Skip to main content

ai_agents_runtime/
runtime.rs

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
13/// Shared lock table used to serialize side-effecting tool calls by canonical resource.
14pub(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
81/// Outcome of processing tool calls within the agent loop.
82enum ToolCallOutcome {
83    /// Tools executed successfully, continue the LLM loop for the next iteration.
84    Continue,
85    /// A state transition fired during tool call handling, continue the loop.
86    TransitionFired,
87    /// HITL rejected a tool call, return this response immediately.
88    Rejected(AgentResponse),
89}
90
91/// Outcome of skill routing — used by `try_skill_route`.
92enum SkillRouteResult {
93    /// No skill matched, continue to normal LLM chat.
94    NoMatch,
95    /// Skill executed successfully.
96    Response(String),
97    /// Skill matched but needs disambiguation first.
98    NeedsClarification(AgentResponse),
99}
100
101/// Result of response-independent parallel transition selection.
102enum ParallelTransitionSelection {
103    /// A transition matched and can be committed before the old-state response.
104    Candidate(TransitionCandidate),
105    /// All eligible transition checks ran and none matched.
106    NoMatch,
107    /// LLM-based route evaluation needed speculative capacity that was unavailable.
108    ReservationExhausted,
109}
110
111/// Outcome of post_loop_processing - drives the caller's next step.
112enum PostLoopResult {
113    /// No transition fired. Content is the LLM response for this turn.
114    NoTransition(String),
115    /// Transition fired. Content is from plain post-transition re-generation.
116    Transitioned(String),
117    /// Transition fired into a state that requires full dispatch.
118    /// Caller re-enters run_loop_internal to apply the correct handler.
119    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/// Host-owned runtime control state shared with active agents.
139#[derive(Debug)]
140struct RuntimeControlState {
141    /// Monotonic version for runtime-control snapshots.
142    version: AtomicU64,
143    /// Emergency switch that denies future calls and is shared with active tool contexts.
144    emergency_deny: Arc<AtomicBool>,
145    /// Optional live replacement for tool security policy.
146    tool_security_override: RwLock<Option<ToolSecurityConfig>>,
147    /// Optional live replacement for the top-level tool scope.
148    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/// Host-only handle for live runtime safety controls.
163#[derive(Clone)]
164pub struct RuntimeControlHandle {
165    state: Arc<RuntimeControlState>,
166}
167
168impl RuntimeControlHandle {
169    /// Returns the current runtime-control version.
170    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    /// Overrides tool security for later tool calls.
179    pub fn set_tool_security(&self, config: ToolSecurityConfig) -> u64 {
180        *self.state.tool_security_override.write() = Some(config);
181        self.bump()
182    }
183
184    /// Clears the live tool security override.
185    pub fn clear_tool_security_override(&self) -> u64 {
186        *self.state.tool_security_override.write() = None;
187        self.bump()
188    }
189
190    /// Overrides the top-level tool scope for later calls.
191    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    /// Clears the live tool scope override.
197    pub fn clear_tool_scope_override(&self) -> u64 {
198        *self.state.tool_scope_override.write() = None;
199        self.bump()
200    }
201
202    /// Enables or disables emergency denial for future tool calls.
203    pub fn set_emergency_deny(&self, enabled: bool) -> u64 {
204        self.state.emergency_deny.store(enabled, Ordering::SeqCst);
205        self.bump()
206    }
207
208    /// Denies future calls and asks active cancellable work to stop.
209    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    /// Structured persona manager for identity, evolution, and secrets.
247    persona_manager: Option<Arc<ai_agents_persona::PersonaManager>>,
248    /// Skill ID that triggered the current pending disambiguation.
249    /// Set by try_skill_route() when skill-level disambiguation triggers clarification.
250    /// Read by run_loop() when clarification resolves to route directly to the skill.
251    pending_skill_id: RwLock<Option<String>>,
252    current_plan: RwLock<Option<Plan>>,
253    /// Tool IDs declared in the top-level `tools:` spec.
254    declared_tool_ids: Option<Vec<String>>,
255    /// Whether the context manager has been initialized (defaults loaded, env resolved, etc.)
256    context_initialized: AtomicBool,
257    /// Spawner for dynamic agent creation (set when YAML has a spawner: section).
258    spawner: Option<Arc<crate::spawner::AgentSpawner>>,
259    /// Registry tracking spawned agents (set when YAML has a spawner: section).
260    spawner_registry: Option<Arc<crate::spawner::AgentRegistry>>,
261    /// Re-dispatch depth for post-transition full dispatch.
262    /// 0 = not re-dispatching. > 0 = user message already in memory, skip re-adding.
263    redispatch_depth: RwLock<u32>,
264    /// Active optimized turn context used to keep root lifecycle state in one place.
265    active_turn_context: RwLock<Option<TurnOptimizationContext>>,
266    /// Tracks whether the root turn already wrote the processed user message.
267    root_user_message_committed: AtomicBool,
268    /// Current actor ID for cross-session memory.
269    actor_id: RwLock<Option<String>>,
270    /// Fact store for managing per-actor extracted facts.
271    fact_store: RwLock<Option<Arc<ai_agents_facts::FactStore>>>,
272    /// Fact extractor for LLM-based fact extraction.
273    /// None when actor_memory is enabled without facts.enabled.
274    fact_extractor: RwLock<Option<Arc<dyn ai_agents_facts::FactExtractor>>>,
275    /// Cached actor facts keyed by actor ID so concurrent or alternating turns do not overwrite one another.
276    actor_facts_cache: Arc<RwLock<HashMap<String, Vec<ai_agents_core::KeyFact>>>>,
277    /// Number of messages since last fact extraction.
278    messages_since_extraction: Arc<RwLock<usize>>,
279    /// Actor memory configuration.
280    actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
281    /// Facts configuration.
282    facts_config: Option<ai_agents_facts::FactsConfig>,
283    /// Session-scoped metadata (tags, ttl, actor roster).
284    session_metadata: RwLock<ai_agents_core::SessionMetadata>,
285    /// Session id currently bound to this runtime instance.
286    current_session_id: RwLock<Option<String>>,
287    /// Relationship manager for actor-scoped social memory.
288    relationship_manager: Option<Arc<RelationshipManager>>,
289    /// Observability manager for traces, metrics, reports, and exports.
290    observability_manager: Option<Arc<ObservabilityManager>>,
291    /// Runtime optimization and maintenance policy.
292    runtime_config: RuntimeConfig,
293    /// Queue for background maintenance tasks.
294    background_maintenance: Arc<BackgroundMaintenanceQueue>,
295    /// Cross-tool locks for side-effecting calls that target the same resource.
296    resource_locks: ToolResourceLocks,
297    /// Host-only runtime control state.
298    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    /// Scopes clarification question generation as disambiguation_clarification.
328    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    /// Scopes clarification response parsing as disambiguation_clarification.
338    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    /// Scopes one process stage using the purpose implied by its stage type.
352    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    /// Attach a relationship manager configured by the builder or host application.
487    pub fn with_relationships(mut self, manager: Arc<RelationshipManager>) -> Self {
488        self.relationship_manager = Some(manager);
489        self
490    }
491
492    /// Attach a shared observability manager for traces, metrics, reports, and exports.
493    pub fn with_observability(mut self, manager: Arc<ObservabilityManager>) -> Self {
494        self.observability_manager = Some(manager);
495        self
496    }
497
498    /// Attach runtime optimization policy and resize the background queue.
499    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    /// Returns the runtime optimization policy.
507    pub fn runtime_config(&self) -> &RuntimeConfig {
508        &self.runtime_config
509    }
510
511    /// Wait for all background maintenance tasks to finish.
512    pub async fn flush_background_tasks(&self) -> Result<()> {
513        self.background_maintenance.flush_all().await
514    }
515
516    /// Wait for background maintenance associated with one actor to finish.
517    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    /// Wait for background maintenance associated with one task kind to finish.
522    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    /// Wait for background maintenance associated with one actor and task kind to finish.
530    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    /// Flush background maintenance before a host shuts down the runtime.
541    pub async fn shutdown_background_tasks(&self) -> Result<()> {
542        self.flush_background_tasks().await
543    }
544
545    /// Returns the configured observability manager for report and export access.
546    pub fn observability(&self) -> Option<Arc<ObservabilityManager>> {
547        self.observability_manager.clone()
548    }
549
550    /// Exports observability files after a turn when export settings request it.
551    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    /// Returns the configured relationship manager, if relationship memory is enabled.
565    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    /// Returns the current session ID or creates one for unsaved observed turns.
605    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    /// Builds the root or child observation context for a chat entry point.
614    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    /// Refreshes task-local context with current runtime labels and a purpose.
631    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    /// Runs a future under a purpose while preserving current trace context.
650    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    /// Runs chat with actor and observation context while avoiding recursive async types.
662    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    /// Run one turn with turn-scoped actor context without mutating the runtime's global actor ID.
687    ///
688    /// The supplied context is available to actor-scoped facts, relationship memory, orchestration, and prompt templates only for the lifetime of this call.
689    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    /// Convenience wrapper around [`Self::chat_with_actor_context`] for a turn whose original actor is known up front.
699    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    /// Ensure the effective actor's relationship is loaded from storage into the relationship manager.
705    pub async fn load_actor_relationship(&self) -> Result<()> {
706        self.maybe_load_actor_relationship().await;
707        Ok(())
708    }
709
710    /// Manually apply a delta to the effective actor's `agent_to_actor` relationship perspective and persist the updated relationship when storage is configured.
711    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    /// Manually apply a delta to a specific relationship perspective for the effective actor.
727    ///
728    /// Use this for two-sided configurations when you need to update `agent_to_actor`, `perceived_actor_to_agent`, or `mutual` explicitly from application logic.
729    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    /// Set only the actor memory and facts configs without creating the store.
775    /// The store and extractor are created lazily in init_storage().
776    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    /// Configure fact store and optional extractor for actor memory.
787    /// Pass `None` for `extractor` to load existing facts without running extraction.
788    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    /// Get the fact store for direct fact manipulation.
803    pub fn fact_store(&self) -> Option<Arc<ai_agents_facts::FactStore>> {
804        self.fact_store.read().clone()
805    }
806
807    /// Get the current actor ID.
808    pub fn actor_id(&self) -> Option<String> {
809        self.actor_id.read().clone()
810    }
811
812    /// Set the current actor ID (player, user, another agent, etc.).
813    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    /// Set the current actor ID. Convenience wrapper around set_actor_id.
826    pub fn set_user_id(&self, user_id: &str) -> ai_agents_core::Result<()> {
827        self.set_actor_id(user_id)
828    }
829
830    /// Load facts for the current actor from storage and cache them for prompt injection.
831    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    /// Load actor memory only when the effective actor has no cached facts yet.
852    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    /// Pre-turn lifecycle shared by streaming and non-streaming paths.
863    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    /// Post-turn lifecycle shared by streaming and non-streaming paths.
876    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    /// Starts root-turn bookkeeping for user-message commit tracking.
885    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    /// Writes the processed user message once for the root turn.
934    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    /// Clears root-turn bookkeeping after final response handling.
953    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    /// Applies freshness policy before rendering the next prompt.
985    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    /// Runs post-turn facts and relationship maintenance according to runtime policy.
1036    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    /// Queues background maintenance or applies the configured overflow behavior.
1269    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    /// Records a skipped background maintenance event when work cannot run.
1404    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    /// Get cached actor facts for the effective actor.
1428    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    /// Returns the formatted relationship prompt text for the effective actor, if relationship injection produced any text for this turn.
1440    pub fn relationship_memory_text(&self) -> Option<String> {
1441        self.format_relationship_for_context().map(|(_, text)| text)
1442    }
1443
1444    /// Manually extract facts from the last N messages.
1445    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        // Save to storage and update the actor-scoped cache when an actor is known.
1489        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                // add_facts now returns the authoritative post-write set.
1495                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    /// Resolve actor_id from context if method is from_context.
1529    /// Supports dotted paths (e.g. "player.id", "user.profile.id").
1530    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                    // get_path resolves dotted paths; get only handles top-level keys.
1544                    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    /// Format actor facts for template injection.
1567    fn format_actor_facts_for_context(&self) -> String {
1568        // Respect inject_in_context: false.
1569        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        // Effective token cap: prefer memory.token_budget.allocation.facts when present,
1594        // otherwise fall back to actor_memory.injection.max_tokens.
1595        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        // Filter by category when injection.mode = category.
1604        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    /// Run auto-extraction after a chat turn.
1876    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        // Complete facts setup now that storage is available.
1950        self.complete_facts_init().await;
1951        Ok(())
1952    }
1953
1954    /// Initialize fact store and extractor from stored config + current storage.
1955    /// Called from init_storage() so facts are ready before the first turn.
1956    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    /// Returns the spawner if configured via a spawner: YAML section.
2025    pub fn spawner(&self) -> Option<&Arc<crate::spawner::AgentSpawner>> {
2026        self.spawner.as_ref()
2027    }
2028
2029    /// Returns the agent registry if configured via a spawner: YAML section.
2030    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    /// Returns the host-only runtime control handle.
2090    pub fn runtime_control(&self) -> RuntimeControlHandle {
2091        RuntimeControlHandle {
2092            state: Arc::clone(&self.runtime_control),
2093        }
2094    }
2095
2096    /// Installs or clears the host question handler used by `ask_user`.
2097    pub fn set_question_handler(&self, handler: Option<Arc<dyn QuestionHandler>>) {
2098        self.tools.set_question_handler(handler);
2099    }
2100
2101    /// Installs the host diagnostics provider used by `diagnostics`.
2102    pub fn set_diagnostics_provider(&self, provider: Arc<dyn DiagnosticsProvider>) {
2103        self.tools.set_diagnostics_provider(provider);
2104    }
2105
2106    /// Installs the host command runner used by `command`.
2107    pub fn set_command_runner(&self, runner: Arc<dyn CommandRunner>) {
2108        self.tools.set_command_runner(runner);
2109    }
2110
2111    /// Returns the current session-local todo list.
2112    pub fn todos(&self) -> Vec<TodoItem> {
2113        self.tools.todos()
2114    }
2115
2116    /// Returns the effective tool security snapshot for the next decision.
2117    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    /// Get a copy of current session metadata.
2198    pub fn session_metadata(&self) -> ai_agents_core::SessionMetadata {
2199        self.session_metadata.read().clone()
2200    }
2201
2202    /// Delete all facts and sessions for an actor, gated by privacy.allow_deletion.
2203    /// Returns Err when actor_memory.privacy.allow_deletion is false.
2204    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    /// Overwrite session metadata (tags, ttl, custom fields).
2230    pub fn set_session_metadata(&self, meta: ai_agents_core::SessionMetadata) {
2231        *self.session_metadata.write() = meta;
2232    }
2233
2234    /// Delete sessions whose TTL has expired. Returns number of sessions removed.
2235    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    /// List sessions matching a filter. Supports actor, tag, and date filters.
2252    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    /// Save state including spawned agents manifest for session persistence.
2295    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                // Fire on_session_created when this session id is first seen on this runtime.
2352                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                // Update metadata before persisting.
2362                {
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                // Persist metadata alongside the snapshot.
2381                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                    // Restore session metadata alongside the snapshot.
2398                    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        // Inject actor_facts for {{ actor_facts }} template variable.
2558        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    /// Returns `Some(tools)` if the current state explicitly declares tools
2644    /// (including `Some([])` for "no tools"), or `None` if the state doesn't
2645    /// specify tools (meaning: fall back to agent-level declared_tool_ids).
2646    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                // Persona always prepended regardless of prompt_mode.
2751                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        // No state machine - prepend persona to base.
2772        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        // Include the current state's prompt text so the detector understands
2854        // what kind of input is expected (e.g., "Ask for the order number").
2855        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        // Extract canonical intent labels from current state's transitions
2871        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        // Try direct JSON parse first
2986        if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(content) {
2987            // Handle JSON array of tool calls (parallel tool calling)
2988            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            // Handle single JSON object
2998            if let Some(tool_call) = self.extract_tool_call_from_value(&parsed) {
2999                return Some(vec![tool_call]);
3000            }
3001        }
3002
3003        // Try to extract JSON from content (handles extra text/braces from LLM)
3004        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                // Handle JSON array of tool calls (parallel tool calling)
3007                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                // Handle single JSON object
3017                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    // Lite models could generate unmatched braces: this function handles such cases
3042    fn extract_json_from_content(&self, content: &str) -> Option<String> {
3043        // Try array first (for parallel tool calls), then single object
3044        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    /// Extract a JSON array `[...]` containing tool calls from mixed content.
3051    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            // Verify it looks like an array of tool calls
3074            if json_str.contains("\"tool\"") {
3075                return Some(json_str.to_string());
3076            }
3077        }
3078
3079        None
3080    }
3081
3082    /// Extract a JSON object `{...}` containing a tool call from mixed content.
3083    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        // Count braces to find the matching closing brace
3088        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            // Verify it looks like a tool call
3107            if json_str.contains("\"tool\"") {
3108                return Some(json_str.to_string());
3109            }
3110        }
3111
3112        None
3113    }
3114
3115    /// Builds a structured record from executor state.
3116    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    /// Sends a finalized tool record to hooks, history, and error handling.
3158    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    /// Invokes a resolved tool once with timeout, cancellation, and actor context.
3181    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    /// Truncates model-facing tool output on a character boundary.
3215    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    /// Acquires a cross-tool lock for side-effecting calls that share a resource.
3229    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    /// Executes a resolved tool with retry policy from recovery settings.
3250    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    /// Executes a tool request through scope, policy, HITL, timeout, recovery, and evidence recording.
3278    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    //
4128    // This method is branch-safe because it only asks the router which skill matches.
4129    // Do not add disambiguation, pending-skill writes, or skill execution here.
4130    //
4131    /// Selects a skill without executing it.
4132    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    //
4159    // This is the commit half of skill routing.
4160    // It may mutate pending skill state, run clarification, and execute skill steps.
4161    //
4162    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    /// Result of skill routing.
4278    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    /// Execute a skill with reasoning and reflection, returning the response string.
4288    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    /// Execute a skill by ID, bypassing the skill router.
4327    /// Used after skill-triggered disambiguation resolves to route directly to the matched skill.
4328    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        // Gate pass against confidence threshold.
4514        // LLM may say PASS but with low confidence - the threshold catches this.
4515        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    /// Process input through the pipeline (state-level override or agent-level).
4538    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    /// Process output through the pipeline (state-level override or agent-level).
4555    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    /// Build a ProcessProcessor from the current state's process config, if any.
4576    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(&current, &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    /// Builds transition context with optional staged writes overlaid.
4633    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    /// Selects a post-response transition without committing side effects.
4647    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, &current_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    /// Selects a guard or resolved-intent transition without an LLM call.
4683    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    /// Commits a selected transition through the shared post-response path.
4725    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    /// Runs state transition approval before any transition side effects.
4731    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    /// Applies transition side effects after approval has succeeded.
4740    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    /// Approves and applies a transition without staged context writes.
4767    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    /// Commits a pre-response transition after approval and before redispatch.
4781    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    /// Handles post-response transition misses with fallback counters.
4804    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    /// Evaluates and commits post-response transitions for the committed response path.
4817    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(&current_state).await
4831    }
4832
4833    /// Attempts deterministic pre-response routing before old-state response generation.
4834    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                &current_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    //
4902    // Speculative branches overlap independent decisions but still commit exactly one path.
4903    // Losing branches must remain data only and must not write memory, run tools, or emit output.
4904    //
4905    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    //
5477    // Pending losers are reported as cancelled because their futures are dropped before completion.
5478    // Completed losers keep failed or discarded status based on their recorded result.
5479    //
5480    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    //
5499    // Finalization is separated from commit so losing branches remain observable.
5500    // This helper must not mutate runtime state other than observability.
5501    //
5502    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    //
5521    // This is only an eligibility check.
5522    // Actual transition selection happens in select_parallel_transition_candidate.
5523    //
5524    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    //
5535    // Parallel transition prompts must not depend on assistant response text.
5536    // Keep this branch response-independent or it can race against invalid context.
5537    //
5538    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            &current_state,
5557            &parallel,
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    /// Re-enters the runtime loop after an optimized transition commits.
5619    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    /// Runs final response hooks and maintenance only for the root dispatch.
5643    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    /// Execute on_exit actions for a state being left.
5656    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    /// Execute on_enter (or on_reenter) actions for a state being entered.
5668    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                // Check if this state was previously visited
5672                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    /// Execute a list of state actions (tool calls, skill invocations, context updates, LLM prompts).
5686    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                            // Render template variables and include conversation context
5761                            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    /// Run context extractors for the current state on the user's input.
5877    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        // Handle overflow AFTER compression, then check warning threshold
5906        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        // Overall budget warning
5923        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        // Per-component warning: summary
5936        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        // Per-component warning: recent_messages
5950        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        // Estimate ~50 tokens per message on average
6020        ((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        // Apply planning-level tool and skill filters.
6175        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        // Build tool descriptions with argument schemas so the planner
6185        // knows how to construct valid args for each step.
6186        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                    // When a tool step has dependency results, ask the LLM to
6350                    // produce the correct arguments given the context and tool schema.
6351                    // This avoids brittle {{stepN}} template substitution and lets
6352                    // the LLM handle type adaptation (e.g. picking the iso field
6353                    // from a datetime result for a downstream format call).
6354                    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                                // Parse the LLM's JSON response, fall back to planner args.
6410                                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        // Set plan status based on whether any steps actually failed.
6526        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        // Synthesize final output from all completed step results.
6542        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        // Synthesize a coherent summary from multiple step results via LLM.
6564        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        // One-shot context initialization: load runtime defaults, resolve env vars,
6624        // populate builtin sources (session, agent), etc.  This must happen before
6625        // the first template render so that {{ context.* }} variables are available.
6626        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        // Clear stale disambiguation context from previous turns.
6635        // This prevents resolved_intent from leaking across turns and causing incorrect deterministic routing on subsequent inputs.
6636        self.clear_disambiguation_context();
6637
6638        // Disambiguation check (before input processing)
6639        if let Some(ref disambiguator) = self.disambiguation_manager {
6640            let disambiguation_context = self.build_disambiguation_context().await?;
6641
6642            // Get state-level disambiguation override
6643            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                    // Routing uses `resolved` (structured, deterministic)
6707                    // This is what makes post-disambiguation routing DETERMINISTIC
6708                    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                    // Check if this clarification was triggered by a skill-level override.
6722                    // If so, route directly to the matched skill instead of going through
6723                    // skill routing again (which might match a different skill).
6724                    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                    // Same skill-id re-check for best-guess path
6738                    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                            // Topic switch: process the user's new input from scratch.
6813                            // The LLM sees full conversation context including the abandoned exchange.
6814                            return self.run_loop_internal(&fresh_input).await;
6815                        }
6816                        None => {
6817                            // Pure abandonment: generate a brief acknowledgment.
6818                            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    /// Generate a localized response using the router LLM
6847    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    /// Clear disambiguation-related keys from the context manager.
6890    ///
6891    /// Render template variables in state action args using the context manager.
6892    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    /// Called at the start of each turn to prevent stale `resolved_intent` from leaking across turns.
6921    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    /// Re-run skill disambiguation on enriched input before executing the skill.
6935    /// After clarification resolves, the enriched input may still be missing required_clarity fields (e.g. "Transfer money to Jane." still lacks amount).
6936    /// This method re-runs the skill's disambiguation pass.
6937    /// If fields are still missing, it returns the new clarification question and keeps pending_skill_id set.
6938    /// If all fields are present (Clear), it executes the skill and returns the response.
6939    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 the skill has disambiguation enabled, re-run it on the enriched input.
6950        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                                // Keep pending_skill_id set (do NOT clear it).
6987                                // The next turn will resolve this new clarification and
6988                                // re-enter this method until all fields are present.
6989                                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                                // Fall through to execute with the further-enriched input.
7020                                *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                                // User abandoned during skill re-check.
7084                                // Clear skill routing state and fall through to normal execution.
7085                                *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        // Skill has no disambiguation or all fields present: execute.
7115        *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    /// Handle skill routing result: output processing, memory, transitions.
7125    /// Returns a fully formed AgentResponse for skill-routed requests.
7126    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    /// Run the Plan-and-Execute flow: generate plan, execute steps, finalize.
7151    /// Supports plan-level reflection with replan loop when configured.
7152    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    /// Inject CoT/ReAct reasoning prompt into the system message (first iteration only).
7245    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    //
7327    // Draft generation must not commit user memory or run tools.
7328    // The current user input is added only as an ephemeral message for this LLM call.
7329    //
7330    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    //
7355    // This is the only place where a winning draft is allowed to become runtime state.
7356    // Parsed tool calls become executable only after this method commits the draft.
7357    //
7358    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    //
7408    // Tool drafts need a committed continuation after function results are written.
7409    // Redispatch depth suppresses duplicate root lifecycle work during that continuation.
7410    //
7411    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    //
7430    // Shared committed text finalization for normal responses and winning text drafts.
7431    // Keep output processing, reflection, transitions, hooks, and maintenance behind this commit boundary.
7432    //
7433    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    //
7480    // Auto reasoning uses this path after the judge wins with a deeper mode.
7481    // It intentionally uses committed message building instead of the draft overlay.
7482    //
7483    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    /// Handle tool calls: check transitions, execute tools in parallel, handle HITL rejection.
7556    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        // Check if a transition should fire before executing the LLM's tool call.
7564        // If a transition fires, on_enter actions handle the tool call correctly
7565        // (with proper URLs from YAML), so skip the LLM's tool call.
7566        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        // Store the assistant's tool-call message so the LLM sees its own decision in conversation history. Without this, the model only sees the tool result and may repeat the same call.
7577        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                    // Check if this is a HITL rejection - if so, break the loop
7592                    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 the rejection message to user, don't continue loop
7600                        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    /// Run the reflection loop on a response, returning (improved_content, reflection_metadata).
7620    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    /// Record the assistant turn, evaluate transitions, and decide what to do next.
7694    /// Returns PostLoopResult so callers can apply_post_loop_result for re-dispatch.
7695    async fn post_loop_processing(
7696        &self,
7697        processed_input: &str,
7698        content: String,
7699    ) -> Result<PostLoopResult> {
7700        // Do NOT add the assistant message to memory yet.
7701        // evaluate_transitions receives content as a direct parameter, so the message does not need to be in memory for transitions to evaluate correctly.
7702        // For NeedsRedispatch we skip adding the stale old-state response entirely, keeping memory clean for the re-dispatched handler.
7703
7704        self.increment_turn();
7705
7706        // Run context extractors so guards can check freshly-extracted values.
7707        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        // Check if we should skip re-generation after this transition.
7720        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        // Check if the new state needs full dispatch.
7729        // Orchestration states (concurrent, group_chat, pipeline, handoff, delegate) need their dedicated handlers.
7730        // Any non-None effective reasoning mode needs CoT/ReAct prompt injection, the plan-and-execute handler, or Auto re-determination - all of which live in run_loop_internal, not here.
7731        if self.needs_redispatch_for_new_state() {
7732            info!("Post-transition NeedsRedispatch: new state requires full dispatch");
7733            // Stale old-state content is NOT added to memory.
7734            // apply_post_loop_result will increment redispatch_depth and call run_loop_internal, which produces the correct response and adds it.
7735            return Ok(PostLoopResult::NeedsRedispatch);
7736        }
7737
7738        // Normal post-transition re-generation (plain LLM with optional tool calls).
7739        // Add the stale response to history so the new LLM sees the conversation.
7740        self.memory
7741            .add_message(ChatMessage::assistant(&content))
7742            .await?;
7743        self.check_memory_compression().await?;
7744
7745        // If a transition fired, on_enter actions already executed (e.g., HTTP calls).
7746        // The current content was generated in the OLD state context and is stale.
7747        // Re-generate in the new state context so the LLM can reference on_enter results.
7748        // If the LLM responds with a tool call, execute it in a mini-loop so the
7749        // result is not returned as raw JSON text.
7750        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            // Check if the post-transition response contains tool calls.
7777            // If so, execute them and loop so the LLM can summarize the result.
7778            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                // Loop to let the LLM see the tool result and produce a text response.
7808                continue;
7809            }
7810
7811            // No tool call - this is the final text response.
7812            self.memory
7813                .add_message(ChatMessage::assistant(&final_content))
7814                .await?;
7815            return Ok(PostLoopResult::Transitioned(final_content));
7816        }
7817
7818        // Exhausted post-transition iterations (unlikely). Return last content.
7819        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    /// Build the final AgentResponse with all metadata.
7828    /// Check whether to re-generate a response after a state transition.
7829    fn should_regenerate_after_transition(&self) -> bool {
7830        if let Some(ref sm) = self.state_machine {
7831            // Global setting
7832            if !sm.config().regenerate_on_transition {
7833                return false;
7834            }
7835            // Per-state override on the new (current) state
7836            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    /// Return true when the new state requires full dispatch via run_loop_internal.
7846    /// Covers orchestration states and any non-None effective reasoning mode.
7847    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                // Any non-None effective reasoning mode requires the main dispatch loop.
7859                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    /// Consume a PostLoopResult. NeedsRedispatch re-enters run_loop_internal.
7869    /// The user message is already in memory - redispatch_depth suppresses re-adding it.
7870    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    // Handle delegation: forward user input to a registry agent.
7951    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        // Prepare input based on delegate_context mode.
7984        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        // Backward-compatible context key.
8007        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        // Structured orchestration context.
8014        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        // post_loop_processing records the assistant turn and evaluates transitions.
8028        // apply_post_loop_result handles NeedsRedispatch by re-entering run_loop_internal.
8029        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    // Handle concurrent execution: run multiple registry agents in parallel and aggregate.
8060    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        // Render input template if provided, otherwise use the raw input.
8074        // Uses direct minijinja rendering (same approach as pipeline) so variables
8075        // are top-level: {{ user_input }}, not {{ context.user_input }}.
8076        // Enrich input with parent conversation history when context_mode is set.
8077        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        // Backward-compatible context key.
8145        let _ = self.context_manager.set(
8146            "concurrent.result",
8147            serde_json::Value::String(result.response.content.clone()),
8148        );
8149
8150        // Build per-agent result data for context and metadata.
8151        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        // Structured orchestration context with per-agent results.
8166        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    // Handle group chat: run a multi-turn multi-agent conversation.
8206    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        // Enrich input with parent conversation history when context_mode is set.
8224        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        // Render input template if provided, otherwise use the raw user message as topic.
8238        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        // Backward-compatible context key.
8264        let _ = self.context_manager.set(
8265            "group_chat.conclusion",
8266            serde_json::Value::String(result.response.content.clone()),
8267        );
8268
8269        // Build transcript data for context and metadata.
8270        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        // Structured orchestration context with full transcript.
8283        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    // Handle pipeline: run agents sequentially with per-stage input templates.
8325    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        // Enrich input with parent conversation history when context_mode is set.
8353        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        // Backward-compatible context key.
8387        let _ = self.context_manager.set(
8388            "pipeline.result",
8389            serde_json::Value::String(result.response.content.clone()),
8390        );
8391
8392        // Build per-stage data for context and metadata.
8393        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        // Structured orchestration context.
8407        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    // Handle handoff: LLM-directed agent-to-agent control transfer.
8445    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        // Enrich input with parent conversation history when context_mode is set.
8466        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        // Render input template if provided, otherwise forward the raw user message.
8480        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        // Backward-compatible context key.
8508        let _ = self.context_manager.set(
8509            "handoff.result",
8510            serde_json::Value::String(result.response.content.clone()),
8511        );
8512
8513        // Build handoff chain data for context and metadata.
8514        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        // Structured orchestration context.
8527        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    // run_loop_internal: blocking (non-streaming) agent pipeline.
8567    async fn run_loop_internal(&self, input: &str) -> Result<AgentResponse> {
8568        self.begin_root_turn();
8569        // Resolve actor_id from context, reload facts if actor changed, bump counter.
8570        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        // Inject process context (detect/extract results) into agent context
8576        // so system prompt templates can use {{ context.detected_language }} etc.
8577        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        // Handle orchestration states (delegate, concurrent, group_chat, pipeline, handoff).
8599        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        //
8630        // The speculative future is boxed to keep the runtime future size manageable.
8631        // Removing the box can overflow small test stacks because this function is recursive through redispatch.
8632        //
8633        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                        // Store the clarification question in memory so the next turn
8657                        // can be handled as a clarification response.
8658                        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 => {} // continue to normal LLM chat
8667        }
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            // When reasoning is active, cap iterations at the reasoning-specific limit.
8697            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            // Try primary LLM (with retry if configured), then apply on_failure policy.
8722            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            // Run reflection (blocking LLM calls for retries)
8815            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            // Post-loop: memory, transitions, post-transition re-generation.
8824            // apply_post_loop_result handles NeedsRedispatch by re-entering
8825            // run_loop_internal so the new state's full dispatch activates.
8826            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    /// Streaming agent pipeline
9111    /// Uses all the same shared helpers as run_loop_internal.
9112    /// The ONLY difference: LLM calls use complete_stream() + yield deltas.
9113    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            // Parity with non-stream: resolve actor from context and load facts if changed.
9123            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            // Inject process context (detect/extract results) into agent context
9135            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                //
9158                // Buffered routing keeps stale stream output hidden until a branch winner is known.
9159                // The boxed future prevents this stream state machine from becoming too large.
9160                //
9161                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            // Handle orchestration states in streaming mode.
9198            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            // Skill routing
9230            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) => {} // no skill matched, continue
9263                Err(e) => {
9264                    yield StreamChunk::error(e.to_string());
9265                    return;
9266                }
9267            }
9268
9269            // Reasoning mode determination
9270            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            // Plan-and-Execute: yield final result as single chunk (not token-by-token)
9287            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                // When reasoning is active, cap iterations at the reasoning-specific limit.
9324                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                // Check if reflection is active — if so, suppress streaming for this LLM call
9357                // because we may need to retry and the user would see a stale first attempt.
9358                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                    // Use blocking call so reflection can retry without partial output
9365                    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                    // Streaming LLM call
9383                    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                    // Construct LLMResponse for hooks
9412                    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                // Tool call handling
9421                if let Some(tool_calls) = self.parse_tool_calls(&content) {
9422                    // Emit tool events for streaming
9423                    // First check transitions (same as blocking path)
9424                    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                    // Store the assistant's tool-call message (same as blocking path)
9445                    let _ = self.memory.add_message(ChatMessage::assistant(&content)).await;
9446
9447                    // Execute tools with streaming events
9448                    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                // Extract thinking, process output
9513                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                // Reflection (uses blocking LLM calls for retries)
9536                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 was active (we used blocking call), yield the final content now
9553                if reflection_active {
9554                    yield StreamChunk::content(&final_content);
9555                }
9556
9557                // Post-loop: memory, transitions, post-transition re-generation.
9558                // For NeedsRedispatch, run_loop_internal handles the new state's full
9559                // dispatch and its result is yielded as a single non-streamed chunk.
9560                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 the post-transition re-generated or re-dispatched content.
9618                    yield StreamChunk::content(&final_content);
9619                }
9620
9621                // Parity with non-stream: finalize the committed response once.
9622                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    /// Streaming entry point with disambiguation
9635    /// Mirrors run_loop but yields StreamChunks instead of AgentResponse.
9636    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            // One-shot context initialization (mirrors run_loop)
9646            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            // Clear stale disambiguation context from previous turns.
9664            self.clear_disambiguation_context();
9665
9666            // Disambiguation check (before input processing)
9667            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                        // Check if this clarification was triggered by a skill-level override.
9753                        // Re-run skill disambiguation to verify all required_clarity fields
9754                        // are present before executing.
9755                        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                        // Forward to internal stream with enriched input
9772                        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                        // Same skill-id re-check for best-guess path
9782                        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                                // Topic switch: forward to internal stream with fresh input.
9889                                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                                // Pure abandonment: generate a brief acknowledgment.
9897                                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            // No disambiguation or Clear result — proceed with internal stream
9929            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    /// Build a context map with language hints from context_manager for HITL message localization.
10000    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    /// Send a HITL check result through the approval flow and return the full ApprovalResult.
10011    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        // Resolve timeout action
10035        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    /// Execute multiple tools in parallel
10086    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    /// Stream a chat response with real-time updates
10139    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
10199//
10200// Render a concurrent input template using direct minijinja.
10201// Same approach as pipeline's render_stage_template so variables are top-level.
10202//
10203// Available variables:
10204//   {{ user_input }}    - the user's actual message
10205//   {{ context.<key> }} - values from the context manager
10206//
10207/// Builds safe runtime tags for background maintenance lifecycle events.
10208fn 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
10250/// Records a background maintenance lifecycle event when observability is enabled.
10251fn 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    // Expose context manager values under {{ context.<key> }}.
10338    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    /// Test hook that counts completed responses.
10368    struct ResponseCountingHooks {
10369        responses: Arc<std::sync::atomic::AtomicUsize>,
10370    }
10371
10372    /// Test tool that returns the execution context it received.
10373    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    /// Test tool that stays active long enough for runtime cancellation.
10423    struct SlowTool;
10424
10425    /// Test tool that fails once and must not be retried for writes.
10426    struct FlakyWriteTool {
10427        calls: Arc<std::sync::atomic::AtomicUsize>,
10428    }
10429
10430    /// Test tool that tracks concurrent execution on one path.
10431    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    // Basic YAML → Build → Chat flow
10606    #[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    // Multi-turn conversation
10621    #[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        // Verify memory accumulated messages
10644        let messages = agent.memory.get_messages(None).await.unwrap();
10645        // 3 user + 3 assistant = 6 messages
10646        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    // Tool execution in chat flow
11139    #[tokio::test]
11140    async fn test_integration_tool_execution() {
11141        // Mock LLM that returns a tool call then a final answer
11142        let mock = mock_with_responses(vec![
11143            // First response: tool call
11144            r#"I'll calculate that for you.
11145[TOOL_CALL: {"name": "calculator", "arguments": {"expression": "2+2"}}]"#,
11146            // After tool result: final answer
11147            "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        // The agent should eventually produce a response
11163        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    // State machine transitions
12215    #[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?", // greeting response
12234            "1",                        // transition evaluator picks first (index 0)
12235            "I'll help you with that.", // support response
12236        ]);
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        // After transition evaluation, state may or may not have changed
12243        // depending on mock evaluator response - the key is that it doesn't crash
12244    }
12245
12246    // State on_enter/on_exit actions
12247    #[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        // The transition evaluator will pick the first transition (index 0)
12271        let mock = mock_with_responses(vec![
12272            "Processing step 1.",
12273            "0", // transition evaluator response: select first transition
12274        ]);
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        // Manually transition to test on_enter/on_exit
12281        agent.transition_to("step2").await.unwrap();
12282
12283        assert_eq!(agent.current_state(), Some("step2".to_string()));
12284
12285        // Verify context was set by on_exit and on_enter actions
12286        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    // Process pipeline transforms input
12292    #[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        // Verify the LLM received the normalized input (trimmed + collapsed whitespace)
12311        let history = mock.call_history();
12312        assert!(!history.is_empty());
12313        // The user message in LLM call should be normalized
12314        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    // ═══════════════════════════════════════════════════════════
12324    // Integration Test 2.1.7: Memory compression triggers
12325    // ═══════════════════════════════════════════════════════════
12326    #[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        // Provide enough responses for compression to trigger
12339        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        // Send enough messages to trigger compression
12345        for i in 0..6 {
12346            let _ = agent.chat(&format!("Message {}", i)).await.unwrap();
12347        }
12348
12349        // Memory should have compressed - verify it didn't crash
12350        // and that messages are bounded
12351        let messages = agent.memory.get_messages(None).await.unwrap();
12352        // With compress_threshold=5 and max_recent_messages=3,
12353        // after 6 turns (12 messages), compression should have run
12354        assert!(messages.len() <= 12); // At most all messages if no compression, fewer if compressed
12355    }
12356
12357    // YAML with multiple LLMs
12358    #[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    // Agent reset clears state
12377    #[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); // user + assistant
12389
12390        agent.reset().await.unwrap();
12391        let messages = agent.memory.get_messages(None).await.unwrap();
12392        assert_eq!(messages.len(), 0);
12393    }
12394
12395    // Process pipeline rejects input
12396    #[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        // Rejected input should produce a rejection response, not call LLM
12430        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, // rejection message is typically short
12436            "Expected rejection response, got: {}",
12437            response.content
12438        );
12439    }
12440
12441    // LLM fallback: primary fails, fallback LLM responds
12442    #[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    // LLM fallback: primary fails, static message returned
12477    #[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    // Tool skip: tool fails, on_failure: skip absorbs the error
12509    #[tokio::test]
12510    async fn test_tool_failure_skip() {
12511        use ai_agents_recovery::{
12512            ErrorRecoveryConfig, ToolFailureAction, ToolRecoveryConfig, ToolRetryConfig,
12513        };
12514
12515        // LLM requests a nonexistent tool, then responds after seeing the skip result
12516        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        // The tool will fail (not found), but on_failure: skip absorbs the error
12540        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}