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, ParallelToolsConfig, StreamChunk, StreamingConfig, ToolCall,
185};
186use crate::optimization::{
187    AwaitBeforeNextTurn, BackgroundMaintenanceQueue, BackgroundOverflowPolicy, MainResponseDraft,
188    MaintenanceMode, MaintenanceSequenceKey, RuntimeBranch, RuntimeBranchResult,
189    RuntimeBranchStatus, RuntimeCommitBehavior, RuntimeConfig, RuntimeOptimizationKind,
190    RuntimeTaskPriority, RuntimeTaskPurpose, ScheduledBranchSet, SkillCandidate,
191    StreamingDraftResult, TransitionCandidate, TurnBranchScheduler, TurnOptimizationContext,
192};
193use crate::spec::StorageConfig;
194
195/// Outcome of processing tool calls within the agent loop.
196enum ToolCallOutcome {
197    /// Tools executed successfully, continue the LLM loop for the next iteration.
198    Continue,
199    /// A state transition fired during tool call handling, continue the loop.
200    TransitionFired,
201    /// HITL rejected a tool call, return this response immediately.
202    Rejected(AgentResponse),
203}
204
205#[derive(Clone)]
206struct MainToolProtocol {
207    choice: Option<ToolChoice>,
208    tool_ids: Vec<String>,
209    definitions: Vec<LLMToolDefinition>,
210}
211
212struct MainProviderResponse {
213    response: LLMResponse,
214    used_native_tools: bool,
215}
216
217//
218// Carries one committed model answer through shared output processing and root-turn finalization.
219//
220struct CommittedTextResponse<'a> {
221    processed_input: &'a str,
222    input_context: &'a HashMap<String, Value>,
223    answer: String,
224    reasoning_mode: ReasoningMode,
225    auto_detected: bool,
226    iterations: u32,
227    thinking_content: Option<String>,
228    all_tool_calls: Vec<ToolCall>,
229}
230
231//
232// Carries the finalized response fields into metadata assembly without changing their ownership.
233//
234struct AgentResponseParts {
235    content: String,
236    all_tool_calls: Vec<ToolCall>,
237    reasoning_mode: ReasoningMode,
238    auto_detected: bool,
239    iterations: u32,
240    thinking: Option<String>,
241    reflection_metadata: Option<ReflectionMetadata>,
242}
243
244#[derive(Clone, Copy)]
245struct DisambiguationOwnership {
246    epoch: u64,
247    state_generation: Option<u64>,
248}
249
250/// Outcome of skill routing — used by `try_skill_route`.
251enum SkillRouteResult {
252    /// No skill matched, continue to normal LLM chat.
253    NoMatch,
254    /// Skill executed successfully.
255    Response { skill_id: String, content: String },
256    /// Skill matched but needs disambiguation first or returned a terminal disambiguation response.
257    NeedsClarification {
258        response: AgentResponse,
259        ownership: Option<DisambiguationOwnership>,
260    },
261}
262
263/// Result of response-independent parallel transition selection.
264enum ParallelTransitionSelection {
265    /// A transition matched and can be committed before the old-state response.
266    Candidate(TransitionCandidate),
267    /// All eligible transition checks ran and none matched.
268    NoMatch,
269    /// LLM-based route evaluation needed speculative capacity that was unavailable.
270    ReservationExhausted,
271}
272
273/// Outcome of post_loop_processing - drives the caller's next step.
274enum PostLoopResult {
275    /// No transition fired. Content is the LLM response for this turn.
276    NoTransition(String),
277    /// Transition fired. Content is from plain post-transition re-generation.
278    Transitioned(String),
279    /// Transition fired into a state that requires full dispatch.
280    /// Caller re-enters run_loop_internal to apply the correct handler.
281    NeedsRedispatch,
282}
283
284struct StateTransitionReservation<'a> {
285    reserved: &'a AtomicBool,
286}
287
288impl Drop for StateTransitionReservation<'_> {
289    fn drop(&mut self) {
290        self.reserved.store(false, Ordering::SeqCst);
291    }
292}
293
294struct RootTurnCleanup<'a> {
295    agent: &'a RuntimeAgent,
296}
297
298impl<'a> RootTurnCleanup<'a> {
299    fn new(agent: &'a RuntimeAgent) -> Self {
300        Self { agent }
301    }
302}
303
304impl Drop for RootTurnCleanup<'_> {
305    fn drop(&mut self) {
306        self.agent.end_root_turn();
307    }
308}
309
310/// Host-owned runtime control state shared with active agents.
311#[derive(Debug)]
312struct RuntimeControlState {
313    /// Serializes control mutations with exact runtime safety snapshots.
314    snapshot_guard: RwLock<()>,
315    /// Monotonic version for runtime-control snapshots.
316    version: AtomicU64,
317    /// Emergency switch that denies future calls and is shared with active tool contexts.
318    emergency_deny: Arc<AtomicBool>,
319    /// Optional live replacement for tool security policy.
320    tool_security_override: RwLock<Option<ToolSecurityEngine>>,
321    /// Optional live narrowing scope applied after the runtime's declared grant.
322    tool_scope_override: RwLock<Option<Vec<String>>>,
323}
324
325impl Default for RuntimeControlState {
326    fn default() -> Self {
327        Self {
328            snapshot_guard: RwLock::new(()),
329            version: AtomicU64::new(1),
330            emergency_deny: Arc::new(AtomicBool::new(false)),
331            tool_security_override: RwLock::new(None),
332            tool_scope_override: RwLock::new(None),
333        }
334    }
335}
336
337/// Host-only handle for live runtime safety controls.
338#[derive(Clone)]
339pub struct RuntimeControlHandle {
340    state: Arc<RuntimeControlState>,
341}
342
343impl RuntimeControlHandle {
344    /// Returns the current runtime-control version.
345    pub fn version(&self) -> u64 {
346        self.state.version.load(Ordering::SeqCst)
347    }
348
349    fn bump(&self) -> u64 {
350        self.state.version.fetch_add(1, Ordering::SeqCst) + 1
351    }
352
353    /// Overrides tool security for later tool calls and panics if host policy is invalid.
354    pub fn set_tool_security(&self, config: ToolSecurityConfig) -> u64 {
355        self.try_set_tool_security(config)
356            .expect("invalid tool security configuration")
357    }
358
359    /// Validates replacement policy before changing the runtime-control generation or active snapshot.
360    pub fn try_set_tool_security(&self, config: ToolSecurityConfig) -> Result<u64> {
361        config.validate()?;
362        let _guard = self.state.snapshot_guard.write();
363        let generation = self.bump();
364        *self.state.tool_security_override.write() = Some(
365            ToolSecurityEngine::new_with_policy_version(config, generation),
366        );
367        Ok(generation)
368    }
369
370    /// Clears the live tool security override.
371    pub fn clear_tool_security_override(&self) -> u64 {
372        let _guard = self.state.snapshot_guard.write();
373        *self.state.tool_security_override.write() = None;
374        self.bump()
375    }
376
377    /// Narrows the runtime's declared tool grant for later calls without adding authority.
378    pub fn set_tool_scope(&self, tool_ids: Vec<String>) -> u64 {
379        let _guard = self.state.snapshot_guard.write();
380        *self.state.tool_scope_override.write() = Some(tool_ids);
381        self.bump()
382    }
383
384    /// Clears live narrowing so later calls return to the runtime's declared grant.
385    pub fn clear_tool_scope_override(&self) -> u64 {
386        let _guard = self.state.snapshot_guard.write();
387        *self.state.tool_scope_override.write() = None;
388        self.bump()
389    }
390
391    /// Enables or disables emergency denial for future tool calls.
392    pub fn set_emergency_deny(&self, enabled: bool) -> u64 {
393        let _guard = self.state.snapshot_guard.write();
394        self.state.emergency_deny.store(enabled, Ordering::SeqCst);
395        self.bump()
396    }
397
398    /// Denies future calls and asks active cancellable work to stop.
399    pub fn cancel_all(&self) -> u64 {
400        self.set_emergency_deny(true)
401    }
402}
403
404pub struct RuntimeAgent {
405    info: AgentInfo,
406    llm_registry: Arc<LLMRegistry>,
407    memory: Arc<dyn Memory>,
408    tools: Arc<ToolRegistry>,
409    skills: Vec<SkillDefinition>,
410    skill_router: Option<SkillRouter>,
411    skill_executor: Option<SkillExecutor>,
412    base_system_prompt: String,
413    max_iterations: u32,
414    iteration_count: RwLock<u32>,
415    max_context_tokens: u32,
416    memory_token_budget: Option<MemoryTokenBudget>,
417    recovery_manager: RecoveryManager,
418    tool_security: ToolSecurityEngine,
419    process_processor: Option<ProcessProcessor>,
420    message_filters: RwLock<HashMap<String, Arc<dyn MessageFilter>>>,
421    state_machine: Option<Arc<StateMachine>>,
422    transition_evaluator: Option<Arc<dyn TransitionEvaluator>>,
423    context_manager: Arc<ContextManager>,
424    template_renderer: TemplateRenderer,
425    tool_call_history: RwLock<Vec<ToolCallRecord>>,
426    parallel_tools: ParallelToolsConfig,
427    streaming: StreamingConfig,
428    hooks: Arc<dyn AgentHooks>,
429    hitl_engine: Option<HITLEngine>,
430    approval_handler: Arc<dyn ApprovalHandler>,
431    storage_config: StorageConfig,
432    storage: RwLock<Option<Arc<dyn AgentStorage>>>,
433    storage_init: tokio::sync::Mutex<()>,
434    reasoning_config: ReasoningConfig,
435    reflection_config: ReflectionConfig,
436    disambiguation_manager: Option<DisambiguationManager>,
437    /// Generation invalidating confirmation work across reset and state changes.
438    disambiguation_epoch: AtomicU64,
439    /// Serializes confirmation redispatch admission with reset and state mutation.
440    disambiguation_admission: tokio::sync::RwLock<()>,
441    /// Reserves one runtime-owned transition before its exit actions can produce side effects.
442    state_transition_reserved: AtomicBool,
443    /// Structured persona manager for identity, evolution, and secrets.
444    persona_manager: Option<Arc<ai_agents_persona::PersonaManager>>,
445    /// Skill ID that triggered the current pending disambiguation.
446    /// Set by try_skill_route() when skill-level disambiguation triggers clarification.
447    /// Read by run_loop() when clarification resolves to route directly to the skill.
448    pending_skill_id: RwLock<Option<String>>,
449    current_plan: RwLock<Option<Plan>>,
450    /// Tool IDs declared in the top-level `tools:` spec.
451    declared_tool_ids: Option<Vec<String>>,
452    /// Whether the context manager has been initialized (defaults loaded, env resolved, etc.)
453    context_initialized: AtomicBool,
454    /// Spawner for dynamic agent creation (set when YAML has a spawner: section).
455    spawner: Option<Arc<crate::spawner::AgentSpawner>>,
456    /// Registry tracking spawned agents (set when YAML has a spawner: section).
457    spawner_registry: Option<Arc<crate::spawner::AgentRegistry>>,
458    /// Re-dispatch depth for post-transition full dispatch.
459    /// 0 = not re-dispatching. > 0 = user message already in memory, skip re-adding.
460    redispatch_depth: RwLock<u32>,
461    /// Active optimized turn context used to keep root lifecycle state in one place.
462    active_turn_context: RwLock<Option<TurnOptimizationContext>>,
463    /// Tracks whether the root turn already wrote the processed user message.
464    root_user_message_committed: AtomicBool,
465    /// Current actor ID for cross-session memory.
466    actor_id: RwLock<Option<String>>,
467    /// Fact store for managing per-actor extracted facts.
468    fact_store: RwLock<Option<Arc<ai_agents_facts::FactStore>>>,
469    /// Fact extractor for LLM-based fact extraction.
470    /// None when actor_memory is enabled without facts.enabled.
471    fact_extractor: RwLock<Option<Arc<dyn ai_agents_facts::FactExtractor>>>,
472    /// Cached actor facts keyed by actor ID so concurrent or alternating turns do not overwrite one another.
473    actor_facts_cache: Arc<RwLock<HashMap<String, Vec<ai_agents_core::KeyFact>>>>,
474    /// Number of messages since last fact extraction.
475    messages_since_extraction: Arc<RwLock<usize>>,
476    /// Actor memory configuration.
477    actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
478    /// Facts configuration.
479    facts_config: Option<ai_agents_facts::FactsConfig>,
480    /// Session-scoped metadata (tags, ttl, actor roster).
481    session_metadata: RwLock<ai_agents_core::SessionMetadata>,
482    /// Session id currently bound to this runtime instance.
483    current_session_id: RwLock<Option<String>>,
484    /// Relationship manager for actor-scoped social memory.
485    relationship_manager: Option<Arc<RelationshipManager>>,
486    /// Observability manager for traces, metrics, reports, and exports.
487    observability_manager: Option<Arc<ObservabilityManager>>,
488    /// Runtime optimization and maintenance policy.
489    runtime_config: RuntimeConfig,
490    /// Queue for background maintenance tasks.
491    background_maintenance: Arc<BackgroundMaintenanceQueue>,
492    /// Cross-tool locks for side-effecting calls that target the same resource.
493    resource_locks: ToolResourceLocks,
494    /// Host-only runtime control state.
495    runtime_control: Arc<RuntimeControlState>,
496}
497
498impl std::fmt::Debug for RuntimeAgent {
499    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
500        f.debug_struct("RuntimeAgent")
501            .field("info", &self.info)
502            .field("base_system_prompt", &self.base_system_prompt)
503            .field("max_iterations", &self.max_iterations)
504            .field("skills_count", &self.skills.len())
505            .field("max_context_tokens", &self.max_context_tokens)
506            .field("has_state_machine", &self.state_machine.is_some())
507            .field("parallel_tools", &self.parallel_tools)
508            .field("streaming", &self.streaming)
509            .field("has_hooks", &true)
510            .field("has_hitl", &self.hitl_engine.is_some())
511            .field("storage_type", &self.storage_config.storage_type())
512            .field("reasoning_mode", &self.reasoning_config.mode)
513            .field("reflection_enabled", &self.reflection_config.enabled)
514            .field("declared_tool_ids", &self.declared_tool_ids)
515            .field("has_persona", &self.persona_manager.is_some())
516            .field("has_observability", &self.observability_manager.is_some())
517            .finish_non_exhaustive()
518    }
519}
520
521struct ObservabilityClarificationObserver;
522
523impl ClarificationObserver for ObservabilityClarificationObserver {
524    /// Scopes clarification question generation as disambiguation_clarification.
525    fn observe_question<'a>(
526        &'a self,
527        future: ClarificationQuestionFuture<'a>,
528    ) -> ClarificationQuestionFuture<'a> {
529        Box::pin(async move {
530            with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
531        })
532    }
533
534    /// Scopes clarification response parsing as disambiguation_clarification.
535    fn observe_parse<'a>(
536        &'a self,
537        future: ClarificationParseFuture<'a>,
538    ) -> ClarificationParseFuture<'a> {
539        Box::pin(async move {
540            with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
541        })
542    }
543
544    /// Scopes semantic confirmation parsing as disambiguation_clarification.
545    fn observe_confirmation_parse<'a>(
546        &'a self,
547        future: ConfirmationParseFuture<'a>,
548    ) -> ConfirmationParseFuture<'a> {
549        Box::pin(async move {
550            with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
551        })
552    }
553}
554
555struct ObservabilityProcessStageObserver;
556
557impl ProcessStageObserver for ObservabilityProcessStageObserver {
558    /// Scopes one process stage using the purpose implied by its stage type.
559    fn observe<'a>(
560        &'a self,
561        hint: ProcessPurposeHint,
562        future: ProcessStageFuture<'a>,
563    ) -> ProcessStageFuture<'a> {
564        Box::pin(async move {
565            with_observation_purpose(observation_purpose_for_process(hint), future).await
566        })
567    }
568}
569
570struct RegistryLLMGetter {
571    registry: Arc<LLMRegistry>,
572}
573
574impl LLMGetter for RegistryLLMGetter {
575    fn get_llm(&self, alias: &str) -> Option<Arc<dyn LLMProvider>> {
576        self.registry.get(alias).ok()
577    }
578}
579
580impl RuntimeAgent {
581    #[allow(clippy::too_many_arguments)]
582    pub fn new(
583        info: AgentInfo,
584        llm_registry: Arc<LLMRegistry>,
585        memory: Arc<dyn Memory>,
586        tools: Arc<ToolRegistry>,
587        skills: Vec<SkillDefinition>,
588        system_prompt: String,
589        max_iterations: u32,
590    ) -> Self {
591        let (skill_router, skill_executor) = if !skills.is_empty() {
592            let router_llm = llm_registry.router().ok();
593            let router = router_llm.map(|llm| SkillRouter::new(llm, skills.clone()));
594            let executor = SkillExecutor::new(llm_registry.clone(), tools.clone());
595            (router, Some(executor))
596        } else {
597            (None, None)
598        };
599
600        let context_manager =
601            ContextManager::new(HashMap::new(), info.name.clone(), info.version.clone());
602
603        Self {
604            info,
605            llm_registry,
606            memory,
607            tools,
608            skills,
609            skill_router,
610            skill_executor,
611            base_system_prompt: system_prompt,
612            max_iterations,
613            iteration_count: RwLock::new(0),
614            max_context_tokens: 128000,
615            memory_token_budget: None,
616            recovery_manager: RecoveryManager::default(),
617            tool_security: ToolSecurityEngine::default(),
618            process_processor: None,
619            message_filters: RwLock::new(HashMap::new()),
620            state_machine: None,
621            transition_evaluator: None,
622            context_manager: Arc::new(context_manager),
623            template_renderer: TemplateRenderer::new(),
624            tool_call_history: RwLock::new(Vec::new()),
625            parallel_tools: ParallelToolsConfig::default(),
626            streaming: StreamingConfig::default(),
627            hooks: Arc::new(NoopHooks),
628            hitl_engine: None,
629            approval_handler: Arc::new(RejectAllHandler::new()),
630            storage_config: StorageConfig::default(),
631            storage: RwLock::new(None),
632            storage_init: tokio::sync::Mutex::new(()),
633            reasoning_config: ReasoningConfig::default(),
634            reflection_config: ReflectionConfig::default(),
635            disambiguation_manager: None,
636            disambiguation_epoch: AtomicU64::new(0),
637            disambiguation_admission: tokio::sync::RwLock::new(()),
638            state_transition_reserved: AtomicBool::new(false),
639            persona_manager: None,
640            pending_skill_id: RwLock::new(None),
641            current_plan: RwLock::new(None),
642            declared_tool_ids: None,
643            context_initialized: AtomicBool::new(false),
644            spawner: None,
645            spawner_registry: None,
646            redispatch_depth: RwLock::new(0),
647            active_turn_context: RwLock::new(None),
648            root_user_message_committed: AtomicBool::new(false),
649            actor_id: RwLock::new(None),
650            fact_store: RwLock::new(None),
651            fact_extractor: RwLock::new(None),
652            actor_facts_cache: Arc::new(RwLock::new(HashMap::new())),
653            messages_since_extraction: Arc::new(RwLock::new(0)),
654            actor_memory_config: None,
655            facts_config: None,
656            session_metadata: RwLock::new(ai_agents_core::SessionMetadata::default()),
657            current_session_id: RwLock::new(None),
658            relationship_manager: None,
659            observability_manager: None,
660            runtime_config: RuntimeConfig::default(),
661            background_maintenance: Arc::new(BackgroundMaintenanceQueue::default()),
662            resource_locks: new_tool_resource_locks(),
663            runtime_control: Arc::new(RuntimeControlState::default()),
664        }
665    }
666
667    pub fn with_declared_tool_ids(mut self, ids: Option<Vec<String>>) -> Self {
668        self.declared_tool_ids = ids;
669        self
670    }
671
672    pub fn with_storage_config(mut self, config: StorageConfig) -> Self {
673        self.storage_config = config;
674        self
675    }
676
677    pub fn with_storage(self, storage: Arc<dyn AgentStorage>) -> Self {
678        *self.storage.write() = Some(storage);
679        self
680    }
681
682    pub(crate) fn with_shared_resource_locks(mut self, locks: ToolResourceLocks) -> Self {
683        self.resource_locks = locks;
684        self
685    }
686
687    pub fn with_reasoning(mut self, config: ReasoningConfig) -> Self {
688        self.reasoning_config = config;
689        self
690    }
691
692    pub fn with_reflection(mut self, config: ReflectionConfig) -> Self {
693        self.reflection_config = config;
694        self
695    }
696
697    /// Attach a relationship manager configured by the builder or host application.
698    pub fn with_relationships(mut self, manager: Arc<RelationshipManager>) -> Self {
699        self.relationship_manager = Some(manager);
700        self
701    }
702
703    /// Attach a shared observability manager for traces, metrics, reports, and exports.
704    pub fn with_observability(mut self, manager: Arc<ObservabilityManager>) -> Self {
705        self.observability_manager = Some(manager);
706        self
707    }
708
709    /// Attach runtime optimization policy and resize the background queue.
710    pub fn with_runtime_config(mut self, config: RuntimeConfig) -> Self {
711        let max_tasks = config.optimization.post_turn.max_background_tasks;
712        self.background_maintenance = Arc::new(BackgroundMaintenanceQueue::new(max_tasks));
713        self.runtime_config = config;
714        self
715    }
716
717    /// Returns the runtime optimization policy.
718    pub fn runtime_config(&self) -> &RuntimeConfig {
719        &self.runtime_config
720    }
721
722    /// Wait for all background maintenance tasks to finish.
723    pub async fn flush_background_tasks(&self) -> Result<()> {
724        self.background_maintenance.flush_all().await
725    }
726
727    /// Wait for background maintenance associated with one actor to finish.
728    pub async fn flush_background_tasks_for_actor(&self, actor_id: &str) -> Result<()> {
729        self.background_maintenance.flush_scope(actor_id).await
730    }
731
732    /// Wait for background maintenance associated with one task kind to finish.
733    pub async fn flush_background_tasks_for_purpose(
734        &self,
735        purpose: RuntimeTaskPurpose,
736    ) -> Result<()> {
737        self.background_maintenance.flush_purpose(purpose).await
738    }
739
740    /// Wait for background maintenance associated with one actor and task kind to finish.
741    pub async fn flush_background_tasks_for_actor_purpose(
742        &self,
743        actor_id: &str,
744        purpose: RuntimeTaskPurpose,
745    ) -> Result<()> {
746        self.background_maintenance
747            .flush_scope_purpose(actor_id, purpose)
748            .await
749    }
750
751    /// Flush background maintenance before a host shuts down the runtime.
752    pub async fn shutdown_background_tasks(&self) -> Result<()> {
753        self.flush_background_tasks().await
754    }
755
756    /// Returns the configured observability manager for report and export access.
757    pub fn observability(&self) -> Option<Arc<ObservabilityManager>> {
758        self.observability_manager.clone()
759    }
760
761    /// Exports observability files after a turn when export settings request it.
762    async fn export_observability_if_configured(&self) {
763        let Some(manager) = self.observability_manager.as_ref() else {
764            return;
765        };
766        let export = &manager.config().export;
767        if !export.write_report && !export.write_raw_events {
768            return;
769        }
770        if let Err(error) = manager.export().await {
771            warn!(error = %error, "Observability export failed");
772        }
773    }
774
775    /// Returns the configured relationship manager, if relationship memory is enabled.
776    pub fn relationship_manager(&self) -> Option<Arc<RelationshipManager>> {
777        self.relationship_manager.clone()
778    }
779
780    fn current_turn_actor_context(&self) -> Option<crate::TurnActorContext> {
781        current_turn_actor_context()
782    }
783
784    fn effective_actor_id(&self) -> Option<String> {
785        self.current_turn_actor_context()
786            .and_then(|ctx| ctx.effective_actor_id().map(|id| id.to_string()))
787            .or_else(|| self.actor_id.read().clone())
788    }
789
790    fn effective_origin_actor_id(&self) -> Option<String> {
791        self.current_turn_actor_context()
792            .and_then(|ctx| ctx.origin_actor_id.clone())
793            .or_else(|| self.actor_id.read().clone())
794    }
795
796    fn record_session_actor_if_needed(&self) {
797        if let Some(actor_id) = self.effective_origin_actor_id() {
798            let mut meta = self.session_metadata.write();
799            meta.actor_id = Some(actor_id.clone());
800            if !meta.actors.iter().any(|a| a == &actor_id) {
801                meta.actors.push(actor_id);
802            }
803        }
804    }
805
806    fn outbound_actor_context(&self) -> crate::TurnActorContext {
807        let mut context = self.current_turn_actor_context().unwrap_or_default();
808        if context.origin_actor_id.is_none() {
809            context.origin_actor_id = self.effective_origin_actor_id();
810        }
811        context.sender_agent_id = Some(self.info.id.clone());
812        context
813    }
814
815    /// Returns the current session ID or creates one for unsaved observed turns.
816    fn observation_session_id(&self) -> Option<String> {
817        let mut current = self.current_session_id.write();
818        if current.is_none() {
819            *current = Some(new_observation_session_id());
820        }
821        current.clone()
822    }
823
824    /// Builds the root or child observation context for a chat entry point.
825    fn build_observation_context(&self, actor_id: Option<String>) -> Option<SpanContext> {
826        let manager = self.observability_manager.as_ref()?;
827        let context = self.build_context_with_overlays();
828        let language = resolve_language_from_context(manager.config(), &context);
829        let context = current_observation_context()
830            .map(|parent| parent.child_for_agent(self.info.id.clone()).with_new_turn())
831            .unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
832        Some(
833            context
834                .with_actor(actor_id.or_else(|| self.effective_actor_id()))
835                .with_session(self.observation_session_id())
836                .with_state(self.current_state())
837                .with_language(Some(language)),
838        )
839    }
840
841    /// Refreshes task-local context with current runtime labels and a purpose.
842    fn current_runtime_observation_context(
843        &self,
844        purpose: ObservationPurpose,
845    ) -> Option<SpanContext> {
846        let manager = self.observability_manager.as_ref()?;
847        let context = self.build_context_with_overlays();
848        let language = resolve_language_from_context(manager.config(), &context);
849        let mut observation = current_observation_context()
850            .unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
851        observation.agent_id = self.info.id.clone();
852        observation.actor_id = self.effective_actor_id();
853        observation.session_id = self.observation_session_id();
854        observation.state = self.current_state();
855        observation.language = Some(language);
856        observation.purpose = purpose;
857        Some(observation)
858    }
859
860    /// Runs a future under a purpose while preserving current trace context.
861    async fn observe_purpose<F, T>(&self, purpose: ObservationPurpose, future: F) -> T
862    where
863        F: Future<Output = T>,
864    {
865        if let Some(context) = self.current_runtime_observation_context(purpose) {
866            with_observation_context(context, future).await
867        } else {
868            future.await
869        }
870    }
871
872    /// Runs chat with actor and observation context while avoiding recursive async types.
873    fn chat_with_actor_context_boxed<'a>(
874        &'a self,
875        input: &'a str,
876        actor_context: crate::TurnActorContext,
877    ) -> Pin<Box<dyn Future<Output = Result<AgentResponse>> + Send + 'a>> {
878        Box::pin(async move {
879            let actor_id = actor_context.effective_actor_id().map(str::to_string);
880            let run = async move {
881                scope_actor_context(
882                    actor_context,
883                    Box::pin(async move { self.run_loop(input).await }),
884                )
885                .await
886            };
887            let result = if let Some(context) = self.build_observation_context(actor_id) {
888                with_observation_context(context, run).await
889            } else {
890                run.await
891            };
892            self.export_observability_if_configured().await;
893            result
894        })
895    }
896
897    /// Run one turn with turn-scoped actor context without mutating the runtime's global actor ID.
898    ///
899    /// The supplied context is available to actor-scoped facts, relationship memory, orchestration, and prompt templates only for the lifetime of this call.
900    pub async fn chat_with_actor_context(
901        &self,
902        input: &str,
903        actor_context: crate::TurnActorContext,
904    ) -> Result<AgentResponse> {
905        self.chat_with_actor_context_boxed(input, actor_context)
906            .await
907    }
908
909    /// Convenience wrapper around [`Self::chat_with_actor_context`] for a turn whose original actor is known up front.
910    pub async fn chat_as_actor(&self, actor_id: &str, input: &str) -> Result<AgentResponse> {
911        let actor_context = crate::TurnActorContext::new().with_origin_actor(actor_id);
912        self.chat_with_actor_context(input, actor_context).await
913    }
914
915    /// Ensure the effective actor's relationship is loaded from storage into the relationship manager.
916    pub async fn load_actor_relationship(&self) -> Result<()> {
917        self.maybe_load_actor_relationship().await;
918        Ok(())
919    }
920
921    /// Manually apply a delta to the effective actor's `agent_to_actor` relationship perspective and persist the updated relationship when storage is configured.
922    pub async fn update_relationship_dimension(
923        &self,
924        dimension: &str,
925        delta: f64,
926        reason: Option<&str>,
927    ) -> Result<ai_agents_relationships::DimensionChange> {
928        self.update_relationship_dimension_for_perspective(
929            ai_agents_relationships::RelationshipPerspective::AgentToActor,
930            dimension,
931            delta,
932            reason,
933        )
934        .await
935    }
936
937    /// Manually apply a delta to a specific relationship perspective for the effective actor.
938    ///
939    /// Use this for two-sided configurations when you need to update `agent_to_actor`, `perceived_actor_to_agent`, or `mutual` explicitly from application logic.
940    pub async fn update_relationship_dimension_for_perspective(
941        &self,
942        perspective: ai_agents_relationships::RelationshipPerspective,
943        dimension: &str,
944        delta: f64,
945        reason: Option<&str>,
946    ) -> Result<ai_agents_relationships::DimensionChange> {
947        let manager = self
948            .relationship_manager
949            .as_ref()
950            .ok_or_else(|| AgentError::Config("Relationship memory is not configured".into()))?;
951        let actor_id = self.effective_actor_id().ok_or_else(|| {
952            AgentError::Config("No actor ID set. Use set_actor_id() first".into())
953        })?;
954        let change = manager.update_dimension_for_perspective(
955            &actor_id,
956            perspective,
957            dimension,
958            delta,
959            1.0,
960            reason.unwrap_or("manual relationship update"),
961        )?;
962        self.persist_actor_relationship(&actor_id).await?;
963        info!(
964            actor_id = %actor_id,
965            perspective = %change.perspective,
966            dimension = %change.dimension,
967            delta = change.delta,
968            current = change.current,
969            "relationship updated manually"
970        );
971        self.hooks
972            .on_relationship_change(&actor_id, std::slice::from_ref(&change))
973            .await;
974        Ok(change)
975    }
976
977    pub fn reasoning_config(&self) -> &ReasoningConfig {
978        &self.reasoning_config
979    }
980
981    pub fn reflection_config(&self) -> &ReflectionConfig {
982        &self.reflection_config
983    }
984
985    /// Set only the actor memory and facts configs without creating the store.
986    /// The store and extractor are created lazily in init_storage().
987    pub fn with_facts_config(
988        mut self,
989        actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
990        facts_config: Option<ai_agents_facts::FactsConfig>,
991    ) -> Self {
992        self.actor_memory_config = actor_memory_config;
993        self.facts_config = facts_config;
994        self
995    }
996
997    /// Configure fact store and optional extractor for actor memory.
998    /// Pass `None` for `extractor` to load existing facts without running extraction.
999    pub fn with_facts(
1000        mut self,
1001        store: Arc<ai_agents_facts::FactStore>,
1002        extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>>,
1003        actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
1004        facts_config: Option<ai_agents_facts::FactsConfig>,
1005    ) -> Self {
1006        *self.fact_store.write() = Some(store);
1007        *self.fact_extractor.write() = extractor;
1008        self.actor_memory_config = actor_memory_config;
1009        self.facts_config = facts_config;
1010        self
1011    }
1012
1013    /// Get the fact store for direct fact manipulation.
1014    pub fn fact_store(&self) -> Option<Arc<ai_agents_facts::FactStore>> {
1015        self.fact_store.read().clone()
1016    }
1017
1018    /// Get the current actor ID.
1019    pub fn actor_id(&self) -> Option<String> {
1020        self.actor_id.read().clone()
1021    }
1022
1023    /// Set the current actor ID (player, user, another agent, etc.).
1024    pub fn set_actor_id(&self, actor_id: &str) -> ai_agents_core::Result<()> {
1025        *self.actor_id.write() = Some(actor_id.to_string());
1026        {
1027            let mut meta = self.session_metadata.write();
1028            meta.actor_id = Some(actor_id.to_string());
1029            if !meta.actors.iter().any(|a| a == actor_id) {
1030                meta.actors.push(actor_id.to_string());
1031            }
1032        }
1033        Ok(())
1034    }
1035
1036    /// Clear the current actor binding while retaining the session actor roster.
1037    pub fn clear_actor_id(&self) {
1038        *self.actor_id.write() = None;
1039        self.session_metadata.write().actor_id = None;
1040    }
1041
1042    /// Set the current actor ID. Convenience wrapper around set_actor_id.
1043    pub fn set_user_id(&self, user_id: &str) -> ai_agents_core::Result<()> {
1044        self.set_actor_id(user_id)
1045    }
1046
1047    /// Load facts for the current actor from storage and cache them for prompt injection.
1048    pub async fn load_actor_memory(&self) -> ai_agents_core::Result<()> {
1049        let actor_id = match self.effective_actor_id() {
1050            Some(id) => id,
1051            None => return Ok(()),
1052        };
1053
1054        let store_opt = self.fact_store.read().clone();
1055        if let Some(store) = store_opt {
1056            let facts = store.get_facts(&actor_id).await?;
1057            let count = facts.len();
1058            self.actor_facts_cache
1059                .write()
1060                .insert(actor_id.clone(), facts);
1061            self.hooks.on_actor_memory_loaded(&actor_id, count).await;
1062            tracing::debug!("loaded {} facts for actor {}", count, actor_id);
1063        }
1064
1065        Ok(())
1066    }
1067
1068    /// Load actor memory only when the effective actor has no cached facts yet.
1069    async fn maybe_load_actor_memory(&self) {
1070        let Some(actor_id) = self.effective_actor_id() else {
1071            return;
1072        };
1073        if self.actor_facts_cache.read().contains_key(&actor_id) {
1074            return;
1075        }
1076        let _ = self.load_actor_memory().await;
1077    }
1078
1079    /// Pre-turn lifecycle shared by streaming and non-streaming paths.
1080    async fn pre_turn_session_lifecycle(&self) {
1081        if *self.redispatch_depth.read() > 0 {
1082            return;
1083        }
1084        self.resolve_actor_id_from_context();
1085        self.await_background_before_next_turn().await;
1086        self.record_session_actor_if_needed();
1087        self.maybe_load_actor_memory().await;
1088        self.maybe_load_actor_relationship().await;
1089        *self.messages_since_extraction.write() += 1;
1090    }
1091
1092    /// Post-turn lifecycle shared by streaming and non-streaming paths.
1093    async fn post_turn_session_lifecycle(&self) -> Result<()> {
1094        if *self.redispatch_depth.read() > 0 {
1095            return Ok(());
1096        }
1097        *self.messages_since_extraction.write() += 1;
1098        self.run_post_turn_maintenance().await
1099    }
1100
1101    /// Starts root-turn bookkeeping for user-message commit tracking.
1102    fn begin_root_turn(&self) {
1103        if *self.redispatch_depth.read() == 0 {
1104            let mut guard = self.active_turn_context.write();
1105            if guard.is_none() {
1106                self.root_user_message_committed
1107                    .store(false, Ordering::SeqCst);
1108                let max_calls = self
1109                    .runtime_config
1110                    .optimization
1111                    .max_speculative_llm_calls_per_turn;
1112                *guard = Some(TurnOptimizationContext::new(
1113                    String::new(),
1114                    HashMap::new(),
1115                    max_calls,
1116                ));
1117            }
1118        }
1119    }
1120
1121    fn update_active_turn_context(
1122        &self,
1123        processed_input: &str,
1124        input_context: HashMap<String, Value>,
1125    ) {
1126        if *self.redispatch_depth.read() > 0 {
1127            return;
1128        }
1129        let max_calls = self
1130            .runtime_config
1131            .optimization
1132            .max_speculative_llm_calls_per_turn;
1133        let mut guard = self.active_turn_context.write();
1134        match guard.as_mut() {
1135            Some(context) => {
1136                context.processed_input = processed_input.to_string();
1137                context.input_context = input_context;
1138                context.max_speculative_llm_calls = max_calls;
1139            }
1140            None => {
1141                *guard = Some(TurnOptimizationContext::new(
1142                    processed_input,
1143                    input_context,
1144                    max_calls,
1145                ));
1146            }
1147        }
1148    }
1149
1150    /// Writes the processed user message once for the root turn.
1151    async fn commit_root_user_message(&self, processed_input: &str) -> Result<()> {
1152        if *self.redispatch_depth.read() > 0 {
1153            return Ok(());
1154        }
1155        if !self
1156            .root_user_message_committed
1157            .swap(true, Ordering::SeqCst)
1158        {
1159            self.memory
1160                .add_message(ChatMessage::user(processed_input))
1161                .await?;
1162            if let Some(context) = self.active_turn_context.write().as_mut() {
1163                context.mark_user_message_committed();
1164            }
1165        }
1166        Ok(())
1167    }
1168
1169    /// Clears root-turn bookkeeping after final response handling.
1170    fn end_root_turn(&self) {
1171        if *self.redispatch_depth.read() == 0 {
1172            self.root_user_message_committed
1173                .store(false, Ordering::SeqCst);
1174            *self.active_turn_context.write() = None;
1175        }
1176    }
1177
1178    fn reserve_active_speculative_llm_call(&self, kind: RuntimeOptimizationKind) -> bool {
1179        self.begin_root_turn();
1180        let mut guard = self.active_turn_context.write();
1181        let Some(context) = guard.as_mut() else {
1182            return false;
1183        };
1184        context.reserve_speculative_llm_call_for(kind)
1185    }
1186
1187    fn branch_context_preview(&self) -> String {
1188        let context = self.build_context_with_overlays();
1189        let mut value = serde_json::to_string_pretty(&context).unwrap_or_else(|_| "{}".to_string());
1190        const MAX_CONTEXT_PREVIEW_CHARS: usize = 2048;
1191        if value.chars().count() > MAX_CONTEXT_PREVIEW_CHARS {
1192            value = value
1193                .chars()
1194                .take(MAX_CONTEXT_PREVIEW_CHARS)
1195                .collect::<String>();
1196            value.push_str("...");
1197        }
1198        value
1199    }
1200
1201    /// Applies freshness policy before rendering the next prompt.
1202    async fn await_background_before_next_turn(&self) {
1203        let optimization = &self.runtime_config.optimization;
1204        if !optimization.enabled {
1205            return;
1206        }
1207        let actor_id = self.effective_actor_id();
1208        let post = &optimization.post_turn;
1209        self.await_background_task(
1210            post.facts.await_before_next_turn,
1211            RuntimeTaskPurpose::PostTurnFacts,
1212            actor_id.as_deref(),
1213            "facts",
1214        )
1215        .await;
1216        self.await_background_task(
1217            post.relationships.await_before_next_turn,
1218            RuntimeTaskPurpose::PostTurnRelationship,
1219            actor_id.as_deref(),
1220            "relationships",
1221        )
1222        .await;
1223    }
1224
1225    async fn await_background_task(
1226        &self,
1227        policy: AwaitBeforeNextTurn,
1228        purpose: RuntimeTaskPurpose,
1229        actor_id: Option<&str>,
1230        label: &str,
1231    ) {
1232        match policy {
1233            AwaitBeforeNextTurn::Never => {}
1234            AwaitBeforeNextTurn::Always => {
1235                if let Err(error) = self.flush_background_tasks_for_purpose(purpose).await {
1236                    warn!(label = label, error = %error, "background maintenance flush failed");
1237                }
1238            }
1239            AwaitBeforeNextTurn::SameActor => {
1240                if let Some(actor_id) = actor_id
1241                    && let Err(error) = self
1242                        .flush_background_tasks_for_actor_purpose(actor_id, purpose)
1243                        .await
1244                {
1245                    warn!(label = label, actor_id = %actor_id, error = %error, "actor background maintenance flush failed");
1246                }
1247            }
1248        }
1249    }
1250
1251    /// Runs post-turn facts and relationship maintenance according to runtime policy.
1252    async fn run_post_turn_maintenance(&self) -> Result<()> {
1253        let optimization = &self.runtime_config.optimization;
1254        if !optimization.enabled {
1255            self.auto_extract_facts().await;
1256            self.auto_update_relationship().await;
1257            return Ok(());
1258        }
1259
1260        let facts_mode = effective_maintenance_mode(
1261            optimization.post_turn.facts.mode,
1262            optimization.parallel_post_turn_memory,
1263        );
1264        let relationships_mode = effective_maintenance_mode(
1265            optimization.post_turn.relationships.mode,
1266            optimization.parallel_post_turn_memory,
1267        );
1268
1269        match (facts_mode, relationships_mode) {
1270            (MaintenanceMode::InlineSerial, MaintenanceMode::InlineSerial) => {
1271                self.auto_extract_facts().await;
1272                self.auto_update_relationship().await;
1273            }
1274            (MaintenanceMode::InlineParallel, MaintenanceMode::InlineParallel) => {
1275                let facts = self.auto_extract_facts();
1276                let relationships = self.auto_update_relationship();
1277                tokio::join!(facts, relationships);
1278            }
1279            (MaintenanceMode::Background, MaintenanceMode::Background) => {
1280                self.schedule_facts_background().await?;
1281                self.schedule_relationship_background().await?;
1282            }
1283            (MaintenanceMode::Background, MaintenanceMode::InlineParallel)
1284            | (MaintenanceMode::Background, MaintenanceMode::InlineSerial) => {
1285                self.schedule_facts_background().await?;
1286                self.auto_update_relationship().await;
1287            }
1288            (MaintenanceMode::InlineParallel, MaintenanceMode::Background)
1289            | (MaintenanceMode::InlineSerial, MaintenanceMode::Background) => {
1290                self.auto_extract_facts().await;
1291                self.schedule_relationship_background().await?;
1292            }
1293            _ => {
1294                self.auto_extract_facts().await;
1295                self.auto_update_relationship().await;
1296            }
1297        }
1298        Ok(())
1299    }
1300
1301    async fn schedule_facts_background(&self) -> Result<()> {
1302        let policy = self.runtime_config.optimization.post_turn.facts.clone();
1303        let should_extract = self
1304            .facts_config
1305            .as_ref()
1306            .map(|c| c.enabled && c.auto_extract)
1307            .unwrap_or(false);
1308        if !should_extract {
1309            return Ok(());
1310        }
1311        let msgs_since = *self.messages_since_extraction.read();
1312        if msgs_since < 2 {
1313            return Ok(());
1314        }
1315        let Some(actor_id) = self.effective_actor_id() else {
1316            self.record_skipped_maintenance(
1317                "facts",
1318                ObservationPurpose::FactsExtraction,
1319                "missing_actor",
1320                Some(&policy),
1321            );
1322            return Ok(());
1323        };
1324        let Some(extractor) = self.fact_extractor.read().clone() else {
1325            return Ok(());
1326        };
1327        let messages = match self.memory.get_messages(None).await {
1328            Ok(messages) => messages,
1329            Err(error) => {
1330                warn!(error = %error, "failed to snapshot messages for fact extraction");
1331                return Ok(());
1332            }
1333        };
1334        let recent: Vec<_> = messages
1335            .iter()
1336            .rev()
1337            .take(msgs_since)
1338            .rev()
1339            .cloned()
1340            .collect();
1341        if recent.is_empty() {
1342            return Ok(());
1343        }
1344        let existing = self
1345            .actor_facts_cache
1346            .read()
1347            .get(&actor_id)
1348            .cloned()
1349            .unwrap_or_default();
1350        let categories = self
1351            .facts_config
1352            .as_ref()
1353            .map(|c| c.custom_categories.clone())
1354            .unwrap_or_default();
1355        let store = self.fact_store.read().clone();
1356        let cache = Arc::clone(&self.actor_facts_cache);
1357        let counter = Arc::clone(&self.messages_since_extraction);
1358        let hooks = Arc::clone(&self.hooks);
1359        let agent_id = self.info.id.clone();
1360        let observation = current_observation_context();
1361        let key = MaintenanceSequenceKey::actor(
1362            agent_id,
1363            actor_id.clone(),
1364            RuntimeTaskPurpose::PostTurnFacts,
1365        );
1366        let actor_for_task = actor_id.clone();
1367        let task = async move {
1368            let run = async move {
1369                let facts = extractor
1370                    .extract(&recent, &existing, Some(&actor_for_task), &categories)
1371                    .await?;
1372                if !facts.is_empty() {
1373                    if let Some(store) = store {
1374                        let authoritative = store.add_facts(&actor_for_task, facts.clone()).await?;
1375                        cache.write().insert(actor_for_task.clone(), authoritative);
1376                    } else {
1377                        cache
1378                            .write()
1379                            .entry(actor_for_task.clone())
1380                            .or_default()
1381                            .extend(facts.clone());
1382                    }
1383                    {
1384                        let mut count = counter.write();
1385                        if *count <= msgs_since {
1386                            *count = 0;
1387                        } else {
1388                            *count -= msgs_since;
1389                        }
1390                    }
1391                    hooks.on_facts_extracted(&actor_for_task, &facts).await;
1392                }
1393                Ok(())
1394            };
1395            if let Some(context) = observation {
1396                with_observation_context(
1397                    context.with_purpose(ObservationPurpose::FactsExtraction),
1398                    run,
1399                )
1400                .await
1401            } else {
1402                run.await
1403            }
1404        };
1405        self.spawn_or_handle_background(Some(key), task, "facts", &policy)
1406            .await
1407    }
1408
1409    async fn schedule_relationship_background(&self) -> Result<()> {
1410        let policy = self
1411            .runtime_config
1412            .optimization
1413            .post_turn
1414            .relationships
1415            .clone();
1416        let Some(manager) = self.relationship_manager.as_ref().cloned() else {
1417            return Ok(());
1418        };
1419        let Some(actor_id) = self.effective_actor_id() else {
1420            self.record_skipped_maintenance(
1421                "relationships",
1422                ObservationPurpose::RelationshipUpdate,
1423                "missing_actor",
1424                Some(&policy),
1425            );
1426            return Ok(());
1427        };
1428        let recent_messages = manager.config().auto_update.recent_messages;
1429        let messages = match self.memory.get_messages(Some(recent_messages)).await {
1430            Ok(messages) => messages,
1431            Err(error) => {
1432                warn!(actor = %actor_id, error = %error, "failed to snapshot messages for relationship update");
1433                return Ok(());
1434            }
1435        };
1436        let storage = self.storage.read().clone();
1437        let hooks = Arc::clone(&self.hooks);
1438        let agent_id = self.info.id.clone();
1439        let observation = current_observation_context();
1440        let key = MaintenanceSequenceKey::actor(
1441            agent_id.clone(),
1442            actor_id.clone(),
1443            RuntimeTaskPurpose::PostTurnRelationship,
1444        );
1445        let actor_for_task = actor_id.clone();
1446        let task = async move {
1447            let run = async move {
1448                if manager.config().auto_update.enabled {
1449                    let update = manager.auto_update(&actor_for_task, &messages).await?;
1450                    if !update.changes.is_empty() {
1451                        hooks
1452                            .on_relationship_change(&actor_for_task, &update.changes)
1453                            .await;
1454                    }
1455                    if let Some(ref event) = update.event {
1456                        hooks.on_notable_event(&actor_for_task, event).await;
1457                    }
1458                }
1459                if manager.config().persistence.enabled
1460                    && let (Some(storage), Some(value)) =
1461                        (storage, manager.relationship_as_value(&actor_for_task)?)
1462                {
1463                    storage
1464                        .save_relationship(&agent_id, &actor_for_task, &value)
1465                        .await?;
1466                }
1467                Ok(())
1468            };
1469            if let Some(context) = observation {
1470                with_observation_context(
1471                    context.with_purpose(ObservationPurpose::RelationshipUpdate),
1472                    run,
1473                )
1474                .await
1475            } else {
1476                run.await
1477            }
1478        };
1479        self.spawn_or_handle_background(Some(key), task, "relationships", &policy)
1480            .await
1481    }
1482
1483    /// Queues background maintenance or applies the configured overflow behavior.
1484    async fn spawn_or_handle_background<F>(
1485        &self,
1486        key: Option<MaintenanceSequenceKey>,
1487        task: F,
1488        label: &'static str,
1489        policy: &crate::optimization::config::MaintenanceTaskPolicy,
1490    ) -> Result<()>
1491    where
1492        F: Future<Output = Result<()>> + Send + 'static,
1493    {
1494        if self.background_maintenance.is_full() {
1495            match self
1496                .runtime_config
1497                .optimization
1498                .post_turn
1499                .on_background_overflow
1500            {
1501                BackgroundOverflowPolicy::RunInline => {
1502                    record_background_maintenance_event(
1503                        self.observability_manager.as_ref(),
1504                        label,
1505                        EventStatus::Success,
1506                        0,
1507                        "inline_overflow",
1508                        None,
1509                        Some(policy),
1510                    );
1511                    let start = Instant::now();
1512                    match task.await {
1513                        Ok(()) => record_background_maintenance_event(
1514                            self.observability_manager.as_ref(),
1515                            label,
1516                            EventStatus::Success,
1517                            start.elapsed().as_millis() as u64,
1518                            "inline_completed",
1519                            None,
1520                            Some(policy),
1521                        ),
1522                        Err(error) => {
1523                            warn!(label = label, error = %error, "inline maintenance fallback failed");
1524                            record_background_maintenance_event(
1525                                self.observability_manager.as_ref(),
1526                                label,
1527                                EventStatus::Error,
1528                                start.elapsed().as_millis() as u64,
1529                                "inline_failed",
1530                                Some(error.to_string()),
1531                                Some(policy),
1532                            );
1533                            return Err(error);
1534                        }
1535                    }
1536                }
1537                BackgroundOverflowPolicy::Drop => {
1538                    self.record_skipped_maintenance(
1539                        label,
1540                        ObservationPurpose::Other(label.to_string()),
1541                        "queue_full",
1542                        Some(policy),
1543                    );
1544                }
1545                BackgroundOverflowPolicy::Error => {
1546                    record_background_maintenance_event(
1547                        self.observability_manager.as_ref(),
1548                        label,
1549                        EventStatus::Error,
1550                        0,
1551                        "queue_full",
1552                        None,
1553                        Some(policy),
1554                    );
1555                    warn!(label = label, "background maintenance queue full");
1556                    return Err(AgentError::Other(format!(
1557                        "background maintenance queue is full for {}",
1558                        label
1559                    )));
1560                }
1561            }
1562            return Ok(());
1563        }
1564
1565        record_background_maintenance_event(
1566            self.observability_manager.as_ref(),
1567            label,
1568            EventStatus::Success,
1569            0,
1570            "scheduled",
1571            None,
1572            Some(policy),
1573        );
1574        let manager = self.observability_manager.clone();
1575        let policy_for_task = policy.clone();
1576        let observed_task = async move {
1577            let start = Instant::now();
1578            let result = task.await;
1579            match &result {
1580                Ok(()) => record_background_maintenance_event(
1581                    manager.as_ref(),
1582                    label,
1583                    EventStatus::Success,
1584                    start.elapsed().as_millis() as u64,
1585                    "completed",
1586                    None,
1587                    Some(&policy_for_task),
1588                ),
1589                Err(error) => record_background_maintenance_event(
1590                    manager.as_ref(),
1591                    label,
1592                    EventStatus::Error,
1593                    start.elapsed().as_millis() as u64,
1594                    "failed",
1595                    Some(error.to_string()),
1596                    Some(&policy_for_task),
1597                ),
1598            }
1599            result
1600        };
1601
1602        if let Err(error) = self.background_maintenance.spawn(key, observed_task) {
1603            record_background_maintenance_event(
1604                self.observability_manager.as_ref(),
1605                label,
1606                EventStatus::Error,
1607                0,
1608                "spawn_failed",
1609                Some(error.to_string()),
1610                Some(policy),
1611            );
1612            warn!(label = label, error = %error, "background maintenance spawn failed");
1613            return Err(error);
1614        }
1615        Ok(())
1616    }
1617
1618    /// Records a skipped background maintenance event when work cannot run.
1619    fn record_skipped_maintenance(
1620        &self,
1621        label: &str,
1622        purpose: ObservationPurpose,
1623        reason: &str,
1624        policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
1625    ) {
1626        if let Some(manager) = self.observability_manager.as_ref() {
1627            let mut tags = background_maintenance_tags(label, "skipped", Some(reason), policy);
1628            tags.insert("runtime.skip_reason".to_string(), reason.to_string());
1629            manager.record_lifecycle_event(
1630                EventType::MemoryOperation {
1631                    operation: format!("{}_maintenance", label),
1632                },
1633                purpose,
1634                EventStatus::Skipped,
1635                0,
1636                tags,
1637                None,
1638            );
1639        }
1640    }
1641
1642    /// Get cached actor facts for the effective actor.
1643    pub fn actor_facts(&self) -> Vec<ai_agents_core::KeyFact> {
1644        let Some(actor_id) = self.effective_actor_id() else {
1645            return Vec::new();
1646        };
1647        self.actor_facts_cache
1648            .read()
1649            .get(&actor_id)
1650            .cloned()
1651            .unwrap_or_default()
1652    }
1653
1654    /// Returns the formatted relationship prompt text for the effective actor, if relationship injection produced any text for this turn.
1655    pub fn relationship_memory_text(&self) -> Option<String> {
1656        self.format_relationship_for_context().map(|(_, text)| text)
1657    }
1658
1659    /// Manually extract facts from the last N messages.
1660    pub async fn extract_facts(
1661        &self,
1662        last_n: usize,
1663    ) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
1664        self.extract_facts_with_source(last_n, "manual").await
1665    }
1666
1667    async fn extract_facts_with_source(
1668        &self,
1669        last_n: usize,
1670        source: &'static str,
1671    ) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
1672        let extractor = match self.fact_extractor.read().clone() {
1673            Some(e) => e,
1674            None => return Ok(vec![]),
1675        };
1676
1677        let messages = self.memory.get_messages(None).await?;
1678        let recent: Vec<_> = messages.iter().rev().take(last_n).rev().cloned().collect();
1679
1680        if recent.is_empty() {
1681            return Ok(vec![]);
1682        }
1683
1684        let actor_id = self.effective_actor_id();
1685        let existing = actor_id
1686            .as_ref()
1687            .and_then(|aid| self.actor_facts_cache.read().get(aid).cloned())
1688            .unwrap_or_default();
1689
1690        let categories = self
1691            .facts_config
1692            .as_ref()
1693            .map(|c| c.custom_categories.clone())
1694            .unwrap_or_default();
1695
1696        let facts = self
1697            .observe_purpose(
1698                ObservationPurpose::FactsExtraction,
1699                extractor.extract(&recent, &existing, actor_id.as_deref(), &categories),
1700            )
1701            .await?;
1702
1703        // Save to storage and update the actor-scoped cache when an actor is known.
1704        if !facts.is_empty() {
1705            let fact_store_opt = self.fact_store.read().clone();
1706            let mut stored_total = 0usize;
1707            let mut cache_updated = false;
1708            if let (Some(store), Some(aid)) = (fact_store_opt, &actor_id) {
1709                // add_facts now returns the authoritative post-write set.
1710                let authoritative = store.add_facts(aid, facts.clone()).await?;
1711                stored_total = authoritative.len();
1712                self.actor_facts_cache
1713                    .write()
1714                    .insert(aid.clone(), authoritative);
1715                cache_updated = true;
1716            } else if let Some(aid) = &actor_id {
1717                let mut cache = self.actor_facts_cache.write();
1718                let entry = cache.entry(aid.clone()).or_default();
1719                entry.extend(facts.clone());
1720                stored_total = entry.len();
1721                cache_updated = true;
1722            }
1723
1724            info!(
1725                actor_id = %actor_id.as_deref().unwrap_or("<none>"),
1726                source = source,
1727                requested_messages = last_n,
1728                message_count = recent.len(),
1729                extracted_count = facts.len(),
1730                cache_updated = cache_updated,
1731                stored_total = stored_total,
1732                "facts extracted"
1733            );
1734
1735            if let Some(ref aid) = actor_id {
1736                self.hooks.on_facts_extracted(aid, &facts).await;
1737            }
1738        }
1739
1740        Ok(facts)
1741    }
1742
1743    /// Resolve actor_id from context if method is from_context.
1744    /// Supports dotted paths (e.g. "player.id", "user.profile.id").
1745    fn resolve_actor_id_from_context(&self) {
1746        if self
1747            .current_turn_actor_context()
1748            .and_then(|ctx| ctx.effective_actor_id().map(str::to_string))
1749            .is_some()
1750        {
1751            return;
1752        }
1753
1754        if let Some(ref am_config) = self.actor_memory_config
1755            && am_config.identification.method == ai_agents_facts::IdentificationMethod::FromContext
1756            && let Some(ref path) = am_config.identification.context_path
1757        {
1758            // get_path resolves dotted paths; get only handles top-level keys.
1759            let val = self
1760                .context_manager
1761                .get_path(path)
1762                .or_else(|| self.context_manager.get(path));
1763            if let Some(val) = val
1764                && let Some(id_str) = val.as_str()
1765            {
1766                let current = self.actor_id.read().clone();
1767                if current.as_deref() != Some(id_str) {
1768                    *self.actor_id.write() = Some(id_str.to_string());
1769                    let mut meta = self.session_metadata.write();
1770                    meta.actor_id = Some(id_str.to_string());
1771                    if !meta.actors.iter().any(|a| a == id_str) {
1772                        meta.actors.push(id_str.to_string());
1773                    }
1774                }
1775            }
1776        }
1777    }
1778
1779    /// Format actor facts for template injection.
1780    fn format_actor_facts_for_context(&self) -> String {
1781        // Respect inject_in_context: false.
1782        let should_inject = self
1783            .facts_config
1784            .as_ref()
1785            .map(|c| c.inject_in_context)
1786            .unwrap_or(true);
1787        if !should_inject {
1788            return String::new();
1789        }
1790
1791        let Some(actor_id) = self.effective_actor_id() else {
1792            return String::new();
1793        };
1794
1795        let facts = self
1796            .actor_facts_cache
1797            .read()
1798            .get(&actor_id)
1799            .cloned()
1800            .unwrap_or_default();
1801        if facts.is_empty() {
1802            return String::new();
1803        }
1804
1805        let am_config = self.actor_memory_config.as_ref();
1806        // Effective token cap: prefer memory.token_budget.allocation.facts when present,
1807        // otherwise fall back to actor_memory.injection.max_tokens.
1808        let facts_budget = self
1809            .memory_token_budget
1810            .as_ref()
1811            .map(|b| b.allocation.facts as usize)
1812            .filter(|n| *n > 0);
1813        let default_max = am_config.map(|c| c.injection.max_tokens).unwrap_or(800);
1814        let max_tokens = facts_budget.unwrap_or(default_max);
1815
1816        // Filter by category when injection.mode = category.
1817        let filtered: Vec<ai_agents_core::KeyFact> = if let Some(cfg) = am_config {
1818            if cfg.injection.mode == ai_agents_facts::InjectionMode::OnDemand {
1819                return String::new();
1820            }
1821            if cfg.injection.mode == ai_agents_facts::InjectionMode::Category
1822                && !cfg.injection.categories.is_empty()
1823            {
1824                facts
1825                    .iter()
1826                    .filter(|f| {
1827                        cfg.injection
1828                            .categories
1829                            .iter()
1830                            .any(|c| f.category.to_string() == *c)
1831                    })
1832                    .cloned()
1833                    .collect()
1834            } else {
1835                facts.clone()
1836            }
1837        } else {
1838            facts.clone()
1839        };
1840
1841        if filtered.is_empty() {
1842            return String::new();
1843        }
1844
1845        if let Some(store) = self.fact_store.read().clone() {
1846            store.format_for_context(&filtered, max_tokens)
1847        } else {
1848            String::new()
1849        }
1850    }
1851
1852    fn build_context_with_staged(&self, staged: &HashMap<String, Value>) -> HashMap<String, Value> {
1853        let context = self.build_context_with_overlays();
1854        let mut root = Value::Object(context.into_iter().collect());
1855        for (path, value) in staged {
1856            if let Ok(updated) = ai_agents_core::set_dot_path(root.clone(), path, value.clone()) {
1857                root = updated;
1858            }
1859        }
1860        match root {
1861            Value::Object(obj) => obj.into_iter().collect(),
1862            _ => HashMap::new(),
1863        }
1864    }
1865
1866    fn build_context_with_overlays(&self) -> HashMap<String, Value> {
1867        let mut context = self.context_manager.get_all();
1868        let mut root = Value::Object(context.clone().into_iter().collect());
1869
1870        if let Some(turn_ctx) = self.current_turn_actor_context() {
1871            if let Some(ref origin_actor_id) = turn_ctx.origin_actor_id
1872                && let Ok(updated) = ai_agents_core::set_dot_path(
1873                    root.clone(),
1874                    "interaction.origin_actor_id",
1875                    serde_json::json!(origin_actor_id),
1876                )
1877            {
1878                root = updated;
1879            }
1880            if let Some(ref sender_agent_id) = turn_ctx.sender_agent_id
1881                && let Ok(updated) = ai_agents_core::set_dot_path(
1882                    root.clone(),
1883                    "interaction.sender_agent_id",
1884                    serde_json::json!(sender_agent_id),
1885                )
1886            {
1887                root = updated;
1888            }
1889        }
1890
1891        if let Some(ref actor_id) = self.effective_actor_id()
1892            && let Ok(updated) = ai_agents_core::set_dot_path(
1893                root.clone(),
1894                "interaction.actor_id",
1895                serde_json::json!(actor_id),
1896            )
1897        {
1898            root = updated;
1899        }
1900
1901        if let Some(manager) = self.relationship_manager.as_ref()
1902            && let Some(actor_id) = self.effective_actor_id()
1903            && let Some(value) = manager.to_context_value(&actor_id)
1904            && let Ok(updated) = ai_agents_core::set_dot_path(
1905                root.clone(),
1906                &manager.config().injection.context_path,
1907                value,
1908            )
1909        {
1910            root = updated;
1911        }
1912
1913        if let Value::Object(obj) = root {
1914            context = obj.into_iter().collect();
1915        }
1916
1917        context
1918    }
1919
1920    fn resolve_actor_name_from_context(&self) -> Option<String> {
1921        for path in ["actor.name", "user.name", "player.name", "customer.name"] {
1922            if let Some(value) = self.context_manager.get_path(path)
1923                && let Some(name) = value.as_str()
1924            {
1925                return Some(name.to_string());
1926            }
1927        }
1928        None
1929    }
1930
1931    async fn maybe_load_actor_relationship(&self) {
1932        let Some(manager) = self.relationship_manager.as_ref() else {
1933            return;
1934        };
1935        let Some(actor_id) = self.effective_actor_id() else {
1936            return;
1937        };
1938
1939        let mut should_fire_loaded = false;
1940        if manager.get(&actor_id).is_none() {
1941            let mut loaded = false;
1942            if manager.config().persistence.enabled {
1943                let storage = self.storage.read().clone();
1944                if let Some(storage) = storage {
1945                    match storage.load_relationship(&self.info.id, &actor_id).await {
1946                        Ok(Some(value)) => match manager.insert_from_value(value) {
1947                            Ok(_) => loaded = true,
1948                            Err(e) => {
1949                                warn!(actor = %actor_id, error = %e, "failed to restore relationship")
1950                            }
1951                        },
1952                        Ok(None) => {}
1953                        Err(e) => {
1954                            warn!(actor = %actor_id, error = %e, "failed to load relationship")
1955                        }
1956                    }
1957                }
1958            }
1959
1960            if !loaded {
1961                manager.get_or_create(&actor_id, self.resolve_actor_name_from_context().as_deref());
1962            }
1963            should_fire_loaded = true;
1964        }
1965
1966        let actor_name = self.resolve_actor_name_from_context();
1967        let relationship = manager.touch_interaction(&actor_id, actor_name.as_deref());
1968        if should_fire_loaded {
1969            self.hooks
1970                .on_relationship_loaded(&actor_id, &relationship)
1971                .await;
1972        }
1973    }
1974
1975    fn format_relationship_for_context(&self) -> Option<(String, String)> {
1976        let manager = self.relationship_manager.as_ref()?;
1977        if !manager.config().injection.enabled {
1978            return None;
1979        }
1980        let actor_id = self.effective_actor_id()?;
1981        let relationship = manager.get(&actor_id)?;
1982        let local_cap = manager.config().injection.max_tokens;
1983        let global_cap = self
1984            .memory_token_budget
1985            .as_ref()
1986            .map(|b| b.allocation.relationships as usize)
1987            .filter(|n| *n > 0);
1988        let max_tokens = global_cap.map(|g| g.min(local_cap)).unwrap_or(local_cap);
1989        let text = ai_agents_relationships::format_relationship(
1990            &relationship,
1991            &manager.config().injection.format,
1992            max_tokens,
1993        );
1994        if text.is_empty() {
1995            None
1996        } else {
1997            Some((manager.config().injection.prompt_variable.clone(), text))
1998        }
1999    }
2000
2001    async fn persist_actor_relationship(&self, actor_id: &str) -> Result<()> {
2002        let Some(manager) = self.relationship_manager.as_ref() else {
2003            return Ok(());
2004        };
2005        if !manager.config().persistence.enabled {
2006            return Ok(());
2007        }
2008        let storage = self.storage.read().clone();
2009        let Some(storage) = storage else {
2010            return Ok(());
2011        };
2012        if let Some(value) = manager.relationship_as_value(actor_id)? {
2013            storage
2014                .save_relationship(&self.info.id, actor_id, &value)
2015                .await?;
2016        }
2017        Ok(())
2018    }
2019
2020    async fn auto_update_relationship(&self) {
2021        let Some(manager) = self.relationship_manager.as_ref() else {
2022            return;
2023        };
2024        let Some(actor_id) = self.effective_actor_id() else {
2025            return;
2026        };
2027        if !manager.config().auto_update.enabled {
2028            let _ = self.persist_actor_relationship(&actor_id).await;
2029            return;
2030        }
2031
2032        let recent_messages = manager.config().auto_update.recent_messages;
2033        let messages = match self.memory.get_messages(Some(recent_messages)).await {
2034            Ok(messages) => messages,
2035            Err(e) => {
2036                warn!(actor = %actor_id, error = %e, "failed to read messages for relationship update");
2037                return;
2038            }
2039        };
2040
2041        match self
2042            .observe_purpose(
2043                ObservationPurpose::RelationshipUpdate,
2044                manager.auto_update(&actor_id, &messages),
2045            )
2046            .await
2047        {
2048            Ok(update) => {
2049                if !update.changes.is_empty() {
2050                    self.hooks
2051                        .on_relationship_change(&actor_id, &update.changes)
2052                        .await;
2053                }
2054                if let Some(ref event) = update.event {
2055                    self.hooks.on_notable_event(&actor_id, event).await;
2056                }
2057                let persisted = match self.persist_actor_relationship(&actor_id).await {
2058                    Ok(()) => true,
2059                    Err(e) => {
2060                        warn!(actor = %actor_id, error = %e, "failed to persist relationship");
2061                        false
2062                    }
2063                };
2064                if !update.changes.is_empty() || update.event.is_some() {
2065                    let changed_dimensions: Vec<String> = update
2066                        .changes
2067                        .iter()
2068                        .map(|change| format!("{}:{}", change.perspective, change.dimension))
2069                        .collect();
2070                    info!(
2071                        actor_id = %actor_id,
2072                        change_count = update.changes.len(),
2073                        changed_dimensions = ?changed_dimensions,
2074                        event_present = update.event.is_some(),
2075                        persisted = persisted,
2076                        "relationship updated"
2077                    );
2078                } else {
2079                    debug!(actor_id = %actor_id, persisted = persisted, "relationship evaluation ran but found no changes");
2080                }
2081            }
2082            Err(e) => warn!(actor = %actor_id, error = %e, "relationship update failed"),
2083        }
2084    }
2085
2086    /// Run auto-extraction after a chat turn.
2087    async fn auto_extract_facts(&self) {
2088        let should_extract = self
2089            .facts_config
2090            .as_ref()
2091            .map(|c| c.enabled && c.auto_extract)
2092            .unwrap_or(false);
2093
2094        if !should_extract {
2095            debug!("fact extraction skipped because auto extraction is disabled");
2096            return;
2097        }
2098
2099        let msgs_since = *self.messages_since_extraction.read();
2100        if msgs_since < 2 {
2101            debug!(
2102                messages_since_extraction = msgs_since,
2103                "fact extraction skipped until threshold is reached"
2104            );
2105            return;
2106        }
2107
2108        match self.extract_facts_with_source(msgs_since, "auto").await {
2109            Ok(facts) => {
2110                if !facts.is_empty() {
2111                    *self.messages_since_extraction.write() = 0;
2112                } else {
2113                    debug!("fact extraction ran but found no new facts");
2114                }
2115            }
2116            Err(e) => {
2117                warn!("fact extraction failed: {}", e);
2118            }
2119        }
2120    }
2121
2122    pub fn with_persona(mut self, manager: Arc<ai_agents_persona::PersonaManager>) -> Self {
2123        self.persona_manager = Some(manager);
2124        self
2125    }
2126
2127    pub fn persona_manager(&self) -> Option<&Arc<ai_agents_persona::PersonaManager>> {
2128        self.persona_manager.as_ref()
2129    }
2130
2131    pub fn with_disambiguation(mut self, config: DisambiguationConfig) -> Self {
2132        if config.is_enabled() {
2133            let manager = DisambiguationManager::new(config, Arc::clone(&self.llm_registry))
2134                .with_clarification_observer(Arc::new(ObservabilityClarificationObserver));
2135            self.disambiguation_manager = Some(manager);
2136        }
2137        self
2138    }
2139
2140    pub fn disambiguation_manager(&self) -> Option<&DisambiguationManager> {
2141        self.disambiguation_manager.as_ref()
2142    }
2143
2144    pub fn has_disambiguation(&self) -> bool {
2145        self.disambiguation_manager
2146            .as_ref()
2147            .is_some_and(|m| m.is_enabled())
2148    }
2149
2150    pub async fn init_storage(&self) -> Result<()> {
2151        //
2152        // One readiness guard prevents concurrent entry points from constructing storage or fact state more than once.
2153        //
2154        let _guard = self.storage_init.lock().await;
2155        let mut storage = self.storage.read().clone();
2156        if storage.is_none() && !self.storage_config.is_none() {
2157            let storage_config = self.convert_storage_config();
2158            storage = create_storage(&storage_config).await?;
2159            *self.storage.write() = storage.clone();
2160        }
2161
2162        self.validate_storage_requirements(storage.as_deref())?;
2163        self.complete_facts_init().await;
2164        Ok(())
2165    }
2166
2167    fn validate_storage_requirements(&self, storage: Option<&dyn AgentStorage>) -> Result<()> {
2168        let facts_required = self
2169            .facts_config
2170            .as_ref()
2171            .is_some_and(|config| config.enabled)
2172            || self
2173                .actor_memory_config
2174                .as_ref()
2175                .is_some_and(|config| config.enabled);
2176        let relationships_required = self
2177            .relationship_manager
2178            .as_ref()
2179            .is_some_and(|manager| manager.config().persistence.enabled);
2180
2181        let Some(storage) = storage else {
2182            let mut requirements = Vec::new();
2183            if facts_required {
2184                requirements.push("actor facts or actor memory");
2185            }
2186            if relationships_required {
2187                requirements.push("persistent relationships");
2188            }
2189            if requirements.is_empty() {
2190                return Ok(());
2191            }
2192            return Err(AgentError::Config(format!(
2193                "Storage is required for enabled {} but none is configured or injected",
2194                requirements.join(" and ")
2195            )));
2196        };
2197
2198        //
2199        // Durable features must fail before execution instead of degrading to volatile state.
2200        //
2201        if facts_required && !storage.supports(StorageCapability::ActorFacts) {
2202            return Err(AgentError::UnsupportedStorageCapability(
2203                StorageCapability::ActorFacts,
2204            ));
2205        }
2206        if relationships_required && !storage.supports(StorageCapability::ActorRelationships) {
2207            return Err(AgentError::UnsupportedStorageCapability(
2208                StorageCapability::ActorRelationships,
2209            ));
2210        }
2211        Ok(())
2212    }
2213
2214    /// Initialize fact store and extractor from stored config and current storage.
2215    /// Called from init_storage() so facts are ready before the first turn.
2216    async fn complete_facts_init(&self) {
2217        if self.fact_store.read().is_some() {
2218            return;
2219        }
2220        let storage = match self.storage.read().clone() {
2221            Some(s) => s,
2222            None => return,
2223        };
2224
2225        let facts_enabled = self
2226            .facts_config
2227            .as_ref()
2228            .map(|f| f.enabled)
2229            .unwrap_or(false);
2230        let actor_memory_enabled = self
2231            .actor_memory_config
2232            .as_ref()
2233            .map(|a| a.enabled)
2234            .unwrap_or(false);
2235
2236        if !facts_enabled && !actor_memory_enabled {
2237            return;
2238        }
2239
2240        let fc = self.facts_config.clone().unwrap_or_default();
2241        let store = Arc::new(ai_agents_facts::FactStore::new(
2242            storage,
2243            self.info.id.clone(),
2244            fc.clone(),
2245        ));
2246
2247        let extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>> = if facts_enabled {
2248            let extractor_llm = fc
2249                .extractor_llm
2250                .as_ref()
2251                .and_then(|alias| self.llm_registry.get(alias).ok())
2252                .or_else(|| self.llm_registry.router().ok())
2253                .or_else(|| self.llm_registry.default().ok());
2254            extractor_llm.map(|llm| {
2255                Arc::new(ai_agents_facts::LLMFactExtractor::new(llm, fc.clone()))
2256                    as Arc<dyn ai_agents_facts::FactExtractor>
2257            })
2258        } else {
2259            None
2260        };
2261
2262        *self.fact_store.write() = Some(store);
2263        *self.fact_extractor.write() = extractor;
2264        debug!(
2265            agent = %self.info.id,
2266            facts_enabled,
2267            actor_memory_enabled,
2268            "facts storage initialized"
2269        );
2270    }
2271
2272    fn convert_storage_config(&self) -> StorageStorageConfig {
2273        crate::spec::storage::to_storage_config(&self.storage_config)
2274    }
2275
2276    pub fn storage(&self) -> Option<Arc<dyn AgentStorage>> {
2277        self.storage.read().clone()
2278    }
2279
2280    pub fn storage_config(&self) -> &StorageConfig {
2281        &self.storage_config
2282    }
2283
2284    /// Returns the spawner if configured via a spawner: YAML section.
2285    pub fn spawner(&self) -> Option<&Arc<crate::spawner::AgentSpawner>> {
2286        self.spawner.as_ref()
2287    }
2288
2289    /// Returns the agent registry if configured via a spawner: YAML section.
2290    pub fn spawner_registry(&self) -> Option<&Arc<crate::spawner::AgentRegistry>> {
2291        self.spawner_registry.as_ref()
2292    }
2293
2294    pub fn has_spawner(&self) -> bool {
2295        self.spawner_registry.is_some()
2296    }
2297
2298    pub fn with_spawner_handles(
2299        mut self,
2300        spawner: Arc<crate::spawner::AgentSpawner>,
2301        registry: Arc<crate::spawner::AgentRegistry>,
2302    ) -> Self {
2303        self.spawner = Some(spawner);
2304        self.spawner_registry = Some(registry);
2305        self
2306    }
2307
2308    pub fn with_hooks(mut self, hooks: Arc<dyn AgentHooks>) -> Self {
2309        self.hooks = hooks;
2310        self
2311    }
2312
2313    pub fn with_parallel_tools(mut self, config: ParallelToolsConfig) -> Self {
2314        self.parallel_tools = config;
2315        self
2316    }
2317
2318    pub fn with_streaming(mut self, config: StreamingConfig) -> Self {
2319        self.streaming = config;
2320        self
2321    }
2322
2323    pub fn with_hitl(mut self, engine: HITLEngine, handler: Arc<dyn ApprovalHandler>) -> Self {
2324        self.hitl_engine = Some(engine);
2325        self.approval_handler = handler;
2326        self
2327    }
2328
2329    pub fn with_max_context_tokens(mut self, tokens: u32) -> Self {
2330        self.max_context_tokens = tokens;
2331        self
2332    }
2333
2334    pub fn with_memory_token_budget(mut self, budget: MemoryTokenBudget) -> Self {
2335        self.memory_token_budget = Some(budget);
2336        self
2337    }
2338
2339    pub fn with_recovery_manager(mut self, manager: RecoveryManager) -> Self {
2340        self.recovery_manager = manager;
2341        self
2342    }
2343
2344    pub fn with_tool_security(mut self, engine: ToolSecurityEngine) -> Self {
2345        self.tool_security = engine;
2346        self
2347    }
2348
2349    /// Returns the host-only runtime control handle.
2350    pub fn runtime_control(&self) -> RuntimeControlHandle {
2351        RuntimeControlHandle {
2352            state: Arc::clone(&self.runtime_control),
2353        }
2354    }
2355
2356    /// Installs or clears the host question handler used by `ask_user`.
2357    pub fn set_question_handler(&self, handler: Option<Arc<dyn QuestionHandler>>) {
2358        self.tools.set_question_handler(handler);
2359    }
2360
2361    /// Installs the host diagnostics provider used by `diagnostics`.
2362    pub fn set_diagnostics_provider(&self, provider: Arc<dyn DiagnosticsProvider>) {
2363        self.tools.set_diagnostics_provider(provider);
2364    }
2365
2366    /// Installs the host command runner used by `command`.
2367    pub fn set_command_runner(&self, runner: Arc<dyn CommandRunner>) {
2368        self.tools.set_command_runner(runner);
2369    }
2370
2371    /// Installs the host web-search provider used by `web_search`.
2372    pub fn set_web_search_provider(&self, provider: Arc<dyn ai_agents_tools::WebSearchProvider>) {
2373        self.tools.set_web_search_provider(provider);
2374    }
2375
2376    /// Returns the current session-local todo list.
2377    pub fn todos(&self) -> Vec<TodoItem> {
2378        self.tools.todos()
2379    }
2380
2381    /// Returns the effective tool security snapshot for the next decision.
2382    fn active_tool_security(&self) -> ToolSecurityEngine {
2383        self.runtime_control
2384            .tool_security_override
2385            .read()
2386            .clone()
2387            .unwrap_or_else(|| self.tool_security.clone())
2388    }
2389
2390    /// Reads policy, scope, emergency state, and generation as one safety snapshot.
2391    fn runtime_safety_snapshot(&self) -> RuntimeSafetySnapshot {
2392        let _guard = self.runtime_control.snapshot_guard.read();
2393        RuntimeSafetySnapshot {
2394            version: self.runtime_control.version.load(Ordering::SeqCst),
2395            emergency_deny: self.runtime_control.emergency_deny.load(Ordering::SeqCst),
2396            tool_security: self
2397                .runtime_control
2398                .tool_security_override
2399                .read()
2400                .clone()
2401                .unwrap_or_else(|| self.tool_security.clone()),
2402            tool_scope_override: self.runtime_control.tool_scope_override.read().clone(),
2403        }
2404    }
2405
2406    /// Performs the single rate admission after locks are held and rejects changed runtime, policy, or state authority.
2407    fn admit_tool_execution(
2408        &self,
2409        expected_runtime_version: u64,
2410        expected_policy_version: u64,
2411        expected_state_generation: Option<u64>,
2412        canonical_id: &str,
2413    ) -> SecurityCheckResult {
2414        let _guard = self.runtime_control.snapshot_guard.read();
2415        if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
2416            return SecurityCheckResult::Block {
2417                reason: "runtime emergency deny is enabled".to_string(),
2418            };
2419        }
2420        let runtime_version = self.runtime_control.version.load(Ordering::SeqCst);
2421        let security_engine = self
2422            .runtime_control
2423            .tool_security_override
2424            .read()
2425            .clone()
2426            .unwrap_or_else(|| self.tool_security.clone());
2427        if runtime_version != expected_runtime_version
2428            || security_engine.policy_version() != expected_policy_version
2429        {
2430            return SecurityCheckResult::Block {
2431                reason: "runtime safety controls changed before admission".to_string(),
2432            };
2433        }
2434        let current_state_generation = self
2435            .state_machine
2436            .as_ref()
2437            .map(|state_machine| state_machine.generation());
2438        if current_state_generation != expected_state_generation {
2439            return SecurityCheckResult::Block {
2440                reason: "state scope changed before admission".to_string(),
2441            };
2442        }
2443        security_engine.admit_tool_execution(canonical_id)
2444    }
2445
2446    pub fn with_process_processor(mut self, processor: ProcessProcessor) -> Self {
2447        let processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
2448        self.process_processor = Some(processor);
2449        self
2450    }
2451
2452    pub fn with_state_machine(
2453        mut self,
2454        state_machine: Arc<StateMachine>,
2455        evaluator: Arc<dyn TransitionEvaluator>,
2456    ) -> Self {
2457        self.state_machine = Some(state_machine);
2458        self.transition_evaluator = Some(evaluator);
2459        self
2460    }
2461
2462    pub fn with_context_manager(mut self, manager: Arc<ContextManager>) -> Self {
2463        self.context_manager = manager;
2464        self
2465    }
2466
2467    pub fn register_message_filter(&self, name: impl Into<String>, filter: Arc<dyn MessageFilter>) {
2468        self.message_filters.write().insert(name.into(), filter);
2469    }
2470
2471    pub fn set_context(&self, key: &str, value: Value) -> Result<()> {
2472        self.context_manager.update(key, value)
2473    }
2474
2475    pub fn update_context(&self, path: &str, value: Value) -> Result<()> {
2476        self.context_manager.update(path, value)
2477    }
2478
2479    pub fn get_context(&self) -> HashMap<String, Value> {
2480        self.build_context_with_overlays()
2481    }
2482
2483    pub fn remove_context(&self, key: &str) -> Option<Value> {
2484        self.context_manager.remove(key)
2485    }
2486
2487    pub async fn refresh_context(&self, key: &str) -> Result<()> {
2488        self.context_manager.refresh(key).await
2489    }
2490
2491    pub fn register_context_provider(&self, name: &str, provider: Arc<dyn ContextProvider>) {
2492        self.context_manager.register_provider(name, provider);
2493    }
2494
2495    pub fn current_state(&self) -> Option<String> {
2496        self.state_machine.as_ref().map(|sm| sm.current())
2497    }
2498
2499    // State changes invalidate only resolved confirmation ownership; ordinary clarification remains manager-owned across layered enablement changes.
2500    async fn invalidate_pending_confirmation(&self, reason: &'static str) {
2501        self.disambiguation_epoch.fetch_add(1, Ordering::SeqCst);
2502        let Some(disambiguator) = self.disambiguation_manager.as_ref() else {
2503            return;
2504        };
2505        if disambiguator.has_pending_confirmation().await {
2506            disambiguator.clear_pending().await;
2507            *self.pending_skill_id.write() = None;
2508            info!(
2509                confirmation_event = "invalidated",
2510                invalidation_reason = reason,
2511                "Runtime invalidated pending confirmation"
2512            );
2513        }
2514    }
2515
2516    // Admission linearizes pending publication or redispatch before reset and state mutation, then rechecks both ownership generations.
2517    async fn admit_disambiguation_redispatch(
2518        &self,
2519        expected_epoch: u64,
2520        expected_state_generation: Option<u64>,
2521    ) -> Result<tokio::sync::RwLockReadGuard<'_, ()>> {
2522        let admission = self.disambiguation_admission.read().await;
2523        let state_generation = self
2524            .state_machine
2525            .as_ref()
2526            .map(|state_machine| state_machine.generation());
2527        if self.disambiguation_epoch.load(Ordering::SeqCst) != expected_epoch
2528            || state_generation != expected_state_generation
2529        {
2530            return Err(AgentError::Other(
2531                "Disambiguation ownership changed before redispatch admission".to_string(),
2532            ));
2533        }
2534        Ok(admission)
2535    }
2536
2537    // A transition reservation prevents losing runtime transitions from running duplicate exit side effects.
2538    fn reserve_state_transition(&self) -> Option<StateTransitionReservation<'_>> {
2539        self.state_transition_reserved
2540            .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
2541            .ok()
2542            .map(|_| StateTransitionReservation {
2543                reserved: &self.state_transition_reserved,
2544            })
2545    }
2546
2547    // Optional ownership is used only for terminal skill responses that came from manager-owned pending state.
2548    async fn admit_optional_disambiguation_ownership(
2549        &self,
2550        ownership: Option<DisambiguationOwnership>,
2551    ) -> Result<Option<tokio::sync::RwLockReadGuard<'_, ()>>> {
2552        match ownership {
2553            Some(ownership) => self
2554                .admit_disambiguation_redispatch(ownership.epoch, ownership.state_generation)
2555                .await
2556                .map(Some),
2557            None => Ok(None),
2558        }
2559    }
2560
2561    /// Applies a manual transition after reserving its exit actions and keeping async lifecycle work outside the commit lock.
2562    pub async fn transition_to(&self, state: &str) -> Result<()> {
2563        let Some(ref sm) = self.state_machine else {
2564            return Ok(());
2565        };
2566        let claim_admission = self.disambiguation_admission.write().await;
2567        let reservation = self.reserve_state_transition().ok_or_else(|| {
2568            AgentError::Other("Another state transition is already in progress".to_string())
2569        })?;
2570        let from_state = sm.current();
2571        let expected_state_generation = sm.generation();
2572        let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
2573        let history_before = sm.history();
2574        drop(claim_admission);
2575
2576        self.execute_state_exit_actions(&from_state).await;
2577
2578        let admission = self.disambiguation_admission.write().await;
2579        if sm.current() != from_state
2580            || sm.generation() != expected_state_generation
2581            || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
2582        {
2583            return Err(AgentError::Other(
2584                "State ownership changed during manual transition preparation".to_string(),
2585            ));
2586        }
2587        sm.transition_to(state, "manual transition")?;
2588        self.invalidate_pending_confirmation("state_transition")
2589            .await;
2590        let entered = sm.current();
2591        let is_reentry = Self::state_was_previously_entered(&entered, &from_state, &history_before);
2592        drop(admission);
2593
2594        self.execute_state_enter_actions(&entered, is_reentry).await;
2595        drop(reservation);
2596        info!(to = %entered, "Manual state transition");
2597        Ok(())
2598    }
2599
2600    pub fn state_history(&self) -> Vec<StateTransitionEvent> {
2601        self.state_machine
2602            .as_ref()
2603            .map(|sm| sm.history())
2604            .unwrap_or_default()
2605    }
2606
2607    /// Get a copy of current session metadata.
2608    pub fn session_metadata(&self) -> ai_agents_core::SessionMetadata {
2609        self.session_metadata.read().clone()
2610    }
2611
2612    /// Delete all facts and sessions for an actor, gated by privacy.allow_deletion.
2613    /// Returns Err when actor_memory.privacy.allow_deletion is false.
2614    pub async fn delete_actor_data(&self, actor_id: &str) -> Result<()> {
2615        let allowed = self
2616            .actor_memory_config
2617            .as_ref()
2618            .map(|c| c.privacy.allow_deletion)
2619            .unwrap_or(true);
2620        if !allowed {
2621            return Err(AgentError::Config(
2622                "privacy.allow_deletion is false; actor data deletion is not permitted".into(),
2623            ));
2624        }
2625        let storage = self.storage.read().clone();
2626        if let Some(storage) = storage {
2627            //
2628            // Composite support is checked before mutation so privacy deletion cannot become partial.
2629            //
2630            if !storage.supports(StorageCapability::ActorDataDeletion) {
2631                return Err(AgentError::UnsupportedStorageCapability(
2632                    StorageCapability::ActorDataDeletion,
2633                ));
2634            }
2635            storage.delete_actor_data(&self.info.id, actor_id).await?;
2636        } else {
2637            //
2638            // Clone the fallback store before awaiting so backend I/O never holds the runtime read guard.
2639            //
2640            let store = { self.fact_store.read().clone() };
2641            if let Some(store) = store {
2642                store.delete_actor_data(actor_id).await?;
2643            }
2644        }
2645        if let Some(manager) = self.relationship_manager.as_ref() {
2646            manager.remove(actor_id);
2647        }
2648        self.actor_facts_cache.write().remove(actor_id);
2649        Ok(())
2650    }
2651
2652    /// Overwrite session metadata (tags, ttl, custom fields).
2653    pub fn set_session_metadata(&self, meta: ai_agents_core::SessionMetadata) {
2654        *self.session_metadata.write() = meta;
2655    }
2656
2657    /// Delete sessions whose TTL has expired. Returns number of sessions removed.
2658    pub async fn cleanup_expired_sessions(&self) -> Result<usize> {
2659        let storage = self.storage.read().clone();
2660        match storage {
2661            Some(s) => {
2662                let count = s.cleanup_expired().await?;
2663                if count > 0 {
2664                    self.hooks.on_sessions_expired(count).await;
2665                }
2666                Ok(count)
2667            }
2668            None => Err(AgentError::Config(
2669                "No storage configured. Use with_storage_config() or with_storage() first".into(),
2670            )),
2671        }
2672    }
2673
2674    /// List sessions matching a filter. Supports actor, tag, and date filters.
2675    pub async fn list_sessions_filtered(
2676        &self,
2677        filter: &ai_agents_core::SessionFilter,
2678    ) -> Result<Vec<ai_agents_core::SessionSummary>> {
2679        let storage = self.storage.read().clone();
2680        match storage {
2681            Some(s) => s.list_sessions_filtered(filter).await,
2682            None => Err(AgentError::Config(
2683                "No storage configured. Use with_storage_config() or with_storage() first".into(),
2684            )),
2685        }
2686    }
2687
2688    pub async fn save_state(&self) -> Result<AgentSnapshot> {
2689        let memory_snapshot = self.memory.snapshot().await?;
2690        let state_machine_snapshot = self.state_machine.as_ref().map(|sm| sm.snapshot());
2691        let context_snapshot = self.context_manager.snapshot();
2692
2693        let mut snapshot = AgentSnapshot::new(self.info.id.clone())
2694            .with_memory(memory_snapshot)
2695            .with_context(context_snapshot)
2696            .with_state_machine(
2697                state_machine_snapshot.unwrap_or_else(|| StateMachineSnapshot {
2698                    current_state: String::new(),
2699                    previous_state: None,
2700                    turn_count: 0,
2701                    no_transition_count: 0,
2702                    history: vec![],
2703                }),
2704            );
2705
2706        if let Some(ref persona) = self.persona_manager {
2707            snapshot.persona = Some(persona.snapshot_as_value()?);
2708        }
2709
2710        if let Some(ref relationships) = self.relationship_manager {
2711            snapshot.relationships = Some(relationships.snapshot_as_value()?);
2712        }
2713
2714        Ok(snapshot)
2715    }
2716
2717    /// Save state including spawned agents manifest for session persistence.
2718    pub async fn save_state_full(&self) -> Result<AgentSnapshot> {
2719        let mut snapshot = self.save_state().await?;
2720        if let Some(ref registry) = self.spawner_registry {
2721            let entries = registry.list_with_specs();
2722            if !entries.is_empty() {
2723                snapshot = snapshot.with_spawned_agents(entries);
2724            }
2725        }
2726        Ok(snapshot)
2727    }
2728
2729    /// Restores persisted state after invalidating any pending confirmation ownership.
2730    pub async fn restore_state(&self, snapshot: AgentSnapshot) -> Result<()> {
2731        let _admission = self.disambiguation_admission.write().await;
2732        if self.state_transition_reserved.load(Ordering::SeqCst) {
2733            return Err(AgentError::Other(
2734                "Cannot restore state while a state transition is in progress".to_string(),
2735            ));
2736        }
2737        self.invalidate_pending_confirmation("state_restore").await;
2738        *self.pending_skill_id.write() = None;
2739        if let Some(disambiguator) = self.disambiguation_manager.as_ref() {
2740            disambiguator.clear_pending().await;
2741        }
2742        self.memory.restore(snapshot.memory).await?;
2743
2744        if let (Some(sm), Some(sm_snapshot)) = (&self.state_machine, snapshot.state_machine)
2745            && !sm_snapshot.current_state.is_empty()
2746        {
2747            sm.restore(sm_snapshot)?;
2748        }
2749
2750        self.context_manager.restore(snapshot.context);
2751
2752        if let (Some(persona_value), Some(persona_manager)) =
2753            (snapshot.persona, &self.persona_manager)
2754        {
2755            persona_manager.restore_from_value(persona_value)?;
2756        }
2757
2758        if let (Some(relationship_value), Some(relationship_manager)) =
2759            (snapshot.relationships, &self.relationship_manager)
2760        {
2761            relationship_manager.restore_from_value(relationship_value)?;
2762        }
2763
2764        info!(agent_id = %snapshot.agent_id, "State restored");
2765        Ok(())
2766    }
2767
2768    pub async fn save_to(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<()> {
2769        let snapshot = self.save_state().await?;
2770        storage.save(session_id, &snapshot).await
2771    }
2772
2773    async fn load_session_restore(
2774        storage: &dyn AgentStorage,
2775        session_id: &str,
2776    ) -> Result<Option<StoredSessionRestore>> {
2777        let Some(snapshot) = storage.load(session_id).await? else {
2778            return Ok(None);
2779        };
2780        //
2781        // Metadata is restored only when the backend explicitly owns that durable contract.
2782        //
2783        let metadata = if storage.supports(StorageCapability::SessionMetadata) {
2784            storage.load_metadata(session_id).await?
2785        } else {
2786            None
2787        };
2788        Ok(Some(StoredSessionRestore { snapshot, metadata }))
2789    }
2790
2791    async fn capture_session_restore_point(&self) -> Result<RuntimeSessionRestorePoint> {
2792        Ok(RuntimeSessionRestorePoint {
2793            snapshot: self.save_state().await?,
2794            metadata: self.session_metadata(),
2795            actor_id: self.actor_id(),
2796            session_id: self.current_session_id.read().clone(),
2797        })
2798    }
2799
2800    async fn apply_session_restore_unchecked(
2801        &self,
2802        session_id: &str,
2803        stored: StoredSessionRestore,
2804    ) -> Result<()> {
2805        self.restore_state(stored.snapshot).await?;
2806        let metadata = stored.metadata.unwrap_or_default();
2807        if let Some(actor_id) = metadata.actor_id.as_deref() {
2808            self.set_actor_id(actor_id)?;
2809        } else {
2810            self.clear_actor_id();
2811        }
2812        self.set_session_metadata(metadata);
2813        *self.current_session_id.write() = Some(session_id.to_string());
2814        Ok(())
2815    }
2816
2817    async fn restore_session_restore_point(
2818        &self,
2819        restore_point: &RuntimeSessionRestorePoint,
2820    ) -> Result<()> {
2821        self.restore_state(restore_point.snapshot.clone()).await?;
2822        if let Some(actor_id) = restore_point.actor_id.as_deref() {
2823            self.set_actor_id(actor_id)?;
2824        } else {
2825            self.clear_actor_id();
2826        }
2827        self.set_session_metadata(restore_point.metadata.clone());
2828        *self.current_session_id.write() = restore_point.session_id.clone();
2829        Ok(())
2830    }
2831
2832    async fn apply_session_restore(
2833        &self,
2834        session_id: &str,
2835        stored: StoredSessionRestore,
2836    ) -> Result<()> {
2837        let before = self.capture_session_restore_point().await?;
2838        if let Err(error) = self
2839            .apply_session_restore_unchecked(session_id, stored)
2840            .await
2841        {
2842            return match self.restore_session_restore_point(&before).await {
2843                Ok(()) => Err(error),
2844                Err(rollback_error) => Err(AgentError::Other(format!(
2845                    "Session restore failed: {error}; rollback failed: {rollback_error}"
2846                ))),
2847            };
2848        }
2849        Ok(())
2850    }
2851
2852    async fn rollback_session_restore_set(
2853        parent: Option<(&RuntimeAgent, &RuntimeSessionRestorePoint)>,
2854        children: &[(String, Arc<RuntimeAgent>, RuntimeSessionRestorePoint)],
2855    ) -> Vec<String> {
2856        let mut errors = Vec::new();
2857        if let Some((agent, restore_point)) = parent
2858            && let Err(error) = agent.restore_session_restore_point(restore_point).await
2859        {
2860            errors.push(format!("parent: {error}"));
2861        }
2862        for (id, agent, restore_point) in children {
2863            if let Err(error) = agent.restore_session_restore_point(restore_point).await {
2864                errors.push(format!("child '{id}': {error}"));
2865            }
2866        }
2867        errors
2868    }
2869
2870    fn restore_failure(error: impl std::fmt::Display, rollback_errors: Vec<String>) -> AgentError {
2871        if rollback_errors.is_empty() {
2872            AgentError::Other(format!(
2873                "Session restore failed: {error}; runtime state was rolled back"
2874            ))
2875        } else {
2876            AgentError::Other(format!(
2877                "Session restore failed: {error}; rollback also failed for {}",
2878                rollback_errors.join(", ")
2879            ))
2880        }
2881    }
2882
2883    pub async fn load_from(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<bool> {
2884        let Some(stored) = Self::load_session_restore(storage, session_id).await? else {
2885            return Ok(false);
2886        };
2887        self.apply_session_restore(session_id, stored).await?;
2888        Ok(true)
2889    }
2890
2891    pub async fn save_session(&self, session_id: &str) -> Result<()> {
2892        let storage = self.storage.read().clone();
2893        match storage {
2894            Some(s) => {
2895                // Fire on_session_created when this session id is first seen on this runtime.
2896                let is_new = {
2897                    let cur = self.current_session_id.read().clone();
2898                    cur.as_deref() != Some(session_id)
2899                };
2900                if is_new {
2901                    *self.current_session_id.write() = Some(session_id.to_string());
2902                    self.hooks.on_session_created(session_id).await;
2903                }
2904
2905                // Update metadata before persisting.
2906                {
2907                    let now = chrono::Utc::now();
2908                    let msg_count = self
2909                        .memory
2910                        .get_messages(None)
2911                        .await
2912                        .map(|v| v.len())
2913                        .unwrap_or(0);
2914                    let mut meta = self.session_metadata.write();
2915                    meta.last_active = now;
2916                    meta.message_count = msg_count;
2917                    if meta.actor_id.is_none() {
2918                        meta.actor_id = self.actor_id.read().clone();
2919                    }
2920                }
2921
2922                let snapshot = self.save_state().await?;
2923                //
2924                // Metadata-capable backends own atomic snapshot and metadata persistence.
2925                //
2926                if s.supports(StorageCapability::SessionMetadata) {
2927                    let metadata = self.session_metadata.read().clone();
2928                    s.save_snapshot_with_metadata(session_id, &snapshot, &metadata)
2929                        .await
2930                } else {
2931                    s.save(session_id, &snapshot).await
2932                }
2933            }
2934            None => Err(AgentError::Config(
2935                "No storage configured. Use with_storage_config() or with_storage() first".into(),
2936            )),
2937        }
2938    }
2939
2940    pub async fn load_session(&self, session_id: &str) -> Result<bool> {
2941        let storage = self.storage.read().clone();
2942        match storage {
2943            Some(storage) => self.load_from(storage.as_ref(), session_id).await,
2944            None => Err(AgentError::Config(
2945                "No storage configured. Use with_storage_config() or with_storage() first".into(),
2946            )),
2947        }
2948    }
2949
2950    /// Restore this runtime and its complete saved child topology from one named session.
2951    pub async fn restore_session_full(&self, session_id: &str) -> Result<usize> {
2952        self.init_storage().await?;
2953        let storage = self.storage.read().clone().ok_or_else(|| {
2954            AgentError::Config(
2955                "No storage configured. Use with_storage_config() or with_storage() first".into(),
2956            )
2957        })?;
2958        let target_parent = Self::load_session_restore(storage.as_ref(), session_id)
2959            .await?
2960            .ok_or_else(|| AgentError::Persistence(format!("Session not found: {session_id}")))?;
2961        let manifest = target_parent
2962            .snapshot
2963            .spawned_agents
2964            .clone()
2965            .unwrap_or_default();
2966
2967        let registry = self.spawner_registry.as_ref().cloned();
2968        let spawner = if manifest.is_empty() {
2969            self.spawner.as_ref().cloned()
2970        } else {
2971            Some(self.spawner.as_ref().cloned().ok_or_else(|| {
2972                AgentError::Config(
2973                    "Saved session contains child agents but this runtime has no spawner".into(),
2974                )
2975            })?)
2976        };
2977        let registry = if manifest.is_empty() {
2978            registry
2979        } else {
2980            Some(registry.ok_or_else(|| {
2981                AgentError::Config(
2982                    "Saved session contains child agents but this runtime has no registry".into(),
2983                )
2984            })?)
2985        };
2986
2987        let mut target_ids = HashSet::with_capacity(manifest.len());
2988        let mut prepared = Vec::with_capacity(manifest.len());
2989        for entry in manifest {
2990            if !target_ids.insert(entry.id.clone()) {
2991                return Err(AgentError::InvalidSpec(format!(
2992                    "Saved child manifest contains duplicate ID: {}",
2993                    entry.id
2994                )));
2995            }
2996            let spec = crate::spec::AgentSpec::from_yaml_strict(&entry.spec_yaml)?;
2997            spawner
2998                .as_ref()
2999                .expect("non-empty manifests require a spawner")
3000                .validate_explicit_child(&entry.id, &spec)?;
3001            prepared.push((entry.id, spec));
3002        }
3003
3004        let current_ids = registry
3005            .as_ref()
3006            .map(|registry| {
3007                registry
3008                    .list()
3009                    .into_iter()
3010                    .map(|info| info.id)
3011                    .collect::<HashSet<_>>()
3012            })
3013            .unwrap_or_default();
3014        let removal_count = current_ids.difference(&target_ids).count();
3015        let additions = prepared
3016            .iter()
3017            .filter(|(id, _)| !current_ids.contains(id))
3018            .cloned()
3019            .collect::<Vec<_>>();
3020
3021        let mut existing = Vec::new();
3022        if let Some(registry) = registry.as_ref() {
3023            for (id, _) in prepared.iter().filter(|(id, _)| current_ids.contains(id)) {
3024                let agent = registry.get(id).ok_or_else(|| {
3025                    AgentError::Config(format!("Retained child disappeared during restore: {id}"))
3026                })?;
3027                let child_storage = agent.storage().ok_or_else(|| {
3028                    AgentError::Config(format!("Child '{id}' has no storage for session restore"))
3029                })?;
3030                let stored = Self::load_session_restore(child_storage.as_ref(), session_id)
3031                    .await?
3032                    .ok_or_else(|| {
3033                        AgentError::Persistence(format!(
3034                            "Child '{id}' has no saved session '{session_id}'"
3035                        ))
3036                    })?;
3037                existing.push((id.clone(), agent, stored));
3038            }
3039        }
3040
3041        let mut staged = Vec::with_capacity(additions.len());
3042        if !additions.is_empty() {
3043            let spawner = spawner
3044                .as_ref()
3045                .expect("restored additions require a spawner");
3046            let reservations = spawner.reserve_restore_capacity(additions.len(), removal_count)?;
3047            for ((id, spec), reservation) in additions.into_iter().zip(reservations) {
3048                let spawned = spawner
3049                    .spawn_with_reserved_capacity(id.clone(), spec, reservation)
3050                    .await?;
3051                let child_storage = spawned.agent.storage().ok_or_else(|| {
3052                    AgentError::Config(format!("Child '{id}' has no storage for session restore"))
3053                })?;
3054                let stored = Self::load_session_restore(child_storage.as_ref(), session_id)
3055                    .await?
3056                    .ok_or_else(|| {
3057                        AgentError::Persistence(format!(
3058                            "Child '{id}' has no saved session '{session_id}'"
3059                        ))
3060                    })?;
3061                staged.push((spawned, stored));
3062            }
3063        } else if let Some(spawner) = spawner.as_ref() {
3064            spawner.reserve_restore_capacity(0, removal_count)?;
3065        }
3066
3067        let parent_before = self.capture_session_restore_point().await?;
3068        let mut existing_before = Vec::with_capacity(existing.len());
3069        for (id, agent, _) in &existing {
3070            existing_before.push((
3071                id.clone(),
3072                Arc::clone(agent),
3073                agent.capture_session_restore_point().await?,
3074            ));
3075        }
3076
3077        //
3078        // Every record is loaded before runtime state changes. Registry replacement is the final topology commit point.
3079        //
3080        for (_, agent, stored) in &existing {
3081            if let Err(error) = agent
3082                .apply_session_restore_unchecked(session_id, stored.clone())
3083                .await
3084            {
3085                drop(staged);
3086                let rollback_errors =
3087                    Self::rollback_session_restore_set(None, &existing_before).await;
3088                return Err(Self::restore_failure(error, rollback_errors));
3089            }
3090        }
3091        for (spawned, stored) in &staged {
3092            if let Err(error) = spawned
3093                .agent
3094                .apply_session_restore_unchecked(session_id, stored.clone())
3095                .await
3096            {
3097                drop(staged);
3098                let rollback_errors =
3099                    Self::rollback_session_restore_set(None, &existing_before).await;
3100                return Err(Self::restore_failure(error, rollback_errors));
3101            }
3102        }
3103        if let Err(error) = self
3104            .apply_session_restore_unchecked(session_id, target_parent)
3105            .await
3106        {
3107            drop(staged);
3108            let rollback_errors =
3109                Self::rollback_session_restore_set(Some((self, &parent_before)), &existing_before)
3110                    .await;
3111            return Err(Self::restore_failure(error, rollback_errors));
3112        }
3113
3114        if let Some(registry) = registry.as_ref()
3115            && let Err(error) = registry
3116                .reconcile(
3117                    &target_ids,
3118                    staged.into_iter().map(|(spawned, _)| spawned).collect(),
3119                )
3120                .await
3121        {
3122            let rollback_errors =
3123                Self::rollback_session_restore_set(Some((self, &parent_before)), &existing_before)
3124                    .await;
3125            return Err(Self::restore_failure(error, rollback_errors));
3126        }
3127
3128        Ok(target_ids.len())
3129    }
3130
3131    pub async fn delete_session(&self, session_id: &str) -> Result<()> {
3132        let storage = self.storage.read().clone();
3133        match storage {
3134            Some(s) => s.delete(session_id).await,
3135            None => Err(AgentError::Config(
3136                "No storage configured. Use with_storage_config() or with_storage() first".into(),
3137            )),
3138        }
3139    }
3140
3141    pub async fn list_sessions(&self) -> Result<Vec<String>> {
3142        let storage = self.storage.read().clone();
3143        match storage {
3144            Some(s) => s.list_sessions().await,
3145            None => Err(AgentError::Config(
3146                "No storage configured. Use with_storage_config() or with_storage() first".into(),
3147            )),
3148        }
3149    }
3150
3151    fn estimate_tokens(&self, text: &str) -> u32 {
3152        (text.len() as f32 / 4.0).ceil() as u32
3153    }
3154
3155    fn estimate_total_tokens(&self, messages: &[ChatMessage]) -> u32 {
3156        messages
3157            .iter()
3158            .map(|m| self.estimate_tokens(&m.content))
3159            .sum()
3160    }
3161
3162    fn truncate_context(&self, messages: &mut Vec<ChatMessage>, keep_recent: usize) {
3163        if messages.len() <= keep_recent + 1 {
3164            return;
3165        }
3166        let system_msg = messages.remove(0);
3167        let to_remove = messages.len().saturating_sub(keep_recent);
3168        messages.drain(..to_remove);
3169        messages.insert(0, system_msg);
3170    }
3171
3172    fn get_filter(&self, config: &FilterConfig) -> Arc<dyn MessageFilter> {
3173        match config {
3174            FilterConfig::KeepRecent(n) => Arc::new(KeepRecentFilter::new(*n)),
3175            FilterConfig::ByRole { keep_roles } => Arc::new(ByRoleFilter::new(keep_roles.clone())),
3176            FilterConfig::SkipPattern { skip_if_contains } => {
3177                Arc::new(SkipPatternFilter::new(skip_if_contains.clone()))
3178            }
3179            FilterConfig::Custom { name } => {
3180                let filters = self.message_filters.read();
3181                filters
3182                    .get(name)
3183                    .cloned()
3184                    .unwrap_or_else(|| Arc::new(KeepRecentFilter::new(10)))
3185            }
3186        }
3187    }
3188
3189    async fn summarize_context(
3190        &self,
3191        messages: &mut Vec<ChatMessage>,
3192        summarizer_llm: Option<&str>,
3193        max_summary_tokens: u32,
3194        custom_prompt: Option<&str>,
3195        keep_recent: usize,
3196        filter: Option<&FilterConfig>,
3197    ) -> Result<()> {
3198        let system_msg = messages.remove(0);
3199
3200        let to_summarize_count = messages.len().saturating_sub(keep_recent);
3201        if to_summarize_count == 0 {
3202            messages.insert(0, system_msg);
3203            return Ok(());
3204        }
3205
3206        let recent_msgs: Vec<ChatMessage> = messages.drain(to_summarize_count..).collect();
3207        let mut to_summarize = std::mem::take(messages);
3208
3209        if let Some(filter_config) = filter {
3210            let filter = self.get_filter(filter_config);
3211            to_summarize = filter.filter(to_summarize);
3212        }
3213
3214        if to_summarize.is_empty() {
3215            *messages = recent_msgs;
3216            messages.insert(0, system_msg);
3217            return Ok(());
3218        }
3219
3220        let conversation_text = to_summarize
3221            .iter()
3222            .map(|m| format!("{:?}: {}", m.role, m.content))
3223            .collect::<Vec<_>>()
3224            .join("\n");
3225
3226        let default_prompt = format!(
3227            "Summarize the following conversation in under {} tokens, preserving key information:\n\n{}",
3228            max_summary_tokens, conversation_text
3229        );
3230
3231        let summary_prompt = custom_prompt
3232            .map(|p| format!("{}\n\n{}", p, conversation_text))
3233            .unwrap_or(default_prompt);
3234
3235        let summarizer = if let Some(alias) = summarizer_llm {
3236            self.llm_registry
3237                .get(alias)
3238                .map_err(|e| AgentError::Config(e.to_string()))?
3239        } else {
3240            self.llm_registry
3241                .router()
3242                .or_else(|_| self.llm_registry.default())
3243                .map_err(|e| AgentError::Config(e.to_string()))?
3244        };
3245
3246        let summary_msgs = vec![ChatMessage::user(&summary_prompt)];
3247        let response = self
3248            .observe_purpose(
3249                ObservationPurpose::Summarization,
3250                summarizer.complete(&summary_msgs, None),
3251            )
3252            .await?;
3253
3254        let summary_message = ChatMessage::system(format!(
3255            "[Previous conversation summary]\n{}",
3256            response.content
3257        ));
3258
3259        *messages = vec![system_msg, summary_message];
3260        messages.extend(recent_msgs);
3261
3262        debug!(
3263            summarized_count = to_summarize_count,
3264            kept_recent = keep_recent,
3265            "Context summarized"
3266        );
3267
3268        Ok(())
3269    }
3270
3271    fn render_system_prompt(&self) -> Result<String> {
3272        let mut context = self.build_context_with_overlays();
3273
3274        // Inject actor_facts for {{ actor_facts }} template variable.
3275        let facts_text = self.format_actor_facts_for_context();
3276        if !facts_text.is_empty() {
3277            context.insert(
3278                "actor_facts".to_string(),
3279                serde_json::Value::String(facts_text),
3280            );
3281        }
3282
3283        if let Some((key, text)) = self.format_relationship_for_context() {
3284            context.insert(key, serde_json::Value::String(text));
3285        }
3286
3287        self.template_renderer
3288            .render(&self.base_system_prompt, &context)
3289    }
3290
3291    /// Canonicalizes IDs once while preserving their first declaration order and ignoring unknown entries.
3292    fn canonical_unique_tool_ids(&self, ids: &[String]) -> Vec<String> {
3293        let mut seen = HashSet::new();
3294        ids.iter()
3295            .filter_map(|id| self.tools.canonical_id(id))
3296            .filter(|canonical_id| seen.insert(canonical_id.clone()))
3297            .collect()
3298    }
3299
3300    /// Applies runtime narrowing to the canonical declared grant without permitting scope-based expansion.
3301    fn get_top_level_tool_ids_for_scope(&self, scope_override: Option<&[String]>) -> Vec<String> {
3302        let Some(declared) = self.declared_tool_ids.as_deref() else {
3303            return Vec::new();
3304        };
3305        let mut effective = self.canonical_unique_tool_ids(declared);
3306        if let Some(scope) = scope_override {
3307            let scope: HashSet<String> =
3308                self.canonical_unique_tool_ids(scope).into_iter().collect();
3309            effective.retain(|canonical_id| scope.contains(canonical_id));
3310        }
3311        effective
3312    }
3313
3314    /// Reads the live runtime narrowing and returns the current deterministic effective tool IDs.
3315    async fn get_available_tool_ids(&self) -> Result<Vec<String>> {
3316        Ok(self.get_available_tool_ids_snapshot().await?.tool_ids)
3317    }
3318
3319    /// Captures the runtime scope used for ordinary availability checks before applying state narrowing.
3320    async fn get_available_tool_ids_snapshot(&self) -> Result<AvailableToolIdsSnapshot> {
3321        let scope_override = self.runtime_control.tool_scope_override.read().clone();
3322        self.get_available_tool_ids_snapshot_for_scope(scope_override.as_deref())
3323            .await
3324    }
3325
3326    /// Applies every explicit ancestor and current-state scope to one runtime snapshot and retains its state generation.
3327    async fn get_available_tool_ids_snapshot_for_scope(
3328        &self,
3329        scope_override: Option<&[String]>,
3330    ) -> Result<AvailableToolIdsSnapshot> {
3331        let mut available = self.get_top_level_tool_ids_for_scope(scope_override);
3332        let (state_generation, state_scopes) = self
3333            .state_machine
3334            .as_ref()
3335            .map(|state_machine| {
3336                let (generation, scopes) = state_machine.current_tool_scope_snapshot();
3337                (Some(generation), scopes)
3338            })
3339            .unwrap_or((None, Vec::new()));
3340
3341        if available.is_empty() || state_scopes.is_empty() {
3342            return Ok(AvailableToolIdsSnapshot {
3343                tool_ids: available,
3344                state_generation,
3345            });
3346        }
3347
3348        let eval_ctx = self.build_evaluation_context().await?;
3349        let llm_getter = RegistryLLMGetter {
3350            registry: self.llm_registry.clone(),
3351        };
3352        let evaluator = ConditionEvaluator::new(llm_getter);
3353
3354        for state_scope in state_scopes {
3355            if state_scope.is_empty() {
3356                available.clear();
3357                break;
3358            }
3359
3360            let mut allowed = HashSet::new();
3361            for tool_ref in &state_scope {
3362                let tool_id = tool_ref.id();
3363                let Some(canonical_id) = self.tools.canonical_id(tool_id) else {
3364                    continue;
3365                };
3366                let condition_matches = if let Some(condition) = tool_ref.condition() {
3367                    match evaluator.evaluate(condition, &eval_ctx).await {
3368                        Ok(matches) => matches,
3369                        Err(error) => {
3370                            warn!(tool = tool_id, error = %error, "Error evaluating tool condition");
3371                            false
3372                        }
3373                    }
3374                } else {
3375                    true
3376                };
3377                if condition_matches {
3378                    allowed.insert(canonical_id);
3379                } else {
3380                    debug!(tool = tool_id, "Tool condition not met, skipping");
3381                }
3382            }
3383            available.retain(|canonical_id| allowed.contains(canonical_id));
3384            if available.is_empty() {
3385                break;
3386            }
3387        }
3388
3389        Ok(AvailableToolIdsSnapshot {
3390            tool_ids: available,
3391            state_generation,
3392        })
3393    }
3394
3395    async fn build_evaluation_context(&self) -> Result<EvaluationContext> {
3396        let context = self.build_context_with_overlays();
3397        let messages = self.memory.get_messages(Some(10)).await?;
3398        let tool_history = self.tool_call_history.read().clone();
3399
3400        let (state_name, turn_count, previous_state) = if let Some(ref sm) = self.state_machine {
3401            (Some(sm.current()), sm.turn_count(), sm.previous())
3402        } else {
3403            (None, 0, None)
3404        };
3405
3406        Ok(EvaluationContext::default()
3407            .with_context(context)
3408            .with_state(state_name, turn_count, previous_state)
3409            .with_called_tools(tool_history)
3410            .with_messages(messages))
3411    }
3412
3413    fn record_tool_call(&self, tool_id: &str, result: Value) {
3414        self.tool_call_history.write().push(ToolCallRecord {
3415            tool_id: tool_id.to_string(),
3416            result,
3417            timestamp: chrono::Utc::now(),
3418        });
3419    }
3420
3421    async fn get_effective_system_prompt_with_persona_hooks(
3422        &self,
3423        fire_persona_hooks: bool,
3424        include_tool_prompt: bool,
3425    ) -> Result<String> {
3426        let rendered_base = self.render_system_prompt()?;
3427
3428        let persona_prefix = if let Some(ref persona) = self.persona_manager {
3429            let context = self.build_context_with_overlays();
3430            if fire_persona_hooks {
3431                let render_result = persona.render_prompt(&context)?;
3432                for content in &render_result.newly_revealed {
3433                    self.hooks.on_secret_revealed(content).await;
3434                }
3435                render_result.prompt
3436            } else {
3437                persona.render_prompt_preview(&context)?
3438            }
3439        } else {
3440            String::new()
3441        };
3442
3443        if let Some(ref sm) = self.state_machine
3444            && let Some(state_def) = sm.current_definition()
3445        {
3446            let state_prompt = if let Some(ref prompt) = state_def.prompt {
3447                let context = self.build_context_with_overlays();
3448                self.template_renderer.render_with_state(
3449                    prompt,
3450                    &context,
3451                    &sm.current(),
3452                    sm.previous().as_deref(),
3453                    sm.turn_count(),
3454                    state_def.max_turns,
3455                )?
3456            } else {
3457                String::new()
3458            };
3459
3460            let combined = match state_def.prompt_mode {
3461                PromptMode::Append => {
3462                    if state_prompt.is_empty() {
3463                        rendered_base
3464                    } else {
3465                        format!(
3466                            "{}\n\n[Current State: {}]\n{}",
3467                            rendered_base,
3468                            sm.current(),
3469                            state_prompt
3470                        )
3471                    }
3472                }
3473                PromptMode::Replace => {
3474                    if state_prompt.is_empty() {
3475                        rendered_base
3476                    } else {
3477                        state_prompt
3478                    }
3479                }
3480                PromptMode::Prepend => {
3481                    if state_prompt.is_empty() {
3482                        rendered_base
3483                    } else {
3484                        format!("{}\n\n{}", state_prompt, rendered_base)
3485                    }
3486                }
3487            };
3488
3489            // Persona always prepended regardless of prompt_mode.
3490            let with_persona = if persona_prefix.is_empty() {
3491                combined
3492            } else {
3493                format!("{}\n\n{}", persona_prefix, combined)
3494            };
3495
3496            if include_tool_prompt {
3497                let available_tool_ids = self.get_available_tool_ids().await?;
3498                if !available_tool_ids.is_empty() {
3499                    let tools_prompt = self.tools.generate_scoped_prompt_with_mode(
3500                        &available_tool_ids,
3501                        None,
3502                        self.parallel_tools.enabled,
3503                        self.runtime_config.tool_schema_prompt_mode,
3504                    );
3505                    if !tools_prompt.is_empty() {
3506                        return Ok(format!("{}\n\n{}", with_persona, tools_prompt));
3507                    }
3508                }
3509            }
3510            return Ok(with_persona);
3511        }
3512
3513        // No state machine - prepend persona to base.
3514        let with_persona = if persona_prefix.is_empty() {
3515            rendered_base
3516        } else {
3517            format!("{}\n\n{}", persona_prefix, rendered_base)
3518        };
3519
3520        if include_tool_prompt {
3521            let available_tool_ids = self.get_available_tool_ids().await?;
3522            let tools_prompt = self.tools.generate_scoped_prompt_with_mode(
3523                &available_tool_ids,
3524                None,
3525                self.parallel_tools.enabled,
3526                self.runtime_config.tool_schema_prompt_mode,
3527            );
3528            if !tools_prompt.is_empty() {
3529                return Ok(format!("{}\n\n{}", with_persona, tools_prompt));
3530            }
3531        }
3532        Ok(with_persona)
3533    }
3534
3535    fn get_state_llm(&self) -> Result<Arc<dyn LLMProvider>> {
3536        if let Some(ref sm) = self.state_machine
3537            && let Some(state_def) = sm.current_definition()
3538            && let Some(ref llm_alias) = state_def.llm
3539        {
3540            return self
3541                .llm_registry
3542                .get(llm_alias)
3543                .map_err(|e| AgentError::Config(e.to_string()));
3544        }
3545        self.llm_registry
3546            .default()
3547            .map_err(|e| AgentError::Config(e.to_string()))
3548    }
3549
3550    fn get_effective_reasoning_config(&self) -> ReasoningConfig {
3551        if let Some(ref sm) = self.state_machine
3552            && let Some(state_def) = sm.current_definition()
3553            && let Some(ref state_reasoning) = state_def.reasoning
3554        {
3555            return state_reasoning.clone();
3556        }
3557        self.reasoning_config.clone()
3558    }
3559
3560    fn get_effective_reflection_config(&self) -> ReflectionConfig {
3561        if let Some(ref sm) = self.state_machine
3562            && let Some(state_def) = sm.current_definition()
3563            && let Some(ref state_reflection) = state_def.reflection
3564        {
3565            return state_reflection.clone();
3566        }
3567        self.reflection_config.clone()
3568    }
3569
3570    fn get_skill_reasoning_config(&self, skill: &SkillDefinition) -> ReasoningConfig {
3571        skill
3572            .reasoning
3573            .clone()
3574            .unwrap_or_else(|| self.get_effective_reasoning_config())
3575    }
3576
3577    fn get_skill_reflection_config(&self, skill: &SkillDefinition) -> ReflectionConfig {
3578        skill
3579            .reflection
3580            .clone()
3581            .unwrap_or_else(|| self.get_effective_reflection_config())
3582    }
3583
3584    async fn build_disambiguation_context(&self) -> Result<DisambiguationContext> {
3585        let recent_messages: Vec<String> = self
3586            .memory
3587            .get_messages(Some(5))
3588            .await?
3589            .iter()
3590            .rev()
3591            .map(|m| format!("{:?}: {}", m.role, m.content))
3592            .collect();
3593
3594        let current_state = self.current_state().map(|s| s.to_string());
3595
3596        // Include the current state's prompt text so the detector understands
3597        // what kind of input is expected (e.g., "Ask for the order number").
3598        let state_prompt: Option<String> = self
3599            .state_machine
3600            .as_ref()
3601            .and_then(|sm| sm.current_definition())
3602            .and_then(|def| def.prompt.clone());
3603
3604        let available_tools: Vec<String> = self
3605            .get_available_tool_ids()
3606            .await
3607            .unwrap_or_else(|_| self.tools.list_ids());
3608
3609        let available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
3610
3611        let mut user_context = self.build_context_with_overlays();
3612        user_context.remove(DISAMBIGUATION_STATE_GENERATION_KEY);
3613        if let Some(state_generation) = self
3614            .state_machine
3615            .as_ref()
3616            .map(|state_machine| state_machine.generation())
3617        {
3618            user_context.insert(
3619                DISAMBIGUATION_STATE_GENERATION_KEY.to_string(),
3620                serde_json::json!(state_generation),
3621            );
3622        }
3623
3624        // Extract canonical intent labels from current state's transitions
3625        let available_intents: Vec<String> = if let Some(ref sm) = self.state_machine {
3626            sm.current_definition()
3627                .map(|def| {
3628                    def.transitions
3629                        .iter()
3630                        .filter_map(|t| t.intent.clone())
3631                        .collect()
3632                })
3633                .unwrap_or_default()
3634        } else {
3635            Vec::new()
3636        };
3637
3638        Ok(DisambiguationContext::from_agent_state(
3639            recent_messages,
3640            current_state,
3641            state_prompt,
3642            available_tools,
3643            available_skills,
3644            available_intents,
3645            user_context,
3646        ))
3647    }
3648
3649    fn get_available_skills(&self) -> Vec<&SkillDefinition> {
3650        if let Some(ref sm) = self.state_machine
3651            && let Some(state_def) = sm.current_definition()
3652        {
3653            let parent_def = sm.get_parent_definition();
3654            let effective_skills = state_def.get_effective_skills(parent_def.as_ref());
3655            if !effective_skills.is_empty() {
3656                return self
3657                    .skills
3658                    .iter()
3659                    .filter(|s| effective_skills.contains(&&s.id))
3660                    .collect();
3661            }
3662        }
3663        self.skills.iter().collect()
3664    }
3665
3666    async fn build_messages(&self) -> Result<Vec<ChatMessage>> {
3667        self.build_messages_internal(true, None, true).await
3668    }
3669
3670    async fn build_messages_for_draft(&self, user_message: &str) -> Result<Vec<ChatMessage>> {
3671        self.build_messages_internal(false, Some(user_message), true)
3672            .await
3673    }
3674
3675    async fn build_messages_internal(
3676        &self,
3677        fire_persona_hooks: bool,
3678        ephemeral_user_message: Option<&str>,
3679        include_tool_prompt: bool,
3680    ) -> Result<Vec<ChatMessage>> {
3681        let system_prompt = self
3682            .get_effective_system_prompt_with_persona_hooks(fire_persona_hooks, include_tool_prompt)
3683            .await?;
3684        let mut messages = vec![ChatMessage::system(&system_prompt)];
3685
3686        let context = self.memory.get_context().await?;
3687        let history = if let Some(ref budget) = self.memory_token_budget {
3688            context.to_llm_messages_with_allocation(&budget.allocation)
3689        } else {
3690            context.to_llm_messages()
3691        };
3692        messages.extend(history);
3693        if let Some(user_message) = ephemeral_user_message {
3694            messages.push(ChatMessage::user(user_message));
3695        }
3696
3697        let total_tokens = self.estimate_total_tokens(&messages);
3698
3699        if total_tokens > self.max_context_tokens {
3700            debug!(
3701                total = total_tokens,
3702                limit = self.max_context_tokens,
3703                "Context overflow"
3704            );
3705
3706            match &self.recovery_manager.config().llm.on_context_overflow {
3707                ContextOverflowAction::Error => {
3708                    return Err(AgentError::LLM(format!(
3709                        "Context overflow: {} tokens > {} limit",
3710                        total_tokens, self.max_context_tokens
3711                    )));
3712                }
3713                ContextOverflowAction::Truncate { keep_recent } => {
3714                    self.truncate_context(&mut messages, *keep_recent);
3715                }
3716                ContextOverflowAction::Summarize {
3717                    summarizer_llm,
3718                    max_summary_tokens,
3719                    custom_prompt,
3720                    keep_recent,
3721                    filter,
3722                } => {
3723                    self.summarize_context(
3724                        &mut messages,
3725                        summarizer_llm.as_deref(),
3726                        *max_summary_tokens,
3727                        custom_prompt.as_deref(),
3728                        *keep_recent,
3729                        filter.as_ref(),
3730                    )
3731                    .await?;
3732                }
3733            }
3734        }
3735
3736        Ok(messages)
3737    }
3738
3739    async fn main_tool_protocol(
3740        &self,
3741        llm: &dyn LLMProvider,
3742        ephemeral_new_turn: bool,
3743    ) -> Result<MainToolProtocol> {
3744        let mut choice = llm.configured_tool_choice();
3745        if matches!(choice.as_ref(), Some(ToolChoice::None)) {
3746            return Ok(MainToolProtocol {
3747                choice,
3748                tool_ids: Vec::new(),
3749                definitions: Vec::new(),
3750            });
3751        }
3752
3753        let mut tool_ids = self.get_available_tool_ids().await?;
3754        tool_ids.sort();
3755        tool_ids.dedup();
3756        if let Some(ToolChoice::Specific(expected)) = choice.as_ref() {
3757            let canonical = self.tools.canonical_id(expected).ok_or_else(|| {
3758                AgentError::Config(format!(
3759                    "specific tool choice '{expected}' is not registered"
3760                ))
3761            })?;
3762            if canonical != *expected {
3763                return Err(AgentError::Config(format!(
3764                    "specific tool choice must use canonical ID '{canonical}', not '{expected}'"
3765                )));
3766            }
3767            if !tool_ids.iter().any(|tool_id| tool_id == expected) {
3768                return Err(AgentError::Config(format!(
3769                    "specific tool choice '{expected}' is outside the effective tool grant"
3770                )));
3771            }
3772        }
3773        if matches!(
3774            choice.as_ref(),
3775            Some(ToolChoice::Required | ToolChoice::Specific(_))
3776        ) && tool_ids.is_empty()
3777        {
3778            return Err(AgentError::Config(
3779                "required tool choice has no tool inside the effective grant".to_string(),
3780            ));
3781        }
3782        if !ephemeral_new_turn
3783            && let Some(configured_choice) = choice.as_ref()
3784            && matches!(
3785                configured_choice,
3786                ToolChoice::Required | ToolChoice::Specific(_)
3787            )
3788            && self
3789                .tool_choice_satisfied_in_current_turn(configured_choice, &tool_ids)
3790                .await?
3791        {
3792            choice = Some(ToolChoice::Auto);
3793        }
3794        if let Some(ToolChoice::Specific(expected)) = choice.as_ref() {
3795            tool_ids.retain(|tool_id| tool_id == expected);
3796        }
3797
3798        let definitions = tool_ids
3799            .iter()
3800            .map(|tool_id| {
3801                let tool = self.tools.get(tool_id).ok_or_else(|| {
3802                    AgentError::Config(format!(
3803                        "effective tool '{tool_id}' disappeared before provider exposure"
3804                    ))
3805                })?;
3806                Ok(LLMToolDefinition {
3807                    name: tool_id.clone(),
3808                    description: tool.description().to_string(),
3809                    input_schema: tool.input_schema(),
3810                })
3811            })
3812            .collect::<Result<Vec<_>>>()?;
3813
3814        //
3815        // Provider-visible definitions are derived only from the current effective grant. Tool choice never expands registration, scope, or authorization.
3816        //
3817        Ok(MainToolProtocol {
3818            choice,
3819            tool_ids,
3820            definitions,
3821        })
3822    }
3823
3824    async fn tool_choice_satisfied_in_current_turn(
3825        &self,
3826        choice: &ToolChoice,
3827        effective_tool_ids: &[String],
3828    ) -> Result<bool> {
3829        let messages = self.memory.get_messages(None).await?;
3830        let mut saw_tool_result = false;
3831        for message in messages.iter().rev() {
3832            match message.role {
3833                ai_agents_core::Role::Tool | ai_agents_core::Role::Function => {
3834                    saw_tool_result = true;
3835                }
3836                ai_agents_core::Role::Assistant if saw_tool_result => {
3837                    let Some(calls) = self.parse_tool_calls(&message.content) else {
3838                        continue;
3839                    };
3840                    let calls_are_effective = !calls.is_empty()
3841                        && calls.iter().all(|call| {
3842                            self.tools
3843                                .canonical_id(&call.name)
3844                                .is_some_and(|canonical| effective_tool_ids.contains(&canonical))
3845                        });
3846                    return Ok(calls_are_effective
3847                        && match choice {
3848                            ToolChoice::Required => true,
3849                            ToolChoice::Specific(expected) => calls.iter().all(|call| {
3850                                self.tools.canonical_id(&call.name).as_deref()
3851                                    == Some(expected.as_str())
3852                            }),
3853                            _ => false,
3854                        });
3855                }
3856                ai_agents_core::Role::User => return Ok(false),
3857                _ => {}
3858            }
3859        }
3860        Ok(false)
3861    }
3862
3863    fn provider_can_use_native_tools(
3864        &self,
3865        llm: &dyn LLMProvider,
3866        protocol: &MainToolProtocol,
3867    ) -> bool {
3868        let Some(choice) = protocol.choice.as_ref() else {
3869            return false;
3870        };
3871        if matches!(choice, ToolChoice::None) || protocol.definitions.is_empty() {
3872            return false;
3873        }
3874        llm.supports_tool_choice(choice)
3875            && protocol.definitions.iter().all(|definition| {
3876                !definition.name.is_empty()
3877                    && definition.name.len() <= 64
3878                    && definition
3879                        .name
3880                        .bytes()
3881                        .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
3882            })
3883    }
3884
3885    fn prompt_messages_for_tool_protocol(
3886        &self,
3887        messages: &[ChatMessage],
3888        protocol: &MainToolProtocol,
3889        corrective: bool,
3890    ) -> Vec<ChatMessage> {
3891        let mut messages = messages.to_vec();
3892        let Some(choice) = protocol.choice.as_ref() else {
3893            return messages;
3894        };
3895        if matches!(choice, ToolChoice::None) || protocol.tool_ids.is_empty() {
3896            return messages;
3897        }
3898
3899        let mut tool_prompt = self.tools.generate_scoped_prompt_with_mode(
3900            &protocol.tool_ids,
3901            None,
3902            self.parallel_tools.enabled,
3903            self.runtime_config.tool_schema_prompt_mode,
3904        );
3905        match choice {
3906            ToolChoice::Required => tool_prompt.push_str(
3907                "\n\nYou must call at least one listed tool before giving a final answer.",
3908            ),
3909            ToolChoice::Specific(tool_id) => tool_prompt.push_str(&format!(
3910                "\n\nYou must call the '{tool_id}' tool before giving a final answer."
3911            )),
3912            ToolChoice::Auto => {}
3913            ToolChoice::None => return messages,
3914            _ => return messages,
3915        }
3916        if let Some(system) = messages
3917            .iter_mut()
3918            .find(|message| message.role == ai_agents_core::Role::System)
3919        {
3920            system.content.push_str("\n\n");
3921            system.content.push_str(&tool_prompt);
3922        } else {
3923            messages.insert(0, ChatMessage::system(tool_prompt));
3924        }
3925        if corrective {
3926            let instruction = match choice {
3927                ToolChoice::Required => {
3928                    "Your previous response did not call a required tool. Call at least one listed tool now and return only the JSON tool call."
3929                }
3930                ToolChoice::Specific(tool_id) => {
3931                    messages.push(ChatMessage::user(format!(
3932                        "Your previous response did not call the required '{tool_id}' tool. Call it now and return only the JSON tool call."
3933                    )));
3934                    return messages;
3935                }
3936                _ => return messages,
3937            };
3938            messages.push(ChatMessage::user(instruction));
3939        }
3940        messages
3941    }
3942
3943    async fn invoke_main_provider(
3944        &self,
3945        llm: Arc<dyn LLMProvider>,
3946        messages: &[ChatMessage],
3947        protocol: &MainToolProtocol,
3948        corrective: bool,
3949    ) -> std::result::Result<MainProviderResponse, LLMError> {
3950        let use_native = self.provider_can_use_native_tools(llm.as_ref(), protocol);
3951        let response = if use_native {
3952            let request = LLMToolRequest {
3953                tools: protocol.definitions.clone(),
3954                choice: protocol
3955                    .choice
3956                    .clone()
3957                    .expect("native tool requests require an explicit choice"),
3958            };
3959            self.observe_purpose(
3960                ObservationPurpose::MainResponse,
3961                llm.complete_with_tools(messages, None, &request),
3962            )
3963            .await?
3964        } else {
3965            let prompt_messages =
3966                self.prompt_messages_for_tool_protocol(messages, protocol, corrective);
3967            self.observe_purpose(
3968                ObservationPurpose::MainResponse,
3969                llm.complete(&prompt_messages, None),
3970            )
3971            .await?
3972        };
3973        Ok(MainProviderResponse {
3974            response,
3975            used_native_tools: use_native,
3976        })
3977    }
3978
3979    async fn complete_main_attempt_with_recovery(
3980        &self,
3981        llm: Arc<dyn LLMProvider>,
3982        messages: &[ChatMessage],
3983        protocol: &MainToolProtocol,
3984        corrective: bool,
3985    ) -> Result<MainProviderResponse> {
3986        let primary_result = if self.recovery_manager.config().default.max_retries > 0 {
3987            self.recovery_manager
3988                .with_retry("llm_call", None, || {
3989                    let llm = Arc::clone(&llm);
3990                    async move {
3991                        self.invoke_main_provider(llm, messages, protocol, corrective)
3992                            .await
3993                            .map_err(|error| error.classify())
3994                    }
3995                })
3996                .await
3997                .map_err(|error| AgentError::LLM(error.to_string()))
3998        } else {
3999            self.invoke_main_provider(Arc::clone(&llm), messages, protocol, corrective)
4000                .await
4001                .map_err(|error| AgentError::LLM(error.to_string()))
4002        };
4003
4004        match primary_result {
4005            Ok(response) => Ok(response),
4006            Err(primary_error) => match &self.recovery_manager.config().llm.on_failure {
4007                LLMFailureAction::FallbackLlm { fallback_llm } => {
4008                    let fallback = self.llm_registry.get(fallback_llm).map_err(|error| {
4009                        AgentError::Config(format!(
4010                            "Fallback LLM '{fallback_llm}' not found: {error}"
4011                        ))
4012                    })?;
4013                    self.invoke_main_provider(fallback, messages, protocol, corrective)
4014                        .await
4015                        .map_err(|error| AgentError::LLM(error.to_string()))
4016                }
4017                LLMFailureAction::FallbackResponse { message } => {
4018                    if matches!(
4019                        protocol.choice.as_ref(),
4020                        Some(ToolChoice::Required | ToolChoice::Specific(_))
4021                    ) {
4022                        Err(AgentError::LLM(format!(
4023                            "Required tool selection failed and cannot be satisfied by a static fallback response: {primary_error}"
4024                        )))
4025                    } else {
4026                        Ok(MainProviderResponse {
4027                            response: LLMResponse::new(message.clone(), FinishReason::Stop),
4028                            used_native_tools: false,
4029                        })
4030                    }
4031                }
4032                LLMFailureAction::Error => Err(primary_error),
4033            },
4034        }
4035    }
4036
4037    fn normalize_main_provider_response(
4038        &self,
4039        mut response: LLMResponse,
4040        protocol: &MainToolProtocol,
4041    ) -> Result<(LLMResponse, bool)> {
4042        let native_calls = response
4043            .tool_calls()
4044            .map_err(|error| AgentError::LLM(error.to_string()))?;
4045        let calls = match native_calls {
4046            Some(calls) => {
4047                let markers = calls
4048                    .iter()
4049                    .map(|call| {
4050                        serde_json::json!({
4051                            "_ai_agents_native_tool_call": true,
4052                            "id": call.id,
4053                            "tool": call.name,
4054                            "arguments": call.arguments,
4055                        })
4056                    })
4057                    .collect::<Vec<_>>();
4058                response.content = if markers.len() == 1 {
4059                    markers[0].to_string()
4060                } else {
4061                    serde_json::Value::Array(markers).to_string()
4062                };
4063                Some(calls)
4064            }
4065            None if !matches!(protocol.choice.as_ref(), Some(ToolChoice::None)) => {
4066                self.parse_tool_calls(response.content.trim())
4067            }
4068            None => None,
4069        };
4070
4071        if protocol.choice.is_some()
4072            && let Some(calls) = calls.as_ref()
4073            && calls.iter().any(|call| {
4074                self.tools
4075                    .canonical_id(&call.name)
4076                    .is_none_or(|canonical| !protocol.tool_ids.contains(&canonical))
4077            })
4078        {
4079            return Err(AgentError::LLM(
4080                "Provider returned a tool call outside the effective grant".to_string(),
4081            ));
4082        }
4083
4084        let compliant = match protocol.choice.as_ref() {
4085            Some(ToolChoice::Required) => calls.as_ref().is_some_and(|calls| !calls.is_empty()),
4086            Some(ToolChoice::Specific(expected)) => calls.as_ref().is_some_and(|calls| {
4087                !calls.is_empty()
4088                    && calls.iter().all(|call| {
4089                        self.tools.canonical_id(&call.name).as_deref() == Some(expected.as_str())
4090                    })
4091            }),
4092            _ => true,
4093        };
4094        Ok((response, compliant))
4095    }
4096
4097    async fn complete_main_llm_with_recovery(
4098        &self,
4099        llm: Arc<dyn LLMProvider>,
4100        messages: &[ChatMessage],
4101        protocol: &MainToolProtocol,
4102    ) -> Result<LLMResponse> {
4103        let first = self
4104            .complete_main_attempt_with_recovery(Arc::clone(&llm), messages, protocol, false)
4105            .await?;
4106        let (response, compliant) =
4107            self.normalize_main_provider_response(first.response, protocol)?;
4108        if compliant {
4109            return Ok(response);
4110        }
4111        if first.used_native_tools {
4112            return Err(AgentError::LLM(
4113                "Provider returned no compliant native call for required tool choice".to_string(),
4114            ));
4115        }
4116
4117        let corrected = self
4118            .complete_main_attempt_with_recovery(llm, messages, protocol, true)
4119            .await?;
4120        let (response, compliant) =
4121            self.normalize_main_provider_response(corrected.response, protocol)?;
4122        if compliant {
4123            return Ok(response);
4124        }
4125        Err(AgentError::LLM(
4126            "Provider returned no compliant tool call after one corrective retry".to_string(),
4127        ))
4128    }
4129
4130    fn is_native_tool_call_content(content: &str) -> bool {
4131        let Ok(value) = serde_json::from_str::<serde_json::Value>(content) else {
4132            return false;
4133        };
4134        match value {
4135            serde_json::Value::Array(values) => {
4136                !values.is_empty()
4137                    && values.iter().all(|value| {
4138                        value
4139                            .get("_ai_agents_native_tool_call")
4140                            .and_then(|marker| marker.as_bool())
4141                            == Some(true)
4142                    })
4143            }
4144            serde_json::Value::Object(map) => {
4145                map.get("_ai_agents_native_tool_call")
4146                    .and_then(|marker| marker.as_bool())
4147                    == Some(true)
4148            }
4149            _ => false,
4150        }
4151    }
4152
4153    fn tool_result_message(
4154        tool_call: &ToolCall,
4155        output: &str,
4156        native_tool_call: bool,
4157    ) -> ChatMessage {
4158        if !native_tool_call {
4159            return ChatMessage::function(&tool_call.name, output);
4160        }
4161        let output = serde_json::from_str::<serde_json::Value>(output)
4162            .unwrap_or_else(|_| serde_json::Value::String(output.to_string()));
4163        ChatMessage::function(
4164            &tool_call.name,
4165            serde_json::json!({
4166                "_ai_agents_native_tool_result": true,
4167                "id": tool_call.id,
4168                "tool": tool_call.name,
4169                "output": output,
4170            })
4171            .to_string(),
4172        )
4173    }
4174
4175    fn parse_main_tool_calls(
4176        &self,
4177        content: &str,
4178        protocol: &MainToolProtocol,
4179    ) -> Option<Vec<ToolCall>> {
4180        if matches!(protocol.choice.as_ref(), Some(ToolChoice::None)) {
4181            None
4182        } else {
4183            self.parse_tool_calls(content)
4184        }
4185    }
4186
4187    fn parse_tool_calls(&self, content: &str) -> Option<Vec<ToolCall>> {
4188        // Try direct JSON parse first
4189        if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(content) {
4190            // Handle JSON array of tool calls (parallel tool calling)
4191            if let Some(arr) = parsed.as_array() {
4192                let calls: Vec<ToolCall> = arr
4193                    .iter()
4194                    .filter_map(|v| self.extract_tool_call_from_value(v))
4195                    .collect();
4196                if !calls.is_empty() {
4197                    return Some(calls);
4198                }
4199            }
4200            // Handle single JSON object
4201            if let Some(tool_call) = self.extract_tool_call_from_value(&parsed) {
4202                return Some(vec![tool_call]);
4203            }
4204        }
4205
4206        // Try to extract JSON from content (handles extra text/braces from LLM)
4207        if let Some(json_str) = self.extract_json_from_content(content)
4208            && let Ok(parsed) = serde_json::from_str::<serde_json::Value>(&json_str)
4209        {
4210            // 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        None
4227    }
4228
4229    fn extract_tool_call_from_value(&self, parsed: &serde_json::Value) -> Option<ToolCall> {
4230        if let Some(tool_name) = parsed.get("tool").and_then(|v| v.as_str()) {
4231            let arguments = parsed
4232                .get("arguments")
4233                .cloned()
4234                .unwrap_or(serde_json::json!({}));
4235            return Some(ToolCall {
4236                id: parsed
4237                    .get("id")
4238                    .and_then(|value| value.as_str())
4239                    .filter(|id| !id.is_empty())
4240                    .map(str::to_string)
4241                    .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
4242                name: tool_name.to_string(),
4243                arguments,
4244            });
4245        }
4246        None
4247    }
4248
4249    // Lite models could generate unmatched braces: this function handles such cases
4250    fn extract_json_from_content(&self, content: &str) -> Option<String> {
4251        // Try array first (for parallel tool calls), then single object
4252        if let Some(result) = self.extract_json_array_from_content(content) {
4253            return Some(result);
4254        }
4255        self.extract_json_object_from_content(content)
4256    }
4257
4258    /// Extract a JSON array `[...]` containing tool calls from mixed content.
4259    fn extract_json_array_from_content(&self, content: &str) -> Option<String> {
4260        let start = content.find('[')?;
4261        let content_from_start = &content[start..];
4262
4263        let mut depth = 0;
4264        let mut end = 0;
4265        for (i, ch) in content_from_start.char_indices() {
4266            match ch {
4267                '[' => depth += 1,
4268                ']' => {
4269                    depth -= 1;
4270                    if depth == 0 {
4271                        end = i + 1;
4272                        break;
4273                    }
4274                }
4275                _ => {}
4276            }
4277        }
4278
4279        if end > 0 {
4280            let json_str = &content_from_start[..end];
4281            // Verify it looks like an array of tool calls
4282            if json_str.contains("\"tool\"") {
4283                return Some(json_str.to_string());
4284            }
4285        }
4286
4287        None
4288    }
4289
4290    /// Extract a JSON object `{...}` containing a tool call from mixed content.
4291    fn extract_json_object_from_content(&self, content: &str) -> Option<String> {
4292        let start = content.find('{')?;
4293        let content_from_start = &content[start..];
4294
4295        // Count braces to find the matching closing brace
4296        let mut depth = 0;
4297        let mut end = 0;
4298        for (i, ch) in content_from_start.char_indices() {
4299            match ch {
4300                '{' => depth += 1,
4301                '}' => {
4302                    depth -= 1;
4303                    if depth == 0 {
4304                        end = i + 1;
4305                        break;
4306                    }
4307                }
4308                _ => {}
4309            }
4310        }
4311
4312        if end > 0 {
4313            let json_str = &content_from_start[..end];
4314            // Verify it looks like a tool call
4315            if json_str.contains("\"tool\"") {
4316                return Some(json_str.to_string());
4317            }
4318        }
4319
4320        None
4321    }
4322
4323    /// Builds a structured record from executor state.
4324    ///
4325    /// The explicit fields preserve one audit boundary for execution, policy, approval, timing, and generation evidence.
4326    #[allow(clippy::too_many_arguments)]
4327    fn record_from_parts(
4328        &self,
4329        request: &ToolExecutionRequest,
4330        canonical_id: String,
4331        executed_arguments: Value,
4332        started_at: chrono::DateTime<chrono::Utc>,
4333        start: Instant,
4334        executed: bool,
4335        success: bool,
4336        output: String,
4337        metadata: HashMap<String, Value>,
4338        policy: ToolPolicyDecisionRecord,
4339        approval: Option<ToolApprovalRecord>,
4340        timed_out: bool,
4341        output_truncated: bool,
4342    ) -> ToolExecutionRecord {
4343        let versions = ToolDecisionVersions {
4344            policy: self.active_tool_security().policy_version(),
4345            registry: self.tools.version(),
4346            runtime_control: self.runtime_control.version.load(Ordering::SeqCst),
4347            state: self
4348                .state_machine
4349                .as_ref()
4350                .map(|state_machine| state_machine.generation()),
4351        };
4352        self.record_from_parts_at(
4353            request,
4354            canonical_id,
4355            executed_arguments,
4356            started_at,
4357            start,
4358            executed,
4359            success,
4360            output,
4361            metadata,
4362            policy,
4363            approval,
4364            timed_out,
4365            output_truncated,
4366            versions,
4367        )
4368    }
4369
4370    /// Builds evidence with the exact generations used by the corresponding decision.
4371    #[allow(clippy::too_many_arguments)]
4372    fn record_from_parts_at(
4373        &self,
4374        request: &ToolExecutionRequest,
4375        canonical_id: String,
4376        executed_arguments: Value,
4377        started_at: chrono::DateTime<chrono::Utc>,
4378        start: Instant,
4379        executed: bool,
4380        success: bool,
4381        output: String,
4382        metadata: HashMap<String, Value>,
4383        policy: ToolPolicyDecisionRecord,
4384        approval: Option<ToolApprovalRecord>,
4385        timed_out: bool,
4386        output_truncated: bool,
4387        versions: ToolDecisionVersions,
4388    ) -> ToolExecutionRecord {
4389        ToolExecutionRecord {
4390            call_id: request.call_id.clone(),
4391            requested_name: request.requested_name.clone(),
4392            canonical_id,
4393            source: request.source.clone(),
4394            arguments: request.arguments.clone(),
4395            executed_arguments,
4396            policy_version: versions.policy,
4397            registry_version: versions.registry,
4398            runtime_config_version: versions.runtime_control,
4399            executed,
4400            success,
4401            output,
4402            metadata,
4403            policy,
4404            approval,
4405            started_at,
4406            duration_ms: start.elapsed().as_millis() as u64,
4407            timed_out,
4408            cancelled: false,
4409            cancellation_reason: None,
4410            output_truncated,
4411        }
4412    }
4413
4414    /// Sends a finalized tool record to hooks, history, and error handling.
4415    async fn finish_tool_record(&self, record: &ToolExecutionRecord) {
4416        let result = ToolResult {
4417            success: record.success,
4418            output: record.model_output_string(),
4419            metadata: if record.metadata.is_empty() {
4420                None
4421            } else {
4422                Some(record.metadata.clone())
4423            },
4424        };
4425        self.hooks
4426            .on_tool_complete(&record.canonical_id, &result, record.duration_ms)
4427            .await;
4428        self.hooks.on_tool_execution_record(record).await;
4429        self.record_tool_call(&record.canonical_id, record.model_output_value());
4430        if !record.success {
4431            self.hooks
4432                .on_error(&AgentError::Tool(record.output.clone()))
4433                .await;
4434        }
4435    }
4436
4437    /// Releases resource locks before hooks run because hooks may invoke another tool.
4438    async fn finish_tool_record_after_resource_guards(
4439        &self,
4440        resource_guards: ToolResourceGuards,
4441        record: &ToolExecutionRecord,
4442    ) {
4443        drop(resource_guards);
4444        self.finish_tool_record(record).await;
4445    }
4446
4447    /// Invokes a resolved tool once with timeout, cancellation, and actor context.
4448    async fn execute_resolved_tool_once(
4449        &self,
4450        tool: Arc<dyn ai_agents_core::Tool>,
4451        args: Value,
4452        ctx: ToolExecutionContext,
4453        timeout_ms: u64,
4454    ) -> Result<(ToolResult, bool, bool, bool)> {
4455        if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4456            return Ok((
4457                ToolResult::error("Tool execution cancelled by runtime control"),
4458                false,
4459                true,
4460                false,
4461            ));
4462        }
4463        //
4464        // Mark invocation inside the future so cancellation before the first poll remains executed false.
4465        //
4466        let invoked = Arc::new(AtomicBool::new(false));
4467        let invoked_by_future = Arc::clone(&invoked);
4468        let actor_context = current_turn_actor_context();
4469        let future = async move {
4470            invoked_by_future.store(true, Ordering::SeqCst);
4471            if let Some(actor_context) = actor_context {
4472                scope_actor_context(actor_context, tool.execute(args, ctx)).await
4473            } else {
4474                tool.execute(args, ctx).await
4475            }
4476        };
4477        tokio::pin!(future);
4478        let timeout = tokio::time::sleep(std::time::Duration::from_millis(timeout_ms));
4479        tokio::pin!(timeout);
4480        let mut cancel_tick = tokio::time::interval(std::time::Duration::from_millis(50));
4481
4482        loop {
4483            tokio::select! {
4484                result = &mut future => return Ok((result, false, false, true)),
4485                _ = &mut timeout => {
4486                    return Ok((
4487                        ToolResult::error("Tool execution timed out"),
4488                        true,
4489                        false,
4490                        invoked.load(Ordering::SeqCst),
4491                    ));
4492                }
4493                _ = cancel_tick.tick() => {
4494                    if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4495                        return Ok((
4496                            ToolResult::error("Tool execution cancelled by runtime control"),
4497                            false,
4498                            true,
4499                            invoked.load(Ordering::SeqCst),
4500                        ));
4501                    }
4502                }
4503            }
4504        }
4505    }
4506
4507    /// Truncates model-facing tool output on a character boundary.
4508    fn truncate_tool_output(output: String, max_chars: Option<usize>) -> (String, bool) {
4509        let Some(max_chars) = max_chars else {
4510            return (output, false);
4511        };
4512        let mut chars = output.chars();
4513        let truncated: String = chars.by_ref().take(max_chars).collect();
4514        if chars.next().is_some() {
4515            (truncated, true)
4516        } else {
4517            (output, false)
4518        }
4519    }
4520
4521    /// Acquires all declared resource locks in stable key order and supports emergency cancellation while waiting.
4522    async fn acquire_tool_resource_locks(&self, keys: &[String]) -> Option<ToolResourceGuards> {
4523        let locks = {
4524            let mut table = self.resource_locks.write();
4525            table.retain(|_, lock| lock.strong_count() > 0);
4526            keys.iter()
4527                .map(|key| {
4528                    if let Some(lock) = table.get(key).and_then(Weak::upgrade) {
4529                        lock
4530                    } else {
4531                        let lock = Arc::new(tokio::sync::Mutex::new(()));
4532                        table.insert(key.clone(), Arc::downgrade(&lock));
4533                        lock
4534                    }
4535                })
4536                .collect::<Vec<_>>()
4537        };
4538        let mut resource_guards = ToolResourceGuards {
4539            guards: Vec::with_capacity(locks.len()),
4540            locks: Arc::clone(&self.resource_locks),
4541        };
4542        let mut locks = locks.into_iter();
4543        while let Some(lock) = locks.next() {
4544            let mut lock = Box::pin(lock.lock_owned());
4545            loop {
4546                tokio::select! {
4547                    guard = &mut lock => {
4548                        resource_guards.guards.push(guard);
4549                        break;
4550                    }
4551                    _ = tokio::time::sleep(std::time::Duration::from_millis(10)) => {
4552                        if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4553                            drop(lock);
4554                            drop(locks);
4555                            drop(resource_guards);
4556                            return None;
4557                        }
4558                    }
4559                }
4560            }
4561        }
4562        Some(resource_guards)
4563    }
4564
4565    /// Executes a resolved tool with retry policy from recovery settings.
4566    async fn run_tool_with_retries(
4567        &self,
4568        canonical_id: &str,
4569        tool: Arc<dyn ai_agents_core::Tool>,
4570        args: Value,
4571        ctx: ToolExecutionContext,
4572        timeout_ms: u64,
4573        max_retries: u32,
4574    ) -> Result<(ToolResult, bool, bool, bool)> {
4575        let max_retries = if ctx.classification.safely_retryable {
4576            max_retries
4577        } else {
4578            0
4579        };
4580        let mut attempts = 0;
4581        let mut invoked = false;
4582        loop {
4583            let (result, timed_out, cancelled, attempt_invoked) = self
4584                .execute_resolved_tool_once(tool.clone(), args.clone(), ctx.clone(), timeout_ms)
4585                .await?;
4586            invoked |= attempt_invoked;
4587            if result.success || timed_out || cancelled || attempts >= max_retries {
4588                return Ok((result, timed_out, cancelled, invoked));
4589            }
4590            attempts += 1;
4591            warn!(tool = %canonical_id, attempt = attempts, error = %result.output, "Retrying failed tool call");
4592        }
4593    }
4594
4595    /// Executes a tool request through scope, policy, HITL, timeout, recovery, and evidence recording.
4596    fn execute_tool_record(
4597        &self,
4598        request: ToolExecutionRequest,
4599    ) -> Pin<Box<dyn Future<Output = Result<ToolExecutionRecord>> + Send + '_>> {
4600        Box::pin(self.execute_tool_record_inner(request))
4601    }
4602
4603    async fn execute_tool_record_inner(
4604        &self,
4605        request: ToolExecutionRequest,
4606    ) -> Result<ToolExecutionRecord> {
4607        let started_at = chrono::Utc::now();
4608        let start = Instant::now();
4609        info!(tool = %request.requested_name, args = %request.arguments, "Executing tool");
4610
4611        if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4612            let record = self.record_from_parts(
4613                &request,
4614                request.requested_name.clone(),
4615                request.arguments.clone(),
4616                started_at,
4617                start,
4618                false,
4619                false,
4620                "Tool execution is disabled by runtime control".to_string(),
4621                HashMap::new(),
4622                ToolPolicyDecisionRecord::deny("runtime emergency deny is enabled"),
4623                None,
4624                false,
4625                false,
4626            );
4627            self.finish_tool_record(&record).await;
4628            return Ok(record);
4629        }
4630
4631        let Some(resolved) = self.tools.resolve(&request.requested_name) else {
4632            let record = self.record_from_parts(
4633                &request,
4634                request.requested_name.clone(),
4635                request.arguments.clone(),
4636                started_at,
4637                start,
4638                false,
4639                false,
4640                format!("Tool '{}' is unavailable", request.requested_name),
4641                HashMap::new(),
4642                ToolPolicyDecisionRecord::unavailable(format!(
4643                    "Tool '{}' is not registered",
4644                    request.requested_name
4645                )),
4646                None,
4647                false,
4648                false,
4649            );
4650            self.finish_tool_record(&record).await;
4651            return Ok(record);
4652        };
4653
4654        let canonical_id = resolved.identity.canonical_id.clone();
4655
4656        let initial_scope_snapshot = self.get_available_tool_ids_snapshot().await?;
4657        if !initial_scope_snapshot
4658            .tool_ids
4659            .iter()
4660            .any(|id| id == &canonical_id)
4661        {
4662            let record = self.record_from_parts(
4663                &request,
4664                canonical_id.clone(),
4665                request.arguments.clone(),
4666                started_at,
4667                start,
4668                false,
4669                false,
4670                format!(
4671                    "Tool '{}' is not available in the current scope",
4672                    canonical_id
4673                ),
4674                HashMap::new(),
4675                ToolPolicyDecisionRecord::deny(format!(
4676                    "Tool '{}' is not granted by the current top-level and state tool scope",
4677                    canonical_id
4678                )),
4679                None,
4680                false,
4681                false,
4682            );
4683            self.finish_tool_record(&record).await;
4684            return Ok(record);
4685        }
4686
4687        let approval_control_snapshot = self.runtime_safety_snapshot();
4688        let security_engine = approval_control_snapshot.tool_security.clone();
4689        let bindings = resolved.tool.policy_bindings();
4690        let mut executed_arguments = security_engine.prepare_tool_arguments_with_bindings(
4691            &canonical_id,
4692            &request.arguments,
4693            &bindings,
4694        );
4695        self.hooks
4696            .on_tool_start(&canonical_id, &executed_arguments)
4697            .await;
4698
4699        let mut metadata = HashMap::new();
4700        let safety = resolved.tool.safety_metadata();
4701        let classification = resolved.tool.classify_call(&executed_arguments);
4702        let limits = security_engine.effective_limits(&canonical_id, &safety, &classification);
4703        metadata.insert(
4704            "classification".to_string(),
4705            serde_json::to_value(&classification).unwrap_or(Value::Null),
4706        );
4707        metadata.insert(
4708            "effective_limits".to_string(),
4709            serde_json::to_value(&limits).unwrap_or(Value::Null),
4710        );
4711        let policy_snapshot = security_engine.policy_snapshot(&canonical_id);
4712        if !policy_snapshot.is_null() {
4713            metadata.insert("policy_snapshot".to_string(), policy_snapshot.clone());
4714        }
4715
4716        let mut approval_record = Some(ToolApprovalRecord {
4717            status: ToolApprovalStatus::NotRequired,
4718            reason: None,
4719            modified_arguments: None,
4720        });
4721
4722        let mut security_result = security_engine
4723            .validate_tool_execution_with_bindings(&canonical_id, &executed_arguments, &bindings)
4724            .await?;
4725        match &security_result {
4726            SecurityCheckResult::Allow => {}
4727            SecurityCheckResult::Warn { message } => {
4728                warn!(tool = %canonical_id, message = %message, "Tool security warning");
4729            }
4730            SecurityCheckResult::Block { reason } => {
4731                let record = self.record_from_parts(
4732                    &request,
4733                    canonical_id,
4734                    executed_arguments,
4735                    started_at,
4736                    start,
4737                    false,
4738                    false,
4739                    format!("Denied: {}", reason),
4740                    metadata,
4741                    ToolPolicyDecisionRecord::deny(reason.clone()),
4742                    approval_record,
4743                    false,
4744                    false,
4745                );
4746                self.finish_tool_record(&record).await;
4747                return Ok(record);
4748            }
4749            SecurityCheckResult::Unavailable { reason } => {
4750                let record = self.record_from_parts(
4751                    &request,
4752                    canonical_id,
4753                    executed_arguments,
4754                    started_at,
4755                    start,
4756                    false,
4757                    false,
4758                    format!("Unavailable: {}", reason),
4759                    metadata,
4760                    ToolPolicyDecisionRecord::unavailable(reason.clone()),
4761                    approval_record,
4762                    false,
4763                    false,
4764                );
4765                self.finish_tool_record(&record).await;
4766                return Ok(record);
4767            }
4768            SecurityCheckResult::RequireConfirmation { message } => {
4769                if self.hitl_engine.is_none() {
4770                    approval_record = Some(ToolApprovalRecord {
4771                        status: ToolApprovalStatus::Unavailable,
4772                        reason: Some("No HITL engine configured".to_string()),
4773                        modified_arguments: None,
4774                    });
4775                    let record = self.record_from_parts(
4776                        &request,
4777                        canonical_id,
4778                        executed_arguments,
4779                        started_at,
4780                        start,
4781                        false,
4782                        false,
4783                        format!("Approval unavailable: {}", message),
4784                        metadata,
4785                        ToolPolicyDecisionRecord::approval(message.clone()),
4786                        approval_record,
4787                        false,
4788                        false,
4789                    );
4790                    self.finish_tool_record(&record).await;
4791                    return Ok(record);
4792                }
4793
4794                let check_result = HITLCheckResult::required(
4795                    ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
4796                    HashMap::new(),
4797                    message.clone(),
4798                    None,
4799                );
4800                match self.request_hitl_approval(check_result).await? {
4801                    ApprovalResult::Approved => {
4802                        merge_approved_record(&mut approval_record);
4803                    }
4804                    ApprovalResult::Modified { changes } => {
4805                        if let Some(obj) = executed_arguments.as_object_mut() {
4806                            for (key, value) in changes {
4807                                obj.insert(key, value);
4808                            }
4809                        }
4810                        security_result = security_engine
4811                            .validate_tool_execution_with_bindings(
4812                                &canonical_id,
4813                                &executed_arguments,
4814                                &bindings,
4815                            )
4816                            .await?;
4817                        if !matches!(
4818                            security_result,
4819                            SecurityCheckResult::Allow
4820                                | SecurityCheckResult::Warn { .. }
4821                                | SecurityCheckResult::RequireConfirmation { .. }
4822                        ) {
4823                            let reason = security_result
4824                                .reason()
4825                                .unwrap_or("modified arguments failed policy")
4826                                .to_string();
4827                            let record = self.record_from_parts(
4828                                &request,
4829                                canonical_id,
4830                                executed_arguments.clone(),
4831                                started_at,
4832                                start,
4833                                false,
4834                                false,
4835                                reason.clone(),
4836                                metadata,
4837                                ToolPolicyDecisionRecord::deny(reason),
4838                                Some(ToolApprovalRecord {
4839                                    status: ToolApprovalStatus::Modified,
4840                                    reason: None,
4841                                    modified_arguments: Some(executed_arguments),
4842                                }),
4843                                false,
4844                                false,
4845                            );
4846                            self.finish_tool_record(&record).await;
4847                            return Ok(record);
4848                        }
4849                        approval_record = Some(ToolApprovalRecord {
4850                            status: ToolApprovalStatus::Modified,
4851                            reason: None,
4852                            modified_arguments: Some(executed_arguments.clone()),
4853                        });
4854                    }
4855                    ApprovalResult::Rejected { reason } => {
4856                        let reason = reason.unwrap_or_else(|| "rejected".to_string());
4857                        approval_record = Some(ToolApprovalRecord {
4858                            status: ToolApprovalStatus::Rejected,
4859                            reason: Some(reason.clone()),
4860                            modified_arguments: None,
4861                        });
4862                        let record = self.record_from_parts(
4863                            &request,
4864                            canonical_id,
4865                            executed_arguments,
4866                            started_at,
4867                            start,
4868                            false,
4869                            false,
4870                            format!("Approval rejected: {}", reason),
4871                            metadata,
4872                            ToolPolicyDecisionRecord::approval(reason),
4873                            approval_record,
4874                            false,
4875                            false,
4876                        );
4877                        self.finish_tool_record(&record).await;
4878                        return Ok(record);
4879                    }
4880                    ApprovalResult::Timeout => {
4881                        approval_record = Some(ToolApprovalRecord {
4882                            status: ToolApprovalStatus::Timeout,
4883                            reason: Some("approval timeout".to_string()),
4884                            modified_arguments: None,
4885                        });
4886                        let record = self.record_from_parts(
4887                            &request,
4888                            canonical_id,
4889                            executed_arguments,
4890                            started_at,
4891                            start,
4892                            false,
4893                            false,
4894                            "Approval timed out".to_string(),
4895                            metadata,
4896                            ToolPolicyDecisionRecord::approval("approval timeout"),
4897                            approval_record,
4898                            false,
4899                            false,
4900                        );
4901                        self.finish_tool_record(&record).await;
4902                        return Ok(record);
4903                    }
4904                }
4905            }
4906        }
4907
4908        if canonical_id == "command" && !self.tools.command_runner_available() {
4909            let record = self.record_from_parts(
4910                &request,
4911                canonical_id.clone(),
4912                executed_arguments.clone(),
4913                started_at,
4914                start,
4915                false,
4916                false,
4917                "Command runner is unavailable".to_string(),
4918                metadata,
4919                ToolPolicyDecisionRecord::unavailable("command runner is unavailable"),
4920                Some(ToolApprovalRecord {
4921                    status: ToolApprovalStatus::Unavailable,
4922                    reason: Some("command runner is unavailable".to_string()),
4923                    modified_arguments: None,
4924                }),
4925                false,
4926                false,
4927            );
4928            self.finish_tool_record(&record).await;
4929            return Ok(record);
4930        }
4931
4932        if approval_record
4933            .as_ref()
4934            .is_some_and(|record| matches!(record.status, ToolApprovalStatus::NotRequired))
4935            && let Some(message) =
4936                security_engine.classification_approval_message(&canonical_id, &classification)
4937        {
4938            if self.hitl_engine.is_none() {
4939                approval_record = Some(ToolApprovalRecord {
4940                    status: ToolApprovalStatus::Unavailable,
4941                    reason: Some("No HITL engine configured".to_string()),
4942                    modified_arguments: None,
4943                });
4944                let record = self.record_from_parts(
4945                    &request,
4946                    canonical_id,
4947                    executed_arguments,
4948                    started_at,
4949                    start,
4950                    false,
4951                    false,
4952                    format!("Approval unavailable: {}", message),
4953                    metadata,
4954                    ToolPolicyDecisionRecord::approval(message),
4955                    approval_record,
4956                    false,
4957                    false,
4958                );
4959                self.finish_tool_record(&record).await;
4960                return Ok(record);
4961            }
4962            let check_result = HITLCheckResult::required(
4963                ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
4964                HashMap::new(),
4965                message.clone(),
4966                None,
4967            );
4968            match self.request_hitl_approval(check_result).await? {
4969                ApprovalResult::Approved => {
4970                    merge_approved_record(&mut approval_record);
4971                }
4972                ApprovalResult::Modified { changes } => {
4973                    if let Some(obj) = executed_arguments.as_object_mut() {
4974                        for (key, value) in changes {
4975                            obj.insert(key, value);
4976                        }
4977                    }
4978                    let modified_security = security_engine
4979                        .validate_tool_execution_with_bindings(
4980                            &canonical_id,
4981                            &executed_arguments,
4982                            &bindings,
4983                        )
4984                        .await?;
4985                    if !matches!(
4986                        modified_security,
4987                        SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
4988                    ) {
4989                        let reason = modified_security
4990                            .reason()
4991                            .unwrap_or("modified arguments failed policy")
4992                            .to_string();
4993                        let record = self.record_from_parts(
4994                            &request,
4995                            canonical_id,
4996                            executed_arguments.clone(),
4997                            started_at,
4998                            start,
4999                            false,
5000                            false,
5001                            reason.clone(),
5002                            metadata,
5003                            ToolPolicyDecisionRecord::deny(reason),
5004                            Some(ToolApprovalRecord {
5005                                status: ToolApprovalStatus::Modified,
5006                                reason: None,
5007                                modified_arguments: Some(executed_arguments),
5008                            }),
5009                            false,
5010                            false,
5011                        );
5012                        self.finish_tool_record(&record).await;
5013                        return Ok(record);
5014                    }
5015                    approval_record = Some(ToolApprovalRecord {
5016                        status: ToolApprovalStatus::Modified,
5017                        reason: None,
5018                        modified_arguments: Some(executed_arguments.clone()),
5019                    });
5020                }
5021                ApprovalResult::Rejected { reason } => {
5022                    let reason = reason.unwrap_or_else(|| "rejected".to_string());
5023                    let record = self.record_from_parts(
5024                        &request,
5025                        canonical_id,
5026                        executed_arguments,
5027                        started_at,
5028                        start,
5029                        false,
5030                        false,
5031                        format!("Approval rejected: {}", reason),
5032                        metadata,
5033                        ToolPolicyDecisionRecord::approval(reason.clone()),
5034                        Some(ToolApprovalRecord {
5035                            status: ToolApprovalStatus::Rejected,
5036                            reason: Some(reason),
5037                            modified_arguments: None,
5038                        }),
5039                        false,
5040                        false,
5041                    );
5042                    self.finish_tool_record(&record).await;
5043                    return Ok(record);
5044                }
5045                ApprovalResult::Timeout => {
5046                    let record = self.record_from_parts(
5047                        &request,
5048                        canonical_id,
5049                        executed_arguments,
5050                        started_at,
5051                        start,
5052                        false,
5053                        false,
5054                        "Approval timed out".to_string(),
5055                        metadata,
5056                        ToolPolicyDecisionRecord::approval("approval timeout"),
5057                        Some(ToolApprovalRecord {
5058                            status: ToolApprovalStatus::Timeout,
5059                            reason: Some("approval timeout".to_string()),
5060                            modified_arguments: None,
5061                        }),
5062                        false,
5063                        false,
5064                    );
5065                    self.finish_tool_record(&record).await;
5066                    return Ok(record);
5067                }
5068            }
5069        }
5070
5071        if canonical_id == "diagnostics" && !self.tools.diagnostics_available() {
5072            let record = self.record_from_parts(
5073                &request,
5074                canonical_id.clone(),
5075                executed_arguments.clone(),
5076                started_at,
5077                start,
5078                false,
5079                false,
5080                "Diagnostics provider is unavailable".to_string(),
5081                metadata,
5082                ToolPolicyDecisionRecord::unavailable("diagnostics provider is unavailable"),
5083                Some(ToolApprovalRecord {
5084                    status: ToolApprovalStatus::Unavailable,
5085                    reason: Some("diagnostics provider is unavailable".to_string()),
5086                    modified_arguments: None,
5087                }),
5088                false,
5089                false,
5090            );
5091            self.finish_tool_record(&record).await;
5092            return Ok(record);
5093        }
5094
5095        if canonical_id == "web_search" && !self.tools.web_search_available() {
5096            let record = self.record_from_parts(
5097                &request,
5098                canonical_id.clone(),
5099                executed_arguments.clone(),
5100                started_at,
5101                start,
5102                false,
5103                false,
5104                "Web search provider is unavailable".to_string(),
5105                metadata,
5106                ToolPolicyDecisionRecord::unavailable("web search provider is unavailable"),
5107                Some(ToolApprovalRecord {
5108                    status: ToolApprovalStatus::Unavailable,
5109                    reason: Some("web search provider is unavailable".to_string()),
5110                    modified_arguments: None,
5111                }),
5112                false,
5113                false,
5114            );
5115            self.finish_tool_record(&record).await;
5116            return Ok(record);
5117        }
5118
5119        let hitl_lang_ctx = self.build_hitl_language_context();
5120        if let Some(ref hitl_engine) = self.hitl_engine {
5121            let check_result = self
5122                .observe_purpose(
5123                    ObservationPurpose::HitlLocalization,
5124                    hitl_engine.check_tool_with_localization(
5125                        &canonical_id,
5126                        &executed_arguments,
5127                        &hitl_lang_ctx,
5128                        self.approval_handler.as_ref(),
5129                        Some(&self.llm_registry),
5130                    ),
5131                )
5132                .await?;
5133            if check_result.is_required() {
5134                match self.request_hitl_approval(check_result).await? {
5135                    ApprovalResult::Approved => {
5136                        merge_approved_record(&mut approval_record);
5137                    }
5138                    ApprovalResult::Modified { changes } => {
5139                        if let Some(obj) = executed_arguments.as_object_mut() {
5140                            for (key, value) in changes {
5141                                obj.insert(key, value);
5142                            }
5143                        }
5144                        let modified_security = security_engine
5145                            .validate_tool_execution_with_bindings(
5146                                &canonical_id,
5147                                &executed_arguments,
5148                                &bindings,
5149                            )
5150                            .await?;
5151                        if !matches!(
5152                            modified_security,
5153                            SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5154                        ) {
5155                            let reason = modified_security
5156                                .reason()
5157                                .unwrap_or("modified arguments failed policy")
5158                                .to_string();
5159                            let record = self.record_from_parts(
5160                                &request,
5161                                canonical_id,
5162                                executed_arguments.clone(),
5163                                started_at,
5164                                start,
5165                                false,
5166                                false,
5167                                reason.clone(),
5168                                metadata,
5169                                ToolPolicyDecisionRecord::deny(reason),
5170                                Some(ToolApprovalRecord {
5171                                    status: ToolApprovalStatus::Modified,
5172                                    reason: None,
5173                                    modified_arguments: Some(executed_arguments),
5174                                }),
5175                                false,
5176                                false,
5177                            );
5178                            self.finish_tool_record(&record).await;
5179                            return Ok(record);
5180                        }
5181                        approval_record = Some(ToolApprovalRecord {
5182                            status: ToolApprovalStatus::Modified,
5183                            reason: None,
5184                            modified_arguments: Some(executed_arguments.clone()),
5185                        });
5186                    }
5187                    ApprovalResult::Rejected { reason } => {
5188                        let reason = reason.unwrap_or_else(|| "rejected".to_string());
5189                        let record = self.record_from_parts(
5190                            &request,
5191                            canonical_id,
5192                            executed_arguments,
5193                            started_at,
5194                            start,
5195                            false,
5196                            false,
5197                            format!("Approval rejected: {}", reason),
5198                            metadata,
5199                            ToolPolicyDecisionRecord::approval(reason.clone()),
5200                            Some(ToolApprovalRecord {
5201                                status: ToolApprovalStatus::Rejected,
5202                                reason: Some(reason),
5203                                modified_arguments: None,
5204                            }),
5205                            false,
5206                            false,
5207                        );
5208                        self.finish_tool_record(&record).await;
5209                        return Ok(record);
5210                    }
5211                    ApprovalResult::Timeout => {
5212                        let record = self.record_from_parts(
5213                            &request,
5214                            canonical_id,
5215                            executed_arguments,
5216                            started_at,
5217                            start,
5218                            false,
5219                            false,
5220                            "Approval timed out".to_string(),
5221                            metadata,
5222                            ToolPolicyDecisionRecord::approval("approval timeout"),
5223                            Some(ToolApprovalRecord {
5224                                status: ToolApprovalStatus::Timeout,
5225                                reason: Some("approval timeout".to_string()),
5226                                modified_arguments: None,
5227                            }),
5228                            false,
5229                            false,
5230                        );
5231                        self.finish_tool_record(&record).await;
5232                        return Ok(record);
5233                    }
5234                }
5235            }
5236
5237            let condition_check = self
5238                .observe_purpose(
5239                    ObservationPurpose::HitlLocalization,
5240                    hitl_engine.check_conditions_with_localization(
5241                        &executed_arguments,
5242                        &hitl_lang_ctx,
5243                        self.approval_handler.as_ref(),
5244                        Some(&self.llm_registry),
5245                    ),
5246                )
5247                .await?;
5248            if condition_check.is_required() {
5249                match self.request_hitl_approval(condition_check).await? {
5250                    ApprovalResult::Approved => {
5251                        merge_approved_record(&mut approval_record);
5252                    }
5253                    ApprovalResult::Modified { changes } => {
5254                        if let Some(obj) = executed_arguments.as_object_mut() {
5255                            for (key, value) in changes {
5256                                obj.insert(key, value);
5257                            }
5258                        }
5259                        let modified_security = security_engine
5260                            .validate_tool_execution_with_bindings(
5261                                &canonical_id,
5262                                &executed_arguments,
5263                                &bindings,
5264                            )
5265                            .await?;
5266                        if !matches!(
5267                            modified_security,
5268                            SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5269                        ) {
5270                            let reason = modified_security
5271                                .reason()
5272                                .unwrap_or("modified arguments failed policy")
5273                                .to_string();
5274                            let record = self.record_from_parts(
5275                                &request,
5276                                canonical_id,
5277                                executed_arguments,
5278                                started_at,
5279                                start,
5280                                false,
5281                                false,
5282                                reason.clone(),
5283                                metadata,
5284                                ToolPolicyDecisionRecord::deny(reason),
5285                                approval_record,
5286                                false,
5287                                false,
5288                            );
5289                            self.finish_tool_record(&record).await;
5290                            return Ok(record);
5291                        }
5292                        approval_record = Some(ToolApprovalRecord {
5293                            status: ToolApprovalStatus::Modified,
5294                            reason: None,
5295                            modified_arguments: Some(executed_arguments.clone()),
5296                        });
5297                    }
5298                    ApprovalResult::Rejected { reason } => {
5299                        let reason = reason.unwrap_or_else(|| "rejected".to_string());
5300                        let record = self.record_from_parts(
5301                            &request,
5302                            canonical_id,
5303                            executed_arguments,
5304                            started_at,
5305                            start,
5306                            false,
5307                            false,
5308                            format!("Approval rejected: {}", reason),
5309                            metadata,
5310                            ToolPolicyDecisionRecord::approval(reason.clone()),
5311                            Some(ToolApprovalRecord {
5312                                status: ToolApprovalStatus::Rejected,
5313                                reason: Some(reason),
5314                                modified_arguments: None,
5315                            }),
5316                            false,
5317                            false,
5318                        );
5319                        self.finish_tool_record(&record).await;
5320                        return Ok(record);
5321                    }
5322                    ApprovalResult::Timeout => {
5323                        let record = self.record_from_parts(
5324                            &request,
5325                            canonical_id,
5326                            executed_arguments,
5327                            started_at,
5328                            start,
5329                            false,
5330                            false,
5331                            "Approval timed out".to_string(),
5332                            metadata,
5333                            ToolPolicyDecisionRecord::approval("approval timeout"),
5334                            Some(ToolApprovalRecord {
5335                                status: ToolApprovalStatus::Timeout,
5336                                reason: Some("approval timeout".to_string()),
5337                                modified_arguments: None,
5338                            }),
5339                            false,
5340                            false,
5341                        );
5342                        self.finish_tool_record(&record).await;
5343                        return Ok(record);
5344                    }
5345                }
5346            }
5347        }
5348
5349        //
5350        // Freeze the action reviewed by HITL after every approved argument change and policy cap.
5351        // A later plain approval preserves any earlier Modified evidence.
5352        //
5353        executed_arguments = security_engine.prepare_tool_arguments_with_bindings(
5354            &canonical_id,
5355            &executed_arguments,
5356            &bindings,
5357        );
5358        if let Some(record) = approval_record.as_mut()
5359            && matches!(record.status, ToolApprovalStatus::Modified)
5360        {
5361            record.modified_arguments = Some(executed_arguments.clone());
5362        }
5363        let binding_security_result = security_engine
5364            .validate_tool_execution_with_bindings(&canonical_id, &executed_arguments, &bindings)
5365            .await?;
5366        let approval_confirmation_required = matches!(
5367            binding_security_result,
5368            SecurityCheckResult::RequireConfirmation { .. }
5369        ) || security_engine
5370            .classification_approval_message(
5371                &canonical_id,
5372                &resolved.tool.classify_call(&executed_arguments),
5373            )
5374            .is_some();
5375        let approval_binding = approval_record.as_ref().and_then(|record| {
5376            matches!(
5377                record.status,
5378                ToolApprovalStatus::Approved | ToolApprovalStatus::Modified
5379            )
5380            .then(|| ToolApprovalBinding {
5381                canonical_id: canonical_id.clone(),
5382                arguments: executed_arguments.clone(),
5383                confirmation_required: approval_confirmation_required,
5384                policy_version: security_engine.policy_version(),
5385                runtime_control_version: approval_control_snapshot.version,
5386                state_generation: initial_scope_snapshot.state_generation,
5387                reviewed_tool: Arc::clone(&resolved.tool),
5388            })
5389        });
5390
5391        //
5392        // Re-resolve once after HITL because registry, scope, and policy may change while approval waits.
5393        // The final Arc must still match the implementation that the approver reviewed.
5394        //
5395        let control_snapshot = self.runtime_safety_snapshot();
5396        let resolved = self.tools.resolve(&request.requested_name);
5397        let registry_version = self.tools.version();
5398        let mut versions = ToolDecisionVersions {
5399            policy: control_snapshot.tool_security.policy_version(),
5400            registry: registry_version,
5401            runtime_control: control_snapshot.version,
5402            state: None,
5403        };
5404        metadata.insert(
5405            "runtime_scope_snapshot".to_string(),
5406            serde_json::to_value(&control_snapshot.tool_scope_override).unwrap_or(Value::Null),
5407        );
5408        let resolved = match resolved {
5409            Some(resolved) => resolved,
5410            None => {
5411                let reason = format!(
5412                    "Tool '{}' became unavailable after approval",
5413                    request.requested_name
5414                );
5415                let record = self.record_from_parts_at(
5416                    &request,
5417                    request.requested_name.clone(),
5418                    executed_arguments,
5419                    started_at,
5420                    start,
5421                    false,
5422                    false,
5423                    reason.clone(),
5424                    metadata,
5425                    ToolPolicyDecisionRecord::unavailable(reason),
5426                    approval_record,
5427                    false,
5428                    false,
5429                    versions,
5430                );
5431                self.finish_tool_record(&record).await;
5432                return Ok(record);
5433            }
5434        };
5435
5436        let canonical_id = resolved.identity.canonical_id.clone();
5437        let bindings = resolved.tool.policy_bindings();
5438        let final_arguments = control_snapshot
5439            .tool_security
5440            .prepare_tool_arguments_with_bindings(&canonical_id, &executed_arguments, &bindings);
5441        if let Some(record) = approval_record.as_mut()
5442            && matches!(record.status, ToolApprovalStatus::Modified)
5443        {
5444            record.modified_arguments = Some(final_arguments.clone());
5445        }
5446        let classification = resolved.tool.classify_call(&final_arguments);
5447        let safety = resolved.tool.safety_metadata();
5448        let security_engine = control_snapshot.tool_security;
5449        let limits = security_engine.effective_limits(&canonical_id, &safety, &classification);
5450        let policy_snapshot = security_engine.policy_snapshot(&canonical_id);
5451        let resource_lock_keys =
5452            tool_resource_lock_keys(&canonical_id, &final_arguments, &bindings, &classification);
5453        metadata.insert(
5454            "classification".to_string(),
5455            serde_json::to_value(&classification).unwrap_or(Value::Null),
5456        );
5457        metadata.insert(
5458            "effective_limits".to_string(),
5459            serde_json::to_value(&limits).unwrap_or(Value::Null),
5460        );
5461        metadata.insert(
5462            "resource_lock_keys".to_string(),
5463            serde_json::to_value(&resource_lock_keys).unwrap_or(Value::Null),
5464        );
5465        if policy_snapshot.is_null() {
5466            metadata.remove("policy_snapshot");
5467        } else {
5468            metadata.insert("policy_snapshot".to_string(), policy_snapshot.clone());
5469        }
5470
5471        let final_denial = |canonical_id: String,
5472                            output: String,
5473                            policy: ToolPolicyDecisionRecord,
5474                            metadata: HashMap<String, Value>,
5475                            decision_versions: ToolDecisionVersions| {
5476            self.record_from_parts_at(
5477                &request,
5478                canonical_id,
5479                final_arguments.clone(),
5480                started_at,
5481                start,
5482                false,
5483                false,
5484                output,
5485                metadata,
5486                policy,
5487                approval_record.clone(),
5488                false,
5489                false,
5490                decision_versions,
5491            )
5492        };
5493
5494        if control_snapshot.emergency_deny {
5495            let reason = "Tool execution is disabled by runtime control".to_string();
5496            let record = final_denial(
5497                canonical_id,
5498                reason.clone(),
5499                ToolPolicyDecisionRecord::deny(reason),
5500                metadata,
5501                versions,
5502            );
5503            self.finish_tool_record(&record).await;
5504            return Ok(record);
5505        }
5506
5507        //
5508        // Final scope evaluation uses the same post-HITL runtime snapshot and captures state authority before locks.
5509        // Admission must reject any later state transition, reset, or restore before consuming rate capacity or invoking the tool.
5510        //
5511        let available_snapshot = self
5512            .get_available_tool_ids_snapshot_for_scope(
5513                control_snapshot.tool_scope_override.as_deref(),
5514            )
5515            .await?;
5516        versions.state = available_snapshot.state_generation;
5517        metadata.insert(
5518            "available_tool_ids_snapshot".to_string(),
5519            serde_json::to_value(&available_snapshot.tool_ids).unwrap_or(Value::Null),
5520        );
5521        metadata.insert(
5522            "state_generation_snapshot".to_string(),
5523            serde_json::to_value(available_snapshot.state_generation).unwrap_or(Value::Null),
5524        );
5525        if !available_snapshot
5526            .tool_ids
5527            .iter()
5528            .any(|tool_id| tool_id == &canonical_id)
5529        {
5530            let reason = format!(
5531                "Tool '{}' is not available in the final runtime scope",
5532                canonical_id
5533            );
5534            let record = final_denial(
5535                canonical_id,
5536                reason.clone(),
5537                ToolPolicyDecisionRecord::deny(reason),
5538                metadata,
5539                versions,
5540            );
5541            self.finish_tool_record(&record).await;
5542            return Ok(record);
5543        }
5544
5545        //
5546        // Validation may run more than once, but it must not consume rate capacity.
5547        // Rate capacity is consumed once after resource locks are held.
5548        //
5549        let final_security_result = security_engine
5550            .validate_tool_execution_with_bindings(&canonical_id, &final_arguments, &bindings)
5551            .await?;
5552        match &final_security_result {
5553            SecurityCheckResult::Block { reason } => {
5554                let record = final_denial(
5555                    canonical_id,
5556                    format!("Denied: {}", reason),
5557                    ToolPolicyDecisionRecord::deny(reason.clone()),
5558                    metadata,
5559                    versions,
5560                );
5561                self.finish_tool_record(&record).await;
5562                return Ok(record);
5563            }
5564            SecurityCheckResult::Unavailable { reason } => {
5565                let record = final_denial(
5566                    canonical_id,
5567                    format!("Unavailable: {}", reason),
5568                    ToolPolicyDecisionRecord::unavailable(reason.clone()),
5569                    metadata,
5570                    versions,
5571                );
5572                self.finish_tool_record(&record).await;
5573                return Ok(record);
5574            }
5575            SecurityCheckResult::Warn { message } => {
5576                warn!(tool = %canonical_id, message = %message, "Tool security warning after approval");
5577            }
5578            SecurityCheckResult::Allow | SecurityCheckResult::RequireConfirmation { .. } => {}
5579        }
5580        let final_confirmation_required = matches!(
5581            final_security_result,
5582            SecurityCheckResult::RequireConfirmation { .. }
5583        ) || security_engine
5584            .classification_approval_message(&canonical_id, &classification)
5585            .is_some();
5586        let stale_approval = approval_binding.as_ref().is_some_and(|binding| {
5587            binding.is_stale(
5588                &canonical_id,
5589                &final_arguments,
5590                final_confirmation_required,
5591                versions,
5592                &resolved.tool,
5593            )
5594        });
5595        if stale_approval {
5596            let reason = "Approval became stale before final admission".to_string();
5597            let record = final_denial(
5598                canonical_id,
5599                reason.clone(),
5600                ToolPolicyDecisionRecord::deny(reason),
5601                metadata,
5602                versions,
5603            );
5604            self.finish_tool_record(&record).await;
5605            return Ok(record);
5606        }
5607        if final_confirmation_required && approval_binding.is_none() {
5608            let reason = "Final policy requires fresh approval".to_string();
5609            let record = final_denial(
5610                canonical_id,
5611                reason.clone(),
5612                ToolPolicyDecisionRecord::approval(reason),
5613                metadata,
5614                versions,
5615            );
5616            self.finish_tool_record(&record).await;
5617            return Ok(record);
5618        }
5619
5620        let unavailable_reason = match canonical_id.as_str() {
5621            "command" if !self.tools.command_runner_available() => {
5622                Some("command runner is unavailable")
5623            }
5624            "diagnostics" if !self.tools.diagnostics_available() => {
5625                Some("diagnostics provider is unavailable")
5626            }
5627            "web_search" if !self.tools.web_search_available() => {
5628                Some("web search provider is unavailable")
5629            }
5630            _ => None,
5631        };
5632        if let Some(reason) = unavailable_reason {
5633            let record = final_denial(
5634                canonical_id,
5635                reason.to_string(),
5636                ToolPolicyDecisionRecord::unavailable(reason),
5637                metadata,
5638                versions,
5639            );
5640            self.finish_tool_record(&record).await;
5641            return Ok(record);
5642        }
5643
5644        //
5645        // Hold conflict locks across final generation admission and invocation.
5646        // Completion hooks and fallback execution run only after these guards are released.
5647        //
5648        let Some(resource_guards) = self.acquire_tool_resource_locks(&resource_lock_keys).await
5649        else {
5650            let reason = "Tool execution cancelled while waiting for resource locks".to_string();
5651            let record = final_denial(
5652                canonical_id,
5653                reason.clone(),
5654                ToolPolicyDecisionRecord::deny(reason),
5655                metadata,
5656                versions,
5657            );
5658            self.finish_tool_record(&record).await;
5659            return Ok(record);
5660        };
5661
5662        //
5663        // The lock-held admission is the last authority boundary. Runtime, policy, and state generations must still match before the one atomic rate admission.
5664        // Updates after admission apply to later calls, while emergency cancellation remains live for this call.
5665        //
5666        let admission = self.admit_tool_execution(
5667            versions.runtime_control,
5668            versions.policy,
5669            versions.state,
5670            &canonical_id,
5671        );
5672        if !matches!(admission, SecurityCheckResult::Allow) {
5673            let latest_control = self.runtime_safety_snapshot();
5674            let reason = admission
5675                .reason()
5676                .unwrap_or("tool admission was denied")
5677                .to_string();
5678            let policy = if admission.is_unavailable() {
5679                ToolPolicyDecisionRecord::unavailable(reason.clone())
5680            } else {
5681                ToolPolicyDecisionRecord::deny(reason.clone())
5682            };
5683            let record = self.record_from_parts_at(
5684                &request,
5685                canonical_id,
5686                final_arguments,
5687                started_at,
5688                start,
5689                false,
5690                false,
5691                reason,
5692                metadata,
5693                policy,
5694                approval_record,
5695                false,
5696                false,
5697                ToolDecisionVersions {
5698                    policy: latest_control.tool_security.policy_version(),
5699                    registry: versions.registry,
5700                    runtime_control: latest_control.version,
5701                    state: self
5702                        .state_machine
5703                        .as_ref()
5704                        .map(|state_machine| state_machine.generation()),
5705                },
5706            );
5707            self.finish_tool_record_after_resource_guards(resource_guards, &record)
5708                .await;
5709            return Ok(record);
5710        }
5711        let executed_arguments = final_arguments;
5712
5713        let tool_config = self.recovery_manager.get_tool_config(&canonical_id);
5714        let timeout_ms = limits
5715            .timeout_ms
5716            .unwrap_or_else(|| security_engine.get_tool_timeout(&canonical_id));
5717        let deadline = Some(started_at + chrono::Duration::milliseconds(timeout_ms as i64));
5718        let turn_actor = current_turn_actor_context();
5719        let actor = ToolActorContext {
5720            actor_id: turn_actor
5721                .as_ref()
5722                .and_then(|context| context.effective_actor_id().map(str::to_string))
5723                .or_else(|| self.actor_id()),
5724            origin_actor_id: turn_actor
5725                .as_ref()
5726                .and_then(|context| context.origin_actor_id.clone()),
5727            sender_agent_id: turn_actor
5728                .as_ref()
5729                .and_then(|context| context.sender_agent_id.clone()),
5730        };
5731        let tool_context = ToolExecutionContext {
5732            requested_name: request.requested_name.clone(),
5733            canonical_id: canonical_id.clone(),
5734            display_name: resolved.identity.display_name.clone(),
5735            provider_id: resolved.identity.provider_id.clone(),
5736            registry_version: versions.registry,
5737            policy_version: versions.policy,
5738            runtime_control_version: versions.runtime_control,
5739            call_id: request.call_id.clone(),
5740            source: request.source.clone(),
5741            actor,
5742            cancellation: ToolCancellationToken::new(
5743                Arc::clone(&self.runtime_control.emergency_deny),
5744                Some("runtime control cancellation".to_string()),
5745            ),
5746            started_at,
5747            deadline,
5748            permission: ToolPolicyDecisionRecord::allow(),
5749            approval: approval_record.clone(),
5750            classification: classification.clone(),
5751            safety,
5752            limits: limits.clone(),
5753            policy_snapshot,
5754            custom_config: security_engine.custom_config(&canonical_id),
5755        };
5756        let (mut result, timed_out, cancelled, invoked) = self
5757            .run_tool_with_retries(
5758                &canonical_id,
5759                resolved.tool.clone(),
5760                executed_arguments.clone(),
5761                tool_context,
5762                timeout_ms,
5763                tool_config.max_retries,
5764            )
5765            .await?;
5766
5767        if !result.success {
5768            match &tool_config.on_failure {
5769                ToolFailureAction::Skip => {
5770                    result = ToolResult::ok(format!(
5771                        "{{\"skipped\": true, \"reason\": \"Tool '{}' was skipped after failure\"}}",
5772                        canonical_id
5773                    ));
5774                }
5775                ToolFailureAction::Fallback { fallback_tool } => {
5776                    drop(resource_guards);
5777                    let fallback_request = ToolExecutionRequest::new(
5778                        request.call_id.clone(),
5779                        fallback_tool.clone(),
5780                        executed_arguments,
5781                        ToolCallSource::Fallback {
5782                            original_tool: canonical_id,
5783                        },
5784                    );
5785                    return Box::pin(self.execute_tool_record(fallback_request)).await;
5786                }
5787                ToolFailureAction::ReportError => {}
5788            }
5789        }
5790
5791        let output_cap = limits.max_output_chars;
5792        let (output, output_truncated) =
5793            Self::truncate_tool_output(result.output.clone(), output_cap);
5794        if let Some(result_metadata) = result.metadata {
5795            metadata.extend(result_metadata);
5796        }
5797        let mut record = self.record_from_parts_at(
5798            &request,
5799            canonical_id,
5800            executed_arguments,
5801            started_at,
5802            start,
5803            invoked,
5804            result.success,
5805            output,
5806            metadata,
5807            ToolPolicyDecisionRecord::allow(),
5808            approval_record,
5809            timed_out,
5810            output_truncated,
5811            versions,
5812        );
5813        record.cancelled = cancelled;
5814        if cancelled {
5815            record.cancellation_reason = Some("runtime control cancellation".to_string());
5816        }
5817        self.finish_tool_record_after_resource_guards(resource_guards, &record)
5818            .await;
5819        Ok(record)
5820    }
5821
5822    #[instrument(skip(self, tool_call), fields(tool = %tool_call.name))]
5823    async fn execute_tool_smart(&self, tool_call: &ToolCall) -> Result<String> {
5824        let record = self
5825            .execute_tool_record(ToolExecutionRequest::new(
5826                tool_call.id.clone(),
5827                tool_call.name.clone(),
5828                tool_call.arguments.clone(),
5829                ToolCallSource::Model,
5830            ))
5831            .await?;
5832        if record.success {
5833            Ok(record.model_output_string())
5834        } else if matches!(record.policy.outcome, PermissionOutcome::RequiresApproval) {
5835            Err(AgentError::HITLRejected(record.model_output_string()))
5836        } else {
5837            Err(AgentError::Tool(record.model_output_string()))
5838        }
5839    }
5840
5841    //
5842    // This method is branch-safe because it only asks the router which skill matches.
5843    // Do not add disambiguation, pending-skill writes, or skill execution here.
5844    //
5845    /// Selects a skill without executing it.
5846    async fn select_skill_candidate(&self, input: &str) -> Result<Option<SkillCandidate>> {
5847        let Some(ref router) = self.skill_router else {
5848            return Ok(None);
5849        };
5850        let available_skills = self.get_available_skills();
5851        if available_skills.is_empty() {
5852            return Ok(None);
5853        }
5854        let skill_ids: Vec<&str> = available_skills.iter().map(|s| s.id.as_str()).collect();
5855        let Some(skill_id) = self
5856            .observe_purpose(
5857                ObservationPurpose::SkillRouting,
5858                router.select_skill_filtered(input, &skill_ids),
5859            )
5860            .await?
5861        else {
5862            return Ok(None);
5863        };
5864        let skill = router
5865            .get_skill(&skill_id)
5866            .cloned()
5867            .ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
5868        info!(skill_id = %skill_id, "Skill selected");
5869        Ok(Some(SkillCandidate::new(skill_id, skill)))
5870    }
5871
5872    //
5873    // This is the commit half of skill routing.
5874    // It may mutate pending skill state, run clarification, and execute skill steps.
5875    //
5876    async fn commit_skill_candidate_route_result(
5877        &self,
5878        candidate: SkillCandidate,
5879        input: &str,
5880    ) -> Result<SkillRouteResult> {
5881        let skill_id = candidate.skill_id;
5882        let skill = candidate.skill;
5883        let expected_state_generation = self
5884            .state_machine
5885            .as_ref()
5886            .map(|state_machine| state_machine.generation());
5887        let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
5888        if let Some(ref skill_disambig) = skill.disambiguation
5889            && skill_disambig.enabled.unwrap_or(false)
5890            && let Some(ref disambiguator) = self.disambiguation_manager
5891        {
5892            let context = self.build_disambiguation_context().await?;
5893            let state_override = self
5894                .state_machine
5895                .as_ref()
5896                .and_then(|sm| sm.current_definition())
5897                .and_then(|def| def.disambiguation.clone());
5898
5899            let disambiguation_result = self
5900                .observe_purpose(
5901                    ObservationPurpose::DisambiguationDetection,
5902                    disambiguator.process_input_with_override(
5903                        input,
5904                        &context,
5905                        state_override.as_ref(),
5906                        Some(skill_disambig),
5907                    ),
5908                )
5909                .await?;
5910            let current_state_generation = self
5911                .state_machine
5912                .as_ref()
5913                .map(|state_machine| state_machine.generation());
5914            if current_state_generation != expected_state_generation
5915                || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
5916            {
5917                disambiguator.clear_pending().await;
5918                *self.pending_skill_id.write() = None;
5919                return Err(AgentError::Other(
5920                    "State or reset ownership changed during skill disambiguation".to_string(),
5921                ));
5922            }
5923            match disambiguation_result {
5924                DisambiguationResult::Clear => {
5925                    debug!(skill_id = %skill_id, "Skill disambiguation: clear");
5926                }
5927                DisambiguationResult::NeedsClarification {
5928                    question,
5929                    detection,
5930                } => {
5931                    let admission = self
5932                        .admit_disambiguation_redispatch(
5933                            expected_disambiguation_epoch,
5934                            expected_state_generation,
5935                        )
5936                        .await?;
5937                    let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
5938                    info!(
5939                        skill_id = %skill_id,
5940                        ambiguity_type = ?detection.ambiguity_type,
5941                        confidence = detection.confidence,
5942                        "Skill requires clarification before execution"
5943                    );
5944                    *self.pending_skill_id.write() = Some(skill_id.clone());
5945                    let response = AgentResponse::new(&question.question).with_metadata(
5946                        "disambiguation",
5947                        serde_json::json!({
5948                            "status": if awaiting_confirmation { "awaiting_confirmation" } else { "awaiting_clarification" },
5949                            "skill_id": skill_id,
5950                            "options": question.options,
5951                            "clarifying": question.clarifying,
5952                            "detection": {
5953                                "type": detection.ambiguity_type,
5954                                "confidence": detection.confidence,
5955                                "what_is_unclear": detection.what_is_unclear,
5956                            }
5957                        }),
5958                    );
5959                    drop(admission);
5960                    return Ok(SkillRouteResult::NeedsClarification {
5961                        response,
5962                        ownership: Some(DisambiguationOwnership {
5963                            epoch: expected_disambiguation_epoch,
5964                            state_generation: expected_state_generation,
5965                        }),
5966                    });
5967                }
5968                DisambiguationResult::Clarified { enriched_input, .. } => {
5969                    info!(skill_id = %skill_id, enriched = %enriched_input, "Skill disambiguation clarified");
5970                    let admission = self
5971                        .admit_disambiguation_redispatch(
5972                            expected_disambiguation_epoch,
5973                            expected_state_generation,
5974                        )
5975                        .await?;
5976                    drop(admission);
5977                    let content = self.execute_skill(&skill, &enriched_input).await?;
5978                    return Ok(SkillRouteResult::Response { skill_id, content });
5979                }
5980                DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
5981                    info!(skill_id = %skill_id, "Skill disambiguation best guess");
5982                    let admission = self
5983                        .admit_disambiguation_redispatch(
5984                            expected_disambiguation_epoch,
5985                            expected_state_generation,
5986                        )
5987                        .await?;
5988                    drop(admission);
5989                    let content = self.execute_skill(&skill, &enriched_input).await?;
5990                    return Ok(SkillRouteResult::Response { skill_id, content });
5991                }
5992                DisambiguationResult::GiveUp { reason } => {
5993                    warn!(skill_id = %skill_id, reason = %reason, "Skill disambiguation gave up");
5994                    let apology = self
5995                                .generate_localized_apology(
5996                                    "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
5997                                    &reason,
5998                                )
5999                                .await
6000                                .unwrap_or_else(|_| {
6001                                    format!("I'm sorry, I couldn't understand your request: {}", reason)
6002                                });
6003                    return Ok(SkillRouteResult::NeedsClarification {
6004                        response: AgentResponse::new(&apology),
6005                        ownership: None,
6006                    });
6007                }
6008                DisambiguationResult::Escalate { reason } => {
6009                    info!(skill_id = %skill_id, reason = %reason, "Skill disambiguation escalating");
6010                    let apology = self
6011                                .generate_localized_apology(
6012                                    "Explain briefly that you're transferring the user to a human agent for help.",
6013                                    &reason,
6014                                )
6015                                .await
6016                                .unwrap_or_else(|_| {
6017                                    format!("I need human assistance to help with your request: {}", reason)
6018                                });
6019                    return Ok(SkillRouteResult::NeedsClarification {
6020                        response: AgentResponse::new(&apology),
6021                        ownership: None,
6022                    });
6023                }
6024                DisambiguationResult::Abandoned { .. } => {
6025                    debug!(skill_id = %skill_id, "Skill disambiguation abandoned");
6026                    return Ok(SkillRouteResult::NoMatch);
6027                }
6028            }
6029        }
6030        let admission = self
6031            .admit_disambiguation_redispatch(
6032                expected_disambiguation_epoch,
6033                expected_state_generation,
6034            )
6035            .await?;
6036        drop(admission);
6037        let content = self.execute_skill(&skill, input).await?;
6038        Ok(SkillRouteResult::Response { skill_id, content })
6039    }
6040
6041    /// Result of skill routing.
6042    async fn try_skill_route(&self, input: &str) -> Result<SkillRouteResult> {
6043        if let Some(candidate) = self.select_skill_candidate(input).await? {
6044            self.commit_skill_candidate_route_result(candidate, input)
6045                .await
6046        } else {
6047            Ok(SkillRouteResult::NoMatch)
6048        }
6049    }
6050
6051    /// Execute a skill with reasoning and reflection, returning the response string.
6052    async fn execute_skill(&self, skill: &SkillDefinition, input: &str) -> Result<String> {
6053        if let Some(ref executor) = self.skill_executor {
6054            let skill_reasoning = self.get_skill_reasoning_config(skill);
6055            let skill_reflection = self.get_skill_reflection_config(skill);
6056
6057            debug!(
6058                skill_id = %skill.id,
6059                reasoning_mode = ?skill_reasoning.mode,
6060                reflection_enabled = ?skill_reflection.enabled,
6061                "Skill reasoning/reflection config"
6062            );
6063
6064            let response = self
6065                .observe_purpose(
6066                    ObservationPurpose::SkillPrompt,
6067                    executor.execute_with_invoker(skill, input, serde_json::json!({}), self),
6068                )
6069                .await?;
6070
6071            if skill_reflection.requires_evaluation() && skill_reflection.is_enabled() {
6072                let should_reflect = self
6073                    .should_reflect_with_config(input, &response, &skill_reflection)
6074                    .await?;
6075                if should_reflect {
6076                    let evaluated = self
6077                        .evaluate_and_retry_with_config(input, response, &skill_reflection)
6078                        .await?;
6079                    return Ok(evaluated);
6080                }
6081            }
6082
6083            return Ok(response);
6084        }
6085        Err(AgentError::Skill(
6086            "No skill executor configured".to_string(),
6087        ))
6088    }
6089
6090    /// Execute a skill by ID, bypassing the skill router.
6091    /// Used after skill-triggered disambiguation resolves to route directly to the matched skill.
6092    async fn execute_skill_by_id(&self, skill_id: &str, input: &str) -> Result<String> {
6093        let skill = self
6094            .skill_router
6095            .as_ref()
6096            .and_then(|r| r.get_skill(skill_id).cloned())
6097            .ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
6098        self.execute_skill(&skill, input).await
6099    }
6100
6101    async fn should_reflect_with_config(
6102        &self,
6103        input: &str,
6104        response: &str,
6105        config: &ReflectionConfig,
6106    ) -> Result<bool> {
6107        if !config.requires_evaluation() {
6108            return Ok(false);
6109        }
6110
6111        if config.is_enabled() {
6112            return Ok(true);
6113        }
6114
6115        let evaluator_llm = config
6116            .evaluator_llm
6117            .as_ref()
6118            .and_then(|alias| self.llm_registry.get(alias).ok())
6119            .or_else(|| self.llm_registry.router().ok())
6120            .or_else(|| self.llm_registry.default().ok());
6121
6122        let Some(llm) = evaluator_llm else {
6123            return Ok(false);
6124        };
6125
6126        let response_preview: String = response.chars().take(500).collect();
6127        let prompt = format!(
6128            r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
6129
6130User query: "{}"
6131Response: "{}"
6132
6133Answer YES or NO only."#,
6134            input, response_preview
6135        );
6136
6137        let messages = vec![ChatMessage::user(&prompt)];
6138        let result = self
6139            .observe_purpose(
6140                ObservationPurpose::ReflectionDecision,
6141                llm.complete(&messages, None),
6142            )
6143            .await;
6144
6145        match result {
6146            Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
6147            Err(_) => Ok(false),
6148        }
6149    }
6150
6151    async fn evaluate_and_retry_with_config(
6152        &self,
6153        input: &str,
6154        mut response: String,
6155        config: &ReflectionConfig,
6156    ) -> Result<String> {
6157        let llm = self.get_state_llm()?;
6158        let mut attempts = 0u32;
6159        let max_retries = config.max_retries;
6160
6161        loop {
6162            let evaluation = self
6163                .evaluate_response_with_config(input, &response, config)
6164                .await?;
6165
6166            if evaluation.passed || attempts >= max_retries {
6167                info!(
6168                    passed = evaluation.passed,
6169                    confidence = evaluation.confidence,
6170                    attempts = attempts + 1,
6171                    "Skill reflection evaluation complete"
6172                );
6173                return Ok(response);
6174            }
6175
6176            debug!(
6177                attempt = attempts + 1,
6178                failed_criteria = evaluation.failed_criteria().count(),
6179                "Skill response did not meet criteria, retrying"
6180            );
6181
6182            let feedback: Vec<String> = evaluation
6183                .failed_criteria()
6184                .map(|c| format!("- {}", c.criterion))
6185                .collect();
6186
6187            let retry_prompt = format!(
6188                "Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response to: {}",
6189                feedback.join("\n"),
6190                input
6191            );
6192
6193            let messages = vec![ChatMessage::user(&retry_prompt)];
6194            let retry_response = self
6195                .observe_purpose(
6196                    ObservationPurpose::ReflectionEvaluation,
6197                    llm.complete(&messages, None),
6198                )
6199                .await
6200                .map_err(|e| AgentError::LLM(e.to_string()))?;
6201
6202            response = retry_response.content.trim().to_string();
6203            attempts += 1;
6204        }
6205    }
6206
6207    async fn evaluate_response_with_config(
6208        &self,
6209        input: &str,
6210        response: &str,
6211        config: &ReflectionConfig,
6212    ) -> Result<EvaluationResult> {
6213        let evaluator_llm = config
6214            .evaluator_llm
6215            .as_ref()
6216            .and_then(|alias| self.llm_registry.get(alias).ok())
6217            .or_else(|| self.llm_registry.router().ok())
6218            .or_else(|| self.llm_registry.default().ok())
6219            .ok_or_else(|| AgentError::Config("No LLM available for evaluation".into()))?;
6220
6221        let criteria = &config.criteria;
6222        let criteria_list = criteria
6223            .iter()
6224            .enumerate()
6225            .map(|(i, c)| format!("{}. {}", i + 1, c))
6226            .collect::<Vec<_>>()
6227            .join("\n");
6228
6229        let prompt = format!(
6230            r#"Evaluate this response against the criteria.
6231
6232User query: "{}"
6233
6234Response to evaluate: "{}"
6235
6236Criteria:
6237{}
6238
6239For each criterion, respond with:
6240- criterion number
6241- PASS or FAIL
6242- brief reason
6243
6244Then provide overall confidence (0.0 to 1.0) and whether it passes overall.
6245
6246Format:
62471. PASS/FAIL - reason
62482. PASS/FAIL - reason
6249...
6250CONFIDENCE: 0.X
6251OVERALL: PASS/FAIL"#,
6252            input, response, criteria_list
6253        );
6254
6255        let messages = vec![ChatMessage::user(&prompt)];
6256        let eval_response = self
6257            .observe_purpose(
6258                ObservationPurpose::ReflectionEvaluation,
6259                evaluator_llm.complete(&messages, None),
6260            )
6261            .await
6262            .map_err(|e| AgentError::LLM(format!("Evaluation failed: {}", e)))?;
6263
6264        let content = eval_response.content.to_uppercase();
6265        let llm_pass = content.contains("OVERALL: PASS");
6266
6267        let confidence = content
6268            .lines()
6269            .find(|l| l.contains("CONFIDENCE:"))
6270            .and_then(|l| {
6271                l.split(':')
6272                    .nth(1)
6273                    .and_then(|v| v.trim().parse::<f32>().ok())
6274            })
6275            .unwrap_or(if llm_pass { 0.8 } else { 0.4 });
6276
6277        // Gate pass against confidence threshold.
6278        // LLM may say PASS but with low confidence - the threshold catches this.
6279        let overall_pass = llm_pass && confidence >= config.pass_threshold;
6280
6281        let mut criteria_results = Vec::new();
6282        for (i, criterion) in criteria.iter().enumerate() {
6283            let line_marker = format!("{}.", i + 1);
6284            let passed = eval_response
6285                .content
6286                .lines()
6287                .find(|l| l.contains(&line_marker))
6288                .map(|l| l.to_uppercase().contains("PASS"))
6289                .unwrap_or(overall_pass);
6290
6291            if passed {
6292                criteria_results.push(CriterionResult::pass(criterion));
6293            } else {
6294                criteria_results.push(CriterionResult::fail(criterion, "Did not meet criterion"));
6295            }
6296        }
6297
6298        Ok(EvaluationResult::new(overall_pass, confidence).with_criteria(criteria_results))
6299    }
6300
6301    /// Process input through the pipeline (state-level override or agent-level).
6302    async fn process_input(&self, input: &str) -> Result<ProcessData> {
6303        if let Some(processor) = self.get_state_process_processor() {
6304            let purpose = observation_purpose_for_process(processor.input_purpose_hint());
6305            return self
6306                .observe_purpose(purpose, processor.process_input(input))
6307                .await;
6308        }
6309        if let Some(ref processor) = self.process_processor {
6310            let purpose = observation_purpose_for_process(processor.input_purpose_hint());
6311            self.observe_purpose(purpose, processor.process_input(input))
6312                .await
6313        } else {
6314            Ok(ProcessData::new(input))
6315        }
6316    }
6317
6318    /// Process output through the pipeline (state-level override or agent-level).
6319    async fn process_output(
6320        &self,
6321        output: &str,
6322        input_context: &std::collections::HashMap<String, serde_json::Value>,
6323    ) -> Result<ProcessData> {
6324        if let Some(processor) = self.get_state_process_processor() {
6325            let purpose = observation_purpose_for_process(processor.output_purpose_hint());
6326            return self
6327                .observe_purpose(purpose, processor.process_output(output, input_context))
6328                .await;
6329        }
6330        if let Some(ref processor) = self.process_processor {
6331            let purpose = observation_purpose_for_process(processor.output_purpose_hint());
6332            self.observe_purpose(purpose, processor.process_output(output, input_context))
6333                .await
6334        } else {
6335            Ok(ProcessData::new(output))
6336        }
6337    }
6338
6339    /// Build a ProcessProcessor from the current state's process config, if any.
6340    fn get_state_process_processor(&self) -> Option<ProcessProcessor> {
6341        let sm = self.state_machine.as_ref()?;
6342        let def = sm.current_definition()?;
6343        let config = def.process.as_ref()?;
6344        let mut processor = ProcessProcessor::new(config.clone());
6345        if let Some(ref registry) = Some(self.llm_registry.clone()) {
6346            processor = processor.with_llm_registry(registry.clone());
6347        }
6348        processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
6349        Some(processor)
6350    }
6351
6352    // Timeout transitions reserve exit actions before committing ownership changes under the short write lock.
6353    async fn check_turn_timeout(&self) -> Result<()> {
6354        let Some(ref sm) = self.state_machine else {
6355            return Ok(());
6356        };
6357        let Some(timeout_state) = sm.check_timeout() else {
6358            return Ok(());
6359        };
6360        let claim_admission = self.disambiguation_admission.write().await;
6361        if sm.check_timeout().as_deref() != Some(timeout_state.as_str()) {
6362            return Ok(());
6363        }
6364        let Some(reservation) = self.reserve_state_transition() else {
6365            return Ok(());
6366        };
6367        let from_state = sm.current();
6368        let expected_state_generation = sm.generation();
6369        let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
6370        let history_before = sm.history();
6371        drop(claim_admission);
6372
6373        self.execute_state_exit_actions(&from_state).await;
6374
6375        let admission = self.disambiguation_admission.write().await;
6376        if sm.current() != from_state
6377            || sm.generation() != expected_state_generation
6378            || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6379            || sm.check_timeout().as_deref() != Some(timeout_state.as_str())
6380        {
6381            return Ok(());
6382        }
6383        sm.transition_to(&timeout_state, "max_turns exceeded")?;
6384        self.invalidate_pending_confirmation("state_timeout").await;
6385        let entered = sm.current();
6386        let is_reentry = Self::state_was_previously_entered(&entered, &from_state, &history_before);
6387        drop(admission);
6388
6389        self.execute_state_enter_actions(&entered, is_reentry).await;
6390        drop(reservation);
6391        info!(to = %entered, "Timeout transition");
6392        Ok(())
6393    }
6394
6395    fn increment_turn(&self) {
6396        if let Some(ref sm) = self.state_machine {
6397            sm.increment_turn();
6398        }
6399    }
6400
6401    fn transitions_available_for_commit(&self) -> Option<(Vec<Transition>, String)> {
6402        let sm = self.state_machine.as_ref()?;
6403        let current = sm.current();
6404        let transitions: Vec<_> = sm
6405            .auto_transitions()
6406            .into_iter()
6407            .filter(|t| match t.cooldown_turns {
6408                Some(cd) if cd > 0 => {
6409                    let resolved = sm.config().resolve_full_path(&current, &t.to);
6410                    !sm.is_on_cooldown(&resolved, cd)
6411                }
6412                _ => true,
6413            })
6414            .collect();
6415        Some((transitions, current))
6416    }
6417
6418    fn transition_reason(transition: &Transition) -> String {
6419        if transition.when.is_empty() {
6420            "guard condition met".to_string()
6421        } else {
6422            transition.when.clone()
6423        }
6424    }
6425
6426    /// Builds transition context with optional staged writes overlaid.
6427    fn build_transition_context(
6428        &self,
6429        user_message: &str,
6430        response: &str,
6431        current_state: &str,
6432        staged: Option<&HashMap<String, Value>>,
6433    ) -> TransitionContext {
6434        let context_map = staged
6435            .map(|writes| self.build_context_with_staged(writes))
6436            .unwrap_or_else(|| self.build_context_with_overlays());
6437        TransitionContext::new(user_message, response, current_state).with_context(context_map)
6438    }
6439
6440    /// Selects a post-response transition without committing side effects.
6441    async fn select_transition_candidate(
6442        &self,
6443        user_message: &str,
6444        response: &str,
6445    ) -> Result<Option<TransitionCandidate>> {
6446        let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6447            return Ok(None);
6448        };
6449        let transitions: Vec<Transition> = transitions
6450            .into_iter()
6451            .filter(|transition| matches!(transition.timing, TransitionTiming::PostResponse))
6452            .collect();
6453        if transitions.is_empty() {
6454            return Ok(None);
6455        }
6456        let Some(evaluator) = self.transition_evaluator.as_ref() else {
6457            return Ok(None);
6458        };
6459        let context = self.build_transition_context(user_message, response, &current_state, None);
6460        let selected = self
6461            .observe_purpose(
6462                ObservationPurpose::StateTransitionEvaluation,
6463                evaluator.select_transition(&transitions, &context),
6464            )
6465            .await?;
6466        Ok(selected.map(|index| {
6467            let transition = transitions[index].clone();
6468            TransitionCandidate::new(
6469                current_state,
6470                transition.clone(),
6471                Self::transition_reason(&transition),
6472            )
6473        }))
6474    }
6475
6476    /// Selects a guard or resolved-intent transition without an LLM call.
6477    fn select_deterministic_transition_candidate(
6478        &self,
6479        user_message: &str,
6480        current_state: &str,
6481        transitions: &[Transition],
6482        staged: &HashMap<String, Value>,
6483    ) -> Option<TransitionCandidate> {
6484        let context = self.build_transition_context(user_message, "", current_state, Some(staged));
6485
6486        for transition in transitions {
6487            if let Some(guard) = transition.guard.as_ref()
6488                && evaluate_guard(guard, &context)
6489            {
6490                return Some(TransitionCandidate::new(
6491                    current_state,
6492                    transition.clone(),
6493                    Self::transition_reason(transition),
6494                ));
6495            }
6496        }
6497
6498        let resolved_intent = context
6499            .context
6500            .get("resolved_intent")
6501            .and_then(Value::as_str)
6502            .filter(|value| !value.is_empty());
6503        if let Some(resolved_intent) = resolved_intent {
6504            for transition in transitions {
6505                if transition.intent.as_deref() == Some(resolved_intent) {
6506                    return Some(TransitionCandidate::new(
6507                        current_state,
6508                        transition.clone(),
6509                        Self::transition_reason(transition),
6510                    ));
6511                }
6512            }
6513        }
6514
6515        None
6516    }
6517
6518    /// Commits a selected transition through the shared post-response path.
6519    async fn commit_transition_candidate(&self, candidate: &TransitionCandidate) -> Result<bool> {
6520        self.commit_transition_target(&candidate.from_state, candidate.target(), &candidate.reason)
6521            .await
6522    }
6523
6524    /// Runs state transition approval before any transition side effects.
6525    async fn approve_transition_target(&self, from_state: &str, target: &str) -> Result<bool> {
6526        let approved = self.check_state_hitl(Some(from_state), target).await?;
6527        if !approved {
6528            info!(to = %target, "State transition rejected by HITL");
6529        }
6530        Ok(approved)
6531    }
6532
6533    /// Applies an approved transition after reserving exit actions and keeping async hooks outside the commit lock.
6534    async fn apply_transition_target(
6535        &self,
6536        from_state: &str,
6537        target: &str,
6538        reason: &str,
6539        staged: Option<&HashMap<String, Value>>,
6540    ) -> Result<bool> {
6541        let Some(ref sm) = self.state_machine else {
6542            return Ok(false);
6543        };
6544        let claim_admission = self.disambiguation_admission.write().await;
6545        if sm.current() != from_state {
6546            return Ok(false);
6547        }
6548        let Some(reservation) = self.reserve_state_transition() else {
6549            return Ok(false);
6550        };
6551        let expected_state_generation = sm.generation();
6552        let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
6553        let history_before = sm.history();
6554        drop(claim_admission);
6555
6556        self.execute_state_exit_actions(from_state).await;
6557
6558        let admission = self.disambiguation_admission.write().await;
6559        if sm.current() != from_state
6560            || sm.generation() != expected_state_generation
6561            || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6562        {
6563            return Ok(false);
6564        }
6565        sm.transition_to(target, reason)?;
6566        self.invalidate_pending_confirmation("state_transition")
6567            .await;
6568        sm.reset_no_transition();
6569        if let Some(staged) = staged {
6570            self.commit_staged_context_writes(staged);
6571        }
6572        let entered = sm.current();
6573        let is_reentry = Self::state_was_previously_entered(&entered, from_state, &history_before);
6574        drop(admission);
6575
6576        self.execute_state_enter_actions(&entered, is_reentry).await;
6577        drop(reservation);
6578        self.hooks
6579            .on_state_transition(Some(from_state), &entered, reason)
6580            .await;
6581        info!(from = %from_state, to = %entered, "State transition");
6582        Ok(true)
6583    }
6584
6585    /// Approves and applies a transition without staged context writes.
6586    async fn commit_transition_target(
6587        &self,
6588        from_state: &str,
6589        target: &str,
6590        reason: &str,
6591    ) -> Result<bool> {
6592        if !self.approve_transition_target(from_state, target).await? {
6593            return Ok(false);
6594        }
6595        self.apply_transition_target(from_state, target, reason, None)
6596            .await
6597    }
6598
6599    /// Applies an approved pre-response transition before redispatch.
6600    async fn apply_pre_response_transition_candidate(
6601        &self,
6602        candidate: &TransitionCandidate,
6603        staged: &HashMap<String, Value>,
6604        processed_input: &str,
6605    ) -> Result<bool> {
6606        self.commit_root_user_message(processed_input).await?;
6607        self.apply_transition_target(
6608            &candidate.from_state,
6609            candidate.target(),
6610            &candidate.reason,
6611            Some(staged),
6612        )
6613        .await
6614    }
6615
6616    /// Commits a pre-response transition after approval and before redispatch.
6617    async fn commit_pre_response_transition_candidate(
6618        &self,
6619        candidate: &TransitionCandidate,
6620        staged: &HashMap<String, Value>,
6621        processed_input: &str,
6622    ) -> Result<bool> {
6623        if !self
6624            .approve_transition_target(&candidate.from_state, candidate.target())
6625            .await?
6626        {
6627            return Ok(false);
6628        }
6629        self.apply_pre_response_transition_candidate(candidate, staged, processed_input)
6630            .await
6631    }
6632
6633    /// Handles post-response transition misses with fallback counters.
6634    async fn handle_transition_miss(&self, current_state: &str) -> Result<bool> {
6635        let Some(ref sm) = self.state_machine else {
6636            return Ok(false);
6637        };
6638        sm.increment_no_transition();
6639        let Some(fallback) = sm.check_fallback() else {
6640            return Ok(false);
6641        };
6642        self.commit_transition_target(current_state, &fallback, "fallback after no transitions")
6643            .await
6644    }
6645
6646    /// Evaluates and commits post-response transitions for the committed response path.
6647    async fn evaluate_transitions(&self, user_message: &str, response: &str) -> Result<bool> {
6648        let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6649            return Ok(false);
6650        };
6651        if transitions.is_empty() {
6652            return Ok(false);
6653        }
6654        if let Some(candidate) = self
6655            .select_transition_candidate(user_message, response)
6656            .await?
6657        {
6658            return self.commit_transition_candidate(&candidate).await;
6659        }
6660        self.handle_transition_miss(&current_state).await
6661    }
6662
6663    /// Attempts deterministic pre-response routing before old-state response generation.
6664    async fn try_pre_response_transition(
6665        &self,
6666        processed_input: &str,
6667    ) -> Result<Option<AgentResponse>> {
6668        let optimization = &self.runtime_config.optimization;
6669        if !optimization.enabled || !optimization.pre_response_deterministic_transitions {
6670            return Ok(None);
6671        }
6672        let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6673            return Ok(None);
6674        };
6675        let eligible: Vec<Transition> = transitions
6676            .into_iter()
6677            .filter(|transition| !transition.requires_response)
6678            .filter(|transition| matches!(transition.timing, TransitionTiming::PreResponse))
6679            .collect();
6680        if eligible.is_empty() {
6681            return Ok(None);
6682        }
6683
6684        let empty_staged = HashMap::new();
6685        let mut extracted_staged: Option<HashMap<String, Value>> = None;
6686        let mut selected: Option<(TransitionCandidate, HashMap<String, Value>)> = None;
6687
6688        for transition in &eligible {
6689            let use_extractors = optimization.pre_response_extractors || transition.run_extractors;
6690            let staged_for_eval = if use_extractors {
6691                if extracted_staged.is_none() {
6692                    extracted_staged =
6693                        Some(self.run_context_extractors_staged(processed_input).await);
6694                }
6695                extracted_staged.as_ref().unwrap_or(&empty_staged)
6696            } else {
6697                &empty_staged
6698            };
6699
6700            if let Some(candidate) = self.select_deterministic_transition_candidate(
6701                processed_input,
6702                &current_state,
6703                std::slice::from_ref(transition),
6704                staged_for_eval,
6705            ) {
6706                let staged_for_commit = if use_extractors {
6707                    staged_for_eval.clone()
6708                } else {
6709                    HashMap::new()
6710                };
6711                selected = Some((candidate, staged_for_commit));
6712                break;
6713            }
6714        }
6715
6716        let Some((candidate, staged)) = selected else {
6717            return Ok(None);
6718        };
6719
6720        if !self
6721            .commit_pre_response_transition_candidate(&candidate, &staged, processed_input)
6722            .await?
6723        {
6724            return Ok(None);
6725        }
6726        self.redispatch_current_state(processed_input)
6727            .await
6728            .map(Some)
6729    }
6730
6731    //
6732    // Speculative branches overlap independent decisions but still commit exactly one path.
6733    // Losing branches must remain data only and must not write memory, run tools, or emit output.
6734    //
6735    async fn try_speculative_branches(
6736        &self,
6737        processed_input: &str,
6738        input_context: &HashMap<String, Value>,
6739    ) -> Result<Option<AgentResponse>> {
6740        let optimization = &self.runtime_config.optimization;
6741        if !optimization.enabled {
6742            return Ok(None);
6743        }
6744
6745        let effective_reasoning_mode = self.get_effective_reasoning_config().mode.clone();
6746        if !matches!(
6747            effective_reasoning_mode,
6748            ReasoningMode::None | ReasoningMode::Auto
6749        ) {
6750            return Ok(None);
6751        }
6752
6753        let mut transition_enabled =
6754            optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
6755        let mut skill_enabled = optimization.speculative_skill_routing
6756            && self.skill_router.is_some()
6757            && self.pending_skill_id.read().is_none();
6758        let mut reasoning_enabled = optimization.speculative_reasoning_auto
6759            && matches!(effective_reasoning_mode, ReasoningMode::Auto);
6760
6761        if matches!(effective_reasoning_mode, ReasoningMode::Auto)
6762            && (!reasoning_enabled || optimization.max_speculative_llm_calls_per_turn < 2)
6763        {
6764            return Ok(None);
6765        }
6766
6767        if !transition_enabled && !skill_enabled && !reasoning_enabled {
6768            return Ok(None);
6769        }
6770
6771        let mut optional_slots = optimization.max_parallel_runtime_tasks.saturating_sub(1);
6772        let mut speculative_call_slots = optimization
6773            .max_speculative_llm_calls_per_turn
6774            .saturating_sub(1);
6775        if reasoning_enabled {
6776            if optional_slots == 0 || speculative_call_slots == 0 {
6777                return Ok(None);
6778            }
6779            optional_slots -= 1;
6780            speculative_call_slots -= 1;
6781        }
6782        if transition_enabled {
6783            if optional_slots == 0 {
6784                transition_enabled = false;
6785            } else {
6786                optional_slots -= 1;
6787            }
6788        }
6789        if skill_enabled && (optional_slots == 0 || speculative_call_slots == 0) {
6790            skill_enabled = false;
6791        }
6792
6793        if !transition_enabled && !skill_enabled && !reasoning_enabled {
6794            return Ok(None);
6795        }
6796
6797        let main_kind = if transition_enabled {
6798            RuntimeOptimizationKind::ParallelStateTransition
6799        } else if skill_enabled {
6800            RuntimeOptimizationKind::SpeculativeSkillRouting
6801        } else {
6802            RuntimeOptimizationKind::SpeculativeReasoningAuto
6803        };
6804        if !self.reserve_active_speculative_llm_call(main_kind) {
6805            return Ok(None);
6806        }
6807
6808        let mut branch_set = ScheduledBranchSet::new(optimization.max_parallel_runtime_tasks)?;
6809        let main_branch = RuntimeBranch::new(
6810            RuntimeTaskPurpose::MainResponse,
6811            main_kind,
6812            RuntimeTaskPriority::Normal,
6813            RuntimeCommitBehavior::FinalResponse,
6814        );
6815        let transition_branch = RuntimeBranch::new(
6816            RuntimeTaskPurpose::StateTransition,
6817            RuntimeOptimizationKind::ParallelStateTransition,
6818            RuntimeTaskPriority::Critical,
6819            RuntimeCommitBehavior::TransitionDecision,
6820        );
6821        let skill_branch = RuntimeBranch::new(
6822            RuntimeTaskPurpose::SkillRouting,
6823            RuntimeOptimizationKind::SpeculativeSkillRouting,
6824            RuntimeTaskPriority::High,
6825            RuntimeCommitBehavior::SkillSelection,
6826        );
6827        let reasoning_branch = RuntimeBranch::new(
6828            RuntimeTaskPurpose::ReasoningJudge,
6829            RuntimeOptimizationKind::SpeculativeReasoningAuto,
6830            RuntimeTaskPriority::Normal,
6831            RuntimeCommitBehavior::ReasoningDecision,
6832        );
6833        let main_id = main_branch.branch_id();
6834        let transition_id = transition_branch.branch_id();
6835        let skill_id = skill_branch.branch_id();
6836        let reasoning_id = reasoning_branch.branch_id();
6837
6838        let main_id_for_future = main_id.clone();
6839        if !branch_set.schedule(
6840            main_branch,
6841            Box::pin(async move {
6842                match crate::optimization::observability::with_branch_observation(
6843                    &main_id_for_future,
6844                    main_kind,
6845                    RuntimeCommitBehavior::FinalResponse,
6846                    self.generate_main_response_draft(processed_input, &ReasoningMode::None),
6847                )
6848                .await
6849                {
6850                    Ok(draft) => RuntimeBranchResult::MainDraft(draft),
6851                    Err(error) => RuntimeBranchResult::Failed(error),
6852                }
6853            }),
6854        ) {
6855            return Ok(None);
6856        }
6857
6858        if transition_enabled {
6859            let transition_id_for_future = transition_id.clone();
6860            if !branch_set.schedule(
6861                transition_branch,
6862                Box::pin(async move {
6863                    match crate::optimization::observability::with_branch_observation(
6864                        &transition_id_for_future,
6865                        RuntimeOptimizationKind::ParallelStateTransition,
6866                        RuntimeCommitBehavior::TransitionDecision,
6867                        self.select_parallel_transition_candidate(processed_input),
6868                    )
6869                    .await
6870                    {
6871                        Ok(ParallelTransitionSelection::Candidate(candidate)) => {
6872                            RuntimeBranchResult::Transition(Some(candidate))
6873                        }
6874                        Ok(ParallelTransitionSelection::NoMatch) => {
6875                            RuntimeBranchResult::Transition(None)
6876                        }
6877                        Ok(ParallelTransitionSelection::ReservationExhausted) => {
6878                            RuntimeBranchResult::Cancelled
6879                        }
6880                        Err(error) => RuntimeBranchResult::Failed(error),
6881                    }
6882                }),
6883            ) {
6884                transition_enabled = false;
6885            }
6886        }
6887
6888        if skill_enabled {
6889            let skill_id_for_future = skill_id.clone();
6890            if !branch_set.schedule(
6891                skill_branch,
6892                Box::pin(async move {
6893                    if !self.reserve_active_speculative_llm_call(
6894                        RuntimeOptimizationKind::SpeculativeSkillRouting,
6895                    ) {
6896                        return RuntimeBranchResult::Cancelled;
6897                    }
6898                    match crate::optimization::observability::with_branch_observation(
6899                        &skill_id_for_future,
6900                        RuntimeOptimizationKind::SpeculativeSkillRouting,
6901                        RuntimeCommitBehavior::SkillSelection,
6902                        self.select_skill_candidate(processed_input),
6903                    )
6904                    .await
6905                    {
6906                        Ok(candidate) => RuntimeBranchResult::Skill(candidate),
6907                        Err(error) => RuntimeBranchResult::Failed(error),
6908                    }
6909                }),
6910            ) {
6911                skill_enabled = false;
6912            }
6913        }
6914
6915        if reasoning_enabled {
6916            let reasoning_id_for_future = reasoning_id.clone();
6917            if !branch_set.schedule(
6918                reasoning_branch,
6919                Box::pin(async move {
6920                    if !self.reserve_active_speculative_llm_call(
6921                        RuntimeOptimizationKind::SpeculativeReasoningAuto,
6922                    ) {
6923                        return RuntimeBranchResult::Cancelled;
6924                    }
6925                    match crate::optimization::observability::with_branch_observation(
6926                        &reasoning_id_for_future,
6927                        RuntimeOptimizationKind::SpeculativeReasoningAuto,
6928                        RuntimeCommitBehavior::ReasoningDecision,
6929                        self.determine_reasoning_mode_strict(processed_input),
6930                    )
6931                    .await
6932                    {
6933                        Ok(mode) => RuntimeBranchResult::Reasoning(mode),
6934                        Err(error) => RuntimeBranchResult::Failed(error),
6935                    }
6936                }),
6937            ) {
6938                reasoning_enabled = false;
6939            }
6940        }
6941
6942        if matches!(effective_reasoning_mode, ReasoningMode::Auto) && !reasoning_enabled {
6943            self.finalize_pending_branches(branch_set.cancel_pending());
6944            return Ok(None);
6945        }
6946
6947        if !transition_enabled && !skill_enabled && !reasoning_enabled {
6948            self.finalize_pending_branches(branch_set.cancel_pending());
6949            return Ok(None);
6950        }
6951
6952        let mut main_pending = true;
6953        let mut skill_pending = skill_enabled;
6954        let mut reasoning_pending = reasoning_enabled;
6955        let mut transition_finalized = !transition_enabled;
6956        let mut skill_finalized = !skill_enabled;
6957        let mut reasoning_finalized = !reasoning_enabled;
6958        let mut main_result: Option<Result<MainResponseDraft>> = None;
6959        let mut transition_candidate: Option<TransitionCandidate> = None;
6960        let mut skill_candidate: Option<SkillCandidate> = None;
6961        let mut reasoning_decision: Option<ReasoningMode> = None;
6962        let mut transition_fallback_required = false;
6963        let mut skill_fallback_required = false;
6964        let mut reasoning_fallback_required = false;
6965
6966        loop {
6967            if let Some(candidate) = transition_candidate.take() {
6968                if self
6969                    .approve_transition_target(&candidate.from_state, candidate.target())
6970                    .await?
6971                {
6972                    // Drop losing provider futures before transition side effects can reuse their shared resources.
6973                    self.finalize_pending_branches(branch_set.cancel_pending());
6974                    if !main_pending {
6975                        self.finalize_branch_loss(
6976                            &main_id,
6977                            main_kind,
6978                            RuntimeCommitBehavior::FinalResponse,
6979                            false,
6980                            main_result.as_ref().map(|result| result.is_err()),
6981                        );
6982                    }
6983                    if skill_enabled && !skill_pending {
6984                        self.finalize_branch_loss(
6985                            &skill_id,
6986                            RuntimeOptimizationKind::SpeculativeSkillRouting,
6987                            RuntimeCommitBehavior::SkillSelection,
6988                            false,
6989                            Some(false),
6990                        );
6991                    }
6992                    if reasoning_enabled && !reasoning_pending {
6993                        self.finalize_branch_loss(
6994                            &reasoning_id,
6995                            RuntimeOptimizationKind::SpeculativeReasoningAuto,
6996                            RuntimeCommitBehavior::ReasoningDecision,
6997                            false,
6998                            Some(false),
6999                        );
7000                    }
7001                    if !self
7002                        .apply_pre_response_transition_candidate(
7003                            &candidate,
7004                            &HashMap::new(),
7005                            processed_input,
7006                        )
7007                        .await?
7008                    {
7009                        self.finalize_optional_branch(
7010                            &transition_id,
7011                            RuntimeOptimizationKind::ParallelStateTransition,
7012                            RuntimeCommitBehavior::TransitionDecision,
7013                            "discarded",
7014                            false,
7015                        );
7016                        return Ok(None);
7017                    }
7018                    self.finalize_optional_branch(
7019                        &transition_id,
7020                        RuntimeOptimizationKind::ParallelStateTransition,
7021                        RuntimeCommitBehavior::TransitionDecision,
7022                        "committed",
7023                        true,
7024                    );
7025                    return self
7026                        .redispatch_current_state(processed_input)
7027                        .await
7028                        .map(Some);
7029                }
7030                self.finalize_optional_branch(
7031                    &transition_id,
7032                    RuntimeOptimizationKind::ParallelStateTransition,
7033                    RuntimeCommitBehavior::TransitionDecision,
7034                    "discarded",
7035                    false,
7036                );
7037                transition_finalized = true;
7038            }
7039
7040            if transition_finalized && skill_candidate.is_some() {
7041                let candidate = skill_candidate.take().unwrap();
7042                self.finalize_optional_branch(
7043                    &skill_id,
7044                    RuntimeOptimizationKind::SpeculativeSkillRouting,
7045                    RuntimeCommitBehavior::SkillSelection,
7046                    "committed",
7047                    true,
7048                );
7049                if !main_pending {
7050                    self.finalize_branch_loss(
7051                        &main_id,
7052                        main_kind,
7053                        RuntimeCommitBehavior::FinalResponse,
7054                        false,
7055                        main_result.as_ref().map(|result| result.is_err()),
7056                    );
7057                }
7058                if reasoning_enabled && !reasoning_pending {
7059                    self.finalize_branch_loss(
7060                        &reasoning_id,
7061                        RuntimeOptimizationKind::SpeculativeReasoningAuto,
7062                        RuntimeCommitBehavior::ReasoningDecision,
7063                        false,
7064                        Some(false),
7065                    );
7066                }
7067                self.finalize_pending_branches(branch_set.cancel_pending());
7068                self.commit_root_user_message(processed_input).await?;
7069                return match self
7070                    .commit_skill_candidate_route_result(candidate, processed_input)
7071                    .await?
7072                {
7073                    SkillRouteResult::Response { skill_id, content } => self
7074                        .handle_skill_response(processed_input, &skill_id, content, input_context)
7075                        .await
7076                        .map(Some),
7077                    SkillRouteResult::NeedsClarification {
7078                        response,
7079                        ownership,
7080                    } => {
7081                        let admission = self
7082                            .admit_optional_disambiguation_ownership(ownership)
7083                            .await?;
7084                        if response
7085                            .metadata
7086                            .as_ref()
7087                            .and_then(|m| m.get("disambiguation"))
7088                            .and_then(|d| d.get("status"))
7089                            .and_then(|s| s.as_str())
7090                            == Some("awaiting_clarification")
7091                        {
7092                            self.memory
7093                                .add_message(ChatMessage::assistant(&response.content))
7094                                .await?;
7095                        }
7096                        drop(admission);
7097                        self.finish_turn_if_root(&response).await?;
7098                        Ok(Some(response))
7099                    }
7100                    SkillRouteResult::NoMatch => Ok(None),
7101                };
7102            }
7103
7104            if transition_finalized
7105                && skill_finalized
7106                && let Some(reasoning_mode) = reasoning_decision.take()
7107            {
7108                if !matches!(reasoning_mode, ReasoningMode::None) {
7109                    self.finalize_optional_branch(
7110                        &reasoning_id,
7111                        RuntimeOptimizationKind::SpeculativeReasoningAuto,
7112                        RuntimeCommitBehavior::ReasoningDecision,
7113                        "committed",
7114                        true,
7115                    );
7116                    if !main_pending {
7117                        self.finalize_branch_loss(
7118                            &main_id,
7119                            main_kind,
7120                            RuntimeCommitBehavior::FinalResponse,
7121                            false,
7122                            main_result.as_ref().map(|result| result.is_err()),
7123                        );
7124                    }
7125                    self.finalize_pending_branches(branch_set.cancel_pending());
7126                    self.commit_root_user_message(processed_input).await?;
7127                    return if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
7128                        self.handle_plan_and_execute(processed_input, input_context, true)
7129                            .await
7130                            .map(Some)
7131                    } else {
7132                        self.run_committed_response_loop_with_reasoning(
7133                            processed_input,
7134                            input_context,
7135                            reasoning_mode,
7136                            true,
7137                        )
7138                        .await
7139                        .map(Some)
7140                    };
7141                }
7142                self.finalize_optional_branch(
7143                    &reasoning_id,
7144                    RuntimeOptimizationKind::SpeculativeReasoningAuto,
7145                    RuntimeCommitBehavior::ReasoningDecision,
7146                    "committed",
7147                    true,
7148                );
7149                reasoning_finalized = true;
7150            }
7151
7152            if transition_finalized && skill_finalized && reasoning_finalized {
7153                if transition_fallback_required
7154                    || skill_fallback_required
7155                    || reasoning_fallback_required
7156                {
7157                    if !main_pending {
7158                        self.finalize_branch_loss(
7159                            &main_id,
7160                            main_kind,
7161                            RuntimeCommitBehavior::FinalResponse,
7162                            false,
7163                            main_result.as_ref().map(|result| result.is_err()),
7164                        );
7165                    }
7166                    self.finalize_pending_branches(branch_set.cancel_pending());
7167                    return Ok(None);
7168                }
7169
7170                if let Some(result) = main_result.take() {
7171                    let draft = match result {
7172                        Ok(draft) => draft,
7173                        Err(error) => {
7174                            self.finalize_optional_branch(
7175                                &main_id,
7176                                main_kind,
7177                                RuntimeCommitBehavior::FinalResponse,
7178                                "failed",
7179                                false,
7180                            );
7181                            self.finalize_pending_branches(branch_set.cancel_pending());
7182                            return Err(error);
7183                        }
7184                    };
7185                    self.finalize_optional_branch(
7186                        &main_id,
7187                        main_kind,
7188                        RuntimeCommitBehavior::FinalResponse,
7189                        "committed",
7190                        true,
7191                    );
7192                    self.finalize_pending_branches(branch_set.cancel_pending());
7193                    return self
7194                        .commit_main_response_draft(
7195                            processed_input,
7196                            input_context,
7197                            draft,
7198                            ReasoningMode::None,
7199                            reasoning_enabled,
7200                        )
7201                        .await
7202                        .map(Some);
7203                }
7204            }
7205
7206            if branch_set.is_empty() {
7207                return Ok(None);
7208            }
7209
7210            let Some(outcome) = branch_set.next_completed().await else {
7211                return Ok(None);
7212            };
7213            let branch_id = outcome.branch.branch_id();
7214            match outcome.result {
7215                RuntimeBranchResult::MainDraft(draft) => {
7216                    main_pending = false;
7217                    main_result = Some(Ok(draft));
7218                }
7219                RuntimeBranchResult::Transition(candidate) => {
7220                    if let Some(candidate) = candidate {
7221                        transition_candidate = Some(candidate);
7222                    } else {
7223                        self.finalize_optional_branch(
7224                            &transition_id,
7225                            RuntimeOptimizationKind::ParallelStateTransition,
7226                            RuntimeCommitBehavior::TransitionDecision,
7227                            "discarded",
7228                            false,
7229                        );
7230                        transition_finalized = true;
7231                    }
7232                }
7233                RuntimeBranchResult::Skill(candidate) => {
7234                    skill_pending = false;
7235                    if let Some(candidate) = candidate {
7236                        skill_candidate = Some(candidate);
7237                    } else {
7238                        self.finalize_optional_branch(
7239                            &skill_id,
7240                            RuntimeOptimizationKind::SpeculativeSkillRouting,
7241                            RuntimeCommitBehavior::SkillSelection,
7242                            "discarded",
7243                            false,
7244                        );
7245                        skill_finalized = true;
7246                    }
7247                }
7248                RuntimeBranchResult::Reasoning(mode) => {
7249                    reasoning_pending = false;
7250                    reasoning_decision = Some(mode);
7251                }
7252                RuntimeBranchResult::Failed(error) => {
7253                    if branch_id == main_id {
7254                        main_pending = false;
7255                        main_result = Some(Err(error));
7256                    } else if branch_id == transition_id {
7257                        self.finalize_optional_branch(
7258                            &transition_id,
7259                            RuntimeOptimizationKind::ParallelStateTransition,
7260                            RuntimeCommitBehavior::TransitionDecision,
7261                            "failed",
7262                            false,
7263                        );
7264                        transition_finalized = true;
7265                    } else if branch_id == skill_id {
7266                        skill_pending = false;
7267                        self.finalize_optional_branch(
7268                            &skill_id,
7269                            RuntimeOptimizationKind::SpeculativeSkillRouting,
7270                            RuntimeCommitBehavior::SkillSelection,
7271                            "failed",
7272                            false,
7273                        );
7274                        skill_finalized = true;
7275                    } else if branch_id == reasoning_id {
7276                        reasoning_pending = false;
7277                        self.finalize_optional_branch(
7278                            &reasoning_id,
7279                            RuntimeOptimizationKind::SpeculativeReasoningAuto,
7280                            RuntimeCommitBehavior::ReasoningDecision,
7281                            "failed",
7282                            false,
7283                        );
7284                        reasoning_finalized = true;
7285                    }
7286                }
7287                RuntimeBranchResult::Cancelled => {
7288                    self.finalize_optional_branch(
7289                        &branch_id,
7290                        outcome.branch.optimization,
7291                        outcome.branch.commit_behavior,
7292                        "cancelled",
7293                        false,
7294                    );
7295                    if branch_id == main_id {
7296                        main_pending = false;
7297                        main_result =
7298                            Some(Err(AgentError::Other("main branch cancelled".to_string())));
7299                    } else if branch_id == transition_id {
7300                        transition_finalized = true;
7301                        transition_fallback_required = true;
7302                    } else if branch_id == skill_id {
7303                        skill_pending = false;
7304                        skill_finalized = true;
7305                        skill_fallback_required = true;
7306                    } else if branch_id == reasoning_id {
7307                        reasoning_pending = false;
7308                        reasoning_finalized = true;
7309                        reasoning_fallback_required = true;
7310                    }
7311                }
7312            }
7313        }
7314    }
7315
7316    fn finalize_pending_branches(&self, branches: Vec<RuntimeBranch>) {
7317        for branch in branches {
7318            self.finalize_optional_branch(
7319                &branch.branch_id(),
7320                branch.optimization,
7321                branch.commit_behavior,
7322                "cancelled",
7323                false,
7324            );
7325        }
7326    }
7327
7328    //
7329    // Pending losers are reported as cancelled because their futures are dropped before completion.
7330    // Completed losers keep failed or discarded status based on their recorded result.
7331    //
7332    fn finalize_branch_loss(
7333        &self,
7334        branch_id: &str,
7335        optimization: RuntimeOptimizationKind,
7336        commit_behavior: RuntimeCommitBehavior,
7337        pending: bool,
7338        completed_failed: Option<bool>,
7339    ) {
7340        let status = if pending {
7341            "cancelled"
7342        } else if completed_failed.unwrap_or(false) {
7343            "failed"
7344        } else {
7345            "discarded"
7346        };
7347        self.finalize_optional_branch(branch_id, optimization, commit_behavior, status, false);
7348    }
7349
7350    //
7351    // Finalization is separated from commit so losing branches remain observable.
7352    // This helper must not mutate runtime state other than observability.
7353    //
7354    fn finalize_optional_branch(
7355        &self,
7356        branch_id: &str,
7357        optimization: RuntimeOptimizationKind,
7358        commit_behavior: RuntimeCommitBehavior,
7359        status: &str,
7360        winner: bool,
7361    ) {
7362        crate::optimization::observability::finalize_branch(
7363            self.observability_manager.as_ref(),
7364            branch_id,
7365            status,
7366            winner,
7367            optimization,
7368            commit_behavior,
7369        );
7370    }
7371
7372    //
7373    // This is only an eligibility check.
7374    // Actual transition selection happens in select_parallel_transition_candidate.
7375    //
7376    fn has_parallel_transition_candidates(&self) -> bool {
7377        self.transitions_available_for_commit()
7378            .map(|(transitions, _)| {
7379                transitions
7380                    .iter()
7381                    .any(|transition| matches!(transition.timing, TransitionTiming::Parallel))
7382            })
7383            .unwrap_or(false)
7384    }
7385
7386    //
7387    // Parallel transition prompts must not depend on assistant response text.
7388    // Keep this branch response-independent or it can race against invalid context.
7389    //
7390    async fn select_parallel_transition_candidate(
7391        &self,
7392        processed_input: &str,
7393    ) -> Result<ParallelTransitionSelection> {
7394        let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
7395            return Ok(ParallelTransitionSelection::NoMatch);
7396        };
7397        let parallel: Vec<Transition> = transitions
7398            .into_iter()
7399            .filter(|transition| matches!(transition.timing, TransitionTiming::Parallel))
7400            .filter(|transition| !transition.requires_response)
7401            .collect();
7402        if parallel.is_empty() {
7403            return Ok(ParallelTransitionSelection::NoMatch);
7404        }
7405        let empty_staged = HashMap::new();
7406        if let Some(candidate) = self.select_deterministic_transition_candidate(
7407            processed_input,
7408            &current_state,
7409            &parallel,
7410            &empty_staged,
7411        ) {
7412            return Ok(ParallelTransitionSelection::Candidate(candidate));
7413        }
7414        let when_transitions: Vec<(usize, &Transition)> = parallel
7415            .iter()
7416            .enumerate()
7417            .filter(|(_, transition)| !transition.when.trim().is_empty())
7418            .collect();
7419        if when_transitions.is_empty() {
7420            return Ok(ParallelTransitionSelection::NoMatch);
7421        }
7422        let llm = self
7423            .llm_registry
7424            .router()
7425            .or_else(|_| self.llm_registry.default())
7426            .map_err(|e| AgentError::Config(e.to_string()))?;
7427        let conditions = when_transitions
7428            .iter()
7429            .enumerate()
7430            .map(|(display_idx, (_, transition))| {
7431                format!("{}. {}", display_idx + 1, transition.when)
7432            })
7433            .collect::<Vec<_>>()
7434            .join("\n");
7435        if !self
7436            .reserve_active_speculative_llm_call(RuntimeOptimizationKind::ParallelStateTransition)
7437        {
7438            return Ok(ParallelTransitionSelection::ReservationExhausted);
7439        }
7440        let context_preview = self.branch_context_preview();
7441        let prompt = format!(
7442            "Based only on the current user message and context, which transition condition is met?\n\nCurrent state: {}\nUser message: {}\nContext:\n{}\n\nConditions:\n{}\n0. None of the above\n\nReply with ONLY the number (0-{}).",
7443            current_state,
7444            processed_input,
7445            context_preview,
7446            conditions,
7447            when_transitions.len()
7448        );
7449        let response = self
7450            .observe_purpose(
7451                ObservationPurpose::StateTransitionEvaluation,
7452                llm.complete(&[ChatMessage::user(prompt)], None),
7453            )
7454            .await
7455            .map_err(|e| AgentError::LLM(e.to_string()))?;
7456        let choice = response.content.trim().parse::<usize>().unwrap_or(0);
7457        if choice == 0 || choice > when_transitions.len() {
7458            return Ok(ParallelTransitionSelection::NoMatch);
7459        }
7460        let transition = when_transitions[choice - 1].1.clone();
7461        Ok(ParallelTransitionSelection::Candidate(
7462            TransitionCandidate::new(
7463                current_state,
7464                transition.clone(),
7465                Self::transition_reason(&transition),
7466            ),
7467        ))
7468    }
7469
7470    /// Re-enters the runtime loop after an optimized transition commits.
7471    async fn redispatch_current_state(&self, processed_input: &str) -> Result<AgentResponse> {
7472        const MAX_REDISPATCH_DEPTH: u32 = 3;
7473        let current_depth = *self.redispatch_depth.read();
7474        if current_depth >= MAX_REDISPATCH_DEPTH {
7475            warn!(depth = current_depth, "Re-dispatch depth limit reached");
7476            let response = AgentResponse::new("");
7477            self.finish_turn_if_root(&response).await?;
7478            return Ok(response);
7479        }
7480        *self.redispatch_depth.write() += 1;
7481        if let Some(context) = self.active_turn_context.write().as_mut() {
7482            context.enter_redispatch();
7483        }
7484        let result = Box::pin(self.run_loop_internal(processed_input)).await;
7485        *self.redispatch_depth.write() -= 1;
7486        if let Some(context) = self.active_turn_context.write().as_mut() {
7487            context.exit_redispatch();
7488        }
7489        let response = result?;
7490        self.finish_turn_if_root(&response).await?;
7491        Ok(response)
7492    }
7493
7494    /// Runs final response hooks and maintenance only for the root dispatch.
7495    async fn finish_turn_if_root(&self, response: &AgentResponse) -> Result<()> {
7496        if *self.redispatch_depth.read() == 0 {
7497            self.post_turn_session_lifecycle().await?;
7498            if let Some(context) = self.active_turn_context.write().as_mut() {
7499                context.mark_post_turn_lifecycle_completed();
7500            }
7501            self.hooks.on_response(response).await;
7502            self.end_root_turn();
7503        }
7504        Ok(())
7505    }
7506
7507    /// Execute on_exit actions for a state being left.
7508    async fn execute_state_exit_actions(&self, state_path: &str) {
7509        if let Some(ref sm) = self.state_machine
7510            && let Some(def) = sm.get_definition(state_path)
7511            && !def.on_exit.is_empty()
7512        {
7513            debug!(state = %state_path, count = def.on_exit.len(), "Executing on_exit actions");
7514            self.execute_state_actions(&def.on_exit).await;
7515        }
7516    }
7517
7518    /// Returns whether a transition target had already been entered before this transition.
7519    fn state_was_previously_entered(
7520        state_path: &str,
7521        from_state: &str,
7522        history_before: &[StateTransitionEvent],
7523    ) -> bool {
7524        state_path == from_state
7525            || history_before
7526                .iter()
7527                .any(|event| event.from == state_path || event.to == state_path)
7528    }
7529
7530    /// Execute on_enter (or on_reenter) actions for a state being entered.
7531    async fn execute_state_enter_actions(&self, state_path: &str, is_reentry: bool) {
7532        if let Some(ref sm) = self.state_machine
7533            && let Some(def) = sm.get_definition(state_path)
7534        {
7535            if is_reentry && !def.on_reenter.is_empty() {
7536                debug!(state = %state_path, count = def.on_reenter.len(), "Executing on_reenter actions");
7537                self.execute_state_actions(&def.on_reenter).await;
7538            } else if !def.on_enter.is_empty() {
7539                debug!(state = %state_path, count = def.on_enter.len(), "Executing on_enter actions");
7540                self.execute_state_actions(&def.on_enter).await;
7541            }
7542        }
7543    }
7544
7545    /// Execute a list of state actions (tool calls, skill invocations, context updates, LLM prompts).
7546    async fn execute_state_actions(&self, actions: &[StateAction]) {
7547        for (action_index, action) in actions.iter().enumerate() {
7548            match action {
7549                StateAction::Tool { tool, args } => {
7550                    let raw_args = args.clone().unwrap_or(Value::Object(Default::default()));
7551                    let args_value = self.render_action_args(&raw_args);
7552                    let state = self.state_machine.as_ref().map(|sm| sm.current());
7553                    let request = ToolExecutionRequest::new(
7554                        uuid::Uuid::new_v4().to_string(),
7555                        tool.clone(),
7556                        args_value,
7557                        ToolCallSource::StateAction {
7558                            state,
7559                            action_index,
7560                        },
7561                    );
7562                    match self.execute_tool_record(request).await {
7563                        Ok(record) if record.success => {
7564                            debug!(tool = %record.canonical_id, "State action: tool executed");
7565                            let _ = self.context_manager.set(
7566                                "last_tool_result",
7567                                serde_json::Value::String(record.model_output_string()),
7568                            );
7569                            let _ = self.context_manager.set(
7570                                "last_tool_record",
7571                                serde_json::to_value(record).unwrap_or(Value::Null),
7572                            );
7573                        }
7574                        Ok(record) => {
7575                            warn!(tool = %record.canonical_id, error = %record.output, "State action: tool failed");
7576                        }
7577                        Err(e) => {
7578                            warn!(tool = %tool, error = %e, "State action: tool failed")
7579                        }
7580                    }
7581                }
7582                StateAction::Skill { skill } => {
7583                    if let Some(ref executor) = self.skill_executor {
7584                        if let Some(def) = self.skills.iter().find(|s| s.id == *skill) {
7585                            match executor
7586                                .execute_with_invoker(def, "", serde_json::json!({}), self)
7587                                .await
7588                            {
7589                                Ok(_) => debug!(skill = %skill, "State action: skill executed"),
7590                                Err(e) => {
7591                                    warn!(skill = %skill, error = %e, "State action: skill failed")
7592                                }
7593                            }
7594                        } else {
7595                            warn!(skill = %skill, "State action: skill not found");
7596                        }
7597                    }
7598                }
7599                StateAction::SetContext { set_context } => {
7600                    for (key, value) in set_context {
7601                        if let Err(e) = self.context_manager.set(key, value.clone()) {
7602                            warn!(key = %key, error = %e, "State action: set_context failed");
7603                        } else {
7604                            debug!(key = %key, "State action: context set");
7605                        }
7606                    }
7607                }
7608                StateAction::Prompt {
7609                    prompt,
7610                    llm,
7611                    store_as,
7612                } => {
7613                    let llm_result = if let Some(alias) = llm {
7614                        self.llm_registry.get(alias)
7615                    } else {
7616                        self.llm_registry.default()
7617                    };
7618                    match llm_result {
7619                        Ok(llm_provider) => {
7620                            // Render template variables and include conversation context
7621                            let context = self.build_context_with_overlays();
7622                            let rendered_prompt = self
7623                                .template_renderer
7624                                .render(prompt, &context)
7625                                .unwrap_or_else(|_| prompt.clone());
7626                            let recent =
7627                                self.memory.get_messages(Some(5)).await.unwrap_or_default();
7628                            let mut messages: Vec<ChatMessage> = recent;
7629                            messages.push(ChatMessage::user(&rendered_prompt));
7630                            match self
7631                                .observe_purpose(
7632                                    ObservationPurpose::StateAction,
7633                                    llm_provider.complete(&messages, None),
7634                                )
7635                                .await
7636                            {
7637                                Ok(response) => {
7638                                    if let Some(key) = store_as {
7639                                        let _ = self
7640                                            .context_manager
7641                                            .set(key, Value::String(response.content));
7642                                        debug!(key = %key, "State action: prompt result stored");
7643                                    }
7644                                }
7645                                Err(e) => {
7646                                    warn!(error = %e, "State action: prompt LLM call failed");
7647                                }
7648                            }
7649                        }
7650                        Err(e) => {
7651                            warn!(error = %e, "State action: LLM not found for prompt");
7652                        }
7653                    }
7654                }
7655            }
7656        }
7657    }
7658
7659    async fn run_context_extractors_staged(&self, user_message: &str) -> HashMap<String, Value> {
7660        let extractors = match &self.state_machine {
7661            Some(sm) => match sm.current_definition() {
7662                Some(def) if !def.extract.is_empty() => def.extract.clone(),
7663                _ => return HashMap::new(),
7664            },
7665            None => return HashMap::new(),
7666        };
7667
7668        let mut staged = HashMap::new();
7669        for extractor in &extractors {
7670            let prompt = if let Some(ref custom) = extractor.llm_extract {
7671                format!(
7672                    "User message:\n\"{}\"\n\nInstruction:\n{}",
7673                    user_message, custom
7674                )
7675            } else if let Some(ref desc) = extractor.description {
7676                format!(
7677                    "From the following message, extract: {}\n\n\
7678                     Message: \"{}\"\n\n\
7679                     If the information is present, return ONLY the extracted value.\n\
7680                     If NOT present, return exactly: __NONE__",
7681                    desc, user_message
7682                )
7683            } else {
7684                continue;
7685            };
7686
7687            let llm = match self
7688                .llm_registry
7689                .get(&extractor.llm)
7690                .or_else(|_| self.llm_registry.get("router"))
7691                .or_else(|_| self.llm_registry.get("default"))
7692            {
7693                Ok(llm) => llm,
7694                Err(e) => {
7695                    warn!(key = %extractor.key, error = %e, "Extractor LLM not found");
7696                    continue;
7697                }
7698            };
7699
7700            let messages = vec![ChatMessage::user(&prompt)];
7701            match self
7702                .observe_purpose(
7703                    ObservationPurpose::ContextExtraction,
7704                    llm.complete(&messages, None),
7705                )
7706                .await
7707            {
7708                Ok(response) => {
7709                    let value = response.content.trim().to_string();
7710                    if value != "__NONE__" && !value.is_empty() {
7711                        staged.insert(
7712                            extractor.key.clone(),
7713                            serde_json::Value::String(value.clone()),
7714                        );
7715                        debug!(key = %extractor.key, value = %value, "Context extracted");
7716                    } else if extractor.required {
7717                        warn!(key = %extractor.key, "Required extraction returned no value");
7718                    }
7719                }
7720                Err(e) => {
7721                    warn!(key = %extractor.key, error = %e, "Context extraction LLM call failed");
7722                }
7723            }
7724        }
7725        staged
7726    }
7727
7728    fn commit_staged_context_writes(&self, staged: &HashMap<String, Value>) {
7729        for (key, value) in staged {
7730            if let Err(error) = self.context_manager.update(key, value.clone()) {
7731                warn!(key = %key, error = %error, "staged context write failed");
7732            }
7733        }
7734    }
7735
7736    /// Run context extractors for the current state on the user's input.
7737    async fn run_context_extractors(&self, user_message: &str) {
7738        let staged = self.run_context_extractors_staged(user_message).await;
7739        self.commit_staged_context_writes(&staged);
7740    }
7741
7742    async fn check_memory_compression(&self) -> Result<()> {
7743        if self.memory.needs_compression() {
7744            let result = self.memory.compress(None).await?;
7745            if let CompressResult::Compressed {
7746                messages_summarized,
7747                new_summary_length,
7748                tokens_saved,
7749            } = result
7750            {
7751                let event = MemoryCompressEvent::new(
7752                    messages_summarized,
7753                    tokens_saved,
7754                    new_summary_length as u32,
7755                );
7756                self.hooks.on_memory_compress(&event).await;
7757                debug!(
7758                    messages = messages_summarized,
7759                    tokens_saved = tokens_saved,
7760                    "Memory compressed"
7761                );
7762            }
7763        }
7764
7765        // Handle overflow AFTER compression, then check warning threshold
7766        self.handle_memory_overflow().await?;
7767        self.check_memory_budget().await;
7768
7769        Ok(())
7770    }
7771
7772    async fn check_memory_budget(&self) {
7773        let Some(ref budget) = self.memory_token_budget else {
7774            return;
7775        };
7776
7777        let context = match self.memory.get_context().await {
7778            Ok(ctx) => ctx,
7779            Err(_) => return,
7780        };
7781
7782        // Overall budget warning
7783        let used_tokens = context.estimated_tokens();
7784        if budget.is_over_warn_threshold(used_tokens) {
7785            let event = MemoryBudgetEvent::new("memory", used_tokens, budget.total);
7786            self.hooks.on_memory_budget_warning(&event).await;
7787            debug!(
7788                used = used_tokens,
7789                total = budget.total,
7790                percent = event.usage_percent,
7791                "Memory budget warning"
7792            );
7793        }
7794
7795        // Per-component warning: summary
7796        if let Some(ref summary) = context.summary {
7797            let summary_tokens = ai_agents_memory::estimate_tokens(summary);
7798            let summary_budget = budget.allocation.summary;
7799            if summary_budget > 0 {
7800                let warn_threshold =
7801                    (summary_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
7802                if summary_tokens >= warn_threshold {
7803                    let event = MemoryBudgetEvent::new("summary", summary_tokens, summary_budget);
7804                    self.hooks.on_memory_budget_warning(&event).await;
7805                }
7806            }
7807        }
7808
7809        // Per-component warning: recent_messages
7810        let recent_tokens: u32 = context
7811            .messages
7812            .iter()
7813            .map(ai_agents_memory::estimate_message_tokens)
7814            .sum();
7815        let recent_budget = budget.allocation.recent_messages;
7816        if recent_budget > 0 {
7817            let warn_threshold =
7818                (recent_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
7819            if recent_tokens >= warn_threshold {
7820                let event = MemoryBudgetEvent::new("recent_messages", recent_tokens, recent_budget);
7821                self.hooks.on_memory_budget_warning(&event).await;
7822            }
7823        }
7824
7825        let relationship_budget = budget.allocation.relationships;
7826        if relationship_budget > 0 {
7827            let relationship_tokens = self
7828                .relationship_memory_text()
7829                .map(|text| ai_agents_memory::estimate_tokens(&text))
7830                .unwrap_or(0);
7831            let warn_threshold =
7832                (relationship_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
7833            if relationship_tokens >= warn_threshold {
7834                let event = MemoryBudgetEvent::new(
7835                    "relationships",
7836                    relationship_tokens,
7837                    relationship_budget,
7838                );
7839                self.hooks.on_memory_budget_warning(&event).await;
7840            }
7841        }
7842    }
7843
7844    async fn handle_memory_overflow(&self) -> Result<()> {
7845        let Some(ref budget) = self.memory_token_budget else {
7846            return Ok(());
7847        };
7848
7849        let context = self.memory.get_context().await?;
7850        let used_tokens = context.estimated_tokens();
7851
7852        if used_tokens <= budget.total {
7853            return Ok(());
7854        }
7855
7856        match budget.overflow_strategy {
7857            OverflowStrategy::TruncateOldest => {
7858                let tokens_to_free = used_tokens - budget.total;
7859                let messages_to_evict = self.calculate_eviction_count(tokens_to_free);
7860                if messages_to_evict > 0 {
7861                    self.evict_messages(messages_to_evict, EvictionReason::TokenBudgetExceeded)
7862                        .await?;
7863                }
7864            }
7865            OverflowStrategy::SummarizeMore => {
7866                let max_attempts = context.total_messages.max(1);
7867                for _ in 0..max_attempts {
7868                    match self.memory.compress(None).await? {
7869                        CompressResult::Compressed {
7870                            messages_summarized,
7871                            ..
7872                        } if messages_summarized > 0 => {
7873                            let context = self.memory.get_context().await?;
7874                            if context.estimated_tokens() <= budget.total {
7875                                return Ok(());
7876                            }
7877                        }
7878                        _ => break,
7879                    }
7880                }
7881                let context = self.memory.get_context().await?;
7882                let used_tokens = context.estimated_tokens();
7883                if used_tokens > budget.total {
7884                    return Err(AgentError::MemoryBudgetExceeded {
7885                        used: used_tokens,
7886                        budget: budget.total,
7887                    });
7888                }
7889            }
7890            OverflowStrategy::Error => {
7891                return Err(AgentError::MemoryBudgetExceeded {
7892                    used: used_tokens,
7893                    budget: budget.total,
7894                });
7895            }
7896        }
7897        Ok(())
7898    }
7899
7900    fn calculate_eviction_count(&self, tokens_to_free: u32) -> usize {
7901        // Estimate ~50 tokens per message on average
7902        ((tokens_to_free as f64 / 50.0).ceil() as usize).max(1)
7903    }
7904
7905    async fn evict_messages(&self, count: usize, reason: EvictionReason) -> Result<()> {
7906        let evicted = self.memory.evict_oldest(count).await?;
7907        if !evicted.is_empty() {
7908            let event = MemoryEvictEvent {
7909                reason,
7910                messages_evicted: evicted.len(),
7911                importance_scores: vec![],
7912            };
7913            self.hooks.on_memory_evict(&event).await;
7914            debug!(count = evicted.len(), "Messages evicted from memory");
7915        }
7916        Ok(())
7917    }
7918
7919    #[instrument(skip(self, input), fields(agent = %self.info.name))]
7920    async fn determine_reasoning_mode(&self, input: &str) -> Result<ReasoningMode> {
7921        match self.determine_reasoning_mode_strict(input).await {
7922            Ok(mode) => Ok(mode),
7923            Err(_) => Ok(ReasoningMode::None),
7924        }
7925    }
7926
7927    async fn determine_reasoning_mode_strict(&self, input: &str) -> Result<ReasoningMode> {
7928        let effective_config = self.get_effective_reasoning_config();
7929
7930        if !matches!(effective_config.mode, ReasoningMode::Auto) {
7931            return Ok(effective_config.mode.clone());
7932        }
7933
7934        let judge_llm = effective_config
7935            .judge_llm
7936            .as_ref()
7937            .and_then(|alias| self.llm_registry.get(alias).ok())
7938            .or_else(|| self.llm_registry.router().ok())
7939            .or_else(|| self.llm_registry.default().ok());
7940
7941        let Some(llm) = judge_llm else {
7942            return Ok(ReasoningMode::None);
7943        };
7944
7945        let prompt = format!(
7946            r#"Analyze this user request and determine the appropriate reasoning mode.
7947
7948User request: "{}"
7949
7950Choose ONE of these modes:
7951- none: Simple queries, greetings, direct answers (fastest)
7952- cot: Complex analysis, multi-step reasoning, math problems
7953- react: Tasks requiring multiple tool calls with observation
7954- plan_and_execute: Complex multi-step tasks requiring coordination
7955
7956Respond with ONLY the mode name (none, cot, react, or plan_and_execute)."#,
7957            input
7958        );
7959
7960        let messages = vec![ChatMessage::user(&prompt)];
7961        let response = self
7962            .observe_purpose(
7963                ObservationPurpose::ReflectionDecision,
7964                llm.complete(&messages, None),
7965            )
7966            .await
7967            .map_err(|e| AgentError::LLM(e.to_string()))?;
7968
7969        let mode_str = response.content.trim().to_lowercase();
7970        Ok(match mode_str.as_str() {
7971            "cot" => ReasoningMode::CoT,
7972            "react" => ReasoningMode::React,
7973            "plan_and_execute" => ReasoningMode::PlanAndExecute,
7974            _ => ReasoningMode::None,
7975        })
7976    }
7977
7978    async fn should_reflect(&self, input: &str, response: &str) -> Result<bool> {
7979        let effective_config = self.get_effective_reflection_config();
7980
7981        if !effective_config.requires_evaluation() {
7982            return Ok(false);
7983        }
7984
7985        if effective_config.is_enabled() {
7986            return Ok(true);
7987        }
7988
7989        let evaluator_llm = effective_config
7990            .evaluator_llm
7991            .as_ref()
7992            .and_then(|alias| self.llm_registry.get(alias).ok())
7993            .or_else(|| self.llm_registry.router().ok())
7994            .or_else(|| self.llm_registry.default().ok());
7995
7996        let Some(llm) = evaluator_llm else {
7997            return Ok(false);
7998        };
7999
8000        let response_preview: String = response.chars().take(500).collect();
8001        let prompt = format!(
8002            r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
8003
8004User query: "{}"
8005Response: "{}"
8006
8007Answer YES or NO only."#,
8008            input, response_preview
8009        );
8010
8011        let messages = vec![ChatMessage::user(&prompt)];
8012        let result = self
8013            .observe_purpose(
8014                ObservationPurpose::ReflectionDecision,
8015                llm.complete(&messages, None),
8016            )
8017            .await;
8018
8019        match result {
8020            Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
8021            Err(_) => Ok(false),
8022        }
8023    }
8024
8025    fn build_cot_system_prompt(&self, base_prompt: &str) -> String {
8026        format!(
8027            "{}\n\n<instruction>\nThink through this step by step before answering:\n1. Understand what is being asked\n2. Break down the problem\n3. Work through each part\n4. Provide your final answer\n\nShow your thinking process, then give your final answer.\n</instruction>",
8028            base_prompt
8029        )
8030    }
8031
8032    fn build_react_system_prompt(&self, base_prompt: &str) -> String {
8033        format!(
8034            "{}\n\n<instruction>\nUse the Reason-Act-Observe pattern:\n1. Thought: Think about what to do\n2. Action: Use a tool if needed\n3. Observation: Analyze the result\n4. Repeat until you have the answer\n\nFormat your response showing Thought/Action/Observation steps.\n</instruction>",
8035            base_prompt
8036        )
8037    }
8038
8039    async fn generate_plan(&self, input: &str) -> Result<Plan> {
8040        let effective = self.get_effective_reasoning_config();
8041        let planning_config = effective.get_planning();
8042
8043        let planner_llm = planning_config
8044            .and_then(|c| c.planner_llm.as_ref())
8045            .and_then(|alias| self.llm_registry.get(alias).ok())
8046            .or_else(|| self.llm_registry.router().ok())
8047            .or_else(|| self.llm_registry.default().ok())
8048            .ok_or_else(|| AgentError::Config("No LLM available for planning".into()))?;
8049
8050        let mut available_tool_ids: Vec<String> = self
8051            .get_available_tool_ids()
8052            .await
8053            .unwrap_or_else(|_| self.tools.list_ids());
8054        let mut available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
8055
8056        // Apply planning-level tool and skill filters.
8057        if let Some(config) = planning_config {
8058            if !config.available.tools.is_all() {
8059                available_tool_ids.retain(|t| config.available.tools.allows(t));
8060            }
8061            if !config.available.skills.is_all() {
8062                available_skills.retain(|s| config.available.skills.allows(s));
8063            }
8064        }
8065
8066        // Build tool descriptions with argument schemas so the planner
8067        // knows how to construct valid args for each step.
8068        let tool_descriptions: Vec<String> = available_tool_ids
8069            .iter()
8070            .filter_map(|id| {
8071                self.tools.get(id).map(|tool| {
8072                    let schema = tool.input_schema();
8073                    let args_desc = schema
8074                        .get("properties")
8075                        .and_then(|p| serde_json::to_string(p).ok())
8076                        .unwrap_or_else(|| "{}".to_string());
8077                    format!(
8078                        "- {} ({}): {}\n  Arguments: {}",
8079                        id,
8080                        tool.name(),
8081                        tool.description(),
8082                        args_desc
8083                    )
8084                })
8085            })
8086            .collect();
8087
8088        let tools_section = if tool_descriptions.is_empty() {
8089            "Available tools: none".to_string()
8090        } else {
8091            format!("Available tools:\n{}", tool_descriptions.join("\n"))
8092        };
8093
8094        let skills_section = if available_skills.is_empty() {
8095            "Available skills: none".to_string()
8096        } else {
8097            format!("Available skills: {}", available_skills.join(", "))
8098        };
8099
8100        let prompt = format!(
8101            r#"Create a step-by-step plan to accomplish this goal.
8102
8103Goal: "{}"
8104
8105{}
8106
8107{}
8108
8109Create a plan with clear steps. For each step, specify:
8110- description: What this step accomplishes
8111- action_type: "tool", "skill", "think", or "respond"
8112- action_target: The tool/skill id (if applicable)
8113- args: The arguments object matching the tool's schema (if action_type is "tool")
8114- dependencies: List of step IDs this depends on (empty if none)
8115
8116Respond in JSON format:
8117{{
8118  "steps": [
8119    {{"id": "step1", "description": "...", "action_type": "tool", "action_target": "tool_id", "args": {{"required_field": "value"}}, "dependencies": []}},
8120    {{"id": "step2", "description": "...", "action_type": "think", "action_target": "...", "dependencies": ["step1"]}}
8121  ]
8122}}"#,
8123            input, tools_section, skills_section,
8124        );
8125
8126        let messages = vec![ChatMessage::user(&prompt)];
8127        let response = self
8128            .observe_purpose(
8129                ObservationPurpose::PlanGeneration,
8130                planner_llm.complete(&messages, None),
8131            )
8132            .await
8133            .map_err(|e| AgentError::LLM(format!("Planning failed: {}", e)))?;
8134
8135        let mut plan = Plan::new(input);
8136
8137        if let Some(json_start) = response.content.find('{')
8138            && let Some(json_end) = response.content.rfind('}')
8139        {
8140            let json_str = &response.content[json_start..=json_end];
8141            if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(json_str)
8142                && let Some(steps) = parsed.get("steps").and_then(|s| s.as_array())
8143            {
8144                for step_value in steps {
8145                    let id = step_value
8146                        .get("id")
8147                        .and_then(|v| v.as_str())
8148                        .unwrap_or("step");
8149                    let desc = step_value
8150                        .get("description")
8151                        .and_then(|v| v.as_str())
8152                        .unwrap_or("");
8153                    let action_type = step_value
8154                        .get("action_type")
8155                        .and_then(|v| v.as_str())
8156                        .unwrap_or("think");
8157                    let action_target = step_value
8158                        .get("action_target")
8159                        .and_then(|v| v.as_str())
8160                        .unwrap_or("");
8161                    let args = step_value
8162                        .get("args")
8163                        .cloned()
8164                        .unwrap_or(serde_json::json!({}));
8165                    let deps: Vec<String> = step_value
8166                        .get("dependencies")
8167                        .and_then(|v| v.as_array())
8168                        .map(|arr| {
8169                            arr.iter()
8170                                .filter_map(|v| v.as_str().map(String::from))
8171                                .collect()
8172                        })
8173                        .unwrap_or_default();
8174
8175                    let action = match action_type {
8176                        "tool" => PlanAction::tool(action_target, args),
8177                        "skill" => PlanAction::skill(action_target),
8178                        "respond" => PlanAction::respond(action_target),
8179                        _ => PlanAction::think(desc),
8180                    };
8181
8182                    let step = PlanStep::new(desc, action)
8183                        .with_id(id)
8184                        .with_dependencies(deps);
8185                    plan.add_step(step);
8186                }
8187            }
8188        }
8189
8190        if plan.steps.is_empty() {
8191            plan.add_step(PlanStep::new(
8192                "Process the request",
8193                PlanAction::think(input),
8194            ));
8195            plan.add_step(PlanStep::new(
8196                "Provide response",
8197                PlanAction::respond("Answer based on analysis"),
8198            ));
8199        }
8200
8201        Ok(plan)
8202    }
8203
8204    async fn execute_plan(&self, plan: &mut Plan) -> Result<String> {
8205        let llm = self.get_state_llm()?;
8206        let mut results: HashMap<String, serde_json::Value> = HashMap::new();
8207        let effective = self.get_effective_reasoning_config();
8208        let max_steps = effective.get_planning().map(|c| c.max_steps).unwrap_or(10);
8209
8210        plan.status = PlanStatus::InProgress;
8211
8212        for step_idx in 0..plan.steps.len().min(max_steps as usize) {
8213            let step = &plan.steps[step_idx];
8214
8215            let deps_satisfied = step.dependencies.iter().all(|dep| {
8216                plan.steps
8217                    .iter()
8218                    .find(|s| &s.id == dep)
8219                    .map(|s| s.status.is_completed())
8220                    .unwrap_or(false)
8221            });
8222
8223            if !deps_satisfied {
8224                continue;
8225            }
8226
8227            plan.steps[step_idx].mark_running();
8228
8229            let result = match &plan.steps[step_idx].action {
8230                PlanAction::Tool { tool, args } => {
8231                    // When a tool step has dependency results, ask the LLM to
8232                    // produce the correct arguments given the context and tool schema.
8233                    // This avoids brittle {{stepN}} template substitution and lets
8234                    // the LLM handle type adaptation (e.g. picking the iso field
8235                    // from a datetime result for a downstream format call).
8236                    let has_dep_results = plan.steps[step_idx]
8237                        .dependencies
8238                        .iter()
8239                        .any(|dep| results.contains_key(dep));
8240
8241                    let final_args = if has_dep_results {
8242                        let dep_context: String = plan.steps[step_idx]
8243                            .dependencies
8244                            .iter()
8245                            .filter_map(|dep| results.get(dep).map(|r| format!("{}: {}", dep, r)))
8246                            .collect::<Vec<_>>()
8247                            .join("\n");
8248
8249                        let tool_schema = self
8250                            .tools
8251                            .get(tool)
8252                            .map(|t| {
8253                                let schema = t.input_schema();
8254                                let props = schema
8255                                    .get("properties")
8256                                    .and_then(|p| serde_json::to_string(p).ok())
8257                                    .unwrap_or_else(|| "{}".to_string());
8258                                format!(
8259                                    "{}: {}\nArguments schema: {}",
8260                                    t.id(),
8261                                    t.description(),
8262                                    props
8263                                )
8264                            })
8265                            .unwrap_or_default();
8266
8267                        let step_desc = &plan.steps[step_idx].description;
8268                        let arg_prompt = format!(
8269                            "Generate the JSON arguments for a tool call.\n\n\
8270                             Tool: {}\n\n\
8271                             Task: {}\n\n\
8272                             Previous step results:\n{}\n\n\
8273                             Planner's draft arguments: {}\n\n\
8274                             Produce ONLY a valid JSON object with the correct argument values.\n\
8275                             Use actual values from the previous step results, not template references.",
8276                            tool_schema,
8277                            step_desc,
8278                            dep_context,
8279                            serde_json::to_string(args).unwrap_or_default()
8280                        );
8281                        let messages = vec![ChatMessage::user(&arg_prompt)];
8282                        match self
8283                            .observe_purpose(
8284                                ObservationPurpose::PlanStep,
8285                                llm.complete(&messages, None),
8286                            )
8287                            .await
8288                        {
8289                            Ok(resp) => {
8290                                let content = resp.content.trim();
8291                                // Parse the LLM's JSON response, fall back to planner args.
8292                                let json_start = content.find('{');
8293                                let json_end = content.rfind('}');
8294                                if let (Some(start), Some(end)) = (json_start, json_end) {
8295                                    serde_json::from_str(&content[start..=end])
8296                                        .unwrap_or_else(|_| args.clone())
8297                                } else {
8298                                    args.clone()
8299                                }
8300                            }
8301                            Err(_) => args.clone(),
8302                        }
8303                    } else {
8304                        args.clone()
8305                    };
8306
8307                    let request = ToolExecutionRequest::new(
8308                        uuid::Uuid::new_v4().to_string(),
8309                        tool.clone(),
8310                        final_args,
8311                        ToolCallSource::Plan {
8312                            step_index: step_idx,
8313                        },
8314                    );
8315                    match self.execute_tool_record(request).await {
8316                        Ok(record) if record.success => {
8317                            serde_json::json!({ "output": record.model_output_string() })
8318                        }
8319                        Ok(record) => {
8320                            plan.steps[step_idx].mark_failed(record.model_output_string());
8321                            continue;
8322                        }
8323                        Err(e) => {
8324                            plan.steps[step_idx].mark_failed(e.to_string());
8325                            continue;
8326                        }
8327                    }
8328                }
8329                PlanAction::Skill { skill } => {
8330                    if let Some(skill_def) = self.skills.iter().find(|s| &s.id == skill) {
8331                        if let Some(ref executor) = self.skill_executor {
8332                            match executor
8333                                .execute_with_invoker(skill_def, "", serde_json::json!({}), self)
8334                                .await
8335                            {
8336                                Ok(output) => serde_json::json!({ "output": output }),
8337                                Err(e) => {
8338                                    plan.steps[step_idx].mark_failed(e.to_string());
8339                                    continue;
8340                                }
8341                            }
8342                        } else {
8343                            serde_json::json!({ "output": "Skill executor not available" })
8344                        }
8345                    } else {
8346                        plan.steps[step_idx].mark_failed("Skill not found");
8347                        continue;
8348                    }
8349                }
8350                PlanAction::Think { prompt } => {
8351                    let context: String = results
8352                        .iter()
8353                        .map(|(k, v)| format!("{}: {}", k, v))
8354                        .collect::<Vec<_>>()
8355                        .join("\n");
8356
8357                    let think_prompt = format!("Context:\n{}\n\nTask: {}", context, prompt);
8358                    let messages = vec![ChatMessage::user(&think_prompt)];
8359
8360                    match self
8361                        .observe_purpose(
8362                            ObservationPurpose::PlanStep,
8363                            llm.complete(&messages, None),
8364                        )
8365                        .await
8366                    {
8367                        Ok(resp) => serde_json::json!({ "output": resp.content }),
8368                        Err(e) => {
8369                            plan.steps[step_idx].mark_failed(e.to_string());
8370                            continue;
8371                        }
8372                    }
8373                }
8374                PlanAction::Respond { template } => {
8375                    let context: String = results
8376                        .iter()
8377                        .map(|(k, v)| format!("{}: {}", k, v))
8378                        .collect::<Vec<_>>()
8379                        .join("\n");
8380
8381                    let respond_prompt = format!(
8382                        "Based on this context:\n{}\n\nGenerate a response following this template/instruction: {}",
8383                        context, template
8384                    );
8385                    let messages = vec![ChatMessage::user(&respond_prompt)];
8386
8387                    match self
8388                        .observe_purpose(
8389                            ObservationPurpose::PlanStep,
8390                            llm.complete(&messages, None),
8391                        )
8392                        .await
8393                    {
8394                        Ok(resp) => serde_json::json!({ "output": resp.content }),
8395                        Err(e) => {
8396                            plan.steps[step_idx].mark_failed(e.to_string());
8397                            continue;
8398                        }
8399                    }
8400                }
8401            };
8402
8403            results.insert(plan.steps[step_idx].id.clone(), result.clone());
8404            plan.steps[step_idx].mark_completed(Some(result));
8405        }
8406
8407        // Set plan status based on whether any steps actually failed.
8408        let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
8409        if has_failures {
8410            let failed_ids: Vec<String> = plan
8411                .steps
8412                .iter()
8413                .filter(|s| s.status.is_failed())
8414                .map(|s| s.id.clone())
8415                .collect();
8416            plan.status = PlanStatus::Failed {
8417                error: format!("Steps failed: {}", failed_ids.join(", ")),
8418            };
8419        } else {
8420            plan.status = PlanStatus::Completed;
8421        }
8422
8423        // Synthesize final output from all completed step results.
8424        let all_outputs: Vec<String> = plan
8425            .steps
8426            .iter()
8427            .filter(|s| s.status.is_completed())
8428            .filter_map(|s| {
8429                s.result
8430                    .as_ref()
8431                    .and_then(|r| r.get("output"))
8432                    .and_then(|o| o.as_str())
8433                    .map(|o| format!("{}: {}", s.description, o))
8434            })
8435            .collect();
8436
8437        if all_outputs.is_empty() {
8438            return Ok("Plan execution completed but produced no results.".to_string());
8439        }
8440
8441        if all_outputs.len() == 1 {
8442            return Ok(all_outputs.into_iter().next().unwrap());
8443        }
8444
8445        // Synthesize a coherent summary from multiple step results via LLM.
8446        let context = all_outputs.join("\n\n");
8447        let prompt = format!(
8448            "You completed a multi-step plan for: \"{}\"\n\nStep results:\n{}\n\nProvide a coherent final response that synthesizes these results.",
8449            plan.goal, context
8450        );
8451        let messages = vec![ChatMessage::user(&prompt)];
8452        match self
8453            .observe_purpose(ObservationPurpose::PlanStep, llm.complete(&messages, None))
8454            .await
8455        {
8456            Ok(resp) => Ok(resp.content.trim().to_string()),
8457            Err(_) => Ok(context),
8458        }
8459    }
8460
8461    async fn evaluate_response(&self, input: &str, response: &str) -> Result<EvaluationResult> {
8462        let effective_config = self.get_effective_reflection_config();
8463        self.evaluate_response_with_config(input, response, &effective_config)
8464            .await
8465    }
8466
8467    fn extract_thinking(&self, content: &str) -> (Option<String>, String) {
8468        if let Some(start) = content.find("<thinking>")
8469            && let Some(end) = content.find("</thinking>")
8470        {
8471            let thinking = content[start + 10..end].trim().to_string();
8472            let answer = content[end + 11..].trim().to_string();
8473            return (Some(thinking), answer);
8474        }
8475        (None, content.to_string())
8476    }
8477
8478    fn format_response_with_thinking(&self, thinking: Option<&str>, answer: &str) -> String {
8479        match self.get_effective_reasoning_config().output {
8480            ReasoningOutput::Hidden => answer.to_string(),
8481            ReasoningOutput::Visible => {
8482                if let Some(t) = thinking {
8483                    format!("Thinking:\n{}\n\nAnswer:\n{}", t, answer)
8484                } else {
8485                    answer.to_string()
8486                }
8487            }
8488            ReasoningOutput::Tagged => {
8489                if let Some(t) = thinking {
8490                    format!("<thinking>{}</thinking>\n{}", t, answer)
8491                } else {
8492                    answer.to_string()
8493                }
8494            }
8495        }
8496    }
8497
8498    //
8499    // Blocking disambiguation returns every clarification or required confirmation as its own root response before redispatch.
8500    // The manager retains resolved input and pending skill ownership until a later turn explicitly confirms it.
8501    //
8502    async fn run_loop(&self, input: &str) -> Result<AgentResponse> {
8503        //
8504        // Blocking execution must fail before a turn starts when required persistence is unavailable.
8505        //
8506        self.init_storage().await?;
8507        self.begin_root_turn();
8508        let _root_cleanup = RootTurnCleanup::new(self);
8509        info!(input_len = input.len(), "Starting chat");
8510
8511        self.hooks.on_message_received(input).await;
8512
8513        // One-shot context initialization: load runtime defaults, resolve env vars,
8514        // populate builtin sources (session, agent), etc.  This must happen before
8515        // the first template render so that {{ context.* }} variables are available.
8516        if !self.context_initialized.swap(true, Ordering::SeqCst) {
8517            self.context_manager.initialize().await?;
8518            debug!("Context manager initialized (defaults, env, builtins)");
8519        }
8520
8521        self.check_turn_timeout().await?;
8522        self.context_manager.refresh_per_turn().await?;
8523
8524        // Clear stale disambiguation context from previous turns.
8525        // This prevents resolved_intent from leaking across turns and causing incorrect deterministic routing on subsequent inputs.
8526        self.clear_disambiguation_context();
8527
8528        // Disambiguation check (before input processing)
8529        if let Some(ref disambiguator) = self.disambiguation_manager {
8530            let disambiguation_context = self.build_disambiguation_context().await?;
8531
8532            // Get state-level disambiguation override
8533            let state_override = self
8534                .state_machine
8535                .as_ref()
8536                .and_then(|sm| sm.current_definition())
8537                .and_then(|def| def.disambiguation.clone());
8538
8539            let state_generation = self
8540                .state_machine
8541                .as_ref()
8542                .map(|state_machine| state_machine.generation());
8543            let disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
8544            let mut disambiguation_result = self
8545                .observe_purpose(
8546                    ObservationPurpose::DisambiguationDetection,
8547                    disambiguator.process_input_with_override(
8548                        input,
8549                        &disambiguation_context,
8550                        state_override.as_ref(),
8551                        None,
8552                    ),
8553                )
8554                .await?;
8555            let current_state_generation = self
8556                .state_machine
8557                .as_ref()
8558                .map(|state_machine| state_machine.generation());
8559            if current_state_generation != state_generation
8560                || self.disambiguation_epoch.load(Ordering::SeqCst) != disambiguation_epoch
8561            {
8562                disambiguator.clear_pending().await;
8563                *self.pending_skill_id.write() = None;
8564                disambiguation_result = DisambiguationResult::Abandoned { new_input: None };
8565                info!(
8566                    confirmation_event = "invalidated",
8567                    invalidation_reason = "state_generation_changed",
8568                    "Disambiguation result invalidated before redispatch"
8569                );
8570            }
8571            match disambiguation_result {
8572                DisambiguationResult::Clear => {
8573                    debug!("Input is clear, proceeding normally");
8574                }
8575                DisambiguationResult::NeedsClarification {
8576                    question,
8577                    detection,
8578                } => {
8579                    let admission = self
8580                        .admit_disambiguation_redispatch(disambiguation_epoch, state_generation)
8581                        .await?;
8582                    let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
8583                    info!(
8584                        ambiguity_type = ?detection.ambiguity_type,
8585                        confidence = detection.confidence,
8586                        "Input requires clarification"
8587                    );
8588
8589                    // This branch also owns post-resolution confirmation questions.
8590                    // Do not clear pending_skill_id or redispatch until the manager returns Clarified on a later turn.
8591                    self.commit_root_user_message(input).await?;
8592                    self.memory
8593                        .add_message(ChatMessage::assistant(&question.question))
8594                        .await?;
8595
8596                    let status = if awaiting_confirmation {
8597                        "awaiting_confirmation"
8598                    } else {
8599                        "awaiting_clarification"
8600                    };
8601                    let response = AgentResponse::new(&question.question).with_metadata(
8602                        "disambiguation",
8603                        serde_json::json!({
8604                            "status": status,
8605                            "options": question.options,
8606                            "clarifying": question.clarifying,
8607                            "detection": {
8608                                "type": detection.ambiguity_type,
8609                                "confidence": detection.confidence,
8610                                "what_is_unclear": detection.what_is_unclear,
8611                            }
8612                        }),
8613                    );
8614                    drop(admission);
8615                    self.finish_turn_if_root(&response).await?;
8616                    return Ok(response);
8617                }
8618                DisambiguationResult::Clarified {
8619                    enriched_input,
8620                    resolved,
8621                    ..
8622                } => {
8623                    let admission = match self
8624                        .admit_disambiguation_redispatch(disambiguation_epoch, state_generation)
8625                        .await
8626                    {
8627                        Ok(admission) => admission,
8628                        Err(error) => {
8629                            *self.pending_skill_id.write() = None;
8630                            return Err(error);
8631                        }
8632                    };
8633                    info!(
8634                        resolved_count = resolved.len(),
8635                        enriched = %enriched_input,
8636                        "Input clarified, injecting resolved intent into context"
8637                    );
8638
8639                    // Routing uses `resolved` (structured, deterministic)
8640                    // This is what makes post-disambiguation routing DETERMINISTIC
8641                    for (key, value) in &resolved {
8642                        let context_key = format!("disambiguation.{}", key);
8643                        let _ = self.context_manager.set(&context_key, value.clone());
8644                    }
8645
8646                    if let Some(intent) = resolved.get("intent") {
8647                        let _ = self.context_manager.set("resolved_intent", intent.clone());
8648                    }
8649
8650                    let _ = self
8651                        .context_manager
8652                        .set("disambiguation.resolved", serde_json::Value::Bool(true));
8653
8654                    // Check if this clarification was triggered by a skill-level override.
8655                    // If so, route directly to the matched skill instead of going through
8656                    // skill routing again (which might match a different skill).
8657                    let skill_id = self.pending_skill_id.read().clone();
8658                    if let Some(skill_id) = skill_id {
8659                        info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input");
8660                        drop(admission);
8661                        return self
8662                            .recheck_skill_disambiguation(
8663                                &skill_id,
8664                                &enriched_input,
8665                                disambiguation_epoch,
8666                                state_generation,
8667                            )
8668                            .await;
8669                    }
8670
8671                    drop(admission);
8672                    return self.run_loop_internal(&enriched_input).await;
8673                }
8674                DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
8675                    info!("Proceeding with best guess interpretation");
8676
8677                    // Same skill-id re-check for best-guess path
8678                    let skill_id = self.pending_skill_id.read().clone();
8679                    if let Some(skill_id) = skill_id {
8680                        info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input");
8681                        return self
8682                            .recheck_skill_disambiguation(
8683                                &skill_id,
8684                                &enriched_input,
8685                                disambiguation_epoch,
8686                                state_generation,
8687                            )
8688                            .await;
8689                    }
8690
8691                    return self.run_loop_internal(&enriched_input).await;
8692                }
8693                DisambiguationResult::GiveUp { reason } => {
8694                    *self.pending_skill_id.write() = None;
8695                    warn!(reason = %reason, "Disambiguation gave up");
8696                    let apology = self
8697                        .generate_localized_apology(
8698                            "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
8699                            &reason,
8700                        )
8701                        .await
8702                        .unwrap_or_else(|_| {
8703                            format!("I'm sorry, I couldn't understand your request: {}", reason)
8704                        });
8705                    let response = AgentResponse::new(&apology);
8706                    self.finish_turn_if_root(&response).await?;
8707                    return Ok(response);
8708                }
8709                DisambiguationResult::Escalate { reason } => {
8710                    *self.pending_skill_id.write() = None;
8711                    info!(reason = %reason, "Escalating to human");
8712                    if let Some(ref hitl) = self.hitl_engine {
8713                        let trigger =
8714                            ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
8715                        let mut context_map = HashMap::new();
8716                        context_map.insert("original_input".to_string(), serde_json::json!(input));
8717                        context_map.insert("reason".to_string(), serde_json::json!(&reason));
8718                        let check_result = HITLCheckResult::required(
8719                            trigger,
8720                            context_map,
8721                            format!("User request needs human assistance: {}", reason),
8722                            Some(hitl.config().default_timeout_seconds),
8723                        );
8724                        let result = self.request_hitl_approval(check_result).await?;
8725                        if matches!(
8726                            result,
8727                            ApprovalResult::Approved | ApprovalResult::Modified { .. }
8728                        ) {
8729                            return self.run_loop_internal(input).await;
8730                        }
8731                    }
8732                    let apology = self
8733                        .generate_localized_apology(
8734                            "Explain briefly that you're transferring the user to a human agent for help.",
8735                            &reason,
8736                        )
8737                        .await
8738                        .unwrap_or_else(|_| {
8739                            format!("I need human assistance to help with your request: {}", reason)
8740                        });
8741                    let response = AgentResponse::new(&apology);
8742                    self.finish_turn_if_root(&response).await?;
8743                    return Ok(response);
8744                }
8745                DisambiguationResult::Abandoned { new_input } => {
8746                    *self.pending_skill_id.write() = None;
8747
8748                    info!(
8749                        has_new_input = new_input.is_some(),
8750                        "Clarification abandoned by user"
8751                    );
8752
8753                    self.commit_root_user_message(input).await?;
8754
8755                    match new_input {
8756                        Some(fresh_input) => {
8757                            // Topic switch: process the user's new input from scratch.
8758                            // The LLM sees full conversation context including the abandoned exchange.
8759                            return self.run_loop_internal(&fresh_input).await;
8760                        }
8761                        None => {
8762                            // Pure abandonment: generate a brief acknowledgment.
8763                            let ack = self
8764                                .generate_localized_apology(
8765                                    "The user changed their mind about their previous request. \
8766                                     Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
8767                                     Do NOT apologize excessively. Be concise.",
8768                                    "User abandoned clarification",
8769                                )
8770                                .await
8771                                .unwrap_or_else(|_| {
8772                                    "OK, no problem. What else can I help with?".to_string()
8773                                });
8774
8775                            self.memory
8776                                .add_message(ChatMessage::assistant(&ack))
8777                                .await?;
8778
8779                            let response = AgentResponse::new(&ack);
8780                            self.finish_turn_if_root(&response).await?;
8781                            return Ok(response);
8782                        }
8783                    }
8784                }
8785            }
8786        }
8787
8788        self.run_loop_internal(input).await
8789    }
8790
8791    /// Generate a localized response using the router LLM
8792    async fn generate_localized_apology(&self, instruction: &str, reason: &str) -> Result<String> {
8793        let llm = self.llm_registry.router().map_err(|e| {
8794            AgentError::LLM(format!(
8795                "Router LLM not available for localized response: {}",
8796                e
8797            ))
8798        })?;
8799
8800        let recent: Vec<String> = self
8801            .memory
8802            .get_messages(Some(3))
8803            .await?
8804            .iter()
8805            .map(|m| m.content.clone())
8806            .collect();
8807
8808        let context_hint = if recent.is_empty() {
8809            String::new()
8810        } else {
8811            format!(
8812                "\nRecent conversation (detect the user's language from this):\n{}\n",
8813                recent.join("\n")
8814            )
8815        };
8816
8817        let prompt = format!(
8818            "{}\nReason: {}\n{}Respond in the same language as the user. Output ONLY the message, nothing else.",
8819            instruction, reason, context_hint
8820        );
8821
8822        let messages = vec![ChatMessage::user(&prompt)];
8823        let response = self
8824            .observe_purpose(
8825                ObservationPurpose::DisambiguationClarification,
8826                llm.complete(&messages, None),
8827            )
8828            .await
8829            .map_err(|e| AgentError::LLM(format!("Localized response generation failed: {}", e)))?;
8830
8831        Ok(response.content.trim().to_string())
8832    }
8833
8834    /// Clear disambiguation-related keys from the context manager.
8835    ///
8836    /// Render template variables in state action args using the context manager.
8837    fn render_action_args(&self, args: &Value) -> Value {
8838        let context = self.build_context_with_overlays();
8839        match args {
8840            Value::Object(map) => {
8841                let mut rendered = serde_json::Map::new();
8842                for (k, v) in map {
8843                    match v {
8844                        Value::String(s) if s.contains("{{") => {
8845                            match self.template_renderer.render(s, &context) {
8846                                Ok(rendered_str) => {
8847                                    rendered.insert(k.clone(), Value::String(rendered_str));
8848                                }
8849                                Err(_) => {
8850                                    rendered.insert(k.clone(), v.clone());
8851                                }
8852                            }
8853                        }
8854                        _ => {
8855                            rendered.insert(k.clone(), v.clone());
8856                        }
8857                    }
8858                }
8859                Value::Object(rendered)
8860            }
8861            _ => args.clone(),
8862        }
8863    }
8864
8865    /// Called at the start of each turn to prevent stale `resolved_intent` from leaking across turns.
8866    fn clear_disambiguation_context(&self) {
8867        let _ = self
8868            .context_manager
8869            .set("resolved_intent", serde_json::Value::Null);
8870
8871        let all = self.context_manager.get_all();
8872        for key in all.keys() {
8873            if key.starts_with("disambiguation.") {
8874                let _ = self.context_manager.set(key, serde_json::Value::Null);
8875            }
8876        }
8877    }
8878
8879    /// Re-run skill disambiguation on enriched input before executing the skill.
8880    /// After clarification resolves, the enriched input may still be missing required_clarity fields (e.g. "Transfer money to Jane." still lacks amount).
8881    /// This method re-runs the skill's disambiguation pass.
8882    /// If fields are still missing, it returns the new clarification question and keeps pending_skill_id set.
8883    /// If all fields are present (Clear), it executes the skill and returns the response.
8884    async fn recheck_skill_disambiguation(
8885        &self,
8886        skill_id: &str,
8887        enriched_input: &str,
8888        expected_disambiguation_epoch: u64,
8889        expected_state_generation: Option<u64>,
8890    ) -> Result<AgentResponse> {
8891        let skill = self
8892            .skill_router
8893            .as_ref()
8894            .and_then(|r| r.get_skill(skill_id).cloned());
8895
8896        // If the skill has disambiguation enabled, re-run it on the enriched input.
8897        if let Some(ref skill) = skill
8898            && let Some(ref skill_disambig) = skill.disambiguation
8899            && skill_disambig.enabled.unwrap_or(false)
8900            && let Some(ref disambiguator) = self.disambiguation_manager
8901        {
8902            let context = self.build_disambiguation_context().await?;
8903            let state_override = self
8904                .state_machine
8905                .as_ref()
8906                .and_then(|sm| sm.current_definition())
8907                .and_then(|def| def.disambiguation.clone());
8908
8909            let disambiguation_result = self
8910                .observe_purpose(
8911                    ObservationPurpose::DisambiguationDetection,
8912                    disambiguator.process_input_with_override(
8913                        enriched_input,
8914                        &context,
8915                        state_override.as_ref(),
8916                        Some(skill_disambig),
8917                    ),
8918                )
8919                .await?;
8920            let current_state_generation = self
8921                .state_machine
8922                .as_ref()
8923                .map(|state_machine| state_machine.generation());
8924            if current_state_generation != expected_state_generation
8925                || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
8926            {
8927                disambiguator.clear_pending().await;
8928                *self.pending_skill_id.write() = None;
8929                return Err(AgentError::Other(
8930                    "State or reset ownership changed during skill disambiguation recheck"
8931                        .to_string(),
8932                ));
8933            }
8934            match disambiguation_result {
8935                DisambiguationResult::Clear => {
8936                    debug!(skill_id = %skill_id, "Skill re-check: all fields present");
8937                }
8938                DisambiguationResult::NeedsClarification {
8939                    question,
8940                    detection,
8941                } => {
8942                    let admission = self
8943                        .admit_disambiguation_redispatch(
8944                            expected_disambiguation_epoch,
8945                            expected_state_generation,
8946                        )
8947                        .await?;
8948                    let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
8949                    info!(
8950                        skill_id = %skill_id,
8951                        ambiguity_type = ?detection.ambiguity_type,
8952                        what_is_unclear = ?detection.what_is_unclear,
8953                        "Skill re-check: still missing fields, asking again"
8954                    );
8955                    // Keep pending_skill_id set (do NOT clear it).
8956                    // The next turn will resolve this new clarification and
8957                    // re-enter this method until all fields are present.
8958                    self.memory
8959                        .add_message(ChatMessage::user(enriched_input))
8960                        .await?;
8961                    self.memory
8962                        .add_message(ChatMessage::assistant(&question.question))
8963                        .await?;
8964
8965                    let response = AgentResponse::new(&question.question).with_metadata(
8966                        "disambiguation",
8967                        serde_json::json!({
8968                            "status": if awaiting_confirmation { "awaiting_confirmation" } else { "awaiting_clarification" },
8969                            "skill_id": skill_id,
8970                            "options": question.options,
8971                            "clarifying": question.clarifying,
8972                            "detection": {
8973                                "type": detection.ambiguity_type,
8974                                "confidence": detection.confidence,
8975                                "what_is_unclear": detection.what_is_unclear,
8976                            }
8977                        }),
8978                    );
8979                    drop(admission);
8980                    self.finish_turn_if_root(&response).await?;
8981                    return Ok(response);
8982                }
8983                DisambiguationResult::Clarified {
8984                    enriched_input: re_enriched,
8985                    ..
8986                } => {
8987                    debug!(skill_id = %skill_id, "Skill re-check: clarified immediately, executing");
8988                    let admission = self
8989                        .admit_disambiguation_redispatch(
8990                            expected_disambiguation_epoch,
8991                            expected_state_generation,
8992                        )
8993                        .await?;
8994                    *self.pending_skill_id.write() = None;
8995                    drop(admission);
8996                    let skill_response = self.execute_skill_by_id(skill_id, &re_enriched).await?;
8997                    self.memory
8998                        .add_message(ChatMessage::user(&re_enriched))
8999                        .await?;
9000                    return self
9001                        .handle_skill_response(
9002                            &re_enriched,
9003                            skill_id,
9004                            skill_response,
9005                            &HashMap::new(),
9006                        )
9007                        .await;
9008                }
9009                DisambiguationResult::ProceedWithBestGuess {
9010                    enriched_input: re_enriched,
9011                } => {
9012                    debug!(skill_id = %skill_id, "Skill re-check: proceeding with best guess");
9013                    let admission = self
9014                        .admit_disambiguation_redispatch(
9015                            expected_disambiguation_epoch,
9016                            expected_state_generation,
9017                        )
9018                        .await?;
9019                    *self.pending_skill_id.write() = None;
9020                    drop(admission);
9021                    let skill_response = self.execute_skill_by_id(skill_id, &re_enriched).await?;
9022                    self.memory
9023                        .add_message(ChatMessage::user(&re_enriched))
9024                        .await?;
9025                    return self
9026                        .handle_skill_response(
9027                            &re_enriched,
9028                            skill_id,
9029                            skill_response,
9030                            &HashMap::new(),
9031                        )
9032                        .await;
9033                }
9034                DisambiguationResult::GiveUp { reason } => {
9035                    *self.pending_skill_id.write() = None;
9036                    let apology = self
9037                                    .generate_localized_apology(
9038                                        "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
9039                                        &reason,
9040                                    )
9041                                    .await
9042                                    .unwrap_or_else(|_| {
9043                                        format!("I'm sorry, I couldn't understand your request: {}", reason)
9044                                    });
9045                    let response = AgentResponse::new(&apology);
9046                    self.finish_turn_if_root(&response).await?;
9047                    return Ok(response);
9048                }
9049                DisambiguationResult::Escalate { reason } => {
9050                    *self.pending_skill_id.write() = None;
9051                    let apology = self
9052                                    .generate_localized_apology(
9053                                        "Explain briefly that you're transferring the user to a human agent for help.",
9054                                        &reason,
9055                                    )
9056                                    .await
9057                                    .unwrap_or_else(|_| {
9058                                        format!("I need human assistance to help with your request: {}", reason)
9059                                    });
9060                    let response = AgentResponse::new(&apology);
9061                    self.finish_turn_if_root(&response).await?;
9062                    return Ok(response);
9063                }
9064                DisambiguationResult::Abandoned { new_input } => {
9065                    // User abandoned during skill re-check.
9066                    // Clear skill routing state and fall through to normal execution.
9067                    *self.pending_skill_id.write() = None;
9068                    debug!(skill_id = %skill_id, "Skill re-check: abandoned by user");
9069                    if let Some(fresh) = new_input {
9070                        return self.run_loop_internal(&fresh).await;
9071                    }
9072                    let ack = self
9073                                    .generate_localized_apology(
9074                                        "The user changed their mind about their previous request. \
9075                                         Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
9076                                         Do NOT apologize excessively. Be concise.",
9077                                        "User abandoned clarification",
9078                                    )
9079                                    .await
9080                                    .unwrap_or_else(|_| {
9081                                        "OK, no problem. What else can I help with?".to_string()
9082                                    });
9083                    self.memory
9084                        .add_message(ChatMessage::assistant(&ack))
9085                        .await?;
9086                    let response = AgentResponse::new(&ack);
9087                    self.finish_turn_if_root(&response).await?;
9088                    return Ok(response);
9089                }
9090            }
9091        }
9092
9093        // Skill execution is admitted before the read guard is released so later invalidation cannot retroactively cancel it.
9094        let admission = self
9095            .admit_disambiguation_redispatch(
9096                expected_disambiguation_epoch,
9097                expected_state_generation,
9098            )
9099            .await?;
9100        *self.pending_skill_id.write() = None;
9101        drop(admission);
9102        let skill_response = self.execute_skill_by_id(skill_id, enriched_input).await?;
9103        self.memory
9104            .add_message(ChatMessage::user(enriched_input))
9105            .await?;
9106        self.handle_skill_response(enriched_input, skill_id, skill_response, &HashMap::new())
9107            .await
9108    }
9109
9110    /// Handle skill routing result: output processing, memory, transitions.
9111    /// Returns a fully formed AgentResponse for skill-routed requests.
9112    async fn handle_skill_response(
9113        &self,
9114        processed_input: &str,
9115        skill_id: &str,
9116        skill_response: String,
9117        input_context: &HashMap<String, Value>,
9118    ) -> Result<AgentResponse> {
9119        let output_data = self.process_output(&skill_response, input_context).await?;
9120        let final_response = output_data.content;
9121
9122        self.memory
9123            .add_message(ChatMessage::assistant(&final_response))
9124            .await?;
9125
9126        self.check_memory_compression().await?;
9127
9128        self.increment_turn();
9129        self.evaluate_transitions(processed_input, &final_response)
9130            .await?;
9131
9132        let response = AgentResponse::new(final_response)
9133            .with_metadata("skill_id", serde_json::json!(skill_id));
9134        self.finish_turn_if_root(&response).await?;
9135        Ok(response)
9136    }
9137
9138    /// Run the Plan-and-Execute flow: generate plan, execute steps, finalize.
9139    /// Supports plan-level reflection with replan loop when configured.
9140    async fn handle_plan_and_execute(
9141        &self,
9142        processed_input: &str,
9143        input_context: &HashMap<String, Value>,
9144        auto_detected: bool,
9145    ) -> Result<AgentResponse> {
9146        let effective = self.get_effective_reasoning_config();
9147        let plan_reflection = effective
9148            .get_planning()
9149            .map(|c| c.reflection.clone())
9150            .unwrap_or_default();
9151
9152        let max_attempts = if plan_reflection.enabled {
9153            1 + plan_reflection.max_replans
9154        } else {
9155            1
9156        };
9157
9158        let mut plan = self.generate_plan(processed_input).await?;
9159        info!(
9160            plan_id = %plan.id,
9161            steps = plan.steps.len(),
9162            "Plan generated"
9163        );
9164
9165        let mut plan_result = String::new();
9166
9167        for attempt in 0..max_attempts {
9168            *self.current_plan.write() = Some(plan.clone());
9169            plan_result = self.execute_plan(&mut plan).await?;
9170
9171            info!(
9172                plan_status = ?plan.status,
9173                completed_steps = plan.completed_steps().count(),
9174                attempt = attempt + 1,
9175                "Plan execution completed"
9176            );
9177
9178            if !plan_reflection.enabled {
9179                break;
9180            }
9181
9182            let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
9183            if !has_failures {
9184                break;
9185            }
9186
9187            if attempt + 1 >= max_attempts {
9188                break;
9189            }
9190
9191            match plan_reflection.on_step_failure {
9192                StepFailureAction::Replan => {
9193                    info!(attempt = attempt + 1, "Plan had failures, replanning");
9194                    plan = self.generate_plan(processed_input).await?;
9195                }
9196                StepFailureAction::Abort => {
9197                    warn!("Plan step failed, aborting");
9198                    break;
9199                }
9200                StepFailureAction::Skip | StepFailureAction::Continue => {
9201                    break;
9202                }
9203            }
9204        }
9205
9206        *self.current_plan.write() = Some(plan);
9207
9208        let output_data = self.process_output(&plan_result, input_context).await?;
9209        let final_content = output_data.content;
9210
9211        self.memory
9212            .add_message(ChatMessage::assistant(&final_content))
9213            .await?;
9214
9215        self.check_memory_compression().await?;
9216        self.increment_turn();
9217        self.evaluate_transitions(processed_input, &final_content)
9218            .await?;
9219
9220        let reasoning_metadata =
9221            ReasoningMetadata::new(ReasoningMode::PlanAndExecute).with_auto_detected(auto_detected);
9222
9223        let response = AgentResponse::new(&final_content).with_metadata(
9224            "reasoning",
9225            serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
9226        );
9227
9228        self.finish_turn_if_root(&response).await?;
9229        Ok(response)
9230    }
9231
9232    /// Inject CoT/ReAct reasoning prompt into the system message (first iteration only).
9233    fn inject_reasoning_prompt(
9234        &self,
9235        messages: &mut [ChatMessage],
9236        reasoning_mode: &ReasoningMode,
9237        is_first_iteration: bool,
9238    ) {
9239        if !is_first_iteration {
9240            return;
9241        }
9242        match reasoning_mode {
9243            ReasoningMode::CoT => {
9244                if let Some(msg) = messages.first_mut()
9245                    && matches!(msg.role, ai_agents_core::Role::System)
9246                {
9247                    msg.content = self.build_cot_system_prompt(&msg.content);
9248                    debug!("Applied Chain-of-Thought system prompt");
9249                }
9250            }
9251            ReasoningMode::React => {
9252                if let Some(msg) = messages.first_mut()
9253                    && matches!(msg.role, ai_agents_core::Role::System)
9254                {
9255                    msg.content = self.build_react_system_prompt(&msg.content);
9256                    debug!("Applied ReAct system prompt");
9257                }
9258            }
9259            _ => {}
9260        }
9261    }
9262
9263    //
9264    // Draft generation must not commit user memory or run tools.
9265    // The current user input is added only as an ephemeral message for this LLM call.
9266    //
9267    async fn generate_main_response_draft(
9268        &self,
9269        processed_input: &str,
9270        reasoning_mode: &ReasoningMode,
9271    ) -> Result<MainResponseDraft> {
9272        let llm = self.get_state_llm()?;
9273        let protocol = self.main_tool_protocol(llm.as_ref(), true).await?;
9274        let mut messages = self
9275            .build_messages_internal(false, Some(processed_input), protocol.choice.is_none())
9276            .await?;
9277        self.inject_reasoning_prompt(&mut messages, reasoning_mode, true);
9278        let response = self
9279            .complete_main_llm_with_recovery(llm, &messages, &protocol)
9280            .await?;
9281        let content = response.content.trim().to_string();
9282        let (thinking, answer) = self.extract_thinking(&content);
9283        if let Some(calls) = self.parse_main_tool_calls(&content, &protocol) {
9284            return Ok(MainResponseDraft::ToolCalls {
9285                raw_content: content,
9286                calls,
9287                thinking,
9288            });
9289        }
9290        Ok(MainResponseDraft::Text {
9291            raw_content: answer,
9292            thinking,
9293        })
9294    }
9295
9296    //
9297    // This is the only place where a winning draft is allowed to become runtime state.
9298    // Prompt and native tool calls remain inert until this method commits the draft into the shared executor path.
9299    //
9300    async fn commit_main_response_draft(
9301        &self,
9302        processed_input: &str,
9303        input_context: &HashMap<String, Value>,
9304        draft: MainResponseDraft,
9305        reasoning_mode: ReasoningMode,
9306        auto_detected: bool,
9307    ) -> Result<AgentResponse> {
9308        self.commit_root_user_message(processed_input).await?;
9309        match draft {
9310            MainResponseDraft::Text {
9311                raw_content,
9312                thinking,
9313            } => {
9314                self.finish_text_response_from_model(CommittedTextResponse {
9315                    processed_input,
9316                    input_context,
9317                    answer: raw_content,
9318                    reasoning_mode,
9319                    auto_detected,
9320                    iterations: 1,
9321                    thinking_content: thinking,
9322                    all_tool_calls: Vec::new(),
9323                })
9324                .await
9325            }
9326            MainResponseDraft::ToolCalls {
9327                raw_content,
9328                calls,
9329                thinking: _,
9330            } => {
9331                let mut all_tool_calls = Vec::new();
9332                match self
9333                    .handle_tool_calls(processed_input, &raw_content, calls, &mut all_tool_calls)
9334                    .await?
9335                {
9336                    ToolCallOutcome::Rejected(response) => {
9337                        self.finish_turn_if_root(&response).await?;
9338                        Ok(response)
9339                    }
9340                    ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => {
9341                        self.continue_after_committed_tool_draft(processed_input)
9342                            .await
9343                    }
9344                }
9345            }
9346        }
9347    }
9348
9349    //
9350    // Tool drafts need a committed continuation after function results are written.
9351    // Redispatch depth suppresses duplicate root lifecycle work during that continuation.
9352    //
9353    async fn continue_after_committed_tool_draft(
9354        &self,
9355        processed_input: &str,
9356    ) -> Result<AgentResponse> {
9357        *self.redispatch_depth.write() += 1;
9358        if let Some(context) = self.active_turn_context.write().as_mut() {
9359            context.enter_redispatch();
9360        }
9361        let result = Box::pin(self.run_loop_internal(processed_input)).await;
9362        *self.redispatch_depth.write() -= 1;
9363        if let Some(context) = self.active_turn_context.write().as_mut() {
9364            context.exit_redispatch();
9365        }
9366        let response = result?;
9367        self.finish_turn_if_root(&response).await?;
9368        Ok(response)
9369    }
9370
9371    //
9372    // Shared committed text finalization for normal responses and winning text drafts.
9373    // Keep output processing, reflection, transitions, hooks, and maintenance behind this commit boundary.
9374    //
9375    async fn finish_text_response_from_model(
9376        &self,
9377        response: CommittedTextResponse<'_>,
9378    ) -> Result<AgentResponse> {
9379        let CommittedTextResponse {
9380            processed_input,
9381            input_context,
9382            answer,
9383            reasoning_mode,
9384            auto_detected,
9385            iterations,
9386            thinking_content,
9387            all_tool_calls,
9388        } = response;
9389        let output_data = self.process_output(&answer, input_context).await?;
9390        let mut final_content = if output_data.metadata.rejected {
9391            output_data
9392                .metadata
9393                .rejection_reason
9394                .unwrap_or_else(|| answer.to_string())
9395        } else {
9396            output_data.content
9397        };
9398        let llm = self.get_state_llm()?;
9399        let reflection_metadata;
9400        (final_content, reflection_metadata) = self
9401            .run_reflection(&*llm, processed_input, final_content)
9402            .await?;
9403        final_content =
9404            self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
9405        let final_content = {
9406            let result = self
9407                .post_loop_processing(processed_input, final_content)
9408                .await?;
9409            self.apply_post_loop_result(processed_input, result).await?
9410        };
9411        let response = self.build_agent_response(AgentResponseParts {
9412            content: final_content,
9413            all_tool_calls,
9414            reasoning_mode,
9415            auto_detected,
9416            iterations,
9417            thinking: thinking_content,
9418            reflection_metadata,
9419        });
9420        self.finish_turn_if_root(&response).await?;
9421        Ok(response)
9422    }
9423
9424    //
9425    // Auto reasoning uses this path after the judge wins with a deeper mode.
9426    // It intentionally uses committed message building instead of the draft overlay.
9427    //
9428    async fn run_committed_response_loop_with_reasoning(
9429        &self,
9430        processed_input: &str,
9431        input_context: &HashMap<String, Value>,
9432        reasoning_mode: ReasoningMode,
9433        auto_detected: bool,
9434    ) -> Result<AgentResponse> {
9435        self.commit_root_user_message(processed_input).await?;
9436        let llm = self.get_state_llm()?;
9437        let mut iterations = 0u32;
9438        let mut all_tool_calls = Vec::new();
9439        let mut thinking_content = None;
9440        loop {
9441            let effective_max = if reasoning_mode != ReasoningMode::None {
9442                let rc = self.get_effective_reasoning_config();
9443                self.max_iterations.min(rc.max_iterations)
9444            } else {
9445                self.max_iterations
9446            };
9447            if iterations >= effective_max {
9448                return Err(AgentError::Other(format!(
9449                    "Max iterations ({}) exceeded",
9450                    effective_max
9451                )));
9452            }
9453            iterations += 1;
9454            *self.iteration_count.write() = iterations;
9455            let protocol = self.main_tool_protocol(llm.as_ref(), false).await?;
9456            let mut messages = self
9457                .build_messages_internal(true, None, protocol.choice.is_none())
9458                .await?;
9459            self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
9460            self.hooks.on_llm_start(&messages).await;
9461            let llm_start = Instant::now();
9462            let response = self
9463                .complete_main_llm_with_recovery(Arc::clone(&llm), &messages, &protocol)
9464                .await?;
9465            let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
9466            self.hooks.on_llm_complete(&response, llm_duration_ms).await;
9467            let content = response.content.trim();
9468            if let Some(tool_calls) = self.parse_main_tool_calls(content, &protocol) {
9469                match self
9470                    .handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
9471                    .await?
9472                {
9473                    ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
9474                    ToolCallOutcome::Rejected(resp) => {
9475                        self.finish_turn_if_root(&resp).await?;
9476                        return Ok(resp);
9477                    }
9478                }
9479            }
9480            let (extracted_thinking, answer) = self.extract_thinking(content);
9481            if extracted_thinking.is_some() {
9482                thinking_content = extracted_thinking;
9483            }
9484            return self
9485                .finish_text_response_from_model(CommittedTextResponse {
9486                    processed_input,
9487                    input_context,
9488                    answer,
9489                    reasoning_mode,
9490                    auto_detected,
9491                    iterations,
9492                    thinking_content,
9493                    all_tool_calls,
9494                })
9495                .await;
9496        }
9497    }
9498
9499    /// Handle tool calls: check transitions, execute tools in parallel, handle HITL rejection.
9500    async fn handle_tool_calls(
9501        &self,
9502        processed_input: &str,
9503        content: &str,
9504        tool_calls: Vec<ToolCall>,
9505        all_tool_calls: &mut Vec<ToolCall>,
9506    ) -> Result<ToolCallOutcome> {
9507        // Check if a transition should fire before executing the LLM's tool call.
9508        // If a transition fires, on_enter actions handle the tool call correctly
9509        // (with proper URLs from YAML), so skip the LLM's tool call.
9510        let transition_fired = self.evaluate_transitions(processed_input, content).await?;
9511        if transition_fired {
9512            self.memory
9513                .add_message(ChatMessage::assistant(
9514                    "(Transitioned to new state — tool call handled by workflow)",
9515                ))
9516                .await?;
9517            return Ok(ToolCallOutcome::TransitionFired);
9518        }
9519
9520        // 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.
9521        self.memory
9522            .add_message(ChatMessage::assistant(content))
9523            .await?;
9524        let native_tool_call = Self::is_native_tool_call_content(content);
9525
9526        let results = self.execute_tools_parallel(&tool_calls).await;
9527
9528        for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
9529            match result {
9530                Ok(output) => {
9531                    self.memory
9532                        .add_message(Self::tool_result_message(
9533                            tool_call,
9534                            &output,
9535                            native_tool_call,
9536                        ))
9537                        .await?;
9538                }
9539                Err(e) => {
9540                    // Check if this is a HITL rejection - if so, break the loop
9541                    if matches!(e, AgentError::HITLRejected(_)) {
9542                        self.memory
9543                            .add_message(ChatMessage::assistant(format!(
9544                                "The operation was rejected by the approver: {}",
9545                                e
9546                            )))
9547                            .await?;
9548                        // Return the rejection message to user, don't continue loop
9549                        return Ok(ToolCallOutcome::Rejected(AgentResponse {
9550                            content: format!("Operation cancelled: {}", e),
9551                            metadata: None,
9552                            tool_calls: Some(all_tool_calls.clone()),
9553                        }));
9554                    }
9555                    self.memory
9556                        .add_message(Self::tool_result_message(
9557                            tool_call,
9558                            &format!("Error: {}", e),
9559                            native_tool_call,
9560                        ))
9561                        .await?;
9562                }
9563            }
9564            all_tool_calls.push(tool_call.clone());
9565        }
9566        Ok(ToolCallOutcome::Continue)
9567    }
9568
9569    /// Run the reflection loop on a response, returning (improved_content, reflection_metadata).
9570    async fn run_reflection(
9571        &self,
9572        llm: &dyn LLMProvider,
9573        processed_input: &str,
9574        mut content: String,
9575    ) -> Result<(String, Option<ReflectionMetadata>)> {
9576        let should_reflect = self.should_reflect(processed_input, &content).await?;
9577        if !should_reflect {
9578            return Ok((content, None));
9579        }
9580
9581        info!("Starting response reflection evaluation");
9582        let mut attempts = 0u32;
9583        let max_retries = self.reflection_config.max_retries;
9584        let mut history: Vec<ReflectionAttempt> = Vec::new();
9585
9586        loop {
9587            let evaluation = self.evaluate_response(processed_input, &content).await?;
9588
9589            if evaluation.passed || attempts >= max_retries {
9590                info!(
9591                    passed = evaluation.passed,
9592                    confidence = evaluation.confidence,
9593                    attempts = attempts + 1,
9594                    "Reflection evaluation complete"
9595                );
9596                let reflection_metadata = Some(
9597                    ReflectionMetadata::new(evaluation)
9598                        .with_attempts(attempts + 1)
9599                        .with_history(history),
9600                );
9601                return Ok((content, reflection_metadata));
9602            }
9603
9604            debug!(
9605                attempt = attempts + 1,
9606                failed_criteria = evaluation.failed_criteria().count(),
9607                "Response did not meet criteria, retrying"
9608            );
9609
9610            history.push(
9611                ReflectionAttempt::new(&content, evaluation.clone())
9612                    .with_feedback("Response did not meet quality criteria"),
9613            );
9614
9615            let feedback: Vec<String> = evaluation
9616                .failed_criteria()
9617                .map(|c| format!("- {}", c.criterion))
9618                .collect();
9619
9620            let retry_prompt = format!(
9621                "Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response.",
9622                feedback.join("\n")
9623            );
9624
9625            self.memory
9626                .add_message(ChatMessage::user(&retry_prompt))
9627                .await?;
9628
9629            let retry_messages = self.build_messages().await?;
9630            let retry_response = self
9631                .observe_purpose(
9632                    ObservationPurpose::ReflectionEvaluation,
9633                    llm.complete(&retry_messages, None),
9634                )
9635                .await
9636                .map_err(|e| AgentError::LLM(e.to_string()))?;
9637
9638            content = retry_response.content.trim().to_string();
9639            attempts += 1;
9640        }
9641    }
9642
9643    /// Record the assistant turn, evaluate transitions, and decide what to do next.
9644    /// Returns PostLoopResult so callers can apply_post_loop_result for re-dispatch.
9645    async fn post_loop_processing(
9646        &self,
9647        processed_input: &str,
9648        content: String,
9649    ) -> Result<PostLoopResult> {
9650        // Do NOT add the assistant message to memory yet.
9651        // evaluate_transitions receives content as a direct parameter, so the message does not need to be in memory for transitions to evaluate correctly.
9652        // For NeedsRedispatch we skip adding the stale old-state response entirely, keeping memory clean for the re-dispatched handler.
9653
9654        self.increment_turn();
9655
9656        // Run context extractors so guards can check freshly-extracted values.
9657        self.run_context_extractors(processed_input).await;
9658
9659        let transitioned = self.evaluate_transitions(processed_input, &content).await?;
9660
9661        if !transitioned {
9662            self.memory
9663                .add_message(ChatMessage::assistant(&content))
9664                .await?;
9665            self.check_memory_compression().await?;
9666            return Ok(PostLoopResult::NoTransition(content));
9667        }
9668
9669        // Check if we should skip re-generation after this transition.
9670        if !self.should_regenerate_after_transition() {
9671            self.memory
9672                .add_message(ChatMessage::assistant(&content))
9673                .await?;
9674            self.check_memory_compression().await?;
9675            return Ok(PostLoopResult::Transitioned(content));
9676        }
9677
9678        // Check if the new state needs full dispatch.
9679        // Orchestration states (concurrent, group_chat, pipeline, handoff, delegate) need their dedicated handlers.
9680        // 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.
9681        if self.needs_redispatch_for_new_state() {
9682            info!("Post-transition NeedsRedispatch: new state requires full dispatch");
9683            // Stale old-state content is NOT added to memory.
9684            // apply_post_loop_result will increment redispatch_depth and call run_loop_internal, which produces the correct response and adds it.
9685            return Ok(PostLoopResult::NeedsRedispatch);
9686        }
9687
9688        // Normal post-transition re-generation (plain LLM with optional tool calls).
9689        // Add the stale response to history so the new LLM sees the conversation.
9690        self.memory
9691            .add_message(ChatMessage::assistant(&content))
9692            .await?;
9693        self.check_memory_compression().await?;
9694
9695        // If a transition fired, on_enter actions already executed (e.g., HTTP calls).
9696        // The current content was generated in the OLD state context and is stale.
9697        // Re-generate in the new state context so the LLM can reference on_enter results.
9698        // If the LLM responds with a tool call, execute it in a mini-loop so the
9699        // result is not returned as raw JSON text.
9700        let new_llm = self.get_state_llm()?;
9701        let mut final_content;
9702
9703        for post_iter in 0..self.max_iterations {
9704            let protocol = self.main_tool_protocol(new_llm.as_ref(), false).await?;
9705            let new_messages = self
9706                .build_messages_internal(true, None, protocol.choice.is_none())
9707                .await?;
9708            if post_iter == 0
9709                && let Some(system_msg) = new_messages.first()
9710                && system_msg.role == ai_agents_core::Role::System
9711            {
9712                debug!(
9713                    prompt_preview =
9714                        &system_msg.content[system_msg.content.len().saturating_sub(200)..],
9715                    "Post-transition system prompt (last 200 chars)"
9716                );
9717            }
9718
9719            let new_response = self
9720                .complete_main_llm_with_recovery(Arc::clone(&new_llm), &new_messages, &protocol)
9721                .await?;
9722            final_content = new_response.content.trim().to_string();
9723
9724            // Check if the post-transition response contains tool calls.
9725            // If so, execute them and loop so the LLM can summarize the result.
9726            if let Some(tool_calls) = self.parse_main_tool_calls(&final_content, &protocol) {
9727                let native_tool_call = Self::is_native_tool_call_content(&final_content);
9728                debug!(
9729                    post_iter = post_iter,
9730                    tools = tool_calls.len(),
9731                    "Post-transition tool call detected, executing"
9732                );
9733
9734                self.memory
9735                    .add_message(ChatMessage::assistant(&final_content))
9736                    .await?;
9737
9738                let results = self.execute_tools_parallel(&tool_calls).await;
9739                for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
9740                    match result {
9741                        Ok(output) => {
9742                            self.memory
9743                                .add_message(Self::tool_result_message(
9744                                    tool_call,
9745                                    &output,
9746                                    native_tool_call,
9747                                ))
9748                                .await?;
9749                        }
9750                        Err(e) => {
9751                            self.memory
9752                                .add_message(Self::tool_result_message(
9753                                    tool_call,
9754                                    &format!("Error: {}", e),
9755                                    native_tool_call,
9756                                ))
9757                                .await?;
9758                        }
9759                    }
9760                }
9761                // Loop to let the LLM see the tool result and produce a text response.
9762                continue;
9763            }
9764
9765            // No tool call - this is the final text response.
9766            self.memory
9767                .add_message(ChatMessage::assistant(&final_content))
9768                .await?;
9769            return Ok(PostLoopResult::Transitioned(final_content));
9770        }
9771
9772        // Exhausted post-transition iterations (unlikely). Return last content.
9773        final_content = "Post-transition processing completed.".to_string();
9774        self.memory
9775            .add_message(ChatMessage::assistant(&final_content))
9776            .await?;
9777
9778        Ok(PostLoopResult::Transitioned(final_content))
9779    }
9780
9781    /// Build the final AgentResponse with all metadata.
9782    /// Check whether to re-generate a response after a state transition.
9783    fn should_regenerate_after_transition(&self) -> bool {
9784        if let Some(ref sm) = self.state_machine {
9785            // Global setting
9786            if !sm.config().regenerate_on_transition {
9787                return false;
9788            }
9789            // Per-state override on the new (current) state
9790            if let Some(def) = sm.current_definition()
9791                && let Some(regen) = def.regenerate_on_enter
9792            {
9793                return regen;
9794            }
9795        }
9796        true
9797    }
9798
9799    /// Return true when the new state requires full dispatch via run_loop_internal.
9800    /// Covers orchestration states and any non-None effective reasoning mode.
9801    fn needs_redispatch_for_new_state(&self) -> bool {
9802        if let Some(ref sm) = self.state_machine
9803            && let Some(def) = sm.current_definition()
9804        {
9805            if def.concurrent.is_some()
9806                || def.group_chat.is_some()
9807                || def.pipeline.is_some()
9808                || def.handoff.is_some()
9809                || def.delegate.is_some()
9810            {
9811                return true;
9812            }
9813            // Any non-None effective reasoning mode requires the main dispatch loop.
9814            let effective = self.get_effective_reasoning_config();
9815            if !matches!(effective.mode, ReasoningMode::None) {
9816                return true;
9817            }
9818        }
9819        false
9820    }
9821
9822    /// Consume a PostLoopResult. NeedsRedispatch re-enters run_loop_internal.
9823    /// The user message is already in memory - redispatch_depth suppresses re-adding it.
9824    async fn apply_post_loop_result(
9825        &self,
9826        processed_input: &str,
9827        result: PostLoopResult,
9828    ) -> Result<String> {
9829        match result {
9830            PostLoopResult::NoTransition(content) | PostLoopResult::Transitioned(content) => {
9831                Ok(content)
9832            }
9833            PostLoopResult::NeedsRedispatch => {
9834                const MAX_REDISPATCH_DEPTH: u32 = 3;
9835                let current_depth = *self.redispatch_depth.read();
9836                if current_depth >= MAX_REDISPATCH_DEPTH {
9837                    warn!(
9838                        depth = current_depth,
9839                        "Post-transition re-dispatch depth limit reached, returning empty response"
9840                    );
9841                    let content = String::new();
9842                    self.memory
9843                        .add_message(ChatMessage::assistant(&content))
9844                        .await?;
9845                    return Ok(content);
9846                }
9847                *self.redispatch_depth.write() += 1;
9848                if let Some(context) = self.active_turn_context.write().as_mut() {
9849                    context.enter_redispatch();
9850                }
9851                info!(
9852                    depth = current_depth + 1,
9853                    "Re-dispatching for new state after transition"
9854                );
9855                let resp = Box::pin(self.run_loop_internal(processed_input)).await;
9856                *self.redispatch_depth.write() -= 1;
9857                if let Some(context) = self.active_turn_context.write().as_mut() {
9858                    context.exit_redispatch();
9859                }
9860                resp.map(|r| r.content)
9861            }
9862        }
9863    }
9864
9865    /// Builds the final response metadata after committed output and transition handling complete.
9866    fn build_agent_response(&self, parts: AgentResponseParts) -> AgentResponse {
9867        let AgentResponseParts {
9868            content,
9869            all_tool_calls,
9870            reasoning_mode,
9871            auto_detected,
9872            iterations,
9873            thinking,
9874            reflection_metadata,
9875        } = parts;
9876        let reasoning_metadata = ReasoningMetadata::new(reasoning_mode.clone())
9877            .with_thinking(thinking.clone().unwrap_or_default())
9878            .with_iterations(iterations)
9879            .with_auto_detected(auto_detected);
9880
9881        let mut response = AgentResponse::new(&content);
9882        if !all_tool_calls.is_empty() {
9883            response = response.with_tool_calls(all_tool_calls);
9884        }
9885
9886        if let Some(state) = self.current_state() {
9887            response = response.with_metadata("current_state", serde_json::json!(state));
9888        }
9889
9890        response = response.with_metadata(
9891            "reasoning",
9892            serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
9893        );
9894
9895        if let Some(ref refl_meta) = reflection_metadata {
9896            response = response.with_metadata(
9897                "reflection",
9898                serde_json::to_value(refl_meta).unwrap_or_default(),
9899            );
9900        }
9901
9902        response
9903    }
9904
9905    // Handle delegation: forward user input to a registry agent.
9906    async fn handle_delegated_state(
9907        &self,
9908        input: &str,
9909        delegate_id: &str,
9910        state_def: &ai_agents_state::StateDefinition,
9911    ) -> Result<AgentResponse> {
9912        use std::time::Instant;
9913
9914        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
9915            AgentError::Config(format!(
9916                "State delegates to '{}' but no agent registry is configured. \
9917                 Add a spawner section with auto_spawn to your YAML.",
9918                delegate_id
9919            ))
9920        })?;
9921
9922        let state_name = self
9923            .state_machine
9924            .as_ref()
9925            .map(|sm| sm.current())
9926            .unwrap_or_else(|| "unknown".to_string());
9927
9928        self.hooks.on_delegate_start(delegate_id, &state_name).await;
9929        let start = Instant::now();
9930
9931        let delegate = registry.get(delegate_id).ok_or_else(|| {
9932            AgentError::Other(format!(
9933                "State '{}' delegates to '{}' but no agent with that ID exists in the registry.",
9934                state_name, delegate_id
9935            ))
9936        })?;
9937
9938        // Prepare input based on delegate_context mode.
9939        let context_mode = state_def.delegate_context.clone().unwrap_or_default();
9940        let effective_input = self
9941            .observe_purpose(
9942                ObservationPurpose::OrchestrationRouting,
9943                crate::orchestration::context::prepare_delegate_input(
9944                    input,
9945                    &context_mode,
9946                    &*self.memory,
9947                    self.llm_registry.get("router").ok().as_deref(),
9948                ),
9949            )
9950            .await?;
9951
9952        let response = delegate
9953            .chat_with_actor_context(&effective_input, self.outbound_actor_context())
9954            .await?;
9955
9956        let duration_ms = start.elapsed().as_millis() as u64;
9957        self.hooks
9958            .on_delegate_complete(delegate_id, &state_name, duration_ms)
9959            .await;
9960
9961        // Backward-compatible context key.
9962        let ctx_key = format!("delegation.{}.last_response", delegate_id);
9963        let _ = self.context_manager.set(
9964            &ctx_key,
9965            serde_json::Value::String(response.content.clone()),
9966        );
9967
9968        // Structured orchestration context.
9969        let _ = self.context_manager.set(
9970            "orchestration",
9971            serde_json::json!({
9972                "type": "delegate",
9973                "agent": delegate_id,
9974                "state": state_name,
9975                "response": response.content,
9976                "duration_ms": duration_ms,
9977            }),
9978        );
9979
9980        self.commit_root_user_message(input).await?;
9981
9982        // post_loop_processing records the assistant turn and evaluates transitions.
9983        // apply_post_loop_result handles NeedsRedispatch by re-entering run_loop_internal.
9984        let post_result = self
9985            .post_loop_processing(
9986                input,
9987                format!("[Delegated to {}]: {}", delegate_id, response.content),
9988            )
9989            .await?;
9990        let final_content = self.apply_post_loop_result(input, post_result).await?;
9991
9992        let mut result = AgentResponse::new(final_content);
9993
9994        let metadata = serde_json::json!({
9995            "orchestration": {
9996                "type": "delegate",
9997                "agent": delegate_id,
9998                "state": state_name,
9999                "response": response.content,
10000                "duration_ms": duration_ms,
10001            }
10002        });
10003        result.metadata = Some(
10004            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10005                metadata,
10006            )
10007            .unwrap_or_default(),
10008        );
10009
10010        self.finish_turn_if_root(&result).await?;
10011        Ok(result)
10012    }
10013
10014    // Handle concurrent execution: run multiple registry agents in parallel and aggregate.
10015    async fn handle_concurrent_state(
10016        &self,
10017        input: &str,
10018        config: &ai_agents_state::ConcurrentStateConfig,
10019    ) -> Result<AgentResponse> {
10020        use std::time::Instant;
10021
10022        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10023            AgentError::Config(
10024                "Concurrent state requires an agent registry. Add a spawner section.".into(),
10025            )
10026        })?;
10027
10028        // Render input template if provided, otherwise use the raw input.
10029        // Uses direct minijinja rendering (same approach as pipeline) so variables
10030        // are top-level: {{ user_input }}, not {{ context.user_input }}.
10031        // Enrich input with parent conversation history when context_mode is set.
10032        let context_mode = config.context_mode.clone().unwrap_or_default();
10033        let context_input = self
10034            .observe_purpose(
10035                ObservationPurpose::OrchestrationRouting,
10036                crate::orchestration::context::prepare_delegate_input(
10037                    input,
10038                    &context_mode,
10039                    &*self.memory,
10040                    self.llm_registry.get("router").ok().as_deref(),
10041                ),
10042            )
10043            .await?;
10044
10045        let effective_input = if let Some(ref tmpl) = config.input {
10046            render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10047                .unwrap_or_else(|_| context_input.clone())
10048        } else {
10049            context_input
10050        };
10051
10052        let start = Instant::now();
10053
10054        let llm_name = config
10055            .aggregation
10056            .synthesizer_llm
10057            .as_deref()
10058            .unwrap_or("router");
10059        let llm_provider = self.llm_registry.get(llm_name).ok();
10060
10061        let vote_parallelism = if self.runtime_config.optimization.enabled
10062            && self
10063                .runtime_config
10064                .optimization
10065                .parallel_orchestration_vote_extraction
10066        {
10067            Some(self.runtime_config.optimization.max_parallel_runtime_tasks)
10068        } else {
10069            None
10070        };
10071
10072        let result = self
10073            .observe_purpose(
10074                ObservationPurpose::OrchestrationAggregation,
10075                scope_actor_context(
10076                    self.outbound_actor_context(),
10077                    crate::orchestration::concurrent(
10078                        registry,
10079                        &effective_input,
10080                        &config.agents,
10081                        &config.aggregation,
10082                        llm_provider.as_deref(),
10083                        config.min_required,
10084                        config.timeout_ms,
10085                        config.on_partial_failure.clone(),
10086                        vote_parallelism,
10087                    ),
10088                ),
10089            )
10090            .await?;
10091
10092        let duration_ms = start.elapsed().as_millis() as u64;
10093        let agent_ids: Vec<String> = config.agents.iter().map(|a| a.id().to_string()).collect();
10094        let strategy = format!("{:?}", config.aggregation.strategy);
10095        self.hooks
10096            .on_concurrent_complete(&agent_ids, &strategy, duration_ms)
10097            .await;
10098
10099        // Backward-compatible context key.
10100        let _ = self.context_manager.set(
10101            "concurrent.result",
10102            serde_json::Value::String(result.response.content.clone()),
10103        );
10104
10105        // Build per-agent result data for context and metadata.
10106        let agents_json: Vec<serde_json::Value> = result
10107            .agent_results
10108            .iter()
10109            .map(|ar| {
10110                serde_json::json!({
10111                    "id": ar.agent_id,
10112                    "response": ar.response.as_ref().map(|r| r.content.as_str()),
10113                    "success": ar.success,
10114                    "error": ar.error,
10115                    "duration_ms": ar.duration_ms,
10116                })
10117            })
10118            .collect();
10119
10120        // Structured orchestration context with per-agent results.
10121        let _ = self.context_manager.set(
10122            "orchestration",
10123            serde_json::json!({
10124                "type": "concurrent",
10125                "result": result.response.content,
10126                "strategy": strategy,
10127                "agents": agents_json,
10128                "duration_ms": duration_ms,
10129            }),
10130        );
10131
10132        self.commit_root_user_message(input).await?;
10133
10134        let post_result = self
10135            .post_loop_processing(input, result.response.content.clone())
10136            .await?;
10137        let final_content = self.apply_post_loop_result(input, post_result).await?;
10138
10139        let mut response = AgentResponse::new(final_content);
10140        let metadata = serde_json::json!({
10141            "orchestration": {
10142                "type": "concurrent",
10143                "result": result.response.content,
10144                "strategy": strategy,
10145                "agents": agents_json,
10146                "duration_ms": duration_ms,
10147            }
10148        });
10149        response.metadata = Some(
10150            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10151                metadata,
10152            )
10153            .unwrap_or_default(),
10154        );
10155
10156        self.finish_turn_if_root(&response).await?;
10157        Ok(response)
10158    }
10159
10160    // Handle group chat: run a multi-turn multi-agent conversation.
10161    async fn handle_group_chat_state(
10162        &self,
10163        input: &str,
10164        config: &ai_agents_state::GroupChatStateConfig,
10165    ) -> Result<AgentResponse> {
10166        use std::time::Instant;
10167
10168        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10169            AgentError::Config(
10170                "Group chat state requires an agent registry. Add a spawner section.".into(),
10171            )
10172        })?;
10173
10174        let start = Instant::now();
10175
10176        let llm_provider = self.llm_registry.get("router").ok();
10177
10178        // Enrich input with parent conversation history when context_mode is set.
10179        let context_mode = config.context_mode.clone().unwrap_or_default();
10180        let context_input = self
10181            .observe_purpose(
10182                ObservationPurpose::OrchestrationRouting,
10183                crate::orchestration::context::prepare_delegate_input(
10184                    input,
10185                    &context_mode,
10186                    &*self.memory,
10187                    self.llm_registry.get("router").ok().as_deref(),
10188                ),
10189            )
10190            .await?;
10191
10192        // Render input template if provided, otherwise use the raw user message as topic.
10193        let effective_topic = if let Some(ref tmpl) = config.input {
10194            render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10195                .unwrap_or_else(|_| context_input.clone())
10196        } else {
10197            context_input
10198        };
10199
10200        let result = self
10201            .observe_purpose(
10202                ObservationPurpose::OrchestrationConversation,
10203                scope_actor_context(
10204                    self.outbound_actor_context(),
10205                    crate::orchestration::group_chat(
10206                        registry,
10207                        &effective_topic,
10208                        config,
10209                        llm_provider.as_deref(),
10210                        Some(&*self.hooks),
10211                    ),
10212                ),
10213            )
10214            .await?;
10215
10216        let duration_ms = start.elapsed().as_millis() as u64;
10217
10218        // Backward-compatible context key.
10219        let _ = self.context_manager.set(
10220            "group_chat.conclusion",
10221            serde_json::Value::String(result.response.content.clone()),
10222        );
10223
10224        // Build transcript data for context and metadata.
10225        let transcript_json: Vec<serde_json::Value> = result
10226            .transcript
10227            .iter()
10228            .map(|t| {
10229                serde_json::json!({
10230                    "speaker": t.speaker,
10231                    "round": t.round,
10232                    "content": t.content,
10233                })
10234            })
10235            .collect();
10236
10237        // Structured orchestration context with full transcript.
10238        let _ = self.context_manager.set(
10239            "orchestration",
10240            serde_json::json!({
10241                "type": "group_chat",
10242                "conclusion": result.response.content,
10243                "transcript": transcript_json,
10244                "rounds": result.rounds_completed,
10245                "termination": result.termination_reason,
10246                "duration_ms": duration_ms,
10247            }),
10248        );
10249
10250        self.commit_root_user_message(input).await?;
10251
10252        let post_result = self
10253            .post_loop_processing(input, result.response.content.clone())
10254            .await?;
10255        let final_content = self.apply_post_loop_result(input, post_result).await?;
10256
10257        let mut response = AgentResponse::new(final_content);
10258        let metadata = serde_json::json!({
10259            "orchestration": {
10260                "type": "group_chat",
10261                "conclusion": result.response.content,
10262                "transcript": transcript_json,
10263                "rounds": result.rounds_completed,
10264                "termination": result.termination_reason,
10265                "duration_ms": duration_ms,
10266            }
10267        });
10268        response.metadata = Some(
10269            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10270                metadata,
10271            )
10272            .unwrap_or_default(),
10273        );
10274
10275        self.finish_turn_if_root(&response).await?;
10276        Ok(response)
10277    }
10278
10279    // Handle pipeline: run agents sequentially with per-stage input templates.
10280    async fn handle_pipeline_state(
10281        &self,
10282        input: &str,
10283        config: &ai_agents_state::PipelineStateConfig,
10284    ) -> Result<AgentResponse> {
10285        use std::time::Instant;
10286
10287        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10288            AgentError::Config(
10289                "Pipeline state requires an agent registry. Add a spawner section.".into(),
10290            )
10291        })?;
10292
10293        let start = Instant::now();
10294
10295        let stages: Vec<crate::orchestration::PipelineStage> = config
10296            .stages
10297            .iter()
10298            .map(|entry| {
10299                let mut stage = crate::orchestration::PipelineStage::id(entry.id());
10300                if let Some(tmpl) = entry.input() {
10301                    stage = stage.with_input(tmpl);
10302                }
10303                stage
10304            })
10305            .collect();
10306
10307        // Enrich input with parent conversation history when context_mode is set.
10308        let context_mode = config.context_mode.clone().unwrap_or_default();
10309        let context_input = self
10310            .observe_purpose(
10311                ObservationPurpose::OrchestrationRouting,
10312                crate::orchestration::context::prepare_delegate_input(
10313                    input,
10314                    &context_mode,
10315                    &*self.memory,
10316                    self.llm_registry.get("router").ok().as_deref(),
10317                ),
10318            )
10319            .await?;
10320
10321        let context_values = self.build_context_with_overlays();
10322        let result = self
10323            .observe_purpose(
10324                ObservationPurpose::OrchestrationRouting,
10325                scope_actor_context(
10326                    self.outbound_actor_context(),
10327                    crate::orchestration::pipeline(
10328                        registry,
10329                        &context_input,
10330                        &stages,
10331                        config.timeout_ms,
10332                        Some(&*self.hooks),
10333                        Some(&context_values),
10334                    ),
10335                ),
10336            )
10337            .await?;
10338
10339        let duration_ms = start.elapsed().as_millis() as u64;
10340
10341        // Backward-compatible context key.
10342        let _ = self.context_manager.set(
10343            "pipeline.result",
10344            serde_json::Value::String(result.response.content.clone()),
10345        );
10346
10347        // Build per-stage data for context and metadata.
10348        let stages_json: Vec<serde_json::Value> = result
10349            .stage_outputs
10350            .iter()
10351            .map(|s| {
10352                serde_json::json!({
10353                    "agent_id": s.agent_id,
10354                    "output": s.output,
10355                    "duration_ms": s.duration_ms,
10356                    "skipped": s.skipped,
10357                })
10358            })
10359            .collect();
10360
10361        // Structured orchestration context.
10362        let _ = self.context_manager.set(
10363            "orchestration",
10364            serde_json::json!({
10365                "type": "pipeline",
10366                "result": result.response.content,
10367                "stages": stages_json,
10368                "duration_ms": duration_ms,
10369            }),
10370        );
10371
10372        self.commit_root_user_message(input).await?;
10373
10374        let post_result = self
10375            .post_loop_processing(input, result.response.content.clone())
10376            .await?;
10377        let final_content = self.apply_post_loop_result(input, post_result).await?;
10378
10379        let mut response = AgentResponse::new(final_content);
10380        let metadata = serde_json::json!({
10381            "orchestration": {
10382                "type": "pipeline",
10383                "result": result.response.content,
10384                "stages": stages_json,
10385                "duration_ms": duration_ms,
10386            }
10387        });
10388        response.metadata = Some(
10389            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10390                metadata,
10391            )
10392            .unwrap_or_default(),
10393        );
10394
10395        self.finish_turn_if_root(&response).await?;
10396        Ok(response)
10397    }
10398
10399    // Handle handoff: LLM-directed agent-to-agent control transfer.
10400    async fn handle_handoff_state(
10401        &self,
10402        input: &str,
10403        config: &ai_agents_state::HandoffStateConfig,
10404    ) -> Result<AgentResponse> {
10405        use std::time::Instant;
10406
10407        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10408            AgentError::Config(
10409                "Handoff state requires an agent registry. Add a spawner section.".into(),
10410            )
10411        })?;
10412
10413        let llm = self
10414            .llm_registry
10415            .get("router")
10416            .map_err(|_| AgentError::Config("Handoff state requires a router LLM.".into()))?;
10417
10418        let start = Instant::now();
10419
10420        // Enrich input with parent conversation history when context_mode is set.
10421        let context_mode = config.context_mode.clone().unwrap_or_default();
10422        let context_input = self
10423            .observe_purpose(
10424                ObservationPurpose::OrchestrationRouting,
10425                crate::orchestration::context::prepare_delegate_input(
10426                    input,
10427                    &context_mode,
10428                    &*self.memory,
10429                    self.llm_registry.get("router").ok().as_deref(),
10430                ),
10431            )
10432            .await?;
10433
10434        // Render input template if provided, otherwise forward the raw user message.
10435        let effective_input = if let Some(ref tmpl) = config.input {
10436            render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10437                .unwrap_or_else(|_| context_input.clone())
10438        } else {
10439            context_input
10440        };
10441
10442        let result = self
10443            .observe_purpose(
10444                ObservationPurpose::OrchestrationRouting,
10445                scope_actor_context(
10446                    self.outbound_actor_context(),
10447                    crate::orchestration::handoff(
10448                        registry,
10449                        &effective_input,
10450                        &config.initial_agent,
10451                        &config.available_agents,
10452                        config.max_handoffs,
10453                        llm.as_ref(),
10454                        Some(&*self.hooks),
10455                    ),
10456                ),
10457            )
10458            .await?;
10459
10460        let duration_ms = start.elapsed().as_millis() as u64;
10461
10462        // Backward-compatible context key.
10463        let _ = self.context_manager.set(
10464            "handoff.result",
10465            serde_json::Value::String(result.response.content.clone()),
10466        );
10467
10468        // Build handoff chain data for context and metadata.
10469        let chain_json: Vec<serde_json::Value> = result
10470            .handoff_chain
10471            .iter()
10472            .map(|h| {
10473                serde_json::json!({
10474                    "from": h.from_agent,
10475                    "to": h.to_agent,
10476                    "reason": h.reason,
10477                })
10478            })
10479            .collect();
10480
10481        // Structured orchestration context.
10482        let _ = self.context_manager.set(
10483            "orchestration",
10484            serde_json::json!({
10485                "type": "handoff",
10486                "result": result.response.content,
10487                "final_agent": result.final_agent,
10488                "handoff_chain": chain_json,
10489                "duration_ms": duration_ms,
10490            }),
10491        );
10492
10493        self.commit_root_user_message(input).await?;
10494
10495        let post_result = self
10496            .post_loop_processing(input, result.response.content.clone())
10497            .await?;
10498        let final_content = self.apply_post_loop_result(input, post_result).await?;
10499
10500        let mut response = AgentResponse::new(final_content);
10501        let metadata = serde_json::json!({
10502            "orchestration": {
10503                "type": "handoff",
10504                "result": result.response.content,
10505                "final_agent": result.final_agent,
10506                "handoff_chain": chain_json,
10507                "duration_ms": duration_ms,
10508            }
10509        });
10510        response.metadata = Some(
10511            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10512                metadata,
10513            )
10514            .unwrap_or_default(),
10515        );
10516
10517        self.finish_turn_if_root(&response).await?;
10518        Ok(response)
10519    }
10520
10521    // run_loop_internal: blocking (non-streaming) agent pipeline.
10522    async fn run_loop_internal(&self, input: &str) -> Result<AgentResponse> {
10523        self.begin_root_turn();
10524        // Resolve actor_id from context, reload facts if actor changed, bump counter.
10525        self.pre_turn_session_lifecycle().await;
10526
10527        let input_data = self.process_input(input).await?;
10528        self.update_active_turn_context(&input_data.content, input_data.context.clone());
10529
10530        // Inject process context (detect/extract results) into agent context
10531        // so system prompt templates can use {{ context.detected_language }} etc.
10532        for (key, value) in &input_data.context {
10533            let _ = self.context_manager.set(key, value.clone());
10534        }
10535
10536        if input_data.metadata.rejected {
10537            let reason = input_data
10538                .metadata
10539                .rejection_reason
10540                .unwrap_or_else(|| "Input rejected".to_string());
10541            warn!(reason = %reason, "Input rejected");
10542            let response = AgentResponse::new(reason);
10543            self.finish_turn_if_root(&response).await?;
10544            return Ok(response);
10545        }
10546
10547        let processed_input = &input_data.content;
10548
10549        if let Some(response) = self.try_pre_response_transition(processed_input).await? {
10550            return Ok(response);
10551        }
10552
10553        // Handle orchestration states (delegate, concurrent, group_chat, pipeline, handoff).
10554        if let Some(ref sm) = self.state_machine
10555            && let Some(def) = sm.current_definition()
10556        {
10557            if let Some(ref delegate_id) = def.delegate {
10558                return self
10559                    .handle_delegated_state(processed_input, delegate_id, &def)
10560                    .await;
10561            }
10562            if let Some(ref concurrent_config) = def.concurrent {
10563                return self
10564                    .handle_concurrent_state(processed_input, concurrent_config)
10565                    .await;
10566            }
10567            if let Some(ref group_chat_config) = def.group_chat {
10568                return self
10569                    .handle_group_chat_state(processed_input, group_chat_config)
10570                    .await;
10571            }
10572            if let Some(ref pipeline_config) = def.pipeline {
10573                return self
10574                    .handle_pipeline_state(processed_input, pipeline_config)
10575                    .await;
10576            }
10577            if let Some(ref handoff_config) = def.handoff {
10578                return self
10579                    .handle_handoff_state(processed_input, handoff_config)
10580                    .await;
10581            }
10582        }
10583
10584        //
10585        // The speculative future is boxed to keep the runtime future size manageable.
10586        // Removing the box can overflow small test stacks because this function is recursive through redispatch.
10587        //
10588        if let Some(response) =
10589            Box::pin(self.try_speculative_branches(processed_input, &input_data.context)).await?
10590        {
10591            return Ok(response);
10592        }
10593
10594        match self.try_skill_route(processed_input).await? {
10595            SkillRouteResult::Response { skill_id, content } => {
10596                self.commit_root_user_message(processed_input).await?;
10597                return self
10598                    .handle_skill_response(processed_input, &skill_id, content, &input_data.context)
10599                    .await;
10600            }
10601            SkillRouteResult::NeedsClarification {
10602                response,
10603                ownership,
10604            } => {
10605                let admission = self
10606                    .admit_optional_disambiguation_ownership(ownership)
10607                    .await?;
10608                self.commit_root_user_message(processed_input).await?;
10609                if let Some(q) = response
10610                    .metadata
10611                    .as_ref()
10612                    .and_then(|m| m.get("disambiguation"))
10613                    .and_then(|d| d.get("status"))
10614                    .and_then(|s| s.as_str())
10615                    && q == "awaiting_clarification"
10616                {
10617                    // Store the clarification question in memory so the next turn
10618                    // can be handled as a clarification response.
10619                    self.memory
10620                        .add_message(ChatMessage::assistant(&response.content))
10621                        .await?;
10622                }
10623                drop(admission);
10624                self.finish_turn_if_root(&response).await?;
10625                return Ok(response);
10626            }
10627            SkillRouteResult::NoMatch => {} // continue to normal LLM chat
10628        }
10629
10630        let effective_reasoning = self.get_effective_reasoning_config();
10631        let reasoning_mode = self.determine_reasoning_mode(processed_input).await?;
10632        let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
10633
10634        info!(
10635            reasoning_mode = ?reasoning_mode,
10636            auto_detected = auto_detected,
10637            reflection_enabled = ?self.reflection_config.enabled,
10638            "Reasoning mode determined"
10639        );
10640
10641        if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
10642            self.commit_root_user_message(processed_input).await?;
10643            return self
10644                .handle_plan_and_execute(processed_input, &input_data.context, auto_detected)
10645                .await;
10646        }
10647
10648        self.commit_root_user_message(processed_input).await?;
10649
10650        let mut iterations = 0u32;
10651        let mut all_tool_calls: Vec<ToolCall> = Vec::new();
10652        let mut thinking_content: Option<String> = None;
10653
10654        let llm = self.get_state_llm()?;
10655
10656        loop {
10657            // When reasoning is active, cap iterations at the reasoning-specific limit.
10658            let effective_max = if reasoning_mode != ReasoningMode::None {
10659                let rc = self.get_effective_reasoning_config();
10660                self.max_iterations.min(rc.max_iterations)
10661            } else {
10662                self.max_iterations
10663            };
10664
10665            if iterations >= effective_max {
10666                let err = AgentError::Other(format!("Max iterations ({}) exceeded", effective_max));
10667                self.hooks.on_error(&err).await;
10668                error!(iterations = iterations, "Max iterations exceeded");
10669                return Err(err);
10670            }
10671            iterations += 1;
10672            *self.iteration_count.write() = iterations;
10673
10674            debug!(iteration = iterations, max = effective_max, "LLM call");
10675
10676            let protocol = self.main_tool_protocol(llm.as_ref(), false).await?;
10677            let mut messages = self
10678                .build_messages_internal(true, None, protocol.choice.is_none())
10679                .await?;
10680            self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
10681
10682            self.hooks.on_llm_start(&messages).await;
10683            let llm_start = Instant::now();
10684            let response = self
10685                .complete_main_llm_with_recovery(Arc::clone(&llm), &messages, &protocol)
10686                .await?;
10687
10688            let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
10689            self.hooks.on_llm_complete(&response, llm_duration_ms).await;
10690
10691            let content = response.content.trim();
10692
10693            if let Some(tool_calls) = self.parse_main_tool_calls(content, &protocol) {
10694                match self
10695                    .handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
10696                    .await?
10697                {
10698                    ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
10699                    ToolCallOutcome::Rejected(resp) => {
10700                        self.finish_turn_if_root(&resp).await?;
10701                        return Ok(resp);
10702                    }
10703                }
10704            }
10705
10706            let (extracted_thinking, answer) = self.extract_thinking(content);
10707            if extracted_thinking.is_some() {
10708                thinking_content = extracted_thinking;
10709            }
10710
10711            let output_data = self.process_output(&answer, &input_data.context).await?;
10712
10713            let mut final_content = if output_data.metadata.rejected {
10714                output_data
10715                    .metadata
10716                    .rejection_reason
10717                    .unwrap_or_else(|| answer.to_string())
10718            } else {
10719                output_data.content
10720            };
10721
10722            // Run reflection (blocking LLM calls for retries)
10723            let reflection_metadata;
10724            (final_content, reflection_metadata) = self
10725                .run_reflection(&*llm, processed_input, final_content)
10726                .await?;
10727
10728            final_content =
10729                self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
10730
10731            // Post-loop: memory, transitions, post-transition re-generation.
10732            // apply_post_loop_result handles NeedsRedispatch by re-entering
10733            // run_loop_internal so the new state's full dispatch activates.
10734            let final_content = {
10735                let result = self
10736                    .post_loop_processing(processed_input, final_content)
10737                    .await?;
10738                self.apply_post_loop_result(processed_input, result).await?
10739            };
10740
10741            let reflected = reflection_metadata.is_some();
10742            let reasoning_mode_debug = format!("{:?}", reasoning_mode);
10743
10744            let response = self.build_agent_response(AgentResponseParts {
10745                content: final_content,
10746                all_tool_calls,
10747                reasoning_mode,
10748                auto_detected,
10749                iterations,
10750                thinking: thinking_content,
10751                reflection_metadata,
10752            });
10753
10754            self.finish_turn_if_root(&response).await?;
10755
10756            let tool_call_count = response.tool_calls.as_ref().map(|tc| tc.len()).unwrap_or(0);
10757            info!(
10758                tool_calls = tool_call_count,
10759                response_len = response.content.len(),
10760                reasoning_mode = %reasoning_mode_debug,
10761                reflected = reflected,
10762                "Chat completed"
10763            );
10764            return Ok(response);
10765        }
10766    }
10767
10768    async fn generate_buffered_streaming_draft(
10769        &self,
10770        processed_input: &str,
10771        routing_resolved: Arc<AtomicBool>,
10772    ) -> Result<StreamingDraftResult> {
10773        let llm = self.get_state_llm()?;
10774        if llm.configured_tool_choice().is_some() {
10775            let draft = self
10776                .generate_main_response_draft(processed_input, &ReasoningMode::None)
10777                .await?;
10778            return Ok(StreamingDraftResult::new(draft, Vec::new()));
10779        }
10780        let messages = self.build_messages_for_draft(processed_input).await?;
10781        let mut stream = self
10782            .observe_purpose(
10783                ObservationPurpose::MainResponse,
10784                llm.complete_stream(&messages, None),
10785            )
10786            .await
10787            .map_err(|e| AgentError::LLM(e.to_string()))?;
10788        let mut buffer = crate::optimization::StreamBranchBuffer::new(self.streaming.buffer_size)?;
10789        let mut chunks = Vec::new();
10790        let mut accumulated = String::new();
10791        while let Some(chunk_result) = stream.next().await {
10792            let chunk = chunk_result.map_err(|e| AgentError::LLM(e.to_string()))?;
10793            accumulated.push_str(&chunk.delta);
10794            let stream_chunk = StreamChunk::content(chunk.delta);
10795            if routing_resolved.load(Ordering::SeqCst) {
10796                chunks.push(stream_chunk);
10797            } else {
10798                buffer.push(stream_chunk)?;
10799            }
10800        }
10801        chunks.splice(0..0, buffer.drain());
10802        let content = accumulated.trim().to_string();
10803        let draft = if let Some(calls) = self.parse_tool_calls(&content) {
10804            MainResponseDraft::ToolCalls {
10805                raw_content: content,
10806                calls,
10807                thinking: None,
10808            }
10809        } else {
10810            MainResponseDraft::Text {
10811                raw_content: content,
10812                thinking: None,
10813            }
10814        };
10815        Ok(StreamingDraftResult::new(draft, chunks))
10816    }
10817
10818    async fn try_buffered_streaming_branches(
10819        &self,
10820        processed_input: &str,
10821        input_context: &HashMap<String, Value>,
10822    ) -> Result<Option<(AgentResponse, Vec<StreamChunk>)>> {
10823        let optimization = &self.runtime_config.optimization;
10824        if !optimization.enabled {
10825            return Ok(None);
10826        }
10827        let transition_enabled =
10828            optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
10829        if !transition_enabled {
10830            return Ok(None);
10831        }
10832        let mut branch_scheduler =
10833            TurnBranchScheduler::new(optimization.max_parallel_runtime_tasks)?;
10834        if !branch_scheduler.reserve_task() {
10835            return Ok(None);
10836        }
10837        if !self
10838            .reserve_active_speculative_llm_call(RuntimeOptimizationKind::BufferedStreamingRouting)
10839        {
10840            branch_scheduler.release_task();
10841            return Ok(None);
10842        }
10843        if !branch_scheduler.reserve_task() {
10844            branch_scheduler.release_task();
10845            return Ok(None);
10846        }
10847        let mut main_branch = RuntimeBranch::new(
10848            RuntimeTaskPurpose::MainResponse,
10849            RuntimeOptimizationKind::BufferedStreamingRouting,
10850            RuntimeTaskPriority::Normal,
10851            RuntimeCommitBehavior::FinalResponse,
10852        );
10853        let mut transition_branch = RuntimeBranch::new(
10854            RuntimeTaskPurpose::StateTransition,
10855            RuntimeOptimizationKind::ParallelStateTransition,
10856            RuntimeTaskPriority::Critical,
10857            RuntimeCommitBehavior::TransitionDecision,
10858        );
10859        let main_id = main_branch.branch_id();
10860        let transition_id = transition_branch.branch_id();
10861        let routing_resolved = Arc::new(AtomicBool::new(false));
10862        let mut main_future =
10863            Box::pin(crate::optimization::observability::with_branch_observation(
10864                &main_id,
10865                RuntimeOptimizationKind::BufferedStreamingRouting,
10866                RuntimeCommitBehavior::FinalResponse,
10867                self.generate_buffered_streaming_draft(
10868                    processed_input,
10869                    Arc::clone(&routing_resolved),
10870                ),
10871            ));
10872        let mut transition_future =
10873            Box::pin(crate::optimization::observability::with_branch_observation(
10874                &transition_id,
10875                RuntimeOptimizationKind::ParallelStateTransition,
10876                RuntimeCommitBehavior::TransitionDecision,
10877                self.select_parallel_transition_candidate(processed_input),
10878            ));
10879        let mut main_pending = true;
10880        let mut transition_pending = true;
10881        let mut main_result: Option<Result<StreamingDraftResult>> = None;
10882        let mut transition_finalized = false;
10883        let mut transition_candidate: Option<TransitionCandidate> = None;
10884        loop {
10885            if let Some(candidate) = transition_candidate.take() {
10886                if self
10887                    .approve_transition_target(&candidate.from_state, candidate.target())
10888                    .await?
10889                {
10890                    // Drop the stale stream future before transition side effects or redispatch reuse the provider.
10891                    drop(main_future);
10892                    drop(transition_future);
10893                    self.finalize_branch_loss(
10894                        &main_id,
10895                        RuntimeOptimizationKind::BufferedStreamingRouting,
10896                        RuntimeCommitBehavior::FinalResponse,
10897                        main_pending,
10898                        main_result.as_ref().map(|result| result.is_err()),
10899                    );
10900                    if !self
10901                        .apply_pre_response_transition_candidate(
10902                            &candidate,
10903                            &HashMap::new(),
10904                            processed_input,
10905                        )
10906                        .await?
10907                    {
10908                        self.finalize_optional_branch(
10909                            &transition_id,
10910                            RuntimeOptimizationKind::ParallelStateTransition,
10911                            RuntimeCommitBehavior::TransitionDecision,
10912                            "discarded",
10913                            false,
10914                        );
10915                        return Ok(None);
10916                    }
10917                    self.finalize_optional_branch(
10918                        &transition_id,
10919                        RuntimeOptimizationKind::ParallelStateTransition,
10920                        RuntimeCommitBehavior::TransitionDecision,
10921                        "committed",
10922                        true,
10923                    );
10924                    let response = self.redispatch_current_state(processed_input).await?;
10925                    return Ok(Some((
10926                        response.clone(),
10927                        vec![StreamChunk::content(response.content)],
10928                    )));
10929                }
10930                self.finalize_optional_branch(
10931                    &transition_id,
10932                    RuntimeOptimizationKind::ParallelStateTransition,
10933                    RuntimeCommitBehavior::TransitionDecision,
10934                    "discarded",
10935                    false,
10936                );
10937                routing_resolved.store(true, Ordering::SeqCst);
10938                transition_finalized = true;
10939            }
10940            if transition_finalized && let Some(result) = main_result.take() {
10941                let stream_draft = match result {
10942                    Ok(stream_draft) => stream_draft,
10943                    Err(error) => {
10944                        self.finalize_optional_branch(
10945                            &main_id,
10946                            RuntimeOptimizationKind::BufferedStreamingRouting,
10947                            RuntimeCommitBehavior::FinalResponse,
10948                            "failed",
10949                            false,
10950                        );
10951                        return Err(error);
10952                    }
10953                };
10954                let raw_draft_content = stream_draft.draft.raw_content().to_string();
10955                let buffered_chunks = stream_draft.chunks;
10956                self.finalize_optional_branch(
10957                    &main_id,
10958                    RuntimeOptimizationKind::BufferedStreamingRouting,
10959                    RuntimeCommitBehavior::FinalResponse,
10960                    "committed",
10961                    true,
10962                );
10963                let response = self
10964                    .commit_main_response_draft(
10965                        processed_input,
10966                        input_context,
10967                        stream_draft.draft,
10968                        ReasoningMode::None,
10969                        false,
10970                    )
10971                    .await?;
10972                let chunks = if response.content == raw_draft_content {
10973                    buffered_chunks
10974                } else {
10975                    vec![StreamChunk::content(response.content.clone())]
10976                };
10977                return Ok(Some((response, chunks)));
10978            }
10979            tokio::select! {
10980                result = &mut main_future, if main_pending => {
10981                    main_pending = false;
10982                    main_branch.transition_to(RuntimeBranchStatus::Completed)?;
10983                    main_result = Some(result);
10984                }
10985                result = &mut transition_future, if transition_pending => {
10986                    transition_pending = false;
10987                    transition_branch.transition_to(RuntimeBranchStatus::Completed)?;
10988                    match result {
10989                        Ok(ParallelTransitionSelection::Candidate(candidate)) => {
10990                            transition_candidate = Some(candidate)
10991                        }
10992                        Ok(ParallelTransitionSelection::NoMatch) => {
10993                            self.finalize_optional_branch(
10994                                &transition_id,
10995                                RuntimeOptimizationKind::ParallelStateTransition,
10996                                RuntimeCommitBehavior::TransitionDecision,
10997                                "discarded",
10998                                false,
10999                            );
11000                            routing_resolved.store(true, Ordering::SeqCst);
11001                            transition_finalized = true;
11002                        }
11003                        Ok(ParallelTransitionSelection::ReservationExhausted) => {
11004                            self.finalize_optional_branch(
11005                                &transition_id,
11006                                RuntimeOptimizationKind::ParallelStateTransition,
11007                                RuntimeCommitBehavior::TransitionDecision,
11008                                "cancelled",
11009                                false,
11010                            );
11011                            routing_resolved.store(true, Ordering::SeqCst);
11012                            self.finalize_branch_loss(
11013                                &main_id,
11014                                RuntimeOptimizationKind::BufferedStreamingRouting,
11015                                RuntimeCommitBehavior::FinalResponse,
11016                                main_pending,
11017                                main_result.as_ref().map(|result| result.is_err()),
11018                            );
11019                            return Ok(None);
11020                        }
11021                        Err(_) => {
11022                            self.finalize_optional_branch(
11023                                &transition_id,
11024                                RuntimeOptimizationKind::ParallelStateTransition,
11025                                RuntimeCommitBehavior::TransitionDecision,
11026                                "failed",
11027                                false,
11028                            );
11029                            routing_resolved.store(true, Ordering::SeqCst);
11030                            transition_finalized = true;
11031                        }
11032                    }
11033                }
11034            }
11035        }
11036    }
11037
11038    /// Streaming agent pipeline
11039    /// Uses all the same shared helpers as run_loop_internal.
11040    /// The ONLY difference: LLM calls use complete_stream() + yield deltas.
11041    fn run_loop_internal_stream<'a>(
11042        &'a self,
11043        input: &'a str,
11044    ) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
11045        let include_tool_events = self.streaming.include_tool_events;
11046        let include_state_events = self.streaming.include_state_events;
11047
11048        Box::pin(async_stream::stream! {
11049            self.begin_root_turn();
11050            // Parity with non-stream: resolve actor from context and load facts if changed.
11051            self.pre_turn_session_lifecycle().await;
11052
11053            let input_data = match self.process_input(input).await {
11054                Ok(data) => data,
11055                Err(e) => {
11056                    yield StreamChunk::error(e.to_string());
11057                    return;
11058                }
11059            };
11060            self.update_active_turn_context(&input_data.content, input_data.context.clone());
11061
11062            // Inject process context (detect/extract results) into agent context
11063            for (key, value) in &input_data.context {
11064                let _ = self.context_manager.set(key, value.clone());
11065            }
11066
11067            if input_data.metadata.rejected {
11068                let reason = input_data
11069                    .metadata
11070                    .rejection_reason
11071                    .unwrap_or_else(|| "Input rejected".to_string());
11072                warn!(reason = %reason, "Input rejected (stream)");
11073                yield StreamChunk::error(reason);
11074                return;
11075            }
11076
11077            let processed_input = &input_data.content;
11078
11079            if self.runtime_config.optimization.enabled
11080                && matches!(
11081                    self.runtime_config.optimization.streaming_policy,
11082                    crate::optimization::StreamingOptimizationPolicy::BufferUntilRoutingDone
11083                )
11084            {
11085                //
11086                // Buffered routing keeps stale stream output hidden until a branch winner is known.
11087                // The boxed future prevents this stream state machine from becoming too large.
11088                //
11089                match Box::pin(self.try_buffered_streaming_branches(processed_input, &input_data.context)).await {
11090                    Ok(Some((_response, chunks))) => {
11091                        for chunk in chunks {
11092                            yield chunk;
11093                        }
11094                        yield StreamChunk::Done {};
11095                        return;
11096                    }
11097                    Ok(None) => {}
11098                    Err(e) => {
11099                        yield StreamChunk::error(e.to_string());
11100                        return;
11101                    }
11102                }
11103            }
11104
11105            if self.runtime_config.optimization.enabled
11106                && matches!(
11107                    self.runtime_config.optimization.streaming_policy,
11108                    crate::optimization::StreamingOptimizationPolicy::PreflightOnly
11109                )
11110            {
11111                match self.try_pre_response_transition(processed_input).await {
11112                    Ok(Some(response)) => {
11113                        yield StreamChunk::content(&response.content);
11114                        yield StreamChunk::Done {};
11115                        return;
11116                    }
11117                    Ok(None) => {}
11118                    Err(e) => {
11119                        yield StreamChunk::error(e.to_string());
11120                        return;
11121                    }
11122                }
11123            }
11124
11125            // Handle orchestration states in streaming mode.
11126            if let Some(ref sm) = self.state_machine
11127                && let Some(def) = sm.current_definition()
11128            {
11129                    let orchestration_result = if let Some(ref delegate_id) = def.delegate {
11130                        Some(self.handle_delegated_state(processed_input, delegate_id, &def).await)
11131                    } else if let Some(ref concurrent_config) = def.concurrent {
11132                        Some(self.handle_concurrent_state(processed_input, concurrent_config).await)
11133                    } else if let Some(ref group_chat_config) = def.group_chat {
11134                        Some(self.handle_group_chat_state(processed_input, group_chat_config).await)
11135                    } else if let Some(ref pipeline_config) = def.pipeline {
11136                        Some(self.handle_pipeline_state(processed_input, pipeline_config).await)
11137                    } else if let Some(ref handoff_config) = def.handoff {
11138                        Some(self.handle_handoff_state(processed_input, handoff_config).await)
11139                    } else {
11140                        None
11141                    };
11142
11143                    if let Some(result) = orchestration_result {
11144                        match result {
11145                            Ok(response) => {
11146                                yield StreamChunk::content(&response.content);
11147                                yield StreamChunk::Done {};
11148                            }
11149                            Err(e) => {
11150                                yield StreamChunk::error(e.to_string());
11151                            }
11152                        }
11153                        return;
11154                    }
11155                }
11156
11157            // Skill routing
11158            match self.try_skill_route(processed_input).await {
11159                Ok(SkillRouteResult::Response { skill_id, content }) => {
11160                    if let Err(e) = self.commit_root_user_message(processed_input).await {
11161                        yield StreamChunk::error(e.to_string());
11162                        return;
11163                    }
11164                    match self.handle_skill_response(processed_input, &skill_id, content, &input_data.context).await {
11165                        Ok(resp) => {
11166                            yield StreamChunk::content(&resp.content);
11167                            yield StreamChunk::Done {};
11168                            return;
11169                        }
11170                        Err(e) => {
11171                            yield StreamChunk::error(e.to_string());
11172                            return;
11173                        }
11174                    }
11175                }
11176                Ok(SkillRouteResult::NeedsClarification {
11177                    response,
11178                    ownership,
11179                }) => {
11180                    let admission = match self
11181                        .admit_optional_disambiguation_ownership(ownership)
11182                        .await
11183                    {
11184                        Ok(admission) => admission,
11185                        Err(e) => {
11186                            yield StreamChunk::error(e.to_string());
11187                            return;
11188                        }
11189                    };
11190                    if let Err(e) = self.commit_root_user_message(processed_input).await {
11191                        yield StreamChunk::error(e.to_string());
11192                        return;
11193                    }
11194                    let _ = self.memory.add_message(ChatMessage::assistant(&response.content)).await;
11195                    drop(admission);
11196                    if let Err(e) = self.finish_turn_if_root(&response).await {
11197                        yield StreamChunk::error(e.to_string());
11198                        return;
11199                    }
11200                    yield StreamChunk::content(&response.content);
11201                    yield StreamChunk::Done {};
11202                    return;
11203                }
11204                Ok(SkillRouteResult::NoMatch) => {} // no skill matched, continue
11205                Err(e) => {
11206                    yield StreamChunk::error(e.to_string());
11207                    return;
11208                }
11209            }
11210
11211            // Reasoning mode determination
11212            let effective_reasoning = self.get_effective_reasoning_config();
11213            let reasoning_mode = match self.determine_reasoning_mode(processed_input).await {
11214                Ok(mode) => mode,
11215                Err(e) => {
11216                    yield StreamChunk::error(e.to_string());
11217                    return;
11218                }
11219            };
11220            let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
11221
11222            info!(
11223                reasoning_mode = ?reasoning_mode,
11224                auto_detected = auto_detected,
11225                "Reasoning mode determined (stream)"
11226            );
11227
11228            // Plan-and-Execute: yield final result as single chunk (not token-by-token)
11229            if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
11230                if let Err(e) = self.commit_root_user_message(processed_input).await {
11231                    yield StreamChunk::error(e.to_string());
11232                    return;
11233                }
11234                match self.handle_plan_and_execute(processed_input, &input_data.context, auto_detected).await {
11235                    Ok(resp) => {
11236                        yield StreamChunk::content(&resp.content);
11237                        yield StreamChunk::Done {};
11238                        return;
11239                    }
11240                    Err(e) => {
11241                        yield StreamChunk::error(e.to_string());
11242                        return;
11243                    }
11244                }
11245            }
11246
11247            if let Err(e) = self.commit_root_user_message(processed_input).await {
11248                yield StreamChunk::error(e.to_string());
11249                return;
11250            }
11251
11252            let llm = match self.get_state_llm() {
11253                Ok(llm) => llm,
11254                Err(e) => {
11255                    yield StreamChunk::error(e.to_string());
11256                    return;
11257                }
11258            };
11259
11260            let mut iterations = 0u32;
11261            let mut all_tool_calls: Vec<ToolCall> = Vec::new();
11262            let mut thinking_content: Option<String> = None;
11263
11264            loop {
11265                // When reasoning is active, cap iterations at the reasoning-specific limit.
11266                let effective_max = if reasoning_mode != ReasoningMode::None {
11267                    let rc = self.get_effective_reasoning_config();
11268                    self.max_iterations.min(rc.max_iterations)
11269                } else {
11270                    self.max_iterations
11271                };
11272
11273                if iterations >= effective_max {
11274                    let err_msg = format!("Max iterations ({}) exceeded", effective_max);
11275                    let err = AgentError::Other(err_msg.clone());
11276                    self.hooks.on_error(&err).await;
11277                    error!(iterations = iterations, "Max iterations exceeded (stream)");
11278                    yield StreamChunk::error(err_msg);
11279                    return;
11280                }
11281                iterations += 1;
11282                *self.iteration_count.write() = iterations;
11283
11284                debug!(iteration = iterations, max = effective_max, "LLM call (stream)");
11285
11286                let protocol = match self.main_tool_protocol(llm.as_ref(), false).await {
11287                    Ok(protocol) => protocol,
11288                    Err(e) => {
11289                        yield StreamChunk::error(e.to_string());
11290                        return;
11291                    }
11292                };
11293                let mut messages = match self
11294                    .build_messages_internal(true, None, protocol.choice.is_none())
11295                    .await
11296                {
11297                    Ok(m) => m,
11298                    Err(e) => {
11299                        yield StreamChunk::error(e.to_string());
11300                        return;
11301                    }
11302                };
11303                self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
11304
11305                self.hooks.on_llm_start(&messages).await;
11306                let llm_start = Instant::now();
11307
11308                // Check if reflection is active — if so, suppress streaming for this LLM call
11309                // because we may need to retry and the user would see a stale first attempt.
11310                let reflection_active = self
11311                    .should_reflect(processed_input, "")
11312                    .await
11313                    .unwrap_or_default();
11314
11315                let buffered_decision = reflection_active || protocol.choice.is_some();
11316                let content = if buffered_decision {
11317                    //
11318                    // Explicit tool choice buffers the provider decision so no text or tool call is visible before the runtime validates and commits it.
11319                    //
11320                    let response = match self
11321                        .complete_main_llm_with_recovery(
11322                            Arc::clone(&llm),
11323                            &messages,
11324                            &protocol,
11325                        )
11326                        .await
11327                    {
11328                        Ok(r) => r,
11329                        Err(e) => {
11330                            yield StreamChunk::error(e.to_string());
11331                            return;
11332                        }
11333                    };
11334                    let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
11335                    self.hooks.on_llm_complete(&response, llm_duration_ms).await;
11336                    response.content.trim().to_string()
11337                } else {
11338                    // Streaming LLM call
11339                    let llm_stream = match self
11340                        .observe_purpose(
11341                            ObservationPurpose::MainResponse,
11342                            llm.complete_stream(&messages, None),
11343                        )
11344                        .await
11345                    {
11346                        Ok(s) => s,
11347                        Err(e) => {
11348                            yield StreamChunk::error(e.to_string());
11349                            return;
11350                        }
11351                    };
11352                    let mut accumulated = String::new();
11353                    let mut stream_inner = llm_stream;
11354                    while let Some(chunk_result) = stream_inner.next().await {
11355                        match chunk_result {
11356                            Ok(chunk) => {
11357                                accumulated.push_str(&chunk.delta);
11358                                yield StreamChunk::content(chunk.delta);
11359                            }
11360                            Err(e) => {
11361                                yield StreamChunk::error(e.to_string());
11362                                return;
11363                            }
11364                        }
11365                    }
11366                    let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
11367                    // Construct LLMResponse for hooks
11368                    let llm_response = ai_agents_core::LLMResponse::new(
11369                        accumulated.trim(),
11370                        ai_agents_core::FinishReason::Stop,
11371                    );
11372                    self.hooks.on_llm_complete(&llm_response, llm_duration_ms).await;
11373                    accumulated.trim().to_string()
11374                };
11375
11376                // Tool call handling
11377                if let Some(tool_calls) = self.parse_main_tool_calls(&content, &protocol) {
11378                    let native_tool_call = Self::is_native_tool_call_content(&content);
11379                    // Emit tool events for streaming
11380                    // First check transitions (same as blocking path)
11381                    let transition_fired = match self.evaluate_transitions(processed_input, &content).await {
11382                        Ok(v) => v,
11383                        Err(e) => {
11384                            yield StreamChunk::error(e.to_string());
11385                            return;
11386                        }
11387                    };
11388                    if transition_fired {
11389                        let _ = self.memory.add_message(ChatMessage::assistant(
11390                            "(Transitioned to new state — tool call handled by workflow)",
11391                        )).await;
11392
11393                        if include_state_events
11394                            && let Some(state) = self.current_state()
11395                        {
11396                                yield StreamChunk::state_transition(None, state);
11397                            }
11398                        continue;
11399                    }
11400
11401                    // Store the assistant's tool-call message (same as blocking path)
11402                    let _ = self.memory.add_message(ChatMessage::assistant(&content)).await;
11403
11404                    // Execute tools with streaming events
11405                    let results = self.execute_tools_parallel(&tool_calls).await;
11406
11407                    for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
11408                        if include_tool_events {
11409                            yield StreamChunk::tool_start(&tool_call.id, &tool_call.name);
11410                        }
11411
11412                        match result {
11413                            Ok(output) => {
11414                                if include_tool_events {
11415                                    yield StreamChunk::tool_result(
11416                                        &tool_call.id,
11417                                        &tool_call.name,
11418                                        &output,
11419                                        true,
11420                                    );
11421                                }
11422                                let _ = self.memory
11423                                    .add_message(Self::tool_result_message(
11424                                        tool_call,
11425                                        &output,
11426                                        native_tool_call,
11427                                    ))
11428                                    .await;
11429                            }
11430                            Err(e) => {
11431                                if matches!(e, AgentError::HITLRejected(_)) {
11432                                    let _ = self.memory.add_message(ChatMessage::assistant(
11433                                        format!("The operation was rejected by the approver: {}", e),
11434                                    )).await;
11435                                    let response = AgentResponse {
11436                                        content: format!("Operation cancelled: {}", e),
11437                                        metadata: None,
11438                                        tool_calls: Some(all_tool_calls.clone()),
11439                                    };
11440                                    if let Err(finalize_error) = self.finish_turn_if_root(&response).await {
11441                                        yield StreamChunk::error(finalize_error.to_string());
11442                                        return;
11443                                    }
11444                                    yield StreamChunk::error(response.content);
11445                                    yield StreamChunk::Done {};
11446                                    return;
11447                                }
11448                                if include_tool_events {
11449                                    yield StreamChunk::tool_result(
11450                                        &tool_call.id,
11451                                        &tool_call.name,
11452                                        e.to_string(),
11453                                        false,
11454                                    );
11455                                }
11456                                let _ = self.memory
11457                                    .add_message(Self::tool_result_message(
11458                                        tool_call,
11459                                        &format!("Error: {}", e),
11460                                        native_tool_call,
11461                                    ))
11462                                    .await;
11463                            }
11464                        }
11465                        all_tool_calls.push(tool_call.clone());
11466
11467                        if include_tool_events {
11468                            yield StreamChunk::tool_end(&tool_call.id);
11469                        }
11470                    }
11471                    continue;
11472                }
11473
11474                // Extract thinking, process output
11475                let (extracted_thinking, answer) = self.extract_thinking(&content);
11476                if extracted_thinking.is_some() {
11477                    thinking_content = extracted_thinking;
11478                }
11479
11480                let output_data = match self.process_output(&answer, &input_data.context).await {
11481                    Ok(d) => d,
11482                    Err(e) => {
11483                        yield StreamChunk::error(e.to_string());
11484                        return;
11485                    }
11486                };
11487
11488                let final_content = if output_data.metadata.rejected {
11489                    output_data
11490                        .metadata
11491                        .rejection_reason
11492                        .unwrap_or_else(|| answer.to_string())
11493                } else {
11494                    output_data.content
11495                };
11496
11497                // Reflection (uses blocking LLM calls for retries)
11498                let (final_content, _reflection_metadata) = match self
11499                    .run_reflection(&*llm, processed_input, final_content)
11500                    .await
11501                {
11502                    Ok(r) => r,
11503                    Err(e) => {
11504                        yield StreamChunk::error(e.to_string());
11505                        return;
11506                    }
11507                };
11508
11509                let final_content = self.format_response_with_thinking(
11510                    thinking_content.as_deref(),
11511                    &final_content,
11512                );
11513
11514                // Buffered decisions emit only the accepted final text.
11515                if buffered_decision {
11516                    yield StreamChunk::content(&final_content);
11517                }
11518
11519                // Post-loop: memory, transitions, post-transition re-generation.
11520                // For NeedsRedispatch, run_loop_internal handles the new state's full
11521                // dispatch and its result is yielded as a single non-streamed chunk.
11522                let post_result = match self
11523                    .post_loop_processing(processed_input, final_content)
11524                    .await
11525                {
11526                    Ok(r) => r,
11527                    Err(e) => {
11528                        yield StreamChunk::error(e.to_string());
11529                        return;
11530                    }
11531                };
11532
11533                let (final_content, transitioned) = match post_result {
11534                    PostLoopResult::NoTransition(content) => (content, false),
11535                    PostLoopResult::Transitioned(content) => (content, true),
11536                    PostLoopResult::NeedsRedispatch => {
11537                        const MAX_REDISPATCH_DEPTH: u32 = 3;
11538                        let current_depth = *self.redispatch_depth.read();
11539                        let content = if current_depth >= MAX_REDISPATCH_DEPTH {
11540                            warn!(
11541                                depth = current_depth,
11542                                "Post-transition re-dispatch depth limit reached (stream)"
11543                            );
11544                            let c = String::new();
11545                            let _ = self.memory.add_message(ChatMessage::assistant(&c)).await;
11546                            c
11547                        } else {
11548                            *self.redispatch_depth.write() += 1;
11549                            if let Some(context) = self.active_turn_context.write().as_mut() {
11550                                context.enter_redispatch();
11551                            }
11552                            info!(
11553                                depth = current_depth + 1,
11554                                "Re-dispatching for new state after transition (stream)"
11555                            );
11556                            let result = self.run_loop_internal(processed_input).await;
11557                            *self.redispatch_depth.write() -= 1;
11558                            if let Some(context) = self.active_turn_context.write().as_mut() {
11559                                context.exit_redispatch();
11560                            }
11561                            match result {
11562                                Ok(resp) => resp.content,
11563                                Err(e) => {
11564                                    yield StreamChunk::error(e.to_string());
11565                                    return;
11566                                }
11567                            }
11568                        };
11569                        (content, true)
11570                    }
11571                };
11572
11573                if transitioned {
11574                    if include_state_events
11575                        && let Some(state) = self.current_state()
11576                    {
11577                            yield StreamChunk::state_transition(None, state);
11578                        }
11579                    // Yield the post-transition re-generated or re-dispatched content.
11580                    yield StreamChunk::content(&final_content);
11581                }
11582
11583                // Parity with non-stream: finalize the committed response once.
11584                let final_response = AgentResponse::new(&final_content);
11585                if let Err(e) = self.finish_turn_if_root(&final_response).await {
11586                    yield StreamChunk::error(e.to_string());
11587                    return;
11588                }
11589
11590                yield StreamChunk::Done {};
11591                return;
11592            }
11593        })
11594    }
11595
11596    /// Streams the root turn while keeping clarification and confirmation questions as terminal responses for their turn.
11597    /// Pending manager and skill ownership must survive until explicit confirmation returns a resolved result for redispatch.
11598    fn run_loop_stream<'a>(
11599        &'a self,
11600        input: &'a str,
11601    ) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
11602        Box::pin(async_stream::stream! {
11603            self.begin_root_turn();
11604            let _root_cleanup = RootTurnCleanup::new(self);
11605            self.hooks.on_message_received(input).await;
11606
11607            // One-shot context initialization (mirrors run_loop)
11608            if !self.context_initialized.swap(true, Ordering::SeqCst) {
11609                if let Err(e) = self.context_manager.initialize().await {
11610                    yield StreamChunk::error(e.to_string());
11611                    return;
11612                }
11613                debug!("Context manager initialized (defaults, env, builtins)");
11614            }
11615
11616            if let Err(e) = self.check_turn_timeout().await {
11617                yield StreamChunk::error(e.to_string());
11618                return;
11619            }
11620            if let Err(e) = self.context_manager.refresh_per_turn().await {
11621                yield StreamChunk::error(e.to_string());
11622                return;
11623            }
11624
11625            // Clear stale disambiguation context from previous turns.
11626            self.clear_disambiguation_context();
11627
11628            // Disambiguation check (before input processing)
11629            if let Some(ref disambiguator) = self.disambiguation_manager {
11630                let disambiguation_context = match self.build_disambiguation_context().await {
11631                    Ok(ctx) => ctx,
11632                    Err(e) => {
11633                        yield StreamChunk::error(e.to_string());
11634                        return;
11635                    }
11636                };
11637
11638                let state_override = self
11639                    .state_machine
11640                    .as_ref()
11641                    .and_then(|sm| sm.current_definition())
11642                    .and_then(|def| def.disambiguation.clone());
11643
11644                let state_generation = self
11645                    .state_machine
11646                    .as_ref()
11647                    .map(|state_machine| state_machine.generation());
11648                let disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
11649                let mut result = match self
11650                    .observe_purpose(
11651                        ObservationPurpose::DisambiguationDetection,
11652                        disambiguator.process_input_with_override(
11653                            input,
11654                            &disambiguation_context,
11655                            state_override.as_ref(),
11656                            None,
11657                        ),
11658                    )
11659                    .await
11660                {
11661                    Ok(r) => r,
11662                    Err(e) => {
11663                        yield StreamChunk::error(e.to_string());
11664                        return;
11665                    }
11666                };
11667                let current_state_generation = self
11668                    .state_machine
11669                    .as_ref()
11670                    .map(|state_machine| state_machine.generation());
11671                if current_state_generation != state_generation
11672                    || self.disambiguation_epoch.load(Ordering::SeqCst) != disambiguation_epoch
11673                {
11674                    disambiguator.clear_pending().await;
11675                    *self.pending_skill_id.write() = None;
11676                    result = DisambiguationResult::Abandoned { new_input: None };
11677                    info!(
11678                        confirmation_event = "invalidated",
11679                        invalidation_reason = "state_generation_changed",
11680                        "Streaming disambiguation result invalidated before redispatch"
11681                    );
11682                }
11683                match result {
11684                    DisambiguationResult::Clear => {
11685                        debug!("Input is clear, proceeding normally (stream)");
11686                    }
11687                    DisambiguationResult::NeedsClarification {
11688                        question,
11689                        detection,
11690                    } => {
11691                        let admission = match self
11692                            .admit_disambiguation_redispatch(
11693                                disambiguation_epoch,
11694                                state_generation,
11695                            )
11696                            .await
11697                        {
11698                            Ok(admission) => admission,
11699                            Err(error) => {
11700                                *self.pending_skill_id.write() = None;
11701                                yield StreamChunk::error(error.to_string());
11702                                return;
11703                            }
11704                        };
11705                        let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
11706                        info!(
11707                            ambiguity_type = ?detection.ambiguity_type,
11708                            confidence = detection.confidence,
11709                            "Input requires clarification (stream)"
11710                        );
11711                        // Confirmation uses the same terminal branch so enriched input cannot stream before explicit agreement.
11712                        // Keep pending_skill_id intact for the later confirmed redispatch.
11713                        if let Err(e) = self.commit_root_user_message(input).await {
11714                            yield StreamChunk::error(e.to_string());
11715                            return;
11716                        }
11717                        let _ = self
11718                            .memory
11719                            .add_message(ChatMessage::assistant(&question.question))
11720                            .await;
11721                        let status = if awaiting_confirmation {
11722                            "awaiting_confirmation"
11723                        } else {
11724                            "awaiting_clarification"
11725                        };
11726                        let response = AgentResponse::new(&question.question).with_metadata(
11727                            "disambiguation",
11728                            serde_json::json!({ "status": status }),
11729                        );
11730                        drop(admission);
11731                        if let Err(e) = self.finish_turn_if_root(&response).await {
11732                            yield StreamChunk::error(e.to_string());
11733                            return;
11734                        }
11735                        yield StreamChunk::content(&question.question);
11736                        yield StreamChunk::Done {};
11737                        return;
11738                    }
11739                    DisambiguationResult::Clarified {
11740                        enriched_input,
11741                        resolved,
11742                        ..
11743                    } => {
11744                        let admission = match self
11745                            .admit_disambiguation_redispatch(
11746                                disambiguation_epoch,
11747                                state_generation,
11748                            )
11749                            .await
11750                        {
11751                            Ok(admission) => admission,
11752                            Err(error) => {
11753                                *self.pending_skill_id.write() = None;
11754                                yield StreamChunk::error(error.to_string());
11755                                return;
11756                            }
11757                        };
11758                        info!(
11759                            resolved_count = resolved.len(),
11760                            enriched = %enriched_input,
11761                            "Input clarified (stream)"
11762                        );
11763                        for (key, value) in &resolved {
11764                            let context_key = format!("disambiguation.{}", key);
11765                            let _ = self.context_manager.set(&context_key, value.clone());
11766                        }
11767                        if let Some(intent) = resolved.get("intent") {
11768                            let _ = self.context_manager.set("resolved_intent", intent.clone());
11769                        }
11770                        let _ = self
11771                            .context_manager
11772                            .set("disambiguation.resolved", serde_json::Value::Bool(true));
11773
11774                        // Check if this clarification was triggered by a skill-level override.
11775                        // Re-run skill disambiguation to verify all required_clarity fields
11776                        // are present before executing.
11777                        let skill_id = self.pending_skill_id.read().clone();
11778                        if let Some(skill_id) = skill_id {
11779                            info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input (stream)");
11780                            drop(admission);
11781                            match self
11782                                .recheck_skill_disambiguation(
11783                                    &skill_id,
11784                                    &enriched_input,
11785                                    disambiguation_epoch,
11786                                    state_generation,
11787                                )
11788                                .await
11789                            {
11790                                Ok(resp) => {
11791                                    yield StreamChunk::content(&resp.content);
11792                                    yield StreamChunk::Done {};
11793                                    return;
11794                                }
11795                                Err(e) => {
11796                                    yield StreamChunk::error(e.to_string());
11797                                    return;
11798                                }
11799                            }
11800                        }
11801
11802                        // Forward to internal stream with enriched input.
11803                        drop(admission);
11804                        let mut inner = self.run_loop_internal_stream(&enriched_input);
11805                        while let Some(chunk) = inner.next().await {
11806                            yield chunk;
11807                        }
11808                        return;
11809                    }
11810                    DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
11811                        info!("Proceeding with best guess (stream)");
11812
11813                        // Same skill-id re-check for best-guess path
11814                        let skill_id = self.pending_skill_id.read().clone();
11815                        if let Some(skill_id) = skill_id {
11816                            info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input (stream)");
11817                            match self
11818                                .recheck_skill_disambiguation(
11819                                    &skill_id,
11820                                    &enriched_input,
11821                                    disambiguation_epoch,
11822                                    state_generation,
11823                                )
11824                                .await
11825                            {
11826                                Ok(resp) => {
11827                                    yield StreamChunk::content(&resp.content);
11828                                    yield StreamChunk::Done {};
11829                                    return;
11830                                }
11831                                Err(e) => {
11832                                    yield StreamChunk::error(e.to_string());
11833                                    return;
11834                                }
11835                            }
11836                        }
11837
11838                        let mut inner = self.run_loop_internal_stream(&enriched_input);
11839                        while let Some(chunk) = inner.next().await {
11840                            yield chunk;
11841                        }
11842                        return;
11843                    }
11844                    DisambiguationResult::GiveUp { reason } => {
11845                        *self.pending_skill_id.write() = None;
11846                        warn!(reason = %reason, "Disambiguation gave up (stream)");
11847                        let apology = self
11848                            .generate_localized_apology(
11849                                "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
11850                                &reason,
11851                            )
11852                            .await
11853                            .unwrap_or_else(|_| {
11854                                format!("I'm sorry, I couldn't understand your request: {}", reason)
11855                            });
11856                        let response = AgentResponse::new(&apology);
11857                        if let Err(e) = self.finish_turn_if_root(&response).await {
11858                            yield StreamChunk::error(e.to_string());
11859                            return;
11860                        }
11861                        yield StreamChunk::content(&apology);
11862                        yield StreamChunk::Done {};
11863                        return;
11864                    }
11865                    DisambiguationResult::Escalate { reason } => {
11866                        *self.pending_skill_id.write() = None;
11867                        info!(reason = %reason, "Escalating to human (stream)");
11868                        if let Some(ref hitl) = self.hitl_engine {
11869                            let trigger =
11870                                ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
11871                            let mut context_map = HashMap::new();
11872                            context_map.insert("original_input".to_string(), serde_json::json!(input));
11873                            context_map.insert("reason".to_string(), serde_json::json!(&reason));
11874                            let check_result = HITLCheckResult::required(
11875                                trigger,
11876                                context_map,
11877                                format!("User request needs human assistance: {}", reason),
11878                                Some(hitl.config().default_timeout_seconds),
11879                            );
11880                            match self.request_hitl_approval(check_result).await {
11881                                Ok(ApprovalResult::Approved | ApprovalResult::Modified { .. }) => {
11882                                    let mut inner = self.run_loop_internal_stream(input);
11883                                    while let Some(chunk) = inner.next().await {
11884                                        yield chunk;
11885                                    }
11886                                    return;
11887                                }
11888                                Ok(_) => {}
11889                                Err(e) => {
11890                                    yield StreamChunk::error(e.to_string());
11891                                    return;
11892                                }
11893                            }
11894                        }
11895                        let apology = self
11896                            .generate_localized_apology(
11897                                "Explain briefly that you're transferring the user to a human agent for help.",
11898                                &reason,
11899                            )
11900                            .await
11901                            .unwrap_or_else(|_| {
11902                                format!("I need human assistance to help with your request: {}", reason)
11903                            });
11904                        let response = AgentResponse::new(&apology);
11905                        if let Err(e) = self.finish_turn_if_root(&response).await {
11906                            yield StreamChunk::error(e.to_string());
11907                            return;
11908                        }
11909                        yield StreamChunk::content(&apology);
11910                        yield StreamChunk::Done {};
11911                        return;
11912                    }
11913                    DisambiguationResult::Abandoned { new_input } => {
11914                        *self.pending_skill_id.write() = None;
11915
11916                        info!(
11917                            has_new_input = new_input.is_some(),
11918                            "Clarification abandoned by user (stream)"
11919                        );
11920
11921                        if let Err(e) = self.commit_root_user_message(input).await {
11922                            yield StreamChunk::error(e.to_string());
11923                            return;
11924                        }
11925
11926                        match new_input {
11927                            Some(fresh_input) => {
11928                                // Topic switch: forward to internal stream with fresh input.
11929                                let mut inner = self.run_loop_internal_stream(&fresh_input);
11930                                while let Some(chunk) = inner.next().await {
11931                                    yield chunk;
11932                                }
11933                                return;
11934                            }
11935                            None => {
11936                                // Pure abandonment: generate a brief acknowledgment.
11937                                let ack = self
11938                                    .generate_localized_apology(
11939                                        "The user changed their mind about their previous request. \
11940                                         Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
11941                                         Do NOT apologize excessively. Be concise.",
11942                                        "User abandoned clarification",
11943                                    )
11944                                    .await
11945                                    .unwrap_or_else(|_| {
11946                                        "OK, no problem. What else can I help with?".to_string()
11947                                    });
11948
11949                                let _ = self
11950                                    .memory
11951                                    .add_message(ChatMessage::assistant(&ack))
11952                                    .await;
11953
11954                                let response = AgentResponse::new(&ack);
11955                                if let Err(e) = self.finish_turn_if_root(&response).await {
11956                                    yield StreamChunk::error(e.to_string());
11957                                    return;
11958                                }
11959                                yield StreamChunk::content(&ack);
11960                                yield StreamChunk::Done {};
11961                                return;
11962                            }
11963                        }
11964                    }
11965                }
11966            }
11967
11968            // No disambiguation or Clear result — proceed with internal stream
11969            let mut inner = self.run_loop_internal_stream(input);
11970            while let Some(chunk) = inner.next().await {
11971                yield chunk;
11972            }
11973        })
11974    }
11975
11976    pub fn info(&self) -> AgentInfo {
11977        self.info.clone()
11978    }
11979
11980    pub fn skills(&self) -> &[SkillDefinition] {
11981        &self.skills
11982    }
11983
11984    /// Clears conversation and pending runtime ownership through one reset contract.
11985    async fn reset_runtime_state(&self) -> Result<()> {
11986        let _admission = self.disambiguation_admission.write().await;
11987        if self.state_transition_reserved.load(Ordering::SeqCst) {
11988            return Err(AgentError::Other(
11989                "Cannot reset while a state transition is in progress".to_string(),
11990            ));
11991        }
11992        self.disambiguation_epoch.fetch_add(1, Ordering::SeqCst);
11993        *self.pending_skill_id.write() = None;
11994        if let Some(disambiguator) = self.disambiguation_manager.as_ref() {
11995            disambiguator.clear_pending().await;
11996        }
11997        self.memory.clear().await?;
11998        *self.iteration_count.write() = 0;
11999        self.tool_call_history.write().clear();
12000        if let Some(ref sm) = self.state_machine {
12001            sm.reset();
12002        }
12003        Ok(())
12004    }
12005
12006    /// Resets the runtime using the same cleanup path as the Agent trait.
12007    pub async fn reset(&self) -> Result<()> {
12008        self.reset_runtime_state().await
12009    }
12010
12011    pub fn max_context_tokens(&self) -> u32 {
12012        self.max_context_tokens
12013    }
12014
12015    pub fn llm_registry(&self) -> &Arc<LLMRegistry> {
12016        &self.llm_registry
12017    }
12018
12019    pub fn state_machine(&self) -> Option<&Arc<StateMachine>> {
12020        self.state_machine.as_ref()
12021    }
12022
12023    pub fn context_manager(&self) -> &Arc<ContextManager> {
12024        &self.context_manager
12025    }
12026
12027    pub fn tool_call_history(&self) -> Vec<ToolCallRecord> {
12028        self.tool_call_history.read().clone()
12029    }
12030
12031    pub fn memory_token_budget(&self) -> Option<&MemoryTokenBudget> {
12032        self.memory_token_budget.as_ref()
12033    }
12034
12035    pub fn parallel_tools_config(&self) -> &ParallelToolsConfig {
12036        &self.parallel_tools
12037    }
12038
12039    pub fn streaming_config(&self) -> &StreamingConfig {
12040        &self.streaming
12041    }
12042
12043    pub fn hooks(&self) -> &Arc<dyn AgentHooks> {
12044        &self.hooks
12045    }
12046
12047    pub fn hitl_engine(&self) -> Option<&HITLEngine> {
12048        self.hitl_engine.as_ref()
12049    }
12050
12051    pub fn approval_handler(&self) -> &Arc<dyn ApprovalHandler> {
12052        &self.approval_handler
12053    }
12054
12055    /// Build a context map with language hints from context_manager for HITL message localization.
12056    fn build_hitl_language_context(&self) -> HashMap<String, Value> {
12057        let mut ctx = HashMap::new();
12058        for key in &["user.language", "input.detected.language", "language"] {
12059            if let Some(val) = self.context_manager.get(key) {
12060                ctx.insert(key.to_string(), val);
12061            }
12062        }
12063        ctx
12064    }
12065
12066    /// Send a HITL check result through the approval flow and return the full ApprovalResult.
12067    async fn request_hitl_approval(&self, check_result: HITLCheckResult) -> Result<ApprovalResult> {
12068        let Some(request) = check_result.into_request() else {
12069            return Ok(ApprovalResult::Approved);
12070        };
12071
12072        self.hooks.on_approval_requested(&request).await;
12073
12074        let timeout = request.timeout;
12075
12076        let raw_result = if let Some(duration) = timeout {
12077            match tokio::time::timeout(
12078                duration,
12079                self.approval_handler.request_approval(request.clone()),
12080            )
12081            .await
12082            {
12083                Ok(result) => result,
12084                Err(_) => ApprovalResult::timeout(),
12085            }
12086        } else {
12087            self.approval_handler
12088                .request_approval(request.clone())
12089                .await
12090        };
12091
12092        self.hooks
12093            .on_approval_result(&request.id, &raw_result)
12094            .await;
12095
12096        let (outcome, effective_result): (ApprovalResolvedOutcome, Result<ApprovalResult>) =
12097            match &raw_result {
12098                ApprovalResult::Approved => (
12099                    ApprovalResolvedOutcome::Approved,
12100                    Ok(ApprovalResult::Approved),
12101                ),
12102                ApprovalResult::Rejected { reason } => (
12103                    ApprovalResolvedOutcome::Rejected {
12104                        reason: reason.clone(),
12105                    },
12106                    Ok(ApprovalResult::Rejected {
12107                        reason: reason.clone(),
12108                    }),
12109                ),
12110                ApprovalResult::Modified { changes } => (
12111                    ApprovalResolvedOutcome::Modified {
12112                        changes: changes.clone(),
12113                    },
12114                    Ok(ApprovalResult::Modified {
12115                        changes: changes.clone(),
12116                    }),
12117                ),
12118                ApprovalResult::Timeout => {
12119                    if let Some(ref engine) = self.hitl_engine {
12120                        match engine.config().on_timeout {
12121                            TimeoutAction::Approve => (
12122                                ApprovalResolvedOutcome::Approved,
12123                                Ok(ApprovalResult::Approved),
12124                            ),
12125                            TimeoutAction::Reject => {
12126                                let reason = Some("Timeout".to_string());
12127                                (
12128                                    ApprovalResolvedOutcome::Rejected {
12129                                        reason: reason.clone(),
12130                                    },
12131                                    Ok(ApprovalResult::Rejected { reason }),
12132                                )
12133                            }
12134                            TimeoutAction::Error => {
12135                                let message = "HITL approval timeout".to_string();
12136                                (
12137                                    ApprovalResolvedOutcome::Error {
12138                                        message: message.clone(),
12139                                    },
12140                                    Err(AgentError::Other(message)),
12141                                )
12142                            }
12143                        }
12144                    } else {
12145                        let reason = Some("Timeout (no engine)".to_string());
12146                        (
12147                            ApprovalResolvedOutcome::Rejected {
12148                                reason: reason.clone(),
12149                            },
12150                            Ok(ApprovalResult::Rejected { reason }),
12151                        )
12152                    }
12153                }
12154            };
12155
12156        self.hooks
12157            .on_approval_resolved(&request, &raw_result, &outcome)
12158            .await;
12159
12160        effective_result
12161    }
12162
12163    pub async fn check_state_hitl(&self, from: Option<&str>, to: &str) -> Result<bool> {
12164        if let Some(ref hitl_engine) = self.hitl_engine {
12165            let hitl_lang_ctx = self.build_hitl_language_context();
12166            let check_result = self
12167                .observe_purpose(
12168                    ObservationPurpose::HitlLocalization,
12169                    hitl_engine.check_state_transition_with_localization(
12170                        from,
12171                        to,
12172                        &hitl_lang_ctx,
12173                        self.approval_handler.as_ref(),
12174                        Some(&self.llm_registry),
12175                    ),
12176                )
12177                .await?;
12178            if check_result.is_required() {
12179                let result = self.request_hitl_approval(check_result).await?;
12180                return Ok(matches!(
12181                    result,
12182                    ApprovalResult::Approved | ApprovalResult::Modified { .. }
12183                ));
12184            }
12185        }
12186        Ok(true)
12187    }
12188
12189    /// Execute multiple tools in parallel
12190    async fn execute_tools_parallel(
12191        &self,
12192        tool_calls: &[ToolCall],
12193    ) -> Vec<(String, Result<String>)> {
12194        let can_run_parallel = tool_calls.iter().all(|tc| {
12195            self.tools
12196                .resolve(&tc.name)
12197                .map(|resolved| resolved.tool.classify_call(&tc.arguments).concurrency_safe)
12198                .unwrap_or(false)
12199        });
12200
12201        if !self.parallel_tools.enabled || tool_calls.len() <= 1 || !can_run_parallel {
12202            let mut results = Vec::new();
12203            for tc in tool_calls {
12204                let result = self
12205                    .observe_purpose(
12206                        current_observation_context()
12207                            .map(|context| context.purpose)
12208                            .unwrap_or_default(),
12209                        self.execute_tool_smart(tc),
12210                    )
12211                    .await;
12212                results.push((tc.id.clone(), result));
12213            }
12214            return results;
12215        }
12216
12217        let chunks: Vec<_> = tool_calls
12218            .chunks(self.parallel_tools.max_parallel)
12219            .collect();
12220
12221        let mut all_results = Vec::new();
12222
12223        for chunk in chunks {
12224            let futures: Vec<_> = chunk
12225                .iter()
12226                .map(|tc| {
12227                    let tc = tc.clone();
12228                    async move {
12229                        let result = self.execute_tool_smart(&tc).await;
12230                        (tc.id.clone(), result)
12231                    }
12232                })
12233                .collect();
12234
12235            let results = futures::future::join_all(futures).await;
12236            all_results.extend(results);
12237        }
12238
12239        all_results
12240    }
12241
12242    /// Stream a chat response with real-time updates
12243    pub async fn chat_stream<'a>(
12244        &'a self,
12245        input: &'a str,
12246    ) -> Result<Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>>> {
12247        //
12248        // Streaming readiness is checked before returning a stream so capability failures cannot appear as partial output.
12249        //
12250        self.init_storage().await?;
12251        info!(input_len = input.len(), "Starting streaming chat");
12252        let inner = self.run_loop_stream(input);
12253        if let Some(context) = self.build_observation_context(None) {
12254            let stream: Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> =
12255                Box::pin(async_stream::stream! {
12256                    let mut inner = inner;
12257                    loop {
12258                        let next = with_observation_context(context.clone(), inner.next()).await;
12259                        match next {
12260                            Some(chunk) => yield chunk,
12261                            None => break,
12262                        }
12263                    }
12264                    self.export_observability_if_configured().await;
12265                });
12266            Ok(stream)
12267        } else {
12268            Ok(inner)
12269        }
12270    }
12271}
12272
12273#[async_trait]
12274impl ToolInvoker for RuntimeAgent {
12275    async fn invoke_tool(&self, request: ToolExecutionRequest) -> Result<ToolExecutionRecord> {
12276        self.execute_tool_record(request).await
12277    }
12278}
12279
12280#[async_trait]
12281impl Agent for RuntimeAgent {
12282    async fn chat(&self, input: &str) -> Result<AgentResponse> {
12283        let result = if let Some(context) = self.build_observation_context(None) {
12284            with_observation_context(context, self.run_loop(input)).await
12285        } else {
12286            self.run_loop(input).await
12287        };
12288        self.export_observability_if_configured().await;
12289        result
12290    }
12291
12292    fn info(&self) -> AgentInfo {
12293        self.info.clone()
12294    }
12295
12296    /// Resets the runtime without leaving clarification or skill ownership behind.
12297    async fn reset(&self) -> Result<()> {
12298        self.reset_runtime_state().await
12299    }
12300}
12301
12302//
12303// Render a concurrent input template using direct minijinja.
12304// Same approach as pipeline's render_stage_template so variables are top-level.
12305//
12306// Available variables:
12307//   {{ user_input }}    - the user's actual message
12308//   {{ context.<key> }} - values from the context manager
12309//
12310/// Builds safe runtime tags for background maintenance lifecycle events.
12311fn background_maintenance_tags(
12312    label: &str,
12313    stage: &str,
12314    reason: Option<&str>,
12315    policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
12316) -> HashMap<String, String> {
12317    let mut tags = HashMap::new();
12318    tags.insert("runtime.background".to_string(), "true".to_string());
12319    tags.insert("runtime.maintenance".to_string(), label.to_string());
12320    tags.insert("runtime.maintenance_stage".to_string(), stage.to_string());
12321    if let Some(policy) = policy {
12322        tags.insert(
12323            "runtime.await_before_next_turn".to_string(),
12324            await_before_next_turn_label(policy.await_before_next_turn).to_string(),
12325        );
12326        tags.insert(
12327            "runtime.maintenance_mode".to_string(),
12328            maintenance_mode_label(policy.mode).to_string(),
12329        );
12330    }
12331    if let Some(reason) = reason {
12332        tags.insert("runtime.reason".to_string(), reason.to_string());
12333    }
12334    tags
12335}
12336
12337fn await_before_next_turn_label(policy: AwaitBeforeNextTurn) -> &'static str {
12338    match policy {
12339        AwaitBeforeNextTurn::Never => "never",
12340        AwaitBeforeNextTurn::SameActor => "same_actor",
12341        AwaitBeforeNextTurn::Always => "always",
12342    }
12343}
12344
12345fn maintenance_mode_label(mode: MaintenanceMode) -> &'static str {
12346    match mode {
12347        MaintenanceMode::InlineSerial => "inline_serial",
12348        MaintenanceMode::InlineParallel => "inline_parallel",
12349        MaintenanceMode::Background => "background",
12350    }
12351}
12352
12353/// Records a background maintenance lifecycle event when observability is enabled.
12354fn record_background_maintenance_event(
12355    manager: Option<&Arc<ObservabilityManager>>,
12356    label: &str,
12357    status: EventStatus,
12358    duration_ms: u64,
12359    stage: &str,
12360    reason: Option<String>,
12361    policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
12362) {
12363    if let Some(manager) = manager {
12364        manager.record_lifecycle_event(
12365            EventType::MemoryOperation {
12366                operation: format!("{}_background_{}", label, stage),
12367            },
12368            ObservationPurpose::Other(format!("{}_maintenance", label)),
12369            status,
12370            duration_ms,
12371            background_maintenance_tags(label, stage, reason.as_deref(), policy),
12372            None,
12373        );
12374    }
12375}
12376
12377fn effective_maintenance_mode(mode: MaintenanceMode, force_parallel: bool) -> MaintenanceMode {
12378    if force_parallel && matches!(mode, MaintenanceMode::InlineSerial) {
12379        MaintenanceMode::InlineParallel
12380    } else {
12381        mode
12382    }
12383}
12384
12385fn observation_purpose_for_process(hint: ProcessPurposeHint) -> ObservationPurpose {
12386    match hint {
12387        ProcessPurposeHint::Detect => ObservationPurpose::ProcessDetect,
12388        ProcessPurposeHint::Extract => ObservationPurpose::ProcessExtract,
12389        ProcessPurposeHint::Validate => ObservationPurpose::ProcessValidate,
12390        ProcessPurposeHint::Transform | ProcessPurposeHint::Other => {
12391            ObservationPurpose::ProcessTransform
12392        }
12393    }
12394}
12395
12396fn new_tool_resource_locks() -> ToolResourceLocks {
12397    Arc::new(RwLock::new(HashMap::new()))
12398}
12399
12400//
12401// Path-bound mutations share one conservative lock so aliases and parent-child paths cannot bypass serialization.
12402// Exact domain keys remain available for non-path side effects, and unbound effects share one fallback lock.
12403//
12404fn tool_resource_lock_keys(
12405    _canonical_id: &str,
12406    args: &Value,
12407    bindings: &ai_agents_core::ToolPolicyBindings,
12408    classification: &ai_agents_core::ToolCallClassification,
12409) -> Vec<String> {
12410    if classification.concurrency_safe {
12411        return Vec::new();
12412    }
12413
12414    let mut keys = Vec::new();
12415    let mut has_path_resource = false;
12416    for binding in &bindings.path_fields {
12417        let value = value_at_argument_path(args, &binding.field)
12418            .cloned()
12419            .or_else(|| {
12420                binding
12421                    .default_path
12422                    .as_ref()
12423                    .map(|path| Value::String(path.clone()))
12424            });
12425        if let Some(value) = value {
12426            collect_resource_strings(&value, |_| {
12427                has_path_resource = true;
12428            });
12429        }
12430    }
12431    for binding in &bindings.domain_fields {
12432        if let Some(value) = value_at_argument_path(args, &binding.field) {
12433            collect_resource_strings(value, |domain| {
12434                let normalized = if binding.is_url {
12435                    normalized_url_resource_key(domain)
12436                } else {
12437                    domain.trim().trim_end_matches('.').to_ascii_lowercase()
12438                };
12439                keys.push(format!("domain:{}", normalized));
12440            });
12441        }
12442    }
12443    for binding in &bindings.command_fields {
12444        if !matches!(binding.kind, ai_agents_core::CommandBindingKind::Cwd) {
12445            continue;
12446        }
12447        if let Some(value) = value_at_argument_path(args, &binding.field) {
12448            collect_resource_strings(value, |_| {
12449                has_path_resource = true;
12450            });
12451        }
12452    }
12453    if has_path_resource {
12454        keys.push("path-mutation:global".to_string());
12455    }
12456    if keys.is_empty() {
12457        keys.push("side-effect:unbound".to_string());
12458    }
12459    keys.sort();
12460    keys.dedup();
12461    keys
12462}
12463
12464fn value_at_argument_path<'a>(value: &'a Value, field: &str) -> Option<&'a Value> {
12465    let mut current = value;
12466    for segment in field.split('.') {
12467        if segment.is_empty() {
12468            return None;
12469        }
12470        current = current.get(segment)?;
12471    }
12472    Some(current)
12473}
12474
12475fn collect_resource_strings(value: &Value, mut collect: impl FnMut(&str)) {
12476    match value {
12477        Value::String(value) => collect(value),
12478        Value::Array(values) => {
12479            for value in values {
12480                if let Some(value) = value.as_str() {
12481                    collect(value);
12482                }
12483            }
12484        }
12485        _ => {}
12486    }
12487}
12488
12489fn normalized_url_resource_key(value: &str) -> String {
12490    let value = value.trim();
12491    let Some((scheme, remainder)) = value.split_once("://") else {
12492        return value.to_ascii_lowercase();
12493    };
12494    let authority_end = remainder.find(['/', '?', '#']).unwrap_or(remainder.len());
12495    let (authority, suffix) = remainder.split_at(authority_end);
12496    format!(
12497        "{}://{}{}",
12498        scheme.to_ascii_lowercase(),
12499        authority.to_ascii_lowercase(),
12500        suffix
12501    )
12502}
12503
12504fn render_concurrent_template(
12505    template: &str,
12506    user_input: &str,
12507    context_values: &std::collections::HashMap<String, serde_json::Value>,
12508) -> Result<String> {
12509    let mut env = minijinja::Environment::new();
12510    env.add_template("concurrent", template)
12511        .map_err(|e| AgentError::Other(format!("Concurrent template parse error: {}", e)))?;
12512
12513    let mut ctx = std::collections::BTreeMap::new();
12514    ctx.insert("user_input".to_string(), minijinja::Value::from(user_input));
12515
12516    // Expose context manager values under {{ context.<key> }}.
12517    let context_obj = minijinja::Value::from_serialize(context_values);
12518    ctx.insert("context".to_string(), context_obj);
12519
12520    let tmpl = env
12521        .get_template("concurrent")
12522        .map_err(|e| AgentError::Other(format!("Concurrent template error: {}", e)))?;
12523
12524    tmpl.render(minijinja::Value::from_serialize(&ctx))
12525        .map_err(|e| AgentError::Other(format!("Concurrent template render error: {}", e)))
12526}
12527
12528#[cfg(test)]
12529mod tests {
12530    use super::*;
12531    use crate::AgentBuilder;
12532    use ai_agents_core::{LLMChunk, LLMConfig, LLMError, LLMFeature, Tool};
12533    use ai_agents_llm::mock::MockLLMProvider;
12534    use ai_agents_skills::{SkillDefinition, SkillStep};
12535    use ai_agents_tools::{
12536        CalculatorTool, CopyPathTool, DeletePathTool, FileWriteTool, MovePathTool,
12537        WebFetchResolver, WebFetchTool, WebFetchTransport, WebFetchTransportRequest,
12538        WebFetchTransportResponse,
12539    };
12540
12541    fn mock_with_response(response: &str) -> MockLLMProvider {
12542        let mut mock = MockLLMProvider::new("test");
12543        mock.set_response(response);
12544        mock
12545    }
12546
12547    fn mock_with_responses(responses: Vec<&str>) -> MockLLMProvider {
12548        let mut mock = MockLLMProvider::new("test");
12549        mock.set_responses(responses.into_iter().map(String::from).collect(), true);
12550        mock
12551    }
12552
12553    /// Builds a two-state fixture so confirmation ownership can be invalidated by transition.
12554    fn disambiguation_state_machine(
12555        state_enabled: Option<bool>,
12556        require_confirmation: bool,
12557    ) -> Arc<StateMachine> {
12558        let definition = ai_agents_state::StateDefinition {
12559            prompt: Some("Handle the resolved request.".to_string()),
12560            disambiguation: Some(ai_agents_disambiguation::StateDisambiguationOverride {
12561                enabled: state_enabled,
12562                require_confirmation,
12563                ..Default::default()
12564            }),
12565            ..Default::default()
12566        };
12567        let review = ai_agents_state::StateDefinition {
12568            prompt: Some("Review a fresh request.".to_string()),
12569            ..Default::default()
12570        };
12571        Arc::new(
12572            StateMachine::new(ai_agents_state::StateConfig {
12573                initial: "active".to_string(),
12574                states: std::collections::HashMap::from([
12575                    ("active".to_string(), definition),
12576                    ("review".to_string(), review),
12577                ]),
12578                global_transitions: Vec::new(),
12579                fallback: None,
12580                max_no_transition: None,
12581                regenerate_on_transition: true,
12582            })
12583            .unwrap(),
12584        )
12585    }
12586
12587    /// Builds a state-aware disambiguation fixture without skills.
12588    fn state_disambiguation_agent(
12589        responses: Vec<&str>,
12590        manager_enabled: bool,
12591        state_enabled: Option<bool>,
12592        require_confirmation: bool,
12593    ) -> (RuntimeAgent, MockLLMProvider) {
12594        state_disambiguation_agent_with_skills(
12595            responses,
12596            manager_enabled,
12597            state_enabled,
12598            require_confirmation,
12599            Vec::new(),
12600        )
12601    }
12602
12603    /// Builds a state-aware disambiguation fixture with optional real skill routing.
12604    fn state_disambiguation_agent_with_skills(
12605        responses: Vec<&str>,
12606        manager_enabled: bool,
12607        state_enabled: Option<bool>,
12608        require_confirmation: bool,
12609        skills: Vec<SkillDefinition>,
12610    ) -> (RuntimeAgent, MockLLMProvider) {
12611        let mut mock = MockLLMProvider::new("state-confirmation");
12612        mock.set_responses(responses.into_iter().map(String::from).collect(), false);
12613        let observed = mock.clone();
12614        let agent = AgentBuilder::new()
12615            .system_prompt("Handle requests.")
12616            .llm(Arc::new(mock.clone()))
12617            .llm_alias("router", Arc::new(mock))
12618            .state_machine(disambiguation_state_machine(
12619                state_enabled,
12620                require_confirmation,
12621            ))
12622            .skills(skills)
12623            .build()
12624            .unwrap()
12625            .with_disambiguation(DisambiguationConfig {
12626                enabled: manager_enabled,
12627                ..Default::default()
12628            });
12629        (agent, observed)
12630    }
12631
12632    /// Defines a prompt skill whose provider call proves committed execution.
12633    fn confirmation_skill() -> SkillDefinition {
12634        SkillDefinition {
12635            id: "send_report".to_string(),
12636            description: "Send a report after clarification".to_string(),
12637            trigger: "When the user asks to send a report".to_string(),
12638            steps: vec![SkillStep::Prompt {
12639                prompt: "Execute confirmed report skill for: {{ input }}".to_string(),
12640                llm: None,
12641            }],
12642            reasoning: None,
12643            reflection: None,
12644            disambiguation: Some(ai_agents_disambiguation::SkillDisambiguationOverride {
12645                enabled: Some(true),
12646                ..Default::default()
12647            }),
12648        }
12649    }
12650
12651    /// Counts committed skill prompt calls without relying on final response wording.
12652    fn confirmation_skill_call_count(observed: &MockLLMProvider) -> usize {
12653        observed
12654            .call_history()
12655            .iter()
12656            .filter(|call| {
12657                call.messages
12658                    .iter()
12659                    .any(|message| message.content.contains("Execute confirmed report skill"))
12660            })
12661            .count()
12662    }
12663
12664    struct BlockingRuntimeConfirmationObserver {
12665        entered: tokio::sync::Barrier,
12666        release: tokio::sync::Notify,
12667    }
12668
12669    impl BlockingRuntimeConfirmationObserver {
12670        fn new() -> Self {
12671            Self {
12672                entered: tokio::sync::Barrier::new(2),
12673                release: tokio::sync::Notify::new(),
12674            }
12675        }
12676    }
12677
12678    struct ResetOnTransitionHooks {
12679        agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
12680        invoked: AtomicBool,
12681    }
12682
12683    #[async_trait]
12684    impl AgentHooks for ResetOnTransitionHooks {
12685        async fn on_state_transition(&self, _from: Option<&str>, _to: &str, _reason: &str) {
12686            if self.invoked.swap(true, Ordering::SeqCst) {
12687                return;
12688            }
12689            let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
12690            if let Some(agent) = agent {
12691                agent.reset().await.unwrap();
12692            }
12693        }
12694    }
12695
12696    impl ClarificationObserver for BlockingRuntimeConfirmationObserver {
12697        fn observe_question<'a>(
12698            &'a self,
12699            future: ClarificationQuestionFuture<'a>,
12700        ) -> ClarificationQuestionFuture<'a> {
12701            future
12702        }
12703
12704        fn observe_parse<'a>(
12705            &'a self,
12706            future: ClarificationParseFuture<'a>,
12707        ) -> ClarificationParseFuture<'a> {
12708            future
12709        }
12710
12711        fn observe_confirmation_parse<'a>(
12712            &'a self,
12713            future: ConfirmationParseFuture<'a>,
12714        ) -> ConfirmationParseFuture<'a> {
12715            Box::pin(async move {
12716                self.entered.wait().await;
12717                self.release.notified().await;
12718                future.await
12719            })
12720        }
12721    }
12722
12723    #[tokio::test]
12724    async fn state_confirmation_blocks_redispatch_until_explicit_agreement() {
12725        let (agent, observed) = state_disambiguation_agent(
12726            vec![
12727                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12728                r#"{"question":"What should I send?","options":null}"#,
12729                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12730                r#"{"question":"Should I send the report to Ada?"}"#,
12731                r#"{"status":"confirmed"}"#,
12732                "Request executed.",
12733            ],
12734            true,
12735            None,
12736            true,
12737        );
12738
12739        let clarification = agent.chat("Send it").await.unwrap();
12740        assert_eq!(clarification.content, "What should I send?");
12741        assert_eq!(observed.call_count(), 2);
12742
12743        let confirmation = agent.chat("The report to Ada").await.unwrap();
12744        assert_eq!(confirmation.content, "Should I send the report to Ada?");
12745        assert_eq!(
12746            confirmation
12747                .metadata
12748                .as_ref()
12749                .and_then(|metadata| metadata.get("disambiguation"))
12750                .and_then(|metadata| metadata.get("status"))
12751                .and_then(Value::as_str),
12752            Some("awaiting_confirmation")
12753        );
12754        assert_eq!(observed.call_count(), 4);
12755
12756        let completed = agent.chat("Yes").await.unwrap();
12757        assert_eq!(completed.content, "Request executed.");
12758        assert_eq!(observed.call_count(), 6);
12759    }
12760
12761    #[tokio::test]
12762    async fn streaming_state_confirmation_ends_the_turn_before_redispatch() {
12763        let (agent, observed) = state_disambiguation_agent(
12764            vec![
12765                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12766                r#"{"question":"What should I send?","options":null}"#,
12767                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12768                r#"{"question":"Should I send the report to Ada?"}"#,
12769                r#"{"status":"confirmed"}"#,
12770                "Request executed.",
12771            ],
12772            true,
12773            None,
12774            true,
12775        );
12776
12777        let mut clarification_stream = agent.chat_stream("Send it").await.unwrap();
12778        let mut clarification = String::new();
12779        while let Some(chunk) = clarification_stream.next().await {
12780            match chunk {
12781                StreamChunk::Content { text } => clarification.push_str(&text),
12782                StreamChunk::Done {} => break,
12783                StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
12784                _ => {}
12785            }
12786        }
12787        assert_eq!(clarification, "What should I send?");
12788        assert_eq!(observed.call_count(), 2);
12789
12790        let mut confirmation_stream = agent.chat_stream("The report to Ada").await.unwrap();
12791        let mut confirmation = String::new();
12792        while let Some(chunk) = confirmation_stream.next().await {
12793            match chunk {
12794                StreamChunk::Content { text } => confirmation.push_str(&text),
12795                StreamChunk::Done {} => break,
12796                StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
12797                _ => {}
12798            }
12799        }
12800        assert_eq!(confirmation, "Should I send the report to Ada?");
12801        assert_eq!(observed.call_count(), 4);
12802
12803        let mut completed_stream = agent.chat_stream("Yes").await.unwrap();
12804        let mut completed = String::new();
12805        while let Some(chunk) = completed_stream.next().await {
12806            match chunk {
12807                StreamChunk::Content { text } => completed.push_str(&text),
12808                StreamChunk::Done {} => break,
12809                StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
12810                _ => {}
12811            }
12812        }
12813        assert_eq!(completed, "Request executed.");
12814        assert_eq!(observed.call_count(), 6);
12815    }
12816
12817    /// Confirms that a pending skill remains inert until confirmation and then runs once.
12818    #[tokio::test]
12819    async fn confirmed_skill_route_executes_exactly_once() {
12820        let (agent, observed) = state_disambiguation_agent_with_skills(
12821            vec![
12822                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
12823                "send_report",
12824                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12825                r#"{"question":"What should I send?","options":null}"#,
12826                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12827                r#"{"question":"Should I send the report to Ada?"}"#,
12828                r#"{"status":"confirmed"}"#,
12829                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"resolved","what_is_unclear":[],"detected_language":"en"}"#,
12830                "Report skill executed.",
12831            ],
12832            true,
12833            None,
12834            true,
12835            vec![confirmation_skill()],
12836        );
12837
12838        let clarification = agent.chat("Send it").await.unwrap();
12839        assert_eq!(clarification.content, "What should I send?");
12840        assert_eq!(confirmation_skill_call_count(&observed), 0);
12841
12842        let confirmation = agent.chat("The report to Ada").await.unwrap();
12843        assert_eq!(confirmation.content, "Should I send the report to Ada?");
12844        assert_eq!(
12845            confirmation
12846                .metadata
12847                .as_ref()
12848                .and_then(|metadata| metadata.get("disambiguation"))
12849                .and_then(|metadata| metadata.get("status"))
12850                .and_then(Value::as_str),
12851            Some("awaiting_confirmation")
12852        );
12853        assert_eq!(confirmation_skill_call_count(&observed), 0);
12854
12855        let completed = agent.chat("Yes").await.unwrap();
12856        assert_eq!(completed.content, "Report skill executed.");
12857        assert_eq!(confirmation_skill_call_count(&observed), 1);
12858        assert!(agent.pending_skill_id.read().is_none());
12859        let messages = agent.memory.get_messages(None).await.unwrap();
12860        assert!(!messages.iter().any(|message| message.content == "Yes"));
12861    }
12862
12863    /// Preserves a second clarification response after confirmation instead of overwriting its metadata.
12864    #[tokio::test]
12865    async fn confirmed_skill_recheck_preserves_new_clarification_metadata() {
12866        let (agent, observed) = state_disambiguation_agent_with_skills(
12867            vec![
12868                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
12869                "send_report",
12870                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12871                r#"{"question":"What should I send?","options":null}"#,
12872                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12873                r#"{"question":"Should I send the report to Ada?"}"#,
12874                r#"{"status":"confirmed"}"#,
12875                r#"{"is_ambiguous":true,"confidence":0.3,"ambiguity_type":"missing_parameters","reasoning":"timing missing","what_is_unclear":["timing"],"detected_language":"en"}"#,
12876                r#"{"question":"When should I send it?","options":null}"#,
12877            ],
12878            true,
12879            None,
12880            true,
12881            vec![confirmation_skill()],
12882        );
12883
12884        agent.chat("Send it").await.unwrap();
12885        agent.chat("The report to Ada").await.unwrap();
12886        let follow_up = agent.chat("Yes").await.unwrap();
12887
12888        assert_eq!(follow_up.content, "When should I send it?");
12889        let metadata = follow_up
12890            .metadata
12891            .as_ref()
12892            .and_then(|metadata| metadata.get("disambiguation"))
12893            .unwrap();
12894        assert_eq!(
12895            metadata.get("status").and_then(Value::as_str),
12896            Some("awaiting_clarification")
12897        );
12898        assert_eq!(
12899            metadata.get("skill_id").and_then(Value::as_str),
12900            Some("send_report")
12901        );
12902        assert!(metadata.get("detection").is_some());
12903        assert_eq!(confirmation_skill_call_count(&observed), 0);
12904    }
12905
12906    /// Confirms rejection clears a pending skill without executing its prompt.
12907    #[tokio::test]
12908    async fn rejected_skill_confirmation_never_executes() {
12909        let (agent, observed) = state_disambiguation_agent_with_skills(
12910            vec![
12911                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
12912                "send_report",
12913                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12914                r#"{"question":"What should I send?","options":null}"#,
12915                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12916                r#"{"question":"Should I send the report to Ada?"}"#,
12917                r#"{"status":"rejected"}"#,
12918                "Confirmation rejected.",
12919            ],
12920            true,
12921            None,
12922            true,
12923            vec![confirmation_skill()],
12924        );
12925
12926        agent.chat("Send it").await.unwrap();
12927        agent.chat("The report to Ada").await.unwrap();
12928        let rejected = agent.chat("No").await.unwrap();
12929
12930        assert_eq!(rejected.content, "Confirmation rejected.");
12931        assert_eq!(confirmation_skill_call_count(&observed), 0);
12932        assert!(agent.pending_skill_id.read().is_none());
12933    }
12934
12935    /// Confirms reset clears manager and skill ownership before a later streaming turn.
12936    #[tokio::test]
12937    async fn reset_invalidates_pending_skill_confirmation_before_streaming_input() {
12938        let (agent, observed) = state_disambiguation_agent_with_skills(
12939            vec![
12940                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
12941                "send_report",
12942                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12943                r#"{"question":"What should I send?","options":null}"#,
12944                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12945                r#"{"question":"Should I send the report to Ada?"}"#,
12946                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"fresh input","what_is_unclear":[],"detected_language":"en"}"#,
12947                "none",
12948                "Fresh response.",
12949            ],
12950            true,
12951            None,
12952            true,
12953            vec![confirmation_skill()],
12954        );
12955
12956        agent.chat("Send it").await.unwrap();
12957        agent.chat("The report to Ada").await.unwrap();
12958        agent.reset().await.unwrap();
12959        assert!(agent.pending_skill_id.read().is_none());
12960        assert!(
12961            !agent
12962                .disambiguation_manager()
12963                .unwrap()
12964                .has_pending_clarification()
12965                .await
12966        );
12967
12968        let mut stream = agent.chat_stream("Yes").await.unwrap();
12969        let mut content = String::new();
12970        while let Some(chunk) = stream.next().await {
12971            match chunk {
12972                StreamChunk::Content { text } => content.push_str(&text),
12973                StreamChunk::Done {} => break,
12974                StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
12975                _ => {}
12976            }
12977        }
12978
12979        assert_eq!(content, "Fresh response.");
12980        assert_eq!(confirmation_skill_call_count(&observed), 0);
12981    }
12982
12983    /// Confirms the Agent trait reset uses the same pending ownership cleanup.
12984    #[tokio::test]
12985    async fn trait_reset_clears_pending_skill_confirmation() {
12986        let (agent, _) = state_disambiguation_agent_with_skills(
12987            vec![
12988                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
12989                "send_report",
12990                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12991                r#"{"question":"What should I send?","options":null}"#,
12992                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12993                r#"{"question":"Should I send the report to Ada?"}"#,
12994            ],
12995            true,
12996            None,
12997            true,
12998            vec![confirmation_skill()],
12999        );
13000
13001        agent.chat("Send it").await.unwrap();
13002        agent.chat("The report to Ada").await.unwrap();
13003        <RuntimeAgent as Agent>::reset(&agent).await.unwrap();
13004
13005        assert!(agent.pending_skill_id.read().is_none());
13006        assert!(
13007            !agent
13008                .disambiguation_manager()
13009                .unwrap()
13010                .has_pending_clarification()
13011                .await
13012        );
13013    }
13014
13015    /// Confirms a state transition cancels stale skill execution before confirmation parsing.
13016    #[tokio::test]
13017    async fn state_change_invalidates_pending_skill_confirmation() {
13018        let (agent, observed) = state_disambiguation_agent_with_skills(
13019            vec![
13020                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13021                "send_report",
13022                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13023                r#"{"question":"What should I send?","options":null}"#,
13024                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13025                r#"{"question":"Should I send the report to Ada?"}"#,
13026                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"fresh input","what_is_unclear":[],"detected_language":"en"}"#,
13027                "none",
13028                "Fresh response.",
13029            ],
13030            true,
13031            None,
13032            true,
13033            vec![confirmation_skill()],
13034        );
13035
13036        agent.chat("Send it").await.unwrap();
13037        agent.chat("The report to Ada").await.unwrap();
13038        agent.transition_to("review").await.unwrap();
13039        let cancelled = agent.chat("Yes").await.unwrap();
13040
13041        assert_eq!(cancelled.content, "Fresh response.");
13042        assert_eq!(confirmation_skill_call_count(&observed), 0);
13043        assert!(agent.pending_skill_id.read().is_none());
13044    }
13045
13046    /// Confirms reset wins over a confirmation result that was already being parsed.
13047    #[tokio::test]
13048    async fn in_flight_confirmation_cannot_redispatch_after_reset() {
13049        let (mut agent, observed) = state_disambiguation_agent_with_skills(
13050            vec![
13051                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13052                "send_report",
13053                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13054                r#"{"question":"What should I send?","options":null}"#,
13055                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13056                r#"{"question":"Should I send the report to Ada?"}"#,
13057                r#"{"status":"confirmed"}"#,
13058                "Confirmation cancelled.",
13059            ],
13060            true,
13061            None,
13062            true,
13063            vec![confirmation_skill()],
13064        );
13065        let observer = Arc::new(BlockingRuntimeConfirmationObserver::new());
13066        let manager = agent
13067            .disambiguation_manager
13068            .take()
13069            .unwrap()
13070            .with_clarification_observer(observer.clone());
13071        agent.disambiguation_manager = Some(manager);
13072        let agent = Arc::new(agent);
13073
13074        agent.chat("Send it").await.unwrap();
13075        agent.chat("The report to Ada").await.unwrap();
13076
13077        let confirming_agent = Arc::clone(&agent);
13078        let confirmation = tokio::spawn(async move { confirming_agent.chat("Yes").await });
13079        observer.entered.wait().await;
13080        agent.reset().await.unwrap();
13081        observer.release.notify_one();
13082
13083        let response = confirmation.await.unwrap().unwrap();
13084        assert_eq!(response.content, "Confirmation cancelled.");
13085        assert_eq!(confirmation_skill_call_count(&observed), 0);
13086        assert!(agent.pending_skill_id.read().is_none());
13087    }
13088
13089    /// Confirms a queued reset prevents an older terminal question from being published afterward.
13090    #[tokio::test]
13091    async fn queued_reset_prevents_stale_confirmation_question_publication() {
13092        let (agent, observed) = state_disambiguation_agent(
13093            vec![
13094                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13095                r#"{"question":"What should I send?","options":null}"#,
13096                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13097                r#"{"question":"Should I send the report to Ada?"}"#,
13098            ],
13099            true,
13100            None,
13101            true,
13102        );
13103        let agent = Arc::new(agent);
13104        agent.chat("Send it").await.unwrap();
13105
13106        let admission = agent.disambiguation_admission.write().await;
13107        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
13108        let resetting_agent = Arc::clone(&agent);
13109        let reset = tokio::spawn(async move {
13110            let _ = started_tx.send(());
13111            resetting_agent.reset().await
13112        });
13113        started_rx.await.unwrap();
13114        tokio::task::yield_now().await;
13115
13116        let responding_agent = Arc::clone(&agent);
13117        let response =
13118            tokio::spawn(async move { responding_agent.chat("The report to Ada").await });
13119        tokio::time::timeout(std::time::Duration::from_secs(2), async {
13120            while observed.call_count() < 4 {
13121                tokio::task::yield_now().await;
13122            }
13123        })
13124        .await
13125        .expect("clarification processing must reach terminal publication");
13126        drop(admission);
13127
13128        reset.await.unwrap().unwrap();
13129        let error = response.await.unwrap().unwrap_err();
13130        assert!(error.to_string().contains("ownership changed"));
13131        assert!(
13132            !agent
13133                .disambiguation_manager()
13134                .unwrap()
13135                .has_pending_clarification()
13136                .await
13137        );
13138        assert!(agent.memory.get_messages(None).await.unwrap().is_empty());
13139    }
13140
13141    /// Confirms skill clarification consumers recheck ownership before publishing returned responses.
13142    #[tokio::test]
13143    async fn queued_reset_prevents_stale_skill_clarification_publication() {
13144        let (agent, observed) = state_disambiguation_agent_with_skills(
13145            vec![
13146                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13147                "send_report",
13148                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13149                r#"{"question":"What should I send?","options":null}"#,
13150            ],
13151            true,
13152            None,
13153            true,
13154            vec![confirmation_skill()],
13155        );
13156        let agent = Arc::new(agent);
13157        let admission = agent.disambiguation_admission.write().await;
13158        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
13159        let resetting_agent = Arc::clone(&agent);
13160        let reset = tokio::spawn(async move {
13161            let _ = started_tx.send(());
13162            resetting_agent.reset().await
13163        });
13164        started_rx.await.unwrap();
13165        tokio::task::yield_now().await;
13166
13167        let responding_agent = Arc::clone(&agent);
13168        let response = tokio::spawn(async move { responding_agent.chat("Send it").await });
13169        tokio::time::timeout(std::time::Duration::from_secs(2), async {
13170            while observed.call_count() < 4 {
13171                tokio::task::yield_now().await;
13172            }
13173        })
13174        .await
13175        .expect("skill clarification must reach terminal publication");
13176        drop(admission);
13177
13178        reset.await.unwrap().unwrap();
13179        let error = response.await.unwrap().unwrap_err();
13180        assert!(error.to_string().contains("ownership changed"));
13181        assert_eq!(confirmation_skill_call_count(&observed), 0);
13182        assert!(agent.pending_skill_id.read().is_none());
13183        assert!(agent.memory.get_messages(None).await.unwrap().is_empty());
13184    }
13185
13186    /// Confirms transition hooks can reenter reset after the state commit lock is released.
13187    #[tokio::test]
13188    async fn transition_hook_can_reset_without_admission_deadlock() {
13189        let hooks = Arc::new(ResetOnTransitionHooks {
13190            agent: parking_lot::Mutex::new(None),
13191            invoked: AtomicBool::new(false),
13192        });
13193        let agent = Arc::new(
13194            AgentBuilder::new()
13195                .system_prompt("Test transition hook reentrancy.")
13196                .llm(Arc::new(mock_with_response("done")))
13197                .state_machine(disambiguation_state_machine(None, false))
13198                .build()
13199                .unwrap()
13200                .with_hooks(hooks.clone()),
13201        );
13202        *hooks.agent.lock() = Some(Arc::downgrade(&agent));
13203
13204        let transitioned = tokio::time::timeout(
13205            std::time::Duration::from_secs(2),
13206            agent.apply_transition_target("active", "review", "test transition", None),
13207        )
13208        .await
13209        .expect("transition hook reset must not deadlock")
13210        .unwrap();
13211
13212        assert!(transitioned);
13213        assert!(hooks.invoked.load(Ordering::SeqCst));
13214        assert_eq!(agent.current_state().as_deref(), Some("active"));
13215    }
13216
13217    /// Confirms only the reserved transition can run source-state exit side effects.
13218    #[tokio::test]
13219    async fn concurrent_transition_cannot_duplicate_exit_actions() {
13220        let gate = PathMutationGate::new();
13221        let active = ai_agents_state::StateDefinition {
13222            on_exit: vec![StateAction::Tool {
13223                tool: "transition_exit".to_string(),
13224                args: Some(serde_json::json!({"path": "./transition-exit.txt"})),
13225            }],
13226            ..Default::default()
13227        };
13228        let state_machine = Arc::new(
13229            StateMachine::new(ai_agents_state::StateConfig {
13230                initial: "active".to_string(),
13231                states: HashMap::from([
13232                    ("active".to_string(), active),
13233                    (
13234                        "review".to_string(),
13235                        ai_agents_state::StateDefinition::default(),
13236                    ),
13237                ]),
13238                global_transitions: Vec::new(),
13239                fallback: None,
13240                max_no_transition: None,
13241                regenerate_on_transition: true,
13242            })
13243            .unwrap(),
13244        );
13245        let agent = Arc::new(
13246            AgentBuilder::new()
13247                .system_prompt("Test transition reservation.")
13248                .llm(Arc::new(mock_with_response("done")))
13249                .tool(Arc::new(BlockingPathMutationTool {
13250                    id: "transition_exit",
13251                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
13252                    gate: gate.clone(),
13253                }))
13254                .state_machine(state_machine)
13255                .build()
13256                .unwrap(),
13257        );
13258
13259        let first_agent = Arc::clone(&agent);
13260        let first = tokio::spawn(async move { first_agent.transition_to("review").await });
13261        tokio::time::timeout(std::time::Duration::from_secs(2), gate.wait_until_entered())
13262            .await
13263            .expect("reserved transition must enter its exit action");
13264
13265        let second = tokio::time::timeout(
13266            std::time::Duration::from_secs(2),
13267            agent.transition_to("review"),
13268        )
13269        .await
13270        .expect("competing transition must fail without waiting for the exit action")
13271        .unwrap_err();
13272        assert!(second.to_string().contains("already in progress"));
13273
13274        gate.release();
13275        first.await.unwrap().unwrap();
13276        assert_eq!(agent.current_state().as_deref(), Some("review"));
13277    }
13278
13279    /// Confirms a later transition cannot overtake the committed state's enter actions.
13280    #[tokio::test]
13281    async fn concurrent_transition_cannot_overtake_enter_actions() {
13282        let gate = PathMutationGate::new();
13283        let review = ai_agents_state::StateDefinition {
13284            on_enter: vec![StateAction::Tool {
13285                tool: "transition_enter".to_string(),
13286                args: Some(serde_json::json!({"path": "./transition-enter.txt"})),
13287            }],
13288            ..Default::default()
13289        };
13290        let state_machine = Arc::new(
13291            StateMachine::new(ai_agents_state::StateConfig {
13292                initial: "active".to_string(),
13293                states: HashMap::from([
13294                    (
13295                        "active".to_string(),
13296                        ai_agents_state::StateDefinition::default(),
13297                    ),
13298                    ("review".to_string(), review),
13299                ]),
13300                global_transitions: Vec::new(),
13301                fallback: None,
13302                max_no_transition: None,
13303                regenerate_on_transition: true,
13304            })
13305            .unwrap(),
13306        );
13307        let agent = Arc::new(
13308            AgentBuilder::new()
13309                .system_prompt("Test transition lifecycle reservation.")
13310                .llm(Arc::new(mock_with_response("done")))
13311                .tool(Arc::new(BlockingPathMutationTool {
13312                    id: "transition_enter",
13313                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
13314                    gate: gate.clone(),
13315                }))
13316                .state_machine(state_machine)
13317                .build()
13318                .unwrap(),
13319        );
13320
13321        let first_agent = Arc::clone(&agent);
13322        let first = tokio::spawn(async move { first_agent.transition_to("review").await });
13323        tokio::time::timeout(std::time::Duration::from_secs(2), gate.wait_until_entered())
13324            .await
13325            .expect("committed transition must enter its destination action");
13326
13327        let second = agent.transition_to("active").await.unwrap_err();
13328        assert!(second.to_string().contains("already in progress"));
13329        assert!(agent.reset().await.is_err());
13330
13331        gate.release();
13332        first.await.unwrap().unwrap();
13333        assert_eq!(agent.current_state().as_deref(), Some("review"));
13334    }
13335
13336    /// Confirms even a same-state restore invalidates manager and skill ownership.
13337    #[tokio::test]
13338    async fn same_state_restore_invalidates_pending_skill_confirmation() {
13339        let (agent, observed) = state_disambiguation_agent_with_skills(
13340            vec![
13341                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13342                "send_report",
13343                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13344                r#"{"question":"What should I send?","options":null}"#,
13345                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13346                r#"{"question":"Should I send the report to Ada?"}"#,
13347            ],
13348            true,
13349            None,
13350            true,
13351            vec![confirmation_skill()],
13352        );
13353
13354        agent.chat("Send it").await.unwrap();
13355        agent.chat("The report to Ada").await.unwrap();
13356        let snapshot = agent.save_state().await.unwrap();
13357        assert_eq!(agent.current_state().as_deref(), Some("active"));
13358
13359        agent.restore_state(snapshot).await.unwrap();
13360
13361        assert_eq!(agent.current_state().as_deref(), Some("active"));
13362        assert!(agent.pending_skill_id.read().is_none());
13363        assert!(
13364            !agent
13365                .disambiguation_manager()
13366                .unwrap()
13367                .has_pending_clarification()
13368                .await
13369        );
13370        assert_eq!(confirmation_skill_call_count(&observed), 0);
13371    }
13372
13373    /// Confirms direct state mutation cannot hide behind a return to the same state path.
13374    #[tokio::test]
13375    async fn direct_state_generation_change_invalidates_confirmation() {
13376        let (agent, observed) = state_disambiguation_agent_with_skills(
13377            vec![
13378                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13379                "send_report",
13380                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13381                r#"{"question":"What should I send?","options":null}"#,
13382                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13383                r#"{"question":"Should I send the report to Ada?"}"#,
13384                "Confirmation cancelled.",
13385            ],
13386            true,
13387            None,
13388            true,
13389            vec![confirmation_skill()],
13390        );
13391
13392        agent.chat("Send it").await.unwrap();
13393        agent.chat("The report to Ada").await.unwrap();
13394        let state_machine = agent.state_machine().unwrap();
13395        state_machine
13396            .transition_to("review", "external test")
13397            .unwrap();
13398        state_machine
13399            .transition_to("active", "external test")
13400            .unwrap();
13401
13402        let response = agent.chat("Yes").await.unwrap();
13403
13404        assert_eq!(response.content, "Confirmation cancelled.");
13405        assert_eq!(confirmation_skill_call_count(&observed), 0);
13406        assert!(agent.pending_skill_id.read().is_none());
13407    }
13408
13409    #[tokio::test]
13410    async fn state_confirmation_does_not_add_a_question_for_clear_input() {
13411        let (agent, observed) = state_disambiguation_agent(
13412            vec![
13413                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"clear","what_is_unclear":[],"detected_language":"en"}"#,
13414                "Request executed.",
13415            ],
13416            true,
13417            None,
13418            true,
13419        );
13420
13421        let response = agent.chat("Send the report to Ada").await.unwrap();
13422
13423        assert_eq!(response.content, "Request executed.");
13424        assert_eq!(observed.call_count(), 2);
13425    }
13426
13427    #[tokio::test]
13428    async fn state_override_cannot_activate_a_disabled_top_level_manager() {
13429        let (agent, observed) =
13430            state_disambiguation_agent(vec!["Request executed."], false, Some(true), true);
13431
13432        assert!(!agent.has_disambiguation());
13433        let response = agent.chat("Send it").await.unwrap();
13434
13435        assert_eq!(response.content, "Request executed.");
13436        assert_eq!(observed.call_count(), 1);
13437    }
13438
13439    #[tokio::test]
13440    async fn native_required_choice_executes_through_the_shared_tool_path() {
13441        let mut mock = MockLLMProvider::new("native-required");
13442        mock.set_tool_choice(Some(ToolChoice::Required));
13443        mock.add_response(
13444            LLMResponse::new("", FinishReason::ToolCall)
13445                .with_tool_calls(vec![ToolCall {
13446                    id: "provider-call-1".to_string(),
13447                    name: "calculator".to_string(),
13448                    arguments: serde_json::json!({"expression": "2 + 2"}),
13449                }])
13450                .unwrap(),
13451        );
13452        mock.add_response(LLMResponse::new("The answer is 4.", FinishReason::Stop));
13453        let observed = mock.clone();
13454        let agent = AgentBuilder::new()
13455            .system_prompt("Use the calculator when needed.")
13456            .llm(Arc::new(mock))
13457            .tool(Arc::new(CalculatorTool::new()))
13458            .build()
13459            .unwrap();
13460
13461        let response = agent.chat("What is 2 + 2?").await.unwrap();
13462
13463        assert_eq!(response.content, "The answer is 4.");
13464        assert_eq!(
13465            response.tool_calls.as_ref().unwrap()[0].id,
13466            "provider-call-1"
13467        );
13468        let calls = observed.call_history();
13469        assert_eq!(calls.len(), 2);
13470        assert!(matches!(
13471            calls[0].request.as_ref().map(|request| &request.choice),
13472            Some(ToolChoice::Required)
13473        ));
13474        assert!(matches!(
13475            calls[1].request.as_ref().map(|request| &request.choice),
13476            Some(ToolChoice::Auto)
13477        ));
13478    }
13479
13480    #[tokio::test]
13481    async fn prompt_fallback_uses_one_corrective_retry() {
13482        let mut mock = MockLLMProvider::new("prompt-required");
13483        mock.set_tool_choice(Some(ToolChoice::Required));
13484        mock.set_native_tool_support(false);
13485        mock.set_responses(
13486            vec![
13487                "I can calculate that.".to_string(),
13488                r#"{"tool":"calculator","arguments":{"expression":"2 + 2"}}"#.to_string(),
13489                "The answer is 4.".to_string(),
13490            ],
13491            false,
13492        );
13493        let observed = mock.clone();
13494        let agent = AgentBuilder::new()
13495            .system_prompt("Use tools.")
13496            .llm(Arc::new(mock))
13497            .tool(Arc::new(CalculatorTool::new()))
13498            .build()
13499            .unwrap();
13500
13501        let response = agent.chat("What is 2 + 2?").await.unwrap();
13502
13503        assert_eq!(response.content, "The answer is 4.");
13504        assert_eq!(observed.call_count(), 3);
13505        let corrective = &observed.call_history()[1].messages;
13506        assert!(
13507            corrective
13508                .last()
13509                .unwrap()
13510                .content
13511                .contains("previous response")
13512        );
13513    }
13514
13515    #[tokio::test]
13516    async fn prompt_fallback_fails_after_one_noncompliant_retry() {
13517        let mut mock = MockLLMProvider::new("prompt-required-failure");
13518        mock.set_tool_choice(Some(ToolChoice::Required));
13519        mock.set_native_tool_support(false);
13520        mock.set_responses(
13521            vec!["No tool.".to_string(), "Still no tool.".to_string()],
13522            false,
13523        );
13524        let observed = mock.clone();
13525        let agent = AgentBuilder::new()
13526            .system_prompt("Use tools.")
13527            .llm(Arc::new(mock))
13528            .tool(Arc::new(CalculatorTool::new()))
13529            .build()
13530            .unwrap();
13531
13532        let error = agent.chat("What is 2 + 2?").await.unwrap_err();
13533
13534        assert!(error.to_string().contains("one corrective retry"));
13535        assert_eq!(observed.call_count(), 2);
13536    }
13537
13538    #[tokio::test]
13539    async fn specific_choice_cannot_widen_the_effective_grant() {
13540        let mut mock = MockLLMProvider::new("specific-outside-grant");
13541        mock.set_tool_choice(Some(ToolChoice::Specific("random".to_string())));
13542        let observed = mock.clone();
13543        let agent = AgentBuilder::new()
13544            .system_prompt("Use tools.")
13545            .llm(Arc::new(mock))
13546            .tool(Arc::new(CalculatorTool::new()))
13547            .build()
13548            .unwrap();
13549
13550        let error = agent.chat("Generate a value.").await.unwrap_err();
13551
13552        assert!(error.to_string().contains("is not registered"));
13553        assert_eq!(observed.call_count(), 0);
13554    }
13555
13556    #[tokio::test]
13557    async fn none_choice_exposes_no_tool_protocol() {
13558        let mut mock = MockLLMProvider::new("no-tools");
13559        mock.set_tool_choice(Some(ToolChoice::None));
13560        mock.set_response(r#"{"tool":"calculator","arguments":{"expression":"2 + 2"}}"#);
13561        let observed = mock.clone();
13562        let agent = AgentBuilder::new()
13563            .system_prompt("Answer directly.")
13564            .llm(Arc::new(mock))
13565            .tool(Arc::new(CalculatorTool::new()))
13566            .build()
13567            .unwrap();
13568
13569        let response = agent.chat("Hello").await.unwrap();
13570
13571        assert!(response.tool_calls.is_none());
13572        assert_eq!(observed.call_count(), 1);
13573        let call = observed.last_call().unwrap();
13574        assert!(call.request.is_none());
13575        assert!(
13576            call.messages
13577                .iter()
13578                .all(|message| !message.content.contains("Available tools:"))
13579        );
13580    }
13581
13582    struct RuntimeStorage {
13583        capabilities: Box<[StorageCapability]>,
13584        snapshots: RwLock<HashMap<String, AgentSnapshot>>,
13585        metadata: RwLock<HashMap<String, ai_agents_core::SessionMetadata>>,
13586        metadata_save_calls: AtomicU64,
13587        metadata_load_calls: AtomicU64,
13588        fail_metadata_save: AtomicBool,
13589        fail_metadata_load: AtomicBool,
13590    }
13591
13592    impl RuntimeStorage {
13593        fn new(capabilities: impl IntoIterator<Item = StorageCapability>) -> Self {
13594            Self {
13595                capabilities: capabilities.into_iter().collect(),
13596                snapshots: RwLock::new(HashMap::new()),
13597                metadata: RwLock::new(HashMap::new()),
13598                metadata_save_calls: AtomicU64::new(0),
13599                metadata_load_calls: AtomicU64::new(0),
13600                fail_metadata_save: AtomicBool::new(false),
13601                fail_metadata_load: AtomicBool::new(false),
13602            }
13603        }
13604    }
13605
13606    #[async_trait]
13607    impl AgentStorage for RuntimeStorage {
13608        fn supports(&self, capability: StorageCapability) -> bool {
13609            self.capabilities.contains(&capability)
13610        }
13611
13612        async fn save(&self, session_id: &str, snapshot: &AgentSnapshot) -> Result<()> {
13613            self.snapshots
13614                .write()
13615                .insert(session_id.to_string(), snapshot.clone());
13616            Ok(())
13617        }
13618
13619        async fn load(&self, session_id: &str) -> Result<Option<AgentSnapshot>> {
13620            Ok(self.snapshots.read().get(session_id).cloned())
13621        }
13622
13623        async fn delete(&self, session_id: &str) -> Result<()> {
13624            self.snapshots.write().remove(session_id);
13625            Ok(())
13626        }
13627
13628        async fn list_sessions(&self) -> Result<Vec<String>> {
13629            Ok(self.snapshots.read().keys().cloned().collect())
13630        }
13631
13632        async fn save_snapshot_with_metadata(
13633            &self,
13634            session_id: &str,
13635            snapshot: &AgentSnapshot,
13636            metadata: &ai_agents_core::SessionMetadata,
13637        ) -> Result<()> {
13638            self.metadata_save_calls.fetch_add(1, Ordering::SeqCst);
13639            if self.fail_metadata_save.load(Ordering::SeqCst) {
13640                return Err(AgentError::Persistence("metadata save failed".into()));
13641            }
13642            self.snapshots
13643                .write()
13644                .insert(session_id.to_string(), snapshot.clone());
13645            self.metadata
13646                .write()
13647                .insert(session_id.to_string(), metadata.clone());
13648            Ok(())
13649        }
13650
13651        async fn save_metadata(
13652            &self,
13653            session_id: &str,
13654            metadata: &ai_agents_core::SessionMetadata,
13655        ) -> Result<()> {
13656            self.metadata_save_calls.fetch_add(1, Ordering::SeqCst);
13657            if self.fail_metadata_save.load(Ordering::SeqCst) {
13658                return Err(AgentError::Persistence("metadata save failed".into()));
13659            }
13660            self.metadata
13661                .write()
13662                .insert(session_id.to_string(), metadata.clone());
13663            Ok(())
13664        }
13665
13666        async fn load_metadata(
13667            &self,
13668            session_id: &str,
13669        ) -> Result<Option<ai_agents_core::SessionMetadata>> {
13670            self.metadata_load_calls.fetch_add(1, Ordering::SeqCst);
13671            if self.fail_metadata_load.load(Ordering::SeqCst) {
13672                return Err(AgentError::Persistence("metadata load failed".into()));
13673            }
13674            Ok(self.metadata.read().get(session_id).cloned())
13675        }
13676    }
13677
13678    fn runtime_storage_agent() -> RuntimeAgent {
13679        AgentBuilder::new()
13680            .system_prompt("Test runtime storage integration.")
13681            .llm(Arc::new(mock_with_response("done")))
13682            .build()
13683            .unwrap()
13684    }
13685
13686    fn restore_spec(id: &str) -> crate::spec::AgentSpec {
13687        crate::spec::AgentSpec {
13688            name: id.to_string(),
13689            system_prompt: format!("Restore child {id}."),
13690            ..crate::spec::AgentSpec::default()
13691        }
13692    }
13693
13694    fn restore_entry(id: &str) -> ai_agents_core::SpawnedAgentEntry {
13695        ai_agents_core::SpawnedAgentEntry {
13696            id: id.to_string(),
13697            name: id.to_string(),
13698            spec_yaml: serde_yaml::to_string(&restore_spec(id)).unwrap(),
13699        }
13700    }
13701
13702    fn restore_spawner(
13703        storage: Arc<RuntimeStorage>,
13704        max_agents: usize,
13705    ) -> (
13706        Arc<crate::spawner::AgentSpawner>,
13707        Arc<crate::spawner::AgentRegistry>,
13708    ) {
13709        let mut llms = LLMRegistry::new();
13710        llms.register("default", Arc::new(mock_with_response("done")));
13711        (
13712            Arc::new(
13713                crate::spawner::AgentSpawner::new()
13714                    .with_shared_llms(llms)
13715                    .with_shared_storage(storage)
13716                    .with_max_agents(max_agents),
13717            ),
13718            Arc::new(crate::spawner::AgentRegistry::new()),
13719        )
13720    }
13721
13722    async fn save_restore_target(
13723        parent: &RuntimeAgent,
13724        storage: &RuntimeStorage,
13725        session_id: &str,
13726        entries: Vec<ai_agents_core::SpawnedAgentEntry>,
13727    ) {
13728        let mut snapshot = parent.save_state().await.unwrap();
13729        snapshot.spawned_agents = Some(entries);
13730        storage.save(session_id, &snapshot).await.unwrap();
13731        storage
13732            .save_metadata(session_id, &ai_agents_core::SessionMetadata::default())
13733            .await
13734            .unwrap();
13735    }
13736
13737    #[tokio::test]
13738    async fn storage_init_requires_storage_for_actor_facts() {
13739        let facts = ai_agents_facts::FactsConfig {
13740            enabled: true,
13741            ..Default::default()
13742        };
13743        let agent = runtime_storage_agent().with_facts_config(None, Some(facts));
13744
13745        let error = agent.init_storage().await.unwrap_err();
13746        assert!(matches!(
13747            error,
13748            AgentError::Config(message)
13749                if message.contains("actor facts or actor memory")
13750                    && message.contains("none is configured or injected")
13751        ));
13752    }
13753
13754    #[tokio::test]
13755    async fn storage_init_validates_actor_facts_capability() {
13756        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
13757        let actor_memory = ai_agents_facts::ActorMemoryConfig {
13758            enabled: true,
13759            ..Default::default()
13760        };
13761        let agent = runtime_storage_agent()
13762            .with_storage(storage)
13763            .with_facts_config(Some(actor_memory), None);
13764
13765        assert!(matches!(
13766            agent.init_storage().await,
13767            Err(AgentError::UnsupportedStorageCapability(
13768                StorageCapability::ActorFacts
13769            ))
13770        ));
13771    }
13772
13773    #[tokio::test]
13774    async fn blocking_chat_rejects_unsupported_required_storage() {
13775        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
13776        let facts = ai_agents_facts::FactsConfig {
13777            enabled: true,
13778            ..Default::default()
13779        };
13780        let agent = runtime_storage_agent()
13781            .with_storage(storage)
13782            .with_facts_config(None, Some(facts));
13783
13784        assert!(matches!(
13785            agent.chat("hello").await,
13786            Err(AgentError::UnsupportedStorageCapability(
13787                StorageCapability::ActorFacts
13788            ))
13789        ));
13790    }
13791
13792    #[tokio::test]
13793    async fn streaming_chat_rejects_unsupported_required_storage_before_stream_creation() {
13794        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
13795        let config = ai_agents_relationships::RelationshipConfig {
13796            enabled: true,
13797            ..Default::default()
13798        };
13799        let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
13800        let agent = runtime_storage_agent()
13801            .with_storage(storage)
13802            .with_relationships(manager);
13803
13804        assert!(matches!(
13805            agent.chat_stream("hello").await,
13806            Err(AgentError::UnsupportedStorageCapability(
13807                StorageCapability::ActorRelationships
13808            ))
13809        ));
13810    }
13811
13812    #[tokio::test]
13813    async fn storage_init_completes_facts_for_injected_storage() {
13814        let storage = Arc::new(RuntimeStorage::new([
13815            StorageCapability::Snapshot,
13816            StorageCapability::ActorFacts,
13817        ]));
13818        let facts = ai_agents_facts::FactsConfig {
13819            enabled: true,
13820            ..Default::default()
13821        };
13822        let agent = runtime_storage_agent()
13823            .with_storage(storage)
13824            .with_facts_config(None, Some(facts));
13825
13826        agent.init_storage().await.unwrap();
13827        assert!(agent.fact_store().is_some());
13828    }
13829
13830    #[tokio::test]
13831    async fn storage_init_requires_storage_for_persistent_relationships() {
13832        let config = ai_agents_relationships::RelationshipConfig {
13833            enabled: true,
13834            ..Default::default()
13835        };
13836        let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
13837        let agent = runtime_storage_agent().with_relationships(manager);
13838
13839        let error = agent.init_storage().await.unwrap_err();
13840        assert!(matches!(
13841            error,
13842            AgentError::Config(message)
13843                if message.contains("persistent relationships")
13844                    && message.contains("none is configured or injected")
13845        ));
13846    }
13847
13848    #[tokio::test]
13849    async fn storage_init_validates_persistent_relationships_capability() {
13850        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
13851        let config = ai_agents_relationships::RelationshipConfig {
13852            enabled: true,
13853            ..Default::default()
13854        };
13855        let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
13856        let agent = runtime_storage_agent()
13857            .with_storage(storage)
13858            .with_relationships(manager);
13859
13860        assert!(matches!(
13861            agent.init_storage().await,
13862            Err(AgentError::UnsupportedStorageCapability(
13863                StorageCapability::ActorRelationships
13864            ))
13865        ));
13866    }
13867
13868    #[tokio::test]
13869    async fn session_restore_updates_identity_and_clears_stale_actor_binding() {
13870        let storage = Arc::new(RuntimeStorage::new([
13871            StorageCapability::Snapshot,
13872            StorageCapability::SessionMetadata,
13873        ]));
13874        let agent = runtime_storage_agent().with_storage(storage.clone());
13875        agent.set_actor_id("old-actor").unwrap();
13876        agent.save_session("old").await.unwrap();
13877        storage
13878            .save("target", &agent.save_state().await.unwrap())
13879            .await
13880            .unwrap();
13881        storage
13882            .save_metadata("target", &ai_agents_core::SessionMetadata::default())
13883            .await
13884            .unwrap();
13885
13886        assert!(agent.load_session("target").await.unwrap());
13887
13888        assert_eq!(agent.current_session_id.read().as_deref(), Some("target"));
13889        assert_eq!(agent.actor_id(), None);
13890    }
13891
13892    #[tokio::test]
13893    async fn complete_restore_reconciles_growth_shrink_and_empty_topologies() {
13894        let storage = Arc::new(RuntimeStorage::new([
13895            StorageCapability::Snapshot,
13896            StorageCapability::SessionMetadata,
13897        ]));
13898        let (spawner, registry) = restore_spawner(storage.clone(), 3);
13899        let parent = runtime_storage_agent()
13900            .with_storage(storage.clone())
13901            .with_spawner_handles(Arc::clone(&spawner), Arc::clone(&registry));
13902
13903        for id in ["a", "b"] {
13904            let spawned = spawner
13905                .spawn_with_id(id.to_string(), restore_spec(id))
13906                .await
13907                .unwrap();
13908            spawned.agent.save_session("grow").await.unwrap();
13909            registry.register(spawned).await.unwrap();
13910        }
13911        let staged_c = crate::spawner::storage::NamespacedStorage::new(storage.clone(), "c");
13912        staged_c
13913            .save("grow", &AgentSnapshot::new("c".into()))
13914            .await
13915            .unwrap();
13916        staged_c
13917            .save_metadata("grow", &ai_agents_core::SessionMetadata::default())
13918            .await
13919            .unwrap();
13920        save_restore_target(
13921            &parent,
13922            storage.as_ref(),
13923            "grow",
13924            vec![restore_entry("a"), restore_entry("b"), restore_entry("c")],
13925        )
13926        .await;
13927
13928        assert_eq!(parent.restore_session_full("grow").await.unwrap(), 3);
13929        assert_eq!(registry.count(), 3);
13930        assert!(registry.contains("c"));
13931        assert_eq!(spawner.spawned_count(), 3);
13932
13933        for id in ["a", "b"] {
13934            registry
13935                .get(id)
13936                .unwrap()
13937                .save_session("shrink")
13938                .await
13939                .unwrap();
13940        }
13941        save_restore_target(
13942            &parent,
13943            storage.as_ref(),
13944            "shrink",
13945            vec![restore_entry("a"), restore_entry("b")],
13946        )
13947        .await;
13948
13949        assert_eq!(parent.restore_session_full("shrink").await.unwrap(), 2);
13950        assert_eq!(registry.count(), 2);
13951        assert!(!registry.contains("c"));
13952        assert_eq!(spawner.spawned_count(), 2);
13953
13954        save_restore_target(&parent, storage.as_ref(), "empty", Vec::new()).await;
13955
13956        assert_eq!(parent.restore_session_full("empty").await.unwrap(), 0);
13957        assert_eq!(registry.count(), 0);
13958        assert_eq!(spawner.spawned_count(), 0);
13959        assert_eq!(parent.current_session_id.read().as_deref(), Some("empty"));
13960    }
13961
13962    #[tokio::test]
13963    async fn storage_session_metadata_is_called_only_when_advertised() {
13964        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
13965        storage.fail_metadata_save.store(true, Ordering::SeqCst);
13966        storage.fail_metadata_load.store(true, Ordering::SeqCst);
13967        let agent = runtime_storage_agent().with_storage(storage.clone());
13968
13969        agent.save_session("session").await.unwrap();
13970        assert!(agent.load_session("session").await.unwrap());
13971        assert_eq!(storage.metadata_save_calls.load(Ordering::SeqCst), 0);
13972        assert_eq!(storage.metadata_load_calls.load(Ordering::SeqCst), 0);
13973    }
13974
13975    #[cfg(feature = "sqlite")]
13976    #[tokio::test]
13977    async fn sqlite_runtime_save_filter_reopen_and_reload_stay_consistent() {
13978        let directory =
13979            std::env::temp_dir().join(format!("ai-agents-runtime-sqlite-{}", uuid::Uuid::new_v4()));
13980        let path = directory.join("sessions.sqlite");
13981        let path_string = path.to_string_lossy().into_owned();
13982        let storage = Arc::new(
13983            ai_agents_storage::SqliteStorage::new(&path_string)
13984                .await
13985                .unwrap(),
13986        );
13987        let agent = runtime_storage_agent().with_storage(storage.clone());
13988        agent.set_session_metadata(ai_agents_core::SessionMetadata {
13989            tags: vec!["initial".into()],
13990            ..Default::default()
13991        });
13992        agent.chat("persist this turn").await.unwrap();
13993        agent.save_session("session").await.unwrap();
13994
13995        agent.set_session_metadata(ai_agents_core::SessionMetadata {
13996            tags: vec!["updated".into()],
13997            ..Default::default()
13998        });
13999        agent.save_session("session").await.unwrap();
14000        assert!(
14001            agent
14002                .list_sessions_filtered(&ai_agents_core::SessionFilter {
14003                    tags: Some(vec!["initial".into()]),
14004                    ..Default::default()
14005                })
14006                .await
14007                .unwrap()
14008                .is_empty()
14009        );
14010        assert_eq!(
14011            agent
14012                .list_sessions_filtered(&ai_agents_core::SessionFilter {
14013                    tags: Some(vec!["updated".into()]),
14014                    ..Default::default()
14015                })
14016                .await
14017                .unwrap()
14018                .len(),
14019            1
14020        );
14021        drop(agent);
14022        storage.close().await;
14023        drop(storage);
14024
14025        let reopened_storage = Arc::new(
14026            ai_agents_storage::SqliteStorage::new(&path_string)
14027                .await
14028                .unwrap(),
14029        );
14030        let restored = runtime_storage_agent().with_storage(reopened_storage.clone());
14031        assert!(restored.load_session("session").await.unwrap());
14032        assert_eq!(restored.session_metadata().tags, vec!["updated"]);
14033        assert_eq!(
14034            restored.current_session_id.read().as_deref(),
14035            Some("session")
14036        );
14037        assert!(restored.save_state().await.unwrap().memory.messages.len() >= 2);
14038        assert_eq!(
14039            restored
14040                .list_sessions_filtered(&ai_agents_core::SessionFilter {
14041                    tags: Some(vec!["updated".into()]),
14042                    ..Default::default()
14043                })
14044                .await
14045                .unwrap()
14046                .len(),
14047            1
14048        );
14049
14050        drop(restored);
14051        reopened_storage.close().await;
14052        drop(reopened_storage);
14053        crate::remove_sqlite_test_directory(&directory)
14054            .await
14055            .unwrap();
14056    }
14057
14058    #[tokio::test]
14059    async fn storage_session_metadata_backend_failures_propagate() {
14060        let storage = Arc::new(RuntimeStorage::new([
14061            StorageCapability::Snapshot,
14062            StorageCapability::SessionMetadata,
14063        ]));
14064        let agent = runtime_storage_agent().with_storage(storage.clone());
14065
14066        agent.save_session("session").await.unwrap();
14067        storage
14068            .save("target", &agent.save_state().await.unwrap())
14069            .await
14070            .unwrap();
14071        storage.fail_metadata_load.store(true, Ordering::SeqCst);
14072        assert!(matches!(
14073            agent.load_session("target").await,
14074            Err(AgentError::Persistence(message)) if message == "metadata load failed"
14075        ));
14076        assert_eq!(agent.current_session_id.read().as_deref(), Some("session"));
14077
14078        storage.fail_metadata_save.store(true, Ordering::SeqCst);
14079        assert!(matches!(
14080            agent.save_session("session").await,
14081            Err(AgentError::Persistence(message)) if message == "metadata save failed"
14082        ));
14083    }
14084
14085    struct ProviderFutureDropSignal {
14086        dropped: Arc<AtomicBool>,
14087    }
14088
14089    impl Drop for ProviderFutureDropSignal {
14090        fn drop(&mut self) {
14091            self.dropped.store(true, Ordering::SeqCst);
14092        }
14093    }
14094
14095    struct BufferedLockingProvider {
14096        lock: Arc<tokio::sync::Mutex<()>>,
14097        stream_started: Arc<tokio::sync::Notify>,
14098        stream_dropped: Arc<AtomicBool>,
14099        committed_after_drop: Arc<AtomicBool>,
14100    }
14101
14102    #[async_trait]
14103    impl LLMProvider for BufferedLockingProvider {
14104        async fn complete(
14105            &self,
14106            _messages: &[ChatMessage],
14107            _config: Option<&LLMConfig>,
14108        ) -> std::result::Result<LLMResponse, LLMError> {
14109            let _guard = self.lock.lock().await;
14110            self.committed_after_drop
14111                .store(self.stream_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14112            Ok(LLMResponse::new(
14113                "Committed technical response.",
14114                FinishReason::Stop,
14115            ))
14116        }
14117
14118        async fn complete_stream(
14119            &self,
14120            _messages: &[ChatMessage],
14121            _config: Option<&LLMConfig>,
14122        ) -> std::result::Result<
14123            Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14124            LLMError,
14125        > {
14126            let _guard = self.lock.lock().await;
14127            let _drop_signal = ProviderFutureDropSignal {
14128                dropped: Arc::clone(&self.stream_dropped),
14129            };
14130            self.stream_started.notify_one();
14131            std::future::pending().await
14132        }
14133
14134        fn provider_name(&self) -> &str {
14135            "buffered-locking"
14136        }
14137
14138        fn supports(&self, _feature: LLMFeature) -> bool {
14139            false
14140        }
14141    }
14142
14143    struct PendingDropStream {
14144        dropped: Arc<AtomicBool>,
14145        dropped_notify: Arc<tokio::sync::Notify>,
14146    }
14147
14148    impl Stream for PendingDropStream {
14149        type Item = std::result::Result<LLMChunk, LLMError>;
14150
14151        fn poll_next(
14152            self: Pin<&mut Self>,
14153            _cx: &mut std::task::Context<'_>,
14154        ) -> std::task::Poll<Option<Self::Item>> {
14155            std::task::Poll::Pending
14156        }
14157    }
14158
14159    impl Drop for PendingDropStream {
14160        fn drop(&mut self) {
14161            self.dropped.store(true, Ordering::SeqCst);
14162            self.dropped_notify.notify_one();
14163        }
14164    }
14165
14166    struct EstablishedStreamProvider {
14167        stream_started: Arc<tokio::sync::Notify>,
14168        stream_dropped: Arc<AtomicBool>,
14169        stream_dropped_notify: Arc<tokio::sync::Notify>,
14170        committed_after_drop: Arc<AtomicBool>,
14171    }
14172
14173    #[async_trait]
14174    impl LLMProvider for EstablishedStreamProvider {
14175        async fn complete(
14176            &self,
14177            _messages: &[ChatMessage],
14178            _config: Option<&LLMConfig>,
14179        ) -> std::result::Result<LLMResponse, LLMError> {
14180            if !self.stream_dropped.load(Ordering::SeqCst) {
14181                self.stream_dropped_notify.notified().await;
14182            }
14183            self.committed_after_drop
14184                .store(self.stream_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14185            Ok(LLMResponse::new(
14186                "Committed technical response.",
14187                FinishReason::Stop,
14188            ))
14189        }
14190
14191        async fn complete_stream(
14192            &self,
14193            _messages: &[ChatMessage],
14194            _config: Option<&LLMConfig>,
14195        ) -> std::result::Result<
14196            Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14197            LLMError,
14198        > {
14199            self.stream_started.notify_one();
14200            Ok(Box::new(PendingDropStream {
14201                dropped: Arc::clone(&self.stream_dropped),
14202                dropped_notify: Arc::clone(&self.stream_dropped_notify),
14203            }))
14204        }
14205
14206        fn provider_name(&self) -> &str {
14207            "established-stream"
14208        }
14209
14210        fn supports(&self, _feature: LLMFeature) -> bool {
14211            false
14212        }
14213    }
14214
14215    struct FirstCallLockingProvider {
14216        lock: Arc<tokio::sync::Mutex<()>>,
14217        first_started: Arc<tokio::sync::Notify>,
14218        first_dropped: Arc<AtomicBool>,
14219        committed_after_drop: Arc<AtomicBool>,
14220        calls: AtomicU64,
14221    }
14222
14223    #[async_trait]
14224    impl LLMProvider for FirstCallLockingProvider {
14225        async fn complete(
14226            &self,
14227            _messages: &[ChatMessage],
14228            _config: Option<&LLMConfig>,
14229        ) -> std::result::Result<LLMResponse, LLMError> {
14230            let _guard = self.lock.lock().await;
14231            let call = self.calls.fetch_add(1, Ordering::SeqCst);
14232            if call == 0 {
14233                let _drop_signal = ProviderFutureDropSignal {
14234                    dropped: Arc::clone(&self.first_dropped),
14235                };
14236                self.first_started.notify_one();
14237                return std::future::pending().await;
14238            }
14239            self.committed_after_drop
14240                .store(self.first_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14241            Ok(LLMResponse::new(
14242                "Committed technical response.",
14243                FinishReason::Stop,
14244            ))
14245        }
14246
14247        async fn complete_stream(
14248            &self,
14249            _messages: &[ChatMessage],
14250            _config: Option<&LLMConfig>,
14251        ) -> std::result::Result<
14252            Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14253            LLMError,
14254        > {
14255            Err(LLMError::Other(
14256                "streaming is not used in this test".to_string(),
14257            ))
14258        }
14259
14260        fn provider_name(&self) -> &str {
14261            "first-call-locking"
14262        }
14263
14264        fn supports(&self, _feature: LLMFeature) -> bool {
14265            false
14266        }
14267    }
14268
14269    struct RoutingAfterProviderStart {
14270        provider_started: Arc<tokio::sync::Notify>,
14271    }
14272
14273    #[async_trait]
14274    impl LLMProvider for RoutingAfterProviderStart {
14275        async fn complete(
14276            &self,
14277            _messages: &[ChatMessage],
14278            _config: Option<&LLMConfig>,
14279        ) -> std::result::Result<LLMResponse, LLMError> {
14280            self.provider_started.notified().await;
14281            Ok(LLMResponse::new("1", FinishReason::Stop))
14282        }
14283
14284        async fn complete_stream(
14285            &self,
14286            _messages: &[ChatMessage],
14287            _config: Option<&LLMConfig>,
14288        ) -> std::result::Result<
14289            Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14290            LLMError,
14291        > {
14292            Err(LLMError::Other(
14293                "streaming is not used in this test".to_string(),
14294            ))
14295        }
14296
14297        fn provider_name(&self) -> &str {
14298            "routing-after-start"
14299        }
14300
14301        fn supports(&self, _feature: LLMFeature) -> bool {
14302            false
14303        }
14304    }
14305
14306    /// Test hook that counts completed responses.
14307    struct ResponseCountingHooks {
14308        responses: Arc<std::sync::atomic::AtomicUsize>,
14309    }
14310
14311    /// Test tool that returns the execution context it received.
14312    struct ContextEchoTool;
14313
14314    #[async_trait]
14315    impl ai_agents_core::Tool for ContextEchoTool {
14316        fn id(&self) -> &str {
14317            "context_echo"
14318        }
14319
14320        fn name(&self) -> &str {
14321            "Context Echo"
14322        }
14323
14324        fn description(&self) -> &str {
14325            "Returns selected execution context fields."
14326        }
14327
14328        fn input_schema(&self) -> Value {
14329            serde_json::json!({"type": "object"})
14330        }
14331
14332        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14333            ai_agents_core::ToolPolicyBindings {
14334                path_fields: vec![ai_agents_core::PathPolicyBinding::read("path")],
14335                result_limit_fields: vec![ai_agents_core::ResultLimitBinding::new(
14336                    "max_results",
14337                    ai_agents_core::ResultLimitKind::MaxResults,
14338                )],
14339                ..Default::default()
14340            }
14341        }
14342
14343        async fn execute(
14344            &self,
14345            _args: Value,
14346            ctx: ai_agents_core::ToolExecutionContext,
14347        ) -> ToolResult {
14348            ToolResult::ok(
14349                serde_json::json!({
14350                    "requested_name": ctx.requested_name,
14351                    "canonical_id": ctx.canonical_id,
14352                    "display_name": ctx.display_name,
14353                    "max_results": ctx.limits.max_results,
14354                    "custom_config": ctx.custom_config,
14355                })
14356                .to_string(),
14357            )
14358        }
14359    }
14360
14361    /// Test tool that stays active long enough for runtime cancellation.
14362    struct SlowTool;
14363
14364    /// Test tool that fails once and must not be retried for writes.
14365    struct FlakyWriteTool {
14366        calls: Arc<std::sync::atomic::AtomicUsize>,
14367    }
14368
14369    /// Test tool that tracks concurrent execution on one path.
14370    struct LockedWriteTool {
14371        active: Arc<std::sync::atomic::AtomicUsize>,
14372        max_active: Arc<std::sync::atomic::AtomicUsize>,
14373    }
14374
14375    struct MultiResourceWriteTool {
14376        active: Arc<std::sync::atomic::AtomicUsize>,
14377        max_active: Arc<std::sync::atomic::AtomicUsize>,
14378    }
14379
14380    #[derive(Clone)]
14381    struct PathMutationGate {
14382        entered: Arc<AtomicBool>,
14383        entered_notify: Arc<tokio::sync::Notify>,
14384        release: Arc<tokio::sync::Notify>,
14385    }
14386
14387    impl PathMutationGate {
14388        fn new() -> Self {
14389            Self {
14390                entered: Arc::new(AtomicBool::new(false)),
14391                entered_notify: Arc::new(tokio::sync::Notify::new()),
14392                release: Arc::new(tokio::sync::Notify::new()),
14393            }
14394        }
14395
14396        async fn wait_until_entered(&self) {
14397            if !self.entered.load(Ordering::SeqCst) {
14398                self.entered_notify.notified().await;
14399            }
14400        }
14401
14402        fn release(&self) {
14403            self.release.notify_one();
14404        }
14405    }
14406
14407    struct BlockingPathMutationTool {
14408        id: &'static str,
14409        path_fields: Vec<ai_agents_core::PathPolicyBinding>,
14410        gate: PathMutationGate,
14411    }
14412
14413    struct NoBindingWriteTool {
14414        active: Arc<std::sync::atomic::AtomicUsize>,
14415        max_active: Arc<std::sync::atomic::AtomicUsize>,
14416    }
14417
14418    struct RecoveryTestTool {
14419        id: String,
14420        succeeds: bool,
14421        calls: Arc<std::sync::atomic::AtomicUsize>,
14422    }
14423
14424    struct BlockingApprovalHandler {
14425        entered: Arc<tokio::sync::Barrier>,
14426        release: Arc<tokio::sync::Notify>,
14427        result: ApprovalResult,
14428    }
14429
14430    struct RuntimeWebFetchTransport {
14431        calls: Arc<std::sync::atomic::AtomicUsize>,
14432    }
14433
14434    struct RuntimeWebFetchResolver;
14435
14436    struct ReentrantToolHooks {
14437        agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
14438        invoked: AtomicBool,
14439        nested_success: AtomicBool,
14440    }
14441
14442    #[async_trait]
14443    impl ai_agents_core::Tool for SlowTool {
14444        fn id(&self) -> &str {
14445            "slow"
14446        }
14447
14448        fn name(&self) -> &str {
14449            "Slow"
14450        }
14451
14452        fn description(&self) -> &str {
14453            "Waits until cancelled or timed out."
14454        }
14455
14456        fn input_schema(&self) -> Value {
14457            serde_json::json!({"type": "object"})
14458        }
14459
14460        async fn execute(
14461            &self,
14462            _args: Value,
14463            _ctx: ai_agents_core::ToolExecutionContext,
14464        ) -> ToolResult {
14465            tokio::time::sleep(std::time::Duration::from_secs(5)).await;
14466            ToolResult::ok("done")
14467        }
14468    }
14469
14470    #[async_trait]
14471    impl ai_agents_core::Tool for FlakyWriteTool {
14472        fn id(&self) -> &str {
14473            "flaky_write"
14474        }
14475
14476        fn name(&self) -> &str {
14477            "Flaky Write"
14478        }
14479
14480        fn description(&self) -> &str {
14481            "Fails on the first write attempt."
14482        }
14483
14484        fn input_schema(&self) -> Value {
14485            serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
14486        }
14487
14488        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14489            ai_agents_core::ToolPolicyBindings {
14490                path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
14491                ..Default::default()
14492            }
14493        }
14494
14495        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14496            ai_agents_core::ToolSafetyMetadata {
14497                read_only: false,
14498                concurrency_safe: false,
14499                operation: ai_agents_core::ToolOperationKind::Write,
14500                side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
14501                requires_network: false,
14502                destructive: false,
14503                open_world: false,
14504                host_dependent: false,
14505                requires_user_interaction: false,
14506                supports_cancellation: true,
14507                default_requires_approval: false,
14508                should_defer_schema: false,
14509                max_output_chars: Some(1024),
14510                max_result_size_chars: Some(1024),
14511            }
14512        }
14513
14514        fn classify_call(&self, _args: &Value) -> ai_agents_core::ToolCallClassification {
14515            let mut classification =
14516                ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
14517            classification.safely_retryable = false;
14518            classification
14519        }
14520
14521        async fn execute(
14522            &self,
14523            _args: Value,
14524            _ctx: ai_agents_core::ToolExecutionContext,
14525        ) -> ToolResult {
14526            let call = self.calls.fetch_add(1, Ordering::SeqCst);
14527            if call == 0 {
14528                ToolResult::error("first failure")
14529            } else {
14530                ToolResult::ok("second success")
14531            }
14532        }
14533    }
14534
14535    #[async_trait]
14536    impl ai_agents_core::Tool for LockedWriteTool {
14537        fn id(&self) -> &str {
14538            "locked_write"
14539        }
14540
14541        fn name(&self) -> &str {
14542            "Locked Write"
14543        }
14544
14545        fn description(&self) -> &str {
14546            "Tracks concurrent execution on one resource."
14547        }
14548
14549        fn input_schema(&self) -> Value {
14550            serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
14551        }
14552
14553        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14554            ai_agents_core::ToolPolicyBindings {
14555                path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
14556                ..Default::default()
14557            }
14558        }
14559
14560        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14561            ai_agents_core::ToolSafetyMetadata {
14562                read_only: false,
14563                concurrency_safe: false,
14564                operation: ai_agents_core::ToolOperationKind::Write,
14565                side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
14566                requires_network: false,
14567                destructive: false,
14568                open_world: false,
14569                host_dependent: false,
14570                requires_user_interaction: false,
14571                supports_cancellation: true,
14572                default_requires_approval: false,
14573                should_defer_schema: false,
14574                max_output_chars: Some(1024),
14575                max_result_size_chars: Some(1024),
14576            }
14577        }
14578
14579        async fn execute(
14580            &self,
14581            _args: Value,
14582            _ctx: ai_agents_core::ToolExecutionContext,
14583        ) -> ToolResult {
14584            let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
14585            loop {
14586                let current_max = self.max_active.load(Ordering::SeqCst);
14587                if active <= current_max {
14588                    break;
14589                }
14590                if self
14591                    .max_active
14592                    .compare_exchange(current_max, active, Ordering::SeqCst, Ordering::SeqCst)
14593                    .is_ok()
14594                {
14595                    break;
14596                }
14597            }
14598            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
14599            self.active.fetch_sub(1, Ordering::SeqCst);
14600            ToolResult::ok("done")
14601        }
14602    }
14603
14604    #[async_trait]
14605    impl ai_agents_core::Tool for MultiResourceWriteTool {
14606        fn id(&self) -> &str {
14607            "multi_resource_write"
14608        }
14609
14610        fn name(&self) -> &str {
14611            "Multi Resource Write"
14612        }
14613
14614        fn description(&self) -> &str {
14615            "Tracks concurrent execution across source and destination resources."
14616        }
14617
14618        fn input_schema(&self) -> Value {
14619            serde_json::json!({"type": "object"})
14620        }
14621
14622        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14623            ai_agents_core::ToolPolicyBindings {
14624                path_fields: vec![
14625                    ai_agents_core::PathPolicyBinding::read_write("source_path"),
14626                    ai_agents_core::PathPolicyBinding::write("destination_path"),
14627                ],
14628                ..Default::default()
14629            }
14630        }
14631
14632        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14633            LockedWriteTool {
14634                active: Arc::clone(&self.active),
14635                max_active: Arc::clone(&self.max_active),
14636            }
14637            .safety_metadata()
14638        }
14639
14640        async fn execute(
14641            &self,
14642            _args: Value,
14643            _ctx: ai_agents_core::ToolExecutionContext,
14644        ) -> ToolResult {
14645            let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
14646            self.max_active.fetch_max(active, Ordering::SeqCst);
14647            tokio::time::sleep(std::time::Duration::from_millis(75)).await;
14648            self.active.fetch_sub(1, Ordering::SeqCst);
14649            ToolResult::ok("done")
14650        }
14651    }
14652
14653    #[async_trait]
14654    impl ai_agents_core::Tool for BlockingPathMutationTool {
14655        fn id(&self) -> &str {
14656            self.id
14657        }
14658
14659        fn name(&self) -> &str {
14660            self.id
14661        }
14662
14663        fn description(&self) -> &str {
14664            "Blocks a path mutation until the test releases it."
14665        }
14666
14667        fn input_schema(&self) -> Value {
14668            serde_json::json!({"type": "object"})
14669        }
14670
14671        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14672            ai_agents_core::ToolPolicyBindings {
14673                path_fields: self.path_fields.clone(),
14674                ..Default::default()
14675            }
14676        }
14677
14678        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14679            ai_agents_core::ToolSafetyMetadata {
14680                read_only: false,
14681                concurrency_safe: false,
14682                operation: ai_agents_core::ToolOperationKind::Write,
14683                side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
14684                requires_network: false,
14685                destructive: false,
14686                open_world: false,
14687                host_dependent: false,
14688                requires_user_interaction: false,
14689                supports_cancellation: true,
14690                default_requires_approval: false,
14691                should_defer_schema: false,
14692                max_output_chars: Some(1024),
14693                max_result_size_chars: Some(1024),
14694            }
14695        }
14696
14697        async fn execute(
14698            &self,
14699            _args: Value,
14700            _ctx: ai_agents_core::ToolExecutionContext,
14701        ) -> ToolResult {
14702            self.gate.entered.store(true, Ordering::SeqCst);
14703            self.gate.entered_notify.notify_one();
14704            self.gate.release.notified().await;
14705            ToolResult::ok("done")
14706        }
14707    }
14708
14709    #[async_trait]
14710    impl ai_agents_core::Tool for NoBindingWriteTool {
14711        fn id(&self) -> &str {
14712            "no_binding_write"
14713        }
14714
14715        fn name(&self) -> &str {
14716            "No Binding Write"
14717        }
14718
14719        fn description(&self) -> &str {
14720            "Tracks concurrent execution without resource bindings."
14721        }
14722
14723        fn input_schema(&self) -> Value {
14724            serde_json::json!({"type": "object"})
14725        }
14726
14727        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14728            LockedWriteTool {
14729                active: Arc::clone(&self.active),
14730                max_active: Arc::clone(&self.max_active),
14731            }
14732            .safety_metadata()
14733        }
14734
14735        async fn execute(
14736            &self,
14737            _args: Value,
14738            _ctx: ai_agents_core::ToolExecutionContext,
14739        ) -> ToolResult {
14740            let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
14741            self.max_active.fetch_max(active, Ordering::SeqCst);
14742            tokio::time::sleep(std::time::Duration::from_millis(75)).await;
14743            self.active.fetch_sub(1, Ordering::SeqCst);
14744            ToolResult::ok("done")
14745        }
14746    }
14747
14748    #[async_trait]
14749    impl ai_agents_core::Tool for RecoveryTestTool {
14750        fn id(&self) -> &str {
14751            &self.id
14752        }
14753
14754        fn name(&self) -> &str {
14755            &self.id
14756        }
14757
14758        fn description(&self) -> &str {
14759            "Records recovery execution and returns a configured result."
14760        }
14761
14762        fn input_schema(&self) -> Value {
14763            serde_json::json!({"type": "object"})
14764        }
14765
14766        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14767            ai_agents_core::ToolPolicyBindings {
14768                path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
14769                ..Default::default()
14770            }
14771        }
14772
14773        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14774            ai_agents_core::ToolSafetyMetadata {
14775                read_only: false,
14776                concurrency_safe: false,
14777                operation: ai_agents_core::ToolOperationKind::Write,
14778                side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
14779                requires_network: false,
14780                destructive: false,
14781                open_world: false,
14782                host_dependent: false,
14783                requires_user_interaction: false,
14784                supports_cancellation: true,
14785                default_requires_approval: false,
14786                should_defer_schema: false,
14787                max_output_chars: Some(1024),
14788                max_result_size_chars: Some(1024),
14789            }
14790        }
14791
14792        async fn execute(
14793            &self,
14794            _args: Value,
14795            _ctx: ai_agents_core::ToolExecutionContext,
14796        ) -> ToolResult {
14797            self.calls.fetch_add(1, Ordering::SeqCst);
14798            if self.succeeds {
14799                ToolResult::ok(format!("{} succeeded", self.id))
14800            } else {
14801                ToolResult::error(format!("{} failed", self.id))
14802            }
14803        }
14804    }
14805
14806    #[async_trait]
14807    impl WebFetchTransport for RuntimeWebFetchTransport {
14808        /// Rejects unvalidated transport calls in the shared-executor fixture.
14809        async fn send(
14810            &self,
14811            _request: WebFetchTransportRequest,
14812        ) -> std::result::Result<WebFetchTransportResponse, String> {
14813            Err("validated addresses are required".to_string())
14814        }
14815
14816        /// Records transport only after runtime and tool policy validation complete.
14817        async fn send_validated(
14818            &self,
14819            _request: WebFetchTransportRequest,
14820            _addresses: &[std::net::SocketAddr],
14821        ) -> std::result::Result<WebFetchTransportResponse, String> {
14822            self.calls.fetch_add(1, Ordering::SeqCst);
14823            Ok(WebFetchTransportResponse {
14824                status: 200,
14825                content_type: Some("text/plain".to_string()),
14826                location: None,
14827                body: b"approved".to_vec(),
14828            })
14829        }
14830    }
14831
14832    #[async_trait]
14833    impl WebFetchResolver for RuntimeWebFetchResolver {
14834        /// Returns one public fixture address without external DNS.
14835        async fn resolve(
14836            &self,
14837            _host: &str,
14838            _port: u16,
14839        ) -> std::result::Result<Vec<std::net::IpAddr>, String> {
14840            Ok(vec![std::net::IpAddr::V4(std::net::Ipv4Addr::new(
14841                93, 184, 216, 34,
14842            ))])
14843        }
14844    }
14845
14846    #[async_trait]
14847    impl ApprovalHandler for BlockingApprovalHandler {
14848        async fn request_approval(
14849            &self,
14850            _request: ai_agents_hitl::ApprovalRequest,
14851        ) -> ApprovalResult {
14852            self.entered.wait().await;
14853            self.release.notified().await;
14854            self.result.clone()
14855        }
14856    }
14857
14858    #[async_trait]
14859    impl AgentHooks for ReentrantToolHooks {
14860        async fn on_tool_complete(&self, tool: &str, _result: &ToolResult, _duration_ms: u64) {
14861            if tool != "reentrant_write" || self.invoked.swap(true, Ordering::SeqCst) {
14862                return;
14863            }
14864            let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
14865            if let Some(agent) = agent {
14866                let result = agent
14867                    .invoke_tool(ToolExecutionRequest::new(
14868                        "nested-hook-call",
14869                        "reentrant_write",
14870                        serde_json::json!({"path": "./hook.txt"}),
14871                        ToolCallSource::Manual,
14872                    ))
14873                    .await;
14874                self.nested_success
14875                    .store(result.is_ok_and(|record| record.success), Ordering::SeqCst);
14876            }
14877        }
14878    }
14879
14880    #[async_trait]
14881    impl AgentHooks for ResponseCountingHooks {
14882        async fn on_response(&self, _response: &AgentResponse) {
14883            self.responses.fetch_add(1, Ordering::SeqCst);
14884        }
14885    }
14886
14887    struct ApprovalRecordingHooks {
14888        events: parking_lot::Mutex<Vec<String>>,
14889    }
14890
14891    impl ApprovalRecordingHooks {
14892        fn new() -> Self {
14893            Self {
14894                events: parking_lot::Mutex::new(Vec::new()),
14895            }
14896        }
14897
14898        fn events(&self) -> Vec<String> {
14899            self.events.lock().clone()
14900        }
14901    }
14902
14903    #[async_trait]
14904    impl AgentHooks for ApprovalRecordingHooks {
14905        async fn on_approval_result(&self, request_id: &str, result: &ApprovalResult) {
14906            self.events.lock().push(format!(
14907                "raw:{}:{}",
14908                request_id,
14909                approval_result_name(result)
14910            ));
14911        }
14912
14913        async fn on_approval_resolved(
14914            &self,
14915            request: &ai_agents_hitl::ApprovalRequest,
14916            raw_result: &ApprovalResult,
14917            outcome: &ApprovalResolvedOutcome,
14918        ) {
14919            self.events.lock().push(format!(
14920                "resolved:{}:{}:{}",
14921                request.id,
14922                approval_result_name(raw_result),
14923                approval_outcome_name(outcome)
14924            ));
14925        }
14926    }
14927
14928    fn approval_result_name(result: &ApprovalResult) -> &'static str {
14929        match result {
14930            ApprovalResult::Approved => "approved",
14931            ApprovalResult::Rejected { .. } => "rejected",
14932            ApprovalResult::Modified { .. } => "modified",
14933            ApprovalResult::Timeout => "timeout",
14934        }
14935    }
14936
14937    fn approval_outcome_name(outcome: &ApprovalResolvedOutcome) -> &'static str {
14938        match outcome {
14939            ApprovalResolvedOutcome::Approved => "approved",
14940            ApprovalResolvedOutcome::Rejected { .. } => "rejected",
14941            ApprovalResolvedOutcome::Modified { .. } => "modified",
14942            ApprovalResolvedOutcome::Error { .. } => "error",
14943        }
14944    }
14945
14946    fn assert_correlated_approval_events(
14947        events: &[String],
14948        raw_status: &str,
14949        outcome_status: &str,
14950    ) {
14951        assert_eq!(events.len(), 2);
14952        let raw: Vec<_> = events[0].split(':').collect();
14953        let resolved: Vec<_> = events[1].split(':').collect();
14954        assert_eq!(raw[0], "raw");
14955        assert_eq!(resolved[0], "resolved");
14956        assert_eq!(raw[1], resolved[1]);
14957        assert_eq!(raw[2], raw_status);
14958        assert_eq!(resolved[2], raw_status);
14959        assert_eq!(resolved[3], outcome_status);
14960    }
14961
14962    fn approval_security_config(policy_enabled: bool) -> ToolSecurityConfig {
14963        let mut security = ToolSecurityConfig {
14964            enabled: true,
14965            fail_closed: true,
14966            ..Default::default()
14967        };
14968        let policy = ai_agents_tools::ToolPolicyConfig {
14969            enabled: policy_enabled,
14970            write_paths: vec![".".to_string()],
14971            require_confirmation: true,
14972            ..Default::default()
14973        };
14974        security.tools.insert("locked_write".to_string(), policy);
14975        security
14976    }
14977
14978    struct MutationTestWorkspace {
14979        root: std::path::PathBuf,
14980    }
14981
14982    impl MutationTestWorkspace {
14983        fn new() -> Self {
14984            let root = std::env::temp_dir().join(format!(
14985                "ai-agents-runtime-mutation-{}",
14986                uuid::Uuid::new_v4()
14987            ));
14988            std::fs::create_dir_all(&root).unwrap();
14989            Self { root }
14990        }
14991    }
14992
14993    impl Drop for MutationTestWorkspace {
14994        fn drop(&mut self) {
14995            let _ = std::fs::remove_dir_all(&self.root);
14996        }
14997    }
14998
14999    async fn wait_for_resource_lock_strong_count(locks: &ToolResourceLocks, minimum: usize) {
15000        tokio::time::timeout(std::time::Duration::from_secs(2), async {
15001            loop {
15002                let strong_count = locks
15003                    .read()
15004                    .get("path-mutation:global")
15005                    .map_or(0, |lock| lock.strong_count());
15006                if strong_count >= minimum {
15007                    break;
15008                }
15009                tokio::task::yield_now().await;
15010            }
15011        })
15012        .await
15013        .expect("path mutation call did not reach the shared lock");
15014    }
15015
15016    async fn assert_path_mutation_pair_serialized(
15017        first_id: &'static str,
15018        first_fields: Vec<ai_agents_core::PathPolicyBinding>,
15019        first_args: Value,
15020        second_id: &'static str,
15021        second_fields: Vec<ai_agents_core::PathPolicyBinding>,
15022        second_args: Value,
15023    ) {
15024        let locks = new_tool_resource_locks();
15025        let first_gate = PathMutationGate::new();
15026        let second_gate = PathMutationGate::new();
15027        second_gate.release();
15028        let agent = Arc::new(
15029            AgentBuilder::new()
15030                .system_prompt("Test global path mutation locking.")
15031                .llm(Arc::new(mock_with_response("done")))
15032                .tool(Arc::new(BlockingPathMutationTool {
15033                    id: first_id,
15034                    path_fields: first_fields,
15035                    gate: first_gate.clone(),
15036                }))
15037                .tool(Arc::new(BlockingPathMutationTool {
15038                    id: second_id,
15039                    path_fields: second_fields,
15040                    gate: second_gate.clone(),
15041                }))
15042                .build()
15043                .unwrap()
15044                .with_shared_resource_locks(Arc::clone(&locks)),
15045        );
15046
15047        let first = {
15048            let agent = Arc::clone(&agent);
15049            tokio::spawn(async move {
15050                agent
15051                    .invoke_tool(ToolExecutionRequest::new(
15052                        format!("{}-first", first_id),
15053                        first_id,
15054                        first_args,
15055                        ToolCallSource::Manual,
15056                    ))
15057                    .await
15058                    .unwrap()
15059            })
15060        };
15061        first_gate.wait_until_entered().await;
15062
15063        let second = {
15064            let agent = Arc::clone(&agent);
15065            tokio::spawn(async move {
15066                agent
15067                    .invoke_tool(ToolExecutionRequest::new(
15068                        format!("{}-second", second_id),
15069                        second_id,
15070                        second_args,
15071                        ToolCallSource::Manual,
15072                    ))
15073                    .await
15074                    .unwrap()
15075            })
15076        };
15077        wait_for_resource_lock_strong_count(&locks, 2).await;
15078        assert!(!second_gate.entered.load(Ordering::SeqCst));
15079        assert!(!second.is_finished());
15080
15081        first_gate.release();
15082        let (first, second) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
15083            tokio::join!(first, second)
15084        })
15085        .await
15086        .expect("serialized path mutation calls did not finish");
15087        assert!(first.unwrap().success);
15088        assert!(second.unwrap().success);
15089        assert!(second_gate.entered.load(Ordering::SeqCst));
15090        assert!(locks.read().is_empty());
15091    }
15092
15093    #[derive(Clone, Copy)]
15094    enum MutationDenial {
15095        Policy,
15096        Approval,
15097    }
15098
15099    fn mutation_denial_security_config(
15100        tool_id: &str,
15101        workspace: &std::path::Path,
15102        denial: MutationDenial,
15103    ) -> ToolSecurityConfig {
15104        let workspace = workspace.to_string_lossy().into_owned();
15105        let mut policy = ai_agents_tools::ToolPolicyConfig {
15106            read_paths: vec![workspace.clone()],
15107            write_paths: vec![workspace.clone()],
15108            ..Default::default()
15109        };
15110        match denial {
15111            MutationDenial::Policy => policy.blocked_paths = vec![workspace],
15112            MutationDenial::Approval => policy.require_confirmation = true,
15113        }
15114
15115        let mut security = ToolSecurityConfig {
15116            enabled: true,
15117            fail_closed: true,
15118            ..Default::default()
15119        };
15120        security.tools.insert(tool_id.to_string(), policy);
15121        security
15122    }
15123
15124    async fn assert_path_mutation_denied(tool: Arc<dyn Tool>, denial: MutationDenial) {
15125        let workspace = MutationTestWorkspace::new();
15126        let tool_id = tool.id().to_string();
15127        let preserved = workspace.root.join(format!("{}-preserved.txt", tool_id));
15128        let destination = workspace.root.join(format!("{}-destination.txt", tool_id));
15129        std::fs::write(&preserved, "preserved").unwrap();
15130        let arguments = match tool_id.as_str() {
15131            "copy_path" | "move_path" => serde_json::json!({
15132                "source_path": preserved.to_string_lossy(),
15133                "destination_path": destination.to_string_lossy(),
15134                "dry_run": false
15135            }),
15136            "delete_path" => serde_json::json!({
15137                "path": preserved.to_string_lossy(),
15138                "recursive": false,
15139                "dry_run": false
15140            }),
15141            _ => panic!("unsupported mutation tool: {}", tool_id),
15142        };
15143        let security = mutation_denial_security_config(&tool_id, &workspace.root, denial);
15144        let builder = AgentBuilder::new()
15145            .system_prompt("Test mutation denial.")
15146            .llm(Arc::new(mock_with_response("done")))
15147            .tool(tool)
15148            .tool_security(ToolSecurityEngine::new(security));
15149        let builder = match denial {
15150            MutationDenial::Policy => builder,
15151            MutationDenial::Approval => builder
15152                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
15153                .approval_handler(Arc::new(RejectAllHandler::new())),
15154        };
15155        let agent = builder.build().unwrap();
15156
15157        let record = agent
15158            .invoke_tool(ToolExecutionRequest::new(
15159                format!("{}-denied", tool_id),
15160                tool_id.clone(),
15161                arguments,
15162                ToolCallSource::Manual,
15163            ))
15164            .await
15165            .unwrap();
15166
15167        assert!(!record.executed, "{} must not be invoked", tool_id);
15168        assert!(!record.success);
15169        match denial {
15170            MutationDenial::Policy => {
15171                assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
15172                assert!(record.approval.as_ref().is_some_and(|approval| matches!(
15173                    &approval.status,
15174                    ToolApprovalStatus::NotRequired
15175                )));
15176            }
15177            MutationDenial::Approval => {
15178                assert_eq!(record.policy.outcome, PermissionOutcome::RequiresApproval);
15179                assert!(record.approval.as_ref().is_some_and(|approval| matches!(
15180                    &approval.status,
15181                    ToolApprovalStatus::Rejected
15182                )));
15183            }
15184        }
15185        assert_eq!(std::fs::read_to_string(&preserved).unwrap(), "preserved");
15186        assert!(!destination.exists());
15187    }
15188
15189    fn recovery_manager_with_fallbacks(
15190        fallbacks: impl IntoIterator<Item = (String, String)>,
15191    ) -> RecoveryManager {
15192        use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
15193
15194        let per_tool = fallbacks
15195            .into_iter()
15196            .map(|(tool, fallback_tool)| {
15197                (
15198                    tool,
15199                    ToolRetryConfig {
15200                        max_retries: 0,
15201                        timeout_ms: Some(1_000),
15202                        on_failure: ToolFailureAction::Fallback { fallback_tool },
15203                    },
15204                )
15205            })
15206            .collect();
15207        RecoveryManager::new(ErrorRecoveryConfig {
15208            tools: ToolRecoveryConfig {
15209                per_tool,
15210                ..Default::default()
15211            },
15212            ..Default::default()
15213        })
15214    }
15215
15216    fn approval_check() -> HITLCheckResult {
15217        HITLCheckResult::required(
15218            ApprovalTrigger::tool("test", serde_json::json!({})),
15219            HashMap::new(),
15220            "Approve?",
15221            None,
15222        )
15223    }
15224
15225    fn agent_with_approval_result(
15226        raw_result: ApprovalResult,
15227        timeout_action: TimeoutAction,
15228        hooks: Arc<ApprovalRecordingHooks>,
15229    ) -> RuntimeAgent {
15230        use ai_agents_hitl::{CallbackHandler, HITLConfig};
15231
15232        let config = HITLConfig {
15233            on_timeout: timeout_action,
15234            ..Default::default()
15235        };
15236        let handler = CallbackHandler::new(move |_| raw_result.clone());
15237        AgentBuilder::new()
15238            .system_prompt("Test HITL hooks.")
15239            .llm(Arc::new(mock_with_response("done")))
15240            .build()
15241            .unwrap()
15242            .with_hooks(hooks)
15243            .with_hitl(HITLEngine::new(config), Arc::new(handler))
15244    }
15245
15246    #[tokio::test]
15247    async fn approval_hooks_expose_direct_effective_decisions_after_raw_results() {
15248        let cases = vec![
15249            (ApprovalResult::Approved, "approved"),
15250            (
15251                ApprovalResult::Rejected {
15252                    reason: Some("denied".to_string()),
15253                },
15254                "rejected",
15255            ),
15256            (
15257                ApprovalResult::Modified {
15258                    changes: HashMap::from([("value".to_string(), serde_json::json!(2))]),
15259                },
15260                "modified",
15261            ),
15262        ];
15263
15264        for (raw_result, expected) in cases {
15265            let hooks = Arc::new(ApprovalRecordingHooks::new());
15266            let agent =
15267                agent_with_approval_result(raw_result, TimeoutAction::Reject, hooks.clone());
15268
15269            let result = agent.request_hitl_approval(approval_check()).await.unwrap();
15270
15271            assert_eq!(approval_result_name(&result), expected);
15272            assert_correlated_approval_events(&hooks.events(), expected, expected);
15273        }
15274    }
15275
15276    #[tokio::test]
15277    async fn approval_hooks_expose_timeout_policy_decisions() {
15278        for (timeout_action, expected) in [
15279            (TimeoutAction::Approve, "approved"),
15280            (TimeoutAction::Reject, "rejected"),
15281        ] {
15282            let hooks = Arc::new(ApprovalRecordingHooks::new());
15283            let agent =
15284                agent_with_approval_result(ApprovalResult::Timeout, timeout_action, hooks.clone());
15285
15286            let result = agent.request_hitl_approval(approval_check()).await.unwrap();
15287
15288            assert_eq!(approval_result_name(&result), expected);
15289            assert_correlated_approval_events(&hooks.events(), "timeout", expected);
15290        }
15291    }
15292
15293    #[tokio::test]
15294    async fn timeout_error_fires_correlated_resolved_error_before_returning() {
15295        let hooks = Arc::new(ApprovalRecordingHooks::new());
15296        let agent = agent_with_approval_result(
15297            ApprovalResult::Timeout,
15298            TimeoutAction::Error,
15299            hooks.clone(),
15300        );
15301
15302        let error = agent
15303            .request_hitl_approval(approval_check())
15304            .await
15305            .unwrap_err();
15306
15307        assert!(error.to_string().contains("HITL approval timeout"));
15308        assert_correlated_approval_events(&hooks.events(), "timeout", "error");
15309    }
15310
15311    // Basic YAML → Build → Chat flow
15312    #[tokio::test]
15313    async fn test_integration_yaml_to_chat_basic() {
15314        let mock = mock_with_response("Hello! How can I help you?");
15315        let agent = AgentBuilder::new()
15316            .system_prompt("You are a test assistant.")
15317            .llm(Arc::new(mock))
15318            .build()
15319            .unwrap();
15320
15321        let response = agent.chat("Hi").await.unwrap();
15322        assert!(!response.content.is_empty());
15323        assert_eq!(response.content, "Hello! How can I help you?");
15324    }
15325
15326    // Multi-turn conversation
15327    #[tokio::test]
15328    async fn test_integration_multi_turn_conversation() {
15329        let mock = mock_with_responses(vec![
15330            "Hello! I'm your assistant.",
15331            "The weather is sunny today.",
15332            "Goodbye!",
15333        ]);
15334        let agent = AgentBuilder::new()
15335            .system_prompt("You are helpful.")
15336            .llm(Arc::new(mock))
15337            .build()
15338            .unwrap();
15339
15340        let r1 = agent.chat("Hi").await.unwrap();
15341        assert_eq!(r1.content, "Hello! I'm your assistant.");
15342
15343        let r2 = agent.chat("What's the weather?").await.unwrap();
15344        assert_eq!(r2.content, "The weather is sunny today.");
15345
15346        let r3 = agent.chat("Bye").await.unwrap();
15347        assert_eq!(r3.content, "Goodbye!");
15348
15349        // Verify memory accumulated messages
15350        let messages = agent.memory.get_messages(None).await.unwrap();
15351        // 3 user + 3 assistant = 6 messages
15352        assert_eq!(messages.len(), 6);
15353    }
15354
15355    #[test]
15356    fn later_approval_preserves_modified_evidence() {
15357        let arguments = serde_json::json!({"dry_run": true});
15358        let mut record = Some(ToolApprovalRecord {
15359            status: ToolApprovalStatus::Modified,
15360            reason: None,
15361            modified_arguments: Some(arguments.clone()),
15362        });
15363
15364        merge_approved_record(&mut record);
15365
15366        let record = record.unwrap();
15367        assert!(matches!(record.status, ToolApprovalStatus::Modified));
15368        assert_eq!(record.modified_arguments, Some(arguments));
15369    }
15370
15371    #[test]
15372    fn approval_binding_rejects_replaced_tool_implementation() {
15373        let reviewed_tool: Arc<dyn ai_agents_core::Tool> = Arc::new(ContextEchoTool);
15374        let same_tool = Arc::clone(&reviewed_tool);
15375        let replacement_tool: Arc<dyn ai_agents_core::Tool> = Arc::new(ContextEchoTool);
15376        let arguments = serde_json::json!({"path": "."});
15377        let versions = ToolDecisionVersions {
15378            policy: 2,
15379            registry: 3,
15380            runtime_control: 4,
15381            state: Some(5),
15382        };
15383        let binding = ToolApprovalBinding {
15384            canonical_id: "context_echo".to_string(),
15385            arguments: arguments.clone(),
15386            confirmation_required: true,
15387            policy_version: versions.policy,
15388            runtime_control_version: versions.runtime_control,
15389            state_generation: versions.state,
15390            reviewed_tool,
15391        };
15392
15393        assert!(!binding.is_stale("context_echo", &arguments, true, versions, &same_tool,));
15394        assert!(binding.is_stale(
15395            "context_echo",
15396            &arguments,
15397            true,
15398            versions,
15399            &replacement_tool,
15400        ));
15401    }
15402
15403    #[tokio::test]
15404    async fn approved_mutation_to_dry_run_remains_executable() {
15405        use ai_agents_hitl::CallbackHandler;
15406
15407        let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
15408            changes: HashMap::from([("dry_run".to_string(), serde_json::json!(true))]),
15409        });
15410        let agent = AgentBuilder::new()
15411            .system_prompt("Test safer approval modifications.")
15412            .llm(Arc::new(mock_with_response("done")))
15413            .tool(Arc::new(ai_agents_tools::FileWriteTool::new()))
15414            .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
15415            .approval_handler(Arc::new(handler))
15416            .build()
15417            .unwrap();
15418
15419        let record = agent
15420            .invoke_tool(ToolExecutionRequest::new(
15421                "approved-dry-run",
15422                "file_write",
15423                serde_json::json!({
15424                    "path": "./approval-dry-run.txt",
15425                    "content": "not written"
15426                }),
15427                ToolCallSource::Manual,
15428            ))
15429            .await
15430            .unwrap();
15431
15432        assert!(record.executed);
15433        assert!(record.success);
15434        assert_eq!(record.executed_arguments["dry_run"], true);
15435        assert!(matches!(
15436            record.approval.as_ref().map(|approval| &approval.status),
15437            Some(ToolApprovalStatus::Modified)
15438        ));
15439        let output: Value = serde_json::from_str(&record.output).unwrap();
15440        assert_eq!(output["mutation_performed"], false);
15441    }
15442
15443    /// Proves shared HITL approval evidence reaches the web fetch implementation unchanged.
15444    #[tokio::test]
15445    async fn shared_executor_approval_reaches_web_fetch_transport() {
15446        use ai_agents_hitl::{CallbackHandler, HITLConfig};
15447        use ai_agents_tools::{DomainPolicyConfig, ToolPolicyConfig};
15448
15449        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15450        let tool = WebFetchTool::with_transport_and_resolver(
15451            Arc::new(RuntimeWebFetchTransport {
15452                calls: Arc::clone(&calls),
15453            }),
15454            Arc::new(RuntimeWebFetchResolver),
15455        );
15456        let mut security = ToolSecurityConfig {
15457            enabled: true,
15458            fail_closed: true,
15459            ..Default::default()
15460        };
15461        security.tools.insert(
15462            "web_fetch".to_string(),
15463            ToolPolicyConfig {
15464                domains: DomainPolicyConfig {
15465                    requires_approval: vec!["approval.test".to_string()],
15466                    ..Default::default()
15467                },
15468                allowed_schemes: vec!["https".to_string()],
15469                allowed_ports: vec![443],
15470                ..Default::default()
15471            },
15472        );
15473        let handler = CallbackHandler::new(|_| ApprovalResult::Approved);
15474        let agent = AgentBuilder::new()
15475            .system_prompt("Test approved web fetch execution.")
15476            .llm(Arc::new(mock_with_response("done")))
15477            .tool(Arc::new(tool))
15478            .tool_security(ToolSecurityEngine::new(security))
15479            .build()
15480            .unwrap()
15481            .with_hitl(HITLEngine::new(HITLConfig::default()), Arc::new(handler));
15482
15483        let record = agent
15484            .invoke_tool(ToolExecutionRequest::new(
15485                "approved-web-fetch",
15486                "web_fetch",
15487                serde_json::json!({
15488                    "url": "https://approval.test/page",
15489                    "cache_ttl_seconds": 0
15490                }),
15491                ToolCallSource::Manual,
15492            ))
15493            .await
15494            .unwrap();
15495
15496        assert!(record.success);
15497        assert!(
15498            record
15499                .approval
15500                .as_ref()
15501                .is_some_and(|approval| matches!(approval.status, ToolApprovalStatus::Approved))
15502        );
15503        assert_eq!(calls.load(Ordering::SeqCst), 1);
15504    }
15505
15506    #[tokio::test]
15507    async fn context_preserves_requested_and_canonical_identity() {
15508        let mock = mock_with_response("hello");
15509        let mut tools = ai_agents_tools::ToolRegistry::new();
15510        tools.register(Arc::new(ContextEchoTool)).unwrap();
15511
15512        let mut security = ToolSecurityConfig {
15513            enabled: true,
15514            fail_closed: true,
15515            ..Default::default()
15516        };
15517        let mut policy = ai_agents_tools::ToolPolicyConfig {
15518            read_paths: vec![".".to_string()],
15519            max_results: Some(7),
15520            ..Default::default()
15521        };
15522        policy
15523            .config
15524            .insert("backend".to_string(), serde_json::json!("memory"));
15525        security.tools.insert("context_echo".to_string(), policy);
15526
15527        let agent = AgentBuilder::new()
15528            .system_prompt("You are helpful.")
15529            .llm(Arc::new(mock))
15530            .tools(tools)
15531            .tool_security(ToolSecurityEngine::new(security))
15532            .build()
15533            .unwrap();
15534
15535        let record = agent
15536            .invoke_tool(ToolExecutionRequest::new(
15537                "ctx-call",
15538                "Context Echo",
15539                serde_json::json!({"path": ".", "max_results": 99}),
15540                ToolCallSource::Manual,
15541            ))
15542            .await
15543            .unwrap();
15544
15545        assert!(record.success);
15546        assert!(matches!(&record.source, ToolCallSource::Manual));
15547        assert_eq!(record.requested_name, "Context Echo");
15548        assert_eq!(record.canonical_id, "context_echo");
15549        assert_eq!(record.policy.outcome, PermissionOutcome::Allow);
15550        assert_eq!(record.executed_arguments["max_results"], 7);
15551        let output: Value = serde_json::from_str(&record.output).unwrap();
15552        assert_eq!(output["requested_name"], "Context Echo");
15553        assert_eq!(output["canonical_id"], "context_echo");
15554        assert_eq!(output["max_results"], 7);
15555        assert_eq!(output["custom_config"]["backend"], "memory");
15556        assert!(record.metadata.contains_key("effective_limits"));
15557        assert!(record.metadata.contains_key("policy_snapshot"));
15558    }
15559
15560    #[tokio::test]
15561    async fn test_runtime_control_cancels_active_tool_call() {
15562        let mock = mock_with_response("hello");
15563        let agent = Arc::new(
15564            AgentBuilder::new()
15565                .system_prompt("You are helpful.")
15566                .llm(Arc::new(mock))
15567                .tool(Arc::new(SlowTool))
15568                .build()
15569                .unwrap(),
15570        );
15571        let control = agent.runtime_control();
15572        let running_agent = Arc::clone(&agent);
15573        let handle = tokio::spawn(async move {
15574            running_agent
15575                .invoke_tool(ToolExecutionRequest::new(
15576                    "slow-call",
15577                    "slow",
15578                    serde_json::json!({}),
15579                    ToolCallSource::Manual,
15580                ))
15581                .await
15582                .unwrap()
15583        });
15584
15585        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
15586        control.cancel_all();
15587        let record = handle.await.unwrap();
15588
15589        assert!(record.executed);
15590        assert!(record.cancelled);
15591        assert!(!record.success);
15592        assert!(record.cancellation_reason.is_some());
15593    }
15594
15595    #[tokio::test]
15596    async fn non_idempotent_tool_calls_are_not_retried() {
15597        use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
15598
15599        let mock = mock_with_response("hello");
15600        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15601        let agent = AgentBuilder::new()
15602            .system_prompt("You are helpful.")
15603            .llm(Arc::new(mock))
15604            .tool(Arc::new(FlakyWriteTool {
15605                calls: Arc::clone(&calls),
15606            }))
15607            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
15608                tools: ToolRecoveryConfig {
15609                    default: ToolRetryConfig {
15610                        max_retries: 2,
15611                        ..Default::default()
15612                    },
15613                    ..Default::default()
15614                },
15615                ..Default::default()
15616            }))
15617            .build()
15618            .unwrap();
15619
15620        let record = agent
15621            .invoke_tool(ToolExecutionRequest::new(
15622                "flaky-call",
15623                "flaky_write",
15624                serde_json::json!({"path": "./tmp.txt"}),
15625                ToolCallSource::Manual,
15626            ))
15627            .await
15628            .unwrap();
15629
15630        assert!(!record.success);
15631        assert_eq!(calls.load(Ordering::SeqCst), 1);
15632    }
15633
15634    #[tokio::test]
15635    async fn side_effecting_tools_are_serialized_per_resource() {
15636        let mock = mock_with_response("hello");
15637        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15638        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15639        let agent = Arc::new(
15640            AgentBuilder::new()
15641                .system_prompt("You are helpful.")
15642                .llm(Arc::new(mock))
15643                .tool(Arc::new(LockedWriteTool {
15644                    active: Arc::clone(&active),
15645                    max_active: Arc::clone(&max_active),
15646                }))
15647                .build()
15648                .unwrap(),
15649        );
15650
15651        let left = {
15652            let agent = Arc::clone(&agent);
15653            tokio::spawn(async move {
15654                agent
15655                    .invoke_tool(ToolExecutionRequest::new(
15656                        "lock-1",
15657                        "locked_write",
15658                        serde_json::json!({"path": "./same.txt"}),
15659                        ToolCallSource::Manual,
15660                    ))
15661                    .await
15662                    .unwrap()
15663            })
15664        };
15665        let right = {
15666            let agent = Arc::clone(&agent);
15667            tokio::spawn(async move {
15668                agent
15669                    .invoke_tool(ToolExecutionRequest::new(
15670                        "lock-2",
15671                        "locked_write",
15672                        serde_json::json!({"path": "./same.txt"}),
15673                        ToolCallSource::Manual,
15674                    ))
15675                    .await
15676                    .unwrap()
15677            })
15678        };
15679
15680        let left = left.await.unwrap();
15681        let right = right.await.unwrap();
15682        assert!(left.success);
15683        assert!(right.success);
15684        assert_eq!(max_active.load(Ordering::SeqCst), 1);
15685    }
15686
15687    #[tokio::test]
15688    async fn path_resources_use_shared_global_lock_and_cleanup() {
15689        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15690        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15691        let bindings = ai_agents_core::ToolPolicyBindings {
15692            path_fields: vec![
15693                ai_agents_core::PathPolicyBinding::read_write("source_path"),
15694                ai_agents_core::PathPolicyBinding::write("destination_path"),
15695            ],
15696            ..Default::default()
15697        };
15698        let classification = ai_agents_core::ToolCallClassification::from_metadata(
15699            &MultiResourceWriteTool {
15700                active: Arc::clone(&active),
15701                max_active: Arc::clone(&max_active),
15702            }
15703            .safety_metadata(),
15704        );
15705        let left_args = serde_json::json!({
15706            "source_path": "./a/../first.txt",
15707            "destination_path": "./second.txt"
15708        });
15709        let right_args = serde_json::json!({
15710            "source_path": "./second.txt",
15711            "destination_path": "./first.txt"
15712        });
15713        let left_keys = tool_resource_lock_keys(
15714            "multi_resource_write",
15715            &left_args,
15716            &bindings,
15717            &classification,
15718        );
15719        let right_keys = tool_resource_lock_keys(
15720            "multi_resource_write",
15721            &right_args,
15722            &bindings,
15723            &classification,
15724        );
15725        assert_eq!(left_keys, right_keys);
15726        assert_eq!(left_keys, vec!["path-mutation:global".to_string()]);
15727
15728        let locks = new_tool_resource_locks();
15729        let build_agent = || {
15730            AgentBuilder::new()
15731                .system_prompt("Test shared resource locks.")
15732                .llm(Arc::new(mock_with_response("done")))
15733                .tool(Arc::new(MultiResourceWriteTool {
15734                    active: Arc::clone(&active),
15735                    max_active: Arc::clone(&max_active),
15736                }))
15737                .build()
15738                .unwrap()
15739                .with_shared_resource_locks(Arc::clone(&locks))
15740        };
15741        let left_agent = Arc::new(build_agent());
15742        let right_agent = Arc::new(build_agent());
15743        let left = tokio::spawn(async move {
15744            left_agent
15745                .invoke_tool(ToolExecutionRequest::new(
15746                    "multi-left",
15747                    "multi_resource_write",
15748                    left_args,
15749                    ToolCallSource::Manual,
15750                ))
15751                .await
15752                .unwrap()
15753        });
15754        let right = tokio::spawn(async move {
15755            right_agent
15756                .invoke_tool(ToolExecutionRequest::new(
15757                    "multi-right",
15758                    "multi_resource_write",
15759                    right_args,
15760                    ToolCallSource::Manual,
15761                ))
15762                .await
15763                .unwrap()
15764        });
15765        let (left, right) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
15766            tokio::join!(left, right)
15767        })
15768        .await
15769        .expect("reversed resource acquisition must not deadlock");
15770
15771        assert!(left.unwrap().success);
15772        assert!(right.unwrap().success);
15773        assert_eq!(max_active.load(Ordering::SeqCst), 1);
15774        assert!(locks.read().is_empty());
15775    }
15776
15777    #[tokio::test]
15778    async fn global_path_lock_serializes_copy_destination_with_file_write() {
15779        assert_path_mutation_pair_serialized(
15780            "copy_path",
15781            CopyPathTool::new().policy_bindings().path_fields,
15782            serde_json::json!({
15783                "source_path": "./source.txt",
15784                "destination_path": "./shared.txt"
15785            }),
15786            "file_write",
15787            FileWriteTool::new().policy_bindings().path_fields,
15788            serde_json::json!({"path": "./shared.txt"}),
15789        )
15790        .await;
15791    }
15792
15793    #[tokio::test]
15794    async fn parent_and_spawned_runtime_share_global_path_lock() {
15795        let workspace = MutationTestWorkspace::new();
15796        let destination = workspace.root.join("spawned.txt");
15797        let parent_gate = PathMutationGate::new();
15798        let parent = Arc::new(
15799            AgentBuilder::from_yaml(
15800                r#"
15801name: LockParent
15802system_prompt: parent
15803llm:
15804  default: default
15805tools:
15806  - parent_path_write
15807spawner:
15808  shared_llms: true
15809"#,
15810            )
15811            .unwrap()
15812            .llm(Arc::new(mock_with_response("done")))
15813            .auto_configure_spawner()
15814            .await
15815            .unwrap()
15816            .tool(Arc::new(BlockingPathMutationTool {
15817                id: "parent_path_write",
15818                path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15819                gate: parent_gate.clone(),
15820            }))
15821            .build()
15822            .unwrap(),
15823        );
15824
15825        let mut child_spec = crate::spec::AgentSpec {
15826            name: "LockChild".to_string(),
15827            system_prompt: "child".to_string(),
15828            tools: Some(vec![crate::spec::ToolEntry::Simple(
15829                "file_write".to_string(),
15830            )]),
15831            ..Default::default()
15832        };
15833        child_spec.tool_security.enabled = true;
15834        child_spec.tool_security.fail_closed = true;
15835        let file_write_policy = ai_agents_tools::ToolPolicyConfig {
15836            write_paths: vec![workspace.root.to_string_lossy().into_owned()],
15837            allow_without_confirmation: true,
15838            ..Default::default()
15839        };
15840        child_spec
15841            .tool_security
15842            .tools
15843            .insert("file_write".to_string(), file_write_policy);
15844        let spawned = parent
15845            .spawner()
15846            .unwrap()
15847            .spawn_from_spec(child_spec)
15848            .await
15849            .unwrap();
15850        assert!(Arc::ptr_eq(
15851            &parent.resource_locks,
15852            &spawned.agent.resource_locks
15853        ));
15854        assert!(!Arc::ptr_eq(
15855            &parent.runtime_control,
15856            &spawned.agent.runtime_control
15857        ));
15858
15859        let parent_call = {
15860            let parent = Arc::clone(&parent);
15861            let destination = destination.clone();
15862            tokio::spawn(async move {
15863                parent
15864                    .invoke_tool(ToolExecutionRequest::new(
15865                        "parent-lock-holder",
15866                        "parent_path_write",
15867                        serde_json::json!({"path": destination}),
15868                        ToolCallSource::Manual,
15869                    ))
15870                    .await
15871                    .unwrap()
15872            })
15873        };
15874        parent_gate.wait_until_entered().await;
15875
15876        let child_call = {
15877            let child = Arc::clone(&spawned.agent);
15878            let destination = destination.clone();
15879            tokio::spawn(async move {
15880                child
15881                    .invoke_tool(ToolExecutionRequest::new(
15882                        "spawned-file-write",
15883                        "file_write",
15884                        serde_json::json!({
15885                            "path": destination,
15886                            "content": "spawned",
15887                            "dry_run": false
15888                        }),
15889                        ToolCallSource::Manual,
15890                    ))
15891                    .await
15892                    .unwrap()
15893            })
15894        };
15895        wait_for_resource_lock_strong_count(&parent.resource_locks, 2).await;
15896        assert!(!child_call.is_finished());
15897
15898        parent_gate.release();
15899        let (parent_record, child_record) =
15900            tokio::time::timeout(std::time::Duration::from_secs(2), async {
15901                tokio::join!(parent_call, child_call)
15902            })
15903            .await
15904            .expect("parent and spawned path mutations did not finish");
15905        assert!(parent_record.unwrap().success);
15906        assert!(child_record.unwrap().success);
15907        assert_eq!(std::fs::read_to_string(destination).unwrap(), "spawned");
15908        assert!(parent.resource_locks.read().is_empty());
15909    }
15910
15911    #[tokio::test]
15912    async fn cancelled_global_path_lock_waiter_does_not_retain_weak_entry() {
15913        let locks = new_tool_resource_locks();
15914        let holder_gate = PathMutationGate::new();
15915        let waiter_gate = PathMutationGate::new();
15916        waiter_gate.release();
15917        let holder = Arc::new(
15918            AgentBuilder::new()
15919                .system_prompt("Hold the global path lock.")
15920                .llm(Arc::new(mock_with_response("done")))
15921                .tool(Arc::new(BlockingPathMutationTool {
15922                    id: "holder_write",
15923                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15924                    gate: holder_gate.clone(),
15925                }))
15926                .build()
15927                .unwrap()
15928                .with_shared_resource_locks(Arc::clone(&locks)),
15929        );
15930        let waiter = Arc::new(
15931            AgentBuilder::new()
15932                .system_prompt("Wait for the global path lock.")
15933                .llm(Arc::new(mock_with_response("done")))
15934                .tool(Arc::new(BlockingPathMutationTool {
15935                    id: "waiter_write",
15936                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15937                    gate: waiter_gate.clone(),
15938                }))
15939                .build()
15940                .unwrap()
15941                .with_shared_resource_locks(Arc::clone(&locks)),
15942        );
15943
15944        let holder_call = {
15945            let holder = Arc::clone(&holder);
15946            tokio::spawn(async move {
15947                holder
15948                    .invoke_tool(ToolExecutionRequest::new(
15949                        "holder-call",
15950                        "holder_write",
15951                        serde_json::json!({"path": "./shared.txt"}),
15952                        ToolCallSource::Manual,
15953                    ))
15954                    .await
15955                    .unwrap()
15956            })
15957        };
15958        holder_gate.wait_until_entered().await;
15959
15960        let waiter_call = {
15961            let waiter = Arc::clone(&waiter);
15962            tokio::spawn(async move {
15963                waiter
15964                    .invoke_tool(ToolExecutionRequest::new(
15965                        "waiter-call",
15966                        "waiter_write",
15967                        serde_json::json!({"path": "./shared.txt"}),
15968                        ToolCallSource::Manual,
15969                    ))
15970                    .await
15971                    .unwrap()
15972            })
15973        };
15974        wait_for_resource_lock_strong_count(&locks, 2).await;
15975        waiter.runtime_control().cancel_all();
15976
15977        let waiter_record = tokio::time::timeout(std::time::Duration::from_secs(2), waiter_call)
15978            .await
15979            .expect("cancelled lock waiter did not finish")
15980            .unwrap();
15981        assert!(!waiter_record.executed);
15982        assert!(!waiter_gate.entered.load(Ordering::SeqCst));
15983        assert_eq!(
15984            locks
15985                .read()
15986                .get("path-mutation:global")
15987                .map_or(0, |lock| lock.strong_count()),
15988            1
15989        );
15990
15991        holder_gate.release();
15992        let holder_record = tokio::time::timeout(std::time::Duration::from_secs(2), holder_call)
15993            .await
15994            .expect("lock holder did not finish")
15995            .unwrap();
15996        assert!(holder_record.success);
15997        assert!(locks.read().is_empty());
15998    }
15999
16000    #[tokio::test]
16001    async fn path_mutation_policy_and_approval_denials_do_not_invoke_tools() {
16002        for denial in [MutationDenial::Policy, MutationDenial::Approval] {
16003            let tools: [Arc<dyn Tool>; 3] = [
16004                Arc::new(CopyPathTool::new()),
16005                Arc::new(MovePathTool::new()),
16006                Arc::new(DeletePathTool::new()),
16007            ];
16008            for tool in tools {
16009                assert_path_mutation_denied(tool, denial).await;
16010            }
16011        }
16012    }
16013
16014    #[tokio::test]
16015    async fn approval_argument_changes_are_rechecked_against_final_scope() {
16016        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16017        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16018        let entered = Arc::new(tokio::sync::Barrier::new(2));
16019        let release = Arc::new(tokio::sync::Notify::new());
16020        let handler = Arc::new(BlockingApprovalHandler {
16021            entered: Arc::clone(&entered),
16022            release: Arc::clone(&release),
16023            result: ApprovalResult::Modified {
16024                changes: HashMap::from([(
16025                    "path".to_string(),
16026                    Value::String("./after-approval.txt".to_string()),
16027                )]),
16028            },
16029        });
16030        let agent = Arc::new(
16031            AgentBuilder::new()
16032                .system_prompt("Test final scope validation.")
16033                .llm(Arc::new(mock_with_response("done")))
16034                .tool(Arc::new(LockedWriteTool {
16035                    active: Arc::clone(&active),
16036                    max_active: Arc::clone(&max_active),
16037                }))
16038                .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
16039                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16040                .approval_handler(handler)
16041                .build()
16042                .unwrap(),
16043        );
16044        let control = agent.runtime_control();
16045        let running = Arc::clone(&agent);
16046        let call = tokio::spawn(async move {
16047            running
16048                .invoke_tool(ToolExecutionRequest::new(
16049                    "approval-scope",
16050                    "locked_write",
16051                    serde_json::json!({"path": "./before-approval.txt"}),
16052                    ToolCallSource::Manual,
16053                ))
16054                .await
16055                .unwrap()
16056        });
16057        entered.wait().await;
16058        let expected_version = control.set_tool_scope(Vec::new());
16059        release.notify_one();
16060        let record = call.await.unwrap();
16061
16062        assert!(!record.executed);
16063        assert!(!record.success);
16064        assert_eq!(record.runtime_config_version, expected_version);
16065        assert_eq!(record.executed_arguments["path"], "./after-approval.txt");
16066        assert_eq!(max_active.load(Ordering::SeqCst), 0);
16067        assert_eq!(
16068            record.metadata["runtime_scope_snapshot"],
16069            serde_json::json!([])
16070        );
16071    }
16072
16073    #[tokio::test]
16074    async fn approval_is_rechecked_against_final_policy_snapshot() {
16075        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16076        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16077        let entered = Arc::new(tokio::sync::Barrier::new(2));
16078        let release = Arc::new(tokio::sync::Notify::new());
16079        let handler = Arc::new(BlockingApprovalHandler {
16080            entered: Arc::clone(&entered),
16081            release: Arc::clone(&release),
16082            result: ApprovalResult::Approved,
16083        });
16084        let agent = Arc::new(
16085            AgentBuilder::new()
16086                .system_prompt("Test final policy validation.")
16087                .llm(Arc::new(mock_with_response("done")))
16088                .tool(Arc::new(LockedWriteTool {
16089                    active: Arc::clone(&active),
16090                    max_active: Arc::clone(&max_active),
16091                }))
16092                .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
16093                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16094                .approval_handler(handler)
16095                .build()
16096                .unwrap(),
16097        );
16098        let control = agent.runtime_control();
16099        let running = Arc::clone(&agent);
16100        let call = tokio::spawn(async move {
16101            running
16102                .invoke_tool(ToolExecutionRequest::new(
16103                    "approval-policy",
16104                    "locked_write",
16105                    serde_json::json!({"path": "./policy.txt"}),
16106                    ToolCallSource::Manual,
16107                ))
16108                .await
16109                .unwrap()
16110        });
16111        entered.wait().await;
16112        let expected_version = control.set_tool_security(approval_security_config(false));
16113        release.notify_one();
16114        let record = call.await.unwrap();
16115
16116        assert!(!record.executed);
16117        assert!(!record.success);
16118        assert_eq!(record.runtime_config_version, expected_version);
16119        assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
16120        assert_eq!(max_active.load(Ordering::SeqCst), 0);
16121        assert!(record.metadata.contains_key("policy_snapshot"));
16122    }
16123
16124    #[test]
16125    fn invalid_live_policy_does_not_replace_snapshot_or_generation() {
16126        let agent = AgentBuilder::new()
16127            .system_prompt("Test runtime policy validation.")
16128            .llm(Arc::new(mock_with_response("done")))
16129            .build()
16130            .unwrap();
16131        let control = agent.runtime_control();
16132        let mut valid = ToolSecurityConfig::default();
16133        valid.tools.insert(
16134            "web_search".to_string(),
16135            ai_agents_tools::ToolPolicyConfig {
16136                max_results: Some(5),
16137                ..Default::default()
16138            },
16139        );
16140        let generation = control.try_set_tool_security(valid).unwrap();
16141
16142        let mut invalid = ToolSecurityConfig::default();
16143        invalid.tools.insert(
16144            "web_search".to_string(),
16145            ai_agents_tools::ToolPolicyConfig {
16146                max_results: Some(0),
16147                ..Default::default()
16148            },
16149        );
16150        let error = control.try_set_tool_security(invalid).unwrap_err();
16151
16152        assert!(
16153            error
16154                .to_string()
16155                .contains("max_results must be greater than 0")
16156        );
16157        assert_eq!(control.version(), generation);
16158        assert_eq!(
16159            control
16160                .state
16161                .tool_security_override
16162                .read()
16163                .as_ref()
16164                .unwrap()
16165                .config()
16166                .tools["web_search"]
16167                .max_results,
16168            Some(5)
16169        );
16170    }
16171
16172    #[tokio::test]
16173    async fn persistent_override_preserves_rate_history_within_generation() {
16174        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16175        let agent = AgentBuilder::new()
16176            .system_prompt("Test persistent policy overrides.")
16177            .llm(Arc::new(mock_with_response("done")))
16178            .tool(Arc::new(RecoveryTestTool {
16179                id: "limited_override".to_string(),
16180                succeeds: true,
16181                calls: Arc::clone(&calls),
16182            }))
16183            .build()
16184            .unwrap();
16185        let mut security = ToolSecurityConfig {
16186            enabled: true,
16187            fail_closed: true,
16188            ..Default::default()
16189        };
16190        let policy = ai_agents_tools::ToolPolicyConfig {
16191            write_paths: vec![".".to_string()],
16192            rate_limit: Some(1),
16193            ..Default::default()
16194        };
16195        security
16196            .tools
16197            .insert("limited_override".to_string(), policy);
16198        let generation = agent.runtime_control().set_tool_security(security);
16199
16200        let first = agent
16201            .invoke_tool(ToolExecutionRequest::new(
16202                "limited-first",
16203                "limited_override",
16204                serde_json::json!({"path": "./limited.txt"}),
16205                ToolCallSource::Manual,
16206            ))
16207            .await
16208            .unwrap();
16209        let second = agent
16210            .invoke_tool(ToolExecutionRequest::new(
16211                "limited-second",
16212                "limited_override",
16213                serde_json::json!({"path": "./limited.txt"}),
16214                ToolCallSource::Manual,
16215            ))
16216            .await
16217            .unwrap();
16218
16219        assert!(first.success);
16220        assert_eq!(first.policy_version, generation);
16221        assert!(!second.executed);
16222        assert!(second.output.contains("Rate limit exceeded"));
16223        assert_eq!(second.policy_version, generation);
16224        assert_eq!(calls.load(Ordering::SeqCst), 1);
16225    }
16226
16227    #[tokio::test]
16228    async fn concurrent_rate_admission_consumes_capacity_atomically() {
16229        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16230        let tool = Arc::new(RecoveryTestTool {
16231            id: "atomic_rate".to_string(),
16232            succeeds: true,
16233            calls: Arc::clone(&calls),
16234        });
16235        let arguments = serde_json::json!({"path": "./atomic-rate.txt"});
16236        let bindings = tool.policy_bindings();
16237        let classification = tool.classify_call(&arguments);
16238        let resource_keys =
16239            tool_resource_lock_keys(tool.id(), &arguments, &bindings, &classification);
16240        let mut security = ToolSecurityConfig {
16241            enabled: true,
16242            fail_closed: true,
16243            ..Default::default()
16244        };
16245        let policy = ai_agents_tools::ToolPolicyConfig {
16246            write_paths: vec![".".to_string()],
16247            rate_limit: Some(1),
16248            ..Default::default()
16249        };
16250        security.tools.insert(tool.id().to_string(), policy);
16251        let agent = Arc::new(
16252            AgentBuilder::new()
16253                .system_prompt("Test atomic rate admission.")
16254                .llm(Arc::new(mock_with_response("done")))
16255                .tool(tool)
16256                .tool_security(ToolSecurityEngine::new(security))
16257                .build()
16258                .unwrap(),
16259        );
16260        let held = agent
16261            .acquire_tool_resource_locks(&resource_keys)
16262            .await
16263            .unwrap();
16264        let left = {
16265            let agent = Arc::clone(&agent);
16266            let arguments = arguments.clone();
16267            tokio::spawn(async move {
16268                agent
16269                    .invoke_tool(ToolExecutionRequest::new(
16270                        "atomic-rate-left",
16271                        "atomic_rate",
16272                        arguments,
16273                        ToolCallSource::Manual,
16274                    ))
16275                    .await
16276                    .unwrap()
16277            })
16278        };
16279        let right = {
16280            let agent = Arc::clone(&agent);
16281            tokio::spawn(async move {
16282                agent
16283                    .invoke_tool(ToolExecutionRequest::new(
16284                        "atomic-rate-right",
16285                        "atomic_rate",
16286                        arguments,
16287                        ToolCallSource::Manual,
16288                    ))
16289                    .await
16290                    .unwrap()
16291            })
16292        };
16293        tokio::time::sleep(std::time::Duration::from_millis(25)).await;
16294        drop(held);
16295        let (left, right) = tokio::join!(left, right);
16296        let records = [left.unwrap(), right.unwrap()];
16297
16298        assert_eq!(records.iter().filter(|record| record.success).count(), 1);
16299        assert_eq!(records.iter().filter(|record| record.executed).count(), 1);
16300        assert!(
16301            records.iter().any(|record| {
16302                !record.executed && record.output.contains("Rate limit exceeded")
16303            })
16304        );
16305        assert_eq!(calls.load(Ordering::SeqCst), 1);
16306    }
16307
16308    #[tokio::test]
16309    async fn changed_policy_generation_invalidates_pending_approval() {
16310        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16311        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16312        let entered = Arc::new(tokio::sync::Barrier::new(2));
16313        let release = Arc::new(tokio::sync::Notify::new());
16314        let handler = Arc::new(BlockingApprovalHandler {
16315            entered: Arc::clone(&entered),
16316            release: Arc::clone(&release),
16317            result: ApprovalResult::Approved,
16318        });
16319        let agent = Arc::new(
16320            AgentBuilder::new()
16321                .system_prompt("Test stale approval denial.")
16322                .llm(Arc::new(mock_with_response("done")))
16323                .tool(Arc::new(LockedWriteTool {
16324                    active: Arc::clone(&active),
16325                    max_active: Arc::clone(&max_active),
16326                }))
16327                .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
16328                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16329                .approval_handler(handler)
16330                .build()
16331                .unwrap(),
16332        );
16333        let running = Arc::clone(&agent);
16334        let call = tokio::spawn(async move {
16335            running
16336                .invoke_tool(ToolExecutionRequest::new(
16337                    "stale-approval",
16338                    "locked_write",
16339                    serde_json::json!({"path": "./stale.txt"}),
16340                    ToolCallSource::Manual,
16341                ))
16342                .await
16343                .unwrap()
16344        });
16345        entered.wait().await;
16346        let generation = agent
16347            .runtime_control()
16348            .set_tool_security(approval_security_config(true));
16349        release.notify_one();
16350        let record = call.await.unwrap();
16351
16352        assert!(!record.executed);
16353        assert!(record.output.contains("Approval became stale"));
16354        assert_eq!(record.policy_version, generation);
16355        assert_eq!(max_active.load(Ordering::SeqCst), 0);
16356    }
16357
16358    #[tokio::test]
16359    async fn final_policy_reapplies_argument_caps_after_approval_changes() {
16360        use ai_agents_hitl::CallbackHandler;
16361
16362        let mut security = ToolSecurityConfig {
16363            enabled: true,
16364            fail_closed: true,
16365            ..Default::default()
16366        };
16367        let policy = ai_agents_tools::ToolPolicyConfig {
16368            read_paths: vec![".".to_string()],
16369            max_results: Some(5),
16370            require_confirmation: true,
16371            ..Default::default()
16372        };
16373        security.tools.insert("context_echo".to_string(), policy);
16374        let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
16375            changes: HashMap::from([("max_results".to_string(), serde_json::json!(99))]),
16376        });
16377        let agent = AgentBuilder::new()
16378            .system_prompt("Test final argument caps.")
16379            .llm(Arc::new(mock_with_response("done")))
16380            .tool(Arc::new(ContextEchoTool))
16381            .tool_security(ToolSecurityEngine::new(security))
16382            .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16383            .approval_handler(Arc::new(handler))
16384            .build()
16385            .unwrap();
16386
16387        let record = agent
16388            .invoke_tool(ToolExecutionRequest::new(
16389                "final-cap",
16390                "context_echo",
16391                serde_json::json!({"path": ".", "max_results": 1}),
16392                ToolCallSource::Manual,
16393            ))
16394            .await
16395            .unwrap();
16396
16397        assert!(record.success);
16398        assert_eq!(record.executed_arguments["max_results"], 5);
16399        assert_eq!(
16400            record.approval.unwrap().modified_arguments.unwrap()["max_results"],
16401            5
16402        );
16403    }
16404
16405    #[tokio::test]
16406    async fn no_binding_writes_use_canonical_fallback_lock() {
16407        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16408        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16409        let agent = Arc::new(
16410            AgentBuilder::new()
16411                .system_prompt("Test fallback resource locks.")
16412                .llm(Arc::new(mock_with_response("done")))
16413                .tool(Arc::new(NoBindingWriteTool {
16414                    active: Arc::clone(&active),
16415                    max_active: Arc::clone(&max_active),
16416                }))
16417                .build()
16418                .unwrap(),
16419        );
16420        let left = {
16421            let agent = Arc::clone(&agent);
16422            tokio::spawn(async move {
16423                agent
16424                    .invoke_tool(ToolExecutionRequest::new(
16425                        "no-binding-left",
16426                        "no_binding_write",
16427                        serde_json::json!({}),
16428                        ToolCallSource::Manual,
16429                    ))
16430                    .await
16431                    .unwrap()
16432            })
16433        };
16434        let right = {
16435            let agent = Arc::clone(&agent);
16436            tokio::spawn(async move {
16437                agent
16438                    .invoke_tool(ToolExecutionRequest::new(
16439                        "no-binding-right",
16440                        "no_binding_write",
16441                        serde_json::json!({}),
16442                        ToolCallSource::Manual,
16443                    ))
16444                    .await
16445                    .unwrap()
16446            })
16447        };
16448        let (left, right) = tokio::join!(left, right);
16449
16450        assert!(left.unwrap().success);
16451        assert!(right.unwrap().success);
16452        assert_eq!(max_active.load(Ordering::SeqCst), 1);
16453    }
16454
16455    #[tokio::test]
16456    async fn parent_and_child_paths_share_a_resource_lock() {
16457        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16458        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16459        let agent = Arc::new(
16460            AgentBuilder::new()
16461                .system_prompt("Test parent child resource locks.")
16462                .llm(Arc::new(mock_with_response("done")))
16463                .tool(Arc::new(LockedWriteTool {
16464                    active: Arc::clone(&active),
16465                    max_active: Arc::clone(&max_active),
16466                }))
16467                .build()
16468                .unwrap(),
16469        );
16470        let parent = format!("./lock-parent-{}", uuid::Uuid::new_v4());
16471        let child = format!("{}/child.txt", parent);
16472        let left = {
16473            let agent = Arc::clone(&agent);
16474            tokio::spawn(async move {
16475                agent
16476                    .invoke_tool(ToolExecutionRequest::new(
16477                        "parent-lock",
16478                        "locked_write",
16479                        serde_json::json!({"path": parent}),
16480                        ToolCallSource::Manual,
16481                    ))
16482                    .await
16483                    .unwrap()
16484            })
16485        };
16486        let right = {
16487            let agent = Arc::clone(&agent);
16488            tokio::spawn(async move {
16489                agent
16490                    .invoke_tool(ToolExecutionRequest::new(
16491                        "child-lock",
16492                        "locked_write",
16493                        serde_json::json!({"path": child}),
16494                        ToolCallSource::Manual,
16495                    ))
16496                    .await
16497                    .unwrap()
16498            })
16499        };
16500        let (left, right) = tokio::join!(left, right);
16501
16502        assert!(left.unwrap().success);
16503        assert!(right.unwrap().success);
16504        assert_eq!(max_active.load(Ordering::SeqCst), 1);
16505    }
16506
16507    #[tokio::test]
16508    async fn tool_hooks_can_reenter_after_resource_guards_are_dropped() {
16509        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16510        let hooks = Arc::new(ReentrantToolHooks {
16511            agent: parking_lot::Mutex::new(None),
16512            invoked: AtomicBool::new(false),
16513            nested_success: AtomicBool::new(false),
16514        });
16515        let agent = Arc::new(
16516            AgentBuilder::new()
16517                .system_prompt("Test hook reentrancy.")
16518                .llm(Arc::new(mock_with_response("done")))
16519                .tool(Arc::new(RecoveryTestTool {
16520                    id: "reentrant_write".to_string(),
16521                    succeeds: true,
16522                    calls: Arc::clone(&calls),
16523                }))
16524                .hooks(hooks.clone())
16525                .build()
16526                .unwrap(),
16527        );
16528        *hooks.agent.lock() = Some(Arc::downgrade(&agent));
16529        let record = tokio::time::timeout(
16530            std::time::Duration::from_secs(2),
16531            agent.invoke_tool(ToolExecutionRequest::new(
16532                "outer-hook-call",
16533                "reentrant_write",
16534                serde_json::json!({"path": "./hook.txt"}),
16535                ToolCallSource::Manual,
16536            )),
16537        )
16538        .await
16539        .expect("tool completion hook must not retain resource guards")
16540        .unwrap();
16541
16542        assert!(record.success);
16543        assert!(hooks.nested_success.load(Ordering::SeqCst));
16544        assert_eq!(calls.load(Ordering::SeqCst), 2);
16545    }
16546
16547    #[tokio::test]
16548    async fn fallback_releases_primary_resource_locks() {
16549        let primary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16550        let fallback_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16551        let agent = AgentBuilder::new()
16552            .system_prompt("Test fallback execution.")
16553            .llm(Arc::new(mock_with_response("done")))
16554            .tool(Arc::new(RecoveryTestTool {
16555                id: "primary".to_string(),
16556                succeeds: false,
16557                calls: Arc::clone(&primary_calls),
16558            }))
16559            .tool(Arc::new(RecoveryTestTool {
16560                id: "fallback".to_string(),
16561                succeeds: true,
16562                calls: Arc::clone(&fallback_calls),
16563            }))
16564            .recovery_manager(recovery_manager_with_fallbacks([(
16565                "primary".to_string(),
16566                "fallback".to_string(),
16567            )]))
16568            .build()
16569            .unwrap();
16570        let record = tokio::time::timeout(
16571            std::time::Duration::from_secs(2),
16572            agent.invoke_tool(ToolExecutionRequest::new(
16573                "fallback-call",
16574                "primary",
16575                serde_json::json!({"path": "./shared.txt"}),
16576                ToolCallSource::Manual,
16577            )),
16578        )
16579        .await
16580        .expect("fallback must not retain the primary resource guard")
16581        .unwrap();
16582
16583        assert!(record.success);
16584        assert_eq!(record.canonical_id, "fallback");
16585        assert_eq!(record.call_id, "fallback-call");
16586        assert!(matches!(record.source, ToolCallSource::Fallback { .. }));
16587        assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
16588        assert_eq!(fallback_calls.load(Ordering::SeqCst), 1);
16589    }
16590
16591    #[tokio::test]
16592    async fn diagnostics_without_provider_records_unavailable_without_execution() {
16593        let mock = mock_with_response("hello");
16594        let yaml = r#"
16595name: DiagnosticsNoProviderAgent
16596system_prompt: "Review diagnostics."
16597tools: [diagnostics]
16598"#;
16599        let agent = AgentBuilder::from_yaml(yaml)
16600            .unwrap()
16601            .llm(Arc::new(mock))
16602            .auto_configure_features()
16603            .unwrap()
16604            .build()
16605            .unwrap();
16606
16607        let record = agent
16608            .invoke_tool(ToolExecutionRequest::new(
16609                "diagnostics-call",
16610                "diagnostics",
16611                serde_json::json!({}),
16612                ToolCallSource::Manual,
16613            ))
16614            .await
16615            .unwrap();
16616
16617        assert!(!record.executed);
16618        assert!(!record.success);
16619        assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
16620    }
16621
16622    #[tokio::test]
16623    async fn web_search_without_provider_records_unavailable_without_execution() {
16624        let mock = mock_with_response("hello");
16625        let yaml = r#"
16626name: WebSearchNoProviderAgent
16627system_prompt: "You search the web."
16628tools: [web_search]
16629"#;
16630        let agent = AgentBuilder::from_yaml(yaml)
16631            .unwrap()
16632            .llm(Arc::new(mock))
16633            .auto_configure_features()
16634            .unwrap()
16635            .build()
16636            .unwrap();
16637
16638        let record = agent
16639            .invoke_tool(ToolExecutionRequest::new(
16640                "web-search-call",
16641                "web_search",
16642                serde_json::json!({"query": "rust async"}),
16643                ToolCallSource::Manual,
16644            ))
16645            .await
16646            .unwrap();
16647
16648        assert!(!record.executed);
16649        assert!(!record.success);
16650        assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
16651    }
16652
16653    #[tokio::test]
16654    async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_omitted() {
16655        let mock = mock_with_response("hello");
16656        let yaml = r#"
16657name: SpawnerNoGrantAgent
16658system_prompt: "You manage agents."
16659spawner:
16660  max_agents: 2
16661"#;
16662        let agent = AgentBuilder::from_yaml(yaml)
16663            .unwrap()
16664            .llm(Arc::new(mock))
16665            .auto_configure_features()
16666            .unwrap()
16667            .auto_configure_spawner()
16668            .await
16669            .unwrap()
16670            .build()
16671            .unwrap();
16672
16673        let available = agent.get_available_tool_ids().await.unwrap();
16674        assert!(available.is_empty());
16675    }
16676
16677    #[tokio::test]
16678    async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_empty() {
16679        let mock = mock_with_response("hello");
16680        let yaml = r#"
16681name: EmptySpawnerNoGrantAgent
16682system_prompt: "You manage agents."
16683tools: []
16684spawner:
16685  max_agents: 2
16686"#;
16687        let agent = AgentBuilder::from_yaml(yaml)
16688            .unwrap()
16689            .llm(Arc::new(mock))
16690            .auto_configure_features()
16691            .unwrap()
16692            .auto_configure_spawner()
16693            .await
16694            .unwrap()
16695            .build()
16696            .unwrap();
16697
16698        let available = agent.get_available_tool_ids().await.unwrap();
16699        assert!(available.is_empty());
16700    }
16701
16702    #[tokio::test]
16703    async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_empty() {
16704        let mock = mock_with_response("hello");
16705        let yaml = r#"
16706name: ManagementGrantAgent
16707system_prompt: "You manage agents."
16708tools: []
16709spawner:
16710  management_tools: true
16711"#;
16712        let agent = AgentBuilder::from_yaml(yaml)
16713            .unwrap()
16714            .llm(Arc::new(mock))
16715            .auto_configure_features()
16716            .unwrap()
16717            .auto_configure_spawner()
16718            .await
16719            .unwrap()
16720            .build()
16721            .unwrap();
16722
16723        let available = agent.get_available_tool_ids().await.unwrap();
16724        assert_eq!(available.len(), 4);
16725        assert!(available.contains(&"spawn_agent".to_string()));
16726        assert!(available.contains(&"send_agent_message".to_string()));
16727        assert!(available.contains(&"list_agents".to_string()));
16728        assert!(available.contains(&"remove_agent".to_string()));
16729    }
16730
16731    #[tokio::test]
16732    async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_omitted() {
16733        let mock = mock_with_response("hello");
16734        let yaml = r#"
16735name: ManagementOmittedToolsGrantAgent
16736system_prompt: "You manage agents."
16737spawner:
16738  management_tools: true
16739"#;
16740        let agent = AgentBuilder::from_yaml(yaml)
16741            .unwrap()
16742            .llm(Arc::new(mock))
16743            .auto_configure_features()
16744            .unwrap()
16745            .auto_configure_spawner()
16746            .await
16747            .unwrap()
16748            .build()
16749            .unwrap();
16750
16751        let available = agent.get_available_tool_ids().await.unwrap();
16752        assert_eq!(available.len(), 4);
16753        assert!(available.contains(&"spawn_agent".to_string()));
16754        assert!(available.contains(&"send_agent_message".to_string()));
16755        assert!(available.contains(&"list_agents".to_string()));
16756        assert!(available.contains(&"remove_agent".to_string()));
16757    }
16758
16759    #[tokio::test]
16760    async fn test_management_tools_selected_grants_only_selected_tools() {
16761        let mock = mock_with_response("hello");
16762        let yaml = r#"
16763name: ManagementSelectedGrantAgent
16764system_prompt: "You manage agents."
16765tools: []
16766spawner:
16767  management_tools:
16768    - spawn_agent
16769    - send_agent_message
16770    - list_agents
16771"#;
16772        let agent = AgentBuilder::from_yaml(yaml)
16773            .unwrap()
16774            .llm(Arc::new(mock))
16775            .auto_configure_features()
16776            .unwrap()
16777            .auto_configure_spawner()
16778            .await
16779            .unwrap()
16780            .build()
16781            .unwrap();
16782
16783        let available = agent.get_available_tool_ids().await.unwrap();
16784        assert_eq!(available.len(), 3);
16785        assert!(available.contains(&"spawn_agent".to_string()));
16786        assert!(available.contains(&"send_agent_message".to_string()));
16787        assert!(available.contains(&"list_agents".to_string()));
16788        assert!(!available.contains(&"remove_agent".to_string()));
16789    }
16790
16791    #[tokio::test]
16792    async fn test_orchestration_tools_flag_grants_tools_when_top_level_tools_empty() {
16793        let mock = mock_with_response("hello");
16794        let yaml = r#"
16795name: OrchestrationGrantAgent
16796system_prompt: "You coordinate agents."
16797llms:
16798  default:
16799    provider: openai
16800    model: gpt-4
16801  router:
16802    provider: openai
16803    model: gpt-4
16804llm:
16805  default: default
16806  router: router
16807tools: []
16808spawner:
16809  orchestration_tools: true
16810"#;
16811        let agent = AgentBuilder::from_yaml(yaml)
16812            .unwrap()
16813            .llm(Arc::new(mock))
16814            .auto_configure_features()
16815            .unwrap()
16816            .auto_configure_spawner()
16817            .await
16818            .unwrap()
16819            .build()
16820            .unwrap();
16821
16822        let available = agent.get_available_tool_ids().await.unwrap();
16823        assert_eq!(available.len(), 5);
16824        assert!(available.contains(&"route_to_agent".to_string()));
16825        assert!(available.contains(&"pipeline_process".to_string()));
16826        assert!(available.contains(&"concurrent_ask".to_string()));
16827        assert!(available.contains(&"group_discussion".to_string()));
16828        assert!(available.contains(&"handoff_conversation".to_string()));
16829    }
16830
16831    #[tokio::test]
16832    async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_empty() {
16833        let mock = mock_with_response("hello");
16834        let yaml = r#"
16835name: PersonaGrantAgent
16836system_prompt: "You can evolve persona."
16837llm:
16838  provider: openai
16839  model: gpt-4
16840tools: []
16841persona:
16842  identity:
16843    name: "Guide"
16844    role: "Helper"
16845  evolution:
16846    enabled: true
16847    allow_llm_evolve: true
16848    mutable_fields:
16849      - traits.personality
16850"#;
16851        let agent = AgentBuilder::from_yaml(yaml)
16852            .unwrap()
16853            .llm(Arc::new(mock))
16854            .build()
16855            .unwrap();
16856
16857        let available = agent.get_available_tool_ids().await.unwrap();
16858        assert_eq!(available, vec!["persona_evolve".to_string()]);
16859    }
16860
16861    #[tokio::test]
16862    async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_omitted() {
16863        let mock = mock_with_response("hello");
16864        let yaml = r#"
16865name: PersonaOmittedToolsGrantAgent
16866system_prompt: "You can evolve persona."
16867llm:
16868  provider: openai
16869  model: gpt-4
16870persona:
16871  identity:
16872    name: "Guide"
16873    role: "Helper"
16874  evolution:
16875    enabled: true
16876    allow_llm_evolve: true
16877    mutable_fields:
16878      - traits.personality
16879"#;
16880        let agent = AgentBuilder::from_yaml(yaml)
16881            .unwrap()
16882            .llm(Arc::new(mock))
16883            .build()
16884            .unwrap();
16885
16886        let available = agent.get_available_tool_ids().await.unwrap();
16887        assert_eq!(available, vec!["persona_evolve".to_string()]);
16888    }
16889
16890    #[tokio::test]
16891    async fn test_omitted_yaml_tools_exposes_no_tools() {
16892        let mock = mock_with_response("hello");
16893        let yaml = r#"
16894name: NoToolsAgent
16895system_prompt: "You are helpful."
16896"#;
16897        let agent = AgentBuilder::from_yaml(yaml)
16898            .unwrap()
16899            .llm(Arc::new(mock))
16900            .auto_configure_features()
16901            .unwrap()
16902            .build()
16903            .unwrap();
16904
16905        let available = agent.get_available_tool_ids().await.unwrap();
16906        assert!(available.is_empty());
16907    }
16908
16909    #[tokio::test]
16910    async fn runtime_scope_cannot_widen_omitted_or_empty_yaml_grants() {
16911        for tools in ["", "tools: []"] {
16912            let yaml = format!(
16913                r#"
16914name: RuntimeScopeNoGrantAgent
16915system_prompt: "No ordinary tools are granted."
16916{tools}
16917"#
16918            );
16919            let agent = AgentBuilder::from_yaml(&yaml)
16920                .unwrap()
16921                .llm(Arc::new(mock_with_response("done")))
16922                .auto_configure_features()
16923                .unwrap()
16924                .build()
16925                .unwrap();
16926
16927            agent
16928                .runtime_control()
16929                .set_tool_scope(vec!["calculator".to_string()]);
16930
16931            assert!(agent.get_available_tool_ids().await.unwrap().is_empty());
16932        }
16933    }
16934
16935    #[tokio::test]
16936    async fn runtime_scope_widening_attempt_keeps_only_declared_tools() {
16937        let yaml = r#"
16938name: RuntimeScopeWideningAgent
16939system_prompt: "Runtime scope cannot add authority."
16940tools: [calculator]
16941"#;
16942        let agent = AgentBuilder::from_yaml(yaml)
16943            .unwrap()
16944            .llm(Arc::new(mock_with_response("done")))
16945            .auto_configure_features()
16946            .unwrap()
16947            .build()
16948            .unwrap();
16949
16950        agent
16951            .runtime_control()
16952            .set_tool_scope(vec!["calculator".to_string(), "datetime".to_string()]);
16953
16954        assert_eq!(
16955            agent.get_available_tool_ids().await.unwrap(),
16956            vec!["calculator".to_string()]
16957        );
16958    }
16959
16960    #[tokio::test]
16961    async fn runtime_scope_is_canonical_unique_ordered_and_clear_restores_declared_grant() {
16962        let yaml = r#"
16963name: RuntimeScopeIntersectionAgent
16964system_prompt: "Use only declared tools."
16965tools: [calculator, datetime]
16966"#;
16967        let agent = AgentBuilder::from_yaml(yaml)
16968            .unwrap()
16969            .llm(Arc::new(mock_with_response("done")))
16970            .auto_configure_features()
16971            .unwrap()
16972            .build()
16973            .unwrap();
16974        let mut aliases = ai_agents_tools::ToolAliases::default();
16975        aliases
16976            .names
16977            .insert("en".to_string(), "calculate_alias".to_string());
16978        agent.tools.set_tool_aliases("calculator", aliases);
16979        let control = agent.runtime_control();
16980
16981        control.set_tool_scope(vec![
16982            "datetime".to_string(),
16983            "calculate_alias".to_string(),
16984            "calculator".to_string(),
16985            "unknown".to_string(),
16986            "datetime".to_string(),
16987        ]);
16988        assert_eq!(
16989            agent.get_available_tool_ids().await.unwrap(),
16990            vec!["calculator".to_string(), "datetime".to_string()]
16991        );
16992
16993        control.set_tool_scope(vec!["datetime".to_string()]);
16994        assert_eq!(
16995            agent.get_available_tool_ids().await.unwrap(),
16996            vec!["datetime".to_string()]
16997        );
16998
16999        control.clear_tool_scope_override();
17000        assert_eq!(
17001            agent.get_available_tool_ids().await.unwrap(),
17002            vec!["calculator".to_string(), "datetime".to_string()]
17003        );
17004    }
17005
17006    #[tokio::test]
17007    async fn runtime_scope_preserves_programmatic_registration_as_declared_grant() {
17008        let agent = AgentBuilder::new()
17009            .system_prompt("Use registered tools.")
17010            .llm(Arc::new(mock_with_response("done")))
17011            .tool(Arc::new(ContextEchoTool))
17012            .tool(Arc::new(SlowTool))
17013            .build()
17014            .unwrap();
17015
17016        agent.runtime_control().set_tool_scope(vec![
17017            "Context Echo".to_string(),
17018            "context_echo".to_string(),
17019            "unknown".to_string(),
17020        ]);
17021
17022        assert_eq!(
17023            agent.get_available_tool_ids().await.unwrap(),
17024            vec!["context_echo".to_string()]
17025        );
17026    }
17027
17028    #[tokio::test]
17029    async fn nested_state_scopes_intersect_every_ancestor_with_aliases() {
17030        let yaml = r#"
17031name: NestedStateScopeAgent
17032system_prompt: "Honor every state scope."
17033tools: [calculator, datetime, echo]
17034states:
17035  initial: root
17036  states:
17037    root:
17038      tools: [calculate_alias, datetime]
17039      initial: middle
17040      states:
17041        middle:
17042          initial: leaf
17043          states:
17044            leaf:
17045              tools: [datetime_alias, echo]
17046"#;
17047        let agent = AgentBuilder::from_yaml(yaml)
17048            .unwrap()
17049            .llm(Arc::new(mock_with_response("done")))
17050            .auto_configure_features()
17051            .unwrap()
17052            .build()
17053            .unwrap();
17054        let mut calculator_aliases = ai_agents_tools::ToolAliases::default();
17055        calculator_aliases
17056            .names
17057            .insert("en".to_string(), "calculate_alias".to_string());
17058        agent
17059            .tools
17060            .set_tool_aliases("calculator", calculator_aliases);
17061        let mut datetime_aliases = ai_agents_tools::ToolAliases::default();
17062        datetime_aliases
17063            .names
17064            .insert("en".to_string(), "datetime_alias".to_string());
17065        agent.tools.set_tool_aliases("datetime", datetime_aliases);
17066        agent.runtime_control().set_tool_scope(vec![
17067            "unknown".to_string(),
17068            "datetime_alias".to_string(),
17069            "calculate_alias".to_string(),
17070            "datetime".to_string(),
17071        ]);
17072
17073        assert_eq!(agent.current_state().as_deref(), Some("root.middle.leaf"));
17074        assert_eq!(
17075            agent.get_available_tool_ids().await.unwrap(),
17076            vec!["datetime".to_string()]
17077        );
17078    }
17079
17080    #[tokio::test]
17081    async fn ancestor_empty_state_scope_denies_omitted_descendants() {
17082        let yaml = r#"
17083name: NestedEmptyStateScopeAgent
17084system_prompt: "An empty ancestor scope denies all tools."
17085tools: [calculator]
17086states:
17087  initial: root
17088  states:
17089    root:
17090      tools: []
17091      initial: middle
17092      states:
17093        middle:
17094          initial: leaf
17095          states:
17096            leaf: {}
17097"#;
17098        let agent = AgentBuilder::from_yaml(yaml)
17099            .unwrap()
17100            .llm(Arc::new(mock_with_response("done")))
17101            .auto_configure_features()
17102            .unwrap()
17103            .build()
17104            .unwrap();
17105
17106        assert!(agent.get_available_tool_ids().await.unwrap().is_empty());
17107    }
17108
17109    #[tokio::test]
17110    async fn state_change_during_approval_invalidates_the_reviewed_authority() {
17111        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17112        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17113        let entered = Arc::new(tokio::sync::Barrier::new(2));
17114        let release = Arc::new(tokio::sync::Notify::new());
17115        let handler = Arc::new(BlockingApprovalHandler {
17116            entered: Arc::clone(&entered),
17117            release: Arc::clone(&release),
17118            result: ApprovalResult::Approved,
17119        });
17120        let yaml = r#"
17121name: ApprovalStateGenerationAgent
17122system_prompt: "State authority may change during approval."
17123tools: [locked_write]
17124states:
17125  initial: first
17126  states:
17127    first:
17128      tools: [locked_write]
17129    second:
17130      tools: [locked_write]
17131"#;
17132        let agent = Arc::new(
17133            AgentBuilder::from_yaml(yaml)
17134                .unwrap()
17135                .llm(Arc::new(mock_with_response("done")))
17136                .tool(Arc::new(LockedWriteTool {
17137                    active: Arc::clone(&active),
17138                    max_active: Arc::clone(&max_active),
17139                }))
17140                .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
17141                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17142                .approval_handler(handler)
17143                .build()
17144                .unwrap(),
17145        );
17146        let running = Arc::clone(&agent);
17147        let call = tokio::spawn(async move {
17148            running
17149                .invoke_tool(ToolExecutionRequest::new(
17150                    "approval-state-generation",
17151                    "locked_write",
17152                    serde_json::json!({"path": "./state-generation.txt"}),
17153                    ToolCallSource::Manual,
17154                ))
17155                .await
17156                .unwrap()
17157        });
17158
17159        entered.wait().await;
17160        agent.transition_to("second").await.unwrap();
17161        release.notify_one();
17162        let record = call.await.unwrap();
17163
17164        assert!(!record.executed);
17165        assert!(record.output.contains("Approval became stale"));
17166        assert_eq!(max_active.load(Ordering::SeqCst), 0);
17167    }
17168
17169    #[tokio::test]
17170    async fn state_change_while_waiting_for_resource_lock_fails_final_admission() {
17171        let holder_gate = PathMutationGate::new();
17172        let waiter_gate = PathMutationGate::new();
17173        let yaml = r#"
17174name: LockedStateGenerationAgent
17175system_prompt: "State authority must remain stable through admission."
17176tools: [state_lock_holder, state_lock_waiter]
17177states:
17178  initial: first
17179  states:
17180    first:
17181      tools: [state_lock_holder, state_lock_waiter]
17182    second:
17183      tools: [state_lock_holder, state_lock_waiter]
17184"#;
17185        let agent = Arc::new(
17186            AgentBuilder::from_yaml(yaml)
17187                .unwrap()
17188                .llm(Arc::new(mock_with_response("done")))
17189                .tool(Arc::new(BlockingPathMutationTool {
17190                    id: "state_lock_holder",
17191                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
17192                    gate: holder_gate.clone(),
17193                }))
17194                .tool(Arc::new(BlockingPathMutationTool {
17195                    id: "state_lock_waiter",
17196                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
17197                    gate: waiter_gate.clone(),
17198                }))
17199                .build()
17200                .unwrap(),
17201        );
17202        let holder_call = {
17203            let agent = Arc::clone(&agent);
17204            tokio::spawn(async move {
17205                agent
17206                    .invoke_tool(ToolExecutionRequest::new(
17207                        "state-lock-holder",
17208                        "state_lock_holder",
17209                        serde_json::json!({"path": "./shared-state-path.txt"}),
17210                        ToolCallSource::Manual,
17211                    ))
17212                    .await
17213                    .unwrap()
17214            })
17215        };
17216        holder_gate.wait_until_entered().await;
17217        let waiter_call = {
17218            let agent = Arc::clone(&agent);
17219            tokio::spawn(async move {
17220                agent
17221                    .invoke_tool(ToolExecutionRequest::new(
17222                        "state-lock-waiter",
17223                        "state_lock_waiter",
17224                        serde_json::json!({"path": "./shared-state-path.txt"}),
17225                        ToolCallSource::Manual,
17226                    ))
17227                    .await
17228                    .unwrap()
17229            })
17230        };
17231
17232        wait_for_resource_lock_strong_count(&agent.resource_locks, 2).await;
17233        agent.transition_to("second").await.unwrap();
17234        holder_gate.release();
17235        let holder_record = holder_call.await.unwrap();
17236        let waiter_record = waiter_call.await.unwrap();
17237
17238        assert!(holder_record.success);
17239        assert!(!waiter_record.executed);
17240        assert!(
17241            waiter_record
17242                .output
17243                .contains("state scope changed before admission")
17244        );
17245        assert!(!waiter_gate.entered.load(Ordering::SeqCst));
17246    }
17247
17248    #[tokio::test]
17249    async fn test_state_tools_cannot_widen_top_level_grant() {
17250        let mock = mock_with_response("hello");
17251        let yaml = r#"
17252name: NarrowToolsAgent
17253system_prompt: "You are helpful."
17254tools:
17255  - calculator
17256states:
17257  initial: current
17258  states:
17259    current:
17260      tools: [datetime]
17261"#;
17262        let agent = AgentBuilder::from_yaml(yaml)
17263            .unwrap()
17264            .llm(Arc::new(mock))
17265            .auto_configure_features()
17266            .unwrap()
17267            .build()
17268            .unwrap();
17269
17270        let available = agent.get_available_tool_ids().await.unwrap();
17271        assert!(available.is_empty());
17272    }
17273
17274    // Tool execution in chat flow
17275    #[tokio::test]
17276    async fn test_integration_tool_execution() {
17277        // Mock LLM that returns a tool call then a final answer
17278        let mock = mock_with_responses(vec![
17279            // First response: tool call
17280            r#"I'll calculate that for you.
17281[TOOL_CALL: {"name": "calculator", "arguments": {"expression": "2+2"}}]"#,
17282            // After tool result: final answer
17283            "The answer is 4.",
17284        ]);
17285        let mut tools = ai_agents_tools::ToolRegistry::new();
17286        tools
17287            .register(Arc::new(ai_agents_tools::CalculatorTool))
17288            .unwrap();
17289
17290        let agent = AgentBuilder::new()
17291            .system_prompt("You are a calculator assistant.")
17292            .llm(Arc::new(mock))
17293            .tools(tools)
17294            .build()
17295            .unwrap();
17296
17297        let response = agent.chat("What is 2+2?").await.unwrap();
17298        // The agent should eventually produce a response
17299        assert!(!response.content.is_empty());
17300    }
17301
17302    #[tokio::test]
17303    async fn test_tool_hitl_rejection_finalizes_blocking_turn() {
17304        let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17305        let hooks = Arc::new(ResponseCountingHooks {
17306            responses: Arc::clone(&responses),
17307        });
17308        let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
17309        let yaml = r#"
17310name: ToolRejectAgent
17311system_prompt: "You use tools when requested."
17312tools:
17313  - echo
17314hitl:
17315  tools:
17316    echo:
17317      require_approval: true
17318      approval_message: "Approve echo?"
17319"#;
17320        let agent = AgentBuilder::from_yaml(yaml)
17321            .unwrap()
17322            .llm(Arc::new(mock))
17323            .auto_configure_features()
17324            .unwrap()
17325            .hooks(hooks)
17326            .build()
17327            .unwrap();
17328
17329        let response = agent.chat("echo hello").await.unwrap();
17330
17331        assert!(
17332            response.content.contains("Operation cancelled"),
17333            "unexpected response: {}",
17334            response.content
17335        );
17336        assert_eq!(responses.load(Ordering::SeqCst), 1);
17337        let messages = agent.memory.get_messages(None).await.unwrap();
17338        assert_eq!(messages.len(), 3);
17339        assert_eq!(messages[0].content, "echo hello");
17340        assert!(messages[1].content.contains("\"tool\":\"echo\""));
17341        assert!(messages[2].content.contains("rejected by the approver"));
17342    }
17343
17344    #[tokio::test]
17345    async fn test_tool_hitl_rejection_finalizes_streaming_turn() {
17346        use futures::StreamExt;
17347
17348        let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17349        let hooks = Arc::new(ResponseCountingHooks {
17350            responses: Arc::clone(&responses),
17351        });
17352        let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
17353        let yaml = r#"
17354name: ToolRejectStreamingAgent
17355system_prompt: "You use tools when requested."
17356tools:
17357  - echo
17358streaming:
17359  enabled: true
17360hitl:
17361  tools:
17362    echo:
17363      require_approval: true
17364      approval_message: "Approve echo?"
17365"#;
17366        let agent = AgentBuilder::from_yaml(yaml)
17367            .unwrap()
17368            .llm(Arc::new(mock))
17369            .auto_configure_features()
17370            .unwrap()
17371            .hooks(hooks)
17372            .build()
17373            .unwrap();
17374
17375        let mut stream = agent.chat_stream("echo hello").await.unwrap();
17376        let mut terminal_error = String::new();
17377        let mut done = false;
17378        while let Some(chunk) = stream.next().await {
17379            match chunk {
17380                StreamChunk::Error { message } => terminal_error = message,
17381                StreamChunk::Done {} => {
17382                    done = true;
17383                    break;
17384                }
17385                _ => {}
17386            }
17387        }
17388
17389        assert!(done);
17390        assert!(
17391            terminal_error.contains("Operation cancelled"),
17392            "unexpected terminal error: {}",
17393            terminal_error
17394        );
17395        assert_eq!(responses.load(Ordering::SeqCst), 1);
17396        let messages = agent.memory.get_messages(None).await.unwrap();
17397        assert_eq!(messages.len(), 3);
17398        assert_eq!(messages[0].content, "echo hello");
17399        assert!(messages[1].content.contains("\"tool\":\"echo\""));
17400        assert!(messages[2].content.contains("rejected by the approver"));
17401    }
17402
17403    #[tokio::test]
17404    async fn test_pre_response_guard_transition_skips_old_state_llm() {
17405        let mock = mock_with_response("Billing state response");
17406        let call_counter = mock.clone();
17407        let yaml = r#"
17408name: OptimizedStateAgent
17409system_prompt: "You route before answering."
17410runtime:
17411  optimization:
17412    enabled: true
17413    pre_response_deterministic_transitions: true
17414states:
17415  initial: greeting
17416  states:
17417    greeting:
17418      prompt: "Old state prompt that should be skipped."
17419      transitions:
17420        - to: billing
17421          guard:
17422            context:
17423              topic:
17424                eq: billing
17425          timing: pre_response
17426    billing:
17427      prompt: "Answer from the billing state."
17428"#;
17429        let agent = AgentBuilder::from_yaml(yaml)
17430            .unwrap()
17431            .llm(Arc::new(mock))
17432            .build()
17433            .unwrap();
17434        agent
17435            .set_context("topic", serde_json::json!("billing"))
17436            .unwrap();
17437
17438        let response = agent.chat("I need billing help").await.unwrap();
17439
17440        assert_eq!(agent.current_state().as_deref(), Some("billing"));
17441        assert_eq!(response.content, "Billing state response");
17442        assert_eq!(call_counter.call_count(), 1);
17443        assert_eq!(agent.actor_facts().len(), 0);
17444    }
17445
17446    #[tokio::test]
17447    async fn test_set_context_supports_dotted_paths_for_pre_response_guards() {
17448        let mock = mock_with_response("Billing state response");
17449        let call_counter = mock.clone();
17450        let yaml = r#"
17451name: OptimizedStateAgent
17452system_prompt: "You route before answering."
17453runtime:
17454  optimization:
17455    enabled: true
17456    pre_response_deterministic_transitions: true
17457context:
17458  request:
17459    type: runtime
17460    default:
17461      topic: general
17462states:
17463  initial: greeting
17464  states:
17465    greeting:
17466      prompt: "Old state prompt that should be skipped."
17467      transitions:
17468        - to: billing
17469          guard:
17470            context:
17471              request.topic:
17472                eq: billing
17473          timing: pre_response
17474    billing:
17475      prompt: "Answer from the billing state."
17476"#;
17477        let agent = AgentBuilder::from_yaml(yaml)
17478            .unwrap()
17479            .llm(Arc::new(mock))
17480            .build()
17481            .unwrap();
17482        agent
17483            .set_context("request.topic", serde_json::json!("billing"))
17484            .unwrap();
17485
17486        let response = agent.chat("I need billing help").await.unwrap();
17487
17488        assert_eq!(agent.current_state().as_deref(), Some("billing"));
17489        assert_eq!(response.content, "Billing state response");
17490        assert_eq!(call_counter.call_count(), 1);
17491        assert_eq!(
17492            agent.get_context().get("request"),
17493            Some(&serde_json::json!({"topic": "billing"}))
17494        );
17495    }
17496
17497    #[tokio::test]
17498    async fn test_pre_response_rejection_does_not_commit_staged_context_or_user() {
17499        let mock = mock_with_response("billing");
17500        let yaml = r#"
17501name: OptimizedStateAgent
17502system_prompt: "You route before answering."
17503runtime:
17504  optimization:
17505    enabled: true
17506    pre_response_deterministic_transitions: true
17507hitl:
17508  states:
17509    billing:
17510      on_enter: require_approval
17511      approval_message: "Approve billing route?"
17512states:
17513  initial: greeting
17514  states:
17515    greeting:
17516      prompt: "Old state prompt."
17517      extract:
17518        - key: topic
17519          description: "Support topic"
17520      transitions:
17521        - to: billing
17522          guard:
17523            context:
17524              topic:
17525                eq: billing
17526          timing: pre_response
17527          run_extractors: true
17528    billing:
17529      prompt: "Billing state."
17530"#;
17531        let agent = AgentBuilder::from_yaml(yaml)
17532            .unwrap()
17533            .llm(Arc::new(mock))
17534            .build()
17535            .unwrap();
17536
17537        let response = agent
17538            .try_pre_response_transition("billing please")
17539            .await
17540            .unwrap();
17541
17542        assert!(response.is_none());
17543        assert_eq!(agent.current_state().as_deref(), Some("greeting"));
17544        assert!(!agent.get_context().contains_key("topic"));
17545        assert_eq!(agent.memory.get_messages(None).await.unwrap().len(), 0);
17546    }
17547
17548    #[tokio::test]
17549    async fn test_pre_response_extractor_commits_context_on_winning_path() {
17550        let mock = mock_with_responses(vec!["billing", "Billing response"]);
17551        let yaml = r#"
17552name: OptimizedStateAgent
17553system_prompt: "You route before answering."
17554runtime:
17555  optimization:
17556    enabled: true
17557    pre_response_deterministic_transitions: true
17558states:
17559  initial: greeting
17560  states:
17561    greeting:
17562      prompt: "Old state prompt."
17563      extract:
17564        - key: topic
17565          description: "Support topic"
17566      transitions:
17567        - to: billing
17568          guard:
17569            context:
17570              topic:
17571                eq: billing
17572          timing: pre_response
17573          run_extractors: true
17574    billing:
17575      prompt: "Billing state."
17576"#;
17577        let agent = AgentBuilder::from_yaml(yaml)
17578            .unwrap()
17579            .llm(Arc::new(mock))
17580            .build()
17581            .unwrap();
17582
17583        let response = agent.chat("billing please").await.unwrap();
17584
17585        assert_eq!(agent.current_state().as_deref(), Some("billing"));
17586        assert_eq!(response.content, "Billing response");
17587        assert_eq!(
17588            agent.get_context().get("topic"),
17589            Some(&serde_json::json!("billing"))
17590        );
17591    }
17592
17593    #[tokio::test]
17594    async fn test_pre_response_extractor_miss_does_not_mutate_context() {
17595        let mock = mock_with_response("__NONE__");
17596        let yaml = r#"
17597name: OptimizedStateAgent
17598system_prompt: "You route before answering."
17599runtime:
17600  optimization:
17601    enabled: true
17602    pre_response_deterministic_transitions: true
17603states:
17604  initial: greeting
17605  states:
17606    greeting:
17607      prompt: "Old state prompt."
17608      extract:
17609        - key: topic
17610          description: "Support topic"
17611      transitions:
17612        - to: billing
17613          guard:
17614            context:
17615              topic:
17616                eq: billing
17617          timing: pre_response
17618          run_extractors: true
17619    billing:
17620      prompt: "Billing state."
17621"#;
17622        let agent = AgentBuilder::from_yaml(yaml)
17623            .unwrap()
17624            .llm(Arc::new(mock))
17625            .build()
17626            .unwrap();
17627
17628        let response = agent.try_pre_response_transition("hello").await.unwrap();
17629
17630        assert!(response.is_none());
17631        assert_eq!(agent.current_state().as_deref(), Some("greeting"));
17632        assert!(!agent.get_context().contains_key("topic"));
17633    }
17634
17635    #[tokio::test]
17636    async fn test_default_guard_transition_stays_post_response() {
17637        let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
17638        let call_counter = mock.clone();
17639        let yaml = r#"
17640name: TimingAgent
17641system_prompt: "You route carefully."
17642runtime:
17643  optimization:
17644    enabled: true
17645    pre_response_deterministic_transitions: true
17646states:
17647  initial: greeting
17648  states:
17649    greeting:
17650      prompt: "Old state prompt."
17651      transitions:
17652        - to: billing
17653          guard:
17654            context:
17655              topic:
17656                eq: billing
17657    billing:
17658      prompt: "Billing state."
17659"#;
17660        let agent = AgentBuilder::from_yaml(yaml)
17661            .unwrap()
17662            .llm(Arc::new(mock))
17663            .build()
17664            .unwrap();
17665        agent
17666            .set_context("topic", serde_json::json!("billing"))
17667            .unwrap();
17668
17669        let response = agent.chat("billing please").await.unwrap();
17670
17671        assert_eq!(agent.current_state().as_deref(), Some("billing"));
17672        assert_eq!(response.content, "Billing response");
17673        assert_eq!(call_counter.call_count(), 2);
17674    }
17675
17676    #[tokio::test]
17677    async fn test_explicit_post_response_guard_transition_stays_post_response() {
17678        let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
17679        let call_counter = mock.clone();
17680        let yaml = r#"
17681name: TimingAgent
17682system_prompt: "You route carefully."
17683runtime:
17684  optimization:
17685    enabled: true
17686    pre_response_deterministic_transitions: true
17687states:
17688  initial: greeting
17689  states:
17690    greeting:
17691      prompt: "Old state prompt."
17692      transitions:
17693        - to: billing
17694          guard:
17695            context:
17696              topic:
17697                eq: billing
17698          timing: post_response
17699    billing:
17700      prompt: "Billing state."
17701"#;
17702        let agent = AgentBuilder::from_yaml(yaml)
17703            .unwrap()
17704            .llm(Arc::new(mock))
17705            .build()
17706            .unwrap();
17707        agent
17708            .set_context("topic", serde_json::json!("billing"))
17709            .unwrap();
17710
17711        let response = agent.chat("billing please").await.unwrap();
17712
17713        assert_eq!(agent.current_state().as_deref(), Some("billing"));
17714        assert_eq!(response.content, "Billing response");
17715        assert_eq!(call_counter.call_count(), 2);
17716    }
17717
17718    #[tokio::test]
17719    async fn test_pre_response_extractors_are_transition_scoped() {
17720        let mock = mock_with_responses(vec!["billing", "Billing response"]);
17721        let yaml = r#"
17722name: ScopedExtractorAgent
17723system_prompt: "You route carefully."
17724runtime:
17725  optimization:
17726    enabled: true
17727    pre_response_deterministic_transitions: true
17728states:
17729  initial: greeting
17730  states:
17731    greeting:
17732      prompt: "Old state prompt."
17733      extract:
17734        - key: topic
17735          description: "Support topic"
17736      transitions:
17737        - to: wrong
17738          guard:
17739            context:
17740              topic:
17741                eq: billing
17742          timing: pre_response
17743        - to: billing
17744          guard:
17745            context:
17746              topic:
17747                eq: billing
17748          timing: pre_response
17749          run_extractors: true
17750    wrong:
17751      prompt: "Wrong state."
17752    billing:
17753      prompt: "Billing state."
17754"#;
17755        let agent = AgentBuilder::from_yaml(yaml)
17756            .unwrap()
17757            .llm(Arc::new(mock))
17758            .build()
17759            .unwrap();
17760
17761        let response = agent.chat("billing please").await.unwrap();
17762
17763        assert_eq!(agent.current_state().as_deref(), Some("billing"));
17764        assert_eq!(response.content, "Billing response");
17765    }
17766
17767    #[tokio::test]
17768    async fn test_pre_response_resolved_intent_routes_early() {
17769        let mock = mock_with_response("Billing response");
17770        let yaml = r#"
17771name: IntentAgent
17772system_prompt: "You route carefully."
17773runtime:
17774  optimization:
17775    enabled: true
17776    pre_response_deterministic_transitions: true
17777states:
17778  initial: greeting
17779  states:
17780    greeting:
17781      prompt: "Old state prompt."
17782      transitions:
17783        - to: billing
17784          intent: billing
17785          timing: pre_response
17786    billing:
17787      prompt: "Billing state."
17788"#;
17789        let agent = AgentBuilder::from_yaml(yaml)
17790            .unwrap()
17791            .llm(Arc::new(mock))
17792            .build()
17793            .unwrap();
17794        agent
17795            .set_context("resolved_intent", serde_json::json!("billing"))
17796            .unwrap();
17797
17798        let response = agent
17799            .try_pre_response_transition("I need billing help")
17800            .await
17801            .unwrap()
17802            .unwrap();
17803
17804        assert_eq!(agent.current_state().as_deref(), Some("billing"));
17805        assert_eq!(response.content, "Billing response");
17806    }
17807
17808    #[tokio::test]
17809    async fn test_background_overflow_error_surfaces() {
17810        let mut config = RuntimeConfig::default();
17811        config.optimization.enabled = true;
17812        config.optimization.post_turn.max_background_tasks = 1;
17813        config.optimization.post_turn.on_background_overflow = BackgroundOverflowPolicy::Error;
17814        let policy = crate::optimization::MaintenanceTaskPolicy {
17815            mode: MaintenanceMode::Background,
17816            await_before_next_turn: AwaitBeforeNextTurn::Always,
17817        };
17818        let agent = AgentBuilder::new()
17819            .system_prompt("You are helpful.")
17820            .llm(Arc::new(mock_with_response("ok")))
17821            .build()
17822            .unwrap()
17823            .with_runtime_config(config);
17824        agent
17825            .background_maintenance
17826            .spawn(None, async { std::future::pending::<Result<()>>().await })
17827            .unwrap();
17828
17829        let result = agent
17830            .spawn_or_handle_background(None, async { Ok(()) }, "facts", &policy)
17831            .await;
17832
17833        assert!(result.is_err());
17834    }
17835
17836    #[tokio::test]
17837    async fn test_speculative_reasoning_low_cap_uses_serial_reasoning() {
17838        let default_mock = mock_with_response("Plain draft response");
17839        let router_mock = mock_with_response("cot");
17840        let router_counter = router_mock.clone();
17841        let yaml = r#"
17842name: ReasoningReservationAgent
17843system_prompt: "You answer plainly unless reasoning wins."
17844llm:
17845  default: default
17846  router: router
17847observability:
17848  enabled: true
17849  export:
17850    write_raw_events: true
17851reasoning:
17852  mode: auto
17853  judge_llm: router
17854runtime:
17855  optimization:
17856    enabled: true
17857    max_speculative_llm_calls_per_turn: 1
17858    speculative_reasoning_auto: true
17859    max_parallel_runtime_tasks: 2
17860"#;
17861        let agent = AgentBuilder::from_yaml(yaml)
17862            .unwrap()
17863            .llm_alias("default", Arc::new(default_mock))
17864            .llm_alias("router", Arc::new(router_mock))
17865            .build()
17866            .unwrap();
17867
17868        let response = agent.chat("hello").await.unwrap();
17869
17870        assert_eq!(response.content, "Plain draft response");
17871        assert_eq!(router_counter.call_count(), 1);
17872        let events = agent.observability().unwrap().raw_events();
17873        assert!(!events.iter().any(|event| {
17874            event.dimensions.get("commit_behavior") == Some(&"reasoning_decision".to_string())
17875        }));
17876    }
17877
17878    #[tokio::test]
17879    async fn test_forced_reasoning_skips_plain_speculative_draft() {
17880        let mock = mock_with_response("Reasoned response");
17881        let yaml = r#"
17882name: ForcedReasoningAgent
17883system_prompt: "You reason before answering."
17884observability:
17885  enabled: true
17886  export:
17887    write_raw_events: true
17888reasoning:
17889  mode: cot
17890runtime:
17891  optimization:
17892    enabled: true
17893    max_speculative_llm_calls_per_turn: 2
17894    speculative_state_transitions: true
17895    max_parallel_runtime_tasks: 2
17896states:
17897  initial: triage
17898  states:
17899    triage:
17900      prompt: "Answer from triage."
17901      transitions:
17902        - to: billing
17903          guard:
17904            context:
17905              route:
17906                eq: billing
17907          timing: parallel
17908    billing:
17909      prompt: "Billing state."
17910"#;
17911        let agent = AgentBuilder::from_yaml(yaml)
17912            .unwrap()
17913            .llm(Arc::new(mock))
17914            .build()
17915            .unwrap();
17916
17917        let response = agent.chat("hello").await.unwrap();
17918
17919        assert_eq!(response.content, "Reasoned response");
17920        let events = agent.observability().unwrap().raw_events();
17921        assert!(
17922            !events
17923                .iter()
17924                .any(|event| event.dimensions.contains_key("branch_status"))
17925        );
17926    }
17927
17928    #[tokio::test]
17929    async fn test_speculative_skill_low_cap_uses_serial_skill_route() {
17930        let default_mock = mock_with_response("Skill committed response");
17931        let router_mock = mock_with_response("helper");
17932        let router_counter = router_mock.clone();
17933        let yaml = r#"
17934name: SkillReservationAgent
17935system_prompt: "Use skills when they match."
17936llm:
17937  default: default
17938  router: router
17939observability:
17940  enabled: true
17941  export:
17942    write_raw_events: true
17943runtime:
17944  optimization:
17945    enabled: true
17946    max_speculative_llm_calls_per_turn: 1
17947    speculative_skill_routing: true
17948    max_parallel_runtime_tasks: 2
17949skills:
17950  - id: helper
17951    description: "Answer helper requests"
17952    trigger: "User asks for helper"
17953    steps:
17954      - prompt: "Answer the helper request: {{ user_input }}"
17955"#;
17956        let agent = AgentBuilder::from_yaml(yaml)
17957            .unwrap()
17958            .llm_alias("default", Arc::new(default_mock))
17959            .llm_alias("router", Arc::new(router_mock))
17960            .build()
17961            .unwrap();
17962
17963        let response = agent.chat("please use helper").await.unwrap();
17964
17965        assert_eq!(response.content, "Skill committed response");
17966        assert_eq!(router_counter.call_count(), 1);
17967        let events = agent.observability().unwrap().raw_events();
17968        assert!(
17969            !events
17970                .iter()
17971                .any(|event| event.dimensions.contains_key("branch_status"))
17972        );
17973    }
17974
17975    #[tokio::test]
17976    async fn test_parallel_transition_low_cap_allows_deterministic_route() {
17977        let mock = mock_with_response("unused");
17978        let call_counter = mock.clone();
17979        let yaml = r#"
17980name: ParallelTransitionLowCapAgent
17981system_prompt: "Route before stale responses when safe."
17982runtime:
17983  optimization:
17984    enabled: true
17985    max_speculative_llm_calls_per_turn: 1
17986    speculative_state_transitions: true
17987    max_parallel_runtime_tasks: 2
17988states:
17989  initial: triage
17990  states:
17991    triage:
17992      prompt: "Triage state."
17993      transitions:
17994        - to: billing
17995          guard:
17996            context:
17997              route:
17998                eq: billing
17999          timing: parallel
18000    billing:
18001      prompt: "Billing state."
18002"#;
18003        let agent = AgentBuilder::from_yaml(yaml)
18004            .unwrap()
18005            .llm(Arc::new(mock))
18006            .build()
18007            .unwrap();
18008        agent
18009            .set_context("route", serde_json::json!("billing"))
18010            .unwrap();
18011        agent.update_active_turn_context("billing help", HashMap::new());
18012        assert!(
18013            agent.reserve_active_speculative_llm_call(
18014                RuntimeOptimizationKind::ParallelStateTransition
18015            )
18016        );
18017
18018        let selection = agent
18019            .select_parallel_transition_candidate("billing help")
18020            .await
18021            .unwrap();
18022        agent.end_root_turn();
18023
18024        match selection {
18025            ParallelTransitionSelection::Candidate(candidate) => {
18026                assert_eq!(candidate.target(), "billing");
18027            }
18028            ParallelTransitionSelection::NoMatch => panic!("deterministic route did not match"),
18029            ParallelTransitionSelection::ReservationExhausted => {
18030                panic!("deterministic route consumed LLM budget")
18031            }
18032        }
18033        assert_eq!(call_counter.call_count(), 0);
18034    }
18035
18036    #[tokio::test]
18037    async fn speculative_transition_drops_loser_before_state_actions() {
18038        let lock = Arc::new(tokio::sync::Mutex::new(()));
18039        let first_started = Arc::new(tokio::sync::Notify::new());
18040        let first_dropped = Arc::new(AtomicBool::new(false));
18041        let committed_after_drop = Arc::new(AtomicBool::new(false));
18042        let default = Arc::new(FirstCallLockingProvider {
18043            lock,
18044            first_started: Arc::clone(&first_started),
18045            first_dropped: Arc::clone(&first_dropped),
18046            committed_after_drop: Arc::clone(&committed_after_drop),
18047            calls: AtomicU64::new(0),
18048        });
18049        let router = Arc::new(RoutingAfterProviderStart {
18050            provider_started: first_started,
18051        });
18052        let yaml = r#"
18053name: SpeculativeCancellationAgent
18054system_prompt: "Route before committed work."
18055llm:
18056  default: default
18057  router: router
18058runtime:
18059  optimization:
18060    enabled: true
18061    max_speculative_llm_calls_per_turn: 2
18062    speculative_state_transitions: true
18063    max_parallel_runtime_tasks: 2
18064states:
18065  initial: triage
18066  states:
18067    triage:
18068      prompt: "Triage state."
18069      transitions:
18070        - to: technical
18071          when: "The request needs technical support"
18072          timing: parallel
18073    technical:
18074      prompt: "Technical state."
18075      on_enter:
18076        - prompt: "Prepare technical context."
18077          llm: default
18078          store_as: preparation
18079"#;
18080        let agent = AgentBuilder::from_yaml(yaml)
18081            .unwrap()
18082            .llm_alias("default", default)
18083            .llm_alias("router", router)
18084            .build()
18085            .unwrap();
18086
18087        let response = tokio::time::timeout(
18088            std::time::Duration::from_secs(2),
18089            agent.chat("I cannot log in because of AUTH-17."),
18090        )
18091        .await
18092        .expect("committed work must not wait on the losing provider future")
18093        .unwrap();
18094
18095        assert_eq!(response.content, "Committed technical response.");
18096        assert_eq!(agent.current_state().as_deref(), Some("technical"));
18097        assert!(first_dropped.load(Ordering::SeqCst));
18098        assert!(committed_after_drop.load(Ordering::SeqCst));
18099    }
18100
18101    #[tokio::test]
18102    async fn buffered_transition_drops_stale_stream_before_redispatch() {
18103        use futures::StreamExt;
18104
18105        let lock = Arc::new(tokio::sync::Mutex::new(()));
18106        let stream_started = Arc::new(tokio::sync::Notify::new());
18107        let stream_dropped = Arc::new(AtomicBool::new(false));
18108        let committed_after_drop = Arc::new(AtomicBool::new(false));
18109        let default = Arc::new(BufferedLockingProvider {
18110            lock,
18111            stream_started: Arc::clone(&stream_started),
18112            stream_dropped: Arc::clone(&stream_dropped),
18113            committed_after_drop: Arc::clone(&committed_after_drop),
18114        });
18115        let router = Arc::new(RoutingAfterProviderStart {
18116            provider_started: stream_started,
18117        });
18118        let yaml = r#"
18119name: BufferedCancellationAgent
18120system_prompt: "Hide stale streamed output."
18121llm:
18122  default: default
18123  router: router
18124streaming:
18125  enabled: true
18126  buffer_size: 8
18127runtime:
18128  optimization:
18129    enabled: true
18130    max_speculative_llm_calls_per_turn: 2
18131    speculative_state_transitions: true
18132    streaming_policy: buffer_until_routing_done
18133    max_parallel_runtime_tasks: 2
18134states:
18135  initial: triage
18136  states:
18137    triage:
18138      prompt: "Triage state."
18139      transitions:
18140        - to: technical
18141          when: "The request needs technical support"
18142          timing: parallel
18143    technical:
18144      prompt: "Technical state."
18145"#;
18146        let agent = AgentBuilder::from_yaml(yaml)
18147            .unwrap()
18148            .llm_alias("default", default)
18149            .llm_alias("router", router)
18150            .build()
18151            .unwrap();
18152
18153        let content = tokio::time::timeout(std::time::Duration::from_secs(2), async {
18154            let mut stream = agent
18155                .chat_stream("AUTH-17 needs technical help.")
18156                .await
18157                .unwrap();
18158            let mut content = String::new();
18159            while let Some(chunk) = stream.next().await {
18160                match chunk {
18161                    StreamChunk::Content { text } => content.push_str(&text),
18162                    StreamChunk::Done {} => break,
18163                    StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
18164                    _ => {}
18165                }
18166            }
18167            content
18168        })
18169        .await
18170        .expect("redispatch must not wait on the stale streaming future");
18171
18172        assert_eq!(content, "Committed technical response.");
18173        assert_eq!(agent.current_state().as_deref(), Some("technical"));
18174        assert!(stream_dropped.load(Ordering::SeqCst));
18175        assert!(committed_after_drop.load(Ordering::SeqCst));
18176    }
18177
18178    #[tokio::test]
18179    async fn buffered_transition_drops_established_stream_before_redispatch() {
18180        use futures::StreamExt;
18181
18182        let stream_started = Arc::new(tokio::sync::Notify::new());
18183        let stream_dropped = Arc::new(AtomicBool::new(false));
18184        let stream_dropped_notify = Arc::new(tokio::sync::Notify::new());
18185        let committed_after_drop = Arc::new(AtomicBool::new(false));
18186        let default = Arc::new(EstablishedStreamProvider {
18187            stream_started: Arc::clone(&stream_started),
18188            stream_dropped: Arc::clone(&stream_dropped),
18189            stream_dropped_notify,
18190            committed_after_drop: Arc::clone(&committed_after_drop),
18191        });
18192        let router = Arc::new(RoutingAfterProviderStart {
18193            provider_started: stream_started,
18194        });
18195        let yaml = r#"
18196name: EstablishedStreamCancellationAgent
18197system_prompt: "Hide stale streamed output."
18198llm:
18199  default: default
18200  router: router
18201streaming:
18202  enabled: true
18203  buffer_size: 8
18204runtime:
18205  optimization:
18206    enabled: true
18207    max_speculative_llm_calls_per_turn: 2
18208    speculative_state_transitions: true
18209    streaming_policy: buffer_until_routing_done
18210    max_parallel_runtime_tasks: 2
18211states:
18212  initial: triage
18213  states:
18214    triage:
18215      prompt: "Triage state."
18216      transitions:
18217        - to: technical
18218          when: "The request needs technical support"
18219          timing: parallel
18220    technical:
18221      prompt: "Technical state."
18222"#;
18223        let agent = AgentBuilder::from_yaml(yaml)
18224            .unwrap()
18225            .llm_alias("default", default)
18226            .llm_alias("router", router)
18227            .build()
18228            .unwrap();
18229
18230        let content = tokio::time::timeout(std::time::Duration::from_secs(2), async {
18231            let mut stream = agent
18232                .chat_stream("AUTH-17 needs technical help.")
18233                .await
18234                .unwrap();
18235            let mut content = String::new();
18236            while let Some(chunk) = stream.next().await {
18237                match chunk {
18238                    StreamChunk::Content { text } => content.push_str(&text),
18239                    StreamChunk::Done {} => break,
18240                    StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
18241                    _ => {}
18242                }
18243            }
18244            content
18245        })
18246        .await
18247        .expect("redispatch must wait for the established stale stream to be dropped");
18248
18249        assert_eq!(content, "Committed technical response.");
18250        assert_eq!(agent.current_state().as_deref(), Some("technical"));
18251        assert!(stream_dropped.load(Ordering::SeqCst));
18252        assert!(committed_after_drop.load(Ordering::SeqCst));
18253    }
18254
18255    #[tokio::test]
18256    async fn test_buffered_streaming_transition_reservation_falls_back() {
18257        use futures::StreamExt;
18258
18259        let mock = mock_with_responses(vec![
18260            "Serial streaming response",
18261            "Serial streaming response",
18262        ]);
18263        let router_mock = mock_with_response("1");
18264        let router_counter = router_mock.clone();
18265        let yaml = r#"
18266name: BufferedReservationFallbackAgent
18267system_prompt: "Stream normally if speculative routing cannot be evaluated."
18268llm:
18269  default: default
18270  router: router
18271observability:
18272  enabled: true
18273  export:
18274    write_raw_events: true
18275streaming:
18276  enabled: true
18277  buffer_size: 8
18278runtime:
18279  optimization:
18280    enabled: true
18281    max_speculative_llm_calls_per_turn: 1
18282    speculative_state_transitions: true
18283    streaming_policy: buffer_until_routing_done
18284    max_parallel_runtime_tasks: 2
18285states:
18286  initial: triage
18287  states:
18288    triage:
18289      prompt: "Triage state."
18290      transitions:
18291        - to: billing
18292          guard:
18293            context:
18294              route:
18295                eq: billing
18296          when: "User asks about billing"
18297          timing: parallel
18298    billing:
18299      prompt: "Billing state."
18300"#;
18301        let agent = AgentBuilder::from_yaml(yaml)
18302            .unwrap()
18303            .llm_alias("default", Arc::new(mock))
18304            .llm_alias("router", Arc::new(router_mock))
18305            .build()
18306            .unwrap();
18307
18308        let mut stream = agent.chat_stream("hello").await.unwrap();
18309        let mut content = String::new();
18310        let mut error = None;
18311        while let Some(chunk) = stream.next().await {
18312            match chunk {
18313                StreamChunk::Content { text } => content.push_str(&text),
18314                StreamChunk::Error { message } => error = Some(message),
18315                StreamChunk::Done {} => break,
18316                _ => {}
18317            }
18318        }
18319
18320        assert_eq!(error, None);
18321        assert_eq!(content, "Serial streaming response");
18322        assert_eq!(router_counter.call_count(), 0);
18323        let events = agent.observability().unwrap().raw_events();
18324        assert!(events.iter().any(|event| {
18325            event.dimensions.get("branch_status") == Some(&"cancelled".to_string())
18326                && event.dimensions.get("commit_behavior")
18327                    == Some(&"transition_decision".to_string())
18328        }));
18329    }
18330
18331    #[tokio::test]
18332    async fn test_blocking_error_cleanup_resets_root_turn_for_next_chat() {
18333        let mut mock = mock_with_response("Recovered response");
18334        mock.set_error("boom");
18335        let mut handle = mock.clone();
18336        let agent = AgentBuilder::new()
18337            .system_prompt("You are helpful.")
18338            .llm(Arc::new(mock))
18339            .build()
18340            .unwrap();
18341
18342        assert!(agent.chat("first").await.is_err());
18343        handle.clear_error();
18344        let response = agent.chat("second").await.unwrap();
18345
18346        assert_eq!(response.content, "Recovered response");
18347        let messages = agent.memory.get_messages(None).await.unwrap();
18348        let user_count = messages
18349            .iter()
18350            .filter(|message| message.role == ai_agents_core::Role::User)
18351            .count();
18352        assert_eq!(user_count, 2);
18353    }
18354
18355    #[tokio::test]
18356    async fn test_streaming_error_cleanup_resets_root_turn_for_next_chat() {
18357        use futures::StreamExt;
18358
18359        let mut mock = mock_with_response("Recovered response");
18360        mock.set_error("stream boom");
18361        let mut handle = mock.clone();
18362        let agent = AgentBuilder::new()
18363            .system_prompt("You are helpful.")
18364            .llm(Arc::new(mock))
18365            .build()
18366            .unwrap();
18367
18368        let mut stream = agent.chat_stream("first").await.unwrap();
18369        let mut saw_error = false;
18370        while let Some(chunk) = stream.next().await {
18371            if matches!(chunk, StreamChunk::Error { .. }) {
18372                saw_error = true;
18373            }
18374        }
18375        assert!(saw_error);
18376
18377        handle.clear_error();
18378        let response = agent.chat("second").await.unwrap();
18379
18380        assert_eq!(response.content, "Recovered response");
18381        let messages = agent.memory.get_messages(None).await.unwrap();
18382        let user_count = messages
18383            .iter()
18384            .filter(|message| message.role == ai_agents_core::Role::User)
18385            .count();
18386        assert_eq!(user_count, 2);
18387    }
18388
18389    #[tokio::test]
18390    async fn test_buffered_streaming_route_miss_releases_buffer_limit() {
18391        use futures::StreamExt;
18392
18393        let mut mock = mock_with_response("one two three");
18394        mock.set_latency(10);
18395        let yaml = r#"
18396name: BufferedMissAgent
18397system_prompt: "You stream safely."
18398llm:
18399  default: default
18400streaming:
18401  enabled: true
18402  buffer_size: 1
18403runtime:
18404  optimization:
18405    enabled: true
18406    max_speculative_llm_calls_per_turn: 2
18407    speculative_state_transitions: true
18408    streaming_policy: buffer_until_routing_done
18409    max_parallel_runtime_tasks: 2
18410states:
18411  initial: triage
18412  states:
18413    triage:
18414      prompt: "Answer from triage."
18415      transitions:
18416        - to: billing
18417          guard:
18418            context:
18419              route:
18420                eq: billing
18421          timing: parallel
18422    billing:
18423      prompt: "Billing state."
18424"#;
18425        let agent = AgentBuilder::from_yaml(yaml)
18426            .unwrap()
18427            .llm_alias("default", Arc::new(mock))
18428            .build()
18429            .unwrap();
18430
18431        let mut stream = agent.chat_stream("hello").await.unwrap();
18432        let mut content = String::new();
18433        let mut error = None;
18434        while let Some(chunk) = stream.next().await {
18435            match chunk {
18436                StreamChunk::Content { text } => content.push_str(&text),
18437                StreamChunk::Error { message } => error = Some(message),
18438                StreamChunk::Done {} => break,
18439                _ => {}
18440            }
18441        }
18442
18443        assert_eq!(error, None);
18444        assert_eq!(content, "one two three");
18445    }
18446
18447    #[tokio::test]
18448    async fn test_buffered_streaming_main_failure_finalizes_branch() {
18449        use futures::StreamExt;
18450
18451        let mock = mock_with_response("one two");
18452        let mut router_mock = mock_with_response("0");
18453        router_mock.set_latency(50);
18454        let yaml = r#"
18455name: BufferedFailureAgent
18456system_prompt: "You stream safely."
18457llm:
18458  default: default
18459  router: router
18460observability:
18461  enabled: true
18462  export:
18463    write_raw_events: true
18464streaming:
18465  enabled: true
18466  buffer_size: 1
18467runtime:
18468  optimization:
18469    enabled: true
18470    max_speculative_llm_calls_per_turn: 2
18471    speculative_state_transitions: true
18472    streaming_policy: buffer_until_routing_done
18473    max_parallel_runtime_tasks: 2
18474states:
18475  initial: triage
18476  states:
18477    triage:
18478      prompt: "Ask for the category."
18479      transitions:
18480        - to: billing
18481          when: "User asks about billing"
18482          timing: parallel
18483    billing:
18484      prompt: "Billing state."
18485"#;
18486        let agent = AgentBuilder::from_yaml(yaml)
18487            .unwrap()
18488            .llm_alias("default", Arc::new(mock))
18489            .llm_alias("router", Arc::new(router_mock))
18490            .build()
18491            .unwrap();
18492
18493        let mut stream = agent.chat_stream("hello").await.unwrap();
18494        let mut error = String::new();
18495        while let Some(chunk) = stream.next().await {
18496            if let StreamChunk::Error { message } = chunk {
18497                error = message;
18498            }
18499        }
18500
18501        assert!(
18502            error.contains("stream buffer filled"),
18503            "unexpected stream error: {}",
18504            error
18505        );
18506        let events = agent.observability().unwrap().raw_events();
18507        assert!(events.iter().any(|event| {
18508            event.dimensions.get("branch_status") == Some(&"failed".to_string())
18509                && event.dimensions.get("commit_behavior") == Some(&"final_response".to_string())
18510                && event.dimensions.get("optimization")
18511                    == Some(&"buffered_streaming_routing".to_string())
18512        }));
18513    }
18514
18515    #[tokio::test]
18516    async fn test_streaming_preflight_does_not_emit_old_state_content() {
18517        use futures::StreamExt;
18518
18519        let mock = mock_with_response("Billing streamed response");
18520        let yaml = r#"
18521name: StreamingOptimizedAgent
18522system_prompt: "You route before streaming."
18523runtime:
18524  optimization:
18525    enabled: true
18526    pre_response_deterministic_transitions: true
18527streaming:
18528  enabled: true
18529states:
18530  initial: greeting
18531  states:
18532    greeting:
18533      prompt: "OLD_STATE_SENTINEL"
18534      transitions:
18535        - to: billing
18536          guard:
18537            context:
18538              topic:
18539                eq: billing
18540          timing: pre_response
18541    billing:
18542      prompt: "Billing state."
18543"#;
18544        let agent = AgentBuilder::from_yaml(yaml)
18545            .unwrap()
18546            .llm(Arc::new(mock))
18547            .build()
18548            .unwrap();
18549        agent
18550            .set_context("topic", serde_json::json!("billing"))
18551            .unwrap();
18552
18553        let mut stream = agent.chat_stream("billing please").await.unwrap();
18554        let mut content = String::new();
18555        while let Some(chunk) = stream.next().await {
18556            match chunk {
18557                StreamChunk::Content { text } => content.push_str(&text),
18558                StreamChunk::Error { message } => panic!("stream error: {}", message),
18559                StreamChunk::Done {} => break,
18560                _ => {}
18561            }
18562        }
18563
18564        assert_eq!(agent.current_state().as_deref(), Some("billing"));
18565        assert!(content.contains("Billing streamed response"));
18566        assert!(!content.contains("OLD_STATE_SENTINEL"));
18567    }
18568
18569    // State machine transitions
18570    #[tokio::test]
18571    async fn test_integration_state_machine_basic() {
18572        let yaml = r#"
18573name: StateAgent
18574system_prompt: "You are a support agent."
18575states:
18576  initial: greeting
18577  states:
18578    greeting:
18579      prompt: "Welcome the user warmly."
18580      transitions:
18581        - to: support
18582          when: "User needs help"
18583          auto: true
18584    support:
18585      prompt: "Help solve the user's problem."
18586"#;
18587        let mock = mock_with_responses(vec![
18588            "Welcome! How can I help?", // greeting response
18589            "1",                        // transition evaluator picks first (index 0)
18590            "I'll help you with that.", // support response
18591        ]);
18592        let builder = AgentBuilder::from_yaml(yaml).unwrap();
18593        let agent = builder.llm(Arc::new(mock)).build().unwrap();
18594
18595        assert_eq!(agent.current_state(), Some("greeting".to_string()));
18596        let _ = agent.chat("I need help").await.unwrap();
18597        // After transition evaluation, state may or may not have changed
18598        // depending on mock evaluator response - the key is that it doesn't crash
18599    }
18600
18601    // State on_enter/on_exit actions
18602    #[tokio::test]
18603    async fn test_integration_state_on_enter_set_context() {
18604        let yaml = r#"
18605name: ActionAgent
18606system_prompt: "You are helpful."
18607states:
18608  initial: step1
18609  states:
18610    step1:
18611      prompt: "Step 1"
18612      on_exit:
18613        - set_context:
18614            step1_exited: true
18615      transitions:
18616        - to: step2
18617          when: "always"
18618          auto: true
18619    step2:
18620      prompt: "Step 2"
18621      on_enter:
18622        - set_context:
18623            step2_entered: true
18624"#;
18625        // The transition evaluator will pick the first transition (index 0)
18626        let mock = mock_with_responses(vec![
18627            "Processing step 1.",
18628            "0", // transition evaluator response: select first transition
18629        ]);
18630        let builder = AgentBuilder::from_yaml(yaml).unwrap();
18631        let agent = builder.llm(Arc::new(mock)).build().unwrap();
18632
18633        assert_eq!(agent.current_state(), Some("step1".to_string()));
18634
18635        // Manually transition to test on_enter/on_exit
18636        agent.transition_to("step2").await.unwrap();
18637
18638        assert_eq!(agent.current_state(), Some("step2".to_string()));
18639
18640        // Verify context was set by on_exit and on_enter actions
18641        let ctx = agent.get_context();
18642        assert_eq!(ctx.get("step1_exited"), Some(&serde_json::json!(true)));
18643        assert_eq!(ctx.get("step2_entered"), Some(&serde_json::json!(true)));
18644    }
18645
18646    #[tokio::test]
18647    async fn state_action_tool_preserves_source_in_stored_record() {
18648        let yaml = r#"
18649name: StateActionToolAgent
18650system_prompt: "You are helpful."
18651tools:
18652  - context_echo
18653states:
18654  initial: idle
18655  states:
18656    idle:
18657      prompt: "Idle"
18658    active:
18659      prompt: "Active"
18660      on_enter:
18661        - set_context:
18662            action_started: true
18663        - tool: context_echo
18664          args: {}
18665"#;
18666        let agent = AgentBuilder::from_yaml(yaml)
18667            .unwrap()
18668            .llm(Arc::new(mock_with_response("unused")))
18669            .tool(Arc::new(ContextEchoTool))
18670            .build()
18671            .unwrap();
18672
18673        agent.transition_to("active").await.unwrap();
18674
18675        let record: ToolExecutionRecord = serde_json::from_value(
18676            agent
18677                .get_context()
18678                .get("last_tool_record")
18679                .cloned()
18680                .expect("successful state action must store its execution record"),
18681        )
18682        .unwrap();
18683        assert!(record.executed);
18684        assert!(record.success);
18685        assert_eq!(record.canonical_id, "context_echo");
18686        assert!(matches!(
18687            &record.source,
18688            ToolCallSource::StateAction {
18689                state: Some(state),
18690                action_index: 1,
18691            } if state == "active"
18692        ));
18693    }
18694
18695    #[tokio::test]
18696    async fn test_ordinary_transition_uses_on_enter_then_on_reenter() {
18697        let yaml = r#"
18698name: OrdinaryLifecycleAgent
18699system_prompt: "You are helpful."
18700states:
18701  initial: intake
18702  regenerate_on_transition: false
18703  states:
18704    intake:
18705      prompt: "Intake"
18706      transitions:
18707        - to: drafting
18708          guard:
18709            context:
18710              route:
18711                eq: drafting
18712    drafting:
18713      prompt: "Drafting"
18714      on_enter:
18715        - set_context:
18716            draft_version: 1
18717      on_reenter:
18718        - set_context:
18719            draft_version: 2
18720      transitions:
18721        - to: review
18722          guard:
18723            context:
18724              route:
18725                eq: review
18726    review:
18727      prompt: "Review"
18728      on_enter:
18729        - set_context:
18730            review_entry: first
18731      transitions:
18732        - to: drafting
18733          guard:
18734            context:
18735              route:
18736                eq: drafting
18737"#;
18738        let agent = AgentBuilder::from_yaml(yaml)
18739            .unwrap()
18740            .llm(Arc::new(mock_with_responses(vec![
18741                "Intake response",
18742                "Draft response",
18743                "Review response",
18744            ])))
18745            .build()
18746            .unwrap();
18747
18748        agent
18749            .set_context("route", serde_json::json!("drafting"))
18750            .unwrap();
18751        agent.chat("Start a draft").await.unwrap();
18752        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
18753        assert_eq!(
18754            agent.get_context().get("draft_version"),
18755            Some(&serde_json::json!(1))
18756        );
18757
18758        agent
18759            .set_context("route", serde_json::json!("review"))
18760            .unwrap();
18761        agent.chat("Review this").await.unwrap();
18762        assert_eq!(agent.current_state().as_deref(), Some("review"));
18763        assert_eq!(
18764            agent.get_context().get("review_entry"),
18765            Some(&serde_json::json!("first"))
18766        );
18767
18768        agent
18769            .set_context("route", serde_json::json!("drafting"))
18770            .unwrap();
18771        agent.chat("Revise this").await.unwrap();
18772        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
18773        assert_eq!(
18774            agent.get_context().get("draft_version"),
18775            Some(&serde_json::json!(2))
18776        );
18777    }
18778
18779    #[tokio::test]
18780    async fn test_manual_transition_uses_on_enter_then_on_reenter() {
18781        let yaml = r#"
18782name: ManualLifecycleAgent
18783system_prompt: "You are helpful."
18784states:
18785  initial: intake
18786  states:
18787    intake:
18788      prompt: "Intake"
18789    drafting:
18790      prompt: "Drafting"
18791      on_enter:
18792        - set_context:
18793            draft_version: 1
18794      on_reenter:
18795        - set_context:
18796            draft_version: 2
18797    review:
18798      prompt: "Review"
18799"#;
18800        let agent = AgentBuilder::from_yaml(yaml)
18801            .unwrap()
18802            .llm(Arc::new(mock_with_response("unused")))
18803            .build()
18804            .unwrap();
18805
18806        assert!(!agent.get_context().contains_key("draft_version"));
18807        agent.transition_to("drafting").await.unwrap();
18808        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
18809        assert_eq!(
18810            agent.get_context().get("draft_version"),
18811            Some(&serde_json::json!(1))
18812        );
18813
18814        agent.transition_to("review").await.unwrap();
18815        agent.transition_to("drafting").await.unwrap();
18816        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
18817        assert_eq!(
18818            agent.get_context().get("draft_version"),
18819            Some(&serde_json::json!(2))
18820        );
18821    }
18822
18823    #[tokio::test]
18824    async fn test_timeout_transition_uses_on_enter_then_on_reenter() {
18825        let yaml = r#"
18826name: TimeoutLifecycleAgent
18827system_prompt: "You are helpful."
18828states:
18829  initial: intake
18830  regenerate_on_transition: false
18831  states:
18832    intake:
18833      prompt: "Intake"
18834      max_turns: 1
18835      timeout_to: drafting
18836    drafting:
18837      prompt: "Drafting"
18838      max_turns: 1
18839      timeout_to: review
18840      on_enter:
18841        - set_context:
18842            draft_version: 1
18843      on_reenter:
18844        - set_context:
18845            draft_version: 2
18846    review:
18847      prompt: "Review"
18848      max_turns: 1
18849      timeout_to: drafting
18850      on_enter:
18851        - set_context:
18852            review_entry: first
18853"#;
18854        let agent = AgentBuilder::from_yaml(yaml)
18855            .unwrap()
18856            .llm(Arc::new(mock_with_responses(vec![
18857                "Intake",
18858                "First draft",
18859                "Review",
18860                "Revised draft",
18861            ])))
18862            .build()
18863            .unwrap();
18864
18865        agent.chat("First turn").await.unwrap();
18866        assert_eq!(agent.current_state().as_deref(), Some("intake"));
18867        assert!(!agent.get_context().contains_key("draft_version"));
18868
18869        agent.chat("Second turn").await.unwrap();
18870        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
18871        assert_eq!(
18872            agent.get_context().get("draft_version"),
18873            Some(&serde_json::json!(1))
18874        );
18875
18876        agent.chat("Third turn").await.unwrap();
18877        assert_eq!(agent.current_state().as_deref(), Some("review"));
18878        assert_eq!(
18879            agent.get_context().get("review_entry"),
18880            Some(&serde_json::json!("first"))
18881        );
18882
18883        agent.chat("Fourth turn").await.unwrap();
18884        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
18885        assert_eq!(
18886            agent.get_context().get("draft_version"),
18887            Some(&serde_json::json!(2))
18888        );
18889    }
18890
18891    // Process pipeline transforms input
18892    #[tokio::test]
18893    async fn test_integration_process_normalize() {
18894        let yaml = r#"
18895name: ProcessAgent
18896system_prompt: "You are helpful."
18897process:
18898  input:
18899    - type: normalize
18900      config:
18901        trim: true
18902        collapse_whitespace: true
18903"#;
18904        let mock = mock_with_response("Got your message.");
18905        let builder = AgentBuilder::from_yaml(yaml).unwrap();
18906        let agent = builder.llm(Arc::new(mock.clone())).build().unwrap();
18907
18908        let _ = agent.chat("  hello   world  ").await.unwrap();
18909
18910        // Verify the LLM received the normalized input (trimmed + collapsed whitespace)
18911        let history = mock.call_history();
18912        assert!(!history.is_empty());
18913        // The user message in LLM call should be normalized
18914        let last_call = history.last().unwrap();
18915        let user_msg = last_call
18916            .messages
18917            .iter()
18918            .find(|m| m.role == ai_agents_core::Role::User)
18919            .unwrap();
18920        assert_eq!(user_msg.content, "hello world");
18921    }
18922
18923    // ═══════════════════════════════════════════════════════════
18924    // Integration Test 2.1.7: Memory compression triggers
18925    // ═══════════════════════════════════════════════════════════
18926    #[tokio::test]
18927    async fn test_integration_memory_compression() {
18928        let yaml = r#"
18929name: MemoryAgent
18930system_prompt: "You are helpful."
18931memory:
18932  type: compacting
18933  max_messages: 100
18934  compress_threshold: 5
18935  max_recent_messages: 3
18936  summarize_batch_size: 2
18937"#;
18938        // Provide enough responses for compression to trigger
18939        let responses: Vec<&str> = (0..8).map(|_| "Response from assistant.").collect();
18940        let mock = mock_with_responses(responses);
18941        let builder = AgentBuilder::from_yaml(yaml).unwrap();
18942        let agent = builder.llm(Arc::new(mock)).build().unwrap();
18943
18944        // Send enough messages to trigger compression
18945        for i in 0..6 {
18946            let _ = agent.chat(&format!("Message {}", i)).await.unwrap();
18947        }
18948
18949        // Memory should have compressed - verify it didn't crash
18950        // and that messages are bounded
18951        let messages = agent.memory.get_messages(None).await.unwrap();
18952        // With compress_threshold=5 and max_recent_messages=3,
18953        // after 6 turns (12 messages), compression should have run
18954        assert!(messages.len() <= 12); // At most all messages if no compression, fewer if compressed
18955    }
18956
18957    // YAML with multiple LLMs
18958    #[tokio::test]
18959    async fn test_integration_multi_llm_registry() {
18960        let mut mock_default = MockLLMProvider::new("default");
18961        mock_default.set_response("Default LLM response.");
18962        let mut mock_router = MockLLMProvider::new("router");
18963        mock_router.set_response("Router response.");
18964
18965        let agent = AgentBuilder::new()
18966            .system_prompt("You are helpful.")
18967            .llm_alias("default", Arc::new(mock_default))
18968            .llm_alias("router", Arc::new(mock_router))
18969            .build()
18970            .unwrap();
18971
18972        let response = agent.chat("Hello").await.unwrap();
18973        assert_eq!(response.content, "Default LLM response.");
18974    }
18975
18976    // Agent reset clears state
18977    #[tokio::test]
18978    async fn test_integration_agent_reset() {
18979        let mock = mock_with_responses(vec!["Hello!", "Hello again!"]);
18980        let agent = AgentBuilder::new()
18981            .system_prompt("You are helpful.")
18982            .llm(Arc::new(mock))
18983            .build()
18984            .unwrap();
18985
18986        let _ = agent.chat("Hi").await.unwrap();
18987        let messages = agent.memory.get_messages(None).await.unwrap();
18988        assert_eq!(messages.len(), 2); // user + assistant
18989
18990        agent.reset().await.unwrap();
18991        let messages = agent.memory.get_messages(None).await.unwrap();
18992        assert_eq!(messages.len(), 0);
18993    }
18994
18995    // Process pipeline rejects input
18996    #[tokio::test]
18997    async fn test_integration_process_validate_reject() {
18998        use ai_agents_process::{ProcessConfig, ProcessProcessor};
18999
19000        let validate_config = ai_agents_process::ValidateStage {
19001            id: Some("length_check".to_string()),
19002            condition: None,
19003            config: ai_agents_process::ValidateConfig {
19004                rules: vec![ai_agents_process::ValidationRule::MinLength {
19005                    min_length: 10,
19006                    on_fail: ai_agents_process::ValidationAction {
19007                        action: ai_agents_process::ValidationActionType::Reject,
19008                        message: None,
19009                    },
19010                }],
19011                ..Default::default()
19012            },
19013        };
19014        let process_config = ProcessConfig {
19015            input: vec![ai_agents_process::ProcessStage::Validate(validate_config)],
19016            ..Default::default()
19017        };
19018        let processor = ProcessProcessor::new(process_config);
19019
19020        let mock = mock_with_response("Should not reach here.");
19021        let agent = AgentBuilder::new()
19022            .system_prompt("You are helpful.")
19023            .llm(Arc::new(mock))
19024            .process_processor(processor)
19025            .build()
19026            .unwrap();
19027
19028        let response = agent.chat("Hi").await.unwrap();
19029        // Rejected input should produce a rejection response, not call LLM
19030        assert!(
19031            response.content.contains("rejected")
19032                || response.content.contains("Input rejected")
19033                || response.content.contains("too short")
19034                || response.content.contains("Too short")
19035                || response.content.len() < 50, // rejection message is typically short
19036            "Expected rejection response, got: {}",
19037            response.content
19038        );
19039    }
19040
19041    // LLM fallback: primary fails, fallback LLM responds
19042    #[tokio::test]
19043    async fn test_llm_fallback_on_failure() {
19044        use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
19045
19046        let mut primary = MockLLMProvider::new("primary");
19047        primary.set_error("Primary LLM is unavailable");
19048
19049        let mut fallback = MockLLMProvider::new("fallback");
19050        fallback.set_response("Fallback response works!");
19051
19052        let agent = AgentBuilder::new()
19053            .system_prompt("You are helpful.")
19054            .llm_alias("default", Arc::new(primary))
19055            .llm_alias("backup", Arc::new(fallback))
19056            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
19057                llm: LLMRecoveryConfig {
19058                    on_failure: LLMFailureAction::FallbackLlm {
19059                        fallback_llm: "backup".to_string(),
19060                    },
19061                    ..Default::default()
19062                },
19063                ..Default::default()
19064            }))
19065            .build()
19066            .unwrap();
19067
19068        let response = agent.chat("Hello").await.unwrap();
19069        assert!(
19070            response.content.contains("Fallback response"),
19071            "Expected fallback response, got: {}",
19072            response.content
19073        );
19074    }
19075
19076    // LLM fallback: primary fails, static message returned
19077    #[tokio::test]
19078    async fn test_llm_fallback_response_static_message() {
19079        use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
19080
19081        let mut primary = MockLLMProvider::new("primary");
19082        primary.set_error("Primary LLM is unavailable");
19083
19084        let agent = AgentBuilder::new()
19085            .system_prompt("You are helpful.")
19086            .llm(Arc::new(primary))
19087            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
19088                llm: LLMRecoveryConfig {
19089                    on_failure: LLMFailureAction::FallbackResponse {
19090                        message: "I am temporarily unavailable. Please try again later."
19091                            .to_string(),
19092                    },
19093                    ..Default::default()
19094                },
19095                ..Default::default()
19096            }))
19097            .build()
19098            .unwrap();
19099
19100        let response = agent.chat("Hello").await.unwrap();
19101        assert!(
19102            response.content.contains("temporarily unavailable"),
19103            "Expected static fallback message, got: {}",
19104            response.content
19105        );
19106    }
19107
19108    // Tool skip: tool fails, on_failure: skip absorbs the error
19109    #[tokio::test]
19110    async fn test_tool_failure_skip() {
19111        use ai_agents_recovery::{
19112            ErrorRecoveryConfig, ToolFailureAction, ToolRecoveryConfig, ToolRetryConfig,
19113        };
19114
19115        // LLM requests a nonexistent tool, then responds after seeing the skip result
19116        let mock = mock_with_responses(vec![
19117            r#"I'll use the nonexistent tool.
19118[TOOL_CALL: {"name": "nonexistent_tool", "arguments": {}}]"#,
19119            "The tool was unavailable, but I can still help you.",
19120        ]);
19121
19122        let agent = AgentBuilder::new()
19123            .system_prompt("You are helpful.")
19124            .llm(Arc::new(mock))
19125            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
19126                tools: ToolRecoveryConfig {
19127                    default: ToolRetryConfig {
19128                        max_retries: 0,
19129                        timeout_ms: None,
19130                        on_failure: ToolFailureAction::Skip,
19131                    },
19132                    ..Default::default()
19133                },
19134                ..Default::default()
19135            }))
19136            .build()
19137            .unwrap();
19138
19139        // The tool will fail (not found), but on_failure: skip absorbs the error
19140        let response = agent.chat("Use the nonexistent tool").await;
19141        assert!(
19142            response.is_ok(),
19143            "Expected Ok with skip policy, got: {:?}",
19144            response
19145        );
19146    }
19147}