1use async_trait::async_trait;
2use futures::stream::{Stream, StreamExt};
3use parking_lot::RwLock;
4use serde_json::Value;
5use std::collections::HashMap;
6use std::future::Future;
7use std::pin::Pin;
8use std::sync::Arc;
9use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
10use std::time::Instant;
11use tracing::{debug, error, info, instrument, warn};
12
13pub(crate) type ToolResourceLocks = Arc<RwLock<HashMap<String, Arc<tokio::sync::Mutex<()>>>>>;
15
16use crate::turn_context::{current_turn_actor_context, scope_actor_context};
17
18use ai_agents_context::{ContextManager, ContextProvider, TemplateRenderer};
19use ai_agents_core::{
20 AgentError, AgentSnapshot, AgentStorage, ChatMessage, FinishReason, LLMProvider, LLMResponse,
21 PermissionOutcome, Result, ToolActorContext, ToolApprovalRecord, ToolApprovalStatus,
22 ToolCallSource, ToolCancellationToken, ToolExecutionContext, ToolExecutionRecord,
23 ToolExecutionRequest, ToolInvoker, ToolPolicyDecisionRecord, ToolResult,
24};
25use ai_agents_disambiguation::{
26 ClarificationObserver, ClarificationParseFuture, ClarificationQuestionFuture,
27 DisambiguationConfig, DisambiguationContext, DisambiguationManager, DisambiguationResult,
28};
29use ai_agents_hitl::{
30 ApprovalHandler, ApprovalResult, ApprovalTrigger, HITLCheckResult, HITLEngine,
31 RejectAllHandler, TimeoutAction,
32};
33use ai_agents_hooks::{AgentHooks, NoopHooks};
34use ai_agents_llm::LLMRegistry;
35use ai_agents_memory::{
36 CompressResult, EvictionReason, Memory, MemoryBudgetEvent, MemoryCompressEvent,
37 MemoryEvictEvent, MemoryTokenBudget, OverflowStrategy,
38};
39use ai_agents_observability::{
40 EventStatus, EventType, ObservabilityManager, ObservationPurpose, SpanContext,
41 current_observation_context, new_session_id as new_observation_session_id,
42 resolve_language_from_context, with_observation_context, with_observation_purpose,
43};
44use ai_agents_process::{
45 ProcessData, ProcessProcessor, ProcessPurposeHint, ProcessStageFuture, ProcessStageObserver,
46};
47use ai_agents_reasoning::{
48 CriterionResult, EvaluationResult, Plan, PlanAction, PlanStatus, PlanStep, ReasoningConfig,
49 ReasoningMetadata, ReasoningMode, ReasoningOutput, ReflectionAttempt, ReflectionConfig,
50 ReflectionMetadata, StepFailureAction,
51};
52use ai_agents_recovery::{
53 ByRoleFilter, ContextOverflowAction, FilterConfig, IntoClassifiedError, KeepRecentFilter,
54 LLMFailureAction, MessageFilter, RecoveryManager, SkipPatternFilter, ToolFailureAction,
55};
56use ai_agents_relationships::RelationshipManager;
57use ai_agents_skills::{SkillDefinition, SkillExecutor, SkillRouter};
58use ai_agents_state::{
59 PromptMode, StateAction, StateMachine, StateMachineSnapshot, StateTransitionEvent, ToolRef,
60 Transition, TransitionContext, TransitionEvaluator, TransitionTiming, evaluate_guard,
61};
62use ai_agents_storage::{StorageConfig as StorageStorageConfig, create_storage};
63use ai_agents_tools::{
64 CommandRunner, ConditionEvaluator, DiagnosticsProvider, EvaluationContext, LLMGetter,
65 QuestionHandler, SecurityCheckResult, TodoItem, ToolCallRecord, ToolRegistry,
66 ToolSecurityConfig, ToolSecurityEngine,
67};
68
69use super::{
70 Agent, AgentInfo, AgentResponse, ParallelToolsConfig, StreamChunk, StreamingConfig, ToolCall,
71};
72use crate::optimization::{
73 AwaitBeforeNextTurn, BackgroundMaintenanceQueue, BackgroundOverflowPolicy, MainResponseDraft,
74 MaintenanceMode, MaintenanceSequenceKey, RuntimeBranch, RuntimeBranchResult,
75 RuntimeBranchStatus, RuntimeCommitBehavior, RuntimeConfig, RuntimeOptimizationKind,
76 RuntimeTaskPriority, RuntimeTaskPurpose, ScheduledBranchSet, SkillCandidate,
77 StreamingDraftResult, TransitionCandidate, TurnBranchScheduler, TurnOptimizationContext,
78};
79use crate::spec::StorageConfig;
80
81enum ToolCallOutcome {
83 Continue,
85 TransitionFired,
87 Rejected(AgentResponse),
89}
90
91enum SkillRouteResult {
93 NoMatch,
95 Response(String),
97 NeedsClarification(AgentResponse),
99}
100
101enum ParallelTransitionSelection {
103 Candidate(TransitionCandidate),
105 NoMatch,
107 ReservationExhausted,
109}
110
111enum PostLoopResult {
113 NoTransition(String),
115 Transitioned(String),
117 NeedsRedispatch,
120}
121
122struct RootTurnCleanup<'a> {
123 agent: &'a RuntimeAgent,
124}
125
126impl<'a> RootTurnCleanup<'a> {
127 fn new(agent: &'a RuntimeAgent) -> Self {
128 Self { agent }
129 }
130}
131
132impl Drop for RootTurnCleanup<'_> {
133 fn drop(&mut self) {
134 self.agent.end_root_turn();
135 }
136}
137
138#[derive(Debug)]
140struct RuntimeControlState {
141 version: AtomicU64,
143 emergency_deny: Arc<AtomicBool>,
145 tool_security_override: RwLock<Option<ToolSecurityConfig>>,
147 tool_scope_override: RwLock<Option<Vec<String>>>,
149}
150
151impl Default for RuntimeControlState {
152 fn default() -> Self {
153 Self {
154 version: AtomicU64::new(1),
155 emergency_deny: Arc::new(AtomicBool::new(false)),
156 tool_security_override: RwLock::new(None),
157 tool_scope_override: RwLock::new(None),
158 }
159 }
160}
161
162#[derive(Clone)]
164pub struct RuntimeControlHandle {
165 state: Arc<RuntimeControlState>,
166}
167
168impl RuntimeControlHandle {
169 pub fn version(&self) -> u64 {
171 self.state.version.load(Ordering::SeqCst)
172 }
173
174 fn bump(&self) -> u64 {
175 self.state.version.fetch_add(1, Ordering::SeqCst) + 1
176 }
177
178 pub fn set_tool_security(&self, config: ToolSecurityConfig) -> u64 {
180 *self.state.tool_security_override.write() = Some(config);
181 self.bump()
182 }
183
184 pub fn clear_tool_security_override(&self) -> u64 {
186 *self.state.tool_security_override.write() = None;
187 self.bump()
188 }
189
190 pub fn set_tool_scope(&self, tool_ids: Vec<String>) -> u64 {
192 *self.state.tool_scope_override.write() = Some(tool_ids);
193 self.bump()
194 }
195
196 pub fn clear_tool_scope_override(&self) -> u64 {
198 *self.state.tool_scope_override.write() = None;
199 self.bump()
200 }
201
202 pub fn set_emergency_deny(&self, enabled: bool) -> u64 {
204 self.state.emergency_deny.store(enabled, Ordering::SeqCst);
205 self.bump()
206 }
207
208 pub fn cancel_all(&self) -> u64 {
210 self.set_emergency_deny(true)
211 }
212}
213
214pub struct RuntimeAgent {
215 info: AgentInfo,
216 llm_registry: Arc<LLMRegistry>,
217 memory: Arc<dyn Memory>,
218 tools: Arc<ToolRegistry>,
219 skills: Vec<SkillDefinition>,
220 skill_router: Option<SkillRouter>,
221 skill_executor: Option<SkillExecutor>,
222 base_system_prompt: String,
223 max_iterations: u32,
224 iteration_count: RwLock<u32>,
225 max_context_tokens: u32,
226 memory_token_budget: Option<MemoryTokenBudget>,
227 recovery_manager: RecoveryManager,
228 tool_security: ToolSecurityEngine,
229 process_processor: Option<ProcessProcessor>,
230 message_filters: RwLock<HashMap<String, Arc<dyn MessageFilter>>>,
231 state_machine: Option<Arc<StateMachine>>,
232 transition_evaluator: Option<Arc<dyn TransitionEvaluator>>,
233 context_manager: Arc<ContextManager>,
234 template_renderer: TemplateRenderer,
235 tool_call_history: RwLock<Vec<ToolCallRecord>>,
236 parallel_tools: ParallelToolsConfig,
237 streaming: StreamingConfig,
238 hooks: Arc<dyn AgentHooks>,
239 hitl_engine: Option<HITLEngine>,
240 approval_handler: Arc<dyn ApprovalHandler>,
241 storage_config: StorageConfig,
242 storage: RwLock<Option<Arc<dyn AgentStorage>>>,
243 reasoning_config: ReasoningConfig,
244 reflection_config: ReflectionConfig,
245 disambiguation_manager: Option<DisambiguationManager>,
246 persona_manager: Option<Arc<ai_agents_persona::PersonaManager>>,
248 pending_skill_id: RwLock<Option<String>>,
252 current_plan: RwLock<Option<Plan>>,
253 declared_tool_ids: Option<Vec<String>>,
255 context_initialized: AtomicBool,
257 spawner: Option<Arc<crate::spawner::AgentSpawner>>,
259 spawner_registry: Option<Arc<crate::spawner::AgentRegistry>>,
261 redispatch_depth: RwLock<u32>,
264 active_turn_context: RwLock<Option<TurnOptimizationContext>>,
266 root_user_message_committed: AtomicBool,
268 actor_id: RwLock<Option<String>>,
270 fact_store: RwLock<Option<Arc<ai_agents_facts::FactStore>>>,
272 fact_extractor: RwLock<Option<Arc<dyn ai_agents_facts::FactExtractor>>>,
275 actor_facts_cache: Arc<RwLock<HashMap<String, Vec<ai_agents_core::KeyFact>>>>,
277 messages_since_extraction: Arc<RwLock<usize>>,
279 actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
281 facts_config: Option<ai_agents_facts::FactsConfig>,
283 session_metadata: RwLock<ai_agents_core::SessionMetadata>,
285 current_session_id: RwLock<Option<String>>,
287 relationship_manager: Option<Arc<RelationshipManager>>,
289 observability_manager: Option<Arc<ObservabilityManager>>,
291 runtime_config: RuntimeConfig,
293 background_maintenance: Arc<BackgroundMaintenanceQueue>,
295 resource_locks: ToolResourceLocks,
297 runtime_control: Arc<RuntimeControlState>,
299}
300
301impl std::fmt::Debug for RuntimeAgent {
302 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
303 f.debug_struct("RuntimeAgent")
304 .field("info", &self.info)
305 .field("base_system_prompt", &self.base_system_prompt)
306 .field("max_iterations", &self.max_iterations)
307 .field("skills_count", &self.skills.len())
308 .field("max_context_tokens", &self.max_context_tokens)
309 .field("has_state_machine", &self.state_machine.is_some())
310 .field("parallel_tools", &self.parallel_tools)
311 .field("streaming", &self.streaming)
312 .field("has_hooks", &true)
313 .field("has_hitl", &self.hitl_engine.is_some())
314 .field("storage_type", &self.storage_config.storage_type())
315 .field("reasoning_mode", &self.reasoning_config.mode)
316 .field("reflection_enabled", &self.reflection_config.enabled)
317 .field("declared_tool_ids", &self.declared_tool_ids)
318 .field("has_persona", &self.persona_manager.is_some())
319 .field("has_observability", &self.observability_manager.is_some())
320 .finish_non_exhaustive()
321 }
322}
323
324struct ObservabilityClarificationObserver;
325
326impl ClarificationObserver for ObservabilityClarificationObserver {
327 fn observe_question<'a>(
329 &'a self,
330 future: ClarificationQuestionFuture<'a>,
331 ) -> ClarificationQuestionFuture<'a> {
332 Box::pin(async move {
333 with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
334 })
335 }
336
337 fn observe_parse<'a>(
339 &'a self,
340 future: ClarificationParseFuture<'a>,
341 ) -> ClarificationParseFuture<'a> {
342 Box::pin(async move {
343 with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
344 })
345 }
346}
347
348struct ObservabilityProcessStageObserver;
349
350impl ProcessStageObserver for ObservabilityProcessStageObserver {
351 fn observe<'a>(
353 &'a self,
354 hint: ProcessPurposeHint,
355 future: ProcessStageFuture<'a>,
356 ) -> ProcessStageFuture<'a> {
357 Box::pin(async move {
358 with_observation_purpose(observation_purpose_for_process(hint), future).await
359 })
360 }
361}
362
363struct RegistryLLMGetter {
364 registry: Arc<LLMRegistry>,
365}
366
367impl LLMGetter for RegistryLLMGetter {
368 fn get_llm(&self, alias: &str) -> Option<Arc<dyn LLMProvider>> {
369 self.registry.get(alias).ok()
370 }
371}
372
373impl RuntimeAgent {
374 #[allow(clippy::too_many_arguments)]
375 pub fn new(
376 info: AgentInfo,
377 llm_registry: Arc<LLMRegistry>,
378 memory: Arc<dyn Memory>,
379 tools: Arc<ToolRegistry>,
380 skills: Vec<SkillDefinition>,
381 system_prompt: String,
382 max_iterations: u32,
383 ) -> Self {
384 let (skill_router, skill_executor) = if !skills.is_empty() {
385 let router_llm = llm_registry.router().ok();
386 let router = router_llm.map(|llm| SkillRouter::new(llm, skills.clone()));
387 let executor = SkillExecutor::new(llm_registry.clone(), tools.clone());
388 (router, Some(executor))
389 } else {
390 (None, None)
391 };
392
393 let context_manager =
394 ContextManager::new(HashMap::new(), info.name.clone(), info.version.clone());
395
396 Self {
397 info,
398 llm_registry,
399 memory,
400 tools,
401 skills,
402 skill_router,
403 skill_executor,
404 base_system_prompt: system_prompt,
405 max_iterations,
406 iteration_count: RwLock::new(0),
407 max_context_tokens: 128000,
408 memory_token_budget: None,
409 recovery_manager: RecoveryManager::default(),
410 tool_security: ToolSecurityEngine::default(),
411 process_processor: None,
412 message_filters: RwLock::new(HashMap::new()),
413 state_machine: None,
414 transition_evaluator: None,
415 context_manager: Arc::new(context_manager),
416 template_renderer: TemplateRenderer::new(),
417 tool_call_history: RwLock::new(Vec::new()),
418 parallel_tools: ParallelToolsConfig::default(),
419 streaming: StreamingConfig::default(),
420 hooks: Arc::new(NoopHooks),
421 hitl_engine: None,
422 approval_handler: Arc::new(RejectAllHandler::new()),
423 storage_config: StorageConfig::default(),
424 storage: RwLock::new(None),
425 reasoning_config: ReasoningConfig::default(),
426 reflection_config: ReflectionConfig::default(),
427 disambiguation_manager: None,
428 persona_manager: None,
429 pending_skill_id: RwLock::new(None),
430 current_plan: RwLock::new(None),
431 declared_tool_ids: None,
432 context_initialized: AtomicBool::new(false),
433 spawner: None,
434 spawner_registry: None,
435 redispatch_depth: RwLock::new(0),
436 active_turn_context: RwLock::new(None),
437 root_user_message_committed: AtomicBool::new(false),
438 actor_id: RwLock::new(None),
439 fact_store: RwLock::new(None),
440 fact_extractor: RwLock::new(None),
441 actor_facts_cache: Arc::new(RwLock::new(HashMap::new())),
442 messages_since_extraction: Arc::new(RwLock::new(0)),
443 actor_memory_config: None,
444 facts_config: None,
445 session_metadata: RwLock::new(ai_agents_core::SessionMetadata::default()),
446 current_session_id: RwLock::new(None),
447 relationship_manager: None,
448 observability_manager: None,
449 runtime_config: RuntimeConfig::default(),
450 background_maintenance: Arc::new(BackgroundMaintenanceQueue::default()),
451 resource_locks: new_tool_resource_locks(),
452 runtime_control: Arc::new(RuntimeControlState::default()),
453 }
454 }
455
456 pub fn with_declared_tool_ids(mut self, ids: Option<Vec<String>>) -> Self {
457 self.declared_tool_ids = ids;
458 self
459 }
460
461 pub fn with_storage_config(mut self, config: StorageConfig) -> Self {
462 self.storage_config = config;
463 self
464 }
465
466 pub fn with_storage(self, storage: Arc<dyn AgentStorage>) -> Self {
467 *self.storage.write() = Some(storage);
468 self
469 }
470
471 pub(crate) fn with_shared_resource_locks(mut self, locks: ToolResourceLocks) -> Self {
472 self.resource_locks = locks;
473 self
474 }
475
476 pub fn with_reasoning(mut self, config: ReasoningConfig) -> Self {
477 self.reasoning_config = config;
478 self
479 }
480
481 pub fn with_reflection(mut self, config: ReflectionConfig) -> Self {
482 self.reflection_config = config;
483 self
484 }
485
486 pub fn with_relationships(mut self, manager: Arc<RelationshipManager>) -> Self {
488 self.relationship_manager = Some(manager);
489 self
490 }
491
492 pub fn with_observability(mut self, manager: Arc<ObservabilityManager>) -> Self {
494 self.observability_manager = Some(manager);
495 self
496 }
497
498 pub fn with_runtime_config(mut self, config: RuntimeConfig) -> Self {
500 let max_tasks = config.optimization.post_turn.max_background_tasks;
501 self.background_maintenance = Arc::new(BackgroundMaintenanceQueue::new(max_tasks));
502 self.runtime_config = config;
503 self
504 }
505
506 pub fn runtime_config(&self) -> &RuntimeConfig {
508 &self.runtime_config
509 }
510
511 pub async fn flush_background_tasks(&self) -> Result<()> {
513 self.background_maintenance.flush_all().await
514 }
515
516 pub async fn flush_background_tasks_for_actor(&self, actor_id: &str) -> Result<()> {
518 self.background_maintenance.flush_scope(actor_id).await
519 }
520
521 pub async fn flush_background_tasks_for_purpose(
523 &self,
524 purpose: RuntimeTaskPurpose,
525 ) -> Result<()> {
526 self.background_maintenance.flush_purpose(purpose).await
527 }
528
529 pub async fn flush_background_tasks_for_actor_purpose(
531 &self,
532 actor_id: &str,
533 purpose: RuntimeTaskPurpose,
534 ) -> Result<()> {
535 self.background_maintenance
536 .flush_scope_purpose(actor_id, purpose)
537 .await
538 }
539
540 pub async fn shutdown_background_tasks(&self) -> Result<()> {
542 self.flush_background_tasks().await
543 }
544
545 pub fn observability(&self) -> Option<Arc<ObservabilityManager>> {
547 self.observability_manager.clone()
548 }
549
550 async fn export_observability_if_configured(&self) {
552 let Some(manager) = self.observability_manager.as_ref() else {
553 return;
554 };
555 let export = &manager.config().export;
556 if !export.write_report && !export.write_raw_events {
557 return;
558 }
559 if let Err(error) = manager.export().await {
560 warn!(error = %error, "Observability export failed");
561 }
562 }
563
564 pub fn relationship_manager(&self) -> Option<Arc<RelationshipManager>> {
566 self.relationship_manager.clone()
567 }
568
569 fn current_turn_actor_context(&self) -> Option<crate::TurnActorContext> {
570 current_turn_actor_context()
571 }
572
573 fn effective_actor_id(&self) -> Option<String> {
574 self.current_turn_actor_context()
575 .and_then(|ctx| ctx.effective_actor_id().map(|id| id.to_string()))
576 .or_else(|| self.actor_id.read().clone())
577 }
578
579 fn effective_origin_actor_id(&self) -> Option<String> {
580 self.current_turn_actor_context()
581 .and_then(|ctx| ctx.origin_actor_id.clone())
582 .or_else(|| self.actor_id.read().clone())
583 }
584
585 fn record_session_actor_if_needed(&self) {
586 if let Some(actor_id) = self.effective_origin_actor_id() {
587 let mut meta = self.session_metadata.write();
588 meta.actor_id = Some(actor_id.clone());
589 if !meta.actors.iter().any(|a| a == &actor_id) {
590 meta.actors.push(actor_id);
591 }
592 }
593 }
594
595 fn outbound_actor_context(&self) -> crate::TurnActorContext {
596 let mut context = self.current_turn_actor_context().unwrap_or_default();
597 if context.origin_actor_id.is_none() {
598 context.origin_actor_id = self.effective_origin_actor_id();
599 }
600 context.sender_agent_id = Some(self.info.id.clone());
601 context
602 }
603
604 fn observation_session_id(&self) -> Option<String> {
606 let mut current = self.current_session_id.write();
607 if current.is_none() {
608 *current = Some(new_observation_session_id());
609 }
610 current.clone()
611 }
612
613 fn build_observation_context(&self, actor_id: Option<String>) -> Option<SpanContext> {
615 let manager = self.observability_manager.as_ref()?;
616 let context = self.build_context_with_overlays();
617 let language = resolve_language_from_context(manager.config(), &context);
618 let context = current_observation_context()
619 .map(|parent| parent.child_for_agent(self.info.id.clone()).with_new_turn())
620 .unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
621 Some(
622 context
623 .with_actor(actor_id.or_else(|| self.effective_actor_id()))
624 .with_session(self.observation_session_id())
625 .with_state(self.current_state())
626 .with_language(Some(language)),
627 )
628 }
629
630 fn current_runtime_observation_context(
632 &self,
633 purpose: ObservationPurpose,
634 ) -> Option<SpanContext> {
635 let manager = self.observability_manager.as_ref()?;
636 let context = self.build_context_with_overlays();
637 let language = resolve_language_from_context(manager.config(), &context);
638 let mut observation = current_observation_context()
639 .unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
640 observation.agent_id = self.info.id.clone();
641 observation.actor_id = self.effective_actor_id();
642 observation.session_id = self.observation_session_id();
643 observation.state = self.current_state();
644 observation.language = Some(language);
645 observation.purpose = purpose;
646 Some(observation)
647 }
648
649 async fn observe_purpose<F, T>(&self, purpose: ObservationPurpose, future: F) -> T
651 where
652 F: Future<Output = T>,
653 {
654 if let Some(context) = self.current_runtime_observation_context(purpose) {
655 with_observation_context(context, future).await
656 } else {
657 future.await
658 }
659 }
660
661 fn chat_with_actor_context_boxed<'a>(
663 &'a self,
664 input: &'a str,
665 actor_context: crate::TurnActorContext,
666 ) -> Pin<Box<dyn Future<Output = Result<AgentResponse>> + Send + 'a>> {
667 Box::pin(async move {
668 let actor_id = actor_context.effective_actor_id().map(str::to_string);
669 let run = async move {
670 scope_actor_context(
671 actor_context,
672 Box::pin(async move { self.run_loop(input).await }),
673 )
674 .await
675 };
676 let result = if let Some(context) = self.build_observation_context(actor_id) {
677 with_observation_context(context, run).await
678 } else {
679 run.await
680 };
681 self.export_observability_if_configured().await;
682 result
683 })
684 }
685
686 pub async fn chat_with_actor_context(
690 &self,
691 input: &str,
692 actor_context: crate::TurnActorContext,
693 ) -> Result<AgentResponse> {
694 self.chat_with_actor_context_boxed(input, actor_context)
695 .await
696 }
697
698 pub async fn chat_as_actor(&self, actor_id: &str, input: &str) -> Result<AgentResponse> {
700 let actor_context = crate::TurnActorContext::new().with_origin_actor(actor_id);
701 self.chat_with_actor_context(input, actor_context).await
702 }
703
704 pub async fn load_actor_relationship(&self) -> Result<()> {
706 self.maybe_load_actor_relationship().await;
707 Ok(())
708 }
709
710 pub async fn update_relationship_dimension(
712 &self,
713 dimension: &str,
714 delta: f64,
715 reason: Option<&str>,
716 ) -> Result<ai_agents_relationships::DimensionChange> {
717 self.update_relationship_dimension_for_perspective(
718 ai_agents_relationships::RelationshipPerspective::AgentToActor,
719 dimension,
720 delta,
721 reason,
722 )
723 .await
724 }
725
726 pub async fn update_relationship_dimension_for_perspective(
730 &self,
731 perspective: ai_agents_relationships::RelationshipPerspective,
732 dimension: &str,
733 delta: f64,
734 reason: Option<&str>,
735 ) -> Result<ai_agents_relationships::DimensionChange> {
736 let manager = self
737 .relationship_manager
738 .as_ref()
739 .ok_or_else(|| AgentError::Config("Relationship memory is not configured".into()))?;
740 let actor_id = self.effective_actor_id().ok_or_else(|| {
741 AgentError::Config("No actor ID set. Use set_actor_id() first".into())
742 })?;
743 let change = manager.update_dimension_for_perspective(
744 &actor_id,
745 perspective,
746 dimension,
747 delta,
748 1.0,
749 reason.unwrap_or("manual relationship update"),
750 )?;
751 self.persist_actor_relationship(&actor_id).await?;
752 info!(
753 actor_id = %actor_id,
754 perspective = %change.perspective,
755 dimension = %change.dimension,
756 delta = change.delta,
757 current = change.current,
758 "relationship updated manually"
759 );
760 self.hooks
761 .on_relationship_change(&actor_id, std::slice::from_ref(&change))
762 .await;
763 Ok(change)
764 }
765
766 pub fn reasoning_config(&self) -> &ReasoningConfig {
767 &self.reasoning_config
768 }
769
770 pub fn reflection_config(&self) -> &ReflectionConfig {
771 &self.reflection_config
772 }
773
774 pub fn with_facts_config(
777 mut self,
778 actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
779 facts_config: Option<ai_agents_facts::FactsConfig>,
780 ) -> Self {
781 self.actor_memory_config = actor_memory_config;
782 self.facts_config = facts_config;
783 self
784 }
785
786 pub fn with_facts(
789 mut self,
790 store: Arc<ai_agents_facts::FactStore>,
791 extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>>,
792 actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
793 facts_config: Option<ai_agents_facts::FactsConfig>,
794 ) -> Self {
795 *self.fact_store.write() = Some(store);
796 *self.fact_extractor.write() = extractor;
797 self.actor_memory_config = actor_memory_config;
798 self.facts_config = facts_config;
799 self
800 }
801
802 pub fn fact_store(&self) -> Option<Arc<ai_agents_facts::FactStore>> {
804 self.fact_store.read().clone()
805 }
806
807 pub fn actor_id(&self) -> Option<String> {
809 self.actor_id.read().clone()
810 }
811
812 pub fn set_actor_id(&self, actor_id: &str) -> ai_agents_core::Result<()> {
814 *self.actor_id.write() = Some(actor_id.to_string());
815 {
816 let mut meta = self.session_metadata.write();
817 meta.actor_id = Some(actor_id.to_string());
818 if !meta.actors.iter().any(|a| a == actor_id) {
819 meta.actors.push(actor_id.to_string());
820 }
821 }
822 Ok(())
823 }
824
825 pub fn set_user_id(&self, user_id: &str) -> ai_agents_core::Result<()> {
827 self.set_actor_id(user_id)
828 }
829
830 pub async fn load_actor_memory(&self) -> ai_agents_core::Result<()> {
832 let actor_id = match self.effective_actor_id() {
833 Some(id) => id,
834 None => return Ok(()),
835 };
836
837 let store_opt = self.fact_store.read().clone();
838 if let Some(store) = store_opt {
839 let facts = store.get_facts(&actor_id).await?;
840 let count = facts.len();
841 self.actor_facts_cache
842 .write()
843 .insert(actor_id.clone(), facts);
844 self.hooks.on_actor_memory_loaded(&actor_id, count).await;
845 tracing::debug!("loaded {} facts for actor {}", count, actor_id);
846 }
847
848 Ok(())
849 }
850
851 async fn maybe_load_actor_memory(&self) {
853 let Some(actor_id) = self.effective_actor_id() else {
854 return;
855 };
856 if self.actor_facts_cache.read().contains_key(&actor_id) {
857 return;
858 }
859 let _ = self.load_actor_memory().await;
860 }
861
862 async fn pre_turn_session_lifecycle(&self) {
864 if *self.redispatch_depth.read() > 0 {
865 return;
866 }
867 self.resolve_actor_id_from_context();
868 self.await_background_before_next_turn().await;
869 self.record_session_actor_if_needed();
870 self.maybe_load_actor_memory().await;
871 self.maybe_load_actor_relationship().await;
872 *self.messages_since_extraction.write() += 1;
873 }
874
875 async fn post_turn_session_lifecycle(&self) -> Result<()> {
877 if *self.redispatch_depth.read() > 0 {
878 return Ok(());
879 }
880 *self.messages_since_extraction.write() += 1;
881 self.run_post_turn_maintenance().await
882 }
883
884 fn begin_root_turn(&self) {
886 if *self.redispatch_depth.read() == 0 {
887 let mut guard = self.active_turn_context.write();
888 if guard.is_none() {
889 self.root_user_message_committed
890 .store(false, Ordering::SeqCst);
891 let max_calls = self
892 .runtime_config
893 .optimization
894 .max_speculative_llm_calls_per_turn;
895 *guard = Some(TurnOptimizationContext::new(
896 String::new(),
897 HashMap::new(),
898 max_calls,
899 ));
900 }
901 }
902 }
903
904 fn update_active_turn_context(
905 &self,
906 processed_input: &str,
907 input_context: HashMap<String, Value>,
908 ) {
909 if *self.redispatch_depth.read() > 0 {
910 return;
911 }
912 let max_calls = self
913 .runtime_config
914 .optimization
915 .max_speculative_llm_calls_per_turn;
916 let mut guard = self.active_turn_context.write();
917 match guard.as_mut() {
918 Some(context) => {
919 context.processed_input = processed_input.to_string();
920 context.input_context = input_context;
921 context.max_speculative_llm_calls = max_calls;
922 }
923 None => {
924 *guard = Some(TurnOptimizationContext::new(
925 processed_input,
926 input_context,
927 max_calls,
928 ));
929 }
930 }
931 }
932
933 async fn commit_root_user_message(&self, processed_input: &str) -> Result<()> {
935 if *self.redispatch_depth.read() > 0 {
936 return Ok(());
937 }
938 if !self
939 .root_user_message_committed
940 .swap(true, Ordering::SeqCst)
941 {
942 self.memory
943 .add_message(ChatMessage::user(processed_input))
944 .await?;
945 if let Some(context) = self.active_turn_context.write().as_mut() {
946 context.mark_user_message_committed();
947 }
948 }
949 Ok(())
950 }
951
952 fn end_root_turn(&self) {
954 if *self.redispatch_depth.read() == 0 {
955 self.root_user_message_committed
956 .store(false, Ordering::SeqCst);
957 *self.active_turn_context.write() = None;
958 }
959 }
960
961 fn reserve_active_speculative_llm_call(&self, kind: RuntimeOptimizationKind) -> bool {
962 self.begin_root_turn();
963 let mut guard = self.active_turn_context.write();
964 let Some(context) = guard.as_mut() else {
965 return false;
966 };
967 context.reserve_speculative_llm_call_for(kind)
968 }
969
970 fn branch_context_preview(&self) -> String {
971 let context = self.build_context_with_overlays();
972 let mut value = serde_json::to_string_pretty(&context).unwrap_or_else(|_| "{}".to_string());
973 const MAX_CONTEXT_PREVIEW_CHARS: usize = 2048;
974 if value.chars().count() > MAX_CONTEXT_PREVIEW_CHARS {
975 value = value
976 .chars()
977 .take(MAX_CONTEXT_PREVIEW_CHARS)
978 .collect::<String>();
979 value.push_str("...");
980 }
981 value
982 }
983
984 async fn await_background_before_next_turn(&self) {
986 let optimization = &self.runtime_config.optimization;
987 if !optimization.enabled {
988 return;
989 }
990 let actor_id = self.effective_actor_id();
991 let post = &optimization.post_turn;
992 self.await_background_task(
993 post.facts.await_before_next_turn,
994 RuntimeTaskPurpose::PostTurnFacts,
995 actor_id.as_deref(),
996 "facts",
997 )
998 .await;
999 self.await_background_task(
1000 post.relationships.await_before_next_turn,
1001 RuntimeTaskPurpose::PostTurnRelationship,
1002 actor_id.as_deref(),
1003 "relationships",
1004 )
1005 .await;
1006 }
1007
1008 async fn await_background_task(
1009 &self,
1010 policy: AwaitBeforeNextTurn,
1011 purpose: RuntimeTaskPurpose,
1012 actor_id: Option<&str>,
1013 label: &str,
1014 ) {
1015 match policy {
1016 AwaitBeforeNextTurn::Never => {}
1017 AwaitBeforeNextTurn::Always => {
1018 if let Err(error) = self.flush_background_tasks_for_purpose(purpose).await {
1019 warn!(label = label, error = %error, "background maintenance flush failed");
1020 }
1021 }
1022 AwaitBeforeNextTurn::SameActor => {
1023 if let Some(actor_id) = actor_id {
1024 if let Err(error) = self
1025 .flush_background_tasks_for_actor_purpose(actor_id, purpose)
1026 .await
1027 {
1028 warn!(label = label, actor_id = %actor_id, error = %error, "actor background maintenance flush failed");
1029 }
1030 }
1031 }
1032 }
1033 }
1034
1035 async fn run_post_turn_maintenance(&self) -> Result<()> {
1037 let optimization = &self.runtime_config.optimization;
1038 if !optimization.enabled {
1039 self.auto_extract_facts().await;
1040 self.auto_update_relationship().await;
1041 return Ok(());
1042 }
1043
1044 let facts_mode = effective_maintenance_mode(
1045 optimization.post_turn.facts.mode,
1046 optimization.parallel_post_turn_memory,
1047 );
1048 let relationships_mode = effective_maintenance_mode(
1049 optimization.post_turn.relationships.mode,
1050 optimization.parallel_post_turn_memory,
1051 );
1052
1053 match (facts_mode, relationships_mode) {
1054 (MaintenanceMode::InlineSerial, MaintenanceMode::InlineSerial) => {
1055 self.auto_extract_facts().await;
1056 self.auto_update_relationship().await;
1057 }
1058 (MaintenanceMode::InlineParallel, MaintenanceMode::InlineParallel) => {
1059 let facts = self.auto_extract_facts();
1060 let relationships = self.auto_update_relationship();
1061 tokio::join!(facts, relationships);
1062 }
1063 (MaintenanceMode::Background, MaintenanceMode::Background) => {
1064 self.schedule_facts_background().await?;
1065 self.schedule_relationship_background().await?;
1066 }
1067 (MaintenanceMode::Background, MaintenanceMode::InlineParallel)
1068 | (MaintenanceMode::Background, MaintenanceMode::InlineSerial) => {
1069 self.schedule_facts_background().await?;
1070 self.auto_update_relationship().await;
1071 }
1072 (MaintenanceMode::InlineParallel, MaintenanceMode::Background)
1073 | (MaintenanceMode::InlineSerial, MaintenanceMode::Background) => {
1074 self.auto_extract_facts().await;
1075 self.schedule_relationship_background().await?;
1076 }
1077 _ => {
1078 self.auto_extract_facts().await;
1079 self.auto_update_relationship().await;
1080 }
1081 }
1082 Ok(())
1083 }
1084
1085 async fn schedule_facts_background(&self) -> Result<()> {
1086 let policy = self.runtime_config.optimization.post_turn.facts.clone();
1087 let should_extract = self
1088 .facts_config
1089 .as_ref()
1090 .map(|c| c.enabled && c.auto_extract)
1091 .unwrap_or(false);
1092 if !should_extract {
1093 return Ok(());
1094 }
1095 let msgs_since = *self.messages_since_extraction.read();
1096 if msgs_since < 2 {
1097 return Ok(());
1098 }
1099 let Some(actor_id) = self.effective_actor_id() else {
1100 self.record_skipped_maintenance(
1101 "facts",
1102 ObservationPurpose::FactsExtraction,
1103 "missing_actor",
1104 Some(&policy),
1105 );
1106 return Ok(());
1107 };
1108 let Some(extractor) = self.fact_extractor.read().clone() else {
1109 return Ok(());
1110 };
1111 let messages = match self.memory.get_messages(None).await {
1112 Ok(messages) => messages,
1113 Err(error) => {
1114 warn!(error = %error, "failed to snapshot messages for fact extraction");
1115 return Ok(());
1116 }
1117 };
1118 let recent: Vec<_> = messages
1119 .iter()
1120 .rev()
1121 .take(msgs_since)
1122 .rev()
1123 .cloned()
1124 .collect();
1125 if recent.is_empty() {
1126 return Ok(());
1127 }
1128 let existing = self
1129 .actor_facts_cache
1130 .read()
1131 .get(&actor_id)
1132 .cloned()
1133 .unwrap_or_default();
1134 let categories = self
1135 .facts_config
1136 .as_ref()
1137 .map(|c| c.custom_categories.clone())
1138 .unwrap_or_default();
1139 let store = self.fact_store.read().clone();
1140 let cache = Arc::clone(&self.actor_facts_cache);
1141 let counter = Arc::clone(&self.messages_since_extraction);
1142 let hooks = Arc::clone(&self.hooks);
1143 let agent_id = self.info.id.clone();
1144 let observation = current_observation_context();
1145 let key = MaintenanceSequenceKey::actor(
1146 agent_id,
1147 actor_id.clone(),
1148 RuntimeTaskPurpose::PostTurnFacts,
1149 );
1150 let actor_for_task = actor_id.clone();
1151 let task = async move {
1152 let run = async move {
1153 let facts = extractor
1154 .extract(&recent, &existing, Some(&actor_for_task), &categories)
1155 .await?;
1156 if !facts.is_empty() {
1157 if let Some(store) = store {
1158 let authoritative = store.add_facts(&actor_for_task, facts.clone()).await?;
1159 cache.write().insert(actor_for_task.clone(), authoritative);
1160 } else {
1161 cache
1162 .write()
1163 .entry(actor_for_task.clone())
1164 .or_default()
1165 .extend(facts.clone());
1166 }
1167 {
1168 let mut count = counter.write();
1169 if *count <= msgs_since {
1170 *count = 0;
1171 } else {
1172 *count -= msgs_since;
1173 }
1174 }
1175 hooks.on_facts_extracted(&actor_for_task, &facts).await;
1176 }
1177 Ok(())
1178 };
1179 if let Some(context) = observation {
1180 with_observation_context(
1181 context.with_purpose(ObservationPurpose::FactsExtraction),
1182 run,
1183 )
1184 .await
1185 } else {
1186 run.await
1187 }
1188 };
1189 self.spawn_or_handle_background(Some(key), task, "facts", &policy)
1190 .await
1191 }
1192
1193 async fn schedule_relationship_background(&self) -> Result<()> {
1194 let policy = self
1195 .runtime_config
1196 .optimization
1197 .post_turn
1198 .relationships
1199 .clone();
1200 let Some(manager) = self.relationship_manager.as_ref().cloned() else {
1201 return Ok(());
1202 };
1203 let Some(actor_id) = self.effective_actor_id() else {
1204 self.record_skipped_maintenance(
1205 "relationships",
1206 ObservationPurpose::RelationshipUpdate,
1207 "missing_actor",
1208 Some(&policy),
1209 );
1210 return Ok(());
1211 };
1212 let recent_messages = manager.config().auto_update.recent_messages;
1213 let messages = match self.memory.get_messages(Some(recent_messages)).await {
1214 Ok(messages) => messages,
1215 Err(error) => {
1216 warn!(actor = %actor_id, error = %error, "failed to snapshot messages for relationship update");
1217 return Ok(());
1218 }
1219 };
1220 let storage = self.storage.read().clone();
1221 let hooks = Arc::clone(&self.hooks);
1222 let agent_id = self.info.id.clone();
1223 let observation = current_observation_context();
1224 let key = MaintenanceSequenceKey::actor(
1225 agent_id.clone(),
1226 actor_id.clone(),
1227 RuntimeTaskPurpose::PostTurnRelationship,
1228 );
1229 let actor_for_task = actor_id.clone();
1230 let task = async move {
1231 let run = async move {
1232 if manager.config().auto_update.enabled {
1233 let update = manager.auto_update(&actor_for_task, &messages).await?;
1234 if !update.changes.is_empty() {
1235 hooks
1236 .on_relationship_change(&actor_for_task, &update.changes)
1237 .await;
1238 }
1239 if let Some(ref event) = update.event {
1240 hooks.on_notable_event(&actor_for_task, event).await;
1241 }
1242 }
1243 if manager.config().persistence.enabled {
1244 if let (Some(storage), Some(value)) =
1245 (storage, manager.relationship_as_value(&actor_for_task)?)
1246 {
1247 storage
1248 .save_relationship(&agent_id, &actor_for_task, &value)
1249 .await?;
1250 }
1251 }
1252 Ok(())
1253 };
1254 if let Some(context) = observation {
1255 with_observation_context(
1256 context.with_purpose(ObservationPurpose::RelationshipUpdate),
1257 run,
1258 )
1259 .await
1260 } else {
1261 run.await
1262 }
1263 };
1264 self.spawn_or_handle_background(Some(key), task, "relationships", &policy)
1265 .await
1266 }
1267
1268 async fn spawn_or_handle_background<F>(
1270 &self,
1271 key: Option<MaintenanceSequenceKey>,
1272 task: F,
1273 label: &'static str,
1274 policy: &crate::optimization::config::MaintenanceTaskPolicy,
1275 ) -> Result<()>
1276 where
1277 F: Future<Output = Result<()>> + Send + 'static,
1278 {
1279 if self.background_maintenance.is_full() {
1280 match self
1281 .runtime_config
1282 .optimization
1283 .post_turn
1284 .on_background_overflow
1285 {
1286 BackgroundOverflowPolicy::RunInline => {
1287 record_background_maintenance_event(
1288 self.observability_manager.as_ref(),
1289 label,
1290 EventStatus::Success,
1291 0,
1292 "inline_overflow",
1293 None,
1294 Some(policy),
1295 );
1296 let start = Instant::now();
1297 match task.await {
1298 Ok(()) => record_background_maintenance_event(
1299 self.observability_manager.as_ref(),
1300 label,
1301 EventStatus::Success,
1302 start.elapsed().as_millis() as u64,
1303 "inline_completed",
1304 None,
1305 Some(policy),
1306 ),
1307 Err(error) => {
1308 warn!(label = label, error = %error, "inline maintenance fallback failed");
1309 record_background_maintenance_event(
1310 self.observability_manager.as_ref(),
1311 label,
1312 EventStatus::Error,
1313 start.elapsed().as_millis() as u64,
1314 "inline_failed",
1315 Some(error.to_string()),
1316 Some(policy),
1317 );
1318 return Err(error);
1319 }
1320 }
1321 }
1322 BackgroundOverflowPolicy::Drop => {
1323 self.record_skipped_maintenance(
1324 label,
1325 ObservationPurpose::Other(label.to_string()),
1326 "queue_full",
1327 Some(policy),
1328 );
1329 }
1330 BackgroundOverflowPolicy::Error => {
1331 record_background_maintenance_event(
1332 self.observability_manager.as_ref(),
1333 label,
1334 EventStatus::Error,
1335 0,
1336 "queue_full",
1337 None,
1338 Some(policy),
1339 );
1340 warn!(label = label, "background maintenance queue full");
1341 return Err(AgentError::Other(format!(
1342 "background maintenance queue is full for {}",
1343 label
1344 )));
1345 }
1346 }
1347 return Ok(());
1348 }
1349
1350 record_background_maintenance_event(
1351 self.observability_manager.as_ref(),
1352 label,
1353 EventStatus::Success,
1354 0,
1355 "scheduled",
1356 None,
1357 Some(policy),
1358 );
1359 let manager = self.observability_manager.clone();
1360 let policy_for_task = policy.clone();
1361 let observed_task = async move {
1362 let start = Instant::now();
1363 let result = task.await;
1364 match &result {
1365 Ok(()) => record_background_maintenance_event(
1366 manager.as_ref(),
1367 label,
1368 EventStatus::Success,
1369 start.elapsed().as_millis() as u64,
1370 "completed",
1371 None,
1372 Some(&policy_for_task),
1373 ),
1374 Err(error) => record_background_maintenance_event(
1375 manager.as_ref(),
1376 label,
1377 EventStatus::Error,
1378 start.elapsed().as_millis() as u64,
1379 "failed",
1380 Some(error.to_string()),
1381 Some(&policy_for_task),
1382 ),
1383 }
1384 result
1385 };
1386
1387 if let Err(error) = self.background_maintenance.spawn(key, observed_task) {
1388 record_background_maintenance_event(
1389 self.observability_manager.as_ref(),
1390 label,
1391 EventStatus::Error,
1392 0,
1393 "spawn_failed",
1394 Some(error.to_string()),
1395 Some(policy),
1396 );
1397 warn!(label = label, error = %error, "background maintenance spawn failed");
1398 return Err(error);
1399 }
1400 Ok(())
1401 }
1402
1403 fn record_skipped_maintenance(
1405 &self,
1406 label: &str,
1407 purpose: ObservationPurpose,
1408 reason: &str,
1409 policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
1410 ) {
1411 if let Some(manager) = self.observability_manager.as_ref() {
1412 let mut tags = background_maintenance_tags(label, "skipped", Some(reason), policy);
1413 tags.insert("runtime.skip_reason".to_string(), reason.to_string());
1414 manager.record_lifecycle_event(
1415 EventType::MemoryOperation {
1416 operation: format!("{}_maintenance", label),
1417 },
1418 purpose,
1419 EventStatus::Skipped,
1420 0,
1421 tags,
1422 None,
1423 );
1424 }
1425 }
1426
1427 pub fn actor_facts(&self) -> Vec<ai_agents_core::KeyFact> {
1429 let Some(actor_id) = self.effective_actor_id() else {
1430 return Vec::new();
1431 };
1432 self.actor_facts_cache
1433 .read()
1434 .get(&actor_id)
1435 .cloned()
1436 .unwrap_or_default()
1437 }
1438
1439 pub fn relationship_memory_text(&self) -> Option<String> {
1441 self.format_relationship_for_context().map(|(_, text)| text)
1442 }
1443
1444 pub async fn extract_facts(
1446 &self,
1447 last_n: usize,
1448 ) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
1449 self.extract_facts_with_source(last_n, "manual").await
1450 }
1451
1452 async fn extract_facts_with_source(
1453 &self,
1454 last_n: usize,
1455 source: &'static str,
1456 ) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
1457 let extractor = match self.fact_extractor.read().clone() {
1458 Some(e) => e,
1459 None => return Ok(vec![]),
1460 };
1461
1462 let messages = self.memory.get_messages(None).await?;
1463 let recent: Vec<_> = messages.iter().rev().take(last_n).rev().cloned().collect();
1464
1465 if recent.is_empty() {
1466 return Ok(vec![]);
1467 }
1468
1469 let actor_id = self.effective_actor_id();
1470 let existing = actor_id
1471 .as_ref()
1472 .and_then(|aid| self.actor_facts_cache.read().get(aid).cloned())
1473 .unwrap_or_default();
1474
1475 let categories = self
1476 .facts_config
1477 .as_ref()
1478 .map(|c| c.custom_categories.clone())
1479 .unwrap_or_default();
1480
1481 let facts = self
1482 .observe_purpose(
1483 ObservationPurpose::FactsExtraction,
1484 extractor.extract(&recent, &existing, actor_id.as_deref(), &categories),
1485 )
1486 .await?;
1487
1488 if !facts.is_empty() {
1490 let fact_store_opt = self.fact_store.read().clone();
1491 let mut stored_total = 0usize;
1492 let mut cache_updated = false;
1493 if let (Some(store), Some(aid)) = (fact_store_opt, &actor_id) {
1494 let authoritative = store.add_facts(aid, facts.clone()).await?;
1496 stored_total = authoritative.len();
1497 self.actor_facts_cache
1498 .write()
1499 .insert(aid.clone(), authoritative);
1500 cache_updated = true;
1501 } else if let Some(aid) = &actor_id {
1502 let mut cache = self.actor_facts_cache.write();
1503 let entry = cache.entry(aid.clone()).or_default();
1504 entry.extend(facts.clone());
1505 stored_total = entry.len();
1506 cache_updated = true;
1507 }
1508
1509 info!(
1510 actor_id = %actor_id.as_deref().unwrap_or("<none>"),
1511 source = source,
1512 requested_messages = last_n,
1513 message_count = recent.len(),
1514 extracted_count = facts.len(),
1515 cache_updated = cache_updated,
1516 stored_total = stored_total,
1517 "facts extracted"
1518 );
1519
1520 if let Some(ref aid) = actor_id {
1521 self.hooks.on_facts_extracted(aid, &facts).await;
1522 }
1523 }
1524
1525 Ok(facts)
1526 }
1527
1528 fn resolve_actor_id_from_context(&self) {
1531 if self
1532 .current_turn_actor_context()
1533 .and_then(|ctx| ctx.effective_actor_id().map(str::to_string))
1534 .is_some()
1535 {
1536 return;
1537 }
1538
1539 if let Some(ref am_config) = self.actor_memory_config {
1540 if am_config.identification.method == ai_agents_facts::IdentificationMethod::FromContext
1541 {
1542 if let Some(ref path) = am_config.identification.context_path {
1543 let val = self
1545 .context_manager
1546 .get_path(path)
1547 .or_else(|| self.context_manager.get(path));
1548 if let Some(val) = val {
1549 if let Some(id_str) = val.as_str() {
1550 let current = self.actor_id.read().clone();
1551 if current.as_deref() != Some(id_str) {
1552 *self.actor_id.write() = Some(id_str.to_string());
1553 let mut meta = self.session_metadata.write();
1554 meta.actor_id = Some(id_str.to_string());
1555 if !meta.actors.iter().any(|a| a == id_str) {
1556 meta.actors.push(id_str.to_string());
1557 }
1558 }
1559 }
1560 }
1561 }
1562 }
1563 }
1564 }
1565
1566 fn format_actor_facts_for_context(&self) -> String {
1568 let should_inject = self
1570 .facts_config
1571 .as_ref()
1572 .map(|c| c.inject_in_context)
1573 .unwrap_or(true);
1574 if !should_inject {
1575 return String::new();
1576 }
1577
1578 let Some(actor_id) = self.effective_actor_id() else {
1579 return String::new();
1580 };
1581
1582 let facts = self
1583 .actor_facts_cache
1584 .read()
1585 .get(&actor_id)
1586 .cloned()
1587 .unwrap_or_default();
1588 if facts.is_empty() {
1589 return String::new();
1590 }
1591
1592 let am_config = self.actor_memory_config.as_ref();
1593 let facts_budget = self
1596 .memory_token_budget
1597 .as_ref()
1598 .map(|b| b.allocation.facts as usize)
1599 .filter(|n| *n > 0);
1600 let default_max = am_config.map(|c| c.injection.max_tokens).unwrap_or(800);
1601 let max_tokens = facts_budget.unwrap_or(default_max);
1602
1603 let filtered: Vec<ai_agents_core::KeyFact> = if let Some(cfg) = am_config {
1605 if cfg.injection.mode == ai_agents_facts::InjectionMode::OnDemand {
1606 return String::new();
1607 }
1608 if cfg.injection.mode == ai_agents_facts::InjectionMode::Category
1609 && !cfg.injection.categories.is_empty()
1610 {
1611 facts
1612 .iter()
1613 .filter(|f| {
1614 cfg.injection
1615 .categories
1616 .iter()
1617 .any(|c| f.category.to_string() == *c)
1618 })
1619 .cloned()
1620 .collect()
1621 } else {
1622 facts.clone()
1623 }
1624 } else {
1625 facts.clone()
1626 };
1627
1628 if filtered.is_empty() {
1629 return String::new();
1630 }
1631
1632 if let Some(store) = self.fact_store.read().clone() {
1633 store.format_for_context(&filtered, max_tokens)
1634 } else {
1635 String::new()
1636 }
1637 }
1638
1639 fn build_context_with_staged(&self, staged: &HashMap<String, Value>) -> HashMap<String, Value> {
1640 let context = self.build_context_with_overlays();
1641 let mut root = Value::Object(context.into_iter().collect());
1642 for (path, value) in staged {
1643 if let Ok(updated) = ai_agents_core::set_dot_path(root.clone(), path, value.clone()) {
1644 root = updated;
1645 }
1646 }
1647 match root {
1648 Value::Object(obj) => obj.into_iter().collect(),
1649 _ => HashMap::new(),
1650 }
1651 }
1652
1653 fn build_context_with_overlays(&self) -> HashMap<String, Value> {
1654 let mut context = self.context_manager.get_all();
1655 let mut root = Value::Object(context.clone().into_iter().collect());
1656
1657 if let Some(turn_ctx) = self.current_turn_actor_context() {
1658 if let Some(ref origin_actor_id) = turn_ctx.origin_actor_id {
1659 if let Ok(updated) = ai_agents_core::set_dot_path(
1660 root.clone(),
1661 "interaction.origin_actor_id",
1662 serde_json::json!(origin_actor_id),
1663 ) {
1664 root = updated;
1665 }
1666 }
1667 if let Some(ref sender_agent_id) = turn_ctx.sender_agent_id {
1668 if let Ok(updated) = ai_agents_core::set_dot_path(
1669 root.clone(),
1670 "interaction.sender_agent_id",
1671 serde_json::json!(sender_agent_id),
1672 ) {
1673 root = updated;
1674 }
1675 }
1676 }
1677
1678 if let Some(ref actor_id) = self.effective_actor_id() {
1679 if let Ok(updated) = ai_agents_core::set_dot_path(
1680 root.clone(),
1681 "interaction.actor_id",
1682 serde_json::json!(actor_id),
1683 ) {
1684 root = updated;
1685 }
1686 }
1687
1688 if let Some(manager) = self.relationship_manager.as_ref() {
1689 if let Some(actor_id) = self.effective_actor_id() {
1690 if let Some(value) = manager.to_context_value(&actor_id) {
1691 if let Ok(updated) = ai_agents_core::set_dot_path(
1692 root.clone(),
1693 &manager.config().injection.context_path,
1694 value,
1695 ) {
1696 root = updated;
1697 }
1698 }
1699 }
1700 }
1701
1702 if let Value::Object(obj) = root {
1703 context = obj.into_iter().collect();
1704 }
1705
1706 context
1707 }
1708
1709 fn resolve_actor_name_from_context(&self) -> Option<String> {
1710 for path in ["actor.name", "user.name", "player.name", "customer.name"] {
1711 if let Some(value) = self.context_manager.get_path(path) {
1712 if let Some(name) = value.as_str() {
1713 return Some(name.to_string());
1714 }
1715 }
1716 }
1717 None
1718 }
1719
1720 async fn maybe_load_actor_relationship(&self) {
1721 let Some(manager) = self.relationship_manager.as_ref() else {
1722 return;
1723 };
1724 let Some(actor_id) = self.effective_actor_id() else {
1725 return;
1726 };
1727
1728 let mut should_fire_loaded = false;
1729 if manager.get(&actor_id).is_none() {
1730 let mut loaded = false;
1731 if manager.config().persistence.enabled {
1732 let storage = self.storage.read().clone();
1733 if let Some(storage) = storage {
1734 match storage.load_relationship(&self.info.id, &actor_id).await {
1735 Ok(Some(value)) => match manager.insert_from_value(value) {
1736 Ok(_) => loaded = true,
1737 Err(e) => {
1738 warn!(actor = %actor_id, error = %e, "failed to restore relationship")
1739 }
1740 },
1741 Ok(None) => {}
1742 Err(e) => {
1743 warn!(actor = %actor_id, error = %e, "failed to load relationship")
1744 }
1745 }
1746 }
1747 }
1748
1749 if !loaded {
1750 manager.get_or_create(&actor_id, self.resolve_actor_name_from_context().as_deref());
1751 }
1752 should_fire_loaded = true;
1753 }
1754
1755 let actor_name = self.resolve_actor_name_from_context();
1756 let relationship = manager.touch_interaction(&actor_id, actor_name.as_deref());
1757 if should_fire_loaded {
1758 self.hooks
1759 .on_relationship_loaded(&actor_id, &relationship)
1760 .await;
1761 }
1762 }
1763
1764 fn format_relationship_for_context(&self) -> Option<(String, String)> {
1765 let manager = self.relationship_manager.as_ref()?;
1766 if !manager.config().injection.enabled {
1767 return None;
1768 }
1769 let actor_id = self.effective_actor_id()?;
1770 let relationship = manager.get(&actor_id)?;
1771 let local_cap = manager.config().injection.max_tokens;
1772 let global_cap = self
1773 .memory_token_budget
1774 .as_ref()
1775 .map(|b| b.allocation.relationships as usize)
1776 .filter(|n| *n > 0);
1777 let max_tokens = global_cap.map(|g| g.min(local_cap)).unwrap_or(local_cap);
1778 let text = ai_agents_relationships::format_relationship(
1779 &relationship,
1780 &manager.config().injection.format,
1781 max_tokens,
1782 );
1783 if text.is_empty() {
1784 None
1785 } else {
1786 Some((manager.config().injection.prompt_variable.clone(), text))
1787 }
1788 }
1789
1790 async fn persist_actor_relationship(&self, actor_id: &str) -> Result<()> {
1791 let Some(manager) = self.relationship_manager.as_ref() else {
1792 return Ok(());
1793 };
1794 if !manager.config().persistence.enabled {
1795 return Ok(());
1796 }
1797 let storage = self.storage.read().clone();
1798 let Some(storage) = storage else {
1799 return Ok(());
1800 };
1801 if let Some(value) = manager.relationship_as_value(actor_id)? {
1802 storage
1803 .save_relationship(&self.info.id, actor_id, &value)
1804 .await?;
1805 }
1806 Ok(())
1807 }
1808
1809 async fn auto_update_relationship(&self) {
1810 let Some(manager) = self.relationship_manager.as_ref() else {
1811 return;
1812 };
1813 let Some(actor_id) = self.effective_actor_id() else {
1814 return;
1815 };
1816 if !manager.config().auto_update.enabled {
1817 let _ = self.persist_actor_relationship(&actor_id).await;
1818 return;
1819 }
1820
1821 let recent_messages = manager.config().auto_update.recent_messages;
1822 let messages = match self.memory.get_messages(Some(recent_messages)).await {
1823 Ok(messages) => messages,
1824 Err(e) => {
1825 warn!(actor = %actor_id, error = %e, "failed to read messages for relationship update");
1826 return;
1827 }
1828 };
1829
1830 match self
1831 .observe_purpose(
1832 ObservationPurpose::RelationshipUpdate,
1833 manager.auto_update(&actor_id, &messages),
1834 )
1835 .await
1836 {
1837 Ok(update) => {
1838 if !update.changes.is_empty() {
1839 self.hooks
1840 .on_relationship_change(&actor_id, &update.changes)
1841 .await;
1842 }
1843 if let Some(ref event) = update.event {
1844 self.hooks.on_notable_event(&actor_id, event).await;
1845 }
1846 let persisted = match self.persist_actor_relationship(&actor_id).await {
1847 Ok(()) => true,
1848 Err(e) => {
1849 warn!(actor = %actor_id, error = %e, "failed to persist relationship");
1850 false
1851 }
1852 };
1853 if !update.changes.is_empty() || update.event.is_some() {
1854 let changed_dimensions: Vec<String> = update
1855 .changes
1856 .iter()
1857 .map(|change| format!("{}:{}", change.perspective, change.dimension))
1858 .collect();
1859 info!(
1860 actor_id = %actor_id,
1861 change_count = update.changes.len(),
1862 changed_dimensions = ?changed_dimensions,
1863 event_present = update.event.is_some(),
1864 persisted = persisted,
1865 "relationship updated"
1866 );
1867 } else {
1868 debug!(actor_id = %actor_id, persisted = persisted, "relationship evaluation ran but found no changes");
1869 }
1870 }
1871 Err(e) => warn!(actor = %actor_id, error = %e, "relationship update failed"),
1872 }
1873 }
1874
1875 async fn auto_extract_facts(&self) {
1877 let should_extract = self
1878 .facts_config
1879 .as_ref()
1880 .map(|c| c.enabled && c.auto_extract)
1881 .unwrap_or(false);
1882
1883 if !should_extract {
1884 debug!("fact extraction skipped because auto extraction is disabled");
1885 return;
1886 }
1887
1888 let msgs_since = *self.messages_since_extraction.read();
1889 if msgs_since < 2 {
1890 debug!(
1891 messages_since_extraction = msgs_since,
1892 "fact extraction skipped until threshold is reached"
1893 );
1894 return;
1895 }
1896
1897 match self.extract_facts_with_source(msgs_since, "auto").await {
1898 Ok(facts) => {
1899 if !facts.is_empty() {
1900 *self.messages_since_extraction.write() = 0;
1901 } else {
1902 debug!("fact extraction ran but found no new facts");
1903 }
1904 }
1905 Err(e) => {
1906 warn!("fact extraction failed: {}", e);
1907 }
1908 }
1909 }
1910
1911 pub fn with_persona(mut self, manager: Arc<ai_agents_persona::PersonaManager>) -> Self {
1912 self.persona_manager = Some(manager);
1913 self
1914 }
1915
1916 pub fn persona_manager(&self) -> Option<&Arc<ai_agents_persona::PersonaManager>> {
1917 self.persona_manager.as_ref()
1918 }
1919
1920 pub fn with_disambiguation(mut self, config: DisambiguationConfig) -> Self {
1921 if config.is_enabled() {
1922 let manager = DisambiguationManager::new(config, Arc::clone(&self.llm_registry))
1923 .with_clarification_observer(Arc::new(ObservabilityClarificationObserver));
1924 self.disambiguation_manager = Some(manager);
1925 }
1926 self
1927 }
1928
1929 pub fn disambiguation_manager(&self) -> Option<&DisambiguationManager> {
1930 self.disambiguation_manager.as_ref()
1931 }
1932
1933 pub fn has_disambiguation(&self) -> bool {
1934 self.disambiguation_manager
1935 .as_ref()
1936 .is_some_and(|m| m.is_enabled())
1937 }
1938
1939 pub async fn init_storage(&self) -> Result<()> {
1940 if self.storage_config.is_none() {
1941 return Ok(());
1942 }
1943 if self.storage.read().is_some() {
1944 return Ok(());
1945 }
1946 let storage_config = self.convert_storage_config();
1947 let storage = create_storage(&storage_config).await?;
1948 *self.storage.write() = storage;
1949 self.complete_facts_init().await;
1951 Ok(())
1952 }
1953
1954 async fn complete_facts_init(&self) {
1957 if self.fact_store.read().is_some() {
1958 return;
1959 }
1960 let storage = match self.storage.read().clone() {
1961 Some(s) => s,
1962 None => return,
1963 };
1964
1965 let facts_enabled = self
1966 .facts_config
1967 .as_ref()
1968 .map(|f| f.enabled)
1969 .unwrap_or(false);
1970 let actor_memory_enabled = self
1971 .actor_memory_config
1972 .as_ref()
1973 .map(|a| a.enabled)
1974 .unwrap_or(false);
1975
1976 if !facts_enabled && !actor_memory_enabled {
1977 return;
1978 }
1979
1980 let fc = self.facts_config.clone().unwrap_or_default();
1981 let store = Arc::new(ai_agents_facts::FactStore::new(
1982 storage,
1983 self.info.id.clone(),
1984 fc.clone(),
1985 ));
1986
1987 let extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>> = if facts_enabled {
1988 let extractor_llm = fc
1989 .extractor_llm
1990 .as_ref()
1991 .and_then(|alias| self.llm_registry.get(alias).ok())
1992 .or_else(|| self.llm_registry.router().ok())
1993 .or_else(|| self.llm_registry.default().ok());
1994 extractor_llm.map(|llm| {
1995 Arc::new(ai_agents_facts::LLMFactExtractor::new(llm, fc.clone()))
1996 as Arc<dyn ai_agents_facts::FactExtractor>
1997 })
1998 } else {
1999 None
2000 };
2001
2002 *self.fact_store.write() = Some(store);
2003 *self.fact_extractor.write() = extractor;
2004 debug!(
2005 agent = %self.info.id,
2006 facts_enabled,
2007 actor_memory_enabled,
2008 "facts storage initialized"
2009 );
2010 }
2011
2012 fn convert_storage_config(&self) -> StorageStorageConfig {
2013 crate::spec::storage::to_storage_config(&self.storage_config)
2014 }
2015
2016 pub fn storage(&self) -> Option<Arc<dyn AgentStorage>> {
2017 self.storage.read().clone()
2018 }
2019
2020 pub fn storage_config(&self) -> &StorageConfig {
2021 &self.storage_config
2022 }
2023
2024 pub fn spawner(&self) -> Option<&Arc<crate::spawner::AgentSpawner>> {
2026 self.spawner.as_ref()
2027 }
2028
2029 pub fn spawner_registry(&self) -> Option<&Arc<crate::spawner::AgentRegistry>> {
2031 self.spawner_registry.as_ref()
2032 }
2033
2034 pub fn has_spawner(&self) -> bool {
2035 self.spawner_registry.is_some()
2036 }
2037
2038 pub fn with_spawner_handles(
2039 mut self,
2040 spawner: Arc<crate::spawner::AgentSpawner>,
2041 registry: Arc<crate::spawner::AgentRegistry>,
2042 ) -> Self {
2043 self.spawner = Some(spawner);
2044 self.spawner_registry = Some(registry);
2045 self
2046 }
2047
2048 pub fn with_hooks(mut self, hooks: Arc<dyn AgentHooks>) -> Self {
2049 self.hooks = hooks;
2050 self
2051 }
2052
2053 pub fn with_parallel_tools(mut self, config: ParallelToolsConfig) -> Self {
2054 self.parallel_tools = config;
2055 self
2056 }
2057
2058 pub fn with_streaming(mut self, config: StreamingConfig) -> Self {
2059 self.streaming = config;
2060 self
2061 }
2062
2063 pub fn with_hitl(mut self, engine: HITLEngine, handler: Arc<dyn ApprovalHandler>) -> Self {
2064 self.hitl_engine = Some(engine);
2065 self.approval_handler = handler;
2066 self
2067 }
2068
2069 pub fn with_max_context_tokens(mut self, tokens: u32) -> Self {
2070 self.max_context_tokens = tokens;
2071 self
2072 }
2073
2074 pub fn with_memory_token_budget(mut self, budget: MemoryTokenBudget) -> Self {
2075 self.memory_token_budget = Some(budget);
2076 self
2077 }
2078
2079 pub fn with_recovery_manager(mut self, manager: RecoveryManager) -> Self {
2080 self.recovery_manager = manager;
2081 self
2082 }
2083
2084 pub fn with_tool_security(mut self, engine: ToolSecurityEngine) -> Self {
2085 self.tool_security = engine;
2086 self
2087 }
2088
2089 pub fn runtime_control(&self) -> RuntimeControlHandle {
2091 RuntimeControlHandle {
2092 state: Arc::clone(&self.runtime_control),
2093 }
2094 }
2095
2096 pub fn set_question_handler(&self, handler: Option<Arc<dyn QuestionHandler>>) {
2098 self.tools.set_question_handler(handler);
2099 }
2100
2101 pub fn set_diagnostics_provider(&self, provider: Arc<dyn DiagnosticsProvider>) {
2103 self.tools.set_diagnostics_provider(provider);
2104 }
2105
2106 pub fn set_command_runner(&self, runner: Arc<dyn CommandRunner>) {
2108 self.tools.set_command_runner(runner);
2109 }
2110
2111 pub fn todos(&self) -> Vec<TodoItem> {
2113 self.tools.todos()
2114 }
2115
2116 fn active_tool_security(&self) -> ToolSecurityEngine {
2118 self.runtime_control
2119 .tool_security_override
2120 .read()
2121 .clone()
2122 .map(ToolSecurityEngine::new)
2123 .unwrap_or_else(|| self.tool_security.clone())
2124 }
2125
2126 pub fn with_process_processor(mut self, processor: ProcessProcessor) -> Self {
2127 let processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
2128 self.process_processor = Some(processor);
2129 self
2130 }
2131
2132 pub fn with_state_machine(
2133 mut self,
2134 state_machine: Arc<StateMachine>,
2135 evaluator: Arc<dyn TransitionEvaluator>,
2136 ) -> Self {
2137 self.state_machine = Some(state_machine);
2138 self.transition_evaluator = Some(evaluator);
2139 self
2140 }
2141
2142 pub fn with_context_manager(mut self, manager: Arc<ContextManager>) -> Self {
2143 self.context_manager = manager;
2144 self
2145 }
2146
2147 pub fn register_message_filter(&self, name: impl Into<String>, filter: Arc<dyn MessageFilter>) {
2148 self.message_filters.write().insert(name.into(), filter);
2149 }
2150
2151 pub fn set_context(&self, key: &str, value: Value) -> Result<()> {
2152 self.context_manager.update(key, value)
2153 }
2154
2155 pub fn update_context(&self, path: &str, value: Value) -> Result<()> {
2156 self.context_manager.update(path, value)
2157 }
2158
2159 pub fn get_context(&self) -> HashMap<String, Value> {
2160 self.build_context_with_overlays()
2161 }
2162
2163 pub fn remove_context(&self, key: &str) -> Option<Value> {
2164 self.context_manager.remove(key)
2165 }
2166
2167 pub async fn refresh_context(&self, key: &str) -> Result<()> {
2168 self.context_manager.refresh(key).await
2169 }
2170
2171 pub fn register_context_provider(&self, name: &str, provider: Arc<dyn ContextProvider>) {
2172 self.context_manager.register_provider(name, provider);
2173 }
2174
2175 pub fn current_state(&self) -> Option<String> {
2176 self.state_machine.as_ref().map(|sm| sm.current())
2177 }
2178
2179 pub async fn transition_to(&self, state: &str) -> Result<()> {
2180 if let Some(ref sm) = self.state_machine {
2181 let from_state = sm.current();
2182 self.execute_state_exit_actions(&from_state).await;
2183 sm.transition_to(state, "manual transition")?;
2184 self.execute_state_enter_actions(state).await;
2185 info!(to = %state, "Manual state transition");
2186 }
2187 Ok(())
2188 }
2189
2190 pub fn state_history(&self) -> Vec<StateTransitionEvent> {
2191 self.state_machine
2192 .as_ref()
2193 .map(|sm| sm.history())
2194 .unwrap_or_default()
2195 }
2196
2197 pub fn session_metadata(&self) -> ai_agents_core::SessionMetadata {
2199 self.session_metadata.read().clone()
2200 }
2201
2202 pub async fn delete_actor_data(&self, actor_id: &str) -> Result<()> {
2205 let allowed = self
2206 .actor_memory_config
2207 .as_ref()
2208 .map(|c| c.privacy.allow_deletion)
2209 .unwrap_or(true);
2210 if !allowed {
2211 return Err(AgentError::Config(
2212 "privacy.allow_deletion is false; actor data deletion is not permitted".into(),
2213 ));
2214 }
2215 let storage = self.storage.read().clone();
2216 if let Some(storage) = storage {
2217 storage.delete_actor_data(&self.info.id, actor_id).await?;
2218 storage.delete_relationship(&self.info.id, actor_id).await?;
2219 } else if let Some(store) = self.fact_store.read().clone() {
2220 store.delete_actor_data(actor_id).await?;
2221 }
2222 if let Some(manager) = self.relationship_manager.as_ref() {
2223 manager.remove(actor_id);
2224 }
2225 self.actor_facts_cache.write().remove(actor_id);
2226 Ok(())
2227 }
2228
2229 pub fn set_session_metadata(&self, meta: ai_agents_core::SessionMetadata) {
2231 *self.session_metadata.write() = meta;
2232 }
2233
2234 pub async fn cleanup_expired_sessions(&self) -> Result<usize> {
2236 let storage = self.storage.read().clone();
2237 match storage {
2238 Some(s) => {
2239 let count = s.cleanup_expired().await?;
2240 if count > 0 {
2241 self.hooks.on_sessions_expired(count).await;
2242 }
2243 Ok(count)
2244 }
2245 None => Err(AgentError::Config(
2246 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2247 )),
2248 }
2249 }
2250
2251 pub async fn list_sessions_filtered(
2253 &self,
2254 filter: &ai_agents_core::SessionFilter,
2255 ) -> Result<Vec<ai_agents_core::SessionSummary>> {
2256 let storage = self.storage.read().clone();
2257 match storage {
2258 Some(s) => s.list_sessions_filtered(filter).await,
2259 None => Err(AgentError::Config(
2260 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2261 )),
2262 }
2263 }
2264
2265 pub async fn save_state(&self) -> Result<AgentSnapshot> {
2266 let memory_snapshot = self.memory.snapshot().await?;
2267 let state_machine_snapshot = self.state_machine.as_ref().map(|sm| sm.snapshot());
2268 let context_snapshot = self.context_manager.snapshot();
2269
2270 let mut snapshot = AgentSnapshot::new(self.info.id.clone())
2271 .with_memory(memory_snapshot)
2272 .with_context(context_snapshot)
2273 .with_state_machine(
2274 state_machine_snapshot.unwrap_or_else(|| StateMachineSnapshot {
2275 current_state: String::new(),
2276 previous_state: None,
2277 turn_count: 0,
2278 no_transition_count: 0,
2279 history: vec![],
2280 }),
2281 );
2282
2283 if let Some(ref persona) = self.persona_manager {
2284 snapshot.persona = Some(persona.snapshot_as_value()?);
2285 }
2286
2287 if let Some(ref relationships) = self.relationship_manager {
2288 snapshot.relationships = Some(relationships.snapshot_as_value()?);
2289 }
2290
2291 Ok(snapshot)
2292 }
2293
2294 pub async fn save_state_full(&self) -> Result<AgentSnapshot> {
2296 let mut snapshot = self.save_state().await?;
2297 if let Some(ref registry) = self.spawner_registry {
2298 let entries = registry.list_with_specs();
2299 if !entries.is_empty() {
2300 snapshot = snapshot.with_spawned_agents(entries);
2301 }
2302 }
2303 Ok(snapshot)
2304 }
2305
2306 pub async fn restore_state(&self, snapshot: AgentSnapshot) -> Result<()> {
2307 self.memory.restore(snapshot.memory).await?;
2308
2309 if let (Some(sm), Some(sm_snapshot)) = (&self.state_machine, snapshot.state_machine) {
2310 if !sm_snapshot.current_state.is_empty() {
2311 sm.restore(sm_snapshot)?;
2312 }
2313 }
2314
2315 self.context_manager.restore(snapshot.context);
2316
2317 if let (Some(persona_value), Some(persona_manager)) =
2318 (snapshot.persona, &self.persona_manager)
2319 {
2320 persona_manager.restore_from_value(persona_value)?;
2321 }
2322
2323 if let (Some(relationship_value), Some(relationship_manager)) =
2324 (snapshot.relationships, &self.relationship_manager)
2325 {
2326 relationship_manager.restore_from_value(relationship_value)?;
2327 }
2328
2329 info!(agent_id = %snapshot.agent_id, "State restored");
2330 Ok(())
2331 }
2332
2333 pub async fn save_to(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<()> {
2334 let snapshot = self.save_state().await?;
2335 storage.save(session_id, &snapshot).await
2336 }
2337
2338 pub async fn load_from(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<bool> {
2339 if let Some(snapshot) = storage.load(session_id).await? {
2340 self.restore_state(snapshot).await?;
2341 Ok(true)
2342 } else {
2343 Ok(false)
2344 }
2345 }
2346
2347 pub async fn save_session(&self, session_id: &str) -> Result<()> {
2348 let storage = self.storage.read().clone();
2349 match storage {
2350 Some(s) => {
2351 let is_new = {
2353 let cur = self.current_session_id.read().clone();
2354 cur.as_deref() != Some(session_id)
2355 };
2356 if is_new {
2357 *self.current_session_id.write() = Some(session_id.to_string());
2358 self.hooks.on_session_created(session_id).await;
2359 }
2360
2361 {
2363 let now = chrono::Utc::now();
2364 let msg_count = self
2365 .memory
2366 .get_messages(None)
2367 .await
2368 .map(|v| v.len())
2369 .unwrap_or(0);
2370 let mut meta = self.session_metadata.write();
2371 meta.last_active = now;
2372 meta.message_count = msg_count;
2373 if meta.actor_id.is_none() {
2374 meta.actor_id = self.actor_id.read().clone();
2375 }
2376 }
2377
2378 self.save_to(s.as_ref(), session_id).await?;
2379
2380 let meta = self.session_metadata.read().clone();
2382 let _ = s.save_metadata(session_id, &meta).await;
2383 Ok(())
2384 }
2385 None => Err(AgentError::Config(
2386 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2387 )),
2388 }
2389 }
2390
2391 pub async fn load_session(&self, session_id: &str) -> Result<bool> {
2392 let storage = self.storage.read().clone();
2393 match storage {
2394 Some(s) => {
2395 let loaded = self.load_from(s.as_ref(), session_id).await?;
2396 if loaded {
2397 if let Ok(Some(meta)) = s.load_metadata(session_id).await {
2399 if let Some(ref aid) = meta.actor_id {
2400 let _ = self.set_actor_id(aid);
2401 }
2402 *self.session_metadata.write() = meta;
2403 }
2404 *self.current_session_id.write() = Some(session_id.to_string());
2405 }
2406 Ok(loaded)
2407 }
2408 None => Err(AgentError::Config(
2409 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2410 )),
2411 }
2412 }
2413
2414 pub async fn delete_session(&self, session_id: &str) -> Result<()> {
2415 let storage = self.storage.read().clone();
2416 match storage {
2417 Some(s) => s.delete(session_id).await,
2418 None => Err(AgentError::Config(
2419 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2420 )),
2421 }
2422 }
2423
2424 pub async fn list_sessions(&self) -> Result<Vec<String>> {
2425 let storage = self.storage.read().clone();
2426 match storage {
2427 Some(s) => s.list_sessions().await,
2428 None => Err(AgentError::Config(
2429 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2430 )),
2431 }
2432 }
2433
2434 fn estimate_tokens(&self, text: &str) -> u32 {
2435 (text.len() as f32 / 4.0).ceil() as u32
2436 }
2437
2438 fn estimate_total_tokens(&self, messages: &[ChatMessage]) -> u32 {
2439 messages
2440 .iter()
2441 .map(|m| self.estimate_tokens(&m.content))
2442 .sum()
2443 }
2444
2445 fn truncate_context(&self, messages: &mut Vec<ChatMessage>, keep_recent: usize) {
2446 if messages.len() <= keep_recent + 1 {
2447 return;
2448 }
2449 let system_msg = messages.remove(0);
2450 let to_remove = messages.len().saturating_sub(keep_recent);
2451 messages.drain(..to_remove);
2452 messages.insert(0, system_msg);
2453 }
2454
2455 fn get_filter(&self, config: &FilterConfig) -> Arc<dyn MessageFilter> {
2456 match config {
2457 FilterConfig::KeepRecent(n) => Arc::new(KeepRecentFilter::new(*n)),
2458 FilterConfig::ByRole { keep_roles } => Arc::new(ByRoleFilter::new(keep_roles.clone())),
2459 FilterConfig::SkipPattern { skip_if_contains } => {
2460 Arc::new(SkipPatternFilter::new(skip_if_contains.clone()))
2461 }
2462 FilterConfig::Custom { name } => {
2463 let filters = self.message_filters.read();
2464 filters
2465 .get(name)
2466 .cloned()
2467 .unwrap_or_else(|| Arc::new(KeepRecentFilter::new(10)))
2468 }
2469 }
2470 }
2471
2472 async fn summarize_context(
2473 &self,
2474 messages: &mut Vec<ChatMessage>,
2475 summarizer_llm: Option<&str>,
2476 max_summary_tokens: u32,
2477 custom_prompt: Option<&str>,
2478 keep_recent: usize,
2479 filter: Option<&FilterConfig>,
2480 ) -> Result<()> {
2481 let system_msg = messages.remove(0);
2482
2483 let to_summarize_count = messages.len().saturating_sub(keep_recent);
2484 if to_summarize_count == 0 {
2485 messages.insert(0, system_msg);
2486 return Ok(());
2487 }
2488
2489 let recent_msgs: Vec<ChatMessage> = messages.drain(to_summarize_count..).collect();
2490 let mut to_summarize = std::mem::take(messages);
2491
2492 if let Some(filter_config) = filter {
2493 let filter = self.get_filter(filter_config);
2494 to_summarize = filter.filter(to_summarize);
2495 }
2496
2497 if to_summarize.is_empty() {
2498 *messages = recent_msgs;
2499 messages.insert(0, system_msg);
2500 return Ok(());
2501 }
2502
2503 let conversation_text = to_summarize
2504 .iter()
2505 .map(|m| format!("{:?}: {}", m.role, m.content))
2506 .collect::<Vec<_>>()
2507 .join("\n");
2508
2509 let default_prompt = format!(
2510 "Summarize the following conversation in under {} tokens, preserving key information:\n\n{}",
2511 max_summary_tokens, conversation_text
2512 );
2513
2514 let summary_prompt = custom_prompt
2515 .map(|p| format!("{}\n\n{}", p, conversation_text))
2516 .unwrap_or(default_prompt);
2517
2518 let summarizer = if let Some(alias) = summarizer_llm {
2519 self.llm_registry
2520 .get(alias)
2521 .map_err(|e| AgentError::Config(e.to_string()))?
2522 } else {
2523 self.llm_registry
2524 .router()
2525 .or_else(|_| self.llm_registry.default())
2526 .map_err(|e| AgentError::Config(e.to_string()))?
2527 };
2528
2529 let summary_msgs = vec![ChatMessage::user(&summary_prompt)];
2530 let response = self
2531 .observe_purpose(
2532 ObservationPurpose::Summarization,
2533 summarizer.complete(&summary_msgs, None),
2534 )
2535 .await?;
2536
2537 let summary_message = ChatMessage::system(&format!(
2538 "[Previous conversation summary]\n{}",
2539 response.content
2540 ));
2541
2542 *messages = vec![system_msg, summary_message];
2543 messages.extend(recent_msgs);
2544
2545 debug!(
2546 summarized_count = to_summarize_count,
2547 kept_recent = keep_recent,
2548 "Context summarized"
2549 );
2550
2551 Ok(())
2552 }
2553
2554 fn render_system_prompt(&self) -> Result<String> {
2555 let mut context = self.build_context_with_overlays();
2556
2557 let facts_text = self.format_actor_facts_for_context();
2559 if !facts_text.is_empty() {
2560 context.insert(
2561 "actor_facts".to_string(),
2562 serde_json::Value::String(facts_text),
2563 );
2564 }
2565
2566 if let Some((key, text)) = self.format_relationship_for_context() {
2567 context.insert(key, serde_json::Value::String(text));
2568 }
2569
2570 self.template_renderer
2571 .render(&self.base_system_prompt, &context)
2572 }
2573
2574 fn get_top_level_tool_ids(&self) -> Vec<String> {
2575 if let Some(ids) = self.runtime_control.tool_scope_override.read().clone() {
2576 return ids
2577 .iter()
2578 .filter_map(|id| self.tools.canonical_id(id))
2579 .collect();
2580 }
2581 self.declared_tool_ids
2582 .as_ref()
2583 .map(|ids| {
2584 ids.iter()
2585 .filter_map(|id| self.tools.canonical_id(id))
2586 .collect::<Vec<_>>()
2587 })
2588 .unwrap_or_default()
2589 }
2590
2591 async fn get_available_tool_ids(&self) -> Result<Vec<String>> {
2592 let top_level = self.get_top_level_tool_ids();
2593 if top_level.is_empty() {
2594 return Ok(Vec::new());
2595 }
2596
2597 match self.get_current_tool_refs() {
2598 Some(tool_refs) => {
2599 if tool_refs.is_empty() {
2600 return Ok(Vec::new());
2601 }
2602
2603 let eval_ctx = self.build_evaluation_context().await?;
2604 let llm_getter = RegistryLLMGetter {
2605 registry: self.llm_registry.clone(),
2606 };
2607 let evaluator = ConditionEvaluator::new(llm_getter);
2608
2609 let mut available = Vec::new();
2610 for tool_ref in &tool_refs {
2611 let tool_id = tool_ref.id();
2612 let Some(canonical_id) = self.tools.canonical_id(tool_id) else {
2613 continue;
2614 };
2615 if !top_level.iter().any(|id| id == &canonical_id) {
2616 debug!(tool = %canonical_id, "Tool not in top-level grant, skipping");
2617 continue;
2618 }
2619
2620 if let Some(condition) = tool_ref.condition() {
2621 match evaluator.evaluate(condition, &eval_ctx).await {
2622 Ok(true) => {
2623 available.push(canonical_id);
2624 }
2625 Ok(false) => {
2626 debug!(tool = tool_id, "Tool condition not met, skipping");
2627 }
2628 Err(e) => {
2629 warn!(tool = tool_id, error = %e, "Error evaluating tool condition");
2630 }
2631 }
2632 } else {
2633 available.push(canonical_id);
2634 }
2635 }
2636
2637 Ok(available)
2638 }
2639 None => Ok(top_level),
2640 }
2641 }
2642
2643 fn get_current_tool_refs(&self) -> Option<Vec<ToolRef>> {
2647 if let Some(ref sm) = self.state_machine {
2648 if let Some(state_def) = sm.current_definition() {
2649 let parent_def = sm.get_parent_definition();
2650 if let Some(effective) = state_def.get_effective_tools(parent_def.as_ref()) {
2651 return Some(effective.into_iter().cloned().collect());
2652 }
2653 }
2654 }
2655 None
2656 }
2657
2658 async fn build_evaluation_context(&self) -> Result<EvaluationContext> {
2659 let context = self.build_context_with_overlays();
2660 let messages = self.memory.get_messages(Some(10)).await?;
2661 let tool_history = self.tool_call_history.read().clone();
2662
2663 let (state_name, turn_count, previous_state) = if let Some(ref sm) = self.state_machine {
2664 (Some(sm.current()), sm.turn_count(), sm.previous())
2665 } else {
2666 (None, 0, None)
2667 };
2668
2669 Ok(EvaluationContext::default()
2670 .with_context(context)
2671 .with_state(state_name, turn_count, previous_state)
2672 .with_called_tools(tool_history)
2673 .with_messages(messages))
2674 }
2675
2676 fn record_tool_call(&self, tool_id: &str, result: Value) {
2677 self.tool_call_history.write().push(ToolCallRecord {
2678 tool_id: tool_id.to_string(),
2679 result,
2680 timestamp: chrono::Utc::now(),
2681 });
2682 }
2683
2684 async fn get_effective_system_prompt_with_persona_hooks(
2685 &self,
2686 fire_persona_hooks: bool,
2687 ) -> Result<String> {
2688 let rendered_base = self.render_system_prompt()?;
2689
2690 let persona_prefix = if let Some(ref persona) = self.persona_manager {
2691 let context = self.build_context_with_overlays();
2692 if fire_persona_hooks {
2693 let render_result = persona.render_prompt(&context)?;
2694 for content in &render_result.newly_revealed {
2695 self.hooks.on_secret_revealed(content).await;
2696 }
2697 render_result.prompt
2698 } else {
2699 persona.render_prompt_preview(&context)?
2700 }
2701 } else {
2702 String::new()
2703 };
2704
2705 if let Some(ref sm) = self.state_machine {
2706 if let Some(state_def) = sm.current_definition() {
2707 let state_prompt = if let Some(ref prompt) = state_def.prompt {
2708 let context = self.build_context_with_overlays();
2709 self.template_renderer.render_with_state(
2710 prompt,
2711 &context,
2712 &sm.current(),
2713 sm.previous().as_deref(),
2714 sm.turn_count(),
2715 state_def.max_turns,
2716 )?
2717 } else {
2718 String::new()
2719 };
2720
2721 let combined = match state_def.prompt_mode {
2722 PromptMode::Append => {
2723 if state_prompt.is_empty() {
2724 rendered_base
2725 } else {
2726 format!(
2727 "{}\n\n[Current State: {}]\n{}",
2728 rendered_base,
2729 sm.current(),
2730 state_prompt
2731 )
2732 }
2733 }
2734 PromptMode::Replace => {
2735 if state_prompt.is_empty() {
2736 rendered_base
2737 } else {
2738 state_prompt
2739 }
2740 }
2741 PromptMode::Prepend => {
2742 if state_prompt.is_empty() {
2743 rendered_base
2744 } else {
2745 format!("{}\n\n{}", state_prompt, rendered_base)
2746 }
2747 }
2748 };
2749
2750 let with_persona = if persona_prefix.is_empty() {
2752 combined
2753 } else {
2754 format!("{}\n\n{}", persona_prefix, combined)
2755 };
2756
2757 let available_tool_ids = self.get_available_tool_ids().await?;
2758 if !available_tool_ids.is_empty() {
2759 let tools_prompt = self.tools.generate_scoped_prompt_with_parallel(
2760 &available_tool_ids,
2761 self.parallel_tools.enabled,
2762 );
2763 if !tools_prompt.is_empty() {
2764 return Ok(format!("{}\n\n{}", with_persona, tools_prompt));
2765 }
2766 }
2767 return Ok(with_persona);
2768 }
2769 }
2770
2771 let with_persona = if persona_prefix.is_empty() {
2773 rendered_base
2774 } else {
2775 format!("{}\n\n{}", persona_prefix, rendered_base)
2776 };
2777
2778 let available_tool_ids = self.get_available_tool_ids().await?;
2779 let tools_prompt = self
2780 .tools
2781 .generate_scoped_prompt_with_parallel(&available_tool_ids, self.parallel_tools.enabled);
2782 if !tools_prompt.is_empty() {
2783 Ok(format!("{}\n\n{}", with_persona, tools_prompt))
2784 } else {
2785 Ok(with_persona)
2786 }
2787 }
2788
2789 fn get_state_llm(&self) -> Result<Arc<dyn LLMProvider>> {
2790 if let Some(ref sm) = self.state_machine {
2791 if let Some(state_def) = sm.current_definition() {
2792 if let Some(ref llm_alias) = state_def.llm {
2793 return self
2794 .llm_registry
2795 .get(llm_alias)
2796 .map_err(|e| AgentError::Config(e.to_string()));
2797 }
2798 }
2799 }
2800 self.llm_registry
2801 .default()
2802 .map_err(|e| AgentError::Config(e.to_string()))
2803 }
2804
2805 fn get_effective_reasoning_config(&self) -> ReasoningConfig {
2806 if let Some(ref sm) = self.state_machine {
2807 if let Some(state_def) = sm.current_definition() {
2808 if let Some(ref state_reasoning) = state_def.reasoning {
2809 return state_reasoning.clone();
2810 }
2811 }
2812 }
2813 self.reasoning_config.clone()
2814 }
2815
2816 fn get_effective_reflection_config(&self) -> ReflectionConfig {
2817 if let Some(ref sm) = self.state_machine {
2818 if let Some(state_def) = sm.current_definition() {
2819 if let Some(ref state_reflection) = state_def.reflection {
2820 return state_reflection.clone();
2821 }
2822 }
2823 }
2824 self.reflection_config.clone()
2825 }
2826
2827 fn get_skill_reasoning_config(&self, skill: &SkillDefinition) -> ReasoningConfig {
2828 skill
2829 .reasoning
2830 .clone()
2831 .unwrap_or_else(|| self.get_effective_reasoning_config())
2832 }
2833
2834 fn get_skill_reflection_config(&self, skill: &SkillDefinition) -> ReflectionConfig {
2835 skill
2836 .reflection
2837 .clone()
2838 .unwrap_or_else(|| self.get_effective_reflection_config())
2839 }
2840
2841 async fn build_disambiguation_context(&self) -> Result<DisambiguationContext> {
2842 let recent_messages: Vec<String> = self
2843 .memory
2844 .get_messages(Some(5))
2845 .await?
2846 .iter()
2847 .rev()
2848 .map(|m| format!("{:?}: {}", m.role, m.content))
2849 .collect();
2850
2851 let current_state = self.current_state().map(|s| s.to_string());
2852
2853 let state_prompt: Option<String> = self
2856 .state_machine
2857 .as_ref()
2858 .and_then(|sm| sm.current_definition())
2859 .and_then(|def| def.prompt.clone());
2860
2861 let available_tools: Vec<String> = self
2862 .get_available_tool_ids()
2863 .await
2864 .unwrap_or_else(|_| self.tools.list_ids());
2865
2866 let available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
2867
2868 let user_context = self.build_context_with_overlays();
2869
2870 let available_intents: Vec<String> = if let Some(ref sm) = self.state_machine {
2872 sm.current_definition()
2873 .map(|def| {
2874 def.transitions
2875 .iter()
2876 .filter_map(|t| t.intent.clone())
2877 .collect()
2878 })
2879 .unwrap_or_default()
2880 } else {
2881 Vec::new()
2882 };
2883
2884 Ok(DisambiguationContext::from_agent_state(
2885 recent_messages,
2886 current_state,
2887 state_prompt,
2888 available_tools,
2889 available_skills,
2890 available_intents,
2891 user_context,
2892 ))
2893 }
2894
2895 fn get_available_skills(&self) -> Vec<&SkillDefinition> {
2896 if let Some(ref sm) = self.state_machine {
2897 if let Some(state_def) = sm.current_definition() {
2898 let parent_def = sm.get_parent_definition();
2899 let effective_skills = state_def.get_effective_skills(parent_def.as_ref());
2900 if !effective_skills.is_empty() {
2901 return self
2902 .skills
2903 .iter()
2904 .filter(|s| effective_skills.contains(&&s.id))
2905 .collect();
2906 }
2907 }
2908 }
2909 self.skills.iter().collect()
2910 }
2911
2912 async fn build_messages(&self) -> Result<Vec<ChatMessage>> {
2913 self.build_messages_internal(true, None).await
2914 }
2915
2916 async fn build_messages_for_draft(&self, user_message: &str) -> Result<Vec<ChatMessage>> {
2917 self.build_messages_internal(false, Some(user_message))
2918 .await
2919 }
2920
2921 async fn build_messages_internal(
2922 &self,
2923 fire_persona_hooks: bool,
2924 ephemeral_user_message: Option<&str>,
2925 ) -> Result<Vec<ChatMessage>> {
2926 let system_prompt = self
2927 .get_effective_system_prompt_with_persona_hooks(fire_persona_hooks)
2928 .await?;
2929 let mut messages = vec![ChatMessage::system(&system_prompt)];
2930
2931 let context = self.memory.get_context().await?;
2932 let history = if let Some(ref budget) = self.memory_token_budget {
2933 context.to_llm_messages_with_allocation(&budget.allocation)
2934 } else {
2935 context.to_llm_messages()
2936 };
2937 messages.extend(history);
2938 if let Some(user_message) = ephemeral_user_message {
2939 messages.push(ChatMessage::user(user_message));
2940 }
2941
2942 let total_tokens = self.estimate_total_tokens(&messages);
2943
2944 if total_tokens > self.max_context_tokens {
2945 debug!(
2946 total = total_tokens,
2947 limit = self.max_context_tokens,
2948 "Context overflow"
2949 );
2950
2951 match &self.recovery_manager.config().llm.on_context_overflow {
2952 ContextOverflowAction::Error => {
2953 return Err(AgentError::LLM(format!(
2954 "Context overflow: {} tokens > {} limit",
2955 total_tokens, self.max_context_tokens
2956 )));
2957 }
2958 ContextOverflowAction::Truncate { keep_recent } => {
2959 self.truncate_context(&mut messages, *keep_recent);
2960 }
2961 ContextOverflowAction::Summarize {
2962 summarizer_llm,
2963 max_summary_tokens,
2964 custom_prompt,
2965 keep_recent,
2966 filter,
2967 } => {
2968 self.summarize_context(
2969 &mut messages,
2970 summarizer_llm.as_deref(),
2971 *max_summary_tokens,
2972 custom_prompt.as_deref(),
2973 *keep_recent,
2974 filter.as_ref(),
2975 )
2976 .await?;
2977 }
2978 }
2979 }
2980
2981 Ok(messages)
2982 }
2983
2984 fn parse_tool_calls(&self, content: &str) -> Option<Vec<ToolCall>> {
2985 if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(content) {
2987 if let Some(arr) = parsed.as_array() {
2989 let calls: Vec<ToolCall> = arr
2990 .iter()
2991 .filter_map(|v| self.extract_tool_call_from_value(v))
2992 .collect();
2993 if !calls.is_empty() {
2994 return Some(calls);
2995 }
2996 }
2997 if let Some(tool_call) = self.extract_tool_call_from_value(&parsed) {
2999 return Some(vec![tool_call]);
3000 }
3001 }
3002
3003 if let Some(json_str) = self.extract_json_from_content(content) {
3005 if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(&json_str) {
3006 if let Some(arr) = parsed.as_array() {
3008 let calls: Vec<ToolCall> = arr
3009 .iter()
3010 .filter_map(|v| self.extract_tool_call_from_value(v))
3011 .collect();
3012 if !calls.is_empty() {
3013 return Some(calls);
3014 }
3015 }
3016 if let Some(tool_call) = self.extract_tool_call_from_value(&parsed) {
3018 return Some(vec![tool_call]);
3019 }
3020 }
3021 }
3022
3023 None
3024 }
3025
3026 fn extract_tool_call_from_value(&self, parsed: &serde_json::Value) -> Option<ToolCall> {
3027 if let Some(tool_name) = parsed.get("tool").and_then(|v| v.as_str()) {
3028 let arguments = parsed
3029 .get("arguments")
3030 .cloned()
3031 .unwrap_or(serde_json::json!({}));
3032 return Some(ToolCall {
3033 id: uuid::Uuid::new_v4().to_string(),
3034 name: tool_name.to_string(),
3035 arguments,
3036 });
3037 }
3038 None
3039 }
3040
3041 fn extract_json_from_content(&self, content: &str) -> Option<String> {
3043 if let Some(result) = self.extract_json_array_from_content(content) {
3045 return Some(result);
3046 }
3047 self.extract_json_object_from_content(content)
3048 }
3049
3050 fn extract_json_array_from_content(&self, content: &str) -> Option<String> {
3052 let start = content.find('[')?;
3053 let content_from_start = &content[start..];
3054
3055 let mut depth = 0;
3056 let mut end = 0;
3057 for (i, ch) in content_from_start.char_indices() {
3058 match ch {
3059 '[' => depth += 1,
3060 ']' => {
3061 depth -= 1;
3062 if depth == 0 {
3063 end = i + 1;
3064 break;
3065 }
3066 }
3067 _ => {}
3068 }
3069 }
3070
3071 if end > 0 {
3072 let json_str = &content_from_start[..end];
3073 if json_str.contains("\"tool\"") {
3075 return Some(json_str.to_string());
3076 }
3077 }
3078
3079 None
3080 }
3081
3082 fn extract_json_object_from_content(&self, content: &str) -> Option<String> {
3084 let start = content.find('{')?;
3085 let content_from_start = &content[start..];
3086
3087 let mut depth = 0;
3089 let mut end = 0;
3090 for (i, ch) in content_from_start.char_indices() {
3091 match ch {
3092 '{' => depth += 1,
3093 '}' => {
3094 depth -= 1;
3095 if depth == 0 {
3096 end = i + 1;
3097 break;
3098 }
3099 }
3100 _ => {}
3101 }
3102 }
3103
3104 if end > 0 {
3105 let json_str = &content_from_start[..end];
3106 if json_str.contains("\"tool\"") {
3108 return Some(json_str.to_string());
3109 }
3110 }
3111
3112 None
3113 }
3114
3115 fn record_from_parts(
3117 &self,
3118 request: &ToolExecutionRequest,
3119 canonical_id: String,
3120 executed_arguments: Value,
3121 started_at: chrono::DateTime<chrono::Utc>,
3122 start: Instant,
3123 executed: bool,
3124 success: bool,
3125 output: String,
3126 metadata: HashMap<String, Value>,
3127 policy: ToolPolicyDecisionRecord,
3128 approval: Option<ToolApprovalRecord>,
3129 timed_out: bool,
3130 output_truncated: bool,
3131 ) -> ToolExecutionRecord {
3132 ToolExecutionRecord {
3133 call_id: request.call_id.clone(),
3134 requested_name: request.requested_name.clone(),
3135 canonical_id,
3136 source: request.source.clone(),
3137 arguments: request.arguments.clone(),
3138 executed_arguments,
3139 policy_version: self.active_tool_security().policy_version(),
3140 registry_version: self.tools.version(),
3141 runtime_config_version: self.runtime_control.version.load(Ordering::SeqCst),
3142 executed,
3143 success,
3144 output,
3145 metadata,
3146 policy,
3147 approval,
3148 started_at,
3149 duration_ms: start.elapsed().as_millis() as u64,
3150 timed_out,
3151 cancelled: false,
3152 cancellation_reason: None,
3153 output_truncated,
3154 }
3155 }
3156
3157 async fn finish_tool_record(&self, record: &ToolExecutionRecord) {
3159 let result = ToolResult {
3160 success: record.success,
3161 output: record.model_output_string(),
3162 metadata: if record.metadata.is_empty() {
3163 None
3164 } else {
3165 Some(record.metadata.clone())
3166 },
3167 };
3168 self.hooks
3169 .on_tool_complete(&record.canonical_id, &result, record.duration_ms)
3170 .await;
3171 self.hooks.on_tool_execution_record(record).await;
3172 self.record_tool_call(&record.canonical_id, record.model_output_value());
3173 if !record.success {
3174 self.hooks
3175 .on_error(&AgentError::Tool(record.output.clone()))
3176 .await;
3177 }
3178 }
3179
3180 async fn execute_resolved_tool_once(
3182 &self,
3183 tool: Arc<dyn ai_agents_core::Tool>,
3184 args: Value,
3185 ctx: ToolExecutionContext,
3186 timeout_ms: u64,
3187 ) -> Result<(ToolResult, bool, bool)> {
3188 let actor_context = current_turn_actor_context();
3189 let future = async move {
3190 if let Some(actor_context) = actor_context {
3191 scope_actor_context(actor_context, tool.execute(args, ctx)).await
3192 } else {
3193 tool.execute(args, ctx).await
3194 }
3195 };
3196 tokio::pin!(future);
3197 let timeout = tokio::time::sleep(std::time::Duration::from_millis(timeout_ms));
3198 tokio::pin!(timeout);
3199 let mut cancel_tick = tokio::time::interval(std::time::Duration::from_millis(50));
3200
3201 loop {
3202 tokio::select! {
3203 result = &mut future => return Ok((result, false, false)),
3204 _ = &mut timeout => return Ok((ToolResult::error("Tool execution timed out"), true, false)),
3205 _ = cancel_tick.tick() => {
3206 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
3207 return Ok((ToolResult::error("Tool execution cancelled by runtime control"), false, true));
3208 }
3209 }
3210 }
3211 }
3212 }
3213
3214 fn truncate_tool_output(output: String, max_chars: Option<usize>) -> (String, bool) {
3216 let Some(max_chars) = max_chars else {
3217 return (output, false);
3218 };
3219 let mut chars = output.chars();
3220 let truncated: String = chars.by_ref().take(max_chars).collect();
3221 if chars.next().is_some() {
3222 (truncated, true)
3223 } else {
3224 (output, false)
3225 }
3226 }
3227
3228 async fn acquire_tool_resource_lock(
3230 &self,
3231 canonical_id: &str,
3232 args: &Value,
3233 classification: &ai_agents_core::ToolCallClassification,
3234 ) -> Option<tokio::sync::OwnedMutexGuard<()>> {
3235 if classification.read_only && classification.concurrency_safe {
3236 return None;
3237 }
3238 let key = tool_resource_lock_key(canonical_id, args);
3239 let lock = {
3240 let mut locks = self.resource_locks.write();
3241 locks
3242 .entry(key)
3243 .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
3244 .clone()
3245 };
3246 Some(lock.lock_owned().await)
3247 }
3248
3249 async fn run_tool_with_retries(
3251 &self,
3252 canonical_id: &str,
3253 tool: Arc<dyn ai_agents_core::Tool>,
3254 args: Value,
3255 ctx: ToolExecutionContext,
3256 timeout_ms: u64,
3257 max_retries: u32,
3258 ) -> Result<(ToolResult, bool, bool)> {
3259 let max_retries = if ctx.classification.safely_retryable {
3260 max_retries
3261 } else {
3262 0
3263 };
3264 let mut attempts = 0;
3265 loop {
3266 let (result, timed_out, cancelled) = self
3267 .execute_resolved_tool_once(tool.clone(), args.clone(), ctx.clone(), timeout_ms)
3268 .await?;
3269 if result.success || timed_out || cancelled || attempts >= max_retries {
3270 return Ok((result, timed_out, cancelled));
3271 }
3272 attempts += 1;
3273 warn!(tool = %canonical_id, attempt = attempts, error = %result.output, "Retrying failed tool call");
3274 }
3275 }
3276
3277 async fn execute_tool_record(
3279 &self,
3280 request: ToolExecutionRequest,
3281 ) -> Result<ToolExecutionRecord> {
3282 let started_at = chrono::Utc::now();
3283 let start = Instant::now();
3284 info!(tool = %request.requested_name, args = %request.arguments, "Executing tool");
3285
3286 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
3287 let record = self.record_from_parts(
3288 &request,
3289 request.requested_name.clone(),
3290 request.arguments.clone(),
3291 started_at,
3292 start,
3293 false,
3294 false,
3295 "Tool execution is disabled by runtime control".to_string(),
3296 HashMap::new(),
3297 ToolPolicyDecisionRecord::deny("runtime emergency deny is enabled"),
3298 None,
3299 false,
3300 false,
3301 );
3302 self.finish_tool_record(&record).await;
3303 return Ok(record);
3304 }
3305
3306 let Some(resolved) = self.tools.resolve(&request.requested_name) else {
3307 let record = self.record_from_parts(
3308 &request,
3309 request.requested_name.clone(),
3310 request.arguments.clone(),
3311 started_at,
3312 start,
3313 false,
3314 false,
3315 format!("Tool '{}' is unavailable", request.requested_name),
3316 HashMap::new(),
3317 ToolPolicyDecisionRecord::unavailable(format!(
3318 "Tool '{}' is not registered",
3319 request.requested_name
3320 )),
3321 None,
3322 false,
3323 false,
3324 );
3325 self.finish_tool_record(&record).await;
3326 return Ok(record);
3327 };
3328
3329 let canonical_id = resolved.identity.canonical_id.clone();
3330 let available_tool_ids = self.get_available_tool_ids().await?;
3331 if !available_tool_ids.iter().any(|id| id == &canonical_id) {
3332 let record = self.record_from_parts(
3333 &request,
3334 canonical_id.clone(),
3335 request.arguments.clone(),
3336 started_at,
3337 start,
3338 false,
3339 false,
3340 format!(
3341 "Tool '{}' is not available in the current scope",
3342 canonical_id
3343 ),
3344 HashMap::new(),
3345 ToolPolicyDecisionRecord::deny(format!(
3346 "Tool '{}' is not granted by the current top-level and state tool scope",
3347 canonical_id
3348 )),
3349 None,
3350 false,
3351 false,
3352 );
3353 self.finish_tool_record(&record).await;
3354 return Ok(record);
3355 }
3356
3357 let security_engine = self.active_tool_security();
3358 let bindings = resolved.tool.policy_bindings();
3359 let mut executed_arguments = security_engine.prepare_tool_arguments_with_bindings(
3360 &canonical_id,
3361 &request.arguments,
3362 &bindings,
3363 );
3364 self.hooks
3365 .on_tool_start(&canonical_id, &executed_arguments)
3366 .await;
3367
3368 let mut metadata = HashMap::new();
3369 let safety = resolved.tool.safety_metadata();
3370 let classification = resolved.tool.classify_call(&executed_arguments);
3371 let limits = security_engine.effective_limits(&canonical_id, &safety, &classification);
3372 metadata.insert(
3373 "classification".to_string(),
3374 serde_json::to_value(&classification).unwrap_or(Value::Null),
3375 );
3376 metadata.insert(
3377 "effective_limits".to_string(),
3378 serde_json::to_value(&limits).unwrap_or(Value::Null),
3379 );
3380 let policy_snapshot = security_engine.policy_snapshot(&canonical_id);
3381 if !policy_snapshot.is_null() {
3382 metadata.insert("policy_snapshot".to_string(), policy_snapshot.clone());
3383 }
3384
3385 let mut approval_record = Some(ToolApprovalRecord {
3386 status: ToolApprovalStatus::NotRequired,
3387 reason: None,
3388 modified_arguments: None,
3389 });
3390
3391 let mut security_result = security_engine
3392 .check_tool_execution_with_bindings(&canonical_id, &executed_arguments, &bindings)
3393 .await?;
3394 match &security_result {
3395 SecurityCheckResult::Allow => {}
3396 SecurityCheckResult::Warn { message } => {
3397 warn!(tool = %canonical_id, message = %message, "Tool security warning");
3398 }
3399 SecurityCheckResult::Block { reason } => {
3400 let record = self.record_from_parts(
3401 &request,
3402 canonical_id,
3403 executed_arguments,
3404 started_at,
3405 start,
3406 false,
3407 false,
3408 format!("Denied: {}", reason),
3409 metadata,
3410 ToolPolicyDecisionRecord::deny(reason.clone()),
3411 approval_record,
3412 false,
3413 false,
3414 );
3415 self.finish_tool_record(&record).await;
3416 return Ok(record);
3417 }
3418 SecurityCheckResult::Unavailable { reason } => {
3419 let record = self.record_from_parts(
3420 &request,
3421 canonical_id,
3422 executed_arguments,
3423 started_at,
3424 start,
3425 false,
3426 false,
3427 format!("Unavailable: {}", reason),
3428 metadata,
3429 ToolPolicyDecisionRecord::unavailable(reason.clone()),
3430 approval_record,
3431 false,
3432 false,
3433 );
3434 self.finish_tool_record(&record).await;
3435 return Ok(record);
3436 }
3437 SecurityCheckResult::RequireConfirmation { message } => {
3438 if self.hitl_engine.is_none() {
3439 approval_record = Some(ToolApprovalRecord {
3440 status: ToolApprovalStatus::Unavailable,
3441 reason: Some("No HITL engine configured".to_string()),
3442 modified_arguments: None,
3443 });
3444 let record = self.record_from_parts(
3445 &request,
3446 canonical_id,
3447 executed_arguments,
3448 started_at,
3449 start,
3450 false,
3451 false,
3452 format!("Approval unavailable: {}", message),
3453 metadata,
3454 ToolPolicyDecisionRecord::approval(message.clone()),
3455 approval_record,
3456 false,
3457 false,
3458 );
3459 self.finish_tool_record(&record).await;
3460 return Ok(record);
3461 }
3462
3463 let check_result = HITLCheckResult::required(
3464 ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
3465 HashMap::new(),
3466 message.clone(),
3467 None,
3468 );
3469 match self.request_hitl_approval(check_result).await? {
3470 ApprovalResult::Approved => {
3471 approval_record = Some(ToolApprovalRecord {
3472 status: ToolApprovalStatus::Approved,
3473 reason: None,
3474 modified_arguments: None,
3475 });
3476 }
3477 ApprovalResult::Modified { changes } => {
3478 if let Some(obj) = executed_arguments.as_object_mut() {
3479 for (key, value) in changes {
3480 obj.insert(key, value);
3481 }
3482 }
3483 security_result = security_engine
3484 .check_tool_execution_with_bindings(
3485 &canonical_id,
3486 &executed_arguments,
3487 &bindings,
3488 )
3489 .await?;
3490 if !matches!(
3491 security_result,
3492 SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
3493 ) {
3494 let reason = security_result
3495 .reason()
3496 .unwrap_or("modified arguments failed policy")
3497 .to_string();
3498 let record = self.record_from_parts(
3499 &request,
3500 canonical_id,
3501 executed_arguments.clone(),
3502 started_at,
3503 start,
3504 false,
3505 false,
3506 reason.clone(),
3507 metadata,
3508 ToolPolicyDecisionRecord::deny(reason),
3509 Some(ToolApprovalRecord {
3510 status: ToolApprovalStatus::Modified,
3511 reason: None,
3512 modified_arguments: Some(executed_arguments),
3513 }),
3514 false,
3515 false,
3516 );
3517 self.finish_tool_record(&record).await;
3518 return Ok(record);
3519 }
3520 approval_record = Some(ToolApprovalRecord {
3521 status: ToolApprovalStatus::Modified,
3522 reason: None,
3523 modified_arguments: Some(executed_arguments.clone()),
3524 });
3525 }
3526 ApprovalResult::Rejected { reason } => {
3527 let reason = reason.unwrap_or_else(|| "rejected".to_string());
3528 approval_record = Some(ToolApprovalRecord {
3529 status: ToolApprovalStatus::Rejected,
3530 reason: Some(reason.clone()),
3531 modified_arguments: None,
3532 });
3533 let record = self.record_from_parts(
3534 &request,
3535 canonical_id,
3536 executed_arguments,
3537 started_at,
3538 start,
3539 false,
3540 false,
3541 format!("Approval rejected: {}", reason),
3542 metadata,
3543 ToolPolicyDecisionRecord::approval(reason),
3544 approval_record,
3545 false,
3546 false,
3547 );
3548 self.finish_tool_record(&record).await;
3549 return Ok(record);
3550 }
3551 ApprovalResult::Timeout => {
3552 approval_record = Some(ToolApprovalRecord {
3553 status: ToolApprovalStatus::Timeout,
3554 reason: Some("approval timeout".to_string()),
3555 modified_arguments: None,
3556 });
3557 let record = self.record_from_parts(
3558 &request,
3559 canonical_id,
3560 executed_arguments,
3561 started_at,
3562 start,
3563 false,
3564 false,
3565 "Approval timed out".to_string(),
3566 metadata,
3567 ToolPolicyDecisionRecord::approval("approval timeout"),
3568 approval_record,
3569 false,
3570 false,
3571 );
3572 self.finish_tool_record(&record).await;
3573 return Ok(record);
3574 }
3575 }
3576 }
3577 }
3578
3579 if canonical_id == "command" && !self.tools.command_runner_available() {
3580 let record = self.record_from_parts(
3581 &request,
3582 canonical_id.clone(),
3583 executed_arguments.clone(),
3584 started_at,
3585 start,
3586 false,
3587 false,
3588 "Command runner is unavailable".to_string(),
3589 metadata,
3590 ToolPolicyDecisionRecord::unavailable("command runner is unavailable"),
3591 Some(ToolApprovalRecord {
3592 status: ToolApprovalStatus::Unavailable,
3593 reason: Some("command runner is unavailable".to_string()),
3594 modified_arguments: None,
3595 }),
3596 false,
3597 false,
3598 );
3599 self.finish_tool_record(&record).await;
3600 return Ok(record);
3601 }
3602
3603 if approval_record
3604 .as_ref()
3605 .is_some_and(|record| matches!(record.status, ToolApprovalStatus::NotRequired))
3606 {
3607 if let Some(message) =
3608 security_engine.classification_approval_message(&canonical_id, &classification)
3609 {
3610 if self.hitl_engine.is_none() {
3611 approval_record = Some(ToolApprovalRecord {
3612 status: ToolApprovalStatus::Unavailable,
3613 reason: Some("No HITL engine configured".to_string()),
3614 modified_arguments: None,
3615 });
3616 let record = self.record_from_parts(
3617 &request,
3618 canonical_id,
3619 executed_arguments,
3620 started_at,
3621 start,
3622 false,
3623 false,
3624 format!("Approval unavailable: {}", message),
3625 metadata,
3626 ToolPolicyDecisionRecord::approval(message),
3627 approval_record,
3628 false,
3629 false,
3630 );
3631 self.finish_tool_record(&record).await;
3632 return Ok(record);
3633 }
3634 let check_result = HITLCheckResult::required(
3635 ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
3636 HashMap::new(),
3637 message.clone(),
3638 None,
3639 );
3640 match self.request_hitl_approval(check_result).await? {
3641 ApprovalResult::Approved => {
3642 approval_record = Some(ToolApprovalRecord {
3643 status: ToolApprovalStatus::Approved,
3644 reason: None,
3645 modified_arguments: None,
3646 });
3647 }
3648 ApprovalResult::Modified { changes } => {
3649 if let Some(obj) = executed_arguments.as_object_mut() {
3650 for (key, value) in changes {
3651 obj.insert(key, value);
3652 }
3653 }
3654 let modified_security = security_engine
3655 .check_tool_execution_with_bindings(
3656 &canonical_id,
3657 &executed_arguments,
3658 &bindings,
3659 )
3660 .await?;
3661 if !matches!(
3662 modified_security,
3663 SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
3664 ) {
3665 let reason = modified_security
3666 .reason()
3667 .unwrap_or("modified arguments failed policy")
3668 .to_string();
3669 let record = self.record_from_parts(
3670 &request,
3671 canonical_id,
3672 executed_arguments.clone(),
3673 started_at,
3674 start,
3675 false,
3676 false,
3677 reason.clone(),
3678 metadata,
3679 ToolPolicyDecisionRecord::deny(reason),
3680 Some(ToolApprovalRecord {
3681 status: ToolApprovalStatus::Modified,
3682 reason: None,
3683 modified_arguments: Some(executed_arguments),
3684 }),
3685 false,
3686 false,
3687 );
3688 self.finish_tool_record(&record).await;
3689 return Ok(record);
3690 }
3691 approval_record = Some(ToolApprovalRecord {
3692 status: ToolApprovalStatus::Modified,
3693 reason: None,
3694 modified_arguments: Some(executed_arguments.clone()),
3695 });
3696 }
3697 ApprovalResult::Rejected { reason } => {
3698 let reason = reason.unwrap_or_else(|| "rejected".to_string());
3699 let record = self.record_from_parts(
3700 &request,
3701 canonical_id,
3702 executed_arguments,
3703 started_at,
3704 start,
3705 false,
3706 false,
3707 format!("Approval rejected: {}", reason),
3708 metadata,
3709 ToolPolicyDecisionRecord::approval(reason.clone()),
3710 Some(ToolApprovalRecord {
3711 status: ToolApprovalStatus::Rejected,
3712 reason: Some(reason),
3713 modified_arguments: None,
3714 }),
3715 false,
3716 false,
3717 );
3718 self.finish_tool_record(&record).await;
3719 return Ok(record);
3720 }
3721 ApprovalResult::Timeout => {
3722 let record = self.record_from_parts(
3723 &request,
3724 canonical_id,
3725 executed_arguments,
3726 started_at,
3727 start,
3728 false,
3729 false,
3730 "Approval timed out".to_string(),
3731 metadata,
3732 ToolPolicyDecisionRecord::approval("approval timeout"),
3733 Some(ToolApprovalRecord {
3734 status: ToolApprovalStatus::Timeout,
3735 reason: Some("approval timeout".to_string()),
3736 modified_arguments: None,
3737 }),
3738 false,
3739 false,
3740 );
3741 self.finish_tool_record(&record).await;
3742 return Ok(record);
3743 }
3744 }
3745 }
3746 }
3747
3748 if canonical_id == "diagnostics" && !self.tools.diagnostics_available() {
3749 let record = self.record_from_parts(
3750 &request,
3751 canonical_id.clone(),
3752 executed_arguments.clone(),
3753 started_at,
3754 start,
3755 false,
3756 false,
3757 "Diagnostics provider is unavailable".to_string(),
3758 metadata,
3759 ToolPolicyDecisionRecord::unavailable("diagnostics provider is unavailable"),
3760 Some(ToolApprovalRecord {
3761 status: ToolApprovalStatus::Unavailable,
3762 reason: Some("diagnostics provider is unavailable".to_string()),
3763 modified_arguments: None,
3764 }),
3765 false,
3766 false,
3767 );
3768 self.finish_tool_record(&record).await;
3769 return Ok(record);
3770 }
3771
3772 let hitl_lang_ctx = self.build_hitl_language_context();
3773 if let Some(ref hitl_engine) = self.hitl_engine {
3774 let check_result = self
3775 .observe_purpose(
3776 ObservationPurpose::HitlLocalization,
3777 hitl_engine.check_tool_with_localization(
3778 &canonical_id,
3779 &executed_arguments,
3780 &hitl_lang_ctx,
3781 self.approval_handler.as_ref(),
3782 Some(&self.llm_registry),
3783 ),
3784 )
3785 .await?;
3786 if check_result.is_required() {
3787 match self.request_hitl_approval(check_result).await? {
3788 ApprovalResult::Approved => {
3789 approval_record = Some(ToolApprovalRecord {
3790 status: ToolApprovalStatus::Approved,
3791 reason: None,
3792 modified_arguments: None,
3793 });
3794 }
3795 ApprovalResult::Modified { changes } => {
3796 if let Some(obj) = executed_arguments.as_object_mut() {
3797 for (key, value) in changes {
3798 obj.insert(key, value);
3799 }
3800 }
3801 let modified_security = security_engine
3802 .check_tool_execution_with_bindings(
3803 &canonical_id,
3804 &executed_arguments,
3805 &bindings,
3806 )
3807 .await?;
3808 if !matches!(
3809 modified_security,
3810 SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
3811 ) {
3812 let reason = modified_security
3813 .reason()
3814 .unwrap_or("modified arguments failed policy")
3815 .to_string();
3816 let record = self.record_from_parts(
3817 &request,
3818 canonical_id,
3819 executed_arguments.clone(),
3820 started_at,
3821 start,
3822 false,
3823 false,
3824 reason.clone(),
3825 metadata,
3826 ToolPolicyDecisionRecord::deny(reason),
3827 Some(ToolApprovalRecord {
3828 status: ToolApprovalStatus::Modified,
3829 reason: None,
3830 modified_arguments: Some(executed_arguments),
3831 }),
3832 false,
3833 false,
3834 );
3835 self.finish_tool_record(&record).await;
3836 return Ok(record);
3837 }
3838 approval_record = Some(ToolApprovalRecord {
3839 status: ToolApprovalStatus::Modified,
3840 reason: None,
3841 modified_arguments: Some(executed_arguments.clone()),
3842 });
3843 }
3844 ApprovalResult::Rejected { reason } => {
3845 let reason = reason.unwrap_or_else(|| "rejected".to_string());
3846 let record = self.record_from_parts(
3847 &request,
3848 canonical_id,
3849 executed_arguments,
3850 started_at,
3851 start,
3852 false,
3853 false,
3854 format!("Approval rejected: {}", reason),
3855 metadata,
3856 ToolPolicyDecisionRecord::approval(reason.clone()),
3857 Some(ToolApprovalRecord {
3858 status: ToolApprovalStatus::Rejected,
3859 reason: Some(reason),
3860 modified_arguments: None,
3861 }),
3862 false,
3863 false,
3864 );
3865 self.finish_tool_record(&record).await;
3866 return Ok(record);
3867 }
3868 ApprovalResult::Timeout => {
3869 let record = self.record_from_parts(
3870 &request,
3871 canonical_id,
3872 executed_arguments,
3873 started_at,
3874 start,
3875 false,
3876 false,
3877 "Approval timed out".to_string(),
3878 metadata,
3879 ToolPolicyDecisionRecord::approval("approval timeout"),
3880 Some(ToolApprovalRecord {
3881 status: ToolApprovalStatus::Timeout,
3882 reason: Some("approval timeout".to_string()),
3883 modified_arguments: None,
3884 }),
3885 false,
3886 false,
3887 );
3888 self.finish_tool_record(&record).await;
3889 return Ok(record);
3890 }
3891 }
3892 }
3893
3894 let condition_check = self
3895 .observe_purpose(
3896 ObservationPurpose::HitlLocalization,
3897 hitl_engine.check_conditions_with_localization(
3898 &executed_arguments,
3899 &hitl_lang_ctx,
3900 self.approval_handler.as_ref(),
3901 Some(&self.llm_registry),
3902 ),
3903 )
3904 .await?;
3905 if condition_check.is_required() {
3906 match self.request_hitl_approval(condition_check).await? {
3907 ApprovalResult::Approved => {}
3908 ApprovalResult::Modified { changes } => {
3909 if let Some(obj) = executed_arguments.as_object_mut() {
3910 for (key, value) in changes {
3911 obj.insert(key, value);
3912 }
3913 }
3914 let modified_security = security_engine
3915 .check_tool_execution_with_bindings(
3916 &canonical_id,
3917 &executed_arguments,
3918 &bindings,
3919 )
3920 .await?;
3921 if !matches!(
3922 modified_security,
3923 SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
3924 ) {
3925 let reason = modified_security
3926 .reason()
3927 .unwrap_or("modified arguments failed policy")
3928 .to_string();
3929 let record = self.record_from_parts(
3930 &request,
3931 canonical_id,
3932 executed_arguments,
3933 started_at,
3934 start,
3935 false,
3936 false,
3937 reason.clone(),
3938 metadata,
3939 ToolPolicyDecisionRecord::deny(reason),
3940 approval_record,
3941 false,
3942 false,
3943 );
3944 self.finish_tool_record(&record).await;
3945 return Ok(record);
3946 }
3947 }
3948 ApprovalResult::Rejected { reason } => {
3949 let reason = reason.unwrap_or_else(|| "rejected".to_string());
3950 let record = self.record_from_parts(
3951 &request,
3952 canonical_id,
3953 executed_arguments,
3954 started_at,
3955 start,
3956 false,
3957 false,
3958 format!("Approval rejected: {}", reason),
3959 metadata,
3960 ToolPolicyDecisionRecord::approval(reason.clone()),
3961 Some(ToolApprovalRecord {
3962 status: ToolApprovalStatus::Rejected,
3963 reason: Some(reason),
3964 modified_arguments: None,
3965 }),
3966 false,
3967 false,
3968 );
3969 self.finish_tool_record(&record).await;
3970 return Ok(record);
3971 }
3972 ApprovalResult::Timeout => {
3973 let record = self.record_from_parts(
3974 &request,
3975 canonical_id,
3976 executed_arguments,
3977 started_at,
3978 start,
3979 false,
3980 false,
3981 "Approval timed out".to_string(),
3982 metadata,
3983 ToolPolicyDecisionRecord::approval("approval timeout"),
3984 Some(ToolApprovalRecord {
3985 status: ToolApprovalStatus::Timeout,
3986 reason: Some("approval timeout".to_string()),
3987 modified_arguments: None,
3988 }),
3989 false,
3990 false,
3991 );
3992 self.finish_tool_record(&record).await;
3993 return Ok(record);
3994 }
3995 }
3996 }
3997 }
3998
3999 let tool_config = self.recovery_manager.get_tool_config(&canonical_id);
4000 let timeout_ms = limits
4001 .timeout_ms
4002 .unwrap_or_else(|| security_engine.get_tool_timeout(&canonical_id));
4003 let deadline = Some(started_at + chrono::Duration::milliseconds(timeout_ms as i64));
4004 let turn_actor = current_turn_actor_context();
4005 let actor = ToolActorContext {
4006 actor_id: turn_actor
4007 .as_ref()
4008 .and_then(|context| context.effective_actor_id().map(str::to_string))
4009 .or_else(|| self.actor_id()),
4010 origin_actor_id: turn_actor
4011 .as_ref()
4012 .and_then(|context| context.origin_actor_id.clone()),
4013 sender_agent_id: turn_actor
4014 .as_ref()
4015 .and_then(|context| context.sender_agent_id.clone()),
4016 };
4017 let tool_context = ToolExecutionContext {
4018 requested_name: request.requested_name.clone(),
4019 canonical_id: canonical_id.clone(),
4020 display_name: resolved.identity.display_name.clone(),
4021 provider_id: resolved.identity.provider_id.clone(),
4022 registry_version: self.tools.version(),
4023 policy_version: security_engine.policy_version(),
4024 runtime_control_version: self.runtime_control.version.load(Ordering::SeqCst),
4025 call_id: request.call_id.clone(),
4026 source: request.source.clone(),
4027 actor,
4028 cancellation: ToolCancellationToken::new(
4029 Arc::clone(&self.runtime_control.emergency_deny),
4030 Some("runtime control cancellation".to_string()),
4031 ),
4032 started_at,
4033 deadline,
4034 permission: ToolPolicyDecisionRecord::allow(),
4035 approval: approval_record.clone(),
4036 classification: classification.clone(),
4037 safety,
4038 limits: limits.clone(),
4039 policy_snapshot,
4040 custom_config: security_engine.custom_config(&canonical_id),
4041 };
4042 let _resource_guard = self
4043 .acquire_tool_resource_lock(&canonical_id, &executed_arguments, &classification)
4044 .await;
4045 let (mut result, timed_out, cancelled) = self
4046 .run_tool_with_retries(
4047 &canonical_id,
4048 resolved.tool.clone(),
4049 executed_arguments.clone(),
4050 tool_context,
4051 timeout_ms,
4052 tool_config.max_retries,
4053 )
4054 .await?;
4055
4056 if !result.success {
4057 match &tool_config.on_failure {
4058 ToolFailureAction::Skip => {
4059 result = ToolResult::ok(format!(
4060 "{{\"skipped\": true, \"reason\": \"Tool '{}' was skipped after failure\"}}",
4061 canonical_id
4062 ));
4063 }
4064 ToolFailureAction::Fallback { fallback_tool } => {
4065 let fallback_request = ToolExecutionRequest::new(
4066 request.call_id.clone(),
4067 fallback_tool.clone(),
4068 executed_arguments.clone(),
4069 ToolCallSource::Fallback {
4070 original_tool: canonical_id.clone(),
4071 },
4072 );
4073 return Box::pin(self.execute_tool_record(fallback_request)).await;
4074 }
4075 ToolFailureAction::ReportError => {}
4076 }
4077 }
4078
4079 let output_cap = limits.max_output_chars;
4080 let (output, output_truncated) =
4081 Self::truncate_tool_output(result.output.clone(), output_cap);
4082 if let Some(result_metadata) = result.metadata {
4083 metadata.extend(result_metadata);
4084 }
4085 let mut record = self.record_from_parts(
4086 &request,
4087 canonical_id,
4088 executed_arguments,
4089 started_at,
4090 start,
4091 true,
4092 result.success,
4093 output,
4094 metadata,
4095 ToolPolicyDecisionRecord::allow(),
4096 approval_record,
4097 timed_out,
4098 output_truncated,
4099 );
4100 record.cancelled = cancelled;
4101 if cancelled {
4102 record.cancellation_reason = Some("runtime control cancellation".to_string());
4103 }
4104 self.finish_tool_record(&record).await;
4105 Ok(record)
4106 }
4107
4108 #[instrument(skip(self, tool_call), fields(tool = %tool_call.name))]
4109 async fn execute_tool_smart(&self, tool_call: &ToolCall) -> Result<String> {
4110 let record = self
4111 .execute_tool_record(ToolExecutionRequest::new(
4112 tool_call.id.clone(),
4113 tool_call.name.clone(),
4114 tool_call.arguments.clone(),
4115 ToolCallSource::Model,
4116 ))
4117 .await?;
4118 if record.success {
4119 Ok(record.model_output_string())
4120 } else if matches!(record.policy.outcome, PermissionOutcome::RequiresApproval) {
4121 Err(AgentError::HITLRejected(record.model_output_string()))
4122 } else {
4123 Err(AgentError::Tool(record.model_output_string()))
4124 }
4125 }
4126
4127 async fn select_skill_candidate(&self, input: &str) -> Result<Option<SkillCandidate>> {
4133 let Some(ref router) = self.skill_router else {
4134 return Ok(None);
4135 };
4136 let available_skills = self.get_available_skills();
4137 if available_skills.is_empty() {
4138 return Ok(None);
4139 }
4140 let skill_ids: Vec<&str> = available_skills.iter().map(|s| s.id.as_str()).collect();
4141 let Some(skill_id) = self
4142 .observe_purpose(
4143 ObservationPurpose::SkillRouting,
4144 router.select_skill_filtered(input, &skill_ids),
4145 )
4146 .await?
4147 else {
4148 return Ok(None);
4149 };
4150 let skill = router
4151 .get_skill(&skill_id)
4152 .cloned()
4153 .ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
4154 info!(skill_id = %skill_id, "Skill selected");
4155 Ok(Some(SkillCandidate::new(skill_id, skill)))
4156 }
4157
4158 async fn commit_skill_candidate_route_result(
4163 &self,
4164 candidate: SkillCandidate,
4165 input: &str,
4166 ) -> Result<SkillRouteResult> {
4167 let skill_id = candidate.skill_id;
4168 let skill = candidate.skill;
4169 if let Some(ref skill_disambig) = skill.disambiguation {
4170 if skill_disambig.enabled.unwrap_or(false) {
4171 if let Some(ref disambiguator) = self.disambiguation_manager {
4172 let context = self.build_disambiguation_context().await?;
4173 let state_override = self
4174 .state_machine
4175 .as_ref()
4176 .and_then(|sm| sm.current_definition())
4177 .and_then(|def| def.disambiguation.clone());
4178
4179 match self
4180 .observe_purpose(
4181 ObservationPurpose::DisambiguationDetection,
4182 disambiguator.process_input_with_override(
4183 input,
4184 &context,
4185 state_override.as_ref(),
4186 Some(skill_disambig),
4187 ),
4188 )
4189 .await?
4190 {
4191 DisambiguationResult::Clear => {
4192 debug!(skill_id = %skill_id, "Skill disambiguation: clear");
4193 }
4194 DisambiguationResult::NeedsClarification {
4195 question,
4196 detection,
4197 } => {
4198 info!(
4199 skill_id = %skill_id,
4200 ambiguity_type = ?detection.ambiguity_type,
4201 confidence = detection.confidence,
4202 "Skill requires clarification before execution"
4203 );
4204 *self.pending_skill_id.write() = Some(skill_id.clone());
4205 return Ok(SkillRouteResult::NeedsClarification(
4206 AgentResponse::new(&question.question).with_metadata(
4207 "disambiguation",
4208 serde_json::json!({
4209 "status": "awaiting_clarification",
4210 "skill_id": skill_id,
4211 "options": question.options,
4212 "clarifying": question.clarifying,
4213 "detection": {
4214 "type": detection.ambiguity_type,
4215 "confidence": detection.confidence,
4216 "what_is_unclear": detection.what_is_unclear,
4217 }
4218 }),
4219 ),
4220 ));
4221 }
4222 DisambiguationResult::Clarified { enriched_input, .. } => {
4223 info!(skill_id = %skill_id, enriched = %enriched_input, "Skill disambiguation clarified");
4224 return Ok(SkillRouteResult::Response(
4225 self.execute_skill(&skill, &enriched_input).await?,
4226 ));
4227 }
4228 DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
4229 info!(skill_id = %skill_id, "Skill disambiguation best guess");
4230 return Ok(SkillRouteResult::Response(
4231 self.execute_skill(&skill, &enriched_input).await?,
4232 ));
4233 }
4234 DisambiguationResult::GiveUp { reason } => {
4235 warn!(skill_id = %skill_id, reason = %reason, "Skill disambiguation gave up");
4236 let apology = self
4237 .generate_localized_apology(
4238 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
4239 &reason,
4240 )
4241 .await
4242 .unwrap_or_else(|_| {
4243 format!("I'm sorry, I couldn't understand your request: {}", reason)
4244 });
4245 return Ok(SkillRouteResult::NeedsClarification(AgentResponse::new(
4246 &apology,
4247 )));
4248 }
4249 DisambiguationResult::Escalate { reason } => {
4250 info!(skill_id = %skill_id, reason = %reason, "Skill disambiguation escalating");
4251 let apology = self
4252 .generate_localized_apology(
4253 "Explain briefly that you're transferring the user to a human agent for help.",
4254 &reason,
4255 )
4256 .await
4257 .unwrap_or_else(|_| {
4258 format!("I need human assistance to help with your request: {}", reason)
4259 });
4260 return Ok(SkillRouteResult::NeedsClarification(AgentResponse::new(
4261 &apology,
4262 )));
4263 }
4264 DisambiguationResult::Abandoned { .. } => {
4265 debug!(skill_id = %skill_id, "Skill disambiguation abandoned");
4266 return Ok(SkillRouteResult::NoMatch);
4267 }
4268 }
4269 }
4270 }
4271 }
4272 Ok(SkillRouteResult::Response(
4273 self.execute_skill(&skill, input).await?,
4274 ))
4275 }
4276
4277 async fn try_skill_route(&self, input: &str) -> Result<SkillRouteResult> {
4279 if let Some(candidate) = self.select_skill_candidate(input).await? {
4280 self.commit_skill_candidate_route_result(candidate, input)
4281 .await
4282 } else {
4283 Ok(SkillRouteResult::NoMatch)
4284 }
4285 }
4286
4287 async fn execute_skill(&self, skill: &SkillDefinition, input: &str) -> Result<String> {
4289 if let Some(ref executor) = self.skill_executor {
4290 let skill_reasoning = self.get_skill_reasoning_config(skill);
4291 let skill_reflection = self.get_skill_reflection_config(skill);
4292
4293 debug!(
4294 skill_id = %skill.id,
4295 reasoning_mode = ?skill_reasoning.mode,
4296 reflection_enabled = ?skill_reflection.enabled,
4297 "Skill reasoning/reflection config"
4298 );
4299
4300 let response = self
4301 .observe_purpose(
4302 ObservationPurpose::SkillPrompt,
4303 executor.execute_with_invoker(skill, input, serde_json::json!({}), self),
4304 )
4305 .await?;
4306
4307 if skill_reflection.requires_evaluation() && skill_reflection.is_enabled() {
4308 let should_reflect = self
4309 .should_reflect_with_config(input, &response, &skill_reflection)
4310 .await?;
4311 if should_reflect {
4312 let evaluated = self
4313 .evaluate_and_retry_with_config(input, response, &skill_reflection)
4314 .await?;
4315 return Ok(evaluated);
4316 }
4317 }
4318
4319 return Ok(response);
4320 }
4321 Err(AgentError::Skill(
4322 "No skill executor configured".to_string(),
4323 ))
4324 }
4325
4326 async fn execute_skill_by_id(&self, skill_id: &str, input: &str) -> Result<String> {
4329 let skill = self
4330 .skill_router
4331 .as_ref()
4332 .and_then(|r| r.get_skill(skill_id).cloned())
4333 .ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
4334 self.execute_skill(&skill, input).await
4335 }
4336
4337 async fn should_reflect_with_config(
4338 &self,
4339 input: &str,
4340 response: &str,
4341 config: &ReflectionConfig,
4342 ) -> Result<bool> {
4343 if !config.requires_evaluation() {
4344 return Ok(false);
4345 }
4346
4347 if config.is_enabled() {
4348 return Ok(true);
4349 }
4350
4351 let evaluator_llm = config
4352 .evaluator_llm
4353 .as_ref()
4354 .and_then(|alias| self.llm_registry.get(alias).ok())
4355 .or_else(|| self.llm_registry.router().ok())
4356 .or_else(|| self.llm_registry.default().ok());
4357
4358 let Some(llm) = evaluator_llm else {
4359 return Ok(false);
4360 };
4361
4362 let response_preview: String = response.chars().take(500).collect();
4363 let prompt = format!(
4364 r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
4365
4366User query: "{}"
4367Response: "{}"
4368
4369Answer YES or NO only."#,
4370 input, response_preview
4371 );
4372
4373 let messages = vec![ChatMessage::user(&prompt)];
4374 let result = self
4375 .observe_purpose(
4376 ObservationPurpose::ReflectionDecision,
4377 llm.complete(&messages, None),
4378 )
4379 .await;
4380
4381 match result {
4382 Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
4383 Err(_) => Ok(false),
4384 }
4385 }
4386
4387 async fn evaluate_and_retry_with_config(
4388 &self,
4389 input: &str,
4390 mut response: String,
4391 config: &ReflectionConfig,
4392 ) -> Result<String> {
4393 let llm = self.get_state_llm()?;
4394 let mut attempts = 0u32;
4395 let max_retries = config.max_retries;
4396
4397 loop {
4398 let evaluation = self
4399 .evaluate_response_with_config(input, &response, config)
4400 .await?;
4401
4402 if evaluation.passed || attempts >= max_retries {
4403 info!(
4404 passed = evaluation.passed,
4405 confidence = evaluation.confidence,
4406 attempts = attempts + 1,
4407 "Skill reflection evaluation complete"
4408 );
4409 return Ok(response);
4410 }
4411
4412 debug!(
4413 attempt = attempts + 1,
4414 failed_criteria = evaluation.failed_criteria().count(),
4415 "Skill response did not meet criteria, retrying"
4416 );
4417
4418 let feedback: Vec<String> = evaluation
4419 .failed_criteria()
4420 .map(|c| format!("- {}", c.criterion))
4421 .collect();
4422
4423 let retry_prompt = format!(
4424 "Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response to: {}",
4425 feedback.join("\n"),
4426 input
4427 );
4428
4429 let messages = vec![ChatMessage::user(&retry_prompt)];
4430 let retry_response = self
4431 .observe_purpose(
4432 ObservationPurpose::ReflectionEvaluation,
4433 llm.complete(&messages, None),
4434 )
4435 .await
4436 .map_err(|e| AgentError::LLM(e.to_string()))?;
4437
4438 response = retry_response.content.trim().to_string();
4439 attempts += 1;
4440 }
4441 }
4442
4443 async fn evaluate_response_with_config(
4444 &self,
4445 input: &str,
4446 response: &str,
4447 config: &ReflectionConfig,
4448 ) -> Result<EvaluationResult> {
4449 let evaluator_llm = config
4450 .evaluator_llm
4451 .as_ref()
4452 .and_then(|alias| self.llm_registry.get(alias).ok())
4453 .or_else(|| self.llm_registry.router().ok())
4454 .or_else(|| self.llm_registry.default().ok())
4455 .ok_or_else(|| AgentError::Config("No LLM available for evaluation".into()))?;
4456
4457 let criteria = &config.criteria;
4458 let criteria_list = criteria
4459 .iter()
4460 .enumerate()
4461 .map(|(i, c)| format!("{}. {}", i + 1, c))
4462 .collect::<Vec<_>>()
4463 .join("\n");
4464
4465 let prompt = format!(
4466 r#"Evaluate this response against the criteria.
4467
4468User query: "{}"
4469
4470Response to evaluate: "{}"
4471
4472Criteria:
4473{}
4474
4475For each criterion, respond with:
4476- criterion number
4477- PASS or FAIL
4478- brief reason
4479
4480Then provide overall confidence (0.0 to 1.0) and whether it passes overall.
4481
4482Format:
44831. PASS/FAIL - reason
44842. PASS/FAIL - reason
4485...
4486CONFIDENCE: 0.X
4487OVERALL: PASS/FAIL"#,
4488 input, response, criteria_list
4489 );
4490
4491 let messages = vec![ChatMessage::user(&prompt)];
4492 let eval_response = self
4493 .observe_purpose(
4494 ObservationPurpose::ReflectionEvaluation,
4495 evaluator_llm.complete(&messages, None),
4496 )
4497 .await
4498 .map_err(|e| AgentError::LLM(format!("Evaluation failed: {}", e)))?;
4499
4500 let content = eval_response.content.to_uppercase();
4501 let llm_pass = content.contains("OVERALL: PASS");
4502
4503 let confidence = content
4504 .lines()
4505 .find(|l| l.contains("CONFIDENCE:"))
4506 .and_then(|l| {
4507 l.split(':')
4508 .nth(1)
4509 .and_then(|v| v.trim().parse::<f32>().ok())
4510 })
4511 .unwrap_or(if llm_pass { 0.8 } else { 0.4 });
4512
4513 let overall_pass = llm_pass && confidence >= config.pass_threshold;
4516
4517 let mut criteria_results = Vec::new();
4518 for (i, criterion) in criteria.iter().enumerate() {
4519 let line_marker = format!("{}.", i + 1);
4520 let passed = eval_response
4521 .content
4522 .lines()
4523 .find(|l| l.contains(&line_marker))
4524 .map(|l| l.to_uppercase().contains("PASS"))
4525 .unwrap_or(overall_pass);
4526
4527 if passed {
4528 criteria_results.push(CriterionResult::pass(criterion));
4529 } else {
4530 criteria_results.push(CriterionResult::fail(criterion, "Did not meet criterion"));
4531 }
4532 }
4533
4534 Ok(EvaluationResult::new(overall_pass, confidence).with_criteria(criteria_results))
4535 }
4536
4537 async fn process_input(&self, input: &str) -> Result<ProcessData> {
4539 if let Some(processor) = self.get_state_process_processor() {
4540 let purpose = observation_purpose_for_process(processor.input_purpose_hint());
4541 return self
4542 .observe_purpose(purpose, processor.process_input(input))
4543 .await;
4544 }
4545 if let Some(ref processor) = self.process_processor {
4546 let purpose = observation_purpose_for_process(processor.input_purpose_hint());
4547 self.observe_purpose(purpose, processor.process_input(input))
4548 .await
4549 } else {
4550 Ok(ProcessData::new(input))
4551 }
4552 }
4553
4554 async fn process_output(
4556 &self,
4557 output: &str,
4558 input_context: &std::collections::HashMap<String, serde_json::Value>,
4559 ) -> Result<ProcessData> {
4560 if let Some(processor) = self.get_state_process_processor() {
4561 let purpose = observation_purpose_for_process(processor.output_purpose_hint());
4562 return self
4563 .observe_purpose(purpose, processor.process_output(output, input_context))
4564 .await;
4565 }
4566 if let Some(ref processor) = self.process_processor {
4567 let purpose = observation_purpose_for_process(processor.output_purpose_hint());
4568 self.observe_purpose(purpose, processor.process_output(output, input_context))
4569 .await
4570 } else {
4571 Ok(ProcessData::new(output))
4572 }
4573 }
4574
4575 fn get_state_process_processor(&self) -> Option<ProcessProcessor> {
4577 let sm = self.state_machine.as_ref()?;
4578 let def = sm.current_definition()?;
4579 let config = def.process.as_ref()?;
4580 let mut processor = ProcessProcessor::new(config.clone());
4581 if let Some(ref registry) = Some(self.llm_registry.clone()) {
4582 processor = processor.with_llm_registry(registry.clone());
4583 }
4584 processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
4585 Some(processor)
4586 }
4587
4588 async fn check_turn_timeout(&self) -> Result<()> {
4589 if let Some(ref sm) = self.state_machine {
4590 if let Some(timeout_state) = sm.check_timeout() {
4591 let from_state = sm.current();
4592 self.execute_state_exit_actions(&from_state).await;
4593 sm.transition_to(&timeout_state, "max_turns exceeded")?;
4594 self.execute_state_enter_actions(&timeout_state).await;
4595 info!(to = %timeout_state, "Timeout transition");
4596 }
4597 }
4598 Ok(())
4599 }
4600
4601 fn increment_turn(&self) {
4602 if let Some(ref sm) = self.state_machine {
4603 sm.increment_turn();
4604 }
4605 }
4606
4607 fn transitions_available_for_commit(&self) -> Option<(Vec<Transition>, String)> {
4608 let sm = self.state_machine.as_ref()?;
4609 let current = sm.current();
4610 let transitions: Vec<_> = sm
4611 .auto_transitions()
4612 .into_iter()
4613 .filter(|t| match t.cooldown_turns {
4614 Some(cd) if cd > 0 => {
4615 let resolved = sm.config().resolve_full_path(¤t, &t.to);
4616 !sm.is_on_cooldown(&resolved, cd)
4617 }
4618 _ => true,
4619 })
4620 .collect();
4621 Some((transitions, current))
4622 }
4623
4624 fn transition_reason(transition: &Transition) -> String {
4625 if transition.when.is_empty() {
4626 "guard condition met".to_string()
4627 } else {
4628 transition.when.clone()
4629 }
4630 }
4631
4632 fn build_transition_context(
4634 &self,
4635 user_message: &str,
4636 response: &str,
4637 current_state: &str,
4638 staged: Option<&HashMap<String, Value>>,
4639 ) -> TransitionContext {
4640 let context_map = staged
4641 .map(|writes| self.build_context_with_staged(writes))
4642 .unwrap_or_else(|| self.build_context_with_overlays());
4643 TransitionContext::new(user_message, response, current_state).with_context(context_map)
4644 }
4645
4646 async fn select_transition_candidate(
4648 &self,
4649 user_message: &str,
4650 response: &str,
4651 ) -> Result<Option<TransitionCandidate>> {
4652 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
4653 return Ok(None);
4654 };
4655 let transitions: Vec<Transition> = transitions
4656 .into_iter()
4657 .filter(|transition| matches!(transition.timing, TransitionTiming::PostResponse))
4658 .collect();
4659 if transitions.is_empty() {
4660 return Ok(None);
4661 }
4662 let Some(evaluator) = self.transition_evaluator.as_ref() else {
4663 return Ok(None);
4664 };
4665 let context = self.build_transition_context(user_message, response, ¤t_state, None);
4666 let selected = self
4667 .observe_purpose(
4668 ObservationPurpose::StateTransitionEvaluation,
4669 evaluator.select_transition(&transitions, &context),
4670 )
4671 .await?;
4672 Ok(selected.map(|index| {
4673 let transition = transitions[index].clone();
4674 TransitionCandidate::new(
4675 current_state,
4676 transition.clone(),
4677 Self::transition_reason(&transition),
4678 )
4679 }))
4680 }
4681
4682 fn select_deterministic_transition_candidate(
4684 &self,
4685 user_message: &str,
4686 current_state: &str,
4687 transitions: &[Transition],
4688 staged: &HashMap<String, Value>,
4689 ) -> Option<TransitionCandidate> {
4690 let context = self.build_transition_context(user_message, "", current_state, Some(staged));
4691
4692 for transition in transitions {
4693 if let Some(guard) = transition.guard.as_ref() {
4694 if evaluate_guard(guard, &context) {
4695 return Some(TransitionCandidate::new(
4696 current_state,
4697 transition.clone(),
4698 Self::transition_reason(transition),
4699 ));
4700 }
4701 }
4702 }
4703
4704 let resolved_intent = context
4705 .context
4706 .get("resolved_intent")
4707 .and_then(Value::as_str)
4708 .filter(|value| !value.is_empty());
4709 if let Some(resolved_intent) = resolved_intent {
4710 for transition in transitions {
4711 if transition.intent.as_deref() == Some(resolved_intent) {
4712 return Some(TransitionCandidate::new(
4713 current_state,
4714 transition.clone(),
4715 Self::transition_reason(transition),
4716 ));
4717 }
4718 }
4719 }
4720
4721 None
4722 }
4723
4724 async fn commit_transition_candidate(&self, candidate: &TransitionCandidate) -> Result<bool> {
4726 self.commit_transition_target(&candidate.from_state, candidate.target(), &candidate.reason)
4727 .await
4728 }
4729
4730 async fn approve_transition_target(&self, from_state: &str, target: &str) -> Result<bool> {
4732 let approved = self.check_state_hitl(Some(from_state), target).await?;
4733 if !approved {
4734 info!(to = %target, "State transition rejected by HITL");
4735 }
4736 Ok(approved)
4737 }
4738
4739 async fn apply_transition_target(
4741 &self,
4742 from_state: &str,
4743 target: &str,
4744 reason: &str,
4745 staged: Option<&HashMap<String, Value>>,
4746 ) -> Result<bool> {
4747 let Some(ref sm) = self.state_machine else {
4748 return Ok(false);
4749 };
4750
4751 self.execute_state_exit_actions(from_state).await;
4752 sm.transition_to(target, reason)?;
4753 sm.reset_no_transition();
4754 if let Some(staged) = staged {
4755 self.commit_staged_context_writes(staged).await;
4756 }
4757 let entered = sm.current();
4758 self.execute_state_enter_actions(&entered).await;
4759 self.hooks
4760 .on_state_transition(Some(from_state), &entered, reason)
4761 .await;
4762 info!(from = %from_state, to = %entered, "State transition");
4763 Ok(true)
4764 }
4765
4766 async fn commit_transition_target(
4768 &self,
4769 from_state: &str,
4770 target: &str,
4771 reason: &str,
4772 ) -> Result<bool> {
4773 if !self.approve_transition_target(from_state, target).await? {
4774 return Ok(false);
4775 }
4776 self.apply_transition_target(from_state, target, reason, None)
4777 .await
4778 }
4779
4780 async fn commit_pre_response_transition_candidate(
4782 &self,
4783 candidate: &TransitionCandidate,
4784 staged: &HashMap<String, Value>,
4785 processed_input: &str,
4786 ) -> Result<bool> {
4787 if !self
4788 .approve_transition_target(&candidate.from_state, candidate.target())
4789 .await?
4790 {
4791 return Ok(false);
4792 }
4793 self.commit_root_user_message(processed_input).await?;
4794 self.apply_transition_target(
4795 &candidate.from_state,
4796 candidate.target(),
4797 &candidate.reason,
4798 Some(staged),
4799 )
4800 .await
4801 }
4802
4803 async fn handle_transition_miss(&self, current_state: &str) -> Result<bool> {
4805 let Some(ref sm) = self.state_machine else {
4806 return Ok(false);
4807 };
4808 sm.increment_no_transition();
4809 let Some(fallback) = sm.check_fallback() else {
4810 return Ok(false);
4811 };
4812 self.commit_transition_target(current_state, &fallback, "fallback after no transitions")
4813 .await
4814 }
4815
4816 async fn evaluate_transitions(&self, user_message: &str, response: &str) -> Result<bool> {
4818 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
4819 return Ok(false);
4820 };
4821 if transitions.is_empty() {
4822 return Ok(false);
4823 }
4824 if let Some(candidate) = self
4825 .select_transition_candidate(user_message, response)
4826 .await?
4827 {
4828 return self.commit_transition_candidate(&candidate).await;
4829 }
4830 self.handle_transition_miss(¤t_state).await
4831 }
4832
4833 async fn try_pre_response_transition(
4835 &self,
4836 processed_input: &str,
4837 ) -> Result<Option<AgentResponse>> {
4838 let optimization = &self.runtime_config.optimization;
4839 if !optimization.enabled || !optimization.pre_response_deterministic_transitions {
4840 return Ok(None);
4841 }
4842 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
4843 return Ok(None);
4844 };
4845 let eligible: Vec<Transition> = transitions
4846 .into_iter()
4847 .filter(|transition| !transition.requires_response)
4848 .filter(|transition| matches!(transition.timing, TransitionTiming::PreResponse))
4849 .collect();
4850 if eligible.is_empty() {
4851 return Ok(None);
4852 }
4853
4854 let empty_staged = HashMap::new();
4855 let mut extracted_staged: Option<HashMap<String, Value>> = None;
4856 let mut selected: Option<(TransitionCandidate, HashMap<String, Value>)> = None;
4857
4858 for transition in &eligible {
4859 let use_extractors = optimization.pre_response_extractors || transition.run_extractors;
4860 let staged_for_eval = if use_extractors {
4861 if extracted_staged.is_none() {
4862 extracted_staged =
4863 Some(self.run_context_extractors_staged(processed_input).await);
4864 }
4865 extracted_staged.as_ref().unwrap_or(&empty_staged)
4866 } else {
4867 &empty_staged
4868 };
4869
4870 if let Some(candidate) = self.select_deterministic_transition_candidate(
4871 processed_input,
4872 ¤t_state,
4873 std::slice::from_ref(transition),
4874 staged_for_eval,
4875 ) {
4876 let staged_for_commit = if use_extractors {
4877 staged_for_eval.clone()
4878 } else {
4879 HashMap::new()
4880 };
4881 selected = Some((candidate, staged_for_commit));
4882 break;
4883 }
4884 }
4885
4886 let Some((candidate, staged)) = selected else {
4887 return Ok(None);
4888 };
4889
4890 if !self
4891 .commit_pre_response_transition_candidate(&candidate, &staged, processed_input)
4892 .await?
4893 {
4894 return Ok(None);
4895 }
4896 self.redispatch_current_state(processed_input)
4897 .await
4898 .map(Some)
4899 }
4900
4901 async fn try_speculative_branches(
4906 &self,
4907 processed_input: &str,
4908 input_context: &HashMap<String, Value>,
4909 ) -> Result<Option<AgentResponse>> {
4910 let optimization = &self.runtime_config.optimization;
4911 if !optimization.enabled {
4912 return Ok(None);
4913 }
4914
4915 let effective_reasoning_mode = self.get_effective_reasoning_config().mode.clone();
4916 if !matches!(
4917 effective_reasoning_mode,
4918 ReasoningMode::None | ReasoningMode::Auto
4919 ) {
4920 return Ok(None);
4921 }
4922
4923 let mut transition_enabled =
4924 optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
4925 let mut skill_enabled = optimization.speculative_skill_routing
4926 && self.skill_router.is_some()
4927 && self.pending_skill_id.read().is_none();
4928 let mut reasoning_enabled = optimization.speculative_reasoning_auto
4929 && matches!(effective_reasoning_mode, ReasoningMode::Auto);
4930
4931 if matches!(effective_reasoning_mode, ReasoningMode::Auto) {
4932 if !reasoning_enabled || optimization.max_speculative_llm_calls_per_turn < 2 {
4933 return Ok(None);
4934 }
4935 }
4936
4937 if !transition_enabled && !skill_enabled && !reasoning_enabled {
4938 return Ok(None);
4939 }
4940
4941 let mut optional_slots = optimization.max_parallel_runtime_tasks.saturating_sub(1);
4942 let mut speculative_call_slots = optimization
4943 .max_speculative_llm_calls_per_turn
4944 .saturating_sub(1);
4945 if reasoning_enabled {
4946 if optional_slots == 0 || speculative_call_slots == 0 {
4947 return Ok(None);
4948 }
4949 optional_slots -= 1;
4950 speculative_call_slots -= 1;
4951 }
4952 if transition_enabled {
4953 if optional_slots == 0 {
4954 transition_enabled = false;
4955 } else {
4956 optional_slots -= 1;
4957 }
4958 }
4959 if skill_enabled && (optional_slots == 0 || speculative_call_slots == 0) {
4960 skill_enabled = false;
4961 }
4962
4963 if !transition_enabled && !skill_enabled && !reasoning_enabled {
4964 return Ok(None);
4965 }
4966
4967 let main_kind = if transition_enabled {
4968 RuntimeOptimizationKind::ParallelStateTransition
4969 } else if skill_enabled {
4970 RuntimeOptimizationKind::SpeculativeSkillRouting
4971 } else {
4972 RuntimeOptimizationKind::SpeculativeReasoningAuto
4973 };
4974 if !self.reserve_active_speculative_llm_call(main_kind) {
4975 return Ok(None);
4976 }
4977
4978 let mut branch_set = ScheduledBranchSet::new(optimization.max_parallel_runtime_tasks)?;
4979 let main_branch = RuntimeBranch::new(
4980 RuntimeTaskPurpose::MainResponse,
4981 main_kind,
4982 RuntimeTaskPriority::Normal,
4983 RuntimeCommitBehavior::FinalResponse,
4984 );
4985 let transition_branch = RuntimeBranch::new(
4986 RuntimeTaskPurpose::StateTransition,
4987 RuntimeOptimizationKind::ParallelStateTransition,
4988 RuntimeTaskPriority::Critical,
4989 RuntimeCommitBehavior::TransitionDecision,
4990 );
4991 let skill_branch = RuntimeBranch::new(
4992 RuntimeTaskPurpose::SkillRouting,
4993 RuntimeOptimizationKind::SpeculativeSkillRouting,
4994 RuntimeTaskPriority::High,
4995 RuntimeCommitBehavior::SkillSelection,
4996 );
4997 let reasoning_branch = RuntimeBranch::new(
4998 RuntimeTaskPurpose::ReasoningJudge,
4999 RuntimeOptimizationKind::SpeculativeReasoningAuto,
5000 RuntimeTaskPriority::Normal,
5001 RuntimeCommitBehavior::ReasoningDecision,
5002 );
5003 let main_id = main_branch.branch_id();
5004 let transition_id = transition_branch.branch_id();
5005 let skill_id = skill_branch.branch_id();
5006 let reasoning_id = reasoning_branch.branch_id();
5007
5008 let main_id_for_future = main_id.clone();
5009 if !branch_set.schedule(
5010 main_branch,
5011 Box::pin(async move {
5012 match crate::optimization::observability::with_branch_observation(
5013 &main_id_for_future,
5014 main_kind,
5015 RuntimeCommitBehavior::FinalResponse,
5016 self.generate_main_response_draft(processed_input, &ReasoningMode::None),
5017 )
5018 .await
5019 {
5020 Ok(draft) => RuntimeBranchResult::MainDraft(draft),
5021 Err(error) => RuntimeBranchResult::Failed(error),
5022 }
5023 }),
5024 ) {
5025 return Ok(None);
5026 }
5027
5028 if transition_enabled {
5029 let transition_id_for_future = transition_id.clone();
5030 if !branch_set.schedule(
5031 transition_branch,
5032 Box::pin(async move {
5033 match crate::optimization::observability::with_branch_observation(
5034 &transition_id_for_future,
5035 RuntimeOptimizationKind::ParallelStateTransition,
5036 RuntimeCommitBehavior::TransitionDecision,
5037 self.select_parallel_transition_candidate(processed_input),
5038 )
5039 .await
5040 {
5041 Ok(ParallelTransitionSelection::Candidate(candidate)) => {
5042 RuntimeBranchResult::Transition(Some(candidate))
5043 }
5044 Ok(ParallelTransitionSelection::NoMatch) => {
5045 RuntimeBranchResult::Transition(None)
5046 }
5047 Ok(ParallelTransitionSelection::ReservationExhausted) => {
5048 RuntimeBranchResult::Cancelled
5049 }
5050 Err(error) => RuntimeBranchResult::Failed(error),
5051 }
5052 }),
5053 ) {
5054 transition_enabled = false;
5055 }
5056 }
5057
5058 if skill_enabled {
5059 let skill_id_for_future = skill_id.clone();
5060 if !branch_set.schedule(
5061 skill_branch,
5062 Box::pin(async move {
5063 if !self.reserve_active_speculative_llm_call(
5064 RuntimeOptimizationKind::SpeculativeSkillRouting,
5065 ) {
5066 return RuntimeBranchResult::Cancelled;
5067 }
5068 match crate::optimization::observability::with_branch_observation(
5069 &skill_id_for_future,
5070 RuntimeOptimizationKind::SpeculativeSkillRouting,
5071 RuntimeCommitBehavior::SkillSelection,
5072 self.select_skill_candidate(processed_input),
5073 )
5074 .await
5075 {
5076 Ok(candidate) => RuntimeBranchResult::Skill(candidate),
5077 Err(error) => RuntimeBranchResult::Failed(error),
5078 }
5079 }),
5080 ) {
5081 skill_enabled = false;
5082 }
5083 }
5084
5085 if reasoning_enabled {
5086 let reasoning_id_for_future = reasoning_id.clone();
5087 if !branch_set.schedule(
5088 reasoning_branch,
5089 Box::pin(async move {
5090 if !self.reserve_active_speculative_llm_call(
5091 RuntimeOptimizationKind::SpeculativeReasoningAuto,
5092 ) {
5093 return RuntimeBranchResult::Cancelled;
5094 }
5095 match crate::optimization::observability::with_branch_observation(
5096 &reasoning_id_for_future,
5097 RuntimeOptimizationKind::SpeculativeReasoningAuto,
5098 RuntimeCommitBehavior::ReasoningDecision,
5099 self.determine_reasoning_mode_strict(processed_input),
5100 )
5101 .await
5102 {
5103 Ok(mode) => RuntimeBranchResult::Reasoning(mode),
5104 Err(error) => RuntimeBranchResult::Failed(error),
5105 }
5106 }),
5107 ) {
5108 reasoning_enabled = false;
5109 }
5110 }
5111
5112 if matches!(effective_reasoning_mode, ReasoningMode::Auto) && !reasoning_enabled {
5113 self.finalize_pending_branches(branch_set.cancel_pending());
5114 return Ok(None);
5115 }
5116
5117 if !transition_enabled && !skill_enabled && !reasoning_enabled {
5118 self.finalize_pending_branches(branch_set.cancel_pending());
5119 return Ok(None);
5120 }
5121
5122 let mut main_pending = true;
5123 let mut skill_pending = skill_enabled;
5124 let mut reasoning_pending = reasoning_enabled;
5125 let mut transition_finalized = !transition_enabled;
5126 let mut skill_finalized = !skill_enabled;
5127 let mut reasoning_finalized = !reasoning_enabled;
5128 let mut main_result: Option<Result<MainResponseDraft>> = None;
5129 let mut transition_candidate: Option<TransitionCandidate> = None;
5130 let mut skill_candidate: Option<SkillCandidate> = None;
5131 let mut reasoning_decision: Option<ReasoningMode> = None;
5132 let mut transition_fallback_required = false;
5133 let mut skill_fallback_required = false;
5134 let mut reasoning_fallback_required = false;
5135
5136 loop {
5137 if let Some(candidate) = transition_candidate.take() {
5138 if self
5139 .commit_pre_response_transition_candidate(
5140 &candidate,
5141 &HashMap::new(),
5142 processed_input,
5143 )
5144 .await?
5145 {
5146 self.finalize_optional_branch(
5147 &transition_id,
5148 RuntimeOptimizationKind::ParallelStateTransition,
5149 RuntimeCommitBehavior::TransitionDecision,
5150 "committed",
5151 true,
5152 );
5153 if !main_pending {
5154 self.finalize_branch_loss(
5155 &main_id,
5156 main_kind,
5157 RuntimeCommitBehavior::FinalResponse,
5158 false,
5159 main_result.as_ref().map(|result| result.is_err()),
5160 );
5161 }
5162 if skill_enabled && !skill_pending {
5163 self.finalize_branch_loss(
5164 &skill_id,
5165 RuntimeOptimizationKind::SpeculativeSkillRouting,
5166 RuntimeCommitBehavior::SkillSelection,
5167 false,
5168 Some(false),
5169 );
5170 }
5171 if reasoning_enabled && !reasoning_pending {
5172 self.finalize_branch_loss(
5173 &reasoning_id,
5174 RuntimeOptimizationKind::SpeculativeReasoningAuto,
5175 RuntimeCommitBehavior::ReasoningDecision,
5176 false,
5177 Some(false),
5178 );
5179 }
5180 self.finalize_pending_branches(branch_set.cancel_pending());
5181 return self
5182 .redispatch_current_state(processed_input)
5183 .await
5184 .map(Some);
5185 }
5186 self.finalize_optional_branch(
5187 &transition_id,
5188 RuntimeOptimizationKind::ParallelStateTransition,
5189 RuntimeCommitBehavior::TransitionDecision,
5190 "discarded",
5191 false,
5192 );
5193 transition_finalized = true;
5194 }
5195
5196 if transition_finalized && skill_candidate.is_some() {
5197 let candidate = skill_candidate.take().unwrap();
5198 self.finalize_optional_branch(
5199 &skill_id,
5200 RuntimeOptimizationKind::SpeculativeSkillRouting,
5201 RuntimeCommitBehavior::SkillSelection,
5202 "committed",
5203 true,
5204 );
5205 if !main_pending {
5206 self.finalize_branch_loss(
5207 &main_id,
5208 main_kind,
5209 RuntimeCommitBehavior::FinalResponse,
5210 false,
5211 main_result.as_ref().map(|result| result.is_err()),
5212 );
5213 }
5214 if reasoning_enabled && !reasoning_pending {
5215 self.finalize_branch_loss(
5216 &reasoning_id,
5217 RuntimeOptimizationKind::SpeculativeReasoningAuto,
5218 RuntimeCommitBehavior::ReasoningDecision,
5219 false,
5220 Some(false),
5221 );
5222 }
5223 self.finalize_pending_branches(branch_set.cancel_pending());
5224 self.commit_root_user_message(processed_input).await?;
5225 return match self
5226 .commit_skill_candidate_route_result(candidate, processed_input)
5227 .await?
5228 {
5229 SkillRouteResult::Response(skill_response) => self
5230 .handle_skill_response(processed_input, skill_response, input_context)
5231 .await
5232 .map(Some),
5233 SkillRouteResult::NeedsClarification(response) => {
5234 if response
5235 .metadata
5236 .as_ref()
5237 .and_then(|m| m.get("disambiguation"))
5238 .and_then(|d| d.get("status"))
5239 .and_then(|s| s.as_str())
5240 == Some("awaiting_clarification")
5241 {
5242 self.memory
5243 .add_message(ChatMessage::assistant(&response.content))
5244 .await?;
5245 }
5246 self.finish_turn_if_root(&response).await?;
5247 Ok(Some(response))
5248 }
5249 SkillRouteResult::NoMatch => Ok(None),
5250 };
5251 }
5252
5253 if transition_finalized && skill_finalized {
5254 if let Some(reasoning_mode) = reasoning_decision.take() {
5255 if !matches!(reasoning_mode, ReasoningMode::None) {
5256 self.finalize_optional_branch(
5257 &reasoning_id,
5258 RuntimeOptimizationKind::SpeculativeReasoningAuto,
5259 RuntimeCommitBehavior::ReasoningDecision,
5260 "committed",
5261 true,
5262 );
5263 if !main_pending {
5264 self.finalize_branch_loss(
5265 &main_id,
5266 main_kind,
5267 RuntimeCommitBehavior::FinalResponse,
5268 false,
5269 main_result.as_ref().map(|result| result.is_err()),
5270 );
5271 }
5272 self.finalize_pending_branches(branch_set.cancel_pending());
5273 self.commit_root_user_message(processed_input).await?;
5274 return if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
5275 self.handle_plan_and_execute(processed_input, input_context, true)
5276 .await
5277 .map(Some)
5278 } else {
5279 self.run_committed_response_loop_with_reasoning(
5280 processed_input,
5281 input_context,
5282 reasoning_mode,
5283 true,
5284 )
5285 .await
5286 .map(Some)
5287 };
5288 }
5289 self.finalize_optional_branch(
5290 &reasoning_id,
5291 RuntimeOptimizationKind::SpeculativeReasoningAuto,
5292 RuntimeCommitBehavior::ReasoningDecision,
5293 "committed",
5294 true,
5295 );
5296 reasoning_finalized = true;
5297 }
5298 }
5299
5300 if transition_finalized && skill_finalized && reasoning_finalized {
5301 if transition_fallback_required
5302 || skill_fallback_required
5303 || reasoning_fallback_required
5304 {
5305 if !main_pending {
5306 self.finalize_branch_loss(
5307 &main_id,
5308 main_kind,
5309 RuntimeCommitBehavior::FinalResponse,
5310 false,
5311 main_result.as_ref().map(|result| result.is_err()),
5312 );
5313 }
5314 self.finalize_pending_branches(branch_set.cancel_pending());
5315 return Ok(None);
5316 }
5317
5318 if let Some(result) = main_result.take() {
5319 let draft = match result {
5320 Ok(draft) => draft,
5321 Err(error) => {
5322 self.finalize_optional_branch(
5323 &main_id,
5324 main_kind,
5325 RuntimeCommitBehavior::FinalResponse,
5326 "failed",
5327 false,
5328 );
5329 self.finalize_pending_branches(branch_set.cancel_pending());
5330 return Err(error);
5331 }
5332 };
5333 self.finalize_optional_branch(
5334 &main_id,
5335 main_kind,
5336 RuntimeCommitBehavior::FinalResponse,
5337 "committed",
5338 true,
5339 );
5340 self.finalize_pending_branches(branch_set.cancel_pending());
5341 return self
5342 .commit_main_response_draft(
5343 processed_input,
5344 input_context,
5345 draft,
5346 ReasoningMode::None,
5347 reasoning_enabled,
5348 )
5349 .await
5350 .map(Some);
5351 }
5352 }
5353
5354 if branch_set.is_empty() {
5355 return Ok(None);
5356 }
5357
5358 let Some(outcome) = branch_set.next_completed().await else {
5359 return Ok(None);
5360 };
5361 let branch_id = outcome.branch.branch_id();
5362 match outcome.result {
5363 RuntimeBranchResult::MainDraft(draft) => {
5364 main_pending = false;
5365 main_result = Some(Ok(draft));
5366 }
5367 RuntimeBranchResult::Transition(candidate) => {
5368 if let Some(candidate) = candidate {
5369 transition_candidate = Some(candidate);
5370 } else {
5371 self.finalize_optional_branch(
5372 &transition_id,
5373 RuntimeOptimizationKind::ParallelStateTransition,
5374 RuntimeCommitBehavior::TransitionDecision,
5375 "discarded",
5376 false,
5377 );
5378 transition_finalized = true;
5379 }
5380 }
5381 RuntimeBranchResult::Skill(candidate) => {
5382 skill_pending = false;
5383 if let Some(candidate) = candidate {
5384 skill_candidate = Some(candidate);
5385 } else {
5386 self.finalize_optional_branch(
5387 &skill_id,
5388 RuntimeOptimizationKind::SpeculativeSkillRouting,
5389 RuntimeCommitBehavior::SkillSelection,
5390 "discarded",
5391 false,
5392 );
5393 skill_finalized = true;
5394 }
5395 }
5396 RuntimeBranchResult::Reasoning(mode) => {
5397 reasoning_pending = false;
5398 reasoning_decision = Some(mode);
5399 }
5400 RuntimeBranchResult::Failed(error) => {
5401 if branch_id == main_id {
5402 main_pending = false;
5403 main_result = Some(Err(error));
5404 } else if branch_id == transition_id {
5405 self.finalize_optional_branch(
5406 &transition_id,
5407 RuntimeOptimizationKind::ParallelStateTransition,
5408 RuntimeCommitBehavior::TransitionDecision,
5409 "failed",
5410 false,
5411 );
5412 transition_finalized = true;
5413 } else if branch_id == skill_id {
5414 skill_pending = false;
5415 self.finalize_optional_branch(
5416 &skill_id,
5417 RuntimeOptimizationKind::SpeculativeSkillRouting,
5418 RuntimeCommitBehavior::SkillSelection,
5419 "failed",
5420 false,
5421 );
5422 skill_finalized = true;
5423 } else if branch_id == reasoning_id {
5424 reasoning_pending = false;
5425 self.finalize_optional_branch(
5426 &reasoning_id,
5427 RuntimeOptimizationKind::SpeculativeReasoningAuto,
5428 RuntimeCommitBehavior::ReasoningDecision,
5429 "failed",
5430 false,
5431 );
5432 reasoning_finalized = true;
5433 }
5434 }
5435 RuntimeBranchResult::Cancelled => {
5436 self.finalize_optional_branch(
5437 &branch_id,
5438 outcome.branch.optimization,
5439 outcome.branch.commit_behavior,
5440 "cancelled",
5441 false,
5442 );
5443 if branch_id == main_id {
5444 main_pending = false;
5445 main_result =
5446 Some(Err(AgentError::Other("main branch cancelled".to_string())));
5447 } else if branch_id == transition_id {
5448 transition_finalized = true;
5449 transition_fallback_required = true;
5450 } else if branch_id == skill_id {
5451 skill_pending = false;
5452 skill_finalized = true;
5453 skill_fallback_required = true;
5454 } else if branch_id == reasoning_id {
5455 reasoning_pending = false;
5456 reasoning_finalized = true;
5457 reasoning_fallback_required = true;
5458 }
5459 }
5460 }
5461 }
5462 }
5463
5464 fn finalize_pending_branches(&self, branches: Vec<RuntimeBranch>) {
5465 for branch in branches {
5466 self.finalize_optional_branch(
5467 &branch.branch_id(),
5468 branch.optimization,
5469 branch.commit_behavior,
5470 "cancelled",
5471 false,
5472 );
5473 }
5474 }
5475
5476 fn finalize_branch_loss(
5481 &self,
5482 branch_id: &str,
5483 optimization: RuntimeOptimizationKind,
5484 commit_behavior: RuntimeCommitBehavior,
5485 pending: bool,
5486 completed_failed: Option<bool>,
5487 ) {
5488 let status = if pending {
5489 "cancelled"
5490 } else if completed_failed.unwrap_or(false) {
5491 "failed"
5492 } else {
5493 "discarded"
5494 };
5495 self.finalize_optional_branch(branch_id, optimization, commit_behavior, status, false);
5496 }
5497
5498 fn finalize_optional_branch(
5503 &self,
5504 branch_id: &str,
5505 optimization: RuntimeOptimizationKind,
5506 commit_behavior: RuntimeCommitBehavior,
5507 status: &str,
5508 winner: bool,
5509 ) {
5510 crate::optimization::observability::finalize_branch(
5511 self.observability_manager.as_ref(),
5512 branch_id,
5513 status,
5514 winner,
5515 optimization,
5516 commit_behavior,
5517 );
5518 }
5519
5520 fn has_parallel_transition_candidates(&self) -> bool {
5525 self.transitions_available_for_commit()
5526 .map(|(transitions, _)| {
5527 transitions
5528 .iter()
5529 .any(|transition| matches!(transition.timing, TransitionTiming::Parallel))
5530 })
5531 .unwrap_or(false)
5532 }
5533
5534 async fn select_parallel_transition_candidate(
5539 &self,
5540 processed_input: &str,
5541 ) -> Result<ParallelTransitionSelection> {
5542 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
5543 return Ok(ParallelTransitionSelection::NoMatch);
5544 };
5545 let parallel: Vec<Transition> = transitions
5546 .into_iter()
5547 .filter(|transition| matches!(transition.timing, TransitionTiming::Parallel))
5548 .filter(|transition| !transition.requires_response)
5549 .collect();
5550 if parallel.is_empty() {
5551 return Ok(ParallelTransitionSelection::NoMatch);
5552 }
5553 let empty_staged = HashMap::new();
5554 if let Some(candidate) = self.select_deterministic_transition_candidate(
5555 processed_input,
5556 ¤t_state,
5557 ¶llel,
5558 &empty_staged,
5559 ) {
5560 return Ok(ParallelTransitionSelection::Candidate(candidate));
5561 }
5562 let when_transitions: Vec<(usize, &Transition)> = parallel
5563 .iter()
5564 .enumerate()
5565 .filter(|(_, transition)| !transition.when.trim().is_empty())
5566 .collect();
5567 if when_transitions.is_empty() {
5568 return Ok(ParallelTransitionSelection::NoMatch);
5569 }
5570 let llm = self
5571 .llm_registry
5572 .router()
5573 .or_else(|_| self.llm_registry.default())
5574 .map_err(|e| AgentError::Config(e.to_string()))?;
5575 let conditions = when_transitions
5576 .iter()
5577 .enumerate()
5578 .map(|(display_idx, (_, transition))| {
5579 format!("{}. {}", display_idx + 1, transition.when)
5580 })
5581 .collect::<Vec<_>>()
5582 .join("\n");
5583 if !self
5584 .reserve_active_speculative_llm_call(RuntimeOptimizationKind::ParallelStateTransition)
5585 {
5586 return Ok(ParallelTransitionSelection::ReservationExhausted);
5587 }
5588 let context_preview = self.branch_context_preview();
5589 let prompt = format!(
5590 "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-{}).",
5591 current_state,
5592 processed_input,
5593 context_preview,
5594 conditions,
5595 when_transitions.len()
5596 );
5597 let response = self
5598 .observe_purpose(
5599 ObservationPurpose::StateTransitionEvaluation,
5600 llm.complete(&[ChatMessage::user(prompt)], None),
5601 )
5602 .await
5603 .map_err(|e| AgentError::LLM(e.to_string()))?;
5604 let choice = response.content.trim().parse::<usize>().unwrap_or(0);
5605 if choice == 0 || choice > when_transitions.len() {
5606 return Ok(ParallelTransitionSelection::NoMatch);
5607 }
5608 let transition = when_transitions[choice - 1].1.clone();
5609 Ok(ParallelTransitionSelection::Candidate(
5610 TransitionCandidate::new(
5611 current_state,
5612 transition.clone(),
5613 Self::transition_reason(&transition),
5614 ),
5615 ))
5616 }
5617
5618 async fn redispatch_current_state(&self, processed_input: &str) -> Result<AgentResponse> {
5620 const MAX_REDISPATCH_DEPTH: u32 = 3;
5621 let current_depth = *self.redispatch_depth.read();
5622 if current_depth >= MAX_REDISPATCH_DEPTH {
5623 warn!(depth = current_depth, "Re-dispatch depth limit reached");
5624 let response = AgentResponse::new("");
5625 self.finish_turn_if_root(&response).await?;
5626 return Ok(response);
5627 }
5628 *self.redispatch_depth.write() += 1;
5629 if let Some(context) = self.active_turn_context.write().as_mut() {
5630 context.enter_redispatch();
5631 }
5632 let result = Box::pin(self.run_loop_internal(processed_input)).await;
5633 *self.redispatch_depth.write() -= 1;
5634 if let Some(context) = self.active_turn_context.write().as_mut() {
5635 context.exit_redispatch();
5636 }
5637 let response = result?;
5638 self.finish_turn_if_root(&response).await?;
5639 Ok(response)
5640 }
5641
5642 async fn finish_turn_if_root(&self, response: &AgentResponse) -> Result<()> {
5644 if *self.redispatch_depth.read() == 0 {
5645 self.post_turn_session_lifecycle().await?;
5646 if let Some(context) = self.active_turn_context.write().as_mut() {
5647 context.mark_post_turn_lifecycle_completed();
5648 }
5649 self.hooks.on_response(response).await;
5650 self.end_root_turn();
5651 }
5652 Ok(())
5653 }
5654
5655 async fn execute_state_exit_actions(&self, state_path: &str) {
5657 if let Some(ref sm) = self.state_machine {
5658 if let Some(def) = sm.get_definition(state_path) {
5659 if !def.on_exit.is_empty() {
5660 debug!(state = %state_path, count = def.on_exit.len(), "Executing on_exit actions");
5661 self.execute_state_actions(&def.on_exit).await;
5662 }
5663 }
5664 }
5665 }
5666
5667 async fn execute_state_enter_actions(&self, state_path: &str) {
5669 if let Some(ref sm) = self.state_machine {
5670 if let Some(def) = sm.get_definition(state_path) {
5671 let is_reentry = sm.history().iter().any(|e| e.to == state_path);
5673
5674 if is_reentry && !def.on_reenter.is_empty() {
5675 debug!(state = %state_path, count = def.on_reenter.len(), "Executing on_reenter actions");
5676 self.execute_state_actions(&def.on_reenter).await;
5677 } else if !def.on_enter.is_empty() {
5678 debug!(state = %state_path, count = def.on_enter.len(), "Executing on_enter actions");
5679 self.execute_state_actions(&def.on_enter).await;
5680 }
5681 }
5682 }
5683 }
5684
5685 async fn execute_state_actions(&self, actions: &[StateAction]) {
5687 for (action_index, action) in actions.iter().enumerate() {
5688 match action {
5689 StateAction::Tool { tool, args } => {
5690 let raw_args = args.clone().unwrap_or(Value::Object(Default::default()));
5691 let args_value = self.render_action_args(&raw_args);
5692 let state = self.state_machine.as_ref().map(|sm| sm.current());
5693 let request = ToolExecutionRequest::new(
5694 uuid::Uuid::new_v4().to_string(),
5695 tool.clone(),
5696 args_value,
5697 ToolCallSource::StateAction {
5698 state,
5699 action_index,
5700 },
5701 );
5702 match self.execute_tool_record(request).await {
5703 Ok(record) if record.success => {
5704 debug!(tool = %record.canonical_id, "State action: tool executed");
5705 let _ = self.context_manager.set(
5706 "last_tool_result",
5707 serde_json::Value::String(record.model_output_string()),
5708 );
5709 let _ = self.context_manager.set(
5710 "last_tool_record",
5711 serde_json::to_value(record).unwrap_or(Value::Null),
5712 );
5713 }
5714 Ok(record) => {
5715 warn!(tool = %record.canonical_id, error = %record.output, "State action: tool failed");
5716 }
5717 Err(e) => {
5718 warn!(tool = %tool, error = %e, "State action: tool failed")
5719 }
5720 }
5721 }
5722 StateAction::Skill { skill } => {
5723 if let Some(ref executor) = self.skill_executor {
5724 if let Some(def) = self.skills.iter().find(|s| s.id == *skill) {
5725 match executor
5726 .execute_with_invoker(def, "", serde_json::json!({}), self)
5727 .await
5728 {
5729 Ok(_) => debug!(skill = %skill, "State action: skill executed"),
5730 Err(e) => {
5731 warn!(skill = %skill, error = %e, "State action: skill failed")
5732 }
5733 }
5734 } else {
5735 warn!(skill = %skill, "State action: skill not found");
5736 }
5737 }
5738 }
5739 StateAction::SetContext { set_context } => {
5740 for (key, value) in set_context {
5741 if let Err(e) = self.context_manager.set(key, value.clone()) {
5742 warn!(key = %key, error = %e, "State action: set_context failed");
5743 } else {
5744 debug!(key = %key, "State action: context set");
5745 }
5746 }
5747 }
5748 StateAction::Prompt {
5749 prompt,
5750 llm,
5751 store_as,
5752 } => {
5753 let llm_result = if let Some(alias) = llm {
5754 self.llm_registry.get(alias)
5755 } else {
5756 self.llm_registry.default()
5757 };
5758 match llm_result {
5759 Ok(llm_provider) => {
5760 let context = self.build_context_with_overlays();
5762 let rendered_prompt = self
5763 .template_renderer
5764 .render(prompt, &context)
5765 .unwrap_or_else(|_| prompt.clone());
5766 let recent =
5767 self.memory.get_messages(Some(5)).await.unwrap_or_default();
5768 let mut messages: Vec<ChatMessage> = recent;
5769 messages.push(ChatMessage::user(&rendered_prompt));
5770 match self
5771 .observe_purpose(
5772 ObservationPurpose::StateAction,
5773 llm_provider.complete(&messages, None),
5774 )
5775 .await
5776 {
5777 Ok(response) => {
5778 if let Some(key) = store_as {
5779 let _ = self
5780 .context_manager
5781 .set(key, Value::String(response.content));
5782 debug!(key = %key, "State action: prompt result stored");
5783 }
5784 }
5785 Err(e) => {
5786 warn!(error = %e, "State action: prompt LLM call failed");
5787 }
5788 }
5789 }
5790 Err(e) => {
5791 warn!(error = %e, "State action: LLM not found for prompt");
5792 }
5793 }
5794 }
5795 }
5796 }
5797 }
5798
5799 async fn run_context_extractors_staged(&self, user_message: &str) -> HashMap<String, Value> {
5800 let extractors = match &self.state_machine {
5801 Some(sm) => match sm.current_definition() {
5802 Some(def) if !def.extract.is_empty() => def.extract.clone(),
5803 _ => return HashMap::new(),
5804 },
5805 None => return HashMap::new(),
5806 };
5807
5808 let mut staged = HashMap::new();
5809 for extractor in &extractors {
5810 let prompt = if let Some(ref custom) = extractor.llm_extract {
5811 format!(
5812 "User message:\n\"{}\"\n\nInstruction:\n{}",
5813 user_message, custom
5814 )
5815 } else if let Some(ref desc) = extractor.description {
5816 format!(
5817 "From the following message, extract: {}\n\n\
5818 Message: \"{}\"\n\n\
5819 If the information is present, return ONLY the extracted value.\n\
5820 If NOT present, return exactly: __NONE__",
5821 desc, user_message
5822 )
5823 } else {
5824 continue;
5825 };
5826
5827 let llm = match self
5828 .llm_registry
5829 .get(&extractor.llm)
5830 .or_else(|_| self.llm_registry.get("router"))
5831 .or_else(|_| self.llm_registry.get("default"))
5832 {
5833 Ok(llm) => llm,
5834 Err(e) => {
5835 warn!(key = %extractor.key, error = %e, "Extractor LLM not found");
5836 continue;
5837 }
5838 };
5839
5840 let messages = vec![ChatMessage::user(&prompt)];
5841 match self
5842 .observe_purpose(
5843 ObservationPurpose::ContextExtraction,
5844 llm.complete(&messages, None),
5845 )
5846 .await
5847 {
5848 Ok(response) => {
5849 let value = response.content.trim().to_string();
5850 if value != "__NONE__" && !value.is_empty() {
5851 staged.insert(
5852 extractor.key.clone(),
5853 serde_json::Value::String(value.clone()),
5854 );
5855 debug!(key = %extractor.key, value = %value, "Context extracted");
5856 } else if extractor.required {
5857 warn!(key = %extractor.key, "Required extraction returned no value");
5858 }
5859 }
5860 Err(e) => {
5861 warn!(key = %extractor.key, error = %e, "Context extraction LLM call failed");
5862 }
5863 }
5864 }
5865 staged
5866 }
5867
5868 async fn commit_staged_context_writes(&self, staged: &HashMap<String, Value>) {
5869 for (key, value) in staged {
5870 if let Err(error) = self.context_manager.update(key, value.clone()) {
5871 warn!(key = %key, error = %error, "staged context write failed");
5872 }
5873 }
5874 }
5875
5876 async fn run_context_extractors(&self, user_message: &str) {
5878 let staged = self.run_context_extractors_staged(user_message).await;
5879 self.commit_staged_context_writes(&staged).await;
5880 }
5881
5882 async fn check_memory_compression(&self) -> Result<()> {
5883 if self.memory.needs_compression() {
5884 let result = self.memory.compress(None).await?;
5885 if let CompressResult::Compressed {
5886 messages_summarized,
5887 new_summary_length,
5888 tokens_saved,
5889 } = result
5890 {
5891 let event = MemoryCompressEvent::new(
5892 messages_summarized,
5893 tokens_saved,
5894 new_summary_length as u32,
5895 );
5896 self.hooks.on_memory_compress(&event).await;
5897 debug!(
5898 messages = messages_summarized,
5899 tokens_saved = tokens_saved,
5900 "Memory compressed"
5901 );
5902 }
5903 }
5904
5905 self.handle_memory_overflow().await?;
5907 self.check_memory_budget().await;
5908
5909 Ok(())
5910 }
5911
5912 async fn check_memory_budget(&self) {
5913 let Some(ref budget) = self.memory_token_budget else {
5914 return;
5915 };
5916
5917 let context = match self.memory.get_context().await {
5918 Ok(ctx) => ctx,
5919 Err(_) => return,
5920 };
5921
5922 let used_tokens = context.estimated_tokens();
5924 if budget.is_over_warn_threshold(used_tokens) {
5925 let event = MemoryBudgetEvent::new("memory", used_tokens, budget.total);
5926 self.hooks.on_memory_budget_warning(&event).await;
5927 debug!(
5928 used = used_tokens,
5929 total = budget.total,
5930 percent = event.usage_percent,
5931 "Memory budget warning"
5932 );
5933 }
5934
5935 if let Some(ref summary) = context.summary {
5937 let summary_tokens = ai_agents_memory::estimate_tokens(summary);
5938 let summary_budget = budget.allocation.summary;
5939 if summary_budget > 0 {
5940 let warn_threshold =
5941 (summary_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
5942 if summary_tokens >= warn_threshold {
5943 let event = MemoryBudgetEvent::new("summary", summary_tokens, summary_budget);
5944 self.hooks.on_memory_budget_warning(&event).await;
5945 }
5946 }
5947 }
5948
5949 let recent_tokens: u32 = context
5951 .messages
5952 .iter()
5953 .map(ai_agents_memory::estimate_message_tokens)
5954 .sum();
5955 let recent_budget = budget.allocation.recent_messages;
5956 if recent_budget > 0 {
5957 let warn_threshold =
5958 (recent_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
5959 if recent_tokens >= warn_threshold {
5960 let event = MemoryBudgetEvent::new("recent_messages", recent_tokens, recent_budget);
5961 self.hooks.on_memory_budget_warning(&event).await;
5962 }
5963 }
5964
5965 let relationship_budget = budget.allocation.relationships;
5966 if relationship_budget > 0 {
5967 let relationship_tokens = self
5968 .relationship_memory_text()
5969 .map(|text| ai_agents_memory::estimate_tokens(&text))
5970 .unwrap_or(0);
5971 let warn_threshold =
5972 (relationship_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
5973 if relationship_tokens >= warn_threshold {
5974 let event = MemoryBudgetEvent::new(
5975 "relationships",
5976 relationship_tokens,
5977 relationship_budget,
5978 );
5979 self.hooks.on_memory_budget_warning(&event).await;
5980 }
5981 }
5982 }
5983
5984 async fn handle_memory_overflow(&self) -> Result<()> {
5985 let Some(ref budget) = self.memory_token_budget else {
5986 return Ok(());
5987 };
5988
5989 let context = self.memory.get_context().await?;
5990 let used_tokens = context.estimated_tokens();
5991
5992 if used_tokens <= budget.total {
5993 return Ok(());
5994 }
5995
5996 match budget.overflow_strategy {
5997 OverflowStrategy::TruncateOldest => {
5998 let tokens_to_free = used_tokens - budget.total;
5999 let messages_to_evict = self.calculate_eviction_count(tokens_to_free);
6000 if messages_to_evict > 0 {
6001 self.evict_messages(messages_to_evict, EvictionReason::TokenBudgetExceeded)
6002 .await?;
6003 }
6004 }
6005 OverflowStrategy::SummarizeMore => {
6006 self.memory.compress(None).await?;
6007 }
6008 OverflowStrategy::Error => {
6009 return Err(AgentError::MemoryBudgetExceeded {
6010 used: used_tokens,
6011 budget: budget.total,
6012 });
6013 }
6014 }
6015 Ok(())
6016 }
6017
6018 fn calculate_eviction_count(&self, tokens_to_free: u32) -> usize {
6019 ((tokens_to_free as f64 / 50.0).ceil() as usize).max(1)
6021 }
6022
6023 async fn evict_messages(&self, count: usize, reason: EvictionReason) -> Result<()> {
6024 let evicted = self.memory.evict_oldest(count).await?;
6025 if !evicted.is_empty() {
6026 let event = MemoryEvictEvent {
6027 reason,
6028 messages_evicted: evicted.len(),
6029 importance_scores: vec![],
6030 };
6031 self.hooks.on_memory_evict(&event).await;
6032 debug!(count = evicted.len(), "Messages evicted from memory");
6033 }
6034 Ok(())
6035 }
6036
6037 #[instrument(skip(self, input), fields(agent = %self.info.name))]
6038 async fn determine_reasoning_mode(&self, input: &str) -> Result<ReasoningMode> {
6039 match self.determine_reasoning_mode_strict(input).await {
6040 Ok(mode) => Ok(mode),
6041 Err(_) => Ok(ReasoningMode::None),
6042 }
6043 }
6044
6045 async fn determine_reasoning_mode_strict(&self, input: &str) -> Result<ReasoningMode> {
6046 let effective_config = self.get_effective_reasoning_config();
6047
6048 if !matches!(effective_config.mode, ReasoningMode::Auto) {
6049 return Ok(effective_config.mode.clone());
6050 }
6051
6052 let judge_llm = effective_config
6053 .judge_llm
6054 .as_ref()
6055 .and_then(|alias| self.llm_registry.get(alias).ok())
6056 .or_else(|| self.llm_registry.router().ok())
6057 .or_else(|| self.llm_registry.default().ok());
6058
6059 let Some(llm) = judge_llm else {
6060 return Ok(ReasoningMode::None);
6061 };
6062
6063 let prompt = format!(
6064 r#"Analyze this user request and determine the appropriate reasoning mode.
6065
6066User request: "{}"
6067
6068Choose ONE of these modes:
6069- none: Simple queries, greetings, direct answers (fastest)
6070- cot: Complex analysis, multi-step reasoning, math problems
6071- react: Tasks requiring multiple tool calls with observation
6072- plan_and_execute: Complex multi-step tasks requiring coordination
6073
6074Respond with ONLY the mode name (none, cot, react, or plan_and_execute)."#,
6075 input
6076 );
6077
6078 let messages = vec![ChatMessage::user(&prompt)];
6079 let response = self
6080 .observe_purpose(
6081 ObservationPurpose::ReflectionDecision,
6082 llm.complete(&messages, None),
6083 )
6084 .await
6085 .map_err(|e| AgentError::LLM(e.to_string()))?;
6086
6087 let mode_str = response.content.trim().to_lowercase();
6088 Ok(match mode_str.as_str() {
6089 "cot" => ReasoningMode::CoT,
6090 "react" => ReasoningMode::React,
6091 "plan_and_execute" => ReasoningMode::PlanAndExecute,
6092 _ => ReasoningMode::None,
6093 })
6094 }
6095
6096 async fn should_reflect(&self, input: &str, response: &str) -> Result<bool> {
6097 let effective_config = self.get_effective_reflection_config();
6098
6099 if !effective_config.requires_evaluation() {
6100 return Ok(false);
6101 }
6102
6103 if effective_config.is_enabled() {
6104 return Ok(true);
6105 }
6106
6107 let evaluator_llm = effective_config
6108 .evaluator_llm
6109 .as_ref()
6110 .and_then(|alias| self.llm_registry.get(alias).ok())
6111 .or_else(|| self.llm_registry.router().ok())
6112 .or_else(|| self.llm_registry.default().ok());
6113
6114 let Some(llm) = evaluator_llm else {
6115 return Ok(false);
6116 };
6117
6118 let response_preview: String = response.chars().take(500).collect();
6119 let prompt = format!(
6120 r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
6121
6122User query: "{}"
6123Response: "{}"
6124
6125Answer YES or NO only."#,
6126 input, response_preview
6127 );
6128
6129 let messages = vec![ChatMessage::user(&prompt)];
6130 let result = self
6131 .observe_purpose(
6132 ObservationPurpose::ReflectionDecision,
6133 llm.complete(&messages, None),
6134 )
6135 .await;
6136
6137 match result {
6138 Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
6139 Err(_) => Ok(false),
6140 }
6141 }
6142
6143 fn build_cot_system_prompt(&self, base_prompt: &str) -> String {
6144 format!(
6145 "{}\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>",
6146 base_prompt
6147 )
6148 }
6149
6150 fn build_react_system_prompt(&self, base_prompt: &str) -> String {
6151 format!(
6152 "{}\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>",
6153 base_prompt
6154 )
6155 }
6156
6157 async fn generate_plan(&self, input: &str) -> Result<Plan> {
6158 let effective = self.get_effective_reasoning_config();
6159 let planning_config = effective.get_planning();
6160
6161 let planner_llm = planning_config
6162 .and_then(|c| c.planner_llm.as_ref())
6163 .and_then(|alias| self.llm_registry.get(alias).ok())
6164 .or_else(|| self.llm_registry.router().ok())
6165 .or_else(|| self.llm_registry.default().ok())
6166 .ok_or_else(|| AgentError::Config("No LLM available for planning".into()))?;
6167
6168 let mut available_tool_ids: Vec<String> = self
6169 .get_available_tool_ids()
6170 .await
6171 .unwrap_or_else(|_| self.tools.list_ids());
6172 let mut available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
6173
6174 if let Some(config) = planning_config {
6176 if !config.available.tools.is_all() {
6177 available_tool_ids.retain(|t| config.available.tools.allows(t));
6178 }
6179 if !config.available.skills.is_all() {
6180 available_skills.retain(|s| config.available.skills.allows(s));
6181 }
6182 }
6183
6184 let tool_descriptions: Vec<String> = available_tool_ids
6187 .iter()
6188 .filter_map(|id| {
6189 self.tools.get(id).map(|tool| {
6190 let schema = tool.input_schema();
6191 let args_desc = schema
6192 .get("properties")
6193 .and_then(|p| serde_json::to_string(p).ok())
6194 .unwrap_or_else(|| "{}".to_string());
6195 format!(
6196 "- {} ({}): {}\n Arguments: {}",
6197 id,
6198 tool.name(),
6199 tool.description(),
6200 args_desc
6201 )
6202 })
6203 })
6204 .collect();
6205
6206 let tools_section = if tool_descriptions.is_empty() {
6207 "Available tools: none".to_string()
6208 } else {
6209 format!("Available tools:\n{}", tool_descriptions.join("\n"))
6210 };
6211
6212 let skills_section = if available_skills.is_empty() {
6213 "Available skills: none".to_string()
6214 } else {
6215 format!("Available skills: {}", available_skills.join(", "))
6216 };
6217
6218 let prompt = format!(
6219 r#"Create a step-by-step plan to accomplish this goal.
6220
6221Goal: "{}"
6222
6223{}
6224
6225{}
6226
6227Create a plan with clear steps. For each step, specify:
6228- description: What this step accomplishes
6229- action_type: "tool", "skill", "think", or "respond"
6230- action_target: The tool/skill id (if applicable)
6231- args: The arguments object matching the tool's schema (if action_type is "tool")
6232- dependencies: List of step IDs this depends on (empty if none)
6233
6234Respond in JSON format:
6235{{
6236 "steps": [
6237 {{"id": "step1", "description": "...", "action_type": "tool", "action_target": "tool_id", "args": {{"required_field": "value"}}, "dependencies": []}},
6238 {{"id": "step2", "description": "...", "action_type": "think", "action_target": "...", "dependencies": ["step1"]}}
6239 ]
6240}}"#,
6241 input, tools_section, skills_section,
6242 );
6243
6244 let messages = vec![ChatMessage::user(&prompt)];
6245 let response = self
6246 .observe_purpose(
6247 ObservationPurpose::PlanGeneration,
6248 planner_llm.complete(&messages, None),
6249 )
6250 .await
6251 .map_err(|e| AgentError::LLM(format!("Planning failed: {}", e)))?;
6252
6253 let mut plan = Plan::new(input);
6254
6255 if let Some(json_start) = response.content.find('{') {
6256 if let Some(json_end) = response.content.rfind('}') {
6257 let json_str = &response.content[json_start..=json_end];
6258 if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(json_str) {
6259 if let Some(steps) = parsed.get("steps").and_then(|s| s.as_array()) {
6260 for step_value in steps {
6261 let id = step_value
6262 .get("id")
6263 .and_then(|v| v.as_str())
6264 .unwrap_or("step");
6265 let desc = step_value
6266 .get("description")
6267 .and_then(|v| v.as_str())
6268 .unwrap_or("");
6269 let action_type = step_value
6270 .get("action_type")
6271 .and_then(|v| v.as_str())
6272 .unwrap_or("think");
6273 let action_target = step_value
6274 .get("action_target")
6275 .and_then(|v| v.as_str())
6276 .unwrap_or("");
6277 let args = step_value
6278 .get("args")
6279 .cloned()
6280 .unwrap_or(serde_json::json!({}));
6281 let deps: Vec<String> = step_value
6282 .get("dependencies")
6283 .and_then(|v| v.as_array())
6284 .map(|arr| {
6285 arr.iter()
6286 .filter_map(|v| v.as_str().map(String::from))
6287 .collect()
6288 })
6289 .unwrap_or_default();
6290
6291 let action = match action_type {
6292 "tool" => PlanAction::tool(action_target, args),
6293 "skill" => PlanAction::skill(action_target),
6294 "respond" => PlanAction::respond(action_target),
6295 _ => PlanAction::think(desc),
6296 };
6297
6298 let step = PlanStep::new(desc, action)
6299 .with_id(id)
6300 .with_dependencies(deps);
6301 plan.add_step(step);
6302 }
6303 }
6304 }
6305 }
6306 }
6307
6308 if plan.steps.is_empty() {
6309 plan.add_step(PlanStep::new(
6310 "Process the request",
6311 PlanAction::think(input),
6312 ));
6313 plan.add_step(PlanStep::new(
6314 "Provide response",
6315 PlanAction::respond("Answer based on analysis"),
6316 ));
6317 }
6318
6319 Ok(plan)
6320 }
6321
6322 async fn execute_plan(&self, plan: &mut Plan) -> Result<String> {
6323 let llm = self.get_state_llm()?;
6324 let mut results: HashMap<String, serde_json::Value> = HashMap::new();
6325 let effective = self.get_effective_reasoning_config();
6326 let max_steps = effective.get_planning().map(|c| c.max_steps).unwrap_or(10);
6327
6328 plan.status = PlanStatus::InProgress;
6329
6330 for step_idx in 0..plan.steps.len().min(max_steps as usize) {
6331 let step = &plan.steps[step_idx];
6332
6333 let deps_satisfied = step.dependencies.iter().all(|dep| {
6334 plan.steps
6335 .iter()
6336 .find(|s| &s.id == dep)
6337 .map(|s| s.status.is_completed())
6338 .unwrap_or(false)
6339 });
6340
6341 if !deps_satisfied {
6342 continue;
6343 }
6344
6345 plan.steps[step_idx].mark_running();
6346
6347 let result = match &plan.steps[step_idx].action {
6348 PlanAction::Tool { tool, args } => {
6349 let has_dep_results = plan.steps[step_idx]
6355 .dependencies
6356 .iter()
6357 .any(|dep| results.contains_key(dep));
6358
6359 let final_args = if has_dep_results {
6360 let dep_context: String = plan.steps[step_idx]
6361 .dependencies
6362 .iter()
6363 .filter_map(|dep| results.get(dep).map(|r| format!("{}: {}", dep, r)))
6364 .collect::<Vec<_>>()
6365 .join("\n");
6366
6367 let tool_schema = self
6368 .tools
6369 .get(tool)
6370 .map(|t| {
6371 let schema = t.input_schema();
6372 let props = schema
6373 .get("properties")
6374 .and_then(|p| serde_json::to_string(p).ok())
6375 .unwrap_or_else(|| "{}".to_string());
6376 format!(
6377 "{}: {}\nArguments schema: {}",
6378 t.id(),
6379 t.description(),
6380 props
6381 )
6382 })
6383 .unwrap_or_default();
6384
6385 let step_desc = &plan.steps[step_idx].description;
6386 let arg_prompt = format!(
6387 "Generate the JSON arguments for a tool call.\n\n\
6388 Tool: {}\n\n\
6389 Task: {}\n\n\
6390 Previous step results:\n{}\n\n\
6391 Planner's draft arguments: {}\n\n\
6392 Produce ONLY a valid JSON object with the correct argument values.\n\
6393 Use actual values from the previous step results, not template references.",
6394 tool_schema,
6395 step_desc,
6396 dep_context,
6397 serde_json::to_string(args).unwrap_or_default()
6398 );
6399 let messages = vec![ChatMessage::user(&arg_prompt)];
6400 match self
6401 .observe_purpose(
6402 ObservationPurpose::PlanStep,
6403 llm.complete(&messages, None),
6404 )
6405 .await
6406 {
6407 Ok(resp) => {
6408 let content = resp.content.trim();
6409 let json_start = content.find('{');
6411 let json_end = content.rfind('}');
6412 if let (Some(start), Some(end)) = (json_start, json_end) {
6413 serde_json::from_str(&content[start..=end])
6414 .unwrap_or_else(|_| args.clone())
6415 } else {
6416 args.clone()
6417 }
6418 }
6419 Err(_) => args.clone(),
6420 }
6421 } else {
6422 args.clone()
6423 };
6424
6425 let request = ToolExecutionRequest::new(
6426 uuid::Uuid::new_v4().to_string(),
6427 tool.clone(),
6428 final_args,
6429 ToolCallSource::Plan {
6430 step_index: step_idx,
6431 },
6432 );
6433 match self.execute_tool_record(request).await {
6434 Ok(record) if record.success => {
6435 serde_json::json!({ "output": record.model_output_string() })
6436 }
6437 Ok(record) => {
6438 plan.steps[step_idx].mark_failed(record.model_output_string());
6439 continue;
6440 }
6441 Err(e) => {
6442 plan.steps[step_idx].mark_failed(e.to_string());
6443 continue;
6444 }
6445 }
6446 }
6447 PlanAction::Skill { skill } => {
6448 if let Some(skill_def) = self.skills.iter().find(|s| &s.id == skill) {
6449 if let Some(ref executor) = self.skill_executor {
6450 match executor
6451 .execute_with_invoker(skill_def, "", serde_json::json!({}), self)
6452 .await
6453 {
6454 Ok(output) => serde_json::json!({ "output": output }),
6455 Err(e) => {
6456 plan.steps[step_idx].mark_failed(e.to_string());
6457 continue;
6458 }
6459 }
6460 } else {
6461 serde_json::json!({ "output": "Skill executor not available" })
6462 }
6463 } else {
6464 plan.steps[step_idx].mark_failed("Skill not found");
6465 continue;
6466 }
6467 }
6468 PlanAction::Think { prompt } => {
6469 let context: String = results
6470 .iter()
6471 .map(|(k, v)| format!("{}: {}", k, v))
6472 .collect::<Vec<_>>()
6473 .join("\n");
6474
6475 let think_prompt = format!("Context:\n{}\n\nTask: {}", context, prompt);
6476 let messages = vec![ChatMessage::user(&think_prompt)];
6477
6478 match self
6479 .observe_purpose(
6480 ObservationPurpose::PlanStep,
6481 llm.complete(&messages, None),
6482 )
6483 .await
6484 {
6485 Ok(resp) => serde_json::json!({ "output": resp.content }),
6486 Err(e) => {
6487 plan.steps[step_idx].mark_failed(e.to_string());
6488 continue;
6489 }
6490 }
6491 }
6492 PlanAction::Respond { template } => {
6493 let context: String = results
6494 .iter()
6495 .map(|(k, v)| format!("{}: {}", k, v))
6496 .collect::<Vec<_>>()
6497 .join("\n");
6498
6499 let respond_prompt = format!(
6500 "Based on this context:\n{}\n\nGenerate a response following this template/instruction: {}",
6501 context, template
6502 );
6503 let messages = vec![ChatMessage::user(&respond_prompt)];
6504
6505 match self
6506 .observe_purpose(
6507 ObservationPurpose::PlanStep,
6508 llm.complete(&messages, None),
6509 )
6510 .await
6511 {
6512 Ok(resp) => serde_json::json!({ "output": resp.content }),
6513 Err(e) => {
6514 plan.steps[step_idx].mark_failed(e.to_string());
6515 continue;
6516 }
6517 }
6518 }
6519 };
6520
6521 results.insert(plan.steps[step_idx].id.clone(), result.clone());
6522 plan.steps[step_idx].mark_completed(Some(result));
6523 }
6524
6525 let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
6527 if has_failures {
6528 let failed_ids: Vec<String> = plan
6529 .steps
6530 .iter()
6531 .filter(|s| s.status.is_failed())
6532 .map(|s| s.id.clone())
6533 .collect();
6534 plan.status = PlanStatus::Failed {
6535 error: format!("Steps failed: {}", failed_ids.join(", ")),
6536 };
6537 } else {
6538 plan.status = PlanStatus::Completed;
6539 }
6540
6541 let all_outputs: Vec<String> = plan
6543 .steps
6544 .iter()
6545 .filter(|s| s.status.is_completed())
6546 .filter_map(|s| {
6547 s.result
6548 .as_ref()
6549 .and_then(|r| r.get("output"))
6550 .and_then(|o| o.as_str())
6551 .map(|o| format!("{}: {}", s.description, o))
6552 })
6553 .collect();
6554
6555 if all_outputs.is_empty() {
6556 return Ok("Plan execution completed but produced no results.".to_string());
6557 }
6558
6559 if all_outputs.len() == 1 {
6560 return Ok(all_outputs.into_iter().next().unwrap());
6561 }
6562
6563 let context = all_outputs.join("\n\n");
6565 let prompt = format!(
6566 "You completed a multi-step plan for: \"{}\"\n\nStep results:\n{}\n\nProvide a coherent final response that synthesizes these results.",
6567 plan.goal, context
6568 );
6569 let messages = vec![ChatMessage::user(&prompt)];
6570 match self
6571 .observe_purpose(ObservationPurpose::PlanStep, llm.complete(&messages, None))
6572 .await
6573 {
6574 Ok(resp) => Ok(resp.content.trim().to_string()),
6575 Err(_) => Ok(context),
6576 }
6577 }
6578
6579 async fn evaluate_response(&self, input: &str, response: &str) -> Result<EvaluationResult> {
6580 let effective_config = self.get_effective_reflection_config();
6581 self.evaluate_response_with_config(input, response, &effective_config)
6582 .await
6583 }
6584
6585 fn extract_thinking(&self, content: &str) -> (Option<String>, String) {
6586 if let Some(start) = content.find("<thinking>") {
6587 if let Some(end) = content.find("</thinking>") {
6588 let thinking = content[start + 10..end].trim().to_string();
6589 let answer = content[end + 11..].trim().to_string();
6590 return (Some(thinking), answer);
6591 }
6592 }
6593 (None, content.to_string())
6594 }
6595
6596 fn format_response_with_thinking(&self, thinking: Option<&str>, answer: &str) -> String {
6597 match self.get_effective_reasoning_config().output {
6598 ReasoningOutput::Hidden => answer.to_string(),
6599 ReasoningOutput::Visible => {
6600 if let Some(t) = thinking {
6601 format!("Thinking:\n{}\n\nAnswer:\n{}", t, answer)
6602 } else {
6603 answer.to_string()
6604 }
6605 }
6606 ReasoningOutput::Tagged => {
6607 if let Some(t) = thinking {
6608 format!("<thinking>{}</thinking>\n{}", t, answer)
6609 } else {
6610 answer.to_string()
6611 }
6612 }
6613 }
6614 }
6615
6616 async fn run_loop(&self, input: &str) -> Result<AgentResponse> {
6617 self.begin_root_turn();
6618 let _root_cleanup = RootTurnCleanup::new(self);
6619 info!(input_len = input.len(), "Starting chat");
6620
6621 self.hooks.on_message_received(input).await;
6622
6623 if !self.context_initialized.swap(true, Ordering::SeqCst) {
6627 self.context_manager.initialize().await?;
6628 debug!("Context manager initialized (defaults, env, builtins)");
6629 }
6630
6631 self.check_turn_timeout().await?;
6632 self.context_manager.refresh_per_turn().await?;
6633
6634 self.clear_disambiguation_context();
6637
6638 if let Some(ref disambiguator) = self.disambiguation_manager {
6640 let disambiguation_context = self.build_disambiguation_context().await?;
6641
6642 let state_override = self
6644 .state_machine
6645 .as_ref()
6646 .and_then(|sm| sm.current_definition())
6647 .and_then(|def| def.disambiguation.clone());
6648
6649 match self
6650 .observe_purpose(
6651 ObservationPurpose::DisambiguationDetection,
6652 disambiguator.process_input_with_override(
6653 input,
6654 &disambiguation_context,
6655 state_override.as_ref(),
6656 None,
6657 ),
6658 )
6659 .await?
6660 {
6661 DisambiguationResult::Clear => {
6662 debug!("Input is clear, proceeding normally");
6663 }
6664 DisambiguationResult::NeedsClarification {
6665 question,
6666 detection,
6667 } => {
6668 info!(
6669 ambiguity_type = ?detection.ambiguity_type,
6670 confidence = detection.confidence,
6671 "Input requires clarification"
6672 );
6673
6674 self.commit_root_user_message(input).await?;
6675 self.memory
6676 .add_message(ChatMessage::assistant(&question.question))
6677 .await?;
6678
6679 let response = AgentResponse::new(&question.question).with_metadata(
6680 "disambiguation",
6681 serde_json::json!({
6682 "status": "awaiting_clarification",
6683 "options": question.options,
6684 "clarifying": question.clarifying,
6685 "detection": {
6686 "type": detection.ambiguity_type,
6687 "confidence": detection.confidence,
6688 "what_is_unclear": detection.what_is_unclear,
6689 }
6690 }),
6691 );
6692 self.finish_turn_if_root(&response).await?;
6693 return Ok(response);
6694 }
6695 DisambiguationResult::Clarified {
6696 enriched_input,
6697 resolved,
6698 ..
6699 } => {
6700 info!(
6701 resolved_count = resolved.len(),
6702 enriched = %enriched_input,
6703 "Input clarified, injecting resolved intent into context"
6704 );
6705
6706 for (key, value) in &resolved {
6709 let context_key = format!("disambiguation.{}", key);
6710 let _ = self.context_manager.set(&context_key, value.clone());
6711 }
6712
6713 if let Some(intent) = resolved.get("intent") {
6714 let _ = self.context_manager.set("resolved_intent", intent.clone());
6715 }
6716
6717 let _ = self
6718 .context_manager
6719 .set("disambiguation.resolved", serde_json::Value::Bool(true));
6720
6721 let skill_id = self.pending_skill_id.read().clone();
6725 if let Some(skill_id) = skill_id {
6726 info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input");
6727 return self
6728 .recheck_skill_disambiguation(&skill_id, &enriched_input)
6729 .await;
6730 }
6731
6732 return self.run_loop_internal(&enriched_input).await;
6733 }
6734 DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
6735 info!("Proceeding with best guess interpretation");
6736
6737 let skill_id = self.pending_skill_id.read().clone();
6739 if let Some(skill_id) = skill_id {
6740 info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input");
6741 return self
6742 .recheck_skill_disambiguation(&skill_id, &enriched_input)
6743 .await;
6744 }
6745
6746 return self.run_loop_internal(&enriched_input).await;
6747 }
6748 DisambiguationResult::GiveUp { reason } => {
6749 *self.pending_skill_id.write() = None;
6750 warn!(reason = %reason, "Disambiguation gave up");
6751 let apology = self
6752 .generate_localized_apology(
6753 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
6754 &reason,
6755 )
6756 .await
6757 .unwrap_or_else(|_| {
6758 format!("I'm sorry, I couldn't understand your request: {}", reason)
6759 });
6760 let response = AgentResponse::new(&apology);
6761 self.finish_turn_if_root(&response).await?;
6762 return Ok(response);
6763 }
6764 DisambiguationResult::Escalate { reason } => {
6765 *self.pending_skill_id.write() = None;
6766 info!(reason = %reason, "Escalating to human");
6767 if let Some(ref hitl) = self.hitl_engine {
6768 let trigger =
6769 ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
6770 let mut context_map = HashMap::new();
6771 context_map.insert("original_input".to_string(), serde_json::json!(input));
6772 context_map.insert("reason".to_string(), serde_json::json!(&reason));
6773 let check_result = HITLCheckResult::required(
6774 trigger,
6775 context_map,
6776 format!("User request needs human assistance: {}", reason),
6777 Some(hitl.config().default_timeout_seconds),
6778 );
6779 let result = self.request_hitl_approval(check_result).await?;
6780 if matches!(
6781 result,
6782 ApprovalResult::Approved | ApprovalResult::Modified { .. }
6783 ) {
6784 return self.run_loop_internal(input).await;
6785 }
6786 }
6787 let apology = self
6788 .generate_localized_apology(
6789 "Explain briefly that you're transferring the user to a human agent for help.",
6790 &reason,
6791 )
6792 .await
6793 .unwrap_or_else(|_| {
6794 format!("I need human assistance to help with your request: {}", reason)
6795 });
6796 let response = AgentResponse::new(&apology);
6797 self.finish_turn_if_root(&response).await?;
6798 return Ok(response);
6799 }
6800 DisambiguationResult::Abandoned { new_input } => {
6801 *self.pending_skill_id.write() = None;
6802
6803 info!(
6804 has_new_input = new_input.is_some(),
6805 "Clarification abandoned by user"
6806 );
6807
6808 self.commit_root_user_message(input).await?;
6809
6810 match new_input {
6811 Some(fresh_input) => {
6812 return self.run_loop_internal(&fresh_input).await;
6815 }
6816 None => {
6817 let ack = self
6819 .generate_localized_apology(
6820 "The user changed their mind about their previous request. \
6821 Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
6822 Do NOT apologize excessively. Be concise.",
6823 "User abandoned clarification",
6824 )
6825 .await
6826 .unwrap_or_else(|_| {
6827 "OK, no problem. What else can I help with?".to_string()
6828 });
6829
6830 self.memory
6831 .add_message(ChatMessage::assistant(&ack))
6832 .await?;
6833
6834 let response = AgentResponse::new(&ack);
6835 self.finish_turn_if_root(&response).await?;
6836 return Ok(response);
6837 }
6838 }
6839 }
6840 }
6841 }
6842
6843 self.run_loop_internal(input).await
6844 }
6845
6846 async fn generate_localized_apology(&self, instruction: &str, reason: &str) -> Result<String> {
6848 let llm = self.llm_registry.router().map_err(|e| {
6849 AgentError::LLM(format!(
6850 "Router LLM not available for localized response: {}",
6851 e
6852 ))
6853 })?;
6854
6855 let recent: Vec<String> = self
6856 .memory
6857 .get_messages(Some(3))
6858 .await?
6859 .iter()
6860 .map(|m| m.content.clone())
6861 .collect();
6862
6863 let context_hint = if recent.is_empty() {
6864 String::new()
6865 } else {
6866 format!(
6867 "\nRecent conversation (detect the user's language from this):\n{}\n",
6868 recent.join("\n")
6869 )
6870 };
6871
6872 let prompt = format!(
6873 "{}\nReason: {}\n{}Respond in the same language as the user. Output ONLY the message, nothing else.",
6874 instruction, reason, context_hint
6875 );
6876
6877 let messages = vec![ChatMessage::user(&prompt)];
6878 let response = self
6879 .observe_purpose(
6880 ObservationPurpose::DisambiguationClarification,
6881 llm.complete(&messages, None),
6882 )
6883 .await
6884 .map_err(|e| AgentError::LLM(format!("Localized response generation failed: {}", e)))?;
6885
6886 Ok(response.content.trim().to_string())
6887 }
6888
6889 fn render_action_args(&self, args: &Value) -> Value {
6893 let context = self.build_context_with_overlays();
6894 match args {
6895 Value::Object(map) => {
6896 let mut rendered = serde_json::Map::new();
6897 for (k, v) in map {
6898 match v {
6899 Value::String(s) if s.contains("{{") => {
6900 match self.template_renderer.render(s, &context) {
6901 Ok(rendered_str) => {
6902 rendered.insert(k.clone(), Value::String(rendered_str));
6903 }
6904 Err(_) => {
6905 rendered.insert(k.clone(), v.clone());
6906 }
6907 }
6908 }
6909 _ => {
6910 rendered.insert(k.clone(), v.clone());
6911 }
6912 }
6913 }
6914 Value::Object(rendered)
6915 }
6916 _ => args.clone(),
6917 }
6918 }
6919
6920 fn clear_disambiguation_context(&self) {
6922 let _ = self
6923 .context_manager
6924 .set("resolved_intent", serde_json::Value::Null);
6925
6926 let all = self.context_manager.get_all();
6927 for key in all.keys() {
6928 if key.starts_with("disambiguation.") {
6929 let _ = self.context_manager.set(key, serde_json::Value::Null);
6930 }
6931 }
6932 }
6933
6934 async fn recheck_skill_disambiguation(
6940 &self,
6941 skill_id: &str,
6942 enriched_input: &str,
6943 ) -> Result<AgentResponse> {
6944 let skill = self
6945 .skill_router
6946 .as_ref()
6947 .and_then(|r| r.get_skill(skill_id).cloned());
6948
6949 if let Some(ref skill) = skill {
6951 if let Some(ref skill_disambig) = skill.disambiguation {
6952 if skill_disambig.enabled.unwrap_or(false) {
6953 if let Some(ref disambiguator) = self.disambiguation_manager {
6954 let context = self.build_disambiguation_context().await?;
6955 let state_override = self
6956 .state_machine
6957 .as_ref()
6958 .and_then(|sm| sm.current_definition())
6959 .and_then(|def| def.disambiguation.clone());
6960
6961 match self
6962 .observe_purpose(
6963 ObservationPurpose::DisambiguationDetection,
6964 disambiguator.process_input_with_override(
6965 enriched_input,
6966 &context,
6967 state_override.as_ref(),
6968 Some(skill_disambig),
6969 ),
6970 )
6971 .await?
6972 {
6973 DisambiguationResult::Clear => {
6974 debug!(skill_id = %skill_id, "Skill re-check: all fields present");
6975 }
6976 DisambiguationResult::NeedsClarification {
6977 question,
6978 detection,
6979 } => {
6980 info!(
6981 skill_id = %skill_id,
6982 ambiguity_type = ?detection.ambiguity_type,
6983 what_is_unclear = ?detection.what_is_unclear,
6984 "Skill re-check: still missing fields, asking again"
6985 );
6986 self.memory
6990 .add_message(ChatMessage::user(enriched_input))
6991 .await?;
6992 self.memory
6993 .add_message(ChatMessage::assistant(&question.question))
6994 .await?;
6995
6996 let response = AgentResponse::new(&question.question)
6997 .with_metadata(
6998 "disambiguation",
6999 serde_json::json!({
7000 "status": "awaiting_clarification",
7001 "skill_id": skill_id,
7002 "options": question.options,
7003 "clarifying": question.clarifying,
7004 "detection": {
7005 "type": detection.ambiguity_type,
7006 "confidence": detection.confidence,
7007 "what_is_unclear": detection.what_is_unclear,
7008 }
7009 }),
7010 );
7011 self.finish_turn_if_root(&response).await?;
7012 return Ok(response);
7013 }
7014 DisambiguationResult::Clarified {
7015 enriched_input: re_enriched,
7016 ..
7017 } => {
7018 debug!(skill_id = %skill_id, "Skill re-check: clarified immediately, executing");
7019 *self.pending_skill_id.write() = None;
7021 let skill_response =
7022 self.execute_skill_by_id(skill_id, &re_enriched).await?;
7023 self.memory
7024 .add_message(ChatMessage::user(&re_enriched))
7025 .await?;
7026 return self
7027 .handle_skill_response(
7028 &re_enriched,
7029 skill_response,
7030 &HashMap::new(),
7031 )
7032 .await;
7033 }
7034 DisambiguationResult::ProceedWithBestGuess {
7035 enriched_input: re_enriched,
7036 } => {
7037 debug!(skill_id = %skill_id, "Skill re-check: proceeding with best guess");
7038 *self.pending_skill_id.write() = None;
7039 let skill_response =
7040 self.execute_skill_by_id(skill_id, &re_enriched).await?;
7041 self.memory
7042 .add_message(ChatMessage::user(&re_enriched))
7043 .await?;
7044 return self
7045 .handle_skill_response(
7046 &re_enriched,
7047 skill_response,
7048 &HashMap::new(),
7049 )
7050 .await;
7051 }
7052 DisambiguationResult::GiveUp { reason } => {
7053 *self.pending_skill_id.write() = None;
7054 let apology = self
7055 .generate_localized_apology(
7056 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
7057 &reason,
7058 )
7059 .await
7060 .unwrap_or_else(|_| {
7061 format!("I'm sorry, I couldn't understand your request: {}", reason)
7062 });
7063 let response = AgentResponse::new(&apology);
7064 self.finish_turn_if_root(&response).await?;
7065 return Ok(response);
7066 }
7067 DisambiguationResult::Escalate { reason } => {
7068 *self.pending_skill_id.write() = None;
7069 let apology = self
7070 .generate_localized_apology(
7071 "Explain briefly that you're transferring the user to a human agent for help.",
7072 &reason,
7073 )
7074 .await
7075 .unwrap_or_else(|_| {
7076 format!("I need human assistance to help with your request: {}", reason)
7077 });
7078 let response = AgentResponse::new(&apology);
7079 self.finish_turn_if_root(&response).await?;
7080 return Ok(response);
7081 }
7082 DisambiguationResult::Abandoned { new_input } => {
7083 *self.pending_skill_id.write() = None;
7086 debug!(skill_id = %skill_id, "Skill re-check: abandoned by user");
7087 if let Some(fresh) = new_input {
7088 return self.run_loop_internal(&fresh).await;
7089 }
7090 let ack = self
7091 .generate_localized_apology(
7092 "The user changed their mind about their previous request. \
7093 Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
7094 Do NOT apologize excessively. Be concise.",
7095 "User abandoned clarification",
7096 )
7097 .await
7098 .unwrap_or_else(|_| {
7099 "OK, no problem. What else can I help with?".to_string()
7100 });
7101 self.memory
7102 .add_message(ChatMessage::assistant(&ack))
7103 .await?;
7104 let response = AgentResponse::new(&ack);
7105 self.finish_turn_if_root(&response).await?;
7106 return Ok(response);
7107 }
7108 }
7109 }
7110 }
7111 }
7112 }
7113
7114 *self.pending_skill_id.write() = None;
7116 let skill_response = self.execute_skill_by_id(skill_id, enriched_input).await?;
7117 self.memory
7118 .add_message(ChatMessage::user(enriched_input))
7119 .await?;
7120 self.handle_skill_response(enriched_input, skill_response, &HashMap::new())
7121 .await
7122 }
7123
7124 async fn handle_skill_response(
7127 &self,
7128 processed_input: &str,
7129 skill_response: String,
7130 input_context: &HashMap<String, Value>,
7131 ) -> Result<AgentResponse> {
7132 let output_data = self.process_output(&skill_response, input_context).await?;
7133 let final_response = output_data.content;
7134
7135 self.memory
7136 .add_message(ChatMessage::assistant(&final_response))
7137 .await?;
7138
7139 self.check_memory_compression().await?;
7140
7141 self.increment_turn();
7142 self.evaluate_transitions(processed_input, &final_response)
7143 .await?;
7144
7145 let response = AgentResponse::new(final_response);
7146 self.finish_turn_if_root(&response).await?;
7147 Ok(response)
7148 }
7149
7150 async fn handle_plan_and_execute(
7153 &self,
7154 processed_input: &str,
7155 input_context: &HashMap<String, Value>,
7156 auto_detected: bool,
7157 ) -> Result<AgentResponse> {
7158 let effective = self.get_effective_reasoning_config();
7159 let plan_reflection = effective
7160 .get_planning()
7161 .map(|c| c.reflection.clone())
7162 .unwrap_or_default();
7163
7164 let max_attempts = if plan_reflection.enabled {
7165 1 + plan_reflection.max_replans
7166 } else {
7167 1
7168 };
7169
7170 let mut plan = self.generate_plan(processed_input).await?;
7171 info!(
7172 plan_id = %plan.id,
7173 steps = plan.steps.len(),
7174 "Plan generated"
7175 );
7176
7177 let mut plan_result = String::new();
7178
7179 for attempt in 0..max_attempts {
7180 *self.current_plan.write() = Some(plan.clone());
7181 plan_result = self.execute_plan(&mut plan).await?;
7182
7183 info!(
7184 plan_status = ?plan.status,
7185 completed_steps = plan.completed_steps().count(),
7186 attempt = attempt + 1,
7187 "Plan execution completed"
7188 );
7189
7190 if !plan_reflection.enabled {
7191 break;
7192 }
7193
7194 let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
7195 if !has_failures {
7196 break;
7197 }
7198
7199 if attempt + 1 >= max_attempts {
7200 break;
7201 }
7202
7203 match plan_reflection.on_step_failure {
7204 StepFailureAction::Replan => {
7205 info!(attempt = attempt + 1, "Plan had failures, replanning");
7206 plan = self.generate_plan(processed_input).await?;
7207 }
7208 StepFailureAction::Abort => {
7209 warn!("Plan step failed, aborting");
7210 break;
7211 }
7212 StepFailureAction::Skip | StepFailureAction::Continue => {
7213 break;
7214 }
7215 }
7216 }
7217
7218 *self.current_plan.write() = Some(plan);
7219
7220 let output_data = self.process_output(&plan_result, input_context).await?;
7221 let final_content = output_data.content;
7222
7223 self.memory
7224 .add_message(ChatMessage::assistant(&final_content))
7225 .await?;
7226
7227 self.check_memory_compression().await?;
7228 self.increment_turn();
7229 self.evaluate_transitions(processed_input, &final_content)
7230 .await?;
7231
7232 let reasoning_metadata =
7233 ReasoningMetadata::new(ReasoningMode::PlanAndExecute).with_auto_detected(auto_detected);
7234
7235 let response = AgentResponse::new(&final_content).with_metadata(
7236 "reasoning",
7237 serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
7238 );
7239
7240 self.finish_turn_if_root(&response).await?;
7241 Ok(response)
7242 }
7243
7244 fn inject_reasoning_prompt(
7246 &self,
7247 messages: &mut [ChatMessage],
7248 reasoning_mode: &ReasoningMode,
7249 is_first_iteration: bool,
7250 ) {
7251 if !is_first_iteration {
7252 return;
7253 }
7254 match reasoning_mode {
7255 ReasoningMode::CoT => {
7256 if let Some(msg) = messages.first_mut() {
7257 if matches!(msg.role, ai_agents_core::Role::System) {
7258 msg.content = self.build_cot_system_prompt(&msg.content);
7259 debug!("Applied Chain-of-Thought system prompt");
7260 }
7261 }
7262 }
7263 ReasoningMode::React => {
7264 if let Some(msg) = messages.first_mut() {
7265 if matches!(msg.role, ai_agents_core::Role::System) {
7266 msg.content = self.build_react_system_prompt(&msg.content);
7267 debug!("Applied ReAct system prompt");
7268 }
7269 }
7270 }
7271 _ => {}
7272 }
7273 }
7274
7275 async fn complete_llm_with_recovery(
7276 &self,
7277 llm: Arc<dyn LLMProvider>,
7278 messages: &[ChatMessage],
7279 ) -> Result<LLMResponse> {
7280 let primary_result = if self.recovery_manager.config().default.max_retries > 0 {
7281 self.recovery_manager
7282 .with_retry("llm_call", None, || async {
7283 self.observe_purpose(
7284 ObservationPurpose::MainResponse,
7285 llm.complete(messages, None),
7286 )
7287 .await
7288 .map_err(|e| e.classify())
7289 })
7290 .await
7291 .map_err(|e| AgentError::LLM(e.to_string()))
7292 } else {
7293 self.observe_purpose(
7294 ObservationPurpose::MainResponse,
7295 llm.complete(messages, None),
7296 )
7297 .await
7298 .map_err(|e| AgentError::LLM(e.to_string()))
7299 };
7300
7301 match primary_result {
7302 Ok(resp) => Ok(resp),
7303 Err(primary_err) => match &self.recovery_manager.config().llm.on_failure {
7304 LLMFailureAction::FallbackLlm { fallback_llm } => {
7305 let fb = self.llm_registry.get(fallback_llm).map_err(|e| {
7306 AgentError::Config(format!(
7307 "Fallback LLM '{}' not found: {}",
7308 fallback_llm, e
7309 ))
7310 })?;
7311 self.observe_purpose(
7312 ObservationPurpose::MainResponse,
7313 fb.complete(messages, None),
7314 )
7315 .await
7316 .map_err(|e| AgentError::LLM(e.to_string()))
7317 }
7318 LLMFailureAction::FallbackResponse { message } => {
7319 Ok(LLMResponse::new(message.clone(), FinishReason::Stop))
7320 }
7321 LLMFailureAction::Error => Err(primary_err),
7322 },
7323 }
7324 }
7325
7326 async fn generate_main_response_draft(
7331 &self,
7332 processed_input: &str,
7333 reasoning_mode: &ReasoningMode,
7334 ) -> Result<MainResponseDraft> {
7335 let llm = self.get_state_llm()?;
7336 let mut messages = self.build_messages_for_draft(processed_input).await?;
7337 self.inject_reasoning_prompt(&mut messages, reasoning_mode, true);
7338 let response = self.complete_llm_with_recovery(llm, &messages).await?;
7339 let content = response.content.trim().to_string();
7340 let (thinking, answer) = self.extract_thinking(&content);
7341 if let Some(calls) = self.parse_tool_calls(&content) {
7342 return Ok(MainResponseDraft::ToolCalls {
7343 raw_content: content,
7344 calls,
7345 thinking,
7346 });
7347 }
7348 Ok(MainResponseDraft::Text {
7349 raw_content: answer,
7350 thinking,
7351 })
7352 }
7353
7354 async fn commit_main_response_draft(
7359 &self,
7360 processed_input: &str,
7361 input_context: &HashMap<String, Value>,
7362 draft: MainResponseDraft,
7363 reasoning_mode: ReasoningMode,
7364 auto_detected: bool,
7365 ) -> Result<AgentResponse> {
7366 self.commit_root_user_message(processed_input).await?;
7367 match draft {
7368 MainResponseDraft::Text {
7369 raw_content,
7370 thinking,
7371 } => {
7372 self.finish_text_response_from_model(
7373 processed_input,
7374 input_context,
7375 raw_content,
7376 reasoning_mode,
7377 auto_detected,
7378 1,
7379 thinking,
7380 Vec::new(),
7381 )
7382 .await
7383 }
7384 MainResponseDraft::ToolCalls {
7385 raw_content,
7386 calls,
7387 thinking: _,
7388 } => {
7389 let mut all_tool_calls = Vec::new();
7390 match self
7391 .handle_tool_calls(processed_input, &raw_content, calls, &mut all_tool_calls)
7392 .await?
7393 {
7394 ToolCallOutcome::Rejected(response) => {
7395 self.finish_turn_if_root(&response).await?;
7396 Ok(response)
7397 }
7398 ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => {
7399 self.continue_after_committed_tool_draft(processed_input)
7400 .await
7401 }
7402 }
7403 }
7404 }
7405 }
7406
7407 async fn continue_after_committed_tool_draft(
7412 &self,
7413 processed_input: &str,
7414 ) -> Result<AgentResponse> {
7415 *self.redispatch_depth.write() += 1;
7416 if let Some(context) = self.active_turn_context.write().as_mut() {
7417 context.enter_redispatch();
7418 }
7419 let result = Box::pin(self.run_loop_internal(processed_input)).await;
7420 *self.redispatch_depth.write() -= 1;
7421 if let Some(context) = self.active_turn_context.write().as_mut() {
7422 context.exit_redispatch();
7423 }
7424 let response = result?;
7425 self.finish_turn_if_root(&response).await?;
7426 Ok(response)
7427 }
7428
7429 async fn finish_text_response_from_model(
7434 &self,
7435 processed_input: &str,
7436 input_context: &HashMap<String, Value>,
7437 answer: String,
7438 reasoning_mode: ReasoningMode,
7439 auto_detected: bool,
7440 iterations: u32,
7441 thinking_content: Option<String>,
7442 all_tool_calls: Vec<ToolCall>,
7443 ) -> Result<AgentResponse> {
7444 let output_data = self.process_output(&answer, input_context).await?;
7445 let mut final_content = if output_data.metadata.rejected {
7446 output_data
7447 .metadata
7448 .rejection_reason
7449 .unwrap_or_else(|| answer.to_string())
7450 } else {
7451 output_data.content
7452 };
7453 let llm = self.get_state_llm()?;
7454 let reflection_metadata;
7455 (final_content, reflection_metadata) = self
7456 .run_reflection(&*llm, processed_input, final_content)
7457 .await?;
7458 final_content =
7459 self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
7460 let final_content = {
7461 let result = self
7462 .post_loop_processing(processed_input, final_content)
7463 .await?;
7464 self.apply_post_loop_result(processed_input, result).await?
7465 };
7466 let response = self.build_agent_response(
7467 final_content,
7468 all_tool_calls,
7469 reasoning_mode,
7470 auto_detected,
7471 iterations,
7472 thinking_content,
7473 reflection_metadata,
7474 );
7475 self.finish_turn_if_root(&response).await?;
7476 Ok(response)
7477 }
7478
7479 async fn run_committed_response_loop_with_reasoning(
7484 &self,
7485 processed_input: &str,
7486 input_context: &HashMap<String, Value>,
7487 reasoning_mode: ReasoningMode,
7488 auto_detected: bool,
7489 ) -> Result<AgentResponse> {
7490 self.commit_root_user_message(processed_input).await?;
7491 let llm = self.get_state_llm()?;
7492 let mut iterations = 0u32;
7493 let mut all_tool_calls = Vec::new();
7494 let mut thinking_content = None;
7495 loop {
7496 let effective_max = if reasoning_mode != ReasoningMode::None {
7497 let rc = self.get_effective_reasoning_config();
7498 self.max_iterations.min(rc.max_iterations)
7499 } else {
7500 self.max_iterations
7501 };
7502 if iterations >= effective_max {
7503 return Err(AgentError::Other(format!(
7504 "Max iterations ({}) exceeded",
7505 effective_max
7506 )));
7507 }
7508 iterations += 1;
7509 *self.iteration_count.write() = iterations;
7510 let mut messages = self.build_messages().await?;
7511 self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
7512 self.hooks.on_llm_start(&messages).await;
7513 let llm_start = Instant::now();
7514 let response = self
7515 .observe_purpose(
7516 ObservationPurpose::MainResponse,
7517 llm.complete(&messages, None),
7518 )
7519 .await
7520 .map_err(|e| AgentError::LLM(e.to_string()))?;
7521 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
7522 self.hooks.on_llm_complete(&response, llm_duration_ms).await;
7523 let content = response.content.trim();
7524 if let Some(tool_calls) = self.parse_tool_calls(content) {
7525 match self
7526 .handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
7527 .await?
7528 {
7529 ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
7530 ToolCallOutcome::Rejected(resp) => {
7531 self.finish_turn_if_root(&resp).await?;
7532 return Ok(resp);
7533 }
7534 }
7535 }
7536 let (extracted_thinking, answer) = self.extract_thinking(content);
7537 if extracted_thinking.is_some() {
7538 thinking_content = extracted_thinking;
7539 }
7540 return self
7541 .finish_text_response_from_model(
7542 processed_input,
7543 input_context,
7544 answer,
7545 reasoning_mode,
7546 auto_detected,
7547 iterations,
7548 thinking_content,
7549 all_tool_calls,
7550 )
7551 .await;
7552 }
7553 }
7554
7555 async fn handle_tool_calls(
7557 &self,
7558 processed_input: &str,
7559 content: &str,
7560 tool_calls: Vec<ToolCall>,
7561 all_tool_calls: &mut Vec<ToolCall>,
7562 ) -> Result<ToolCallOutcome> {
7563 let transition_fired = self.evaluate_transitions(processed_input, content).await?;
7567 if transition_fired {
7568 self.memory
7569 .add_message(ChatMessage::assistant(
7570 "(Transitioned to new state — tool call handled by workflow)",
7571 ))
7572 .await?;
7573 return Ok(ToolCallOutcome::TransitionFired);
7574 }
7575
7576 self.memory
7578 .add_message(ChatMessage::assistant(content))
7579 .await?;
7580
7581 let results = self.execute_tools_parallel(&tool_calls).await;
7582
7583 for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
7584 match result {
7585 Ok(output) => {
7586 self.memory
7587 .add_message(ChatMessage::function(&tool_call.name, &output))
7588 .await?;
7589 }
7590 Err(e) => {
7591 if matches!(e, AgentError::HITLRejected(_)) {
7593 self.memory
7594 .add_message(ChatMessage::assistant(&format!(
7595 "The operation was rejected by the approver: {}",
7596 e
7597 )))
7598 .await?;
7599 return Ok(ToolCallOutcome::Rejected(AgentResponse {
7601 content: format!("Operation cancelled: {}", e),
7602 metadata: None,
7603 tool_calls: Some(all_tool_calls.clone()),
7604 }));
7605 }
7606 self.memory
7607 .add_message(ChatMessage::function(
7608 &tool_call.name,
7609 &format!("Error: {}", e),
7610 ))
7611 .await?;
7612 }
7613 }
7614 all_tool_calls.push(tool_call.clone());
7615 }
7616 Ok(ToolCallOutcome::Continue)
7617 }
7618
7619 async fn run_reflection(
7621 &self,
7622 llm: &dyn LLMProvider,
7623 processed_input: &str,
7624 mut content: String,
7625 ) -> Result<(String, Option<ReflectionMetadata>)> {
7626 let should_reflect = self.should_reflect(processed_input, &content).await?;
7627 if !should_reflect {
7628 return Ok((content, None));
7629 }
7630
7631 info!("Starting response reflection evaluation");
7632 let mut attempts = 0u32;
7633 let max_retries = self.reflection_config.max_retries;
7634 let mut history: Vec<ReflectionAttempt> = Vec::new();
7635
7636 loop {
7637 let evaluation = self.evaluate_response(processed_input, &content).await?;
7638
7639 if evaluation.passed || attempts >= max_retries {
7640 info!(
7641 passed = evaluation.passed,
7642 confidence = evaluation.confidence,
7643 attempts = attempts + 1,
7644 "Reflection evaluation complete"
7645 );
7646 let reflection_metadata = Some(
7647 ReflectionMetadata::new(evaluation)
7648 .with_attempts(attempts + 1)
7649 .with_history(history),
7650 );
7651 return Ok((content, reflection_metadata));
7652 }
7653
7654 debug!(
7655 attempt = attempts + 1,
7656 failed_criteria = evaluation.failed_criteria().count(),
7657 "Response did not meet criteria, retrying"
7658 );
7659
7660 history.push(
7661 ReflectionAttempt::new(&content, evaluation.clone())
7662 .with_feedback("Response did not meet quality criteria"),
7663 );
7664
7665 let feedback: Vec<String> = evaluation
7666 .failed_criteria()
7667 .map(|c| format!("- {}", c.criterion))
7668 .collect();
7669
7670 let retry_prompt = format!(
7671 "Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response.",
7672 feedback.join("\n")
7673 );
7674
7675 self.memory
7676 .add_message(ChatMessage::user(&retry_prompt))
7677 .await?;
7678
7679 let retry_messages = self.build_messages().await?;
7680 let retry_response = self
7681 .observe_purpose(
7682 ObservationPurpose::ReflectionEvaluation,
7683 llm.complete(&retry_messages, None),
7684 )
7685 .await
7686 .map_err(|e| AgentError::LLM(e.to_string()))?;
7687
7688 content = retry_response.content.trim().to_string();
7689 attempts += 1;
7690 }
7691 }
7692
7693 async fn post_loop_processing(
7696 &self,
7697 processed_input: &str,
7698 content: String,
7699 ) -> Result<PostLoopResult> {
7700 self.increment_turn();
7705
7706 self.run_context_extractors(processed_input).await;
7708
7709 let transitioned = self.evaluate_transitions(processed_input, &content).await?;
7710
7711 if !transitioned {
7712 self.memory
7713 .add_message(ChatMessage::assistant(&content))
7714 .await?;
7715 self.check_memory_compression().await?;
7716 return Ok(PostLoopResult::NoTransition(content));
7717 }
7718
7719 if !self.should_regenerate_after_transition() {
7721 self.memory
7722 .add_message(ChatMessage::assistant(&content))
7723 .await?;
7724 self.check_memory_compression().await?;
7725 return Ok(PostLoopResult::Transitioned(content));
7726 }
7727
7728 if self.needs_redispatch_for_new_state() {
7732 info!("Post-transition NeedsRedispatch: new state requires full dispatch");
7733 return Ok(PostLoopResult::NeedsRedispatch);
7736 }
7737
7738 self.memory
7741 .add_message(ChatMessage::assistant(&content))
7742 .await?;
7743 self.check_memory_compression().await?;
7744
7745 let new_llm = self.get_state_llm()?;
7751 let mut final_content;
7752
7753 for post_iter in 0..self.max_iterations {
7754 let new_messages = self.build_messages().await?;
7755 if post_iter == 0 {
7756 if let Some(system_msg) = new_messages.first() {
7757 if system_msg.role == ai_agents_core::Role::System {
7758 debug!(
7759 prompt_preview =
7760 &system_msg.content[system_msg.content.len().saturating_sub(200)..],
7761 "Post-transition system prompt (last 200 chars)"
7762 );
7763 }
7764 }
7765 }
7766
7767 let new_response = self
7768 .observe_purpose(
7769 ObservationPurpose::MainResponse,
7770 new_llm.complete(&new_messages, None),
7771 )
7772 .await
7773 .map_err(|e| AgentError::LLM(e.to_string()))?;
7774 final_content = new_response.content.trim().to_string();
7775
7776 if let Some(tool_calls) = self.parse_tool_calls(&final_content) {
7779 debug!(
7780 post_iter = post_iter,
7781 tools = tool_calls.len(),
7782 "Post-transition tool call detected, executing"
7783 );
7784
7785 self.memory
7786 .add_message(ChatMessage::assistant(&final_content))
7787 .await?;
7788
7789 let results = self.execute_tools_parallel(&tool_calls).await;
7790 for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
7791 match result {
7792 Ok(output) => {
7793 self.memory
7794 .add_message(ChatMessage::function(&tool_call.name, &output))
7795 .await?;
7796 }
7797 Err(e) => {
7798 self.memory
7799 .add_message(ChatMessage::function(
7800 &tool_call.name,
7801 &format!("Error: {}", e),
7802 ))
7803 .await?;
7804 }
7805 }
7806 }
7807 continue;
7809 }
7810
7811 self.memory
7813 .add_message(ChatMessage::assistant(&final_content))
7814 .await?;
7815 return Ok(PostLoopResult::Transitioned(final_content));
7816 }
7817
7818 final_content = "Post-transition processing completed.".to_string();
7820 self.memory
7821 .add_message(ChatMessage::assistant(&final_content))
7822 .await?;
7823
7824 Ok(PostLoopResult::Transitioned(final_content))
7825 }
7826
7827 fn should_regenerate_after_transition(&self) -> bool {
7830 if let Some(ref sm) = self.state_machine {
7831 if !sm.config().regenerate_on_transition {
7833 return false;
7834 }
7835 if let Some(def) = sm.current_definition() {
7837 if let Some(regen) = def.regenerate_on_enter {
7838 return regen;
7839 }
7840 }
7841 }
7842 true
7843 }
7844
7845 fn needs_redispatch_for_new_state(&self) -> bool {
7848 if let Some(ref sm) = self.state_machine {
7849 if let Some(def) = sm.current_definition() {
7850 if def.concurrent.is_some()
7851 || def.group_chat.is_some()
7852 || def.pipeline.is_some()
7853 || def.handoff.is_some()
7854 || def.delegate.is_some()
7855 {
7856 return true;
7857 }
7858 let effective = self.get_effective_reasoning_config();
7860 if !matches!(effective.mode, ReasoningMode::None) {
7861 return true;
7862 }
7863 }
7864 }
7865 false
7866 }
7867
7868 async fn apply_post_loop_result(
7871 &self,
7872 processed_input: &str,
7873 result: PostLoopResult,
7874 ) -> Result<String> {
7875 match result {
7876 PostLoopResult::NoTransition(content) | PostLoopResult::Transitioned(content) => {
7877 Ok(content)
7878 }
7879 PostLoopResult::NeedsRedispatch => {
7880 const MAX_REDISPATCH_DEPTH: u32 = 3;
7881 let current_depth = *self.redispatch_depth.read();
7882 if current_depth >= MAX_REDISPATCH_DEPTH {
7883 warn!(
7884 depth = current_depth,
7885 "Post-transition re-dispatch depth limit reached, returning empty response"
7886 );
7887 let content = String::new();
7888 self.memory
7889 .add_message(ChatMessage::assistant(&content))
7890 .await?;
7891 return Ok(content);
7892 }
7893 *self.redispatch_depth.write() += 1;
7894 if let Some(context) = self.active_turn_context.write().as_mut() {
7895 context.enter_redispatch();
7896 }
7897 info!(
7898 depth = current_depth + 1,
7899 "Re-dispatching for new state after transition"
7900 );
7901 let resp = Box::pin(self.run_loop_internal(processed_input)).await;
7902 *self.redispatch_depth.write() -= 1;
7903 if let Some(context) = self.active_turn_context.write().as_mut() {
7904 context.exit_redispatch();
7905 }
7906 resp.map(|r| r.content)
7907 }
7908 }
7909 }
7910
7911 fn build_agent_response(
7912 &self,
7913 content: String,
7914 all_tool_calls: Vec<ToolCall>,
7915 reasoning_mode: ReasoningMode,
7916 auto_detected: bool,
7917 iterations: u32,
7918 thinking: Option<String>,
7919 reflection_metadata: Option<ReflectionMetadata>,
7920 ) -> AgentResponse {
7921 let reasoning_metadata = ReasoningMetadata::new(reasoning_mode.clone())
7922 .with_thinking(thinking.clone().unwrap_or_default())
7923 .with_iterations(iterations)
7924 .with_auto_detected(auto_detected);
7925
7926 let mut response = AgentResponse::new(&content);
7927 if !all_tool_calls.is_empty() {
7928 response = response.with_tool_calls(all_tool_calls);
7929 }
7930
7931 if let Some(state) = self.current_state() {
7932 response = response.with_metadata("current_state", serde_json::json!(state));
7933 }
7934
7935 response = response.with_metadata(
7936 "reasoning",
7937 serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
7938 );
7939
7940 if let Some(ref refl_meta) = reflection_metadata {
7941 response = response.with_metadata(
7942 "reflection",
7943 serde_json::to_value(refl_meta).unwrap_or_default(),
7944 );
7945 }
7946
7947 response
7948 }
7949
7950 async fn handle_delegated_state(
7952 &self,
7953 input: &str,
7954 delegate_id: &str,
7955 state_def: &ai_agents_state::StateDefinition,
7956 ) -> Result<AgentResponse> {
7957 use std::time::Instant;
7958
7959 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
7960 AgentError::Config(format!(
7961 "State delegates to '{}' but no agent registry is configured. \
7962 Add a spawner section with auto_spawn to your YAML.",
7963 delegate_id
7964 ))
7965 })?;
7966
7967 let state_name = self
7968 .state_machine
7969 .as_ref()
7970 .map(|sm| sm.current())
7971 .unwrap_or_else(|| "unknown".to_string());
7972
7973 self.hooks.on_delegate_start(delegate_id, &state_name).await;
7974 let start = Instant::now();
7975
7976 let delegate = registry.get(delegate_id).ok_or_else(|| {
7977 AgentError::Other(format!(
7978 "State '{}' delegates to '{}' but no agent with that ID exists in the registry.",
7979 state_name, delegate_id
7980 ))
7981 })?;
7982
7983 let context_mode = state_def.delegate_context.clone().unwrap_or_default();
7985 let effective_input = self
7986 .observe_purpose(
7987 ObservationPurpose::OrchestrationRouting,
7988 crate::orchestration::context::prepare_delegate_input(
7989 input,
7990 &context_mode,
7991 &*self.memory,
7992 self.llm_registry.get("router").ok().as_deref(),
7993 ),
7994 )
7995 .await?;
7996
7997 let response = delegate
7998 .chat_with_actor_context(&effective_input, self.outbound_actor_context())
7999 .await?;
8000
8001 let duration_ms = start.elapsed().as_millis() as u64;
8002 self.hooks
8003 .on_delegate_complete(delegate_id, &state_name, duration_ms)
8004 .await;
8005
8006 let ctx_key = format!("delegation.{}.last_response", delegate_id);
8008 let _ = self.context_manager.set(
8009 &ctx_key,
8010 serde_json::Value::String(response.content.clone()),
8011 );
8012
8013 let _ = self.context_manager.set(
8015 "orchestration",
8016 serde_json::json!({
8017 "type": "delegate",
8018 "agent": delegate_id,
8019 "state": state_name,
8020 "response": response.content,
8021 "duration_ms": duration_ms,
8022 }),
8023 );
8024
8025 self.commit_root_user_message(input).await?;
8026
8027 let post_result = self
8030 .post_loop_processing(
8031 input,
8032 format!("[Delegated to {}]: {}", delegate_id, response.content),
8033 )
8034 .await?;
8035 let final_content = self.apply_post_loop_result(input, post_result).await?;
8036
8037 let mut result = AgentResponse::new(final_content);
8038
8039 let metadata = serde_json::json!({
8040 "orchestration": {
8041 "type": "delegate",
8042 "agent": delegate_id,
8043 "state": state_name,
8044 "response": response.content,
8045 "duration_ms": duration_ms,
8046 }
8047 });
8048 result.metadata = Some(
8049 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
8050 metadata,
8051 )
8052 .unwrap_or_default(),
8053 );
8054
8055 self.finish_turn_if_root(&result).await?;
8056 Ok(result)
8057 }
8058
8059 async fn handle_concurrent_state(
8061 &self,
8062 input: &str,
8063 config: &ai_agents_state::ConcurrentStateConfig,
8064 ) -> Result<AgentResponse> {
8065 use std::time::Instant;
8066
8067 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
8068 AgentError::Config(
8069 "Concurrent state requires an agent registry. Add a spawner section.".into(),
8070 )
8071 })?;
8072
8073 let context_mode = config.context_mode.clone().unwrap_or_default();
8078 let context_input = self
8079 .observe_purpose(
8080 ObservationPurpose::OrchestrationRouting,
8081 crate::orchestration::context::prepare_delegate_input(
8082 input,
8083 &context_mode,
8084 &*self.memory,
8085 self.llm_registry.get("router").ok().as_deref(),
8086 ),
8087 )
8088 .await?;
8089
8090 let effective_input = if let Some(ref tmpl) = config.input {
8091 render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
8092 .unwrap_or_else(|_| context_input.clone())
8093 } else {
8094 context_input
8095 };
8096
8097 let start = Instant::now();
8098
8099 let llm_name = config
8100 .aggregation
8101 .synthesizer_llm
8102 .as_deref()
8103 .unwrap_or("router");
8104 let llm_provider = self.llm_registry.get(llm_name).ok();
8105
8106 let vote_parallelism = if self.runtime_config.optimization.enabled
8107 && self
8108 .runtime_config
8109 .optimization
8110 .parallel_orchestration_vote_extraction
8111 {
8112 Some(self.runtime_config.optimization.max_parallel_runtime_tasks)
8113 } else {
8114 None
8115 };
8116
8117 let result = self
8118 .observe_purpose(
8119 ObservationPurpose::OrchestrationAggregation,
8120 scope_actor_context(
8121 self.outbound_actor_context(),
8122 crate::orchestration::concurrent(
8123 registry,
8124 &effective_input,
8125 &config.agents,
8126 &config.aggregation,
8127 llm_provider.as_deref(),
8128 config.min_required,
8129 config.timeout_ms,
8130 config.on_partial_failure.clone(),
8131 vote_parallelism,
8132 ),
8133 ),
8134 )
8135 .await?;
8136
8137 let duration_ms = start.elapsed().as_millis() as u64;
8138 let agent_ids: Vec<String> = config.agents.iter().map(|a| a.id().to_string()).collect();
8139 let strategy = format!("{:?}", config.aggregation.strategy);
8140 self.hooks
8141 .on_concurrent_complete(&agent_ids, &strategy, duration_ms)
8142 .await;
8143
8144 let _ = self.context_manager.set(
8146 "concurrent.result",
8147 serde_json::Value::String(result.response.content.clone()),
8148 );
8149
8150 let agents_json: Vec<serde_json::Value> = result
8152 .agent_results
8153 .iter()
8154 .map(|ar| {
8155 serde_json::json!({
8156 "id": ar.agent_id,
8157 "response": ar.response.as_ref().map(|r| r.content.as_str()),
8158 "success": ar.success,
8159 "error": ar.error,
8160 "duration_ms": ar.duration_ms,
8161 })
8162 })
8163 .collect();
8164
8165 let _ = self.context_manager.set(
8167 "orchestration",
8168 serde_json::json!({
8169 "type": "concurrent",
8170 "result": result.response.content,
8171 "strategy": strategy,
8172 "agents": agents_json,
8173 "duration_ms": duration_ms,
8174 }),
8175 );
8176
8177 self.commit_root_user_message(input).await?;
8178
8179 let post_result = self
8180 .post_loop_processing(input, result.response.content.clone())
8181 .await?;
8182 let final_content = self.apply_post_loop_result(input, post_result).await?;
8183
8184 let mut response = AgentResponse::new(final_content);
8185 let metadata = serde_json::json!({
8186 "orchestration": {
8187 "type": "concurrent",
8188 "result": result.response.content,
8189 "strategy": strategy,
8190 "agents": agents_json,
8191 "duration_ms": duration_ms,
8192 }
8193 });
8194 response.metadata = Some(
8195 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
8196 metadata,
8197 )
8198 .unwrap_or_default(),
8199 );
8200
8201 self.finish_turn_if_root(&response).await?;
8202 Ok(response)
8203 }
8204
8205 async fn handle_group_chat_state(
8207 &self,
8208 input: &str,
8209 config: &ai_agents_state::GroupChatStateConfig,
8210 ) -> Result<AgentResponse> {
8211 use std::time::Instant;
8212
8213 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
8214 AgentError::Config(
8215 "Group chat state requires an agent registry. Add a spawner section.".into(),
8216 )
8217 })?;
8218
8219 let start = Instant::now();
8220
8221 let llm_provider = self.llm_registry.get("router").ok();
8222
8223 let context_mode = config.context_mode.clone().unwrap_or_default();
8225 let context_input = self
8226 .observe_purpose(
8227 ObservationPurpose::OrchestrationRouting,
8228 crate::orchestration::context::prepare_delegate_input(
8229 input,
8230 &context_mode,
8231 &*self.memory,
8232 self.llm_registry.get("router").ok().as_deref(),
8233 ),
8234 )
8235 .await?;
8236
8237 let effective_topic = if let Some(ref tmpl) = config.input {
8239 render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
8240 .unwrap_or_else(|_| context_input.clone())
8241 } else {
8242 context_input
8243 };
8244
8245 let result = self
8246 .observe_purpose(
8247 ObservationPurpose::OrchestrationConversation,
8248 scope_actor_context(
8249 self.outbound_actor_context(),
8250 crate::orchestration::group_chat(
8251 registry,
8252 &effective_topic,
8253 config,
8254 llm_provider.as_deref(),
8255 Some(&*self.hooks),
8256 ),
8257 ),
8258 )
8259 .await?;
8260
8261 let duration_ms = start.elapsed().as_millis() as u64;
8262
8263 let _ = self.context_manager.set(
8265 "group_chat.conclusion",
8266 serde_json::Value::String(result.response.content.clone()),
8267 );
8268
8269 let transcript_json: Vec<serde_json::Value> = result
8271 .transcript
8272 .iter()
8273 .map(|t| {
8274 serde_json::json!({
8275 "speaker": t.speaker,
8276 "round": t.round,
8277 "content": t.content,
8278 })
8279 })
8280 .collect();
8281
8282 let _ = self.context_manager.set(
8284 "orchestration",
8285 serde_json::json!({
8286 "type": "group_chat",
8287 "conclusion": result.response.content,
8288 "transcript": transcript_json,
8289 "rounds": result.rounds_completed,
8290 "termination": result.termination_reason,
8291 "duration_ms": duration_ms,
8292 }),
8293 );
8294
8295 self.commit_root_user_message(input).await?;
8296
8297 let post_result = self
8298 .post_loop_processing(input, result.response.content.clone())
8299 .await?;
8300 let final_content = self.apply_post_loop_result(input, post_result).await?;
8301
8302 let mut response = AgentResponse::new(final_content);
8303 let metadata = serde_json::json!({
8304 "orchestration": {
8305 "type": "group_chat",
8306 "conclusion": result.response.content,
8307 "transcript": transcript_json,
8308 "rounds": result.rounds_completed,
8309 "termination": result.termination_reason,
8310 "duration_ms": duration_ms,
8311 }
8312 });
8313 response.metadata = Some(
8314 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
8315 metadata,
8316 )
8317 .unwrap_or_default(),
8318 );
8319
8320 self.finish_turn_if_root(&response).await?;
8321 Ok(response)
8322 }
8323
8324 async fn handle_pipeline_state(
8326 &self,
8327 input: &str,
8328 config: &ai_agents_state::PipelineStateConfig,
8329 ) -> Result<AgentResponse> {
8330 use std::time::Instant;
8331
8332 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
8333 AgentError::Config(
8334 "Pipeline state requires an agent registry. Add a spawner section.".into(),
8335 )
8336 })?;
8337
8338 let start = Instant::now();
8339
8340 let stages: Vec<crate::orchestration::PipelineStage> = config
8341 .stages
8342 .iter()
8343 .map(|entry| {
8344 let mut stage = crate::orchestration::PipelineStage::id(entry.id());
8345 if let Some(tmpl) = entry.input() {
8346 stage = stage.with_input(tmpl);
8347 }
8348 stage
8349 })
8350 .collect();
8351
8352 let context_mode = config.context_mode.clone().unwrap_or_default();
8354 let context_input = self
8355 .observe_purpose(
8356 ObservationPurpose::OrchestrationRouting,
8357 crate::orchestration::context::prepare_delegate_input(
8358 input,
8359 &context_mode,
8360 &*self.memory,
8361 self.llm_registry.get("router").ok().as_deref(),
8362 ),
8363 )
8364 .await?;
8365
8366 let context_values = self.build_context_with_overlays();
8367 let result = self
8368 .observe_purpose(
8369 ObservationPurpose::OrchestrationRouting,
8370 scope_actor_context(
8371 self.outbound_actor_context(),
8372 crate::orchestration::pipeline(
8373 registry,
8374 &context_input,
8375 &stages,
8376 config.timeout_ms,
8377 Some(&*self.hooks),
8378 Some(&context_values),
8379 ),
8380 ),
8381 )
8382 .await?;
8383
8384 let duration_ms = start.elapsed().as_millis() as u64;
8385
8386 let _ = self.context_manager.set(
8388 "pipeline.result",
8389 serde_json::Value::String(result.response.content.clone()),
8390 );
8391
8392 let stages_json: Vec<serde_json::Value> = result
8394 .stage_outputs
8395 .iter()
8396 .map(|s| {
8397 serde_json::json!({
8398 "agent_id": s.agent_id,
8399 "output": s.output,
8400 "duration_ms": s.duration_ms,
8401 "skipped": s.skipped,
8402 })
8403 })
8404 .collect();
8405
8406 let _ = self.context_manager.set(
8408 "orchestration",
8409 serde_json::json!({
8410 "type": "pipeline",
8411 "result": result.response.content,
8412 "stages": stages_json,
8413 "duration_ms": duration_ms,
8414 }),
8415 );
8416
8417 self.commit_root_user_message(input).await?;
8418
8419 let post_result = self
8420 .post_loop_processing(input, result.response.content.clone())
8421 .await?;
8422 let final_content = self.apply_post_loop_result(input, post_result).await?;
8423
8424 let mut response = AgentResponse::new(final_content);
8425 let metadata = serde_json::json!({
8426 "orchestration": {
8427 "type": "pipeline",
8428 "result": result.response.content,
8429 "stages": stages_json,
8430 "duration_ms": duration_ms,
8431 }
8432 });
8433 response.metadata = Some(
8434 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
8435 metadata,
8436 )
8437 .unwrap_or_default(),
8438 );
8439
8440 self.finish_turn_if_root(&response).await?;
8441 Ok(response)
8442 }
8443
8444 async fn handle_handoff_state(
8446 &self,
8447 input: &str,
8448 config: &ai_agents_state::HandoffStateConfig,
8449 ) -> Result<AgentResponse> {
8450 use std::time::Instant;
8451
8452 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
8453 AgentError::Config(
8454 "Handoff state requires an agent registry. Add a spawner section.".into(),
8455 )
8456 })?;
8457
8458 let llm = self
8459 .llm_registry
8460 .get("router")
8461 .map_err(|_| AgentError::Config("Handoff state requires a router LLM.".into()))?;
8462
8463 let start = Instant::now();
8464
8465 let context_mode = config.context_mode.clone().unwrap_or_default();
8467 let context_input = self
8468 .observe_purpose(
8469 ObservationPurpose::OrchestrationRouting,
8470 crate::orchestration::context::prepare_delegate_input(
8471 input,
8472 &context_mode,
8473 &*self.memory,
8474 self.llm_registry.get("router").ok().as_deref(),
8475 ),
8476 )
8477 .await?;
8478
8479 let effective_input = if let Some(ref tmpl) = config.input {
8481 render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
8482 .unwrap_or_else(|_| context_input.clone())
8483 } else {
8484 context_input
8485 };
8486
8487 let result = self
8488 .observe_purpose(
8489 ObservationPurpose::OrchestrationRouting,
8490 scope_actor_context(
8491 self.outbound_actor_context(),
8492 crate::orchestration::handoff(
8493 registry,
8494 &effective_input,
8495 &config.initial_agent,
8496 &config.available_agents,
8497 config.max_handoffs,
8498 llm.as_ref(),
8499 Some(&*self.hooks),
8500 ),
8501 ),
8502 )
8503 .await?;
8504
8505 let duration_ms = start.elapsed().as_millis() as u64;
8506
8507 let _ = self.context_manager.set(
8509 "handoff.result",
8510 serde_json::Value::String(result.response.content.clone()),
8511 );
8512
8513 let chain_json: Vec<serde_json::Value> = result
8515 .handoff_chain
8516 .iter()
8517 .map(|h| {
8518 serde_json::json!({
8519 "from": h.from_agent,
8520 "to": h.to_agent,
8521 "reason": h.reason,
8522 })
8523 })
8524 .collect();
8525
8526 let _ = self.context_manager.set(
8528 "orchestration",
8529 serde_json::json!({
8530 "type": "handoff",
8531 "result": result.response.content,
8532 "final_agent": result.final_agent,
8533 "handoff_chain": chain_json,
8534 "duration_ms": duration_ms,
8535 }),
8536 );
8537
8538 self.commit_root_user_message(input).await?;
8539
8540 let post_result = self
8541 .post_loop_processing(input, result.response.content.clone())
8542 .await?;
8543 let final_content = self.apply_post_loop_result(input, post_result).await?;
8544
8545 let mut response = AgentResponse::new(final_content);
8546 let metadata = serde_json::json!({
8547 "orchestration": {
8548 "type": "handoff",
8549 "result": result.response.content,
8550 "final_agent": result.final_agent,
8551 "handoff_chain": chain_json,
8552 "duration_ms": duration_ms,
8553 }
8554 });
8555 response.metadata = Some(
8556 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
8557 metadata,
8558 )
8559 .unwrap_or_default(),
8560 );
8561
8562 self.finish_turn_if_root(&response).await?;
8563 Ok(response)
8564 }
8565
8566 async fn run_loop_internal(&self, input: &str) -> Result<AgentResponse> {
8568 self.begin_root_turn();
8569 self.pre_turn_session_lifecycle().await;
8571
8572 let input_data = self.process_input(input).await?;
8573 self.update_active_turn_context(&input_data.content, input_data.context.clone());
8574
8575 for (key, value) in &input_data.context {
8578 let _ = self.context_manager.set(key, value.clone());
8579 }
8580
8581 if input_data.metadata.rejected {
8582 let reason = input_data
8583 .metadata
8584 .rejection_reason
8585 .unwrap_or_else(|| "Input rejected".to_string());
8586 warn!(reason = %reason, "Input rejected");
8587 let response = AgentResponse::new(reason);
8588 self.finish_turn_if_root(&response).await?;
8589 return Ok(response);
8590 }
8591
8592 let processed_input = &input_data.content;
8593
8594 if let Some(response) = self.try_pre_response_transition(processed_input).await? {
8595 return Ok(response);
8596 }
8597
8598 if let Some(ref sm) = self.state_machine {
8600 if let Some(def) = sm.current_definition() {
8601 if let Some(ref delegate_id) = def.delegate {
8602 return self
8603 .handle_delegated_state(processed_input, delegate_id, &def)
8604 .await;
8605 }
8606 if let Some(ref concurrent_config) = def.concurrent {
8607 return self
8608 .handle_concurrent_state(processed_input, concurrent_config)
8609 .await;
8610 }
8611 if let Some(ref group_chat_config) = def.group_chat {
8612 return self
8613 .handle_group_chat_state(processed_input, group_chat_config)
8614 .await;
8615 }
8616 if let Some(ref pipeline_config) = def.pipeline {
8617 return self
8618 .handle_pipeline_state(processed_input, pipeline_config)
8619 .await;
8620 }
8621 if let Some(ref handoff_config) = def.handoff {
8622 return self
8623 .handle_handoff_state(processed_input, handoff_config)
8624 .await;
8625 }
8626 }
8627 }
8628
8629 if let Some(response) =
8634 Box::pin(self.try_speculative_branches(processed_input, &input_data.context)).await?
8635 {
8636 return Ok(response);
8637 }
8638
8639 match self.try_skill_route(processed_input).await? {
8640 SkillRouteResult::Response(skill_response) => {
8641 self.commit_root_user_message(processed_input).await?;
8642 return self
8643 .handle_skill_response(processed_input, skill_response, &input_data.context)
8644 .await;
8645 }
8646 SkillRouteResult::NeedsClarification(response) => {
8647 self.commit_root_user_message(processed_input).await?;
8648 if let Some(q) = response
8649 .metadata
8650 .as_ref()
8651 .and_then(|m| m.get("disambiguation"))
8652 .and_then(|d| d.get("status"))
8653 .and_then(|s| s.as_str())
8654 {
8655 if q == "awaiting_clarification" {
8656 self.memory
8659 .add_message(ChatMessage::assistant(&response.content))
8660 .await?;
8661 }
8662 }
8663 self.finish_turn_if_root(&response).await?;
8664 return Ok(response);
8665 }
8666 SkillRouteResult::NoMatch => {} }
8668
8669 let effective_reasoning = self.get_effective_reasoning_config();
8670 let reasoning_mode = self.determine_reasoning_mode(processed_input).await?;
8671 let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
8672
8673 info!(
8674 reasoning_mode = ?reasoning_mode,
8675 auto_detected = auto_detected,
8676 reflection_enabled = ?self.reflection_config.enabled,
8677 "Reasoning mode determined"
8678 );
8679
8680 if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
8681 self.commit_root_user_message(processed_input).await?;
8682 return self
8683 .handle_plan_and_execute(processed_input, &input_data.context, auto_detected)
8684 .await;
8685 }
8686
8687 self.commit_root_user_message(processed_input).await?;
8688
8689 let mut iterations = 0u32;
8690 let mut all_tool_calls: Vec<ToolCall> = Vec::new();
8691 let mut thinking_content: Option<String> = None;
8692
8693 let llm = self.get_state_llm()?;
8694
8695 loop {
8696 let effective_max = if reasoning_mode != ReasoningMode::None {
8698 let rc = self.get_effective_reasoning_config();
8699 self.max_iterations.min(rc.max_iterations)
8700 } else {
8701 self.max_iterations
8702 };
8703
8704 if iterations >= effective_max {
8705 let err = AgentError::Other(format!("Max iterations ({}) exceeded", effective_max));
8706 self.hooks.on_error(&err).await;
8707 error!(iterations = iterations, "Max iterations exceeded");
8708 return Err(err);
8709 }
8710 iterations += 1;
8711 *self.iteration_count.write() = iterations;
8712
8713 debug!(iteration = iterations, max = effective_max, "LLM call");
8714
8715 let mut messages = self.build_messages().await?;
8716 self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
8717
8718 self.hooks.on_llm_start(&messages).await;
8719 let llm_start = Instant::now();
8720
8721 let response = {
8723 let primary_result = if self.recovery_manager.config().default.max_retries > 0 {
8724 self.recovery_manager
8725 .with_retry("llm_call", None, || async {
8726 self.observe_purpose(
8727 ObservationPurpose::MainResponse,
8728 llm.complete(&messages, None),
8729 )
8730 .await
8731 .map_err(|e| e.classify())
8732 })
8733 .await
8734 .map_err(|e| AgentError::LLM(e.to_string()))
8735 } else {
8736 self.observe_purpose(
8737 ObservationPurpose::MainResponse,
8738 llm.complete(&messages, None),
8739 )
8740 .await
8741 .map_err(|e| AgentError::LLM(e.to_string()))
8742 };
8743
8744 match primary_result {
8745 Ok(resp) => resp,
8746 Err(primary_err) => match &self.recovery_manager.config().llm.on_failure {
8747 LLMFailureAction::FallbackLlm { fallback_llm } => {
8748 warn!(
8749 fallback = %fallback_llm,
8750 error = %primary_err,
8751 "Primary LLM failed, attempting fallback LLM"
8752 );
8753 let fb = self.llm_registry.get(fallback_llm).map_err(|e| {
8754 AgentError::Config(format!(
8755 "Fallback LLM '{}' not found: {}",
8756 fallback_llm, e
8757 ))
8758 })?;
8759 self.observe_purpose(
8760 ObservationPurpose::MainResponse,
8761 fb.complete(&messages, None),
8762 )
8763 .await
8764 .map_err(|e| AgentError::LLM(e.to_string()))?
8765 }
8766 LLMFailureAction::FallbackResponse { message } => {
8767 warn!(
8768 error = %primary_err,
8769 "Primary LLM failed, using static fallback response"
8770 );
8771 LLMResponse::new(message.clone(), FinishReason::Stop)
8772 }
8773 LLMFailureAction::Error => {
8774 return Err(primary_err);
8775 }
8776 },
8777 }
8778 };
8779
8780 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
8781 self.hooks.on_llm_complete(&response, llm_duration_ms).await;
8782
8783 let content = response.content.trim();
8784
8785 if let Some(tool_calls) = self.parse_tool_calls(content) {
8786 match self
8787 .handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
8788 .await?
8789 {
8790 ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
8791 ToolCallOutcome::Rejected(resp) => {
8792 self.finish_turn_if_root(&resp).await?;
8793 return Ok(resp);
8794 }
8795 }
8796 }
8797
8798 let (extracted_thinking, answer) = self.extract_thinking(content);
8799 if extracted_thinking.is_some() {
8800 thinking_content = extracted_thinking;
8801 }
8802
8803 let output_data = self.process_output(&answer, &input_data.context).await?;
8804
8805 let mut final_content = if output_data.metadata.rejected {
8806 output_data
8807 .metadata
8808 .rejection_reason
8809 .unwrap_or_else(|| answer.to_string())
8810 } else {
8811 output_data.content
8812 };
8813
8814 let reflection_metadata;
8816 (final_content, reflection_metadata) = self
8817 .run_reflection(&*llm, processed_input, final_content)
8818 .await?;
8819
8820 final_content =
8821 self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
8822
8823 let final_content = {
8827 let result = self
8828 .post_loop_processing(processed_input, final_content)
8829 .await?;
8830 self.apply_post_loop_result(processed_input, result).await?
8831 };
8832
8833 let reflected = reflection_metadata.is_some();
8834 let reasoning_mode_debug = format!("{:?}", reasoning_mode);
8835
8836 let response = self.build_agent_response(
8837 final_content,
8838 all_tool_calls,
8839 reasoning_mode,
8840 auto_detected,
8841 iterations,
8842 thinking_content,
8843 reflection_metadata,
8844 );
8845
8846 self.finish_turn_if_root(&response).await?;
8847
8848 let tool_call_count = response.tool_calls.as_ref().map(|tc| tc.len()).unwrap_or(0);
8849 info!(
8850 tool_calls = tool_call_count,
8851 response_len = response.content.len(),
8852 reasoning_mode = %reasoning_mode_debug,
8853 reflected = reflected,
8854 "Chat completed"
8855 );
8856 return Ok(response);
8857 }
8858 }
8859
8860 async fn generate_buffered_streaming_draft(
8861 &self,
8862 processed_input: &str,
8863 routing_resolved: Arc<AtomicBool>,
8864 ) -> Result<StreamingDraftResult> {
8865 let llm = self.get_state_llm()?;
8866 let messages = self.build_messages_for_draft(processed_input).await?;
8867 let mut stream = self
8868 .observe_purpose(
8869 ObservationPurpose::MainResponse,
8870 llm.complete_stream(&messages, None),
8871 )
8872 .await
8873 .map_err(|e| AgentError::LLM(e.to_string()))?;
8874 let mut buffer = crate::optimization::StreamBranchBuffer::new(self.streaming.buffer_size)?;
8875 let mut chunks = Vec::new();
8876 let mut accumulated = String::new();
8877 while let Some(chunk_result) = stream.next().await {
8878 let chunk = chunk_result.map_err(|e| AgentError::LLM(e.to_string()))?;
8879 accumulated.push_str(&chunk.delta);
8880 let stream_chunk = StreamChunk::content(chunk.delta);
8881 if routing_resolved.load(Ordering::SeqCst) {
8882 chunks.push(stream_chunk);
8883 } else {
8884 buffer.push(stream_chunk)?;
8885 }
8886 }
8887 chunks.splice(0..0, buffer.drain());
8888 let content = accumulated.trim().to_string();
8889 let draft = if let Some(calls) = self.parse_tool_calls(&content) {
8890 MainResponseDraft::ToolCalls {
8891 raw_content: content,
8892 calls,
8893 thinking: None,
8894 }
8895 } else {
8896 MainResponseDraft::Text {
8897 raw_content: content,
8898 thinking: None,
8899 }
8900 };
8901 Ok(StreamingDraftResult::new(draft, chunks))
8902 }
8903
8904 async fn try_buffered_streaming_branches(
8905 &self,
8906 processed_input: &str,
8907 input_context: &HashMap<String, Value>,
8908 ) -> Result<Option<(AgentResponse, Vec<StreamChunk>)>> {
8909 let optimization = &self.runtime_config.optimization;
8910 if !optimization.enabled {
8911 return Ok(None);
8912 }
8913 let transition_enabled =
8914 optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
8915 if !transition_enabled {
8916 return Ok(None);
8917 }
8918 let mut branch_scheduler =
8919 TurnBranchScheduler::new(optimization.max_parallel_runtime_tasks)?;
8920 if !branch_scheduler.reserve_task() {
8921 return Ok(None);
8922 }
8923 if !self
8924 .reserve_active_speculative_llm_call(RuntimeOptimizationKind::BufferedStreamingRouting)
8925 {
8926 branch_scheduler.release_task();
8927 return Ok(None);
8928 }
8929 if !branch_scheduler.reserve_task() {
8930 branch_scheduler.release_task();
8931 return Ok(None);
8932 }
8933 let mut main_branch = RuntimeBranch::new(
8934 RuntimeTaskPurpose::MainResponse,
8935 RuntimeOptimizationKind::BufferedStreamingRouting,
8936 RuntimeTaskPriority::Normal,
8937 RuntimeCommitBehavior::FinalResponse,
8938 );
8939 let mut transition_branch = RuntimeBranch::new(
8940 RuntimeTaskPurpose::StateTransition,
8941 RuntimeOptimizationKind::ParallelStateTransition,
8942 RuntimeTaskPriority::Critical,
8943 RuntimeCommitBehavior::TransitionDecision,
8944 );
8945 let main_id = main_branch.branch_id();
8946 let transition_id = transition_branch.branch_id();
8947 let routing_resolved = Arc::new(AtomicBool::new(false));
8948 let mut main_future =
8949 Box::pin(crate::optimization::observability::with_branch_observation(
8950 &main_id,
8951 RuntimeOptimizationKind::BufferedStreamingRouting,
8952 RuntimeCommitBehavior::FinalResponse,
8953 self.generate_buffered_streaming_draft(
8954 processed_input,
8955 Arc::clone(&routing_resolved),
8956 ),
8957 ));
8958 let mut transition_future =
8959 Box::pin(crate::optimization::observability::with_branch_observation(
8960 &transition_id,
8961 RuntimeOptimizationKind::ParallelStateTransition,
8962 RuntimeCommitBehavior::TransitionDecision,
8963 self.select_parallel_transition_candidate(processed_input),
8964 ));
8965 let mut main_pending = true;
8966 let mut transition_pending = true;
8967 let mut main_result: Option<Result<StreamingDraftResult>> = None;
8968 let mut transition_finalized = false;
8969 let mut transition_candidate: Option<TransitionCandidate> = None;
8970 loop {
8971 if let Some(candidate) = transition_candidate.take() {
8972 if self
8973 .commit_pre_response_transition_candidate(
8974 &candidate,
8975 &HashMap::new(),
8976 processed_input,
8977 )
8978 .await?
8979 {
8980 self.finalize_optional_branch(
8981 &transition_id,
8982 RuntimeOptimizationKind::ParallelStateTransition,
8983 RuntimeCommitBehavior::TransitionDecision,
8984 "committed",
8985 true,
8986 );
8987 self.finalize_branch_loss(
8988 &main_id,
8989 RuntimeOptimizationKind::BufferedStreamingRouting,
8990 RuntimeCommitBehavior::FinalResponse,
8991 main_pending,
8992 main_result.as_ref().map(|result| result.is_err()),
8993 );
8994 let response = self.redispatch_current_state(processed_input).await?;
8995 return Ok(Some((
8996 response.clone(),
8997 vec![StreamChunk::content(response.content)],
8998 )));
8999 }
9000 self.finalize_optional_branch(
9001 &transition_id,
9002 RuntimeOptimizationKind::ParallelStateTransition,
9003 RuntimeCommitBehavior::TransitionDecision,
9004 "discarded",
9005 false,
9006 );
9007 routing_resolved.store(true, Ordering::SeqCst);
9008 transition_finalized = true;
9009 }
9010 if transition_finalized {
9011 if let Some(result) = main_result.take() {
9012 let stream_draft = match result {
9013 Ok(stream_draft) => stream_draft,
9014 Err(error) => {
9015 self.finalize_optional_branch(
9016 &main_id,
9017 RuntimeOptimizationKind::BufferedStreamingRouting,
9018 RuntimeCommitBehavior::FinalResponse,
9019 "failed",
9020 false,
9021 );
9022 return Err(error);
9023 }
9024 };
9025 let raw_draft_content = stream_draft.draft.raw_content().to_string();
9026 let buffered_chunks = stream_draft.chunks;
9027 self.finalize_optional_branch(
9028 &main_id,
9029 RuntimeOptimizationKind::BufferedStreamingRouting,
9030 RuntimeCommitBehavior::FinalResponse,
9031 "committed",
9032 true,
9033 );
9034 let response = self
9035 .commit_main_response_draft(
9036 processed_input,
9037 input_context,
9038 stream_draft.draft,
9039 ReasoningMode::None,
9040 false,
9041 )
9042 .await?;
9043 let chunks = if response.content == raw_draft_content {
9044 buffered_chunks
9045 } else {
9046 vec![StreamChunk::content(response.content.clone())]
9047 };
9048 return Ok(Some((response, chunks)));
9049 }
9050 }
9051 tokio::select! {
9052 result = &mut main_future, if main_pending => {
9053 main_pending = false;
9054 main_branch.transition_to(RuntimeBranchStatus::Completed)?;
9055 main_result = Some(result);
9056 }
9057 result = &mut transition_future, if transition_pending => {
9058 transition_pending = false;
9059 transition_branch.transition_to(RuntimeBranchStatus::Completed)?;
9060 match result {
9061 Ok(ParallelTransitionSelection::Candidate(candidate)) => {
9062 transition_candidate = Some(candidate)
9063 }
9064 Ok(ParallelTransitionSelection::NoMatch) => {
9065 self.finalize_optional_branch(
9066 &transition_id,
9067 RuntimeOptimizationKind::ParallelStateTransition,
9068 RuntimeCommitBehavior::TransitionDecision,
9069 "discarded",
9070 false,
9071 );
9072 routing_resolved.store(true, Ordering::SeqCst);
9073 transition_finalized = true;
9074 }
9075 Ok(ParallelTransitionSelection::ReservationExhausted) => {
9076 self.finalize_optional_branch(
9077 &transition_id,
9078 RuntimeOptimizationKind::ParallelStateTransition,
9079 RuntimeCommitBehavior::TransitionDecision,
9080 "cancelled",
9081 false,
9082 );
9083 routing_resolved.store(true, Ordering::SeqCst);
9084 self.finalize_branch_loss(
9085 &main_id,
9086 RuntimeOptimizationKind::BufferedStreamingRouting,
9087 RuntimeCommitBehavior::FinalResponse,
9088 main_pending,
9089 main_result.as_ref().map(|result| result.is_err()),
9090 );
9091 return Ok(None);
9092 }
9093 Err(_) => {
9094 self.finalize_optional_branch(
9095 &transition_id,
9096 RuntimeOptimizationKind::ParallelStateTransition,
9097 RuntimeCommitBehavior::TransitionDecision,
9098 "failed",
9099 false,
9100 );
9101 routing_resolved.store(true, Ordering::SeqCst);
9102 transition_finalized = true;
9103 }
9104 }
9105 }
9106 }
9107 }
9108 }
9109
9110 fn run_loop_internal_stream<'a>(
9114 &'a self,
9115 input: &'a str,
9116 ) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
9117 let include_tool_events = self.streaming.include_tool_events;
9118 let include_state_events = self.streaming.include_state_events;
9119
9120 Box::pin(async_stream::stream! {
9121 self.begin_root_turn();
9122 self.pre_turn_session_lifecycle().await;
9124
9125 let input_data = match self.process_input(input).await {
9126 Ok(data) => data,
9127 Err(e) => {
9128 yield StreamChunk::error(e.to_string());
9129 return;
9130 }
9131 };
9132 self.update_active_turn_context(&input_data.content, input_data.context.clone());
9133
9134 for (key, value) in &input_data.context {
9136 let _ = self.context_manager.set(key, value.clone());
9137 }
9138
9139 if input_data.metadata.rejected {
9140 let reason = input_data
9141 .metadata
9142 .rejection_reason
9143 .unwrap_or_else(|| "Input rejected".to_string());
9144 warn!(reason = %reason, "Input rejected (stream)");
9145 yield StreamChunk::error(reason);
9146 return;
9147 }
9148
9149 let processed_input = &input_data.content;
9150
9151 if self.runtime_config.optimization.enabled
9152 && matches!(
9153 self.runtime_config.optimization.streaming_policy,
9154 crate::optimization::StreamingOptimizationPolicy::BufferUntilRoutingDone
9155 )
9156 {
9157 match Box::pin(self.try_buffered_streaming_branches(processed_input, &input_data.context)).await {
9162 Ok(Some((_response, chunks))) => {
9163 for chunk in chunks {
9164 yield chunk;
9165 }
9166 yield StreamChunk::Done {};
9167 return;
9168 }
9169 Ok(None) => {}
9170 Err(e) => {
9171 yield StreamChunk::error(e.to_string());
9172 return;
9173 }
9174 }
9175 }
9176
9177 if self.runtime_config.optimization.enabled
9178 && matches!(
9179 self.runtime_config.optimization.streaming_policy,
9180 crate::optimization::StreamingOptimizationPolicy::PreflightOnly
9181 )
9182 {
9183 match self.try_pre_response_transition(processed_input).await {
9184 Ok(Some(response)) => {
9185 yield StreamChunk::content(&response.content);
9186 yield StreamChunk::Done {};
9187 return;
9188 }
9189 Ok(None) => {}
9190 Err(e) => {
9191 yield StreamChunk::error(e.to_string());
9192 return;
9193 }
9194 }
9195 }
9196
9197 if let Some(ref sm) = self.state_machine {
9199 if let Some(def) = sm.current_definition() {
9200 let orchestration_result = if let Some(ref delegate_id) = def.delegate {
9201 Some(self.handle_delegated_state(processed_input, delegate_id, &def).await)
9202 } else if let Some(ref concurrent_config) = def.concurrent {
9203 Some(self.handle_concurrent_state(processed_input, concurrent_config).await)
9204 } else if let Some(ref group_chat_config) = def.group_chat {
9205 Some(self.handle_group_chat_state(processed_input, group_chat_config).await)
9206 } else if let Some(ref pipeline_config) = def.pipeline {
9207 Some(self.handle_pipeline_state(processed_input, pipeline_config).await)
9208 } else if let Some(ref handoff_config) = def.handoff {
9209 Some(self.handle_handoff_state(processed_input, handoff_config).await)
9210 } else {
9211 None
9212 };
9213
9214 if let Some(result) = orchestration_result {
9215 match result {
9216 Ok(response) => {
9217 yield StreamChunk::content(&response.content);
9218 yield StreamChunk::Done {};
9219 }
9220 Err(e) => {
9221 yield StreamChunk::error(e.to_string());
9222 }
9223 }
9224 return;
9225 }
9226 }
9227 }
9228
9229 match self.try_skill_route(processed_input).await {
9231 Ok(SkillRouteResult::Response(skill_response)) => {
9232 if let Err(e) = self.commit_root_user_message(processed_input).await {
9233 yield StreamChunk::error(e.to_string());
9234 return;
9235 }
9236 match self.handle_skill_response(processed_input, skill_response, &input_data.context).await {
9237 Ok(resp) => {
9238 yield StreamChunk::content(&resp.content);
9239 yield StreamChunk::Done {};
9240 return;
9241 }
9242 Err(e) => {
9243 yield StreamChunk::error(e.to_string());
9244 return;
9245 }
9246 }
9247 }
9248 Ok(SkillRouteResult::NeedsClarification(response)) => {
9249 if let Err(e) = self.commit_root_user_message(processed_input).await {
9250 yield StreamChunk::error(e.to_string());
9251 return;
9252 }
9253 let _ = self.memory.add_message(ChatMessage::assistant(&response.content)).await;
9254 if let Err(e) = self.finish_turn_if_root(&response).await {
9255 yield StreamChunk::error(e.to_string());
9256 return;
9257 }
9258 yield StreamChunk::content(&response.content);
9259 yield StreamChunk::Done {};
9260 return;
9261 }
9262 Ok(SkillRouteResult::NoMatch) => {} Err(e) => {
9264 yield StreamChunk::error(e.to_string());
9265 return;
9266 }
9267 }
9268
9269 let effective_reasoning = self.get_effective_reasoning_config();
9271 let reasoning_mode = match self.determine_reasoning_mode(processed_input).await {
9272 Ok(mode) => mode,
9273 Err(e) => {
9274 yield StreamChunk::error(e.to_string());
9275 return;
9276 }
9277 };
9278 let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
9279
9280 info!(
9281 reasoning_mode = ?reasoning_mode,
9282 auto_detected = auto_detected,
9283 "Reasoning mode determined (stream)"
9284 );
9285
9286 if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
9288 if let Err(e) = self.commit_root_user_message(processed_input).await {
9289 yield StreamChunk::error(e.to_string());
9290 return;
9291 }
9292 match self.handle_plan_and_execute(processed_input, &input_data.context, auto_detected).await {
9293 Ok(resp) => {
9294 yield StreamChunk::content(&resp.content);
9295 yield StreamChunk::Done {};
9296 return;
9297 }
9298 Err(e) => {
9299 yield StreamChunk::error(e.to_string());
9300 return;
9301 }
9302 }
9303 }
9304
9305 if let Err(e) = self.commit_root_user_message(processed_input).await {
9306 yield StreamChunk::error(e.to_string());
9307 return;
9308 }
9309
9310 let llm = match self.get_state_llm() {
9311 Ok(llm) => llm,
9312 Err(e) => {
9313 yield StreamChunk::error(e.to_string());
9314 return;
9315 }
9316 };
9317
9318 let mut iterations = 0u32;
9319 let mut all_tool_calls: Vec<ToolCall> = Vec::new();
9320 let mut thinking_content: Option<String> = None;
9321
9322 loop {
9323 let effective_max = if reasoning_mode != ReasoningMode::None {
9325 let rc = self.get_effective_reasoning_config();
9326 self.max_iterations.min(rc.max_iterations)
9327 } else {
9328 self.max_iterations
9329 };
9330
9331 if iterations >= effective_max {
9332 let err_msg = format!("Max iterations ({}) exceeded", effective_max);
9333 let err = AgentError::Other(err_msg.clone());
9334 self.hooks.on_error(&err).await;
9335 error!(iterations = iterations, "Max iterations exceeded (stream)");
9336 yield StreamChunk::error(err_msg);
9337 return;
9338 }
9339 iterations += 1;
9340 *self.iteration_count.write() = iterations;
9341
9342 debug!(iteration = iterations, max = effective_max, "LLM call (stream)");
9343
9344 let mut messages = match self.build_messages().await {
9345 Ok(m) => m,
9346 Err(e) => {
9347 yield StreamChunk::error(e.to_string());
9348 return;
9349 }
9350 };
9351 self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
9352
9353 self.hooks.on_llm_start(&messages).await;
9354 let llm_start = Instant::now();
9355
9356 let reflection_active = match self.should_reflect(processed_input, "").await {
9359 Ok(v) => v,
9360 Err(_) => false,
9361 };
9362
9363 let content = if reflection_active {
9364 let response = match self
9366 .observe_purpose(
9367 ObservationPurpose::MainResponse,
9368 llm.complete(&messages, None),
9369 )
9370 .await
9371 {
9372 Ok(r) => r,
9373 Err(e) => {
9374 yield StreamChunk::error(e.to_string());
9375 return;
9376 }
9377 };
9378 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
9379 self.hooks.on_llm_complete(&response, llm_duration_ms).await;
9380 response.content.trim().to_string()
9381 } else {
9382 let llm_stream = match self
9384 .observe_purpose(
9385 ObservationPurpose::MainResponse,
9386 llm.complete_stream(&messages, None),
9387 )
9388 .await
9389 {
9390 Ok(s) => s,
9391 Err(e) => {
9392 yield StreamChunk::error(e.to_string());
9393 return;
9394 }
9395 };
9396 let mut accumulated = String::new();
9397 let mut stream_inner = llm_stream;
9398 while let Some(chunk_result) = stream_inner.next().await {
9399 match chunk_result {
9400 Ok(chunk) => {
9401 accumulated.push_str(&chunk.delta);
9402 yield StreamChunk::content(chunk.delta);
9403 }
9404 Err(e) => {
9405 yield StreamChunk::error(e.to_string());
9406 return;
9407 }
9408 }
9409 }
9410 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
9411 let llm_response = ai_agents_core::LLMResponse::new(
9413 accumulated.trim(),
9414 ai_agents_core::FinishReason::Stop,
9415 );
9416 self.hooks.on_llm_complete(&llm_response, llm_duration_ms).await;
9417 accumulated.trim().to_string()
9418 };
9419
9420 if let Some(tool_calls) = self.parse_tool_calls(&content) {
9422 let transition_fired = match self.evaluate_transitions(processed_input, &content).await {
9425 Ok(v) => v,
9426 Err(e) => {
9427 yield StreamChunk::error(e.to_string());
9428 return;
9429 }
9430 };
9431 if transition_fired {
9432 let _ = self.memory.add_message(ChatMessage::assistant(
9433 "(Transitioned to new state — tool call handled by workflow)",
9434 )).await;
9435
9436 if include_state_events {
9437 if let Some(state) = self.current_state() {
9438 yield StreamChunk::state_transition(None, state);
9439 }
9440 }
9441 continue;
9442 }
9443
9444 let _ = self.memory.add_message(ChatMessage::assistant(&content)).await;
9446
9447 let results = self.execute_tools_parallel(&tool_calls).await;
9449
9450 for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
9451 if include_tool_events {
9452 yield StreamChunk::tool_start(&tool_call.id, &tool_call.name);
9453 }
9454
9455 match result {
9456 Ok(output) => {
9457 if include_tool_events {
9458 yield StreamChunk::tool_result(
9459 &tool_call.id,
9460 &tool_call.name,
9461 &output,
9462 true,
9463 );
9464 }
9465 let _ = self.memory
9466 .add_message(ChatMessage::function(&tool_call.name, &output))
9467 .await;
9468 }
9469 Err(e) => {
9470 if matches!(e, AgentError::HITLRejected(_)) {
9471 let _ = self.memory.add_message(ChatMessage::assistant(
9472 &format!("The operation was rejected by the approver: {}", e),
9473 )).await;
9474 let response = AgentResponse {
9475 content: format!("Operation cancelled: {}", e),
9476 metadata: None,
9477 tool_calls: Some(all_tool_calls.clone()),
9478 };
9479 if let Err(finalize_error) = self.finish_turn_if_root(&response).await {
9480 yield StreamChunk::error(finalize_error.to_string());
9481 return;
9482 }
9483 yield StreamChunk::error(response.content);
9484 yield StreamChunk::Done {};
9485 return;
9486 }
9487 if include_tool_events {
9488 yield StreamChunk::tool_result(
9489 &tool_call.id,
9490 &tool_call.name,
9491 &e.to_string(),
9492 false,
9493 );
9494 }
9495 let _ = self.memory
9496 .add_message(ChatMessage::function(
9497 &tool_call.name,
9498 &format!("Error: {}", e),
9499 ))
9500 .await;
9501 }
9502 }
9503 all_tool_calls.push(tool_call.clone());
9504
9505 if include_tool_events {
9506 yield StreamChunk::tool_end(&tool_call.id);
9507 }
9508 }
9509 continue;
9510 }
9511
9512 let (extracted_thinking, answer) = self.extract_thinking(&content);
9514 if extracted_thinking.is_some() {
9515 thinking_content = extracted_thinking;
9516 }
9517
9518 let output_data = match self.process_output(&answer, &input_data.context).await {
9519 Ok(d) => d,
9520 Err(e) => {
9521 yield StreamChunk::error(e.to_string());
9522 return;
9523 }
9524 };
9525
9526 let final_content = if output_data.metadata.rejected {
9527 output_data
9528 .metadata
9529 .rejection_reason
9530 .unwrap_or_else(|| answer.to_string())
9531 } else {
9532 output_data.content
9533 };
9534
9535 let (final_content, _reflection_metadata) = match self
9537 .run_reflection(&*llm, processed_input, final_content)
9538 .await
9539 {
9540 Ok(r) => r,
9541 Err(e) => {
9542 yield StreamChunk::error(e.to_string());
9543 return;
9544 }
9545 };
9546
9547 let final_content = self.format_response_with_thinking(
9548 thinking_content.as_deref(),
9549 &final_content,
9550 );
9551
9552 if reflection_active {
9554 yield StreamChunk::content(&final_content);
9555 }
9556
9557 let post_result = match self
9561 .post_loop_processing(processed_input, final_content)
9562 .await
9563 {
9564 Ok(r) => r,
9565 Err(e) => {
9566 yield StreamChunk::error(e.to_string());
9567 return;
9568 }
9569 };
9570
9571 let (final_content, transitioned) = match post_result {
9572 PostLoopResult::NoTransition(content) => (content, false),
9573 PostLoopResult::Transitioned(content) => (content, true),
9574 PostLoopResult::NeedsRedispatch => {
9575 const MAX_REDISPATCH_DEPTH: u32 = 3;
9576 let current_depth = *self.redispatch_depth.read();
9577 let content = if current_depth >= MAX_REDISPATCH_DEPTH {
9578 warn!(
9579 depth = current_depth,
9580 "Post-transition re-dispatch depth limit reached (stream)"
9581 );
9582 let c = String::new();
9583 let _ = self.memory.add_message(ChatMessage::assistant(&c)).await;
9584 c
9585 } else {
9586 *self.redispatch_depth.write() += 1;
9587 if let Some(context) = self.active_turn_context.write().as_mut() {
9588 context.enter_redispatch();
9589 }
9590 info!(
9591 depth = current_depth + 1,
9592 "Re-dispatching for new state after transition (stream)"
9593 );
9594 let result = self.run_loop_internal(processed_input).await;
9595 *self.redispatch_depth.write() -= 1;
9596 if let Some(context) = self.active_turn_context.write().as_mut() {
9597 context.exit_redispatch();
9598 }
9599 match result {
9600 Ok(resp) => resp.content,
9601 Err(e) => {
9602 yield StreamChunk::error(e.to_string());
9603 return;
9604 }
9605 }
9606 };
9607 (content, true)
9608 }
9609 };
9610
9611 if transitioned {
9612 if include_state_events {
9613 if let Some(state) = self.current_state() {
9614 yield StreamChunk::state_transition(None, state);
9615 }
9616 }
9617 yield StreamChunk::content(&final_content);
9619 }
9620
9621 let final_response = AgentResponse::new(&final_content);
9623 if let Err(e) = self.finish_turn_if_root(&final_response).await {
9624 yield StreamChunk::error(e.to_string());
9625 return;
9626 }
9627
9628 yield StreamChunk::Done {};
9629 return;
9630 }
9631 })
9632 }
9633
9634 fn run_loop_stream<'a>(
9637 &'a self,
9638 input: &'a str,
9639 ) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
9640 Box::pin(async_stream::stream! {
9641 self.begin_root_turn();
9642 let _root_cleanup = RootTurnCleanup::new(self);
9643 self.hooks.on_message_received(input).await;
9644
9645 if !self.context_initialized.swap(true, Ordering::SeqCst) {
9647 if let Err(e) = self.context_manager.initialize().await {
9648 yield StreamChunk::error(e.to_string());
9649 return;
9650 }
9651 debug!("Context manager initialized (defaults, env, builtins)");
9652 }
9653
9654 if let Err(e) = self.check_turn_timeout().await {
9655 yield StreamChunk::error(e.to_string());
9656 return;
9657 }
9658 if let Err(e) = self.context_manager.refresh_per_turn().await {
9659 yield StreamChunk::error(e.to_string());
9660 return;
9661 }
9662
9663 self.clear_disambiguation_context();
9665
9666 if let Some(ref disambiguator) = self.disambiguation_manager {
9668 let disambiguation_context = match self.build_disambiguation_context().await {
9669 Ok(ctx) => ctx,
9670 Err(e) => {
9671 yield StreamChunk::error(e.to_string());
9672 return;
9673 }
9674 };
9675
9676 let state_override = self
9677 .state_machine
9678 .as_ref()
9679 .and_then(|sm| sm.current_definition())
9680 .and_then(|def| def.disambiguation.clone());
9681
9682 let result = match self
9683 .observe_purpose(
9684 ObservationPurpose::DisambiguationDetection,
9685 disambiguator.process_input_with_override(
9686 input,
9687 &disambiguation_context,
9688 state_override.as_ref(),
9689 None,
9690 ),
9691 )
9692 .await
9693 {
9694 Ok(r) => r,
9695 Err(e) => {
9696 yield StreamChunk::error(e.to_string());
9697 return;
9698 }
9699 };
9700
9701 match result {
9702 DisambiguationResult::Clear => {
9703 debug!("Input is clear, proceeding normally (stream)");
9704 }
9705 DisambiguationResult::NeedsClarification {
9706 question,
9707 detection,
9708 } => {
9709 info!(
9710 ambiguity_type = ?detection.ambiguity_type,
9711 confidence = detection.confidence,
9712 "Input requires clarification (stream)"
9713 );
9714 if let Err(e) = self.commit_root_user_message(input).await {
9715 yield StreamChunk::error(e.to_string());
9716 return;
9717 }
9718 let _ = self
9719 .memory
9720 .add_message(ChatMessage::assistant(&question.question))
9721 .await;
9722 let response = AgentResponse::new(&question.question);
9723 if let Err(e) = self.finish_turn_if_root(&response).await {
9724 yield StreamChunk::error(e.to_string());
9725 return;
9726 }
9727 yield StreamChunk::content(&question.question);
9728 yield StreamChunk::Done {};
9729 return;
9730 }
9731 DisambiguationResult::Clarified {
9732 enriched_input,
9733 resolved,
9734 ..
9735 } => {
9736 info!(
9737 resolved_count = resolved.len(),
9738 enriched = %enriched_input,
9739 "Input clarified (stream)"
9740 );
9741 for (key, value) in &resolved {
9742 let context_key = format!("disambiguation.{}", key);
9743 let _ = self.context_manager.set(&context_key, value.clone());
9744 }
9745 if let Some(intent) = resolved.get("intent") {
9746 let _ = self.context_manager.set("resolved_intent", intent.clone());
9747 }
9748 let _ = self
9749 .context_manager
9750 .set("disambiguation.resolved", serde_json::Value::Bool(true));
9751
9752 let skill_id = self.pending_skill_id.read().clone();
9756 if let Some(skill_id) = skill_id {
9757 info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input (stream)");
9758 match self.recheck_skill_disambiguation(&skill_id, &enriched_input).await {
9759 Ok(resp) => {
9760 yield StreamChunk::content(&resp.content);
9761 yield StreamChunk::Done {};
9762 return;
9763 }
9764 Err(e) => {
9765 yield StreamChunk::error(e.to_string());
9766 return;
9767 }
9768 }
9769 }
9770
9771 let mut inner = self.run_loop_internal_stream(&enriched_input);
9773 while let Some(chunk) = inner.next().await {
9774 yield chunk;
9775 }
9776 return;
9777 }
9778 DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
9779 info!("Proceeding with best guess (stream)");
9780
9781 let skill_id = self.pending_skill_id.read().clone();
9783 if let Some(skill_id) = skill_id {
9784 info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input (stream)");
9785 match self.recheck_skill_disambiguation(&skill_id, &enriched_input).await {
9786 Ok(resp) => {
9787 yield StreamChunk::content(&resp.content);
9788 yield StreamChunk::Done {};
9789 return;
9790 }
9791 Err(e) => {
9792 yield StreamChunk::error(e.to_string());
9793 return;
9794 }
9795 }
9796 }
9797
9798 let mut inner = self.run_loop_internal_stream(&enriched_input);
9799 while let Some(chunk) = inner.next().await {
9800 yield chunk;
9801 }
9802 return;
9803 }
9804 DisambiguationResult::GiveUp { reason } => {
9805 *self.pending_skill_id.write() = None;
9806 warn!(reason = %reason, "Disambiguation gave up (stream)");
9807 let apology = self
9808 .generate_localized_apology(
9809 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
9810 &reason,
9811 )
9812 .await
9813 .unwrap_or_else(|_| {
9814 format!("I'm sorry, I couldn't understand your request: {}", reason)
9815 });
9816 let response = AgentResponse::new(&apology);
9817 if let Err(e) = self.finish_turn_if_root(&response).await {
9818 yield StreamChunk::error(e.to_string());
9819 return;
9820 }
9821 yield StreamChunk::content(&apology);
9822 yield StreamChunk::Done {};
9823 return;
9824 }
9825 DisambiguationResult::Escalate { reason } => {
9826 *self.pending_skill_id.write() = None;
9827 info!(reason = %reason, "Escalating to human (stream)");
9828 if let Some(ref hitl) = self.hitl_engine {
9829 let trigger =
9830 ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
9831 let mut context_map = HashMap::new();
9832 context_map.insert("original_input".to_string(), serde_json::json!(input));
9833 context_map.insert("reason".to_string(), serde_json::json!(&reason));
9834 let check_result = HITLCheckResult::required(
9835 trigger,
9836 context_map,
9837 format!("User request needs human assistance: {}", reason),
9838 Some(hitl.config().default_timeout_seconds),
9839 );
9840 match self.request_hitl_approval(check_result).await {
9841 Ok(result) if matches!(result, ApprovalResult::Approved | ApprovalResult::Modified { .. }) => {
9842 let mut inner = self.run_loop_internal_stream(input);
9843 while let Some(chunk) = inner.next().await {
9844 yield chunk;
9845 }
9846 return;
9847 }
9848 Ok(_) => {}
9849 Err(e) => {
9850 yield StreamChunk::error(e.to_string());
9851 return;
9852 }
9853 }
9854 }
9855 let apology = self
9856 .generate_localized_apology(
9857 "Explain briefly that you're transferring the user to a human agent for help.",
9858 &reason,
9859 )
9860 .await
9861 .unwrap_or_else(|_| {
9862 format!("I need human assistance to help with your request: {}", reason)
9863 });
9864 let response = AgentResponse::new(&apology);
9865 if let Err(e) = self.finish_turn_if_root(&response).await {
9866 yield StreamChunk::error(e.to_string());
9867 return;
9868 }
9869 yield StreamChunk::content(&apology);
9870 yield StreamChunk::Done {};
9871 return;
9872 }
9873 DisambiguationResult::Abandoned { new_input } => {
9874 *self.pending_skill_id.write() = None;
9875
9876 info!(
9877 has_new_input = new_input.is_some(),
9878 "Clarification abandoned by user (stream)"
9879 );
9880
9881 if let Err(e) = self.commit_root_user_message(input).await {
9882 yield StreamChunk::error(e.to_string());
9883 return;
9884 }
9885
9886 match new_input {
9887 Some(fresh_input) => {
9888 let mut inner = self.run_loop_internal_stream(&fresh_input);
9890 while let Some(chunk) = inner.next().await {
9891 yield chunk;
9892 }
9893 return;
9894 }
9895 None => {
9896 let ack = self
9898 .generate_localized_apology(
9899 "The user changed their mind about their previous request. \
9900 Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
9901 Do NOT apologize excessively. Be concise.",
9902 "User abandoned clarification",
9903 )
9904 .await
9905 .unwrap_or_else(|_| {
9906 "OK, no problem. What else can I help with?".to_string()
9907 });
9908
9909 let _ = self
9910 .memory
9911 .add_message(ChatMessage::assistant(&ack))
9912 .await;
9913
9914 let response = AgentResponse::new(&ack);
9915 if let Err(e) = self.finish_turn_if_root(&response).await {
9916 yield StreamChunk::error(e.to_string());
9917 return;
9918 }
9919 yield StreamChunk::content(&ack);
9920 yield StreamChunk::Done {};
9921 return;
9922 }
9923 }
9924 }
9925 }
9926 }
9927
9928 let mut inner = self.run_loop_internal_stream(input);
9930 while let Some(chunk) = inner.next().await {
9931 yield chunk;
9932 }
9933 })
9934 }
9935
9936 pub fn info(&self) -> AgentInfo {
9937 self.info.clone()
9938 }
9939
9940 pub fn skills(&self) -> &[SkillDefinition] {
9941 &self.skills
9942 }
9943
9944 pub async fn reset(&self) -> Result<()> {
9945 self.memory.clear().await?;
9946 *self.iteration_count.write() = 0;
9947 self.tool_call_history.write().clear();
9948 *self.pending_skill_id.write() = None;
9949 if let Some(ref sm) = self.state_machine {
9950 sm.reset();
9951 }
9952 Ok(())
9953 }
9954
9955 pub fn max_context_tokens(&self) -> u32 {
9956 self.max_context_tokens
9957 }
9958
9959 pub fn llm_registry(&self) -> &Arc<LLMRegistry> {
9960 &self.llm_registry
9961 }
9962
9963 pub fn state_machine(&self) -> Option<&Arc<StateMachine>> {
9964 self.state_machine.as_ref()
9965 }
9966
9967 pub fn context_manager(&self) -> &Arc<ContextManager> {
9968 &self.context_manager
9969 }
9970
9971 pub fn tool_call_history(&self) -> Vec<ToolCallRecord> {
9972 self.tool_call_history.read().clone()
9973 }
9974
9975 pub fn memory_token_budget(&self) -> Option<&MemoryTokenBudget> {
9976 self.memory_token_budget.as_ref()
9977 }
9978
9979 pub fn parallel_tools_config(&self) -> &ParallelToolsConfig {
9980 &self.parallel_tools
9981 }
9982
9983 pub fn streaming_config(&self) -> &StreamingConfig {
9984 &self.streaming
9985 }
9986
9987 pub fn hooks(&self) -> &Arc<dyn AgentHooks> {
9988 &self.hooks
9989 }
9990
9991 pub fn hitl_engine(&self) -> Option<&HITLEngine> {
9992 self.hitl_engine.as_ref()
9993 }
9994
9995 pub fn approval_handler(&self) -> &Arc<dyn ApprovalHandler> {
9996 &self.approval_handler
9997 }
9998
9999 fn build_hitl_language_context(&self) -> HashMap<String, Value> {
10001 let mut ctx = HashMap::new();
10002 for key in &["user.language", "input.detected.language", "language"] {
10003 if let Some(val) = self.context_manager.get(key) {
10004 ctx.insert(key.to_string(), val);
10005 }
10006 }
10007 ctx
10008 }
10009
10010 async fn request_hitl_approval(&self, check_result: HITLCheckResult) -> Result<ApprovalResult> {
10012 let Some(request) = check_result.into_request() else {
10013 return Ok(ApprovalResult::Approved);
10014 };
10015
10016 self.hooks.on_approval_requested(&request).await;
10017
10018 let request_id = request.id.clone();
10019 let timeout = request.timeout;
10020
10021 let result = if let Some(duration) = timeout {
10022 match tokio::time::timeout(duration, self.approval_handler.request_approval(request))
10023 .await
10024 {
10025 Ok(result) => result,
10026 Err(_) => ApprovalResult::timeout(),
10027 }
10028 } else {
10029 self.approval_handler.request_approval(request).await
10030 };
10031
10032 self.hooks.on_approval_result(&request_id, &result).await;
10033
10034 let result = match result {
10036 ApprovalResult::Timeout => {
10037 if let Some(ref engine) = self.hitl_engine {
10038 match engine.config().on_timeout {
10039 TimeoutAction::Approve => ApprovalResult::Approved,
10040 TimeoutAction::Reject => ApprovalResult::Rejected {
10041 reason: Some("Timeout".to_string()),
10042 },
10043 TimeoutAction::Error => {
10044 return Err(AgentError::Other("HITL approval timeout".to_string()));
10045 }
10046 }
10047 } else {
10048 ApprovalResult::Rejected {
10049 reason: Some("Timeout (no engine)".to_string()),
10050 }
10051 }
10052 }
10053 other => other,
10054 };
10055
10056 Ok(result)
10057 }
10058
10059 pub async fn check_state_hitl(&self, from: Option<&str>, to: &str) -> Result<bool> {
10060 if let Some(ref hitl_engine) = self.hitl_engine {
10061 let hitl_lang_ctx = self.build_hitl_language_context();
10062 let check_result = self
10063 .observe_purpose(
10064 ObservationPurpose::HitlLocalization,
10065 hitl_engine.check_state_transition_with_localization(
10066 from,
10067 to,
10068 &hitl_lang_ctx,
10069 self.approval_handler.as_ref(),
10070 Some(&self.llm_registry),
10071 ),
10072 )
10073 .await?;
10074 if check_result.is_required() {
10075 let result = self.request_hitl_approval(check_result).await?;
10076 return Ok(matches!(
10077 result,
10078 ApprovalResult::Approved | ApprovalResult::Modified { .. }
10079 ));
10080 }
10081 }
10082 Ok(true)
10083 }
10084
10085 async fn execute_tools_parallel(
10087 &self,
10088 tool_calls: &[ToolCall],
10089 ) -> Vec<(String, Result<String>)> {
10090 let can_run_parallel = tool_calls.iter().all(|tc| {
10091 self.tools
10092 .resolve(&tc.name)
10093 .map(|resolved| resolved.tool.classify_call(&tc.arguments).concurrency_safe)
10094 .unwrap_or(false)
10095 });
10096
10097 if !self.parallel_tools.enabled || tool_calls.len() <= 1 || !can_run_parallel {
10098 let mut results = Vec::new();
10099 for tc in tool_calls {
10100 let result = self
10101 .observe_purpose(
10102 current_observation_context()
10103 .map(|context| context.purpose)
10104 .unwrap_or_default(),
10105 self.execute_tool_smart(tc),
10106 )
10107 .await;
10108 results.push((tc.id.clone(), result));
10109 }
10110 return results;
10111 }
10112
10113 let chunks: Vec<_> = tool_calls
10114 .chunks(self.parallel_tools.max_parallel)
10115 .collect();
10116
10117 let mut all_results = Vec::new();
10118
10119 for chunk in chunks {
10120 let futures: Vec<_> = chunk
10121 .iter()
10122 .map(|tc| {
10123 let tc = tc.clone();
10124 async move {
10125 let result = self.execute_tool_smart(&tc).await;
10126 (tc.id.clone(), result)
10127 }
10128 })
10129 .collect();
10130
10131 let results = futures::future::join_all(futures).await;
10132 all_results.extend(results);
10133 }
10134
10135 all_results
10136 }
10137
10138 pub async fn chat_stream<'a>(
10140 &'a self,
10141 input: &'a str,
10142 ) -> Result<Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>>> {
10143 info!(input_len = input.len(), "Starting streaming chat");
10144 let inner = self.run_loop_stream(input);
10145 if let Some(context) = self.build_observation_context(None) {
10146 let stream: Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> =
10147 Box::pin(async_stream::stream! {
10148 let mut inner = inner;
10149 loop {
10150 let next = with_observation_context(context.clone(), inner.next()).await;
10151 match next {
10152 Some(chunk) => yield chunk,
10153 None => break,
10154 }
10155 }
10156 self.export_observability_if_configured().await;
10157 });
10158 Ok(stream)
10159 } else {
10160 Ok(inner)
10161 }
10162 }
10163}
10164
10165#[async_trait]
10166impl ToolInvoker for RuntimeAgent {
10167 async fn invoke_tool(&self, request: ToolExecutionRequest) -> Result<ToolExecutionRecord> {
10168 self.execute_tool_record(request).await
10169 }
10170}
10171
10172#[async_trait]
10173impl Agent for RuntimeAgent {
10174 async fn chat(&self, input: &str) -> Result<AgentResponse> {
10175 let result = if let Some(context) = self.build_observation_context(None) {
10176 with_observation_context(context, self.run_loop(input)).await
10177 } else {
10178 self.run_loop(input).await
10179 };
10180 self.export_observability_if_configured().await;
10181 result
10182 }
10183
10184 fn info(&self) -> AgentInfo {
10185 self.info.clone()
10186 }
10187
10188 async fn reset(&self) -> Result<()> {
10189 self.memory.clear().await?;
10190 *self.iteration_count.write() = 0;
10191 self.tool_call_history.write().clear();
10192 if let Some(ref sm) = self.state_machine {
10193 sm.reset();
10194 }
10195 Ok(())
10196 }
10197}
10198
10199fn background_maintenance_tags(
10209 label: &str,
10210 stage: &str,
10211 reason: Option<&str>,
10212 policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
10213) -> HashMap<String, String> {
10214 let mut tags = HashMap::new();
10215 tags.insert("runtime.background".to_string(), "true".to_string());
10216 tags.insert("runtime.maintenance".to_string(), label.to_string());
10217 tags.insert("runtime.maintenance_stage".to_string(), stage.to_string());
10218 if let Some(policy) = policy {
10219 tags.insert(
10220 "runtime.await_before_next_turn".to_string(),
10221 await_before_next_turn_label(policy.await_before_next_turn).to_string(),
10222 );
10223 tags.insert(
10224 "runtime.maintenance_mode".to_string(),
10225 maintenance_mode_label(policy.mode).to_string(),
10226 );
10227 }
10228 if let Some(reason) = reason {
10229 tags.insert("runtime.reason".to_string(), reason.to_string());
10230 }
10231 tags
10232}
10233
10234fn await_before_next_turn_label(policy: AwaitBeforeNextTurn) -> &'static str {
10235 match policy {
10236 AwaitBeforeNextTurn::Never => "never",
10237 AwaitBeforeNextTurn::SameActor => "same_actor",
10238 AwaitBeforeNextTurn::Always => "always",
10239 }
10240}
10241
10242fn maintenance_mode_label(mode: MaintenanceMode) -> &'static str {
10243 match mode {
10244 MaintenanceMode::InlineSerial => "inline_serial",
10245 MaintenanceMode::InlineParallel => "inline_parallel",
10246 MaintenanceMode::Background => "background",
10247 }
10248}
10249
10250fn record_background_maintenance_event(
10252 manager: Option<&Arc<ObservabilityManager>>,
10253 label: &str,
10254 status: EventStatus,
10255 duration_ms: u64,
10256 stage: &str,
10257 reason: Option<String>,
10258 policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
10259) {
10260 if let Some(manager) = manager {
10261 manager.record_lifecycle_event(
10262 EventType::MemoryOperation {
10263 operation: format!("{}_background_{}", label, stage),
10264 },
10265 ObservationPurpose::Other(format!("{}_maintenance", label)),
10266 status,
10267 duration_ms,
10268 background_maintenance_tags(label, stage, reason.as_deref(), policy),
10269 None,
10270 );
10271 }
10272}
10273
10274fn effective_maintenance_mode(mode: MaintenanceMode, force_parallel: bool) -> MaintenanceMode {
10275 if force_parallel && matches!(mode, MaintenanceMode::InlineSerial) {
10276 MaintenanceMode::InlineParallel
10277 } else {
10278 mode
10279 }
10280}
10281
10282fn observation_purpose_for_process(hint: ProcessPurposeHint) -> ObservationPurpose {
10283 match hint {
10284 ProcessPurposeHint::Detect => ObservationPurpose::ProcessDetect,
10285 ProcessPurposeHint::Extract => ObservationPurpose::ProcessExtract,
10286 ProcessPurposeHint::Validate => ObservationPurpose::ProcessValidate,
10287 ProcessPurposeHint::Transform | ProcessPurposeHint::Other => {
10288 ObservationPurpose::ProcessTransform
10289 }
10290 }
10291}
10292
10293fn new_tool_resource_locks() -> ToolResourceLocks {
10294 Arc::new(RwLock::new(HashMap::new()))
10295}
10296
10297fn tool_resource_lock_key(canonical_id: &str, args: &Value) -> String {
10298 let resource = ["path", "base_path", "cwd", "url"]
10299 .into_iter()
10300 .find_map(|field| args.get(field).and_then(Value::as_str))
10301 .map(normalized_resource_key);
10302 match resource {
10303 Some(value) => format!("resource:{}", value),
10304 None => format!("tool:{}", canonical_id),
10305 }
10306}
10307
10308fn normalized_resource_key(value: &str) -> String {
10309 if value.starts_with("http://") || value.starts_with("https://") {
10310 return value.to_ascii_lowercase();
10311 }
10312 let path = std::path::Path::new(value);
10313 let path = if path.is_absolute() {
10314 path.components().collect::<std::path::PathBuf>()
10315 } else {
10316 std::env::current_dir()
10317 .unwrap_or_else(|_| std::path::PathBuf::from("."))
10318 .join(path)
10319 .components()
10320 .collect()
10321 };
10322 path.to_string_lossy().to_string()
10323}
10324
10325fn render_concurrent_template(
10326 template: &str,
10327 user_input: &str,
10328 context_values: &std::collections::HashMap<String, serde_json::Value>,
10329) -> Result<String> {
10330 let mut env = minijinja::Environment::new();
10331 env.add_template("concurrent", template)
10332 .map_err(|e| AgentError::Other(format!("Concurrent template parse error: {}", e)))?;
10333
10334 let mut ctx = std::collections::BTreeMap::new();
10335 ctx.insert("user_input".to_string(), minijinja::Value::from(user_input));
10336
10337 let context_obj = minijinja::Value::from_serialize(context_values);
10339 ctx.insert("context".to_string(), context_obj);
10340
10341 let tmpl = env
10342 .get_template("concurrent")
10343 .map_err(|e| AgentError::Other(format!("Concurrent template error: {}", e)))?;
10344
10345 tmpl.render(minijinja::Value::from_serialize(&ctx))
10346 .map_err(|e| AgentError::Other(format!("Concurrent template render error: {}", e)))
10347}
10348
10349#[cfg(test)]
10350mod tests {
10351 use super::*;
10352 use crate::AgentBuilder;
10353 use ai_agents_llm::mock::MockLLMProvider;
10354
10355 fn mock_with_response(response: &str) -> MockLLMProvider {
10356 let mut mock = MockLLMProvider::new("test");
10357 mock.set_response(response);
10358 mock
10359 }
10360
10361 fn mock_with_responses(responses: Vec<&str>) -> MockLLMProvider {
10362 let mut mock = MockLLMProvider::new("test");
10363 mock.set_responses(responses.into_iter().map(String::from).collect(), true);
10364 mock
10365 }
10366
10367 struct ResponseCountingHooks {
10369 responses: Arc<std::sync::atomic::AtomicUsize>,
10370 }
10371
10372 struct ContextEchoTool;
10374
10375 #[async_trait]
10376 impl ai_agents_core::Tool for ContextEchoTool {
10377 fn id(&self) -> &str {
10378 "context_echo"
10379 }
10380
10381 fn name(&self) -> &str {
10382 "Context Echo"
10383 }
10384
10385 fn description(&self) -> &str {
10386 "Returns selected execution context fields."
10387 }
10388
10389 fn input_schema(&self) -> Value {
10390 serde_json::json!({"type": "object"})
10391 }
10392
10393 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
10394 ai_agents_core::ToolPolicyBindings {
10395 path_fields: vec![ai_agents_core::PathPolicyBinding::read("path")],
10396 result_limit_fields: vec![ai_agents_core::ResultLimitBinding::new(
10397 "max_results",
10398 ai_agents_core::ResultLimitKind::MaxResults,
10399 )],
10400 ..Default::default()
10401 }
10402 }
10403
10404 async fn execute(
10405 &self,
10406 _args: Value,
10407 ctx: ai_agents_core::ToolExecutionContext,
10408 ) -> ToolResult {
10409 ToolResult::ok(
10410 serde_json::json!({
10411 "requested_name": ctx.requested_name,
10412 "canonical_id": ctx.canonical_id,
10413 "display_name": ctx.display_name,
10414 "max_results": ctx.limits.max_results,
10415 "custom_config": ctx.custom_config,
10416 })
10417 .to_string(),
10418 )
10419 }
10420 }
10421
10422 struct SlowTool;
10424
10425 struct FlakyWriteTool {
10427 calls: Arc<std::sync::atomic::AtomicUsize>,
10428 }
10429
10430 struct LockedWriteTool {
10432 active: Arc<std::sync::atomic::AtomicUsize>,
10433 max_active: Arc<std::sync::atomic::AtomicUsize>,
10434 }
10435
10436 #[async_trait]
10437 impl ai_agents_core::Tool for SlowTool {
10438 fn id(&self) -> &str {
10439 "slow"
10440 }
10441
10442 fn name(&self) -> &str {
10443 "Slow"
10444 }
10445
10446 fn description(&self) -> &str {
10447 "Waits until cancelled or timed out."
10448 }
10449
10450 fn input_schema(&self) -> Value {
10451 serde_json::json!({"type": "object"})
10452 }
10453
10454 async fn execute(
10455 &self,
10456 _args: Value,
10457 _ctx: ai_agents_core::ToolExecutionContext,
10458 ) -> ToolResult {
10459 tokio::time::sleep(std::time::Duration::from_secs(5)).await;
10460 ToolResult::ok("done")
10461 }
10462 }
10463
10464 #[async_trait]
10465 impl ai_agents_core::Tool for FlakyWriteTool {
10466 fn id(&self) -> &str {
10467 "flaky_write"
10468 }
10469
10470 fn name(&self) -> &str {
10471 "Flaky Write"
10472 }
10473
10474 fn description(&self) -> &str {
10475 "Fails on the first write attempt."
10476 }
10477
10478 fn input_schema(&self) -> Value {
10479 serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
10480 }
10481
10482 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
10483 ai_agents_core::ToolPolicyBindings {
10484 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
10485 ..Default::default()
10486 }
10487 }
10488
10489 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
10490 ai_agents_core::ToolSafetyMetadata {
10491 read_only: false,
10492 concurrency_safe: false,
10493 operation: ai_agents_core::ToolOperationKind::Write,
10494 side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
10495 requires_network: false,
10496 destructive: false,
10497 open_world: false,
10498 host_dependent: false,
10499 requires_user_interaction: false,
10500 supports_cancellation: true,
10501 default_requires_approval: false,
10502 should_defer_schema: false,
10503 max_output_chars: Some(1024),
10504 max_result_size_chars: Some(1024),
10505 }
10506 }
10507
10508 fn classify_call(&self, _args: &Value) -> ai_agents_core::ToolCallClassification {
10509 let mut classification =
10510 ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
10511 classification.safely_retryable = false;
10512 classification
10513 }
10514
10515 async fn execute(
10516 &self,
10517 _args: Value,
10518 _ctx: ai_agents_core::ToolExecutionContext,
10519 ) -> ToolResult {
10520 let call = self.calls.fetch_add(1, Ordering::SeqCst);
10521 if call == 0 {
10522 ToolResult::error("first failure")
10523 } else {
10524 ToolResult::ok("second success")
10525 }
10526 }
10527 }
10528
10529 #[async_trait]
10530 impl ai_agents_core::Tool for LockedWriteTool {
10531 fn id(&self) -> &str {
10532 "locked_write"
10533 }
10534
10535 fn name(&self) -> &str {
10536 "Locked Write"
10537 }
10538
10539 fn description(&self) -> &str {
10540 "Tracks concurrent execution on one resource."
10541 }
10542
10543 fn input_schema(&self) -> Value {
10544 serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
10545 }
10546
10547 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
10548 ai_agents_core::ToolPolicyBindings {
10549 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
10550 ..Default::default()
10551 }
10552 }
10553
10554 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
10555 ai_agents_core::ToolSafetyMetadata {
10556 read_only: false,
10557 concurrency_safe: false,
10558 operation: ai_agents_core::ToolOperationKind::Write,
10559 side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
10560 requires_network: false,
10561 destructive: false,
10562 open_world: false,
10563 host_dependent: false,
10564 requires_user_interaction: false,
10565 supports_cancellation: true,
10566 default_requires_approval: false,
10567 should_defer_schema: false,
10568 max_output_chars: Some(1024),
10569 max_result_size_chars: Some(1024),
10570 }
10571 }
10572
10573 async fn execute(
10574 &self,
10575 _args: Value,
10576 _ctx: ai_agents_core::ToolExecutionContext,
10577 ) -> ToolResult {
10578 let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
10579 loop {
10580 let current_max = self.max_active.load(Ordering::SeqCst);
10581 if active <= current_max {
10582 break;
10583 }
10584 if self
10585 .max_active
10586 .compare_exchange(current_max, active, Ordering::SeqCst, Ordering::SeqCst)
10587 .is_ok()
10588 {
10589 break;
10590 }
10591 }
10592 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
10593 self.active.fetch_sub(1, Ordering::SeqCst);
10594 ToolResult::ok("done")
10595 }
10596 }
10597
10598 #[async_trait]
10599 impl AgentHooks for ResponseCountingHooks {
10600 async fn on_response(&self, _response: &AgentResponse) {
10601 self.responses.fetch_add(1, Ordering::SeqCst);
10602 }
10603 }
10604
10605 #[tokio::test]
10607 async fn test_integration_yaml_to_chat_basic() {
10608 let mock = mock_with_response("Hello! How can I help you?");
10609 let agent = AgentBuilder::new()
10610 .system_prompt("You are a test assistant.")
10611 .llm(Arc::new(mock))
10612 .build()
10613 .unwrap();
10614
10615 let response = agent.chat("Hi").await.unwrap();
10616 assert!(!response.content.is_empty());
10617 assert_eq!(response.content, "Hello! How can I help you?");
10618 }
10619
10620 #[tokio::test]
10622 async fn test_integration_multi_turn_conversation() {
10623 let mock = mock_with_responses(vec![
10624 "Hello! I'm your assistant.",
10625 "The weather is sunny today.",
10626 "Goodbye!",
10627 ]);
10628 let agent = AgentBuilder::new()
10629 .system_prompt("You are helpful.")
10630 .llm(Arc::new(mock))
10631 .build()
10632 .unwrap();
10633
10634 let r1 = agent.chat("Hi").await.unwrap();
10635 assert_eq!(r1.content, "Hello! I'm your assistant.");
10636
10637 let r2 = agent.chat("What's the weather?").await.unwrap();
10638 assert_eq!(r2.content, "The weather is sunny today.");
10639
10640 let r3 = agent.chat("Bye").await.unwrap();
10641 assert_eq!(r3.content, "Goodbye!");
10642
10643 let messages = agent.memory.get_messages(None).await.unwrap();
10645 assert_eq!(messages.len(), 6);
10647 }
10648
10649 #[tokio::test]
10650 async fn context_preserves_requested_and_canonical_identity() {
10651 let mock = mock_with_response("hello");
10652 let mut tools = ai_agents_tools::ToolRegistry::new();
10653 tools.register(Arc::new(ContextEchoTool)).unwrap();
10654
10655 let mut security = ToolSecurityConfig::default();
10656 security.enabled = true;
10657 security.fail_closed = true;
10658 let mut policy = ai_agents_tools::ToolPolicyConfig::default();
10659 policy.read_paths = vec![".".to_string()];
10660 policy.max_results = Some(7);
10661 policy
10662 .config
10663 .insert("backend".to_string(), serde_json::json!("memory"));
10664 security.tools.insert("context_echo".to_string(), policy);
10665
10666 let agent = AgentBuilder::new()
10667 .system_prompt("You are helpful.")
10668 .llm(Arc::new(mock))
10669 .tools(tools)
10670 .tool_security(ToolSecurityEngine::new(security))
10671 .build()
10672 .unwrap();
10673
10674 let record = agent
10675 .invoke_tool(ToolExecutionRequest::new(
10676 "ctx-call",
10677 "Context Echo",
10678 serde_json::json!({"path": ".", "max_results": 99}),
10679 ToolCallSource::Manual,
10680 ))
10681 .await
10682 .unwrap();
10683
10684 assert!(record.success);
10685 assert_eq!(record.requested_name, "Context Echo");
10686 assert_eq!(record.canonical_id, "context_echo");
10687 assert_eq!(record.policy.outcome, PermissionOutcome::Allow);
10688 assert_eq!(record.executed_arguments["max_results"], 7);
10689 let output: Value = serde_json::from_str(&record.output).unwrap();
10690 assert_eq!(output["requested_name"], "Context Echo");
10691 assert_eq!(output["canonical_id"], "context_echo");
10692 assert_eq!(output["max_results"], 7);
10693 assert_eq!(output["custom_config"]["backend"], "memory");
10694 assert!(record.metadata.contains_key("effective_limits"));
10695 assert!(record.metadata.contains_key("policy_snapshot"));
10696 }
10697
10698 #[tokio::test]
10699 async fn test_runtime_control_cancels_active_tool_call() {
10700 let mock = mock_with_response("hello");
10701 let agent = Arc::new(
10702 AgentBuilder::new()
10703 .system_prompt("You are helpful.")
10704 .llm(Arc::new(mock))
10705 .tool(Arc::new(SlowTool))
10706 .build()
10707 .unwrap(),
10708 );
10709 let control = agent.runtime_control();
10710 let running_agent = Arc::clone(&agent);
10711 let handle = tokio::spawn(async move {
10712 running_agent
10713 .invoke_tool(ToolExecutionRequest::new(
10714 "slow-call",
10715 "slow",
10716 serde_json::json!({}),
10717 ToolCallSource::Manual,
10718 ))
10719 .await
10720 .unwrap()
10721 });
10722
10723 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
10724 control.cancel_all();
10725 let record = handle.await.unwrap();
10726
10727 assert!(record.executed);
10728 assert!(record.cancelled);
10729 assert!(!record.success);
10730 assert!(record.cancellation_reason.is_some());
10731 }
10732
10733 #[tokio::test]
10734 async fn non_idempotent_tool_calls_are_not_retried() {
10735 use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
10736
10737 let mock = mock_with_response("hello");
10738 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
10739 let agent = AgentBuilder::new()
10740 .system_prompt("You are helpful.")
10741 .llm(Arc::new(mock))
10742 .tool(Arc::new(FlakyWriteTool {
10743 calls: Arc::clone(&calls),
10744 }))
10745 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
10746 tools: ToolRecoveryConfig {
10747 default: ToolRetryConfig {
10748 max_retries: 2,
10749 ..Default::default()
10750 },
10751 ..Default::default()
10752 },
10753 ..Default::default()
10754 }))
10755 .build()
10756 .unwrap();
10757
10758 let record = agent
10759 .invoke_tool(ToolExecutionRequest::new(
10760 "flaky-call",
10761 "flaky_write",
10762 serde_json::json!({"path": "./tmp.txt"}),
10763 ToolCallSource::Manual,
10764 ))
10765 .await
10766 .unwrap();
10767
10768 assert!(!record.success);
10769 assert_eq!(calls.load(Ordering::SeqCst), 1);
10770 }
10771
10772 #[tokio::test]
10773 async fn side_effecting_tools_are_serialized_per_resource() {
10774 let mock = mock_with_response("hello");
10775 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
10776 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
10777 let agent = Arc::new(
10778 AgentBuilder::new()
10779 .system_prompt("You are helpful.")
10780 .llm(Arc::new(mock))
10781 .tool(Arc::new(LockedWriteTool {
10782 active: Arc::clone(&active),
10783 max_active: Arc::clone(&max_active),
10784 }))
10785 .build()
10786 .unwrap(),
10787 );
10788
10789 let left = {
10790 let agent = Arc::clone(&agent);
10791 tokio::spawn(async move {
10792 agent
10793 .invoke_tool(ToolExecutionRequest::new(
10794 "lock-1",
10795 "locked_write",
10796 serde_json::json!({"path": "./same.txt"}),
10797 ToolCallSource::Manual,
10798 ))
10799 .await
10800 .unwrap()
10801 })
10802 };
10803 let right = {
10804 let agent = Arc::clone(&agent);
10805 tokio::spawn(async move {
10806 agent
10807 .invoke_tool(ToolExecutionRequest::new(
10808 "lock-2",
10809 "locked_write",
10810 serde_json::json!({"path": "./same.txt"}),
10811 ToolCallSource::Manual,
10812 ))
10813 .await
10814 .unwrap()
10815 })
10816 };
10817
10818 let left = left.await.unwrap();
10819 let right = right.await.unwrap();
10820 assert!(left.success);
10821 assert!(right.success);
10822 assert_eq!(max_active.load(Ordering::SeqCst), 1);
10823 }
10824
10825 #[tokio::test]
10826 async fn diagnostics_without_provider_records_unavailable_without_execution() {
10827 let mock = mock_with_response("hello");
10828 let yaml = r#"
10829name: DiagnosticsNoProviderAgent
10830system_prompt: "Review diagnostics."
10831tools: [diagnostics]
10832"#;
10833 let agent = AgentBuilder::from_yaml(yaml)
10834 .unwrap()
10835 .llm(Arc::new(mock))
10836 .auto_configure_features()
10837 .unwrap()
10838 .build()
10839 .unwrap();
10840
10841 let record = agent
10842 .invoke_tool(ToolExecutionRequest::new(
10843 "diagnostics-call",
10844 "diagnostics",
10845 serde_json::json!({}),
10846 ToolCallSource::Manual,
10847 ))
10848 .await
10849 .unwrap();
10850
10851 assert!(!record.executed);
10852 assert!(!record.success);
10853 assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
10854 }
10855
10856 #[tokio::test]
10857 async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_omitted() {
10858 let mock = mock_with_response("hello");
10859 let yaml = r#"
10860name: SpawnerNoGrantAgent
10861system_prompt: "You manage agents."
10862spawner:
10863 max_agents: 2
10864"#;
10865 let agent = AgentBuilder::from_yaml(yaml)
10866 .unwrap()
10867 .llm(Arc::new(mock))
10868 .auto_configure_features()
10869 .unwrap()
10870 .auto_configure_spawner()
10871 .await
10872 .unwrap()
10873 .build()
10874 .unwrap();
10875
10876 let available = agent.get_available_tool_ids().await.unwrap();
10877 assert!(available.is_empty());
10878 }
10879
10880 #[tokio::test]
10881 async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_empty() {
10882 let mock = mock_with_response("hello");
10883 let yaml = r#"
10884name: EmptySpawnerNoGrantAgent
10885system_prompt: "You manage agents."
10886tools: []
10887spawner:
10888 max_agents: 2
10889"#;
10890 let agent = AgentBuilder::from_yaml(yaml)
10891 .unwrap()
10892 .llm(Arc::new(mock))
10893 .auto_configure_features()
10894 .unwrap()
10895 .auto_configure_spawner()
10896 .await
10897 .unwrap()
10898 .build()
10899 .unwrap();
10900
10901 let available = agent.get_available_tool_ids().await.unwrap();
10902 assert!(available.is_empty());
10903 }
10904
10905 #[tokio::test]
10906 async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_empty() {
10907 let mock = mock_with_response("hello");
10908 let yaml = r#"
10909name: ManagementGrantAgent
10910system_prompt: "You manage agents."
10911tools: []
10912spawner:
10913 management_tools: true
10914"#;
10915 let agent = AgentBuilder::from_yaml(yaml)
10916 .unwrap()
10917 .llm(Arc::new(mock))
10918 .auto_configure_features()
10919 .unwrap()
10920 .auto_configure_spawner()
10921 .await
10922 .unwrap()
10923 .build()
10924 .unwrap();
10925
10926 let available = agent.get_available_tool_ids().await.unwrap();
10927 assert_eq!(available.len(), 4);
10928 assert!(available.contains(&"spawn_agent".to_string()));
10929 assert!(available.contains(&"send_agent_message".to_string()));
10930 assert!(available.contains(&"list_agents".to_string()));
10931 assert!(available.contains(&"remove_agent".to_string()));
10932 }
10933
10934 #[tokio::test]
10935 async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_omitted() {
10936 let mock = mock_with_response("hello");
10937 let yaml = r#"
10938name: ManagementOmittedToolsGrantAgent
10939system_prompt: "You manage agents."
10940spawner:
10941 management_tools: true
10942"#;
10943 let agent = AgentBuilder::from_yaml(yaml)
10944 .unwrap()
10945 .llm(Arc::new(mock))
10946 .auto_configure_features()
10947 .unwrap()
10948 .auto_configure_spawner()
10949 .await
10950 .unwrap()
10951 .build()
10952 .unwrap();
10953
10954 let available = agent.get_available_tool_ids().await.unwrap();
10955 assert_eq!(available.len(), 4);
10956 assert!(available.contains(&"spawn_agent".to_string()));
10957 assert!(available.contains(&"send_agent_message".to_string()));
10958 assert!(available.contains(&"list_agents".to_string()));
10959 assert!(available.contains(&"remove_agent".to_string()));
10960 }
10961
10962 #[tokio::test]
10963 async fn test_management_tools_selected_grants_only_selected_tools() {
10964 let mock = mock_with_response("hello");
10965 let yaml = r#"
10966name: ManagementSelectedGrantAgent
10967system_prompt: "You manage agents."
10968tools: []
10969spawner:
10970 management_tools:
10971 - spawn_agent
10972 - send_agent_message
10973 - list_agents
10974"#;
10975 let agent = AgentBuilder::from_yaml(yaml)
10976 .unwrap()
10977 .llm(Arc::new(mock))
10978 .auto_configure_features()
10979 .unwrap()
10980 .auto_configure_spawner()
10981 .await
10982 .unwrap()
10983 .build()
10984 .unwrap();
10985
10986 let available = agent.get_available_tool_ids().await.unwrap();
10987 assert_eq!(available.len(), 3);
10988 assert!(available.contains(&"spawn_agent".to_string()));
10989 assert!(available.contains(&"send_agent_message".to_string()));
10990 assert!(available.contains(&"list_agents".to_string()));
10991 assert!(!available.contains(&"remove_agent".to_string()));
10992 }
10993
10994 #[tokio::test]
10995 async fn test_orchestration_tools_flag_grants_tools_when_top_level_tools_empty() {
10996 let mock = mock_with_response("hello");
10997 let yaml = r#"
10998name: OrchestrationGrantAgent
10999system_prompt: "You coordinate agents."
11000llms:
11001 default:
11002 provider: openai
11003 model: gpt-4
11004 router:
11005 provider: openai
11006 model: gpt-4
11007llm:
11008 default: default
11009 router: router
11010tools: []
11011spawner:
11012 orchestration_tools: true
11013"#;
11014 let agent = AgentBuilder::from_yaml(yaml)
11015 .unwrap()
11016 .llm(Arc::new(mock))
11017 .auto_configure_features()
11018 .unwrap()
11019 .auto_configure_spawner()
11020 .await
11021 .unwrap()
11022 .build()
11023 .unwrap();
11024
11025 let available = agent.get_available_tool_ids().await.unwrap();
11026 assert_eq!(available.len(), 5);
11027 assert!(available.contains(&"route_to_agent".to_string()));
11028 assert!(available.contains(&"pipeline_process".to_string()));
11029 assert!(available.contains(&"concurrent_ask".to_string()));
11030 assert!(available.contains(&"group_discussion".to_string()));
11031 assert!(available.contains(&"handoff_conversation".to_string()));
11032 }
11033
11034 #[tokio::test]
11035 async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_empty() {
11036 let mock = mock_with_response("hello");
11037 let yaml = r#"
11038name: PersonaGrantAgent
11039system_prompt: "You can evolve persona."
11040llm:
11041 provider: openai
11042 model: gpt-4
11043tools: []
11044persona:
11045 identity:
11046 name: "Guide"
11047 role: "Helper"
11048 evolution:
11049 enabled: true
11050 allow_llm_evolve: true
11051 mutable_fields:
11052 - traits.personality
11053"#;
11054 let agent = AgentBuilder::from_yaml(yaml)
11055 .unwrap()
11056 .llm(Arc::new(mock))
11057 .build()
11058 .unwrap();
11059
11060 let available = agent.get_available_tool_ids().await.unwrap();
11061 assert_eq!(available, vec!["persona_evolve".to_string()]);
11062 }
11063
11064 #[tokio::test]
11065 async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_omitted() {
11066 let mock = mock_with_response("hello");
11067 let yaml = r#"
11068name: PersonaOmittedToolsGrantAgent
11069system_prompt: "You can evolve persona."
11070llm:
11071 provider: openai
11072 model: gpt-4
11073persona:
11074 identity:
11075 name: "Guide"
11076 role: "Helper"
11077 evolution:
11078 enabled: true
11079 allow_llm_evolve: true
11080 mutable_fields:
11081 - traits.personality
11082"#;
11083 let agent = AgentBuilder::from_yaml(yaml)
11084 .unwrap()
11085 .llm(Arc::new(mock))
11086 .build()
11087 .unwrap();
11088
11089 let available = agent.get_available_tool_ids().await.unwrap();
11090 assert_eq!(available, vec!["persona_evolve".to_string()]);
11091 }
11092
11093 #[tokio::test]
11094 async fn test_omitted_yaml_tools_exposes_no_tools() {
11095 let mock = mock_with_response("hello");
11096 let yaml = r#"
11097name: NoToolsAgent
11098system_prompt: "You are helpful."
11099"#;
11100 let agent = AgentBuilder::from_yaml(yaml)
11101 .unwrap()
11102 .llm(Arc::new(mock))
11103 .auto_configure_features()
11104 .unwrap()
11105 .build()
11106 .unwrap();
11107
11108 let available = agent.get_available_tool_ids().await.unwrap();
11109 assert!(available.is_empty());
11110 }
11111
11112 #[tokio::test]
11113 async fn test_state_tools_cannot_widen_top_level_grant() {
11114 let mock = mock_with_response("hello");
11115 let yaml = r#"
11116name: NarrowToolsAgent
11117system_prompt: "You are helpful."
11118tools:
11119 - calculator
11120states:
11121 initial: current
11122 states:
11123 current:
11124 tools: [datetime]
11125"#;
11126 let agent = AgentBuilder::from_yaml(yaml)
11127 .unwrap()
11128 .llm(Arc::new(mock))
11129 .auto_configure_features()
11130 .unwrap()
11131 .build()
11132 .unwrap();
11133
11134 let available = agent.get_available_tool_ids().await.unwrap();
11135 assert!(available.is_empty());
11136 }
11137
11138 #[tokio::test]
11140 async fn test_integration_tool_execution() {
11141 let mock = mock_with_responses(vec![
11143 r#"I'll calculate that for you.
11145[TOOL_CALL: {"name": "calculator", "arguments": {"expression": "2+2"}}]"#,
11146 "The answer is 4.",
11148 ]);
11149 let mut tools = ai_agents_tools::ToolRegistry::new();
11150 tools
11151 .register(Arc::new(ai_agents_tools::CalculatorTool))
11152 .unwrap();
11153
11154 let agent = AgentBuilder::new()
11155 .system_prompt("You are a calculator assistant.")
11156 .llm(Arc::new(mock))
11157 .tools(tools)
11158 .build()
11159 .unwrap();
11160
11161 let response = agent.chat("What is 2+2?").await.unwrap();
11162 assert!(!response.content.is_empty());
11164 }
11165
11166 #[tokio::test]
11167 async fn test_tool_hitl_rejection_finalizes_blocking_turn() {
11168 let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
11169 let hooks = Arc::new(ResponseCountingHooks {
11170 responses: Arc::clone(&responses),
11171 });
11172 let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
11173 let yaml = r#"
11174name: ToolRejectAgent
11175system_prompt: "You use tools when requested."
11176tools:
11177 - echo
11178hitl:
11179 tools:
11180 echo:
11181 require_approval: true
11182 approval_message: "Approve echo?"
11183"#;
11184 let agent = AgentBuilder::from_yaml(yaml)
11185 .unwrap()
11186 .llm(Arc::new(mock))
11187 .auto_configure_features()
11188 .unwrap()
11189 .hooks(hooks)
11190 .build()
11191 .unwrap();
11192
11193 let response = agent.chat("echo hello").await.unwrap();
11194
11195 assert!(
11196 response.content.contains("Operation cancelled"),
11197 "unexpected response: {}",
11198 response.content
11199 );
11200 assert_eq!(responses.load(Ordering::SeqCst), 1);
11201 let messages = agent.memory.get_messages(None).await.unwrap();
11202 assert_eq!(messages.len(), 3);
11203 assert_eq!(messages[0].content, "echo hello");
11204 assert!(messages[1].content.contains("\"tool\":\"echo\""));
11205 assert!(messages[2].content.contains("rejected by the approver"));
11206 }
11207
11208 #[tokio::test]
11209 async fn test_tool_hitl_rejection_finalizes_streaming_turn() {
11210 use futures::StreamExt;
11211
11212 let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
11213 let hooks = Arc::new(ResponseCountingHooks {
11214 responses: Arc::clone(&responses),
11215 });
11216 let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
11217 let yaml = r#"
11218name: ToolRejectStreamingAgent
11219system_prompt: "You use tools when requested."
11220tools:
11221 - echo
11222streaming:
11223 enabled: true
11224hitl:
11225 tools:
11226 echo:
11227 require_approval: true
11228 approval_message: "Approve echo?"
11229"#;
11230 let agent = AgentBuilder::from_yaml(yaml)
11231 .unwrap()
11232 .llm(Arc::new(mock))
11233 .auto_configure_features()
11234 .unwrap()
11235 .hooks(hooks)
11236 .build()
11237 .unwrap();
11238
11239 let mut stream = agent.chat_stream("echo hello").await.unwrap();
11240 let mut terminal_error = String::new();
11241 let mut done = false;
11242 while let Some(chunk) = stream.next().await {
11243 match chunk {
11244 StreamChunk::Error { message } => terminal_error = message,
11245 StreamChunk::Done {} => {
11246 done = true;
11247 break;
11248 }
11249 _ => {}
11250 }
11251 }
11252
11253 assert!(done);
11254 assert!(
11255 terminal_error.contains("Operation cancelled"),
11256 "unexpected terminal error: {}",
11257 terminal_error
11258 );
11259 assert_eq!(responses.load(Ordering::SeqCst), 1);
11260 let messages = agent.memory.get_messages(None).await.unwrap();
11261 assert_eq!(messages.len(), 3);
11262 assert_eq!(messages[0].content, "echo hello");
11263 assert!(messages[1].content.contains("\"tool\":\"echo\""));
11264 assert!(messages[2].content.contains("rejected by the approver"));
11265 }
11266
11267 #[tokio::test]
11268 async fn test_pre_response_guard_transition_skips_old_state_llm() {
11269 let mock = mock_with_response("Billing state response");
11270 let call_counter = mock.clone();
11271 let yaml = r#"
11272name: OptimizedStateAgent
11273system_prompt: "You route before answering."
11274runtime:
11275 optimization:
11276 enabled: true
11277 pre_response_deterministic_transitions: true
11278states:
11279 initial: greeting
11280 states:
11281 greeting:
11282 prompt: "Old state prompt that should be skipped."
11283 transitions:
11284 - to: billing
11285 guard:
11286 context:
11287 topic:
11288 eq: billing
11289 timing: pre_response
11290 billing:
11291 prompt: "Answer from the billing state."
11292"#;
11293 let agent = AgentBuilder::from_yaml(yaml)
11294 .unwrap()
11295 .llm(Arc::new(mock))
11296 .build()
11297 .unwrap();
11298 agent
11299 .set_context("topic", serde_json::json!("billing"))
11300 .unwrap();
11301
11302 let response = agent.chat("I need billing help").await.unwrap();
11303
11304 assert_eq!(agent.current_state().as_deref(), Some("billing"));
11305 assert_eq!(response.content, "Billing state response");
11306 assert_eq!(call_counter.call_count(), 1);
11307 assert_eq!(agent.actor_facts().len(), 0);
11308 }
11309
11310 #[tokio::test]
11311 async fn test_set_context_supports_dotted_paths_for_pre_response_guards() {
11312 let mock = mock_with_response("Billing state response");
11313 let call_counter = mock.clone();
11314 let yaml = r#"
11315name: OptimizedStateAgent
11316system_prompt: "You route before answering."
11317runtime:
11318 optimization:
11319 enabled: true
11320 pre_response_deterministic_transitions: true
11321context:
11322 request:
11323 type: runtime
11324 default:
11325 topic: general
11326states:
11327 initial: greeting
11328 states:
11329 greeting:
11330 prompt: "Old state prompt that should be skipped."
11331 transitions:
11332 - to: billing
11333 guard:
11334 context:
11335 request.topic:
11336 eq: billing
11337 timing: pre_response
11338 billing:
11339 prompt: "Answer from the billing state."
11340"#;
11341 let agent = AgentBuilder::from_yaml(yaml)
11342 .unwrap()
11343 .llm(Arc::new(mock))
11344 .build()
11345 .unwrap();
11346 agent
11347 .set_context("request.topic", serde_json::json!("billing"))
11348 .unwrap();
11349
11350 let response = agent.chat("I need billing help").await.unwrap();
11351
11352 assert_eq!(agent.current_state().as_deref(), Some("billing"));
11353 assert_eq!(response.content, "Billing state response");
11354 assert_eq!(call_counter.call_count(), 1);
11355 assert_eq!(
11356 agent.get_context().get("request"),
11357 Some(&serde_json::json!({"topic": "billing"}))
11358 );
11359 }
11360
11361 #[tokio::test]
11362 async fn test_pre_response_rejection_does_not_commit_staged_context_or_user() {
11363 let mock = mock_with_response("billing");
11364 let yaml = r#"
11365name: OptimizedStateAgent
11366system_prompt: "You route before answering."
11367runtime:
11368 optimization:
11369 enabled: true
11370 pre_response_deterministic_transitions: true
11371hitl:
11372 states:
11373 billing:
11374 on_enter: require_approval
11375 approval_message: "Approve billing route?"
11376states:
11377 initial: greeting
11378 states:
11379 greeting:
11380 prompt: "Old state prompt."
11381 extract:
11382 - key: topic
11383 description: "Support topic"
11384 transitions:
11385 - to: billing
11386 guard:
11387 context:
11388 topic:
11389 eq: billing
11390 timing: pre_response
11391 run_extractors: true
11392 billing:
11393 prompt: "Billing state."
11394"#;
11395 let agent = AgentBuilder::from_yaml(yaml)
11396 .unwrap()
11397 .llm(Arc::new(mock))
11398 .build()
11399 .unwrap();
11400
11401 let response = agent
11402 .try_pre_response_transition("billing please")
11403 .await
11404 .unwrap();
11405
11406 assert!(response.is_none());
11407 assert_eq!(agent.current_state().as_deref(), Some("greeting"));
11408 assert!(!agent.get_context().contains_key("topic"));
11409 assert_eq!(agent.memory.get_messages(None).await.unwrap().len(), 0);
11410 }
11411
11412 #[tokio::test]
11413 async fn test_pre_response_extractor_commits_context_on_winning_path() {
11414 let mock = mock_with_responses(vec!["billing", "Billing response"]);
11415 let yaml = r#"
11416name: OptimizedStateAgent
11417system_prompt: "You route before answering."
11418runtime:
11419 optimization:
11420 enabled: true
11421 pre_response_deterministic_transitions: true
11422states:
11423 initial: greeting
11424 states:
11425 greeting:
11426 prompt: "Old state prompt."
11427 extract:
11428 - key: topic
11429 description: "Support topic"
11430 transitions:
11431 - to: billing
11432 guard:
11433 context:
11434 topic:
11435 eq: billing
11436 timing: pre_response
11437 run_extractors: true
11438 billing:
11439 prompt: "Billing state."
11440"#;
11441 let agent = AgentBuilder::from_yaml(yaml)
11442 .unwrap()
11443 .llm(Arc::new(mock))
11444 .build()
11445 .unwrap();
11446
11447 let response = agent.chat("billing please").await.unwrap();
11448
11449 assert_eq!(agent.current_state().as_deref(), Some("billing"));
11450 assert_eq!(response.content, "Billing response");
11451 assert_eq!(
11452 agent.get_context().get("topic"),
11453 Some(&serde_json::json!("billing"))
11454 );
11455 }
11456
11457 #[tokio::test]
11458 async fn test_pre_response_extractor_miss_does_not_mutate_context() {
11459 let mock = mock_with_response("__NONE__");
11460 let yaml = r#"
11461name: OptimizedStateAgent
11462system_prompt: "You route before answering."
11463runtime:
11464 optimization:
11465 enabled: true
11466 pre_response_deterministic_transitions: true
11467states:
11468 initial: greeting
11469 states:
11470 greeting:
11471 prompt: "Old state prompt."
11472 extract:
11473 - key: topic
11474 description: "Support topic"
11475 transitions:
11476 - to: billing
11477 guard:
11478 context:
11479 topic:
11480 eq: billing
11481 timing: pre_response
11482 run_extractors: true
11483 billing:
11484 prompt: "Billing state."
11485"#;
11486 let agent = AgentBuilder::from_yaml(yaml)
11487 .unwrap()
11488 .llm(Arc::new(mock))
11489 .build()
11490 .unwrap();
11491
11492 let response = agent.try_pre_response_transition("hello").await.unwrap();
11493
11494 assert!(response.is_none());
11495 assert_eq!(agent.current_state().as_deref(), Some("greeting"));
11496 assert!(!agent.get_context().contains_key("topic"));
11497 }
11498
11499 #[tokio::test]
11500 async fn test_default_guard_transition_stays_post_response() {
11501 let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
11502 let call_counter = mock.clone();
11503 let yaml = r#"
11504name: TimingAgent
11505system_prompt: "You route carefully."
11506runtime:
11507 optimization:
11508 enabled: true
11509 pre_response_deterministic_transitions: true
11510states:
11511 initial: greeting
11512 states:
11513 greeting:
11514 prompt: "Old state prompt."
11515 transitions:
11516 - to: billing
11517 guard:
11518 context:
11519 topic:
11520 eq: billing
11521 billing:
11522 prompt: "Billing state."
11523"#;
11524 let agent = AgentBuilder::from_yaml(yaml)
11525 .unwrap()
11526 .llm(Arc::new(mock))
11527 .build()
11528 .unwrap();
11529 agent
11530 .set_context("topic", serde_json::json!("billing"))
11531 .unwrap();
11532
11533 let response = agent.chat("billing please").await.unwrap();
11534
11535 assert_eq!(agent.current_state().as_deref(), Some("billing"));
11536 assert_eq!(response.content, "Billing response");
11537 assert_eq!(call_counter.call_count(), 2);
11538 }
11539
11540 #[tokio::test]
11541 async fn test_explicit_post_response_guard_transition_stays_post_response() {
11542 let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
11543 let call_counter = mock.clone();
11544 let yaml = r#"
11545name: TimingAgent
11546system_prompt: "You route carefully."
11547runtime:
11548 optimization:
11549 enabled: true
11550 pre_response_deterministic_transitions: true
11551states:
11552 initial: greeting
11553 states:
11554 greeting:
11555 prompt: "Old state prompt."
11556 transitions:
11557 - to: billing
11558 guard:
11559 context:
11560 topic:
11561 eq: billing
11562 timing: post_response
11563 billing:
11564 prompt: "Billing state."
11565"#;
11566 let agent = AgentBuilder::from_yaml(yaml)
11567 .unwrap()
11568 .llm(Arc::new(mock))
11569 .build()
11570 .unwrap();
11571 agent
11572 .set_context("topic", serde_json::json!("billing"))
11573 .unwrap();
11574
11575 let response = agent.chat("billing please").await.unwrap();
11576
11577 assert_eq!(agent.current_state().as_deref(), Some("billing"));
11578 assert_eq!(response.content, "Billing response");
11579 assert_eq!(call_counter.call_count(), 2);
11580 }
11581
11582 #[tokio::test]
11583 async fn test_pre_response_extractors_are_transition_scoped() {
11584 let mock = mock_with_responses(vec!["billing", "Billing response"]);
11585 let yaml = r#"
11586name: ScopedExtractorAgent
11587system_prompt: "You route carefully."
11588runtime:
11589 optimization:
11590 enabled: true
11591 pre_response_deterministic_transitions: true
11592states:
11593 initial: greeting
11594 states:
11595 greeting:
11596 prompt: "Old state prompt."
11597 extract:
11598 - key: topic
11599 description: "Support topic"
11600 transitions:
11601 - to: wrong
11602 guard:
11603 context:
11604 topic:
11605 eq: billing
11606 timing: pre_response
11607 - to: billing
11608 guard:
11609 context:
11610 topic:
11611 eq: billing
11612 timing: pre_response
11613 run_extractors: true
11614 wrong:
11615 prompt: "Wrong state."
11616 billing:
11617 prompt: "Billing state."
11618"#;
11619 let agent = AgentBuilder::from_yaml(yaml)
11620 .unwrap()
11621 .llm(Arc::new(mock))
11622 .build()
11623 .unwrap();
11624
11625 let response = agent.chat("billing please").await.unwrap();
11626
11627 assert_eq!(agent.current_state().as_deref(), Some("billing"));
11628 assert_eq!(response.content, "Billing response");
11629 }
11630
11631 #[tokio::test]
11632 async fn test_pre_response_resolved_intent_routes_early() {
11633 let mock = mock_with_response("Billing response");
11634 let yaml = r#"
11635name: IntentAgent
11636system_prompt: "You route carefully."
11637runtime:
11638 optimization:
11639 enabled: true
11640 pre_response_deterministic_transitions: true
11641states:
11642 initial: greeting
11643 states:
11644 greeting:
11645 prompt: "Old state prompt."
11646 transitions:
11647 - to: billing
11648 intent: billing
11649 timing: pre_response
11650 billing:
11651 prompt: "Billing state."
11652"#;
11653 let agent = AgentBuilder::from_yaml(yaml)
11654 .unwrap()
11655 .llm(Arc::new(mock))
11656 .build()
11657 .unwrap();
11658 agent
11659 .set_context("resolved_intent", serde_json::json!("billing"))
11660 .unwrap();
11661
11662 let response = agent
11663 .try_pre_response_transition("I need billing help")
11664 .await
11665 .unwrap()
11666 .unwrap();
11667
11668 assert_eq!(agent.current_state().as_deref(), Some("billing"));
11669 assert_eq!(response.content, "Billing response");
11670 }
11671
11672 #[tokio::test]
11673 async fn test_background_overflow_error_surfaces() {
11674 let mut config = RuntimeConfig::default();
11675 config.optimization.enabled = true;
11676 config.optimization.post_turn.max_background_tasks = 1;
11677 config.optimization.post_turn.on_background_overflow = BackgroundOverflowPolicy::Error;
11678 let policy = crate::optimization::MaintenanceTaskPolicy {
11679 mode: MaintenanceMode::Background,
11680 await_before_next_turn: AwaitBeforeNextTurn::Always,
11681 };
11682 let agent = AgentBuilder::new()
11683 .system_prompt("You are helpful.")
11684 .llm(Arc::new(mock_with_response("ok")))
11685 .build()
11686 .unwrap()
11687 .with_runtime_config(config);
11688 agent
11689 .background_maintenance
11690 .spawn(None, async { std::future::pending::<Result<()>>().await })
11691 .unwrap();
11692
11693 let result = agent
11694 .spawn_or_handle_background(None, async { Ok(()) }, "facts", &policy)
11695 .await;
11696
11697 assert!(result.is_err());
11698 }
11699
11700 #[tokio::test]
11701 async fn test_speculative_reasoning_low_cap_uses_serial_reasoning() {
11702 let default_mock = mock_with_response("Plain draft response");
11703 let router_mock = mock_with_response("cot");
11704 let router_counter = router_mock.clone();
11705 let yaml = r#"
11706name: ReasoningReservationAgent
11707system_prompt: "You answer plainly unless reasoning wins."
11708llm:
11709 default: default
11710 router: router
11711observability:
11712 enabled: true
11713 export:
11714 write_raw_events: true
11715reasoning:
11716 mode: auto
11717 judge_llm: router
11718runtime:
11719 optimization:
11720 enabled: true
11721 max_speculative_llm_calls_per_turn: 1
11722 speculative_reasoning_auto: true
11723 max_parallel_runtime_tasks: 2
11724"#;
11725 let agent = AgentBuilder::from_yaml(yaml)
11726 .unwrap()
11727 .llm_alias("default", Arc::new(default_mock))
11728 .llm_alias("router", Arc::new(router_mock))
11729 .build()
11730 .unwrap();
11731
11732 let response = agent.chat("hello").await.unwrap();
11733
11734 assert_eq!(response.content, "Plain draft response");
11735 assert_eq!(router_counter.call_count(), 1);
11736 let events = agent.observability().unwrap().raw_events();
11737 assert!(!events.iter().any(|event| {
11738 event.dimensions.get("commit_behavior") == Some(&"reasoning_decision".to_string())
11739 }));
11740 }
11741
11742 #[tokio::test]
11743 async fn test_forced_reasoning_skips_plain_speculative_draft() {
11744 let mock = mock_with_response("Reasoned response");
11745 let yaml = r#"
11746name: ForcedReasoningAgent
11747system_prompt: "You reason before answering."
11748observability:
11749 enabled: true
11750 export:
11751 write_raw_events: true
11752reasoning:
11753 mode: cot
11754runtime:
11755 optimization:
11756 enabled: true
11757 max_speculative_llm_calls_per_turn: 2
11758 speculative_state_transitions: true
11759 max_parallel_runtime_tasks: 2
11760states:
11761 initial: triage
11762 states:
11763 triage:
11764 prompt: "Answer from triage."
11765 transitions:
11766 - to: billing
11767 guard:
11768 context:
11769 route:
11770 eq: billing
11771 timing: parallel
11772 billing:
11773 prompt: "Billing state."
11774"#;
11775 let agent = AgentBuilder::from_yaml(yaml)
11776 .unwrap()
11777 .llm(Arc::new(mock))
11778 .build()
11779 .unwrap();
11780
11781 let response = agent.chat("hello").await.unwrap();
11782
11783 assert_eq!(response.content, "Reasoned response");
11784 let events = agent.observability().unwrap().raw_events();
11785 assert!(
11786 !events
11787 .iter()
11788 .any(|event| event.dimensions.contains_key("branch_status"))
11789 );
11790 }
11791
11792 #[tokio::test]
11793 async fn test_speculative_skill_low_cap_uses_serial_skill_route() {
11794 let default_mock = mock_with_response("Skill committed response");
11795 let router_mock = mock_with_response("helper");
11796 let router_counter = router_mock.clone();
11797 let yaml = r#"
11798name: SkillReservationAgent
11799system_prompt: "Use skills when they match."
11800llm:
11801 default: default
11802 router: router
11803observability:
11804 enabled: true
11805 export:
11806 write_raw_events: true
11807runtime:
11808 optimization:
11809 enabled: true
11810 max_speculative_llm_calls_per_turn: 1
11811 speculative_skill_routing: true
11812 max_parallel_runtime_tasks: 2
11813skills:
11814 - id: helper
11815 description: "Answer helper requests"
11816 trigger: "User asks for helper"
11817 steps:
11818 - prompt: "Answer the helper request: {{ user_input }}"
11819"#;
11820 let agent = AgentBuilder::from_yaml(yaml)
11821 .unwrap()
11822 .llm_alias("default", Arc::new(default_mock))
11823 .llm_alias("router", Arc::new(router_mock))
11824 .build()
11825 .unwrap();
11826
11827 let response = agent.chat("please use helper").await.unwrap();
11828
11829 assert_eq!(response.content, "Skill committed response");
11830 assert_eq!(router_counter.call_count(), 1);
11831 let events = agent.observability().unwrap().raw_events();
11832 assert!(
11833 !events
11834 .iter()
11835 .any(|event| event.dimensions.contains_key("branch_status"))
11836 );
11837 }
11838
11839 #[tokio::test]
11840 async fn test_parallel_transition_low_cap_allows_deterministic_route() {
11841 let mock = mock_with_response("unused");
11842 let call_counter = mock.clone();
11843 let yaml = r#"
11844name: ParallelTransitionLowCapAgent
11845system_prompt: "Route before stale responses when safe."
11846runtime:
11847 optimization:
11848 enabled: true
11849 max_speculative_llm_calls_per_turn: 1
11850 speculative_state_transitions: true
11851 max_parallel_runtime_tasks: 2
11852states:
11853 initial: triage
11854 states:
11855 triage:
11856 prompt: "Triage state."
11857 transitions:
11858 - to: billing
11859 guard:
11860 context:
11861 route:
11862 eq: billing
11863 timing: parallel
11864 billing:
11865 prompt: "Billing state."
11866"#;
11867 let agent = AgentBuilder::from_yaml(yaml)
11868 .unwrap()
11869 .llm(Arc::new(mock))
11870 .build()
11871 .unwrap();
11872 agent
11873 .set_context("route", serde_json::json!("billing"))
11874 .unwrap();
11875 agent.update_active_turn_context("billing help", HashMap::new());
11876 assert!(
11877 agent.reserve_active_speculative_llm_call(
11878 RuntimeOptimizationKind::ParallelStateTransition
11879 )
11880 );
11881
11882 let selection = agent
11883 .select_parallel_transition_candidate("billing help")
11884 .await
11885 .unwrap();
11886 agent.end_root_turn();
11887
11888 match selection {
11889 ParallelTransitionSelection::Candidate(candidate) => {
11890 assert_eq!(candidate.target(), "billing");
11891 }
11892 ParallelTransitionSelection::NoMatch => panic!("deterministic route did not match"),
11893 ParallelTransitionSelection::ReservationExhausted => {
11894 panic!("deterministic route consumed LLM budget")
11895 }
11896 }
11897 assert_eq!(call_counter.call_count(), 0);
11898 }
11899
11900 #[tokio::test]
11901 async fn test_buffered_streaming_transition_reservation_falls_back() {
11902 use futures::StreamExt;
11903
11904 let mock = mock_with_responses(vec![
11905 "Serial streaming response",
11906 "Serial streaming response",
11907 ]);
11908 let router_mock = mock_with_response("1");
11909 let router_counter = router_mock.clone();
11910 let yaml = r#"
11911name: BufferedReservationFallbackAgent
11912system_prompt: "Stream normally if speculative routing cannot be evaluated."
11913llm:
11914 default: default
11915 router: router
11916observability:
11917 enabled: true
11918 export:
11919 write_raw_events: true
11920streaming:
11921 enabled: true
11922 buffer_size: 8
11923runtime:
11924 optimization:
11925 enabled: true
11926 max_speculative_llm_calls_per_turn: 1
11927 speculative_state_transitions: true
11928 streaming_policy: buffer_until_routing_done
11929 max_parallel_runtime_tasks: 2
11930states:
11931 initial: triage
11932 states:
11933 triage:
11934 prompt: "Triage state."
11935 transitions:
11936 - to: billing
11937 guard:
11938 context:
11939 route:
11940 eq: billing
11941 when: "User asks about billing"
11942 timing: parallel
11943 billing:
11944 prompt: "Billing state."
11945"#;
11946 let agent = AgentBuilder::from_yaml(yaml)
11947 .unwrap()
11948 .llm_alias("default", Arc::new(mock))
11949 .llm_alias("router", Arc::new(router_mock))
11950 .build()
11951 .unwrap();
11952
11953 let mut stream = agent.chat_stream("hello").await.unwrap();
11954 let mut content = String::new();
11955 let mut error = None;
11956 while let Some(chunk) = stream.next().await {
11957 match chunk {
11958 StreamChunk::Content { text } => content.push_str(&text),
11959 StreamChunk::Error { message } => error = Some(message),
11960 StreamChunk::Done {} => break,
11961 _ => {}
11962 }
11963 }
11964
11965 assert_eq!(error, None);
11966 assert_eq!(content, "Serial streaming response");
11967 assert_eq!(router_counter.call_count(), 0);
11968 let events = agent.observability().unwrap().raw_events();
11969 assert!(events.iter().any(|event| {
11970 event.dimensions.get("branch_status") == Some(&"cancelled".to_string())
11971 && event.dimensions.get("commit_behavior")
11972 == Some(&"transition_decision".to_string())
11973 }));
11974 }
11975
11976 #[tokio::test]
11977 async fn test_blocking_error_cleanup_resets_root_turn_for_next_chat() {
11978 let mut mock = mock_with_response("Recovered response");
11979 mock.set_error("boom");
11980 let mut handle = mock.clone();
11981 let agent = AgentBuilder::new()
11982 .system_prompt("You are helpful.")
11983 .llm(Arc::new(mock))
11984 .build()
11985 .unwrap();
11986
11987 assert!(agent.chat("first").await.is_err());
11988 handle.clear_error();
11989 let response = agent.chat("second").await.unwrap();
11990
11991 assert_eq!(response.content, "Recovered response");
11992 let messages = agent.memory.get_messages(None).await.unwrap();
11993 let user_count = messages
11994 .iter()
11995 .filter(|message| message.role == ai_agents_core::Role::User)
11996 .count();
11997 assert_eq!(user_count, 2);
11998 }
11999
12000 #[tokio::test]
12001 async fn test_streaming_error_cleanup_resets_root_turn_for_next_chat() {
12002 use futures::StreamExt;
12003
12004 let mut mock = mock_with_response("Recovered response");
12005 mock.set_error("stream boom");
12006 let mut handle = mock.clone();
12007 let agent = AgentBuilder::new()
12008 .system_prompt("You are helpful.")
12009 .llm(Arc::new(mock))
12010 .build()
12011 .unwrap();
12012
12013 let mut stream = agent.chat_stream("first").await.unwrap();
12014 let mut saw_error = false;
12015 while let Some(chunk) = stream.next().await {
12016 if matches!(chunk, StreamChunk::Error { .. }) {
12017 saw_error = true;
12018 }
12019 }
12020 assert!(saw_error);
12021
12022 handle.clear_error();
12023 let response = agent.chat("second").await.unwrap();
12024
12025 assert_eq!(response.content, "Recovered response");
12026 let messages = agent.memory.get_messages(None).await.unwrap();
12027 let user_count = messages
12028 .iter()
12029 .filter(|message| message.role == ai_agents_core::Role::User)
12030 .count();
12031 assert_eq!(user_count, 2);
12032 }
12033
12034 #[tokio::test]
12035 async fn test_buffered_streaming_route_miss_releases_buffer_limit() {
12036 use futures::StreamExt;
12037
12038 let mut mock = mock_with_response("one two three");
12039 mock.set_latency(10);
12040 let yaml = r#"
12041name: BufferedMissAgent
12042system_prompt: "You stream safely."
12043llm:
12044 default: default
12045streaming:
12046 enabled: true
12047 buffer_size: 1
12048runtime:
12049 optimization:
12050 enabled: true
12051 max_speculative_llm_calls_per_turn: 2
12052 speculative_state_transitions: true
12053 streaming_policy: buffer_until_routing_done
12054 max_parallel_runtime_tasks: 2
12055states:
12056 initial: triage
12057 states:
12058 triage:
12059 prompt: "Answer from triage."
12060 transitions:
12061 - to: billing
12062 guard:
12063 context:
12064 route:
12065 eq: billing
12066 timing: parallel
12067 billing:
12068 prompt: "Billing state."
12069"#;
12070 let agent = AgentBuilder::from_yaml(yaml)
12071 .unwrap()
12072 .llm_alias("default", Arc::new(mock))
12073 .build()
12074 .unwrap();
12075
12076 let mut stream = agent.chat_stream("hello").await.unwrap();
12077 let mut content = String::new();
12078 let mut error = None;
12079 while let Some(chunk) = stream.next().await {
12080 match chunk {
12081 StreamChunk::Content { text } => content.push_str(&text),
12082 StreamChunk::Error { message } => error = Some(message),
12083 StreamChunk::Done {} => break,
12084 _ => {}
12085 }
12086 }
12087
12088 assert_eq!(error, None);
12089 assert_eq!(content, "one two three");
12090 }
12091
12092 #[tokio::test]
12093 async fn test_buffered_streaming_main_failure_finalizes_branch() {
12094 use futures::StreamExt;
12095
12096 let mock = mock_with_response("one two");
12097 let mut router_mock = mock_with_response("0");
12098 router_mock.set_latency(50);
12099 let yaml = r#"
12100name: BufferedFailureAgent
12101system_prompt: "You stream safely."
12102llm:
12103 default: default
12104 router: router
12105observability:
12106 enabled: true
12107 export:
12108 write_raw_events: true
12109streaming:
12110 enabled: true
12111 buffer_size: 1
12112runtime:
12113 optimization:
12114 enabled: true
12115 max_speculative_llm_calls_per_turn: 2
12116 speculative_state_transitions: true
12117 streaming_policy: buffer_until_routing_done
12118 max_parallel_runtime_tasks: 2
12119states:
12120 initial: triage
12121 states:
12122 triage:
12123 prompt: "Ask for the category."
12124 transitions:
12125 - to: billing
12126 when: "User asks about billing"
12127 timing: parallel
12128 billing:
12129 prompt: "Billing state."
12130"#;
12131 let agent = AgentBuilder::from_yaml(yaml)
12132 .unwrap()
12133 .llm_alias("default", Arc::new(mock))
12134 .llm_alias("router", Arc::new(router_mock))
12135 .build()
12136 .unwrap();
12137
12138 let mut stream = agent.chat_stream("hello").await.unwrap();
12139 let mut error = String::new();
12140 while let Some(chunk) = stream.next().await {
12141 if let StreamChunk::Error { message } = chunk {
12142 error = message;
12143 }
12144 }
12145
12146 assert!(
12147 error.contains("stream buffer filled"),
12148 "unexpected stream error: {}",
12149 error
12150 );
12151 let events = agent.observability().unwrap().raw_events();
12152 assert!(events.iter().any(|event| {
12153 event.dimensions.get("branch_status") == Some(&"failed".to_string())
12154 && event.dimensions.get("commit_behavior") == Some(&"final_response".to_string())
12155 && event.dimensions.get("optimization")
12156 == Some(&"buffered_streaming_routing".to_string())
12157 }));
12158 }
12159
12160 #[tokio::test]
12161 async fn test_streaming_preflight_does_not_emit_old_state_content() {
12162 use futures::StreamExt;
12163
12164 let mock = mock_with_response("Billing streamed response");
12165 let yaml = r#"
12166name: StreamingOptimizedAgent
12167system_prompt: "You route before streaming."
12168runtime:
12169 optimization:
12170 enabled: true
12171 pre_response_deterministic_transitions: true
12172streaming:
12173 enabled: true
12174states:
12175 initial: greeting
12176 states:
12177 greeting:
12178 prompt: "OLD_STATE_SENTINEL"
12179 transitions:
12180 - to: billing
12181 guard:
12182 context:
12183 topic:
12184 eq: billing
12185 timing: pre_response
12186 billing:
12187 prompt: "Billing state."
12188"#;
12189 let agent = AgentBuilder::from_yaml(yaml)
12190 .unwrap()
12191 .llm(Arc::new(mock))
12192 .build()
12193 .unwrap();
12194 agent
12195 .set_context("topic", serde_json::json!("billing"))
12196 .unwrap();
12197
12198 let mut stream = agent.chat_stream("billing please").await.unwrap();
12199 let mut content = String::new();
12200 while let Some(chunk) = stream.next().await {
12201 match chunk {
12202 StreamChunk::Content { text } => content.push_str(&text),
12203 StreamChunk::Error { message } => panic!("stream error: {}", message),
12204 StreamChunk::Done {} => break,
12205 _ => {}
12206 }
12207 }
12208
12209 assert_eq!(agent.current_state().as_deref(), Some("billing"));
12210 assert!(content.contains("Billing streamed response"));
12211 assert!(!content.contains("OLD_STATE_SENTINEL"));
12212 }
12213
12214 #[tokio::test]
12216 async fn test_integration_state_machine_basic() {
12217 let yaml = r#"
12218name: StateAgent
12219system_prompt: "You are a support agent."
12220states:
12221 initial: greeting
12222 states:
12223 greeting:
12224 prompt: "Welcome the user warmly."
12225 transitions:
12226 - to: support
12227 when: "User needs help"
12228 auto: true
12229 support:
12230 prompt: "Help solve the user's problem."
12231"#;
12232 let mock = mock_with_responses(vec![
12233 "Welcome! How can I help?", "1", "I'll help you with that.", ]);
12237 let builder = AgentBuilder::from_yaml(yaml).unwrap();
12238 let agent = builder.llm(Arc::new(mock)).build().unwrap();
12239
12240 assert_eq!(agent.current_state(), Some("greeting".to_string()));
12241 let _ = agent.chat("I need help").await.unwrap();
12242 }
12245
12246 #[tokio::test]
12248 async fn test_integration_state_on_enter_set_context() {
12249 let yaml = r#"
12250name: ActionAgent
12251system_prompt: "You are helpful."
12252states:
12253 initial: step1
12254 states:
12255 step1:
12256 prompt: "Step 1"
12257 on_exit:
12258 - set_context:
12259 step1_exited: true
12260 transitions:
12261 - to: step2
12262 when: "always"
12263 auto: true
12264 step2:
12265 prompt: "Step 2"
12266 on_enter:
12267 - set_context:
12268 step2_entered: true
12269"#;
12270 let mock = mock_with_responses(vec![
12272 "Processing step 1.",
12273 "0", ]);
12275 let builder = AgentBuilder::from_yaml(yaml).unwrap();
12276 let agent = builder.llm(Arc::new(mock)).build().unwrap();
12277
12278 assert_eq!(agent.current_state(), Some("step1".to_string()));
12279
12280 agent.transition_to("step2").await.unwrap();
12282
12283 assert_eq!(agent.current_state(), Some("step2".to_string()));
12284
12285 let ctx = agent.get_context();
12287 assert_eq!(ctx.get("step1_exited"), Some(&serde_json::json!(true)));
12288 assert_eq!(ctx.get("step2_entered"), Some(&serde_json::json!(true)));
12289 }
12290
12291 #[tokio::test]
12293 async fn test_integration_process_normalize() {
12294 let yaml = r#"
12295name: ProcessAgent
12296system_prompt: "You are helpful."
12297process:
12298 input:
12299 - type: normalize
12300 config:
12301 trim: true
12302 collapse_whitespace: true
12303"#;
12304 let mock = mock_with_response("Got your message.");
12305 let builder = AgentBuilder::from_yaml(yaml).unwrap();
12306 let agent = builder.llm(Arc::new(mock.clone())).build().unwrap();
12307
12308 let _ = agent.chat(" hello world ").await.unwrap();
12309
12310 let history = mock.call_history();
12312 assert!(!history.is_empty());
12313 let last_call = history.last().unwrap();
12315 let user_msg = last_call
12316 .messages
12317 .iter()
12318 .find(|m| m.role == ai_agents_core::Role::User)
12319 .unwrap();
12320 assert_eq!(user_msg.content, "hello world");
12321 }
12322
12323 #[tokio::test]
12327 async fn test_integration_memory_compression() {
12328 let yaml = r#"
12329name: MemoryAgent
12330system_prompt: "You are helpful."
12331memory:
12332 type: compacting
12333 max_messages: 100
12334 compress_threshold: 5
12335 max_recent_messages: 3
12336 summarize_batch_size: 2
12337"#;
12338 let responses: Vec<&str> = (0..8).map(|_| "Response from assistant.").collect();
12340 let mock = mock_with_responses(responses);
12341 let builder = AgentBuilder::from_yaml(yaml).unwrap();
12342 let agent = builder.llm(Arc::new(mock)).build().unwrap();
12343
12344 for i in 0..6 {
12346 let _ = agent.chat(&format!("Message {}", i)).await.unwrap();
12347 }
12348
12349 let messages = agent.memory.get_messages(None).await.unwrap();
12352 assert!(messages.len() <= 12); }
12356
12357 #[tokio::test]
12359 async fn test_integration_multi_llm_registry() {
12360 let mut mock_default = MockLLMProvider::new("default");
12361 mock_default.set_response("Default LLM response.");
12362 let mut mock_router = MockLLMProvider::new("router");
12363 mock_router.set_response("Router response.");
12364
12365 let agent = AgentBuilder::new()
12366 .system_prompt("You are helpful.")
12367 .llm_alias("default", Arc::new(mock_default))
12368 .llm_alias("router", Arc::new(mock_router))
12369 .build()
12370 .unwrap();
12371
12372 let response = agent.chat("Hello").await.unwrap();
12373 assert_eq!(response.content, "Default LLM response.");
12374 }
12375
12376 #[tokio::test]
12378 async fn test_integration_agent_reset() {
12379 let mock = mock_with_responses(vec!["Hello!", "Hello again!"]);
12380 let agent = AgentBuilder::new()
12381 .system_prompt("You are helpful.")
12382 .llm(Arc::new(mock))
12383 .build()
12384 .unwrap();
12385
12386 let _ = agent.chat("Hi").await.unwrap();
12387 let messages = agent.memory.get_messages(None).await.unwrap();
12388 assert_eq!(messages.len(), 2); agent.reset().await.unwrap();
12391 let messages = agent.memory.get_messages(None).await.unwrap();
12392 assert_eq!(messages.len(), 0);
12393 }
12394
12395 #[tokio::test]
12397 async fn test_integration_process_validate_reject() {
12398 use ai_agents_process::{ProcessConfig, ProcessProcessor};
12399
12400 let validate_config = ai_agents_process::ValidateStage {
12401 id: Some("length_check".to_string()),
12402 condition: None,
12403 config: ai_agents_process::ValidateConfig {
12404 rules: vec![ai_agents_process::ValidationRule::MinLength {
12405 min_length: 10,
12406 on_fail: ai_agents_process::ValidationAction {
12407 action: ai_agents_process::ValidationActionType::Reject,
12408 message: None,
12409 },
12410 }],
12411 ..Default::default()
12412 },
12413 };
12414 let process_config = ProcessConfig {
12415 input: vec![ai_agents_process::ProcessStage::Validate(validate_config)],
12416 ..Default::default()
12417 };
12418 let processor = ProcessProcessor::new(process_config);
12419
12420 let mock = mock_with_response("Should not reach here.");
12421 let agent = AgentBuilder::new()
12422 .system_prompt("You are helpful.")
12423 .llm(Arc::new(mock))
12424 .process_processor(processor)
12425 .build()
12426 .unwrap();
12427
12428 let response = agent.chat("Hi").await.unwrap();
12429 assert!(
12431 response.content.contains("rejected")
12432 || response.content.contains("Input rejected")
12433 || response.content.contains("too short")
12434 || response.content.contains("Too short")
12435 || response.content.len() < 50, "Expected rejection response, got: {}",
12437 response.content
12438 );
12439 }
12440
12441 #[tokio::test]
12443 async fn test_llm_fallback_on_failure() {
12444 use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
12445
12446 let mut primary = MockLLMProvider::new("primary");
12447 primary.set_error("Primary LLM is unavailable");
12448
12449 let mut fallback = MockLLMProvider::new("fallback");
12450 fallback.set_response("Fallback response works!");
12451
12452 let agent = AgentBuilder::new()
12453 .system_prompt("You are helpful.")
12454 .llm_alias("default", Arc::new(primary))
12455 .llm_alias("backup", Arc::new(fallback))
12456 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
12457 llm: LLMRecoveryConfig {
12458 on_failure: LLMFailureAction::FallbackLlm {
12459 fallback_llm: "backup".to_string(),
12460 },
12461 ..Default::default()
12462 },
12463 ..Default::default()
12464 }))
12465 .build()
12466 .unwrap();
12467
12468 let response = agent.chat("Hello").await.unwrap();
12469 assert!(
12470 response.content.contains("Fallback response"),
12471 "Expected fallback response, got: {}",
12472 response.content
12473 );
12474 }
12475
12476 #[tokio::test]
12478 async fn test_llm_fallback_response_static_message() {
12479 use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
12480
12481 let mut primary = MockLLMProvider::new("primary");
12482 primary.set_error("Primary LLM is unavailable");
12483
12484 let agent = AgentBuilder::new()
12485 .system_prompt("You are helpful.")
12486 .llm(Arc::new(primary))
12487 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
12488 llm: LLMRecoveryConfig {
12489 on_failure: LLMFailureAction::FallbackResponse {
12490 message: "I am temporarily unavailable. Please try again later."
12491 .to_string(),
12492 },
12493 ..Default::default()
12494 },
12495 ..Default::default()
12496 }))
12497 .build()
12498 .unwrap();
12499
12500 let response = agent.chat("Hello").await.unwrap();
12501 assert!(
12502 response.content.contains("temporarily unavailable"),
12503 "Expected static fallback message, got: {}",
12504 response.content
12505 );
12506 }
12507
12508 #[tokio::test]
12510 async fn test_tool_failure_skip() {
12511 use ai_agents_recovery::{
12512 ErrorRecoveryConfig, ToolFailureAction, ToolRecoveryConfig, ToolRetryConfig,
12513 };
12514
12515 let mock = mock_with_responses(vec![
12517 r#"I'll use the nonexistent tool.
12518[TOOL_CALL: {"name": "nonexistent_tool", "arguments": {}}]"#,
12519 "The tool was unavailable, but I can still help you.",
12520 ]);
12521
12522 let agent = AgentBuilder::new()
12523 .system_prompt("You are helpful.")
12524 .llm(Arc::new(mock))
12525 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
12526 tools: ToolRecoveryConfig {
12527 default: ToolRetryConfig {
12528 max_retries: 0,
12529 timeout_ms: None,
12530 on_failure: ToolFailureAction::Skip,
12531 },
12532 ..Default::default()
12533 },
12534 ..Default::default()
12535 }))
12536 .build()
12537 .unwrap();
12538
12539 let response = agent.chat("Use the nonexistent tool").await;
12541 assert!(
12542 response.is_ok(),
12543 "Expected Ok with skip policy, got: {:?}",
12544 response
12545 );
12546 }
12547}