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, HashSet};
6use std::future::Future;
7use std::pin::Pin;
8use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
9use std::sync::{Arc, Weak};
10use std::time::Instant;
11use tracing::{debug, error, info, instrument, warn};
12
13const DISAMBIGUATION_STATE_GENERATION_KEY: &str = "_runtime.disambiguation_state_generation";
14
15/// Shared lock table used to serialize side-effecting tool calls by canonical resource.
16pub(crate) type ToolResourceLocks = Arc<RwLock<HashMap<String, Weak<tokio::sync::Mutex<()>>>>>;
17
18//
19// Owns every lock acquired for one tool call and removes dead weak entries after release.
20//
21struct ToolResourceGuards {
22    guards: Vec<tokio::sync::OwnedMutexGuard<()>>,
23    locks: ToolResourceLocks,
24}
25
26#[derive(Clone)]
27struct StoredSessionRestore {
28    snapshot: AgentSnapshot,
29    metadata: Option<ai_agents_core::SessionMetadata>,
30}
31
32struct RuntimeSessionRestorePoint {
33    snapshot: AgentSnapshot,
34    metadata: ai_agents_core::SessionMetadata,
35    actor_id: Option<String>,
36    session_id: Option<String>,
37}
38
39impl Drop for ToolResourceGuards {
40    fn drop(&mut self) {
41        self.guards.clear();
42        self.locks.write().retain(|_, lock| lock.strong_count() > 0);
43    }
44}
45
46//
47// Captures runtime controls under one read guard so policy and scope belong to the same generation.
48//
49#[derive(Clone)]
50struct RuntimeSafetySnapshot {
51    version: u64,
52    emergency_deny: bool,
53    tool_security: ToolSecurityEngine,
54    tool_scope_override: Option<Vec<String>>,
55}
56
57//
58// Records the generations used by the final authorization decision.
59//
60#[derive(Clone, Copy)]
61struct ToolDecisionVersions {
62    policy: u64,
63    registry: u64,
64    runtime_control: u64,
65    state: Option<u64>,
66}
67
68//
69// Carries deterministic effective tools with the state generation that authorized their scopes.
70//
71struct AvailableToolIdsSnapshot {
72    tool_ids: Vec<String>,
73    state_generation: Option<u64>,
74}
75
76//
77// Binds human approval to the reviewed action and exact tool implementation.
78//
79#[derive(Clone)]
80struct ToolApprovalBinding {
81    canonical_id: String,
82    arguments: Value,
83    confirmation_required: bool,
84    policy_version: u64,
85    runtime_control_version: u64,
86    state_generation: Option<u64>,
87    reviewed_tool: Arc<dyn ai_agents_core::Tool>,
88}
89
90//
91// A later plain approval must not erase arguments modified by an earlier approval stage.
92//
93fn merge_approved_record(record: &mut Option<ToolApprovalRecord>) {
94    if record
95        .as_ref()
96        .is_some_and(|record| matches!(record.status, ToolApprovalStatus::Modified))
97    {
98        return;
99    }
100    *record = Some(ToolApprovalRecord {
101        status: ToolApprovalStatus::Approved,
102        reason: None,
103        modified_arguments: None,
104    });
105}
106
107impl ToolApprovalBinding {
108    /// Returns true when final authorization no longer matches the reviewed action.
109    fn is_stale(
110        &self,
111        canonical_id: &str,
112        arguments: &Value,
113        confirmation_required: bool,
114        versions: ToolDecisionVersions,
115        resolved_tool: &Arc<dyn ai_agents_core::Tool>,
116    ) -> bool {
117        self.canonical_id != canonical_id
118            || self.arguments != *arguments
119            || self.confirmation_required != confirmation_required
120            || self.policy_version != versions.policy
121            || self.runtime_control_version != versions.runtime_control
122            || self.state_generation != versions.state
123            || !Arc::ptr_eq(&self.reviewed_tool, resolved_tool)
124    }
125}
126
127use crate::turn_context::{current_turn_actor_context, scope_actor_context};
128
129use ai_agents_context::{ContextManager, ContextProvider, TemplateRenderer};
130use ai_agents_core::traits::storage::StorageCapability;
131use ai_agents_core::{
132    AgentError, AgentSnapshot, AgentStorage, ChatMessage, FinishReason, LLMError, LLMProvider,
133    LLMResponse, LLMToolDefinition, LLMToolRequest, PermissionOutcome, Result, ToolActorContext,
134    ToolApprovalRecord, ToolApprovalStatus, ToolCallSource, ToolCancellationToken, ToolChoice,
135    ToolExecutionContext, ToolExecutionRecord, ToolExecutionRequest, ToolInvoker,
136    ToolPolicyDecisionRecord, ToolResult,
137};
138use ai_agents_disambiguation::{
139    ClarificationObserver, ClarificationParseFuture, ClarificationQuestionFuture,
140    ConfirmationParseFuture, DisambiguationConfig, DisambiguationContext, DisambiguationManager,
141    DisambiguationResult,
142};
143use ai_agents_hitl::{
144    ApprovalHandler, ApprovalResolvedOutcome, ApprovalResult, ApprovalTrigger, HITLCheckResult,
145    HITLEngine, RejectAllHandler, TimeoutAction,
146};
147use ai_agents_hooks::{AgentHooks, NoopHooks};
148use ai_agents_llm::LLMRegistry;
149use ai_agents_memory::{
150    CompressResult, EvictionReason, Memory, MemoryBudgetEvent, MemoryCompressEvent,
151    MemoryEvictEvent, MemoryTokenBudget, OverflowStrategy,
152};
153use ai_agents_observability::{
154    EventStatus, EventType, ObservabilityManager, ObservationPurpose, SpanContext,
155    current_observation_context, new_session_id as new_observation_session_id,
156    resolve_language_from_context, with_observation_context, with_observation_purpose,
157};
158use ai_agents_process::{
159    ProcessData, ProcessProcessor, ProcessPurposeHint, ProcessStageFuture, ProcessStageObserver,
160};
161use ai_agents_reasoning::{
162    CriterionResult, EvaluationResult, Plan, PlanAction, PlanStatus, PlanStep, ReasoningConfig,
163    ReasoningMetadata, ReasoningMode, ReasoningOutput, ReflectionAttempt, ReflectionConfig,
164    ReflectionMetadata, StepFailureAction,
165};
166use ai_agents_recovery::{
167    ByRoleFilter, ContextOverflowAction, FilterConfig, IntoClassifiedError, KeepRecentFilter,
168    LLMFailureAction, MessageFilter, RecoveryManager, SkipPatternFilter, ToolFailureAction,
169};
170use ai_agents_relationships::RelationshipManager;
171use ai_agents_skills::{SkillDefinition, SkillExecutor, SkillRouter};
172use ai_agents_state::{
173    PromptMode, StateAction, StateMachine, StateMachineSnapshot, StateTransitionEvent, Transition,
174    TransitionContext, TransitionEvaluator, TransitionTiming, evaluate_guard,
175};
176use ai_agents_storage::{StorageConfig as StorageStorageConfig, create_storage};
177use ai_agents_tools::{
178    CommandRunner, ConditionEvaluator, DiagnosticsProvider, EvaluationContext, LLMGetter,
179    QuestionHandler, SecurityCheckResult, TodoItem, ToolCallRecord, ToolRegistry,
180    ToolSecurityConfig, ToolSecurityEngine,
181};
182
183use super::{
184    Agent, AgentInfo, AgentResponse, AgentStreamEvent, ParallelToolsConfig, StreamChunk,
185    StreamingConfig, ToolCall,
186};
187use crate::optimization::{
188    AwaitBeforeNextTurn, BackgroundMaintenanceQueue, BackgroundOverflowPolicy, MainResponseDraft,
189    MaintenanceMode, MaintenanceSequenceKey, RuntimeBranch, RuntimeBranchResult,
190    RuntimeBranchStatus, RuntimeCommitBehavior, RuntimeConfig, RuntimeOptimizationKind,
191    RuntimeTaskPriority, RuntimeTaskPurpose, ScheduledBranchSet, SkillCandidate,
192    StreamingDraftResult, TransitionCandidate, TurnBranchScheduler, TurnOptimizationContext,
193};
194use crate::spec::StorageConfig;
195
196/// Outcome of processing tool calls within the agent loop.
197enum ToolCallOutcome {
198    /// Tools executed successfully, continue the LLM loop for the next iteration.
199    Continue,
200    /// A state transition fired during tool call handling, continue the loop.
201    TransitionFired,
202    /// HITL rejected a tool call, return this response immediately.
203    Rejected(AgentResponse),
204}
205
206#[derive(Clone)]
207struct MainToolProtocol {
208    choice: Option<ToolChoice>,
209    tool_ids: Vec<String>,
210    definitions: Vec<LLMToolDefinition>,
211}
212
213struct MainProviderResponse {
214    response: LLMResponse,
215    used_native_tools: bool,
216}
217
218//
219// Carries one committed model answer through shared output processing and root-turn finalization.
220//
221struct CommittedTextResponse<'a> {
222    processed_input: &'a str,
223    input_context: &'a HashMap<String, Value>,
224    answer: String,
225    reasoning_mode: ReasoningMode,
226    auto_detected: bool,
227    iterations: u32,
228    thinking_content: Option<String>,
229    all_tool_calls: Vec<ToolCall>,
230}
231
232//
233// Carries the finalized response fields into metadata assembly without changing their ownership.
234//
235struct AgentResponseParts {
236    content: String,
237    all_tool_calls: Vec<ToolCall>,
238    reasoning_mode: ReasoningMode,
239    auto_detected: bool,
240    iterations: u32,
241    thinking: Option<String>,
242    reflection_metadata: Option<ReflectionMetadata>,
243}
244
245//
246// Owns the finalized response paired with a legacy Done chunk without changing provisional chunk timing.
247//
248type RuntimeStreamTerminalSlot = Arc<RwLock<Option<AgentResponse>>>;
249
250//
251// Creates isolated terminal ownership for one public stream.
252//
253fn new_runtime_stream_terminal_slot() -> RuntimeStreamTerminalSlot {
254    Arc::new(RwLock::new(None))
255}
256
257//
258// Records the exact response that already completed root-turn finalization before Done is emitted.
259//
260fn record_runtime_stream_final(slot: &RuntimeStreamTerminalSlot, response: AgentResponse) {
261    *slot.write() = Some(response);
262}
263
264#[derive(Clone, Copy)]
265struct DisambiguationOwnership {
266    epoch: u64,
267    state_generation: Option<u64>,
268}
269
270/// Outcome of skill routing — used by `try_skill_route`.
271enum SkillRouteResult {
272    /// No skill matched, continue to normal LLM chat.
273    NoMatch,
274    /// Skill executed successfully.
275    Response { skill_id: String, content: String },
276    /// Skill matched but needs disambiguation first or returned a terminal disambiguation response.
277    NeedsClarification {
278        response: AgentResponse,
279        ownership: Option<DisambiguationOwnership>,
280    },
281}
282
283/// Result of response-independent parallel transition selection.
284enum ParallelTransitionSelection {
285    /// A transition matched and can be committed before the old-state response.
286    Candidate(TransitionCandidate),
287    /// All eligible transition checks ran and none matched.
288    NoMatch,
289    /// LLM-based route evaluation needed speculative capacity that was unavailable.
290    ReservationExhausted,
291}
292
293/// Outcome of post_loop_processing - drives the caller's next step.
294enum PostLoopResult {
295    /// No transition fired. Content is the LLM response for this turn.
296    NoTransition(String),
297    /// Transition fired. Content is from plain post-transition re-generation.
298    Transitioned(String),
299    /// Transition fired into a state that requires full dispatch.
300    /// Caller re-enters run_loop_internal to apply the correct handler.
301    NeedsRedispatch,
302}
303
304struct StateTransitionReservation<'a> {
305    reserved: &'a AtomicBool,
306}
307
308impl Drop for StateTransitionReservation<'_> {
309    fn drop(&mut self) {
310        self.reserved.store(false, Ordering::SeqCst);
311    }
312}
313
314struct RootTurnCleanup<'a> {
315    agent: &'a RuntimeAgent,
316}
317
318impl<'a> RootTurnCleanup<'a> {
319    fn new(agent: &'a RuntimeAgent) -> Self {
320        Self { agent }
321    }
322}
323
324impl Drop for RootTurnCleanup<'_> {
325    fn drop(&mut self) {
326        self.agent.end_root_turn();
327    }
328}
329
330/// Host-owned runtime control state shared with active agents.
331#[derive(Debug)]
332struct RuntimeControlState {
333    /// Serializes control mutations with exact runtime safety snapshots.
334    snapshot_guard: RwLock<()>,
335    /// Monotonic version for runtime-control snapshots.
336    version: AtomicU64,
337    /// Emergency switch that denies future calls and is shared with active tool contexts.
338    emergency_deny: Arc<AtomicBool>,
339    /// Optional live replacement for tool security policy.
340    tool_security_override: RwLock<Option<ToolSecurityEngine>>,
341    /// Optional live narrowing scope applied after the runtime's declared grant.
342    tool_scope_override: RwLock<Option<Vec<String>>>,
343}
344
345impl Default for RuntimeControlState {
346    fn default() -> Self {
347        Self {
348            snapshot_guard: RwLock::new(()),
349            version: AtomicU64::new(1),
350            emergency_deny: Arc::new(AtomicBool::new(false)),
351            tool_security_override: RwLock::new(None),
352            tool_scope_override: RwLock::new(None),
353        }
354    }
355}
356
357/// Host-only handle for live runtime safety controls.
358#[derive(Clone)]
359pub struct RuntimeControlHandle {
360    state: Arc<RuntimeControlState>,
361}
362
363impl RuntimeControlHandle {
364    /// Returns the current runtime-control version.
365    pub fn version(&self) -> u64 {
366        self.state.version.load(Ordering::SeqCst)
367    }
368
369    fn bump(&self) -> u64 {
370        self.state.version.fetch_add(1, Ordering::SeqCst) + 1
371    }
372
373    /// Overrides tool security for later tool calls and panics if host policy is invalid.
374    pub fn set_tool_security(&self, config: ToolSecurityConfig) -> u64 {
375        self.try_set_tool_security(config)
376            .expect("invalid tool security configuration")
377    }
378
379    /// Validates replacement policy before changing the runtime-control generation or active snapshot.
380    pub fn try_set_tool_security(&self, config: ToolSecurityConfig) -> Result<u64> {
381        config.validate()?;
382        let _guard = self.state.snapshot_guard.write();
383        let generation = self.bump();
384        *self.state.tool_security_override.write() = Some(
385            ToolSecurityEngine::new_with_policy_version(config, generation),
386        );
387        Ok(generation)
388    }
389
390    /// Clears the live tool security override.
391    pub fn clear_tool_security_override(&self) -> u64 {
392        let _guard = self.state.snapshot_guard.write();
393        *self.state.tool_security_override.write() = None;
394        self.bump()
395    }
396
397    /// Narrows the runtime's declared tool grant for later calls without adding authority.
398    pub fn set_tool_scope(&self, tool_ids: Vec<String>) -> u64 {
399        let _guard = self.state.snapshot_guard.write();
400        *self.state.tool_scope_override.write() = Some(tool_ids);
401        self.bump()
402    }
403
404    /// Clears live narrowing so later calls return to the runtime's declared grant.
405    pub fn clear_tool_scope_override(&self) -> u64 {
406        let _guard = self.state.snapshot_guard.write();
407        *self.state.tool_scope_override.write() = None;
408        self.bump()
409    }
410
411    /// Enables or disables emergency denial for future tool calls.
412    pub fn set_emergency_deny(&self, enabled: bool) -> u64 {
413        let _guard = self.state.snapshot_guard.write();
414        self.state.emergency_deny.store(enabled, Ordering::SeqCst);
415        self.bump()
416    }
417
418    /// Denies future calls and asks active cancellable work to stop.
419    pub fn cancel_all(&self) -> u64 {
420        self.set_emergency_deny(true)
421    }
422}
423
424pub struct RuntimeAgent {
425    info: AgentInfo,
426    llm_registry: Arc<LLMRegistry>,
427    memory: Arc<dyn Memory>,
428    tools: Arc<ToolRegistry>,
429    skills: Vec<SkillDefinition>,
430    skill_router: Option<SkillRouter>,
431    skill_executor: Option<SkillExecutor>,
432    base_system_prompt: String,
433    max_iterations: u32,
434    iteration_count: RwLock<u32>,
435    max_context_tokens: u32,
436    memory_token_budget: Option<MemoryTokenBudget>,
437    recovery_manager: RecoveryManager,
438    tool_security: ToolSecurityEngine,
439    process_processor: Option<ProcessProcessor>,
440    message_filters: RwLock<HashMap<String, Arc<dyn MessageFilter>>>,
441    state_machine: Option<Arc<StateMachine>>,
442    transition_evaluator: Option<Arc<dyn TransitionEvaluator>>,
443    context_manager: Arc<ContextManager>,
444    template_renderer: TemplateRenderer,
445    tool_call_history: RwLock<Vec<ToolCallRecord>>,
446    parallel_tools: ParallelToolsConfig,
447    streaming: StreamingConfig,
448    hooks: Arc<dyn AgentHooks>,
449    hitl_engine: Option<HITLEngine>,
450    approval_handler: Arc<dyn ApprovalHandler>,
451    storage_config: StorageConfig,
452    storage: RwLock<Option<Arc<dyn AgentStorage>>>,
453    storage_init: tokio::sync::Mutex<()>,
454    reasoning_config: ReasoningConfig,
455    reflection_config: ReflectionConfig,
456    disambiguation_manager: Option<DisambiguationManager>,
457    /// Generation invalidating confirmation work across reset and state changes.
458    disambiguation_epoch: AtomicU64,
459    /// Serializes confirmation redispatch admission with reset and state mutation.
460    disambiguation_admission: tokio::sync::RwLock<()>,
461    /// Reserves one runtime-owned transition before its exit actions can produce side effects.
462    state_transition_reserved: AtomicBool,
463    /// Structured persona manager for identity, evolution, and secrets.
464    persona_manager: Option<Arc<ai_agents_persona::PersonaManager>>,
465    /// Skill ID that triggered the current pending disambiguation.
466    /// Set by try_skill_route() when skill-level disambiguation triggers clarification.
467    /// Read by run_loop() when clarification resolves to route directly to the skill.
468    pending_skill_id: RwLock<Option<String>>,
469    current_plan: RwLock<Option<Plan>>,
470    /// Tool IDs declared in the top-level `tools:` spec.
471    declared_tool_ids: Option<Vec<String>>,
472    /// Whether the context manager has been initialized (defaults loaded, env resolved, etc.)
473    context_initialized: AtomicBool,
474    /// Spawner for dynamic agent creation (set when YAML has a spawner: section).
475    spawner: Option<Arc<crate::spawner::AgentSpawner>>,
476    /// Registry tracking spawned agents (set when YAML has a spawner: section).
477    spawner_registry: Option<Arc<crate::spawner::AgentRegistry>>,
478    /// Re-dispatch depth for post-transition full dispatch.
479    /// 0 = not re-dispatching. > 0 = user message already in memory, skip re-adding.
480    redispatch_depth: RwLock<u32>,
481    /// Active optimized turn context used to keep root lifecycle state in one place.
482    active_turn_context: RwLock<Option<TurnOptimizationContext>>,
483    /// Tracks whether the root turn already wrote the processed user message.
484    root_user_message_committed: AtomicBool,
485    /// Current actor ID for cross-session memory.
486    actor_id: RwLock<Option<String>>,
487    /// Fact store for managing per-actor extracted facts.
488    fact_store: RwLock<Option<Arc<ai_agents_facts::FactStore>>>,
489    /// Fact extractor for LLM-based fact extraction.
490    /// None when actor_memory is enabled without facts.enabled.
491    fact_extractor: RwLock<Option<Arc<dyn ai_agents_facts::FactExtractor>>>,
492    /// Cached actor facts keyed by actor ID so concurrent or alternating turns do not overwrite one another.
493    actor_facts_cache: Arc<RwLock<HashMap<String, Vec<ai_agents_core::KeyFact>>>>,
494    /// Number of messages since last fact extraction.
495    messages_since_extraction: Arc<RwLock<usize>>,
496    /// Actor memory configuration.
497    actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
498    /// Facts configuration.
499    facts_config: Option<ai_agents_facts::FactsConfig>,
500    /// Session-scoped metadata (tags, ttl, actor roster).
501    session_metadata: RwLock<ai_agents_core::SessionMetadata>,
502    /// Session id currently bound to this runtime instance.
503    current_session_id: RwLock<Option<String>>,
504    /// Relationship manager for actor-scoped social memory.
505    relationship_manager: Option<Arc<RelationshipManager>>,
506    /// Observability manager for traces, metrics, reports, and exports.
507    observability_manager: Option<Arc<ObservabilityManager>>,
508    /// Runtime optimization and maintenance policy.
509    runtime_config: RuntimeConfig,
510    /// Queue for background maintenance tasks.
511    background_maintenance: Arc<BackgroundMaintenanceQueue>,
512    /// Cross-tool locks for side-effecting calls that target the same resource.
513    resource_locks: ToolResourceLocks,
514    /// Host-only runtime control state.
515    runtime_control: Arc<RuntimeControlState>,
516}
517
518impl std::fmt::Debug for RuntimeAgent {
519    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
520        f.debug_struct("RuntimeAgent")
521            .field("info", &self.info)
522            .field("base_system_prompt", &self.base_system_prompt)
523            .field("max_iterations", &self.max_iterations)
524            .field("skills_count", &self.skills.len())
525            .field("max_context_tokens", &self.max_context_tokens)
526            .field("has_state_machine", &self.state_machine.is_some())
527            .field("parallel_tools", &self.parallel_tools)
528            .field("streaming", &self.streaming)
529            .field("has_hooks", &true)
530            .field("has_hitl", &self.hitl_engine.is_some())
531            .field("storage_type", &self.storage_config.storage_type())
532            .field("reasoning_mode", &self.reasoning_config.mode)
533            .field("reflection_enabled", &self.reflection_config.enabled)
534            .field("declared_tool_ids", &self.declared_tool_ids)
535            .field("has_persona", &self.persona_manager.is_some())
536            .field("has_observability", &self.observability_manager.is_some())
537            .finish_non_exhaustive()
538    }
539}
540
541struct ObservabilityClarificationObserver;
542
543impl ClarificationObserver for ObservabilityClarificationObserver {
544    /// Scopes clarification question generation as disambiguation_clarification.
545    fn observe_question<'a>(
546        &'a self,
547        future: ClarificationQuestionFuture<'a>,
548    ) -> ClarificationQuestionFuture<'a> {
549        Box::pin(async move {
550            with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
551        })
552    }
553
554    /// Scopes clarification response parsing as disambiguation_clarification.
555    fn observe_parse<'a>(
556        &'a self,
557        future: ClarificationParseFuture<'a>,
558    ) -> ClarificationParseFuture<'a> {
559        Box::pin(async move {
560            with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
561        })
562    }
563
564    /// Scopes semantic confirmation parsing as disambiguation_clarification.
565    fn observe_confirmation_parse<'a>(
566        &'a self,
567        future: ConfirmationParseFuture<'a>,
568    ) -> ConfirmationParseFuture<'a> {
569        Box::pin(async move {
570            with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
571        })
572    }
573}
574
575struct ObservabilityProcessStageObserver;
576
577impl ProcessStageObserver for ObservabilityProcessStageObserver {
578    /// Scopes one process stage using the purpose implied by its stage type.
579    fn observe<'a>(
580        &'a self,
581        hint: ProcessPurposeHint,
582        future: ProcessStageFuture<'a>,
583    ) -> ProcessStageFuture<'a> {
584        Box::pin(async move {
585            with_observation_purpose(observation_purpose_for_process(hint), future).await
586        })
587    }
588}
589
590struct RegistryLLMGetter {
591    registry: Arc<LLMRegistry>,
592}
593
594impl LLMGetter for RegistryLLMGetter {
595    fn get_llm(&self, alias: &str) -> Option<Arc<dyn LLMProvider>> {
596        self.registry.get(alias).ok()
597    }
598}
599
600impl RuntimeAgent {
601    #[allow(clippy::too_many_arguments)]
602    pub fn new(
603        info: AgentInfo,
604        llm_registry: Arc<LLMRegistry>,
605        memory: Arc<dyn Memory>,
606        tools: Arc<ToolRegistry>,
607        skills: Vec<SkillDefinition>,
608        system_prompt: String,
609        max_iterations: u32,
610    ) -> Self {
611        let (skill_router, skill_executor) = if !skills.is_empty() {
612            let router_llm = llm_registry.router().ok();
613            let router = router_llm.map(|llm| SkillRouter::new(llm, skills.clone()));
614            let executor = SkillExecutor::new(llm_registry.clone(), tools.clone());
615            (router, Some(executor))
616        } else {
617            (None, None)
618        };
619
620        let context_manager =
621            ContextManager::new(HashMap::new(), info.name.clone(), info.version.clone());
622
623        Self {
624            info,
625            llm_registry,
626            memory,
627            tools,
628            skills,
629            skill_router,
630            skill_executor,
631            base_system_prompt: system_prompt,
632            max_iterations,
633            iteration_count: RwLock::new(0),
634            max_context_tokens: 128000,
635            memory_token_budget: None,
636            recovery_manager: RecoveryManager::default(),
637            tool_security: ToolSecurityEngine::default(),
638            process_processor: None,
639            message_filters: RwLock::new(HashMap::new()),
640            state_machine: None,
641            transition_evaluator: None,
642            context_manager: Arc::new(context_manager),
643            template_renderer: TemplateRenderer::new(),
644            tool_call_history: RwLock::new(Vec::new()),
645            parallel_tools: ParallelToolsConfig::default(),
646            streaming: StreamingConfig::default(),
647            hooks: Arc::new(NoopHooks),
648            hitl_engine: None,
649            approval_handler: Arc::new(RejectAllHandler::new()),
650            storage_config: StorageConfig::default(),
651            storage: RwLock::new(None),
652            storage_init: tokio::sync::Mutex::new(()),
653            reasoning_config: ReasoningConfig::default(),
654            reflection_config: ReflectionConfig::default(),
655            disambiguation_manager: None,
656            disambiguation_epoch: AtomicU64::new(0),
657            disambiguation_admission: tokio::sync::RwLock::new(()),
658            state_transition_reserved: AtomicBool::new(false),
659            persona_manager: None,
660            pending_skill_id: RwLock::new(None),
661            current_plan: RwLock::new(None),
662            declared_tool_ids: None,
663            context_initialized: AtomicBool::new(false),
664            spawner: None,
665            spawner_registry: None,
666            redispatch_depth: RwLock::new(0),
667            active_turn_context: RwLock::new(None),
668            root_user_message_committed: AtomicBool::new(false),
669            actor_id: RwLock::new(None),
670            fact_store: RwLock::new(None),
671            fact_extractor: RwLock::new(None),
672            actor_facts_cache: Arc::new(RwLock::new(HashMap::new())),
673            messages_since_extraction: Arc::new(RwLock::new(0)),
674            actor_memory_config: None,
675            facts_config: None,
676            session_metadata: RwLock::new(ai_agents_core::SessionMetadata::default()),
677            current_session_id: RwLock::new(None),
678            relationship_manager: None,
679            observability_manager: None,
680            runtime_config: RuntimeConfig::default(),
681            background_maintenance: Arc::new(BackgroundMaintenanceQueue::default()),
682            resource_locks: new_tool_resource_locks(),
683            runtime_control: Arc::new(RuntimeControlState::default()),
684        }
685    }
686
687    pub fn with_declared_tool_ids(mut self, ids: Option<Vec<String>>) -> Self {
688        self.declared_tool_ids = ids;
689        self
690    }
691
692    pub fn with_storage_config(mut self, config: StorageConfig) -> Self {
693        self.storage_config = config;
694        self
695    }
696
697    pub fn with_storage(self, storage: Arc<dyn AgentStorage>) -> Self {
698        *self.storage.write() = Some(storage);
699        self
700    }
701
702    pub(crate) fn with_shared_resource_locks(mut self, locks: ToolResourceLocks) -> Self {
703        self.resource_locks = locks;
704        self
705    }
706
707    pub fn with_reasoning(mut self, config: ReasoningConfig) -> Self {
708        self.reasoning_config = config;
709        self
710    }
711
712    pub fn with_reflection(mut self, config: ReflectionConfig) -> Self {
713        self.reflection_config = config;
714        self
715    }
716
717    /// Attach a relationship manager configured by the builder or host application.
718    pub fn with_relationships(mut self, manager: Arc<RelationshipManager>) -> Self {
719        self.relationship_manager = Some(manager);
720        self
721    }
722
723    /// Attach a shared observability manager for traces, metrics, reports, and exports.
724    pub fn with_observability(mut self, manager: Arc<ObservabilityManager>) -> Self {
725        self.observability_manager = Some(manager);
726        self
727    }
728
729    /// Attach runtime optimization policy and resize the background queue.
730    pub fn with_runtime_config(mut self, config: RuntimeConfig) -> Self {
731        let max_tasks = config.optimization.post_turn.max_background_tasks;
732        self.background_maintenance = Arc::new(BackgroundMaintenanceQueue::new(max_tasks));
733        self.runtime_config = config;
734        self
735    }
736
737    /// Returns the runtime optimization policy.
738    pub fn runtime_config(&self) -> &RuntimeConfig {
739        &self.runtime_config
740    }
741
742    /// Wait for all background maintenance tasks to finish.
743    pub async fn flush_background_tasks(&self) -> Result<()> {
744        self.background_maintenance.flush_all().await
745    }
746
747    /// Wait for background maintenance associated with one actor to finish.
748    pub async fn flush_background_tasks_for_actor(&self, actor_id: &str) -> Result<()> {
749        self.background_maintenance.flush_scope(actor_id).await
750    }
751
752    /// Wait for background maintenance associated with one task kind to finish.
753    pub async fn flush_background_tasks_for_purpose(
754        &self,
755        purpose: RuntimeTaskPurpose,
756    ) -> Result<()> {
757        self.background_maintenance.flush_purpose(purpose).await
758    }
759
760    /// Wait for background maintenance associated with one actor and task kind to finish.
761    pub async fn flush_background_tasks_for_actor_purpose(
762        &self,
763        actor_id: &str,
764        purpose: RuntimeTaskPurpose,
765    ) -> Result<()> {
766        self.background_maintenance
767            .flush_scope_purpose(actor_id, purpose)
768            .await
769    }
770
771    /// Flush background maintenance before a host shuts down the runtime.
772    pub async fn shutdown_background_tasks(&self) -> Result<()> {
773        self.flush_background_tasks().await
774    }
775
776    /// Returns the configured observability manager for report and export access.
777    pub fn observability(&self) -> Option<Arc<ObservabilityManager>> {
778        self.observability_manager.clone()
779    }
780
781    /// Exports observability files after a turn when export settings request it.
782    async fn export_observability_if_configured(&self) {
783        let Some(manager) = self.observability_manager.as_ref() else {
784            return;
785        };
786        let export = &manager.config().export;
787        if !export.write_report && !export.write_raw_events {
788            return;
789        }
790        if let Err(error) = manager.export().await {
791            warn!(error = %error, "Observability export failed");
792        }
793    }
794
795    /// Returns the configured relationship manager, if relationship memory is enabled.
796    pub fn relationship_manager(&self) -> Option<Arc<RelationshipManager>> {
797        self.relationship_manager.clone()
798    }
799
800    fn current_turn_actor_context(&self) -> Option<crate::TurnActorContext> {
801        current_turn_actor_context()
802    }
803
804    fn effective_actor_id(&self) -> Option<String> {
805        self.current_turn_actor_context()
806            .and_then(|ctx| ctx.effective_actor_id().map(|id| id.to_string()))
807            .or_else(|| self.actor_id.read().clone())
808    }
809
810    fn effective_origin_actor_id(&self) -> Option<String> {
811        self.current_turn_actor_context()
812            .and_then(|ctx| ctx.origin_actor_id.clone())
813            .or_else(|| self.actor_id.read().clone())
814    }
815
816    fn record_session_actor_if_needed(&self) {
817        if let Some(actor_id) = self.effective_origin_actor_id() {
818            let mut meta = self.session_metadata.write();
819            meta.actor_id = Some(actor_id.clone());
820            if !meta.actors.iter().any(|a| a == &actor_id) {
821                meta.actors.push(actor_id);
822            }
823        }
824    }
825
826    fn outbound_actor_context(&self) -> crate::TurnActorContext {
827        let mut context = self.current_turn_actor_context().unwrap_or_default();
828        if context.origin_actor_id.is_none() {
829            context.origin_actor_id = self.effective_origin_actor_id();
830        }
831        context.sender_agent_id = Some(self.info.id.clone());
832        context
833    }
834
835    /// Returns the current session ID or creates one for unsaved observed turns.
836    fn observation_session_id(&self) -> Option<String> {
837        let mut current = self.current_session_id.write();
838        if current.is_none() {
839            *current = Some(new_observation_session_id());
840        }
841        current.clone()
842    }
843
844    /// Builds the root or child observation context for a chat entry point.
845    fn build_observation_context(&self, actor_id: Option<String>) -> Option<SpanContext> {
846        let manager = self.observability_manager.as_ref()?;
847        let context = self.build_context_with_overlays();
848        let language = resolve_language_from_context(manager.config(), &context);
849        let context = current_observation_context()
850            .map(|parent| parent.child_for_agent(self.info.id.clone()).with_new_turn())
851            .unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
852        Some(
853            context
854                .with_actor(actor_id.or_else(|| self.effective_actor_id()))
855                .with_session(self.observation_session_id())
856                .with_state(self.current_state())
857                .with_language(Some(language)),
858        )
859    }
860
861    /// Refreshes task-local context with current runtime labels and a purpose.
862    fn current_runtime_observation_context(
863        &self,
864        purpose: ObservationPurpose,
865    ) -> Option<SpanContext> {
866        let manager = self.observability_manager.as_ref()?;
867        let context = self.build_context_with_overlays();
868        let language = resolve_language_from_context(manager.config(), &context);
869        let mut observation = current_observation_context()
870            .unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
871        observation.agent_id = self.info.id.clone();
872        observation.actor_id = self.effective_actor_id();
873        observation.session_id = self.observation_session_id();
874        observation.state = self.current_state();
875        observation.language = Some(language);
876        observation.purpose = purpose;
877        Some(observation)
878    }
879
880    /// Runs a future under a purpose while preserving current trace context.
881    async fn observe_purpose<F, T>(&self, purpose: ObservationPurpose, future: F) -> T
882    where
883        F: Future<Output = T>,
884    {
885        if let Some(context) = self.current_runtime_observation_context(purpose) {
886            with_observation_context(context, future).await
887        } else {
888            future.await
889        }
890    }
891
892    /// Runs chat with actor and observation context while avoiding recursive async types.
893    fn chat_with_actor_context_boxed<'a>(
894        &'a self,
895        input: &'a str,
896        actor_context: crate::TurnActorContext,
897    ) -> Pin<Box<dyn Future<Output = Result<AgentResponse>> + Send + 'a>> {
898        Box::pin(async move {
899            let actor_id = actor_context.effective_actor_id().map(str::to_string);
900            let run = async move {
901                scope_actor_context(
902                    actor_context,
903                    Box::pin(async move { self.run_loop(input).await }),
904                )
905                .await
906            };
907            let result = if let Some(context) = self.build_observation_context(actor_id) {
908                with_observation_context(context, run).await
909            } else {
910                run.await
911            };
912            self.export_observability_if_configured().await;
913            result
914        })
915    }
916
917    /// Run one turn with turn-scoped actor context without mutating the runtime's global actor ID.
918    ///
919    /// The supplied context is available to actor-scoped facts, relationship memory, orchestration, and prompt templates only for the lifetime of this call.
920    pub async fn chat_with_actor_context(
921        &self,
922        input: &str,
923        actor_context: crate::TurnActorContext,
924    ) -> Result<AgentResponse> {
925        self.chat_with_actor_context_boxed(input, actor_context)
926            .await
927    }
928
929    /// Convenience wrapper around [`Self::chat_with_actor_context`] for a turn whose original actor is known up front.
930    pub async fn chat_as_actor(&self, actor_id: &str, input: &str) -> Result<AgentResponse> {
931        let actor_context = crate::TurnActorContext::new().with_origin_actor(actor_id);
932        self.chat_with_actor_context(input, actor_context).await
933    }
934
935    /// Ensure the effective actor's relationship is loaded from storage into the relationship manager.
936    pub async fn load_actor_relationship(&self) -> Result<()> {
937        self.maybe_load_actor_relationship().await;
938        Ok(())
939    }
940
941    /// Manually apply a delta to the effective actor's `agent_to_actor` relationship perspective and persist the updated relationship when storage is configured.
942    pub async fn update_relationship_dimension(
943        &self,
944        dimension: &str,
945        delta: f64,
946        reason: Option<&str>,
947    ) -> Result<ai_agents_relationships::DimensionChange> {
948        self.update_relationship_dimension_for_perspective(
949            ai_agents_relationships::RelationshipPerspective::AgentToActor,
950            dimension,
951            delta,
952            reason,
953        )
954        .await
955    }
956
957    /// Manually apply a delta to a specific relationship perspective for the effective actor.
958    ///
959    /// Use this for two-sided configurations when you need to update `agent_to_actor`, `perceived_actor_to_agent`, or `mutual` explicitly from application logic.
960    pub async fn update_relationship_dimension_for_perspective(
961        &self,
962        perspective: ai_agents_relationships::RelationshipPerspective,
963        dimension: &str,
964        delta: f64,
965        reason: Option<&str>,
966    ) -> Result<ai_agents_relationships::DimensionChange> {
967        let manager = self
968            .relationship_manager
969            .as_ref()
970            .ok_or_else(|| AgentError::Config("Relationship memory is not configured".into()))?;
971        let actor_id = self.effective_actor_id().ok_or_else(|| {
972            AgentError::Config("No actor ID set. Use set_actor_id() first".into())
973        })?;
974        let change = manager.update_dimension_for_perspective(
975            &actor_id,
976            perspective,
977            dimension,
978            delta,
979            1.0,
980            reason.unwrap_or("manual relationship update"),
981        )?;
982        self.persist_actor_relationship(&actor_id).await?;
983        info!(
984            actor_id = %actor_id,
985            perspective = %change.perspective,
986            dimension = %change.dimension,
987            delta = change.delta,
988            current = change.current,
989            "relationship updated manually"
990        );
991        self.hooks
992            .on_relationship_change(&actor_id, std::slice::from_ref(&change))
993            .await;
994        Ok(change)
995    }
996
997    pub fn reasoning_config(&self) -> &ReasoningConfig {
998        &self.reasoning_config
999    }
1000
1001    pub fn reflection_config(&self) -> &ReflectionConfig {
1002        &self.reflection_config
1003    }
1004
1005    /// Set only the actor memory and facts configs without creating the store.
1006    /// The store and extractor are created lazily in init_storage().
1007    pub fn with_facts_config(
1008        mut self,
1009        actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
1010        facts_config: Option<ai_agents_facts::FactsConfig>,
1011    ) -> Self {
1012        self.actor_memory_config = actor_memory_config;
1013        self.facts_config = facts_config;
1014        self
1015    }
1016
1017    /// Configure fact store and optional extractor for actor memory.
1018    /// Pass `None` for `extractor` to load existing facts without running extraction.
1019    pub fn with_facts(
1020        mut self,
1021        store: Arc<ai_agents_facts::FactStore>,
1022        extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>>,
1023        actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
1024        facts_config: Option<ai_agents_facts::FactsConfig>,
1025    ) -> Self {
1026        *self.fact_store.write() = Some(store);
1027        *self.fact_extractor.write() = extractor;
1028        self.actor_memory_config = actor_memory_config;
1029        self.facts_config = facts_config;
1030        self
1031    }
1032
1033    /// Get the fact store for direct fact manipulation.
1034    pub fn fact_store(&self) -> Option<Arc<ai_agents_facts::FactStore>> {
1035        self.fact_store.read().clone()
1036    }
1037
1038    /// Get the current actor ID.
1039    pub fn actor_id(&self) -> Option<String> {
1040        self.actor_id.read().clone()
1041    }
1042
1043    /// Set the current actor ID (player, user, another agent, etc.).
1044    pub fn set_actor_id(&self, actor_id: &str) -> ai_agents_core::Result<()> {
1045        *self.actor_id.write() = Some(actor_id.to_string());
1046        {
1047            let mut meta = self.session_metadata.write();
1048            meta.actor_id = Some(actor_id.to_string());
1049            if !meta.actors.iter().any(|a| a == actor_id) {
1050                meta.actors.push(actor_id.to_string());
1051            }
1052        }
1053        Ok(())
1054    }
1055
1056    /// Clear the current actor binding while retaining the session actor roster.
1057    pub fn clear_actor_id(&self) {
1058        *self.actor_id.write() = None;
1059        self.session_metadata.write().actor_id = None;
1060    }
1061
1062    /// Set the current actor ID. Convenience wrapper around set_actor_id.
1063    pub fn set_user_id(&self, user_id: &str) -> ai_agents_core::Result<()> {
1064        self.set_actor_id(user_id)
1065    }
1066
1067    /// Load facts for the current actor from storage and cache them for prompt injection.
1068    pub async fn load_actor_memory(&self) -> ai_agents_core::Result<()> {
1069        let actor_id = match self.effective_actor_id() {
1070            Some(id) => id,
1071            None => return Ok(()),
1072        };
1073
1074        let store_opt = self.fact_store.read().clone();
1075        if let Some(store) = store_opt {
1076            let facts = store.get_facts(&actor_id).await?;
1077            let count = facts.len();
1078            self.actor_facts_cache
1079                .write()
1080                .insert(actor_id.clone(), facts);
1081            self.hooks.on_actor_memory_loaded(&actor_id, count).await;
1082            tracing::debug!("loaded {} facts for actor {}", count, actor_id);
1083        }
1084
1085        Ok(())
1086    }
1087
1088    /// Load actor memory only when the effective actor has no cached facts yet.
1089    async fn maybe_load_actor_memory(&self) {
1090        let Some(actor_id) = self.effective_actor_id() else {
1091            return;
1092        };
1093        if self.actor_facts_cache.read().contains_key(&actor_id) {
1094            return;
1095        }
1096        let _ = self.load_actor_memory().await;
1097    }
1098
1099    /// Pre-turn lifecycle shared by streaming and non-streaming paths.
1100    async fn pre_turn_session_lifecycle(&self) {
1101        if *self.redispatch_depth.read() > 0 {
1102            return;
1103        }
1104        self.resolve_actor_id_from_context();
1105        self.await_background_before_next_turn().await;
1106        self.record_session_actor_if_needed();
1107        self.maybe_load_actor_memory().await;
1108        self.maybe_load_actor_relationship().await;
1109        *self.messages_since_extraction.write() += 1;
1110    }
1111
1112    /// Post-turn lifecycle shared by streaming and non-streaming paths.
1113    async fn post_turn_session_lifecycle(&self) -> Result<()> {
1114        if *self.redispatch_depth.read() > 0 {
1115            return Ok(());
1116        }
1117        *self.messages_since_extraction.write() += 1;
1118        self.run_post_turn_maintenance().await
1119    }
1120
1121    /// Starts root-turn bookkeeping for user-message commit tracking.
1122    fn begin_root_turn(&self) {
1123        if *self.redispatch_depth.read() == 0 {
1124            let mut guard = self.active_turn_context.write();
1125            if guard.is_none() {
1126                self.root_user_message_committed
1127                    .store(false, Ordering::SeqCst);
1128                let max_calls = self
1129                    .runtime_config
1130                    .optimization
1131                    .max_speculative_llm_calls_per_turn;
1132                *guard = Some(TurnOptimizationContext::new(
1133                    String::new(),
1134                    HashMap::new(),
1135                    max_calls,
1136                ));
1137            }
1138        }
1139    }
1140
1141    fn update_active_turn_context(
1142        &self,
1143        processed_input: &str,
1144        input_context: HashMap<String, Value>,
1145    ) {
1146        if *self.redispatch_depth.read() > 0 {
1147            return;
1148        }
1149        let max_calls = self
1150            .runtime_config
1151            .optimization
1152            .max_speculative_llm_calls_per_turn;
1153        let mut guard = self.active_turn_context.write();
1154        match guard.as_mut() {
1155            Some(context) => {
1156                context.processed_input = processed_input.to_string();
1157                context.input_context = input_context;
1158                context.max_speculative_llm_calls = max_calls;
1159            }
1160            None => {
1161                *guard = Some(TurnOptimizationContext::new(
1162                    processed_input,
1163                    input_context,
1164                    max_calls,
1165                ));
1166            }
1167        }
1168    }
1169
1170    /// Writes the processed user message once for the root turn.
1171    async fn commit_root_user_message(&self, processed_input: &str) -> Result<()> {
1172        if *self.redispatch_depth.read() > 0 {
1173            return Ok(());
1174        }
1175        if !self
1176            .root_user_message_committed
1177            .swap(true, Ordering::SeqCst)
1178        {
1179            self.memory
1180                .add_message(ChatMessage::user(processed_input))
1181                .await?;
1182            if let Some(context) = self.active_turn_context.write().as_mut() {
1183                context.mark_user_message_committed();
1184            }
1185        }
1186        Ok(())
1187    }
1188
1189    /// Clears root-turn bookkeeping after final response handling.
1190    fn end_root_turn(&self) {
1191        if *self.redispatch_depth.read() == 0 {
1192            self.root_user_message_committed
1193                .store(false, Ordering::SeqCst);
1194            *self.active_turn_context.write() = None;
1195        }
1196    }
1197
1198    fn reserve_active_speculative_llm_call(&self, kind: RuntimeOptimizationKind) -> bool {
1199        self.begin_root_turn();
1200        let mut guard = self.active_turn_context.write();
1201        let Some(context) = guard.as_mut() else {
1202            return false;
1203        };
1204        context.reserve_speculative_llm_call_for(kind)
1205    }
1206
1207    fn branch_context_preview(&self) -> String {
1208        let context = self.build_context_with_overlays();
1209        let mut value = serde_json::to_string_pretty(&context).unwrap_or_else(|_| "{}".to_string());
1210        const MAX_CONTEXT_PREVIEW_CHARS: usize = 2048;
1211        if value.chars().count() > MAX_CONTEXT_PREVIEW_CHARS {
1212            value = value
1213                .chars()
1214                .take(MAX_CONTEXT_PREVIEW_CHARS)
1215                .collect::<String>();
1216            value.push_str("...");
1217        }
1218        value
1219    }
1220
1221    /// Applies freshness policy before rendering the next prompt.
1222    async fn await_background_before_next_turn(&self) {
1223        let optimization = &self.runtime_config.optimization;
1224        if !optimization.enabled {
1225            return;
1226        }
1227        let actor_id = self.effective_actor_id();
1228        let post = &optimization.post_turn;
1229        self.await_background_task(
1230            post.facts.await_before_next_turn,
1231            RuntimeTaskPurpose::PostTurnFacts,
1232            actor_id.as_deref(),
1233            "facts",
1234        )
1235        .await;
1236        self.await_background_task(
1237            post.relationships.await_before_next_turn,
1238            RuntimeTaskPurpose::PostTurnRelationship,
1239            actor_id.as_deref(),
1240            "relationships",
1241        )
1242        .await;
1243    }
1244
1245    async fn await_background_task(
1246        &self,
1247        policy: AwaitBeforeNextTurn,
1248        purpose: RuntimeTaskPurpose,
1249        actor_id: Option<&str>,
1250        label: &str,
1251    ) {
1252        match policy {
1253            AwaitBeforeNextTurn::Never => {}
1254            AwaitBeforeNextTurn::Always => {
1255                if let Err(error) = self.flush_background_tasks_for_purpose(purpose).await {
1256                    warn!(label = label, error = %error, "background maintenance flush failed");
1257                }
1258            }
1259            AwaitBeforeNextTurn::SameActor => {
1260                if let Some(actor_id) = actor_id
1261                    && let Err(error) = self
1262                        .flush_background_tasks_for_actor_purpose(actor_id, purpose)
1263                        .await
1264                {
1265                    warn!(label = label, actor_id = %actor_id, error = %error, "actor background maintenance flush failed");
1266                }
1267            }
1268        }
1269    }
1270
1271    /// Runs post-turn facts and relationship maintenance according to runtime policy.
1272    async fn run_post_turn_maintenance(&self) -> Result<()> {
1273        let optimization = &self.runtime_config.optimization;
1274        if !optimization.enabled {
1275            self.auto_extract_facts().await;
1276            self.auto_update_relationship().await;
1277            return Ok(());
1278        }
1279
1280        let facts_mode = effective_maintenance_mode(
1281            optimization.post_turn.facts.mode,
1282            optimization.parallel_post_turn_memory,
1283        );
1284        let relationships_mode = effective_maintenance_mode(
1285            optimization.post_turn.relationships.mode,
1286            optimization.parallel_post_turn_memory,
1287        );
1288
1289        match (facts_mode, relationships_mode) {
1290            (MaintenanceMode::InlineSerial, MaintenanceMode::InlineSerial) => {
1291                self.auto_extract_facts().await;
1292                self.auto_update_relationship().await;
1293            }
1294            (MaintenanceMode::InlineParallel, MaintenanceMode::InlineParallel) => {
1295                let facts = self.auto_extract_facts();
1296                let relationships = self.auto_update_relationship();
1297                tokio::join!(facts, relationships);
1298            }
1299            (MaintenanceMode::Background, MaintenanceMode::Background) => {
1300                self.schedule_facts_background().await?;
1301                self.schedule_relationship_background().await?;
1302            }
1303            (MaintenanceMode::Background, MaintenanceMode::InlineParallel)
1304            | (MaintenanceMode::Background, MaintenanceMode::InlineSerial) => {
1305                self.schedule_facts_background().await?;
1306                self.auto_update_relationship().await;
1307            }
1308            (MaintenanceMode::InlineParallel, MaintenanceMode::Background)
1309            | (MaintenanceMode::InlineSerial, MaintenanceMode::Background) => {
1310                self.auto_extract_facts().await;
1311                self.schedule_relationship_background().await?;
1312            }
1313            _ => {
1314                self.auto_extract_facts().await;
1315                self.auto_update_relationship().await;
1316            }
1317        }
1318        Ok(())
1319    }
1320
1321    async fn schedule_facts_background(&self) -> Result<()> {
1322        let policy = self.runtime_config.optimization.post_turn.facts.clone();
1323        let should_extract = self
1324            .facts_config
1325            .as_ref()
1326            .map(|c| c.enabled && c.auto_extract)
1327            .unwrap_or(false);
1328        if !should_extract {
1329            return Ok(());
1330        }
1331        let msgs_since = *self.messages_since_extraction.read();
1332        if msgs_since < 2 {
1333            return Ok(());
1334        }
1335        let Some(actor_id) = self.effective_actor_id() else {
1336            self.record_skipped_maintenance(
1337                "facts",
1338                ObservationPurpose::FactsExtraction,
1339                "missing_actor",
1340                Some(&policy),
1341            );
1342            return Ok(());
1343        };
1344        let Some(extractor) = self.fact_extractor.read().clone() else {
1345            return Ok(());
1346        };
1347        let messages = match self.memory.get_messages(None).await {
1348            Ok(messages) => messages,
1349            Err(error) => {
1350                warn!(error = %error, "failed to snapshot messages for fact extraction");
1351                return Ok(());
1352            }
1353        };
1354        let recent: Vec<_> = messages
1355            .iter()
1356            .rev()
1357            .take(msgs_since)
1358            .rev()
1359            .cloned()
1360            .collect();
1361        if recent.is_empty() {
1362            return Ok(());
1363        }
1364        let existing = self
1365            .actor_facts_cache
1366            .read()
1367            .get(&actor_id)
1368            .cloned()
1369            .unwrap_or_default();
1370        let categories = self
1371            .facts_config
1372            .as_ref()
1373            .map(|c| c.custom_categories.clone())
1374            .unwrap_or_default();
1375        let store = self.fact_store.read().clone();
1376        let cache = Arc::clone(&self.actor_facts_cache);
1377        let counter = Arc::clone(&self.messages_since_extraction);
1378        let hooks = Arc::clone(&self.hooks);
1379        let agent_id = self.info.id.clone();
1380        let observation = current_observation_context();
1381        let key = MaintenanceSequenceKey::actor(
1382            agent_id,
1383            actor_id.clone(),
1384            RuntimeTaskPurpose::PostTurnFacts,
1385        );
1386        let actor_for_task = actor_id.clone();
1387        let task = async move {
1388            let run = async move {
1389                let facts = extractor
1390                    .extract(&recent, &existing, Some(&actor_for_task), &categories)
1391                    .await?;
1392                if !facts.is_empty() {
1393                    if let Some(store) = store {
1394                        let authoritative = store.add_facts(&actor_for_task, facts.clone()).await?;
1395                        cache.write().insert(actor_for_task.clone(), authoritative);
1396                    } else {
1397                        cache
1398                            .write()
1399                            .entry(actor_for_task.clone())
1400                            .or_default()
1401                            .extend(facts.clone());
1402                    }
1403                    {
1404                        let mut count = counter.write();
1405                        if *count <= msgs_since {
1406                            *count = 0;
1407                        } else {
1408                            *count -= msgs_since;
1409                        }
1410                    }
1411                    hooks.on_facts_extracted(&actor_for_task, &facts).await;
1412                }
1413                Ok(())
1414            };
1415            if let Some(context) = observation {
1416                with_observation_context(
1417                    context.with_purpose(ObservationPurpose::FactsExtraction),
1418                    run,
1419                )
1420                .await
1421            } else {
1422                run.await
1423            }
1424        };
1425        self.spawn_or_handle_background(Some(key), task, "facts", &policy)
1426            .await
1427    }
1428
1429    async fn schedule_relationship_background(&self) -> Result<()> {
1430        let policy = self
1431            .runtime_config
1432            .optimization
1433            .post_turn
1434            .relationships
1435            .clone();
1436        let Some(manager) = self.relationship_manager.as_ref().cloned() else {
1437            return Ok(());
1438        };
1439        let Some(actor_id) = self.effective_actor_id() else {
1440            self.record_skipped_maintenance(
1441                "relationships",
1442                ObservationPurpose::RelationshipUpdate,
1443                "missing_actor",
1444                Some(&policy),
1445            );
1446            return Ok(());
1447        };
1448        let recent_messages = manager.config().auto_update.recent_messages;
1449        let messages = match self.memory.get_messages(Some(recent_messages)).await {
1450            Ok(messages) => messages,
1451            Err(error) => {
1452                warn!(actor = %actor_id, error = %error, "failed to snapshot messages for relationship update");
1453                return Ok(());
1454            }
1455        };
1456        let storage = self.storage.read().clone();
1457        let hooks = Arc::clone(&self.hooks);
1458        let agent_id = self.info.id.clone();
1459        let observation = current_observation_context();
1460        let key = MaintenanceSequenceKey::actor(
1461            agent_id.clone(),
1462            actor_id.clone(),
1463            RuntimeTaskPurpose::PostTurnRelationship,
1464        );
1465        let actor_for_task = actor_id.clone();
1466        let task = async move {
1467            let run = async move {
1468                if manager.config().auto_update.enabled {
1469                    let update = manager.auto_update(&actor_for_task, &messages).await?;
1470                    if !update.changes.is_empty() {
1471                        hooks
1472                            .on_relationship_change(&actor_for_task, &update.changes)
1473                            .await;
1474                    }
1475                    if let Some(ref event) = update.event {
1476                        hooks.on_notable_event(&actor_for_task, event).await;
1477                    }
1478                }
1479                if manager.config().persistence.enabled
1480                    && let (Some(storage), Some(value)) =
1481                        (storage, manager.relationship_as_value(&actor_for_task)?)
1482                {
1483                    storage
1484                        .save_relationship(&agent_id, &actor_for_task, &value)
1485                        .await?;
1486                }
1487                Ok(())
1488            };
1489            if let Some(context) = observation {
1490                with_observation_context(
1491                    context.with_purpose(ObservationPurpose::RelationshipUpdate),
1492                    run,
1493                )
1494                .await
1495            } else {
1496                run.await
1497            }
1498        };
1499        self.spawn_or_handle_background(Some(key), task, "relationships", &policy)
1500            .await
1501    }
1502
1503    /// Queues background maintenance or applies the configured overflow behavior.
1504    async fn spawn_or_handle_background<F>(
1505        &self,
1506        key: Option<MaintenanceSequenceKey>,
1507        task: F,
1508        label: &'static str,
1509        policy: &crate::optimization::config::MaintenanceTaskPolicy,
1510    ) -> Result<()>
1511    where
1512        F: Future<Output = Result<()>> + Send + 'static,
1513    {
1514        if self.background_maintenance.is_full() {
1515            match self
1516                .runtime_config
1517                .optimization
1518                .post_turn
1519                .on_background_overflow
1520            {
1521                BackgroundOverflowPolicy::RunInline => {
1522                    record_background_maintenance_event(
1523                        self.observability_manager.as_ref(),
1524                        label,
1525                        EventStatus::Success,
1526                        0,
1527                        "inline_overflow",
1528                        None,
1529                        Some(policy),
1530                    );
1531                    let start = Instant::now();
1532                    match task.await {
1533                        Ok(()) => record_background_maintenance_event(
1534                            self.observability_manager.as_ref(),
1535                            label,
1536                            EventStatus::Success,
1537                            start.elapsed().as_millis() as u64,
1538                            "inline_completed",
1539                            None,
1540                            Some(policy),
1541                        ),
1542                        Err(error) => {
1543                            warn!(label = label, error = %error, "inline maintenance fallback failed");
1544                            record_background_maintenance_event(
1545                                self.observability_manager.as_ref(),
1546                                label,
1547                                EventStatus::Error,
1548                                start.elapsed().as_millis() as u64,
1549                                "inline_failed",
1550                                Some(error.to_string()),
1551                                Some(policy),
1552                            );
1553                            return Err(error);
1554                        }
1555                    }
1556                }
1557                BackgroundOverflowPolicy::Drop => {
1558                    self.record_skipped_maintenance(
1559                        label,
1560                        ObservationPurpose::Other(label.to_string()),
1561                        "queue_full",
1562                        Some(policy),
1563                    );
1564                }
1565                BackgroundOverflowPolicy::Error => {
1566                    record_background_maintenance_event(
1567                        self.observability_manager.as_ref(),
1568                        label,
1569                        EventStatus::Error,
1570                        0,
1571                        "queue_full",
1572                        None,
1573                        Some(policy),
1574                    );
1575                    warn!(label = label, "background maintenance queue full");
1576                    return Err(AgentError::Other(format!(
1577                        "background maintenance queue is full for {}",
1578                        label
1579                    )));
1580                }
1581            }
1582            return Ok(());
1583        }
1584
1585        record_background_maintenance_event(
1586            self.observability_manager.as_ref(),
1587            label,
1588            EventStatus::Success,
1589            0,
1590            "scheduled",
1591            None,
1592            Some(policy),
1593        );
1594        let manager = self.observability_manager.clone();
1595        let policy_for_task = policy.clone();
1596        let observed_task = async move {
1597            let start = Instant::now();
1598            let result = task.await;
1599            match &result {
1600                Ok(()) => record_background_maintenance_event(
1601                    manager.as_ref(),
1602                    label,
1603                    EventStatus::Success,
1604                    start.elapsed().as_millis() as u64,
1605                    "completed",
1606                    None,
1607                    Some(&policy_for_task),
1608                ),
1609                Err(error) => record_background_maintenance_event(
1610                    manager.as_ref(),
1611                    label,
1612                    EventStatus::Error,
1613                    start.elapsed().as_millis() as u64,
1614                    "failed",
1615                    Some(error.to_string()),
1616                    Some(&policy_for_task),
1617                ),
1618            }
1619            result
1620        };
1621
1622        if let Err(error) = self.background_maintenance.spawn(key, observed_task) {
1623            record_background_maintenance_event(
1624                self.observability_manager.as_ref(),
1625                label,
1626                EventStatus::Error,
1627                0,
1628                "spawn_failed",
1629                Some(error.to_string()),
1630                Some(policy),
1631            );
1632            warn!(label = label, error = %error, "background maintenance spawn failed");
1633            return Err(error);
1634        }
1635        Ok(())
1636    }
1637
1638    /// Records a skipped background maintenance event when work cannot run.
1639    fn record_skipped_maintenance(
1640        &self,
1641        label: &str,
1642        purpose: ObservationPurpose,
1643        reason: &str,
1644        policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
1645    ) {
1646        if let Some(manager) = self.observability_manager.as_ref() {
1647            let mut tags = background_maintenance_tags(label, "skipped", Some(reason), policy);
1648            tags.insert("runtime.skip_reason".to_string(), reason.to_string());
1649            manager.record_lifecycle_event(
1650                EventType::MemoryOperation {
1651                    operation: format!("{}_maintenance", label),
1652                },
1653                purpose,
1654                EventStatus::Skipped,
1655                0,
1656                tags,
1657                None,
1658            );
1659        }
1660    }
1661
1662    /// Get cached actor facts for the effective actor.
1663    pub fn actor_facts(&self) -> Vec<ai_agents_core::KeyFact> {
1664        let Some(actor_id) = self.effective_actor_id() else {
1665            return Vec::new();
1666        };
1667        self.actor_facts_cache
1668            .read()
1669            .get(&actor_id)
1670            .cloned()
1671            .unwrap_or_default()
1672    }
1673
1674    /// Returns the formatted relationship prompt text for the effective actor, if relationship injection produced any text for this turn.
1675    pub fn relationship_memory_text(&self) -> Option<String> {
1676        self.format_relationship_for_context().map(|(_, text)| text)
1677    }
1678
1679    /// Manually extract facts from the last N messages.
1680    pub async fn extract_facts(
1681        &self,
1682        last_n: usize,
1683    ) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
1684        self.extract_facts_with_source(last_n, "manual").await
1685    }
1686
1687    async fn extract_facts_with_source(
1688        &self,
1689        last_n: usize,
1690        source: &'static str,
1691    ) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
1692        let extractor = match self.fact_extractor.read().clone() {
1693            Some(e) => e,
1694            None => return Ok(vec![]),
1695        };
1696
1697        let messages = self.memory.get_messages(None).await?;
1698        let recent: Vec<_> = messages.iter().rev().take(last_n).rev().cloned().collect();
1699
1700        if recent.is_empty() {
1701            return Ok(vec![]);
1702        }
1703
1704        let actor_id = self.effective_actor_id();
1705        let existing = actor_id
1706            .as_ref()
1707            .and_then(|aid| self.actor_facts_cache.read().get(aid).cloned())
1708            .unwrap_or_default();
1709
1710        let categories = self
1711            .facts_config
1712            .as_ref()
1713            .map(|c| c.custom_categories.clone())
1714            .unwrap_or_default();
1715
1716        let facts = self
1717            .observe_purpose(
1718                ObservationPurpose::FactsExtraction,
1719                extractor.extract(&recent, &existing, actor_id.as_deref(), &categories),
1720            )
1721            .await?;
1722
1723        // Save to storage and update the actor-scoped cache when an actor is known.
1724        if !facts.is_empty() {
1725            let fact_store_opt = self.fact_store.read().clone();
1726            let mut stored_total = 0usize;
1727            let mut cache_updated = false;
1728            if let (Some(store), Some(aid)) = (fact_store_opt, &actor_id) {
1729                // add_facts now returns the authoritative post-write set.
1730                let authoritative = store.add_facts(aid, facts.clone()).await?;
1731                stored_total = authoritative.len();
1732                self.actor_facts_cache
1733                    .write()
1734                    .insert(aid.clone(), authoritative);
1735                cache_updated = true;
1736            } else if let Some(aid) = &actor_id {
1737                let mut cache = self.actor_facts_cache.write();
1738                let entry = cache.entry(aid.clone()).or_default();
1739                entry.extend(facts.clone());
1740                stored_total = entry.len();
1741                cache_updated = true;
1742            }
1743
1744            info!(
1745                actor_id = %actor_id.as_deref().unwrap_or("<none>"),
1746                source = source,
1747                requested_messages = last_n,
1748                message_count = recent.len(),
1749                extracted_count = facts.len(),
1750                cache_updated = cache_updated,
1751                stored_total = stored_total,
1752                "facts extracted"
1753            );
1754
1755            if let Some(ref aid) = actor_id {
1756                self.hooks.on_facts_extracted(aid, &facts).await;
1757            }
1758        }
1759
1760        Ok(facts)
1761    }
1762
1763    /// Resolve actor_id from context if method is from_context.
1764    /// Supports dotted paths (e.g. "player.id", "user.profile.id").
1765    fn resolve_actor_id_from_context(&self) {
1766        if self
1767            .current_turn_actor_context()
1768            .and_then(|ctx| ctx.effective_actor_id().map(str::to_string))
1769            .is_some()
1770        {
1771            return;
1772        }
1773
1774        if let Some(ref am_config) = self.actor_memory_config
1775            && am_config.identification.method == ai_agents_facts::IdentificationMethod::FromContext
1776            && let Some(ref path) = am_config.identification.context_path
1777        {
1778            // get_path resolves dotted paths; get only handles top-level keys.
1779            let val = self
1780                .context_manager
1781                .get_path(path)
1782                .or_else(|| self.context_manager.get(path));
1783            if let Some(val) = val
1784                && let Some(id_str) = val.as_str()
1785            {
1786                let current = self.actor_id.read().clone();
1787                if current.as_deref() != Some(id_str) {
1788                    *self.actor_id.write() = Some(id_str.to_string());
1789                    let mut meta = self.session_metadata.write();
1790                    meta.actor_id = Some(id_str.to_string());
1791                    if !meta.actors.iter().any(|a| a == id_str) {
1792                        meta.actors.push(id_str.to_string());
1793                    }
1794                }
1795            }
1796        }
1797    }
1798
1799    /// Format actor facts for template injection.
1800    fn format_actor_facts_for_context(&self) -> String {
1801        // Respect inject_in_context: false.
1802        let should_inject = self
1803            .facts_config
1804            .as_ref()
1805            .map(|c| c.inject_in_context)
1806            .unwrap_or(true);
1807        if !should_inject {
1808            return String::new();
1809        }
1810
1811        let Some(actor_id) = self.effective_actor_id() else {
1812            return String::new();
1813        };
1814
1815        let facts = self
1816            .actor_facts_cache
1817            .read()
1818            .get(&actor_id)
1819            .cloned()
1820            .unwrap_or_default();
1821        if facts.is_empty() {
1822            return String::new();
1823        }
1824
1825        let am_config = self.actor_memory_config.as_ref();
1826        // Effective token cap: prefer memory.token_budget.allocation.facts when present,
1827        // otherwise fall back to actor_memory.injection.max_tokens.
1828        let facts_budget = self
1829            .memory_token_budget
1830            .as_ref()
1831            .map(|b| b.allocation.facts as usize)
1832            .filter(|n| *n > 0);
1833        let default_max = am_config.map(|c| c.injection.max_tokens).unwrap_or(800);
1834        let max_tokens = facts_budget.unwrap_or(default_max);
1835
1836        // Filter by category when injection.mode = category.
1837        let filtered: Vec<ai_agents_core::KeyFact> = if let Some(cfg) = am_config {
1838            if cfg.injection.mode == ai_agents_facts::InjectionMode::OnDemand {
1839                return String::new();
1840            }
1841            if cfg.injection.mode == ai_agents_facts::InjectionMode::Category
1842                && !cfg.injection.categories.is_empty()
1843            {
1844                facts
1845                    .iter()
1846                    .filter(|f| {
1847                        cfg.injection
1848                            .categories
1849                            .iter()
1850                            .any(|c| f.category.to_string() == *c)
1851                    })
1852                    .cloned()
1853                    .collect()
1854            } else {
1855                facts.clone()
1856            }
1857        } else {
1858            facts.clone()
1859        };
1860
1861        if filtered.is_empty() {
1862            return String::new();
1863        }
1864
1865        if let Some(store) = self.fact_store.read().clone() {
1866            store.format_for_context(&filtered, max_tokens)
1867        } else {
1868            String::new()
1869        }
1870    }
1871
1872    fn build_context_with_staged(&self, staged: &HashMap<String, Value>) -> HashMap<String, Value> {
1873        let context = self.build_context_with_overlays();
1874        let mut root = Value::Object(context.into_iter().collect());
1875        for (path, value) in staged {
1876            if let Ok(updated) = ai_agents_core::set_dot_path(root.clone(), path, value.clone()) {
1877                root = updated;
1878            }
1879        }
1880        match root {
1881            Value::Object(obj) => obj.into_iter().collect(),
1882            _ => HashMap::new(),
1883        }
1884    }
1885
1886    fn build_context_with_overlays(&self) -> HashMap<String, Value> {
1887        let mut context = self.context_manager.get_all();
1888        let mut root = Value::Object(context.clone().into_iter().collect());
1889
1890        if let Some(turn_ctx) = self.current_turn_actor_context() {
1891            if let Some(ref origin_actor_id) = turn_ctx.origin_actor_id
1892                && let Ok(updated) = ai_agents_core::set_dot_path(
1893                    root.clone(),
1894                    "interaction.origin_actor_id",
1895                    serde_json::json!(origin_actor_id),
1896                )
1897            {
1898                root = updated;
1899            }
1900            if let Some(ref sender_agent_id) = turn_ctx.sender_agent_id
1901                && let Ok(updated) = ai_agents_core::set_dot_path(
1902                    root.clone(),
1903                    "interaction.sender_agent_id",
1904                    serde_json::json!(sender_agent_id),
1905                )
1906            {
1907                root = updated;
1908            }
1909        }
1910
1911        if let Some(ref actor_id) = self.effective_actor_id()
1912            && let Ok(updated) = ai_agents_core::set_dot_path(
1913                root.clone(),
1914                "interaction.actor_id",
1915                serde_json::json!(actor_id),
1916            )
1917        {
1918            root = updated;
1919        }
1920
1921        if let Some(manager) = self.relationship_manager.as_ref()
1922            && let Some(actor_id) = self.effective_actor_id()
1923            && let Some(value) = manager.to_context_value(&actor_id)
1924            && let Ok(updated) = ai_agents_core::set_dot_path(
1925                root.clone(),
1926                &manager.config().injection.context_path,
1927                value,
1928            )
1929        {
1930            root = updated;
1931        }
1932
1933        if let Value::Object(obj) = root {
1934            context = obj.into_iter().collect();
1935        }
1936
1937        context
1938    }
1939
1940    fn resolve_actor_name_from_context(&self) -> Option<String> {
1941        for path in ["actor.name", "user.name", "player.name", "customer.name"] {
1942            if let Some(value) = self.context_manager.get_path(path)
1943                && let Some(name) = value.as_str()
1944            {
1945                return Some(name.to_string());
1946            }
1947        }
1948        None
1949    }
1950
1951    async fn maybe_load_actor_relationship(&self) {
1952        let Some(manager) = self.relationship_manager.as_ref() else {
1953            return;
1954        };
1955        let Some(actor_id) = self.effective_actor_id() else {
1956            return;
1957        };
1958
1959        let mut should_fire_loaded = false;
1960        if manager.get(&actor_id).is_none() {
1961            let mut loaded = false;
1962            if manager.config().persistence.enabled {
1963                let storage = self.storage.read().clone();
1964                if let Some(storage) = storage {
1965                    match storage.load_relationship(&self.info.id, &actor_id).await {
1966                        Ok(Some(value)) => match manager.insert_from_value(value) {
1967                            Ok(_) => loaded = true,
1968                            Err(e) => {
1969                                warn!(actor = %actor_id, error = %e, "failed to restore relationship")
1970                            }
1971                        },
1972                        Ok(None) => {}
1973                        Err(e) => {
1974                            warn!(actor = %actor_id, error = %e, "failed to load relationship")
1975                        }
1976                    }
1977                }
1978            }
1979
1980            if !loaded {
1981                manager.get_or_create(&actor_id, self.resolve_actor_name_from_context().as_deref());
1982            }
1983            should_fire_loaded = true;
1984        }
1985
1986        let actor_name = self.resolve_actor_name_from_context();
1987        let relationship = manager.touch_interaction(&actor_id, actor_name.as_deref());
1988        if should_fire_loaded {
1989            self.hooks
1990                .on_relationship_loaded(&actor_id, &relationship)
1991                .await;
1992        }
1993    }
1994
1995    fn format_relationship_for_context(&self) -> Option<(String, String)> {
1996        let manager = self.relationship_manager.as_ref()?;
1997        if !manager.config().injection.enabled {
1998            return None;
1999        }
2000        let actor_id = self.effective_actor_id()?;
2001        let relationship = manager.get(&actor_id)?;
2002        let local_cap = manager.config().injection.max_tokens;
2003        let global_cap = self
2004            .memory_token_budget
2005            .as_ref()
2006            .map(|b| b.allocation.relationships as usize)
2007            .filter(|n| *n > 0);
2008        let max_tokens = global_cap.map(|g| g.min(local_cap)).unwrap_or(local_cap);
2009        let text = ai_agents_relationships::format_relationship(
2010            &relationship,
2011            &manager.config().injection.format,
2012            max_tokens,
2013        );
2014        if text.is_empty() {
2015            None
2016        } else {
2017            Some((manager.config().injection.prompt_variable.clone(), text))
2018        }
2019    }
2020
2021    async fn persist_actor_relationship(&self, actor_id: &str) -> Result<()> {
2022        let Some(manager) = self.relationship_manager.as_ref() else {
2023            return Ok(());
2024        };
2025        if !manager.config().persistence.enabled {
2026            return Ok(());
2027        }
2028        let storage = self.storage.read().clone();
2029        let Some(storage) = storage else {
2030            return Ok(());
2031        };
2032        if let Some(value) = manager.relationship_as_value(actor_id)? {
2033            storage
2034                .save_relationship(&self.info.id, actor_id, &value)
2035                .await?;
2036        }
2037        Ok(())
2038    }
2039
2040    async fn auto_update_relationship(&self) {
2041        let Some(manager) = self.relationship_manager.as_ref() else {
2042            return;
2043        };
2044        let Some(actor_id) = self.effective_actor_id() else {
2045            return;
2046        };
2047        if !manager.config().auto_update.enabled {
2048            let _ = self.persist_actor_relationship(&actor_id).await;
2049            return;
2050        }
2051
2052        let recent_messages = manager.config().auto_update.recent_messages;
2053        let messages = match self.memory.get_messages(Some(recent_messages)).await {
2054            Ok(messages) => messages,
2055            Err(e) => {
2056                warn!(actor = %actor_id, error = %e, "failed to read messages for relationship update");
2057                return;
2058            }
2059        };
2060
2061        match self
2062            .observe_purpose(
2063                ObservationPurpose::RelationshipUpdate,
2064                manager.auto_update(&actor_id, &messages),
2065            )
2066            .await
2067        {
2068            Ok(update) => {
2069                if !update.changes.is_empty() {
2070                    self.hooks
2071                        .on_relationship_change(&actor_id, &update.changes)
2072                        .await;
2073                }
2074                if let Some(ref event) = update.event {
2075                    self.hooks.on_notable_event(&actor_id, event).await;
2076                }
2077                let persisted = match self.persist_actor_relationship(&actor_id).await {
2078                    Ok(()) => true,
2079                    Err(e) => {
2080                        warn!(actor = %actor_id, error = %e, "failed to persist relationship");
2081                        false
2082                    }
2083                };
2084                if !update.changes.is_empty() || update.event.is_some() {
2085                    let changed_dimensions: Vec<String> = update
2086                        .changes
2087                        .iter()
2088                        .map(|change| format!("{}:{}", change.perspective, change.dimension))
2089                        .collect();
2090                    info!(
2091                        actor_id = %actor_id,
2092                        change_count = update.changes.len(),
2093                        changed_dimensions = ?changed_dimensions,
2094                        event_present = update.event.is_some(),
2095                        persisted = persisted,
2096                        "relationship updated"
2097                    );
2098                } else {
2099                    debug!(actor_id = %actor_id, persisted = persisted, "relationship evaluation ran but found no changes");
2100                }
2101            }
2102            Err(e) => warn!(actor = %actor_id, error = %e, "relationship update failed"),
2103        }
2104    }
2105
2106    /// Run auto-extraction after a chat turn.
2107    async fn auto_extract_facts(&self) {
2108        let should_extract = self
2109            .facts_config
2110            .as_ref()
2111            .map(|c| c.enabled && c.auto_extract)
2112            .unwrap_or(false);
2113
2114        if !should_extract {
2115            debug!("fact extraction skipped because auto extraction is disabled");
2116            return;
2117        }
2118
2119        let msgs_since = *self.messages_since_extraction.read();
2120        if msgs_since < 2 {
2121            debug!(
2122                messages_since_extraction = msgs_since,
2123                "fact extraction skipped until threshold is reached"
2124            );
2125            return;
2126        }
2127
2128        match self.extract_facts_with_source(msgs_since, "auto").await {
2129            Ok(facts) => {
2130                if !facts.is_empty() {
2131                    *self.messages_since_extraction.write() = 0;
2132                } else {
2133                    debug!("fact extraction ran but found no new facts");
2134                }
2135            }
2136            Err(e) => {
2137                warn!("fact extraction failed: {}", e);
2138            }
2139        }
2140    }
2141
2142    pub fn with_persona(mut self, manager: Arc<ai_agents_persona::PersonaManager>) -> Self {
2143        self.persona_manager = Some(manager);
2144        self
2145    }
2146
2147    pub fn persona_manager(&self) -> Option<&Arc<ai_agents_persona::PersonaManager>> {
2148        self.persona_manager.as_ref()
2149    }
2150
2151    pub fn with_disambiguation(mut self, config: DisambiguationConfig) -> Self {
2152        if config.is_enabled() {
2153            let manager = DisambiguationManager::new(config, Arc::clone(&self.llm_registry))
2154                .with_clarification_observer(Arc::new(ObservabilityClarificationObserver));
2155            self.disambiguation_manager = Some(manager);
2156        }
2157        self
2158    }
2159
2160    pub fn disambiguation_manager(&self) -> Option<&DisambiguationManager> {
2161        self.disambiguation_manager.as_ref()
2162    }
2163
2164    pub fn has_disambiguation(&self) -> bool {
2165        self.disambiguation_manager
2166            .as_ref()
2167            .is_some_and(|m| m.is_enabled())
2168    }
2169
2170    pub async fn init_storage(&self) -> Result<()> {
2171        //
2172        // One readiness guard prevents concurrent entry points from constructing storage or fact state more than once.
2173        //
2174        let _guard = self.storage_init.lock().await;
2175        let mut storage = self.storage.read().clone();
2176        if storage.is_none() && !self.storage_config.is_none() {
2177            let storage_config = self.convert_storage_config();
2178            storage = create_storage(&storage_config).await?;
2179            *self.storage.write() = storage.clone();
2180        }
2181
2182        self.validate_storage_requirements(storage.as_deref())?;
2183        self.complete_facts_init().await;
2184        Ok(())
2185    }
2186
2187    fn validate_storage_requirements(&self, storage: Option<&dyn AgentStorage>) -> Result<()> {
2188        let facts_required = self
2189            .facts_config
2190            .as_ref()
2191            .is_some_and(|config| config.enabled)
2192            || self
2193                .actor_memory_config
2194                .as_ref()
2195                .is_some_and(|config| config.enabled);
2196        let relationships_required = self
2197            .relationship_manager
2198            .as_ref()
2199            .is_some_and(|manager| manager.config().persistence.enabled);
2200
2201        let Some(storage) = storage else {
2202            let mut requirements = Vec::new();
2203            if facts_required {
2204                requirements.push("actor facts or actor memory");
2205            }
2206            if relationships_required {
2207                requirements.push("persistent relationships");
2208            }
2209            if requirements.is_empty() {
2210                return Ok(());
2211            }
2212            return Err(AgentError::Config(format!(
2213                "Storage is required for enabled {} but none is configured or injected",
2214                requirements.join(" and ")
2215            )));
2216        };
2217
2218        //
2219        // Durable features must fail before execution instead of degrading to volatile state.
2220        //
2221        if facts_required && !storage.supports(StorageCapability::ActorFacts) {
2222            return Err(AgentError::UnsupportedStorageCapability(
2223                StorageCapability::ActorFacts,
2224            ));
2225        }
2226        if relationships_required && !storage.supports(StorageCapability::ActorRelationships) {
2227            return Err(AgentError::UnsupportedStorageCapability(
2228                StorageCapability::ActorRelationships,
2229            ));
2230        }
2231        Ok(())
2232    }
2233
2234    /// Initialize fact store and extractor from stored config and current storage.
2235    /// Called from init_storage() so facts are ready before the first turn.
2236    async fn complete_facts_init(&self) {
2237        if self.fact_store.read().is_some() {
2238            return;
2239        }
2240        let storage = match self.storage.read().clone() {
2241            Some(s) => s,
2242            None => return,
2243        };
2244
2245        let facts_enabled = self
2246            .facts_config
2247            .as_ref()
2248            .map(|f| f.enabled)
2249            .unwrap_or(false);
2250        let actor_memory_enabled = self
2251            .actor_memory_config
2252            .as_ref()
2253            .map(|a| a.enabled)
2254            .unwrap_or(false);
2255
2256        if !facts_enabled && !actor_memory_enabled {
2257            return;
2258        }
2259
2260        let fc = self.facts_config.clone().unwrap_or_default();
2261        let store = Arc::new(ai_agents_facts::FactStore::new(
2262            storage,
2263            self.info.id.clone(),
2264            fc.clone(),
2265        ));
2266
2267        let extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>> = if facts_enabled {
2268            let extractor_llm = fc
2269                .extractor_llm
2270                .as_ref()
2271                .and_then(|alias| self.llm_registry.get(alias).ok())
2272                .or_else(|| self.llm_registry.router().ok())
2273                .or_else(|| self.llm_registry.default().ok());
2274            extractor_llm.map(|llm| {
2275                Arc::new(ai_agents_facts::LLMFactExtractor::new(llm, fc.clone()))
2276                    as Arc<dyn ai_agents_facts::FactExtractor>
2277            })
2278        } else {
2279            None
2280        };
2281
2282        *self.fact_store.write() = Some(store);
2283        *self.fact_extractor.write() = extractor;
2284        debug!(
2285            agent = %self.info.id,
2286            facts_enabled,
2287            actor_memory_enabled,
2288            "facts storage initialized"
2289        );
2290    }
2291
2292    fn convert_storage_config(&self) -> StorageStorageConfig {
2293        crate::spec::storage::to_storage_config(&self.storage_config)
2294    }
2295
2296    pub fn storage(&self) -> Option<Arc<dyn AgentStorage>> {
2297        self.storage.read().clone()
2298    }
2299
2300    pub fn storage_config(&self) -> &StorageConfig {
2301        &self.storage_config
2302    }
2303
2304    /// Returns the spawner if configured via a spawner: YAML section.
2305    pub fn spawner(&self) -> Option<&Arc<crate::spawner::AgentSpawner>> {
2306        self.spawner.as_ref()
2307    }
2308
2309    /// Returns the agent registry if configured via a spawner: YAML section.
2310    pub fn spawner_registry(&self) -> Option<&Arc<crate::spawner::AgentRegistry>> {
2311        self.spawner_registry.as_ref()
2312    }
2313
2314    pub fn has_spawner(&self) -> bool {
2315        self.spawner_registry.is_some()
2316    }
2317
2318    pub fn with_spawner_handles(
2319        mut self,
2320        spawner: Arc<crate::spawner::AgentSpawner>,
2321        registry: Arc<crate::spawner::AgentRegistry>,
2322    ) -> Self {
2323        self.spawner = Some(spawner);
2324        self.spawner_registry = Some(registry);
2325        self
2326    }
2327
2328    pub fn with_hooks(mut self, hooks: Arc<dyn AgentHooks>) -> Self {
2329        self.hooks = hooks;
2330        self
2331    }
2332
2333    pub fn with_parallel_tools(mut self, config: ParallelToolsConfig) -> Self {
2334        self.parallel_tools = config;
2335        self
2336    }
2337
2338    pub fn with_streaming(mut self, config: StreamingConfig) -> Self {
2339        self.streaming = config;
2340        self
2341    }
2342
2343    pub fn with_hitl(mut self, engine: HITLEngine, handler: Arc<dyn ApprovalHandler>) -> Self {
2344        self.hitl_engine = Some(engine);
2345        self.approval_handler = handler;
2346        self
2347    }
2348
2349    pub fn with_max_context_tokens(mut self, tokens: u32) -> Self {
2350        self.max_context_tokens = tokens;
2351        self
2352    }
2353
2354    pub fn with_memory_token_budget(mut self, budget: MemoryTokenBudget) -> Self {
2355        self.memory_token_budget = Some(budget);
2356        self
2357    }
2358
2359    pub fn with_recovery_manager(mut self, manager: RecoveryManager) -> Self {
2360        self.recovery_manager = manager;
2361        self
2362    }
2363
2364    pub fn with_tool_security(mut self, engine: ToolSecurityEngine) -> Self {
2365        self.tool_security = engine;
2366        self
2367    }
2368
2369    /// Returns the host-only runtime control handle.
2370    pub fn runtime_control(&self) -> RuntimeControlHandle {
2371        RuntimeControlHandle {
2372            state: Arc::clone(&self.runtime_control),
2373        }
2374    }
2375
2376    /// Installs or clears the host question handler used by `ask_user`.
2377    pub fn set_question_handler(&self, handler: Option<Arc<dyn QuestionHandler>>) {
2378        self.tools.set_question_handler(handler);
2379    }
2380
2381    /// Installs the host diagnostics provider used by `diagnostics`.
2382    pub fn set_diagnostics_provider(&self, provider: Arc<dyn DiagnosticsProvider>) {
2383        self.tools.set_diagnostics_provider(provider);
2384    }
2385
2386    /// Installs the host command runner used by `command`.
2387    pub fn set_command_runner(&self, runner: Arc<dyn CommandRunner>) {
2388        self.tools.set_command_runner(runner);
2389    }
2390
2391    /// Installs the host web-search provider used by `web_search`.
2392    pub fn set_web_search_provider(&self, provider: Arc<dyn ai_agents_tools::WebSearchProvider>) {
2393        self.tools.set_web_search_provider(provider);
2394    }
2395
2396    /// Returns the current session-local todo list.
2397    pub fn todos(&self) -> Vec<TodoItem> {
2398        self.tools.todos()
2399    }
2400
2401    /// Returns the effective tool security snapshot for the next decision.
2402    fn active_tool_security(&self) -> ToolSecurityEngine {
2403        self.runtime_control
2404            .tool_security_override
2405            .read()
2406            .clone()
2407            .unwrap_or_else(|| self.tool_security.clone())
2408    }
2409
2410    /// Reads policy, scope, emergency state, and generation as one safety snapshot.
2411    fn runtime_safety_snapshot(&self) -> RuntimeSafetySnapshot {
2412        let _guard = self.runtime_control.snapshot_guard.read();
2413        RuntimeSafetySnapshot {
2414            version: self.runtime_control.version.load(Ordering::SeqCst),
2415            emergency_deny: self.runtime_control.emergency_deny.load(Ordering::SeqCst),
2416            tool_security: self
2417                .runtime_control
2418                .tool_security_override
2419                .read()
2420                .clone()
2421                .unwrap_or_else(|| self.tool_security.clone()),
2422            tool_scope_override: self.runtime_control.tool_scope_override.read().clone(),
2423        }
2424    }
2425
2426    /// Performs the single rate admission after locks are held and rejects changed runtime, policy, or state authority.
2427    fn admit_tool_execution(
2428        &self,
2429        expected_runtime_version: u64,
2430        expected_policy_version: u64,
2431        expected_state_generation: Option<u64>,
2432        canonical_id: &str,
2433    ) -> SecurityCheckResult {
2434        let _guard = self.runtime_control.snapshot_guard.read();
2435        if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
2436            return SecurityCheckResult::Block {
2437                reason: "runtime emergency deny is enabled".to_string(),
2438            };
2439        }
2440        let runtime_version = self.runtime_control.version.load(Ordering::SeqCst);
2441        let security_engine = self
2442            .runtime_control
2443            .tool_security_override
2444            .read()
2445            .clone()
2446            .unwrap_or_else(|| self.tool_security.clone());
2447        if runtime_version != expected_runtime_version
2448            || security_engine.policy_version() != expected_policy_version
2449        {
2450            return SecurityCheckResult::Block {
2451                reason: "runtime safety controls changed before admission".to_string(),
2452            };
2453        }
2454        let current_state_generation = self
2455            .state_machine
2456            .as_ref()
2457            .map(|state_machine| state_machine.generation());
2458        if current_state_generation != expected_state_generation {
2459            return SecurityCheckResult::Block {
2460                reason: "state scope changed before admission".to_string(),
2461            };
2462        }
2463        security_engine.admit_tool_execution(canonical_id)
2464    }
2465
2466    pub fn with_process_processor(mut self, processor: ProcessProcessor) -> Self {
2467        let processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
2468        self.process_processor = Some(processor);
2469        self
2470    }
2471
2472    pub fn with_state_machine(
2473        mut self,
2474        state_machine: Arc<StateMachine>,
2475        evaluator: Arc<dyn TransitionEvaluator>,
2476    ) -> Self {
2477        self.state_machine = Some(state_machine);
2478        self.transition_evaluator = Some(evaluator);
2479        self
2480    }
2481
2482    pub fn with_context_manager(mut self, manager: Arc<ContextManager>) -> Self {
2483        self.context_manager = manager;
2484        self
2485    }
2486
2487    pub fn register_message_filter(&self, name: impl Into<String>, filter: Arc<dyn MessageFilter>) {
2488        self.message_filters.write().insert(name.into(), filter);
2489    }
2490
2491    pub fn set_context(&self, key: &str, value: Value) -> Result<()> {
2492        self.context_manager.update(key, value)
2493    }
2494
2495    pub fn update_context(&self, path: &str, value: Value) -> Result<()> {
2496        self.context_manager.update(path, value)
2497    }
2498
2499    pub fn get_context(&self) -> HashMap<String, Value> {
2500        self.build_context_with_overlays()
2501    }
2502
2503    pub fn remove_context(&self, key: &str) -> Option<Value> {
2504        self.context_manager.remove(key)
2505    }
2506
2507    pub async fn refresh_context(&self, key: &str) -> Result<()> {
2508        self.context_manager.refresh(key).await
2509    }
2510
2511    pub fn register_context_provider(&self, name: &str, provider: Arc<dyn ContextProvider>) {
2512        self.context_manager.register_provider(name, provider);
2513    }
2514
2515    pub fn current_state(&self) -> Option<String> {
2516        self.state_machine.as_ref().map(|sm| sm.current())
2517    }
2518
2519    // State changes invalidate only resolved confirmation ownership; ordinary clarification remains manager-owned across layered enablement changes.
2520    async fn invalidate_pending_confirmation(&self, reason: &'static str) {
2521        self.disambiguation_epoch.fetch_add(1, Ordering::SeqCst);
2522        let Some(disambiguator) = self.disambiguation_manager.as_ref() else {
2523            return;
2524        };
2525        if disambiguator.has_pending_confirmation().await {
2526            disambiguator.clear_pending().await;
2527            *self.pending_skill_id.write() = None;
2528            info!(
2529                confirmation_event = "invalidated",
2530                invalidation_reason = reason,
2531                "Runtime invalidated pending confirmation"
2532            );
2533        }
2534    }
2535
2536    // Admission linearizes pending publication or redispatch before reset and state mutation, then rechecks both ownership generations.
2537    async fn admit_disambiguation_redispatch(
2538        &self,
2539        expected_epoch: u64,
2540        expected_state_generation: Option<u64>,
2541    ) -> Result<tokio::sync::RwLockReadGuard<'_, ()>> {
2542        let admission = self.disambiguation_admission.read().await;
2543        let state_generation = self
2544            .state_machine
2545            .as_ref()
2546            .map(|state_machine| state_machine.generation());
2547        if self.disambiguation_epoch.load(Ordering::SeqCst) != expected_epoch
2548            || state_generation != expected_state_generation
2549        {
2550            return Err(AgentError::Other(
2551                "Disambiguation ownership changed before redispatch admission".to_string(),
2552            ));
2553        }
2554        Ok(admission)
2555    }
2556
2557    // A transition reservation prevents losing runtime transitions from running duplicate exit side effects.
2558    fn reserve_state_transition(&self) -> Option<StateTransitionReservation<'_>> {
2559        self.state_transition_reserved
2560            .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
2561            .ok()
2562            .map(|_| StateTransitionReservation {
2563                reserved: &self.state_transition_reserved,
2564            })
2565    }
2566
2567    // Optional ownership is used only for terminal skill responses that came from manager-owned pending state.
2568    async fn admit_optional_disambiguation_ownership(
2569        &self,
2570        ownership: Option<DisambiguationOwnership>,
2571    ) -> Result<Option<tokio::sync::RwLockReadGuard<'_, ()>>> {
2572        match ownership {
2573            Some(ownership) => self
2574                .admit_disambiguation_redispatch(ownership.epoch, ownership.state_generation)
2575                .await
2576                .map(Some),
2577            None => Ok(None),
2578        }
2579    }
2580
2581    /// Applies a manual transition after reserving its exit actions and keeping async lifecycle work outside the commit lock.
2582    pub async fn transition_to(&self, state: &str) -> Result<()> {
2583        let Some(ref sm) = self.state_machine else {
2584            return Ok(());
2585        };
2586        let claim_admission = self.disambiguation_admission.write().await;
2587        let reservation = self.reserve_state_transition().ok_or_else(|| {
2588            AgentError::Other("Another state transition is already in progress".to_string())
2589        })?;
2590        let from_state = sm.current();
2591        let expected_state_generation = sm.generation();
2592        let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
2593        let history_before = sm.history();
2594        drop(claim_admission);
2595
2596        self.execute_state_exit_actions(&from_state).await;
2597
2598        let admission = self.disambiguation_admission.write().await;
2599        if sm.current() != from_state
2600            || sm.generation() != expected_state_generation
2601            || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
2602        {
2603            return Err(AgentError::Other(
2604                "State ownership changed during manual transition preparation".to_string(),
2605            ));
2606        }
2607        sm.transition_to(state, "manual transition")?;
2608        self.invalidate_pending_confirmation("state_transition")
2609            .await;
2610        let entered = sm.current();
2611        let is_reentry = Self::state_was_previously_entered(&entered, &from_state, &history_before);
2612        drop(admission);
2613
2614        self.execute_state_enter_actions(&entered, is_reentry).await;
2615        drop(reservation);
2616        info!(to = %entered, "Manual state transition");
2617        Ok(())
2618    }
2619
2620    pub fn state_history(&self) -> Vec<StateTransitionEvent> {
2621        self.state_machine
2622            .as_ref()
2623            .map(|sm| sm.history())
2624            .unwrap_or_default()
2625    }
2626
2627    /// Get a copy of current session metadata.
2628    pub fn session_metadata(&self) -> ai_agents_core::SessionMetadata {
2629        self.session_metadata.read().clone()
2630    }
2631
2632    /// Delete all facts and sessions for an actor, gated by privacy.allow_deletion.
2633    /// Returns Err when actor_memory.privacy.allow_deletion is false.
2634    pub async fn delete_actor_data(&self, actor_id: &str) -> Result<()> {
2635        let allowed = self
2636            .actor_memory_config
2637            .as_ref()
2638            .map(|c| c.privacy.allow_deletion)
2639            .unwrap_or(true);
2640        if !allowed {
2641            return Err(AgentError::Config(
2642                "privacy.allow_deletion is false; actor data deletion is not permitted".into(),
2643            ));
2644        }
2645        let storage = self.storage.read().clone();
2646        if let Some(storage) = storage {
2647            //
2648            // Composite support is checked before mutation so privacy deletion cannot become partial.
2649            //
2650            if !storage.supports(StorageCapability::ActorDataDeletion) {
2651                return Err(AgentError::UnsupportedStorageCapability(
2652                    StorageCapability::ActorDataDeletion,
2653                ));
2654            }
2655            storage.delete_actor_data(&self.info.id, actor_id).await?;
2656        } else {
2657            //
2658            // Clone the fallback store before awaiting so backend I/O never holds the runtime read guard.
2659            //
2660            let store = { self.fact_store.read().clone() };
2661            if let Some(store) = store {
2662                store.delete_actor_data(actor_id).await?;
2663            }
2664        }
2665        if let Some(manager) = self.relationship_manager.as_ref() {
2666            manager.remove(actor_id);
2667        }
2668        self.actor_facts_cache.write().remove(actor_id);
2669        Ok(())
2670    }
2671
2672    /// Overwrite session metadata (tags, ttl, custom fields).
2673    pub fn set_session_metadata(&self, meta: ai_agents_core::SessionMetadata) {
2674        *self.session_metadata.write() = meta;
2675    }
2676
2677    /// Delete sessions whose TTL has expired. Returns number of sessions removed.
2678    pub async fn cleanup_expired_sessions(&self) -> Result<usize> {
2679        let storage = self.storage.read().clone();
2680        match storage {
2681            Some(s) => {
2682                let count = s.cleanup_expired().await?;
2683                if count > 0 {
2684                    self.hooks.on_sessions_expired(count).await;
2685                }
2686                Ok(count)
2687            }
2688            None => Err(AgentError::Config(
2689                "No storage configured. Use with_storage_config() or with_storage() first".into(),
2690            )),
2691        }
2692    }
2693
2694    /// List sessions matching a filter. Supports actor, tag, and date filters.
2695    pub async fn list_sessions_filtered(
2696        &self,
2697        filter: &ai_agents_core::SessionFilter,
2698    ) -> Result<Vec<ai_agents_core::SessionSummary>> {
2699        let storage = self.storage.read().clone();
2700        match storage {
2701            Some(s) => s.list_sessions_filtered(filter).await,
2702            None => Err(AgentError::Config(
2703                "No storage configured. Use with_storage_config() or with_storage() first".into(),
2704            )),
2705        }
2706    }
2707
2708    pub async fn save_state(&self) -> Result<AgentSnapshot> {
2709        let memory_snapshot = self.memory.snapshot().await?;
2710        let state_machine_snapshot = self.state_machine.as_ref().map(|sm| sm.snapshot());
2711        let context_snapshot = self.context_manager.snapshot();
2712
2713        let mut snapshot = AgentSnapshot::new(self.info.id.clone())
2714            .with_memory(memory_snapshot)
2715            .with_context(context_snapshot)
2716            .with_state_machine(
2717                state_machine_snapshot.unwrap_or_else(|| StateMachineSnapshot {
2718                    current_state: String::new(),
2719                    previous_state: None,
2720                    turn_count: 0,
2721                    no_transition_count: 0,
2722                    history: vec![],
2723                }),
2724            );
2725
2726        if let Some(ref persona) = self.persona_manager {
2727            snapshot.persona = Some(persona.snapshot_as_value()?);
2728        }
2729
2730        if let Some(ref relationships) = self.relationship_manager {
2731            snapshot.relationships = Some(relationships.snapshot_as_value()?);
2732        }
2733
2734        Ok(snapshot)
2735    }
2736
2737    /// Save state including spawned agents manifest for session persistence.
2738    pub async fn save_state_full(&self) -> Result<AgentSnapshot> {
2739        let mut snapshot = self.save_state().await?;
2740        if let Some(ref registry) = self.spawner_registry {
2741            let entries = registry.list_with_specs();
2742            if !entries.is_empty() {
2743                snapshot = snapshot.with_spawned_agents(entries);
2744            }
2745        }
2746        Ok(snapshot)
2747    }
2748
2749    /// Restores persisted state after invalidating any pending confirmation ownership.
2750    pub async fn restore_state(&self, snapshot: AgentSnapshot) -> Result<()> {
2751        let _admission = self.disambiguation_admission.write().await;
2752        if self.state_transition_reserved.load(Ordering::SeqCst) {
2753            return Err(AgentError::Other(
2754                "Cannot restore state while a state transition is in progress".to_string(),
2755            ));
2756        }
2757        self.invalidate_pending_confirmation("state_restore").await;
2758        *self.pending_skill_id.write() = None;
2759        if let Some(disambiguator) = self.disambiguation_manager.as_ref() {
2760            disambiguator.clear_pending().await;
2761        }
2762        self.memory.restore(snapshot.memory).await?;
2763
2764        if let (Some(sm), Some(sm_snapshot)) = (&self.state_machine, snapshot.state_machine)
2765            && !sm_snapshot.current_state.is_empty()
2766        {
2767            sm.restore(sm_snapshot)?;
2768        }
2769
2770        self.context_manager.restore(snapshot.context);
2771
2772        if let (Some(persona_value), Some(persona_manager)) =
2773            (snapshot.persona, &self.persona_manager)
2774        {
2775            persona_manager.restore_from_value(persona_value)?;
2776        }
2777
2778        if let (Some(relationship_value), Some(relationship_manager)) =
2779            (snapshot.relationships, &self.relationship_manager)
2780        {
2781            relationship_manager.restore_from_value(relationship_value)?;
2782        }
2783
2784        info!(agent_id = %snapshot.agent_id, "State restored");
2785        Ok(())
2786    }
2787
2788    pub async fn save_to(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<()> {
2789        let snapshot = self.save_state().await?;
2790        storage.save(session_id, &snapshot).await
2791    }
2792
2793    async fn load_session_restore(
2794        storage: &dyn AgentStorage,
2795        session_id: &str,
2796    ) -> Result<Option<StoredSessionRestore>> {
2797        let Some(snapshot) = storage.load(session_id).await? else {
2798            return Ok(None);
2799        };
2800        //
2801        // Metadata is restored only when the backend explicitly owns that durable contract.
2802        //
2803        let metadata = if storage.supports(StorageCapability::SessionMetadata) {
2804            storage.load_metadata(session_id).await?
2805        } else {
2806            None
2807        };
2808        Ok(Some(StoredSessionRestore { snapshot, metadata }))
2809    }
2810
2811    async fn capture_session_restore_point(&self) -> Result<RuntimeSessionRestorePoint> {
2812        Ok(RuntimeSessionRestorePoint {
2813            snapshot: self.save_state().await?,
2814            metadata: self.session_metadata(),
2815            actor_id: self.actor_id(),
2816            session_id: self.current_session_id.read().clone(),
2817        })
2818    }
2819
2820    async fn apply_session_restore_unchecked(
2821        &self,
2822        session_id: &str,
2823        stored: StoredSessionRestore,
2824    ) -> Result<()> {
2825        self.restore_state(stored.snapshot).await?;
2826        let metadata = stored.metadata.unwrap_or_default();
2827        if let Some(actor_id) = metadata.actor_id.as_deref() {
2828            self.set_actor_id(actor_id)?;
2829        } else {
2830            self.clear_actor_id();
2831        }
2832        self.set_session_metadata(metadata);
2833        *self.current_session_id.write() = Some(session_id.to_string());
2834        Ok(())
2835    }
2836
2837    async fn restore_session_restore_point(
2838        &self,
2839        restore_point: &RuntimeSessionRestorePoint,
2840    ) -> Result<()> {
2841        self.restore_state(restore_point.snapshot.clone()).await?;
2842        if let Some(actor_id) = restore_point.actor_id.as_deref() {
2843            self.set_actor_id(actor_id)?;
2844        } else {
2845            self.clear_actor_id();
2846        }
2847        self.set_session_metadata(restore_point.metadata.clone());
2848        *self.current_session_id.write() = restore_point.session_id.clone();
2849        Ok(())
2850    }
2851
2852    async fn apply_session_restore(
2853        &self,
2854        session_id: &str,
2855        stored: StoredSessionRestore,
2856    ) -> Result<()> {
2857        let before = self.capture_session_restore_point().await?;
2858        if let Err(error) = self
2859            .apply_session_restore_unchecked(session_id, stored)
2860            .await
2861        {
2862            return match self.restore_session_restore_point(&before).await {
2863                Ok(()) => Err(error),
2864                Err(rollback_error) => Err(AgentError::Other(format!(
2865                    "Session restore failed: {error}; rollback failed: {rollback_error}"
2866                ))),
2867            };
2868        }
2869        Ok(())
2870    }
2871
2872    async fn rollback_session_restore_set(
2873        parent: Option<(&RuntimeAgent, &RuntimeSessionRestorePoint)>,
2874        children: &[(String, Arc<RuntimeAgent>, RuntimeSessionRestorePoint)],
2875    ) -> Vec<String> {
2876        let mut errors = Vec::new();
2877        if let Some((agent, restore_point)) = parent
2878            && let Err(error) = agent.restore_session_restore_point(restore_point).await
2879        {
2880            errors.push(format!("parent: {error}"));
2881        }
2882        for (id, agent, restore_point) in children {
2883            if let Err(error) = agent.restore_session_restore_point(restore_point).await {
2884                errors.push(format!("child '{id}': {error}"));
2885            }
2886        }
2887        errors
2888    }
2889
2890    fn restore_failure(error: impl std::fmt::Display, rollback_errors: Vec<String>) -> AgentError {
2891        if rollback_errors.is_empty() {
2892            AgentError::Other(format!(
2893                "Session restore failed: {error}; runtime state was rolled back"
2894            ))
2895        } else {
2896            AgentError::Other(format!(
2897                "Session restore failed: {error}; rollback also failed for {}",
2898                rollback_errors.join(", ")
2899            ))
2900        }
2901    }
2902
2903    pub async fn load_from(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<bool> {
2904        let Some(stored) = Self::load_session_restore(storage, session_id).await? else {
2905            return Ok(false);
2906        };
2907        self.apply_session_restore(session_id, stored).await?;
2908        Ok(true)
2909    }
2910
2911    pub async fn save_session(&self, session_id: &str) -> Result<()> {
2912        let storage = self.storage.read().clone();
2913        match storage {
2914            Some(s) => {
2915                // Fire on_session_created when this session id is first seen on this runtime.
2916                let is_new = {
2917                    let cur = self.current_session_id.read().clone();
2918                    cur.as_deref() != Some(session_id)
2919                };
2920                if is_new {
2921                    *self.current_session_id.write() = Some(session_id.to_string());
2922                    self.hooks.on_session_created(session_id).await;
2923                }
2924
2925                // Update metadata before persisting.
2926                {
2927                    let now = chrono::Utc::now();
2928                    let msg_count = self
2929                        .memory
2930                        .get_messages(None)
2931                        .await
2932                        .map(|v| v.len())
2933                        .unwrap_or(0);
2934                    let mut meta = self.session_metadata.write();
2935                    meta.last_active = now;
2936                    meta.message_count = msg_count;
2937                    if meta.actor_id.is_none() {
2938                        meta.actor_id = self.actor_id.read().clone();
2939                    }
2940                }
2941
2942                let snapshot = self.save_state().await?;
2943                //
2944                // Metadata-capable backends own atomic snapshot and metadata persistence.
2945                //
2946                if s.supports(StorageCapability::SessionMetadata) {
2947                    let metadata = self.session_metadata.read().clone();
2948                    s.save_snapshot_with_metadata(session_id, &snapshot, &metadata)
2949                        .await
2950                } else {
2951                    s.save(session_id, &snapshot).await
2952                }
2953            }
2954            None => Err(AgentError::Config(
2955                "No storage configured. Use with_storage_config() or with_storage() first".into(),
2956            )),
2957        }
2958    }
2959
2960    pub async fn load_session(&self, session_id: &str) -> Result<bool> {
2961        let storage = self.storage.read().clone();
2962        match storage {
2963            Some(storage) => self.load_from(storage.as_ref(), session_id).await,
2964            None => Err(AgentError::Config(
2965                "No storage configured. Use with_storage_config() or with_storage() first".into(),
2966            )),
2967        }
2968    }
2969
2970    /// Restore this runtime and its complete saved child topology from one named session.
2971    pub async fn restore_session_full(&self, session_id: &str) -> Result<usize> {
2972        self.init_storage().await?;
2973        let storage = self.storage.read().clone().ok_or_else(|| {
2974            AgentError::Config(
2975                "No storage configured. Use with_storage_config() or with_storage() first".into(),
2976            )
2977        })?;
2978        let target_parent = Self::load_session_restore(storage.as_ref(), session_id)
2979            .await?
2980            .ok_or_else(|| AgentError::Persistence(format!("Session not found: {session_id}")))?;
2981        let manifest = target_parent
2982            .snapshot
2983            .spawned_agents
2984            .clone()
2985            .unwrap_or_default();
2986
2987        let registry = self.spawner_registry.as_ref().cloned();
2988        let spawner = if manifest.is_empty() {
2989            self.spawner.as_ref().cloned()
2990        } else {
2991            Some(self.spawner.as_ref().cloned().ok_or_else(|| {
2992                AgentError::Config(
2993                    "Saved session contains child agents but this runtime has no spawner".into(),
2994                )
2995            })?)
2996        };
2997        let registry = if manifest.is_empty() {
2998            registry
2999        } else {
3000            Some(registry.ok_or_else(|| {
3001                AgentError::Config(
3002                    "Saved session contains child agents but this runtime has no registry".into(),
3003                )
3004            })?)
3005        };
3006
3007        let mut target_ids = HashSet::with_capacity(manifest.len());
3008        let mut prepared = Vec::with_capacity(manifest.len());
3009        for entry in manifest {
3010            if !target_ids.insert(entry.id.clone()) {
3011                return Err(AgentError::InvalidSpec(format!(
3012                    "Saved child manifest contains duplicate ID: {}",
3013                    entry.id
3014                )));
3015            }
3016            let spec = crate::spec::AgentSpec::from_yaml_strict(&entry.spec_yaml)?;
3017            spawner
3018                .as_ref()
3019                .expect("non-empty manifests require a spawner")
3020                .validate_explicit_child(&entry.id, &spec)?;
3021            prepared.push((entry.id, spec));
3022        }
3023
3024        let current_ids = registry
3025            .as_ref()
3026            .map(|registry| {
3027                registry
3028                    .list()
3029                    .into_iter()
3030                    .map(|info| info.id)
3031                    .collect::<HashSet<_>>()
3032            })
3033            .unwrap_or_default();
3034        let removal_count = current_ids.difference(&target_ids).count();
3035        let additions = prepared
3036            .iter()
3037            .filter(|(id, _)| !current_ids.contains(id))
3038            .cloned()
3039            .collect::<Vec<_>>();
3040
3041        let mut existing = Vec::new();
3042        if let Some(registry) = registry.as_ref() {
3043            for (id, _) in prepared.iter().filter(|(id, _)| current_ids.contains(id)) {
3044                let agent = registry.get(id).ok_or_else(|| {
3045                    AgentError::Config(format!("Retained child disappeared during restore: {id}"))
3046                })?;
3047                let child_storage = agent.storage().ok_or_else(|| {
3048                    AgentError::Config(format!("Child '{id}' has no storage for session restore"))
3049                })?;
3050                let stored = Self::load_session_restore(child_storage.as_ref(), session_id)
3051                    .await?
3052                    .ok_or_else(|| {
3053                        AgentError::Persistence(format!(
3054                            "Child '{id}' has no saved session '{session_id}'"
3055                        ))
3056                    })?;
3057                existing.push((id.clone(), agent, stored));
3058            }
3059        }
3060
3061        let mut staged = Vec::with_capacity(additions.len());
3062        if !additions.is_empty() {
3063            let spawner = spawner
3064                .as_ref()
3065                .expect("restored additions require a spawner");
3066            let reservations = spawner.reserve_restore_capacity(additions.len(), removal_count)?;
3067            for ((id, spec), reservation) in additions.into_iter().zip(reservations) {
3068                let spawned = spawner
3069                    .spawn_with_reserved_capacity(id.clone(), spec, reservation)
3070                    .await?;
3071                let child_storage = spawned.agent.storage().ok_or_else(|| {
3072                    AgentError::Config(format!("Child '{id}' has no storage for session restore"))
3073                })?;
3074                let stored = Self::load_session_restore(child_storage.as_ref(), session_id)
3075                    .await?
3076                    .ok_or_else(|| {
3077                        AgentError::Persistence(format!(
3078                            "Child '{id}' has no saved session '{session_id}'"
3079                        ))
3080                    })?;
3081                staged.push((spawned, stored));
3082            }
3083        } else if let Some(spawner) = spawner.as_ref() {
3084            spawner.reserve_restore_capacity(0, removal_count)?;
3085        }
3086
3087        let parent_before = self.capture_session_restore_point().await?;
3088        let mut existing_before = Vec::with_capacity(existing.len());
3089        for (id, agent, _) in &existing {
3090            existing_before.push((
3091                id.clone(),
3092                Arc::clone(agent),
3093                agent.capture_session_restore_point().await?,
3094            ));
3095        }
3096
3097        //
3098        // Every record is loaded before runtime state changes. Registry replacement is the final topology commit point.
3099        //
3100        for (_, agent, stored) in &existing {
3101            if let Err(error) = agent
3102                .apply_session_restore_unchecked(session_id, stored.clone())
3103                .await
3104            {
3105                drop(staged);
3106                let rollback_errors =
3107                    Self::rollback_session_restore_set(None, &existing_before).await;
3108                return Err(Self::restore_failure(error, rollback_errors));
3109            }
3110        }
3111        for (spawned, stored) in &staged {
3112            if let Err(error) = spawned
3113                .agent
3114                .apply_session_restore_unchecked(session_id, stored.clone())
3115                .await
3116            {
3117                drop(staged);
3118                let rollback_errors =
3119                    Self::rollback_session_restore_set(None, &existing_before).await;
3120                return Err(Self::restore_failure(error, rollback_errors));
3121            }
3122        }
3123        if let Err(error) = self
3124            .apply_session_restore_unchecked(session_id, target_parent)
3125            .await
3126        {
3127            drop(staged);
3128            let rollback_errors =
3129                Self::rollback_session_restore_set(Some((self, &parent_before)), &existing_before)
3130                    .await;
3131            return Err(Self::restore_failure(error, rollback_errors));
3132        }
3133
3134        if let Some(registry) = registry.as_ref()
3135            && let Err(error) = registry
3136                .reconcile(
3137                    &target_ids,
3138                    staged.into_iter().map(|(spawned, _)| spawned).collect(),
3139                )
3140                .await
3141        {
3142            let rollback_errors =
3143                Self::rollback_session_restore_set(Some((self, &parent_before)), &existing_before)
3144                    .await;
3145            return Err(Self::restore_failure(error, rollback_errors));
3146        }
3147
3148        Ok(target_ids.len())
3149    }
3150
3151    pub async fn delete_session(&self, session_id: &str) -> Result<()> {
3152        let storage = self.storage.read().clone();
3153        match storage {
3154            Some(s) => s.delete(session_id).await,
3155            None => Err(AgentError::Config(
3156                "No storage configured. Use with_storage_config() or with_storage() first".into(),
3157            )),
3158        }
3159    }
3160
3161    pub async fn list_sessions(&self) -> Result<Vec<String>> {
3162        let storage = self.storage.read().clone();
3163        match storage {
3164            Some(s) => s.list_sessions().await,
3165            None => Err(AgentError::Config(
3166                "No storage configured. Use with_storage_config() or with_storage() first".into(),
3167            )),
3168        }
3169    }
3170
3171    fn estimate_tokens(&self, text: &str) -> u32 {
3172        (text.len() as f32 / 4.0).ceil() as u32
3173    }
3174
3175    fn estimate_total_tokens(&self, messages: &[ChatMessage]) -> u32 {
3176        messages
3177            .iter()
3178            .map(|m| self.estimate_tokens(&m.content))
3179            .sum()
3180    }
3181
3182    fn truncate_context(&self, messages: &mut Vec<ChatMessage>, keep_recent: usize) {
3183        if messages.len() <= keep_recent + 1 {
3184            return;
3185        }
3186        let system_msg = messages.remove(0);
3187        let to_remove = messages.len().saturating_sub(keep_recent);
3188        messages.drain(..to_remove);
3189        messages.insert(0, system_msg);
3190    }
3191
3192    fn get_filter(&self, config: &FilterConfig) -> Arc<dyn MessageFilter> {
3193        match config {
3194            FilterConfig::KeepRecent(n) => Arc::new(KeepRecentFilter::new(*n)),
3195            FilterConfig::ByRole { keep_roles } => Arc::new(ByRoleFilter::new(keep_roles.clone())),
3196            FilterConfig::SkipPattern { skip_if_contains } => {
3197                Arc::new(SkipPatternFilter::new(skip_if_contains.clone()))
3198            }
3199            FilterConfig::Custom { name } => {
3200                let filters = self.message_filters.read();
3201                filters
3202                    .get(name)
3203                    .cloned()
3204                    .unwrap_or_else(|| Arc::new(KeepRecentFilter::new(10)))
3205            }
3206        }
3207    }
3208
3209    async fn summarize_context(
3210        &self,
3211        messages: &mut Vec<ChatMessage>,
3212        summarizer_llm: Option<&str>,
3213        max_summary_tokens: u32,
3214        custom_prompt: Option<&str>,
3215        keep_recent: usize,
3216        filter: Option<&FilterConfig>,
3217    ) -> Result<()> {
3218        let system_msg = messages.remove(0);
3219
3220        let to_summarize_count = messages.len().saturating_sub(keep_recent);
3221        if to_summarize_count == 0 {
3222            messages.insert(0, system_msg);
3223            return Ok(());
3224        }
3225
3226        let recent_msgs: Vec<ChatMessage> = messages.drain(to_summarize_count..).collect();
3227        let mut to_summarize = std::mem::take(messages);
3228
3229        if let Some(filter_config) = filter {
3230            let filter = self.get_filter(filter_config);
3231            to_summarize = filter.filter(to_summarize);
3232        }
3233
3234        if to_summarize.is_empty() {
3235            *messages = recent_msgs;
3236            messages.insert(0, system_msg);
3237            return Ok(());
3238        }
3239
3240        let conversation_text = to_summarize
3241            .iter()
3242            .map(|m| format!("{:?}: {}", m.role, m.content))
3243            .collect::<Vec<_>>()
3244            .join("\n");
3245
3246        let default_prompt = format!(
3247            "Summarize the following conversation in under {} tokens, preserving key information:\n\n{}",
3248            max_summary_tokens, conversation_text
3249        );
3250
3251        let summary_prompt = custom_prompt
3252            .map(|p| format!("{}\n\n{}", p, conversation_text))
3253            .unwrap_or(default_prompt);
3254
3255        let summarizer = if let Some(alias) = summarizer_llm {
3256            self.llm_registry
3257                .get(alias)
3258                .map_err(|e| AgentError::Config(e.to_string()))?
3259        } else {
3260            self.llm_registry
3261                .router()
3262                .or_else(|_| self.llm_registry.default())
3263                .map_err(|e| AgentError::Config(e.to_string()))?
3264        };
3265
3266        let summary_msgs = vec![ChatMessage::user(&summary_prompt)];
3267        let response = self
3268            .observe_purpose(
3269                ObservationPurpose::Summarization,
3270                summarizer.complete(&summary_msgs, None),
3271            )
3272            .await?;
3273
3274        let summary_message = ChatMessage::system(format!(
3275            "[Previous conversation summary]\n{}",
3276            response.content
3277        ));
3278
3279        *messages = vec![system_msg, summary_message];
3280        messages.extend(recent_msgs);
3281
3282        debug!(
3283            summarized_count = to_summarize_count,
3284            kept_recent = keep_recent,
3285            "Context summarized"
3286        );
3287
3288        Ok(())
3289    }
3290
3291    fn render_system_prompt(&self) -> Result<String> {
3292        let mut context = self.build_context_with_overlays();
3293
3294        // Inject actor_facts for {{ actor_facts }} template variable.
3295        let facts_text = self.format_actor_facts_for_context();
3296        if !facts_text.is_empty() {
3297            context.insert(
3298                "actor_facts".to_string(),
3299                serde_json::Value::String(facts_text),
3300            );
3301        }
3302
3303        if let Some((key, text)) = self.format_relationship_for_context() {
3304            context.insert(key, serde_json::Value::String(text));
3305        }
3306
3307        self.template_renderer
3308            .render(&self.base_system_prompt, &context)
3309    }
3310
3311    /// Canonicalizes IDs once while preserving their first declaration order and ignoring unknown entries.
3312    fn canonical_unique_tool_ids(&self, ids: &[String]) -> Vec<String> {
3313        let mut seen = HashSet::new();
3314        ids.iter()
3315            .filter_map(|id| self.tools.canonical_id(id))
3316            .filter(|canonical_id| seen.insert(canonical_id.clone()))
3317            .collect()
3318    }
3319
3320    /// Applies runtime narrowing to the canonical declared grant without permitting scope-based expansion.
3321    fn get_top_level_tool_ids_for_scope(&self, scope_override: Option<&[String]>) -> Vec<String> {
3322        let Some(declared) = self.declared_tool_ids.as_deref() else {
3323            return Vec::new();
3324        };
3325        let mut effective = self.canonical_unique_tool_ids(declared);
3326        if let Some(scope) = scope_override {
3327            let scope: HashSet<String> =
3328                self.canonical_unique_tool_ids(scope).into_iter().collect();
3329            effective.retain(|canonical_id| scope.contains(canonical_id));
3330        }
3331        effective
3332    }
3333
3334    /// Reads the live runtime narrowing and returns the current deterministic effective tool IDs.
3335    async fn get_available_tool_ids(&self) -> Result<Vec<String>> {
3336        Ok(self.get_available_tool_ids_snapshot().await?.tool_ids)
3337    }
3338
3339    /// Captures the runtime scope used for ordinary availability checks before applying state narrowing.
3340    async fn get_available_tool_ids_snapshot(&self) -> Result<AvailableToolIdsSnapshot> {
3341        let scope_override = self.runtime_control.tool_scope_override.read().clone();
3342        self.get_available_tool_ids_snapshot_for_scope(scope_override.as_deref())
3343            .await
3344    }
3345
3346    /// Applies every explicit ancestor and current-state scope to one runtime snapshot and retains its state generation.
3347    async fn get_available_tool_ids_snapshot_for_scope(
3348        &self,
3349        scope_override: Option<&[String]>,
3350    ) -> Result<AvailableToolIdsSnapshot> {
3351        let mut available = self.get_top_level_tool_ids_for_scope(scope_override);
3352        let (state_generation, state_scopes) = self
3353            .state_machine
3354            .as_ref()
3355            .map(|state_machine| {
3356                let (generation, scopes) = state_machine.current_tool_scope_snapshot();
3357                (Some(generation), scopes)
3358            })
3359            .unwrap_or((None, Vec::new()));
3360
3361        if available.is_empty() || state_scopes.is_empty() {
3362            return Ok(AvailableToolIdsSnapshot {
3363                tool_ids: available,
3364                state_generation,
3365            });
3366        }
3367
3368        let eval_ctx = self.build_evaluation_context().await?;
3369        let llm_getter = RegistryLLMGetter {
3370            registry: self.llm_registry.clone(),
3371        };
3372        let evaluator = ConditionEvaluator::new(llm_getter);
3373
3374        for state_scope in state_scopes {
3375            if state_scope.is_empty() {
3376                available.clear();
3377                break;
3378            }
3379
3380            let mut allowed = HashSet::new();
3381            for tool_ref in &state_scope {
3382                let tool_id = tool_ref.id();
3383                let Some(canonical_id) = self.tools.canonical_id(tool_id) else {
3384                    continue;
3385                };
3386                let condition_matches = if let Some(condition) = tool_ref.condition() {
3387                    match evaluator.evaluate(condition, &eval_ctx).await {
3388                        Ok(matches) => matches,
3389                        Err(error) => {
3390                            warn!(tool = tool_id, error = %error, "Error evaluating tool condition");
3391                            false
3392                        }
3393                    }
3394                } else {
3395                    true
3396                };
3397                if condition_matches {
3398                    allowed.insert(canonical_id);
3399                } else {
3400                    debug!(tool = tool_id, "Tool condition not met, skipping");
3401                }
3402            }
3403            available.retain(|canonical_id| allowed.contains(canonical_id));
3404            if available.is_empty() {
3405                break;
3406            }
3407        }
3408
3409        Ok(AvailableToolIdsSnapshot {
3410            tool_ids: available,
3411            state_generation,
3412        })
3413    }
3414
3415    async fn build_evaluation_context(&self) -> Result<EvaluationContext> {
3416        let context = self.build_context_with_overlays();
3417        let messages = self.memory.get_messages(Some(10)).await?;
3418        let tool_history = self.tool_call_history.read().clone();
3419
3420        let (state_name, turn_count, previous_state) = if let Some(ref sm) = self.state_machine {
3421            (Some(sm.current()), sm.turn_count(), sm.previous())
3422        } else {
3423            (None, 0, None)
3424        };
3425
3426        Ok(EvaluationContext::default()
3427            .with_context(context)
3428            .with_state(state_name, turn_count, previous_state)
3429            .with_called_tools(tool_history)
3430            .with_messages(messages))
3431    }
3432
3433    fn record_tool_call(&self, tool_id: &str, result: Value) {
3434        self.tool_call_history.write().push(ToolCallRecord {
3435            tool_id: tool_id.to_string(),
3436            result,
3437            timestamp: chrono::Utc::now(),
3438        });
3439    }
3440
3441    async fn get_effective_system_prompt_with_persona_hooks(
3442        &self,
3443        fire_persona_hooks: bool,
3444        include_tool_prompt: bool,
3445    ) -> Result<String> {
3446        let rendered_base = self.render_system_prompt()?;
3447
3448        let persona_prefix = if let Some(ref persona) = self.persona_manager {
3449            let context = self.build_context_with_overlays();
3450            if fire_persona_hooks {
3451                let render_result = persona.render_prompt(&context)?;
3452                for content in &render_result.newly_revealed {
3453                    self.hooks.on_secret_revealed(content).await;
3454                }
3455                render_result.prompt
3456            } else {
3457                persona.render_prompt_preview(&context)?
3458            }
3459        } else {
3460            String::new()
3461        };
3462
3463        if let Some(ref sm) = self.state_machine
3464            && let Some(state_def) = sm.current_definition()
3465        {
3466            let state_prompt = if let Some(ref prompt) = state_def.prompt {
3467                let context = self.build_context_with_overlays();
3468                self.template_renderer.render_with_state(
3469                    prompt,
3470                    &context,
3471                    &sm.current(),
3472                    sm.previous().as_deref(),
3473                    sm.turn_count(),
3474                    state_def.max_turns,
3475                )?
3476            } else {
3477                String::new()
3478            };
3479
3480            let combined = match state_def.prompt_mode {
3481                PromptMode::Append => {
3482                    if state_prompt.is_empty() {
3483                        rendered_base
3484                    } else {
3485                        format!(
3486                            "{}\n\n[Current State: {}]\n{}",
3487                            rendered_base,
3488                            sm.current(),
3489                            state_prompt
3490                        )
3491                    }
3492                }
3493                PromptMode::Replace => {
3494                    if state_prompt.is_empty() {
3495                        rendered_base
3496                    } else {
3497                        state_prompt
3498                    }
3499                }
3500                PromptMode::Prepend => {
3501                    if state_prompt.is_empty() {
3502                        rendered_base
3503                    } else {
3504                        format!("{}\n\n{}", state_prompt, rendered_base)
3505                    }
3506                }
3507            };
3508
3509            // Persona always prepended regardless of prompt_mode.
3510            let with_persona = if persona_prefix.is_empty() {
3511                combined
3512            } else {
3513                format!("{}\n\n{}", persona_prefix, combined)
3514            };
3515
3516            if include_tool_prompt {
3517                let available_tool_ids = self.get_available_tool_ids().await?;
3518                if !available_tool_ids.is_empty() {
3519                    let tools_prompt = self.tools.generate_scoped_prompt_with_mode(
3520                        &available_tool_ids,
3521                        None,
3522                        self.parallel_tools.enabled,
3523                        self.runtime_config.tool_schema_prompt_mode,
3524                    );
3525                    if !tools_prompt.is_empty() {
3526                        return Ok(format!("{}\n\n{}", with_persona, tools_prompt));
3527                    }
3528                }
3529            }
3530            return Ok(with_persona);
3531        }
3532
3533        // No state machine - prepend persona to base.
3534        let with_persona = if persona_prefix.is_empty() {
3535            rendered_base
3536        } else {
3537            format!("{}\n\n{}", persona_prefix, rendered_base)
3538        };
3539
3540        if include_tool_prompt {
3541            let available_tool_ids = self.get_available_tool_ids().await?;
3542            let tools_prompt = self.tools.generate_scoped_prompt_with_mode(
3543                &available_tool_ids,
3544                None,
3545                self.parallel_tools.enabled,
3546                self.runtime_config.tool_schema_prompt_mode,
3547            );
3548            if !tools_prompt.is_empty() {
3549                return Ok(format!("{}\n\n{}", with_persona, tools_prompt));
3550            }
3551        }
3552        Ok(with_persona)
3553    }
3554
3555    fn get_state_llm(&self) -> Result<Arc<dyn LLMProvider>> {
3556        if let Some(ref sm) = self.state_machine
3557            && let Some(state_def) = sm.current_definition()
3558            && let Some(ref llm_alias) = state_def.llm
3559        {
3560            return self
3561                .llm_registry
3562                .get(llm_alias)
3563                .map_err(|e| AgentError::Config(e.to_string()));
3564        }
3565        self.llm_registry
3566            .default()
3567            .map_err(|e| AgentError::Config(e.to_string()))
3568    }
3569
3570    fn get_effective_reasoning_config(&self) -> ReasoningConfig {
3571        if let Some(ref sm) = self.state_machine
3572            && let Some(state_def) = sm.current_definition()
3573            && let Some(ref state_reasoning) = state_def.reasoning
3574        {
3575            return state_reasoning.clone();
3576        }
3577        self.reasoning_config.clone()
3578    }
3579
3580    fn get_effective_reflection_config(&self) -> ReflectionConfig {
3581        if let Some(ref sm) = self.state_machine
3582            && let Some(state_def) = sm.current_definition()
3583            && let Some(ref state_reflection) = state_def.reflection
3584        {
3585            return state_reflection.clone();
3586        }
3587        self.reflection_config.clone()
3588    }
3589
3590    fn get_skill_reasoning_config(&self, skill: &SkillDefinition) -> ReasoningConfig {
3591        skill
3592            .reasoning
3593            .clone()
3594            .unwrap_or_else(|| self.get_effective_reasoning_config())
3595    }
3596
3597    fn get_skill_reflection_config(&self, skill: &SkillDefinition) -> ReflectionConfig {
3598        skill
3599            .reflection
3600            .clone()
3601            .unwrap_or_else(|| self.get_effective_reflection_config())
3602    }
3603
3604    async fn build_disambiguation_context(&self) -> Result<DisambiguationContext> {
3605        let recent_messages: Vec<String> = self
3606            .memory
3607            .get_messages(Some(5))
3608            .await?
3609            .iter()
3610            .rev()
3611            .map(|m| format!("{:?}: {}", m.role, m.content))
3612            .collect();
3613
3614        let current_state = self.current_state().map(|s| s.to_string());
3615
3616        // Include the current state's prompt text so the detector understands
3617        // what kind of input is expected (e.g., "Ask for the order number").
3618        let state_prompt: Option<String> = self
3619            .state_machine
3620            .as_ref()
3621            .and_then(|sm| sm.current_definition())
3622            .and_then(|def| def.prompt.clone());
3623
3624        let available_tools: Vec<String> = self
3625            .get_available_tool_ids()
3626            .await
3627            .unwrap_or_else(|_| self.tools.list_ids());
3628
3629        let available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
3630
3631        let mut user_context = self.build_context_with_overlays();
3632        user_context.remove(DISAMBIGUATION_STATE_GENERATION_KEY);
3633        if let Some(state_generation) = self
3634            .state_machine
3635            .as_ref()
3636            .map(|state_machine| state_machine.generation())
3637        {
3638            user_context.insert(
3639                DISAMBIGUATION_STATE_GENERATION_KEY.to_string(),
3640                serde_json::json!(state_generation),
3641            );
3642        }
3643
3644        // Extract canonical intent labels from current state's transitions
3645        let available_intents: Vec<String> = if let Some(ref sm) = self.state_machine {
3646            sm.current_definition()
3647                .map(|def| {
3648                    def.transitions
3649                        .iter()
3650                        .filter_map(|t| t.intent.clone())
3651                        .collect()
3652                })
3653                .unwrap_or_default()
3654        } else {
3655            Vec::new()
3656        };
3657
3658        Ok(DisambiguationContext::from_agent_state(
3659            recent_messages,
3660            current_state,
3661            state_prompt,
3662            available_tools,
3663            available_skills,
3664            available_intents,
3665            user_context,
3666        ))
3667    }
3668
3669    fn get_available_skills(&self) -> Vec<&SkillDefinition> {
3670        if let Some(ref sm) = self.state_machine
3671            && let Some(state_def) = sm.current_definition()
3672        {
3673            let parent_def = sm.get_parent_definition();
3674            let effective_skills = state_def.get_effective_skills(parent_def.as_ref());
3675            if !effective_skills.is_empty() {
3676                return self
3677                    .skills
3678                    .iter()
3679                    .filter(|s| effective_skills.contains(&&s.id))
3680                    .collect();
3681            }
3682        }
3683        self.skills.iter().collect()
3684    }
3685
3686    async fn build_messages(&self) -> Result<Vec<ChatMessage>> {
3687        self.build_messages_internal(true, None, true).await
3688    }
3689
3690    async fn build_messages_for_draft(&self, user_message: &str) -> Result<Vec<ChatMessage>> {
3691        self.build_messages_internal(false, Some(user_message), true)
3692            .await
3693    }
3694
3695    async fn build_messages_internal(
3696        &self,
3697        fire_persona_hooks: bool,
3698        ephemeral_user_message: Option<&str>,
3699        include_tool_prompt: bool,
3700    ) -> Result<Vec<ChatMessage>> {
3701        let system_prompt = self
3702            .get_effective_system_prompt_with_persona_hooks(fire_persona_hooks, include_tool_prompt)
3703            .await?;
3704        let mut messages = vec![ChatMessage::system(&system_prompt)];
3705
3706        let context = self.memory.get_context().await?;
3707        let history = if let Some(ref budget) = self.memory_token_budget {
3708            context.to_llm_messages_with_allocation(&budget.allocation)
3709        } else {
3710            context.to_llm_messages()
3711        };
3712        messages.extend(history);
3713        if let Some(user_message) = ephemeral_user_message {
3714            messages.push(ChatMessage::user(user_message));
3715        }
3716
3717        let total_tokens = self.estimate_total_tokens(&messages);
3718
3719        if total_tokens > self.max_context_tokens {
3720            debug!(
3721                total = total_tokens,
3722                limit = self.max_context_tokens,
3723                "Context overflow"
3724            );
3725
3726            match &self.recovery_manager.config().llm.on_context_overflow {
3727                ContextOverflowAction::Error => {
3728                    return Err(AgentError::LLM(format!(
3729                        "Context overflow: {} tokens > {} limit",
3730                        total_tokens, self.max_context_tokens
3731                    )));
3732                }
3733                ContextOverflowAction::Truncate { keep_recent } => {
3734                    self.truncate_context(&mut messages, *keep_recent);
3735                }
3736                ContextOverflowAction::Summarize {
3737                    summarizer_llm,
3738                    max_summary_tokens,
3739                    custom_prompt,
3740                    keep_recent,
3741                    filter,
3742                } => {
3743                    self.summarize_context(
3744                        &mut messages,
3745                        summarizer_llm.as_deref(),
3746                        *max_summary_tokens,
3747                        custom_prompt.as_deref(),
3748                        *keep_recent,
3749                        filter.as_ref(),
3750                    )
3751                    .await?;
3752                }
3753            }
3754        }
3755
3756        Ok(messages)
3757    }
3758
3759    async fn main_tool_protocol(
3760        &self,
3761        llm: &dyn LLMProvider,
3762        ephemeral_new_turn: bool,
3763    ) -> Result<MainToolProtocol> {
3764        let mut choice = llm.configured_tool_choice();
3765        if matches!(choice.as_ref(), Some(ToolChoice::None)) {
3766            return Ok(MainToolProtocol {
3767                choice,
3768                tool_ids: Vec::new(),
3769                definitions: Vec::new(),
3770            });
3771        }
3772
3773        let mut tool_ids = self.get_available_tool_ids().await?;
3774        tool_ids.sort();
3775        tool_ids.dedup();
3776        if let Some(ToolChoice::Specific(expected)) = choice.as_ref() {
3777            let canonical = self.tools.canonical_id(expected).ok_or_else(|| {
3778                AgentError::Config(format!(
3779                    "specific tool choice '{expected}' is not registered"
3780                ))
3781            })?;
3782            if canonical != *expected {
3783                return Err(AgentError::Config(format!(
3784                    "specific tool choice must use canonical ID '{canonical}', not '{expected}'"
3785                )));
3786            }
3787            if !tool_ids.iter().any(|tool_id| tool_id == expected) {
3788                return Err(AgentError::Config(format!(
3789                    "specific tool choice '{expected}' is outside the effective tool grant"
3790                )));
3791            }
3792        }
3793        if matches!(
3794            choice.as_ref(),
3795            Some(ToolChoice::Required | ToolChoice::Specific(_))
3796        ) && tool_ids.is_empty()
3797        {
3798            return Err(AgentError::Config(
3799                "required tool choice has no tool inside the effective grant".to_string(),
3800            ));
3801        }
3802        if !ephemeral_new_turn
3803            && let Some(configured_choice) = choice.as_ref()
3804            && matches!(
3805                configured_choice,
3806                ToolChoice::Required | ToolChoice::Specific(_)
3807            )
3808            && self
3809                .tool_choice_satisfied_in_current_turn(configured_choice, &tool_ids)
3810                .await?
3811        {
3812            choice = Some(ToolChoice::Auto);
3813        }
3814        if let Some(ToolChoice::Specific(expected)) = choice.as_ref() {
3815            tool_ids.retain(|tool_id| tool_id == expected);
3816        }
3817
3818        let definitions = tool_ids
3819            .iter()
3820            .map(|tool_id| {
3821                let tool = self.tools.get(tool_id).ok_or_else(|| {
3822                    AgentError::Config(format!(
3823                        "effective tool '{tool_id}' disappeared before provider exposure"
3824                    ))
3825                })?;
3826                Ok(LLMToolDefinition {
3827                    name: tool_id.clone(),
3828                    description: tool.description().to_string(),
3829                    input_schema: tool.input_schema(),
3830                })
3831            })
3832            .collect::<Result<Vec<_>>>()?;
3833
3834        //
3835        // Provider-visible definitions are derived only from the current effective grant. Tool choice never expands registration, scope, or authorization.
3836        //
3837        Ok(MainToolProtocol {
3838            choice,
3839            tool_ids,
3840            definitions,
3841        })
3842    }
3843
3844    async fn tool_choice_satisfied_in_current_turn(
3845        &self,
3846        choice: &ToolChoice,
3847        effective_tool_ids: &[String],
3848    ) -> Result<bool> {
3849        let messages = self.memory.get_messages(None).await?;
3850        let mut saw_tool_result = false;
3851        for message in messages.iter().rev() {
3852            match message.role {
3853                ai_agents_core::Role::Tool | ai_agents_core::Role::Function => {
3854                    saw_tool_result = true;
3855                }
3856                ai_agents_core::Role::Assistant if saw_tool_result => {
3857                    let Some(calls) = self.parse_tool_calls(&message.content) else {
3858                        continue;
3859                    };
3860                    let calls_are_effective = !calls.is_empty()
3861                        && calls.iter().all(|call| {
3862                            self.tools
3863                                .canonical_id(&call.name)
3864                                .is_some_and(|canonical| effective_tool_ids.contains(&canonical))
3865                        });
3866                    return Ok(calls_are_effective
3867                        && match choice {
3868                            ToolChoice::Required => true,
3869                            ToolChoice::Specific(expected) => calls.iter().all(|call| {
3870                                self.tools.canonical_id(&call.name).as_deref()
3871                                    == Some(expected.as_str())
3872                            }),
3873                            _ => false,
3874                        });
3875                }
3876                ai_agents_core::Role::User => return Ok(false),
3877                _ => {}
3878            }
3879        }
3880        Ok(false)
3881    }
3882
3883    fn provider_can_use_native_tools(
3884        &self,
3885        llm: &dyn LLMProvider,
3886        protocol: &MainToolProtocol,
3887    ) -> bool {
3888        let Some(choice) = protocol.choice.as_ref() else {
3889            return false;
3890        };
3891        if matches!(choice, ToolChoice::None) || protocol.definitions.is_empty() {
3892            return false;
3893        }
3894        llm.supports_tool_choice(choice)
3895            && protocol.definitions.iter().all(|definition| {
3896                !definition.name.is_empty()
3897                    && definition.name.len() <= 64
3898                    && definition
3899                        .name
3900                        .bytes()
3901                        .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
3902            })
3903    }
3904
3905    fn prompt_messages_for_tool_protocol(
3906        &self,
3907        messages: &[ChatMessage],
3908        protocol: &MainToolProtocol,
3909        corrective: bool,
3910    ) -> Vec<ChatMessage> {
3911        let mut messages = messages.to_vec();
3912        let Some(choice) = protocol.choice.as_ref() else {
3913            return messages;
3914        };
3915        if matches!(choice, ToolChoice::None) || protocol.tool_ids.is_empty() {
3916            return messages;
3917        }
3918
3919        let mut tool_prompt = self.tools.generate_scoped_prompt_with_mode(
3920            &protocol.tool_ids,
3921            None,
3922            self.parallel_tools.enabled,
3923            self.runtime_config.tool_schema_prompt_mode,
3924        );
3925        match choice {
3926            ToolChoice::Required => tool_prompt.push_str(
3927                "\n\nYou must call at least one listed tool before giving a final answer.",
3928            ),
3929            ToolChoice::Specific(tool_id) => tool_prompt.push_str(&format!(
3930                "\n\nYou must call the '{tool_id}' tool before giving a final answer."
3931            )),
3932            ToolChoice::Auto => {}
3933            ToolChoice::None => return messages,
3934            _ => return messages,
3935        }
3936        if let Some(system) = messages
3937            .iter_mut()
3938            .find(|message| message.role == ai_agents_core::Role::System)
3939        {
3940            system.content.push_str("\n\n");
3941            system.content.push_str(&tool_prompt);
3942        } else {
3943            messages.insert(0, ChatMessage::system(tool_prompt));
3944        }
3945        if corrective {
3946            let instruction = match choice {
3947                ToolChoice::Required => {
3948                    "Your previous response did not call a required tool. Call at least one listed tool now and return only the JSON tool call."
3949                }
3950                ToolChoice::Specific(tool_id) => {
3951                    messages.push(ChatMessage::user(format!(
3952                        "Your previous response did not call the required '{tool_id}' tool. Call it now and return only the JSON tool call."
3953                    )));
3954                    return messages;
3955                }
3956                _ => return messages,
3957            };
3958            messages.push(ChatMessage::user(instruction));
3959        }
3960        messages
3961    }
3962
3963    async fn invoke_main_provider(
3964        &self,
3965        llm: Arc<dyn LLMProvider>,
3966        messages: &[ChatMessage],
3967        protocol: &MainToolProtocol,
3968        corrective: bool,
3969    ) -> std::result::Result<MainProviderResponse, LLMError> {
3970        let use_native = self.provider_can_use_native_tools(llm.as_ref(), protocol);
3971        let response = if use_native {
3972            let request = LLMToolRequest {
3973                tools: protocol.definitions.clone(),
3974                choice: protocol
3975                    .choice
3976                    .clone()
3977                    .expect("native tool requests require an explicit choice"),
3978            };
3979            self.observe_purpose(
3980                ObservationPurpose::MainResponse,
3981                llm.complete_with_tools(messages, None, &request),
3982            )
3983            .await?
3984        } else {
3985            let prompt_messages =
3986                self.prompt_messages_for_tool_protocol(messages, protocol, corrective);
3987            self.observe_purpose(
3988                ObservationPurpose::MainResponse,
3989                llm.complete(&prompt_messages, None),
3990            )
3991            .await?
3992        };
3993        Ok(MainProviderResponse {
3994            response,
3995            used_native_tools: use_native,
3996        })
3997    }
3998
3999    async fn complete_main_attempt_with_recovery(
4000        &self,
4001        llm: Arc<dyn LLMProvider>,
4002        messages: &[ChatMessage],
4003        protocol: &MainToolProtocol,
4004        corrective: bool,
4005    ) -> Result<MainProviderResponse> {
4006        let primary_result = if self.recovery_manager.config().default.max_retries > 0 {
4007            self.recovery_manager
4008                .with_retry("llm_call", None, || {
4009                    let llm = Arc::clone(&llm);
4010                    async move {
4011                        self.invoke_main_provider(llm, messages, protocol, corrective)
4012                            .await
4013                            .map_err(|error| error.classify())
4014                    }
4015                })
4016                .await
4017                .map_err(|error| AgentError::LLM(error.to_string()))
4018        } else {
4019            self.invoke_main_provider(Arc::clone(&llm), messages, protocol, corrective)
4020                .await
4021                .map_err(|error| AgentError::LLM(error.to_string()))
4022        };
4023
4024        match primary_result {
4025            Ok(response) => Ok(response),
4026            Err(primary_error) => match &self.recovery_manager.config().llm.on_failure {
4027                LLMFailureAction::FallbackLlm { fallback_llm } => {
4028                    let fallback = self.llm_registry.get(fallback_llm).map_err(|error| {
4029                        AgentError::Config(format!(
4030                            "Fallback LLM '{fallback_llm}' not found: {error}"
4031                        ))
4032                    })?;
4033                    self.invoke_main_provider(fallback, messages, protocol, corrective)
4034                        .await
4035                        .map_err(|error| AgentError::LLM(error.to_string()))
4036                }
4037                LLMFailureAction::FallbackResponse { message } => {
4038                    if matches!(
4039                        protocol.choice.as_ref(),
4040                        Some(ToolChoice::Required | ToolChoice::Specific(_))
4041                    ) {
4042                        Err(AgentError::LLM(format!(
4043                            "Required tool selection failed and cannot be satisfied by a static fallback response: {primary_error}"
4044                        )))
4045                    } else {
4046                        Ok(MainProviderResponse {
4047                            response: LLMResponse::new(message.clone(), FinishReason::Stop),
4048                            used_native_tools: false,
4049                        })
4050                    }
4051                }
4052                LLMFailureAction::Error => Err(primary_error),
4053            },
4054        }
4055    }
4056
4057    fn normalize_main_provider_response(
4058        &self,
4059        mut response: LLMResponse,
4060        protocol: &MainToolProtocol,
4061    ) -> Result<(LLMResponse, bool)> {
4062        let native_calls = response
4063            .tool_calls()
4064            .map_err(|error| AgentError::LLM(error.to_string()))?;
4065        let calls = match native_calls {
4066            Some(calls) => {
4067                let markers = calls
4068                    .iter()
4069                    .map(|call| {
4070                        serde_json::json!({
4071                            "_ai_agents_native_tool_call": true,
4072                            "id": call.id,
4073                            "tool": call.name,
4074                            "arguments": call.arguments,
4075                        })
4076                    })
4077                    .collect::<Vec<_>>();
4078                response.content = if markers.len() == 1 {
4079                    markers[0].to_string()
4080                } else {
4081                    serde_json::Value::Array(markers).to_string()
4082                };
4083                Some(calls)
4084            }
4085            None if !matches!(protocol.choice.as_ref(), Some(ToolChoice::None)) => {
4086                self.parse_tool_calls(response.content.trim())
4087            }
4088            None => None,
4089        };
4090
4091        if protocol.choice.is_some()
4092            && let Some(calls) = calls.as_ref()
4093            && calls.iter().any(|call| {
4094                self.tools
4095                    .canonical_id(&call.name)
4096                    .is_none_or(|canonical| !protocol.tool_ids.contains(&canonical))
4097            })
4098        {
4099            return Err(AgentError::LLM(
4100                "Provider returned a tool call outside the effective grant".to_string(),
4101            ));
4102        }
4103
4104        let compliant = match protocol.choice.as_ref() {
4105            Some(ToolChoice::Required) => calls.as_ref().is_some_and(|calls| !calls.is_empty()),
4106            Some(ToolChoice::Specific(expected)) => calls.as_ref().is_some_and(|calls| {
4107                !calls.is_empty()
4108                    && calls.iter().all(|call| {
4109                        self.tools.canonical_id(&call.name).as_deref() == Some(expected.as_str())
4110                    })
4111            }),
4112            _ => true,
4113        };
4114        Ok((response, compliant))
4115    }
4116
4117    async fn complete_main_llm_with_recovery(
4118        &self,
4119        llm: Arc<dyn LLMProvider>,
4120        messages: &[ChatMessage],
4121        protocol: &MainToolProtocol,
4122    ) -> Result<LLMResponse> {
4123        let first = self
4124            .complete_main_attempt_with_recovery(Arc::clone(&llm), messages, protocol, false)
4125            .await?;
4126        let (response, compliant) =
4127            self.normalize_main_provider_response(first.response, protocol)?;
4128        if compliant {
4129            return Ok(response);
4130        }
4131        if first.used_native_tools {
4132            return Err(AgentError::LLM(
4133                "Provider returned no compliant native call for required tool choice".to_string(),
4134            ));
4135        }
4136
4137        let corrected = self
4138            .complete_main_attempt_with_recovery(llm, messages, protocol, true)
4139            .await?;
4140        let (response, compliant) =
4141            self.normalize_main_provider_response(corrected.response, protocol)?;
4142        if compliant {
4143            return Ok(response);
4144        }
4145        Err(AgentError::LLM(
4146            "Provider returned no compliant tool call after one corrective retry".to_string(),
4147        ))
4148    }
4149
4150    fn is_native_tool_call_content(content: &str) -> bool {
4151        let Ok(value) = serde_json::from_str::<serde_json::Value>(content) else {
4152            return false;
4153        };
4154        match value {
4155            serde_json::Value::Array(values) => {
4156                !values.is_empty()
4157                    && values.iter().all(|value| {
4158                        value
4159                            .get("_ai_agents_native_tool_call")
4160                            .and_then(|marker| marker.as_bool())
4161                            == Some(true)
4162                    })
4163            }
4164            serde_json::Value::Object(map) => {
4165                map.get("_ai_agents_native_tool_call")
4166                    .and_then(|marker| marker.as_bool())
4167                    == Some(true)
4168            }
4169            _ => false,
4170        }
4171    }
4172
4173    fn tool_result_message(
4174        tool_call: &ToolCall,
4175        output: &str,
4176        native_tool_call: bool,
4177    ) -> ChatMessage {
4178        if !native_tool_call {
4179            return ChatMessage::function(&tool_call.name, output);
4180        }
4181        let output = serde_json::from_str::<serde_json::Value>(output)
4182            .unwrap_or_else(|_| serde_json::Value::String(output.to_string()));
4183        ChatMessage::function(
4184            &tool_call.name,
4185            serde_json::json!({
4186                "_ai_agents_native_tool_result": true,
4187                "id": tool_call.id,
4188                "tool": tool_call.name,
4189                "output": output,
4190            })
4191            .to_string(),
4192        )
4193    }
4194
4195    fn parse_main_tool_calls(
4196        &self,
4197        content: &str,
4198        protocol: &MainToolProtocol,
4199    ) -> Option<Vec<ToolCall>> {
4200        if matches!(protocol.choice.as_ref(), Some(ToolChoice::None)) {
4201            None
4202        } else {
4203            self.parse_tool_calls(content)
4204        }
4205    }
4206
4207    fn parse_tool_calls(&self, content: &str) -> Option<Vec<ToolCall>> {
4208        // Try direct JSON parse first
4209        if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(content) {
4210            // Handle JSON array of tool calls (parallel tool calling)
4211            if let Some(arr) = parsed.as_array() {
4212                let calls: Vec<ToolCall> = arr
4213                    .iter()
4214                    .filter_map(|v| self.extract_tool_call_from_value(v))
4215                    .collect();
4216                if !calls.is_empty() {
4217                    return Some(calls);
4218                }
4219            }
4220            // Handle single JSON object
4221            if let Some(tool_call) = self.extract_tool_call_from_value(&parsed) {
4222                return Some(vec![tool_call]);
4223            }
4224        }
4225
4226        // Try to extract JSON from content (handles extra text/braces from LLM)
4227        if let Some(json_str) = self.extract_json_from_content(content)
4228            && let Ok(parsed) = serde_json::from_str::<serde_json::Value>(&json_str)
4229        {
4230            // Handle JSON array of tool calls (parallel tool calling)
4231            if let Some(arr) = parsed.as_array() {
4232                let calls: Vec<ToolCall> = arr
4233                    .iter()
4234                    .filter_map(|v| self.extract_tool_call_from_value(v))
4235                    .collect();
4236                if !calls.is_empty() {
4237                    return Some(calls);
4238                }
4239            }
4240            // Handle single JSON object
4241            if let Some(tool_call) = self.extract_tool_call_from_value(&parsed) {
4242                return Some(vec![tool_call]);
4243            }
4244        }
4245
4246        None
4247    }
4248
4249    fn extract_tool_call_from_value(&self, parsed: &serde_json::Value) -> Option<ToolCall> {
4250        if let Some(tool_name) = parsed.get("tool").and_then(|v| v.as_str()) {
4251            let arguments = parsed
4252                .get("arguments")
4253                .cloned()
4254                .unwrap_or(serde_json::json!({}));
4255            return Some(ToolCall {
4256                id: parsed
4257                    .get("id")
4258                    .and_then(|value| value.as_str())
4259                    .filter(|id| !id.is_empty())
4260                    .map(str::to_string)
4261                    .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
4262                name: tool_name.to_string(),
4263                arguments,
4264            });
4265        }
4266        None
4267    }
4268
4269    // Lite models could generate unmatched braces: this function handles such cases
4270    fn extract_json_from_content(&self, content: &str) -> Option<String> {
4271        // Try array first (for parallel tool calls), then single object
4272        if let Some(result) = self.extract_json_array_from_content(content) {
4273            return Some(result);
4274        }
4275        self.extract_json_object_from_content(content)
4276    }
4277
4278    /// Extract a JSON array `[...]` containing tool calls from mixed content.
4279    fn extract_json_array_from_content(&self, content: &str) -> Option<String> {
4280        let start = content.find('[')?;
4281        let content_from_start = &content[start..];
4282
4283        let mut depth = 0;
4284        let mut end = 0;
4285        for (i, ch) in content_from_start.char_indices() {
4286            match ch {
4287                '[' => depth += 1,
4288                ']' => {
4289                    depth -= 1;
4290                    if depth == 0 {
4291                        end = i + 1;
4292                        break;
4293                    }
4294                }
4295                _ => {}
4296            }
4297        }
4298
4299        if end > 0 {
4300            let json_str = &content_from_start[..end];
4301            // Verify it looks like an array of tool calls
4302            if json_str.contains("\"tool\"") {
4303                return Some(json_str.to_string());
4304            }
4305        }
4306
4307        None
4308    }
4309
4310    /// Extract a JSON object `{...}` containing a tool call from mixed content.
4311    fn extract_json_object_from_content(&self, content: &str) -> Option<String> {
4312        let start = content.find('{')?;
4313        let content_from_start = &content[start..];
4314
4315        // Count braces to find the matching closing brace
4316        let mut depth = 0;
4317        let mut end = 0;
4318        for (i, ch) in content_from_start.char_indices() {
4319            match ch {
4320                '{' => depth += 1,
4321                '}' => {
4322                    depth -= 1;
4323                    if depth == 0 {
4324                        end = i + 1;
4325                        break;
4326                    }
4327                }
4328                _ => {}
4329            }
4330        }
4331
4332        if end > 0 {
4333            let json_str = &content_from_start[..end];
4334            // Verify it looks like a tool call
4335            if json_str.contains("\"tool\"") {
4336                return Some(json_str.to_string());
4337            }
4338        }
4339
4340        None
4341    }
4342
4343    /// Builds a structured record from executor state.
4344    ///
4345    /// The explicit fields preserve one audit boundary for execution, policy, approval, timing, and generation evidence.
4346    #[allow(clippy::too_many_arguments)]
4347    fn record_from_parts(
4348        &self,
4349        request: &ToolExecutionRequest,
4350        canonical_id: String,
4351        executed_arguments: Value,
4352        started_at: chrono::DateTime<chrono::Utc>,
4353        start: Instant,
4354        executed: bool,
4355        success: bool,
4356        output: String,
4357        metadata: HashMap<String, Value>,
4358        policy: ToolPolicyDecisionRecord,
4359        approval: Option<ToolApprovalRecord>,
4360        timed_out: bool,
4361        output_truncated: bool,
4362    ) -> ToolExecutionRecord {
4363        let versions = ToolDecisionVersions {
4364            policy: self.active_tool_security().policy_version(),
4365            registry: self.tools.version(),
4366            runtime_control: self.runtime_control.version.load(Ordering::SeqCst),
4367            state: self
4368                .state_machine
4369                .as_ref()
4370                .map(|state_machine| state_machine.generation()),
4371        };
4372        self.record_from_parts_at(
4373            request,
4374            canonical_id,
4375            executed_arguments,
4376            started_at,
4377            start,
4378            executed,
4379            success,
4380            output,
4381            metadata,
4382            policy,
4383            approval,
4384            timed_out,
4385            output_truncated,
4386            versions,
4387        )
4388    }
4389
4390    /// Builds evidence with the exact generations used by the corresponding decision.
4391    #[allow(clippy::too_many_arguments)]
4392    fn record_from_parts_at(
4393        &self,
4394        request: &ToolExecutionRequest,
4395        canonical_id: String,
4396        executed_arguments: Value,
4397        started_at: chrono::DateTime<chrono::Utc>,
4398        start: Instant,
4399        executed: bool,
4400        success: bool,
4401        output: String,
4402        metadata: HashMap<String, Value>,
4403        policy: ToolPolicyDecisionRecord,
4404        approval: Option<ToolApprovalRecord>,
4405        timed_out: bool,
4406        output_truncated: bool,
4407        versions: ToolDecisionVersions,
4408    ) -> ToolExecutionRecord {
4409        ToolExecutionRecord {
4410            call_id: request.call_id.clone(),
4411            requested_name: request.requested_name.clone(),
4412            canonical_id,
4413            source: request.source.clone(),
4414            arguments: request.arguments.clone(),
4415            executed_arguments,
4416            policy_version: versions.policy,
4417            registry_version: versions.registry,
4418            runtime_config_version: versions.runtime_control,
4419            executed,
4420            success,
4421            output,
4422            metadata,
4423            policy,
4424            approval,
4425            started_at,
4426            duration_ms: start.elapsed().as_millis() as u64,
4427            timed_out,
4428            cancelled: false,
4429            cancellation_reason: None,
4430            output_truncated,
4431        }
4432    }
4433
4434    /// Sends a finalized tool record to hooks, history, and error handling.
4435    async fn finish_tool_record(&self, record: &ToolExecutionRecord) {
4436        let result = ToolResult {
4437            success: record.success,
4438            output: record.model_output_string(),
4439            metadata: if record.metadata.is_empty() {
4440                None
4441            } else {
4442                Some(record.metadata.clone())
4443            },
4444        };
4445        self.hooks
4446            .on_tool_complete(&record.canonical_id, &result, record.duration_ms)
4447            .await;
4448        self.hooks.on_tool_execution_record(record).await;
4449        self.record_tool_call(&record.canonical_id, record.model_output_value());
4450        if !record.success {
4451            self.hooks
4452                .on_error(&AgentError::Tool(record.output.clone()))
4453                .await;
4454        }
4455    }
4456
4457    /// Releases resource locks before hooks run because hooks may invoke another tool.
4458    async fn finish_tool_record_after_resource_guards(
4459        &self,
4460        resource_guards: ToolResourceGuards,
4461        record: &ToolExecutionRecord,
4462    ) {
4463        drop(resource_guards);
4464        self.finish_tool_record(record).await;
4465    }
4466
4467    /// Invokes a resolved tool once with timeout, cancellation, and actor context.
4468    async fn execute_resolved_tool_once(
4469        &self,
4470        tool: Arc<dyn ai_agents_core::Tool>,
4471        args: Value,
4472        ctx: ToolExecutionContext,
4473        timeout_ms: u64,
4474    ) -> Result<(ToolResult, bool, bool, bool)> {
4475        if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4476            return Ok((
4477                ToolResult::error("Tool execution cancelled by runtime control"),
4478                false,
4479                true,
4480                false,
4481            ));
4482        }
4483        //
4484        // Mark invocation inside the future so cancellation before the first poll remains executed false.
4485        //
4486        let invoked = Arc::new(AtomicBool::new(false));
4487        let invoked_by_future = Arc::clone(&invoked);
4488        let actor_context = current_turn_actor_context();
4489        let future = async move {
4490            invoked_by_future.store(true, Ordering::SeqCst);
4491            if let Some(actor_context) = actor_context {
4492                scope_actor_context(actor_context, tool.execute(args, ctx)).await
4493            } else {
4494                tool.execute(args, ctx).await
4495            }
4496        };
4497        tokio::pin!(future);
4498        let timeout = tokio::time::sleep(std::time::Duration::from_millis(timeout_ms));
4499        tokio::pin!(timeout);
4500        let mut cancel_tick = tokio::time::interval(std::time::Duration::from_millis(50));
4501
4502        loop {
4503            tokio::select! {
4504                result = &mut future => return Ok((result, false, false, true)),
4505                _ = &mut timeout => {
4506                    return Ok((
4507                        ToolResult::error("Tool execution timed out"),
4508                        true,
4509                        false,
4510                        invoked.load(Ordering::SeqCst),
4511                    ));
4512                }
4513                _ = cancel_tick.tick() => {
4514                    if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4515                        return Ok((
4516                            ToolResult::error("Tool execution cancelled by runtime control"),
4517                            false,
4518                            true,
4519                            invoked.load(Ordering::SeqCst),
4520                        ));
4521                    }
4522                }
4523            }
4524        }
4525    }
4526
4527    /// Truncates model-facing tool output on a character boundary.
4528    fn truncate_tool_output(output: String, max_chars: Option<usize>) -> (String, bool) {
4529        let Some(max_chars) = max_chars else {
4530            return (output, false);
4531        };
4532        let mut chars = output.chars();
4533        let truncated: String = chars.by_ref().take(max_chars).collect();
4534        if chars.next().is_some() {
4535            (truncated, true)
4536        } else {
4537            (output, false)
4538        }
4539    }
4540
4541    /// Acquires all declared resource locks in stable key order and supports emergency cancellation while waiting.
4542    async fn acquire_tool_resource_locks(&self, keys: &[String]) -> Option<ToolResourceGuards> {
4543        let locks = {
4544            let mut table = self.resource_locks.write();
4545            table.retain(|_, lock| lock.strong_count() > 0);
4546            keys.iter()
4547                .map(|key| {
4548                    if let Some(lock) = table.get(key).and_then(Weak::upgrade) {
4549                        lock
4550                    } else {
4551                        let lock = Arc::new(tokio::sync::Mutex::new(()));
4552                        table.insert(key.clone(), Arc::downgrade(&lock));
4553                        lock
4554                    }
4555                })
4556                .collect::<Vec<_>>()
4557        };
4558        let mut resource_guards = ToolResourceGuards {
4559            guards: Vec::with_capacity(locks.len()),
4560            locks: Arc::clone(&self.resource_locks),
4561        };
4562        let mut locks = locks.into_iter();
4563        while let Some(lock) = locks.next() {
4564            let mut lock = Box::pin(lock.lock_owned());
4565            loop {
4566                tokio::select! {
4567                    guard = &mut lock => {
4568                        resource_guards.guards.push(guard);
4569                        break;
4570                    }
4571                    _ = tokio::time::sleep(std::time::Duration::from_millis(10)) => {
4572                        if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4573                            drop(lock);
4574                            drop(locks);
4575                            drop(resource_guards);
4576                            return None;
4577                        }
4578                    }
4579                }
4580            }
4581        }
4582        Some(resource_guards)
4583    }
4584
4585    /// Executes a resolved tool with retry policy from recovery settings.
4586    async fn run_tool_with_retries(
4587        &self,
4588        canonical_id: &str,
4589        tool: Arc<dyn ai_agents_core::Tool>,
4590        args: Value,
4591        ctx: ToolExecutionContext,
4592        timeout_ms: u64,
4593        max_retries: u32,
4594    ) -> Result<(ToolResult, bool, bool, bool)> {
4595        let max_retries = if ctx.classification.safely_retryable {
4596            max_retries
4597        } else {
4598            0
4599        };
4600        let mut attempts = 0;
4601        let mut invoked = false;
4602        loop {
4603            let (result, timed_out, cancelled, attempt_invoked) = self
4604                .execute_resolved_tool_once(tool.clone(), args.clone(), ctx.clone(), timeout_ms)
4605                .await?;
4606            invoked |= attempt_invoked;
4607            if result.success || timed_out || cancelled || attempts >= max_retries {
4608                return Ok((result, timed_out, cancelled, invoked));
4609            }
4610            attempts += 1;
4611            warn!(tool = %canonical_id, attempt = attempts, error = %result.output, "Retrying failed tool call");
4612        }
4613    }
4614
4615    /// Executes a tool request through scope, policy, HITL, timeout, recovery, and evidence recording.
4616    fn execute_tool_record(
4617        &self,
4618        request: ToolExecutionRequest,
4619    ) -> Pin<Box<dyn Future<Output = Result<ToolExecutionRecord>> + Send + '_>> {
4620        Box::pin(self.execute_tool_record_inner(request))
4621    }
4622
4623    async fn execute_tool_record_inner(
4624        &self,
4625        request: ToolExecutionRequest,
4626    ) -> Result<ToolExecutionRecord> {
4627        let started_at = chrono::Utc::now();
4628        let start = Instant::now();
4629        info!(tool = %request.requested_name, args = %request.arguments, "Executing tool");
4630
4631        if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4632            let record = self.record_from_parts(
4633                &request,
4634                request.requested_name.clone(),
4635                request.arguments.clone(),
4636                started_at,
4637                start,
4638                false,
4639                false,
4640                "Tool execution is disabled by runtime control".to_string(),
4641                HashMap::new(),
4642                ToolPolicyDecisionRecord::deny("runtime emergency deny is enabled"),
4643                None,
4644                false,
4645                false,
4646            );
4647            self.finish_tool_record(&record).await;
4648            return Ok(record);
4649        }
4650
4651        let Some(resolved) = self.tools.resolve(&request.requested_name) else {
4652            let record = self.record_from_parts(
4653                &request,
4654                request.requested_name.clone(),
4655                request.arguments.clone(),
4656                started_at,
4657                start,
4658                false,
4659                false,
4660                format!("Tool '{}' is unavailable", request.requested_name),
4661                HashMap::new(),
4662                ToolPolicyDecisionRecord::unavailable(format!(
4663                    "Tool '{}' is not registered",
4664                    request.requested_name
4665                )),
4666                None,
4667                false,
4668                false,
4669            );
4670            self.finish_tool_record(&record).await;
4671            return Ok(record);
4672        };
4673
4674        let canonical_id = resolved.identity.canonical_id.clone();
4675
4676        let initial_scope_snapshot = self.get_available_tool_ids_snapshot().await?;
4677        if !initial_scope_snapshot
4678            .tool_ids
4679            .iter()
4680            .any(|id| id == &canonical_id)
4681        {
4682            let record = self.record_from_parts(
4683                &request,
4684                canonical_id.clone(),
4685                request.arguments.clone(),
4686                started_at,
4687                start,
4688                false,
4689                false,
4690                format!(
4691                    "Tool '{}' is not available in the current scope",
4692                    canonical_id
4693                ),
4694                HashMap::new(),
4695                ToolPolicyDecisionRecord::deny(format!(
4696                    "Tool '{}' is not granted by the current top-level and state tool scope",
4697                    canonical_id
4698                )),
4699                None,
4700                false,
4701                false,
4702            );
4703            self.finish_tool_record(&record).await;
4704            return Ok(record);
4705        }
4706
4707        let approval_control_snapshot = self.runtime_safety_snapshot();
4708        let security_engine = approval_control_snapshot.tool_security.clone();
4709        let bindings = resolved.tool.policy_bindings();
4710        let mut executed_arguments = security_engine.prepare_tool_arguments_with_bindings(
4711            &canonical_id,
4712            &request.arguments,
4713            &bindings,
4714        );
4715        self.hooks
4716            .on_tool_start(&canonical_id, &executed_arguments)
4717            .await;
4718
4719        let mut metadata = HashMap::new();
4720        let safety = resolved.tool.safety_metadata();
4721        let classification = resolved.tool.classify_call(&executed_arguments);
4722        let limits = security_engine.effective_limits(&canonical_id, &safety, &classification);
4723        metadata.insert(
4724            "classification".to_string(),
4725            serde_json::to_value(&classification).unwrap_or(Value::Null),
4726        );
4727        metadata.insert(
4728            "effective_limits".to_string(),
4729            serde_json::to_value(&limits).unwrap_or(Value::Null),
4730        );
4731        let policy_snapshot = security_engine.policy_snapshot(&canonical_id);
4732        if !policy_snapshot.is_null() {
4733            metadata.insert("policy_snapshot".to_string(), policy_snapshot.clone());
4734        }
4735
4736        let mut approval_record = Some(ToolApprovalRecord {
4737            status: ToolApprovalStatus::NotRequired,
4738            reason: None,
4739            modified_arguments: None,
4740        });
4741
4742        let mut security_result = security_engine
4743            .validate_tool_execution_with_bindings(&canonical_id, &executed_arguments, &bindings)
4744            .await?;
4745        match &security_result {
4746            SecurityCheckResult::Allow => {}
4747            SecurityCheckResult::Warn { message } => {
4748                warn!(tool = %canonical_id, message = %message, "Tool security warning");
4749            }
4750            SecurityCheckResult::Block { reason } => {
4751                let record = self.record_from_parts(
4752                    &request,
4753                    canonical_id,
4754                    executed_arguments,
4755                    started_at,
4756                    start,
4757                    false,
4758                    false,
4759                    format!("Denied: {}", reason),
4760                    metadata,
4761                    ToolPolicyDecisionRecord::deny(reason.clone()),
4762                    approval_record,
4763                    false,
4764                    false,
4765                );
4766                self.finish_tool_record(&record).await;
4767                return Ok(record);
4768            }
4769            SecurityCheckResult::Unavailable { reason } => {
4770                let record = self.record_from_parts(
4771                    &request,
4772                    canonical_id,
4773                    executed_arguments,
4774                    started_at,
4775                    start,
4776                    false,
4777                    false,
4778                    format!("Unavailable: {}", reason),
4779                    metadata,
4780                    ToolPolicyDecisionRecord::unavailable(reason.clone()),
4781                    approval_record,
4782                    false,
4783                    false,
4784                );
4785                self.finish_tool_record(&record).await;
4786                return Ok(record);
4787            }
4788            SecurityCheckResult::RequireConfirmation { message } => {
4789                if self.hitl_engine.is_none() {
4790                    approval_record = Some(ToolApprovalRecord {
4791                        status: ToolApprovalStatus::Unavailable,
4792                        reason: Some("No HITL engine configured".to_string()),
4793                        modified_arguments: None,
4794                    });
4795                    let record = self.record_from_parts(
4796                        &request,
4797                        canonical_id,
4798                        executed_arguments,
4799                        started_at,
4800                        start,
4801                        false,
4802                        false,
4803                        format!("Approval unavailable: {}", message),
4804                        metadata,
4805                        ToolPolicyDecisionRecord::approval(message.clone()),
4806                        approval_record,
4807                        false,
4808                        false,
4809                    );
4810                    self.finish_tool_record(&record).await;
4811                    return Ok(record);
4812                }
4813
4814                let check_result = HITLCheckResult::required(
4815                    ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
4816                    HashMap::new(),
4817                    message.clone(),
4818                    None,
4819                );
4820                match self.request_hitl_approval(check_result).await? {
4821                    ApprovalResult::Approved => {
4822                        merge_approved_record(&mut approval_record);
4823                    }
4824                    ApprovalResult::Modified { changes } => {
4825                        if let Some(obj) = executed_arguments.as_object_mut() {
4826                            for (key, value) in changes {
4827                                obj.insert(key, value);
4828                            }
4829                        }
4830                        security_result = security_engine
4831                            .validate_tool_execution_with_bindings(
4832                                &canonical_id,
4833                                &executed_arguments,
4834                                &bindings,
4835                            )
4836                            .await?;
4837                        if !matches!(
4838                            security_result,
4839                            SecurityCheckResult::Allow
4840                                | SecurityCheckResult::Warn { .. }
4841                                | SecurityCheckResult::RequireConfirmation { .. }
4842                        ) {
4843                            let reason = security_result
4844                                .reason()
4845                                .unwrap_or("modified arguments failed policy")
4846                                .to_string();
4847                            let record = self.record_from_parts(
4848                                &request,
4849                                canonical_id,
4850                                executed_arguments.clone(),
4851                                started_at,
4852                                start,
4853                                false,
4854                                false,
4855                                reason.clone(),
4856                                metadata,
4857                                ToolPolicyDecisionRecord::deny(reason),
4858                                Some(ToolApprovalRecord {
4859                                    status: ToolApprovalStatus::Modified,
4860                                    reason: None,
4861                                    modified_arguments: Some(executed_arguments),
4862                                }),
4863                                false,
4864                                false,
4865                            );
4866                            self.finish_tool_record(&record).await;
4867                            return Ok(record);
4868                        }
4869                        approval_record = Some(ToolApprovalRecord {
4870                            status: ToolApprovalStatus::Modified,
4871                            reason: None,
4872                            modified_arguments: Some(executed_arguments.clone()),
4873                        });
4874                    }
4875                    ApprovalResult::Rejected { reason } => {
4876                        let reason = reason.unwrap_or_else(|| "rejected".to_string());
4877                        approval_record = Some(ToolApprovalRecord {
4878                            status: ToolApprovalStatus::Rejected,
4879                            reason: Some(reason.clone()),
4880                            modified_arguments: None,
4881                        });
4882                        let record = self.record_from_parts(
4883                            &request,
4884                            canonical_id,
4885                            executed_arguments,
4886                            started_at,
4887                            start,
4888                            false,
4889                            false,
4890                            format!("Approval rejected: {}", reason),
4891                            metadata,
4892                            ToolPolicyDecisionRecord::approval(reason),
4893                            approval_record,
4894                            false,
4895                            false,
4896                        );
4897                        self.finish_tool_record(&record).await;
4898                        return Ok(record);
4899                    }
4900                    ApprovalResult::Timeout => {
4901                        approval_record = Some(ToolApprovalRecord {
4902                            status: ToolApprovalStatus::Timeout,
4903                            reason: Some("approval timeout".to_string()),
4904                            modified_arguments: None,
4905                        });
4906                        let record = self.record_from_parts(
4907                            &request,
4908                            canonical_id,
4909                            executed_arguments,
4910                            started_at,
4911                            start,
4912                            false,
4913                            false,
4914                            "Approval timed out".to_string(),
4915                            metadata,
4916                            ToolPolicyDecisionRecord::approval("approval timeout"),
4917                            approval_record,
4918                            false,
4919                            false,
4920                        );
4921                        self.finish_tool_record(&record).await;
4922                        return Ok(record);
4923                    }
4924                }
4925            }
4926        }
4927
4928        if canonical_id == "command" && !self.tools.command_runner_available() {
4929            let record = self.record_from_parts(
4930                &request,
4931                canonical_id.clone(),
4932                executed_arguments.clone(),
4933                started_at,
4934                start,
4935                false,
4936                false,
4937                "Command runner is unavailable".to_string(),
4938                metadata,
4939                ToolPolicyDecisionRecord::unavailable("command runner is unavailable"),
4940                Some(ToolApprovalRecord {
4941                    status: ToolApprovalStatus::Unavailable,
4942                    reason: Some("command runner is unavailable".to_string()),
4943                    modified_arguments: None,
4944                }),
4945                false,
4946                false,
4947            );
4948            self.finish_tool_record(&record).await;
4949            return Ok(record);
4950        }
4951
4952        if approval_record
4953            .as_ref()
4954            .is_some_and(|record| matches!(record.status, ToolApprovalStatus::NotRequired))
4955            && let Some(message) =
4956                security_engine.classification_approval_message(&canonical_id, &classification)
4957        {
4958            if self.hitl_engine.is_none() {
4959                approval_record = Some(ToolApprovalRecord {
4960                    status: ToolApprovalStatus::Unavailable,
4961                    reason: Some("No HITL engine configured".to_string()),
4962                    modified_arguments: None,
4963                });
4964                let record = self.record_from_parts(
4965                    &request,
4966                    canonical_id,
4967                    executed_arguments,
4968                    started_at,
4969                    start,
4970                    false,
4971                    false,
4972                    format!("Approval unavailable: {}", message),
4973                    metadata,
4974                    ToolPolicyDecisionRecord::approval(message),
4975                    approval_record,
4976                    false,
4977                    false,
4978                );
4979                self.finish_tool_record(&record).await;
4980                return Ok(record);
4981            }
4982            let check_result = HITLCheckResult::required(
4983                ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
4984                HashMap::new(),
4985                message.clone(),
4986                None,
4987            );
4988            match self.request_hitl_approval(check_result).await? {
4989                ApprovalResult::Approved => {
4990                    merge_approved_record(&mut approval_record);
4991                }
4992                ApprovalResult::Modified { changes } => {
4993                    if let Some(obj) = executed_arguments.as_object_mut() {
4994                        for (key, value) in changes {
4995                            obj.insert(key, value);
4996                        }
4997                    }
4998                    let modified_security = security_engine
4999                        .validate_tool_execution_with_bindings(
5000                            &canonical_id,
5001                            &executed_arguments,
5002                            &bindings,
5003                        )
5004                        .await?;
5005                    if !matches!(
5006                        modified_security,
5007                        SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5008                    ) {
5009                        let reason = modified_security
5010                            .reason()
5011                            .unwrap_or("modified arguments failed policy")
5012                            .to_string();
5013                        let record = self.record_from_parts(
5014                            &request,
5015                            canonical_id,
5016                            executed_arguments.clone(),
5017                            started_at,
5018                            start,
5019                            false,
5020                            false,
5021                            reason.clone(),
5022                            metadata,
5023                            ToolPolicyDecisionRecord::deny(reason),
5024                            Some(ToolApprovalRecord {
5025                                status: ToolApprovalStatus::Modified,
5026                                reason: None,
5027                                modified_arguments: Some(executed_arguments),
5028                            }),
5029                            false,
5030                            false,
5031                        );
5032                        self.finish_tool_record(&record).await;
5033                        return Ok(record);
5034                    }
5035                    approval_record = Some(ToolApprovalRecord {
5036                        status: ToolApprovalStatus::Modified,
5037                        reason: None,
5038                        modified_arguments: Some(executed_arguments.clone()),
5039                    });
5040                }
5041                ApprovalResult::Rejected { reason } => {
5042                    let reason = reason.unwrap_or_else(|| "rejected".to_string());
5043                    let record = self.record_from_parts(
5044                        &request,
5045                        canonical_id,
5046                        executed_arguments,
5047                        started_at,
5048                        start,
5049                        false,
5050                        false,
5051                        format!("Approval rejected: {}", reason),
5052                        metadata,
5053                        ToolPolicyDecisionRecord::approval(reason.clone()),
5054                        Some(ToolApprovalRecord {
5055                            status: ToolApprovalStatus::Rejected,
5056                            reason: Some(reason),
5057                            modified_arguments: None,
5058                        }),
5059                        false,
5060                        false,
5061                    );
5062                    self.finish_tool_record(&record).await;
5063                    return Ok(record);
5064                }
5065                ApprovalResult::Timeout => {
5066                    let record = self.record_from_parts(
5067                        &request,
5068                        canonical_id,
5069                        executed_arguments,
5070                        started_at,
5071                        start,
5072                        false,
5073                        false,
5074                        "Approval timed out".to_string(),
5075                        metadata,
5076                        ToolPolicyDecisionRecord::approval("approval timeout"),
5077                        Some(ToolApprovalRecord {
5078                            status: ToolApprovalStatus::Timeout,
5079                            reason: Some("approval timeout".to_string()),
5080                            modified_arguments: None,
5081                        }),
5082                        false,
5083                        false,
5084                    );
5085                    self.finish_tool_record(&record).await;
5086                    return Ok(record);
5087                }
5088            }
5089        }
5090
5091        if canonical_id == "diagnostics" && !self.tools.diagnostics_available() {
5092            let record = self.record_from_parts(
5093                &request,
5094                canonical_id.clone(),
5095                executed_arguments.clone(),
5096                started_at,
5097                start,
5098                false,
5099                false,
5100                "Diagnostics provider is unavailable".to_string(),
5101                metadata,
5102                ToolPolicyDecisionRecord::unavailable("diagnostics provider is unavailable"),
5103                Some(ToolApprovalRecord {
5104                    status: ToolApprovalStatus::Unavailable,
5105                    reason: Some("diagnostics provider is unavailable".to_string()),
5106                    modified_arguments: None,
5107                }),
5108                false,
5109                false,
5110            );
5111            self.finish_tool_record(&record).await;
5112            return Ok(record);
5113        }
5114
5115        if canonical_id == "web_search" && !self.tools.web_search_available() {
5116            let record = self.record_from_parts(
5117                &request,
5118                canonical_id.clone(),
5119                executed_arguments.clone(),
5120                started_at,
5121                start,
5122                false,
5123                false,
5124                "Web search provider is unavailable".to_string(),
5125                metadata,
5126                ToolPolicyDecisionRecord::unavailable("web search provider is unavailable"),
5127                Some(ToolApprovalRecord {
5128                    status: ToolApprovalStatus::Unavailable,
5129                    reason: Some("web search provider is unavailable".to_string()),
5130                    modified_arguments: None,
5131                }),
5132                false,
5133                false,
5134            );
5135            self.finish_tool_record(&record).await;
5136            return Ok(record);
5137        }
5138
5139        let hitl_lang_ctx = self.build_hitl_language_context();
5140        if let Some(ref hitl_engine) = self.hitl_engine {
5141            let check_result = self
5142                .observe_purpose(
5143                    ObservationPurpose::HitlLocalization,
5144                    hitl_engine.check_tool_with_localization(
5145                        &canonical_id,
5146                        &executed_arguments,
5147                        &hitl_lang_ctx,
5148                        self.approval_handler.as_ref(),
5149                        Some(&self.llm_registry),
5150                    ),
5151                )
5152                .await?;
5153            if check_result.is_required() {
5154                match self.request_hitl_approval(check_result).await? {
5155                    ApprovalResult::Approved => {
5156                        merge_approved_record(&mut approval_record);
5157                    }
5158                    ApprovalResult::Modified { changes } => {
5159                        if let Some(obj) = executed_arguments.as_object_mut() {
5160                            for (key, value) in changes {
5161                                obj.insert(key, value);
5162                            }
5163                        }
5164                        let modified_security = security_engine
5165                            .validate_tool_execution_with_bindings(
5166                                &canonical_id,
5167                                &executed_arguments,
5168                                &bindings,
5169                            )
5170                            .await?;
5171                        if !matches!(
5172                            modified_security,
5173                            SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5174                        ) {
5175                            let reason = modified_security
5176                                .reason()
5177                                .unwrap_or("modified arguments failed policy")
5178                                .to_string();
5179                            let record = self.record_from_parts(
5180                                &request,
5181                                canonical_id,
5182                                executed_arguments.clone(),
5183                                started_at,
5184                                start,
5185                                false,
5186                                false,
5187                                reason.clone(),
5188                                metadata,
5189                                ToolPolicyDecisionRecord::deny(reason),
5190                                Some(ToolApprovalRecord {
5191                                    status: ToolApprovalStatus::Modified,
5192                                    reason: None,
5193                                    modified_arguments: Some(executed_arguments),
5194                                }),
5195                                false,
5196                                false,
5197                            );
5198                            self.finish_tool_record(&record).await;
5199                            return Ok(record);
5200                        }
5201                        approval_record = Some(ToolApprovalRecord {
5202                            status: ToolApprovalStatus::Modified,
5203                            reason: None,
5204                            modified_arguments: Some(executed_arguments.clone()),
5205                        });
5206                    }
5207                    ApprovalResult::Rejected { reason } => {
5208                        let reason = reason.unwrap_or_else(|| "rejected".to_string());
5209                        let record = self.record_from_parts(
5210                            &request,
5211                            canonical_id,
5212                            executed_arguments,
5213                            started_at,
5214                            start,
5215                            false,
5216                            false,
5217                            format!("Approval rejected: {}", reason),
5218                            metadata,
5219                            ToolPolicyDecisionRecord::approval(reason.clone()),
5220                            Some(ToolApprovalRecord {
5221                                status: ToolApprovalStatus::Rejected,
5222                                reason: Some(reason),
5223                                modified_arguments: None,
5224                            }),
5225                            false,
5226                            false,
5227                        );
5228                        self.finish_tool_record(&record).await;
5229                        return Ok(record);
5230                    }
5231                    ApprovalResult::Timeout => {
5232                        let record = self.record_from_parts(
5233                            &request,
5234                            canonical_id,
5235                            executed_arguments,
5236                            started_at,
5237                            start,
5238                            false,
5239                            false,
5240                            "Approval timed out".to_string(),
5241                            metadata,
5242                            ToolPolicyDecisionRecord::approval("approval timeout"),
5243                            Some(ToolApprovalRecord {
5244                                status: ToolApprovalStatus::Timeout,
5245                                reason: Some("approval timeout".to_string()),
5246                                modified_arguments: None,
5247                            }),
5248                            false,
5249                            false,
5250                        );
5251                        self.finish_tool_record(&record).await;
5252                        return Ok(record);
5253                    }
5254                }
5255            }
5256
5257            let condition_check = self
5258                .observe_purpose(
5259                    ObservationPurpose::HitlLocalization,
5260                    hitl_engine.check_conditions_with_localization(
5261                        &executed_arguments,
5262                        &hitl_lang_ctx,
5263                        self.approval_handler.as_ref(),
5264                        Some(&self.llm_registry),
5265                    ),
5266                )
5267                .await?;
5268            if condition_check.is_required() {
5269                match self.request_hitl_approval(condition_check).await? {
5270                    ApprovalResult::Approved => {
5271                        merge_approved_record(&mut approval_record);
5272                    }
5273                    ApprovalResult::Modified { changes } => {
5274                        if let Some(obj) = executed_arguments.as_object_mut() {
5275                            for (key, value) in changes {
5276                                obj.insert(key, value);
5277                            }
5278                        }
5279                        let modified_security = security_engine
5280                            .validate_tool_execution_with_bindings(
5281                                &canonical_id,
5282                                &executed_arguments,
5283                                &bindings,
5284                            )
5285                            .await?;
5286                        if !matches!(
5287                            modified_security,
5288                            SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5289                        ) {
5290                            let reason = modified_security
5291                                .reason()
5292                                .unwrap_or("modified arguments failed policy")
5293                                .to_string();
5294                            let record = self.record_from_parts(
5295                                &request,
5296                                canonical_id,
5297                                executed_arguments,
5298                                started_at,
5299                                start,
5300                                false,
5301                                false,
5302                                reason.clone(),
5303                                metadata,
5304                                ToolPolicyDecisionRecord::deny(reason),
5305                                approval_record,
5306                                false,
5307                                false,
5308                            );
5309                            self.finish_tool_record(&record).await;
5310                            return Ok(record);
5311                        }
5312                        approval_record = Some(ToolApprovalRecord {
5313                            status: ToolApprovalStatus::Modified,
5314                            reason: None,
5315                            modified_arguments: Some(executed_arguments.clone()),
5316                        });
5317                    }
5318                    ApprovalResult::Rejected { reason } => {
5319                        let reason = reason.unwrap_or_else(|| "rejected".to_string());
5320                        let record = self.record_from_parts(
5321                            &request,
5322                            canonical_id,
5323                            executed_arguments,
5324                            started_at,
5325                            start,
5326                            false,
5327                            false,
5328                            format!("Approval rejected: {}", reason),
5329                            metadata,
5330                            ToolPolicyDecisionRecord::approval(reason.clone()),
5331                            Some(ToolApprovalRecord {
5332                                status: ToolApprovalStatus::Rejected,
5333                                reason: Some(reason),
5334                                modified_arguments: None,
5335                            }),
5336                            false,
5337                            false,
5338                        );
5339                        self.finish_tool_record(&record).await;
5340                        return Ok(record);
5341                    }
5342                    ApprovalResult::Timeout => {
5343                        let record = self.record_from_parts(
5344                            &request,
5345                            canonical_id,
5346                            executed_arguments,
5347                            started_at,
5348                            start,
5349                            false,
5350                            false,
5351                            "Approval timed out".to_string(),
5352                            metadata,
5353                            ToolPolicyDecisionRecord::approval("approval timeout"),
5354                            Some(ToolApprovalRecord {
5355                                status: ToolApprovalStatus::Timeout,
5356                                reason: Some("approval timeout".to_string()),
5357                                modified_arguments: None,
5358                            }),
5359                            false,
5360                            false,
5361                        );
5362                        self.finish_tool_record(&record).await;
5363                        return Ok(record);
5364                    }
5365                }
5366            }
5367        }
5368
5369        //
5370        // Freeze the action reviewed by HITL after every approved argument change and policy cap.
5371        // A later plain approval preserves any earlier Modified evidence.
5372        //
5373        executed_arguments = security_engine.prepare_tool_arguments_with_bindings(
5374            &canonical_id,
5375            &executed_arguments,
5376            &bindings,
5377        );
5378        if let Some(record) = approval_record.as_mut()
5379            && matches!(record.status, ToolApprovalStatus::Modified)
5380        {
5381            record.modified_arguments = Some(executed_arguments.clone());
5382        }
5383        let binding_security_result = security_engine
5384            .validate_tool_execution_with_bindings(&canonical_id, &executed_arguments, &bindings)
5385            .await?;
5386        let approval_confirmation_required = matches!(
5387            binding_security_result,
5388            SecurityCheckResult::RequireConfirmation { .. }
5389        ) || security_engine
5390            .classification_approval_message(
5391                &canonical_id,
5392                &resolved.tool.classify_call(&executed_arguments),
5393            )
5394            .is_some();
5395        let approval_binding = approval_record.as_ref().and_then(|record| {
5396            matches!(
5397                record.status,
5398                ToolApprovalStatus::Approved | ToolApprovalStatus::Modified
5399            )
5400            .then(|| ToolApprovalBinding {
5401                canonical_id: canonical_id.clone(),
5402                arguments: executed_arguments.clone(),
5403                confirmation_required: approval_confirmation_required,
5404                policy_version: security_engine.policy_version(),
5405                runtime_control_version: approval_control_snapshot.version,
5406                state_generation: initial_scope_snapshot.state_generation,
5407                reviewed_tool: Arc::clone(&resolved.tool),
5408            })
5409        });
5410
5411        //
5412        // Re-resolve once after HITL because registry, scope, and policy may change while approval waits.
5413        // The final Arc must still match the implementation that the approver reviewed.
5414        //
5415        let control_snapshot = self.runtime_safety_snapshot();
5416        let resolved = self.tools.resolve(&request.requested_name);
5417        let registry_version = self.tools.version();
5418        let mut versions = ToolDecisionVersions {
5419            policy: control_snapshot.tool_security.policy_version(),
5420            registry: registry_version,
5421            runtime_control: control_snapshot.version,
5422            state: None,
5423        };
5424        metadata.insert(
5425            "runtime_scope_snapshot".to_string(),
5426            serde_json::to_value(&control_snapshot.tool_scope_override).unwrap_or(Value::Null),
5427        );
5428        let resolved = match resolved {
5429            Some(resolved) => resolved,
5430            None => {
5431                let reason = format!(
5432                    "Tool '{}' became unavailable after approval",
5433                    request.requested_name
5434                );
5435                let record = self.record_from_parts_at(
5436                    &request,
5437                    request.requested_name.clone(),
5438                    executed_arguments,
5439                    started_at,
5440                    start,
5441                    false,
5442                    false,
5443                    reason.clone(),
5444                    metadata,
5445                    ToolPolicyDecisionRecord::unavailable(reason),
5446                    approval_record,
5447                    false,
5448                    false,
5449                    versions,
5450                );
5451                self.finish_tool_record(&record).await;
5452                return Ok(record);
5453            }
5454        };
5455
5456        let canonical_id = resolved.identity.canonical_id.clone();
5457        let bindings = resolved.tool.policy_bindings();
5458        let final_arguments = control_snapshot
5459            .tool_security
5460            .prepare_tool_arguments_with_bindings(&canonical_id, &executed_arguments, &bindings);
5461        if let Some(record) = approval_record.as_mut()
5462            && matches!(record.status, ToolApprovalStatus::Modified)
5463        {
5464            record.modified_arguments = Some(final_arguments.clone());
5465        }
5466        let classification = resolved.tool.classify_call(&final_arguments);
5467        let safety = resolved.tool.safety_metadata();
5468        let security_engine = control_snapshot.tool_security;
5469        let limits = security_engine.effective_limits(&canonical_id, &safety, &classification);
5470        let policy_snapshot = security_engine.policy_snapshot(&canonical_id);
5471        let resource_lock_keys =
5472            tool_resource_lock_keys(&canonical_id, &final_arguments, &bindings, &classification);
5473        metadata.insert(
5474            "classification".to_string(),
5475            serde_json::to_value(&classification).unwrap_or(Value::Null),
5476        );
5477        metadata.insert(
5478            "effective_limits".to_string(),
5479            serde_json::to_value(&limits).unwrap_or(Value::Null),
5480        );
5481        metadata.insert(
5482            "resource_lock_keys".to_string(),
5483            serde_json::to_value(&resource_lock_keys).unwrap_or(Value::Null),
5484        );
5485        if policy_snapshot.is_null() {
5486            metadata.remove("policy_snapshot");
5487        } else {
5488            metadata.insert("policy_snapshot".to_string(), policy_snapshot.clone());
5489        }
5490
5491        let final_denial = |canonical_id: String,
5492                            output: String,
5493                            policy: ToolPolicyDecisionRecord,
5494                            metadata: HashMap<String, Value>,
5495                            decision_versions: ToolDecisionVersions| {
5496            self.record_from_parts_at(
5497                &request,
5498                canonical_id,
5499                final_arguments.clone(),
5500                started_at,
5501                start,
5502                false,
5503                false,
5504                output,
5505                metadata,
5506                policy,
5507                approval_record.clone(),
5508                false,
5509                false,
5510                decision_versions,
5511            )
5512        };
5513
5514        if control_snapshot.emergency_deny {
5515            let reason = "Tool execution is disabled by runtime control".to_string();
5516            let record = final_denial(
5517                canonical_id,
5518                reason.clone(),
5519                ToolPolicyDecisionRecord::deny(reason),
5520                metadata,
5521                versions,
5522            );
5523            self.finish_tool_record(&record).await;
5524            return Ok(record);
5525        }
5526
5527        //
5528        // Final scope evaluation uses the same post-HITL runtime snapshot and captures state authority before locks.
5529        // Admission must reject any later state transition, reset, or restore before consuming rate capacity or invoking the tool.
5530        //
5531        let available_snapshot = self
5532            .get_available_tool_ids_snapshot_for_scope(
5533                control_snapshot.tool_scope_override.as_deref(),
5534            )
5535            .await?;
5536        versions.state = available_snapshot.state_generation;
5537        metadata.insert(
5538            "available_tool_ids_snapshot".to_string(),
5539            serde_json::to_value(&available_snapshot.tool_ids).unwrap_or(Value::Null),
5540        );
5541        metadata.insert(
5542            "state_generation_snapshot".to_string(),
5543            serde_json::to_value(available_snapshot.state_generation).unwrap_or(Value::Null),
5544        );
5545        if !available_snapshot
5546            .tool_ids
5547            .iter()
5548            .any(|tool_id| tool_id == &canonical_id)
5549        {
5550            let reason = format!(
5551                "Tool '{}' is not available in the final runtime scope",
5552                canonical_id
5553            );
5554            let record = final_denial(
5555                canonical_id,
5556                reason.clone(),
5557                ToolPolicyDecisionRecord::deny(reason),
5558                metadata,
5559                versions,
5560            );
5561            self.finish_tool_record(&record).await;
5562            return Ok(record);
5563        }
5564
5565        //
5566        // Validation may run more than once, but it must not consume rate capacity.
5567        // Rate capacity is consumed once after resource locks are held.
5568        //
5569        let final_security_result = security_engine
5570            .validate_tool_execution_with_bindings(&canonical_id, &final_arguments, &bindings)
5571            .await?;
5572        match &final_security_result {
5573            SecurityCheckResult::Block { reason } => {
5574                let record = final_denial(
5575                    canonical_id,
5576                    format!("Denied: {}", reason),
5577                    ToolPolicyDecisionRecord::deny(reason.clone()),
5578                    metadata,
5579                    versions,
5580                );
5581                self.finish_tool_record(&record).await;
5582                return Ok(record);
5583            }
5584            SecurityCheckResult::Unavailable { reason } => {
5585                let record = final_denial(
5586                    canonical_id,
5587                    format!("Unavailable: {}", reason),
5588                    ToolPolicyDecisionRecord::unavailable(reason.clone()),
5589                    metadata,
5590                    versions,
5591                );
5592                self.finish_tool_record(&record).await;
5593                return Ok(record);
5594            }
5595            SecurityCheckResult::Warn { message } => {
5596                warn!(tool = %canonical_id, message = %message, "Tool security warning after approval");
5597            }
5598            SecurityCheckResult::Allow | SecurityCheckResult::RequireConfirmation { .. } => {}
5599        }
5600        let final_confirmation_required = matches!(
5601            final_security_result,
5602            SecurityCheckResult::RequireConfirmation { .. }
5603        ) || security_engine
5604            .classification_approval_message(&canonical_id, &classification)
5605            .is_some();
5606        let stale_approval = approval_binding.as_ref().is_some_and(|binding| {
5607            binding.is_stale(
5608                &canonical_id,
5609                &final_arguments,
5610                final_confirmation_required,
5611                versions,
5612                &resolved.tool,
5613            )
5614        });
5615        if stale_approval {
5616            let reason = "Approval became stale before final admission".to_string();
5617            let record = final_denial(
5618                canonical_id,
5619                reason.clone(),
5620                ToolPolicyDecisionRecord::deny(reason),
5621                metadata,
5622                versions,
5623            );
5624            self.finish_tool_record(&record).await;
5625            return Ok(record);
5626        }
5627        if final_confirmation_required && approval_binding.is_none() {
5628            let reason = "Final policy requires fresh approval".to_string();
5629            let record = final_denial(
5630                canonical_id,
5631                reason.clone(),
5632                ToolPolicyDecisionRecord::approval(reason),
5633                metadata,
5634                versions,
5635            );
5636            self.finish_tool_record(&record).await;
5637            return Ok(record);
5638        }
5639
5640        let unavailable_reason = match canonical_id.as_str() {
5641            "command" if !self.tools.command_runner_available() => {
5642                Some("command runner is unavailable")
5643            }
5644            "diagnostics" if !self.tools.diagnostics_available() => {
5645                Some("diagnostics provider is unavailable")
5646            }
5647            "web_search" if !self.tools.web_search_available() => {
5648                Some("web search provider is unavailable")
5649            }
5650            _ => None,
5651        };
5652        if let Some(reason) = unavailable_reason {
5653            let record = final_denial(
5654                canonical_id,
5655                reason.to_string(),
5656                ToolPolicyDecisionRecord::unavailable(reason),
5657                metadata,
5658                versions,
5659            );
5660            self.finish_tool_record(&record).await;
5661            return Ok(record);
5662        }
5663
5664        //
5665        // Hold conflict locks across final generation admission and invocation.
5666        // Completion hooks and fallback execution run only after these guards are released.
5667        //
5668        let Some(resource_guards) = self.acquire_tool_resource_locks(&resource_lock_keys).await
5669        else {
5670            let reason = "Tool execution cancelled while waiting for resource locks".to_string();
5671            let record = final_denial(
5672                canonical_id,
5673                reason.clone(),
5674                ToolPolicyDecisionRecord::deny(reason),
5675                metadata,
5676                versions,
5677            );
5678            self.finish_tool_record(&record).await;
5679            return Ok(record);
5680        };
5681
5682        //
5683        // The lock-held admission is the last authority boundary. Runtime, policy, and state generations must still match before the one atomic rate admission.
5684        // Updates after admission apply to later calls, while emergency cancellation remains live for this call.
5685        //
5686        let admission = self.admit_tool_execution(
5687            versions.runtime_control,
5688            versions.policy,
5689            versions.state,
5690            &canonical_id,
5691        );
5692        if !matches!(admission, SecurityCheckResult::Allow) {
5693            let latest_control = self.runtime_safety_snapshot();
5694            let reason = admission
5695                .reason()
5696                .unwrap_or("tool admission was denied")
5697                .to_string();
5698            let policy = if admission.is_unavailable() {
5699                ToolPolicyDecisionRecord::unavailable(reason.clone())
5700            } else {
5701                ToolPolicyDecisionRecord::deny(reason.clone())
5702            };
5703            let record = self.record_from_parts_at(
5704                &request,
5705                canonical_id,
5706                final_arguments,
5707                started_at,
5708                start,
5709                false,
5710                false,
5711                reason,
5712                metadata,
5713                policy,
5714                approval_record,
5715                false,
5716                false,
5717                ToolDecisionVersions {
5718                    policy: latest_control.tool_security.policy_version(),
5719                    registry: versions.registry,
5720                    runtime_control: latest_control.version,
5721                    state: self
5722                        .state_machine
5723                        .as_ref()
5724                        .map(|state_machine| state_machine.generation()),
5725                },
5726            );
5727            self.finish_tool_record_after_resource_guards(resource_guards, &record)
5728                .await;
5729            return Ok(record);
5730        }
5731        let executed_arguments = final_arguments;
5732
5733        let tool_config = self.recovery_manager.get_tool_config(&canonical_id);
5734        let timeout_ms = limits
5735            .timeout_ms
5736            .unwrap_or_else(|| security_engine.get_tool_timeout(&canonical_id));
5737        let deadline = Some(started_at + chrono::Duration::milliseconds(timeout_ms as i64));
5738        let turn_actor = current_turn_actor_context();
5739        let actor = ToolActorContext {
5740            actor_id: turn_actor
5741                .as_ref()
5742                .and_then(|context| context.effective_actor_id().map(str::to_string))
5743                .or_else(|| self.actor_id()),
5744            origin_actor_id: turn_actor
5745                .as_ref()
5746                .and_then(|context| context.origin_actor_id.clone()),
5747            sender_agent_id: turn_actor
5748                .as_ref()
5749                .and_then(|context| context.sender_agent_id.clone()),
5750        };
5751        let tool_context = ToolExecutionContext {
5752            requested_name: request.requested_name.clone(),
5753            canonical_id: canonical_id.clone(),
5754            display_name: resolved.identity.display_name.clone(),
5755            provider_id: resolved.identity.provider_id.clone(),
5756            registry_version: versions.registry,
5757            policy_version: versions.policy,
5758            runtime_control_version: versions.runtime_control,
5759            call_id: request.call_id.clone(),
5760            source: request.source.clone(),
5761            actor,
5762            cancellation: ToolCancellationToken::new(
5763                Arc::clone(&self.runtime_control.emergency_deny),
5764                Some("runtime control cancellation".to_string()),
5765            ),
5766            started_at,
5767            deadline,
5768            permission: ToolPolicyDecisionRecord::allow(),
5769            approval: approval_record.clone(),
5770            classification: classification.clone(),
5771            safety,
5772            limits: limits.clone(),
5773            policy_snapshot,
5774            custom_config: security_engine.custom_config(&canonical_id),
5775        };
5776        let (mut result, timed_out, cancelled, invoked) = self
5777            .run_tool_with_retries(
5778                &canonical_id,
5779                resolved.tool.clone(),
5780                executed_arguments.clone(),
5781                tool_context,
5782                timeout_ms,
5783                tool_config.max_retries,
5784            )
5785            .await?;
5786
5787        if !result.success {
5788            match &tool_config.on_failure {
5789                ToolFailureAction::Skip => {
5790                    result = ToolResult::ok(format!(
5791                        "{{\"skipped\": true, \"reason\": \"Tool '{}' was skipped after failure\"}}",
5792                        canonical_id
5793                    ));
5794                }
5795                ToolFailureAction::Fallback { fallback_tool } => {
5796                    drop(resource_guards);
5797                    let fallback_request = ToolExecutionRequest::new(
5798                        request.call_id.clone(),
5799                        fallback_tool.clone(),
5800                        executed_arguments,
5801                        ToolCallSource::Fallback {
5802                            original_tool: canonical_id,
5803                        },
5804                    );
5805                    return Box::pin(self.execute_tool_record(fallback_request)).await;
5806                }
5807                ToolFailureAction::ReportError => {}
5808            }
5809        }
5810
5811        let output_cap = limits.max_output_chars;
5812        let (output, output_truncated) =
5813            Self::truncate_tool_output(result.output.clone(), output_cap);
5814        if let Some(result_metadata) = result.metadata {
5815            metadata.extend(result_metadata);
5816        }
5817        let mut record = self.record_from_parts_at(
5818            &request,
5819            canonical_id,
5820            executed_arguments,
5821            started_at,
5822            start,
5823            invoked,
5824            result.success,
5825            output,
5826            metadata,
5827            ToolPolicyDecisionRecord::allow(),
5828            approval_record,
5829            timed_out,
5830            output_truncated,
5831            versions,
5832        );
5833        record.cancelled = cancelled;
5834        if cancelled {
5835            record.cancellation_reason = Some("runtime control cancellation".to_string());
5836        }
5837        self.finish_tool_record_after_resource_guards(resource_guards, &record)
5838            .await;
5839        Ok(record)
5840    }
5841
5842    #[instrument(skip(self, tool_call), fields(tool = %tool_call.name))]
5843    async fn execute_tool_smart(&self, tool_call: &ToolCall) -> Result<String> {
5844        let record = self
5845            .execute_tool_record(ToolExecutionRequest::new(
5846                tool_call.id.clone(),
5847                tool_call.name.clone(),
5848                tool_call.arguments.clone(),
5849                ToolCallSource::Model,
5850            ))
5851            .await?;
5852        if record.success {
5853            Ok(record.model_output_string())
5854        } else if matches!(record.policy.outcome, PermissionOutcome::RequiresApproval) {
5855            Err(AgentError::HITLRejected(record.model_output_string()))
5856        } else {
5857            Err(AgentError::Tool(record.model_output_string()))
5858        }
5859    }
5860
5861    //
5862    // This method is branch-safe because it only asks the router which skill matches.
5863    // Do not add disambiguation, pending-skill writes, or skill execution here.
5864    //
5865    /// Selects a skill without executing it.
5866    async fn select_skill_candidate(&self, input: &str) -> Result<Option<SkillCandidate>> {
5867        let Some(ref router) = self.skill_router else {
5868            return Ok(None);
5869        };
5870        let available_skills = self.get_available_skills();
5871        if available_skills.is_empty() {
5872            return Ok(None);
5873        }
5874        let skill_ids: Vec<&str> = available_skills.iter().map(|s| s.id.as_str()).collect();
5875        let Some(skill_id) = self
5876            .observe_purpose(
5877                ObservationPurpose::SkillRouting,
5878                router.select_skill_filtered(input, &skill_ids),
5879            )
5880            .await?
5881        else {
5882            return Ok(None);
5883        };
5884        let skill = router
5885            .get_skill(&skill_id)
5886            .cloned()
5887            .ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
5888        info!(skill_id = %skill_id, "Skill selected");
5889        Ok(Some(SkillCandidate::new(skill_id, skill)))
5890    }
5891
5892    //
5893    // This is the commit half of skill routing.
5894    // It may mutate pending skill state, run clarification, and execute skill steps.
5895    //
5896    async fn commit_skill_candidate_route_result(
5897        &self,
5898        candidate: SkillCandidate,
5899        input: &str,
5900    ) -> Result<SkillRouteResult> {
5901        let skill_id = candidate.skill_id;
5902        let skill = candidate.skill;
5903        let expected_state_generation = self
5904            .state_machine
5905            .as_ref()
5906            .map(|state_machine| state_machine.generation());
5907        let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
5908        if let Some(ref skill_disambig) = skill.disambiguation
5909            && skill_disambig.enabled.unwrap_or(false)
5910            && let Some(ref disambiguator) = self.disambiguation_manager
5911        {
5912            let context = self.build_disambiguation_context().await?;
5913            let state_override = self
5914                .state_machine
5915                .as_ref()
5916                .and_then(|sm| sm.current_definition())
5917                .and_then(|def| def.disambiguation.clone());
5918
5919            let disambiguation_result = self
5920                .observe_purpose(
5921                    ObservationPurpose::DisambiguationDetection,
5922                    disambiguator.process_input_with_override(
5923                        input,
5924                        &context,
5925                        state_override.as_ref(),
5926                        Some(skill_disambig),
5927                    ),
5928                )
5929                .await?;
5930            let current_state_generation = self
5931                .state_machine
5932                .as_ref()
5933                .map(|state_machine| state_machine.generation());
5934            if current_state_generation != expected_state_generation
5935                || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
5936            {
5937                disambiguator.clear_pending().await;
5938                *self.pending_skill_id.write() = None;
5939                return Err(AgentError::Other(
5940                    "State or reset ownership changed during skill disambiguation".to_string(),
5941                ));
5942            }
5943            match disambiguation_result {
5944                DisambiguationResult::Clear => {
5945                    debug!(skill_id = %skill_id, "Skill disambiguation: clear");
5946                }
5947                DisambiguationResult::NeedsClarification {
5948                    question,
5949                    detection,
5950                } => {
5951                    let admission = self
5952                        .admit_disambiguation_redispatch(
5953                            expected_disambiguation_epoch,
5954                            expected_state_generation,
5955                        )
5956                        .await?;
5957                    let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
5958                    info!(
5959                        skill_id = %skill_id,
5960                        ambiguity_type = ?detection.ambiguity_type,
5961                        confidence = detection.confidence,
5962                        "Skill requires clarification before execution"
5963                    );
5964                    *self.pending_skill_id.write() = Some(skill_id.clone());
5965                    let response = AgentResponse::new(&question.question).with_metadata(
5966                        "disambiguation",
5967                        serde_json::json!({
5968                            "status": if awaiting_confirmation { "awaiting_confirmation" } else { "awaiting_clarification" },
5969                            "skill_id": skill_id,
5970                            "options": question.options,
5971                            "clarifying": question.clarifying,
5972                            "detection": {
5973                                "type": detection.ambiguity_type,
5974                                "confidence": detection.confidence,
5975                                "what_is_unclear": detection.what_is_unclear,
5976                            }
5977                        }),
5978                    );
5979                    drop(admission);
5980                    return Ok(SkillRouteResult::NeedsClarification {
5981                        response,
5982                        ownership: Some(DisambiguationOwnership {
5983                            epoch: expected_disambiguation_epoch,
5984                            state_generation: expected_state_generation,
5985                        }),
5986                    });
5987                }
5988                DisambiguationResult::Clarified { enriched_input, .. } => {
5989                    info!(skill_id = %skill_id, enriched = %enriched_input, "Skill disambiguation clarified");
5990                    let admission = self
5991                        .admit_disambiguation_redispatch(
5992                            expected_disambiguation_epoch,
5993                            expected_state_generation,
5994                        )
5995                        .await?;
5996                    drop(admission);
5997                    let content = self.execute_skill(&skill, &enriched_input).await?;
5998                    return Ok(SkillRouteResult::Response { skill_id, content });
5999                }
6000                DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
6001                    info!(skill_id = %skill_id, "Skill disambiguation best guess");
6002                    let admission = self
6003                        .admit_disambiguation_redispatch(
6004                            expected_disambiguation_epoch,
6005                            expected_state_generation,
6006                        )
6007                        .await?;
6008                    drop(admission);
6009                    let content = self.execute_skill(&skill, &enriched_input).await?;
6010                    return Ok(SkillRouteResult::Response { skill_id, content });
6011                }
6012                DisambiguationResult::GiveUp { reason } => {
6013                    warn!(skill_id = %skill_id, reason = %reason, "Skill disambiguation gave up");
6014                    let apology = self
6015                                .generate_localized_apology(
6016                                    "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
6017                                    &reason,
6018                                )
6019                                .await
6020                                .unwrap_or_else(|_| {
6021                                    format!("I'm sorry, I couldn't understand your request: {}", reason)
6022                                });
6023                    return Ok(SkillRouteResult::NeedsClarification {
6024                        response: AgentResponse::new(&apology),
6025                        ownership: None,
6026                    });
6027                }
6028                DisambiguationResult::Escalate { reason } => {
6029                    info!(skill_id = %skill_id, reason = %reason, "Skill disambiguation escalating");
6030                    let apology = self
6031                                .generate_localized_apology(
6032                                    "Explain briefly that you're transferring the user to a human agent for help.",
6033                                    &reason,
6034                                )
6035                                .await
6036                                .unwrap_or_else(|_| {
6037                                    format!("I need human assistance to help with your request: {}", reason)
6038                                });
6039                    return Ok(SkillRouteResult::NeedsClarification {
6040                        response: AgentResponse::new(&apology),
6041                        ownership: None,
6042                    });
6043                }
6044                DisambiguationResult::Abandoned { .. } => {
6045                    debug!(skill_id = %skill_id, "Skill disambiguation abandoned");
6046                    return Ok(SkillRouteResult::NoMatch);
6047                }
6048            }
6049        }
6050        let admission = self
6051            .admit_disambiguation_redispatch(
6052                expected_disambiguation_epoch,
6053                expected_state_generation,
6054            )
6055            .await?;
6056        drop(admission);
6057        let content = self.execute_skill(&skill, input).await?;
6058        Ok(SkillRouteResult::Response { skill_id, content })
6059    }
6060
6061    /// Result of skill routing.
6062    async fn try_skill_route(&self, input: &str) -> Result<SkillRouteResult> {
6063        if let Some(candidate) = self.select_skill_candidate(input).await? {
6064            self.commit_skill_candidate_route_result(candidate, input)
6065                .await
6066        } else {
6067            Ok(SkillRouteResult::NoMatch)
6068        }
6069    }
6070
6071    /// Execute a skill with reasoning and reflection, returning the response string.
6072    async fn execute_skill(&self, skill: &SkillDefinition, input: &str) -> Result<String> {
6073        if let Some(ref executor) = self.skill_executor {
6074            let skill_reasoning = self.get_skill_reasoning_config(skill);
6075            let skill_reflection = self.get_skill_reflection_config(skill);
6076
6077            debug!(
6078                skill_id = %skill.id,
6079                reasoning_mode = ?skill_reasoning.mode,
6080                reflection_enabled = ?skill_reflection.enabled,
6081                "Skill reasoning/reflection config"
6082            );
6083
6084            let response = self
6085                .observe_purpose(
6086                    ObservationPurpose::SkillPrompt,
6087                    executor.execute_with_invoker(skill, input, serde_json::json!({}), self),
6088                )
6089                .await?;
6090
6091            if skill_reflection.requires_evaluation() && skill_reflection.is_enabled() {
6092                let should_reflect = self
6093                    .should_reflect_with_config(input, &response, &skill_reflection)
6094                    .await?;
6095                if should_reflect {
6096                    let evaluated = self
6097                        .evaluate_and_retry_with_config(input, response, &skill_reflection)
6098                        .await?;
6099                    return Ok(evaluated);
6100                }
6101            }
6102
6103            return Ok(response);
6104        }
6105        Err(AgentError::Skill(
6106            "No skill executor configured".to_string(),
6107        ))
6108    }
6109
6110    /// Execute a skill by ID, bypassing the skill router.
6111    /// Used after skill-triggered disambiguation resolves to route directly to the matched skill.
6112    async fn execute_skill_by_id(&self, skill_id: &str, input: &str) -> Result<String> {
6113        let skill = self
6114            .skill_router
6115            .as_ref()
6116            .and_then(|r| r.get_skill(skill_id).cloned())
6117            .ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
6118        self.execute_skill(&skill, input).await
6119    }
6120
6121    async fn should_reflect_with_config(
6122        &self,
6123        input: &str,
6124        response: &str,
6125        config: &ReflectionConfig,
6126    ) -> Result<bool> {
6127        if !config.requires_evaluation() {
6128            return Ok(false);
6129        }
6130
6131        if config.is_enabled() {
6132            return Ok(true);
6133        }
6134
6135        let evaluator_llm = config
6136            .evaluator_llm
6137            .as_ref()
6138            .and_then(|alias| self.llm_registry.get(alias).ok())
6139            .or_else(|| self.llm_registry.router().ok())
6140            .or_else(|| self.llm_registry.default().ok());
6141
6142        let Some(llm) = evaluator_llm else {
6143            return Ok(false);
6144        };
6145
6146        let response_preview: String = response.chars().take(500).collect();
6147        let prompt = format!(
6148            r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
6149
6150User query: "{}"
6151Response: "{}"
6152
6153Answer YES or NO only."#,
6154            input, response_preview
6155        );
6156
6157        let messages = vec![ChatMessage::user(&prompt)];
6158        let result = self
6159            .observe_purpose(
6160                ObservationPurpose::ReflectionDecision,
6161                llm.complete(&messages, None),
6162            )
6163            .await;
6164
6165        match result {
6166            Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
6167            Err(_) => Ok(false),
6168        }
6169    }
6170
6171    async fn evaluate_and_retry_with_config(
6172        &self,
6173        input: &str,
6174        mut response: String,
6175        config: &ReflectionConfig,
6176    ) -> Result<String> {
6177        let llm = self.get_state_llm()?;
6178        let mut attempts = 0u32;
6179        let max_retries = config.max_retries;
6180
6181        loop {
6182            let evaluation = self
6183                .evaluate_response_with_config(input, &response, config)
6184                .await?;
6185
6186            if evaluation.passed || attempts >= max_retries {
6187                info!(
6188                    passed = evaluation.passed,
6189                    confidence = evaluation.confidence,
6190                    attempts = attempts + 1,
6191                    "Skill reflection evaluation complete"
6192                );
6193                return Ok(response);
6194            }
6195
6196            debug!(
6197                attempt = attempts + 1,
6198                failed_criteria = evaluation.failed_criteria().count(),
6199                "Skill response did not meet criteria, retrying"
6200            );
6201
6202            let feedback: Vec<String> = evaluation
6203                .failed_criteria()
6204                .map(|c| format!("- {}", c.criterion))
6205                .collect();
6206
6207            let retry_prompt = format!(
6208                "Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response to: {}",
6209                feedback.join("\n"),
6210                input
6211            );
6212
6213            let messages = vec![ChatMessage::user(&retry_prompt)];
6214            let retry_response = self
6215                .observe_purpose(
6216                    ObservationPurpose::ReflectionEvaluation,
6217                    llm.complete(&messages, None),
6218                )
6219                .await
6220                .map_err(|e| AgentError::LLM(e.to_string()))?;
6221
6222            response = retry_response.content.trim().to_string();
6223            attempts += 1;
6224        }
6225    }
6226
6227    async fn evaluate_response_with_config(
6228        &self,
6229        input: &str,
6230        response: &str,
6231        config: &ReflectionConfig,
6232    ) -> Result<EvaluationResult> {
6233        let evaluator_llm = config
6234            .evaluator_llm
6235            .as_ref()
6236            .and_then(|alias| self.llm_registry.get(alias).ok())
6237            .or_else(|| self.llm_registry.router().ok())
6238            .or_else(|| self.llm_registry.default().ok())
6239            .ok_or_else(|| AgentError::Config("No LLM available for evaluation".into()))?;
6240
6241        let criteria = &config.criteria;
6242        let criteria_list = criteria
6243            .iter()
6244            .enumerate()
6245            .map(|(i, c)| format!("{}. {}", i + 1, c))
6246            .collect::<Vec<_>>()
6247            .join("\n");
6248
6249        let prompt = format!(
6250            r#"Evaluate this response against the criteria.
6251
6252User query: "{}"
6253
6254Response to evaluate: "{}"
6255
6256Criteria:
6257{}
6258
6259For each criterion, respond with:
6260- criterion number
6261- PASS or FAIL
6262- brief reason
6263
6264Then provide overall confidence (0.0 to 1.0) and whether it passes overall.
6265
6266Format:
62671. PASS/FAIL - reason
62682. PASS/FAIL - reason
6269...
6270CONFIDENCE: 0.X
6271OVERALL: PASS/FAIL"#,
6272            input, response, criteria_list
6273        );
6274
6275        let messages = vec![ChatMessage::user(&prompt)];
6276        let eval_response = self
6277            .observe_purpose(
6278                ObservationPurpose::ReflectionEvaluation,
6279                evaluator_llm.complete(&messages, None),
6280            )
6281            .await
6282            .map_err(|e| AgentError::LLM(format!("Evaluation failed: {}", e)))?;
6283
6284        let content = eval_response.content.to_uppercase();
6285        let llm_pass = content.contains("OVERALL: PASS");
6286
6287        let confidence = content
6288            .lines()
6289            .find(|l| l.contains("CONFIDENCE:"))
6290            .and_then(|l| {
6291                l.split(':')
6292                    .nth(1)
6293                    .and_then(|v| v.trim().parse::<f32>().ok())
6294            })
6295            .unwrap_or(if llm_pass { 0.8 } else { 0.4 });
6296
6297        // Gate pass against confidence threshold.
6298        // LLM may say PASS but with low confidence - the threshold catches this.
6299        let overall_pass = llm_pass && confidence >= config.pass_threshold;
6300
6301        let mut criteria_results = Vec::new();
6302        for (i, criterion) in criteria.iter().enumerate() {
6303            let line_marker = format!("{}.", i + 1);
6304            let passed = eval_response
6305                .content
6306                .lines()
6307                .find(|l| l.contains(&line_marker))
6308                .map(|l| l.to_uppercase().contains("PASS"))
6309                .unwrap_or(overall_pass);
6310
6311            if passed {
6312                criteria_results.push(CriterionResult::pass(criterion));
6313            } else {
6314                criteria_results.push(CriterionResult::fail(criterion, "Did not meet criterion"));
6315            }
6316        }
6317
6318        Ok(EvaluationResult::new(overall_pass, confidence).with_criteria(criteria_results))
6319    }
6320
6321    /// Process input through the pipeline (state-level override or agent-level).
6322    async fn process_input(&self, input: &str) -> Result<ProcessData> {
6323        if let Some(processor) = self.get_state_process_processor() {
6324            let purpose = observation_purpose_for_process(processor.input_purpose_hint());
6325            return self
6326                .observe_purpose(purpose, processor.process_input(input))
6327                .await;
6328        }
6329        if let Some(ref processor) = self.process_processor {
6330            let purpose = observation_purpose_for_process(processor.input_purpose_hint());
6331            self.observe_purpose(purpose, processor.process_input(input))
6332                .await
6333        } else {
6334            Ok(ProcessData::new(input))
6335        }
6336    }
6337
6338    /// Process output through the pipeline (state-level override or agent-level).
6339    async fn process_output(
6340        &self,
6341        output: &str,
6342        input_context: &std::collections::HashMap<String, serde_json::Value>,
6343    ) -> Result<ProcessData> {
6344        if let Some(processor) = self.get_state_process_processor() {
6345            let purpose = observation_purpose_for_process(processor.output_purpose_hint());
6346            return self
6347                .observe_purpose(purpose, processor.process_output(output, input_context))
6348                .await;
6349        }
6350        if let Some(ref processor) = self.process_processor {
6351            let purpose = observation_purpose_for_process(processor.output_purpose_hint());
6352            self.observe_purpose(purpose, processor.process_output(output, input_context))
6353                .await
6354        } else {
6355            Ok(ProcessData::new(output))
6356        }
6357    }
6358
6359    /// Build a ProcessProcessor from the current state's process config, if any.
6360    fn get_state_process_processor(&self) -> Option<ProcessProcessor> {
6361        let sm = self.state_machine.as_ref()?;
6362        let def = sm.current_definition()?;
6363        let config = def.process.as_ref()?;
6364        let mut processor = ProcessProcessor::new(config.clone());
6365        if let Some(ref registry) = Some(self.llm_registry.clone()) {
6366            processor = processor.with_llm_registry(registry.clone());
6367        }
6368        processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
6369        Some(processor)
6370    }
6371
6372    // Timeout transitions reserve exit actions before committing ownership changes under the short write lock.
6373    async fn check_turn_timeout(&self) -> Result<()> {
6374        let Some(ref sm) = self.state_machine else {
6375            return Ok(());
6376        };
6377        let Some(timeout_state) = sm.check_timeout() else {
6378            return Ok(());
6379        };
6380        let claim_admission = self.disambiguation_admission.write().await;
6381        if sm.check_timeout().as_deref() != Some(timeout_state.as_str()) {
6382            return Ok(());
6383        }
6384        let Some(reservation) = self.reserve_state_transition() else {
6385            return Ok(());
6386        };
6387        let from_state = sm.current();
6388        let expected_state_generation = sm.generation();
6389        let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
6390        let history_before = sm.history();
6391        drop(claim_admission);
6392
6393        self.execute_state_exit_actions(&from_state).await;
6394
6395        let admission = self.disambiguation_admission.write().await;
6396        if sm.current() != from_state
6397            || sm.generation() != expected_state_generation
6398            || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6399            || sm.check_timeout().as_deref() != Some(timeout_state.as_str())
6400        {
6401            return Ok(());
6402        }
6403        sm.transition_to(&timeout_state, "max_turns exceeded")?;
6404        self.invalidate_pending_confirmation("state_timeout").await;
6405        let entered = sm.current();
6406        let is_reentry = Self::state_was_previously_entered(&entered, &from_state, &history_before);
6407        drop(admission);
6408
6409        self.execute_state_enter_actions(&entered, is_reentry).await;
6410        drop(reservation);
6411        info!(to = %entered, "Timeout transition");
6412        Ok(())
6413    }
6414
6415    fn increment_turn(&self) {
6416        if let Some(ref sm) = self.state_machine {
6417            sm.increment_turn();
6418        }
6419    }
6420
6421    fn transitions_available_for_commit(&self) -> Option<(Vec<Transition>, String)> {
6422        let sm = self.state_machine.as_ref()?;
6423        let current = sm.current();
6424        let transitions: Vec<_> = sm
6425            .auto_transitions()
6426            .into_iter()
6427            .filter(|t| match t.cooldown_turns {
6428                Some(cd) if cd > 0 => {
6429                    let resolved = sm.config().resolve_full_path(&current, &t.to);
6430                    !sm.is_on_cooldown(&resolved, cd)
6431                }
6432                _ => true,
6433            })
6434            .collect();
6435        Some((transitions, current))
6436    }
6437
6438    fn transition_reason(transition: &Transition) -> String {
6439        if transition.when.is_empty() {
6440            "guard condition met".to_string()
6441        } else {
6442            transition.when.clone()
6443        }
6444    }
6445
6446    /// Builds transition context with optional staged writes overlaid.
6447    fn build_transition_context(
6448        &self,
6449        user_message: &str,
6450        response: &str,
6451        current_state: &str,
6452        staged: Option<&HashMap<String, Value>>,
6453    ) -> TransitionContext {
6454        let context_map = staged
6455            .map(|writes| self.build_context_with_staged(writes))
6456            .unwrap_or_else(|| self.build_context_with_overlays());
6457        TransitionContext::new(user_message, response, current_state).with_context(context_map)
6458    }
6459
6460    /// Selects a post-response transition without committing side effects.
6461    async fn select_transition_candidate(
6462        &self,
6463        user_message: &str,
6464        response: &str,
6465    ) -> Result<Option<TransitionCandidate>> {
6466        let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6467            return Ok(None);
6468        };
6469        let transitions: Vec<Transition> = transitions
6470            .into_iter()
6471            .filter(|transition| matches!(transition.timing, TransitionTiming::PostResponse))
6472            .collect();
6473        if transitions.is_empty() {
6474            return Ok(None);
6475        }
6476        let Some(evaluator) = self.transition_evaluator.as_ref() else {
6477            return Ok(None);
6478        };
6479        let context = self.build_transition_context(user_message, response, &current_state, None);
6480        let selected = self
6481            .observe_purpose(
6482                ObservationPurpose::StateTransitionEvaluation,
6483                evaluator.select_transition(&transitions, &context),
6484            )
6485            .await?;
6486        Ok(selected.map(|index| {
6487            let transition = transitions[index].clone();
6488            TransitionCandidate::new(
6489                current_state,
6490                transition.clone(),
6491                Self::transition_reason(&transition),
6492            )
6493        }))
6494    }
6495
6496    /// Selects a guard or resolved-intent transition without an LLM call.
6497    fn select_deterministic_transition_candidate(
6498        &self,
6499        user_message: &str,
6500        current_state: &str,
6501        transitions: &[Transition],
6502        staged: &HashMap<String, Value>,
6503    ) -> Option<TransitionCandidate> {
6504        let context = self.build_transition_context(user_message, "", current_state, Some(staged));
6505
6506        for transition in transitions {
6507            if let Some(guard) = transition.guard.as_ref()
6508                && evaluate_guard(guard, &context)
6509            {
6510                return Some(TransitionCandidate::new(
6511                    current_state,
6512                    transition.clone(),
6513                    Self::transition_reason(transition),
6514                ));
6515            }
6516        }
6517
6518        let resolved_intent = context
6519            .context
6520            .get("resolved_intent")
6521            .and_then(Value::as_str)
6522            .filter(|value| !value.is_empty());
6523        if let Some(resolved_intent) = resolved_intent {
6524            for transition in transitions {
6525                if transition.intent.as_deref() == Some(resolved_intent) {
6526                    return Some(TransitionCandidate::new(
6527                        current_state,
6528                        transition.clone(),
6529                        Self::transition_reason(transition),
6530                    ));
6531                }
6532            }
6533        }
6534
6535        None
6536    }
6537
6538    /// Commits a selected transition through the shared post-response path.
6539    async fn commit_transition_candidate(&self, candidate: &TransitionCandidate) -> Result<bool> {
6540        self.commit_transition_target(&candidate.from_state, candidate.target(), &candidate.reason)
6541            .await
6542    }
6543
6544    /// Runs state transition approval before any transition side effects.
6545    async fn approve_transition_target(&self, from_state: &str, target: &str) -> Result<bool> {
6546        let approved = self.check_state_hitl(Some(from_state), target).await?;
6547        if !approved {
6548            info!(to = %target, "State transition rejected by HITL");
6549        }
6550        Ok(approved)
6551    }
6552
6553    /// Applies an approved transition after reserving exit actions and keeping async hooks outside the commit lock.
6554    async fn apply_transition_target(
6555        &self,
6556        from_state: &str,
6557        target: &str,
6558        reason: &str,
6559        staged: Option<&HashMap<String, Value>>,
6560    ) -> Result<bool> {
6561        let Some(ref sm) = self.state_machine else {
6562            return Ok(false);
6563        };
6564        let claim_admission = self.disambiguation_admission.write().await;
6565        if sm.current() != from_state {
6566            return Ok(false);
6567        }
6568        let Some(reservation) = self.reserve_state_transition() else {
6569            return Ok(false);
6570        };
6571        let expected_state_generation = sm.generation();
6572        let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
6573        let history_before = sm.history();
6574        drop(claim_admission);
6575
6576        self.execute_state_exit_actions(from_state).await;
6577
6578        let admission = self.disambiguation_admission.write().await;
6579        if sm.current() != from_state
6580            || sm.generation() != expected_state_generation
6581            || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6582        {
6583            return Ok(false);
6584        }
6585        sm.transition_to(target, reason)?;
6586        self.invalidate_pending_confirmation("state_transition")
6587            .await;
6588        sm.reset_no_transition();
6589        if let Some(staged) = staged {
6590            self.commit_staged_context_writes(staged);
6591        }
6592        let entered = sm.current();
6593        let is_reentry = Self::state_was_previously_entered(&entered, from_state, &history_before);
6594        drop(admission);
6595
6596        self.execute_state_enter_actions(&entered, is_reentry).await;
6597        drop(reservation);
6598        self.hooks
6599            .on_state_transition(Some(from_state), &entered, reason)
6600            .await;
6601        info!(from = %from_state, to = %entered, "State transition");
6602        Ok(true)
6603    }
6604
6605    /// Approves and applies a transition without staged context writes.
6606    async fn commit_transition_target(
6607        &self,
6608        from_state: &str,
6609        target: &str,
6610        reason: &str,
6611    ) -> Result<bool> {
6612        if !self.approve_transition_target(from_state, target).await? {
6613            return Ok(false);
6614        }
6615        self.apply_transition_target(from_state, target, reason, None)
6616            .await
6617    }
6618
6619    /// Applies an approved pre-response transition before redispatch.
6620    async fn apply_pre_response_transition_candidate(
6621        &self,
6622        candidate: &TransitionCandidate,
6623        staged: &HashMap<String, Value>,
6624        processed_input: &str,
6625    ) -> Result<bool> {
6626        self.commit_root_user_message(processed_input).await?;
6627        self.apply_transition_target(
6628            &candidate.from_state,
6629            candidate.target(),
6630            &candidate.reason,
6631            Some(staged),
6632        )
6633        .await
6634    }
6635
6636    /// Commits a pre-response transition after approval and before redispatch.
6637    async fn commit_pre_response_transition_candidate(
6638        &self,
6639        candidate: &TransitionCandidate,
6640        staged: &HashMap<String, Value>,
6641        processed_input: &str,
6642    ) -> Result<bool> {
6643        if !self
6644            .approve_transition_target(&candidate.from_state, candidate.target())
6645            .await?
6646        {
6647            return Ok(false);
6648        }
6649        self.apply_pre_response_transition_candidate(candidate, staged, processed_input)
6650            .await
6651    }
6652
6653    /// Handles post-response transition misses with fallback counters.
6654    async fn handle_transition_miss(&self, current_state: &str) -> Result<bool> {
6655        let Some(ref sm) = self.state_machine else {
6656            return Ok(false);
6657        };
6658        sm.increment_no_transition();
6659        let Some(fallback) = sm.check_fallback() else {
6660            return Ok(false);
6661        };
6662        self.commit_transition_target(current_state, &fallback, "fallback after no transitions")
6663            .await
6664    }
6665
6666    /// Evaluates and commits post-response transitions for the committed response path.
6667    async fn evaluate_transitions(&self, user_message: &str, response: &str) -> Result<bool> {
6668        let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6669            return Ok(false);
6670        };
6671        if transitions.is_empty() {
6672            return Ok(false);
6673        }
6674        if let Some(candidate) = self
6675            .select_transition_candidate(user_message, response)
6676            .await?
6677        {
6678            return self.commit_transition_candidate(&candidate).await;
6679        }
6680        self.handle_transition_miss(&current_state).await
6681    }
6682
6683    /// Attempts deterministic pre-response routing before old-state response generation.
6684    async fn try_pre_response_transition(
6685        &self,
6686        processed_input: &str,
6687    ) -> Result<Option<AgentResponse>> {
6688        let optimization = &self.runtime_config.optimization;
6689        if !optimization.enabled || !optimization.pre_response_deterministic_transitions {
6690            return Ok(None);
6691        }
6692        let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6693            return Ok(None);
6694        };
6695        let eligible: Vec<Transition> = transitions
6696            .into_iter()
6697            .filter(|transition| !transition.requires_response)
6698            .filter(|transition| matches!(transition.timing, TransitionTiming::PreResponse))
6699            .collect();
6700        if eligible.is_empty() {
6701            return Ok(None);
6702        }
6703
6704        let empty_staged = HashMap::new();
6705        let mut extracted_staged: Option<HashMap<String, Value>> = None;
6706        let mut selected: Option<(TransitionCandidate, HashMap<String, Value>)> = None;
6707
6708        for transition in &eligible {
6709            let use_extractors = optimization.pre_response_extractors || transition.run_extractors;
6710            let staged_for_eval = if use_extractors {
6711                if extracted_staged.is_none() {
6712                    extracted_staged =
6713                        Some(self.run_context_extractors_staged(processed_input).await);
6714                }
6715                extracted_staged.as_ref().unwrap_or(&empty_staged)
6716            } else {
6717                &empty_staged
6718            };
6719
6720            if let Some(candidate) = self.select_deterministic_transition_candidate(
6721                processed_input,
6722                &current_state,
6723                std::slice::from_ref(transition),
6724                staged_for_eval,
6725            ) {
6726                let staged_for_commit = if use_extractors {
6727                    staged_for_eval.clone()
6728                } else {
6729                    HashMap::new()
6730                };
6731                selected = Some((candidate, staged_for_commit));
6732                break;
6733            }
6734        }
6735
6736        let Some((candidate, staged)) = selected else {
6737            return Ok(None);
6738        };
6739
6740        if !self
6741            .commit_pre_response_transition_candidate(&candidate, &staged, processed_input)
6742            .await?
6743        {
6744            return Ok(None);
6745        }
6746        self.redispatch_current_state(processed_input)
6747            .await
6748            .map(Some)
6749    }
6750
6751    //
6752    // Speculative branches overlap independent decisions but still commit exactly one path.
6753    // Losing branches must remain data only and must not write memory, run tools, or emit output.
6754    //
6755    async fn try_speculative_branches(
6756        &self,
6757        processed_input: &str,
6758        input_context: &HashMap<String, Value>,
6759    ) -> Result<Option<AgentResponse>> {
6760        let optimization = &self.runtime_config.optimization;
6761        if !optimization.enabled {
6762            return Ok(None);
6763        }
6764
6765        let effective_reasoning_mode = self.get_effective_reasoning_config().mode.clone();
6766        if !matches!(
6767            effective_reasoning_mode,
6768            ReasoningMode::None | ReasoningMode::Auto
6769        ) {
6770            return Ok(None);
6771        }
6772
6773        let mut transition_enabled =
6774            optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
6775        let mut skill_enabled = optimization.speculative_skill_routing
6776            && self.skill_router.is_some()
6777            && self.pending_skill_id.read().is_none();
6778        let mut reasoning_enabled = optimization.speculative_reasoning_auto
6779            && matches!(effective_reasoning_mode, ReasoningMode::Auto);
6780
6781        if matches!(effective_reasoning_mode, ReasoningMode::Auto)
6782            && (!reasoning_enabled || optimization.max_speculative_llm_calls_per_turn < 2)
6783        {
6784            return Ok(None);
6785        }
6786
6787        if !transition_enabled && !skill_enabled && !reasoning_enabled {
6788            return Ok(None);
6789        }
6790
6791        let mut optional_slots = optimization.max_parallel_runtime_tasks.saturating_sub(1);
6792        let mut speculative_call_slots = optimization
6793            .max_speculative_llm_calls_per_turn
6794            .saturating_sub(1);
6795        if reasoning_enabled {
6796            if optional_slots == 0 || speculative_call_slots == 0 {
6797                return Ok(None);
6798            }
6799            optional_slots -= 1;
6800            speculative_call_slots -= 1;
6801        }
6802        if transition_enabled {
6803            if optional_slots == 0 {
6804                transition_enabled = false;
6805            } else {
6806                optional_slots -= 1;
6807            }
6808        }
6809        if skill_enabled && (optional_slots == 0 || speculative_call_slots == 0) {
6810            skill_enabled = false;
6811        }
6812
6813        if !transition_enabled && !skill_enabled && !reasoning_enabled {
6814            return Ok(None);
6815        }
6816
6817        let main_kind = if transition_enabled {
6818            RuntimeOptimizationKind::ParallelStateTransition
6819        } else if skill_enabled {
6820            RuntimeOptimizationKind::SpeculativeSkillRouting
6821        } else {
6822            RuntimeOptimizationKind::SpeculativeReasoningAuto
6823        };
6824        if !self.reserve_active_speculative_llm_call(main_kind) {
6825            return Ok(None);
6826        }
6827
6828        let mut branch_set = ScheduledBranchSet::new(optimization.max_parallel_runtime_tasks)?;
6829        let main_branch = RuntimeBranch::new(
6830            RuntimeTaskPurpose::MainResponse,
6831            main_kind,
6832            RuntimeTaskPriority::Normal,
6833            RuntimeCommitBehavior::FinalResponse,
6834        );
6835        let transition_branch = RuntimeBranch::new(
6836            RuntimeTaskPurpose::StateTransition,
6837            RuntimeOptimizationKind::ParallelStateTransition,
6838            RuntimeTaskPriority::Critical,
6839            RuntimeCommitBehavior::TransitionDecision,
6840        );
6841        let skill_branch = RuntimeBranch::new(
6842            RuntimeTaskPurpose::SkillRouting,
6843            RuntimeOptimizationKind::SpeculativeSkillRouting,
6844            RuntimeTaskPriority::High,
6845            RuntimeCommitBehavior::SkillSelection,
6846        );
6847        let reasoning_branch = RuntimeBranch::new(
6848            RuntimeTaskPurpose::ReasoningJudge,
6849            RuntimeOptimizationKind::SpeculativeReasoningAuto,
6850            RuntimeTaskPriority::Normal,
6851            RuntimeCommitBehavior::ReasoningDecision,
6852        );
6853        let main_id = main_branch.branch_id();
6854        let transition_id = transition_branch.branch_id();
6855        let skill_id = skill_branch.branch_id();
6856        let reasoning_id = reasoning_branch.branch_id();
6857
6858        let main_id_for_future = main_id.clone();
6859        if !branch_set.schedule(
6860            main_branch,
6861            Box::pin(async move {
6862                match crate::optimization::observability::with_branch_observation(
6863                    &main_id_for_future,
6864                    main_kind,
6865                    RuntimeCommitBehavior::FinalResponse,
6866                    self.generate_main_response_draft(processed_input, &ReasoningMode::None),
6867                )
6868                .await
6869                {
6870                    Ok(draft) => RuntimeBranchResult::MainDraft(draft),
6871                    Err(error) => RuntimeBranchResult::Failed(error),
6872                }
6873            }),
6874        ) {
6875            return Ok(None);
6876        }
6877
6878        if transition_enabled {
6879            let transition_id_for_future = transition_id.clone();
6880            if !branch_set.schedule(
6881                transition_branch,
6882                Box::pin(async move {
6883                    match crate::optimization::observability::with_branch_observation(
6884                        &transition_id_for_future,
6885                        RuntimeOptimizationKind::ParallelStateTransition,
6886                        RuntimeCommitBehavior::TransitionDecision,
6887                        self.select_parallel_transition_candidate(processed_input),
6888                    )
6889                    .await
6890                    {
6891                        Ok(ParallelTransitionSelection::Candidate(candidate)) => {
6892                            RuntimeBranchResult::Transition(Some(candidate))
6893                        }
6894                        Ok(ParallelTransitionSelection::NoMatch) => {
6895                            RuntimeBranchResult::Transition(None)
6896                        }
6897                        Ok(ParallelTransitionSelection::ReservationExhausted) => {
6898                            RuntimeBranchResult::Cancelled
6899                        }
6900                        Err(error) => RuntimeBranchResult::Failed(error),
6901                    }
6902                }),
6903            ) {
6904                transition_enabled = false;
6905            }
6906        }
6907
6908        if skill_enabled {
6909            let skill_id_for_future = skill_id.clone();
6910            if !branch_set.schedule(
6911                skill_branch,
6912                Box::pin(async move {
6913                    if !self.reserve_active_speculative_llm_call(
6914                        RuntimeOptimizationKind::SpeculativeSkillRouting,
6915                    ) {
6916                        return RuntimeBranchResult::Cancelled;
6917                    }
6918                    match crate::optimization::observability::with_branch_observation(
6919                        &skill_id_for_future,
6920                        RuntimeOptimizationKind::SpeculativeSkillRouting,
6921                        RuntimeCommitBehavior::SkillSelection,
6922                        self.select_skill_candidate(processed_input),
6923                    )
6924                    .await
6925                    {
6926                        Ok(candidate) => RuntimeBranchResult::Skill(candidate),
6927                        Err(error) => RuntimeBranchResult::Failed(error),
6928                    }
6929                }),
6930            ) {
6931                skill_enabled = false;
6932            }
6933        }
6934
6935        if reasoning_enabled {
6936            let reasoning_id_for_future = reasoning_id.clone();
6937            if !branch_set.schedule(
6938                reasoning_branch,
6939                Box::pin(async move {
6940                    if !self.reserve_active_speculative_llm_call(
6941                        RuntimeOptimizationKind::SpeculativeReasoningAuto,
6942                    ) {
6943                        return RuntimeBranchResult::Cancelled;
6944                    }
6945                    match crate::optimization::observability::with_branch_observation(
6946                        &reasoning_id_for_future,
6947                        RuntimeOptimizationKind::SpeculativeReasoningAuto,
6948                        RuntimeCommitBehavior::ReasoningDecision,
6949                        self.determine_reasoning_mode_strict(processed_input),
6950                    )
6951                    .await
6952                    {
6953                        Ok(mode) => RuntimeBranchResult::Reasoning(mode),
6954                        Err(error) => RuntimeBranchResult::Failed(error),
6955                    }
6956                }),
6957            ) {
6958                reasoning_enabled = false;
6959            }
6960        }
6961
6962        if matches!(effective_reasoning_mode, ReasoningMode::Auto) && !reasoning_enabled {
6963            self.finalize_pending_branches(branch_set.cancel_pending());
6964            return Ok(None);
6965        }
6966
6967        if !transition_enabled && !skill_enabled && !reasoning_enabled {
6968            self.finalize_pending_branches(branch_set.cancel_pending());
6969            return Ok(None);
6970        }
6971
6972        let mut main_pending = true;
6973        let mut skill_pending = skill_enabled;
6974        let mut reasoning_pending = reasoning_enabled;
6975        let mut transition_finalized = !transition_enabled;
6976        let mut skill_finalized = !skill_enabled;
6977        let mut reasoning_finalized = !reasoning_enabled;
6978        let mut main_result: Option<Result<MainResponseDraft>> = None;
6979        let mut transition_candidate: Option<TransitionCandidate> = None;
6980        let mut skill_candidate: Option<SkillCandidate> = None;
6981        let mut reasoning_decision: Option<ReasoningMode> = None;
6982        let mut transition_fallback_required = false;
6983        let mut skill_fallback_required = false;
6984        let mut reasoning_fallback_required = false;
6985
6986        loop {
6987            if let Some(candidate) = transition_candidate.take() {
6988                if self
6989                    .approve_transition_target(&candidate.from_state, candidate.target())
6990                    .await?
6991                {
6992                    // Drop losing provider futures before transition side effects can reuse their shared resources.
6993                    self.finalize_pending_branches(branch_set.cancel_pending());
6994                    if !main_pending {
6995                        self.finalize_branch_loss(
6996                            &main_id,
6997                            main_kind,
6998                            RuntimeCommitBehavior::FinalResponse,
6999                            false,
7000                            main_result.as_ref().map(|result| result.is_err()),
7001                        );
7002                    }
7003                    if skill_enabled && !skill_pending {
7004                        self.finalize_branch_loss(
7005                            &skill_id,
7006                            RuntimeOptimizationKind::SpeculativeSkillRouting,
7007                            RuntimeCommitBehavior::SkillSelection,
7008                            false,
7009                            Some(false),
7010                        );
7011                    }
7012                    if reasoning_enabled && !reasoning_pending {
7013                        self.finalize_branch_loss(
7014                            &reasoning_id,
7015                            RuntimeOptimizationKind::SpeculativeReasoningAuto,
7016                            RuntimeCommitBehavior::ReasoningDecision,
7017                            false,
7018                            Some(false),
7019                        );
7020                    }
7021                    if !self
7022                        .apply_pre_response_transition_candidate(
7023                            &candidate,
7024                            &HashMap::new(),
7025                            processed_input,
7026                        )
7027                        .await?
7028                    {
7029                        self.finalize_optional_branch(
7030                            &transition_id,
7031                            RuntimeOptimizationKind::ParallelStateTransition,
7032                            RuntimeCommitBehavior::TransitionDecision,
7033                            "discarded",
7034                            false,
7035                        );
7036                        return Ok(None);
7037                    }
7038                    self.finalize_optional_branch(
7039                        &transition_id,
7040                        RuntimeOptimizationKind::ParallelStateTransition,
7041                        RuntimeCommitBehavior::TransitionDecision,
7042                        "committed",
7043                        true,
7044                    );
7045                    return self
7046                        .redispatch_current_state(processed_input)
7047                        .await
7048                        .map(Some);
7049                }
7050                self.finalize_optional_branch(
7051                    &transition_id,
7052                    RuntimeOptimizationKind::ParallelStateTransition,
7053                    RuntimeCommitBehavior::TransitionDecision,
7054                    "discarded",
7055                    false,
7056                );
7057                transition_finalized = true;
7058            }
7059
7060            if transition_finalized && skill_candidate.is_some() {
7061                let candidate = skill_candidate.take().unwrap();
7062                self.finalize_optional_branch(
7063                    &skill_id,
7064                    RuntimeOptimizationKind::SpeculativeSkillRouting,
7065                    RuntimeCommitBehavior::SkillSelection,
7066                    "committed",
7067                    true,
7068                );
7069                if !main_pending {
7070                    self.finalize_branch_loss(
7071                        &main_id,
7072                        main_kind,
7073                        RuntimeCommitBehavior::FinalResponse,
7074                        false,
7075                        main_result.as_ref().map(|result| result.is_err()),
7076                    );
7077                }
7078                if reasoning_enabled && !reasoning_pending {
7079                    self.finalize_branch_loss(
7080                        &reasoning_id,
7081                        RuntimeOptimizationKind::SpeculativeReasoningAuto,
7082                        RuntimeCommitBehavior::ReasoningDecision,
7083                        false,
7084                        Some(false),
7085                    );
7086                }
7087                self.finalize_pending_branches(branch_set.cancel_pending());
7088                self.commit_root_user_message(processed_input).await?;
7089                return match self
7090                    .commit_skill_candidate_route_result(candidate, processed_input)
7091                    .await?
7092                {
7093                    SkillRouteResult::Response { skill_id, content } => self
7094                        .handle_skill_response(processed_input, &skill_id, content, input_context)
7095                        .await
7096                        .map(Some),
7097                    SkillRouteResult::NeedsClarification {
7098                        response,
7099                        ownership,
7100                    } => {
7101                        let admission = self
7102                            .admit_optional_disambiguation_ownership(ownership)
7103                            .await?;
7104                        if response
7105                            .metadata
7106                            .as_ref()
7107                            .and_then(|m| m.get("disambiguation"))
7108                            .and_then(|d| d.get("status"))
7109                            .and_then(|s| s.as_str())
7110                            == Some("awaiting_clarification")
7111                        {
7112                            self.memory
7113                                .add_message(ChatMessage::assistant(&response.content))
7114                                .await?;
7115                        }
7116                        drop(admission);
7117                        self.finish_turn_if_root(&response).await?;
7118                        Ok(Some(response))
7119                    }
7120                    SkillRouteResult::NoMatch => Ok(None),
7121                };
7122            }
7123
7124            if transition_finalized
7125                && skill_finalized
7126                && let Some(reasoning_mode) = reasoning_decision.take()
7127            {
7128                if !matches!(reasoning_mode, ReasoningMode::None) {
7129                    self.finalize_optional_branch(
7130                        &reasoning_id,
7131                        RuntimeOptimizationKind::SpeculativeReasoningAuto,
7132                        RuntimeCommitBehavior::ReasoningDecision,
7133                        "committed",
7134                        true,
7135                    );
7136                    if !main_pending {
7137                        self.finalize_branch_loss(
7138                            &main_id,
7139                            main_kind,
7140                            RuntimeCommitBehavior::FinalResponse,
7141                            false,
7142                            main_result.as_ref().map(|result| result.is_err()),
7143                        );
7144                    }
7145                    self.finalize_pending_branches(branch_set.cancel_pending());
7146                    self.commit_root_user_message(processed_input).await?;
7147                    return if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
7148                        self.handle_plan_and_execute(processed_input, input_context, true)
7149                            .await
7150                            .map(Some)
7151                    } else {
7152                        self.run_committed_response_loop_with_reasoning(
7153                            processed_input,
7154                            input_context,
7155                            reasoning_mode,
7156                            true,
7157                        )
7158                        .await
7159                        .map(Some)
7160                    };
7161                }
7162                self.finalize_optional_branch(
7163                    &reasoning_id,
7164                    RuntimeOptimizationKind::SpeculativeReasoningAuto,
7165                    RuntimeCommitBehavior::ReasoningDecision,
7166                    "committed",
7167                    true,
7168                );
7169                reasoning_finalized = true;
7170            }
7171
7172            if transition_finalized && skill_finalized && reasoning_finalized {
7173                if transition_fallback_required
7174                    || skill_fallback_required
7175                    || reasoning_fallback_required
7176                {
7177                    if !main_pending {
7178                        self.finalize_branch_loss(
7179                            &main_id,
7180                            main_kind,
7181                            RuntimeCommitBehavior::FinalResponse,
7182                            false,
7183                            main_result.as_ref().map(|result| result.is_err()),
7184                        );
7185                    }
7186                    self.finalize_pending_branches(branch_set.cancel_pending());
7187                    return Ok(None);
7188                }
7189
7190                if let Some(result) = main_result.take() {
7191                    let draft = match result {
7192                        Ok(draft) => draft,
7193                        Err(error) => {
7194                            self.finalize_optional_branch(
7195                                &main_id,
7196                                main_kind,
7197                                RuntimeCommitBehavior::FinalResponse,
7198                                "failed",
7199                                false,
7200                            );
7201                            self.finalize_pending_branches(branch_set.cancel_pending());
7202                            return Err(error);
7203                        }
7204                    };
7205                    self.finalize_optional_branch(
7206                        &main_id,
7207                        main_kind,
7208                        RuntimeCommitBehavior::FinalResponse,
7209                        "committed",
7210                        true,
7211                    );
7212                    self.finalize_pending_branches(branch_set.cancel_pending());
7213                    return self
7214                        .commit_main_response_draft(
7215                            processed_input,
7216                            input_context,
7217                            draft,
7218                            ReasoningMode::None,
7219                            reasoning_enabled,
7220                        )
7221                        .await
7222                        .map(Some);
7223                }
7224            }
7225
7226            if branch_set.is_empty() {
7227                return Ok(None);
7228            }
7229
7230            let Some(outcome) = branch_set.next_completed().await else {
7231                return Ok(None);
7232            };
7233            let branch_id = outcome.branch.branch_id();
7234            match outcome.result {
7235                RuntimeBranchResult::MainDraft(draft) => {
7236                    main_pending = false;
7237                    main_result = Some(Ok(draft));
7238                }
7239                RuntimeBranchResult::Transition(candidate) => {
7240                    if let Some(candidate) = candidate {
7241                        transition_candidate = Some(candidate);
7242                    } else {
7243                        self.finalize_optional_branch(
7244                            &transition_id,
7245                            RuntimeOptimizationKind::ParallelStateTransition,
7246                            RuntimeCommitBehavior::TransitionDecision,
7247                            "discarded",
7248                            false,
7249                        );
7250                        transition_finalized = true;
7251                    }
7252                }
7253                RuntimeBranchResult::Skill(candidate) => {
7254                    skill_pending = false;
7255                    if let Some(candidate) = candidate {
7256                        skill_candidate = Some(candidate);
7257                    } else {
7258                        self.finalize_optional_branch(
7259                            &skill_id,
7260                            RuntimeOptimizationKind::SpeculativeSkillRouting,
7261                            RuntimeCommitBehavior::SkillSelection,
7262                            "discarded",
7263                            false,
7264                        );
7265                        skill_finalized = true;
7266                    }
7267                }
7268                RuntimeBranchResult::Reasoning(mode) => {
7269                    reasoning_pending = false;
7270                    reasoning_decision = Some(mode);
7271                }
7272                RuntimeBranchResult::Failed(error) => {
7273                    if branch_id == main_id {
7274                        main_pending = false;
7275                        main_result = Some(Err(error));
7276                    } else if branch_id == transition_id {
7277                        self.finalize_optional_branch(
7278                            &transition_id,
7279                            RuntimeOptimizationKind::ParallelStateTransition,
7280                            RuntimeCommitBehavior::TransitionDecision,
7281                            "failed",
7282                            false,
7283                        );
7284                        transition_finalized = true;
7285                    } else if branch_id == skill_id {
7286                        skill_pending = false;
7287                        self.finalize_optional_branch(
7288                            &skill_id,
7289                            RuntimeOptimizationKind::SpeculativeSkillRouting,
7290                            RuntimeCommitBehavior::SkillSelection,
7291                            "failed",
7292                            false,
7293                        );
7294                        skill_finalized = true;
7295                    } else if branch_id == reasoning_id {
7296                        reasoning_pending = false;
7297                        self.finalize_optional_branch(
7298                            &reasoning_id,
7299                            RuntimeOptimizationKind::SpeculativeReasoningAuto,
7300                            RuntimeCommitBehavior::ReasoningDecision,
7301                            "failed",
7302                            false,
7303                        );
7304                        reasoning_finalized = true;
7305                    }
7306                }
7307                RuntimeBranchResult::Cancelled => {
7308                    self.finalize_optional_branch(
7309                        &branch_id,
7310                        outcome.branch.optimization,
7311                        outcome.branch.commit_behavior,
7312                        "cancelled",
7313                        false,
7314                    );
7315                    if branch_id == main_id {
7316                        main_pending = false;
7317                        main_result =
7318                            Some(Err(AgentError::Other("main branch cancelled".to_string())));
7319                    } else if branch_id == transition_id {
7320                        transition_finalized = true;
7321                        transition_fallback_required = true;
7322                    } else if branch_id == skill_id {
7323                        skill_pending = false;
7324                        skill_finalized = true;
7325                        skill_fallback_required = true;
7326                    } else if branch_id == reasoning_id {
7327                        reasoning_pending = false;
7328                        reasoning_finalized = true;
7329                        reasoning_fallback_required = true;
7330                    }
7331                }
7332            }
7333        }
7334    }
7335
7336    fn finalize_pending_branches(&self, branches: Vec<RuntimeBranch>) {
7337        for branch in branches {
7338            self.finalize_optional_branch(
7339                &branch.branch_id(),
7340                branch.optimization,
7341                branch.commit_behavior,
7342                "cancelled",
7343                false,
7344            );
7345        }
7346    }
7347
7348    //
7349    // Pending losers are reported as cancelled because their futures are dropped before completion.
7350    // Completed losers keep failed or discarded status based on their recorded result.
7351    //
7352    fn finalize_branch_loss(
7353        &self,
7354        branch_id: &str,
7355        optimization: RuntimeOptimizationKind,
7356        commit_behavior: RuntimeCommitBehavior,
7357        pending: bool,
7358        completed_failed: Option<bool>,
7359    ) {
7360        let status = if pending {
7361            "cancelled"
7362        } else if completed_failed.unwrap_or(false) {
7363            "failed"
7364        } else {
7365            "discarded"
7366        };
7367        self.finalize_optional_branch(branch_id, optimization, commit_behavior, status, false);
7368    }
7369
7370    //
7371    // Finalization is separated from commit so losing branches remain observable.
7372    // This helper must not mutate runtime state other than observability.
7373    //
7374    fn finalize_optional_branch(
7375        &self,
7376        branch_id: &str,
7377        optimization: RuntimeOptimizationKind,
7378        commit_behavior: RuntimeCommitBehavior,
7379        status: &str,
7380        winner: bool,
7381    ) {
7382        crate::optimization::observability::finalize_branch(
7383            self.observability_manager.as_ref(),
7384            branch_id,
7385            status,
7386            winner,
7387            optimization,
7388            commit_behavior,
7389        );
7390    }
7391
7392    //
7393    // This is only an eligibility check.
7394    // Actual transition selection happens in select_parallel_transition_candidate.
7395    //
7396    fn has_parallel_transition_candidates(&self) -> bool {
7397        self.transitions_available_for_commit()
7398            .map(|(transitions, _)| {
7399                transitions
7400                    .iter()
7401                    .any(|transition| matches!(transition.timing, TransitionTiming::Parallel))
7402            })
7403            .unwrap_or(false)
7404    }
7405
7406    //
7407    // Parallel transition prompts must not depend on assistant response text.
7408    // Keep this branch response-independent or it can race against invalid context.
7409    //
7410    async fn select_parallel_transition_candidate(
7411        &self,
7412        processed_input: &str,
7413    ) -> Result<ParallelTransitionSelection> {
7414        let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
7415            return Ok(ParallelTransitionSelection::NoMatch);
7416        };
7417        let parallel: Vec<Transition> = transitions
7418            .into_iter()
7419            .filter(|transition| matches!(transition.timing, TransitionTiming::Parallel))
7420            .filter(|transition| !transition.requires_response)
7421            .collect();
7422        if parallel.is_empty() {
7423            return Ok(ParallelTransitionSelection::NoMatch);
7424        }
7425        let empty_staged = HashMap::new();
7426        if let Some(candidate) = self.select_deterministic_transition_candidate(
7427            processed_input,
7428            &current_state,
7429            &parallel,
7430            &empty_staged,
7431        ) {
7432            return Ok(ParallelTransitionSelection::Candidate(candidate));
7433        }
7434        let when_transitions: Vec<(usize, &Transition)> = parallel
7435            .iter()
7436            .enumerate()
7437            .filter(|(_, transition)| !transition.when.trim().is_empty())
7438            .collect();
7439        if when_transitions.is_empty() {
7440            return Ok(ParallelTransitionSelection::NoMatch);
7441        }
7442        let llm = self
7443            .llm_registry
7444            .router()
7445            .or_else(|_| self.llm_registry.default())
7446            .map_err(|e| AgentError::Config(e.to_string()))?;
7447        let conditions = when_transitions
7448            .iter()
7449            .enumerate()
7450            .map(|(display_idx, (_, transition))| {
7451                format!("{}. {}", display_idx + 1, transition.when)
7452            })
7453            .collect::<Vec<_>>()
7454            .join("\n");
7455        if !self
7456            .reserve_active_speculative_llm_call(RuntimeOptimizationKind::ParallelStateTransition)
7457        {
7458            return Ok(ParallelTransitionSelection::ReservationExhausted);
7459        }
7460        let context_preview = self.branch_context_preview();
7461        let prompt = format!(
7462            "Based only on the current user message and context, which transition condition is met?\n\nCurrent state: {}\nUser message: {}\nContext:\n{}\n\nConditions:\n{}\n0. None of the above\n\nReply with ONLY the number (0-{}).",
7463            current_state,
7464            processed_input,
7465            context_preview,
7466            conditions,
7467            when_transitions.len()
7468        );
7469        let response = self
7470            .observe_purpose(
7471                ObservationPurpose::StateTransitionEvaluation,
7472                llm.complete(&[ChatMessage::user(prompt)], None),
7473            )
7474            .await
7475            .map_err(|e| AgentError::LLM(e.to_string()))?;
7476        let choice = response.content.trim().parse::<usize>().unwrap_or(0);
7477        if choice == 0 || choice > when_transitions.len() {
7478            return Ok(ParallelTransitionSelection::NoMatch);
7479        }
7480        let transition = when_transitions[choice - 1].1.clone();
7481        Ok(ParallelTransitionSelection::Candidate(
7482            TransitionCandidate::new(
7483                current_state,
7484                transition.clone(),
7485                Self::transition_reason(&transition),
7486            ),
7487        ))
7488    }
7489
7490    /// Re-enters the runtime loop after an optimized transition commits.
7491    async fn redispatch_current_state(&self, processed_input: &str) -> Result<AgentResponse> {
7492        const MAX_REDISPATCH_DEPTH: u32 = 3;
7493        let current_depth = *self.redispatch_depth.read();
7494        if current_depth >= MAX_REDISPATCH_DEPTH {
7495            warn!(depth = current_depth, "Re-dispatch depth limit reached");
7496            let response = AgentResponse::new("");
7497            self.finish_turn_if_root(&response).await?;
7498            return Ok(response);
7499        }
7500        *self.redispatch_depth.write() += 1;
7501        if let Some(context) = self.active_turn_context.write().as_mut() {
7502            context.enter_redispatch();
7503        }
7504        let result = Box::pin(self.run_loop_internal(processed_input)).await;
7505        *self.redispatch_depth.write() -= 1;
7506        if let Some(context) = self.active_turn_context.write().as_mut() {
7507            context.exit_redispatch();
7508        }
7509        let response = result?;
7510        self.finish_turn_if_root(&response).await?;
7511        Ok(response)
7512    }
7513
7514    /// Runs final response hooks and maintenance only for the root dispatch.
7515    async fn finish_turn_if_root(&self, response: &AgentResponse) -> Result<()> {
7516        if *self.redispatch_depth.read() == 0 {
7517            self.post_turn_session_lifecycle().await?;
7518            if let Some(context) = self.active_turn_context.write().as_mut() {
7519                context.mark_post_turn_lifecycle_completed();
7520            }
7521            self.hooks.on_response(response).await;
7522            self.end_root_turn();
7523        }
7524        Ok(())
7525    }
7526
7527    /// Execute on_exit actions for a state being left.
7528    async fn execute_state_exit_actions(&self, state_path: &str) {
7529        if let Some(ref sm) = self.state_machine
7530            && let Some(def) = sm.get_definition(state_path)
7531            && !def.on_exit.is_empty()
7532        {
7533            debug!(state = %state_path, count = def.on_exit.len(), "Executing on_exit actions");
7534            self.execute_state_actions(&def.on_exit).await;
7535        }
7536    }
7537
7538    /// Returns whether a transition target had already been entered before this transition.
7539    fn state_was_previously_entered(
7540        state_path: &str,
7541        from_state: &str,
7542        history_before: &[StateTransitionEvent],
7543    ) -> bool {
7544        state_path == from_state
7545            || history_before
7546                .iter()
7547                .any(|event| event.from == state_path || event.to == state_path)
7548    }
7549
7550    /// Execute on_enter (or on_reenter) actions for a state being entered.
7551    async fn execute_state_enter_actions(&self, state_path: &str, is_reentry: bool) {
7552        if let Some(ref sm) = self.state_machine
7553            && let Some(def) = sm.get_definition(state_path)
7554        {
7555            if is_reentry && !def.on_reenter.is_empty() {
7556                debug!(state = %state_path, count = def.on_reenter.len(), "Executing on_reenter actions");
7557                self.execute_state_actions(&def.on_reenter).await;
7558            } else if !def.on_enter.is_empty() {
7559                debug!(state = %state_path, count = def.on_enter.len(), "Executing on_enter actions");
7560                self.execute_state_actions(&def.on_enter).await;
7561            }
7562        }
7563    }
7564
7565    /// Execute a list of state actions (tool calls, skill invocations, context updates, LLM prompts).
7566    async fn execute_state_actions(&self, actions: &[StateAction]) {
7567        for (action_index, action) in actions.iter().enumerate() {
7568            match action {
7569                StateAction::Tool { tool, args } => {
7570                    let raw_args = args.clone().unwrap_or(Value::Object(Default::default()));
7571                    let args_value = self.render_action_args(&raw_args);
7572                    let state = self.state_machine.as_ref().map(|sm| sm.current());
7573                    let request = ToolExecutionRequest::new(
7574                        uuid::Uuid::new_v4().to_string(),
7575                        tool.clone(),
7576                        args_value,
7577                        ToolCallSource::StateAction {
7578                            state,
7579                            action_index,
7580                        },
7581                    );
7582                    match self.execute_tool_record(request).await {
7583                        Ok(record) if record.success => {
7584                            debug!(tool = %record.canonical_id, "State action: tool executed");
7585                            let _ = self.context_manager.set(
7586                                "last_tool_result",
7587                                serde_json::Value::String(record.model_output_string()),
7588                            );
7589                            let _ = self.context_manager.set(
7590                                "last_tool_record",
7591                                serde_json::to_value(record).unwrap_or(Value::Null),
7592                            );
7593                        }
7594                        Ok(record) => {
7595                            warn!(tool = %record.canonical_id, error = %record.output, "State action: tool failed");
7596                        }
7597                        Err(e) => {
7598                            warn!(tool = %tool, error = %e, "State action: tool failed")
7599                        }
7600                    }
7601                }
7602                StateAction::Skill { skill } => {
7603                    if let Some(ref executor) = self.skill_executor {
7604                        if let Some(def) = self.skills.iter().find(|s| s.id == *skill) {
7605                            match executor
7606                                .execute_with_invoker(def, "", serde_json::json!({}), self)
7607                                .await
7608                            {
7609                                Ok(_) => debug!(skill = %skill, "State action: skill executed"),
7610                                Err(e) => {
7611                                    warn!(skill = %skill, error = %e, "State action: skill failed")
7612                                }
7613                            }
7614                        } else {
7615                            warn!(skill = %skill, "State action: skill not found");
7616                        }
7617                    }
7618                }
7619                StateAction::SetContext { set_context } => {
7620                    for (key, value) in set_context {
7621                        if let Err(e) = self.context_manager.set(key, value.clone()) {
7622                            warn!(key = %key, error = %e, "State action: set_context failed");
7623                        } else {
7624                            debug!(key = %key, "State action: context set");
7625                        }
7626                    }
7627                }
7628                StateAction::Prompt {
7629                    prompt,
7630                    llm,
7631                    store_as,
7632                } => {
7633                    let llm_result = if let Some(alias) = llm {
7634                        self.llm_registry.get(alias)
7635                    } else {
7636                        self.llm_registry.default()
7637                    };
7638                    match llm_result {
7639                        Ok(llm_provider) => {
7640                            // Render template variables and include conversation context
7641                            let context = self.build_context_with_overlays();
7642                            let rendered_prompt = self
7643                                .template_renderer
7644                                .render(prompt, &context)
7645                                .unwrap_or_else(|_| prompt.clone());
7646                            let recent =
7647                                self.memory.get_messages(Some(5)).await.unwrap_or_default();
7648                            let mut messages: Vec<ChatMessage> = recent;
7649                            messages.push(ChatMessage::user(&rendered_prompt));
7650                            match self
7651                                .observe_purpose(
7652                                    ObservationPurpose::StateAction,
7653                                    llm_provider.complete(&messages, None),
7654                                )
7655                                .await
7656                            {
7657                                Ok(response) => {
7658                                    if let Some(key) = store_as {
7659                                        let _ = self
7660                                            .context_manager
7661                                            .set(key, Value::String(response.content));
7662                                        debug!(key = %key, "State action: prompt result stored");
7663                                    }
7664                                }
7665                                Err(e) => {
7666                                    warn!(error = %e, "State action: prompt LLM call failed");
7667                                }
7668                            }
7669                        }
7670                        Err(e) => {
7671                            warn!(error = %e, "State action: LLM not found for prompt");
7672                        }
7673                    }
7674                }
7675            }
7676        }
7677    }
7678
7679    async fn run_context_extractors_staged(&self, user_message: &str) -> HashMap<String, Value> {
7680        let extractors = match &self.state_machine {
7681            Some(sm) => match sm.current_definition() {
7682                Some(def) if !def.extract.is_empty() => def.extract.clone(),
7683                _ => return HashMap::new(),
7684            },
7685            None => return HashMap::new(),
7686        };
7687
7688        let mut staged = HashMap::new();
7689        for extractor in &extractors {
7690            let prompt = if let Some(ref custom) = extractor.llm_extract {
7691                format!(
7692                    "User message:\n\"{}\"\n\nInstruction:\n{}",
7693                    user_message, custom
7694                )
7695            } else if let Some(ref desc) = extractor.description {
7696                format!(
7697                    "From the following message, extract: {}\n\n\
7698                     Message: \"{}\"\n\n\
7699                     If the information is present, return ONLY the extracted value.\n\
7700                     If NOT present, return exactly: __NONE__",
7701                    desc, user_message
7702                )
7703            } else {
7704                continue;
7705            };
7706
7707            let llm = match self
7708                .llm_registry
7709                .get(&extractor.llm)
7710                .or_else(|_| self.llm_registry.get("router"))
7711                .or_else(|_| self.llm_registry.get("default"))
7712            {
7713                Ok(llm) => llm,
7714                Err(e) => {
7715                    warn!(key = %extractor.key, error = %e, "Extractor LLM not found");
7716                    continue;
7717                }
7718            };
7719
7720            let messages = vec![ChatMessage::user(&prompt)];
7721            match self
7722                .observe_purpose(
7723                    ObservationPurpose::ContextExtraction,
7724                    llm.complete(&messages, None),
7725                )
7726                .await
7727            {
7728                Ok(response) => {
7729                    let value = response.content.trim().to_string();
7730                    if value != "__NONE__" && !value.is_empty() {
7731                        staged.insert(
7732                            extractor.key.clone(),
7733                            serde_json::Value::String(value.clone()),
7734                        );
7735                        debug!(key = %extractor.key, value = %value, "Context extracted");
7736                    } else if extractor.required {
7737                        warn!(key = %extractor.key, "Required extraction returned no value");
7738                    }
7739                }
7740                Err(e) => {
7741                    warn!(key = %extractor.key, error = %e, "Context extraction LLM call failed");
7742                }
7743            }
7744        }
7745        staged
7746    }
7747
7748    fn commit_staged_context_writes(&self, staged: &HashMap<String, Value>) {
7749        for (key, value) in staged {
7750            if let Err(error) = self.context_manager.update(key, value.clone()) {
7751                warn!(key = %key, error = %error, "staged context write failed");
7752            }
7753        }
7754    }
7755
7756    /// Run context extractors for the current state on the user's input.
7757    async fn run_context_extractors(&self, user_message: &str) {
7758        let staged = self.run_context_extractors_staged(user_message).await;
7759        self.commit_staged_context_writes(&staged);
7760    }
7761
7762    async fn check_memory_compression(&self) -> Result<()> {
7763        if self.memory.needs_compression() {
7764            let result = self.memory.compress(None).await?;
7765            if let CompressResult::Compressed {
7766                messages_summarized,
7767                new_summary_length,
7768                tokens_saved,
7769            } = result
7770            {
7771                let event = MemoryCompressEvent::new(
7772                    messages_summarized,
7773                    tokens_saved,
7774                    new_summary_length as u32,
7775                );
7776                self.hooks.on_memory_compress(&event).await;
7777                debug!(
7778                    messages = messages_summarized,
7779                    tokens_saved = tokens_saved,
7780                    "Memory compressed"
7781                );
7782            }
7783        }
7784
7785        // Handle overflow AFTER compression, then check warning threshold
7786        self.handle_memory_overflow().await?;
7787        self.check_memory_budget().await;
7788
7789        Ok(())
7790    }
7791
7792    async fn check_memory_budget(&self) {
7793        let Some(ref budget) = self.memory_token_budget else {
7794            return;
7795        };
7796
7797        let context = match self.memory.get_context().await {
7798            Ok(ctx) => ctx,
7799            Err(_) => return,
7800        };
7801
7802        // Overall budget warning
7803        let used_tokens = context.estimated_tokens();
7804        if budget.is_over_warn_threshold(used_tokens) {
7805            let event = MemoryBudgetEvent::new("memory", used_tokens, budget.total);
7806            self.hooks.on_memory_budget_warning(&event).await;
7807            debug!(
7808                used = used_tokens,
7809                total = budget.total,
7810                percent = event.usage_percent,
7811                "Memory budget warning"
7812            );
7813        }
7814
7815        // Per-component warning: summary
7816        if let Some(ref summary) = context.summary {
7817            let summary_tokens = ai_agents_memory::estimate_tokens(summary);
7818            let summary_budget = budget.allocation.summary;
7819            if summary_budget > 0 {
7820                let warn_threshold =
7821                    (summary_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
7822                if summary_tokens >= warn_threshold {
7823                    let event = MemoryBudgetEvent::new("summary", summary_tokens, summary_budget);
7824                    self.hooks.on_memory_budget_warning(&event).await;
7825                }
7826            }
7827        }
7828
7829        // Per-component warning: recent_messages
7830        let recent_tokens: u32 = context
7831            .messages
7832            .iter()
7833            .map(ai_agents_memory::estimate_message_tokens)
7834            .sum();
7835        let recent_budget = budget.allocation.recent_messages;
7836        if recent_budget > 0 {
7837            let warn_threshold =
7838                (recent_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
7839            if recent_tokens >= warn_threshold {
7840                let event = MemoryBudgetEvent::new("recent_messages", recent_tokens, recent_budget);
7841                self.hooks.on_memory_budget_warning(&event).await;
7842            }
7843        }
7844
7845        let relationship_budget = budget.allocation.relationships;
7846        if relationship_budget > 0 {
7847            let relationship_tokens = self
7848                .relationship_memory_text()
7849                .map(|text| ai_agents_memory::estimate_tokens(&text))
7850                .unwrap_or(0);
7851            let warn_threshold =
7852                (relationship_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
7853            if relationship_tokens >= warn_threshold {
7854                let event = MemoryBudgetEvent::new(
7855                    "relationships",
7856                    relationship_tokens,
7857                    relationship_budget,
7858                );
7859                self.hooks.on_memory_budget_warning(&event).await;
7860            }
7861        }
7862    }
7863
7864    async fn handle_memory_overflow(&self) -> Result<()> {
7865        let Some(ref budget) = self.memory_token_budget else {
7866            return Ok(());
7867        };
7868
7869        let context = self.memory.get_context().await?;
7870        let used_tokens = context.estimated_tokens();
7871
7872        if used_tokens <= budget.total {
7873            return Ok(());
7874        }
7875
7876        match budget.overflow_strategy {
7877            OverflowStrategy::TruncateOldest => {
7878                let tokens_to_free = used_tokens - budget.total;
7879                let messages_to_evict = self.calculate_eviction_count(tokens_to_free);
7880                if messages_to_evict > 0 {
7881                    self.evict_messages(messages_to_evict, EvictionReason::TokenBudgetExceeded)
7882                        .await?;
7883                }
7884            }
7885            OverflowStrategy::SummarizeMore => {
7886                let max_attempts = context.total_messages.max(1);
7887                for _ in 0..max_attempts {
7888                    match self.memory.compress(None).await? {
7889                        CompressResult::Compressed {
7890                            messages_summarized,
7891                            ..
7892                        } if messages_summarized > 0 => {
7893                            let context = self.memory.get_context().await?;
7894                            if context.estimated_tokens() <= budget.total {
7895                                return Ok(());
7896                            }
7897                        }
7898                        _ => break,
7899                    }
7900                }
7901                let context = self.memory.get_context().await?;
7902                let used_tokens = context.estimated_tokens();
7903                if used_tokens > budget.total {
7904                    return Err(AgentError::MemoryBudgetExceeded {
7905                        used: used_tokens,
7906                        budget: budget.total,
7907                    });
7908                }
7909            }
7910            OverflowStrategy::Error => {
7911                return Err(AgentError::MemoryBudgetExceeded {
7912                    used: used_tokens,
7913                    budget: budget.total,
7914                });
7915            }
7916        }
7917        Ok(())
7918    }
7919
7920    fn calculate_eviction_count(&self, tokens_to_free: u32) -> usize {
7921        // Estimate ~50 tokens per message on average
7922        ((tokens_to_free as f64 / 50.0).ceil() as usize).max(1)
7923    }
7924
7925    async fn evict_messages(&self, count: usize, reason: EvictionReason) -> Result<()> {
7926        let evicted = self.memory.evict_oldest(count).await?;
7927        if !evicted.is_empty() {
7928            let event = MemoryEvictEvent {
7929                reason,
7930                messages_evicted: evicted.len(),
7931                importance_scores: vec![],
7932            };
7933            self.hooks.on_memory_evict(&event).await;
7934            debug!(count = evicted.len(), "Messages evicted from memory");
7935        }
7936        Ok(())
7937    }
7938
7939    #[instrument(skip(self, input), fields(agent = %self.info.name))]
7940    async fn determine_reasoning_mode(&self, input: &str) -> Result<ReasoningMode> {
7941        match self.determine_reasoning_mode_strict(input).await {
7942            Ok(mode) => Ok(mode),
7943            Err(_) => Ok(ReasoningMode::None),
7944        }
7945    }
7946
7947    async fn determine_reasoning_mode_strict(&self, input: &str) -> Result<ReasoningMode> {
7948        let effective_config = self.get_effective_reasoning_config();
7949
7950        if !matches!(effective_config.mode, ReasoningMode::Auto) {
7951            return Ok(effective_config.mode.clone());
7952        }
7953
7954        let judge_llm = effective_config
7955            .judge_llm
7956            .as_ref()
7957            .and_then(|alias| self.llm_registry.get(alias).ok())
7958            .or_else(|| self.llm_registry.router().ok())
7959            .or_else(|| self.llm_registry.default().ok());
7960
7961        let Some(llm) = judge_llm else {
7962            return Ok(ReasoningMode::None);
7963        };
7964
7965        let prompt = format!(
7966            r#"Analyze this user request and determine the appropriate reasoning mode.
7967
7968User request: "{}"
7969
7970Choose ONE of these modes:
7971- none: Simple queries, greetings, direct answers (fastest)
7972- cot: Complex analysis, multi-step reasoning, math problems
7973- react: Tasks requiring multiple tool calls with observation
7974- plan_and_execute: Complex multi-step tasks requiring coordination
7975
7976Respond with ONLY the mode name (none, cot, react, or plan_and_execute)."#,
7977            input
7978        );
7979
7980        let messages = vec![ChatMessage::user(&prompt)];
7981        let response = self
7982            .observe_purpose(
7983                ObservationPurpose::ReflectionDecision,
7984                llm.complete(&messages, None),
7985            )
7986            .await
7987            .map_err(|e| AgentError::LLM(e.to_string()))?;
7988
7989        let mode_str = response.content.trim().to_lowercase();
7990        Ok(match mode_str.as_str() {
7991            "cot" => ReasoningMode::CoT,
7992            "react" => ReasoningMode::React,
7993            "plan_and_execute" => ReasoningMode::PlanAndExecute,
7994            _ => ReasoningMode::None,
7995        })
7996    }
7997
7998    async fn should_reflect(&self, input: &str, response: &str) -> Result<bool> {
7999        let effective_config = self.get_effective_reflection_config();
8000
8001        if !effective_config.requires_evaluation() {
8002            return Ok(false);
8003        }
8004
8005        if effective_config.is_enabled() {
8006            return Ok(true);
8007        }
8008
8009        let evaluator_llm = effective_config
8010            .evaluator_llm
8011            .as_ref()
8012            .and_then(|alias| self.llm_registry.get(alias).ok())
8013            .or_else(|| self.llm_registry.router().ok())
8014            .or_else(|| self.llm_registry.default().ok());
8015
8016        let Some(llm) = evaluator_llm else {
8017            return Ok(false);
8018        };
8019
8020        let response_preview: String = response.chars().take(500).collect();
8021        let prompt = format!(
8022            r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
8023
8024User query: "{}"
8025Response: "{}"
8026
8027Answer YES or NO only."#,
8028            input, response_preview
8029        );
8030
8031        let messages = vec![ChatMessage::user(&prompt)];
8032        let result = self
8033            .observe_purpose(
8034                ObservationPurpose::ReflectionDecision,
8035                llm.complete(&messages, None),
8036            )
8037            .await;
8038
8039        match result {
8040            Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
8041            Err(_) => Ok(false),
8042        }
8043    }
8044
8045    fn build_cot_system_prompt(&self, base_prompt: &str) -> String {
8046        format!(
8047            "{}\n\n<instruction>\nThink through this step by step before answering:\n1. Understand what is being asked\n2. Break down the problem\n3. Work through each part\n4. Provide your final answer\n\nShow your thinking process, then give your final answer.\n</instruction>",
8048            base_prompt
8049        )
8050    }
8051
8052    fn build_react_system_prompt(&self, base_prompt: &str) -> String {
8053        format!(
8054            "{}\n\n<instruction>\nUse the Reason-Act-Observe pattern:\n1. Thought: Think about what to do\n2. Action: Use a tool if needed\n3. Observation: Analyze the result\n4. Repeat until you have the answer\n\nFormat your response showing Thought/Action/Observation steps.\n</instruction>",
8055            base_prompt
8056        )
8057    }
8058
8059    async fn generate_plan(&self, input: &str) -> Result<Plan> {
8060        let effective = self.get_effective_reasoning_config();
8061        let planning_config = effective.get_planning();
8062
8063        let planner_llm = planning_config
8064            .and_then(|c| c.planner_llm.as_ref())
8065            .and_then(|alias| self.llm_registry.get(alias).ok())
8066            .or_else(|| self.llm_registry.router().ok())
8067            .or_else(|| self.llm_registry.default().ok())
8068            .ok_or_else(|| AgentError::Config("No LLM available for planning".into()))?;
8069
8070        let mut available_tool_ids: Vec<String> = self
8071            .get_available_tool_ids()
8072            .await
8073            .unwrap_or_else(|_| self.tools.list_ids());
8074        let mut available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
8075
8076        // Apply planning-level tool and skill filters.
8077        if let Some(config) = planning_config {
8078            if !config.available.tools.is_all() {
8079                available_tool_ids.retain(|t| config.available.tools.allows(t));
8080            }
8081            if !config.available.skills.is_all() {
8082                available_skills.retain(|s| config.available.skills.allows(s));
8083            }
8084        }
8085
8086        // Build tool descriptions with argument schemas so the planner
8087        // knows how to construct valid args for each step.
8088        let tool_descriptions: Vec<String> = available_tool_ids
8089            .iter()
8090            .filter_map(|id| {
8091                self.tools.get(id).map(|tool| {
8092                    let schema = tool.input_schema();
8093                    let args_desc = schema
8094                        .get("properties")
8095                        .and_then(|p| serde_json::to_string(p).ok())
8096                        .unwrap_or_else(|| "{}".to_string());
8097                    format!(
8098                        "- {} ({}): {}\n  Arguments: {}",
8099                        id,
8100                        tool.name(),
8101                        tool.description(),
8102                        args_desc
8103                    )
8104                })
8105            })
8106            .collect();
8107
8108        let tools_section = if tool_descriptions.is_empty() {
8109            "Available tools: none".to_string()
8110        } else {
8111            format!("Available tools:\n{}", tool_descriptions.join("\n"))
8112        };
8113
8114        let skills_section = if available_skills.is_empty() {
8115            "Available skills: none".to_string()
8116        } else {
8117            format!("Available skills: {}", available_skills.join(", "))
8118        };
8119
8120        let prompt = format!(
8121            r#"Create a step-by-step plan to accomplish this goal.
8122
8123Goal: "{}"
8124
8125{}
8126
8127{}
8128
8129Create a plan with clear steps. For each step, specify:
8130- description: What this step accomplishes
8131- action_type: "tool", "skill", "think", or "respond"
8132- action_target: The tool/skill id (if applicable)
8133- args: The arguments object matching the tool's schema (if action_type is "tool")
8134- dependencies: List of step IDs this depends on (empty if none)
8135
8136Respond in JSON format:
8137{{
8138  "steps": [
8139    {{"id": "step1", "description": "...", "action_type": "tool", "action_target": "tool_id", "args": {{"required_field": "value"}}, "dependencies": []}},
8140    {{"id": "step2", "description": "...", "action_type": "think", "action_target": "...", "dependencies": ["step1"]}}
8141  ]
8142}}"#,
8143            input, tools_section, skills_section,
8144        );
8145
8146        let messages = vec![ChatMessage::user(&prompt)];
8147        let response = self
8148            .observe_purpose(
8149                ObservationPurpose::PlanGeneration,
8150                planner_llm.complete(&messages, None),
8151            )
8152            .await
8153            .map_err(|e| AgentError::LLM(format!("Planning failed: {}", e)))?;
8154
8155        let mut plan = Plan::new(input);
8156
8157        if let Some(json_start) = response.content.find('{')
8158            && let Some(json_end) = response.content.rfind('}')
8159        {
8160            let json_str = &response.content[json_start..=json_end];
8161            if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(json_str)
8162                && let Some(steps) = parsed.get("steps").and_then(|s| s.as_array())
8163            {
8164                for step_value in steps {
8165                    let id = step_value
8166                        .get("id")
8167                        .and_then(|v| v.as_str())
8168                        .unwrap_or("step");
8169                    let desc = step_value
8170                        .get("description")
8171                        .and_then(|v| v.as_str())
8172                        .unwrap_or("");
8173                    let action_type = step_value
8174                        .get("action_type")
8175                        .and_then(|v| v.as_str())
8176                        .unwrap_or("think");
8177                    let action_target = step_value
8178                        .get("action_target")
8179                        .and_then(|v| v.as_str())
8180                        .unwrap_or("");
8181                    let args = step_value
8182                        .get("args")
8183                        .cloned()
8184                        .unwrap_or(serde_json::json!({}));
8185                    let deps: Vec<String> = step_value
8186                        .get("dependencies")
8187                        .and_then(|v| v.as_array())
8188                        .map(|arr| {
8189                            arr.iter()
8190                                .filter_map(|v| v.as_str().map(String::from))
8191                                .collect()
8192                        })
8193                        .unwrap_or_default();
8194
8195                    let action = match action_type {
8196                        "tool" => PlanAction::tool(action_target, args),
8197                        "skill" => PlanAction::skill(action_target),
8198                        "respond" => PlanAction::respond(action_target),
8199                        _ => PlanAction::think(desc),
8200                    };
8201
8202                    let step = PlanStep::new(desc, action)
8203                        .with_id(id)
8204                        .with_dependencies(deps);
8205                    plan.add_step(step);
8206                }
8207            }
8208        }
8209
8210        if plan.steps.is_empty() {
8211            plan.add_step(PlanStep::new(
8212                "Process the request",
8213                PlanAction::think(input),
8214            ));
8215            plan.add_step(PlanStep::new(
8216                "Provide response",
8217                PlanAction::respond("Answer based on analysis"),
8218            ));
8219        }
8220
8221        Ok(plan)
8222    }
8223
8224    async fn execute_plan(&self, plan: &mut Plan) -> Result<String> {
8225        let llm = self.get_state_llm()?;
8226        let mut results: HashMap<String, serde_json::Value> = HashMap::new();
8227        let effective = self.get_effective_reasoning_config();
8228        let max_steps = effective.get_planning().map(|c| c.max_steps).unwrap_or(10);
8229
8230        plan.status = PlanStatus::InProgress;
8231
8232        for step_idx in 0..plan.steps.len().min(max_steps as usize) {
8233            let step = &plan.steps[step_idx];
8234
8235            let deps_satisfied = step.dependencies.iter().all(|dep| {
8236                plan.steps
8237                    .iter()
8238                    .find(|s| &s.id == dep)
8239                    .map(|s| s.status.is_completed())
8240                    .unwrap_or(false)
8241            });
8242
8243            if !deps_satisfied {
8244                continue;
8245            }
8246
8247            plan.steps[step_idx].mark_running();
8248
8249            let result = match &plan.steps[step_idx].action {
8250                PlanAction::Tool { tool, args } => {
8251                    // When a tool step has dependency results, ask the LLM to
8252                    // produce the correct arguments given the context and tool schema.
8253                    // This avoids brittle {{stepN}} template substitution and lets
8254                    // the LLM handle type adaptation (e.g. picking the iso field
8255                    // from a datetime result for a downstream format call).
8256                    let has_dep_results = plan.steps[step_idx]
8257                        .dependencies
8258                        .iter()
8259                        .any(|dep| results.contains_key(dep));
8260
8261                    let final_args = if has_dep_results {
8262                        let dep_context: String = plan.steps[step_idx]
8263                            .dependencies
8264                            .iter()
8265                            .filter_map(|dep| results.get(dep).map(|r| format!("{}: {}", dep, r)))
8266                            .collect::<Vec<_>>()
8267                            .join("\n");
8268
8269                        let tool_schema = self
8270                            .tools
8271                            .get(tool)
8272                            .map(|t| {
8273                                let schema = t.input_schema();
8274                                let props = schema
8275                                    .get("properties")
8276                                    .and_then(|p| serde_json::to_string(p).ok())
8277                                    .unwrap_or_else(|| "{}".to_string());
8278                                format!(
8279                                    "{}: {}\nArguments schema: {}",
8280                                    t.id(),
8281                                    t.description(),
8282                                    props
8283                                )
8284                            })
8285                            .unwrap_or_default();
8286
8287                        let step_desc = &plan.steps[step_idx].description;
8288                        let arg_prompt = format!(
8289                            "Generate the JSON arguments for a tool call.\n\n\
8290                             Tool: {}\n\n\
8291                             Task: {}\n\n\
8292                             Previous step results:\n{}\n\n\
8293                             Planner's draft arguments: {}\n\n\
8294                             Produce ONLY a valid JSON object with the correct argument values.\n\
8295                             Use actual values from the previous step results, not template references.",
8296                            tool_schema,
8297                            step_desc,
8298                            dep_context,
8299                            serde_json::to_string(args).unwrap_or_default()
8300                        );
8301                        let messages = vec![ChatMessage::user(&arg_prompt)];
8302                        match self
8303                            .observe_purpose(
8304                                ObservationPurpose::PlanStep,
8305                                llm.complete(&messages, None),
8306                            )
8307                            .await
8308                        {
8309                            Ok(resp) => {
8310                                let content = resp.content.trim();
8311                                // Parse the LLM's JSON response, fall back to planner args.
8312                                let json_start = content.find('{');
8313                                let json_end = content.rfind('}');
8314                                if let (Some(start), Some(end)) = (json_start, json_end) {
8315                                    serde_json::from_str(&content[start..=end])
8316                                        .unwrap_or_else(|_| args.clone())
8317                                } else {
8318                                    args.clone()
8319                                }
8320                            }
8321                            Err(_) => args.clone(),
8322                        }
8323                    } else {
8324                        args.clone()
8325                    };
8326
8327                    let request = ToolExecutionRequest::new(
8328                        uuid::Uuid::new_v4().to_string(),
8329                        tool.clone(),
8330                        final_args,
8331                        ToolCallSource::Plan {
8332                            step_index: step_idx,
8333                        },
8334                    );
8335                    match self.execute_tool_record(request).await {
8336                        Ok(record) if record.success => {
8337                            serde_json::json!({ "output": record.model_output_string() })
8338                        }
8339                        Ok(record) => {
8340                            plan.steps[step_idx].mark_failed(record.model_output_string());
8341                            continue;
8342                        }
8343                        Err(e) => {
8344                            plan.steps[step_idx].mark_failed(e.to_string());
8345                            continue;
8346                        }
8347                    }
8348                }
8349                PlanAction::Skill { skill } => {
8350                    if let Some(skill_def) = self.skills.iter().find(|s| &s.id == skill) {
8351                        if let Some(ref executor) = self.skill_executor {
8352                            match executor
8353                                .execute_with_invoker(skill_def, "", serde_json::json!({}), self)
8354                                .await
8355                            {
8356                                Ok(output) => serde_json::json!({ "output": output }),
8357                                Err(e) => {
8358                                    plan.steps[step_idx].mark_failed(e.to_string());
8359                                    continue;
8360                                }
8361                            }
8362                        } else {
8363                            serde_json::json!({ "output": "Skill executor not available" })
8364                        }
8365                    } else {
8366                        plan.steps[step_idx].mark_failed("Skill not found");
8367                        continue;
8368                    }
8369                }
8370                PlanAction::Think { prompt } => {
8371                    let context: String = results
8372                        .iter()
8373                        .map(|(k, v)| format!("{}: {}", k, v))
8374                        .collect::<Vec<_>>()
8375                        .join("\n");
8376
8377                    let think_prompt = format!("Context:\n{}\n\nTask: {}", context, prompt);
8378                    let messages = vec![ChatMessage::user(&think_prompt)];
8379
8380                    match self
8381                        .observe_purpose(
8382                            ObservationPurpose::PlanStep,
8383                            llm.complete(&messages, None),
8384                        )
8385                        .await
8386                    {
8387                        Ok(resp) => serde_json::json!({ "output": resp.content }),
8388                        Err(e) => {
8389                            plan.steps[step_idx].mark_failed(e.to_string());
8390                            continue;
8391                        }
8392                    }
8393                }
8394                PlanAction::Respond { template } => {
8395                    let context: String = results
8396                        .iter()
8397                        .map(|(k, v)| format!("{}: {}", k, v))
8398                        .collect::<Vec<_>>()
8399                        .join("\n");
8400
8401                    let respond_prompt = format!(
8402                        "Based on this context:\n{}\n\nGenerate a response following this template/instruction: {}",
8403                        context, template
8404                    );
8405                    let messages = vec![ChatMessage::user(&respond_prompt)];
8406
8407                    match self
8408                        .observe_purpose(
8409                            ObservationPurpose::PlanStep,
8410                            llm.complete(&messages, None),
8411                        )
8412                        .await
8413                    {
8414                        Ok(resp) => serde_json::json!({ "output": resp.content }),
8415                        Err(e) => {
8416                            plan.steps[step_idx].mark_failed(e.to_string());
8417                            continue;
8418                        }
8419                    }
8420                }
8421            };
8422
8423            results.insert(plan.steps[step_idx].id.clone(), result.clone());
8424            plan.steps[step_idx].mark_completed(Some(result));
8425        }
8426
8427        // Set plan status based on whether any steps actually failed.
8428        let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
8429        if has_failures {
8430            let failed_ids: Vec<String> = plan
8431                .steps
8432                .iter()
8433                .filter(|s| s.status.is_failed())
8434                .map(|s| s.id.clone())
8435                .collect();
8436            plan.status = PlanStatus::Failed {
8437                error: format!("Steps failed: {}", failed_ids.join(", ")),
8438            };
8439        } else {
8440            plan.status = PlanStatus::Completed;
8441        }
8442
8443        // Synthesize final output from all completed step results.
8444        let all_outputs: Vec<String> = plan
8445            .steps
8446            .iter()
8447            .filter(|s| s.status.is_completed())
8448            .filter_map(|s| {
8449                s.result
8450                    .as_ref()
8451                    .and_then(|r| r.get("output"))
8452                    .and_then(|o| o.as_str())
8453                    .map(|o| format!("{}: {}", s.description, o))
8454            })
8455            .collect();
8456
8457        if all_outputs.is_empty() {
8458            return Ok("Plan execution completed but produced no results.".to_string());
8459        }
8460
8461        if all_outputs.len() == 1 {
8462            return Ok(all_outputs.into_iter().next().unwrap());
8463        }
8464
8465        // Synthesize a coherent summary from multiple step results via LLM.
8466        let context = all_outputs.join("\n\n");
8467        let prompt = format!(
8468            "You completed a multi-step plan for: \"{}\"\n\nStep results:\n{}\n\nProvide a coherent final response that synthesizes these results.",
8469            plan.goal, context
8470        );
8471        let messages = vec![ChatMessage::user(&prompt)];
8472        match self
8473            .observe_purpose(ObservationPurpose::PlanStep, llm.complete(&messages, None))
8474            .await
8475        {
8476            Ok(resp) => Ok(resp.content.trim().to_string()),
8477            Err(_) => Ok(context),
8478        }
8479    }
8480
8481    async fn evaluate_response(&self, input: &str, response: &str) -> Result<EvaluationResult> {
8482        let effective_config = self.get_effective_reflection_config();
8483        self.evaluate_response_with_config(input, response, &effective_config)
8484            .await
8485    }
8486
8487    fn extract_thinking(&self, content: &str) -> (Option<String>, String) {
8488        if let Some(start) = content.find("<thinking>")
8489            && let Some(end) = content.find("</thinking>")
8490        {
8491            let thinking = content[start + 10..end].trim().to_string();
8492            let answer = content[end + 11..].trim().to_string();
8493            return (Some(thinking), answer);
8494        }
8495        (None, content.to_string())
8496    }
8497
8498    fn format_response_with_thinking(&self, thinking: Option<&str>, answer: &str) -> String {
8499        match self.get_effective_reasoning_config().output {
8500            ReasoningOutput::Hidden => answer.to_string(),
8501            ReasoningOutput::Visible => {
8502                if let Some(t) = thinking {
8503                    format!("Thinking:\n{}\n\nAnswer:\n{}", t, answer)
8504                } else {
8505                    answer.to_string()
8506                }
8507            }
8508            ReasoningOutput::Tagged => {
8509                if let Some(t) = thinking {
8510                    format!("<thinking>{}</thinking>\n{}", t, answer)
8511                } else {
8512                    answer.to_string()
8513                }
8514            }
8515        }
8516    }
8517
8518    //
8519    // Blocking disambiguation returns every clarification or required confirmation as its own root response before redispatch.
8520    // The manager retains resolved input and pending skill ownership until a later turn explicitly confirms it.
8521    //
8522    async fn run_loop(&self, input: &str) -> Result<AgentResponse> {
8523        //
8524        // Blocking execution must fail before a turn starts when required persistence is unavailable.
8525        //
8526        self.init_storage().await?;
8527        self.begin_root_turn();
8528        let _root_cleanup = RootTurnCleanup::new(self);
8529        info!(input_len = input.len(), "Starting chat");
8530
8531        self.hooks.on_message_received(input).await;
8532
8533        // One-shot context initialization: load runtime defaults, resolve env vars,
8534        // populate builtin sources (session, agent), etc.  This must happen before
8535        // the first template render so that {{ context.* }} variables are available.
8536        if !self.context_initialized.swap(true, Ordering::SeqCst) {
8537            self.context_manager.initialize().await?;
8538            debug!("Context manager initialized (defaults, env, builtins)");
8539        }
8540
8541        self.check_turn_timeout().await?;
8542        self.context_manager.refresh_per_turn().await?;
8543
8544        // Clear stale disambiguation context from previous turns.
8545        // This prevents resolved_intent from leaking across turns and causing incorrect deterministic routing on subsequent inputs.
8546        self.clear_disambiguation_context();
8547
8548        // Disambiguation check (before input processing)
8549        if let Some(ref disambiguator) = self.disambiguation_manager {
8550            let disambiguation_context = self.build_disambiguation_context().await?;
8551
8552            // Get state-level disambiguation override
8553            let state_override = self
8554                .state_machine
8555                .as_ref()
8556                .and_then(|sm| sm.current_definition())
8557                .and_then(|def| def.disambiguation.clone());
8558
8559            let state_generation = self
8560                .state_machine
8561                .as_ref()
8562                .map(|state_machine| state_machine.generation());
8563            let disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
8564            let mut disambiguation_result = self
8565                .observe_purpose(
8566                    ObservationPurpose::DisambiguationDetection,
8567                    disambiguator.process_input_with_override(
8568                        input,
8569                        &disambiguation_context,
8570                        state_override.as_ref(),
8571                        None,
8572                    ),
8573                )
8574                .await?;
8575            let current_state_generation = self
8576                .state_machine
8577                .as_ref()
8578                .map(|state_machine| state_machine.generation());
8579            if current_state_generation != state_generation
8580                || self.disambiguation_epoch.load(Ordering::SeqCst) != disambiguation_epoch
8581            {
8582                disambiguator.clear_pending().await;
8583                *self.pending_skill_id.write() = None;
8584                disambiguation_result = DisambiguationResult::Abandoned { new_input: None };
8585                info!(
8586                    confirmation_event = "invalidated",
8587                    invalidation_reason = "state_generation_changed",
8588                    "Disambiguation result invalidated before redispatch"
8589                );
8590            }
8591            match disambiguation_result {
8592                DisambiguationResult::Clear => {
8593                    debug!("Input is clear, proceeding normally");
8594                }
8595                DisambiguationResult::NeedsClarification {
8596                    question,
8597                    detection,
8598                } => {
8599                    let admission = self
8600                        .admit_disambiguation_redispatch(disambiguation_epoch, state_generation)
8601                        .await?;
8602                    let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
8603                    info!(
8604                        ambiguity_type = ?detection.ambiguity_type,
8605                        confidence = detection.confidence,
8606                        "Input requires clarification"
8607                    );
8608
8609                    // This branch also owns post-resolution confirmation questions.
8610                    // Do not clear pending_skill_id or redispatch until the manager returns Clarified on a later turn.
8611                    self.commit_root_user_message(input).await?;
8612                    self.memory
8613                        .add_message(ChatMessage::assistant(&question.question))
8614                        .await?;
8615
8616                    let status = if awaiting_confirmation {
8617                        "awaiting_confirmation"
8618                    } else {
8619                        "awaiting_clarification"
8620                    };
8621                    let response = AgentResponse::new(&question.question).with_metadata(
8622                        "disambiguation",
8623                        serde_json::json!({
8624                            "status": status,
8625                            "options": question.options,
8626                            "clarifying": question.clarifying,
8627                            "detection": {
8628                                "type": detection.ambiguity_type,
8629                                "confidence": detection.confidence,
8630                                "what_is_unclear": detection.what_is_unclear,
8631                            }
8632                        }),
8633                    );
8634                    drop(admission);
8635                    self.finish_turn_if_root(&response).await?;
8636                    return Ok(response);
8637                }
8638                DisambiguationResult::Clarified {
8639                    enriched_input,
8640                    resolved,
8641                    ..
8642                } => {
8643                    let admission = match self
8644                        .admit_disambiguation_redispatch(disambiguation_epoch, state_generation)
8645                        .await
8646                    {
8647                        Ok(admission) => admission,
8648                        Err(error) => {
8649                            *self.pending_skill_id.write() = None;
8650                            return Err(error);
8651                        }
8652                    };
8653                    info!(
8654                        resolved_count = resolved.len(),
8655                        enriched = %enriched_input,
8656                        "Input clarified, injecting resolved intent into context"
8657                    );
8658
8659                    // Routing uses `resolved` (structured, deterministic)
8660                    // This is what makes post-disambiguation routing DETERMINISTIC
8661                    for (key, value) in &resolved {
8662                        let context_key = format!("disambiguation.{}", key);
8663                        let _ = self.context_manager.set(&context_key, value.clone());
8664                    }
8665
8666                    if let Some(intent) = resolved.get("intent") {
8667                        let _ = self.context_manager.set("resolved_intent", intent.clone());
8668                    }
8669
8670                    let _ = self
8671                        .context_manager
8672                        .set("disambiguation.resolved", serde_json::Value::Bool(true));
8673
8674                    // Check if this clarification was triggered by a skill-level override.
8675                    // If so, route directly to the matched skill instead of going through
8676                    // skill routing again (which might match a different skill).
8677                    let skill_id = self.pending_skill_id.read().clone();
8678                    if let Some(skill_id) = skill_id {
8679                        info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input");
8680                        drop(admission);
8681                        return self
8682                            .recheck_skill_disambiguation(
8683                                &skill_id,
8684                                &enriched_input,
8685                                disambiguation_epoch,
8686                                state_generation,
8687                            )
8688                            .await;
8689                    }
8690
8691                    drop(admission);
8692                    return self.run_loop_internal(&enriched_input).await;
8693                }
8694                DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
8695                    info!("Proceeding with best guess interpretation");
8696
8697                    // Same skill-id re-check for best-guess path
8698                    let skill_id = self.pending_skill_id.read().clone();
8699                    if let Some(skill_id) = skill_id {
8700                        info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input");
8701                        return self
8702                            .recheck_skill_disambiguation(
8703                                &skill_id,
8704                                &enriched_input,
8705                                disambiguation_epoch,
8706                                state_generation,
8707                            )
8708                            .await;
8709                    }
8710
8711                    return self.run_loop_internal(&enriched_input).await;
8712                }
8713                DisambiguationResult::GiveUp { reason } => {
8714                    *self.pending_skill_id.write() = None;
8715                    warn!(reason = %reason, "Disambiguation gave up");
8716                    let apology = self
8717                        .generate_localized_apology(
8718                            "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
8719                            &reason,
8720                        )
8721                        .await
8722                        .unwrap_or_else(|_| {
8723                            format!("I'm sorry, I couldn't understand your request: {}", reason)
8724                        });
8725                    let response = AgentResponse::new(&apology);
8726                    self.finish_turn_if_root(&response).await?;
8727                    return Ok(response);
8728                }
8729                DisambiguationResult::Escalate { reason } => {
8730                    *self.pending_skill_id.write() = None;
8731                    info!(reason = %reason, "Escalating to human");
8732                    if let Some(ref hitl) = self.hitl_engine {
8733                        let trigger =
8734                            ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
8735                        let mut context_map = HashMap::new();
8736                        context_map.insert("original_input".to_string(), serde_json::json!(input));
8737                        context_map.insert("reason".to_string(), serde_json::json!(&reason));
8738                        let check_result = HITLCheckResult::required(
8739                            trigger,
8740                            context_map,
8741                            format!("User request needs human assistance: {}", reason),
8742                            Some(hitl.config().default_timeout_seconds),
8743                        );
8744                        let result = self.request_hitl_approval(check_result).await?;
8745                        if matches!(
8746                            result,
8747                            ApprovalResult::Approved | ApprovalResult::Modified { .. }
8748                        ) {
8749                            return self.run_loop_internal(input).await;
8750                        }
8751                    }
8752                    let apology = self
8753                        .generate_localized_apology(
8754                            "Explain briefly that you're transferring the user to a human agent for help.",
8755                            &reason,
8756                        )
8757                        .await
8758                        .unwrap_or_else(|_| {
8759                            format!("I need human assistance to help with your request: {}", reason)
8760                        });
8761                    let response = AgentResponse::new(&apology);
8762                    self.finish_turn_if_root(&response).await?;
8763                    return Ok(response);
8764                }
8765                DisambiguationResult::Abandoned { new_input } => {
8766                    *self.pending_skill_id.write() = None;
8767
8768                    info!(
8769                        has_new_input = new_input.is_some(),
8770                        "Clarification abandoned by user"
8771                    );
8772
8773                    self.commit_root_user_message(input).await?;
8774
8775                    match new_input {
8776                        Some(fresh_input) => {
8777                            // Topic switch: process the user's new input from scratch.
8778                            // The LLM sees full conversation context including the abandoned exchange.
8779                            return self.run_loop_internal(&fresh_input).await;
8780                        }
8781                        None => {
8782                            // Pure abandonment: generate a brief acknowledgment.
8783                            let ack = self
8784                                .generate_localized_apology(
8785                                    "The user changed their mind about their previous request. \
8786                                     Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
8787                                     Do NOT apologize excessively. Be concise.",
8788                                    "User abandoned clarification",
8789                                )
8790                                .await
8791                                .unwrap_or_else(|_| {
8792                                    "OK, no problem. What else can I help with?".to_string()
8793                                });
8794
8795                            self.memory
8796                                .add_message(ChatMessage::assistant(&ack))
8797                                .await?;
8798
8799                            let response = AgentResponse::new(&ack);
8800                            self.finish_turn_if_root(&response).await?;
8801                            return Ok(response);
8802                        }
8803                    }
8804                }
8805            }
8806        }
8807
8808        self.run_loop_internal(input).await
8809    }
8810
8811    /// Generate a localized response using the router LLM
8812    async fn generate_localized_apology(&self, instruction: &str, reason: &str) -> Result<String> {
8813        let llm = self.llm_registry.router().map_err(|e| {
8814            AgentError::LLM(format!(
8815                "Router LLM not available for localized response: {}",
8816                e
8817            ))
8818        })?;
8819
8820        let recent: Vec<String> = self
8821            .memory
8822            .get_messages(Some(3))
8823            .await?
8824            .iter()
8825            .map(|m| m.content.clone())
8826            .collect();
8827
8828        let context_hint = if recent.is_empty() {
8829            String::new()
8830        } else {
8831            format!(
8832                "\nRecent conversation (detect the user's language from this):\n{}\n",
8833                recent.join("\n")
8834            )
8835        };
8836
8837        let prompt = format!(
8838            "{}\nReason: {}\n{}Respond in the same language as the user. Output ONLY the message, nothing else.",
8839            instruction, reason, context_hint
8840        );
8841
8842        let messages = vec![ChatMessage::user(&prompt)];
8843        let response = self
8844            .observe_purpose(
8845                ObservationPurpose::DisambiguationClarification,
8846                llm.complete(&messages, None),
8847            )
8848            .await
8849            .map_err(|e| AgentError::LLM(format!("Localized response generation failed: {}", e)))?;
8850
8851        Ok(response.content.trim().to_string())
8852    }
8853
8854    /// Clear disambiguation-related keys from the context manager.
8855    ///
8856    /// Render template variables in state action args using the context manager.
8857    fn render_action_args(&self, args: &Value) -> Value {
8858        let context = self.build_context_with_overlays();
8859        match args {
8860            Value::Object(map) => {
8861                let mut rendered = serde_json::Map::new();
8862                for (k, v) in map {
8863                    match v {
8864                        Value::String(s) if s.contains("{{") => {
8865                            match self.template_renderer.render(s, &context) {
8866                                Ok(rendered_str) => {
8867                                    rendered.insert(k.clone(), Value::String(rendered_str));
8868                                }
8869                                Err(_) => {
8870                                    rendered.insert(k.clone(), v.clone());
8871                                }
8872                            }
8873                        }
8874                        _ => {
8875                            rendered.insert(k.clone(), v.clone());
8876                        }
8877                    }
8878                }
8879                Value::Object(rendered)
8880            }
8881            _ => args.clone(),
8882        }
8883    }
8884
8885    /// Called at the start of each turn to prevent stale `resolved_intent` from leaking across turns.
8886    fn clear_disambiguation_context(&self) {
8887        let _ = self
8888            .context_manager
8889            .set("resolved_intent", serde_json::Value::Null);
8890
8891        let all = self.context_manager.get_all();
8892        for key in all.keys() {
8893            if key.starts_with("disambiguation.") {
8894                let _ = self.context_manager.set(key, serde_json::Value::Null);
8895            }
8896        }
8897    }
8898
8899    /// Re-run skill disambiguation on enriched input before executing the skill.
8900    /// After clarification resolves, the enriched input may still be missing required_clarity fields (e.g. "Transfer money to Jane." still lacks amount).
8901    /// This method re-runs the skill's disambiguation pass.
8902    /// If fields are still missing, it returns the new clarification question and keeps pending_skill_id set.
8903    /// If all fields are present (Clear), it executes the skill and returns the response.
8904    async fn recheck_skill_disambiguation(
8905        &self,
8906        skill_id: &str,
8907        enriched_input: &str,
8908        expected_disambiguation_epoch: u64,
8909        expected_state_generation: Option<u64>,
8910    ) -> Result<AgentResponse> {
8911        let skill = self
8912            .skill_router
8913            .as_ref()
8914            .and_then(|r| r.get_skill(skill_id).cloned());
8915
8916        // If the skill has disambiguation enabled, re-run it on the enriched input.
8917        if let Some(ref skill) = skill
8918            && let Some(ref skill_disambig) = skill.disambiguation
8919            && skill_disambig.enabled.unwrap_or(false)
8920            && let Some(ref disambiguator) = self.disambiguation_manager
8921        {
8922            let context = self.build_disambiguation_context().await?;
8923            let state_override = self
8924                .state_machine
8925                .as_ref()
8926                .and_then(|sm| sm.current_definition())
8927                .and_then(|def| def.disambiguation.clone());
8928
8929            let disambiguation_result = self
8930                .observe_purpose(
8931                    ObservationPurpose::DisambiguationDetection,
8932                    disambiguator.process_input_with_override(
8933                        enriched_input,
8934                        &context,
8935                        state_override.as_ref(),
8936                        Some(skill_disambig),
8937                    ),
8938                )
8939                .await?;
8940            let current_state_generation = self
8941                .state_machine
8942                .as_ref()
8943                .map(|state_machine| state_machine.generation());
8944            if current_state_generation != expected_state_generation
8945                || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
8946            {
8947                disambiguator.clear_pending().await;
8948                *self.pending_skill_id.write() = None;
8949                return Err(AgentError::Other(
8950                    "State or reset ownership changed during skill disambiguation recheck"
8951                        .to_string(),
8952                ));
8953            }
8954            match disambiguation_result {
8955                DisambiguationResult::Clear => {
8956                    debug!(skill_id = %skill_id, "Skill re-check: all fields present");
8957                }
8958                DisambiguationResult::NeedsClarification {
8959                    question,
8960                    detection,
8961                } => {
8962                    let admission = self
8963                        .admit_disambiguation_redispatch(
8964                            expected_disambiguation_epoch,
8965                            expected_state_generation,
8966                        )
8967                        .await?;
8968                    let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
8969                    info!(
8970                        skill_id = %skill_id,
8971                        ambiguity_type = ?detection.ambiguity_type,
8972                        what_is_unclear = ?detection.what_is_unclear,
8973                        "Skill re-check: still missing fields, asking again"
8974                    );
8975                    // Keep pending_skill_id set (do NOT clear it).
8976                    // The next turn will resolve this new clarification and
8977                    // re-enter this method until all fields are present.
8978                    self.memory
8979                        .add_message(ChatMessage::user(enriched_input))
8980                        .await?;
8981                    self.memory
8982                        .add_message(ChatMessage::assistant(&question.question))
8983                        .await?;
8984
8985                    let response = AgentResponse::new(&question.question).with_metadata(
8986                        "disambiguation",
8987                        serde_json::json!({
8988                            "status": if awaiting_confirmation { "awaiting_confirmation" } else { "awaiting_clarification" },
8989                            "skill_id": skill_id,
8990                            "options": question.options,
8991                            "clarifying": question.clarifying,
8992                            "detection": {
8993                                "type": detection.ambiguity_type,
8994                                "confidence": detection.confidence,
8995                                "what_is_unclear": detection.what_is_unclear,
8996                            }
8997                        }),
8998                    );
8999                    drop(admission);
9000                    self.finish_turn_if_root(&response).await?;
9001                    return Ok(response);
9002                }
9003                DisambiguationResult::Clarified {
9004                    enriched_input: re_enriched,
9005                    ..
9006                } => {
9007                    debug!(skill_id = %skill_id, "Skill re-check: clarified immediately, executing");
9008                    let admission = self
9009                        .admit_disambiguation_redispatch(
9010                            expected_disambiguation_epoch,
9011                            expected_state_generation,
9012                        )
9013                        .await?;
9014                    *self.pending_skill_id.write() = None;
9015                    drop(admission);
9016                    let skill_response = self.execute_skill_by_id(skill_id, &re_enriched).await?;
9017                    self.memory
9018                        .add_message(ChatMessage::user(&re_enriched))
9019                        .await?;
9020                    return self
9021                        .handle_skill_response(
9022                            &re_enriched,
9023                            skill_id,
9024                            skill_response,
9025                            &HashMap::new(),
9026                        )
9027                        .await;
9028                }
9029                DisambiguationResult::ProceedWithBestGuess {
9030                    enriched_input: re_enriched,
9031                } => {
9032                    debug!(skill_id = %skill_id, "Skill re-check: proceeding with best guess");
9033                    let admission = self
9034                        .admit_disambiguation_redispatch(
9035                            expected_disambiguation_epoch,
9036                            expected_state_generation,
9037                        )
9038                        .await?;
9039                    *self.pending_skill_id.write() = None;
9040                    drop(admission);
9041                    let skill_response = self.execute_skill_by_id(skill_id, &re_enriched).await?;
9042                    self.memory
9043                        .add_message(ChatMessage::user(&re_enriched))
9044                        .await?;
9045                    return self
9046                        .handle_skill_response(
9047                            &re_enriched,
9048                            skill_id,
9049                            skill_response,
9050                            &HashMap::new(),
9051                        )
9052                        .await;
9053                }
9054                DisambiguationResult::GiveUp { reason } => {
9055                    *self.pending_skill_id.write() = None;
9056                    let apology = self
9057                                    .generate_localized_apology(
9058                                        "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
9059                                        &reason,
9060                                    )
9061                                    .await
9062                                    .unwrap_or_else(|_| {
9063                                        format!("I'm sorry, I couldn't understand your request: {}", reason)
9064                                    });
9065                    let response = AgentResponse::new(&apology);
9066                    self.finish_turn_if_root(&response).await?;
9067                    return Ok(response);
9068                }
9069                DisambiguationResult::Escalate { reason } => {
9070                    *self.pending_skill_id.write() = None;
9071                    let apology = self
9072                                    .generate_localized_apology(
9073                                        "Explain briefly that you're transferring the user to a human agent for help.",
9074                                        &reason,
9075                                    )
9076                                    .await
9077                                    .unwrap_or_else(|_| {
9078                                        format!("I need human assistance to help with your request: {}", reason)
9079                                    });
9080                    let response = AgentResponse::new(&apology);
9081                    self.finish_turn_if_root(&response).await?;
9082                    return Ok(response);
9083                }
9084                DisambiguationResult::Abandoned { new_input } => {
9085                    // User abandoned during skill re-check.
9086                    // Clear skill routing state and fall through to normal execution.
9087                    *self.pending_skill_id.write() = None;
9088                    debug!(skill_id = %skill_id, "Skill re-check: abandoned by user");
9089                    if let Some(fresh) = new_input {
9090                        return self.run_loop_internal(&fresh).await;
9091                    }
9092                    let ack = self
9093                                    .generate_localized_apology(
9094                                        "The user changed their mind about their previous request. \
9095                                         Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
9096                                         Do NOT apologize excessively. Be concise.",
9097                                        "User abandoned clarification",
9098                                    )
9099                                    .await
9100                                    .unwrap_or_else(|_| {
9101                                        "OK, no problem. What else can I help with?".to_string()
9102                                    });
9103                    self.memory
9104                        .add_message(ChatMessage::assistant(&ack))
9105                        .await?;
9106                    let response = AgentResponse::new(&ack);
9107                    self.finish_turn_if_root(&response).await?;
9108                    return Ok(response);
9109                }
9110            }
9111        }
9112
9113        // Skill execution is admitted before the read guard is released so later invalidation cannot retroactively cancel it.
9114        let admission = self
9115            .admit_disambiguation_redispatch(
9116                expected_disambiguation_epoch,
9117                expected_state_generation,
9118            )
9119            .await?;
9120        *self.pending_skill_id.write() = None;
9121        drop(admission);
9122        let skill_response = self.execute_skill_by_id(skill_id, enriched_input).await?;
9123        self.memory
9124            .add_message(ChatMessage::user(enriched_input))
9125            .await?;
9126        self.handle_skill_response(enriched_input, skill_id, skill_response, &HashMap::new())
9127            .await
9128    }
9129
9130    /// Handle skill routing result: output processing, memory, transitions.
9131    /// Returns a fully formed AgentResponse for skill-routed requests.
9132    async fn handle_skill_response(
9133        &self,
9134        processed_input: &str,
9135        skill_id: &str,
9136        skill_response: String,
9137        input_context: &HashMap<String, Value>,
9138    ) -> Result<AgentResponse> {
9139        let output_data = self.process_output(&skill_response, input_context).await?;
9140        let final_response = output_data.content;
9141
9142        self.memory
9143            .add_message(ChatMessage::assistant(&final_response))
9144            .await?;
9145
9146        self.check_memory_compression().await?;
9147
9148        self.increment_turn();
9149        self.evaluate_transitions(processed_input, &final_response)
9150            .await?;
9151
9152        let response = AgentResponse::new(final_response)
9153            .with_metadata("skill_id", serde_json::json!(skill_id));
9154        self.finish_turn_if_root(&response).await?;
9155        Ok(response)
9156    }
9157
9158    /// Run the Plan-and-Execute flow: generate plan, execute steps, finalize.
9159    /// Supports plan-level reflection with replan loop when configured.
9160    async fn handle_plan_and_execute(
9161        &self,
9162        processed_input: &str,
9163        input_context: &HashMap<String, Value>,
9164        auto_detected: bool,
9165    ) -> Result<AgentResponse> {
9166        let effective = self.get_effective_reasoning_config();
9167        let plan_reflection = effective
9168            .get_planning()
9169            .map(|c| c.reflection.clone())
9170            .unwrap_or_default();
9171
9172        let max_attempts = if plan_reflection.enabled {
9173            1 + plan_reflection.max_replans
9174        } else {
9175            1
9176        };
9177
9178        let mut plan = self.generate_plan(processed_input).await?;
9179        info!(
9180            plan_id = %plan.id,
9181            steps = plan.steps.len(),
9182            "Plan generated"
9183        );
9184
9185        let mut plan_result = String::new();
9186
9187        for attempt in 0..max_attempts {
9188            *self.current_plan.write() = Some(plan.clone());
9189            plan_result = self.execute_plan(&mut plan).await?;
9190
9191            info!(
9192                plan_status = ?plan.status,
9193                completed_steps = plan.completed_steps().count(),
9194                attempt = attempt + 1,
9195                "Plan execution completed"
9196            );
9197
9198            if !plan_reflection.enabled {
9199                break;
9200            }
9201
9202            let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
9203            if !has_failures {
9204                break;
9205            }
9206
9207            if attempt + 1 >= max_attempts {
9208                break;
9209            }
9210
9211            match plan_reflection.on_step_failure {
9212                StepFailureAction::Replan => {
9213                    info!(attempt = attempt + 1, "Plan had failures, replanning");
9214                    plan = self.generate_plan(processed_input).await?;
9215                }
9216                StepFailureAction::Abort => {
9217                    warn!("Plan step failed, aborting");
9218                    break;
9219                }
9220                StepFailureAction::Skip | StepFailureAction::Continue => {
9221                    break;
9222                }
9223            }
9224        }
9225
9226        *self.current_plan.write() = Some(plan);
9227
9228        let output_data = self.process_output(&plan_result, input_context).await?;
9229        let final_content = output_data.content;
9230
9231        self.memory
9232            .add_message(ChatMessage::assistant(&final_content))
9233            .await?;
9234
9235        self.check_memory_compression().await?;
9236        self.increment_turn();
9237        self.evaluate_transitions(processed_input, &final_content)
9238            .await?;
9239
9240        let reasoning_metadata =
9241            ReasoningMetadata::new(ReasoningMode::PlanAndExecute).with_auto_detected(auto_detected);
9242
9243        let response = AgentResponse::new(&final_content).with_metadata(
9244            "reasoning",
9245            serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
9246        );
9247
9248        self.finish_turn_if_root(&response).await?;
9249        Ok(response)
9250    }
9251
9252    /// Inject CoT/ReAct reasoning prompt into the system message (first iteration only).
9253    fn inject_reasoning_prompt(
9254        &self,
9255        messages: &mut [ChatMessage],
9256        reasoning_mode: &ReasoningMode,
9257        is_first_iteration: bool,
9258    ) {
9259        if !is_first_iteration {
9260            return;
9261        }
9262        match reasoning_mode {
9263            ReasoningMode::CoT => {
9264                if let Some(msg) = messages.first_mut()
9265                    && matches!(msg.role, ai_agents_core::Role::System)
9266                {
9267                    msg.content = self.build_cot_system_prompt(&msg.content);
9268                    debug!("Applied Chain-of-Thought system prompt");
9269                }
9270            }
9271            ReasoningMode::React => {
9272                if let Some(msg) = messages.first_mut()
9273                    && matches!(msg.role, ai_agents_core::Role::System)
9274                {
9275                    msg.content = self.build_react_system_prompt(&msg.content);
9276                    debug!("Applied ReAct system prompt");
9277                }
9278            }
9279            _ => {}
9280        }
9281    }
9282
9283    //
9284    // Draft generation must not commit user memory or run tools.
9285    // The current user input is added only as an ephemeral message for this LLM call.
9286    //
9287    async fn generate_main_response_draft(
9288        &self,
9289        processed_input: &str,
9290        reasoning_mode: &ReasoningMode,
9291    ) -> Result<MainResponseDraft> {
9292        let llm = self.get_state_llm()?;
9293        let protocol = self.main_tool_protocol(llm.as_ref(), true).await?;
9294        let mut messages = self
9295            .build_messages_internal(false, Some(processed_input), protocol.choice.is_none())
9296            .await?;
9297        self.inject_reasoning_prompt(&mut messages, reasoning_mode, true);
9298        let response = self
9299            .complete_main_llm_with_recovery(llm, &messages, &protocol)
9300            .await?;
9301        let content = response.content.trim().to_string();
9302        let (thinking, answer) = self.extract_thinking(&content);
9303        if let Some(calls) = self.parse_main_tool_calls(&content, &protocol) {
9304            return Ok(MainResponseDraft::ToolCalls {
9305                raw_content: content,
9306                calls,
9307                thinking,
9308            });
9309        }
9310        Ok(MainResponseDraft::Text {
9311            raw_content: answer,
9312            thinking,
9313        })
9314    }
9315
9316    //
9317    // This is the only place where a winning draft is allowed to become runtime state.
9318    // Prompt and native tool calls remain inert until this method commits the draft into the shared executor path.
9319    //
9320    async fn commit_main_response_draft(
9321        &self,
9322        processed_input: &str,
9323        input_context: &HashMap<String, Value>,
9324        draft: MainResponseDraft,
9325        reasoning_mode: ReasoningMode,
9326        auto_detected: bool,
9327    ) -> Result<AgentResponse> {
9328        self.commit_root_user_message(processed_input).await?;
9329        match draft {
9330            MainResponseDraft::Text {
9331                raw_content,
9332                thinking,
9333            } => {
9334                self.finish_text_response_from_model(CommittedTextResponse {
9335                    processed_input,
9336                    input_context,
9337                    answer: raw_content,
9338                    reasoning_mode,
9339                    auto_detected,
9340                    iterations: 1,
9341                    thinking_content: thinking,
9342                    all_tool_calls: Vec::new(),
9343                })
9344                .await
9345            }
9346            MainResponseDraft::ToolCalls {
9347                raw_content,
9348                calls,
9349                thinking: _,
9350            } => {
9351                let mut all_tool_calls = Vec::new();
9352                match self
9353                    .handle_tool_calls(processed_input, &raw_content, calls, &mut all_tool_calls)
9354                    .await?
9355                {
9356                    ToolCallOutcome::Rejected(response) => {
9357                        self.finish_turn_if_root(&response).await?;
9358                        Ok(response)
9359                    }
9360                    ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => {
9361                        self.continue_after_committed_tool_draft(processed_input)
9362                            .await
9363                    }
9364                }
9365            }
9366        }
9367    }
9368
9369    //
9370    // Tool drafts need a committed continuation after function results are written.
9371    // Redispatch depth suppresses duplicate root lifecycle work during that continuation.
9372    //
9373    async fn continue_after_committed_tool_draft(
9374        &self,
9375        processed_input: &str,
9376    ) -> Result<AgentResponse> {
9377        *self.redispatch_depth.write() += 1;
9378        if let Some(context) = self.active_turn_context.write().as_mut() {
9379            context.enter_redispatch();
9380        }
9381        let result = Box::pin(self.run_loop_internal(processed_input)).await;
9382        *self.redispatch_depth.write() -= 1;
9383        if let Some(context) = self.active_turn_context.write().as_mut() {
9384            context.exit_redispatch();
9385        }
9386        let response = result?;
9387        self.finish_turn_if_root(&response).await?;
9388        Ok(response)
9389    }
9390
9391    //
9392    // Shared committed text finalization for normal responses and winning text drafts.
9393    // Keep output processing, reflection, transitions, hooks, and maintenance behind this commit boundary.
9394    //
9395    async fn finish_text_response_from_model(
9396        &self,
9397        response: CommittedTextResponse<'_>,
9398    ) -> Result<AgentResponse> {
9399        let CommittedTextResponse {
9400            processed_input,
9401            input_context,
9402            answer,
9403            reasoning_mode,
9404            auto_detected,
9405            iterations,
9406            thinking_content,
9407            all_tool_calls,
9408        } = response;
9409        let output_data = self.process_output(&answer, input_context).await?;
9410        let mut final_content = if output_data.metadata.rejected {
9411            output_data
9412                .metadata
9413                .rejection_reason
9414                .unwrap_or_else(|| answer.to_string())
9415        } else {
9416            output_data.content
9417        };
9418        let llm = self.get_state_llm()?;
9419        let reflection_metadata;
9420        (final_content, reflection_metadata) = self
9421            .run_reflection(&*llm, processed_input, final_content)
9422            .await?;
9423        final_content =
9424            self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
9425        let final_content = {
9426            let result = self
9427                .post_loop_processing(processed_input, final_content)
9428                .await?;
9429            self.apply_post_loop_result(processed_input, result).await?
9430        };
9431        let response = self.build_agent_response(AgentResponseParts {
9432            content: final_content,
9433            all_tool_calls,
9434            reasoning_mode,
9435            auto_detected,
9436            iterations,
9437            thinking: thinking_content,
9438            reflection_metadata,
9439        });
9440        self.finish_turn_if_root(&response).await?;
9441        Ok(response)
9442    }
9443
9444    //
9445    // Auto reasoning uses this path after the judge wins with a deeper mode.
9446    // It intentionally uses committed message building instead of the draft overlay.
9447    //
9448    async fn run_committed_response_loop_with_reasoning(
9449        &self,
9450        processed_input: &str,
9451        input_context: &HashMap<String, Value>,
9452        reasoning_mode: ReasoningMode,
9453        auto_detected: bool,
9454    ) -> Result<AgentResponse> {
9455        self.commit_root_user_message(processed_input).await?;
9456        let llm = self.get_state_llm()?;
9457        let mut iterations = 0u32;
9458        let mut all_tool_calls = Vec::new();
9459        let mut thinking_content = None;
9460        loop {
9461            let effective_max = if reasoning_mode != ReasoningMode::None {
9462                let rc = self.get_effective_reasoning_config();
9463                self.max_iterations.min(rc.max_iterations)
9464            } else {
9465                self.max_iterations
9466            };
9467            if iterations >= effective_max {
9468                return Err(AgentError::Other(format!(
9469                    "Max iterations ({}) exceeded",
9470                    effective_max
9471                )));
9472            }
9473            iterations += 1;
9474            *self.iteration_count.write() = iterations;
9475            let protocol = self.main_tool_protocol(llm.as_ref(), false).await?;
9476            let mut messages = self
9477                .build_messages_internal(true, None, protocol.choice.is_none())
9478                .await?;
9479            self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
9480            self.hooks.on_llm_start(&messages).await;
9481            let llm_start = Instant::now();
9482            let response = self
9483                .complete_main_llm_with_recovery(Arc::clone(&llm), &messages, &protocol)
9484                .await?;
9485            let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
9486            self.hooks.on_llm_complete(&response, llm_duration_ms).await;
9487            let content = response.content.trim();
9488            if let Some(tool_calls) = self.parse_main_tool_calls(content, &protocol) {
9489                match self
9490                    .handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
9491                    .await?
9492                {
9493                    ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
9494                    ToolCallOutcome::Rejected(resp) => {
9495                        self.finish_turn_if_root(&resp).await?;
9496                        return Ok(resp);
9497                    }
9498                }
9499            }
9500            let (extracted_thinking, answer) = self.extract_thinking(content);
9501            if extracted_thinking.is_some() {
9502                thinking_content = extracted_thinking;
9503            }
9504            return self
9505                .finish_text_response_from_model(CommittedTextResponse {
9506                    processed_input,
9507                    input_context,
9508                    answer,
9509                    reasoning_mode,
9510                    auto_detected,
9511                    iterations,
9512                    thinking_content,
9513                    all_tool_calls,
9514                })
9515                .await;
9516        }
9517    }
9518
9519    /// Handle tool calls: check transitions, execute tools in parallel, handle HITL rejection.
9520    async fn handle_tool_calls(
9521        &self,
9522        processed_input: &str,
9523        content: &str,
9524        tool_calls: Vec<ToolCall>,
9525        all_tool_calls: &mut Vec<ToolCall>,
9526    ) -> Result<ToolCallOutcome> {
9527        // Check if a transition should fire before executing the LLM's tool call.
9528        // If a transition fires, on_enter actions handle the tool call correctly
9529        // (with proper URLs from YAML), so skip the LLM's tool call.
9530        let transition_fired = self.evaluate_transitions(processed_input, content).await?;
9531        if transition_fired {
9532            self.memory
9533                .add_message(ChatMessage::assistant(
9534                    "(Transitioned to new state — tool call handled by workflow)",
9535                ))
9536                .await?;
9537            return Ok(ToolCallOutcome::TransitionFired);
9538        }
9539
9540        // 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.
9541        self.memory
9542            .add_message(ChatMessage::assistant(content))
9543            .await?;
9544        let native_tool_call = Self::is_native_tool_call_content(content);
9545
9546        let results = self.execute_tools_parallel(&tool_calls).await;
9547
9548        for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
9549            match result {
9550                Ok(output) => {
9551                    self.memory
9552                        .add_message(Self::tool_result_message(
9553                            tool_call,
9554                            &output,
9555                            native_tool_call,
9556                        ))
9557                        .await?;
9558                }
9559                Err(e) => {
9560                    // Check if this is a HITL rejection - if so, break the loop
9561                    if matches!(e, AgentError::HITLRejected(_)) {
9562                        self.memory
9563                            .add_message(ChatMessage::assistant(format!(
9564                                "The operation was rejected by the approver: {}",
9565                                e
9566                            )))
9567                            .await?;
9568                        // Return the rejection message to user, don't continue loop
9569                        return Ok(ToolCallOutcome::Rejected(AgentResponse {
9570                            content: format!("Operation cancelled: {}", e),
9571                            metadata: None,
9572                            tool_calls: Some(all_tool_calls.clone()),
9573                        }));
9574                    }
9575                    self.memory
9576                        .add_message(Self::tool_result_message(
9577                            tool_call,
9578                            &format!("Error: {}", e),
9579                            native_tool_call,
9580                        ))
9581                        .await?;
9582                }
9583            }
9584            all_tool_calls.push(tool_call.clone());
9585        }
9586        Ok(ToolCallOutcome::Continue)
9587    }
9588
9589    /// Run the reflection loop on a response, returning (improved_content, reflection_metadata).
9590    async fn run_reflection(
9591        &self,
9592        llm: &dyn LLMProvider,
9593        processed_input: &str,
9594        mut content: String,
9595    ) -> Result<(String, Option<ReflectionMetadata>)> {
9596        let should_reflect = self.should_reflect(processed_input, &content).await?;
9597        if !should_reflect {
9598            return Ok((content, None));
9599        }
9600
9601        info!("Starting response reflection evaluation");
9602        let mut attempts = 0u32;
9603        let max_retries = self.reflection_config.max_retries;
9604        let mut history: Vec<ReflectionAttempt> = Vec::new();
9605
9606        loop {
9607            let evaluation = self.evaluate_response(processed_input, &content).await?;
9608
9609            if evaluation.passed || attempts >= max_retries {
9610                info!(
9611                    passed = evaluation.passed,
9612                    confidence = evaluation.confidence,
9613                    attempts = attempts + 1,
9614                    "Reflection evaluation complete"
9615                );
9616                let reflection_metadata = Some(
9617                    ReflectionMetadata::new(evaluation)
9618                        .with_attempts(attempts + 1)
9619                        .with_history(history),
9620                );
9621                return Ok((content, reflection_metadata));
9622            }
9623
9624            debug!(
9625                attempt = attempts + 1,
9626                failed_criteria = evaluation.failed_criteria().count(),
9627                "Response did not meet criteria, retrying"
9628            );
9629
9630            history.push(
9631                ReflectionAttempt::new(&content, evaluation.clone())
9632                    .with_feedback("Response did not meet quality criteria"),
9633            );
9634
9635            let feedback: Vec<String> = evaluation
9636                .failed_criteria()
9637                .map(|c| format!("- {}", c.criterion))
9638                .collect();
9639
9640            let retry_prompt = format!(
9641                "Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response.",
9642                feedback.join("\n")
9643            );
9644
9645            self.memory
9646                .add_message(ChatMessage::user(&retry_prompt))
9647                .await?;
9648
9649            let retry_messages = self.build_messages().await?;
9650            let retry_response = self
9651                .observe_purpose(
9652                    ObservationPurpose::ReflectionEvaluation,
9653                    llm.complete(&retry_messages, None),
9654                )
9655                .await
9656                .map_err(|e| AgentError::LLM(e.to_string()))?;
9657
9658            content = retry_response.content.trim().to_string();
9659            attempts += 1;
9660        }
9661    }
9662
9663    /// Record the assistant turn, evaluate transitions, and decide what to do next.
9664    /// Returns PostLoopResult so callers can apply_post_loop_result for re-dispatch.
9665    async fn post_loop_processing(
9666        &self,
9667        processed_input: &str,
9668        content: String,
9669    ) -> Result<PostLoopResult> {
9670        // Do NOT add the assistant message to memory yet.
9671        // evaluate_transitions receives content as a direct parameter, so the message does not need to be in memory for transitions to evaluate correctly.
9672        // For NeedsRedispatch we skip adding the stale old-state response entirely, keeping memory clean for the re-dispatched handler.
9673
9674        self.increment_turn();
9675
9676        // Run context extractors so guards can check freshly-extracted values.
9677        self.run_context_extractors(processed_input).await;
9678
9679        let transitioned = self.evaluate_transitions(processed_input, &content).await?;
9680
9681        if !transitioned {
9682            self.memory
9683                .add_message(ChatMessage::assistant(&content))
9684                .await?;
9685            self.check_memory_compression().await?;
9686            return Ok(PostLoopResult::NoTransition(content));
9687        }
9688
9689        // Check if we should skip re-generation after this transition.
9690        if !self.should_regenerate_after_transition() {
9691            self.memory
9692                .add_message(ChatMessage::assistant(&content))
9693                .await?;
9694            self.check_memory_compression().await?;
9695            return Ok(PostLoopResult::Transitioned(content));
9696        }
9697
9698        // Check if the new state needs full dispatch.
9699        // Orchestration states (concurrent, group_chat, pipeline, handoff, delegate) need their dedicated handlers.
9700        // 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.
9701        if self.needs_redispatch_for_new_state() {
9702            info!("Post-transition NeedsRedispatch: new state requires full dispatch");
9703            // Stale old-state content is NOT added to memory.
9704            // apply_post_loop_result will increment redispatch_depth and call run_loop_internal, which produces the correct response and adds it.
9705            return Ok(PostLoopResult::NeedsRedispatch);
9706        }
9707
9708        // Normal post-transition re-generation (plain LLM with optional tool calls).
9709        // Add the stale response to history so the new LLM sees the conversation.
9710        self.memory
9711            .add_message(ChatMessage::assistant(&content))
9712            .await?;
9713        self.check_memory_compression().await?;
9714
9715        // If a transition fired, on_enter actions already executed (e.g., HTTP calls).
9716        // The current content was generated in the OLD state context and is stale.
9717        // Re-generate in the new state context so the LLM can reference on_enter results.
9718        // If the LLM responds with a tool call, execute it in a mini-loop so the
9719        // result is not returned as raw JSON text.
9720        let new_llm = self.get_state_llm()?;
9721        let mut final_content;
9722
9723        for post_iter in 0..self.max_iterations {
9724            let protocol = self.main_tool_protocol(new_llm.as_ref(), false).await?;
9725            let new_messages = self
9726                .build_messages_internal(true, None, protocol.choice.is_none())
9727                .await?;
9728            if post_iter == 0
9729                && let Some(system_msg) = new_messages.first()
9730                && system_msg.role == ai_agents_core::Role::System
9731            {
9732                debug!(
9733                    prompt_preview =
9734                        &system_msg.content[system_msg.content.len().saturating_sub(200)..],
9735                    "Post-transition system prompt (last 200 chars)"
9736                );
9737            }
9738
9739            let new_response = self
9740                .complete_main_llm_with_recovery(Arc::clone(&new_llm), &new_messages, &protocol)
9741                .await?;
9742            final_content = new_response.content.trim().to_string();
9743
9744            // Check if the post-transition response contains tool calls.
9745            // If so, execute them and loop so the LLM can summarize the result.
9746            if let Some(tool_calls) = self.parse_main_tool_calls(&final_content, &protocol) {
9747                let native_tool_call = Self::is_native_tool_call_content(&final_content);
9748                debug!(
9749                    post_iter = post_iter,
9750                    tools = tool_calls.len(),
9751                    "Post-transition tool call detected, executing"
9752                );
9753
9754                self.memory
9755                    .add_message(ChatMessage::assistant(&final_content))
9756                    .await?;
9757
9758                let results = self.execute_tools_parallel(&tool_calls).await;
9759                for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
9760                    match result {
9761                        Ok(output) => {
9762                            self.memory
9763                                .add_message(Self::tool_result_message(
9764                                    tool_call,
9765                                    &output,
9766                                    native_tool_call,
9767                                ))
9768                                .await?;
9769                        }
9770                        Err(e) => {
9771                            self.memory
9772                                .add_message(Self::tool_result_message(
9773                                    tool_call,
9774                                    &format!("Error: {}", e),
9775                                    native_tool_call,
9776                                ))
9777                                .await?;
9778                        }
9779                    }
9780                }
9781                // Loop to let the LLM see the tool result and produce a text response.
9782                continue;
9783            }
9784
9785            // No tool call - this is the final text response.
9786            self.memory
9787                .add_message(ChatMessage::assistant(&final_content))
9788                .await?;
9789            return Ok(PostLoopResult::Transitioned(final_content));
9790        }
9791
9792        // Exhausted post-transition iterations (unlikely). Return last content.
9793        final_content = "Post-transition processing completed.".to_string();
9794        self.memory
9795            .add_message(ChatMessage::assistant(&final_content))
9796            .await?;
9797
9798        Ok(PostLoopResult::Transitioned(final_content))
9799    }
9800
9801    /// Build the final AgentResponse with all metadata.
9802    /// Check whether to re-generate a response after a state transition.
9803    fn should_regenerate_after_transition(&self) -> bool {
9804        if let Some(ref sm) = self.state_machine {
9805            // Global setting
9806            if !sm.config().regenerate_on_transition {
9807                return false;
9808            }
9809            // Per-state override on the new (current) state
9810            if let Some(def) = sm.current_definition()
9811                && let Some(regen) = def.regenerate_on_enter
9812            {
9813                return regen;
9814            }
9815        }
9816        true
9817    }
9818
9819    /// Return true when the new state requires full dispatch via run_loop_internal.
9820    /// Covers orchestration states and any non-None effective reasoning mode.
9821    fn needs_redispatch_for_new_state(&self) -> bool {
9822        if let Some(ref sm) = self.state_machine
9823            && let Some(def) = sm.current_definition()
9824        {
9825            if def.concurrent.is_some()
9826                || def.group_chat.is_some()
9827                || def.pipeline.is_some()
9828                || def.handoff.is_some()
9829                || def.delegate.is_some()
9830            {
9831                return true;
9832            }
9833            // Any non-None effective reasoning mode requires the main dispatch loop.
9834            let effective = self.get_effective_reasoning_config();
9835            if !matches!(effective.mode, ReasoningMode::None) {
9836                return true;
9837            }
9838        }
9839        false
9840    }
9841
9842    /// Consume a PostLoopResult. NeedsRedispatch re-enters run_loop_internal.
9843    /// The user message is already in memory - redispatch_depth suppresses re-adding it.
9844    async fn apply_post_loop_result(
9845        &self,
9846        processed_input: &str,
9847        result: PostLoopResult,
9848    ) -> Result<String> {
9849        match result {
9850            PostLoopResult::NoTransition(content) | PostLoopResult::Transitioned(content) => {
9851                Ok(content)
9852            }
9853            PostLoopResult::NeedsRedispatch => {
9854                const MAX_REDISPATCH_DEPTH: u32 = 3;
9855                let current_depth = *self.redispatch_depth.read();
9856                if current_depth >= MAX_REDISPATCH_DEPTH {
9857                    warn!(
9858                        depth = current_depth,
9859                        "Post-transition re-dispatch depth limit reached, returning empty response"
9860                    );
9861                    let content = String::new();
9862                    self.memory
9863                        .add_message(ChatMessage::assistant(&content))
9864                        .await?;
9865                    return Ok(content);
9866                }
9867                *self.redispatch_depth.write() += 1;
9868                if let Some(context) = self.active_turn_context.write().as_mut() {
9869                    context.enter_redispatch();
9870                }
9871                info!(
9872                    depth = current_depth + 1,
9873                    "Re-dispatching for new state after transition"
9874                );
9875                let resp = Box::pin(self.run_loop_internal(processed_input)).await;
9876                *self.redispatch_depth.write() -= 1;
9877                if let Some(context) = self.active_turn_context.write().as_mut() {
9878                    context.exit_redispatch();
9879                }
9880                resp.map(|r| r.content)
9881            }
9882        }
9883    }
9884
9885    /// Builds the final response metadata after committed output and transition handling complete.
9886    fn build_agent_response(&self, parts: AgentResponseParts) -> AgentResponse {
9887        let AgentResponseParts {
9888            content,
9889            all_tool_calls,
9890            reasoning_mode,
9891            auto_detected,
9892            iterations,
9893            thinking,
9894            reflection_metadata,
9895        } = parts;
9896        let reasoning_metadata = ReasoningMetadata::new(reasoning_mode.clone())
9897            .with_thinking(thinking.clone().unwrap_or_default())
9898            .with_iterations(iterations)
9899            .with_auto_detected(auto_detected);
9900
9901        let mut response = AgentResponse::new(&content);
9902        if !all_tool_calls.is_empty() {
9903            response = response.with_tool_calls(all_tool_calls);
9904        }
9905
9906        if let Some(state) = self.current_state() {
9907            response = response.with_metadata("current_state", serde_json::json!(state));
9908        }
9909
9910        response = response.with_metadata(
9911            "reasoning",
9912            serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
9913        );
9914
9915        if let Some(ref refl_meta) = reflection_metadata {
9916            response = response.with_metadata(
9917                "reflection",
9918                serde_json::to_value(refl_meta).unwrap_or_default(),
9919            );
9920        }
9921
9922        response
9923    }
9924
9925    // Handle delegation: forward user input to a registry agent.
9926    async fn handle_delegated_state(
9927        &self,
9928        input: &str,
9929        delegate_id: &str,
9930        state_def: &ai_agents_state::StateDefinition,
9931    ) -> Result<AgentResponse> {
9932        use std::time::Instant;
9933
9934        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
9935            AgentError::Config(format!(
9936                "State delegates to '{}' but no agent registry is configured. \
9937                 Add a spawner section with auto_spawn to your YAML.",
9938                delegate_id
9939            ))
9940        })?;
9941
9942        let state_name = self
9943            .state_machine
9944            .as_ref()
9945            .map(|sm| sm.current())
9946            .unwrap_or_else(|| "unknown".to_string());
9947
9948        self.hooks.on_delegate_start(delegate_id, &state_name).await;
9949        let start = Instant::now();
9950
9951        let delegate = registry.get(delegate_id).ok_or_else(|| {
9952            AgentError::Other(format!(
9953                "State '{}' delegates to '{}' but no agent with that ID exists in the registry.",
9954                state_name, delegate_id
9955            ))
9956        })?;
9957
9958        // Prepare input based on delegate_context mode.
9959        let context_mode = state_def.delegate_context.clone().unwrap_or_default();
9960        let effective_input = self
9961            .observe_purpose(
9962                ObservationPurpose::OrchestrationRouting,
9963                crate::orchestration::context::prepare_delegate_input(
9964                    input,
9965                    &context_mode,
9966                    &*self.memory,
9967                    self.llm_registry.get("router").ok().as_deref(),
9968                ),
9969            )
9970            .await?;
9971
9972        let response = delegate
9973            .chat_with_actor_context(&effective_input, self.outbound_actor_context())
9974            .await?;
9975
9976        let duration_ms = start.elapsed().as_millis() as u64;
9977        self.hooks
9978            .on_delegate_complete(delegate_id, &state_name, duration_ms)
9979            .await;
9980
9981        // Backward-compatible context key.
9982        let ctx_key = format!("delegation.{}.last_response", delegate_id);
9983        let _ = self.context_manager.set(
9984            &ctx_key,
9985            serde_json::Value::String(response.content.clone()),
9986        );
9987
9988        // Structured orchestration context.
9989        let _ = self.context_manager.set(
9990            "orchestration",
9991            serde_json::json!({
9992                "type": "delegate",
9993                "agent": delegate_id,
9994                "state": state_name,
9995                "response": response.content,
9996                "duration_ms": duration_ms,
9997            }),
9998        );
9999
10000        self.commit_root_user_message(input).await?;
10001
10002        // post_loop_processing records the assistant turn and evaluates transitions.
10003        // apply_post_loop_result handles NeedsRedispatch by re-entering run_loop_internal.
10004        let post_result = self
10005            .post_loop_processing(
10006                input,
10007                format!("[Delegated to {}]: {}", delegate_id, response.content),
10008            )
10009            .await?;
10010        let final_content = self.apply_post_loop_result(input, post_result).await?;
10011
10012        let mut result = AgentResponse::new(final_content);
10013
10014        let metadata = serde_json::json!({
10015            "orchestration": {
10016                "type": "delegate",
10017                "agent": delegate_id,
10018                "state": state_name,
10019                "response": response.content,
10020                "duration_ms": duration_ms,
10021            }
10022        });
10023        result.metadata = Some(
10024            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10025                metadata,
10026            )
10027            .unwrap_or_default(),
10028        );
10029
10030        self.finish_turn_if_root(&result).await?;
10031        Ok(result)
10032    }
10033
10034    // Handle concurrent execution: run multiple registry agents in parallel and aggregate.
10035    async fn handle_concurrent_state(
10036        &self,
10037        input: &str,
10038        config: &ai_agents_state::ConcurrentStateConfig,
10039    ) -> Result<AgentResponse> {
10040        use std::time::Instant;
10041
10042        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10043            AgentError::Config(
10044                "Concurrent state requires an agent registry. Add a spawner section.".into(),
10045            )
10046        })?;
10047
10048        // Render input template if provided, otherwise use the raw input.
10049        // Uses direct minijinja rendering (same approach as pipeline) so variables
10050        // are top-level: {{ user_input }}, not {{ context.user_input }}.
10051        // Enrich input with parent conversation history when context_mode is set.
10052        let context_mode = config.context_mode.clone().unwrap_or_default();
10053        let context_input = self
10054            .observe_purpose(
10055                ObservationPurpose::OrchestrationRouting,
10056                crate::orchestration::context::prepare_delegate_input(
10057                    input,
10058                    &context_mode,
10059                    &*self.memory,
10060                    self.llm_registry.get("router").ok().as_deref(),
10061                ),
10062            )
10063            .await?;
10064
10065        let effective_input = if let Some(ref tmpl) = config.input {
10066            render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10067                .unwrap_or_else(|_| context_input.clone())
10068        } else {
10069            context_input
10070        };
10071
10072        let start = Instant::now();
10073
10074        let llm_name = config
10075            .aggregation
10076            .synthesizer_llm
10077            .as_deref()
10078            .unwrap_or("router");
10079        let llm_provider = self.llm_registry.get(llm_name).ok();
10080
10081        let vote_parallelism = if self.runtime_config.optimization.enabled
10082            && self
10083                .runtime_config
10084                .optimization
10085                .parallel_orchestration_vote_extraction
10086        {
10087            Some(self.runtime_config.optimization.max_parallel_runtime_tasks)
10088        } else {
10089            None
10090        };
10091
10092        let result = self
10093            .observe_purpose(
10094                ObservationPurpose::OrchestrationAggregation,
10095                scope_actor_context(
10096                    self.outbound_actor_context(),
10097                    crate::orchestration::concurrent(
10098                        registry,
10099                        &effective_input,
10100                        &config.agents,
10101                        &config.aggregation,
10102                        llm_provider.as_deref(),
10103                        config.min_required,
10104                        config.timeout_ms,
10105                        config.on_partial_failure.clone(),
10106                        vote_parallelism,
10107                    ),
10108                ),
10109            )
10110            .await?;
10111
10112        let duration_ms = start.elapsed().as_millis() as u64;
10113        let agent_ids: Vec<String> = config.agents.iter().map(|a| a.id().to_string()).collect();
10114        let strategy = format!("{:?}", config.aggregation.strategy);
10115        self.hooks
10116            .on_concurrent_complete(&agent_ids, &strategy, duration_ms)
10117            .await;
10118
10119        // Backward-compatible context key.
10120        let _ = self.context_manager.set(
10121            "concurrent.result",
10122            serde_json::Value::String(result.response.content.clone()),
10123        );
10124
10125        // Build per-agent result data for context and metadata.
10126        let agents_json: Vec<serde_json::Value> = result
10127            .agent_results
10128            .iter()
10129            .map(|ar| {
10130                serde_json::json!({
10131                    "id": ar.agent_id,
10132                    "response": ar.response.as_ref().map(|r| r.content.as_str()),
10133                    "success": ar.success,
10134                    "error": ar.error,
10135                    "duration_ms": ar.duration_ms,
10136                })
10137            })
10138            .collect();
10139
10140        // Structured orchestration context with per-agent results.
10141        let _ = self.context_manager.set(
10142            "orchestration",
10143            serde_json::json!({
10144                "type": "concurrent",
10145                "result": result.response.content,
10146                "strategy": strategy,
10147                "agents": agents_json,
10148                "duration_ms": duration_ms,
10149            }),
10150        );
10151
10152        self.commit_root_user_message(input).await?;
10153
10154        let post_result = self
10155            .post_loop_processing(input, result.response.content.clone())
10156            .await?;
10157        let final_content = self.apply_post_loop_result(input, post_result).await?;
10158
10159        let mut response = AgentResponse::new(final_content);
10160        let metadata = serde_json::json!({
10161            "orchestration": {
10162                "type": "concurrent",
10163                "result": result.response.content,
10164                "strategy": strategy,
10165                "agents": agents_json,
10166                "duration_ms": duration_ms,
10167            }
10168        });
10169        response.metadata = Some(
10170            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10171                metadata,
10172            )
10173            .unwrap_or_default(),
10174        );
10175
10176        self.finish_turn_if_root(&response).await?;
10177        Ok(response)
10178    }
10179
10180    // Handle group chat: run a multi-turn multi-agent conversation.
10181    async fn handle_group_chat_state(
10182        &self,
10183        input: &str,
10184        config: &ai_agents_state::GroupChatStateConfig,
10185    ) -> Result<AgentResponse> {
10186        use std::time::Instant;
10187
10188        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10189            AgentError::Config(
10190                "Group chat state requires an agent registry. Add a spawner section.".into(),
10191            )
10192        })?;
10193
10194        let start = Instant::now();
10195
10196        let llm_provider = self.llm_registry.get("router").ok();
10197
10198        // Enrich input with parent conversation history when context_mode is set.
10199        let context_mode = config.context_mode.clone().unwrap_or_default();
10200        let context_input = self
10201            .observe_purpose(
10202                ObservationPurpose::OrchestrationRouting,
10203                crate::orchestration::context::prepare_delegate_input(
10204                    input,
10205                    &context_mode,
10206                    &*self.memory,
10207                    self.llm_registry.get("router").ok().as_deref(),
10208                ),
10209            )
10210            .await?;
10211
10212        // Render input template if provided, otherwise use the raw user message as topic.
10213        let effective_topic = if let Some(ref tmpl) = config.input {
10214            render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10215                .unwrap_or_else(|_| context_input.clone())
10216        } else {
10217            context_input
10218        };
10219
10220        let result = self
10221            .observe_purpose(
10222                ObservationPurpose::OrchestrationConversation,
10223                scope_actor_context(
10224                    self.outbound_actor_context(),
10225                    crate::orchestration::group_chat(
10226                        registry,
10227                        &effective_topic,
10228                        config,
10229                        llm_provider.as_deref(),
10230                        Some(&*self.hooks),
10231                    ),
10232                ),
10233            )
10234            .await?;
10235
10236        let duration_ms = start.elapsed().as_millis() as u64;
10237
10238        // Backward-compatible context key.
10239        let _ = self.context_manager.set(
10240            "group_chat.conclusion",
10241            serde_json::Value::String(result.response.content.clone()),
10242        );
10243
10244        // Build transcript data for context and metadata.
10245        let transcript_json: Vec<serde_json::Value> = result
10246            .transcript
10247            .iter()
10248            .map(|t| {
10249                serde_json::json!({
10250                    "speaker": t.speaker,
10251                    "round": t.round,
10252                    "content": t.content,
10253                })
10254            })
10255            .collect();
10256
10257        // Structured orchestration context with full transcript.
10258        let _ = self.context_manager.set(
10259            "orchestration",
10260            serde_json::json!({
10261                "type": "group_chat",
10262                "conclusion": result.response.content,
10263                "transcript": transcript_json,
10264                "rounds": result.rounds_completed,
10265                "termination": result.termination_reason,
10266                "duration_ms": duration_ms,
10267            }),
10268        );
10269
10270        self.commit_root_user_message(input).await?;
10271
10272        let post_result = self
10273            .post_loop_processing(input, result.response.content.clone())
10274            .await?;
10275        let final_content = self.apply_post_loop_result(input, post_result).await?;
10276
10277        let mut response = AgentResponse::new(final_content);
10278        let metadata = serde_json::json!({
10279            "orchestration": {
10280                "type": "group_chat",
10281                "conclusion": result.response.content,
10282                "transcript": transcript_json,
10283                "rounds": result.rounds_completed,
10284                "termination": result.termination_reason,
10285                "duration_ms": duration_ms,
10286            }
10287        });
10288        response.metadata = Some(
10289            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10290                metadata,
10291            )
10292            .unwrap_or_default(),
10293        );
10294
10295        self.finish_turn_if_root(&response).await?;
10296        Ok(response)
10297    }
10298
10299    // Handle pipeline: run agents sequentially with per-stage input templates.
10300    async fn handle_pipeline_state(
10301        &self,
10302        input: &str,
10303        config: &ai_agents_state::PipelineStateConfig,
10304    ) -> Result<AgentResponse> {
10305        use std::time::Instant;
10306
10307        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10308            AgentError::Config(
10309                "Pipeline state requires an agent registry. Add a spawner section.".into(),
10310            )
10311        })?;
10312
10313        let start = Instant::now();
10314
10315        let stages: Vec<crate::orchestration::PipelineStage> = config
10316            .stages
10317            .iter()
10318            .map(|entry| {
10319                let mut stage = crate::orchestration::PipelineStage::id(entry.id());
10320                if let Some(tmpl) = entry.input() {
10321                    stage = stage.with_input(tmpl);
10322                }
10323                stage
10324            })
10325            .collect();
10326
10327        // Enrich input with parent conversation history when context_mode is set.
10328        let context_mode = config.context_mode.clone().unwrap_or_default();
10329        let context_input = self
10330            .observe_purpose(
10331                ObservationPurpose::OrchestrationRouting,
10332                crate::orchestration::context::prepare_delegate_input(
10333                    input,
10334                    &context_mode,
10335                    &*self.memory,
10336                    self.llm_registry.get("router").ok().as_deref(),
10337                ),
10338            )
10339            .await?;
10340
10341        let context_values = self.build_context_with_overlays();
10342        let result = self
10343            .observe_purpose(
10344                ObservationPurpose::OrchestrationRouting,
10345                scope_actor_context(
10346                    self.outbound_actor_context(),
10347                    crate::orchestration::pipeline(
10348                        registry,
10349                        &context_input,
10350                        &stages,
10351                        config.timeout_ms,
10352                        Some(&*self.hooks),
10353                        Some(&context_values),
10354                    ),
10355                ),
10356            )
10357            .await?;
10358
10359        let duration_ms = start.elapsed().as_millis() as u64;
10360
10361        // Backward-compatible context key.
10362        let _ = self.context_manager.set(
10363            "pipeline.result",
10364            serde_json::Value::String(result.response.content.clone()),
10365        );
10366
10367        // Build per-stage data for context and metadata.
10368        let stages_json: Vec<serde_json::Value> = result
10369            .stage_outputs
10370            .iter()
10371            .map(|s| {
10372                serde_json::json!({
10373                    "agent_id": s.agent_id,
10374                    "output": s.output,
10375                    "duration_ms": s.duration_ms,
10376                    "skipped": s.skipped,
10377                })
10378            })
10379            .collect();
10380
10381        // Structured orchestration context.
10382        let _ = self.context_manager.set(
10383            "orchestration",
10384            serde_json::json!({
10385                "type": "pipeline",
10386                "result": result.response.content,
10387                "stages": stages_json,
10388                "duration_ms": duration_ms,
10389            }),
10390        );
10391
10392        self.commit_root_user_message(input).await?;
10393
10394        let post_result = self
10395            .post_loop_processing(input, result.response.content.clone())
10396            .await?;
10397        let final_content = self.apply_post_loop_result(input, post_result).await?;
10398
10399        let mut response = AgentResponse::new(final_content);
10400        let metadata = serde_json::json!({
10401            "orchestration": {
10402                "type": "pipeline",
10403                "result": result.response.content,
10404                "stages": stages_json,
10405                "duration_ms": duration_ms,
10406            }
10407        });
10408        response.metadata = Some(
10409            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10410                metadata,
10411            )
10412            .unwrap_or_default(),
10413        );
10414
10415        self.finish_turn_if_root(&response).await?;
10416        Ok(response)
10417    }
10418
10419    // Handle handoff: LLM-directed agent-to-agent control transfer.
10420    async fn handle_handoff_state(
10421        &self,
10422        input: &str,
10423        config: &ai_agents_state::HandoffStateConfig,
10424    ) -> Result<AgentResponse> {
10425        use std::time::Instant;
10426
10427        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10428            AgentError::Config(
10429                "Handoff state requires an agent registry. Add a spawner section.".into(),
10430            )
10431        })?;
10432
10433        let llm = self
10434            .llm_registry
10435            .get("router")
10436            .map_err(|_| AgentError::Config("Handoff state requires a router LLM.".into()))?;
10437
10438        let start = Instant::now();
10439
10440        // Enrich input with parent conversation history when context_mode is set.
10441        let context_mode = config.context_mode.clone().unwrap_or_default();
10442        let context_input = self
10443            .observe_purpose(
10444                ObservationPurpose::OrchestrationRouting,
10445                crate::orchestration::context::prepare_delegate_input(
10446                    input,
10447                    &context_mode,
10448                    &*self.memory,
10449                    self.llm_registry.get("router").ok().as_deref(),
10450                ),
10451            )
10452            .await?;
10453
10454        // Render input template if provided, otherwise forward the raw user message.
10455        let effective_input = if let Some(ref tmpl) = config.input {
10456            render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10457                .unwrap_or_else(|_| context_input.clone())
10458        } else {
10459            context_input
10460        };
10461
10462        let result = self
10463            .observe_purpose(
10464                ObservationPurpose::OrchestrationRouting,
10465                scope_actor_context(
10466                    self.outbound_actor_context(),
10467                    crate::orchestration::handoff(
10468                        registry,
10469                        &effective_input,
10470                        &config.initial_agent,
10471                        &config.available_agents,
10472                        config.max_handoffs,
10473                        llm.as_ref(),
10474                        Some(&*self.hooks),
10475                    ),
10476                ),
10477            )
10478            .await?;
10479
10480        let duration_ms = start.elapsed().as_millis() as u64;
10481
10482        // Backward-compatible context key.
10483        let _ = self.context_manager.set(
10484            "handoff.result",
10485            serde_json::Value::String(result.response.content.clone()),
10486        );
10487
10488        // Build handoff chain data for context and metadata.
10489        let chain_json: Vec<serde_json::Value> = result
10490            .handoff_chain
10491            .iter()
10492            .map(|h| {
10493                serde_json::json!({
10494                    "from": h.from_agent,
10495                    "to": h.to_agent,
10496                    "reason": h.reason,
10497                })
10498            })
10499            .collect();
10500
10501        // Structured orchestration context.
10502        let _ = self.context_manager.set(
10503            "orchestration",
10504            serde_json::json!({
10505                "type": "handoff",
10506                "result": result.response.content,
10507                "final_agent": result.final_agent,
10508                "handoff_chain": chain_json,
10509                "duration_ms": duration_ms,
10510            }),
10511        );
10512
10513        self.commit_root_user_message(input).await?;
10514
10515        let post_result = self
10516            .post_loop_processing(input, result.response.content.clone())
10517            .await?;
10518        let final_content = self.apply_post_loop_result(input, post_result).await?;
10519
10520        let mut response = AgentResponse::new(final_content);
10521        let metadata = serde_json::json!({
10522            "orchestration": {
10523                "type": "handoff",
10524                "result": result.response.content,
10525                "final_agent": result.final_agent,
10526                "handoff_chain": chain_json,
10527                "duration_ms": duration_ms,
10528            }
10529        });
10530        response.metadata = Some(
10531            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10532                metadata,
10533            )
10534            .unwrap_or_default(),
10535        );
10536
10537        self.finish_turn_if_root(&response).await?;
10538        Ok(response)
10539    }
10540
10541    // run_loop_internal: blocking (non-streaming) agent pipeline.
10542    async fn run_loop_internal(&self, input: &str) -> Result<AgentResponse> {
10543        self.begin_root_turn();
10544        // Resolve actor_id from context, reload facts if actor changed, bump counter.
10545        self.pre_turn_session_lifecycle().await;
10546
10547        let input_data = self.process_input(input).await?;
10548        self.update_active_turn_context(&input_data.content, input_data.context.clone());
10549
10550        // Inject process context (detect/extract results) into agent context
10551        // so system prompt templates can use {{ context.detected_language }} etc.
10552        for (key, value) in &input_data.context {
10553            let _ = self.context_manager.set(key, value.clone());
10554        }
10555
10556        if input_data.metadata.rejected {
10557            let reason = input_data
10558                .metadata
10559                .rejection_reason
10560                .unwrap_or_else(|| "Input rejected".to_string());
10561            warn!(reason = %reason, "Input rejected");
10562            let response = AgentResponse::new(reason);
10563            self.finish_turn_if_root(&response).await?;
10564            return Ok(response);
10565        }
10566
10567        let processed_input = &input_data.content;
10568
10569        if let Some(response) = self.try_pre_response_transition(processed_input).await? {
10570            return Ok(response);
10571        }
10572
10573        // Handle orchestration states (delegate, concurrent, group_chat, pipeline, handoff).
10574        if let Some(ref sm) = self.state_machine
10575            && let Some(def) = sm.current_definition()
10576        {
10577            if let Some(ref delegate_id) = def.delegate {
10578                return self
10579                    .handle_delegated_state(processed_input, delegate_id, &def)
10580                    .await;
10581            }
10582            if let Some(ref concurrent_config) = def.concurrent {
10583                return self
10584                    .handle_concurrent_state(processed_input, concurrent_config)
10585                    .await;
10586            }
10587            if let Some(ref group_chat_config) = def.group_chat {
10588                return self
10589                    .handle_group_chat_state(processed_input, group_chat_config)
10590                    .await;
10591            }
10592            if let Some(ref pipeline_config) = def.pipeline {
10593                return self
10594                    .handle_pipeline_state(processed_input, pipeline_config)
10595                    .await;
10596            }
10597            if let Some(ref handoff_config) = def.handoff {
10598                return self
10599                    .handle_handoff_state(processed_input, handoff_config)
10600                    .await;
10601            }
10602        }
10603
10604        //
10605        // The speculative future is boxed to keep the runtime future size manageable.
10606        // Removing the box can overflow small test stacks because this function is recursive through redispatch.
10607        //
10608        if let Some(response) =
10609            Box::pin(self.try_speculative_branches(processed_input, &input_data.context)).await?
10610        {
10611            return Ok(response);
10612        }
10613
10614        match self.try_skill_route(processed_input).await? {
10615            SkillRouteResult::Response { skill_id, content } => {
10616                self.commit_root_user_message(processed_input).await?;
10617                return self
10618                    .handle_skill_response(processed_input, &skill_id, content, &input_data.context)
10619                    .await;
10620            }
10621            SkillRouteResult::NeedsClarification {
10622                response,
10623                ownership,
10624            } => {
10625                let admission = self
10626                    .admit_optional_disambiguation_ownership(ownership)
10627                    .await?;
10628                self.commit_root_user_message(processed_input).await?;
10629                if let Some(q) = response
10630                    .metadata
10631                    .as_ref()
10632                    .and_then(|m| m.get("disambiguation"))
10633                    .and_then(|d| d.get("status"))
10634                    .and_then(|s| s.as_str())
10635                    && q == "awaiting_clarification"
10636                {
10637                    // Store the clarification question in memory so the next turn
10638                    // can be handled as a clarification response.
10639                    self.memory
10640                        .add_message(ChatMessage::assistant(&response.content))
10641                        .await?;
10642                }
10643                drop(admission);
10644                self.finish_turn_if_root(&response).await?;
10645                return Ok(response);
10646            }
10647            SkillRouteResult::NoMatch => {} // continue to normal LLM chat
10648        }
10649
10650        let effective_reasoning = self.get_effective_reasoning_config();
10651        let reasoning_mode = self.determine_reasoning_mode(processed_input).await?;
10652        let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
10653
10654        info!(
10655            reasoning_mode = ?reasoning_mode,
10656            auto_detected = auto_detected,
10657            reflection_enabled = ?self.reflection_config.enabled,
10658            "Reasoning mode determined"
10659        );
10660
10661        if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
10662            self.commit_root_user_message(processed_input).await?;
10663            return self
10664                .handle_plan_and_execute(processed_input, &input_data.context, auto_detected)
10665                .await;
10666        }
10667
10668        self.commit_root_user_message(processed_input).await?;
10669
10670        let mut iterations = 0u32;
10671        let mut all_tool_calls: Vec<ToolCall> = Vec::new();
10672        let mut thinking_content: Option<String> = None;
10673
10674        let llm = self.get_state_llm()?;
10675
10676        loop {
10677            // When reasoning is active, cap iterations at the reasoning-specific limit.
10678            let effective_max = if reasoning_mode != ReasoningMode::None {
10679                let rc = self.get_effective_reasoning_config();
10680                self.max_iterations.min(rc.max_iterations)
10681            } else {
10682                self.max_iterations
10683            };
10684
10685            if iterations >= effective_max {
10686                let err = AgentError::Other(format!("Max iterations ({}) exceeded", effective_max));
10687                self.hooks.on_error(&err).await;
10688                error!(iterations = iterations, "Max iterations exceeded");
10689                return Err(err);
10690            }
10691            iterations += 1;
10692            *self.iteration_count.write() = iterations;
10693
10694            debug!(iteration = iterations, max = effective_max, "LLM call");
10695
10696            let protocol = self.main_tool_protocol(llm.as_ref(), false).await?;
10697            let mut messages = self
10698                .build_messages_internal(true, None, protocol.choice.is_none())
10699                .await?;
10700            self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
10701
10702            self.hooks.on_llm_start(&messages).await;
10703            let llm_start = Instant::now();
10704            let response = self
10705                .complete_main_llm_with_recovery(Arc::clone(&llm), &messages, &protocol)
10706                .await?;
10707
10708            let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
10709            self.hooks.on_llm_complete(&response, llm_duration_ms).await;
10710
10711            let content = response.content.trim();
10712
10713            if let Some(tool_calls) = self.parse_main_tool_calls(content, &protocol) {
10714                match self
10715                    .handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
10716                    .await?
10717                {
10718                    ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
10719                    ToolCallOutcome::Rejected(resp) => {
10720                        self.finish_turn_if_root(&resp).await?;
10721                        return Ok(resp);
10722                    }
10723                }
10724            }
10725
10726            let (extracted_thinking, answer) = self.extract_thinking(content);
10727            if extracted_thinking.is_some() {
10728                thinking_content = extracted_thinking;
10729            }
10730
10731            let output_data = self.process_output(&answer, &input_data.context).await?;
10732
10733            let mut final_content = if output_data.metadata.rejected {
10734                output_data
10735                    .metadata
10736                    .rejection_reason
10737                    .unwrap_or_else(|| answer.to_string())
10738            } else {
10739                output_data.content
10740            };
10741
10742            // Run reflection (blocking LLM calls for retries)
10743            let reflection_metadata;
10744            (final_content, reflection_metadata) = self
10745                .run_reflection(&*llm, processed_input, final_content)
10746                .await?;
10747
10748            final_content =
10749                self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
10750
10751            // Post-loop: memory, transitions, post-transition re-generation.
10752            // apply_post_loop_result handles NeedsRedispatch by re-entering
10753            // run_loop_internal so the new state's full dispatch activates.
10754            let final_content = {
10755                let result = self
10756                    .post_loop_processing(processed_input, final_content)
10757                    .await?;
10758                self.apply_post_loop_result(processed_input, result).await?
10759            };
10760
10761            let reflected = reflection_metadata.is_some();
10762            let reasoning_mode_debug = format!("{:?}", reasoning_mode);
10763
10764            let response = self.build_agent_response(AgentResponseParts {
10765                content: final_content,
10766                all_tool_calls,
10767                reasoning_mode,
10768                auto_detected,
10769                iterations,
10770                thinking: thinking_content,
10771                reflection_metadata,
10772            });
10773
10774            self.finish_turn_if_root(&response).await?;
10775
10776            let tool_call_count = response.tool_calls.as_ref().map(|tc| tc.len()).unwrap_or(0);
10777            info!(
10778                tool_calls = tool_call_count,
10779                response_len = response.content.len(),
10780                reasoning_mode = %reasoning_mode_debug,
10781                reflected = reflected,
10782                "Chat completed"
10783            );
10784            return Ok(response);
10785        }
10786    }
10787
10788    async fn generate_buffered_streaming_draft(
10789        &self,
10790        processed_input: &str,
10791        routing_resolved: Arc<AtomicBool>,
10792    ) -> Result<StreamingDraftResult> {
10793        let llm = self.get_state_llm()?;
10794        if llm.configured_tool_choice().is_some() {
10795            let draft = self
10796                .generate_main_response_draft(processed_input, &ReasoningMode::None)
10797                .await?;
10798            return Ok(StreamingDraftResult::new(draft, Vec::new()));
10799        }
10800        let messages = self.build_messages_for_draft(processed_input).await?;
10801        let mut stream = self
10802            .observe_purpose(
10803                ObservationPurpose::MainResponse,
10804                llm.complete_stream(&messages, None),
10805            )
10806            .await
10807            .map_err(|e| AgentError::LLM(e.to_string()))?;
10808        let mut buffer = crate::optimization::StreamBranchBuffer::new(self.streaming.buffer_size)?;
10809        let mut chunks = Vec::new();
10810        let mut accumulated = String::new();
10811        while let Some(chunk_result) = stream.next().await {
10812            let chunk = chunk_result.map_err(|e| AgentError::LLM(e.to_string()))?;
10813            accumulated.push_str(&chunk.delta);
10814            let stream_chunk = StreamChunk::content(chunk.delta);
10815            if routing_resolved.load(Ordering::SeqCst) {
10816                chunks.push(stream_chunk);
10817            } else {
10818                buffer.push(stream_chunk)?;
10819            }
10820        }
10821        chunks.splice(0..0, buffer.drain());
10822        let content = accumulated.trim().to_string();
10823        let draft = if let Some(calls) = self.parse_tool_calls(&content) {
10824            MainResponseDraft::ToolCalls {
10825                raw_content: content,
10826                calls,
10827                thinking: None,
10828            }
10829        } else {
10830            MainResponseDraft::Text {
10831                raw_content: content,
10832                thinking: None,
10833            }
10834        };
10835        Ok(StreamingDraftResult::new(draft, chunks))
10836    }
10837
10838    async fn try_buffered_streaming_branches(
10839        &self,
10840        processed_input: &str,
10841        input_context: &HashMap<String, Value>,
10842    ) -> Result<Option<(AgentResponse, Vec<StreamChunk>)>> {
10843        let optimization = &self.runtime_config.optimization;
10844        if !optimization.enabled {
10845            return Ok(None);
10846        }
10847        let transition_enabled =
10848            optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
10849        if !transition_enabled {
10850            return Ok(None);
10851        }
10852        let mut branch_scheduler =
10853            TurnBranchScheduler::new(optimization.max_parallel_runtime_tasks)?;
10854        if !branch_scheduler.reserve_task() {
10855            return Ok(None);
10856        }
10857        if !self
10858            .reserve_active_speculative_llm_call(RuntimeOptimizationKind::BufferedStreamingRouting)
10859        {
10860            branch_scheduler.release_task();
10861            return Ok(None);
10862        }
10863        if !branch_scheduler.reserve_task() {
10864            branch_scheduler.release_task();
10865            return Ok(None);
10866        }
10867        let mut main_branch = RuntimeBranch::new(
10868            RuntimeTaskPurpose::MainResponse,
10869            RuntimeOptimizationKind::BufferedStreamingRouting,
10870            RuntimeTaskPriority::Normal,
10871            RuntimeCommitBehavior::FinalResponse,
10872        );
10873        let mut transition_branch = RuntimeBranch::new(
10874            RuntimeTaskPurpose::StateTransition,
10875            RuntimeOptimizationKind::ParallelStateTransition,
10876            RuntimeTaskPriority::Critical,
10877            RuntimeCommitBehavior::TransitionDecision,
10878        );
10879        let main_id = main_branch.branch_id();
10880        let transition_id = transition_branch.branch_id();
10881        let routing_resolved = Arc::new(AtomicBool::new(false));
10882        let mut main_future =
10883            Box::pin(crate::optimization::observability::with_branch_observation(
10884                &main_id,
10885                RuntimeOptimizationKind::BufferedStreamingRouting,
10886                RuntimeCommitBehavior::FinalResponse,
10887                self.generate_buffered_streaming_draft(
10888                    processed_input,
10889                    Arc::clone(&routing_resolved),
10890                ),
10891            ));
10892        let mut transition_future =
10893            Box::pin(crate::optimization::observability::with_branch_observation(
10894                &transition_id,
10895                RuntimeOptimizationKind::ParallelStateTransition,
10896                RuntimeCommitBehavior::TransitionDecision,
10897                self.select_parallel_transition_candidate(processed_input),
10898            ));
10899        let mut main_pending = true;
10900        let mut transition_pending = true;
10901        let mut main_result: Option<Result<StreamingDraftResult>> = None;
10902        let mut transition_finalized = false;
10903        let mut transition_candidate: Option<TransitionCandidate> = None;
10904        loop {
10905            if let Some(candidate) = transition_candidate.take() {
10906                if self
10907                    .approve_transition_target(&candidate.from_state, candidate.target())
10908                    .await?
10909                {
10910                    // Drop the stale stream future before transition side effects or redispatch reuse the provider.
10911                    drop(main_future);
10912                    drop(transition_future);
10913                    self.finalize_branch_loss(
10914                        &main_id,
10915                        RuntimeOptimizationKind::BufferedStreamingRouting,
10916                        RuntimeCommitBehavior::FinalResponse,
10917                        main_pending,
10918                        main_result.as_ref().map(|result| result.is_err()),
10919                    );
10920                    if !self
10921                        .apply_pre_response_transition_candidate(
10922                            &candidate,
10923                            &HashMap::new(),
10924                            processed_input,
10925                        )
10926                        .await?
10927                    {
10928                        self.finalize_optional_branch(
10929                            &transition_id,
10930                            RuntimeOptimizationKind::ParallelStateTransition,
10931                            RuntimeCommitBehavior::TransitionDecision,
10932                            "discarded",
10933                            false,
10934                        );
10935                        return Ok(None);
10936                    }
10937                    self.finalize_optional_branch(
10938                        &transition_id,
10939                        RuntimeOptimizationKind::ParallelStateTransition,
10940                        RuntimeCommitBehavior::TransitionDecision,
10941                        "committed",
10942                        true,
10943                    );
10944                    let response = self.redispatch_current_state(processed_input).await?;
10945                    return Ok(Some((
10946                        response.clone(),
10947                        vec![StreamChunk::content(response.content)],
10948                    )));
10949                }
10950                self.finalize_optional_branch(
10951                    &transition_id,
10952                    RuntimeOptimizationKind::ParallelStateTransition,
10953                    RuntimeCommitBehavior::TransitionDecision,
10954                    "discarded",
10955                    false,
10956                );
10957                routing_resolved.store(true, Ordering::SeqCst);
10958                transition_finalized = true;
10959            }
10960            if transition_finalized && let Some(result) = main_result.take() {
10961                let stream_draft = match result {
10962                    Ok(stream_draft) => stream_draft,
10963                    Err(error) => {
10964                        self.finalize_optional_branch(
10965                            &main_id,
10966                            RuntimeOptimizationKind::BufferedStreamingRouting,
10967                            RuntimeCommitBehavior::FinalResponse,
10968                            "failed",
10969                            false,
10970                        );
10971                        return Err(error);
10972                    }
10973                };
10974                let raw_draft_content = stream_draft.draft.raw_content().to_string();
10975                let buffered_chunks = stream_draft.chunks;
10976                self.finalize_optional_branch(
10977                    &main_id,
10978                    RuntimeOptimizationKind::BufferedStreamingRouting,
10979                    RuntimeCommitBehavior::FinalResponse,
10980                    "committed",
10981                    true,
10982                );
10983                let response = self
10984                    .commit_main_response_draft(
10985                        processed_input,
10986                        input_context,
10987                        stream_draft.draft,
10988                        ReasoningMode::None,
10989                        false,
10990                    )
10991                    .await?;
10992                let chunks = if response.content == raw_draft_content {
10993                    buffered_chunks
10994                } else {
10995                    vec![StreamChunk::content(response.content.clone())]
10996                };
10997                return Ok(Some((response, chunks)));
10998            }
10999            tokio::select! {
11000                result = &mut main_future, if main_pending => {
11001                    main_pending = false;
11002                    main_branch.transition_to(RuntimeBranchStatus::Completed)?;
11003                    main_result = Some(result);
11004                }
11005                result = &mut transition_future, if transition_pending => {
11006                    transition_pending = false;
11007                    transition_branch.transition_to(RuntimeBranchStatus::Completed)?;
11008                    match result {
11009                        Ok(ParallelTransitionSelection::Candidate(candidate)) => {
11010                            transition_candidate = Some(candidate)
11011                        }
11012                        Ok(ParallelTransitionSelection::NoMatch) => {
11013                            self.finalize_optional_branch(
11014                                &transition_id,
11015                                RuntimeOptimizationKind::ParallelStateTransition,
11016                                RuntimeCommitBehavior::TransitionDecision,
11017                                "discarded",
11018                                false,
11019                            );
11020                            routing_resolved.store(true, Ordering::SeqCst);
11021                            transition_finalized = true;
11022                        }
11023                        Ok(ParallelTransitionSelection::ReservationExhausted) => {
11024                            self.finalize_optional_branch(
11025                                &transition_id,
11026                                RuntimeOptimizationKind::ParallelStateTransition,
11027                                RuntimeCommitBehavior::TransitionDecision,
11028                                "cancelled",
11029                                false,
11030                            );
11031                            routing_resolved.store(true, Ordering::SeqCst);
11032                            self.finalize_branch_loss(
11033                                &main_id,
11034                                RuntimeOptimizationKind::BufferedStreamingRouting,
11035                                RuntimeCommitBehavior::FinalResponse,
11036                                main_pending,
11037                                main_result.as_ref().map(|result| result.is_err()),
11038                            );
11039                            return Ok(None);
11040                        }
11041                        Err(_) => {
11042                            self.finalize_optional_branch(
11043                                &transition_id,
11044                                RuntimeOptimizationKind::ParallelStateTransition,
11045                                RuntimeCommitBehavior::TransitionDecision,
11046                                "failed",
11047                                false,
11048                            );
11049                            routing_resolved.store(true, Ordering::SeqCst);
11050                            transition_finalized = true;
11051                        }
11052                    }
11053                }
11054            }
11055        }
11056    }
11057
11058    /// Streaming agent pipeline
11059    /// Uses all the same shared helpers as run_loop_internal.
11060    /// The ONLY difference: LLM calls use complete_stream() + yield deltas.
11061    fn run_loop_internal_stream<'a>(
11062        &'a self,
11063        input: &'a str,
11064        terminal: RuntimeStreamTerminalSlot,
11065    ) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
11066        let include_tool_events = self.streaming.include_tool_events;
11067        let include_state_events = self.streaming.include_state_events;
11068
11069        Box::pin(async_stream::stream! {
11070            self.begin_root_turn();
11071            // Parity with non-stream: resolve actor from context and load facts if changed.
11072            self.pre_turn_session_lifecycle().await;
11073
11074            let input_data = match self.process_input(input).await {
11075                Ok(data) => data,
11076                Err(e) => {
11077                    yield StreamChunk::error(e.to_string());
11078                    return;
11079                }
11080            };
11081            self.update_active_turn_context(&input_data.content, input_data.context.clone());
11082
11083            // Inject process context (detect/extract results) into agent context
11084            for (key, value) in &input_data.context {
11085                let _ = self.context_manager.set(key, value.clone());
11086            }
11087
11088            if input_data.metadata.rejected {
11089                let reason = input_data
11090                    .metadata
11091                    .rejection_reason
11092                    .unwrap_or_else(|| "Input rejected".to_string());
11093                warn!(reason = %reason, "Input rejected (stream)");
11094                yield StreamChunk::error(reason);
11095                return;
11096            }
11097
11098            let processed_input = &input_data.content;
11099
11100            if self.runtime_config.optimization.enabled
11101                && matches!(
11102                    self.runtime_config.optimization.streaming_policy,
11103                    crate::optimization::StreamingOptimizationPolicy::BufferUntilRoutingDone
11104                )
11105            {
11106                //
11107                // Buffered routing keeps stale stream output hidden until a branch winner is known.
11108                // The boxed future prevents this stream state machine from becoming too large.
11109                //
11110                match Box::pin(self.try_buffered_streaming_branches(processed_input, &input_data.context)).await {
11111                    Ok(Some((response, chunks))) => {
11112                        for chunk in chunks {
11113                            yield chunk;
11114                        }
11115                        record_runtime_stream_final(&terminal, response);
11116                        yield StreamChunk::Done {};
11117                        return;
11118                    }
11119                    Ok(None) => {}
11120                    Err(e) => {
11121                        yield StreamChunk::error(e.to_string());
11122                        return;
11123                    }
11124                }
11125            }
11126
11127            if self.runtime_config.optimization.enabled
11128                && matches!(
11129                    self.runtime_config.optimization.streaming_policy,
11130                    crate::optimization::StreamingOptimizationPolicy::PreflightOnly
11131                )
11132            {
11133                match self.try_pre_response_transition(processed_input).await {
11134                    Ok(Some(response)) => {
11135                        yield StreamChunk::content(&response.content);
11136                        record_runtime_stream_final(&terminal, response);
11137                        yield StreamChunk::Done {};
11138                        return;
11139                    }
11140                    Ok(None) => {}
11141                    Err(e) => {
11142                        yield StreamChunk::error(e.to_string());
11143                        return;
11144                    }
11145                }
11146            }
11147
11148            // Handle orchestration states in streaming mode.
11149            if let Some(ref sm) = self.state_machine
11150                && let Some(def) = sm.current_definition()
11151            {
11152                    let orchestration_result = if let Some(ref delegate_id) = def.delegate {
11153                        Some(self.handle_delegated_state(processed_input, delegate_id, &def).await)
11154                    } else if let Some(ref concurrent_config) = def.concurrent {
11155                        Some(self.handle_concurrent_state(processed_input, concurrent_config).await)
11156                    } else if let Some(ref group_chat_config) = def.group_chat {
11157                        Some(self.handle_group_chat_state(processed_input, group_chat_config).await)
11158                    } else if let Some(ref pipeline_config) = def.pipeline {
11159                        Some(self.handle_pipeline_state(processed_input, pipeline_config).await)
11160                    } else if let Some(ref handoff_config) = def.handoff {
11161                        Some(self.handle_handoff_state(processed_input, handoff_config).await)
11162                    } else {
11163                        None
11164                    };
11165
11166                    if let Some(result) = orchestration_result {
11167                        match result {
11168                            Ok(response) => {
11169                                yield StreamChunk::content(&response.content);
11170                                record_runtime_stream_final(&terminal, response);
11171                                yield StreamChunk::Done {};
11172                            }
11173                            Err(e) => {
11174                                yield StreamChunk::error(e.to_string());
11175                            }
11176                        }
11177                        return;
11178                    }
11179                }
11180
11181            // Skill routing
11182            match self.try_skill_route(processed_input).await {
11183                Ok(SkillRouteResult::Response { skill_id, content }) => {
11184                    if let Err(e) = self.commit_root_user_message(processed_input).await {
11185                        yield StreamChunk::error(e.to_string());
11186                        return;
11187                    }
11188                    match self.handle_skill_response(processed_input, &skill_id, content, &input_data.context).await {
11189                        Ok(resp) => {
11190                            yield StreamChunk::content(&resp.content);
11191                            record_runtime_stream_final(&terminal, resp);
11192                            yield StreamChunk::Done {};
11193                            return;
11194                        }
11195                        Err(e) => {
11196                            yield StreamChunk::error(e.to_string());
11197                            return;
11198                        }
11199                    }
11200                }
11201                Ok(SkillRouteResult::NeedsClarification {
11202                    response,
11203                    ownership,
11204                }) => {
11205                    let admission = match self
11206                        .admit_optional_disambiguation_ownership(ownership)
11207                        .await
11208                    {
11209                        Ok(admission) => admission,
11210                        Err(e) => {
11211                            yield StreamChunk::error(e.to_string());
11212                            return;
11213                        }
11214                    };
11215                    if let Err(e) = self.commit_root_user_message(processed_input).await {
11216                        yield StreamChunk::error(e.to_string());
11217                        return;
11218                    }
11219                    let _ = self.memory.add_message(ChatMessage::assistant(&response.content)).await;
11220                    drop(admission);
11221                    if let Err(e) = self.finish_turn_if_root(&response).await {
11222                        yield StreamChunk::error(e.to_string());
11223                        return;
11224                    }
11225                    yield StreamChunk::content(&response.content);
11226                    record_runtime_stream_final(&terminal, response);
11227                    yield StreamChunk::Done {};
11228                    return;
11229                }
11230                Ok(SkillRouteResult::NoMatch) => {} // no skill matched, continue
11231                Err(e) => {
11232                    yield StreamChunk::error(e.to_string());
11233                    return;
11234                }
11235            }
11236
11237            // Reasoning mode determination
11238            let effective_reasoning = self.get_effective_reasoning_config();
11239            let reasoning_mode = match self.determine_reasoning_mode(processed_input).await {
11240                Ok(mode) => mode,
11241                Err(e) => {
11242                    yield StreamChunk::error(e.to_string());
11243                    return;
11244                }
11245            };
11246            let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
11247
11248            info!(
11249                reasoning_mode = ?reasoning_mode,
11250                auto_detected = auto_detected,
11251                "Reasoning mode determined (stream)"
11252            );
11253
11254            // Plan-and-Execute: yield final result as single chunk (not token-by-token)
11255            if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
11256                if let Err(e) = self.commit_root_user_message(processed_input).await {
11257                    yield StreamChunk::error(e.to_string());
11258                    return;
11259                }
11260                match self.handle_plan_and_execute(processed_input, &input_data.context, auto_detected).await {
11261                    Ok(resp) => {
11262                        yield StreamChunk::content(&resp.content);
11263                        record_runtime_stream_final(&terminal, resp);
11264                        yield StreamChunk::Done {};
11265                        return;
11266                    }
11267                    Err(e) => {
11268                        yield StreamChunk::error(e.to_string());
11269                        return;
11270                    }
11271                }
11272            }
11273
11274            if let Err(e) = self.commit_root_user_message(processed_input).await {
11275                yield StreamChunk::error(e.to_string());
11276                return;
11277            }
11278
11279            let llm = match self.get_state_llm() {
11280                Ok(llm) => llm,
11281                Err(e) => {
11282                    yield StreamChunk::error(e.to_string());
11283                    return;
11284                }
11285            };
11286
11287            let mut iterations = 0u32;
11288            let mut all_tool_calls: Vec<ToolCall> = Vec::new();
11289            let mut thinking_content: Option<String> = None;
11290
11291            loop {
11292                // When reasoning is active, cap iterations at the reasoning-specific limit.
11293                let effective_max = if reasoning_mode != ReasoningMode::None {
11294                    let rc = self.get_effective_reasoning_config();
11295                    self.max_iterations.min(rc.max_iterations)
11296                } else {
11297                    self.max_iterations
11298                };
11299
11300                if iterations >= effective_max {
11301                    let err_msg = format!("Max iterations ({}) exceeded", effective_max);
11302                    let err = AgentError::Other(err_msg.clone());
11303                    self.hooks.on_error(&err).await;
11304                    error!(iterations = iterations, "Max iterations exceeded (stream)");
11305                    yield StreamChunk::error(err_msg);
11306                    return;
11307                }
11308                iterations += 1;
11309                *self.iteration_count.write() = iterations;
11310
11311                debug!(iteration = iterations, max = effective_max, "LLM call (stream)");
11312
11313                let protocol = match self.main_tool_protocol(llm.as_ref(), false).await {
11314                    Ok(protocol) => protocol,
11315                    Err(e) => {
11316                        yield StreamChunk::error(e.to_string());
11317                        return;
11318                    }
11319                };
11320                let mut messages = match self
11321                    .build_messages_internal(true, None, protocol.choice.is_none())
11322                    .await
11323                {
11324                    Ok(m) => m,
11325                    Err(e) => {
11326                        yield StreamChunk::error(e.to_string());
11327                        return;
11328                    }
11329                };
11330                self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
11331
11332                self.hooks.on_llm_start(&messages).await;
11333                let llm_start = Instant::now();
11334
11335                // Check if reflection is active — if so, suppress streaming for this LLM call
11336                // because we may need to retry and the user would see a stale first attempt.
11337                let reflection_active = self
11338                    .should_reflect(processed_input, "")
11339                    .await
11340                    .unwrap_or_default();
11341
11342                let buffered_decision = reflection_active || protocol.choice.is_some();
11343                let content = if buffered_decision {
11344                    //
11345                    // Explicit tool choice buffers the provider decision so no text or tool call is visible before the runtime validates and commits it.
11346                    //
11347                    let response = match self
11348                        .complete_main_llm_with_recovery(
11349                            Arc::clone(&llm),
11350                            &messages,
11351                            &protocol,
11352                        )
11353                        .await
11354                    {
11355                        Ok(r) => r,
11356                        Err(e) => {
11357                            yield StreamChunk::error(e.to_string());
11358                            return;
11359                        }
11360                    };
11361                    let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
11362                    self.hooks.on_llm_complete(&response, llm_duration_ms).await;
11363                    response.content.trim().to_string()
11364                } else {
11365                    // Streaming LLM call
11366                    let llm_stream = match self
11367                        .observe_purpose(
11368                            ObservationPurpose::MainResponse,
11369                            llm.complete_stream(&messages, None),
11370                        )
11371                        .await
11372                    {
11373                        Ok(s) => s,
11374                        Err(e) => {
11375                            yield StreamChunk::error(e.to_string());
11376                            return;
11377                        }
11378                    };
11379                    let mut accumulated = String::new();
11380                    let mut stream_inner = llm_stream;
11381                    while let Some(chunk_result) = stream_inner.next().await {
11382                        match chunk_result {
11383                            Ok(chunk) => {
11384                                accumulated.push_str(&chunk.delta);
11385                                yield StreamChunk::content(chunk.delta);
11386                            }
11387                            Err(e) => {
11388                                yield StreamChunk::error(e.to_string());
11389                                return;
11390                            }
11391                        }
11392                    }
11393                    let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
11394                    // Construct LLMResponse for hooks
11395                    let llm_response = ai_agents_core::LLMResponse::new(
11396                        accumulated.trim(),
11397                        ai_agents_core::FinishReason::Stop,
11398                    );
11399                    self.hooks.on_llm_complete(&llm_response, llm_duration_ms).await;
11400                    accumulated.trim().to_string()
11401                };
11402
11403                // Tool call handling
11404                if let Some(tool_calls) = self.parse_main_tool_calls(&content, &protocol) {
11405                    let native_tool_call = Self::is_native_tool_call_content(&content);
11406                    // Emit tool events for streaming
11407                    // First check transitions (same as blocking path)
11408                    let transition_fired = match self.evaluate_transitions(processed_input, &content).await {
11409                        Ok(v) => v,
11410                        Err(e) => {
11411                            yield StreamChunk::error(e.to_string());
11412                            return;
11413                        }
11414                    };
11415                    if transition_fired {
11416                        let _ = self.memory.add_message(ChatMessage::assistant(
11417                            "(Transitioned to new state — tool call handled by workflow)",
11418                        )).await;
11419
11420                        if include_state_events
11421                            && let Some(state) = self.current_state()
11422                        {
11423                                yield StreamChunk::state_transition(None, state);
11424                            }
11425                        continue;
11426                    }
11427
11428                    // Store the assistant's tool-call message (same as blocking path)
11429                    let _ = self.memory.add_message(ChatMessage::assistant(&content)).await;
11430
11431                    // Execute tools with streaming events
11432                    let results = self.execute_tools_parallel(&tool_calls).await;
11433
11434                    for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
11435                        if include_tool_events {
11436                            yield StreamChunk::tool_start(&tool_call.id, &tool_call.name);
11437                        }
11438
11439                        match result {
11440                            Ok(output) => {
11441                                if include_tool_events {
11442                                    yield StreamChunk::tool_result(
11443                                        &tool_call.id,
11444                                        &tool_call.name,
11445                                        &output,
11446                                        true,
11447                                    );
11448                                }
11449                                let _ = self.memory
11450                                    .add_message(Self::tool_result_message(
11451                                        tool_call,
11452                                        &output,
11453                                        native_tool_call,
11454                                    ))
11455                                    .await;
11456                            }
11457                            Err(e) => {
11458                                if matches!(e, AgentError::HITLRejected(_)) {
11459                                    let _ = self.memory.add_message(ChatMessage::assistant(
11460                                        format!("The operation was rejected by the approver: {}", e),
11461                                    )).await;
11462                                    let response = AgentResponse {
11463                                        content: format!("Operation cancelled: {}", e),
11464                                        metadata: None,
11465                                        tool_calls: Some(all_tool_calls.clone()),
11466                                    };
11467                                    if let Err(finalize_error) = self.finish_turn_if_root(&response).await {
11468                                        yield StreamChunk::error(finalize_error.to_string());
11469                                        return;
11470                                    }
11471                                    let legacy_error = response.content.clone();
11472                                    record_runtime_stream_final(&terminal, response);
11473                                    yield StreamChunk::error(legacy_error);
11474                                    yield StreamChunk::Done {};
11475                                    return;
11476                                }
11477                                if include_tool_events {
11478                                    yield StreamChunk::tool_result(
11479                                        &tool_call.id,
11480                                        &tool_call.name,
11481                                        e.to_string(),
11482                                        false,
11483                                    );
11484                                }
11485                                let _ = self.memory
11486                                    .add_message(Self::tool_result_message(
11487                                        tool_call,
11488                                        &format!("Error: {}", e),
11489                                        native_tool_call,
11490                                    ))
11491                                    .await;
11492                            }
11493                        }
11494                        all_tool_calls.push(tool_call.clone());
11495
11496                        if include_tool_events {
11497                            yield StreamChunk::tool_end(&tool_call.id);
11498                        }
11499                    }
11500                    continue;
11501                }
11502
11503                // Extract thinking, process output
11504                let (extracted_thinking, answer) = self.extract_thinking(&content);
11505                if extracted_thinking.is_some() {
11506                    thinking_content = extracted_thinking;
11507                }
11508
11509                let output_data = match self.process_output(&answer, &input_data.context).await {
11510                    Ok(d) => d,
11511                    Err(e) => {
11512                        yield StreamChunk::error(e.to_string());
11513                        return;
11514                    }
11515                };
11516
11517                let final_content = if output_data.metadata.rejected {
11518                    output_data
11519                        .metadata
11520                        .rejection_reason
11521                        .unwrap_or_else(|| answer.to_string())
11522                } else {
11523                    output_data.content
11524                };
11525
11526                // Reflection (uses blocking LLM calls for retries)
11527                let (final_content, reflection_metadata) = match self
11528                    .run_reflection(&*llm, processed_input, final_content)
11529                    .await
11530                {
11531                    Ok(r) => r,
11532                    Err(e) => {
11533                        yield StreamChunk::error(e.to_string());
11534                        return;
11535                    }
11536                };
11537
11538                let final_content = self.format_response_with_thinking(
11539                    thinking_content.as_deref(),
11540                    &final_content,
11541                );
11542
11543                // Buffered decisions emit only the accepted final text.
11544                if buffered_decision {
11545                    yield StreamChunk::content(&final_content);
11546                }
11547
11548                // Post-loop: memory, transitions, post-transition re-generation.
11549                // For NeedsRedispatch, run_loop_internal handles the new state's full
11550                // dispatch and its result is yielded as a single non-streamed chunk.
11551                let post_result = match self
11552                    .post_loop_processing(processed_input, final_content)
11553                    .await
11554                {
11555                    Ok(r) => r,
11556                    Err(e) => {
11557                        yield StreamChunk::error(e.to_string());
11558                        return;
11559                    }
11560                };
11561
11562                let (final_content, transitioned) = match post_result {
11563                    PostLoopResult::NoTransition(content) => (content, false),
11564                    PostLoopResult::Transitioned(content) => (content, true),
11565                    PostLoopResult::NeedsRedispatch => {
11566                        const MAX_REDISPATCH_DEPTH: u32 = 3;
11567                        let current_depth = *self.redispatch_depth.read();
11568                        let content = if current_depth >= MAX_REDISPATCH_DEPTH {
11569                            warn!(
11570                                depth = current_depth,
11571                                "Post-transition re-dispatch depth limit reached (stream)"
11572                            );
11573                            let c = String::new();
11574                            let _ = self.memory.add_message(ChatMessage::assistant(&c)).await;
11575                            c
11576                        } else {
11577                            *self.redispatch_depth.write() += 1;
11578                            if let Some(context) = self.active_turn_context.write().as_mut() {
11579                                context.enter_redispatch();
11580                            }
11581                            info!(
11582                                depth = current_depth + 1,
11583                                "Re-dispatching for new state after transition (stream)"
11584                            );
11585                            let result = self.run_loop_internal(processed_input).await;
11586                            *self.redispatch_depth.write() -= 1;
11587                            if let Some(context) = self.active_turn_context.write().as_mut() {
11588                                context.exit_redispatch();
11589                            }
11590                            match result {
11591                                Ok(resp) => resp.content,
11592                                Err(e) => {
11593                                    yield StreamChunk::error(e.to_string());
11594                                    return;
11595                                }
11596                            }
11597                        };
11598                        (content, true)
11599                    }
11600                };
11601
11602                if transitioned {
11603                    if include_state_events
11604                        && let Some(state) = self.current_state()
11605                    {
11606                            yield StreamChunk::state_transition(None, state);
11607                        }
11608                    // Yield the post-transition re-generated or re-dispatched content.
11609                    yield StreamChunk::content(&final_content);
11610                }
11611
11612                // Build and finalize the same authoritative response shape before exposing the terminal event.
11613                let final_response = self.build_agent_response(AgentResponseParts {
11614                    content: final_content,
11615                    all_tool_calls,
11616                    reasoning_mode,
11617                    auto_detected,
11618                    iterations,
11619                    thinking: thinking_content,
11620                    reflection_metadata,
11621                });
11622                if let Err(e) = self.finish_turn_if_root(&final_response).await {
11623                    yield StreamChunk::error(e.to_string());
11624                    return;
11625                }
11626
11627                record_runtime_stream_final(&terminal, final_response);
11628                yield StreamChunk::Done {};
11629                return;
11630            }
11631        })
11632    }
11633
11634    /// Streams the root turn while keeping clarification and confirmation questions as terminal responses for their turn.
11635    /// Pending manager and skill ownership must survive until explicit confirmation returns a resolved result for redispatch.
11636    fn run_loop_stream<'a>(
11637        &'a self,
11638        input: &'a str,
11639        terminal: RuntimeStreamTerminalSlot,
11640    ) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
11641        Box::pin(async_stream::stream! {
11642            self.begin_root_turn();
11643            let _root_cleanup = RootTurnCleanup::new(self);
11644            self.hooks.on_message_received(input).await;
11645
11646            // One-shot context initialization (mirrors run_loop)
11647            if !self.context_initialized.swap(true, Ordering::SeqCst) {
11648                if let Err(e) = self.context_manager.initialize().await {
11649                    yield StreamChunk::error(e.to_string());
11650                    return;
11651                }
11652                debug!("Context manager initialized (defaults, env, builtins)");
11653            }
11654
11655            if let Err(e) = self.check_turn_timeout().await {
11656                yield StreamChunk::error(e.to_string());
11657                return;
11658            }
11659            if let Err(e) = self.context_manager.refresh_per_turn().await {
11660                yield StreamChunk::error(e.to_string());
11661                return;
11662            }
11663
11664            // Clear stale disambiguation context from previous turns.
11665            self.clear_disambiguation_context();
11666
11667            // Disambiguation check (before input processing)
11668            if let Some(ref disambiguator) = self.disambiguation_manager {
11669                let disambiguation_context = match self.build_disambiguation_context().await {
11670                    Ok(ctx) => ctx,
11671                    Err(e) => {
11672                        yield StreamChunk::error(e.to_string());
11673                        return;
11674                    }
11675                };
11676
11677                let state_override = self
11678                    .state_machine
11679                    .as_ref()
11680                    .and_then(|sm| sm.current_definition())
11681                    .and_then(|def| def.disambiguation.clone());
11682
11683                let state_generation = self
11684                    .state_machine
11685                    .as_ref()
11686                    .map(|state_machine| state_machine.generation());
11687                let disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
11688                let mut result = match self
11689                    .observe_purpose(
11690                        ObservationPurpose::DisambiguationDetection,
11691                        disambiguator.process_input_with_override(
11692                            input,
11693                            &disambiguation_context,
11694                            state_override.as_ref(),
11695                            None,
11696                        ),
11697                    )
11698                    .await
11699                {
11700                    Ok(r) => r,
11701                    Err(e) => {
11702                        yield StreamChunk::error(e.to_string());
11703                        return;
11704                    }
11705                };
11706                let current_state_generation = self
11707                    .state_machine
11708                    .as_ref()
11709                    .map(|state_machine| state_machine.generation());
11710                if current_state_generation != state_generation
11711                    || self.disambiguation_epoch.load(Ordering::SeqCst) != disambiguation_epoch
11712                {
11713                    disambiguator.clear_pending().await;
11714                    *self.pending_skill_id.write() = None;
11715                    result = DisambiguationResult::Abandoned { new_input: None };
11716                    info!(
11717                        confirmation_event = "invalidated",
11718                        invalidation_reason = "state_generation_changed",
11719                        "Streaming disambiguation result invalidated before redispatch"
11720                    );
11721                }
11722                match result {
11723                    DisambiguationResult::Clear => {
11724                        debug!("Input is clear, proceeding normally (stream)");
11725                    }
11726                    DisambiguationResult::NeedsClarification {
11727                        question,
11728                        detection,
11729                    } => {
11730                        let admission = match self
11731                            .admit_disambiguation_redispatch(
11732                                disambiguation_epoch,
11733                                state_generation,
11734                            )
11735                            .await
11736                        {
11737                            Ok(admission) => admission,
11738                            Err(error) => {
11739                                *self.pending_skill_id.write() = None;
11740                                yield StreamChunk::error(error.to_string());
11741                                return;
11742                            }
11743                        };
11744                        let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
11745                        info!(
11746                            ambiguity_type = ?detection.ambiguity_type,
11747                            confidence = detection.confidence,
11748                            "Input requires clarification (stream)"
11749                        );
11750                        // Confirmation uses the same terminal branch so enriched input cannot stream before explicit agreement.
11751                        // Keep pending_skill_id intact for the later confirmed redispatch.
11752                        if let Err(e) = self.commit_root_user_message(input).await {
11753                            yield StreamChunk::error(e.to_string());
11754                            return;
11755                        }
11756                        let _ = self
11757                            .memory
11758                            .add_message(ChatMessage::assistant(&question.question))
11759                            .await;
11760                        let status = if awaiting_confirmation {
11761                            "awaiting_confirmation"
11762                        } else {
11763                            "awaiting_clarification"
11764                        };
11765                        let response = AgentResponse::new(&question.question).with_metadata(
11766                            "disambiguation",
11767                            serde_json::json!({ "status": status }),
11768                        );
11769                        drop(admission);
11770                        if let Err(e) = self.finish_turn_if_root(&response).await {
11771                            yield StreamChunk::error(e.to_string());
11772                            return;
11773                        }
11774                        yield StreamChunk::content(&question.question);
11775                        record_runtime_stream_final(&terminal, response);
11776                        yield StreamChunk::Done {};
11777                        return;
11778                    }
11779                    DisambiguationResult::Clarified {
11780                        enriched_input,
11781                        resolved,
11782                        ..
11783                    } => {
11784                        let admission = match self
11785                            .admit_disambiguation_redispatch(
11786                                disambiguation_epoch,
11787                                state_generation,
11788                            )
11789                            .await
11790                        {
11791                            Ok(admission) => admission,
11792                            Err(error) => {
11793                                *self.pending_skill_id.write() = None;
11794                                yield StreamChunk::error(error.to_string());
11795                                return;
11796                            }
11797                        };
11798                        info!(
11799                            resolved_count = resolved.len(),
11800                            enriched = %enriched_input,
11801                            "Input clarified (stream)"
11802                        );
11803                        for (key, value) in &resolved {
11804                            let context_key = format!("disambiguation.{}", key);
11805                            let _ = self.context_manager.set(&context_key, value.clone());
11806                        }
11807                        if let Some(intent) = resolved.get("intent") {
11808                            let _ = self.context_manager.set("resolved_intent", intent.clone());
11809                        }
11810                        let _ = self
11811                            .context_manager
11812                            .set("disambiguation.resolved", serde_json::Value::Bool(true));
11813
11814                        // Check if this clarification was triggered by a skill-level override.
11815                        // Re-run skill disambiguation to verify all required_clarity fields
11816                        // are present before executing.
11817                        let skill_id = self.pending_skill_id.read().clone();
11818                        if let Some(skill_id) = skill_id {
11819                            info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input (stream)");
11820                            drop(admission);
11821                            match self
11822                                .recheck_skill_disambiguation(
11823                                    &skill_id,
11824                                    &enriched_input,
11825                                    disambiguation_epoch,
11826                                    state_generation,
11827                                )
11828                                .await
11829                            {
11830                                Ok(resp) => {
11831                                    yield StreamChunk::content(&resp.content);
11832                                    record_runtime_stream_final(&terminal, resp);
11833                                    yield StreamChunk::Done {};
11834                                    return;
11835                                }
11836                                Err(e) => {
11837                                    yield StreamChunk::error(e.to_string());
11838                                    return;
11839                                }
11840                            }
11841                        }
11842
11843                        // Forward to internal stream with enriched input.
11844                        drop(admission);
11845                        let mut inner = self.run_loop_internal_stream(
11846                            &enriched_input,
11847                            Arc::clone(&terminal),
11848                        );
11849                        while let Some(chunk) = inner.next().await {
11850                            yield chunk;
11851                        }
11852                        return;
11853                    }
11854                    DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
11855                        info!("Proceeding with best guess (stream)");
11856
11857                        // Same skill-id re-check for best-guess path
11858                        let skill_id = self.pending_skill_id.read().clone();
11859                        if let Some(skill_id) = skill_id {
11860                            info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input (stream)");
11861                            match self
11862                                .recheck_skill_disambiguation(
11863                                    &skill_id,
11864                                    &enriched_input,
11865                                    disambiguation_epoch,
11866                                    state_generation,
11867                                )
11868                                .await
11869                            {
11870                                Ok(resp) => {
11871                                    yield StreamChunk::content(&resp.content);
11872                                    record_runtime_stream_final(&terminal, resp);
11873                                    yield StreamChunk::Done {};
11874                                    return;
11875                                }
11876                                Err(e) => {
11877                                    yield StreamChunk::error(e.to_string());
11878                                    return;
11879                                }
11880                            }
11881                        }
11882
11883                        let mut inner = self.run_loop_internal_stream(
11884                            &enriched_input,
11885                            Arc::clone(&terminal),
11886                        );
11887                        while let Some(chunk) = inner.next().await {
11888                            yield chunk;
11889                        }
11890                        return;
11891                    }
11892                    DisambiguationResult::GiveUp { reason } => {
11893                        *self.pending_skill_id.write() = None;
11894                        warn!(reason = %reason, "Disambiguation gave up (stream)");
11895                        let apology = self
11896                            .generate_localized_apology(
11897                                "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
11898                                &reason,
11899                            )
11900                            .await
11901                            .unwrap_or_else(|_| {
11902                                format!("I'm sorry, I couldn't understand your request: {}", reason)
11903                            });
11904                        let response = AgentResponse::new(&apology);
11905                        if let Err(e) = self.finish_turn_if_root(&response).await {
11906                            yield StreamChunk::error(e.to_string());
11907                            return;
11908                        }
11909                        yield StreamChunk::content(&apology);
11910                        record_runtime_stream_final(&terminal, response);
11911                        yield StreamChunk::Done {};
11912                        return;
11913                    }
11914                    DisambiguationResult::Escalate { reason } => {
11915                        *self.pending_skill_id.write() = None;
11916                        info!(reason = %reason, "Escalating to human (stream)");
11917                        if let Some(ref hitl) = self.hitl_engine {
11918                            let trigger =
11919                                ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
11920                            let mut context_map = HashMap::new();
11921                            context_map.insert("original_input".to_string(), serde_json::json!(input));
11922                            context_map.insert("reason".to_string(), serde_json::json!(&reason));
11923                            let check_result = HITLCheckResult::required(
11924                                trigger,
11925                                context_map,
11926                                format!("User request needs human assistance: {}", reason),
11927                                Some(hitl.config().default_timeout_seconds),
11928                            );
11929                            match self.request_hitl_approval(check_result).await {
11930                                Ok(ApprovalResult::Approved | ApprovalResult::Modified { .. }) => {
11931                                    let mut inner = self.run_loop_internal_stream(
11932                                        input,
11933                                        Arc::clone(&terminal),
11934                                    );
11935                                    while let Some(chunk) = inner.next().await {
11936                                        yield chunk;
11937                                    }
11938                                    return;
11939                                }
11940                                Ok(_) => {}
11941                                Err(e) => {
11942                                    yield StreamChunk::error(e.to_string());
11943                                    return;
11944                                }
11945                            }
11946                        }
11947                        let apology = self
11948                            .generate_localized_apology(
11949                                "Explain briefly that you're transferring the user to a human agent for help.",
11950                                &reason,
11951                            )
11952                            .await
11953                            .unwrap_or_else(|_| {
11954                                format!("I need human assistance to help with your request: {}", reason)
11955                            });
11956                        let response = AgentResponse::new(&apology);
11957                        if let Err(e) = self.finish_turn_if_root(&response).await {
11958                            yield StreamChunk::error(e.to_string());
11959                            return;
11960                        }
11961                        yield StreamChunk::content(&apology);
11962                        record_runtime_stream_final(&terminal, response);
11963                        yield StreamChunk::Done {};
11964                        return;
11965                    }
11966                    DisambiguationResult::Abandoned { new_input } => {
11967                        *self.pending_skill_id.write() = None;
11968
11969                        info!(
11970                            has_new_input = new_input.is_some(),
11971                            "Clarification abandoned by user (stream)"
11972                        );
11973
11974                        if let Err(e) = self.commit_root_user_message(input).await {
11975                            yield StreamChunk::error(e.to_string());
11976                            return;
11977                        }
11978
11979                        match new_input {
11980                            Some(fresh_input) => {
11981                                // Topic switch: forward to internal stream with fresh input.
11982                                let mut inner = self.run_loop_internal_stream(
11983                                    &fresh_input,
11984                                    Arc::clone(&terminal),
11985                                );
11986                                while let Some(chunk) = inner.next().await {
11987                                    yield chunk;
11988                                }
11989                                return;
11990                            }
11991                            None => {
11992                                // Pure abandonment: generate a brief acknowledgment.
11993                                let ack = self
11994                                    .generate_localized_apology(
11995                                        "The user changed their mind about their previous request. \
11996                                         Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
11997                                         Do NOT apologize excessively. Be concise.",
11998                                        "User abandoned clarification",
11999                                    )
12000                                    .await
12001                                    .unwrap_or_else(|_| {
12002                                        "OK, no problem. What else can I help with?".to_string()
12003                                    });
12004
12005                                let _ = self
12006                                    .memory
12007                                    .add_message(ChatMessage::assistant(&ack))
12008                                    .await;
12009
12010                                let response = AgentResponse::new(&ack);
12011                                if let Err(e) = self.finish_turn_if_root(&response).await {
12012                                    yield StreamChunk::error(e.to_string());
12013                                    return;
12014                                }
12015                                yield StreamChunk::content(&ack);
12016                                record_runtime_stream_final(&terminal, response);
12017                                yield StreamChunk::Done {};
12018                                return;
12019                            }
12020                        }
12021                    }
12022                }
12023            }
12024
12025            // No disambiguation or Clear result — proceed with internal stream
12026            let mut inner = self.run_loop_internal_stream(input, Arc::clone(&terminal));
12027            while let Some(chunk) = inner.next().await {
12028                yield chunk;
12029            }
12030        })
12031    }
12032
12033    pub fn info(&self) -> AgentInfo {
12034        self.info.clone()
12035    }
12036
12037    pub fn skills(&self) -> &[SkillDefinition] {
12038        &self.skills
12039    }
12040
12041    /// Clears conversation and pending runtime ownership through one reset contract.
12042    async fn reset_runtime_state(&self) -> Result<()> {
12043        let _admission = self.disambiguation_admission.write().await;
12044        if self.state_transition_reserved.load(Ordering::SeqCst) {
12045            return Err(AgentError::Other(
12046                "Cannot reset while a state transition is in progress".to_string(),
12047            ));
12048        }
12049        self.disambiguation_epoch.fetch_add(1, Ordering::SeqCst);
12050        *self.pending_skill_id.write() = None;
12051        if let Some(disambiguator) = self.disambiguation_manager.as_ref() {
12052            disambiguator.clear_pending().await;
12053        }
12054        self.memory.clear().await?;
12055        *self.iteration_count.write() = 0;
12056        self.tool_call_history.write().clear();
12057        if let Some(ref sm) = self.state_machine {
12058            sm.reset();
12059        }
12060        Ok(())
12061    }
12062
12063    /// Resets the runtime using the same cleanup path as the Agent trait.
12064    pub async fn reset(&self) -> Result<()> {
12065        self.reset_runtime_state().await
12066    }
12067
12068    pub fn max_context_tokens(&self) -> u32 {
12069        self.max_context_tokens
12070    }
12071
12072    pub fn llm_registry(&self) -> &Arc<LLMRegistry> {
12073        &self.llm_registry
12074    }
12075
12076    pub fn state_machine(&self) -> Option<&Arc<StateMachine>> {
12077        self.state_machine.as_ref()
12078    }
12079
12080    pub fn context_manager(&self) -> &Arc<ContextManager> {
12081        &self.context_manager
12082    }
12083
12084    pub fn tool_call_history(&self) -> Vec<ToolCallRecord> {
12085        self.tool_call_history.read().clone()
12086    }
12087
12088    pub fn memory_token_budget(&self) -> Option<&MemoryTokenBudget> {
12089        self.memory_token_budget.as_ref()
12090    }
12091
12092    pub fn parallel_tools_config(&self) -> &ParallelToolsConfig {
12093        &self.parallel_tools
12094    }
12095
12096    pub fn streaming_config(&self) -> &StreamingConfig {
12097        &self.streaming
12098    }
12099
12100    pub fn hooks(&self) -> &Arc<dyn AgentHooks> {
12101        &self.hooks
12102    }
12103
12104    pub fn hitl_engine(&self) -> Option<&HITLEngine> {
12105        self.hitl_engine.as_ref()
12106    }
12107
12108    pub fn approval_handler(&self) -> &Arc<dyn ApprovalHandler> {
12109        &self.approval_handler
12110    }
12111
12112    /// Build a context map with language hints from context_manager for HITL message localization.
12113    fn build_hitl_language_context(&self) -> HashMap<String, Value> {
12114        let mut ctx = HashMap::new();
12115        for key in &["user.language", "input.detected.language", "language"] {
12116            if let Some(val) = self.context_manager.get(key) {
12117                ctx.insert(key.to_string(), val);
12118            }
12119        }
12120        ctx
12121    }
12122
12123    /// Send a HITL check result through the approval flow and return the full ApprovalResult.
12124    async fn request_hitl_approval(&self, check_result: HITLCheckResult) -> Result<ApprovalResult> {
12125        let Some(request) = check_result.into_request() else {
12126            return Ok(ApprovalResult::Approved);
12127        };
12128
12129        self.hooks.on_approval_requested(&request).await;
12130
12131        let timeout = request.timeout;
12132
12133        let raw_result = if let Some(duration) = timeout {
12134            match tokio::time::timeout(
12135                duration,
12136                self.approval_handler.request_approval(request.clone()),
12137            )
12138            .await
12139            {
12140                Ok(result) => result,
12141                Err(_) => ApprovalResult::timeout(),
12142            }
12143        } else {
12144            self.approval_handler
12145                .request_approval(request.clone())
12146                .await
12147        };
12148
12149        self.hooks
12150            .on_approval_result(&request.id, &raw_result)
12151            .await;
12152
12153        let (outcome, effective_result): (ApprovalResolvedOutcome, Result<ApprovalResult>) =
12154            match &raw_result {
12155                ApprovalResult::Approved => (
12156                    ApprovalResolvedOutcome::Approved,
12157                    Ok(ApprovalResult::Approved),
12158                ),
12159                ApprovalResult::Rejected { reason } => (
12160                    ApprovalResolvedOutcome::Rejected {
12161                        reason: reason.clone(),
12162                    },
12163                    Ok(ApprovalResult::Rejected {
12164                        reason: reason.clone(),
12165                    }),
12166                ),
12167                ApprovalResult::Modified { changes } => (
12168                    ApprovalResolvedOutcome::Modified {
12169                        changes: changes.clone(),
12170                    },
12171                    Ok(ApprovalResult::Modified {
12172                        changes: changes.clone(),
12173                    }),
12174                ),
12175                ApprovalResult::Timeout => {
12176                    if let Some(ref engine) = self.hitl_engine {
12177                        match engine.config().on_timeout {
12178                            TimeoutAction::Approve => (
12179                                ApprovalResolvedOutcome::Approved,
12180                                Ok(ApprovalResult::Approved),
12181                            ),
12182                            TimeoutAction::Reject => {
12183                                let reason = Some("Timeout".to_string());
12184                                (
12185                                    ApprovalResolvedOutcome::Rejected {
12186                                        reason: reason.clone(),
12187                                    },
12188                                    Ok(ApprovalResult::Rejected { reason }),
12189                                )
12190                            }
12191                            TimeoutAction::Error => {
12192                                let message = "HITL approval timeout".to_string();
12193                                (
12194                                    ApprovalResolvedOutcome::Error {
12195                                        message: message.clone(),
12196                                    },
12197                                    Err(AgentError::Other(message)),
12198                                )
12199                            }
12200                        }
12201                    } else {
12202                        let reason = Some("Timeout (no engine)".to_string());
12203                        (
12204                            ApprovalResolvedOutcome::Rejected {
12205                                reason: reason.clone(),
12206                            },
12207                            Ok(ApprovalResult::Rejected { reason }),
12208                        )
12209                    }
12210                }
12211            };
12212
12213        self.hooks
12214            .on_approval_resolved(&request, &raw_result, &outcome)
12215            .await;
12216
12217        effective_result
12218    }
12219
12220    pub async fn check_state_hitl(&self, from: Option<&str>, to: &str) -> Result<bool> {
12221        if let Some(ref hitl_engine) = self.hitl_engine {
12222            let hitl_lang_ctx = self.build_hitl_language_context();
12223            let check_result = self
12224                .observe_purpose(
12225                    ObservationPurpose::HitlLocalization,
12226                    hitl_engine.check_state_transition_with_localization(
12227                        from,
12228                        to,
12229                        &hitl_lang_ctx,
12230                        self.approval_handler.as_ref(),
12231                        Some(&self.llm_registry),
12232                    ),
12233                )
12234                .await?;
12235            if check_result.is_required() {
12236                let result = self.request_hitl_approval(check_result).await?;
12237                return Ok(matches!(
12238                    result,
12239                    ApprovalResult::Approved | ApprovalResult::Modified { .. }
12240                ));
12241            }
12242        }
12243        Ok(true)
12244    }
12245
12246    /// Execute multiple tools in parallel
12247    async fn execute_tools_parallel(
12248        &self,
12249        tool_calls: &[ToolCall],
12250    ) -> Vec<(String, Result<String>)> {
12251        let can_run_parallel = tool_calls.iter().all(|tc| {
12252            self.tools
12253                .resolve(&tc.name)
12254                .map(|resolved| resolved.tool.classify_call(&tc.arguments).concurrency_safe)
12255                .unwrap_or(false)
12256        });
12257
12258        if !self.parallel_tools.enabled || tool_calls.len() <= 1 || !can_run_parallel {
12259            let mut results = Vec::new();
12260            for tc in tool_calls {
12261                let result = self
12262                    .observe_purpose(
12263                        current_observation_context()
12264                            .map(|context| context.purpose)
12265                            .unwrap_or_default(),
12266                        self.execute_tool_smart(tc),
12267                    )
12268                    .await;
12269                results.push((tc.id.clone(), result));
12270            }
12271            return results;
12272        }
12273
12274        let chunks: Vec<_> = tool_calls
12275            .chunks(self.parallel_tools.max_parallel)
12276            .collect();
12277
12278        let mut all_results = Vec::new();
12279
12280        for chunk in chunks {
12281            let futures: Vec<_> = chunk
12282                .iter()
12283                .map(|tc| {
12284                    let tc = tc.clone();
12285                    async move {
12286                        let result = self.execute_tool_smart(&tc).await;
12287                        (tc.id.clone(), result)
12288                    }
12289                })
12290                .collect();
12291
12292            let results = futures::future::join_all(futures).await;
12293            all_results.extend(results);
12294        }
12295
12296        all_results
12297    }
12298
12299    /// Streams provisional chunks and preserves the legacy Done terminal contract.
12300    pub async fn chat_stream<'a>(
12301        &'a self,
12302        input: &'a str,
12303    ) -> Result<Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>>> {
12304        //
12305        // Streaming readiness is checked before returning a stream so capability failures cannot appear as partial output.
12306        //
12307        self.init_storage().await?;
12308        info!(input_len = input.len(), "Starting streaming chat");
12309        let terminal = new_runtime_stream_terminal_slot();
12310        let inner = self.run_loop_stream(input, terminal);
12311        if let Some(context) = self.build_observation_context(None) {
12312            let stream: Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> =
12313                Box::pin(async_stream::stream! {
12314                    let mut inner = inner;
12315                    loop {
12316                        let next = with_observation_context(context.clone(), inner.next()).await;
12317                        match next {
12318                            Some(chunk) => yield chunk,
12319                            None => break,
12320                        }
12321                    }
12322                    self.export_observability_if_configured().await;
12323                });
12324            Ok(stream)
12325        } else {
12326            Ok(inner)
12327        }
12328    }
12329
12330    /// Streams provisional chunks and emits one authoritative response after successful root-turn finalization.
12331    pub async fn chat_stream_events<'a>(
12332        &'a self,
12333        input: &'a str,
12334    ) -> Result<Pin<Box<dyn Stream<Item = AgentStreamEvent> + Send + 'a>>> {
12335        //
12336        // The event API shares the legacy execution stream so chunk timing, cancellation, and side effects cannot diverge.
12337        //
12338        self.init_storage().await?;
12339        info!(input_len = input.len(), "Starting streaming chat events");
12340        let terminal = new_runtime_stream_terminal_slot();
12341        let mut inner = self.run_loop_stream(input, Arc::clone(&terminal));
12342        let observation_context = self.build_observation_context(None);
12343        let stream: Pin<Box<dyn Stream<Item = AgentStreamEvent> + Send + 'a>> =
12344            Box::pin(async_stream::stream! {
12345                loop {
12346                    let next = if let Some(context) = observation_context.as_ref() {
12347                        with_observation_context(context.clone(), inner.next()).await
12348                    } else {
12349                        inner.next().await
12350                    };
12351                    match next {
12352                        Some(StreamChunk::Done {}) => {
12353                            let terminal_event = { terminal.write().take() };
12354                            if let Some(response) = terminal_event {
12355                                while if let Some(context) = observation_context.as_ref() {
12356                                    with_observation_context(context.clone(), inner.next())
12357                                        .await
12358                                        .is_some()
12359                                } else {
12360                                    inner.next().await.is_some()
12361                                } {}
12362                                if observation_context.is_some() {
12363                                    self.export_observability_if_configured().await;
12364                                }
12365                                yield AgentStreamEvent::Final(response);
12366                                return;
12367                            }
12368                        }
12369                        Some(StreamChunk::Error { message }) => {
12370                            let finalized = { terminal.read().is_some() };
12371                            if finalized {
12372                                continue;
12373                            }
12374                            while if let Some(context) = observation_context.as_ref() {
12375                                with_observation_context(context.clone(), inner.next())
12376                                    .await
12377                                    .is_some()
12378                            } else {
12379                                inner.next().await.is_some()
12380                            } {}
12381                            if observation_context.is_some() {
12382                                self.export_observability_if_configured().await;
12383                            }
12384                            yield AgentStreamEvent::Chunk(StreamChunk::Error { message });
12385                            return;
12386                        }
12387                        Some(chunk) => yield AgentStreamEvent::Chunk(chunk),
12388                        None => {
12389                            if observation_context.is_some() {
12390                                self.export_observability_if_configured().await;
12391                            }
12392                            break;
12393                        }
12394                    }
12395                }
12396            });
12397        Ok(stream)
12398    }
12399}
12400
12401#[async_trait]
12402impl ToolInvoker for RuntimeAgent {
12403    async fn invoke_tool(&self, request: ToolExecutionRequest) -> Result<ToolExecutionRecord> {
12404        self.execute_tool_record(request).await
12405    }
12406}
12407
12408#[async_trait]
12409impl Agent for RuntimeAgent {
12410    async fn chat(&self, input: &str) -> Result<AgentResponse> {
12411        let result = if let Some(context) = self.build_observation_context(None) {
12412            with_observation_context(context, self.run_loop(input)).await
12413        } else {
12414            self.run_loop(input).await
12415        };
12416        self.export_observability_if_configured().await;
12417        result
12418    }
12419
12420    fn info(&self) -> AgentInfo {
12421        self.info.clone()
12422    }
12423
12424    /// Resets the runtime without leaving clarification or skill ownership behind.
12425    async fn reset(&self) -> Result<()> {
12426        self.reset_runtime_state().await
12427    }
12428}
12429
12430//
12431// Render a concurrent input template using direct minijinja.
12432// Same approach as pipeline's render_stage_template so variables are top-level.
12433//
12434// Available variables:
12435//   {{ user_input }}    - the user's actual message
12436//   {{ context.<key> }} - values from the context manager
12437//
12438/// Builds safe runtime tags for background maintenance lifecycle events.
12439fn background_maintenance_tags(
12440    label: &str,
12441    stage: &str,
12442    reason: Option<&str>,
12443    policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
12444) -> HashMap<String, String> {
12445    let mut tags = HashMap::new();
12446    tags.insert("runtime.background".to_string(), "true".to_string());
12447    tags.insert("runtime.maintenance".to_string(), label.to_string());
12448    tags.insert("runtime.maintenance_stage".to_string(), stage.to_string());
12449    if let Some(policy) = policy {
12450        tags.insert(
12451            "runtime.await_before_next_turn".to_string(),
12452            await_before_next_turn_label(policy.await_before_next_turn).to_string(),
12453        );
12454        tags.insert(
12455            "runtime.maintenance_mode".to_string(),
12456            maintenance_mode_label(policy.mode).to_string(),
12457        );
12458    }
12459    if let Some(reason) = reason {
12460        tags.insert("runtime.reason".to_string(), reason.to_string());
12461    }
12462    tags
12463}
12464
12465fn await_before_next_turn_label(policy: AwaitBeforeNextTurn) -> &'static str {
12466    match policy {
12467        AwaitBeforeNextTurn::Never => "never",
12468        AwaitBeforeNextTurn::SameActor => "same_actor",
12469        AwaitBeforeNextTurn::Always => "always",
12470    }
12471}
12472
12473fn maintenance_mode_label(mode: MaintenanceMode) -> &'static str {
12474    match mode {
12475        MaintenanceMode::InlineSerial => "inline_serial",
12476        MaintenanceMode::InlineParallel => "inline_parallel",
12477        MaintenanceMode::Background => "background",
12478    }
12479}
12480
12481/// Records a background maintenance lifecycle event when observability is enabled.
12482fn record_background_maintenance_event(
12483    manager: Option<&Arc<ObservabilityManager>>,
12484    label: &str,
12485    status: EventStatus,
12486    duration_ms: u64,
12487    stage: &str,
12488    reason: Option<String>,
12489    policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
12490) {
12491    if let Some(manager) = manager {
12492        manager.record_lifecycle_event(
12493            EventType::MemoryOperation {
12494                operation: format!("{}_background_{}", label, stage),
12495            },
12496            ObservationPurpose::Other(format!("{}_maintenance", label)),
12497            status,
12498            duration_ms,
12499            background_maintenance_tags(label, stage, reason.as_deref(), policy),
12500            None,
12501        );
12502    }
12503}
12504
12505fn effective_maintenance_mode(mode: MaintenanceMode, force_parallel: bool) -> MaintenanceMode {
12506    if force_parallel && matches!(mode, MaintenanceMode::InlineSerial) {
12507        MaintenanceMode::InlineParallel
12508    } else {
12509        mode
12510    }
12511}
12512
12513fn observation_purpose_for_process(hint: ProcessPurposeHint) -> ObservationPurpose {
12514    match hint {
12515        ProcessPurposeHint::Detect => ObservationPurpose::ProcessDetect,
12516        ProcessPurposeHint::Extract => ObservationPurpose::ProcessExtract,
12517        ProcessPurposeHint::Validate => ObservationPurpose::ProcessValidate,
12518        ProcessPurposeHint::Transform | ProcessPurposeHint::Other => {
12519            ObservationPurpose::ProcessTransform
12520        }
12521    }
12522}
12523
12524fn new_tool_resource_locks() -> ToolResourceLocks {
12525    Arc::new(RwLock::new(HashMap::new()))
12526}
12527
12528//
12529// Path-bound mutations share one conservative lock so aliases and parent-child paths cannot bypass serialization.
12530// Exact domain keys remain available for non-path side effects, and unbound effects share one fallback lock.
12531//
12532fn tool_resource_lock_keys(
12533    _canonical_id: &str,
12534    args: &Value,
12535    bindings: &ai_agents_core::ToolPolicyBindings,
12536    classification: &ai_agents_core::ToolCallClassification,
12537) -> Vec<String> {
12538    if classification.concurrency_safe {
12539        return Vec::new();
12540    }
12541
12542    let mut keys = Vec::new();
12543    let mut has_path_resource = false;
12544    for binding in &bindings.path_fields {
12545        let value = value_at_argument_path(args, &binding.field)
12546            .cloned()
12547            .or_else(|| {
12548                binding
12549                    .default_path
12550                    .as_ref()
12551                    .map(|path| Value::String(path.clone()))
12552            });
12553        if let Some(value) = value {
12554            collect_resource_strings(&value, |_| {
12555                has_path_resource = true;
12556            });
12557        }
12558    }
12559    for binding in &bindings.domain_fields {
12560        if let Some(value) = value_at_argument_path(args, &binding.field) {
12561            collect_resource_strings(value, |domain| {
12562                let normalized = if binding.is_url {
12563                    normalized_url_resource_key(domain)
12564                } else {
12565                    domain.trim().trim_end_matches('.').to_ascii_lowercase()
12566                };
12567                keys.push(format!("domain:{}", normalized));
12568            });
12569        }
12570    }
12571    for binding in &bindings.command_fields {
12572        if !matches!(binding.kind, ai_agents_core::CommandBindingKind::Cwd) {
12573            continue;
12574        }
12575        if let Some(value) = value_at_argument_path(args, &binding.field) {
12576            collect_resource_strings(value, |_| {
12577                has_path_resource = true;
12578            });
12579        }
12580    }
12581    if has_path_resource {
12582        keys.push("path-mutation:global".to_string());
12583    }
12584    if keys.is_empty() {
12585        keys.push("side-effect:unbound".to_string());
12586    }
12587    keys.sort();
12588    keys.dedup();
12589    keys
12590}
12591
12592fn value_at_argument_path<'a>(value: &'a Value, field: &str) -> Option<&'a Value> {
12593    let mut current = value;
12594    for segment in field.split('.') {
12595        if segment.is_empty() {
12596            return None;
12597        }
12598        current = current.get(segment)?;
12599    }
12600    Some(current)
12601}
12602
12603fn collect_resource_strings(value: &Value, mut collect: impl FnMut(&str)) {
12604    match value {
12605        Value::String(value) => collect(value),
12606        Value::Array(values) => {
12607            for value in values {
12608                if let Some(value) = value.as_str() {
12609                    collect(value);
12610                }
12611            }
12612        }
12613        _ => {}
12614    }
12615}
12616
12617fn normalized_url_resource_key(value: &str) -> String {
12618    let value = value.trim();
12619    let Some((scheme, remainder)) = value.split_once("://") else {
12620        return value.to_ascii_lowercase();
12621    };
12622    let authority_end = remainder.find(['/', '?', '#']).unwrap_or(remainder.len());
12623    let (authority, suffix) = remainder.split_at(authority_end);
12624    format!(
12625        "{}://{}{}",
12626        scheme.to_ascii_lowercase(),
12627        authority.to_ascii_lowercase(),
12628        suffix
12629    )
12630}
12631
12632fn render_concurrent_template(
12633    template: &str,
12634    user_input: &str,
12635    context_values: &std::collections::HashMap<String, serde_json::Value>,
12636) -> Result<String> {
12637    let mut env = minijinja::Environment::new();
12638    env.add_template("concurrent", template)
12639        .map_err(|e| AgentError::Other(format!("Concurrent template parse error: {}", e)))?;
12640
12641    let mut ctx = std::collections::BTreeMap::new();
12642    ctx.insert("user_input".to_string(), minijinja::Value::from(user_input));
12643
12644    // Expose context manager values under {{ context.<key> }}.
12645    let context_obj = minijinja::Value::from_serialize(context_values);
12646    ctx.insert("context".to_string(), context_obj);
12647
12648    let tmpl = env
12649        .get_template("concurrent")
12650        .map_err(|e| AgentError::Other(format!("Concurrent template error: {}", e)))?;
12651
12652    tmpl.render(minijinja::Value::from_serialize(&ctx))
12653        .map_err(|e| AgentError::Other(format!("Concurrent template render error: {}", e)))
12654}
12655
12656#[cfg(test)]
12657mod tests {
12658    use super::*;
12659    use crate::AgentBuilder;
12660    use ai_agents_core::{LLMChunk, LLMConfig, LLMError, LLMFeature, Tool};
12661    use ai_agents_llm::mock::MockLLMProvider;
12662    use ai_agents_skills::{SkillDefinition, SkillStep};
12663    use ai_agents_tools::{
12664        CalculatorTool, CopyPathTool, DeletePathTool, FileWriteTool, MovePathTool,
12665        WebFetchResolver, WebFetchTool, WebFetchTransport, WebFetchTransportRequest,
12666        WebFetchTransportResponse,
12667    };
12668
12669    fn mock_with_response(response: &str) -> MockLLMProvider {
12670        let mut mock = MockLLMProvider::new("test");
12671        mock.set_response(response);
12672        mock
12673    }
12674
12675    fn mock_with_responses(responses: Vec<&str>) -> MockLLMProvider {
12676        let mut mock = MockLLMProvider::new("test");
12677        mock.set_responses(responses.into_iter().map(String::from).collect(), true);
12678        mock
12679    }
12680
12681    /// Builds a two-state fixture so confirmation ownership can be invalidated by transition.
12682    fn disambiguation_state_machine(
12683        state_enabled: Option<bool>,
12684        require_confirmation: bool,
12685    ) -> Arc<StateMachine> {
12686        let definition = ai_agents_state::StateDefinition {
12687            prompt: Some("Handle the resolved request.".to_string()),
12688            disambiguation: Some(ai_agents_disambiguation::StateDisambiguationOverride {
12689                enabled: state_enabled,
12690                require_confirmation,
12691                ..Default::default()
12692            }),
12693            ..Default::default()
12694        };
12695        let review = ai_agents_state::StateDefinition {
12696            prompt: Some("Review a fresh request.".to_string()),
12697            ..Default::default()
12698        };
12699        Arc::new(
12700            StateMachine::new(ai_agents_state::StateConfig {
12701                initial: "active".to_string(),
12702                states: std::collections::HashMap::from([
12703                    ("active".to_string(), definition),
12704                    ("review".to_string(), review),
12705                ]),
12706                global_transitions: Vec::new(),
12707                fallback: None,
12708                max_no_transition: None,
12709                regenerate_on_transition: true,
12710            })
12711            .unwrap(),
12712        )
12713    }
12714
12715    /// Builds a state-aware disambiguation fixture without skills.
12716    fn state_disambiguation_agent(
12717        responses: Vec<&str>,
12718        manager_enabled: bool,
12719        state_enabled: Option<bool>,
12720        require_confirmation: bool,
12721    ) -> (RuntimeAgent, MockLLMProvider) {
12722        state_disambiguation_agent_with_skills(
12723            responses,
12724            manager_enabled,
12725            state_enabled,
12726            require_confirmation,
12727            Vec::new(),
12728        )
12729    }
12730
12731    /// Builds a state-aware disambiguation fixture with optional real skill routing.
12732    fn state_disambiguation_agent_with_skills(
12733        responses: Vec<&str>,
12734        manager_enabled: bool,
12735        state_enabled: Option<bool>,
12736        require_confirmation: bool,
12737        skills: Vec<SkillDefinition>,
12738    ) -> (RuntimeAgent, MockLLMProvider) {
12739        let mut mock = MockLLMProvider::new("state-confirmation");
12740        mock.set_responses(responses.into_iter().map(String::from).collect(), false);
12741        let observed = mock.clone();
12742        let agent = AgentBuilder::new()
12743            .system_prompt("Handle requests.")
12744            .llm(Arc::new(mock.clone()))
12745            .llm_alias("router", Arc::new(mock))
12746            .state_machine(disambiguation_state_machine(
12747                state_enabled,
12748                require_confirmation,
12749            ))
12750            .skills(skills)
12751            .build()
12752            .unwrap()
12753            .with_disambiguation(DisambiguationConfig {
12754                enabled: manager_enabled,
12755                ..Default::default()
12756            });
12757        (agent, observed)
12758    }
12759
12760    /// Defines a prompt skill whose provider call proves committed execution.
12761    fn confirmation_skill() -> SkillDefinition {
12762        SkillDefinition {
12763            id: "send_report".to_string(),
12764            description: "Send a report after clarification".to_string(),
12765            trigger: "When the user asks to send a report".to_string(),
12766            steps: vec![SkillStep::Prompt {
12767                prompt: "Execute confirmed report skill for: {{ input }}".to_string(),
12768                llm: None,
12769            }],
12770            reasoning: None,
12771            reflection: None,
12772            disambiguation: Some(ai_agents_disambiguation::SkillDisambiguationOverride {
12773                enabled: Some(true),
12774                ..Default::default()
12775            }),
12776        }
12777    }
12778
12779    /// Counts committed skill prompt calls without relying on final response wording.
12780    fn confirmation_skill_call_count(observed: &MockLLMProvider) -> usize {
12781        observed
12782            .call_history()
12783            .iter()
12784            .filter(|call| {
12785                call.messages
12786                    .iter()
12787                    .any(|message| message.content.contains("Execute confirmed report skill"))
12788            })
12789            .count()
12790    }
12791
12792    struct BlockingRuntimeConfirmationObserver {
12793        entered: tokio::sync::Barrier,
12794        release: tokio::sync::Notify,
12795    }
12796
12797    impl BlockingRuntimeConfirmationObserver {
12798        fn new() -> Self {
12799            Self {
12800                entered: tokio::sync::Barrier::new(2),
12801                release: tokio::sync::Notify::new(),
12802            }
12803        }
12804    }
12805
12806    struct ResetOnTransitionHooks {
12807        agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
12808        invoked: AtomicBool,
12809    }
12810
12811    #[async_trait]
12812    impl AgentHooks for ResetOnTransitionHooks {
12813        async fn on_state_transition(&self, _from: Option<&str>, _to: &str, _reason: &str) {
12814            if self.invoked.swap(true, Ordering::SeqCst) {
12815                return;
12816            }
12817            let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
12818            if let Some(agent) = agent {
12819                agent.reset().await.unwrap();
12820            }
12821        }
12822    }
12823
12824    impl ClarificationObserver for BlockingRuntimeConfirmationObserver {
12825        fn observe_question<'a>(
12826            &'a self,
12827            future: ClarificationQuestionFuture<'a>,
12828        ) -> ClarificationQuestionFuture<'a> {
12829            future
12830        }
12831
12832        fn observe_parse<'a>(
12833            &'a self,
12834            future: ClarificationParseFuture<'a>,
12835        ) -> ClarificationParseFuture<'a> {
12836            future
12837        }
12838
12839        fn observe_confirmation_parse<'a>(
12840            &'a self,
12841            future: ConfirmationParseFuture<'a>,
12842        ) -> ConfirmationParseFuture<'a> {
12843            Box::pin(async move {
12844                self.entered.wait().await;
12845                self.release.notified().await;
12846                future.await
12847            })
12848        }
12849    }
12850
12851    #[tokio::test]
12852    async fn state_confirmation_blocks_redispatch_until_explicit_agreement() {
12853        let (agent, observed) = state_disambiguation_agent(
12854            vec![
12855                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12856                r#"{"question":"What should I send?","options":null}"#,
12857                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12858                r#"{"question":"Should I send the report to Ada?"}"#,
12859                r#"{"status":"confirmed"}"#,
12860                "Request executed.",
12861            ],
12862            true,
12863            None,
12864            true,
12865        );
12866
12867        let clarification = agent.chat("Send it").await.unwrap();
12868        assert_eq!(clarification.content, "What should I send?");
12869        assert_eq!(observed.call_count(), 2);
12870
12871        let confirmation = agent.chat("The report to Ada").await.unwrap();
12872        assert_eq!(confirmation.content, "Should I send the report to Ada?");
12873        assert_eq!(
12874            confirmation
12875                .metadata
12876                .as_ref()
12877                .and_then(|metadata| metadata.get("disambiguation"))
12878                .and_then(|metadata| metadata.get("status"))
12879                .and_then(Value::as_str),
12880            Some("awaiting_confirmation")
12881        );
12882        assert_eq!(observed.call_count(), 4);
12883
12884        let completed = agent.chat("Yes").await.unwrap();
12885        assert_eq!(completed.content, "Request executed.");
12886        assert_eq!(observed.call_count(), 6);
12887    }
12888
12889    #[tokio::test]
12890    async fn streaming_state_confirmation_ends_the_turn_before_redispatch() {
12891        let (agent, observed) = state_disambiguation_agent(
12892            vec![
12893                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12894                r#"{"question":"What should I send?","options":null}"#,
12895                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12896                r#"{"question":"Should I send the report to Ada?"}"#,
12897                r#"{"status":"confirmed"}"#,
12898                "Request executed.",
12899            ],
12900            true,
12901            None,
12902            true,
12903        );
12904
12905        let mut clarification_stream = agent.chat_stream("Send it").await.unwrap();
12906        let mut clarification = String::new();
12907        while let Some(chunk) = clarification_stream.next().await {
12908            match chunk {
12909                StreamChunk::Content { text } => clarification.push_str(&text),
12910                StreamChunk::Done {} => break,
12911                StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
12912                _ => {}
12913            }
12914        }
12915        assert_eq!(clarification, "What should I send?");
12916        assert_eq!(observed.call_count(), 2);
12917
12918        let mut confirmation_stream = agent.chat_stream_events("The report to Ada").await.unwrap();
12919        let mut confirmation = None;
12920        while let Some(event) = confirmation_stream.next().await {
12921            match event {
12922                AgentStreamEvent::Final(response) => confirmation = Some(response),
12923                AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
12924                    panic!("unexpected stream error: {message}")
12925                }
12926                AgentStreamEvent::Chunk(_) => {}
12927            }
12928        }
12929        let confirmation = confirmation.expect("confirmation must finalize");
12930        assert_eq!(confirmation.content, "Should I send the report to Ada?");
12931        assert_eq!(
12932            confirmation
12933                .metadata
12934                .as_ref()
12935                .and_then(|metadata| metadata.get("disambiguation"))
12936                .and_then(|metadata| metadata.get("status"))
12937                .and_then(Value::as_str),
12938            Some("awaiting_confirmation")
12939        );
12940        assert_eq!(observed.call_count(), 4);
12941
12942        let mut completed_stream = agent.chat_stream("Yes").await.unwrap();
12943        let mut completed = String::new();
12944        while let Some(chunk) = completed_stream.next().await {
12945            match chunk {
12946                StreamChunk::Content { text } => completed.push_str(&text),
12947                StreamChunk::Done {} => break,
12948                StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
12949                _ => {}
12950            }
12951        }
12952        assert_eq!(completed, "Request executed.");
12953        assert_eq!(observed.call_count(), 6);
12954    }
12955
12956    /// Confirms that a pending skill remains inert until confirmation and then runs once.
12957    #[tokio::test]
12958    async fn confirmed_skill_route_executes_exactly_once() {
12959        let (agent, observed) = state_disambiguation_agent_with_skills(
12960            vec![
12961                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
12962                "send_report",
12963                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12964                r#"{"question":"What should I send?","options":null}"#,
12965                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12966                r#"{"question":"Should I send the report to Ada?"}"#,
12967                r#"{"status":"confirmed"}"#,
12968                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"resolved","what_is_unclear":[],"detected_language":"en"}"#,
12969                "Report skill executed.",
12970            ],
12971            true,
12972            None,
12973            true,
12974            vec![confirmation_skill()],
12975        );
12976
12977        let clarification = agent.chat("Send it").await.unwrap();
12978        assert_eq!(clarification.content, "What should I send?");
12979        assert_eq!(confirmation_skill_call_count(&observed), 0);
12980
12981        let confirmation = agent.chat("The report to Ada").await.unwrap();
12982        assert_eq!(confirmation.content, "Should I send the report to Ada?");
12983        assert_eq!(
12984            confirmation
12985                .metadata
12986                .as_ref()
12987                .and_then(|metadata| metadata.get("disambiguation"))
12988                .and_then(|metadata| metadata.get("status"))
12989                .and_then(Value::as_str),
12990            Some("awaiting_confirmation")
12991        );
12992        assert_eq!(confirmation_skill_call_count(&observed), 0);
12993
12994        let completed = agent.chat("Yes").await.unwrap();
12995        assert_eq!(completed.content, "Report skill executed.");
12996        assert_eq!(confirmation_skill_call_count(&observed), 1);
12997        assert!(agent.pending_skill_id.read().is_none());
12998        let messages = agent.memory.get_messages(None).await.unwrap();
12999        assert!(!messages.iter().any(|message| message.content == "Yes"));
13000    }
13001
13002    /// Preserves a second clarification response after confirmation instead of overwriting its metadata.
13003    #[tokio::test]
13004    async fn confirmed_skill_recheck_preserves_new_clarification_metadata() {
13005        let (agent, observed) = state_disambiguation_agent_with_skills(
13006            vec![
13007                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13008                "send_report",
13009                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13010                r#"{"question":"What should I send?","options":null}"#,
13011                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13012                r#"{"question":"Should I send the report to Ada?"}"#,
13013                r#"{"status":"confirmed"}"#,
13014                r#"{"is_ambiguous":true,"confidence":0.3,"ambiguity_type":"missing_parameters","reasoning":"timing missing","what_is_unclear":["timing"],"detected_language":"en"}"#,
13015                r#"{"question":"When should I send it?","options":null}"#,
13016            ],
13017            true,
13018            None,
13019            true,
13020            vec![confirmation_skill()],
13021        );
13022
13023        agent.chat("Send it").await.unwrap();
13024        agent.chat("The report to Ada").await.unwrap();
13025        let follow_up = agent.chat("Yes").await.unwrap();
13026
13027        assert_eq!(follow_up.content, "When should I send it?");
13028        let metadata = follow_up
13029            .metadata
13030            .as_ref()
13031            .and_then(|metadata| metadata.get("disambiguation"))
13032            .unwrap();
13033        assert_eq!(
13034            metadata.get("status").and_then(Value::as_str),
13035            Some("awaiting_clarification")
13036        );
13037        assert_eq!(
13038            metadata.get("skill_id").and_then(Value::as_str),
13039            Some("send_report")
13040        );
13041        assert!(metadata.get("detection").is_some());
13042        assert_eq!(confirmation_skill_call_count(&observed), 0);
13043    }
13044
13045    /// Confirms rejection clears a pending skill without executing its prompt.
13046    #[tokio::test]
13047    async fn rejected_skill_confirmation_never_executes() {
13048        let (agent, observed) = state_disambiguation_agent_with_skills(
13049            vec![
13050                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13051                "send_report",
13052                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13053                r#"{"question":"What should I send?","options":null}"#,
13054                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13055                r#"{"question":"Should I send the report to Ada?"}"#,
13056                r#"{"status":"rejected"}"#,
13057                "Confirmation rejected.",
13058            ],
13059            true,
13060            None,
13061            true,
13062            vec![confirmation_skill()],
13063        );
13064
13065        agent.chat("Send it").await.unwrap();
13066        agent.chat("The report to Ada").await.unwrap();
13067        let rejected = agent.chat("No").await.unwrap();
13068
13069        assert_eq!(rejected.content, "Confirmation rejected.");
13070        assert_eq!(confirmation_skill_call_count(&observed), 0);
13071        assert!(agent.pending_skill_id.read().is_none());
13072    }
13073
13074    /// Confirms reset clears manager and skill ownership before a later streaming turn.
13075    #[tokio::test]
13076    async fn reset_invalidates_pending_skill_confirmation_before_streaming_input() {
13077        let (agent, observed) = state_disambiguation_agent_with_skills(
13078            vec![
13079                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13080                "send_report",
13081                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13082                r#"{"question":"What should I send?","options":null}"#,
13083                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13084                r#"{"question":"Should I send the report to Ada?"}"#,
13085                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"fresh input","what_is_unclear":[],"detected_language":"en"}"#,
13086                "none",
13087                "Fresh response.",
13088            ],
13089            true,
13090            None,
13091            true,
13092            vec![confirmation_skill()],
13093        );
13094
13095        agent.chat("Send it").await.unwrap();
13096        agent.chat("The report to Ada").await.unwrap();
13097        agent.reset().await.unwrap();
13098        assert!(agent.pending_skill_id.read().is_none());
13099        assert!(
13100            !agent
13101                .disambiguation_manager()
13102                .unwrap()
13103                .has_pending_clarification()
13104                .await
13105        );
13106
13107        let mut stream = agent.chat_stream("Yes").await.unwrap();
13108        let mut content = String::new();
13109        while let Some(chunk) = stream.next().await {
13110            match chunk {
13111                StreamChunk::Content { text } => content.push_str(&text),
13112                StreamChunk::Done {} => break,
13113                StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
13114                _ => {}
13115            }
13116        }
13117
13118        assert_eq!(content, "Fresh response.");
13119        assert_eq!(confirmation_skill_call_count(&observed), 0);
13120    }
13121
13122    /// Confirms the Agent trait reset uses the same pending ownership cleanup.
13123    #[tokio::test]
13124    async fn trait_reset_clears_pending_skill_confirmation() {
13125        let (agent, _) = state_disambiguation_agent_with_skills(
13126            vec![
13127                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13128                "send_report",
13129                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13130                r#"{"question":"What should I send?","options":null}"#,
13131                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13132                r#"{"question":"Should I send the report to Ada?"}"#,
13133            ],
13134            true,
13135            None,
13136            true,
13137            vec![confirmation_skill()],
13138        );
13139
13140        agent.chat("Send it").await.unwrap();
13141        agent.chat("The report to Ada").await.unwrap();
13142        <RuntimeAgent as Agent>::reset(&agent).await.unwrap();
13143
13144        assert!(agent.pending_skill_id.read().is_none());
13145        assert!(
13146            !agent
13147                .disambiguation_manager()
13148                .unwrap()
13149                .has_pending_clarification()
13150                .await
13151        );
13152    }
13153
13154    /// Confirms a state transition cancels stale skill execution before confirmation parsing.
13155    #[tokio::test]
13156    async fn state_change_invalidates_pending_skill_confirmation() {
13157        let (agent, observed) = state_disambiguation_agent_with_skills(
13158            vec![
13159                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13160                "send_report",
13161                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13162                r#"{"question":"What should I send?","options":null}"#,
13163                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13164                r#"{"question":"Should I send the report to Ada?"}"#,
13165                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"fresh input","what_is_unclear":[],"detected_language":"en"}"#,
13166                "none",
13167                "Fresh response.",
13168            ],
13169            true,
13170            None,
13171            true,
13172            vec![confirmation_skill()],
13173        );
13174
13175        agent.chat("Send it").await.unwrap();
13176        agent.chat("The report to Ada").await.unwrap();
13177        agent.transition_to("review").await.unwrap();
13178        let cancelled = agent.chat("Yes").await.unwrap();
13179
13180        assert_eq!(cancelled.content, "Fresh response.");
13181        assert_eq!(confirmation_skill_call_count(&observed), 0);
13182        assert!(agent.pending_skill_id.read().is_none());
13183    }
13184
13185    /// Confirms reset wins over a confirmation result that was already being parsed.
13186    #[tokio::test]
13187    async fn in_flight_confirmation_cannot_redispatch_after_reset() {
13188        let (mut agent, observed) = state_disambiguation_agent_with_skills(
13189            vec![
13190                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13191                "send_report",
13192                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13193                r#"{"question":"What should I send?","options":null}"#,
13194                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13195                r#"{"question":"Should I send the report to Ada?"}"#,
13196                r#"{"status":"confirmed"}"#,
13197                "Confirmation cancelled.",
13198            ],
13199            true,
13200            None,
13201            true,
13202            vec![confirmation_skill()],
13203        );
13204        let observer = Arc::new(BlockingRuntimeConfirmationObserver::new());
13205        let manager = agent
13206            .disambiguation_manager
13207            .take()
13208            .unwrap()
13209            .with_clarification_observer(observer.clone());
13210        agent.disambiguation_manager = Some(manager);
13211        let agent = Arc::new(agent);
13212
13213        agent.chat("Send it").await.unwrap();
13214        agent.chat("The report to Ada").await.unwrap();
13215
13216        let confirming_agent = Arc::clone(&agent);
13217        let confirmation = tokio::spawn(async move { confirming_agent.chat("Yes").await });
13218        observer.entered.wait().await;
13219        agent.reset().await.unwrap();
13220        observer.release.notify_one();
13221
13222        let response = confirmation.await.unwrap().unwrap();
13223        assert_eq!(response.content, "Confirmation cancelled.");
13224        assert_eq!(confirmation_skill_call_count(&observed), 0);
13225        assert!(agent.pending_skill_id.read().is_none());
13226    }
13227
13228    /// Confirms a queued reset prevents an older terminal question from being published afterward.
13229    #[tokio::test]
13230    async fn queued_reset_prevents_stale_confirmation_question_publication() {
13231        let (agent, observed) = state_disambiguation_agent(
13232            vec![
13233                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13234                r#"{"question":"What should I send?","options":null}"#,
13235                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13236                r#"{"question":"Should I send the report to Ada?"}"#,
13237            ],
13238            true,
13239            None,
13240            true,
13241        );
13242        let agent = Arc::new(agent);
13243        agent.chat("Send it").await.unwrap();
13244
13245        let admission = agent.disambiguation_admission.write().await;
13246        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
13247        let resetting_agent = Arc::clone(&agent);
13248        let reset = tokio::spawn(async move {
13249            let _ = started_tx.send(());
13250            resetting_agent.reset().await
13251        });
13252        started_rx.await.unwrap();
13253        tokio::task::yield_now().await;
13254
13255        let responding_agent = Arc::clone(&agent);
13256        let response =
13257            tokio::spawn(async move { responding_agent.chat("The report to Ada").await });
13258        tokio::time::timeout(std::time::Duration::from_secs(2), async {
13259            while observed.call_count() < 4 {
13260                tokio::task::yield_now().await;
13261            }
13262        })
13263        .await
13264        .expect("clarification processing must reach terminal publication");
13265        drop(admission);
13266
13267        reset.await.unwrap().unwrap();
13268        let error = response.await.unwrap().unwrap_err();
13269        assert!(error.to_string().contains("ownership changed"));
13270        assert!(
13271            !agent
13272                .disambiguation_manager()
13273                .unwrap()
13274                .has_pending_clarification()
13275                .await
13276        );
13277        assert!(agent.memory.get_messages(None).await.unwrap().is_empty());
13278    }
13279
13280    /// Confirms skill clarification consumers recheck ownership before publishing returned responses.
13281    #[tokio::test]
13282    async fn queued_reset_prevents_stale_skill_clarification_publication() {
13283        let (agent, observed) = state_disambiguation_agent_with_skills(
13284            vec![
13285                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13286                "send_report",
13287                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13288                r#"{"question":"What should I send?","options":null}"#,
13289            ],
13290            true,
13291            None,
13292            true,
13293            vec![confirmation_skill()],
13294        );
13295        let agent = Arc::new(agent);
13296        let admission = agent.disambiguation_admission.write().await;
13297        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
13298        let resetting_agent = Arc::clone(&agent);
13299        let reset = tokio::spawn(async move {
13300            let _ = started_tx.send(());
13301            resetting_agent.reset().await
13302        });
13303        started_rx.await.unwrap();
13304        tokio::task::yield_now().await;
13305
13306        let responding_agent = Arc::clone(&agent);
13307        let response = tokio::spawn(async move { responding_agent.chat("Send it").await });
13308        tokio::time::timeout(std::time::Duration::from_secs(2), async {
13309            while observed.call_count() < 4 {
13310                tokio::task::yield_now().await;
13311            }
13312        })
13313        .await
13314        .expect("skill clarification must reach terminal publication");
13315        drop(admission);
13316
13317        reset.await.unwrap().unwrap();
13318        let error = response.await.unwrap().unwrap_err();
13319        assert!(error.to_string().contains("ownership changed"));
13320        assert_eq!(confirmation_skill_call_count(&observed), 0);
13321        assert!(agent.pending_skill_id.read().is_none());
13322        assert!(agent.memory.get_messages(None).await.unwrap().is_empty());
13323    }
13324
13325    /// Confirms transition hooks can reenter reset after the state commit lock is released.
13326    #[tokio::test]
13327    async fn transition_hook_can_reset_without_admission_deadlock() {
13328        let hooks = Arc::new(ResetOnTransitionHooks {
13329            agent: parking_lot::Mutex::new(None),
13330            invoked: AtomicBool::new(false),
13331        });
13332        let agent = Arc::new(
13333            AgentBuilder::new()
13334                .system_prompt("Test transition hook reentrancy.")
13335                .llm(Arc::new(mock_with_response("done")))
13336                .state_machine(disambiguation_state_machine(None, false))
13337                .build()
13338                .unwrap()
13339                .with_hooks(hooks.clone()),
13340        );
13341        *hooks.agent.lock() = Some(Arc::downgrade(&agent));
13342
13343        let transitioned = tokio::time::timeout(
13344            std::time::Duration::from_secs(2),
13345            agent.apply_transition_target("active", "review", "test transition", None),
13346        )
13347        .await
13348        .expect("transition hook reset must not deadlock")
13349        .unwrap();
13350
13351        assert!(transitioned);
13352        assert!(hooks.invoked.load(Ordering::SeqCst));
13353        assert_eq!(agent.current_state().as_deref(), Some("active"));
13354    }
13355
13356    /// Confirms only the reserved transition can run source-state exit side effects.
13357    #[tokio::test]
13358    async fn concurrent_transition_cannot_duplicate_exit_actions() {
13359        let gate = PathMutationGate::new();
13360        let active = ai_agents_state::StateDefinition {
13361            on_exit: vec![StateAction::Tool {
13362                tool: "transition_exit".to_string(),
13363                args: Some(serde_json::json!({"path": "./transition-exit.txt"})),
13364            }],
13365            ..Default::default()
13366        };
13367        let state_machine = Arc::new(
13368            StateMachine::new(ai_agents_state::StateConfig {
13369                initial: "active".to_string(),
13370                states: HashMap::from([
13371                    ("active".to_string(), active),
13372                    (
13373                        "review".to_string(),
13374                        ai_agents_state::StateDefinition::default(),
13375                    ),
13376                ]),
13377                global_transitions: Vec::new(),
13378                fallback: None,
13379                max_no_transition: None,
13380                regenerate_on_transition: true,
13381            })
13382            .unwrap(),
13383        );
13384        let agent = Arc::new(
13385            AgentBuilder::new()
13386                .system_prompt("Test transition reservation.")
13387                .llm(Arc::new(mock_with_response("done")))
13388                .tool(Arc::new(BlockingPathMutationTool {
13389                    id: "transition_exit",
13390                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
13391                    gate: gate.clone(),
13392                }))
13393                .state_machine(state_machine)
13394                .build()
13395                .unwrap(),
13396        );
13397
13398        let first_agent = Arc::clone(&agent);
13399        let first = tokio::spawn(async move { first_agent.transition_to("review").await });
13400        tokio::time::timeout(std::time::Duration::from_secs(2), gate.wait_until_entered())
13401            .await
13402            .expect("reserved transition must enter its exit action");
13403
13404        let second = tokio::time::timeout(
13405            std::time::Duration::from_secs(2),
13406            agent.transition_to("review"),
13407        )
13408        .await
13409        .expect("competing transition must fail without waiting for the exit action")
13410        .unwrap_err();
13411        assert!(second.to_string().contains("already in progress"));
13412
13413        gate.release();
13414        first.await.unwrap().unwrap();
13415        assert_eq!(agent.current_state().as_deref(), Some("review"));
13416    }
13417
13418    /// Confirms a later transition cannot overtake the committed state's enter actions.
13419    #[tokio::test]
13420    async fn concurrent_transition_cannot_overtake_enter_actions() {
13421        let gate = PathMutationGate::new();
13422        let review = ai_agents_state::StateDefinition {
13423            on_enter: vec![StateAction::Tool {
13424                tool: "transition_enter".to_string(),
13425                args: Some(serde_json::json!({"path": "./transition-enter.txt"})),
13426            }],
13427            ..Default::default()
13428        };
13429        let state_machine = Arc::new(
13430            StateMachine::new(ai_agents_state::StateConfig {
13431                initial: "active".to_string(),
13432                states: HashMap::from([
13433                    (
13434                        "active".to_string(),
13435                        ai_agents_state::StateDefinition::default(),
13436                    ),
13437                    ("review".to_string(), review),
13438                ]),
13439                global_transitions: Vec::new(),
13440                fallback: None,
13441                max_no_transition: None,
13442                regenerate_on_transition: true,
13443            })
13444            .unwrap(),
13445        );
13446        let agent = Arc::new(
13447            AgentBuilder::new()
13448                .system_prompt("Test transition lifecycle reservation.")
13449                .llm(Arc::new(mock_with_response("done")))
13450                .tool(Arc::new(BlockingPathMutationTool {
13451                    id: "transition_enter",
13452                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
13453                    gate: gate.clone(),
13454                }))
13455                .state_machine(state_machine)
13456                .build()
13457                .unwrap(),
13458        );
13459
13460        let first_agent = Arc::clone(&agent);
13461        let first = tokio::spawn(async move { first_agent.transition_to("review").await });
13462        tokio::time::timeout(std::time::Duration::from_secs(2), gate.wait_until_entered())
13463            .await
13464            .expect("committed transition must enter its destination action");
13465
13466        let second = agent.transition_to("active").await.unwrap_err();
13467        assert!(second.to_string().contains("already in progress"));
13468        assert!(agent.reset().await.is_err());
13469
13470        gate.release();
13471        first.await.unwrap().unwrap();
13472        assert_eq!(agent.current_state().as_deref(), Some("review"));
13473    }
13474
13475    /// Confirms even a same-state restore invalidates manager and skill ownership.
13476    #[tokio::test]
13477    async fn same_state_restore_invalidates_pending_skill_confirmation() {
13478        let (agent, observed) = state_disambiguation_agent_with_skills(
13479            vec![
13480                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13481                "send_report",
13482                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13483                r#"{"question":"What should I send?","options":null}"#,
13484                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13485                r#"{"question":"Should I send the report to Ada?"}"#,
13486            ],
13487            true,
13488            None,
13489            true,
13490            vec![confirmation_skill()],
13491        );
13492
13493        agent.chat("Send it").await.unwrap();
13494        agent.chat("The report to Ada").await.unwrap();
13495        let snapshot = agent.save_state().await.unwrap();
13496        assert_eq!(agent.current_state().as_deref(), Some("active"));
13497
13498        agent.restore_state(snapshot).await.unwrap();
13499
13500        assert_eq!(agent.current_state().as_deref(), Some("active"));
13501        assert!(agent.pending_skill_id.read().is_none());
13502        assert!(
13503            !agent
13504                .disambiguation_manager()
13505                .unwrap()
13506                .has_pending_clarification()
13507                .await
13508        );
13509        assert_eq!(confirmation_skill_call_count(&observed), 0);
13510    }
13511
13512    /// Confirms direct state mutation cannot hide behind a return to the same state path.
13513    #[tokio::test]
13514    async fn direct_state_generation_change_invalidates_confirmation() {
13515        let (agent, observed) = state_disambiguation_agent_with_skills(
13516            vec![
13517                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13518                "send_report",
13519                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13520                r#"{"question":"What should I send?","options":null}"#,
13521                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13522                r#"{"question":"Should I send the report to Ada?"}"#,
13523                "Confirmation cancelled.",
13524            ],
13525            true,
13526            None,
13527            true,
13528            vec![confirmation_skill()],
13529        );
13530
13531        agent.chat("Send it").await.unwrap();
13532        agent.chat("The report to Ada").await.unwrap();
13533        let state_machine = agent.state_machine().unwrap();
13534        state_machine
13535            .transition_to("review", "external test")
13536            .unwrap();
13537        state_machine
13538            .transition_to("active", "external test")
13539            .unwrap();
13540
13541        let response = agent.chat("Yes").await.unwrap();
13542
13543        assert_eq!(response.content, "Confirmation cancelled.");
13544        assert_eq!(confirmation_skill_call_count(&observed), 0);
13545        assert!(agent.pending_skill_id.read().is_none());
13546    }
13547
13548    #[tokio::test]
13549    async fn state_confirmation_does_not_add_a_question_for_clear_input() {
13550        let (agent, observed) = state_disambiguation_agent(
13551            vec![
13552                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"clear","what_is_unclear":[],"detected_language":"en"}"#,
13553                "Request executed.",
13554            ],
13555            true,
13556            None,
13557            true,
13558        );
13559
13560        let response = agent.chat("Send the report to Ada").await.unwrap();
13561
13562        assert_eq!(response.content, "Request executed.");
13563        assert_eq!(observed.call_count(), 2);
13564    }
13565
13566    #[tokio::test]
13567    async fn state_override_cannot_activate_a_disabled_top_level_manager() {
13568        let (agent, observed) =
13569            state_disambiguation_agent(vec!["Request executed."], false, Some(true), true);
13570
13571        assert!(!agent.has_disambiguation());
13572        let response = agent.chat("Send it").await.unwrap();
13573
13574        assert_eq!(response.content, "Request executed.");
13575        assert_eq!(observed.call_count(), 1);
13576    }
13577
13578    #[tokio::test]
13579    async fn native_required_choice_executes_through_the_shared_tool_path() {
13580        let mut mock = MockLLMProvider::new("native-required");
13581        mock.set_tool_choice(Some(ToolChoice::Required));
13582        mock.add_response(
13583            LLMResponse::new("", FinishReason::ToolCall)
13584                .with_tool_calls(vec![ToolCall {
13585                    id: "provider-call-1".to_string(),
13586                    name: "calculator".to_string(),
13587                    arguments: serde_json::json!({"expression": "2 + 2"}),
13588                }])
13589                .unwrap(),
13590        );
13591        mock.add_response(LLMResponse::new("The answer is 4.", FinishReason::Stop));
13592        let observed = mock.clone();
13593        let agent = AgentBuilder::new()
13594            .system_prompt("Use the calculator when needed.")
13595            .llm(Arc::new(mock))
13596            .tool(Arc::new(CalculatorTool::new()))
13597            .build()
13598            .unwrap();
13599
13600        let response = agent.chat("What is 2 + 2?").await.unwrap();
13601
13602        assert_eq!(response.content, "The answer is 4.");
13603        assert_eq!(
13604            response.tool_calls.as_ref().unwrap()[0].id,
13605            "provider-call-1"
13606        );
13607        let calls = observed.call_history();
13608        assert_eq!(calls.len(), 2);
13609        assert!(matches!(
13610            calls[0].request.as_ref().map(|request| &request.choice),
13611            Some(ToolChoice::Required)
13612        ));
13613        assert!(matches!(
13614            calls[1].request.as_ref().map(|request| &request.choice),
13615            Some(ToolChoice::Auto)
13616        ));
13617    }
13618
13619    #[tokio::test]
13620    async fn prompt_fallback_uses_one_corrective_retry() {
13621        let mut mock = MockLLMProvider::new("prompt-required");
13622        mock.set_tool_choice(Some(ToolChoice::Required));
13623        mock.set_native_tool_support(false);
13624        mock.set_responses(
13625            vec![
13626                "I can calculate that.".to_string(),
13627                r#"{"tool":"calculator","arguments":{"expression":"2 + 2"}}"#.to_string(),
13628                "The answer is 4.".to_string(),
13629            ],
13630            false,
13631        );
13632        let observed = mock.clone();
13633        let agent = AgentBuilder::new()
13634            .system_prompt("Use tools.")
13635            .llm(Arc::new(mock))
13636            .tool(Arc::new(CalculatorTool::new()))
13637            .build()
13638            .unwrap();
13639
13640        let response = agent.chat("What is 2 + 2?").await.unwrap();
13641
13642        assert_eq!(response.content, "The answer is 4.");
13643        assert_eq!(observed.call_count(), 3);
13644        let corrective = &observed.call_history()[1].messages;
13645        assert!(
13646            corrective
13647                .last()
13648                .unwrap()
13649                .content
13650                .contains("previous response")
13651        );
13652    }
13653
13654    #[tokio::test]
13655    async fn prompt_fallback_fails_after_one_noncompliant_retry() {
13656        let mut mock = MockLLMProvider::new("prompt-required-failure");
13657        mock.set_tool_choice(Some(ToolChoice::Required));
13658        mock.set_native_tool_support(false);
13659        mock.set_responses(
13660            vec!["No tool.".to_string(), "Still no tool.".to_string()],
13661            false,
13662        );
13663        let observed = mock.clone();
13664        let agent = AgentBuilder::new()
13665            .system_prompt("Use tools.")
13666            .llm(Arc::new(mock))
13667            .tool(Arc::new(CalculatorTool::new()))
13668            .build()
13669            .unwrap();
13670
13671        let error = agent.chat("What is 2 + 2?").await.unwrap_err();
13672
13673        assert!(error.to_string().contains("one corrective retry"));
13674        assert_eq!(observed.call_count(), 2);
13675    }
13676
13677    #[tokio::test]
13678    async fn specific_choice_cannot_widen_the_effective_grant() {
13679        let mut mock = MockLLMProvider::new("specific-outside-grant");
13680        mock.set_tool_choice(Some(ToolChoice::Specific("random".to_string())));
13681        let observed = mock.clone();
13682        let agent = AgentBuilder::new()
13683            .system_prompt("Use tools.")
13684            .llm(Arc::new(mock))
13685            .tool(Arc::new(CalculatorTool::new()))
13686            .build()
13687            .unwrap();
13688
13689        let error = agent.chat("Generate a value.").await.unwrap_err();
13690
13691        assert!(error.to_string().contains("is not registered"));
13692        assert_eq!(observed.call_count(), 0);
13693    }
13694
13695    #[tokio::test]
13696    async fn none_choice_exposes_no_tool_protocol() {
13697        let mut mock = MockLLMProvider::new("no-tools");
13698        mock.set_tool_choice(Some(ToolChoice::None));
13699        mock.set_response(r#"{"tool":"calculator","arguments":{"expression":"2 + 2"}}"#);
13700        let observed = mock.clone();
13701        let agent = AgentBuilder::new()
13702            .system_prompt("Answer directly.")
13703            .llm(Arc::new(mock))
13704            .tool(Arc::new(CalculatorTool::new()))
13705            .build()
13706            .unwrap();
13707
13708        let response = agent.chat("Hello").await.unwrap();
13709
13710        assert!(response.tool_calls.is_none());
13711        assert_eq!(observed.call_count(), 1);
13712        let call = observed.last_call().unwrap();
13713        assert!(call.request.is_none());
13714        assert!(
13715            call.messages
13716                .iter()
13717                .all(|message| !message.content.contains("Available tools:"))
13718        );
13719    }
13720
13721    struct RuntimeStorage {
13722        capabilities: Box<[StorageCapability]>,
13723        snapshots: RwLock<HashMap<String, AgentSnapshot>>,
13724        metadata: RwLock<HashMap<String, ai_agents_core::SessionMetadata>>,
13725        metadata_save_calls: AtomicU64,
13726        metadata_load_calls: AtomicU64,
13727        fail_metadata_save: AtomicBool,
13728        fail_metadata_load: AtomicBool,
13729    }
13730
13731    impl RuntimeStorage {
13732        fn new(capabilities: impl IntoIterator<Item = StorageCapability>) -> Self {
13733            Self {
13734                capabilities: capabilities.into_iter().collect(),
13735                snapshots: RwLock::new(HashMap::new()),
13736                metadata: RwLock::new(HashMap::new()),
13737                metadata_save_calls: AtomicU64::new(0),
13738                metadata_load_calls: AtomicU64::new(0),
13739                fail_metadata_save: AtomicBool::new(false),
13740                fail_metadata_load: AtomicBool::new(false),
13741            }
13742        }
13743    }
13744
13745    #[async_trait]
13746    impl AgentStorage for RuntimeStorage {
13747        fn supports(&self, capability: StorageCapability) -> bool {
13748            self.capabilities.contains(&capability)
13749        }
13750
13751        async fn save(&self, session_id: &str, snapshot: &AgentSnapshot) -> Result<()> {
13752            self.snapshots
13753                .write()
13754                .insert(session_id.to_string(), snapshot.clone());
13755            Ok(())
13756        }
13757
13758        async fn load(&self, session_id: &str) -> Result<Option<AgentSnapshot>> {
13759            Ok(self.snapshots.read().get(session_id).cloned())
13760        }
13761
13762        async fn delete(&self, session_id: &str) -> Result<()> {
13763            self.snapshots.write().remove(session_id);
13764            Ok(())
13765        }
13766
13767        async fn list_sessions(&self) -> Result<Vec<String>> {
13768            Ok(self.snapshots.read().keys().cloned().collect())
13769        }
13770
13771        async fn save_snapshot_with_metadata(
13772            &self,
13773            session_id: &str,
13774            snapshot: &AgentSnapshot,
13775            metadata: &ai_agents_core::SessionMetadata,
13776        ) -> Result<()> {
13777            self.metadata_save_calls.fetch_add(1, Ordering::SeqCst);
13778            if self.fail_metadata_save.load(Ordering::SeqCst) {
13779                return Err(AgentError::Persistence("metadata save failed".into()));
13780            }
13781            self.snapshots
13782                .write()
13783                .insert(session_id.to_string(), snapshot.clone());
13784            self.metadata
13785                .write()
13786                .insert(session_id.to_string(), metadata.clone());
13787            Ok(())
13788        }
13789
13790        async fn save_metadata(
13791            &self,
13792            session_id: &str,
13793            metadata: &ai_agents_core::SessionMetadata,
13794        ) -> Result<()> {
13795            self.metadata_save_calls.fetch_add(1, Ordering::SeqCst);
13796            if self.fail_metadata_save.load(Ordering::SeqCst) {
13797                return Err(AgentError::Persistence("metadata save failed".into()));
13798            }
13799            self.metadata
13800                .write()
13801                .insert(session_id.to_string(), metadata.clone());
13802            Ok(())
13803        }
13804
13805        async fn load_metadata(
13806            &self,
13807            session_id: &str,
13808        ) -> Result<Option<ai_agents_core::SessionMetadata>> {
13809            self.metadata_load_calls.fetch_add(1, Ordering::SeqCst);
13810            if self.fail_metadata_load.load(Ordering::SeqCst) {
13811                return Err(AgentError::Persistence("metadata load failed".into()));
13812            }
13813            Ok(self.metadata.read().get(session_id).cloned())
13814        }
13815    }
13816
13817    fn runtime_storage_agent() -> RuntimeAgent {
13818        AgentBuilder::new()
13819            .system_prompt("Test runtime storage integration.")
13820            .llm(Arc::new(mock_with_response("done")))
13821            .build()
13822            .unwrap()
13823    }
13824
13825    fn restore_spec(id: &str) -> crate::spec::AgentSpec {
13826        crate::spec::AgentSpec {
13827            name: id.to_string(),
13828            system_prompt: format!("Restore child {id}."),
13829            ..crate::spec::AgentSpec::default()
13830        }
13831    }
13832
13833    fn restore_entry(id: &str) -> ai_agents_core::SpawnedAgentEntry {
13834        ai_agents_core::SpawnedAgentEntry {
13835            id: id.to_string(),
13836            name: id.to_string(),
13837            spec_yaml: serde_yaml::to_string(&restore_spec(id)).unwrap(),
13838        }
13839    }
13840
13841    fn restore_spawner(
13842        storage: Arc<RuntimeStorage>,
13843        max_agents: usize,
13844    ) -> (
13845        Arc<crate::spawner::AgentSpawner>,
13846        Arc<crate::spawner::AgentRegistry>,
13847    ) {
13848        let mut llms = LLMRegistry::new();
13849        llms.register("default", Arc::new(mock_with_response("done")));
13850        (
13851            Arc::new(
13852                crate::spawner::AgentSpawner::new()
13853                    .with_shared_llms(llms)
13854                    .with_shared_storage(storage)
13855                    .with_max_agents(max_agents),
13856            ),
13857            Arc::new(crate::spawner::AgentRegistry::new()),
13858        )
13859    }
13860
13861    async fn save_restore_target(
13862        parent: &RuntimeAgent,
13863        storage: &RuntimeStorage,
13864        session_id: &str,
13865        entries: Vec<ai_agents_core::SpawnedAgentEntry>,
13866    ) {
13867        let mut snapshot = parent.save_state().await.unwrap();
13868        snapshot.spawned_agents = Some(entries);
13869        storage.save(session_id, &snapshot).await.unwrap();
13870        storage
13871            .save_metadata(session_id, &ai_agents_core::SessionMetadata::default())
13872            .await
13873            .unwrap();
13874    }
13875
13876    #[tokio::test]
13877    async fn storage_init_requires_storage_for_actor_facts() {
13878        let facts = ai_agents_facts::FactsConfig {
13879            enabled: true,
13880            ..Default::default()
13881        };
13882        let agent = runtime_storage_agent().with_facts_config(None, Some(facts));
13883
13884        let error = agent.init_storage().await.unwrap_err();
13885        assert!(matches!(
13886            error,
13887            AgentError::Config(message)
13888                if message.contains("actor facts or actor memory")
13889                    && message.contains("none is configured or injected")
13890        ));
13891    }
13892
13893    #[tokio::test]
13894    async fn storage_init_validates_actor_facts_capability() {
13895        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
13896        let actor_memory = ai_agents_facts::ActorMemoryConfig {
13897            enabled: true,
13898            ..Default::default()
13899        };
13900        let agent = runtime_storage_agent()
13901            .with_storage(storage)
13902            .with_facts_config(Some(actor_memory), None);
13903
13904        assert!(matches!(
13905            agent.init_storage().await,
13906            Err(AgentError::UnsupportedStorageCapability(
13907                StorageCapability::ActorFacts
13908            ))
13909        ));
13910    }
13911
13912    #[tokio::test]
13913    async fn blocking_chat_rejects_unsupported_required_storage() {
13914        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
13915        let facts = ai_agents_facts::FactsConfig {
13916            enabled: true,
13917            ..Default::default()
13918        };
13919        let agent = runtime_storage_agent()
13920            .with_storage(storage)
13921            .with_facts_config(None, Some(facts));
13922
13923        assert!(matches!(
13924            agent.chat("hello").await,
13925            Err(AgentError::UnsupportedStorageCapability(
13926                StorageCapability::ActorFacts
13927            ))
13928        ));
13929    }
13930
13931    #[tokio::test]
13932    async fn streaming_chat_rejects_unsupported_required_storage_before_stream_creation() {
13933        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
13934        let config = ai_agents_relationships::RelationshipConfig {
13935            enabled: true,
13936            ..Default::default()
13937        };
13938        let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
13939        let agent = runtime_storage_agent()
13940            .with_storage(storage)
13941            .with_relationships(manager);
13942
13943        assert!(matches!(
13944            agent.chat_stream("hello").await,
13945            Err(AgentError::UnsupportedStorageCapability(
13946                StorageCapability::ActorRelationships
13947            ))
13948        ));
13949    }
13950
13951    #[tokio::test]
13952    async fn storage_init_completes_facts_for_injected_storage() {
13953        let storage = Arc::new(RuntimeStorage::new([
13954            StorageCapability::Snapshot,
13955            StorageCapability::ActorFacts,
13956        ]));
13957        let facts = ai_agents_facts::FactsConfig {
13958            enabled: true,
13959            ..Default::default()
13960        };
13961        let agent = runtime_storage_agent()
13962            .with_storage(storage)
13963            .with_facts_config(None, Some(facts));
13964
13965        agent.init_storage().await.unwrap();
13966        assert!(agent.fact_store().is_some());
13967    }
13968
13969    #[tokio::test]
13970    async fn storage_init_requires_storage_for_persistent_relationships() {
13971        let config = ai_agents_relationships::RelationshipConfig {
13972            enabled: true,
13973            ..Default::default()
13974        };
13975        let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
13976        let agent = runtime_storage_agent().with_relationships(manager);
13977
13978        let error = agent.init_storage().await.unwrap_err();
13979        assert!(matches!(
13980            error,
13981            AgentError::Config(message)
13982                if message.contains("persistent relationships")
13983                    && message.contains("none is configured or injected")
13984        ));
13985    }
13986
13987    #[tokio::test]
13988    async fn storage_init_validates_persistent_relationships_capability() {
13989        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
13990        let config = ai_agents_relationships::RelationshipConfig {
13991            enabled: true,
13992            ..Default::default()
13993        };
13994        let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
13995        let agent = runtime_storage_agent()
13996            .with_storage(storage)
13997            .with_relationships(manager);
13998
13999        assert!(matches!(
14000            agent.init_storage().await,
14001            Err(AgentError::UnsupportedStorageCapability(
14002                StorageCapability::ActorRelationships
14003            ))
14004        ));
14005    }
14006
14007    #[tokio::test]
14008    async fn session_restore_updates_identity_and_clears_stale_actor_binding() {
14009        let storage = Arc::new(RuntimeStorage::new([
14010            StorageCapability::Snapshot,
14011            StorageCapability::SessionMetadata,
14012        ]));
14013        let agent = runtime_storage_agent().with_storage(storage.clone());
14014        agent.set_actor_id("old-actor").unwrap();
14015        agent.save_session("old").await.unwrap();
14016        storage
14017            .save("target", &agent.save_state().await.unwrap())
14018            .await
14019            .unwrap();
14020        storage
14021            .save_metadata("target", &ai_agents_core::SessionMetadata::default())
14022            .await
14023            .unwrap();
14024
14025        assert!(agent.load_session("target").await.unwrap());
14026
14027        assert_eq!(agent.current_session_id.read().as_deref(), Some("target"));
14028        assert_eq!(agent.actor_id(), None);
14029    }
14030
14031    #[tokio::test]
14032    async fn complete_restore_reconciles_growth_shrink_and_empty_topologies() {
14033        let storage = Arc::new(RuntimeStorage::new([
14034            StorageCapability::Snapshot,
14035            StorageCapability::SessionMetadata,
14036        ]));
14037        let (spawner, registry) = restore_spawner(storage.clone(), 3);
14038        let parent = runtime_storage_agent()
14039            .with_storage(storage.clone())
14040            .with_spawner_handles(Arc::clone(&spawner), Arc::clone(&registry));
14041
14042        for id in ["a", "b"] {
14043            let spawned = spawner
14044                .spawn_with_id(id.to_string(), restore_spec(id))
14045                .await
14046                .unwrap();
14047            spawned.agent.save_session("grow").await.unwrap();
14048            registry.register(spawned).await.unwrap();
14049        }
14050        let staged_c = crate::spawner::storage::NamespacedStorage::new(storage.clone(), "c");
14051        staged_c
14052            .save("grow", &AgentSnapshot::new("c".into()))
14053            .await
14054            .unwrap();
14055        staged_c
14056            .save_metadata("grow", &ai_agents_core::SessionMetadata::default())
14057            .await
14058            .unwrap();
14059        save_restore_target(
14060            &parent,
14061            storage.as_ref(),
14062            "grow",
14063            vec![restore_entry("a"), restore_entry("b"), restore_entry("c")],
14064        )
14065        .await;
14066
14067        assert_eq!(parent.restore_session_full("grow").await.unwrap(), 3);
14068        assert_eq!(registry.count(), 3);
14069        assert!(registry.contains("c"));
14070        assert_eq!(spawner.spawned_count(), 3);
14071
14072        for id in ["a", "b"] {
14073            registry
14074                .get(id)
14075                .unwrap()
14076                .save_session("shrink")
14077                .await
14078                .unwrap();
14079        }
14080        save_restore_target(
14081            &parent,
14082            storage.as_ref(),
14083            "shrink",
14084            vec![restore_entry("a"), restore_entry("b")],
14085        )
14086        .await;
14087
14088        assert_eq!(parent.restore_session_full("shrink").await.unwrap(), 2);
14089        assert_eq!(registry.count(), 2);
14090        assert!(!registry.contains("c"));
14091        assert_eq!(spawner.spawned_count(), 2);
14092
14093        save_restore_target(&parent, storage.as_ref(), "empty", Vec::new()).await;
14094
14095        assert_eq!(parent.restore_session_full("empty").await.unwrap(), 0);
14096        assert_eq!(registry.count(), 0);
14097        assert_eq!(spawner.spawned_count(), 0);
14098        assert_eq!(parent.current_session_id.read().as_deref(), Some("empty"));
14099    }
14100
14101    #[tokio::test]
14102    async fn storage_session_metadata_is_called_only_when_advertised() {
14103        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14104        storage.fail_metadata_save.store(true, Ordering::SeqCst);
14105        storage.fail_metadata_load.store(true, Ordering::SeqCst);
14106        let agent = runtime_storage_agent().with_storage(storage.clone());
14107
14108        agent.save_session("session").await.unwrap();
14109        assert!(agent.load_session("session").await.unwrap());
14110        assert_eq!(storage.metadata_save_calls.load(Ordering::SeqCst), 0);
14111        assert_eq!(storage.metadata_load_calls.load(Ordering::SeqCst), 0);
14112    }
14113
14114    #[cfg(feature = "sqlite")]
14115    #[tokio::test]
14116    async fn sqlite_runtime_save_filter_reopen_and_reload_stay_consistent() {
14117        let directory =
14118            std::env::temp_dir().join(format!("ai-agents-runtime-sqlite-{}", uuid::Uuid::new_v4()));
14119        let path = directory.join("sessions.sqlite");
14120        let path_string = path.to_string_lossy().into_owned();
14121        let storage = Arc::new(
14122            ai_agents_storage::SqliteStorage::new(&path_string)
14123                .await
14124                .unwrap(),
14125        );
14126        let agent = runtime_storage_agent().with_storage(storage.clone());
14127        agent.set_session_metadata(ai_agents_core::SessionMetadata {
14128            tags: vec!["initial".into()],
14129            ..Default::default()
14130        });
14131        agent.chat("persist this turn").await.unwrap();
14132        agent.save_session("session").await.unwrap();
14133
14134        agent.set_session_metadata(ai_agents_core::SessionMetadata {
14135            tags: vec!["updated".into()],
14136            ..Default::default()
14137        });
14138        agent.save_session("session").await.unwrap();
14139        assert!(
14140            agent
14141                .list_sessions_filtered(&ai_agents_core::SessionFilter {
14142                    tags: Some(vec!["initial".into()]),
14143                    ..Default::default()
14144                })
14145                .await
14146                .unwrap()
14147                .is_empty()
14148        );
14149        assert_eq!(
14150            agent
14151                .list_sessions_filtered(&ai_agents_core::SessionFilter {
14152                    tags: Some(vec!["updated".into()]),
14153                    ..Default::default()
14154                })
14155                .await
14156                .unwrap()
14157                .len(),
14158            1
14159        );
14160        drop(agent);
14161        storage.close().await;
14162        drop(storage);
14163
14164        let reopened_storage = Arc::new(
14165            ai_agents_storage::SqliteStorage::new(&path_string)
14166                .await
14167                .unwrap(),
14168        );
14169        let restored = runtime_storage_agent().with_storage(reopened_storage.clone());
14170        assert!(restored.load_session("session").await.unwrap());
14171        assert_eq!(restored.session_metadata().tags, vec!["updated"]);
14172        assert_eq!(
14173            restored.current_session_id.read().as_deref(),
14174            Some("session")
14175        );
14176        assert!(restored.save_state().await.unwrap().memory.messages.len() >= 2);
14177        assert_eq!(
14178            restored
14179                .list_sessions_filtered(&ai_agents_core::SessionFilter {
14180                    tags: Some(vec!["updated".into()]),
14181                    ..Default::default()
14182                })
14183                .await
14184                .unwrap()
14185                .len(),
14186            1
14187        );
14188
14189        drop(restored);
14190        reopened_storage.close().await;
14191        drop(reopened_storage);
14192        crate::remove_sqlite_test_directory(&directory)
14193            .await
14194            .unwrap();
14195    }
14196
14197    #[tokio::test]
14198    async fn storage_session_metadata_backend_failures_propagate() {
14199        let storage = Arc::new(RuntimeStorage::new([
14200            StorageCapability::Snapshot,
14201            StorageCapability::SessionMetadata,
14202        ]));
14203        let agent = runtime_storage_agent().with_storage(storage.clone());
14204
14205        agent.save_session("session").await.unwrap();
14206        storage
14207            .save("target", &agent.save_state().await.unwrap())
14208            .await
14209            .unwrap();
14210        storage.fail_metadata_load.store(true, Ordering::SeqCst);
14211        assert!(matches!(
14212            agent.load_session("target").await,
14213            Err(AgentError::Persistence(message)) if message == "metadata load failed"
14214        ));
14215        assert_eq!(agent.current_session_id.read().as_deref(), Some("session"));
14216
14217        storage.fail_metadata_save.store(true, Ordering::SeqCst);
14218        assert!(matches!(
14219            agent.save_session("session").await,
14220            Err(AgentError::Persistence(message)) if message == "metadata save failed"
14221        ));
14222    }
14223
14224    struct ProviderFutureDropSignal {
14225        dropped: Arc<AtomicBool>,
14226    }
14227
14228    impl Drop for ProviderFutureDropSignal {
14229        fn drop(&mut self) {
14230            self.dropped.store(true, Ordering::SeqCst);
14231        }
14232    }
14233
14234    struct BufferedLockingProvider {
14235        lock: Arc<tokio::sync::Mutex<()>>,
14236        stream_started: Arc<tokio::sync::Notify>,
14237        stream_dropped: Arc<AtomicBool>,
14238        committed_after_drop: Arc<AtomicBool>,
14239    }
14240
14241    #[async_trait]
14242    impl LLMProvider for BufferedLockingProvider {
14243        async fn complete(
14244            &self,
14245            _messages: &[ChatMessage],
14246            _config: Option<&LLMConfig>,
14247        ) -> std::result::Result<LLMResponse, LLMError> {
14248            let _guard = self.lock.lock().await;
14249            self.committed_after_drop
14250                .store(self.stream_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14251            Ok(LLMResponse::new(
14252                "Committed technical response.",
14253                FinishReason::Stop,
14254            ))
14255        }
14256
14257        async fn complete_stream(
14258            &self,
14259            _messages: &[ChatMessage],
14260            _config: Option<&LLMConfig>,
14261        ) -> std::result::Result<
14262            Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14263            LLMError,
14264        > {
14265            let _guard = self.lock.lock().await;
14266            let _drop_signal = ProviderFutureDropSignal {
14267                dropped: Arc::clone(&self.stream_dropped),
14268            };
14269            self.stream_started.notify_one();
14270            std::future::pending().await
14271        }
14272
14273        fn provider_name(&self) -> &str {
14274            "buffered-locking"
14275        }
14276
14277        fn supports(&self, _feature: LLMFeature) -> bool {
14278            false
14279        }
14280    }
14281
14282    struct PendingDropStream {
14283        dropped: Arc<AtomicBool>,
14284        dropped_notify: Arc<tokio::sync::Notify>,
14285    }
14286
14287    impl Stream for PendingDropStream {
14288        type Item = std::result::Result<LLMChunk, LLMError>;
14289
14290        fn poll_next(
14291            self: Pin<&mut Self>,
14292            _cx: &mut std::task::Context<'_>,
14293        ) -> std::task::Poll<Option<Self::Item>> {
14294            std::task::Poll::Pending
14295        }
14296    }
14297
14298    impl Drop for PendingDropStream {
14299        fn drop(&mut self) {
14300            self.dropped.store(true, Ordering::SeqCst);
14301            self.dropped_notify.notify_one();
14302        }
14303    }
14304
14305    struct EstablishedStreamProvider {
14306        stream_started: Arc<tokio::sync::Notify>,
14307        stream_dropped: Arc<AtomicBool>,
14308        stream_dropped_notify: Arc<tokio::sync::Notify>,
14309        committed_after_drop: Arc<AtomicBool>,
14310    }
14311
14312    #[async_trait]
14313    impl LLMProvider for EstablishedStreamProvider {
14314        async fn complete(
14315            &self,
14316            _messages: &[ChatMessage],
14317            _config: Option<&LLMConfig>,
14318        ) -> std::result::Result<LLMResponse, LLMError> {
14319            if !self.stream_dropped.load(Ordering::SeqCst) {
14320                self.stream_dropped_notify.notified().await;
14321            }
14322            self.committed_after_drop
14323                .store(self.stream_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14324            Ok(LLMResponse::new(
14325                "Committed technical response.",
14326                FinishReason::Stop,
14327            ))
14328        }
14329
14330        async fn complete_stream(
14331            &self,
14332            _messages: &[ChatMessage],
14333            _config: Option<&LLMConfig>,
14334        ) -> std::result::Result<
14335            Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14336            LLMError,
14337        > {
14338            self.stream_started.notify_one();
14339            Ok(Box::new(PendingDropStream {
14340                dropped: Arc::clone(&self.stream_dropped),
14341                dropped_notify: Arc::clone(&self.stream_dropped_notify),
14342            }))
14343        }
14344
14345        fn provider_name(&self) -> &str {
14346            "established-stream"
14347        }
14348
14349        fn supports(&self, _feature: LLMFeature) -> bool {
14350            false
14351        }
14352    }
14353
14354    struct FirstCallLockingProvider {
14355        lock: Arc<tokio::sync::Mutex<()>>,
14356        first_started: Arc<tokio::sync::Notify>,
14357        first_dropped: Arc<AtomicBool>,
14358        committed_after_drop: Arc<AtomicBool>,
14359        calls: AtomicU64,
14360    }
14361
14362    #[async_trait]
14363    impl LLMProvider for FirstCallLockingProvider {
14364        async fn complete(
14365            &self,
14366            _messages: &[ChatMessage],
14367            _config: Option<&LLMConfig>,
14368        ) -> std::result::Result<LLMResponse, LLMError> {
14369            let _guard = self.lock.lock().await;
14370            let call = self.calls.fetch_add(1, Ordering::SeqCst);
14371            if call == 0 {
14372                let _drop_signal = ProviderFutureDropSignal {
14373                    dropped: Arc::clone(&self.first_dropped),
14374                };
14375                self.first_started.notify_one();
14376                return std::future::pending().await;
14377            }
14378            self.committed_after_drop
14379                .store(self.first_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14380            Ok(LLMResponse::new(
14381                "Committed technical response.",
14382                FinishReason::Stop,
14383            ))
14384        }
14385
14386        async fn complete_stream(
14387            &self,
14388            _messages: &[ChatMessage],
14389            _config: Option<&LLMConfig>,
14390        ) -> std::result::Result<
14391            Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14392            LLMError,
14393        > {
14394            Err(LLMError::Other(
14395                "streaming is not used in this test".to_string(),
14396            ))
14397        }
14398
14399        fn provider_name(&self) -> &str {
14400            "first-call-locking"
14401        }
14402
14403        fn supports(&self, _feature: LLMFeature) -> bool {
14404            false
14405        }
14406    }
14407
14408    struct RoutingAfterProviderStart {
14409        provider_started: Arc<tokio::sync::Notify>,
14410    }
14411
14412    #[async_trait]
14413    impl LLMProvider for RoutingAfterProviderStart {
14414        async fn complete(
14415            &self,
14416            _messages: &[ChatMessage],
14417            _config: Option<&LLMConfig>,
14418        ) -> std::result::Result<LLMResponse, LLMError> {
14419            self.provider_started.notified().await;
14420            Ok(LLMResponse::new("1", FinishReason::Stop))
14421        }
14422
14423        async fn complete_stream(
14424            &self,
14425            _messages: &[ChatMessage],
14426            _config: Option<&LLMConfig>,
14427        ) -> std::result::Result<
14428            Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14429            LLMError,
14430        > {
14431            Err(LLMError::Other(
14432                "streaming is not used in this test".to_string(),
14433            ))
14434        }
14435
14436        fn provider_name(&self) -> &str {
14437            "routing-after-start"
14438        }
14439
14440        fn supports(&self, _feature: LLMFeature) -> bool {
14441            false
14442        }
14443    }
14444
14445    /// Test hook that counts completed responses.
14446    struct ResponseCountingHooks {
14447        responses: Arc<std::sync::atomic::AtomicUsize>,
14448    }
14449
14450    /// Test tool that returns the execution context it received.
14451    struct ContextEchoTool;
14452
14453    #[async_trait]
14454    impl ai_agents_core::Tool for ContextEchoTool {
14455        fn id(&self) -> &str {
14456            "context_echo"
14457        }
14458
14459        fn name(&self) -> &str {
14460            "Context Echo"
14461        }
14462
14463        fn description(&self) -> &str {
14464            "Returns selected execution context fields."
14465        }
14466
14467        fn input_schema(&self) -> Value {
14468            serde_json::json!({"type": "object"})
14469        }
14470
14471        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14472            ai_agents_core::ToolPolicyBindings {
14473                path_fields: vec![ai_agents_core::PathPolicyBinding::read("path")],
14474                result_limit_fields: vec![ai_agents_core::ResultLimitBinding::new(
14475                    "max_results",
14476                    ai_agents_core::ResultLimitKind::MaxResults,
14477                )],
14478                ..Default::default()
14479            }
14480        }
14481
14482        async fn execute(
14483            &self,
14484            _args: Value,
14485            ctx: ai_agents_core::ToolExecutionContext,
14486        ) -> ToolResult {
14487            ToolResult::ok(
14488                serde_json::json!({
14489                    "requested_name": ctx.requested_name,
14490                    "canonical_id": ctx.canonical_id,
14491                    "display_name": ctx.display_name,
14492                    "max_results": ctx.limits.max_results,
14493                    "custom_config": ctx.custom_config,
14494                })
14495                .to_string(),
14496            )
14497        }
14498    }
14499
14500    /// Test tool that stays active long enough for runtime cancellation.
14501    struct SlowTool;
14502
14503    /// Test tool that fails once and must not be retried for writes.
14504    struct FlakyWriteTool {
14505        calls: Arc<std::sync::atomic::AtomicUsize>,
14506    }
14507
14508    /// Test tool that tracks concurrent execution on one path.
14509    struct LockedWriteTool {
14510        active: Arc<std::sync::atomic::AtomicUsize>,
14511        max_active: Arc<std::sync::atomic::AtomicUsize>,
14512    }
14513
14514    struct MultiResourceWriteTool {
14515        active: Arc<std::sync::atomic::AtomicUsize>,
14516        max_active: Arc<std::sync::atomic::AtomicUsize>,
14517    }
14518
14519    #[derive(Clone)]
14520    struct PathMutationGate {
14521        entered: Arc<AtomicBool>,
14522        entered_notify: Arc<tokio::sync::Notify>,
14523        release: Arc<tokio::sync::Notify>,
14524    }
14525
14526    impl PathMutationGate {
14527        fn new() -> Self {
14528            Self {
14529                entered: Arc::new(AtomicBool::new(false)),
14530                entered_notify: Arc::new(tokio::sync::Notify::new()),
14531                release: Arc::new(tokio::sync::Notify::new()),
14532            }
14533        }
14534
14535        async fn wait_until_entered(&self) {
14536            if !self.entered.load(Ordering::SeqCst) {
14537                self.entered_notify.notified().await;
14538            }
14539        }
14540
14541        fn release(&self) {
14542            self.release.notify_one();
14543        }
14544    }
14545
14546    struct BlockingPathMutationTool {
14547        id: &'static str,
14548        path_fields: Vec<ai_agents_core::PathPolicyBinding>,
14549        gate: PathMutationGate,
14550    }
14551
14552    struct NoBindingWriteTool {
14553        active: Arc<std::sync::atomic::AtomicUsize>,
14554        max_active: Arc<std::sync::atomic::AtomicUsize>,
14555    }
14556
14557    struct RecoveryTestTool {
14558        id: String,
14559        succeeds: bool,
14560        calls: Arc<std::sync::atomic::AtomicUsize>,
14561    }
14562
14563    struct BlockingApprovalHandler {
14564        entered: Arc<tokio::sync::Barrier>,
14565        release: Arc<tokio::sync::Notify>,
14566        result: ApprovalResult,
14567    }
14568
14569    struct RuntimeWebFetchTransport {
14570        calls: Arc<std::sync::atomic::AtomicUsize>,
14571    }
14572
14573    struct RuntimeWebFetchResolver;
14574
14575    struct ReentrantToolHooks {
14576        agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
14577        invoked: AtomicBool,
14578        nested_success: AtomicBool,
14579    }
14580
14581    #[async_trait]
14582    impl ai_agents_core::Tool for SlowTool {
14583        fn id(&self) -> &str {
14584            "slow"
14585        }
14586
14587        fn name(&self) -> &str {
14588            "Slow"
14589        }
14590
14591        fn description(&self) -> &str {
14592            "Waits until cancelled or timed out."
14593        }
14594
14595        fn input_schema(&self) -> Value {
14596            serde_json::json!({"type": "object"})
14597        }
14598
14599        async fn execute(
14600            &self,
14601            _args: Value,
14602            _ctx: ai_agents_core::ToolExecutionContext,
14603        ) -> ToolResult {
14604            tokio::time::sleep(std::time::Duration::from_secs(5)).await;
14605            ToolResult::ok("done")
14606        }
14607    }
14608
14609    #[async_trait]
14610    impl ai_agents_core::Tool for FlakyWriteTool {
14611        fn id(&self) -> &str {
14612            "flaky_write"
14613        }
14614
14615        fn name(&self) -> &str {
14616            "Flaky Write"
14617        }
14618
14619        fn description(&self) -> &str {
14620            "Fails on the first write attempt."
14621        }
14622
14623        fn input_schema(&self) -> Value {
14624            serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
14625        }
14626
14627        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14628            ai_agents_core::ToolPolicyBindings {
14629                path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
14630                ..Default::default()
14631            }
14632        }
14633
14634        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14635            ai_agents_core::ToolSafetyMetadata {
14636                read_only: false,
14637                concurrency_safe: false,
14638                operation: ai_agents_core::ToolOperationKind::Write,
14639                side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
14640                requires_network: false,
14641                destructive: false,
14642                open_world: false,
14643                host_dependent: false,
14644                requires_user_interaction: false,
14645                supports_cancellation: true,
14646                default_requires_approval: false,
14647                should_defer_schema: false,
14648                max_output_chars: Some(1024),
14649                max_result_size_chars: Some(1024),
14650            }
14651        }
14652
14653        fn classify_call(&self, _args: &Value) -> ai_agents_core::ToolCallClassification {
14654            let mut classification =
14655                ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
14656            classification.safely_retryable = false;
14657            classification
14658        }
14659
14660        async fn execute(
14661            &self,
14662            _args: Value,
14663            _ctx: ai_agents_core::ToolExecutionContext,
14664        ) -> ToolResult {
14665            let call = self.calls.fetch_add(1, Ordering::SeqCst);
14666            if call == 0 {
14667                ToolResult::error("first failure")
14668            } else {
14669                ToolResult::ok("second success")
14670            }
14671        }
14672    }
14673
14674    #[async_trait]
14675    impl ai_agents_core::Tool for LockedWriteTool {
14676        fn id(&self) -> &str {
14677            "locked_write"
14678        }
14679
14680        fn name(&self) -> &str {
14681            "Locked Write"
14682        }
14683
14684        fn description(&self) -> &str {
14685            "Tracks concurrent execution on one resource."
14686        }
14687
14688        fn input_schema(&self) -> Value {
14689            serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
14690        }
14691
14692        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14693            ai_agents_core::ToolPolicyBindings {
14694                path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
14695                ..Default::default()
14696            }
14697        }
14698
14699        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14700            ai_agents_core::ToolSafetyMetadata {
14701                read_only: false,
14702                concurrency_safe: false,
14703                operation: ai_agents_core::ToolOperationKind::Write,
14704                side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
14705                requires_network: false,
14706                destructive: false,
14707                open_world: false,
14708                host_dependent: false,
14709                requires_user_interaction: false,
14710                supports_cancellation: true,
14711                default_requires_approval: false,
14712                should_defer_schema: false,
14713                max_output_chars: Some(1024),
14714                max_result_size_chars: Some(1024),
14715            }
14716        }
14717
14718        async fn execute(
14719            &self,
14720            _args: Value,
14721            _ctx: ai_agents_core::ToolExecutionContext,
14722        ) -> ToolResult {
14723            let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
14724            loop {
14725                let current_max = self.max_active.load(Ordering::SeqCst);
14726                if active <= current_max {
14727                    break;
14728                }
14729                if self
14730                    .max_active
14731                    .compare_exchange(current_max, active, Ordering::SeqCst, Ordering::SeqCst)
14732                    .is_ok()
14733                {
14734                    break;
14735                }
14736            }
14737            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
14738            self.active.fetch_sub(1, Ordering::SeqCst);
14739            ToolResult::ok("done")
14740        }
14741    }
14742
14743    #[async_trait]
14744    impl ai_agents_core::Tool for MultiResourceWriteTool {
14745        fn id(&self) -> &str {
14746            "multi_resource_write"
14747        }
14748
14749        fn name(&self) -> &str {
14750            "Multi Resource Write"
14751        }
14752
14753        fn description(&self) -> &str {
14754            "Tracks concurrent execution across source and destination resources."
14755        }
14756
14757        fn input_schema(&self) -> Value {
14758            serde_json::json!({"type": "object"})
14759        }
14760
14761        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14762            ai_agents_core::ToolPolicyBindings {
14763                path_fields: vec![
14764                    ai_agents_core::PathPolicyBinding::read_write("source_path"),
14765                    ai_agents_core::PathPolicyBinding::write("destination_path"),
14766                ],
14767                ..Default::default()
14768            }
14769        }
14770
14771        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14772            LockedWriteTool {
14773                active: Arc::clone(&self.active),
14774                max_active: Arc::clone(&self.max_active),
14775            }
14776            .safety_metadata()
14777        }
14778
14779        async fn execute(
14780            &self,
14781            _args: Value,
14782            _ctx: ai_agents_core::ToolExecutionContext,
14783        ) -> ToolResult {
14784            let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
14785            self.max_active.fetch_max(active, Ordering::SeqCst);
14786            tokio::time::sleep(std::time::Duration::from_millis(75)).await;
14787            self.active.fetch_sub(1, Ordering::SeqCst);
14788            ToolResult::ok("done")
14789        }
14790    }
14791
14792    #[async_trait]
14793    impl ai_agents_core::Tool for BlockingPathMutationTool {
14794        fn id(&self) -> &str {
14795            self.id
14796        }
14797
14798        fn name(&self) -> &str {
14799            self.id
14800        }
14801
14802        fn description(&self) -> &str {
14803            "Blocks a path mutation until the test releases it."
14804        }
14805
14806        fn input_schema(&self) -> Value {
14807            serde_json::json!({"type": "object"})
14808        }
14809
14810        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14811            ai_agents_core::ToolPolicyBindings {
14812                path_fields: self.path_fields.clone(),
14813                ..Default::default()
14814            }
14815        }
14816
14817        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14818            ai_agents_core::ToolSafetyMetadata {
14819                read_only: false,
14820                concurrency_safe: false,
14821                operation: ai_agents_core::ToolOperationKind::Write,
14822                side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
14823                requires_network: false,
14824                destructive: false,
14825                open_world: false,
14826                host_dependent: false,
14827                requires_user_interaction: false,
14828                supports_cancellation: true,
14829                default_requires_approval: false,
14830                should_defer_schema: false,
14831                max_output_chars: Some(1024),
14832                max_result_size_chars: Some(1024),
14833            }
14834        }
14835
14836        async fn execute(
14837            &self,
14838            _args: Value,
14839            _ctx: ai_agents_core::ToolExecutionContext,
14840        ) -> ToolResult {
14841            self.gate.entered.store(true, Ordering::SeqCst);
14842            self.gate.entered_notify.notify_one();
14843            self.gate.release.notified().await;
14844            ToolResult::ok("done")
14845        }
14846    }
14847
14848    #[async_trait]
14849    impl ai_agents_core::Tool for NoBindingWriteTool {
14850        fn id(&self) -> &str {
14851            "no_binding_write"
14852        }
14853
14854        fn name(&self) -> &str {
14855            "No Binding Write"
14856        }
14857
14858        fn description(&self) -> &str {
14859            "Tracks concurrent execution without resource bindings."
14860        }
14861
14862        fn input_schema(&self) -> Value {
14863            serde_json::json!({"type": "object"})
14864        }
14865
14866        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14867            LockedWriteTool {
14868                active: Arc::clone(&self.active),
14869                max_active: Arc::clone(&self.max_active),
14870            }
14871            .safety_metadata()
14872        }
14873
14874        async fn execute(
14875            &self,
14876            _args: Value,
14877            _ctx: ai_agents_core::ToolExecutionContext,
14878        ) -> ToolResult {
14879            let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
14880            self.max_active.fetch_max(active, Ordering::SeqCst);
14881            tokio::time::sleep(std::time::Duration::from_millis(75)).await;
14882            self.active.fetch_sub(1, Ordering::SeqCst);
14883            ToolResult::ok("done")
14884        }
14885    }
14886
14887    #[async_trait]
14888    impl ai_agents_core::Tool for RecoveryTestTool {
14889        fn id(&self) -> &str {
14890            &self.id
14891        }
14892
14893        fn name(&self) -> &str {
14894            &self.id
14895        }
14896
14897        fn description(&self) -> &str {
14898            "Records recovery execution and returns a configured result."
14899        }
14900
14901        fn input_schema(&self) -> Value {
14902            serde_json::json!({"type": "object"})
14903        }
14904
14905        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14906            ai_agents_core::ToolPolicyBindings {
14907                path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
14908                ..Default::default()
14909            }
14910        }
14911
14912        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14913            ai_agents_core::ToolSafetyMetadata {
14914                read_only: false,
14915                concurrency_safe: false,
14916                operation: ai_agents_core::ToolOperationKind::Write,
14917                side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
14918                requires_network: false,
14919                destructive: false,
14920                open_world: false,
14921                host_dependent: false,
14922                requires_user_interaction: false,
14923                supports_cancellation: true,
14924                default_requires_approval: false,
14925                should_defer_schema: false,
14926                max_output_chars: Some(1024),
14927                max_result_size_chars: Some(1024),
14928            }
14929        }
14930
14931        async fn execute(
14932            &self,
14933            _args: Value,
14934            _ctx: ai_agents_core::ToolExecutionContext,
14935        ) -> ToolResult {
14936            self.calls.fetch_add(1, Ordering::SeqCst);
14937            if self.succeeds {
14938                ToolResult::ok(format!("{} succeeded", self.id))
14939            } else {
14940                ToolResult::error(format!("{} failed", self.id))
14941            }
14942        }
14943    }
14944
14945    #[async_trait]
14946    impl WebFetchTransport for RuntimeWebFetchTransport {
14947        /// Rejects unvalidated transport calls in the shared-executor fixture.
14948        async fn send(
14949            &self,
14950            _request: WebFetchTransportRequest,
14951        ) -> std::result::Result<WebFetchTransportResponse, String> {
14952            Err("validated addresses are required".to_string())
14953        }
14954
14955        /// Records transport only after runtime and tool policy validation complete.
14956        async fn send_validated(
14957            &self,
14958            _request: WebFetchTransportRequest,
14959            _addresses: &[std::net::SocketAddr],
14960        ) -> std::result::Result<WebFetchTransportResponse, String> {
14961            self.calls.fetch_add(1, Ordering::SeqCst);
14962            Ok(WebFetchTransportResponse {
14963                status: 200,
14964                content_type: Some("text/plain".to_string()),
14965                location: None,
14966                body: b"approved".to_vec(),
14967            })
14968        }
14969    }
14970
14971    #[async_trait]
14972    impl WebFetchResolver for RuntimeWebFetchResolver {
14973        /// Returns one public fixture address without external DNS.
14974        async fn resolve(
14975            &self,
14976            _host: &str,
14977            _port: u16,
14978        ) -> std::result::Result<Vec<std::net::IpAddr>, String> {
14979            Ok(vec![std::net::IpAddr::V4(std::net::Ipv4Addr::new(
14980                93, 184, 216, 34,
14981            ))])
14982        }
14983    }
14984
14985    #[async_trait]
14986    impl ApprovalHandler for BlockingApprovalHandler {
14987        async fn request_approval(
14988            &self,
14989            _request: ai_agents_hitl::ApprovalRequest,
14990        ) -> ApprovalResult {
14991            self.entered.wait().await;
14992            self.release.notified().await;
14993            self.result.clone()
14994        }
14995    }
14996
14997    #[async_trait]
14998    impl AgentHooks for ReentrantToolHooks {
14999        async fn on_tool_complete(&self, tool: &str, _result: &ToolResult, _duration_ms: u64) {
15000            if tool != "reentrant_write" || self.invoked.swap(true, Ordering::SeqCst) {
15001                return;
15002            }
15003            let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
15004            if let Some(agent) = agent {
15005                let result = agent
15006                    .invoke_tool(ToolExecutionRequest::new(
15007                        "nested-hook-call",
15008                        "reentrant_write",
15009                        serde_json::json!({"path": "./hook.txt"}),
15010                        ToolCallSource::Manual,
15011                    ))
15012                    .await;
15013                self.nested_success
15014                    .store(result.is_ok_and(|record| record.success), Ordering::SeqCst);
15015            }
15016        }
15017    }
15018
15019    #[async_trait]
15020    impl AgentHooks for ResponseCountingHooks {
15021        async fn on_response(&self, _response: &AgentResponse) {
15022            self.responses.fetch_add(1, Ordering::SeqCst);
15023        }
15024    }
15025
15026    struct ApprovalRecordingHooks {
15027        events: parking_lot::Mutex<Vec<String>>,
15028    }
15029
15030    impl ApprovalRecordingHooks {
15031        fn new() -> Self {
15032            Self {
15033                events: parking_lot::Mutex::new(Vec::new()),
15034            }
15035        }
15036
15037        fn events(&self) -> Vec<String> {
15038            self.events.lock().clone()
15039        }
15040    }
15041
15042    #[async_trait]
15043    impl AgentHooks for ApprovalRecordingHooks {
15044        async fn on_approval_result(&self, request_id: &str, result: &ApprovalResult) {
15045            self.events.lock().push(format!(
15046                "raw:{}:{}",
15047                request_id,
15048                approval_result_name(result)
15049            ));
15050        }
15051
15052        async fn on_approval_resolved(
15053            &self,
15054            request: &ai_agents_hitl::ApprovalRequest,
15055            raw_result: &ApprovalResult,
15056            outcome: &ApprovalResolvedOutcome,
15057        ) {
15058            self.events.lock().push(format!(
15059                "resolved:{}:{}:{}",
15060                request.id,
15061                approval_result_name(raw_result),
15062                approval_outcome_name(outcome)
15063            ));
15064        }
15065    }
15066
15067    fn approval_result_name(result: &ApprovalResult) -> &'static str {
15068        match result {
15069            ApprovalResult::Approved => "approved",
15070            ApprovalResult::Rejected { .. } => "rejected",
15071            ApprovalResult::Modified { .. } => "modified",
15072            ApprovalResult::Timeout => "timeout",
15073        }
15074    }
15075
15076    fn approval_outcome_name(outcome: &ApprovalResolvedOutcome) -> &'static str {
15077        match outcome {
15078            ApprovalResolvedOutcome::Approved => "approved",
15079            ApprovalResolvedOutcome::Rejected { .. } => "rejected",
15080            ApprovalResolvedOutcome::Modified { .. } => "modified",
15081            ApprovalResolvedOutcome::Error { .. } => "error",
15082        }
15083    }
15084
15085    fn assert_correlated_approval_events(
15086        events: &[String],
15087        raw_status: &str,
15088        outcome_status: &str,
15089    ) {
15090        assert_eq!(events.len(), 2);
15091        let raw: Vec<_> = events[0].split(':').collect();
15092        let resolved: Vec<_> = events[1].split(':').collect();
15093        assert_eq!(raw[0], "raw");
15094        assert_eq!(resolved[0], "resolved");
15095        assert_eq!(raw[1], resolved[1]);
15096        assert_eq!(raw[2], raw_status);
15097        assert_eq!(resolved[2], raw_status);
15098        assert_eq!(resolved[3], outcome_status);
15099    }
15100
15101    fn approval_security_config(policy_enabled: bool) -> ToolSecurityConfig {
15102        let mut security = ToolSecurityConfig {
15103            enabled: true,
15104            fail_closed: true,
15105            ..Default::default()
15106        };
15107        let policy = ai_agents_tools::ToolPolicyConfig {
15108            enabled: policy_enabled,
15109            write_paths: vec![".".to_string()],
15110            require_confirmation: true,
15111            ..Default::default()
15112        };
15113        security.tools.insert("locked_write".to_string(), policy);
15114        security
15115    }
15116
15117    struct MutationTestWorkspace {
15118        root: std::path::PathBuf,
15119    }
15120
15121    impl MutationTestWorkspace {
15122        fn new() -> Self {
15123            let root = std::env::temp_dir().join(format!(
15124                "ai-agents-runtime-mutation-{}",
15125                uuid::Uuid::new_v4()
15126            ));
15127            std::fs::create_dir_all(&root).unwrap();
15128            Self { root }
15129        }
15130    }
15131
15132    impl Drop for MutationTestWorkspace {
15133        fn drop(&mut self) {
15134            let _ = std::fs::remove_dir_all(&self.root);
15135        }
15136    }
15137
15138    async fn wait_for_resource_lock_strong_count(locks: &ToolResourceLocks, minimum: usize) {
15139        tokio::time::timeout(std::time::Duration::from_secs(2), async {
15140            loop {
15141                let strong_count = locks
15142                    .read()
15143                    .get("path-mutation:global")
15144                    .map_or(0, |lock| lock.strong_count());
15145                if strong_count >= minimum {
15146                    break;
15147                }
15148                tokio::task::yield_now().await;
15149            }
15150        })
15151        .await
15152        .expect("path mutation call did not reach the shared lock");
15153    }
15154
15155    async fn assert_path_mutation_pair_serialized(
15156        first_id: &'static str,
15157        first_fields: Vec<ai_agents_core::PathPolicyBinding>,
15158        first_args: Value,
15159        second_id: &'static str,
15160        second_fields: Vec<ai_agents_core::PathPolicyBinding>,
15161        second_args: Value,
15162    ) {
15163        let locks = new_tool_resource_locks();
15164        let first_gate = PathMutationGate::new();
15165        let second_gate = PathMutationGate::new();
15166        second_gate.release();
15167        let agent = Arc::new(
15168            AgentBuilder::new()
15169                .system_prompt("Test global path mutation locking.")
15170                .llm(Arc::new(mock_with_response("done")))
15171                .tool(Arc::new(BlockingPathMutationTool {
15172                    id: first_id,
15173                    path_fields: first_fields,
15174                    gate: first_gate.clone(),
15175                }))
15176                .tool(Arc::new(BlockingPathMutationTool {
15177                    id: second_id,
15178                    path_fields: second_fields,
15179                    gate: second_gate.clone(),
15180                }))
15181                .build()
15182                .unwrap()
15183                .with_shared_resource_locks(Arc::clone(&locks)),
15184        );
15185
15186        let first = {
15187            let agent = Arc::clone(&agent);
15188            tokio::spawn(async move {
15189                agent
15190                    .invoke_tool(ToolExecutionRequest::new(
15191                        format!("{}-first", first_id),
15192                        first_id,
15193                        first_args,
15194                        ToolCallSource::Manual,
15195                    ))
15196                    .await
15197                    .unwrap()
15198            })
15199        };
15200        first_gate.wait_until_entered().await;
15201
15202        let second = {
15203            let agent = Arc::clone(&agent);
15204            tokio::spawn(async move {
15205                agent
15206                    .invoke_tool(ToolExecutionRequest::new(
15207                        format!("{}-second", second_id),
15208                        second_id,
15209                        second_args,
15210                        ToolCallSource::Manual,
15211                    ))
15212                    .await
15213                    .unwrap()
15214            })
15215        };
15216        wait_for_resource_lock_strong_count(&locks, 2).await;
15217        assert!(!second_gate.entered.load(Ordering::SeqCst));
15218        assert!(!second.is_finished());
15219
15220        first_gate.release();
15221        let (first, second) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
15222            tokio::join!(first, second)
15223        })
15224        .await
15225        .expect("serialized path mutation calls did not finish");
15226        assert!(first.unwrap().success);
15227        assert!(second.unwrap().success);
15228        assert!(second_gate.entered.load(Ordering::SeqCst));
15229        assert!(locks.read().is_empty());
15230    }
15231
15232    #[derive(Clone, Copy)]
15233    enum MutationDenial {
15234        Policy,
15235        Approval,
15236    }
15237
15238    fn mutation_denial_security_config(
15239        tool_id: &str,
15240        workspace: &std::path::Path,
15241        denial: MutationDenial,
15242    ) -> ToolSecurityConfig {
15243        let workspace = workspace.to_string_lossy().into_owned();
15244        let mut policy = ai_agents_tools::ToolPolicyConfig {
15245            read_paths: vec![workspace.clone()],
15246            write_paths: vec![workspace.clone()],
15247            ..Default::default()
15248        };
15249        match denial {
15250            MutationDenial::Policy => policy.blocked_paths = vec![workspace],
15251            MutationDenial::Approval => policy.require_confirmation = true,
15252        }
15253
15254        let mut security = ToolSecurityConfig {
15255            enabled: true,
15256            fail_closed: true,
15257            ..Default::default()
15258        };
15259        security.tools.insert(tool_id.to_string(), policy);
15260        security
15261    }
15262
15263    async fn assert_path_mutation_denied(tool: Arc<dyn Tool>, denial: MutationDenial) {
15264        let workspace = MutationTestWorkspace::new();
15265        let tool_id = tool.id().to_string();
15266        let preserved = workspace.root.join(format!("{}-preserved.txt", tool_id));
15267        let destination = workspace.root.join(format!("{}-destination.txt", tool_id));
15268        std::fs::write(&preserved, "preserved").unwrap();
15269        let arguments = match tool_id.as_str() {
15270            "copy_path" | "move_path" => serde_json::json!({
15271                "source_path": preserved.to_string_lossy(),
15272                "destination_path": destination.to_string_lossy(),
15273                "dry_run": false
15274            }),
15275            "delete_path" => serde_json::json!({
15276                "path": preserved.to_string_lossy(),
15277                "recursive": false,
15278                "dry_run": false
15279            }),
15280            _ => panic!("unsupported mutation tool: {}", tool_id),
15281        };
15282        let security = mutation_denial_security_config(&tool_id, &workspace.root, denial);
15283        let builder = AgentBuilder::new()
15284            .system_prompt("Test mutation denial.")
15285            .llm(Arc::new(mock_with_response("done")))
15286            .tool(tool)
15287            .tool_security(ToolSecurityEngine::new(security));
15288        let builder = match denial {
15289            MutationDenial::Policy => builder,
15290            MutationDenial::Approval => builder
15291                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
15292                .approval_handler(Arc::new(RejectAllHandler::new())),
15293        };
15294        let agent = builder.build().unwrap();
15295
15296        let record = agent
15297            .invoke_tool(ToolExecutionRequest::new(
15298                format!("{}-denied", tool_id),
15299                tool_id.clone(),
15300                arguments,
15301                ToolCallSource::Manual,
15302            ))
15303            .await
15304            .unwrap();
15305
15306        assert!(!record.executed, "{} must not be invoked", tool_id);
15307        assert!(!record.success);
15308        match denial {
15309            MutationDenial::Policy => {
15310                assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
15311                assert!(record.approval.as_ref().is_some_and(|approval| matches!(
15312                    &approval.status,
15313                    ToolApprovalStatus::NotRequired
15314                )));
15315            }
15316            MutationDenial::Approval => {
15317                assert_eq!(record.policy.outcome, PermissionOutcome::RequiresApproval);
15318                assert!(record.approval.as_ref().is_some_and(|approval| matches!(
15319                    &approval.status,
15320                    ToolApprovalStatus::Rejected
15321                )));
15322            }
15323        }
15324        assert_eq!(std::fs::read_to_string(&preserved).unwrap(), "preserved");
15325        assert!(!destination.exists());
15326    }
15327
15328    fn recovery_manager_with_fallbacks(
15329        fallbacks: impl IntoIterator<Item = (String, String)>,
15330    ) -> RecoveryManager {
15331        use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
15332
15333        let per_tool = fallbacks
15334            .into_iter()
15335            .map(|(tool, fallback_tool)| {
15336                (
15337                    tool,
15338                    ToolRetryConfig {
15339                        max_retries: 0,
15340                        timeout_ms: Some(1_000),
15341                        on_failure: ToolFailureAction::Fallback { fallback_tool },
15342                    },
15343                )
15344            })
15345            .collect();
15346        RecoveryManager::new(ErrorRecoveryConfig {
15347            tools: ToolRecoveryConfig {
15348                per_tool,
15349                ..Default::default()
15350            },
15351            ..Default::default()
15352        })
15353    }
15354
15355    fn approval_check() -> HITLCheckResult {
15356        HITLCheckResult::required(
15357            ApprovalTrigger::tool("test", serde_json::json!({})),
15358            HashMap::new(),
15359            "Approve?",
15360            None,
15361        )
15362    }
15363
15364    fn agent_with_approval_result(
15365        raw_result: ApprovalResult,
15366        timeout_action: TimeoutAction,
15367        hooks: Arc<ApprovalRecordingHooks>,
15368    ) -> RuntimeAgent {
15369        use ai_agents_hitl::{CallbackHandler, HITLConfig};
15370
15371        let config = HITLConfig {
15372            on_timeout: timeout_action,
15373            ..Default::default()
15374        };
15375        let handler = CallbackHandler::new(move |_| raw_result.clone());
15376        AgentBuilder::new()
15377            .system_prompt("Test HITL hooks.")
15378            .llm(Arc::new(mock_with_response("done")))
15379            .build()
15380            .unwrap()
15381            .with_hooks(hooks)
15382            .with_hitl(HITLEngine::new(config), Arc::new(handler))
15383    }
15384
15385    #[tokio::test]
15386    async fn approval_hooks_expose_direct_effective_decisions_after_raw_results() {
15387        let cases = vec![
15388            (ApprovalResult::Approved, "approved"),
15389            (
15390                ApprovalResult::Rejected {
15391                    reason: Some("denied".to_string()),
15392                },
15393                "rejected",
15394            ),
15395            (
15396                ApprovalResult::Modified {
15397                    changes: HashMap::from([("value".to_string(), serde_json::json!(2))]),
15398                },
15399                "modified",
15400            ),
15401        ];
15402
15403        for (raw_result, expected) in cases {
15404            let hooks = Arc::new(ApprovalRecordingHooks::new());
15405            let agent =
15406                agent_with_approval_result(raw_result, TimeoutAction::Reject, hooks.clone());
15407
15408            let result = agent.request_hitl_approval(approval_check()).await.unwrap();
15409
15410            assert_eq!(approval_result_name(&result), expected);
15411            assert_correlated_approval_events(&hooks.events(), expected, expected);
15412        }
15413    }
15414
15415    #[tokio::test]
15416    async fn approval_hooks_expose_timeout_policy_decisions() {
15417        for (timeout_action, expected) in [
15418            (TimeoutAction::Approve, "approved"),
15419            (TimeoutAction::Reject, "rejected"),
15420        ] {
15421            let hooks = Arc::new(ApprovalRecordingHooks::new());
15422            let agent =
15423                agent_with_approval_result(ApprovalResult::Timeout, timeout_action, hooks.clone());
15424
15425            let result = agent.request_hitl_approval(approval_check()).await.unwrap();
15426
15427            assert_eq!(approval_result_name(&result), expected);
15428            assert_correlated_approval_events(&hooks.events(), "timeout", expected);
15429        }
15430    }
15431
15432    #[tokio::test]
15433    async fn timeout_error_fires_correlated_resolved_error_before_returning() {
15434        let hooks = Arc::new(ApprovalRecordingHooks::new());
15435        let agent = agent_with_approval_result(
15436            ApprovalResult::Timeout,
15437            TimeoutAction::Error,
15438            hooks.clone(),
15439        );
15440
15441        let error = agent
15442            .request_hitl_approval(approval_check())
15443            .await
15444            .unwrap_err();
15445
15446        assert!(error.to_string().contains("HITL approval timeout"));
15447        assert_correlated_approval_events(&hooks.events(), "timeout", "error");
15448    }
15449
15450    // Basic YAML → Build → Chat flow
15451    #[tokio::test]
15452    async fn test_integration_yaml_to_chat_basic() {
15453        let mock = mock_with_response("Hello! How can I help you?");
15454        let agent = AgentBuilder::new()
15455            .system_prompt("You are a test assistant.")
15456            .llm(Arc::new(mock))
15457            .build()
15458            .unwrap();
15459
15460        let response = agent.chat("Hi").await.unwrap();
15461        assert!(!response.content.is_empty());
15462        assert_eq!(response.content, "Hello! How can I help you?");
15463    }
15464
15465    #[tokio::test]
15466    async fn stream_events_emit_one_authoritative_final_without_legacy_done() {
15467        let agent = AgentBuilder::new()
15468            .system_prompt("You are a test assistant.")
15469            .llm(Arc::new(mock_with_response(
15470                "Hello from the final response.",
15471            )))
15472            .build()
15473            .unwrap();
15474
15475        let mut stream = agent.chat_stream_events("Hi").await.unwrap();
15476        let mut final_responses = Vec::new();
15477        let mut legacy_done = 0;
15478        while let Some(event) = stream.next().await {
15479            match event {
15480                AgentStreamEvent::Chunk(StreamChunk::Done {}) => legacy_done += 1,
15481                AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
15482                    panic!("unexpected stream error: {message}")
15483                }
15484                AgentStreamEvent::Final(response) => final_responses.push(response),
15485                AgentStreamEvent::Chunk(_) => {}
15486            }
15487        }
15488
15489        assert_eq!(legacy_done, 0);
15490        assert_eq!(final_responses.len(), 1);
15491        let response = final_responses.pop().unwrap();
15492        assert_eq!(response.content, "Hello from the final response.");
15493        assert!(
15494            response
15495                .metadata
15496                .as_ref()
15497                .is_some_and(|metadata| { metadata.contains_key("reasoning") })
15498        );
15499    }
15500
15501    #[tokio::test]
15502    async fn stream_final_content_includes_output_processing_after_provisional_chunks() {
15503        let yaml = r#"
15504name: ProcessedStreamAgent
15505system_prompt: "Answer directly."
15506process:
15507  output:
15508    - type: format
15509      config:
15510        template: "{{ response }} [finalized]"
15511streaming:
15512  enabled: true
15513"#;
15514        let agent = AgentBuilder::from_yaml(yaml)
15515            .unwrap()
15516            .llm(Arc::new(mock_with_response("provisional answer")))
15517            .auto_configure_features()
15518            .unwrap()
15519            .build()
15520            .unwrap();
15521
15522        let mut stream = agent.chat_stream_events("Hi").await.unwrap();
15523        let mut provisional = String::new();
15524        let mut final_content = None;
15525        while let Some(event) = stream.next().await {
15526            match event {
15527                AgentStreamEvent::Chunk(StreamChunk::Content { text }) => {
15528                    provisional.push_str(&text)
15529                }
15530                AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
15531                    panic!("unexpected stream error: {message}")
15532                }
15533                AgentStreamEvent::Final(response) => final_content = Some(response.content),
15534                AgentStreamEvent::Chunk(_) => {}
15535            }
15536        }
15537
15538        assert_eq!(provisional, "provisional answer");
15539        assert_eq!(
15540            final_content.as_deref(),
15541            Some("provisional answer [finalized]")
15542        );
15543    }
15544
15545    #[tokio::test]
15546    async fn stream_events_preserve_tool_progress_and_final_tool_calls() {
15547        let agent = AgentBuilder::new()
15548            .system_prompt("Use the echo tool once, then answer.")
15549            .llm(Arc::new(mock_with_responses(vec![
15550                r#"{"tool":"echo","arguments":{"message":"hello"}}"#,
15551                "Echo completed.",
15552            ])))
15553            .tool(Arc::new(ai_agents_tools::EchoTool::new()))
15554            .build()
15555            .unwrap();
15556
15557        let mut stream = agent.chat_stream_events("echo hello").await.unwrap();
15558        let mut starts = 0;
15559        let mut results = 0;
15560        let mut ends = 0;
15561        let mut final_response = None;
15562        while let Some(event) = stream.next().await {
15563            match event {
15564                AgentStreamEvent::Chunk(StreamChunk::ToolCallStart { name, .. }) => {
15565                    assert_eq!(name, "echo");
15566                    starts += 1;
15567                }
15568                AgentStreamEvent::Chunk(StreamChunk::ToolResult { name, success, .. }) => {
15569                    assert_eq!(name, "echo");
15570                    assert!(success);
15571                    results += 1;
15572                }
15573                AgentStreamEvent::Chunk(StreamChunk::ToolCallEnd { .. }) => ends += 1,
15574                AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
15575                    panic!("unexpected stream error: {message}")
15576                }
15577                AgentStreamEvent::Final(response) => final_response = Some(response),
15578                AgentStreamEvent::Chunk(_) => {}
15579            }
15580        }
15581
15582        assert_eq!((starts, results, ends), (1, 1, 1));
15583        let response = final_response.expect("tool stream must finalize");
15584        assert_eq!(response.content, "Echo completed.");
15585        assert_eq!(
15586            response.tool_calls.as_ref().map(|calls| calls
15587                .iter()
15588                .map(|call| call.name.as_str())
15589                .collect::<Vec<_>>()),
15590            Some(vec!["echo"])
15591        );
15592    }
15593
15594    #[tokio::test]
15595    async fn legacy_stream_still_emits_one_done_chunk() {
15596        let agent = AgentBuilder::new()
15597            .system_prompt("You are a test assistant.")
15598            .llm(Arc::new(mock_with_response(
15599                "Hello from the legacy stream.",
15600            )))
15601            .build()
15602            .unwrap();
15603
15604        let mut stream = agent.chat_stream("Hi").await.unwrap();
15605        let mut done = 0;
15606        while let Some(chunk) = stream.next().await {
15607            match chunk {
15608                StreamChunk::Done {} => done += 1,
15609                StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
15610                _ => {}
15611            }
15612        }
15613
15614        assert_eq!(done, 1);
15615    }
15616
15617    // Multi-turn conversation
15618    #[tokio::test]
15619    async fn test_integration_multi_turn_conversation() {
15620        let mock = mock_with_responses(vec![
15621            "Hello! I'm your assistant.",
15622            "The weather is sunny today.",
15623            "Goodbye!",
15624        ]);
15625        let agent = AgentBuilder::new()
15626            .system_prompt("You are helpful.")
15627            .llm(Arc::new(mock))
15628            .build()
15629            .unwrap();
15630
15631        let r1 = agent.chat("Hi").await.unwrap();
15632        assert_eq!(r1.content, "Hello! I'm your assistant.");
15633
15634        let r2 = agent.chat("What's the weather?").await.unwrap();
15635        assert_eq!(r2.content, "The weather is sunny today.");
15636
15637        let r3 = agent.chat("Bye").await.unwrap();
15638        assert_eq!(r3.content, "Goodbye!");
15639
15640        // Verify memory accumulated messages
15641        let messages = agent.memory.get_messages(None).await.unwrap();
15642        // 3 user + 3 assistant = 6 messages
15643        assert_eq!(messages.len(), 6);
15644    }
15645
15646    #[test]
15647    fn later_approval_preserves_modified_evidence() {
15648        let arguments = serde_json::json!({"dry_run": true});
15649        let mut record = Some(ToolApprovalRecord {
15650            status: ToolApprovalStatus::Modified,
15651            reason: None,
15652            modified_arguments: Some(arguments.clone()),
15653        });
15654
15655        merge_approved_record(&mut record);
15656
15657        let record = record.unwrap();
15658        assert!(matches!(record.status, ToolApprovalStatus::Modified));
15659        assert_eq!(record.modified_arguments, Some(arguments));
15660    }
15661
15662    #[test]
15663    fn approval_binding_rejects_replaced_tool_implementation() {
15664        let reviewed_tool: Arc<dyn ai_agents_core::Tool> = Arc::new(ContextEchoTool);
15665        let same_tool = Arc::clone(&reviewed_tool);
15666        let replacement_tool: Arc<dyn ai_agents_core::Tool> = Arc::new(ContextEchoTool);
15667        let arguments = serde_json::json!({"path": "."});
15668        let versions = ToolDecisionVersions {
15669            policy: 2,
15670            registry: 3,
15671            runtime_control: 4,
15672            state: Some(5),
15673        };
15674        let binding = ToolApprovalBinding {
15675            canonical_id: "context_echo".to_string(),
15676            arguments: arguments.clone(),
15677            confirmation_required: true,
15678            policy_version: versions.policy,
15679            runtime_control_version: versions.runtime_control,
15680            state_generation: versions.state,
15681            reviewed_tool,
15682        };
15683
15684        assert!(!binding.is_stale("context_echo", &arguments, true, versions, &same_tool,));
15685        assert!(binding.is_stale(
15686            "context_echo",
15687            &arguments,
15688            true,
15689            versions,
15690            &replacement_tool,
15691        ));
15692    }
15693
15694    #[tokio::test]
15695    async fn approved_mutation_to_dry_run_remains_executable() {
15696        use ai_agents_hitl::CallbackHandler;
15697
15698        let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
15699            changes: HashMap::from([("dry_run".to_string(), serde_json::json!(true))]),
15700        });
15701        let agent = AgentBuilder::new()
15702            .system_prompt("Test safer approval modifications.")
15703            .llm(Arc::new(mock_with_response("done")))
15704            .tool(Arc::new(ai_agents_tools::FileWriteTool::new()))
15705            .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
15706            .approval_handler(Arc::new(handler))
15707            .build()
15708            .unwrap();
15709
15710        let record = agent
15711            .invoke_tool(ToolExecutionRequest::new(
15712                "approved-dry-run",
15713                "file_write",
15714                serde_json::json!({
15715                    "path": "./approval-dry-run.txt",
15716                    "content": "not written"
15717                }),
15718                ToolCallSource::Manual,
15719            ))
15720            .await
15721            .unwrap();
15722
15723        assert!(record.executed);
15724        assert!(record.success);
15725        assert_eq!(record.executed_arguments["dry_run"], true);
15726        assert!(matches!(
15727            record.approval.as_ref().map(|approval| &approval.status),
15728            Some(ToolApprovalStatus::Modified)
15729        ));
15730        let output: Value = serde_json::from_str(&record.output).unwrap();
15731        assert_eq!(output["mutation_performed"], false);
15732    }
15733
15734    /// Proves shared HITL approval evidence reaches the web fetch implementation unchanged.
15735    #[tokio::test]
15736    async fn shared_executor_approval_reaches_web_fetch_transport() {
15737        use ai_agents_hitl::{CallbackHandler, HITLConfig};
15738        use ai_agents_tools::{DomainPolicyConfig, ToolPolicyConfig};
15739
15740        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15741        let tool = WebFetchTool::with_transport_and_resolver(
15742            Arc::new(RuntimeWebFetchTransport {
15743                calls: Arc::clone(&calls),
15744            }),
15745            Arc::new(RuntimeWebFetchResolver),
15746        );
15747        let mut security = ToolSecurityConfig {
15748            enabled: true,
15749            fail_closed: true,
15750            ..Default::default()
15751        };
15752        security.tools.insert(
15753            "web_fetch".to_string(),
15754            ToolPolicyConfig {
15755                domains: DomainPolicyConfig {
15756                    requires_approval: vec!["approval.test".to_string()],
15757                    ..Default::default()
15758                },
15759                allowed_schemes: vec!["https".to_string()],
15760                allowed_ports: vec![443],
15761                ..Default::default()
15762            },
15763        );
15764        let handler = CallbackHandler::new(|_| ApprovalResult::Approved);
15765        let agent = AgentBuilder::new()
15766            .system_prompt("Test approved web fetch execution.")
15767            .llm(Arc::new(mock_with_response("done")))
15768            .tool(Arc::new(tool))
15769            .tool_security(ToolSecurityEngine::new(security))
15770            .build()
15771            .unwrap()
15772            .with_hitl(HITLEngine::new(HITLConfig::default()), Arc::new(handler));
15773
15774        let record = agent
15775            .invoke_tool(ToolExecutionRequest::new(
15776                "approved-web-fetch",
15777                "web_fetch",
15778                serde_json::json!({
15779                    "url": "https://approval.test/page",
15780                    "cache_ttl_seconds": 0
15781                }),
15782                ToolCallSource::Manual,
15783            ))
15784            .await
15785            .unwrap();
15786
15787        assert!(record.success);
15788        assert!(
15789            record
15790                .approval
15791                .as_ref()
15792                .is_some_and(|approval| matches!(approval.status, ToolApprovalStatus::Approved))
15793        );
15794        assert_eq!(calls.load(Ordering::SeqCst), 1);
15795    }
15796
15797    #[tokio::test]
15798    async fn context_preserves_requested_and_canonical_identity() {
15799        let mock = mock_with_response("hello");
15800        let mut tools = ai_agents_tools::ToolRegistry::new();
15801        tools.register(Arc::new(ContextEchoTool)).unwrap();
15802
15803        let mut security = ToolSecurityConfig {
15804            enabled: true,
15805            fail_closed: true,
15806            ..Default::default()
15807        };
15808        let mut policy = ai_agents_tools::ToolPolicyConfig {
15809            read_paths: vec![".".to_string()],
15810            max_results: Some(7),
15811            ..Default::default()
15812        };
15813        policy
15814            .config
15815            .insert("backend".to_string(), serde_json::json!("memory"));
15816        security.tools.insert("context_echo".to_string(), policy);
15817
15818        let agent = AgentBuilder::new()
15819            .system_prompt("You are helpful.")
15820            .llm(Arc::new(mock))
15821            .tools(tools)
15822            .tool_security(ToolSecurityEngine::new(security))
15823            .build()
15824            .unwrap();
15825
15826        let record = agent
15827            .invoke_tool(ToolExecutionRequest::new(
15828                "ctx-call",
15829                "Context Echo",
15830                serde_json::json!({"path": ".", "max_results": 99}),
15831                ToolCallSource::Manual,
15832            ))
15833            .await
15834            .unwrap();
15835
15836        assert!(record.success);
15837        assert!(matches!(&record.source, ToolCallSource::Manual));
15838        assert_eq!(record.requested_name, "Context Echo");
15839        assert_eq!(record.canonical_id, "context_echo");
15840        assert_eq!(record.policy.outcome, PermissionOutcome::Allow);
15841        assert_eq!(record.executed_arguments["max_results"], 7);
15842        let output: Value = serde_json::from_str(&record.output).unwrap();
15843        assert_eq!(output["requested_name"], "Context Echo");
15844        assert_eq!(output["canonical_id"], "context_echo");
15845        assert_eq!(output["max_results"], 7);
15846        assert_eq!(output["custom_config"]["backend"], "memory");
15847        assert!(record.metadata.contains_key("effective_limits"));
15848        assert!(record.metadata.contains_key("policy_snapshot"));
15849    }
15850
15851    #[tokio::test]
15852    async fn test_runtime_control_cancels_active_tool_call() {
15853        let mock = mock_with_response("hello");
15854        let agent = Arc::new(
15855            AgentBuilder::new()
15856                .system_prompt("You are helpful.")
15857                .llm(Arc::new(mock))
15858                .tool(Arc::new(SlowTool))
15859                .build()
15860                .unwrap(),
15861        );
15862        let control = agent.runtime_control();
15863        let running_agent = Arc::clone(&agent);
15864        let handle = tokio::spawn(async move {
15865            running_agent
15866                .invoke_tool(ToolExecutionRequest::new(
15867                    "slow-call",
15868                    "slow",
15869                    serde_json::json!({}),
15870                    ToolCallSource::Manual,
15871                ))
15872                .await
15873                .unwrap()
15874        });
15875
15876        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
15877        control.cancel_all();
15878        let record = handle.await.unwrap();
15879
15880        assert!(record.executed);
15881        assert!(record.cancelled);
15882        assert!(!record.success);
15883        assert!(record.cancellation_reason.is_some());
15884    }
15885
15886    #[tokio::test]
15887    async fn non_idempotent_tool_calls_are_not_retried() {
15888        use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
15889
15890        let mock = mock_with_response("hello");
15891        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15892        let agent = AgentBuilder::new()
15893            .system_prompt("You are helpful.")
15894            .llm(Arc::new(mock))
15895            .tool(Arc::new(FlakyWriteTool {
15896                calls: Arc::clone(&calls),
15897            }))
15898            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
15899                tools: ToolRecoveryConfig {
15900                    default: ToolRetryConfig {
15901                        max_retries: 2,
15902                        ..Default::default()
15903                    },
15904                    ..Default::default()
15905                },
15906                ..Default::default()
15907            }))
15908            .build()
15909            .unwrap();
15910
15911        let record = agent
15912            .invoke_tool(ToolExecutionRequest::new(
15913                "flaky-call",
15914                "flaky_write",
15915                serde_json::json!({"path": "./tmp.txt"}),
15916                ToolCallSource::Manual,
15917            ))
15918            .await
15919            .unwrap();
15920
15921        assert!(!record.success);
15922        assert_eq!(calls.load(Ordering::SeqCst), 1);
15923    }
15924
15925    #[tokio::test]
15926    async fn side_effecting_tools_are_serialized_per_resource() {
15927        let mock = mock_with_response("hello");
15928        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15929        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15930        let agent = Arc::new(
15931            AgentBuilder::new()
15932                .system_prompt("You are helpful.")
15933                .llm(Arc::new(mock))
15934                .tool(Arc::new(LockedWriteTool {
15935                    active: Arc::clone(&active),
15936                    max_active: Arc::clone(&max_active),
15937                }))
15938                .build()
15939                .unwrap(),
15940        );
15941
15942        let left = {
15943            let agent = Arc::clone(&agent);
15944            tokio::spawn(async move {
15945                agent
15946                    .invoke_tool(ToolExecutionRequest::new(
15947                        "lock-1",
15948                        "locked_write",
15949                        serde_json::json!({"path": "./same.txt"}),
15950                        ToolCallSource::Manual,
15951                    ))
15952                    .await
15953                    .unwrap()
15954            })
15955        };
15956        let right = {
15957            let agent = Arc::clone(&agent);
15958            tokio::spawn(async move {
15959                agent
15960                    .invoke_tool(ToolExecutionRequest::new(
15961                        "lock-2",
15962                        "locked_write",
15963                        serde_json::json!({"path": "./same.txt"}),
15964                        ToolCallSource::Manual,
15965                    ))
15966                    .await
15967                    .unwrap()
15968            })
15969        };
15970
15971        let left = left.await.unwrap();
15972        let right = right.await.unwrap();
15973        assert!(left.success);
15974        assert!(right.success);
15975        assert_eq!(max_active.load(Ordering::SeqCst), 1);
15976    }
15977
15978    #[tokio::test]
15979    async fn path_resources_use_shared_global_lock_and_cleanup() {
15980        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15981        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15982        let bindings = ai_agents_core::ToolPolicyBindings {
15983            path_fields: vec![
15984                ai_agents_core::PathPolicyBinding::read_write("source_path"),
15985                ai_agents_core::PathPolicyBinding::write("destination_path"),
15986            ],
15987            ..Default::default()
15988        };
15989        let classification = ai_agents_core::ToolCallClassification::from_metadata(
15990            &MultiResourceWriteTool {
15991                active: Arc::clone(&active),
15992                max_active: Arc::clone(&max_active),
15993            }
15994            .safety_metadata(),
15995        );
15996        let left_args = serde_json::json!({
15997            "source_path": "./a/../first.txt",
15998            "destination_path": "./second.txt"
15999        });
16000        let right_args = serde_json::json!({
16001            "source_path": "./second.txt",
16002            "destination_path": "./first.txt"
16003        });
16004        let left_keys = tool_resource_lock_keys(
16005            "multi_resource_write",
16006            &left_args,
16007            &bindings,
16008            &classification,
16009        );
16010        let right_keys = tool_resource_lock_keys(
16011            "multi_resource_write",
16012            &right_args,
16013            &bindings,
16014            &classification,
16015        );
16016        assert_eq!(left_keys, right_keys);
16017        assert_eq!(left_keys, vec!["path-mutation:global".to_string()]);
16018
16019        let locks = new_tool_resource_locks();
16020        let build_agent = || {
16021            AgentBuilder::new()
16022                .system_prompt("Test shared resource locks.")
16023                .llm(Arc::new(mock_with_response("done")))
16024                .tool(Arc::new(MultiResourceWriteTool {
16025                    active: Arc::clone(&active),
16026                    max_active: Arc::clone(&max_active),
16027                }))
16028                .build()
16029                .unwrap()
16030                .with_shared_resource_locks(Arc::clone(&locks))
16031        };
16032        let left_agent = Arc::new(build_agent());
16033        let right_agent = Arc::new(build_agent());
16034        let left = tokio::spawn(async move {
16035            left_agent
16036                .invoke_tool(ToolExecutionRequest::new(
16037                    "multi-left",
16038                    "multi_resource_write",
16039                    left_args,
16040                    ToolCallSource::Manual,
16041                ))
16042                .await
16043                .unwrap()
16044        });
16045        let right = tokio::spawn(async move {
16046            right_agent
16047                .invoke_tool(ToolExecutionRequest::new(
16048                    "multi-right",
16049                    "multi_resource_write",
16050                    right_args,
16051                    ToolCallSource::Manual,
16052                ))
16053                .await
16054                .unwrap()
16055        });
16056        let (left, right) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
16057            tokio::join!(left, right)
16058        })
16059        .await
16060        .expect("reversed resource acquisition must not deadlock");
16061
16062        assert!(left.unwrap().success);
16063        assert!(right.unwrap().success);
16064        assert_eq!(max_active.load(Ordering::SeqCst), 1);
16065        assert!(locks.read().is_empty());
16066    }
16067
16068    #[tokio::test]
16069    async fn global_path_lock_serializes_copy_destination_with_file_write() {
16070        assert_path_mutation_pair_serialized(
16071            "copy_path",
16072            CopyPathTool::new().policy_bindings().path_fields,
16073            serde_json::json!({
16074                "source_path": "./source.txt",
16075                "destination_path": "./shared.txt"
16076            }),
16077            "file_write",
16078            FileWriteTool::new().policy_bindings().path_fields,
16079            serde_json::json!({"path": "./shared.txt"}),
16080        )
16081        .await;
16082    }
16083
16084    #[tokio::test]
16085    async fn parent_and_spawned_runtime_share_global_path_lock() {
16086        let workspace = MutationTestWorkspace::new();
16087        let destination = workspace.root.join("spawned.txt");
16088        let parent_gate = PathMutationGate::new();
16089        let parent = Arc::new(
16090            AgentBuilder::from_yaml(
16091                r#"
16092name: LockParent
16093system_prompt: parent
16094llm:
16095  default: default
16096tools:
16097  - parent_path_write
16098spawner:
16099  shared_llms: true
16100"#,
16101            )
16102            .unwrap()
16103            .llm(Arc::new(mock_with_response("done")))
16104            .auto_configure_spawner()
16105            .await
16106            .unwrap()
16107            .tool(Arc::new(BlockingPathMutationTool {
16108                id: "parent_path_write",
16109                path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
16110                gate: parent_gate.clone(),
16111            }))
16112            .build()
16113            .unwrap(),
16114        );
16115
16116        let mut child_spec = crate::spec::AgentSpec {
16117            name: "LockChild".to_string(),
16118            system_prompt: "child".to_string(),
16119            tools: Some(vec![crate::spec::ToolEntry::Simple(
16120                "file_write".to_string(),
16121            )]),
16122            ..Default::default()
16123        };
16124        child_spec.tool_security.enabled = true;
16125        child_spec.tool_security.fail_closed = true;
16126        let file_write_policy = ai_agents_tools::ToolPolicyConfig {
16127            write_paths: vec![workspace.root.to_string_lossy().into_owned()],
16128            allow_without_confirmation: true,
16129            ..Default::default()
16130        };
16131        child_spec
16132            .tool_security
16133            .tools
16134            .insert("file_write".to_string(), file_write_policy);
16135        let spawned = parent
16136            .spawner()
16137            .unwrap()
16138            .spawn_from_spec(child_spec)
16139            .await
16140            .unwrap();
16141        assert!(Arc::ptr_eq(
16142            &parent.resource_locks,
16143            &spawned.agent.resource_locks
16144        ));
16145        assert!(!Arc::ptr_eq(
16146            &parent.runtime_control,
16147            &spawned.agent.runtime_control
16148        ));
16149
16150        let parent_call = {
16151            let parent = Arc::clone(&parent);
16152            let destination = destination.clone();
16153            tokio::spawn(async move {
16154                parent
16155                    .invoke_tool(ToolExecutionRequest::new(
16156                        "parent-lock-holder",
16157                        "parent_path_write",
16158                        serde_json::json!({"path": destination}),
16159                        ToolCallSource::Manual,
16160                    ))
16161                    .await
16162                    .unwrap()
16163            })
16164        };
16165        parent_gate.wait_until_entered().await;
16166
16167        let child_call = {
16168            let child = Arc::clone(&spawned.agent);
16169            let destination = destination.clone();
16170            tokio::spawn(async move {
16171                child
16172                    .invoke_tool(ToolExecutionRequest::new(
16173                        "spawned-file-write",
16174                        "file_write",
16175                        serde_json::json!({
16176                            "path": destination,
16177                            "content": "spawned",
16178                            "dry_run": false
16179                        }),
16180                        ToolCallSource::Manual,
16181                    ))
16182                    .await
16183                    .unwrap()
16184            })
16185        };
16186        wait_for_resource_lock_strong_count(&parent.resource_locks, 2).await;
16187        assert!(!child_call.is_finished());
16188
16189        parent_gate.release();
16190        let (parent_record, child_record) =
16191            tokio::time::timeout(std::time::Duration::from_secs(2), async {
16192                tokio::join!(parent_call, child_call)
16193            })
16194            .await
16195            .expect("parent and spawned path mutations did not finish");
16196        assert!(parent_record.unwrap().success);
16197        assert!(child_record.unwrap().success);
16198        assert_eq!(std::fs::read_to_string(destination).unwrap(), "spawned");
16199        assert!(parent.resource_locks.read().is_empty());
16200    }
16201
16202    #[tokio::test]
16203    async fn cancelled_global_path_lock_waiter_does_not_retain_weak_entry() {
16204        let locks = new_tool_resource_locks();
16205        let holder_gate = PathMutationGate::new();
16206        let waiter_gate = PathMutationGate::new();
16207        waiter_gate.release();
16208        let holder = Arc::new(
16209            AgentBuilder::new()
16210                .system_prompt("Hold the global path lock.")
16211                .llm(Arc::new(mock_with_response("done")))
16212                .tool(Arc::new(BlockingPathMutationTool {
16213                    id: "holder_write",
16214                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
16215                    gate: holder_gate.clone(),
16216                }))
16217                .build()
16218                .unwrap()
16219                .with_shared_resource_locks(Arc::clone(&locks)),
16220        );
16221        let waiter = Arc::new(
16222            AgentBuilder::new()
16223                .system_prompt("Wait for the global path lock.")
16224                .llm(Arc::new(mock_with_response("done")))
16225                .tool(Arc::new(BlockingPathMutationTool {
16226                    id: "waiter_write",
16227                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
16228                    gate: waiter_gate.clone(),
16229                }))
16230                .build()
16231                .unwrap()
16232                .with_shared_resource_locks(Arc::clone(&locks)),
16233        );
16234
16235        let holder_call = {
16236            let holder = Arc::clone(&holder);
16237            tokio::spawn(async move {
16238                holder
16239                    .invoke_tool(ToolExecutionRequest::new(
16240                        "holder-call",
16241                        "holder_write",
16242                        serde_json::json!({"path": "./shared.txt"}),
16243                        ToolCallSource::Manual,
16244                    ))
16245                    .await
16246                    .unwrap()
16247            })
16248        };
16249        holder_gate.wait_until_entered().await;
16250
16251        let waiter_call = {
16252            let waiter = Arc::clone(&waiter);
16253            tokio::spawn(async move {
16254                waiter
16255                    .invoke_tool(ToolExecutionRequest::new(
16256                        "waiter-call",
16257                        "waiter_write",
16258                        serde_json::json!({"path": "./shared.txt"}),
16259                        ToolCallSource::Manual,
16260                    ))
16261                    .await
16262                    .unwrap()
16263            })
16264        };
16265        wait_for_resource_lock_strong_count(&locks, 2).await;
16266        waiter.runtime_control().cancel_all();
16267
16268        let waiter_record = tokio::time::timeout(std::time::Duration::from_secs(2), waiter_call)
16269            .await
16270            .expect("cancelled lock waiter did not finish")
16271            .unwrap();
16272        assert!(!waiter_record.executed);
16273        assert!(!waiter_gate.entered.load(Ordering::SeqCst));
16274        assert_eq!(
16275            locks
16276                .read()
16277                .get("path-mutation:global")
16278                .map_or(0, |lock| lock.strong_count()),
16279            1
16280        );
16281
16282        holder_gate.release();
16283        let holder_record = tokio::time::timeout(std::time::Duration::from_secs(2), holder_call)
16284            .await
16285            .expect("lock holder did not finish")
16286            .unwrap();
16287        assert!(holder_record.success);
16288        assert!(locks.read().is_empty());
16289    }
16290
16291    #[tokio::test]
16292    async fn path_mutation_policy_and_approval_denials_do_not_invoke_tools() {
16293        for denial in [MutationDenial::Policy, MutationDenial::Approval] {
16294            let tools: [Arc<dyn Tool>; 3] = [
16295                Arc::new(CopyPathTool::new()),
16296                Arc::new(MovePathTool::new()),
16297                Arc::new(DeletePathTool::new()),
16298            ];
16299            for tool in tools {
16300                assert_path_mutation_denied(tool, denial).await;
16301            }
16302        }
16303    }
16304
16305    #[tokio::test]
16306    async fn approval_argument_changes_are_rechecked_against_final_scope() {
16307        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16308        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16309        let entered = Arc::new(tokio::sync::Barrier::new(2));
16310        let release = Arc::new(tokio::sync::Notify::new());
16311        let handler = Arc::new(BlockingApprovalHandler {
16312            entered: Arc::clone(&entered),
16313            release: Arc::clone(&release),
16314            result: ApprovalResult::Modified {
16315                changes: HashMap::from([(
16316                    "path".to_string(),
16317                    Value::String("./after-approval.txt".to_string()),
16318                )]),
16319            },
16320        });
16321        let agent = Arc::new(
16322            AgentBuilder::new()
16323                .system_prompt("Test final scope validation.")
16324                .llm(Arc::new(mock_with_response("done")))
16325                .tool(Arc::new(LockedWriteTool {
16326                    active: Arc::clone(&active),
16327                    max_active: Arc::clone(&max_active),
16328                }))
16329                .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
16330                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16331                .approval_handler(handler)
16332                .build()
16333                .unwrap(),
16334        );
16335        let control = agent.runtime_control();
16336        let running = Arc::clone(&agent);
16337        let call = tokio::spawn(async move {
16338            running
16339                .invoke_tool(ToolExecutionRequest::new(
16340                    "approval-scope",
16341                    "locked_write",
16342                    serde_json::json!({"path": "./before-approval.txt"}),
16343                    ToolCallSource::Manual,
16344                ))
16345                .await
16346                .unwrap()
16347        });
16348        entered.wait().await;
16349        let expected_version = control.set_tool_scope(Vec::new());
16350        release.notify_one();
16351        let record = call.await.unwrap();
16352
16353        assert!(!record.executed);
16354        assert!(!record.success);
16355        assert_eq!(record.runtime_config_version, expected_version);
16356        assert_eq!(record.executed_arguments["path"], "./after-approval.txt");
16357        assert_eq!(max_active.load(Ordering::SeqCst), 0);
16358        assert_eq!(
16359            record.metadata["runtime_scope_snapshot"],
16360            serde_json::json!([])
16361        );
16362    }
16363
16364    #[tokio::test]
16365    async fn approval_is_rechecked_against_final_policy_snapshot() {
16366        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16367        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16368        let entered = Arc::new(tokio::sync::Barrier::new(2));
16369        let release = Arc::new(tokio::sync::Notify::new());
16370        let handler = Arc::new(BlockingApprovalHandler {
16371            entered: Arc::clone(&entered),
16372            release: Arc::clone(&release),
16373            result: ApprovalResult::Approved,
16374        });
16375        let agent = Arc::new(
16376            AgentBuilder::new()
16377                .system_prompt("Test final policy validation.")
16378                .llm(Arc::new(mock_with_response("done")))
16379                .tool(Arc::new(LockedWriteTool {
16380                    active: Arc::clone(&active),
16381                    max_active: Arc::clone(&max_active),
16382                }))
16383                .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
16384                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16385                .approval_handler(handler)
16386                .build()
16387                .unwrap(),
16388        );
16389        let control = agent.runtime_control();
16390        let running = Arc::clone(&agent);
16391        let call = tokio::spawn(async move {
16392            running
16393                .invoke_tool(ToolExecutionRequest::new(
16394                    "approval-policy",
16395                    "locked_write",
16396                    serde_json::json!({"path": "./policy.txt"}),
16397                    ToolCallSource::Manual,
16398                ))
16399                .await
16400                .unwrap()
16401        });
16402        entered.wait().await;
16403        let expected_version = control.set_tool_security(approval_security_config(false));
16404        release.notify_one();
16405        let record = call.await.unwrap();
16406
16407        assert!(!record.executed);
16408        assert!(!record.success);
16409        assert_eq!(record.runtime_config_version, expected_version);
16410        assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
16411        assert_eq!(max_active.load(Ordering::SeqCst), 0);
16412        assert!(record.metadata.contains_key("policy_snapshot"));
16413    }
16414
16415    #[test]
16416    fn invalid_live_policy_does_not_replace_snapshot_or_generation() {
16417        let agent = AgentBuilder::new()
16418            .system_prompt("Test runtime policy validation.")
16419            .llm(Arc::new(mock_with_response("done")))
16420            .build()
16421            .unwrap();
16422        let control = agent.runtime_control();
16423        let mut valid = ToolSecurityConfig::default();
16424        valid.tools.insert(
16425            "web_search".to_string(),
16426            ai_agents_tools::ToolPolicyConfig {
16427                max_results: Some(5),
16428                ..Default::default()
16429            },
16430        );
16431        let generation = control.try_set_tool_security(valid).unwrap();
16432
16433        let mut invalid = ToolSecurityConfig::default();
16434        invalid.tools.insert(
16435            "web_search".to_string(),
16436            ai_agents_tools::ToolPolicyConfig {
16437                max_results: Some(0),
16438                ..Default::default()
16439            },
16440        );
16441        let error = control.try_set_tool_security(invalid).unwrap_err();
16442
16443        assert!(
16444            error
16445                .to_string()
16446                .contains("max_results must be greater than 0")
16447        );
16448        assert_eq!(control.version(), generation);
16449        assert_eq!(
16450            control
16451                .state
16452                .tool_security_override
16453                .read()
16454                .as_ref()
16455                .unwrap()
16456                .config()
16457                .tools["web_search"]
16458                .max_results,
16459            Some(5)
16460        );
16461    }
16462
16463    #[tokio::test]
16464    async fn persistent_override_preserves_rate_history_within_generation() {
16465        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16466        let agent = AgentBuilder::new()
16467            .system_prompt("Test persistent policy overrides.")
16468            .llm(Arc::new(mock_with_response("done")))
16469            .tool(Arc::new(RecoveryTestTool {
16470                id: "limited_override".to_string(),
16471                succeeds: true,
16472                calls: Arc::clone(&calls),
16473            }))
16474            .build()
16475            .unwrap();
16476        let mut security = ToolSecurityConfig {
16477            enabled: true,
16478            fail_closed: true,
16479            ..Default::default()
16480        };
16481        let policy = ai_agents_tools::ToolPolicyConfig {
16482            write_paths: vec![".".to_string()],
16483            rate_limit: Some(1),
16484            ..Default::default()
16485        };
16486        security
16487            .tools
16488            .insert("limited_override".to_string(), policy);
16489        let generation = agent.runtime_control().set_tool_security(security);
16490
16491        let first = agent
16492            .invoke_tool(ToolExecutionRequest::new(
16493                "limited-first",
16494                "limited_override",
16495                serde_json::json!({"path": "./limited.txt"}),
16496                ToolCallSource::Manual,
16497            ))
16498            .await
16499            .unwrap();
16500        let second = agent
16501            .invoke_tool(ToolExecutionRequest::new(
16502                "limited-second",
16503                "limited_override",
16504                serde_json::json!({"path": "./limited.txt"}),
16505                ToolCallSource::Manual,
16506            ))
16507            .await
16508            .unwrap();
16509
16510        assert!(first.success);
16511        assert_eq!(first.policy_version, generation);
16512        assert!(!second.executed);
16513        assert!(second.output.contains("Rate limit exceeded"));
16514        assert_eq!(second.policy_version, generation);
16515        assert_eq!(calls.load(Ordering::SeqCst), 1);
16516    }
16517
16518    #[tokio::test]
16519    async fn concurrent_rate_admission_consumes_capacity_atomically() {
16520        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16521        let tool = Arc::new(RecoveryTestTool {
16522            id: "atomic_rate".to_string(),
16523            succeeds: true,
16524            calls: Arc::clone(&calls),
16525        });
16526        let arguments = serde_json::json!({"path": "./atomic-rate.txt"});
16527        let bindings = tool.policy_bindings();
16528        let classification = tool.classify_call(&arguments);
16529        let resource_keys =
16530            tool_resource_lock_keys(tool.id(), &arguments, &bindings, &classification);
16531        let mut security = ToolSecurityConfig {
16532            enabled: true,
16533            fail_closed: true,
16534            ..Default::default()
16535        };
16536        let policy = ai_agents_tools::ToolPolicyConfig {
16537            write_paths: vec![".".to_string()],
16538            rate_limit: Some(1),
16539            ..Default::default()
16540        };
16541        security.tools.insert(tool.id().to_string(), policy);
16542        let agent = Arc::new(
16543            AgentBuilder::new()
16544                .system_prompt("Test atomic rate admission.")
16545                .llm(Arc::new(mock_with_response("done")))
16546                .tool(tool)
16547                .tool_security(ToolSecurityEngine::new(security))
16548                .build()
16549                .unwrap(),
16550        );
16551        let held = agent
16552            .acquire_tool_resource_locks(&resource_keys)
16553            .await
16554            .unwrap();
16555        let left = {
16556            let agent = Arc::clone(&agent);
16557            let arguments = arguments.clone();
16558            tokio::spawn(async move {
16559                agent
16560                    .invoke_tool(ToolExecutionRequest::new(
16561                        "atomic-rate-left",
16562                        "atomic_rate",
16563                        arguments,
16564                        ToolCallSource::Manual,
16565                    ))
16566                    .await
16567                    .unwrap()
16568            })
16569        };
16570        let right = {
16571            let agent = Arc::clone(&agent);
16572            tokio::spawn(async move {
16573                agent
16574                    .invoke_tool(ToolExecutionRequest::new(
16575                        "atomic-rate-right",
16576                        "atomic_rate",
16577                        arguments,
16578                        ToolCallSource::Manual,
16579                    ))
16580                    .await
16581                    .unwrap()
16582            })
16583        };
16584        tokio::time::sleep(std::time::Duration::from_millis(25)).await;
16585        drop(held);
16586        let (left, right) = tokio::join!(left, right);
16587        let records = [left.unwrap(), right.unwrap()];
16588
16589        assert_eq!(records.iter().filter(|record| record.success).count(), 1);
16590        assert_eq!(records.iter().filter(|record| record.executed).count(), 1);
16591        assert!(
16592            records.iter().any(|record| {
16593                !record.executed && record.output.contains("Rate limit exceeded")
16594            })
16595        );
16596        assert_eq!(calls.load(Ordering::SeqCst), 1);
16597    }
16598
16599    #[tokio::test]
16600    async fn changed_policy_generation_invalidates_pending_approval() {
16601        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16602        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16603        let entered = Arc::new(tokio::sync::Barrier::new(2));
16604        let release = Arc::new(tokio::sync::Notify::new());
16605        let handler = Arc::new(BlockingApprovalHandler {
16606            entered: Arc::clone(&entered),
16607            release: Arc::clone(&release),
16608            result: ApprovalResult::Approved,
16609        });
16610        let agent = Arc::new(
16611            AgentBuilder::new()
16612                .system_prompt("Test stale approval denial.")
16613                .llm(Arc::new(mock_with_response("done")))
16614                .tool(Arc::new(LockedWriteTool {
16615                    active: Arc::clone(&active),
16616                    max_active: Arc::clone(&max_active),
16617                }))
16618                .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
16619                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16620                .approval_handler(handler)
16621                .build()
16622                .unwrap(),
16623        );
16624        let running = Arc::clone(&agent);
16625        let call = tokio::spawn(async move {
16626            running
16627                .invoke_tool(ToolExecutionRequest::new(
16628                    "stale-approval",
16629                    "locked_write",
16630                    serde_json::json!({"path": "./stale.txt"}),
16631                    ToolCallSource::Manual,
16632                ))
16633                .await
16634                .unwrap()
16635        });
16636        entered.wait().await;
16637        let generation = agent
16638            .runtime_control()
16639            .set_tool_security(approval_security_config(true));
16640        release.notify_one();
16641        let record = call.await.unwrap();
16642
16643        assert!(!record.executed);
16644        assert!(record.output.contains("Approval became stale"));
16645        assert_eq!(record.policy_version, generation);
16646        assert_eq!(max_active.load(Ordering::SeqCst), 0);
16647    }
16648
16649    #[tokio::test]
16650    async fn final_policy_reapplies_argument_caps_after_approval_changes() {
16651        use ai_agents_hitl::CallbackHandler;
16652
16653        let mut security = ToolSecurityConfig {
16654            enabled: true,
16655            fail_closed: true,
16656            ..Default::default()
16657        };
16658        let policy = ai_agents_tools::ToolPolicyConfig {
16659            read_paths: vec![".".to_string()],
16660            max_results: Some(5),
16661            require_confirmation: true,
16662            ..Default::default()
16663        };
16664        security.tools.insert("context_echo".to_string(), policy);
16665        let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
16666            changes: HashMap::from([("max_results".to_string(), serde_json::json!(99))]),
16667        });
16668        let agent = AgentBuilder::new()
16669            .system_prompt("Test final argument caps.")
16670            .llm(Arc::new(mock_with_response("done")))
16671            .tool(Arc::new(ContextEchoTool))
16672            .tool_security(ToolSecurityEngine::new(security))
16673            .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16674            .approval_handler(Arc::new(handler))
16675            .build()
16676            .unwrap();
16677
16678        let record = agent
16679            .invoke_tool(ToolExecutionRequest::new(
16680                "final-cap",
16681                "context_echo",
16682                serde_json::json!({"path": ".", "max_results": 1}),
16683                ToolCallSource::Manual,
16684            ))
16685            .await
16686            .unwrap();
16687
16688        assert!(record.success);
16689        assert_eq!(record.executed_arguments["max_results"], 5);
16690        assert_eq!(
16691            record.approval.unwrap().modified_arguments.unwrap()["max_results"],
16692            5
16693        );
16694    }
16695
16696    #[tokio::test]
16697    async fn no_binding_writes_use_canonical_fallback_lock() {
16698        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16699        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16700        let agent = Arc::new(
16701            AgentBuilder::new()
16702                .system_prompt("Test fallback resource locks.")
16703                .llm(Arc::new(mock_with_response("done")))
16704                .tool(Arc::new(NoBindingWriteTool {
16705                    active: Arc::clone(&active),
16706                    max_active: Arc::clone(&max_active),
16707                }))
16708                .build()
16709                .unwrap(),
16710        );
16711        let left = {
16712            let agent = Arc::clone(&agent);
16713            tokio::spawn(async move {
16714                agent
16715                    .invoke_tool(ToolExecutionRequest::new(
16716                        "no-binding-left",
16717                        "no_binding_write",
16718                        serde_json::json!({}),
16719                        ToolCallSource::Manual,
16720                    ))
16721                    .await
16722                    .unwrap()
16723            })
16724        };
16725        let right = {
16726            let agent = Arc::clone(&agent);
16727            tokio::spawn(async move {
16728                agent
16729                    .invoke_tool(ToolExecutionRequest::new(
16730                        "no-binding-right",
16731                        "no_binding_write",
16732                        serde_json::json!({}),
16733                        ToolCallSource::Manual,
16734                    ))
16735                    .await
16736                    .unwrap()
16737            })
16738        };
16739        let (left, right) = tokio::join!(left, right);
16740
16741        assert!(left.unwrap().success);
16742        assert!(right.unwrap().success);
16743        assert_eq!(max_active.load(Ordering::SeqCst), 1);
16744    }
16745
16746    #[tokio::test]
16747    async fn parent_and_child_paths_share_a_resource_lock() {
16748        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16749        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16750        let agent = Arc::new(
16751            AgentBuilder::new()
16752                .system_prompt("Test parent child resource locks.")
16753                .llm(Arc::new(mock_with_response("done")))
16754                .tool(Arc::new(LockedWriteTool {
16755                    active: Arc::clone(&active),
16756                    max_active: Arc::clone(&max_active),
16757                }))
16758                .build()
16759                .unwrap(),
16760        );
16761        let parent = format!("./lock-parent-{}", uuid::Uuid::new_v4());
16762        let child = format!("{}/child.txt", parent);
16763        let left = {
16764            let agent = Arc::clone(&agent);
16765            tokio::spawn(async move {
16766                agent
16767                    .invoke_tool(ToolExecutionRequest::new(
16768                        "parent-lock",
16769                        "locked_write",
16770                        serde_json::json!({"path": parent}),
16771                        ToolCallSource::Manual,
16772                    ))
16773                    .await
16774                    .unwrap()
16775            })
16776        };
16777        let right = {
16778            let agent = Arc::clone(&agent);
16779            tokio::spawn(async move {
16780                agent
16781                    .invoke_tool(ToolExecutionRequest::new(
16782                        "child-lock",
16783                        "locked_write",
16784                        serde_json::json!({"path": child}),
16785                        ToolCallSource::Manual,
16786                    ))
16787                    .await
16788                    .unwrap()
16789            })
16790        };
16791        let (left, right) = tokio::join!(left, right);
16792
16793        assert!(left.unwrap().success);
16794        assert!(right.unwrap().success);
16795        assert_eq!(max_active.load(Ordering::SeqCst), 1);
16796    }
16797
16798    #[tokio::test]
16799    async fn tool_hooks_can_reenter_after_resource_guards_are_dropped() {
16800        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16801        let hooks = Arc::new(ReentrantToolHooks {
16802            agent: parking_lot::Mutex::new(None),
16803            invoked: AtomicBool::new(false),
16804            nested_success: AtomicBool::new(false),
16805        });
16806        let agent = Arc::new(
16807            AgentBuilder::new()
16808                .system_prompt("Test hook reentrancy.")
16809                .llm(Arc::new(mock_with_response("done")))
16810                .tool(Arc::new(RecoveryTestTool {
16811                    id: "reentrant_write".to_string(),
16812                    succeeds: true,
16813                    calls: Arc::clone(&calls),
16814                }))
16815                .hooks(hooks.clone())
16816                .build()
16817                .unwrap(),
16818        );
16819        *hooks.agent.lock() = Some(Arc::downgrade(&agent));
16820        let record = tokio::time::timeout(
16821            std::time::Duration::from_secs(2),
16822            agent.invoke_tool(ToolExecutionRequest::new(
16823                "outer-hook-call",
16824                "reentrant_write",
16825                serde_json::json!({"path": "./hook.txt"}),
16826                ToolCallSource::Manual,
16827            )),
16828        )
16829        .await
16830        .expect("tool completion hook must not retain resource guards")
16831        .unwrap();
16832
16833        assert!(record.success);
16834        assert!(hooks.nested_success.load(Ordering::SeqCst));
16835        assert_eq!(calls.load(Ordering::SeqCst), 2);
16836    }
16837
16838    #[tokio::test]
16839    async fn fallback_releases_primary_resource_locks() {
16840        let primary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16841        let fallback_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16842        let agent = AgentBuilder::new()
16843            .system_prompt("Test fallback execution.")
16844            .llm(Arc::new(mock_with_response("done")))
16845            .tool(Arc::new(RecoveryTestTool {
16846                id: "primary".to_string(),
16847                succeeds: false,
16848                calls: Arc::clone(&primary_calls),
16849            }))
16850            .tool(Arc::new(RecoveryTestTool {
16851                id: "fallback".to_string(),
16852                succeeds: true,
16853                calls: Arc::clone(&fallback_calls),
16854            }))
16855            .recovery_manager(recovery_manager_with_fallbacks([(
16856                "primary".to_string(),
16857                "fallback".to_string(),
16858            )]))
16859            .build()
16860            .unwrap();
16861        let record = tokio::time::timeout(
16862            std::time::Duration::from_secs(2),
16863            agent.invoke_tool(ToolExecutionRequest::new(
16864                "fallback-call",
16865                "primary",
16866                serde_json::json!({"path": "./shared.txt"}),
16867                ToolCallSource::Manual,
16868            )),
16869        )
16870        .await
16871        .expect("fallback must not retain the primary resource guard")
16872        .unwrap();
16873
16874        assert!(record.success);
16875        assert_eq!(record.canonical_id, "fallback");
16876        assert_eq!(record.call_id, "fallback-call");
16877        assert!(matches!(record.source, ToolCallSource::Fallback { .. }));
16878        assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
16879        assert_eq!(fallback_calls.load(Ordering::SeqCst), 1);
16880    }
16881
16882    #[tokio::test]
16883    async fn diagnostics_without_provider_records_unavailable_without_execution() {
16884        let mock = mock_with_response("hello");
16885        let yaml = r#"
16886name: DiagnosticsNoProviderAgent
16887system_prompt: "Review diagnostics."
16888tools: [diagnostics]
16889"#;
16890        let agent = AgentBuilder::from_yaml(yaml)
16891            .unwrap()
16892            .llm(Arc::new(mock))
16893            .auto_configure_features()
16894            .unwrap()
16895            .build()
16896            .unwrap();
16897
16898        let record = agent
16899            .invoke_tool(ToolExecutionRequest::new(
16900                "diagnostics-call",
16901                "diagnostics",
16902                serde_json::json!({}),
16903                ToolCallSource::Manual,
16904            ))
16905            .await
16906            .unwrap();
16907
16908        assert!(!record.executed);
16909        assert!(!record.success);
16910        assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
16911    }
16912
16913    #[tokio::test]
16914    async fn web_search_without_provider_records_unavailable_without_execution() {
16915        let mock = mock_with_response("hello");
16916        let yaml = r#"
16917name: WebSearchNoProviderAgent
16918system_prompt: "You search the web."
16919tools: [web_search]
16920"#;
16921        let agent = AgentBuilder::from_yaml(yaml)
16922            .unwrap()
16923            .llm(Arc::new(mock))
16924            .auto_configure_features()
16925            .unwrap()
16926            .build()
16927            .unwrap();
16928
16929        let record = agent
16930            .invoke_tool(ToolExecutionRequest::new(
16931                "web-search-call",
16932                "web_search",
16933                serde_json::json!({"query": "rust async"}),
16934                ToolCallSource::Manual,
16935            ))
16936            .await
16937            .unwrap();
16938
16939        assert!(!record.executed);
16940        assert!(!record.success);
16941        assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
16942    }
16943
16944    #[tokio::test]
16945    async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_omitted() {
16946        let mock = mock_with_response("hello");
16947        let yaml = r#"
16948name: SpawnerNoGrantAgent
16949system_prompt: "You manage agents."
16950spawner:
16951  max_agents: 2
16952"#;
16953        let agent = AgentBuilder::from_yaml(yaml)
16954            .unwrap()
16955            .llm(Arc::new(mock))
16956            .auto_configure_features()
16957            .unwrap()
16958            .auto_configure_spawner()
16959            .await
16960            .unwrap()
16961            .build()
16962            .unwrap();
16963
16964        let available = agent.get_available_tool_ids().await.unwrap();
16965        assert!(available.is_empty());
16966    }
16967
16968    #[tokio::test]
16969    async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_empty() {
16970        let mock = mock_with_response("hello");
16971        let yaml = r#"
16972name: EmptySpawnerNoGrantAgent
16973system_prompt: "You manage agents."
16974tools: []
16975spawner:
16976  max_agents: 2
16977"#;
16978        let agent = AgentBuilder::from_yaml(yaml)
16979            .unwrap()
16980            .llm(Arc::new(mock))
16981            .auto_configure_features()
16982            .unwrap()
16983            .auto_configure_spawner()
16984            .await
16985            .unwrap()
16986            .build()
16987            .unwrap();
16988
16989        let available = agent.get_available_tool_ids().await.unwrap();
16990        assert!(available.is_empty());
16991    }
16992
16993    #[tokio::test]
16994    async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_empty() {
16995        let mock = mock_with_response("hello");
16996        let yaml = r#"
16997name: ManagementGrantAgent
16998system_prompt: "You manage agents."
16999tools: []
17000spawner:
17001  management_tools: true
17002"#;
17003        let agent = AgentBuilder::from_yaml(yaml)
17004            .unwrap()
17005            .llm(Arc::new(mock))
17006            .auto_configure_features()
17007            .unwrap()
17008            .auto_configure_spawner()
17009            .await
17010            .unwrap()
17011            .build()
17012            .unwrap();
17013
17014        let available = agent.get_available_tool_ids().await.unwrap();
17015        assert_eq!(available.len(), 4);
17016        assert!(available.contains(&"spawn_agent".to_string()));
17017        assert!(available.contains(&"send_agent_message".to_string()));
17018        assert!(available.contains(&"list_agents".to_string()));
17019        assert!(available.contains(&"remove_agent".to_string()));
17020    }
17021
17022    #[tokio::test]
17023    async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_omitted() {
17024        let mock = mock_with_response("hello");
17025        let yaml = r#"
17026name: ManagementOmittedToolsGrantAgent
17027system_prompt: "You manage agents."
17028spawner:
17029  management_tools: true
17030"#;
17031        let agent = AgentBuilder::from_yaml(yaml)
17032            .unwrap()
17033            .llm(Arc::new(mock))
17034            .auto_configure_features()
17035            .unwrap()
17036            .auto_configure_spawner()
17037            .await
17038            .unwrap()
17039            .build()
17040            .unwrap();
17041
17042        let available = agent.get_available_tool_ids().await.unwrap();
17043        assert_eq!(available.len(), 4);
17044        assert!(available.contains(&"spawn_agent".to_string()));
17045        assert!(available.contains(&"send_agent_message".to_string()));
17046        assert!(available.contains(&"list_agents".to_string()));
17047        assert!(available.contains(&"remove_agent".to_string()));
17048    }
17049
17050    #[tokio::test]
17051    async fn test_management_tools_selected_grants_only_selected_tools() {
17052        let mock = mock_with_response("hello");
17053        let yaml = r#"
17054name: ManagementSelectedGrantAgent
17055system_prompt: "You manage agents."
17056tools: []
17057spawner:
17058  management_tools:
17059    - spawn_agent
17060    - send_agent_message
17061    - list_agents
17062"#;
17063        let agent = AgentBuilder::from_yaml(yaml)
17064            .unwrap()
17065            .llm(Arc::new(mock))
17066            .auto_configure_features()
17067            .unwrap()
17068            .auto_configure_spawner()
17069            .await
17070            .unwrap()
17071            .build()
17072            .unwrap();
17073
17074        let available = agent.get_available_tool_ids().await.unwrap();
17075        assert_eq!(available.len(), 3);
17076        assert!(available.contains(&"spawn_agent".to_string()));
17077        assert!(available.contains(&"send_agent_message".to_string()));
17078        assert!(available.contains(&"list_agents".to_string()));
17079        assert!(!available.contains(&"remove_agent".to_string()));
17080    }
17081
17082    #[tokio::test]
17083    async fn test_orchestration_tools_flag_grants_tools_when_top_level_tools_empty() {
17084        let mock = mock_with_response("hello");
17085        let yaml = r#"
17086name: OrchestrationGrantAgent
17087system_prompt: "You coordinate agents."
17088llms:
17089  default:
17090    provider: openai
17091    model: gpt-4
17092  router:
17093    provider: openai
17094    model: gpt-4
17095llm:
17096  default: default
17097  router: router
17098tools: []
17099spawner:
17100  orchestration_tools: true
17101"#;
17102        let agent = AgentBuilder::from_yaml(yaml)
17103            .unwrap()
17104            .llm(Arc::new(mock))
17105            .auto_configure_features()
17106            .unwrap()
17107            .auto_configure_spawner()
17108            .await
17109            .unwrap()
17110            .build()
17111            .unwrap();
17112
17113        let available = agent.get_available_tool_ids().await.unwrap();
17114        assert_eq!(available.len(), 5);
17115        assert!(available.contains(&"route_to_agent".to_string()));
17116        assert!(available.contains(&"pipeline_process".to_string()));
17117        assert!(available.contains(&"concurrent_ask".to_string()));
17118        assert!(available.contains(&"group_discussion".to_string()));
17119        assert!(available.contains(&"handoff_conversation".to_string()));
17120    }
17121
17122    #[tokio::test]
17123    async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_empty() {
17124        let mock = mock_with_response("hello");
17125        let yaml = r#"
17126name: PersonaGrantAgent
17127system_prompt: "You can evolve persona."
17128llm:
17129  provider: openai
17130  model: gpt-4
17131tools: []
17132persona:
17133  identity:
17134    name: "Guide"
17135    role: "Helper"
17136  evolution:
17137    enabled: true
17138    allow_llm_evolve: true
17139    mutable_fields:
17140      - traits.personality
17141"#;
17142        let agent = AgentBuilder::from_yaml(yaml)
17143            .unwrap()
17144            .llm(Arc::new(mock))
17145            .build()
17146            .unwrap();
17147
17148        let available = agent.get_available_tool_ids().await.unwrap();
17149        assert_eq!(available, vec!["persona_evolve".to_string()]);
17150    }
17151
17152    #[tokio::test]
17153    async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_omitted() {
17154        let mock = mock_with_response("hello");
17155        let yaml = r#"
17156name: PersonaOmittedToolsGrantAgent
17157system_prompt: "You can evolve persona."
17158llm:
17159  provider: openai
17160  model: gpt-4
17161persona:
17162  identity:
17163    name: "Guide"
17164    role: "Helper"
17165  evolution:
17166    enabled: true
17167    allow_llm_evolve: true
17168    mutable_fields:
17169      - traits.personality
17170"#;
17171        let agent = AgentBuilder::from_yaml(yaml)
17172            .unwrap()
17173            .llm(Arc::new(mock))
17174            .build()
17175            .unwrap();
17176
17177        let available = agent.get_available_tool_ids().await.unwrap();
17178        assert_eq!(available, vec!["persona_evolve".to_string()]);
17179    }
17180
17181    #[tokio::test]
17182    async fn test_omitted_yaml_tools_exposes_no_tools() {
17183        let mock = mock_with_response("hello");
17184        let yaml = r#"
17185name: NoToolsAgent
17186system_prompt: "You are helpful."
17187"#;
17188        let agent = AgentBuilder::from_yaml(yaml)
17189            .unwrap()
17190            .llm(Arc::new(mock))
17191            .auto_configure_features()
17192            .unwrap()
17193            .build()
17194            .unwrap();
17195
17196        let available = agent.get_available_tool_ids().await.unwrap();
17197        assert!(available.is_empty());
17198    }
17199
17200    #[tokio::test]
17201    async fn runtime_scope_cannot_widen_omitted_or_empty_yaml_grants() {
17202        for tools in ["", "tools: []"] {
17203            let yaml = format!(
17204                r#"
17205name: RuntimeScopeNoGrantAgent
17206system_prompt: "No ordinary tools are granted."
17207{tools}
17208"#
17209            );
17210            let agent = AgentBuilder::from_yaml(&yaml)
17211                .unwrap()
17212                .llm(Arc::new(mock_with_response("done")))
17213                .auto_configure_features()
17214                .unwrap()
17215                .build()
17216                .unwrap();
17217
17218            agent
17219                .runtime_control()
17220                .set_tool_scope(vec!["calculator".to_string()]);
17221
17222            assert!(agent.get_available_tool_ids().await.unwrap().is_empty());
17223        }
17224    }
17225
17226    #[tokio::test]
17227    async fn runtime_scope_widening_attempt_keeps_only_declared_tools() {
17228        let yaml = r#"
17229name: RuntimeScopeWideningAgent
17230system_prompt: "Runtime scope cannot add authority."
17231tools: [calculator]
17232"#;
17233        let agent = AgentBuilder::from_yaml(yaml)
17234            .unwrap()
17235            .llm(Arc::new(mock_with_response("done")))
17236            .auto_configure_features()
17237            .unwrap()
17238            .build()
17239            .unwrap();
17240
17241        agent
17242            .runtime_control()
17243            .set_tool_scope(vec!["calculator".to_string(), "datetime".to_string()]);
17244
17245        assert_eq!(
17246            agent.get_available_tool_ids().await.unwrap(),
17247            vec!["calculator".to_string()]
17248        );
17249    }
17250
17251    #[tokio::test]
17252    async fn runtime_scope_is_canonical_unique_ordered_and_clear_restores_declared_grant() {
17253        let yaml = r#"
17254name: RuntimeScopeIntersectionAgent
17255system_prompt: "Use only declared tools."
17256tools: [calculator, datetime]
17257"#;
17258        let agent = AgentBuilder::from_yaml(yaml)
17259            .unwrap()
17260            .llm(Arc::new(mock_with_response("done")))
17261            .auto_configure_features()
17262            .unwrap()
17263            .build()
17264            .unwrap();
17265        let mut aliases = ai_agents_tools::ToolAliases::default();
17266        aliases
17267            .names
17268            .insert("en".to_string(), "calculate_alias".to_string());
17269        agent.tools.set_tool_aliases("calculator", aliases);
17270        let control = agent.runtime_control();
17271
17272        control.set_tool_scope(vec![
17273            "datetime".to_string(),
17274            "calculate_alias".to_string(),
17275            "calculator".to_string(),
17276            "unknown".to_string(),
17277            "datetime".to_string(),
17278        ]);
17279        assert_eq!(
17280            agent.get_available_tool_ids().await.unwrap(),
17281            vec!["calculator".to_string(), "datetime".to_string()]
17282        );
17283
17284        control.set_tool_scope(vec!["datetime".to_string()]);
17285        assert_eq!(
17286            agent.get_available_tool_ids().await.unwrap(),
17287            vec!["datetime".to_string()]
17288        );
17289
17290        control.clear_tool_scope_override();
17291        assert_eq!(
17292            agent.get_available_tool_ids().await.unwrap(),
17293            vec!["calculator".to_string(), "datetime".to_string()]
17294        );
17295    }
17296
17297    #[tokio::test]
17298    async fn runtime_scope_preserves_programmatic_registration_as_declared_grant() {
17299        let agent = AgentBuilder::new()
17300            .system_prompt("Use registered tools.")
17301            .llm(Arc::new(mock_with_response("done")))
17302            .tool(Arc::new(ContextEchoTool))
17303            .tool(Arc::new(SlowTool))
17304            .build()
17305            .unwrap();
17306
17307        agent.runtime_control().set_tool_scope(vec![
17308            "Context Echo".to_string(),
17309            "context_echo".to_string(),
17310            "unknown".to_string(),
17311        ]);
17312
17313        assert_eq!(
17314            agent.get_available_tool_ids().await.unwrap(),
17315            vec!["context_echo".to_string()]
17316        );
17317    }
17318
17319    #[tokio::test]
17320    async fn nested_state_scopes_intersect_every_ancestor_with_aliases() {
17321        let yaml = r#"
17322name: NestedStateScopeAgent
17323system_prompt: "Honor every state scope."
17324tools: [calculator, datetime, echo]
17325states:
17326  initial: root
17327  states:
17328    root:
17329      tools: [calculate_alias, datetime]
17330      initial: middle
17331      states:
17332        middle:
17333          initial: leaf
17334          states:
17335            leaf:
17336              tools: [datetime_alias, echo]
17337"#;
17338        let agent = AgentBuilder::from_yaml(yaml)
17339            .unwrap()
17340            .llm(Arc::new(mock_with_response("done")))
17341            .auto_configure_features()
17342            .unwrap()
17343            .build()
17344            .unwrap();
17345        let mut calculator_aliases = ai_agents_tools::ToolAliases::default();
17346        calculator_aliases
17347            .names
17348            .insert("en".to_string(), "calculate_alias".to_string());
17349        agent
17350            .tools
17351            .set_tool_aliases("calculator", calculator_aliases);
17352        let mut datetime_aliases = ai_agents_tools::ToolAliases::default();
17353        datetime_aliases
17354            .names
17355            .insert("en".to_string(), "datetime_alias".to_string());
17356        agent.tools.set_tool_aliases("datetime", datetime_aliases);
17357        agent.runtime_control().set_tool_scope(vec![
17358            "unknown".to_string(),
17359            "datetime_alias".to_string(),
17360            "calculate_alias".to_string(),
17361            "datetime".to_string(),
17362        ]);
17363
17364        assert_eq!(agent.current_state().as_deref(), Some("root.middle.leaf"));
17365        assert_eq!(
17366            agent.get_available_tool_ids().await.unwrap(),
17367            vec!["datetime".to_string()]
17368        );
17369    }
17370
17371    #[tokio::test]
17372    async fn ancestor_empty_state_scope_denies_omitted_descendants() {
17373        let yaml = r#"
17374name: NestedEmptyStateScopeAgent
17375system_prompt: "An empty ancestor scope denies all tools."
17376tools: [calculator]
17377states:
17378  initial: root
17379  states:
17380    root:
17381      tools: []
17382      initial: middle
17383      states:
17384        middle:
17385          initial: leaf
17386          states:
17387            leaf: {}
17388"#;
17389        let agent = AgentBuilder::from_yaml(yaml)
17390            .unwrap()
17391            .llm(Arc::new(mock_with_response("done")))
17392            .auto_configure_features()
17393            .unwrap()
17394            .build()
17395            .unwrap();
17396
17397        assert!(agent.get_available_tool_ids().await.unwrap().is_empty());
17398    }
17399
17400    #[tokio::test]
17401    async fn state_change_during_approval_invalidates_the_reviewed_authority() {
17402        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17403        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17404        let entered = Arc::new(tokio::sync::Barrier::new(2));
17405        let release = Arc::new(tokio::sync::Notify::new());
17406        let handler = Arc::new(BlockingApprovalHandler {
17407            entered: Arc::clone(&entered),
17408            release: Arc::clone(&release),
17409            result: ApprovalResult::Approved,
17410        });
17411        let yaml = r#"
17412name: ApprovalStateGenerationAgent
17413system_prompt: "State authority may change during approval."
17414tools: [locked_write]
17415states:
17416  initial: first
17417  states:
17418    first:
17419      tools: [locked_write]
17420    second:
17421      tools: [locked_write]
17422"#;
17423        let agent = Arc::new(
17424            AgentBuilder::from_yaml(yaml)
17425                .unwrap()
17426                .llm(Arc::new(mock_with_response("done")))
17427                .tool(Arc::new(LockedWriteTool {
17428                    active: Arc::clone(&active),
17429                    max_active: Arc::clone(&max_active),
17430                }))
17431                .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
17432                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17433                .approval_handler(handler)
17434                .build()
17435                .unwrap(),
17436        );
17437        let running = Arc::clone(&agent);
17438        let call = tokio::spawn(async move {
17439            running
17440                .invoke_tool(ToolExecutionRequest::new(
17441                    "approval-state-generation",
17442                    "locked_write",
17443                    serde_json::json!({"path": "./state-generation.txt"}),
17444                    ToolCallSource::Manual,
17445                ))
17446                .await
17447                .unwrap()
17448        });
17449
17450        entered.wait().await;
17451        agent.transition_to("second").await.unwrap();
17452        release.notify_one();
17453        let record = call.await.unwrap();
17454
17455        assert!(!record.executed);
17456        assert!(record.output.contains("Approval became stale"));
17457        assert_eq!(max_active.load(Ordering::SeqCst), 0);
17458    }
17459
17460    #[tokio::test]
17461    async fn state_change_while_waiting_for_resource_lock_fails_final_admission() {
17462        let holder_gate = PathMutationGate::new();
17463        let waiter_gate = PathMutationGate::new();
17464        let yaml = r#"
17465name: LockedStateGenerationAgent
17466system_prompt: "State authority must remain stable through admission."
17467tools: [state_lock_holder, state_lock_waiter]
17468states:
17469  initial: first
17470  states:
17471    first:
17472      tools: [state_lock_holder, state_lock_waiter]
17473    second:
17474      tools: [state_lock_holder, state_lock_waiter]
17475"#;
17476        let agent = Arc::new(
17477            AgentBuilder::from_yaml(yaml)
17478                .unwrap()
17479                .llm(Arc::new(mock_with_response("done")))
17480                .tool(Arc::new(BlockingPathMutationTool {
17481                    id: "state_lock_holder",
17482                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
17483                    gate: holder_gate.clone(),
17484                }))
17485                .tool(Arc::new(BlockingPathMutationTool {
17486                    id: "state_lock_waiter",
17487                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
17488                    gate: waiter_gate.clone(),
17489                }))
17490                .build()
17491                .unwrap(),
17492        );
17493        let holder_call = {
17494            let agent = Arc::clone(&agent);
17495            tokio::spawn(async move {
17496                agent
17497                    .invoke_tool(ToolExecutionRequest::new(
17498                        "state-lock-holder",
17499                        "state_lock_holder",
17500                        serde_json::json!({"path": "./shared-state-path.txt"}),
17501                        ToolCallSource::Manual,
17502                    ))
17503                    .await
17504                    .unwrap()
17505            })
17506        };
17507        holder_gate.wait_until_entered().await;
17508        let waiter_call = {
17509            let agent = Arc::clone(&agent);
17510            tokio::spawn(async move {
17511                agent
17512                    .invoke_tool(ToolExecutionRequest::new(
17513                        "state-lock-waiter",
17514                        "state_lock_waiter",
17515                        serde_json::json!({"path": "./shared-state-path.txt"}),
17516                        ToolCallSource::Manual,
17517                    ))
17518                    .await
17519                    .unwrap()
17520            })
17521        };
17522
17523        wait_for_resource_lock_strong_count(&agent.resource_locks, 2).await;
17524        agent.transition_to("second").await.unwrap();
17525        holder_gate.release();
17526        let holder_record = holder_call.await.unwrap();
17527        let waiter_record = waiter_call.await.unwrap();
17528
17529        assert!(holder_record.success);
17530        assert!(!waiter_record.executed);
17531        assert!(
17532            waiter_record
17533                .output
17534                .contains("state scope changed before admission")
17535        );
17536        assert!(!waiter_gate.entered.load(Ordering::SeqCst));
17537    }
17538
17539    #[tokio::test]
17540    async fn test_state_tools_cannot_widen_top_level_grant() {
17541        let mock = mock_with_response("hello");
17542        let yaml = r#"
17543name: NarrowToolsAgent
17544system_prompt: "You are helpful."
17545tools:
17546  - calculator
17547states:
17548  initial: current
17549  states:
17550    current:
17551      tools: [datetime]
17552"#;
17553        let agent = AgentBuilder::from_yaml(yaml)
17554            .unwrap()
17555            .llm(Arc::new(mock))
17556            .auto_configure_features()
17557            .unwrap()
17558            .build()
17559            .unwrap();
17560
17561        let available = agent.get_available_tool_ids().await.unwrap();
17562        assert!(available.is_empty());
17563    }
17564
17565    // Tool execution in chat flow
17566    #[tokio::test]
17567    async fn test_integration_tool_execution() {
17568        // Mock LLM that returns a tool call then a final answer
17569        let mock = mock_with_responses(vec![
17570            // First response: tool call
17571            r#"I'll calculate that for you.
17572[TOOL_CALL: {"name": "calculator", "arguments": {"expression": "2+2"}}]"#,
17573            // After tool result: final answer
17574            "The answer is 4.",
17575        ]);
17576        let mut tools = ai_agents_tools::ToolRegistry::new();
17577        tools
17578            .register(Arc::new(ai_agents_tools::CalculatorTool))
17579            .unwrap();
17580
17581        let agent = AgentBuilder::new()
17582            .system_prompt("You are a calculator assistant.")
17583            .llm(Arc::new(mock))
17584            .tools(tools)
17585            .build()
17586            .unwrap();
17587
17588        let response = agent.chat("What is 2+2?").await.unwrap();
17589        // The agent should eventually produce a response
17590        assert!(!response.content.is_empty());
17591    }
17592
17593    #[tokio::test]
17594    async fn test_tool_hitl_rejection_finalizes_blocking_turn() {
17595        let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17596        let hooks = Arc::new(ResponseCountingHooks {
17597            responses: Arc::clone(&responses),
17598        });
17599        let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
17600        let yaml = r#"
17601name: ToolRejectAgent
17602system_prompt: "You use tools when requested."
17603tools:
17604  - echo
17605hitl:
17606  tools:
17607    echo:
17608      require_approval: true
17609      approval_message: "Approve echo?"
17610"#;
17611        let agent = AgentBuilder::from_yaml(yaml)
17612            .unwrap()
17613            .llm(Arc::new(mock))
17614            .auto_configure_features()
17615            .unwrap()
17616            .hooks(hooks)
17617            .build()
17618            .unwrap();
17619
17620        let response = agent.chat("echo hello").await.unwrap();
17621
17622        assert!(
17623            response.content.contains("Operation cancelled"),
17624            "unexpected response: {}",
17625            response.content
17626        );
17627        assert_eq!(responses.load(Ordering::SeqCst), 1);
17628        let messages = agent.memory.get_messages(None).await.unwrap();
17629        assert_eq!(messages.len(), 3);
17630        assert_eq!(messages[0].content, "echo hello");
17631        assert!(messages[1].content.contains("\"tool\":\"echo\""));
17632        assert!(messages[2].content.contains("rejected by the approver"));
17633    }
17634
17635    #[tokio::test]
17636    async fn test_tool_hitl_rejection_finalizes_streaming_turn() {
17637        use futures::StreamExt;
17638
17639        let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17640        let hooks = Arc::new(ResponseCountingHooks {
17641            responses: Arc::clone(&responses),
17642        });
17643        let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
17644        let yaml = r#"
17645name: ToolRejectStreamingAgent
17646system_prompt: "You use tools when requested."
17647tools:
17648  - echo
17649streaming:
17650  enabled: true
17651hitl:
17652  tools:
17653    echo:
17654      require_approval: true
17655      approval_message: "Approve echo?"
17656"#;
17657        let agent = AgentBuilder::from_yaml(yaml)
17658            .unwrap()
17659            .llm(Arc::new(mock))
17660            .auto_configure_features()
17661            .unwrap()
17662            .hooks(hooks)
17663            .build()
17664            .unwrap();
17665
17666        let mut stream = agent.chat_stream("echo hello").await.unwrap();
17667        let mut terminal_error = String::new();
17668        let mut done = false;
17669        while let Some(chunk) = stream.next().await {
17670            match chunk {
17671                StreamChunk::Error { message } => terminal_error = message,
17672                StreamChunk::Done {} => {
17673                    done = true;
17674                    break;
17675                }
17676                _ => {}
17677            }
17678        }
17679
17680        assert!(done);
17681        assert!(
17682            terminal_error.contains("Operation cancelled"),
17683            "unexpected terminal error: {}",
17684            terminal_error
17685        );
17686        assert_eq!(responses.load(Ordering::SeqCst), 1);
17687        let messages = agent.memory.get_messages(None).await.unwrap();
17688        assert_eq!(messages.len(), 3);
17689        assert_eq!(messages[0].content, "echo hello");
17690        assert!(messages[1].content.contains("\"tool\":\"echo\""));
17691        assert!(messages[2].content.contains("rejected by the approver"));
17692    }
17693
17694    #[tokio::test]
17695    async fn tool_hitl_rejection_preserves_legacy_error_but_finalizes_event_stream() {
17696        let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
17697        let yaml = r#"
17698name: ToolRejectEventAgent
17699system_prompt: "You use tools when requested."
17700tools:
17701  - echo
17702streaming:
17703  enabled: true
17704hitl:
17705  tools:
17706    echo:
17707      require_approval: true
17708      approval_message: "Approve echo?"
17709"#;
17710        let agent = AgentBuilder::from_yaml(yaml)
17711            .unwrap()
17712            .llm(Arc::new(mock))
17713            .auto_configure_features()
17714            .unwrap()
17715            .build()
17716            .unwrap();
17717
17718        let mut stream = agent.chat_stream_events("echo hello").await.unwrap();
17719        let mut error_seen = false;
17720        let mut final_response = None;
17721        while let Some(event) = stream.next().await {
17722            match event {
17723                AgentStreamEvent::Chunk(StreamChunk::Error { .. }) => error_seen = true,
17724                AgentStreamEvent::Final(response) => final_response = Some(response),
17725                AgentStreamEvent::Chunk(_) => {}
17726            }
17727        }
17728
17729        assert!(!error_seen);
17730        assert!(
17731            final_response
17732                .is_some_and(|response| { response.content.contains("Operation cancelled") })
17733        );
17734    }
17735
17736    #[tokio::test]
17737    async fn test_pre_response_guard_transition_skips_old_state_llm() {
17738        let mock = mock_with_response("Billing state response");
17739        let call_counter = mock.clone();
17740        let yaml = r#"
17741name: OptimizedStateAgent
17742system_prompt: "You route before answering."
17743runtime:
17744  optimization:
17745    enabled: true
17746    pre_response_deterministic_transitions: true
17747states:
17748  initial: greeting
17749  states:
17750    greeting:
17751      prompt: "Old state prompt that should be skipped."
17752      transitions:
17753        - to: billing
17754          guard:
17755            context:
17756              topic:
17757                eq: billing
17758          timing: pre_response
17759    billing:
17760      prompt: "Answer from the billing state."
17761"#;
17762        let agent = AgentBuilder::from_yaml(yaml)
17763            .unwrap()
17764            .llm(Arc::new(mock))
17765            .build()
17766            .unwrap();
17767        agent
17768            .set_context("topic", serde_json::json!("billing"))
17769            .unwrap();
17770
17771        let response = agent.chat("I need billing help").await.unwrap();
17772
17773        assert_eq!(agent.current_state().as_deref(), Some("billing"));
17774        assert_eq!(response.content, "Billing state response");
17775        assert_eq!(call_counter.call_count(), 1);
17776        assert_eq!(agent.actor_facts().len(), 0);
17777    }
17778
17779    #[tokio::test]
17780    async fn test_set_context_supports_dotted_paths_for_pre_response_guards() {
17781        let mock = mock_with_response("Billing state response");
17782        let call_counter = mock.clone();
17783        let yaml = r#"
17784name: OptimizedStateAgent
17785system_prompt: "You route before answering."
17786runtime:
17787  optimization:
17788    enabled: true
17789    pre_response_deterministic_transitions: true
17790context:
17791  request:
17792    type: runtime
17793    default:
17794      topic: general
17795states:
17796  initial: greeting
17797  states:
17798    greeting:
17799      prompt: "Old state prompt that should be skipped."
17800      transitions:
17801        - to: billing
17802          guard:
17803            context:
17804              request.topic:
17805                eq: billing
17806          timing: pre_response
17807    billing:
17808      prompt: "Answer from the billing state."
17809"#;
17810        let agent = AgentBuilder::from_yaml(yaml)
17811            .unwrap()
17812            .llm(Arc::new(mock))
17813            .build()
17814            .unwrap();
17815        agent
17816            .set_context("request.topic", serde_json::json!("billing"))
17817            .unwrap();
17818
17819        let response = agent.chat("I need billing help").await.unwrap();
17820
17821        assert_eq!(agent.current_state().as_deref(), Some("billing"));
17822        assert_eq!(response.content, "Billing state response");
17823        assert_eq!(call_counter.call_count(), 1);
17824        assert_eq!(
17825            agent.get_context().get("request"),
17826            Some(&serde_json::json!({"topic": "billing"}))
17827        );
17828    }
17829
17830    #[tokio::test]
17831    async fn test_pre_response_rejection_does_not_commit_staged_context_or_user() {
17832        let mock = mock_with_response("billing");
17833        let yaml = r#"
17834name: OptimizedStateAgent
17835system_prompt: "You route before answering."
17836runtime:
17837  optimization:
17838    enabled: true
17839    pre_response_deterministic_transitions: true
17840hitl:
17841  states:
17842    billing:
17843      on_enter: require_approval
17844      approval_message: "Approve billing route?"
17845states:
17846  initial: greeting
17847  states:
17848    greeting:
17849      prompt: "Old state prompt."
17850      extract:
17851        - key: topic
17852          description: "Support topic"
17853      transitions:
17854        - to: billing
17855          guard:
17856            context:
17857              topic:
17858                eq: billing
17859          timing: pre_response
17860          run_extractors: true
17861    billing:
17862      prompt: "Billing state."
17863"#;
17864        let agent = AgentBuilder::from_yaml(yaml)
17865            .unwrap()
17866            .llm(Arc::new(mock))
17867            .build()
17868            .unwrap();
17869
17870        let response = agent
17871            .try_pre_response_transition("billing please")
17872            .await
17873            .unwrap();
17874
17875        assert!(response.is_none());
17876        assert_eq!(agent.current_state().as_deref(), Some("greeting"));
17877        assert!(!agent.get_context().contains_key("topic"));
17878        assert_eq!(agent.memory.get_messages(None).await.unwrap().len(), 0);
17879    }
17880
17881    #[tokio::test]
17882    async fn test_pre_response_extractor_commits_context_on_winning_path() {
17883        let mock = mock_with_responses(vec!["billing", "Billing response"]);
17884        let yaml = r#"
17885name: OptimizedStateAgent
17886system_prompt: "You route before answering."
17887runtime:
17888  optimization:
17889    enabled: true
17890    pre_response_deterministic_transitions: true
17891states:
17892  initial: greeting
17893  states:
17894    greeting:
17895      prompt: "Old state prompt."
17896      extract:
17897        - key: topic
17898          description: "Support topic"
17899      transitions:
17900        - to: billing
17901          guard:
17902            context:
17903              topic:
17904                eq: billing
17905          timing: pre_response
17906          run_extractors: true
17907    billing:
17908      prompt: "Billing state."
17909"#;
17910        let agent = AgentBuilder::from_yaml(yaml)
17911            .unwrap()
17912            .llm(Arc::new(mock))
17913            .build()
17914            .unwrap();
17915
17916        let response = agent.chat("billing please").await.unwrap();
17917
17918        assert_eq!(agent.current_state().as_deref(), Some("billing"));
17919        assert_eq!(response.content, "Billing response");
17920        assert_eq!(
17921            agent.get_context().get("topic"),
17922            Some(&serde_json::json!("billing"))
17923        );
17924    }
17925
17926    #[tokio::test]
17927    async fn test_pre_response_extractor_miss_does_not_mutate_context() {
17928        let mock = mock_with_response("__NONE__");
17929        let yaml = r#"
17930name: OptimizedStateAgent
17931system_prompt: "You route before answering."
17932runtime:
17933  optimization:
17934    enabled: true
17935    pre_response_deterministic_transitions: true
17936states:
17937  initial: greeting
17938  states:
17939    greeting:
17940      prompt: "Old state prompt."
17941      extract:
17942        - key: topic
17943          description: "Support topic"
17944      transitions:
17945        - to: billing
17946          guard:
17947            context:
17948              topic:
17949                eq: billing
17950          timing: pre_response
17951          run_extractors: true
17952    billing:
17953      prompt: "Billing state."
17954"#;
17955        let agent = AgentBuilder::from_yaml(yaml)
17956            .unwrap()
17957            .llm(Arc::new(mock))
17958            .build()
17959            .unwrap();
17960
17961        let response = agent.try_pre_response_transition("hello").await.unwrap();
17962
17963        assert!(response.is_none());
17964        assert_eq!(agent.current_state().as_deref(), Some("greeting"));
17965        assert!(!agent.get_context().contains_key("topic"));
17966    }
17967
17968    #[tokio::test]
17969    async fn test_default_guard_transition_stays_post_response() {
17970        let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
17971        let call_counter = mock.clone();
17972        let yaml = r#"
17973name: TimingAgent
17974system_prompt: "You route carefully."
17975runtime:
17976  optimization:
17977    enabled: true
17978    pre_response_deterministic_transitions: true
17979states:
17980  initial: greeting
17981  states:
17982    greeting:
17983      prompt: "Old state prompt."
17984      transitions:
17985        - to: billing
17986          guard:
17987            context:
17988              topic:
17989                eq: billing
17990    billing:
17991      prompt: "Billing state."
17992"#;
17993        let agent = AgentBuilder::from_yaml(yaml)
17994            .unwrap()
17995            .llm(Arc::new(mock))
17996            .build()
17997            .unwrap();
17998        agent
17999            .set_context("topic", serde_json::json!("billing"))
18000            .unwrap();
18001
18002        let response = agent.chat("billing please").await.unwrap();
18003
18004        assert_eq!(agent.current_state().as_deref(), Some("billing"));
18005        assert_eq!(response.content, "Billing response");
18006        assert_eq!(call_counter.call_count(), 2);
18007    }
18008
18009    #[tokio::test]
18010    async fn test_explicit_post_response_guard_transition_stays_post_response() {
18011        let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
18012        let call_counter = mock.clone();
18013        let yaml = r#"
18014name: TimingAgent
18015system_prompt: "You route carefully."
18016runtime:
18017  optimization:
18018    enabled: true
18019    pre_response_deterministic_transitions: true
18020states:
18021  initial: greeting
18022  states:
18023    greeting:
18024      prompt: "Old state prompt."
18025      transitions:
18026        - to: billing
18027          guard:
18028            context:
18029              topic:
18030                eq: billing
18031          timing: post_response
18032    billing:
18033      prompt: "Billing state."
18034"#;
18035        let agent = AgentBuilder::from_yaml(yaml)
18036            .unwrap()
18037            .llm(Arc::new(mock))
18038            .build()
18039            .unwrap();
18040        agent
18041            .set_context("topic", serde_json::json!("billing"))
18042            .unwrap();
18043
18044        let response = agent.chat("billing please").await.unwrap();
18045
18046        assert_eq!(agent.current_state().as_deref(), Some("billing"));
18047        assert_eq!(response.content, "Billing response");
18048        assert_eq!(call_counter.call_count(), 2);
18049    }
18050
18051    #[tokio::test]
18052    async fn test_pre_response_extractors_are_transition_scoped() {
18053        let mock = mock_with_responses(vec!["billing", "Billing response"]);
18054        let yaml = r#"
18055name: ScopedExtractorAgent
18056system_prompt: "You route carefully."
18057runtime:
18058  optimization:
18059    enabled: true
18060    pre_response_deterministic_transitions: true
18061states:
18062  initial: greeting
18063  states:
18064    greeting:
18065      prompt: "Old state prompt."
18066      extract:
18067        - key: topic
18068          description: "Support topic"
18069      transitions:
18070        - to: wrong
18071          guard:
18072            context:
18073              topic:
18074                eq: billing
18075          timing: pre_response
18076        - to: billing
18077          guard:
18078            context:
18079              topic:
18080                eq: billing
18081          timing: pre_response
18082          run_extractors: true
18083    wrong:
18084      prompt: "Wrong state."
18085    billing:
18086      prompt: "Billing state."
18087"#;
18088        let agent = AgentBuilder::from_yaml(yaml)
18089            .unwrap()
18090            .llm(Arc::new(mock))
18091            .build()
18092            .unwrap();
18093
18094        let response = agent.chat("billing please").await.unwrap();
18095
18096        assert_eq!(agent.current_state().as_deref(), Some("billing"));
18097        assert_eq!(response.content, "Billing response");
18098    }
18099
18100    #[tokio::test]
18101    async fn test_pre_response_resolved_intent_routes_early() {
18102        let mock = mock_with_response("Billing response");
18103        let yaml = r#"
18104name: IntentAgent
18105system_prompt: "You route carefully."
18106runtime:
18107  optimization:
18108    enabled: true
18109    pre_response_deterministic_transitions: true
18110states:
18111  initial: greeting
18112  states:
18113    greeting:
18114      prompt: "Old state prompt."
18115      transitions:
18116        - to: billing
18117          intent: billing
18118          timing: pre_response
18119    billing:
18120      prompt: "Billing state."
18121"#;
18122        let agent = AgentBuilder::from_yaml(yaml)
18123            .unwrap()
18124            .llm(Arc::new(mock))
18125            .build()
18126            .unwrap();
18127        agent
18128            .set_context("resolved_intent", serde_json::json!("billing"))
18129            .unwrap();
18130
18131        let response = agent
18132            .try_pre_response_transition("I need billing help")
18133            .await
18134            .unwrap()
18135            .unwrap();
18136
18137        assert_eq!(agent.current_state().as_deref(), Some("billing"));
18138        assert_eq!(response.content, "Billing response");
18139    }
18140
18141    #[tokio::test]
18142    async fn test_background_overflow_error_surfaces() {
18143        let mut config = RuntimeConfig::default();
18144        config.optimization.enabled = true;
18145        config.optimization.post_turn.max_background_tasks = 1;
18146        config.optimization.post_turn.on_background_overflow = BackgroundOverflowPolicy::Error;
18147        let policy = crate::optimization::MaintenanceTaskPolicy {
18148            mode: MaintenanceMode::Background,
18149            await_before_next_turn: AwaitBeforeNextTurn::Always,
18150        };
18151        let agent = AgentBuilder::new()
18152            .system_prompt("You are helpful.")
18153            .llm(Arc::new(mock_with_response("ok")))
18154            .build()
18155            .unwrap()
18156            .with_runtime_config(config);
18157        agent
18158            .background_maintenance
18159            .spawn(None, async { std::future::pending::<Result<()>>().await })
18160            .unwrap();
18161
18162        let result = agent
18163            .spawn_or_handle_background(None, async { Ok(()) }, "facts", &policy)
18164            .await;
18165
18166        assert!(result.is_err());
18167    }
18168
18169    #[tokio::test]
18170    async fn test_speculative_reasoning_low_cap_uses_serial_reasoning() {
18171        let default_mock = mock_with_response("Plain draft response");
18172        let router_mock = mock_with_response("cot");
18173        let router_counter = router_mock.clone();
18174        let yaml = r#"
18175name: ReasoningReservationAgent
18176system_prompt: "You answer plainly unless reasoning wins."
18177llm:
18178  default: default
18179  router: router
18180observability:
18181  enabled: true
18182  export:
18183    write_raw_events: true
18184reasoning:
18185  mode: auto
18186  judge_llm: router
18187runtime:
18188  optimization:
18189    enabled: true
18190    max_speculative_llm_calls_per_turn: 1
18191    speculative_reasoning_auto: true
18192    max_parallel_runtime_tasks: 2
18193"#;
18194        let agent = AgentBuilder::from_yaml(yaml)
18195            .unwrap()
18196            .llm_alias("default", Arc::new(default_mock))
18197            .llm_alias("router", Arc::new(router_mock))
18198            .build()
18199            .unwrap();
18200
18201        let response = agent.chat("hello").await.unwrap();
18202
18203        assert_eq!(response.content, "Plain draft response");
18204        assert_eq!(router_counter.call_count(), 1);
18205        let events = agent.observability().unwrap().raw_events();
18206        assert!(!events.iter().any(|event| {
18207            event.dimensions.get("commit_behavior") == Some(&"reasoning_decision".to_string())
18208        }));
18209    }
18210
18211    #[tokio::test]
18212    async fn test_forced_reasoning_skips_plain_speculative_draft() {
18213        let mock = mock_with_response("Reasoned response");
18214        let yaml = r#"
18215name: ForcedReasoningAgent
18216system_prompt: "You reason before answering."
18217observability:
18218  enabled: true
18219  export:
18220    write_raw_events: true
18221reasoning:
18222  mode: cot
18223runtime:
18224  optimization:
18225    enabled: true
18226    max_speculative_llm_calls_per_turn: 2
18227    speculative_state_transitions: true
18228    max_parallel_runtime_tasks: 2
18229states:
18230  initial: triage
18231  states:
18232    triage:
18233      prompt: "Answer from triage."
18234      transitions:
18235        - to: billing
18236          guard:
18237            context:
18238              route:
18239                eq: billing
18240          timing: parallel
18241    billing:
18242      prompt: "Billing state."
18243"#;
18244        let agent = AgentBuilder::from_yaml(yaml)
18245            .unwrap()
18246            .llm(Arc::new(mock))
18247            .build()
18248            .unwrap();
18249
18250        let response = agent.chat("hello").await.unwrap();
18251
18252        assert_eq!(response.content, "Reasoned response");
18253        let events = agent.observability().unwrap().raw_events();
18254        assert!(
18255            !events
18256                .iter()
18257                .any(|event| event.dimensions.contains_key("branch_status"))
18258        );
18259    }
18260
18261    #[tokio::test]
18262    async fn test_speculative_skill_low_cap_uses_serial_skill_route() {
18263        let default_mock = mock_with_response("Skill committed response");
18264        let router_mock = mock_with_response("helper");
18265        let router_counter = router_mock.clone();
18266        let yaml = r#"
18267name: SkillReservationAgent
18268system_prompt: "Use skills when they match."
18269llm:
18270  default: default
18271  router: router
18272observability:
18273  enabled: true
18274  export:
18275    write_raw_events: true
18276runtime:
18277  optimization:
18278    enabled: true
18279    max_speculative_llm_calls_per_turn: 1
18280    speculative_skill_routing: true
18281    max_parallel_runtime_tasks: 2
18282skills:
18283  - id: helper
18284    description: "Answer helper requests"
18285    trigger: "User asks for helper"
18286    steps:
18287      - prompt: "Answer the helper request: {{ user_input }}"
18288"#;
18289        let agent = AgentBuilder::from_yaml(yaml)
18290            .unwrap()
18291            .llm_alias("default", Arc::new(default_mock))
18292            .llm_alias("router", Arc::new(router_mock))
18293            .build()
18294            .unwrap();
18295
18296        let response = agent.chat("please use helper").await.unwrap();
18297
18298        assert_eq!(response.content, "Skill committed response");
18299        assert_eq!(router_counter.call_count(), 1);
18300        let events = agent.observability().unwrap().raw_events();
18301        assert!(
18302            !events
18303                .iter()
18304                .any(|event| event.dimensions.contains_key("branch_status"))
18305        );
18306    }
18307
18308    #[tokio::test]
18309    async fn test_parallel_transition_low_cap_allows_deterministic_route() {
18310        let mock = mock_with_response("unused");
18311        let call_counter = mock.clone();
18312        let yaml = r#"
18313name: ParallelTransitionLowCapAgent
18314system_prompt: "Route before stale responses when safe."
18315runtime:
18316  optimization:
18317    enabled: true
18318    max_speculative_llm_calls_per_turn: 1
18319    speculative_state_transitions: true
18320    max_parallel_runtime_tasks: 2
18321states:
18322  initial: triage
18323  states:
18324    triage:
18325      prompt: "Triage state."
18326      transitions:
18327        - to: billing
18328          guard:
18329            context:
18330              route:
18331                eq: billing
18332          timing: parallel
18333    billing:
18334      prompt: "Billing state."
18335"#;
18336        let agent = AgentBuilder::from_yaml(yaml)
18337            .unwrap()
18338            .llm(Arc::new(mock))
18339            .build()
18340            .unwrap();
18341        agent
18342            .set_context("route", serde_json::json!("billing"))
18343            .unwrap();
18344        agent.update_active_turn_context("billing help", HashMap::new());
18345        assert!(
18346            agent.reserve_active_speculative_llm_call(
18347                RuntimeOptimizationKind::ParallelStateTransition
18348            )
18349        );
18350
18351        let selection = agent
18352            .select_parallel_transition_candidate("billing help")
18353            .await
18354            .unwrap();
18355        agent.end_root_turn();
18356
18357        match selection {
18358            ParallelTransitionSelection::Candidate(candidate) => {
18359                assert_eq!(candidate.target(), "billing");
18360            }
18361            ParallelTransitionSelection::NoMatch => panic!("deterministic route did not match"),
18362            ParallelTransitionSelection::ReservationExhausted => {
18363                panic!("deterministic route consumed LLM budget")
18364            }
18365        }
18366        assert_eq!(call_counter.call_count(), 0);
18367    }
18368
18369    #[tokio::test]
18370    async fn speculative_transition_drops_loser_before_state_actions() {
18371        let lock = Arc::new(tokio::sync::Mutex::new(()));
18372        let first_started = Arc::new(tokio::sync::Notify::new());
18373        let first_dropped = Arc::new(AtomicBool::new(false));
18374        let committed_after_drop = Arc::new(AtomicBool::new(false));
18375        let default = Arc::new(FirstCallLockingProvider {
18376            lock,
18377            first_started: Arc::clone(&first_started),
18378            first_dropped: Arc::clone(&first_dropped),
18379            committed_after_drop: Arc::clone(&committed_after_drop),
18380            calls: AtomicU64::new(0),
18381        });
18382        let router = Arc::new(RoutingAfterProviderStart {
18383            provider_started: first_started,
18384        });
18385        let yaml = r#"
18386name: SpeculativeCancellationAgent
18387system_prompt: "Route before committed work."
18388llm:
18389  default: default
18390  router: router
18391runtime:
18392  optimization:
18393    enabled: true
18394    max_speculative_llm_calls_per_turn: 2
18395    speculative_state_transitions: true
18396    max_parallel_runtime_tasks: 2
18397states:
18398  initial: triage
18399  states:
18400    triage:
18401      prompt: "Triage state."
18402      transitions:
18403        - to: technical
18404          when: "The request needs technical support"
18405          timing: parallel
18406    technical:
18407      prompt: "Technical state."
18408      on_enter:
18409        - prompt: "Prepare technical context."
18410          llm: default
18411          store_as: preparation
18412"#;
18413        let agent = AgentBuilder::from_yaml(yaml)
18414            .unwrap()
18415            .llm_alias("default", default)
18416            .llm_alias("router", router)
18417            .build()
18418            .unwrap();
18419
18420        let response = tokio::time::timeout(
18421            std::time::Duration::from_secs(2),
18422            agent.chat("I cannot log in because of AUTH-17."),
18423        )
18424        .await
18425        .expect("committed work must not wait on the losing provider future")
18426        .unwrap();
18427
18428        assert_eq!(response.content, "Committed technical response.");
18429        assert_eq!(agent.current_state().as_deref(), Some("technical"));
18430        assert!(first_dropped.load(Ordering::SeqCst));
18431        assert!(committed_after_drop.load(Ordering::SeqCst));
18432    }
18433
18434    #[tokio::test]
18435    async fn buffered_transition_drops_stale_stream_before_redispatch() {
18436        use futures::StreamExt;
18437
18438        let lock = Arc::new(tokio::sync::Mutex::new(()));
18439        let stream_started = Arc::new(tokio::sync::Notify::new());
18440        let stream_dropped = Arc::new(AtomicBool::new(false));
18441        let committed_after_drop = Arc::new(AtomicBool::new(false));
18442        let default = Arc::new(BufferedLockingProvider {
18443            lock,
18444            stream_started: Arc::clone(&stream_started),
18445            stream_dropped: Arc::clone(&stream_dropped),
18446            committed_after_drop: Arc::clone(&committed_after_drop),
18447        });
18448        let router = Arc::new(RoutingAfterProviderStart {
18449            provider_started: stream_started,
18450        });
18451        let yaml = r#"
18452name: BufferedCancellationAgent
18453system_prompt: "Hide stale streamed output."
18454llm:
18455  default: default
18456  router: router
18457streaming:
18458  enabled: true
18459  buffer_size: 8
18460runtime:
18461  optimization:
18462    enabled: true
18463    max_speculative_llm_calls_per_turn: 2
18464    speculative_state_transitions: true
18465    streaming_policy: buffer_until_routing_done
18466    max_parallel_runtime_tasks: 2
18467states:
18468  initial: triage
18469  states:
18470    triage:
18471      prompt: "Triage state."
18472      transitions:
18473        - to: technical
18474          when: "The request needs technical support"
18475          timing: parallel
18476    technical:
18477      prompt: "Technical state."
18478"#;
18479        let agent = AgentBuilder::from_yaml(yaml)
18480            .unwrap()
18481            .llm_alias("default", default)
18482            .llm_alias("router", router)
18483            .build()
18484            .unwrap();
18485
18486        let content = tokio::time::timeout(std::time::Duration::from_secs(2), async {
18487            let mut stream = agent
18488                .chat_stream("AUTH-17 needs technical help.")
18489                .await
18490                .unwrap();
18491            let mut content = String::new();
18492            while let Some(chunk) = stream.next().await {
18493                match chunk {
18494                    StreamChunk::Content { text } => content.push_str(&text),
18495                    StreamChunk::Done {} => break,
18496                    StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
18497                    _ => {}
18498                }
18499            }
18500            content
18501        })
18502        .await
18503        .expect("redispatch must not wait on the stale streaming future");
18504
18505        assert_eq!(content, "Committed technical response.");
18506        assert_eq!(agent.current_state().as_deref(), Some("technical"));
18507        assert!(stream_dropped.load(Ordering::SeqCst));
18508        assert!(committed_after_drop.load(Ordering::SeqCst));
18509    }
18510
18511    #[tokio::test]
18512    async fn buffered_transition_drops_established_stream_before_redispatch() {
18513        use futures::StreamExt;
18514
18515        let stream_started = Arc::new(tokio::sync::Notify::new());
18516        let stream_dropped = Arc::new(AtomicBool::new(false));
18517        let stream_dropped_notify = Arc::new(tokio::sync::Notify::new());
18518        let committed_after_drop = Arc::new(AtomicBool::new(false));
18519        let default = Arc::new(EstablishedStreamProvider {
18520            stream_started: Arc::clone(&stream_started),
18521            stream_dropped: Arc::clone(&stream_dropped),
18522            stream_dropped_notify,
18523            committed_after_drop: Arc::clone(&committed_after_drop),
18524        });
18525        let router = Arc::new(RoutingAfterProviderStart {
18526            provider_started: stream_started,
18527        });
18528        let yaml = r#"
18529name: EstablishedStreamCancellationAgent
18530system_prompt: "Hide stale streamed output."
18531llm:
18532  default: default
18533  router: router
18534streaming:
18535  enabled: true
18536  buffer_size: 8
18537runtime:
18538  optimization:
18539    enabled: true
18540    max_speculative_llm_calls_per_turn: 2
18541    speculative_state_transitions: true
18542    streaming_policy: buffer_until_routing_done
18543    max_parallel_runtime_tasks: 2
18544states:
18545  initial: triage
18546  states:
18547    triage:
18548      prompt: "Triage state."
18549      transitions:
18550        - to: technical
18551          when: "The request needs technical support"
18552          timing: parallel
18553    technical:
18554      prompt: "Technical state."
18555"#;
18556        let agent = AgentBuilder::from_yaml(yaml)
18557            .unwrap()
18558            .llm_alias("default", default)
18559            .llm_alias("router", router)
18560            .build()
18561            .unwrap();
18562
18563        let content = tokio::time::timeout(std::time::Duration::from_secs(2), async {
18564            let mut stream = agent
18565                .chat_stream("AUTH-17 needs technical help.")
18566                .await
18567                .unwrap();
18568            let mut content = String::new();
18569            while let Some(chunk) = stream.next().await {
18570                match chunk {
18571                    StreamChunk::Content { text } => content.push_str(&text),
18572                    StreamChunk::Done {} => break,
18573                    StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
18574                    _ => {}
18575                }
18576            }
18577            content
18578        })
18579        .await
18580        .expect("redispatch must wait for the established stale stream to be dropped");
18581
18582        assert_eq!(content, "Committed technical response.");
18583        assert_eq!(agent.current_state().as_deref(), Some("technical"));
18584        assert!(stream_dropped.load(Ordering::SeqCst));
18585        assert!(committed_after_drop.load(Ordering::SeqCst));
18586    }
18587
18588    #[tokio::test]
18589    async fn test_buffered_streaming_transition_reservation_falls_back() {
18590        use futures::StreamExt;
18591
18592        let mock = mock_with_responses(vec![
18593            "Serial streaming response",
18594            "Serial streaming response",
18595        ]);
18596        let router_mock = mock_with_response("1");
18597        let router_counter = router_mock.clone();
18598        let yaml = r#"
18599name: BufferedReservationFallbackAgent
18600system_prompt: "Stream normally if speculative routing cannot be evaluated."
18601llm:
18602  default: default
18603  router: router
18604observability:
18605  enabled: true
18606  export:
18607    write_raw_events: true
18608streaming:
18609  enabled: true
18610  buffer_size: 8
18611runtime:
18612  optimization:
18613    enabled: true
18614    max_speculative_llm_calls_per_turn: 1
18615    speculative_state_transitions: true
18616    streaming_policy: buffer_until_routing_done
18617    max_parallel_runtime_tasks: 2
18618states:
18619  initial: triage
18620  states:
18621    triage:
18622      prompt: "Triage state."
18623      transitions:
18624        - to: billing
18625          guard:
18626            context:
18627              route:
18628                eq: billing
18629          when: "User asks about billing"
18630          timing: parallel
18631    billing:
18632      prompt: "Billing state."
18633"#;
18634        let agent = AgentBuilder::from_yaml(yaml)
18635            .unwrap()
18636            .llm_alias("default", Arc::new(mock))
18637            .llm_alias("router", Arc::new(router_mock))
18638            .build()
18639            .unwrap();
18640
18641        let mut stream = agent.chat_stream("hello").await.unwrap();
18642        let mut content = String::new();
18643        let mut error = None;
18644        while let Some(chunk) = stream.next().await {
18645            match chunk {
18646                StreamChunk::Content { text } => content.push_str(&text),
18647                StreamChunk::Error { message } => error = Some(message),
18648                StreamChunk::Done {} => break,
18649                _ => {}
18650            }
18651        }
18652
18653        assert_eq!(error, None);
18654        assert_eq!(content, "Serial streaming response");
18655        assert_eq!(router_counter.call_count(), 0);
18656        let events = agent.observability().unwrap().raw_events();
18657        assert!(events.iter().any(|event| {
18658            event.dimensions.get("branch_status") == Some(&"cancelled".to_string())
18659                && event.dimensions.get("commit_behavior")
18660                    == Some(&"transition_decision".to_string())
18661        }));
18662    }
18663
18664    #[tokio::test]
18665    async fn test_blocking_error_cleanup_resets_root_turn_for_next_chat() {
18666        let mut mock = mock_with_response("Recovered response");
18667        mock.set_error("boom");
18668        let mut handle = mock.clone();
18669        let agent = AgentBuilder::new()
18670            .system_prompt("You are helpful.")
18671            .llm(Arc::new(mock))
18672            .build()
18673            .unwrap();
18674
18675        assert!(agent.chat("first").await.is_err());
18676        handle.clear_error();
18677        let response = agent.chat("second").await.unwrap();
18678
18679        assert_eq!(response.content, "Recovered response");
18680        let messages = agent.memory.get_messages(None).await.unwrap();
18681        let user_count = messages
18682            .iter()
18683            .filter(|message| message.role == ai_agents_core::Role::User)
18684            .count();
18685        assert_eq!(user_count, 2);
18686    }
18687
18688    #[tokio::test]
18689    async fn test_streaming_error_cleanup_resets_root_turn_for_next_chat() {
18690        use futures::StreamExt;
18691
18692        let mut mock = mock_with_response("Recovered response");
18693        mock.set_error("stream boom");
18694        let mut handle = mock.clone();
18695        let agent = AgentBuilder::new()
18696            .system_prompt("You are helpful.")
18697            .llm(Arc::new(mock))
18698            .build()
18699            .unwrap();
18700
18701        let mut stream = agent.chat_stream("first").await.unwrap();
18702        let mut saw_error = false;
18703        while let Some(chunk) = stream.next().await {
18704            if matches!(chunk, StreamChunk::Error { .. }) {
18705                saw_error = true;
18706            }
18707        }
18708        assert!(saw_error);
18709
18710        handle.clear_error();
18711        let response = agent.chat("second").await.unwrap();
18712
18713        assert_eq!(response.content, "Recovered response");
18714        let messages = agent.memory.get_messages(None).await.unwrap();
18715        let user_count = messages
18716            .iter()
18717            .filter(|message| message.role == ai_agents_core::Role::User)
18718            .count();
18719        assert_eq!(user_count, 2);
18720    }
18721
18722    #[tokio::test]
18723    async fn test_buffered_streaming_route_miss_releases_buffer_limit() {
18724        use futures::StreamExt;
18725
18726        let mut mock = mock_with_response("one two three");
18727        mock.set_latency(10);
18728        let yaml = r#"
18729name: BufferedMissAgent
18730system_prompt: "You stream safely."
18731llm:
18732  default: default
18733streaming:
18734  enabled: true
18735  buffer_size: 1
18736runtime:
18737  optimization:
18738    enabled: true
18739    max_speculative_llm_calls_per_turn: 2
18740    speculative_state_transitions: true
18741    streaming_policy: buffer_until_routing_done
18742    max_parallel_runtime_tasks: 2
18743states:
18744  initial: triage
18745  states:
18746    triage:
18747      prompt: "Answer from triage."
18748      transitions:
18749        - to: billing
18750          guard:
18751            context:
18752              route:
18753                eq: billing
18754          timing: parallel
18755    billing:
18756      prompt: "Billing state."
18757"#;
18758        let agent = AgentBuilder::from_yaml(yaml)
18759            .unwrap()
18760            .llm_alias("default", Arc::new(mock))
18761            .build()
18762            .unwrap();
18763
18764        let mut stream = agent.chat_stream("hello").await.unwrap();
18765        let mut content = String::new();
18766        let mut error = None;
18767        while let Some(chunk) = stream.next().await {
18768            match chunk {
18769                StreamChunk::Content { text } => content.push_str(&text),
18770                StreamChunk::Error { message } => error = Some(message),
18771                StreamChunk::Done {} => break,
18772                _ => {}
18773            }
18774        }
18775
18776        assert_eq!(error, None);
18777        assert_eq!(content, "one two three");
18778    }
18779
18780    #[tokio::test]
18781    async fn test_buffered_streaming_main_failure_finalizes_branch() {
18782        use futures::StreamExt;
18783
18784        let mock = mock_with_response("one two");
18785        let mut router_mock = mock_with_response("0");
18786        router_mock.set_latency(50);
18787        let yaml = r#"
18788name: BufferedFailureAgent
18789system_prompt: "You stream safely."
18790llm:
18791  default: default
18792  router: router
18793observability:
18794  enabled: true
18795  export:
18796    write_raw_events: true
18797streaming:
18798  enabled: true
18799  buffer_size: 1
18800runtime:
18801  optimization:
18802    enabled: true
18803    max_speculative_llm_calls_per_turn: 2
18804    speculative_state_transitions: true
18805    streaming_policy: buffer_until_routing_done
18806    max_parallel_runtime_tasks: 2
18807states:
18808  initial: triage
18809  states:
18810    triage:
18811      prompt: "Ask for the category."
18812      transitions:
18813        - to: billing
18814          when: "User asks about billing"
18815          timing: parallel
18816    billing:
18817      prompt: "Billing state."
18818"#;
18819        let agent = AgentBuilder::from_yaml(yaml)
18820            .unwrap()
18821            .llm_alias("default", Arc::new(mock))
18822            .llm_alias("router", Arc::new(router_mock))
18823            .build()
18824            .unwrap();
18825
18826        let mut stream = agent.chat_stream("hello").await.unwrap();
18827        let mut error = String::new();
18828        while let Some(chunk) = stream.next().await {
18829            if let StreamChunk::Error { message } = chunk {
18830                error = message;
18831            }
18832        }
18833
18834        assert!(
18835            error.contains("stream buffer filled"),
18836            "unexpected stream error: {}",
18837            error
18838        );
18839        let events = agent.observability().unwrap().raw_events();
18840        assert!(events.iter().any(|event| {
18841            event.dimensions.get("branch_status") == Some(&"failed".to_string())
18842                && event.dimensions.get("commit_behavior") == Some(&"final_response".to_string())
18843                && event.dimensions.get("optimization")
18844                    == Some(&"buffered_streaming_routing".to_string())
18845        }));
18846    }
18847
18848    #[tokio::test]
18849    async fn test_streaming_preflight_does_not_emit_old_state_content() {
18850        use futures::StreamExt;
18851
18852        let mock = mock_with_response("Billing streamed response");
18853        let yaml = r#"
18854name: StreamingOptimizedAgent
18855system_prompt: "You route before streaming."
18856runtime:
18857  optimization:
18858    enabled: true
18859    pre_response_deterministic_transitions: true
18860streaming:
18861  enabled: true
18862states:
18863  initial: greeting
18864  states:
18865    greeting:
18866      prompt: "OLD_STATE_SENTINEL"
18867      transitions:
18868        - to: billing
18869          guard:
18870            context:
18871              topic:
18872                eq: billing
18873          timing: pre_response
18874    billing:
18875      prompt: "Billing state."
18876"#;
18877        let agent = AgentBuilder::from_yaml(yaml)
18878            .unwrap()
18879            .llm(Arc::new(mock))
18880            .build()
18881            .unwrap();
18882        agent
18883            .set_context("topic", serde_json::json!("billing"))
18884            .unwrap();
18885
18886        let mut stream = agent.chat_stream("billing please").await.unwrap();
18887        let mut content = String::new();
18888        while let Some(chunk) = stream.next().await {
18889            match chunk {
18890                StreamChunk::Content { text } => content.push_str(&text),
18891                StreamChunk::Error { message } => panic!("stream error: {}", message),
18892                StreamChunk::Done {} => break,
18893                _ => {}
18894            }
18895        }
18896
18897        assert_eq!(agent.current_state().as_deref(), Some("billing"));
18898        assert!(content.contains("Billing streamed response"));
18899        assert!(!content.contains("OLD_STATE_SENTINEL"));
18900    }
18901
18902    // State machine transitions
18903    #[tokio::test]
18904    async fn test_integration_state_machine_basic() {
18905        let yaml = r#"
18906name: StateAgent
18907system_prompt: "You are a support agent."
18908states:
18909  initial: greeting
18910  states:
18911    greeting:
18912      prompt: "Welcome the user warmly."
18913      transitions:
18914        - to: support
18915          when: "User needs help"
18916          auto: true
18917    support:
18918      prompt: "Help solve the user's problem."
18919"#;
18920        let mock = mock_with_responses(vec![
18921            "Welcome! How can I help?", // greeting response
18922            "1",                        // transition evaluator picks first (index 0)
18923            "I'll help you with that.", // support response
18924        ]);
18925        let builder = AgentBuilder::from_yaml(yaml).unwrap();
18926        let agent = builder.llm(Arc::new(mock)).build().unwrap();
18927
18928        assert_eq!(agent.current_state(), Some("greeting".to_string()));
18929        let _ = agent.chat("I need help").await.unwrap();
18930        // After transition evaluation, state may or may not have changed
18931        // depending on mock evaluator response - the key is that it doesn't crash
18932    }
18933
18934    // State on_enter/on_exit actions
18935    #[tokio::test]
18936    async fn test_integration_state_on_enter_set_context() {
18937        let yaml = r#"
18938name: ActionAgent
18939system_prompt: "You are helpful."
18940states:
18941  initial: step1
18942  states:
18943    step1:
18944      prompt: "Step 1"
18945      on_exit:
18946        - set_context:
18947            step1_exited: true
18948      transitions:
18949        - to: step2
18950          when: "always"
18951          auto: true
18952    step2:
18953      prompt: "Step 2"
18954      on_enter:
18955        - set_context:
18956            step2_entered: true
18957"#;
18958        // The transition evaluator will pick the first transition (index 0)
18959        let mock = mock_with_responses(vec![
18960            "Processing step 1.",
18961            "0", // transition evaluator response: select first transition
18962        ]);
18963        let builder = AgentBuilder::from_yaml(yaml).unwrap();
18964        let agent = builder.llm(Arc::new(mock)).build().unwrap();
18965
18966        assert_eq!(agent.current_state(), Some("step1".to_string()));
18967
18968        // Manually transition to test on_enter/on_exit
18969        agent.transition_to("step2").await.unwrap();
18970
18971        assert_eq!(agent.current_state(), Some("step2".to_string()));
18972
18973        // Verify context was set by on_exit and on_enter actions
18974        let ctx = agent.get_context();
18975        assert_eq!(ctx.get("step1_exited"), Some(&serde_json::json!(true)));
18976        assert_eq!(ctx.get("step2_entered"), Some(&serde_json::json!(true)));
18977    }
18978
18979    #[tokio::test]
18980    async fn state_action_tool_preserves_source_in_stored_record() {
18981        let yaml = r#"
18982name: StateActionToolAgent
18983system_prompt: "You are helpful."
18984tools:
18985  - context_echo
18986states:
18987  initial: idle
18988  states:
18989    idle:
18990      prompt: "Idle"
18991    active:
18992      prompt: "Active"
18993      on_enter:
18994        - set_context:
18995            action_started: true
18996        - tool: context_echo
18997          args: {}
18998"#;
18999        let agent = AgentBuilder::from_yaml(yaml)
19000            .unwrap()
19001            .llm(Arc::new(mock_with_response("unused")))
19002            .tool(Arc::new(ContextEchoTool))
19003            .build()
19004            .unwrap();
19005
19006        agent.transition_to("active").await.unwrap();
19007
19008        let record: ToolExecutionRecord = serde_json::from_value(
19009            agent
19010                .get_context()
19011                .get("last_tool_record")
19012                .cloned()
19013                .expect("successful state action must store its execution record"),
19014        )
19015        .unwrap();
19016        assert!(record.executed);
19017        assert!(record.success);
19018        assert_eq!(record.canonical_id, "context_echo");
19019        assert!(matches!(
19020            &record.source,
19021            ToolCallSource::StateAction {
19022                state: Some(state),
19023                action_index: 1,
19024            } if state == "active"
19025        ));
19026    }
19027
19028    #[tokio::test]
19029    async fn test_ordinary_transition_uses_on_enter_then_on_reenter() {
19030        let yaml = r#"
19031name: OrdinaryLifecycleAgent
19032system_prompt: "You are helpful."
19033states:
19034  initial: intake
19035  regenerate_on_transition: false
19036  states:
19037    intake:
19038      prompt: "Intake"
19039      transitions:
19040        - to: drafting
19041          guard:
19042            context:
19043              route:
19044                eq: drafting
19045    drafting:
19046      prompt: "Drafting"
19047      on_enter:
19048        - set_context:
19049            draft_version: 1
19050      on_reenter:
19051        - set_context:
19052            draft_version: 2
19053      transitions:
19054        - to: review
19055          guard:
19056            context:
19057              route:
19058                eq: review
19059    review:
19060      prompt: "Review"
19061      on_enter:
19062        - set_context:
19063            review_entry: first
19064      transitions:
19065        - to: drafting
19066          guard:
19067            context:
19068              route:
19069                eq: drafting
19070"#;
19071        let agent = AgentBuilder::from_yaml(yaml)
19072            .unwrap()
19073            .llm(Arc::new(mock_with_responses(vec![
19074                "Intake response",
19075                "Draft response",
19076                "Review response",
19077            ])))
19078            .build()
19079            .unwrap();
19080
19081        agent
19082            .set_context("route", serde_json::json!("drafting"))
19083            .unwrap();
19084        agent.chat("Start a draft").await.unwrap();
19085        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19086        assert_eq!(
19087            agent.get_context().get("draft_version"),
19088            Some(&serde_json::json!(1))
19089        );
19090
19091        agent
19092            .set_context("route", serde_json::json!("review"))
19093            .unwrap();
19094        agent.chat("Review this").await.unwrap();
19095        assert_eq!(agent.current_state().as_deref(), Some("review"));
19096        assert_eq!(
19097            agent.get_context().get("review_entry"),
19098            Some(&serde_json::json!("first"))
19099        );
19100
19101        agent
19102            .set_context("route", serde_json::json!("drafting"))
19103            .unwrap();
19104        agent.chat("Revise this").await.unwrap();
19105        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19106        assert_eq!(
19107            agent.get_context().get("draft_version"),
19108            Some(&serde_json::json!(2))
19109        );
19110    }
19111
19112    #[tokio::test]
19113    async fn test_manual_transition_uses_on_enter_then_on_reenter() {
19114        let yaml = r#"
19115name: ManualLifecycleAgent
19116system_prompt: "You are helpful."
19117states:
19118  initial: intake
19119  states:
19120    intake:
19121      prompt: "Intake"
19122    drafting:
19123      prompt: "Drafting"
19124      on_enter:
19125        - set_context:
19126            draft_version: 1
19127      on_reenter:
19128        - set_context:
19129            draft_version: 2
19130    review:
19131      prompt: "Review"
19132"#;
19133        let agent = AgentBuilder::from_yaml(yaml)
19134            .unwrap()
19135            .llm(Arc::new(mock_with_response("unused")))
19136            .build()
19137            .unwrap();
19138
19139        assert!(!agent.get_context().contains_key("draft_version"));
19140        agent.transition_to("drafting").await.unwrap();
19141        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19142        assert_eq!(
19143            agent.get_context().get("draft_version"),
19144            Some(&serde_json::json!(1))
19145        );
19146
19147        agent.transition_to("review").await.unwrap();
19148        agent.transition_to("drafting").await.unwrap();
19149        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19150        assert_eq!(
19151            agent.get_context().get("draft_version"),
19152            Some(&serde_json::json!(2))
19153        );
19154    }
19155
19156    #[tokio::test]
19157    async fn test_timeout_transition_uses_on_enter_then_on_reenter() {
19158        let yaml = r#"
19159name: TimeoutLifecycleAgent
19160system_prompt: "You are helpful."
19161states:
19162  initial: intake
19163  regenerate_on_transition: false
19164  states:
19165    intake:
19166      prompt: "Intake"
19167      max_turns: 1
19168      timeout_to: drafting
19169    drafting:
19170      prompt: "Drafting"
19171      max_turns: 1
19172      timeout_to: review
19173      on_enter:
19174        - set_context:
19175            draft_version: 1
19176      on_reenter:
19177        - set_context:
19178            draft_version: 2
19179    review:
19180      prompt: "Review"
19181      max_turns: 1
19182      timeout_to: drafting
19183      on_enter:
19184        - set_context:
19185            review_entry: first
19186"#;
19187        let agent = AgentBuilder::from_yaml(yaml)
19188            .unwrap()
19189            .llm(Arc::new(mock_with_responses(vec![
19190                "Intake",
19191                "First draft",
19192                "Review",
19193                "Revised draft",
19194            ])))
19195            .build()
19196            .unwrap();
19197
19198        agent.chat("First turn").await.unwrap();
19199        assert_eq!(agent.current_state().as_deref(), Some("intake"));
19200        assert!(!agent.get_context().contains_key("draft_version"));
19201
19202        agent.chat("Second turn").await.unwrap();
19203        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19204        assert_eq!(
19205            agent.get_context().get("draft_version"),
19206            Some(&serde_json::json!(1))
19207        );
19208
19209        agent.chat("Third turn").await.unwrap();
19210        assert_eq!(agent.current_state().as_deref(), Some("review"));
19211        assert_eq!(
19212            agent.get_context().get("review_entry"),
19213            Some(&serde_json::json!("first"))
19214        );
19215
19216        agent.chat("Fourth turn").await.unwrap();
19217        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19218        assert_eq!(
19219            agent.get_context().get("draft_version"),
19220            Some(&serde_json::json!(2))
19221        );
19222    }
19223
19224    // Process pipeline transforms input
19225    #[tokio::test]
19226    async fn test_integration_process_normalize() {
19227        let yaml = r#"
19228name: ProcessAgent
19229system_prompt: "You are helpful."
19230process:
19231  input:
19232    - type: normalize
19233      config:
19234        trim: true
19235        collapse_whitespace: true
19236"#;
19237        let mock = mock_with_response("Got your message.");
19238        let builder = AgentBuilder::from_yaml(yaml).unwrap();
19239        let agent = builder.llm(Arc::new(mock.clone())).build().unwrap();
19240
19241        let _ = agent.chat("  hello   world  ").await.unwrap();
19242
19243        // Verify the LLM received the normalized input (trimmed + collapsed whitespace)
19244        let history = mock.call_history();
19245        assert!(!history.is_empty());
19246        // The user message in LLM call should be normalized
19247        let last_call = history.last().unwrap();
19248        let user_msg = last_call
19249            .messages
19250            .iter()
19251            .find(|m| m.role == ai_agents_core::Role::User)
19252            .unwrap();
19253        assert_eq!(user_msg.content, "hello world");
19254    }
19255
19256    // ═══════════════════════════════════════════════════════════
19257    // Integration Test 2.1.7: Memory compression triggers
19258    // ═══════════════════════════════════════════════════════════
19259    #[tokio::test]
19260    async fn test_integration_memory_compression() {
19261        let yaml = r#"
19262name: MemoryAgent
19263system_prompt: "You are helpful."
19264memory:
19265  type: compacting
19266  max_messages: 100
19267  compress_threshold: 5
19268  max_recent_messages: 3
19269  summarize_batch_size: 2
19270"#;
19271        // Provide enough responses for compression to trigger
19272        let responses: Vec<&str> = (0..8).map(|_| "Response from assistant.").collect();
19273        let mock = mock_with_responses(responses);
19274        let builder = AgentBuilder::from_yaml(yaml).unwrap();
19275        let agent = builder.llm(Arc::new(mock)).build().unwrap();
19276
19277        // Send enough messages to trigger compression
19278        for i in 0..6 {
19279            let _ = agent.chat(&format!("Message {}", i)).await.unwrap();
19280        }
19281
19282        // Memory should have compressed - verify it didn't crash
19283        // and that messages are bounded
19284        let messages = agent.memory.get_messages(None).await.unwrap();
19285        // With compress_threshold=5 and max_recent_messages=3,
19286        // after 6 turns (12 messages), compression should have run
19287        assert!(messages.len() <= 12); // At most all messages if no compression, fewer if compressed
19288    }
19289
19290    // YAML with multiple LLMs
19291    #[tokio::test]
19292    async fn test_integration_multi_llm_registry() {
19293        let mut mock_default = MockLLMProvider::new("default");
19294        mock_default.set_response("Default LLM response.");
19295        let mut mock_router = MockLLMProvider::new("router");
19296        mock_router.set_response("Router response.");
19297
19298        let agent = AgentBuilder::new()
19299            .system_prompt("You are helpful.")
19300            .llm_alias("default", Arc::new(mock_default))
19301            .llm_alias("router", Arc::new(mock_router))
19302            .build()
19303            .unwrap();
19304
19305        let response = agent.chat("Hello").await.unwrap();
19306        assert_eq!(response.content, "Default LLM response.");
19307    }
19308
19309    // Agent reset clears state
19310    #[tokio::test]
19311    async fn test_integration_agent_reset() {
19312        let mock = mock_with_responses(vec!["Hello!", "Hello again!"]);
19313        let agent = AgentBuilder::new()
19314            .system_prompt("You are helpful.")
19315            .llm(Arc::new(mock))
19316            .build()
19317            .unwrap();
19318
19319        let _ = agent.chat("Hi").await.unwrap();
19320        let messages = agent.memory.get_messages(None).await.unwrap();
19321        assert_eq!(messages.len(), 2); // user + assistant
19322
19323        agent.reset().await.unwrap();
19324        let messages = agent.memory.get_messages(None).await.unwrap();
19325        assert_eq!(messages.len(), 0);
19326    }
19327
19328    // Process pipeline rejects input
19329    #[tokio::test]
19330    async fn test_integration_process_validate_reject() {
19331        use ai_agents_process::{ProcessConfig, ProcessProcessor};
19332
19333        let validate_config = ai_agents_process::ValidateStage {
19334            id: Some("length_check".to_string()),
19335            condition: None,
19336            config: ai_agents_process::ValidateConfig {
19337                rules: vec![ai_agents_process::ValidationRule::MinLength {
19338                    min_length: 10,
19339                    on_fail: ai_agents_process::ValidationAction {
19340                        action: ai_agents_process::ValidationActionType::Reject,
19341                        message: None,
19342                    },
19343                }],
19344                ..Default::default()
19345            },
19346        };
19347        let process_config = ProcessConfig {
19348            input: vec![ai_agents_process::ProcessStage::Validate(validate_config)],
19349            ..Default::default()
19350        };
19351        let processor = ProcessProcessor::new(process_config);
19352
19353        let mock = mock_with_response("Should not reach here.");
19354        let agent = AgentBuilder::new()
19355            .system_prompt("You are helpful.")
19356            .llm(Arc::new(mock))
19357            .process_processor(processor)
19358            .build()
19359            .unwrap();
19360
19361        let response = agent.chat("Hi").await.unwrap();
19362        // Rejected input should produce a rejection response, not call LLM
19363        assert!(
19364            response.content.contains("rejected")
19365                || response.content.contains("Input rejected")
19366                || response.content.contains("too short")
19367                || response.content.contains("Too short")
19368                || response.content.len() < 50, // rejection message is typically short
19369            "Expected rejection response, got: {}",
19370            response.content
19371        );
19372    }
19373
19374    // LLM fallback: primary fails, fallback LLM responds
19375    #[tokio::test]
19376    async fn test_llm_fallback_on_failure() {
19377        use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
19378
19379        let mut primary = MockLLMProvider::new("primary");
19380        primary.set_error("Primary LLM is unavailable");
19381
19382        let mut fallback = MockLLMProvider::new("fallback");
19383        fallback.set_response("Fallback response works!");
19384
19385        let agent = AgentBuilder::new()
19386            .system_prompt("You are helpful.")
19387            .llm_alias("default", Arc::new(primary))
19388            .llm_alias("backup", Arc::new(fallback))
19389            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
19390                llm: LLMRecoveryConfig {
19391                    on_failure: LLMFailureAction::FallbackLlm {
19392                        fallback_llm: "backup".to_string(),
19393                    },
19394                    ..Default::default()
19395                },
19396                ..Default::default()
19397            }))
19398            .build()
19399            .unwrap();
19400
19401        let response = agent.chat("Hello").await.unwrap();
19402        assert!(
19403            response.content.contains("Fallback response"),
19404            "Expected fallback response, got: {}",
19405            response.content
19406        );
19407    }
19408
19409    // LLM fallback: primary fails, static message returned
19410    #[tokio::test]
19411    async fn test_llm_fallback_response_static_message() {
19412        use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
19413
19414        let mut primary = MockLLMProvider::new("primary");
19415        primary.set_error("Primary LLM is unavailable");
19416
19417        let agent = AgentBuilder::new()
19418            .system_prompt("You are helpful.")
19419            .llm(Arc::new(primary))
19420            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
19421                llm: LLMRecoveryConfig {
19422                    on_failure: LLMFailureAction::FallbackResponse {
19423                        message: "I am temporarily unavailable. Please try again later."
19424                            .to_string(),
19425                    },
19426                    ..Default::default()
19427                },
19428                ..Default::default()
19429            }))
19430            .build()
19431            .unwrap();
19432
19433        let response = agent.chat("Hello").await.unwrap();
19434        assert!(
19435            response.content.contains("temporarily unavailable"),
19436            "Expected static fallback message, got: {}",
19437            response.content
19438        );
19439    }
19440
19441    // Tool skip: tool fails, on_failure: skip absorbs the error
19442    #[tokio::test]
19443    async fn test_tool_failure_skip() {
19444        use ai_agents_recovery::{
19445            ErrorRecoveryConfig, ToolFailureAction, ToolRecoveryConfig, ToolRetryConfig,
19446        };
19447
19448        // LLM requests a nonexistent tool, then responds after seeing the skip result
19449        let mock = mock_with_responses(vec![
19450            r#"I'll use the nonexistent tool.
19451[TOOL_CALL: {"name": "nonexistent_tool", "arguments": {}}]"#,
19452            "The tool was unavailable, but I can still help you.",
19453        ]);
19454
19455        let agent = AgentBuilder::new()
19456            .system_prompt("You are helpful.")
19457            .llm(Arc::new(mock))
19458            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
19459                tools: ToolRecoveryConfig {
19460                    default: ToolRetryConfig {
19461                        max_retries: 0,
19462                        timeout_ms: None,
19463                        on_failure: ToolFailureAction::Skip,
19464                    },
19465                    ..Default::default()
19466                },
19467                ..Default::default()
19468            }))
19469            .build()
19470            .unwrap();
19471
19472        // The tool will fail (not found), but on_failure: skip absorbs the error
19473        let response = agent.chat("Use the nonexistent tool").await;
19474        assert!(
19475            response.is_ok(),
19476            "Expected Ok with skip policy, got: {:?}",
19477            response
19478        );
19479    }
19480}