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::{Duration, Instant};
11use tracing::{debug, error, info, instrument, warn};
12
13const DISAMBIGUATION_STATE_GENERATION_KEY: &str = "_runtime.disambiguation_state_generation";
14const MAX_TOOL_FALLBACK_HOPS: usize = 16;
15
16/// Keeps the strong identity of one non-reentrant root-turn gate.
17pub(crate) type RootTurnGate = Arc<tokio::sync::Mutex<()>>;
18
19/// Owns one immutable root-gate ancestry snapshot that can cross polls and spawned tasks by cloning one `Arc`.
20pub(crate) type RootTurnGateIdentityStack = Arc<[RootTurnGate]>;
21
22tokio::task_local! {
23    static RUNTIME_GATE_IDENTITY_STACK: RootTurnGateIdentityStack;
24}
25
26/// Captures the immutable gate ownership chain for propagation and returns an empty chain outside any root-turn scope.
27pub(crate) fn current_runtime_gate_identity_stack() -> RootTurnGateIdentityStack {
28    RUNTIME_GATE_IDENTITY_STACK
29        .try_with(Arc::clone)
30        .unwrap_or_default()
31}
32
33/// Polls a future with an explicit gate ownership chain while cloning only its immutable owner instead of rebuilding the ancestry.
34pub(crate) async fn scope_runtime_gate_identity_stack<F, T>(
35    identity_stack: &RootTurnGateIdentityStack,
36    future: F,
37) -> T
38where
39    F: Future<Output = T>,
40{
41    RUNTIME_GATE_IDENTITY_STACK
42        .scope(Arc::clone(identity_stack), future)
43        .await
44}
45
46/// Shared lock table used to serialize side-effecting tool calls by canonical resource.
47pub(crate) type ToolResourceLocks = Arc<RwLock<HashMap<String, Weak<tokio::sync::Mutex<()>>>>>;
48
49//
50// Owns every lock acquired for one tool call and removes dead weak entries after release.
51//
52struct ToolResourceGuards {
53    guards: Vec<tokio::sync::OwnedMutexGuard<()>>,
54    locks: ToolResourceLocks,
55}
56
57//
58// Couples one acquired root-turn gate with the immutable ancestry owner that must remain active through hooks, orchestration, and stream polling.
59//
60struct RootTurnAdmission {
61    guard: tokio::sync::OwnedMutexGuard<()>,
62    identity_stack: RootTurnGateIdentityStack,
63}
64
65#[derive(Clone)]
66struct StoredSessionRestore {
67    snapshot: AgentSnapshot,
68    metadata: Option<ai_agents_core::SessionMetadata>,
69}
70
71struct RuntimeSessionRestorePoint {
72    snapshot: AgentSnapshot,
73    metadata: ai_agents_core::SessionMetadata,
74    actor_id: Option<String>,
75    session_id: Option<String>,
76}
77
78impl Drop for ToolResourceGuards {
79    fn drop(&mut self) {
80        self.guards.clear();
81        self.locks.write().retain(|_, lock| lock.strong_count() > 0);
82    }
83}
84
85//
86// Captures runtime controls under one read guard so policy and scope belong to the same generation.
87//
88#[derive(Clone)]
89struct RuntimeSafetySnapshot {
90    version: u64,
91    emergency_deny: bool,
92    tool_security: ToolSecurityEngine,
93    tool_scope_override: Option<Vec<String>>,
94}
95
96//
97// Records the generations used by the final authorization decision.
98//
99#[derive(Clone, Copy)]
100struct ToolDecisionVersions {
101    policy: u64,
102    registry: u64,
103    runtime_control: u64,
104    state: Option<u64>,
105}
106
107//
108// Carries canonical fallback ancestry across separate shared-executor requests so aliases cannot hide cycles and an acyclic configuration cannot recurse without a fixed bound.
109//
110#[derive(Clone, Debug, Default)]
111struct ToolFallbackState {
112    visited_canonical_ids: Vec<String>,
113}
114
115impl ToolFallbackState {
116    //
117    // Rejects before the start hook, approval, locks, or side effects when the resolved target repeats or would exceed the stable fallback-hop bound; ordinary terminal completion evidence still runs.
118    //
119    fn rejection_reason(&self, canonical_id: &str) -> Option<String> {
120        if self
121            .visited_canonical_ids
122            .iter()
123            .any(|visited| visited == canonical_id)
124        {
125            return Some(format!(
126                "Tool fallback cycle detected at '{canonical_id}' after [{}]",
127                self.visited_canonical_ids.join(" -> ")
128            ));
129        }
130        if self.visited_canonical_ids.len() > MAX_TOOL_FALLBACK_HOPS {
131            return Some(format!(
132                "Tool fallback chain exceeds the maximum of {MAX_TOOL_FALLBACK_HOPS} hops"
133            ));
134        }
135        None
136    }
137
138    //
139    // Extends ancestry only after the current canonical request passes cycle and hop admission.
140    //
141    fn with_current(mut self, canonical_id: String) -> Self {
142        self.visited_canonical_ids.push(canonical_id);
143        self
144    }
145
146    //
147    // Rejects canonical drift after an async approval boundary because a changed final target did not pass the initial fallback admission and may re-enter an ancestor through a refreshed alias or provider.
148    //
149    fn final_rejection_reason(
150        &self,
151        admitted_canonical_id: &str,
152        final_canonical_id: &str,
153    ) -> Option<String> {
154        if admitted_canonical_id == final_canonical_id {
155            return None;
156        }
157        if self
158            .visited_canonical_ids
159            .iter()
160            .any(|visited| visited == final_canonical_id)
161        {
162            return Some(format!(
163                "Tool fallback cycle detected after final resolution changed '{admitted_canonical_id}' to '{final_canonical_id}'"
164            ));
165        }
166        Some(format!(
167            "Tool canonical target changed after initial admission from '{admitted_canonical_id}' to '{final_canonical_id}'"
168        ))
169    }
170}
171
172//
173// Carries one checked timeout representation into every retry attempt so deadline and timer cannot diverge.
174//
175#[derive(Clone, Copy, Debug)]
176struct ValidatedToolTimeout {
177    timer: Duration,
178    deadline_delta: chrono::Duration,
179}
180
181//
182// Carries deterministic effective tools with the state generation that authorized their scopes.
183//
184struct AvailableToolIdsSnapshot {
185    tool_ids: Vec<String>,
186    state_generation: Option<u64>,
187}
188
189//
190// Binds human approval to the reviewed action and exact tool implementation.
191//
192#[derive(Clone)]
193struct ToolApprovalBinding {
194    canonical_id: String,
195    arguments: Value,
196    confirmation_required: bool,
197    policy_version: u64,
198    runtime_control_version: u64,
199    state_generation: Option<u64>,
200    reviewed_tool: Arc<dyn ai_agents_core::Tool>,
201}
202
203//
204// A later plain approval must not erase arguments modified by an earlier approval stage.
205//
206fn merge_approved_record(record: &mut Option<ToolApprovalRecord>) {
207    if record
208        .as_ref()
209        .is_some_and(|record| matches!(record.status, ToolApprovalStatus::Modified))
210    {
211        return;
212    }
213    *record = Some(ToolApprovalRecord {
214        status: ToolApprovalStatus::Approved,
215        reason: None,
216        modified_arguments: None,
217    });
218}
219
220impl ToolApprovalBinding {
221    /// Returns true when final authorization no longer matches the reviewed action.
222    fn is_stale(
223        &self,
224        canonical_id: &str,
225        arguments: &Value,
226        confirmation_required: bool,
227        versions: ToolDecisionVersions,
228        resolved_tool: &Arc<dyn ai_agents_core::Tool>,
229    ) -> bool {
230        self.canonical_id != canonical_id
231            || self.arguments != *arguments
232            || self.confirmation_required != confirmation_required
233            || self.policy_version != versions.policy
234            || self.runtime_control_version != versions.runtime_control
235            || self.state_generation != versions.state
236            || !Arc::ptr_eq(&self.reviewed_tool, resolved_tool)
237    }
238}
239
240use crate::turn_context::{current_turn_actor_context, scope_actor_context};
241
242use ai_agents_context::{ContextManager, ContextProvider, TemplateRenderer};
243use ai_agents_core::traits::storage::StorageCapability;
244use ai_agents_core::{
245    AgentError, AgentSnapshot, AgentStorage, ChatMessage, FinishReason, LLMError, LLMProvider,
246    LLMResponse, LLMToolDefinition, LLMToolRequest, PermissionOutcome, Result, ToolActorContext,
247    ToolApprovalRecord, ToolApprovalStatus, ToolCallClassification, ToolCallSource,
248    ToolCancellationToken, ToolChoice, ToolExecutionContext, ToolExecutionLimits,
249    ToolExecutionRecord, ToolExecutionRequest, ToolInvoker, ToolPolicyDecisionRecord, ToolResult,
250    ToolSafetyMetadata,
251};
252use ai_agents_disambiguation::{
253    ClarificationObserver, ClarificationParseFuture, ClarificationQuestionFuture,
254    ConfirmationParseFuture, DisambiguationConfig, DisambiguationContext, DisambiguationManager,
255    DisambiguationResult,
256};
257use ai_agents_hitl::{
258    ApprovalHandler, ApprovalResolvedOutcome, ApprovalResult, ApprovalTrigger, HITLCheckResult,
259    HITLEngine, RejectAllHandler, TimeoutAction,
260};
261use ai_agents_hooks::{AgentHooks, NoopHooks};
262use ai_agents_llm::LLMRegistry;
263use ai_agents_memory::{
264    CompressResult, EvictionReason, Memory, MemoryBudgetEvent, MemoryCompressEvent,
265    MemoryEvictEvent, MemoryTokenBudget, OverflowStrategy,
266};
267use ai_agents_observability::{
268    EventStatus, EventType, ObservabilityManager, ObservationPurpose, SpanContext,
269    current_observation_context, new_session_id as new_observation_session_id,
270    resolve_language_from_context, with_observation_context, with_observation_purpose,
271};
272use ai_agents_process::{
273    ProcessData, ProcessProcessor, ProcessPurposeHint, ProcessStageFuture, ProcessStageObserver,
274};
275use ai_agents_reasoning::{
276    CriterionResult, EvaluationResult, Plan, PlanAction, PlanStatus, PlanStep, ReasoningConfig,
277    ReasoningMetadata, ReasoningMode, ReasoningOutput, ReflectionAttempt, ReflectionConfig,
278    ReflectionMetadata, StepFailureAction,
279};
280use ai_agents_recovery::{
281    ByRoleFilter, ContextOverflowAction, FilterConfig, IntoClassifiedError, KeepRecentFilter,
282    LLMFailureAction, MessageFilter, RecoveryManager, SkipPatternFilter, ToolFailureAction,
283};
284use ai_agents_relationships::RelationshipManager;
285use ai_agents_skills::{SkillDefinition, SkillExecutor, SkillRouter};
286use ai_agents_state::{
287    PromptMode, StateAction, StateMachine, StateMachineSnapshot, StateTransitionEvent, Transition,
288    TransitionContext, TransitionEvaluator, TransitionTiming, evaluate_guard,
289};
290use ai_agents_storage::{StorageConfig as StorageStorageConfig, create_storage};
291use ai_agents_tools::{
292    CommandRunner, ConditionEvaluator, DiagnosticsProvider, EvaluationContext, LLMGetter,
293    MAX_TOOL_TIMEOUT_MS, QuestionHandler, SecurityCheckResult, TodoItem, ToolCallRecord,
294    ToolRegistry, ToolSecurityConfig, ToolSecurityEngine,
295};
296
297use super::{
298    Agent, AgentInfo, AgentResponse, AgentStreamEvent, ParallelToolsConfig, StreamChunk,
299    StreamingConfig, ToolCall,
300};
301use crate::optimization::{
302    AwaitBeforeNextTurn, BackgroundMaintenanceQueue, BackgroundOverflowPolicy, MainResponseDraft,
303    MaintenanceMode, MaintenanceSequenceKey, RuntimeBranch, RuntimeBranchResult,
304    RuntimeBranchStatus, RuntimeCommitBehavior, RuntimeConfig, RuntimeOptimizationKind,
305    RuntimeTaskPriority, RuntimeTaskPurpose, ScheduledBranchSet, SkillCandidate,
306    StreamingDraftResult, TransitionCandidate, TurnBranchScheduler, TurnOptimizationContext,
307};
308use crate::spec::StorageConfig;
309
310/// Outcome of processing tool calls within the agent loop.
311enum ToolCallOutcome {
312    /// Tools executed successfully, continue the LLM loop for the next iteration.
313    Continue,
314    /// A state transition fired during tool call handling, continue the loop.
315    TransitionFired,
316    /// HITL rejected a tool call, return this response immediately.
317    Rejected(AgentResponse),
318}
319
320#[derive(Clone)]
321struct MainToolProtocol {
322    choice: Option<ToolChoice>,
323    tool_ids: Vec<String>,
324    definitions: Vec<LLMToolDefinition>,
325}
326
327struct MainProviderResponse {
328    response: LLMResponse,
329    used_native_tools: bool,
330}
331
332//
333// Carries one committed model answer through shared output processing and root-turn finalization.
334//
335struct CommittedTextResponse<'a> {
336    processed_input: &'a str,
337    input_context: &'a HashMap<String, Value>,
338    answer: String,
339    reasoning_mode: ReasoningMode,
340    auto_detected: bool,
341    iterations: u32,
342    thinking_content: Option<String>,
343    all_tool_calls: Vec<ToolCall>,
344}
345
346//
347// Carries the finalized response fields into metadata assembly without changing their ownership.
348//
349struct AgentResponseParts {
350    content: String,
351    all_tool_calls: Vec<ToolCall>,
352    reasoning_mode: ReasoningMode,
353    auto_detected: bool,
354    iterations: u32,
355    thinking: Option<String>,
356    reflection_metadata: Option<ReflectionMetadata>,
357}
358
359//
360// Owns the finalized response paired with a legacy Done chunk without changing provisional chunk timing.
361//
362type RuntimeStreamTerminalSlot = Arc<RwLock<Option<AgentResponse>>>;
363
364//
365// Creates isolated terminal ownership for one public stream.
366//
367fn new_runtime_stream_terminal_slot() -> RuntimeStreamTerminalSlot {
368    Arc::new(RwLock::new(None))
369}
370
371//
372// Records the exact response that already completed root-turn finalization before Done is emitted.
373//
374fn record_runtime_stream_final(slot: &RuntimeStreamTerminalSlot, response: AgentResponse) {
375    *slot.write() = Some(response);
376}
377
378#[derive(Clone, Copy)]
379struct DisambiguationOwnership {
380    epoch: u64,
381    state_generation: Option<u64>,
382}
383
384/// Outcome of skill routing — used by `try_skill_route`.
385enum SkillRouteResult {
386    /// No skill matched, continue to normal LLM chat.
387    NoMatch,
388    /// Skill executed successfully.
389    Response { skill_id: String, content: String },
390    /// Skill matched but needs disambiguation first or returned a terminal disambiguation response.
391    NeedsClarification {
392        response: AgentResponse,
393        ownership: Option<DisambiguationOwnership>,
394    },
395}
396
397/// Result of response-independent parallel transition selection.
398enum ParallelTransitionSelection {
399    /// A transition matched and can be committed before the old-state response.
400    Candidate(TransitionCandidate),
401    /// All eligible transition checks ran and none matched.
402    NoMatch,
403    /// LLM-based route evaluation needed speculative capacity that was unavailable.
404    ReservationExhausted,
405}
406
407/// Outcome of post_loop_processing - drives the caller's next step.
408enum PostLoopResult {
409    /// No transition fired. Content is the LLM response for this turn.
410    NoTransition(String),
411    /// Transition fired. Content is from plain post-transition re-generation.
412    Transitioned(String),
413    /// Transition fired into a state that requires full dispatch.
414    /// Caller re-enters run_loop_internal to apply the correct handler.
415    NeedsRedispatch,
416}
417
418struct StateTransitionReservation<'a> {
419    reserved: &'a AtomicBool,
420}
421
422impl Drop for StateTransitionReservation<'_> {
423    fn drop(&mut self) {
424        self.reserved.store(false, Ordering::SeqCst);
425    }
426}
427
428struct RootTurnCleanup<'a> {
429    agent: &'a RuntimeAgent,
430}
431
432impl<'a> RootTurnCleanup<'a> {
433    fn new(agent: &'a RuntimeAgent) -> Self {
434        Self { agent }
435    }
436}
437
438impl Drop for RootTurnCleanup<'_> {
439    fn drop(&mut self) {
440        self.agent.end_root_turn();
441    }
442}
443
444/// Host-owned runtime control state shared with active agents.
445#[derive(Debug)]
446struct RuntimeControlState {
447    /// Serializes control mutations with exact runtime safety snapshots.
448    snapshot_guard: RwLock<()>,
449    /// Monotonic version for runtime-control snapshots.
450    version: AtomicU64,
451    /// Emergency switch that denies future calls and is shared with active tool contexts.
452    emergency_deny: Arc<AtomicBool>,
453    /// Optional live replacement for tool security policy.
454    tool_security_override: RwLock<Option<ToolSecurityEngine>>,
455    /// Optional live narrowing scope applied after the runtime's declared grant.
456    tool_scope_override: RwLock<Option<Vec<String>>>,
457}
458
459impl Default for RuntimeControlState {
460    fn default() -> Self {
461        Self {
462            snapshot_guard: RwLock::new(()),
463            version: AtomicU64::new(1),
464            emergency_deny: Arc::new(AtomicBool::new(false)),
465            tool_security_override: RwLock::new(None),
466            tool_scope_override: RwLock::new(None),
467        }
468    }
469}
470
471/// Host-only handle for live runtime safety controls.
472#[derive(Clone)]
473pub struct RuntimeControlHandle {
474    state: Arc<RuntimeControlState>,
475}
476
477impl RuntimeControlHandle {
478    /// Returns the current runtime-control version.
479    pub fn version(&self) -> u64 {
480        self.state.version.load(Ordering::SeqCst)
481    }
482
483    fn bump(&self) -> u64 {
484        self.state.version.fetch_add(1, Ordering::SeqCst) + 1
485    }
486
487    /// Overrides tool security for later tool calls and panics if host policy is invalid.
488    pub fn set_tool_security(&self, config: ToolSecurityConfig) -> u64 {
489        self.try_set_tool_security(config)
490            .expect("invalid tool security configuration")
491    }
492
493    /// Validates replacement policy before changing the runtime-control generation or active snapshot.
494    pub fn try_set_tool_security(&self, config: ToolSecurityConfig) -> Result<u64> {
495        config.validate()?;
496        let _guard = self.state.snapshot_guard.write();
497        let generation = self.bump();
498        *self.state.tool_security_override.write() = Some(
499            ToolSecurityEngine::new_with_policy_version(config, generation),
500        );
501        Ok(generation)
502    }
503
504    /// Clears the live tool security override.
505    pub fn clear_tool_security_override(&self) -> u64 {
506        let _guard = self.state.snapshot_guard.write();
507        *self.state.tool_security_override.write() = None;
508        self.bump()
509    }
510
511    /// Narrows the runtime's declared tool grant for later calls without adding authority.
512    pub fn set_tool_scope(&self, tool_ids: Vec<String>) -> u64 {
513        let _guard = self.state.snapshot_guard.write();
514        *self.state.tool_scope_override.write() = Some(tool_ids);
515        self.bump()
516    }
517
518    /// Clears live narrowing so later calls return to the runtime's declared grant.
519    pub fn clear_tool_scope_override(&self) -> u64 {
520        let _guard = self.state.snapshot_guard.write();
521        *self.state.tool_scope_override.write() = None;
522        self.bump()
523    }
524
525    /// Enables or disables emergency denial for future tool calls.
526    pub fn set_emergency_deny(&self, enabled: bool) -> u64 {
527        let _guard = self.state.snapshot_guard.write();
528        self.state.emergency_deny.store(enabled, Ordering::SeqCst);
529        self.bump()
530    }
531
532    /// Denies future calls and asks active cancellable work to stop.
533    pub fn cancel_all(&self) -> u64 {
534        self.set_emergency_deny(true)
535    }
536}
537
538pub struct RuntimeAgent {
539    info: AgentInfo,
540    llm_registry: Arc<LLMRegistry>,
541    memory: Arc<dyn Memory>,
542    tools: Arc<ToolRegistry>,
543    skills: Vec<SkillDefinition>,
544    skill_router: Option<SkillRouter>,
545    skill_executor: Option<SkillExecutor>,
546    base_system_prompt: String,
547    max_iterations: u32,
548    iteration_count: RwLock<u32>,
549    max_context_tokens: u32,
550    memory_token_budget: Option<MemoryTokenBudget>,
551    recovery_manager: RecoveryManager,
552    tool_security: ToolSecurityEngine,
553    process_processor: Option<ProcessProcessor>,
554    message_filters: RwLock<HashMap<String, Arc<dyn MessageFilter>>>,
555    state_machine: Option<Arc<StateMachine>>,
556    transition_evaluator: Option<Arc<dyn TransitionEvaluator>>,
557    context_manager: Arc<ContextManager>,
558    template_renderer: TemplateRenderer,
559    tool_call_history: RwLock<Vec<ToolCallRecord>>,
560    parallel_tools: ParallelToolsConfig,
561    streaming: StreamingConfig,
562    hooks: Arc<dyn AgentHooks>,
563    hitl_engine: Option<HITLEngine>,
564    approval_handler: Arc<dyn ApprovalHandler>,
565    storage_config: StorageConfig,
566    storage: RwLock<Option<Arc<dyn AgentStorage>>>,
567    storage_init: tokio::sync::Mutex<()>,
568    reasoning_config: ReasoningConfig,
569    reflection_config: ReflectionConfig,
570    disambiguation_manager: Option<DisambiguationManager>,
571    /// Generation invalidating confirmation work across reset and state changes.
572    disambiguation_epoch: AtomicU64,
573    /// Serializes confirmation redispatch admission with reset and state mutation.
574    disambiguation_admission: tokio::sync::RwLock<()>,
575    /// Reserves one runtime-owned transition before its exit actions can produce side effects.
576    state_transition_reserved: AtomicBool,
577    /// Structured persona manager for identity, evolution, and secrets.
578    persona_manager: Option<Arc<ai_agents_persona::PersonaManager>>,
579    /// Skill ID that triggered the current pending disambiguation.
580    /// Set by try_skill_route() when skill-level disambiguation triggers clarification.
581    /// Read by run_loop() when clarification resolves to route directly to the skill.
582    pending_skill_id: RwLock<Option<String>>,
583    current_plan: RwLock<Option<Plan>>,
584    /// Tool IDs declared in the top-level `tools:` spec.
585    declared_tool_ids: Option<Vec<String>>,
586    /// Whether the context manager has been initialized (defaults loaded, env resolved, etc.)
587    context_initialized: AtomicBool,
588    /// Spawner for dynamic agent creation (set when YAML has a spawner: section).
589    spawner: Option<Arc<crate::spawner::AgentSpawner>>,
590    /// Registry tracking spawned agents (set when YAML has a spawner: section).
591    spawner_registry: Option<Arc<crate::spawner::AgentRegistry>>,
592    /// Re-dispatch depth for post-transition full dispatch.
593    /// 0 = not re-dispatching. > 0 = user message already in memory, skip re-adding.
594    redispatch_depth: RwLock<u32>,
595    /// Active optimized turn context used to keep root lifecycle state in one place.
596    active_turn_context: RwLock<Option<TurnOptimizationContext>>,
597    /// Tracks whether the root turn already wrote the processed user message.
598    root_user_message_committed: AtomicBool,
599    /// Current actor ID for cross-session memory.
600    actor_id: RwLock<Option<String>>,
601    /// Fact store for managing per-actor extracted facts.
602    fact_store: RwLock<Option<Arc<ai_agents_facts::FactStore>>>,
603    /// Fact extractor for LLM-based fact extraction.
604    /// None when actor_memory is enabled without facts.enabled.
605    fact_extractor: RwLock<Option<Arc<dyn ai_agents_facts::FactExtractor>>>,
606    /// Cached actor facts keyed by actor ID so concurrent or alternating turns do not overwrite one another.
607    actor_facts_cache: Arc<RwLock<HashMap<String, Vec<ai_agents_core::KeyFact>>>>,
608    /// Number of messages since last fact extraction.
609    messages_since_extraction: Arc<RwLock<usize>>,
610    /// Actor memory configuration.
611    actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
612    /// Facts configuration.
613    facts_config: Option<ai_agents_facts::FactsConfig>,
614    /// Session-scoped metadata (tags, ttl, actor roster).
615    session_metadata: RwLock<ai_agents_core::SessionMetadata>,
616    /// Session id currently bound to this runtime instance.
617    current_session_id: RwLock<Option<String>>,
618    /// Relationship manager for actor-scoped social memory.
619    relationship_manager: Option<Arc<RelationshipManager>>,
620    /// Observability manager for traces, metrics, reports, and exports.
621    observability_manager: Option<Arc<ObservabilityManager>>,
622    /// Runtime optimization and maintenance policy.
623    runtime_config: RuntimeConfig,
624    /// Queue for background maintenance tasks.
625    background_maintenance: Arc<BackgroundMaintenanceQueue>,
626    /// Cross-tool locks for side-effecting calls that target the same resource.
627    resource_locks: ToolResourceLocks,
628    /// Host-only runtime control state.
629    runtime_control: Arc<RuntimeControlState>,
630    /// Serializes independent externally initiated root turns across blocking and streaming APIs.
631    root_turn_gate: RootTurnGate,
632}
633
634impl std::fmt::Debug for RuntimeAgent {
635    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
636        f.debug_struct("RuntimeAgent")
637            .field("info", &self.info)
638            .field("base_system_prompt", &self.base_system_prompt)
639            .field("max_iterations", &self.max_iterations)
640            .field("skills_count", &self.skills.len())
641            .field("max_context_tokens", &self.max_context_tokens)
642            .field("has_state_machine", &self.state_machine.is_some())
643            .field("parallel_tools", &self.parallel_tools)
644            .field("streaming", &self.streaming)
645            .field("has_hooks", &true)
646            .field("has_hitl", &self.hitl_engine.is_some())
647            .field("storage_type", &self.storage_config.storage_type())
648            .field("reasoning_mode", &self.reasoning_config.mode)
649            .field("reflection_enabled", &self.reflection_config.enabled)
650            .field("declared_tool_ids", &self.declared_tool_ids)
651            .field("has_persona", &self.persona_manager.is_some())
652            .field("has_observability", &self.observability_manager.is_some())
653            .finish_non_exhaustive()
654    }
655}
656
657struct ObservabilityClarificationObserver;
658
659impl ClarificationObserver for ObservabilityClarificationObserver {
660    /// Scopes clarification question generation as disambiguation_clarification.
661    fn observe_question<'a>(
662        &'a self,
663        future: ClarificationQuestionFuture<'a>,
664    ) -> ClarificationQuestionFuture<'a> {
665        Box::pin(async move {
666            with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
667        })
668    }
669
670    /// Scopes clarification response parsing as disambiguation_clarification.
671    fn observe_parse<'a>(
672        &'a self,
673        future: ClarificationParseFuture<'a>,
674    ) -> ClarificationParseFuture<'a> {
675        Box::pin(async move {
676            with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
677        })
678    }
679
680    /// Scopes semantic confirmation parsing as disambiguation_clarification.
681    fn observe_confirmation_parse<'a>(
682        &'a self,
683        future: ConfirmationParseFuture<'a>,
684    ) -> ConfirmationParseFuture<'a> {
685        Box::pin(async move {
686            with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
687        })
688    }
689}
690
691struct ObservabilityProcessStageObserver;
692
693impl ProcessStageObserver for ObservabilityProcessStageObserver {
694    /// Scopes one process stage using the purpose implied by its stage type.
695    fn observe<'a>(
696        &'a self,
697        hint: ProcessPurposeHint,
698        future: ProcessStageFuture<'a>,
699    ) -> ProcessStageFuture<'a> {
700        Box::pin(async move {
701            with_observation_purpose(observation_purpose_for_process(hint), future).await
702        })
703    }
704}
705
706struct RegistryLLMGetter {
707    registry: Arc<LLMRegistry>,
708}
709
710impl LLMGetter for RegistryLLMGetter {
711    fn get_llm(&self, alias: &str) -> Option<Arc<dyn LLMProvider>> {
712        self.registry.get(alias).ok()
713    }
714}
715
716impl RuntimeAgent {
717    /// Constructs a runtime with one gate shared by every external root-turn entry point.
718    #[allow(clippy::too_many_arguments)]
719    pub fn new(
720        info: AgentInfo,
721        llm_registry: Arc<LLMRegistry>,
722        memory: Arc<dyn Memory>,
723        tools: Arc<ToolRegistry>,
724        skills: Vec<SkillDefinition>,
725        system_prompt: String,
726        max_iterations: u32,
727    ) -> Self {
728        let (skill_router, skill_executor) = if !skills.is_empty() {
729            let router_llm = llm_registry.router().ok();
730            let router = router_llm.map(|llm| SkillRouter::new(llm, skills.clone()));
731            let executor = SkillExecutor::new(llm_registry.clone(), tools.clone());
732            (router, Some(executor))
733        } else {
734            (None, None)
735        };
736
737        let context_manager =
738            ContextManager::new(HashMap::new(), info.name.clone(), info.version.clone());
739
740        Self {
741            info,
742            llm_registry,
743            memory,
744            tools,
745            skills,
746            skill_router,
747            skill_executor,
748            base_system_prompt: system_prompt,
749            max_iterations,
750            iteration_count: RwLock::new(0),
751            max_context_tokens: 128000,
752            memory_token_budget: None,
753            recovery_manager: RecoveryManager::default(),
754            tool_security: ToolSecurityEngine::default(),
755            process_processor: None,
756            message_filters: RwLock::new(HashMap::new()),
757            state_machine: None,
758            transition_evaluator: None,
759            context_manager: Arc::new(context_manager),
760            template_renderer: TemplateRenderer::new(),
761            tool_call_history: RwLock::new(Vec::new()),
762            parallel_tools: ParallelToolsConfig::default(),
763            streaming: StreamingConfig::default(),
764            hooks: Arc::new(NoopHooks),
765            hitl_engine: None,
766            approval_handler: Arc::new(RejectAllHandler::new()),
767            storage_config: StorageConfig::default(),
768            storage: RwLock::new(None),
769            storage_init: tokio::sync::Mutex::new(()),
770            reasoning_config: ReasoningConfig::default(),
771            reflection_config: ReflectionConfig::default(),
772            disambiguation_manager: None,
773            disambiguation_epoch: AtomicU64::new(0),
774            disambiguation_admission: tokio::sync::RwLock::new(()),
775            state_transition_reserved: AtomicBool::new(false),
776            persona_manager: None,
777            pending_skill_id: RwLock::new(None),
778            current_plan: RwLock::new(None),
779            declared_tool_ids: None,
780            context_initialized: AtomicBool::new(false),
781            spawner: None,
782            spawner_registry: None,
783            redispatch_depth: RwLock::new(0),
784            active_turn_context: RwLock::new(None),
785            root_user_message_committed: AtomicBool::new(false),
786            actor_id: RwLock::new(None),
787            fact_store: RwLock::new(None),
788            fact_extractor: RwLock::new(None),
789            actor_facts_cache: Arc::new(RwLock::new(HashMap::new())),
790            messages_since_extraction: Arc::new(RwLock::new(0)),
791            actor_memory_config: None,
792            facts_config: None,
793            session_metadata: RwLock::new(ai_agents_core::SessionMetadata::default()),
794            current_session_id: RwLock::new(None),
795            relationship_manager: None,
796            observability_manager: None,
797            runtime_config: RuntimeConfig::default(),
798            background_maintenance: Arc::new(BackgroundMaintenanceQueue::default()),
799            resource_locks: new_tool_resource_locks(),
800            runtime_control: Arc::new(RuntimeControlState::default()),
801            root_turn_gate: Arc::new(tokio::sync::Mutex::new(())),
802        }
803    }
804
805    pub fn with_declared_tool_ids(mut self, ids: Option<Vec<String>>) -> Self {
806        self.declared_tool_ids = ids;
807        self
808    }
809
810    pub fn with_storage_config(mut self, config: StorageConfig) -> Self {
811        self.storage_config = config;
812        self
813    }
814
815    pub fn with_storage(self, storage: Arc<dyn AgentStorage>) -> Self {
816        *self.storage.write() = Some(storage);
817        self
818    }
819
820    pub(crate) fn with_shared_resource_locks(mut self, locks: ToolResourceLocks) -> Self {
821        self.resource_locks = locks;
822        self
823    }
824
825    pub fn with_reasoning(mut self, config: ReasoningConfig) -> Self {
826        self.reasoning_config = config;
827        self
828    }
829
830    pub fn with_reflection(mut self, config: ReflectionConfig) -> Self {
831        self.reflection_config = config;
832        self
833    }
834
835    /// Attach a relationship manager configured by the builder or host application.
836    pub fn with_relationships(mut self, manager: Arc<RelationshipManager>) -> Self {
837        self.relationship_manager = Some(manager);
838        self
839    }
840
841    /// Attach a shared observability manager for traces, metrics, reports, and exports.
842    pub fn with_observability(mut self, manager: Arc<ObservabilityManager>) -> Self {
843        self.observability_manager = Some(manager);
844        self
845    }
846
847    /// Attach runtime optimization policy and resize the background queue.
848    pub fn with_runtime_config(mut self, config: RuntimeConfig) -> Self {
849        let max_tasks = config.optimization.post_turn.max_background_tasks;
850        self.background_maintenance = Arc::new(BackgroundMaintenanceQueue::new(max_tasks));
851        self.runtime_config = config;
852        self
853    }
854
855    /// Returns the runtime optimization policy.
856    pub fn runtime_config(&self) -> &RuntimeConfig {
857        &self.runtime_config
858    }
859
860    /// Wait for all background maintenance tasks to finish.
861    pub async fn flush_background_tasks(&self) -> Result<()> {
862        self.background_maintenance.flush_all().await
863    }
864
865    /// Wait for background maintenance associated with one actor to finish.
866    pub async fn flush_background_tasks_for_actor(&self, actor_id: &str) -> Result<()> {
867        self.background_maintenance.flush_scope(actor_id).await
868    }
869
870    /// Wait for background maintenance associated with one task kind to finish.
871    pub async fn flush_background_tasks_for_purpose(
872        &self,
873        purpose: RuntimeTaskPurpose,
874    ) -> Result<()> {
875        self.background_maintenance.flush_purpose(purpose).await
876    }
877
878    /// Wait for background maintenance associated with one actor and task kind to finish.
879    pub async fn flush_background_tasks_for_actor_purpose(
880        &self,
881        actor_id: &str,
882        purpose: RuntimeTaskPurpose,
883    ) -> Result<()> {
884        self.background_maintenance
885            .flush_scope_purpose(actor_id, purpose)
886            .await
887    }
888
889    /// Flush background maintenance before a host shuts down the runtime.
890    pub async fn shutdown_background_tasks(&self) -> Result<()> {
891        self.flush_background_tasks().await
892    }
893
894    /// Returns the configured observability manager for report and export access.
895    pub fn observability(&self) -> Option<Arc<ObservabilityManager>> {
896        self.observability_manager.clone()
897    }
898
899    /// Exports observability files after a turn when export settings request it.
900    async fn export_observability_if_configured(&self) {
901        let Some(manager) = self.observability_manager.as_ref() else {
902            return;
903        };
904        let export = &manager.config().export;
905        if !export.write_report && !export.write_raw_events {
906            return;
907        }
908        if let Err(error) = manager.export().await {
909            warn!(error = %error, "Observability export failed");
910        }
911    }
912
913    /// Returns the configured relationship manager, if relationship memory is enabled.
914    pub fn relationship_manager(&self) -> Option<Arc<RelationshipManager>> {
915        self.relationship_manager.clone()
916    }
917
918    fn current_turn_actor_context(&self) -> Option<crate::TurnActorContext> {
919        current_turn_actor_context()
920    }
921
922    fn effective_actor_id(&self) -> Option<String> {
923        self.current_turn_actor_context()
924            .and_then(|ctx| ctx.effective_actor_id().map(|id| id.to_string()))
925            .or_else(|| self.actor_id.read().clone())
926    }
927
928    fn effective_origin_actor_id(&self) -> Option<String> {
929        self.current_turn_actor_context()
930            .and_then(|ctx| ctx.origin_actor_id.clone())
931            .or_else(|| self.actor_id.read().clone())
932    }
933
934    fn record_session_actor_if_needed(&self) {
935        if let Some(actor_id) = self.effective_origin_actor_id() {
936            let mut meta = self.session_metadata.write();
937            meta.actor_id = Some(actor_id.clone());
938            if !meta.actors.iter().any(|a| a == &actor_id) {
939                meta.actors.push(actor_id);
940            }
941        }
942    }
943
944    fn outbound_actor_context(&self) -> crate::TurnActorContext {
945        let mut context = self.current_turn_actor_context().unwrap_or_default();
946        if context.origin_actor_id.is_none() {
947            context.origin_actor_id = self.effective_origin_actor_id();
948        }
949        context.sender_agent_id = Some(self.info.id.clone());
950        context
951    }
952
953    /// Returns the current session ID or creates one for unsaved observed turns.
954    fn observation_session_id(&self) -> Option<String> {
955        let mut current = self.current_session_id.write();
956        if current.is_none() {
957            *current = Some(new_observation_session_id());
958        }
959        current.clone()
960    }
961
962    /// Builds the root or child observation context for a chat entry point.
963    fn build_observation_context(&self, actor_id: Option<String>) -> Option<SpanContext> {
964        let manager = self.observability_manager.as_ref()?;
965        let context = self.build_context_with_overlays();
966        let language = resolve_language_from_context(manager.config(), &context);
967        let context = current_observation_context()
968            .map(|parent| parent.child_for_agent(self.info.id.clone()).with_new_turn())
969            .unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
970        Some(
971            context
972                .with_actor(actor_id.or_else(|| self.effective_actor_id()))
973                .with_session(self.observation_session_id())
974                .with_state(self.current_state())
975                .with_language(Some(language)),
976        )
977    }
978
979    /// Refreshes task-local context with current runtime labels and a purpose.
980    fn current_runtime_observation_context(
981        &self,
982        purpose: ObservationPurpose,
983    ) -> Option<SpanContext> {
984        let manager = self.observability_manager.as_ref()?;
985        let context = self.build_context_with_overlays();
986        let language = resolve_language_from_context(manager.config(), &context);
987        let mut observation = current_observation_context()
988            .unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
989        observation.agent_id = self.info.id.clone();
990        observation.actor_id = self.effective_actor_id();
991        observation.session_id = self.observation_session_id();
992        observation.state = self.current_state();
993        observation.language = Some(language);
994        observation.purpose = purpose;
995        Some(observation)
996    }
997
998    /// Runs a future under a purpose while preserving current trace context.
999    async fn observe_purpose<F, T>(&self, purpose: ObservationPurpose, future: F) -> T
1000    where
1001        F: Future<Output = T>,
1002    {
1003        if let Some(context) = self.current_runtime_observation_context(purpose) {
1004            with_observation_context(context, future).await
1005        } else {
1006            future.await
1007        }
1008    }
1009
1010    /// Runs one actor-scoped external root turn with gate ownership visible through finalization, hooks, orchestration, and export.
1011    ///
1012    /// Internal redispatch calls `run_loop` directly and therefore keeps the original guard and identity stack instead of acquiring this non-reentrant gate again.
1013    fn chat_with_actor_context_boxed<'a>(
1014        &'a self,
1015        input: &'a str,
1016        actor_context: crate::TurnActorContext,
1017    ) -> Pin<Box<dyn Future<Output = Result<AgentResponse>> + Send + 'a>> {
1018        Box::pin(async move {
1019            let RootTurnAdmission {
1020                guard,
1021                identity_stack,
1022            } = self.acquire_root_turn().await?;
1023            let result = scope_runtime_gate_identity_stack(&identity_stack, async move {
1024                let actor_id = actor_context.effective_actor_id().map(str::to_string);
1025                let run = async move {
1026                    scope_actor_context(
1027                        actor_context,
1028                        Box::pin(async move { self.run_loop(input).await }),
1029                    )
1030                    .await
1031                };
1032                let result = if let Some(context) = self.build_observation_context(actor_id) {
1033                    with_observation_context(context, run).await
1034                } else {
1035                    run.await
1036                };
1037                self.export_observability_if_configured().await;
1038                result
1039            })
1040            .await;
1041            drop(guard);
1042            result
1043        })
1044    }
1045
1046    /// Rejects recursive ownership before waiting, then acquires the external root-turn gate and builds one immutable extended ancestry snapshot.
1047    async fn acquire_root_turn(&self) -> Result<RootTurnAdmission> {
1048        let gate_identity = Arc::clone(&self.root_turn_gate);
1049        let current_identity_stack = current_runtime_gate_identity_stack();
1050        if current_identity_stack
1051            .iter()
1052            .any(|owned_gate| Arc::ptr_eq(owned_gate, &gate_identity))
1053        {
1054            return Err(AgentError::Other(format!(
1055                "RuntimeAgent '{}' rejected reentrant root turn ownership",
1056                self.info.id
1057            )));
1058        }
1059        let guard = Arc::clone(&gate_identity).lock_owned().await;
1060        //
1061        // Root admission is the only place that extends ancestry, so later scopes can preserve the complete chain with an `Arc` clone.
1062        //
1063        let mut identity_stack = Vec::with_capacity(current_identity_stack.len() + 1);
1064        identity_stack.extend(current_identity_stack.iter().cloned());
1065        identity_stack.push(gate_identity);
1066        Ok(RootTurnAdmission {
1067            guard,
1068            identity_stack: identity_stack.into(),
1069        })
1070    }
1071
1072    /// Run one serialized turn with turn-scoped actor context without mutating the runtime's global actor ID.
1073    ///
1074    /// The supplied context is available to actor-scoped facts, relationship memory, orchestration, and prompt templates only for the lifetime of this call.
1075    pub async fn chat_with_actor_context(
1076        &self,
1077        input: &str,
1078        actor_context: crate::TurnActorContext,
1079    ) -> Result<AgentResponse> {
1080        self.chat_with_actor_context_boxed(input, actor_context)
1081            .await
1082    }
1083
1084    /// Convenience wrapper around [`Self::chat_with_actor_context`] for a turn whose original actor is known up front.
1085    pub async fn chat_as_actor(&self, actor_id: &str, input: &str) -> Result<AgentResponse> {
1086        let actor_context = crate::TurnActorContext::new().with_origin_actor(actor_id);
1087        self.chat_with_actor_context(input, actor_context).await
1088    }
1089
1090    /// Ensure the effective actor's relationship is loaded from storage into the relationship manager.
1091    pub async fn load_actor_relationship(&self) -> Result<()> {
1092        self.maybe_load_actor_relationship().await;
1093        Ok(())
1094    }
1095
1096    /// Manually apply a delta to the effective actor's `agent_to_actor` relationship perspective and persist the updated relationship when storage is configured.
1097    pub async fn update_relationship_dimension(
1098        &self,
1099        dimension: &str,
1100        delta: f64,
1101        reason: Option<&str>,
1102    ) -> Result<ai_agents_relationships::DimensionChange> {
1103        self.update_relationship_dimension_for_perspective(
1104            ai_agents_relationships::RelationshipPerspective::AgentToActor,
1105            dimension,
1106            delta,
1107            reason,
1108        )
1109        .await
1110    }
1111
1112    /// Manually apply a delta to a specific relationship perspective for the effective actor.
1113    ///
1114    /// Use this for two-sided configurations when you need to update `agent_to_actor`, `perceived_actor_to_agent`, or `mutual` explicitly from application logic.
1115    pub async fn update_relationship_dimension_for_perspective(
1116        &self,
1117        perspective: ai_agents_relationships::RelationshipPerspective,
1118        dimension: &str,
1119        delta: f64,
1120        reason: Option<&str>,
1121    ) -> Result<ai_agents_relationships::DimensionChange> {
1122        let manager = self
1123            .relationship_manager
1124            .as_ref()
1125            .ok_or_else(|| AgentError::Config("Relationship memory is not configured".into()))?;
1126        let actor_id = self.effective_actor_id().ok_or_else(|| {
1127            AgentError::Config("No actor ID set. Use set_actor_id() first".into())
1128        })?;
1129        let change = manager.update_dimension_for_perspective(
1130            &actor_id,
1131            perspective,
1132            dimension,
1133            delta,
1134            1.0,
1135            reason.unwrap_or("manual relationship update"),
1136        )?;
1137        self.persist_actor_relationship(&actor_id).await?;
1138        info!(
1139            actor_id = %actor_id,
1140            perspective = %change.perspective,
1141            dimension = %change.dimension,
1142            delta = change.delta,
1143            current = change.current,
1144            "relationship updated manually"
1145        );
1146        self.hooks
1147            .on_relationship_change(&actor_id, std::slice::from_ref(&change))
1148            .await;
1149        Ok(change)
1150    }
1151
1152    pub fn reasoning_config(&self) -> &ReasoningConfig {
1153        &self.reasoning_config
1154    }
1155
1156    pub fn reflection_config(&self) -> &ReflectionConfig {
1157        &self.reflection_config
1158    }
1159
1160    /// Set only the actor memory and facts configs without creating the store.
1161    /// The store and extractor are created lazily in init_storage().
1162    pub fn with_facts_config(
1163        mut self,
1164        actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
1165        facts_config: Option<ai_agents_facts::FactsConfig>,
1166    ) -> Self {
1167        self.actor_memory_config = actor_memory_config;
1168        self.facts_config = facts_config;
1169        self
1170    }
1171
1172    /// Configure fact store and optional extractor for actor memory.
1173    /// Pass `None` for `extractor` to load existing facts without running extraction.
1174    pub fn with_facts(
1175        mut self,
1176        store: Arc<ai_agents_facts::FactStore>,
1177        extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>>,
1178        actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
1179        facts_config: Option<ai_agents_facts::FactsConfig>,
1180    ) -> Self {
1181        *self.fact_store.write() = Some(store);
1182        *self.fact_extractor.write() = extractor;
1183        self.actor_memory_config = actor_memory_config;
1184        self.facts_config = facts_config;
1185        self
1186    }
1187
1188    /// Get the fact store for direct fact manipulation.
1189    pub fn fact_store(&self) -> Option<Arc<ai_agents_facts::FactStore>> {
1190        self.fact_store.read().clone()
1191    }
1192
1193    /// Get the current actor ID.
1194    pub fn actor_id(&self) -> Option<String> {
1195        self.actor_id.read().clone()
1196    }
1197
1198    /// Set the current actor ID (player, user, another agent, etc.).
1199    pub fn set_actor_id(&self, actor_id: &str) -> ai_agents_core::Result<()> {
1200        *self.actor_id.write() = Some(actor_id.to_string());
1201        {
1202            let mut meta = self.session_metadata.write();
1203            meta.actor_id = Some(actor_id.to_string());
1204            if !meta.actors.iter().any(|a| a == actor_id) {
1205                meta.actors.push(actor_id.to_string());
1206            }
1207        }
1208        Ok(())
1209    }
1210
1211    /// Clear the current actor binding while retaining the session actor roster.
1212    pub fn clear_actor_id(&self) {
1213        *self.actor_id.write() = None;
1214        self.session_metadata.write().actor_id = None;
1215    }
1216
1217    /// Set the current actor ID. Convenience wrapper around set_actor_id.
1218    pub fn set_user_id(&self, user_id: &str) -> ai_agents_core::Result<()> {
1219        self.set_actor_id(user_id)
1220    }
1221
1222    /// Load facts for the current actor from storage and cache them for prompt injection.
1223    pub async fn load_actor_memory(&self) -> ai_agents_core::Result<()> {
1224        let actor_id = match self.effective_actor_id() {
1225            Some(id) => id,
1226            None => return Ok(()),
1227        };
1228
1229        let store_opt = self.fact_store.read().clone();
1230        if let Some(store) = store_opt {
1231            let facts = store.get_facts(&actor_id).await?;
1232            let count = facts.len();
1233            self.actor_facts_cache
1234                .write()
1235                .insert(actor_id.clone(), facts);
1236            self.hooks.on_actor_memory_loaded(&actor_id, count).await;
1237            tracing::debug!("loaded {} facts for actor {}", count, actor_id);
1238        }
1239
1240        Ok(())
1241    }
1242
1243    /// Load actor memory only when the effective actor has no cached facts yet.
1244    async fn maybe_load_actor_memory(&self) {
1245        let Some(actor_id) = self.effective_actor_id() else {
1246            return;
1247        };
1248        if self.actor_facts_cache.read().contains_key(&actor_id) {
1249            return;
1250        }
1251        let _ = self.load_actor_memory().await;
1252    }
1253
1254    /// Pre-turn lifecycle shared by streaming and non-streaming paths.
1255    async fn pre_turn_session_lifecycle(&self) {
1256        if *self.redispatch_depth.read() > 0 {
1257            return;
1258        }
1259        self.resolve_actor_id_from_context();
1260        self.await_background_before_next_turn().await;
1261        self.record_session_actor_if_needed();
1262        self.maybe_load_actor_memory().await;
1263        self.maybe_load_actor_relationship().await;
1264        *self.messages_since_extraction.write() += 1;
1265    }
1266
1267    /// Post-turn lifecycle shared by streaming and non-streaming paths.
1268    async fn post_turn_session_lifecycle(&self) -> Result<()> {
1269        if *self.redispatch_depth.read() > 0 {
1270            return Ok(());
1271        }
1272        *self.messages_since_extraction.write() += 1;
1273        self.run_post_turn_maintenance().await
1274    }
1275
1276    /// Starts root-turn bookkeeping for user-message commit tracking.
1277    fn begin_root_turn(&self) {
1278        if *self.redispatch_depth.read() == 0 {
1279            let mut guard = self.active_turn_context.write();
1280            if guard.is_none() {
1281                self.root_user_message_committed
1282                    .store(false, Ordering::SeqCst);
1283                let max_calls = self
1284                    .runtime_config
1285                    .optimization
1286                    .max_speculative_llm_calls_per_turn;
1287                *guard = Some(TurnOptimizationContext::new(
1288                    String::new(),
1289                    HashMap::new(),
1290                    max_calls,
1291                ));
1292            }
1293        }
1294    }
1295
1296    fn update_active_turn_context(
1297        &self,
1298        processed_input: &str,
1299        input_context: HashMap<String, Value>,
1300    ) {
1301        if *self.redispatch_depth.read() > 0 {
1302            return;
1303        }
1304        let max_calls = self
1305            .runtime_config
1306            .optimization
1307            .max_speculative_llm_calls_per_turn;
1308        let mut guard = self.active_turn_context.write();
1309        match guard.as_mut() {
1310            Some(context) => {
1311                context.processed_input = processed_input.to_string();
1312                context.input_context = input_context;
1313                context.max_speculative_llm_calls = max_calls;
1314            }
1315            None => {
1316                *guard = Some(TurnOptimizationContext::new(
1317                    processed_input,
1318                    input_context,
1319                    max_calls,
1320                ));
1321            }
1322        }
1323    }
1324
1325    /// Writes the processed user message once for the root turn.
1326    async fn commit_root_user_message(&self, processed_input: &str) -> Result<()> {
1327        if *self.redispatch_depth.read() > 0 {
1328            return Ok(());
1329        }
1330        if !self
1331            .root_user_message_committed
1332            .swap(true, Ordering::SeqCst)
1333        {
1334            self.memory
1335                .add_message(ChatMessage::user(processed_input))
1336                .await?;
1337            if let Some(context) = self.active_turn_context.write().as_mut() {
1338                context.mark_user_message_committed();
1339            }
1340        }
1341        Ok(())
1342    }
1343
1344    /// Clears root-turn bookkeeping after final response handling.
1345    fn end_root_turn(&self) {
1346        if *self.redispatch_depth.read() == 0 {
1347            self.root_user_message_committed
1348                .store(false, Ordering::SeqCst);
1349            *self.active_turn_context.write() = None;
1350        }
1351    }
1352
1353    fn reserve_active_speculative_llm_call(&self, kind: RuntimeOptimizationKind) -> bool {
1354        self.begin_root_turn();
1355        let mut guard = self.active_turn_context.write();
1356        let Some(context) = guard.as_mut() else {
1357            return false;
1358        };
1359        context.reserve_speculative_llm_call_for(kind)
1360    }
1361
1362    fn branch_context_preview(&self) -> String {
1363        let context = self.build_context_with_overlays();
1364        let mut value = serde_json::to_string_pretty(&context).unwrap_or_else(|_| "{}".to_string());
1365        const MAX_CONTEXT_PREVIEW_CHARS: usize = 2048;
1366        if value.chars().count() > MAX_CONTEXT_PREVIEW_CHARS {
1367            value = value
1368                .chars()
1369                .take(MAX_CONTEXT_PREVIEW_CHARS)
1370                .collect::<String>();
1371            value.push_str("...");
1372        }
1373        value
1374    }
1375
1376    /// Applies freshness policy before rendering the next prompt.
1377    async fn await_background_before_next_turn(&self) {
1378        let optimization = &self.runtime_config.optimization;
1379        if !optimization.enabled {
1380            return;
1381        }
1382        let actor_id = self.effective_actor_id();
1383        let post = &optimization.post_turn;
1384        self.await_background_task(
1385            post.facts.await_before_next_turn,
1386            RuntimeTaskPurpose::PostTurnFacts,
1387            actor_id.as_deref(),
1388            "facts",
1389        )
1390        .await;
1391        self.await_background_task(
1392            post.relationships.await_before_next_turn,
1393            RuntimeTaskPurpose::PostTurnRelationship,
1394            actor_id.as_deref(),
1395            "relationships",
1396        )
1397        .await;
1398    }
1399
1400    async fn await_background_task(
1401        &self,
1402        policy: AwaitBeforeNextTurn,
1403        purpose: RuntimeTaskPurpose,
1404        actor_id: Option<&str>,
1405        label: &str,
1406    ) {
1407        match policy {
1408            AwaitBeforeNextTurn::Never => {}
1409            AwaitBeforeNextTurn::Always => {
1410                if let Err(error) = self.flush_background_tasks_for_purpose(purpose).await {
1411                    warn!(label = label, error = %error, "background maintenance flush failed");
1412                }
1413            }
1414            AwaitBeforeNextTurn::SameActor => {
1415                if let Some(actor_id) = actor_id
1416                    && let Err(error) = self
1417                        .flush_background_tasks_for_actor_purpose(actor_id, purpose)
1418                        .await
1419                {
1420                    warn!(label = label, actor_id = %actor_id, error = %error, "actor background maintenance flush failed");
1421                }
1422            }
1423        }
1424    }
1425
1426    /// Runs post-turn facts and relationship maintenance according to runtime policy.
1427    async fn run_post_turn_maintenance(&self) -> Result<()> {
1428        let optimization = &self.runtime_config.optimization;
1429        if !optimization.enabled {
1430            self.auto_extract_facts().await;
1431            self.auto_update_relationship().await;
1432            return Ok(());
1433        }
1434
1435        let facts_mode = effective_maintenance_mode(
1436            optimization.post_turn.facts.mode,
1437            optimization.parallel_post_turn_memory,
1438        );
1439        let relationships_mode = effective_maintenance_mode(
1440            optimization.post_turn.relationships.mode,
1441            optimization.parallel_post_turn_memory,
1442        );
1443
1444        match (facts_mode, relationships_mode) {
1445            (MaintenanceMode::InlineSerial, MaintenanceMode::InlineSerial) => {
1446                self.auto_extract_facts().await;
1447                self.auto_update_relationship().await;
1448            }
1449            (MaintenanceMode::InlineParallel, MaintenanceMode::InlineParallel) => {
1450                let facts = self.auto_extract_facts();
1451                let relationships = self.auto_update_relationship();
1452                tokio::join!(facts, relationships);
1453            }
1454            (MaintenanceMode::Background, MaintenanceMode::Background) => {
1455                self.schedule_facts_background().await?;
1456                self.schedule_relationship_background().await?;
1457            }
1458            (MaintenanceMode::Background, MaintenanceMode::InlineParallel)
1459            | (MaintenanceMode::Background, MaintenanceMode::InlineSerial) => {
1460                self.schedule_facts_background().await?;
1461                self.auto_update_relationship().await;
1462            }
1463            (MaintenanceMode::InlineParallel, MaintenanceMode::Background)
1464            | (MaintenanceMode::InlineSerial, MaintenanceMode::Background) => {
1465                self.auto_extract_facts().await;
1466                self.schedule_relationship_background().await?;
1467            }
1468            _ => {
1469                self.auto_extract_facts().await;
1470                self.auto_update_relationship().await;
1471            }
1472        }
1473        Ok(())
1474    }
1475
1476    async fn schedule_facts_background(&self) -> Result<()> {
1477        let policy = self.runtime_config.optimization.post_turn.facts.clone();
1478        let should_extract = self
1479            .facts_config
1480            .as_ref()
1481            .map(|c| c.enabled && c.auto_extract)
1482            .unwrap_or(false);
1483        if !should_extract {
1484            return Ok(());
1485        }
1486        let msgs_since = *self.messages_since_extraction.read();
1487        if msgs_since < 2 {
1488            return Ok(());
1489        }
1490        let Some(actor_id) = self.effective_actor_id() else {
1491            self.record_skipped_maintenance(
1492                "facts",
1493                ObservationPurpose::FactsExtraction,
1494                "missing_actor",
1495                Some(&policy),
1496            );
1497            return Ok(());
1498        };
1499        let Some(extractor) = self.fact_extractor.read().clone() else {
1500            return Ok(());
1501        };
1502        let messages = match self.memory.get_messages(None).await {
1503            Ok(messages) => messages,
1504            Err(error) => {
1505                warn!(error = %error, "failed to snapshot messages for fact extraction");
1506                return Ok(());
1507            }
1508        };
1509        let recent: Vec<_> = messages
1510            .iter()
1511            .rev()
1512            .take(msgs_since)
1513            .rev()
1514            .cloned()
1515            .collect();
1516        if recent.is_empty() {
1517            return Ok(());
1518        }
1519        let existing = self
1520            .actor_facts_cache
1521            .read()
1522            .get(&actor_id)
1523            .cloned()
1524            .unwrap_or_default();
1525        let categories = self
1526            .facts_config
1527            .as_ref()
1528            .map(|c| c.custom_categories.clone())
1529            .unwrap_or_default();
1530        let store = self.fact_store.read().clone();
1531        let cache = Arc::clone(&self.actor_facts_cache);
1532        let counter = Arc::clone(&self.messages_since_extraction);
1533        let hooks = Arc::clone(&self.hooks);
1534        let agent_id = self.info.id.clone();
1535        let observation = current_observation_context();
1536        let key = MaintenanceSequenceKey::actor(
1537            agent_id,
1538            actor_id.clone(),
1539            RuntimeTaskPurpose::PostTurnFacts,
1540        );
1541        let actor_for_task = actor_id.clone();
1542        let task = async move {
1543            let run = async move {
1544                let facts = extractor
1545                    .extract(&recent, &existing, Some(&actor_for_task), &categories)
1546                    .await?;
1547                if !facts.is_empty() {
1548                    if let Some(store) = store {
1549                        let authoritative = store.add_facts(&actor_for_task, facts.clone()).await?;
1550                        cache.write().insert(actor_for_task.clone(), authoritative);
1551                    } else {
1552                        cache
1553                            .write()
1554                            .entry(actor_for_task.clone())
1555                            .or_default()
1556                            .extend(facts.clone());
1557                    }
1558                    {
1559                        let mut count = counter.write();
1560                        if *count <= msgs_since {
1561                            *count = 0;
1562                        } else {
1563                            *count -= msgs_since;
1564                        }
1565                    }
1566                    hooks.on_facts_extracted(&actor_for_task, &facts).await;
1567                }
1568                Ok(())
1569            };
1570            if let Some(context) = observation {
1571                with_observation_context(
1572                    context.with_purpose(ObservationPurpose::FactsExtraction),
1573                    run,
1574                )
1575                .await
1576            } else {
1577                run.await
1578            }
1579        };
1580        self.spawn_or_handle_background(Some(key), task, "facts", &policy)
1581            .await
1582    }
1583
1584    async fn schedule_relationship_background(&self) -> Result<()> {
1585        let policy = self
1586            .runtime_config
1587            .optimization
1588            .post_turn
1589            .relationships
1590            .clone();
1591        let Some(manager) = self.relationship_manager.as_ref().cloned() else {
1592            return Ok(());
1593        };
1594        let Some(actor_id) = self.effective_actor_id() else {
1595            self.record_skipped_maintenance(
1596                "relationships",
1597                ObservationPurpose::RelationshipUpdate,
1598                "missing_actor",
1599                Some(&policy),
1600            );
1601            return Ok(());
1602        };
1603        let recent_messages = manager.config().auto_update.recent_messages;
1604        let messages = match self.memory.get_messages(Some(recent_messages)).await {
1605            Ok(messages) => messages,
1606            Err(error) => {
1607                warn!(actor = %actor_id, error = %error, "failed to snapshot messages for relationship update");
1608                return Ok(());
1609            }
1610        };
1611        let storage = self.storage.read().clone();
1612        let hooks = Arc::clone(&self.hooks);
1613        let agent_id = self.info.id.clone();
1614        let observation = current_observation_context();
1615        let key = MaintenanceSequenceKey::actor(
1616            agent_id.clone(),
1617            actor_id.clone(),
1618            RuntimeTaskPurpose::PostTurnRelationship,
1619        );
1620        let actor_for_task = actor_id.clone();
1621        let task = async move {
1622            let run = async move {
1623                if manager.config().auto_update.enabled {
1624                    let update = manager.auto_update(&actor_for_task, &messages).await?;
1625                    if !update.changes.is_empty() {
1626                        hooks
1627                            .on_relationship_change(&actor_for_task, &update.changes)
1628                            .await;
1629                    }
1630                    if let Some(ref event) = update.event {
1631                        hooks.on_notable_event(&actor_for_task, event).await;
1632                    }
1633                }
1634                if manager.config().persistence.enabled
1635                    && let (Some(storage), Some(value)) =
1636                        (storage, manager.relationship_as_value(&actor_for_task)?)
1637                {
1638                    storage
1639                        .save_relationship(&agent_id, &actor_for_task, &value)
1640                        .await?;
1641                }
1642                Ok(())
1643            };
1644            if let Some(context) = observation {
1645                with_observation_context(
1646                    context.with_purpose(ObservationPurpose::RelationshipUpdate),
1647                    run,
1648                )
1649                .await
1650            } else {
1651                run.await
1652            }
1653        };
1654        self.spawn_or_handle_background(Some(key), task, "relationships", &policy)
1655            .await
1656    }
1657
1658    /// Queues background maintenance or applies the configured overflow behavior.
1659    async fn spawn_or_handle_background<F>(
1660        &self,
1661        key: Option<MaintenanceSequenceKey>,
1662        task: F,
1663        label: &'static str,
1664        policy: &crate::optimization::config::MaintenanceTaskPolicy,
1665    ) -> Result<()>
1666    where
1667        F: Future<Output = Result<()>> + Send + 'static,
1668    {
1669        if self.background_maintenance.is_full() {
1670            match self
1671                .runtime_config
1672                .optimization
1673                .post_turn
1674                .on_background_overflow
1675            {
1676                BackgroundOverflowPolicy::RunInline => {
1677                    record_background_maintenance_event(
1678                        self.observability_manager.as_ref(),
1679                        label,
1680                        EventStatus::Success,
1681                        0,
1682                        "inline_overflow",
1683                        None,
1684                        Some(policy),
1685                    );
1686                    let start = Instant::now();
1687                    match task.await {
1688                        Ok(()) => record_background_maintenance_event(
1689                            self.observability_manager.as_ref(),
1690                            label,
1691                            EventStatus::Success,
1692                            start.elapsed().as_millis() as u64,
1693                            "inline_completed",
1694                            None,
1695                            Some(policy),
1696                        ),
1697                        Err(error) => {
1698                            warn!(label = label, error = %error, "inline maintenance fallback failed");
1699                            record_background_maintenance_event(
1700                                self.observability_manager.as_ref(),
1701                                label,
1702                                EventStatus::Error,
1703                                start.elapsed().as_millis() as u64,
1704                                "inline_failed",
1705                                Some(error.to_string()),
1706                                Some(policy),
1707                            );
1708                            return Err(error);
1709                        }
1710                    }
1711                }
1712                BackgroundOverflowPolicy::Drop => {
1713                    self.record_skipped_maintenance(
1714                        label,
1715                        ObservationPurpose::Other(label.to_string()),
1716                        "queue_full",
1717                        Some(policy),
1718                    );
1719                }
1720                BackgroundOverflowPolicy::Error => {
1721                    record_background_maintenance_event(
1722                        self.observability_manager.as_ref(),
1723                        label,
1724                        EventStatus::Error,
1725                        0,
1726                        "queue_full",
1727                        None,
1728                        Some(policy),
1729                    );
1730                    warn!(label = label, "background maintenance queue full");
1731                    return Err(AgentError::Other(format!(
1732                        "background maintenance queue is full for {}",
1733                        label
1734                    )));
1735                }
1736            }
1737            return Ok(());
1738        }
1739
1740        record_background_maintenance_event(
1741            self.observability_manager.as_ref(),
1742            label,
1743            EventStatus::Success,
1744            0,
1745            "scheduled",
1746            None,
1747            Some(policy),
1748        );
1749        let manager = self.observability_manager.clone();
1750        let policy_for_task = policy.clone();
1751        let observed_task = async move {
1752            let start = Instant::now();
1753            let result = task.await;
1754            match &result {
1755                Ok(()) => record_background_maintenance_event(
1756                    manager.as_ref(),
1757                    label,
1758                    EventStatus::Success,
1759                    start.elapsed().as_millis() as u64,
1760                    "completed",
1761                    None,
1762                    Some(&policy_for_task),
1763                ),
1764                Err(error) => record_background_maintenance_event(
1765                    manager.as_ref(),
1766                    label,
1767                    EventStatus::Error,
1768                    start.elapsed().as_millis() as u64,
1769                    "failed",
1770                    Some(error.to_string()),
1771                    Some(&policy_for_task),
1772                ),
1773            }
1774            result
1775        };
1776
1777        if let Err(error) = self.background_maintenance.spawn(key, observed_task) {
1778            record_background_maintenance_event(
1779                self.observability_manager.as_ref(),
1780                label,
1781                EventStatus::Error,
1782                0,
1783                "spawn_failed",
1784                Some(error.to_string()),
1785                Some(policy),
1786            );
1787            warn!(label = label, error = %error, "background maintenance spawn failed");
1788            return Err(error);
1789        }
1790        Ok(())
1791    }
1792
1793    /// Records a skipped background maintenance event when work cannot run.
1794    fn record_skipped_maintenance(
1795        &self,
1796        label: &str,
1797        purpose: ObservationPurpose,
1798        reason: &str,
1799        policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
1800    ) {
1801        if let Some(manager) = self.observability_manager.as_ref() {
1802            let mut tags = background_maintenance_tags(label, "skipped", Some(reason), policy);
1803            tags.insert("runtime.skip_reason".to_string(), reason.to_string());
1804            manager.record_lifecycle_event(
1805                EventType::MemoryOperation {
1806                    operation: format!("{}_maintenance", label),
1807                },
1808                purpose,
1809                EventStatus::Skipped,
1810                0,
1811                tags,
1812                None,
1813            );
1814        }
1815    }
1816
1817    /// Get cached actor facts for the effective actor.
1818    pub fn actor_facts(&self) -> Vec<ai_agents_core::KeyFact> {
1819        let Some(actor_id) = self.effective_actor_id() else {
1820            return Vec::new();
1821        };
1822        self.actor_facts_cache
1823            .read()
1824            .get(&actor_id)
1825            .cloned()
1826            .unwrap_or_default()
1827    }
1828
1829    /// Returns the formatted relationship prompt text for the effective actor, if relationship injection produced any text for this turn.
1830    pub fn relationship_memory_text(&self) -> Option<String> {
1831        self.format_relationship_for_context().map(|(_, text)| text)
1832    }
1833
1834    /// Manually extract facts from the last N messages.
1835    pub async fn extract_facts(
1836        &self,
1837        last_n: usize,
1838    ) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
1839        self.extract_facts_with_source(last_n, "manual").await
1840    }
1841
1842    async fn extract_facts_with_source(
1843        &self,
1844        last_n: usize,
1845        source: &'static str,
1846    ) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
1847        let extractor = match self.fact_extractor.read().clone() {
1848            Some(e) => e,
1849            None => return Ok(vec![]),
1850        };
1851
1852        let messages = self.memory.get_messages(None).await?;
1853        let recent: Vec<_> = messages.iter().rev().take(last_n).rev().cloned().collect();
1854
1855        if recent.is_empty() {
1856            return Ok(vec![]);
1857        }
1858
1859        let actor_id = self.effective_actor_id();
1860        let existing = actor_id
1861            .as_ref()
1862            .and_then(|aid| self.actor_facts_cache.read().get(aid).cloned())
1863            .unwrap_or_default();
1864
1865        let categories = self
1866            .facts_config
1867            .as_ref()
1868            .map(|c| c.custom_categories.clone())
1869            .unwrap_or_default();
1870
1871        let facts = self
1872            .observe_purpose(
1873                ObservationPurpose::FactsExtraction,
1874                extractor.extract(&recent, &existing, actor_id.as_deref(), &categories),
1875            )
1876            .await?;
1877
1878        // Save to storage and update the actor-scoped cache when an actor is known.
1879        if !facts.is_empty() {
1880            let fact_store_opt = self.fact_store.read().clone();
1881            let mut stored_total = 0usize;
1882            let mut cache_updated = false;
1883            if let (Some(store), Some(aid)) = (fact_store_opt, &actor_id) {
1884                // add_facts now returns the authoritative post-write set.
1885                let authoritative = store.add_facts(aid, facts.clone()).await?;
1886                stored_total = authoritative.len();
1887                self.actor_facts_cache
1888                    .write()
1889                    .insert(aid.clone(), authoritative);
1890                cache_updated = true;
1891            } else if let Some(aid) = &actor_id {
1892                let mut cache = self.actor_facts_cache.write();
1893                let entry = cache.entry(aid.clone()).or_default();
1894                entry.extend(facts.clone());
1895                stored_total = entry.len();
1896                cache_updated = true;
1897            }
1898
1899            info!(
1900                actor_id = %actor_id.as_deref().unwrap_or("<none>"),
1901                source = source,
1902                requested_messages = last_n,
1903                message_count = recent.len(),
1904                extracted_count = facts.len(),
1905                cache_updated = cache_updated,
1906                stored_total = stored_total,
1907                "facts extracted"
1908            );
1909
1910            if let Some(ref aid) = actor_id {
1911                self.hooks.on_facts_extracted(aid, &facts).await;
1912            }
1913        }
1914
1915        Ok(facts)
1916    }
1917
1918    /// Resolve actor_id from context if method is from_context.
1919    /// Supports dotted paths (e.g. "player.id", "user.profile.id").
1920    fn resolve_actor_id_from_context(&self) {
1921        if self
1922            .current_turn_actor_context()
1923            .and_then(|ctx| ctx.effective_actor_id().map(str::to_string))
1924            .is_some()
1925        {
1926            return;
1927        }
1928
1929        if let Some(ref am_config) = self.actor_memory_config
1930            && am_config.identification.method == ai_agents_facts::IdentificationMethod::FromContext
1931            && let Some(ref path) = am_config.identification.context_path
1932        {
1933            // get_path resolves dotted paths; get only handles top-level keys.
1934            let val = self
1935                .context_manager
1936                .get_path(path)
1937                .or_else(|| self.context_manager.get(path));
1938            if let Some(val) = val
1939                && let Some(id_str) = val.as_str()
1940            {
1941                let current = self.actor_id.read().clone();
1942                if current.as_deref() != Some(id_str) {
1943                    *self.actor_id.write() = Some(id_str.to_string());
1944                    let mut meta = self.session_metadata.write();
1945                    meta.actor_id = Some(id_str.to_string());
1946                    if !meta.actors.iter().any(|a| a == id_str) {
1947                        meta.actors.push(id_str.to_string());
1948                    }
1949                }
1950            }
1951        }
1952    }
1953
1954    /// Format actor facts for template injection.
1955    fn format_actor_facts_for_context(&self) -> String {
1956        // Respect inject_in_context: false.
1957        let should_inject = self
1958            .facts_config
1959            .as_ref()
1960            .map(|c| c.inject_in_context)
1961            .unwrap_or(true);
1962        if !should_inject {
1963            return String::new();
1964        }
1965
1966        let Some(actor_id) = self.effective_actor_id() else {
1967            return String::new();
1968        };
1969
1970        let facts = self
1971            .actor_facts_cache
1972            .read()
1973            .get(&actor_id)
1974            .cloned()
1975            .unwrap_or_default();
1976        if facts.is_empty() {
1977            return String::new();
1978        }
1979
1980        let am_config = self.actor_memory_config.as_ref();
1981        // Effective token cap: prefer memory.token_budget.allocation.facts when present,
1982        // otherwise fall back to actor_memory.injection.max_tokens.
1983        let facts_budget = self
1984            .memory_token_budget
1985            .as_ref()
1986            .map(|b| b.allocation.facts as usize)
1987            .filter(|n| *n > 0);
1988        let default_max = am_config.map(|c| c.injection.max_tokens).unwrap_or(800);
1989        let max_tokens = facts_budget.unwrap_or(default_max);
1990
1991        // Filter by category when injection.mode = category.
1992        let filtered: Vec<ai_agents_core::KeyFact> = if let Some(cfg) = am_config {
1993            if cfg.injection.mode == ai_agents_facts::InjectionMode::OnDemand {
1994                return String::new();
1995            }
1996            if cfg.injection.mode == ai_agents_facts::InjectionMode::Category
1997                && !cfg.injection.categories.is_empty()
1998            {
1999                facts
2000                    .iter()
2001                    .filter(|f| {
2002                        cfg.injection
2003                            .categories
2004                            .iter()
2005                            .any(|c| f.category.to_string() == *c)
2006                    })
2007                    .cloned()
2008                    .collect()
2009            } else {
2010                facts.clone()
2011            }
2012        } else {
2013            facts.clone()
2014        };
2015
2016        if filtered.is_empty() {
2017            return String::new();
2018        }
2019
2020        if let Some(store) = self.fact_store.read().clone() {
2021            store.format_for_context(&filtered, max_tokens)
2022        } else {
2023            String::new()
2024        }
2025    }
2026
2027    fn build_context_with_staged(&self, staged: &HashMap<String, Value>) -> HashMap<String, Value> {
2028        let context = self.build_context_with_overlays();
2029        let mut root = Value::Object(context.into_iter().collect());
2030        for (path, value) in staged {
2031            if let Ok(updated) = ai_agents_core::set_dot_path(root.clone(), path, value.clone()) {
2032                root = updated;
2033            }
2034        }
2035        match root {
2036            Value::Object(obj) => obj.into_iter().collect(),
2037            _ => HashMap::new(),
2038        }
2039    }
2040
2041    fn build_context_with_overlays(&self) -> HashMap<String, Value> {
2042        let mut context = self.context_manager.get_all();
2043        let mut root = Value::Object(context.clone().into_iter().collect());
2044
2045        if let Some(turn_ctx) = self.current_turn_actor_context() {
2046            if let Some(ref origin_actor_id) = turn_ctx.origin_actor_id
2047                && let Ok(updated) = ai_agents_core::set_dot_path(
2048                    root.clone(),
2049                    "interaction.origin_actor_id",
2050                    serde_json::json!(origin_actor_id),
2051                )
2052            {
2053                root = updated;
2054            }
2055            if let Some(ref sender_agent_id) = turn_ctx.sender_agent_id
2056                && let Ok(updated) = ai_agents_core::set_dot_path(
2057                    root.clone(),
2058                    "interaction.sender_agent_id",
2059                    serde_json::json!(sender_agent_id),
2060                )
2061            {
2062                root = updated;
2063            }
2064        }
2065
2066        if let Some(ref actor_id) = self.effective_actor_id()
2067            && let Ok(updated) = ai_agents_core::set_dot_path(
2068                root.clone(),
2069                "interaction.actor_id",
2070                serde_json::json!(actor_id),
2071            )
2072        {
2073            root = updated;
2074        }
2075
2076        if let Some(manager) = self.relationship_manager.as_ref()
2077            && let Some(actor_id) = self.effective_actor_id()
2078            && let Some(value) = manager.to_context_value(&actor_id)
2079            && let Ok(updated) = ai_agents_core::set_dot_path(
2080                root.clone(),
2081                &manager.config().injection.context_path,
2082                value,
2083            )
2084        {
2085            root = updated;
2086        }
2087
2088        if let Value::Object(obj) = root {
2089            context = obj.into_iter().collect();
2090        }
2091
2092        context
2093    }
2094
2095    fn resolve_actor_name_from_context(&self) -> Option<String> {
2096        for path in ["actor.name", "user.name", "player.name", "customer.name"] {
2097            if let Some(value) = self.context_manager.get_path(path)
2098                && let Some(name) = value.as_str()
2099            {
2100                return Some(name.to_string());
2101            }
2102        }
2103        None
2104    }
2105
2106    async fn maybe_load_actor_relationship(&self) {
2107        let Some(manager) = self.relationship_manager.as_ref() else {
2108            return;
2109        };
2110        let Some(actor_id) = self.effective_actor_id() else {
2111            return;
2112        };
2113
2114        let mut should_fire_loaded = false;
2115        if manager.get(&actor_id).is_none() {
2116            let mut loaded = false;
2117            if manager.config().persistence.enabled {
2118                let storage = self.storage.read().clone();
2119                if let Some(storage) = storage {
2120                    match storage.load_relationship(&self.info.id, &actor_id).await {
2121                        Ok(Some(value)) => match manager.insert_from_value(value) {
2122                            Ok(_) => loaded = true,
2123                            Err(e) => {
2124                                warn!(actor = %actor_id, error = %e, "failed to restore relationship")
2125                            }
2126                        },
2127                        Ok(None) => {}
2128                        Err(e) => {
2129                            warn!(actor = %actor_id, error = %e, "failed to load relationship")
2130                        }
2131                    }
2132                }
2133            }
2134
2135            if !loaded {
2136                manager.get_or_create(&actor_id, self.resolve_actor_name_from_context().as_deref());
2137            }
2138            should_fire_loaded = true;
2139        }
2140
2141        let actor_name = self.resolve_actor_name_from_context();
2142        let relationship = manager.touch_interaction(&actor_id, actor_name.as_deref());
2143        if should_fire_loaded {
2144            self.hooks
2145                .on_relationship_loaded(&actor_id, &relationship)
2146                .await;
2147        }
2148    }
2149
2150    fn format_relationship_for_context(&self) -> Option<(String, String)> {
2151        let manager = self.relationship_manager.as_ref()?;
2152        if !manager.config().injection.enabled {
2153            return None;
2154        }
2155        let actor_id = self.effective_actor_id()?;
2156        let relationship = manager.get(&actor_id)?;
2157        let local_cap = manager.config().injection.max_tokens;
2158        let global_cap = self
2159            .memory_token_budget
2160            .as_ref()
2161            .map(|b| b.allocation.relationships as usize)
2162            .filter(|n| *n > 0);
2163        let max_tokens = global_cap.map(|g| g.min(local_cap)).unwrap_or(local_cap);
2164        let text = ai_agents_relationships::format_relationship(
2165            &relationship,
2166            &manager.config().injection.format,
2167            max_tokens,
2168        );
2169        if text.is_empty() {
2170            None
2171        } else {
2172            Some((manager.config().injection.prompt_variable.clone(), text))
2173        }
2174    }
2175
2176    async fn persist_actor_relationship(&self, actor_id: &str) -> Result<()> {
2177        let Some(manager) = self.relationship_manager.as_ref() else {
2178            return Ok(());
2179        };
2180        if !manager.config().persistence.enabled {
2181            return Ok(());
2182        }
2183        let storage = self.storage.read().clone();
2184        let Some(storage) = storage else {
2185            return Ok(());
2186        };
2187        if let Some(value) = manager.relationship_as_value(actor_id)? {
2188            storage
2189                .save_relationship(&self.info.id, actor_id, &value)
2190                .await?;
2191        }
2192        Ok(())
2193    }
2194
2195    async fn auto_update_relationship(&self) {
2196        let Some(manager) = self.relationship_manager.as_ref() else {
2197            return;
2198        };
2199        let Some(actor_id) = self.effective_actor_id() else {
2200            return;
2201        };
2202        if !manager.config().auto_update.enabled {
2203            let _ = self.persist_actor_relationship(&actor_id).await;
2204            return;
2205        }
2206
2207        let recent_messages = manager.config().auto_update.recent_messages;
2208        let messages = match self.memory.get_messages(Some(recent_messages)).await {
2209            Ok(messages) => messages,
2210            Err(e) => {
2211                warn!(actor = %actor_id, error = %e, "failed to read messages for relationship update");
2212                return;
2213            }
2214        };
2215
2216        match self
2217            .observe_purpose(
2218                ObservationPurpose::RelationshipUpdate,
2219                manager.auto_update(&actor_id, &messages),
2220            )
2221            .await
2222        {
2223            Ok(update) => {
2224                if !update.changes.is_empty() {
2225                    self.hooks
2226                        .on_relationship_change(&actor_id, &update.changes)
2227                        .await;
2228                }
2229                if let Some(ref event) = update.event {
2230                    self.hooks.on_notable_event(&actor_id, event).await;
2231                }
2232                let persisted = match self.persist_actor_relationship(&actor_id).await {
2233                    Ok(()) => true,
2234                    Err(e) => {
2235                        warn!(actor = %actor_id, error = %e, "failed to persist relationship");
2236                        false
2237                    }
2238                };
2239                if !update.changes.is_empty() || update.event.is_some() {
2240                    let changed_dimensions: Vec<String> = update
2241                        .changes
2242                        .iter()
2243                        .map(|change| format!("{}:{}", change.perspective, change.dimension))
2244                        .collect();
2245                    info!(
2246                        actor_id = %actor_id,
2247                        change_count = update.changes.len(),
2248                        changed_dimensions = ?changed_dimensions,
2249                        event_present = update.event.is_some(),
2250                        persisted = persisted,
2251                        "relationship updated"
2252                    );
2253                } else {
2254                    debug!(actor_id = %actor_id, persisted = persisted, "relationship evaluation ran but found no changes");
2255                }
2256            }
2257            Err(e) => warn!(actor = %actor_id, error = %e, "relationship update failed"),
2258        }
2259    }
2260
2261    /// Run auto-extraction after a chat turn.
2262    async fn auto_extract_facts(&self) {
2263        let should_extract = self
2264            .facts_config
2265            .as_ref()
2266            .map(|c| c.enabled && c.auto_extract)
2267            .unwrap_or(false);
2268
2269        if !should_extract {
2270            debug!("fact extraction skipped because auto extraction is disabled");
2271            return;
2272        }
2273
2274        let msgs_since = *self.messages_since_extraction.read();
2275        if msgs_since < 2 {
2276            debug!(
2277                messages_since_extraction = msgs_since,
2278                "fact extraction skipped until threshold is reached"
2279            );
2280            return;
2281        }
2282
2283        match self.extract_facts_with_source(msgs_since, "auto").await {
2284            Ok(facts) => {
2285                if !facts.is_empty() {
2286                    *self.messages_since_extraction.write() = 0;
2287                } else {
2288                    debug!("fact extraction ran but found no new facts");
2289                }
2290            }
2291            Err(e) => {
2292                warn!("fact extraction failed: {}", e);
2293            }
2294        }
2295    }
2296
2297    pub fn with_persona(mut self, manager: Arc<ai_agents_persona::PersonaManager>) -> Self {
2298        self.persona_manager = Some(manager);
2299        self
2300    }
2301
2302    pub fn persona_manager(&self) -> Option<&Arc<ai_agents_persona::PersonaManager>> {
2303        self.persona_manager.as_ref()
2304    }
2305
2306    pub fn with_disambiguation(mut self, config: DisambiguationConfig) -> Self {
2307        if config.is_enabled() {
2308            let manager = DisambiguationManager::new(config, Arc::clone(&self.llm_registry))
2309                .with_clarification_observer(Arc::new(ObservabilityClarificationObserver));
2310            self.disambiguation_manager = Some(manager);
2311        }
2312        self
2313    }
2314
2315    pub fn disambiguation_manager(&self) -> Option<&DisambiguationManager> {
2316        self.disambiguation_manager.as_ref()
2317    }
2318
2319    pub fn has_disambiguation(&self) -> bool {
2320        self.disambiguation_manager
2321            .as_ref()
2322            .is_some_and(|m| m.is_enabled())
2323    }
2324
2325    pub async fn init_storage(&self) -> Result<()> {
2326        //
2327        // One readiness guard prevents concurrent entry points from constructing storage or fact state more than once.
2328        //
2329        let _guard = self.storage_init.lock().await;
2330        let mut storage = self.storage.read().clone();
2331        if storage.is_none() && !self.storage_config.is_none() {
2332            let storage_config = self.convert_storage_config();
2333            storage = create_storage(&storage_config).await?;
2334            *self.storage.write() = storage.clone();
2335        }
2336
2337        self.validate_storage_requirements(storage.as_deref())?;
2338        self.complete_facts_init().await;
2339        Ok(())
2340    }
2341
2342    fn validate_storage_requirements(&self, storage: Option<&dyn AgentStorage>) -> Result<()> {
2343        let facts_required = self
2344            .facts_config
2345            .as_ref()
2346            .is_some_and(|config| config.enabled)
2347            || self
2348                .actor_memory_config
2349                .as_ref()
2350                .is_some_and(|config| config.enabled);
2351        let relationships_required = self
2352            .relationship_manager
2353            .as_ref()
2354            .is_some_and(|manager| manager.config().persistence.enabled);
2355
2356        let Some(storage) = storage else {
2357            let mut requirements = Vec::new();
2358            if facts_required {
2359                requirements.push("actor facts or actor memory");
2360            }
2361            if relationships_required {
2362                requirements.push("persistent relationships");
2363            }
2364            if requirements.is_empty() {
2365                return Ok(());
2366            }
2367            return Err(AgentError::Config(format!(
2368                "Storage is required for enabled {} but none is configured or injected",
2369                requirements.join(" and ")
2370            )));
2371        };
2372
2373        //
2374        // Durable features must fail before execution instead of degrading to volatile state.
2375        //
2376        if facts_required && !storage.supports(StorageCapability::ActorFacts) {
2377            return Err(AgentError::UnsupportedStorageCapability(
2378                StorageCapability::ActorFacts,
2379            ));
2380        }
2381        if relationships_required && !storage.supports(StorageCapability::ActorRelationships) {
2382            return Err(AgentError::UnsupportedStorageCapability(
2383                StorageCapability::ActorRelationships,
2384            ));
2385        }
2386        Ok(())
2387    }
2388
2389    /// Initialize fact store and extractor from stored config and current storage.
2390    /// Called from init_storage() so facts are ready before the first turn.
2391    async fn complete_facts_init(&self) {
2392        if self.fact_store.read().is_some() {
2393            return;
2394        }
2395        let storage = match self.storage.read().clone() {
2396            Some(s) => s,
2397            None => return,
2398        };
2399
2400        let facts_enabled = self
2401            .facts_config
2402            .as_ref()
2403            .map(|f| f.enabled)
2404            .unwrap_or(false);
2405        let actor_memory_enabled = self
2406            .actor_memory_config
2407            .as_ref()
2408            .map(|a| a.enabled)
2409            .unwrap_or(false);
2410
2411        if !facts_enabled && !actor_memory_enabled {
2412            return;
2413        }
2414
2415        let fc = self.facts_config.clone().unwrap_or_default();
2416        let store = Arc::new(ai_agents_facts::FactStore::new(
2417            storage,
2418            self.info.id.clone(),
2419            fc.clone(),
2420        ));
2421
2422        let extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>> = if facts_enabled {
2423            let extractor_llm = fc
2424                .extractor_llm
2425                .as_ref()
2426                .and_then(|alias| self.llm_registry.get(alias).ok())
2427                .or_else(|| self.llm_registry.router().ok())
2428                .or_else(|| self.llm_registry.default().ok());
2429            extractor_llm.map(|llm| {
2430                Arc::new(ai_agents_facts::LLMFactExtractor::new(llm, fc.clone()))
2431                    as Arc<dyn ai_agents_facts::FactExtractor>
2432            })
2433        } else {
2434            None
2435        };
2436
2437        *self.fact_store.write() = Some(store);
2438        *self.fact_extractor.write() = extractor;
2439        debug!(
2440            agent = %self.info.id,
2441            facts_enabled,
2442            actor_memory_enabled,
2443            "facts storage initialized"
2444        );
2445    }
2446
2447    fn convert_storage_config(&self) -> StorageStorageConfig {
2448        crate::spec::storage::to_storage_config(&self.storage_config)
2449    }
2450
2451    pub fn storage(&self) -> Option<Arc<dyn AgentStorage>> {
2452        self.storage.read().clone()
2453    }
2454
2455    pub fn storage_config(&self) -> &StorageConfig {
2456        &self.storage_config
2457    }
2458
2459    /// Returns the spawner if configured via a spawner: YAML section.
2460    pub fn spawner(&self) -> Option<&Arc<crate::spawner::AgentSpawner>> {
2461        self.spawner.as_ref()
2462    }
2463
2464    /// Returns the agent registry if configured via a spawner: YAML section.
2465    pub fn spawner_registry(&self) -> Option<&Arc<crate::spawner::AgentRegistry>> {
2466        self.spawner_registry.as_ref()
2467    }
2468
2469    pub fn has_spawner(&self) -> bool {
2470        self.spawner_registry.is_some()
2471    }
2472
2473    pub fn with_spawner_handles(
2474        mut self,
2475        spawner: Arc<crate::spawner::AgentSpawner>,
2476        registry: Arc<crate::spawner::AgentRegistry>,
2477    ) -> Self {
2478        self.spawner = Some(spawner);
2479        self.spawner_registry = Some(registry);
2480        self
2481    }
2482
2483    pub fn with_hooks(mut self, hooks: Arc<dyn AgentHooks>) -> Self {
2484        self.hooks = hooks;
2485        self
2486    }
2487
2488    pub fn with_parallel_tools(mut self, config: ParallelToolsConfig) -> Self {
2489        self.parallel_tools = config;
2490        self
2491    }
2492
2493    pub fn with_streaming(mut self, config: StreamingConfig) -> Self {
2494        self.streaming = config;
2495        self
2496    }
2497
2498    pub fn with_hitl(mut self, engine: HITLEngine, handler: Arc<dyn ApprovalHandler>) -> Self {
2499        self.hitl_engine = Some(engine);
2500        self.approval_handler = handler;
2501        self
2502    }
2503
2504    pub fn with_max_context_tokens(mut self, tokens: u32) -> Self {
2505        self.max_context_tokens = tokens;
2506        self
2507    }
2508
2509    pub fn with_memory_token_budget(mut self, budget: MemoryTokenBudget) -> Self {
2510        self.memory_token_budget = Some(budget);
2511        self
2512    }
2513
2514    pub fn with_recovery_manager(mut self, manager: RecoveryManager) -> Self {
2515        self.recovery_manager = manager;
2516        self
2517    }
2518
2519    pub fn with_tool_security(mut self, engine: ToolSecurityEngine) -> Self {
2520        self.tool_security = engine;
2521        self
2522    }
2523
2524    /// Returns the host-only runtime control handle.
2525    pub fn runtime_control(&self) -> RuntimeControlHandle {
2526        RuntimeControlHandle {
2527            state: Arc::clone(&self.runtime_control),
2528        }
2529    }
2530
2531    /// Installs or clears the host question handler used by `ask_user`.
2532    pub fn set_question_handler(&self, handler: Option<Arc<dyn QuestionHandler>>) {
2533        self.tools.set_question_handler(handler);
2534    }
2535
2536    /// Installs the host diagnostics provider used by `diagnostics`.
2537    pub fn set_diagnostics_provider(&self, provider: Arc<dyn DiagnosticsProvider>) {
2538        self.tools.set_diagnostics_provider(provider);
2539    }
2540
2541    /// Installs the host command runner used by `command`.
2542    pub fn set_command_runner(&self, runner: Arc<dyn CommandRunner>) {
2543        self.tools.set_command_runner(runner);
2544    }
2545
2546    /// Installs the host web-search provider used by `web_search`.
2547    pub fn set_web_search_provider(&self, provider: Arc<dyn ai_agents_tools::WebSearchProvider>) {
2548        self.tools.set_web_search_provider(provider);
2549    }
2550
2551    /// Returns the current session-local todo list.
2552    pub fn todos(&self) -> Vec<TodoItem> {
2553        self.tools.todos()
2554    }
2555
2556    /// Returns the effective tool security snapshot for the next decision.
2557    fn active_tool_security(&self) -> ToolSecurityEngine {
2558        self.runtime_control
2559            .tool_security_override
2560            .read()
2561            .clone()
2562            .unwrap_or_else(|| self.tool_security.clone())
2563    }
2564
2565    /// Reads policy, scope, emergency state, and generation as one safety snapshot.
2566    fn runtime_safety_snapshot(&self) -> RuntimeSafetySnapshot {
2567        let _guard = self.runtime_control.snapshot_guard.read();
2568        RuntimeSafetySnapshot {
2569            version: self.runtime_control.version.load(Ordering::SeqCst),
2570            emergency_deny: self.runtime_control.emergency_deny.load(Ordering::SeqCst),
2571            tool_security: self
2572                .runtime_control
2573                .tool_security_override
2574                .read()
2575                .clone()
2576                .unwrap_or_else(|| self.tool_security.clone()),
2577            tool_scope_override: self.runtime_control.tool_scope_override.read().clone(),
2578        }
2579    }
2580
2581    /// Performs the single rate admission after locks are held and rejects changed runtime, policy, or state authority.
2582    fn admit_tool_execution(
2583        &self,
2584        expected_runtime_version: u64,
2585        expected_policy_version: u64,
2586        expected_state_generation: Option<u64>,
2587        canonical_id: &str,
2588    ) -> SecurityCheckResult {
2589        let _guard = self.runtime_control.snapshot_guard.read();
2590        if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
2591            return SecurityCheckResult::Block {
2592                reason: "runtime emergency deny is enabled".to_string(),
2593            };
2594        }
2595        let runtime_version = self.runtime_control.version.load(Ordering::SeqCst);
2596        let security_engine = self
2597            .runtime_control
2598            .tool_security_override
2599            .read()
2600            .clone()
2601            .unwrap_or_else(|| self.tool_security.clone());
2602        if runtime_version != expected_runtime_version
2603            || security_engine.policy_version() != expected_policy_version
2604        {
2605            return SecurityCheckResult::Block {
2606                reason: "runtime safety controls changed before admission".to_string(),
2607            };
2608        }
2609        let current_state_generation = self
2610            .state_machine
2611            .as_ref()
2612            .map(|state_machine| state_machine.generation());
2613        if current_state_generation != expected_state_generation {
2614            return SecurityCheckResult::Block {
2615                reason: "state scope changed before admission".to_string(),
2616            };
2617        }
2618        security_engine.admit_tool_execution(canonical_id)
2619    }
2620
2621    pub fn with_process_processor(mut self, processor: ProcessProcessor) -> Self {
2622        let processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
2623        self.process_processor = Some(processor);
2624        self
2625    }
2626
2627    pub fn with_state_machine(
2628        mut self,
2629        state_machine: Arc<StateMachine>,
2630        evaluator: Arc<dyn TransitionEvaluator>,
2631    ) -> Self {
2632        self.state_machine = Some(state_machine);
2633        self.transition_evaluator = Some(evaluator);
2634        self
2635    }
2636
2637    pub fn with_context_manager(mut self, manager: Arc<ContextManager>) -> Self {
2638        self.context_manager = manager;
2639        self
2640    }
2641
2642    pub fn register_message_filter(&self, name: impl Into<String>, filter: Arc<dyn MessageFilter>) {
2643        self.message_filters.write().insert(name.into(), filter);
2644    }
2645
2646    pub fn set_context(&self, key: &str, value: Value) -> Result<()> {
2647        self.context_manager.update(key, value)
2648    }
2649
2650    pub fn update_context(&self, path: &str, value: Value) -> Result<()> {
2651        self.context_manager.update(path, value)
2652    }
2653
2654    pub fn get_context(&self) -> HashMap<String, Value> {
2655        self.build_context_with_overlays()
2656    }
2657
2658    pub fn remove_context(&self, key: &str) -> Option<Value> {
2659        self.context_manager.remove(key)
2660    }
2661
2662    pub async fn refresh_context(&self, key: &str) -> Result<()> {
2663        self.context_manager.refresh(key).await
2664    }
2665
2666    pub fn register_context_provider(&self, name: &str, provider: Arc<dyn ContextProvider>) {
2667        self.context_manager.register_provider(name, provider);
2668    }
2669
2670    pub fn current_state(&self) -> Option<String> {
2671        self.state_machine.as_ref().map(|sm| sm.current())
2672    }
2673
2674    // State changes invalidate only resolved confirmation ownership; ordinary clarification remains manager-owned across layered enablement changes.
2675    async fn invalidate_pending_confirmation(&self, reason: &'static str) {
2676        self.disambiguation_epoch.fetch_add(1, Ordering::SeqCst);
2677        let Some(disambiguator) = self.disambiguation_manager.as_ref() else {
2678            return;
2679        };
2680        if disambiguator.has_pending_confirmation().await {
2681            disambiguator.clear_pending().await;
2682            *self.pending_skill_id.write() = None;
2683            info!(
2684                confirmation_event = "invalidated",
2685                invalidation_reason = reason,
2686                "Runtime invalidated pending confirmation"
2687            );
2688        }
2689    }
2690
2691    // Admission linearizes pending publication or redispatch before reset and state mutation, then rechecks both ownership generations.
2692    async fn admit_disambiguation_redispatch(
2693        &self,
2694        expected_epoch: u64,
2695        expected_state_generation: Option<u64>,
2696    ) -> Result<tokio::sync::RwLockReadGuard<'_, ()>> {
2697        let admission = self.disambiguation_admission.read().await;
2698        let state_generation = self
2699            .state_machine
2700            .as_ref()
2701            .map(|state_machine| state_machine.generation());
2702        if self.disambiguation_epoch.load(Ordering::SeqCst) != expected_epoch
2703            || state_generation != expected_state_generation
2704        {
2705            return Err(AgentError::Other(
2706                "Disambiguation ownership changed before redispatch admission".to_string(),
2707            ));
2708        }
2709        Ok(admission)
2710    }
2711
2712    // A transition reservation prevents losing runtime transitions from running duplicate exit side effects.
2713    fn reserve_state_transition(&self) -> Option<StateTransitionReservation<'_>> {
2714        self.state_transition_reserved
2715            .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
2716            .ok()
2717            .map(|_| StateTransitionReservation {
2718                reserved: &self.state_transition_reserved,
2719            })
2720    }
2721
2722    // Optional ownership is used only for terminal skill responses that came from manager-owned pending state.
2723    async fn admit_optional_disambiguation_ownership(
2724        &self,
2725        ownership: Option<DisambiguationOwnership>,
2726    ) -> Result<Option<tokio::sync::RwLockReadGuard<'_, ()>>> {
2727        match ownership {
2728            Some(ownership) => self
2729                .admit_disambiguation_redispatch(ownership.epoch, ownership.state_generation)
2730                .await
2731                .map(Some),
2732            None => Ok(None),
2733        }
2734    }
2735
2736    /// Applies a manual transition after reserving its exit actions and keeping async lifecycle work outside the commit lock.
2737    pub async fn transition_to(&self, state: &str) -> Result<()> {
2738        let Some(ref sm) = self.state_machine else {
2739            return Ok(());
2740        };
2741        let claim_admission = self.disambiguation_admission.write().await;
2742        let reservation = self.reserve_state_transition().ok_or_else(|| {
2743            AgentError::Other("Another state transition is already in progress".to_string())
2744        })?;
2745        let from_state = sm.current();
2746        let expected_state_generation = sm.generation();
2747        let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
2748        let history_before = sm.history();
2749        drop(claim_admission);
2750
2751        self.execute_state_exit_actions(&from_state).await;
2752
2753        let admission = self.disambiguation_admission.write().await;
2754        if sm.current() != from_state
2755            || sm.generation() != expected_state_generation
2756            || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
2757        {
2758            return Err(AgentError::Other(
2759                "State ownership changed during manual transition preparation".to_string(),
2760            ));
2761        }
2762        sm.transition_to(state, "manual transition")?;
2763        self.invalidate_pending_confirmation("state_transition")
2764            .await;
2765        let entered = sm.current();
2766        let is_reentry = Self::state_was_previously_entered(&entered, &from_state, &history_before);
2767        drop(admission);
2768
2769        self.execute_state_enter_actions(&entered, is_reentry).await;
2770        drop(reservation);
2771        info!(to = %entered, "Manual state transition");
2772        Ok(())
2773    }
2774
2775    pub fn state_history(&self) -> Vec<StateTransitionEvent> {
2776        self.state_machine
2777            .as_ref()
2778            .map(|sm| sm.history())
2779            .unwrap_or_default()
2780    }
2781
2782    /// Get a copy of current session metadata.
2783    pub fn session_metadata(&self) -> ai_agents_core::SessionMetadata {
2784        self.session_metadata.read().clone()
2785    }
2786
2787    /// Delete all facts and sessions for an actor, gated by privacy.allow_deletion.
2788    /// Returns Err when actor_memory.privacy.allow_deletion is false.
2789    pub async fn delete_actor_data(&self, actor_id: &str) -> Result<()> {
2790        let allowed = self
2791            .actor_memory_config
2792            .as_ref()
2793            .map(|c| c.privacy.allow_deletion)
2794            .unwrap_or(true);
2795        if !allowed {
2796            return Err(AgentError::Config(
2797                "privacy.allow_deletion is false; actor data deletion is not permitted".into(),
2798            ));
2799        }
2800        let storage = self.storage.read().clone();
2801        if let Some(storage) = storage {
2802            //
2803            // Composite support is checked before mutation so privacy deletion cannot become partial.
2804            //
2805            if !storage.supports(StorageCapability::ActorDataDeletion) {
2806                return Err(AgentError::UnsupportedStorageCapability(
2807                    StorageCapability::ActorDataDeletion,
2808                ));
2809            }
2810            storage.delete_actor_data(&self.info.id, actor_id).await?;
2811        } else {
2812            //
2813            // Clone the fallback store before awaiting so backend I/O never holds the runtime read guard.
2814            //
2815            let store = { self.fact_store.read().clone() };
2816            if let Some(store) = store {
2817                store.delete_actor_data(actor_id).await?;
2818            }
2819        }
2820        if let Some(manager) = self.relationship_manager.as_ref() {
2821            manager.remove(actor_id);
2822        }
2823        self.actor_facts_cache.write().remove(actor_id);
2824        Ok(())
2825    }
2826
2827    /// Overwrite session metadata (tags, ttl, custom fields).
2828    pub fn set_session_metadata(&self, meta: ai_agents_core::SessionMetadata) {
2829        *self.session_metadata.write() = meta;
2830    }
2831
2832    /// Delete sessions whose TTL has expired. Returns number of sessions removed.
2833    pub async fn cleanup_expired_sessions(&self) -> Result<usize> {
2834        let storage = self.storage.read().clone();
2835        match storage {
2836            Some(s) => {
2837                let count = s.cleanup_expired().await?;
2838                if count > 0 {
2839                    self.hooks.on_sessions_expired(count).await;
2840                }
2841                Ok(count)
2842            }
2843            None => Err(AgentError::Config(
2844                "No storage configured. Use with_storage_config() or with_storage() first".into(),
2845            )),
2846        }
2847    }
2848
2849    /// List sessions matching a filter. Supports actor, tag, and date filters.
2850    pub async fn list_sessions_filtered(
2851        &self,
2852        filter: &ai_agents_core::SessionFilter,
2853    ) -> Result<Vec<ai_agents_core::SessionSummary>> {
2854        let storage = self.storage.read().clone();
2855        match storage {
2856            Some(s) => s.list_sessions_filtered(filter).await,
2857            None => Err(AgentError::Config(
2858                "No storage configured. Use with_storage_config() or with_storage() first".into(),
2859            )),
2860        }
2861    }
2862
2863    pub async fn save_state(&self) -> Result<AgentSnapshot> {
2864        let memory_snapshot = self.memory.snapshot().await?;
2865        let state_machine_snapshot = self.state_machine.as_ref().map(|sm| sm.snapshot());
2866        let context_snapshot = self.context_manager.snapshot();
2867
2868        let mut snapshot = AgentSnapshot::new(self.info.id.clone())
2869            .with_memory(memory_snapshot)
2870            .with_context(context_snapshot)
2871            .with_state_machine(
2872                state_machine_snapshot.unwrap_or_else(|| StateMachineSnapshot {
2873                    current_state: String::new(),
2874                    previous_state: None,
2875                    turn_count: 0,
2876                    no_transition_count: 0,
2877                    history: vec![],
2878                }),
2879            );
2880
2881        if let Some(ref persona) = self.persona_manager {
2882            snapshot.persona = Some(persona.snapshot_as_value()?);
2883        }
2884
2885        if let Some(ref relationships) = self.relationship_manager {
2886            snapshot.relationships = Some(relationships.snapshot_as_value()?);
2887        }
2888
2889        Ok(snapshot)
2890    }
2891
2892    /// Save state including spawned agents manifest for session persistence.
2893    pub async fn save_state_full(&self) -> Result<AgentSnapshot> {
2894        let mut snapshot = self.save_state().await?;
2895        if let Some(ref registry) = self.spawner_registry {
2896            let entries = registry.list_with_specs();
2897            if !entries.is_empty() {
2898                snapshot = snapshot.with_spawned_agents(entries);
2899            }
2900        }
2901        Ok(snapshot)
2902    }
2903
2904    /// Restores persisted state after invalidating any pending confirmation ownership.
2905    pub async fn restore_state(&self, snapshot: AgentSnapshot) -> Result<()> {
2906        let _admission = self.disambiguation_admission.write().await;
2907        if self.state_transition_reserved.load(Ordering::SeqCst) {
2908            return Err(AgentError::Other(
2909                "Cannot restore state while a state transition is in progress".to_string(),
2910            ));
2911        }
2912        self.invalidate_pending_confirmation("state_restore").await;
2913        *self.pending_skill_id.write() = None;
2914        if let Some(disambiguator) = self.disambiguation_manager.as_ref() {
2915            disambiguator.clear_pending().await;
2916        }
2917        self.memory.restore(snapshot.memory).await?;
2918
2919        if let (Some(sm), Some(sm_snapshot)) = (&self.state_machine, snapshot.state_machine)
2920            && !sm_snapshot.current_state.is_empty()
2921        {
2922            sm.restore(sm_snapshot)?;
2923        }
2924
2925        self.context_manager.restore(snapshot.context);
2926
2927        if let (Some(persona_value), Some(persona_manager)) =
2928            (snapshot.persona, &self.persona_manager)
2929        {
2930            persona_manager.restore_from_value(persona_value)?;
2931        }
2932
2933        if let (Some(relationship_value), Some(relationship_manager)) =
2934            (snapshot.relationships, &self.relationship_manager)
2935        {
2936            relationship_manager.restore_from_value(relationship_value)?;
2937        }
2938
2939        info!(agent_id = %snapshot.agent_id, "State restored");
2940        Ok(())
2941    }
2942
2943    pub async fn save_to(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<()> {
2944        let snapshot = self.save_state().await?;
2945        storage.save(session_id, &snapshot).await
2946    }
2947
2948    async fn load_session_restore(
2949        storage: &dyn AgentStorage,
2950        session_id: &str,
2951    ) -> Result<Option<StoredSessionRestore>> {
2952        let Some(snapshot) = storage.load(session_id).await? else {
2953            return Ok(None);
2954        };
2955        //
2956        // Metadata is restored only when the backend explicitly owns that durable contract.
2957        //
2958        let metadata = if storage.supports(StorageCapability::SessionMetadata) {
2959            storage.load_metadata(session_id).await?
2960        } else {
2961            None
2962        };
2963        Ok(Some(StoredSessionRestore { snapshot, metadata }))
2964    }
2965
2966    async fn capture_session_restore_point(&self) -> Result<RuntimeSessionRestorePoint> {
2967        Ok(RuntimeSessionRestorePoint {
2968            snapshot: self.save_state().await?,
2969            metadata: self.session_metadata(),
2970            actor_id: self.actor_id(),
2971            session_id: self.current_session_id.read().clone(),
2972        })
2973    }
2974
2975    async fn apply_session_restore_unchecked(
2976        &self,
2977        session_id: &str,
2978        stored: StoredSessionRestore,
2979    ) -> Result<()> {
2980        self.restore_state(stored.snapshot).await?;
2981        let metadata = stored.metadata.unwrap_or_default();
2982        if let Some(actor_id) = metadata.actor_id.as_deref() {
2983            self.set_actor_id(actor_id)?;
2984        } else {
2985            self.clear_actor_id();
2986        }
2987        self.set_session_metadata(metadata);
2988        *self.current_session_id.write() = Some(session_id.to_string());
2989        Ok(())
2990    }
2991
2992    async fn restore_session_restore_point(
2993        &self,
2994        restore_point: &RuntimeSessionRestorePoint,
2995    ) -> Result<()> {
2996        self.restore_state(restore_point.snapshot.clone()).await?;
2997        if let Some(actor_id) = restore_point.actor_id.as_deref() {
2998            self.set_actor_id(actor_id)?;
2999        } else {
3000            self.clear_actor_id();
3001        }
3002        self.set_session_metadata(restore_point.metadata.clone());
3003        *self.current_session_id.write() = restore_point.session_id.clone();
3004        Ok(())
3005    }
3006
3007    async fn apply_session_restore(
3008        &self,
3009        session_id: &str,
3010        stored: StoredSessionRestore,
3011    ) -> Result<()> {
3012        let before = self.capture_session_restore_point().await?;
3013        if let Err(error) = self
3014            .apply_session_restore_unchecked(session_id, stored)
3015            .await
3016        {
3017            return match self.restore_session_restore_point(&before).await {
3018                Ok(()) => Err(error),
3019                Err(rollback_error) => Err(AgentError::Other(format!(
3020                    "Session restore failed: {error}; rollback failed: {rollback_error}"
3021                ))),
3022            };
3023        }
3024        Ok(())
3025    }
3026
3027    async fn rollback_session_restore_set(
3028        parent: Option<(&RuntimeAgent, &RuntimeSessionRestorePoint)>,
3029        children: &[(String, Arc<RuntimeAgent>, RuntimeSessionRestorePoint)],
3030    ) -> Vec<String> {
3031        let mut errors = Vec::new();
3032        if let Some((agent, restore_point)) = parent
3033            && let Err(error) = agent.restore_session_restore_point(restore_point).await
3034        {
3035            errors.push(format!("parent: {error}"));
3036        }
3037        for (id, agent, restore_point) in children {
3038            if let Err(error) = agent.restore_session_restore_point(restore_point).await {
3039                errors.push(format!("child '{id}': {error}"));
3040            }
3041        }
3042        errors
3043    }
3044
3045    fn restore_failure(error: impl std::fmt::Display, rollback_errors: Vec<String>) -> AgentError {
3046        if rollback_errors.is_empty() {
3047            AgentError::Other(format!(
3048                "Session restore failed: {error}; runtime state was rolled back"
3049            ))
3050        } else {
3051            AgentError::Other(format!(
3052                "Session restore failed: {error}; rollback also failed for {}",
3053                rollback_errors.join(", ")
3054            ))
3055        }
3056    }
3057
3058    pub async fn load_from(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<bool> {
3059        let Some(stored) = Self::load_session_restore(storage, session_id).await? else {
3060            return Ok(false);
3061        };
3062        self.apply_session_restore(session_id, stored).await?;
3063        Ok(true)
3064    }
3065
3066    pub async fn save_session(&self, session_id: &str) -> Result<()> {
3067        let storage = self.storage.read().clone();
3068        match storage {
3069            Some(s) => {
3070                // Fire on_session_created when this session id is first seen on this runtime.
3071                let is_new = {
3072                    let cur = self.current_session_id.read().clone();
3073                    cur.as_deref() != Some(session_id)
3074                };
3075                if is_new {
3076                    *self.current_session_id.write() = Some(session_id.to_string());
3077                    self.hooks.on_session_created(session_id).await;
3078                }
3079
3080                // Update metadata before persisting.
3081                {
3082                    let now = chrono::Utc::now();
3083                    let msg_count = self
3084                        .memory
3085                        .get_messages(None)
3086                        .await
3087                        .map(|v| v.len())
3088                        .unwrap_or(0);
3089                    let mut meta = self.session_metadata.write();
3090                    meta.last_active = now;
3091                    meta.message_count = msg_count;
3092                    if meta.actor_id.is_none() {
3093                        meta.actor_id = self.actor_id.read().clone();
3094                    }
3095                }
3096
3097                let snapshot = self.save_state().await?;
3098                //
3099                // Metadata-capable backends own atomic snapshot and metadata persistence.
3100                //
3101                if s.supports(StorageCapability::SessionMetadata) {
3102                    let metadata = self.session_metadata.read().clone();
3103                    s.save_snapshot_with_metadata(session_id, &snapshot, &metadata)
3104                        .await
3105                } else {
3106                    s.save(session_id, &snapshot).await
3107                }
3108            }
3109            None => Err(AgentError::Config(
3110                "No storage configured. Use with_storage_config() or with_storage() first".into(),
3111            )),
3112        }
3113    }
3114
3115    pub async fn load_session(&self, session_id: &str) -> Result<bool> {
3116        let storage = self.storage.read().clone();
3117        match storage {
3118            Some(storage) => self.load_from(storage.as_ref(), session_id).await,
3119            None => Err(AgentError::Config(
3120                "No storage configured. Use with_storage_config() or with_storage() first".into(),
3121            )),
3122        }
3123    }
3124
3125    /// Restore this runtime and its complete saved child topology from one named session.
3126    pub async fn restore_session_full(&self, session_id: &str) -> Result<usize> {
3127        self.init_storage().await?;
3128        let storage = self.storage.read().clone().ok_or_else(|| {
3129            AgentError::Config(
3130                "No storage configured. Use with_storage_config() or with_storage() first".into(),
3131            )
3132        })?;
3133        let target_parent = Self::load_session_restore(storage.as_ref(), session_id)
3134            .await?
3135            .ok_or_else(|| AgentError::Persistence(format!("Session not found: {session_id}")))?;
3136        let manifest = target_parent
3137            .snapshot
3138            .spawned_agents
3139            .clone()
3140            .unwrap_or_default();
3141
3142        let registry = self.spawner_registry.as_ref().cloned();
3143        let spawner = if manifest.is_empty() {
3144            self.spawner.as_ref().cloned()
3145        } else {
3146            Some(self.spawner.as_ref().cloned().ok_or_else(|| {
3147                AgentError::Config(
3148                    "Saved session contains child agents but this runtime has no spawner".into(),
3149                )
3150            })?)
3151        };
3152        let registry = if manifest.is_empty() {
3153            registry
3154        } else {
3155            Some(registry.ok_or_else(|| {
3156                AgentError::Config(
3157                    "Saved session contains child agents but this runtime has no registry".into(),
3158                )
3159            })?)
3160        };
3161
3162        let mut target_ids = HashSet::with_capacity(manifest.len());
3163        let mut prepared = Vec::with_capacity(manifest.len());
3164        for entry in manifest {
3165            if !target_ids.insert(entry.id.clone()) {
3166                return Err(AgentError::InvalidSpec(format!(
3167                    "Saved child manifest contains duplicate ID: {}",
3168                    entry.id
3169                )));
3170            }
3171            let spec = crate::spec::AgentSpec::from_yaml_strict(&entry.spec_yaml)?;
3172            spawner
3173                .as_ref()
3174                .expect("non-empty manifests require a spawner")
3175                .validate_explicit_child(&entry.id, &spec)?;
3176            prepared.push((entry.id, spec));
3177        }
3178
3179        let current_ids = registry
3180            .as_ref()
3181            .map(|registry| {
3182                registry
3183                    .list()
3184                    .into_iter()
3185                    .map(|info| info.id)
3186                    .collect::<HashSet<_>>()
3187            })
3188            .unwrap_or_default();
3189        let removal_count = current_ids.difference(&target_ids).count();
3190        let additions = prepared
3191            .iter()
3192            .filter(|(id, _)| !current_ids.contains(id))
3193            .cloned()
3194            .collect::<Vec<_>>();
3195
3196        let mut existing = Vec::new();
3197        if let Some(registry) = registry.as_ref() {
3198            for (id, _) in prepared.iter().filter(|(id, _)| current_ids.contains(id)) {
3199                let agent = registry.get(id).ok_or_else(|| {
3200                    AgentError::Config(format!("Retained child disappeared during restore: {id}"))
3201                })?;
3202                let child_storage = agent.storage().ok_or_else(|| {
3203                    AgentError::Config(format!("Child '{id}' has no storage for session restore"))
3204                })?;
3205                let stored = Self::load_session_restore(child_storage.as_ref(), session_id)
3206                    .await?
3207                    .ok_or_else(|| {
3208                        AgentError::Persistence(format!(
3209                            "Child '{id}' has no saved session '{session_id}'"
3210                        ))
3211                    })?;
3212                existing.push((id.clone(), agent, stored));
3213            }
3214        }
3215
3216        let mut staged = Vec::with_capacity(additions.len());
3217        if !additions.is_empty() {
3218            let spawner = spawner
3219                .as_ref()
3220                .expect("restored additions require a spawner");
3221            let reservations = spawner.reserve_restore_capacity(additions.len(), removal_count)?;
3222            for ((id, spec), reservation) in additions.into_iter().zip(reservations) {
3223                let spawned = spawner
3224                    .spawn_with_reserved_capacity(id.clone(), spec, reservation)
3225                    .await?;
3226                let child_storage = spawned.agent.storage().ok_or_else(|| {
3227                    AgentError::Config(format!("Child '{id}' has no storage for session restore"))
3228                })?;
3229                let stored = Self::load_session_restore(child_storage.as_ref(), session_id)
3230                    .await?
3231                    .ok_or_else(|| {
3232                        AgentError::Persistence(format!(
3233                            "Child '{id}' has no saved session '{session_id}'"
3234                        ))
3235                    })?;
3236                staged.push((spawned, stored));
3237            }
3238        } else if let Some(spawner) = spawner.as_ref() {
3239            spawner.reserve_restore_capacity(0, removal_count)?;
3240        }
3241
3242        let parent_before = self.capture_session_restore_point().await?;
3243        let mut existing_before = Vec::with_capacity(existing.len());
3244        for (id, agent, _) in &existing {
3245            existing_before.push((
3246                id.clone(),
3247                Arc::clone(agent),
3248                agent.capture_session_restore_point().await?,
3249            ));
3250        }
3251
3252        //
3253        // Every record is loaded before runtime state changes. Registry replacement is the final topology commit point.
3254        //
3255        for (_, agent, stored) in &existing {
3256            if let Err(error) = agent
3257                .apply_session_restore_unchecked(session_id, stored.clone())
3258                .await
3259            {
3260                drop(staged);
3261                let rollback_errors =
3262                    Self::rollback_session_restore_set(None, &existing_before).await;
3263                return Err(Self::restore_failure(error, rollback_errors));
3264            }
3265        }
3266        for (spawned, stored) in &staged {
3267            if let Err(error) = spawned
3268                .agent
3269                .apply_session_restore_unchecked(session_id, stored.clone())
3270                .await
3271            {
3272                drop(staged);
3273                let rollback_errors =
3274                    Self::rollback_session_restore_set(None, &existing_before).await;
3275                return Err(Self::restore_failure(error, rollback_errors));
3276            }
3277        }
3278        if let Err(error) = self
3279            .apply_session_restore_unchecked(session_id, target_parent)
3280            .await
3281        {
3282            drop(staged);
3283            let rollback_errors =
3284                Self::rollback_session_restore_set(Some((self, &parent_before)), &existing_before)
3285                    .await;
3286            return Err(Self::restore_failure(error, rollback_errors));
3287        }
3288
3289        if let Some(registry) = registry.as_ref()
3290            && let Err(error) = registry
3291                .reconcile(
3292                    &target_ids,
3293                    staged.into_iter().map(|(spawned, _)| spawned).collect(),
3294                )
3295                .await
3296        {
3297            let rollback_errors =
3298                Self::rollback_session_restore_set(Some((self, &parent_before)), &existing_before)
3299                    .await;
3300            return Err(Self::restore_failure(error, rollback_errors));
3301        }
3302
3303        Ok(target_ids.len())
3304    }
3305
3306    pub async fn delete_session(&self, session_id: &str) -> Result<()> {
3307        let storage = self.storage.read().clone();
3308        match storage {
3309            Some(s) => s.delete(session_id).await,
3310            None => Err(AgentError::Config(
3311                "No storage configured. Use with_storage_config() or with_storage() first".into(),
3312            )),
3313        }
3314    }
3315
3316    pub async fn list_sessions(&self) -> Result<Vec<String>> {
3317        let storage = self.storage.read().clone();
3318        match storage {
3319            Some(s) => s.list_sessions().await,
3320            None => Err(AgentError::Config(
3321                "No storage configured. Use with_storage_config() or with_storage() first".into(),
3322            )),
3323        }
3324    }
3325
3326    fn estimate_tokens(&self, text: &str) -> u32 {
3327        (text.len() as f32 / 4.0).ceil() as u32
3328    }
3329
3330    fn estimate_total_tokens(&self, messages: &[ChatMessage]) -> u32 {
3331        messages
3332            .iter()
3333            .map(|m| self.estimate_tokens(&m.content))
3334            .sum()
3335    }
3336
3337    fn truncate_context(&self, messages: &mut Vec<ChatMessage>, keep_recent: usize) {
3338        if messages.len() <= keep_recent + 1 {
3339            return;
3340        }
3341        let system_msg = messages.remove(0);
3342        let to_remove = messages.len().saturating_sub(keep_recent);
3343        messages.drain(..to_remove);
3344        messages.insert(0, system_msg);
3345    }
3346
3347    fn get_filter(&self, config: &FilterConfig) -> Arc<dyn MessageFilter> {
3348        match config {
3349            FilterConfig::KeepRecent(n) => Arc::new(KeepRecentFilter::new(*n)),
3350            FilterConfig::ByRole { keep_roles } => Arc::new(ByRoleFilter::new(keep_roles.clone())),
3351            FilterConfig::SkipPattern { skip_if_contains } => {
3352                Arc::new(SkipPatternFilter::new(skip_if_contains.clone()))
3353            }
3354            FilterConfig::Custom { name } => {
3355                let filters = self.message_filters.read();
3356                filters
3357                    .get(name)
3358                    .cloned()
3359                    .unwrap_or_else(|| Arc::new(KeepRecentFilter::new(10)))
3360            }
3361        }
3362    }
3363
3364    async fn summarize_context(
3365        &self,
3366        messages: &mut Vec<ChatMessage>,
3367        summarizer_llm: Option<&str>,
3368        max_summary_tokens: u32,
3369        custom_prompt: Option<&str>,
3370        keep_recent: usize,
3371        filter: Option<&FilterConfig>,
3372    ) -> Result<()> {
3373        let system_msg = messages.remove(0);
3374
3375        let to_summarize_count = messages.len().saturating_sub(keep_recent);
3376        if to_summarize_count == 0 {
3377            messages.insert(0, system_msg);
3378            return Ok(());
3379        }
3380
3381        let recent_msgs: Vec<ChatMessage> = messages.drain(to_summarize_count..).collect();
3382        let mut to_summarize = std::mem::take(messages);
3383
3384        if let Some(filter_config) = filter {
3385            let filter = self.get_filter(filter_config);
3386            to_summarize = filter.filter(to_summarize);
3387        }
3388
3389        if to_summarize.is_empty() {
3390            *messages = recent_msgs;
3391            messages.insert(0, system_msg);
3392            return Ok(());
3393        }
3394
3395        let conversation_text = to_summarize
3396            .iter()
3397            .map(|m| format!("{:?}: {}", m.role, m.content))
3398            .collect::<Vec<_>>()
3399            .join("\n");
3400
3401        let default_prompt = format!(
3402            "Summarize the following conversation in under {} tokens, preserving key information:\n\n{}",
3403            max_summary_tokens, conversation_text
3404        );
3405
3406        let summary_prompt = custom_prompt
3407            .map(|p| format!("{}\n\n{}", p, conversation_text))
3408            .unwrap_or(default_prompt);
3409
3410        let summarizer = if let Some(alias) = summarizer_llm {
3411            self.llm_registry
3412                .get(alias)
3413                .map_err(|e| AgentError::Config(e.to_string()))?
3414        } else {
3415            self.llm_registry
3416                .router()
3417                .or_else(|_| self.llm_registry.default())
3418                .map_err(|e| AgentError::Config(e.to_string()))?
3419        };
3420
3421        let summary_msgs = vec![ChatMessage::user(&summary_prompt)];
3422        let response = self
3423            .observe_purpose(
3424                ObservationPurpose::Summarization,
3425                summarizer.complete(&summary_msgs, None),
3426            )
3427            .await?;
3428
3429        let summary_message = ChatMessage::system(format!(
3430            "[Previous conversation summary]\n{}",
3431            response.content
3432        ));
3433
3434        *messages = vec![system_msg, summary_message];
3435        messages.extend(recent_msgs);
3436
3437        debug!(
3438            summarized_count = to_summarize_count,
3439            kept_recent = keep_recent,
3440            "Context summarized"
3441        );
3442
3443        Ok(())
3444    }
3445
3446    fn render_system_prompt(&self) -> Result<String> {
3447        let mut context = self.build_context_with_overlays();
3448
3449        // Inject actor_facts for {{ actor_facts }} template variable.
3450        let facts_text = self.format_actor_facts_for_context();
3451        if !facts_text.is_empty() {
3452            context.insert(
3453                "actor_facts".to_string(),
3454                serde_json::Value::String(facts_text),
3455            );
3456        }
3457
3458        if let Some((key, text)) = self.format_relationship_for_context() {
3459            context.insert(key, serde_json::Value::String(text));
3460        }
3461
3462        self.template_renderer
3463            .render(&self.base_system_prompt, &context)
3464    }
3465
3466    /// Canonicalizes IDs once while preserving their first declaration order and ignoring unknown entries.
3467    fn canonical_unique_tool_ids(&self, ids: &[String]) -> Vec<String> {
3468        let mut seen = HashSet::new();
3469        ids.iter()
3470            .filter_map(|id| self.tools.canonical_id(id))
3471            .filter(|canonical_id| seen.insert(canonical_id.clone()))
3472            .collect()
3473    }
3474
3475    /// Applies runtime narrowing to the canonical declared grant without permitting scope-based expansion.
3476    fn get_top_level_tool_ids_for_scope(&self, scope_override: Option<&[String]>) -> Vec<String> {
3477        let Some(declared) = self.declared_tool_ids.as_deref() else {
3478            return Vec::new();
3479        };
3480        let mut effective = self.canonical_unique_tool_ids(declared);
3481        if let Some(scope) = scope_override {
3482            let scope: HashSet<String> =
3483                self.canonical_unique_tool_ids(scope).into_iter().collect();
3484            effective.retain(|canonical_id| scope.contains(canonical_id));
3485        }
3486        effective
3487    }
3488
3489    /// Reads the live runtime narrowing and returns the current deterministic effective tool IDs.
3490    async fn get_available_tool_ids(&self) -> Result<Vec<String>> {
3491        Ok(self.get_available_tool_ids_snapshot().await?.tool_ids)
3492    }
3493
3494    /// Captures the runtime scope used for ordinary availability checks before applying state narrowing.
3495    async fn get_available_tool_ids_snapshot(&self) -> Result<AvailableToolIdsSnapshot> {
3496        let scope_override = self.runtime_control.tool_scope_override.read().clone();
3497        self.get_available_tool_ids_snapshot_for_scope(scope_override.as_deref())
3498            .await
3499    }
3500
3501    /// Applies every explicit ancestor and current-state scope to one runtime snapshot and retains its state generation.
3502    async fn get_available_tool_ids_snapshot_for_scope(
3503        &self,
3504        scope_override: Option<&[String]>,
3505    ) -> Result<AvailableToolIdsSnapshot> {
3506        let mut available = self.get_top_level_tool_ids_for_scope(scope_override);
3507        let (state_generation, state_scopes) = self
3508            .state_machine
3509            .as_ref()
3510            .map(|state_machine| {
3511                let (generation, scopes) = state_machine.current_tool_scope_snapshot();
3512                (Some(generation), scopes)
3513            })
3514            .unwrap_or((None, Vec::new()));
3515
3516        if available.is_empty() || state_scopes.is_empty() {
3517            return Ok(AvailableToolIdsSnapshot {
3518                tool_ids: available,
3519                state_generation,
3520            });
3521        }
3522
3523        let eval_ctx = self.build_evaluation_context().await?;
3524        let llm_getter = RegistryLLMGetter {
3525            registry: self.llm_registry.clone(),
3526        };
3527        let evaluator = ConditionEvaluator::new(llm_getter);
3528
3529        for state_scope in state_scopes {
3530            if state_scope.is_empty() {
3531                available.clear();
3532                break;
3533            }
3534
3535            let mut allowed = HashSet::new();
3536            for tool_ref in &state_scope {
3537                let tool_id = tool_ref.id();
3538                let Some(canonical_id) = self.tools.canonical_id(tool_id) else {
3539                    continue;
3540                };
3541                let condition_matches = if let Some(condition) = tool_ref.condition() {
3542                    match evaluator.evaluate(condition, &eval_ctx).await {
3543                        Ok(matches) => matches,
3544                        Err(error) => {
3545                            warn!(tool = tool_id, error = %error, "Error evaluating tool condition");
3546                            false
3547                        }
3548                    }
3549                } else {
3550                    true
3551                };
3552                if condition_matches {
3553                    allowed.insert(canonical_id);
3554                } else {
3555                    debug!(tool = tool_id, "Tool condition not met, skipping");
3556                }
3557            }
3558            available.retain(|canonical_id| allowed.contains(canonical_id));
3559            if available.is_empty() {
3560                break;
3561            }
3562        }
3563
3564        Ok(AvailableToolIdsSnapshot {
3565            tool_ids: available,
3566            state_generation,
3567        })
3568    }
3569
3570    async fn build_evaluation_context(&self) -> Result<EvaluationContext> {
3571        let context = self.build_context_with_overlays();
3572        let messages = self.memory.get_messages(Some(10)).await?;
3573        let tool_history = self.tool_call_history.read().clone();
3574
3575        let (state_name, turn_count, previous_state) = if let Some(ref sm) = self.state_machine {
3576            (Some(sm.current()), sm.turn_count(), sm.previous())
3577        } else {
3578            (None, 0, None)
3579        };
3580
3581        Ok(EvaluationContext::default()
3582            .with_context(context)
3583            .with_state(state_name, turn_count, previous_state)
3584            .with_called_tools(tool_history)
3585            .with_messages(messages))
3586    }
3587
3588    fn record_tool_call(&self, tool_id: &str, result: Value) {
3589        self.tool_call_history.write().push(ToolCallRecord {
3590            tool_id: tool_id.to_string(),
3591            result,
3592            timestamp: chrono::Utc::now(),
3593        });
3594    }
3595
3596    async fn get_effective_system_prompt_with_persona_hooks(
3597        &self,
3598        fire_persona_hooks: bool,
3599        include_tool_prompt: bool,
3600    ) -> Result<String> {
3601        let rendered_base = self.render_system_prompt()?;
3602
3603        let persona_prefix = if let Some(ref persona) = self.persona_manager {
3604            let context = self.build_context_with_overlays();
3605            if fire_persona_hooks {
3606                let render_result = persona.render_prompt(&context)?;
3607                for content in &render_result.newly_revealed {
3608                    self.hooks.on_secret_revealed(content).await;
3609                }
3610                render_result.prompt
3611            } else {
3612                persona.render_prompt_preview(&context)?
3613            }
3614        } else {
3615            String::new()
3616        };
3617
3618        if let Some(ref sm) = self.state_machine
3619            && let Some(state_def) = sm.current_definition()
3620        {
3621            let state_prompt = if let Some(ref prompt) = state_def.prompt {
3622                let context = self.build_context_with_overlays();
3623                self.template_renderer.render_with_state(
3624                    prompt,
3625                    &context,
3626                    &sm.current(),
3627                    sm.previous().as_deref(),
3628                    sm.turn_count(),
3629                    state_def.max_turns,
3630                )?
3631            } else {
3632                String::new()
3633            };
3634
3635            let combined = match state_def.prompt_mode {
3636                PromptMode::Append => {
3637                    if state_prompt.is_empty() {
3638                        rendered_base
3639                    } else {
3640                        format!(
3641                            "{}\n\n[Current State: {}]\n{}",
3642                            rendered_base,
3643                            sm.current(),
3644                            state_prompt
3645                        )
3646                    }
3647                }
3648                PromptMode::Replace => {
3649                    if state_prompt.is_empty() {
3650                        rendered_base
3651                    } else {
3652                        state_prompt
3653                    }
3654                }
3655                PromptMode::Prepend => {
3656                    if state_prompt.is_empty() {
3657                        rendered_base
3658                    } else {
3659                        format!("{}\n\n{}", state_prompt, rendered_base)
3660                    }
3661                }
3662            };
3663
3664            // Persona always prepended regardless of prompt_mode.
3665            let with_persona = if persona_prefix.is_empty() {
3666                combined
3667            } else {
3668                format!("{}\n\n{}", persona_prefix, combined)
3669            };
3670
3671            if include_tool_prompt {
3672                let available_tool_ids = self.get_available_tool_ids().await?;
3673                if !available_tool_ids.is_empty() {
3674                    let tools_prompt = self.tools.generate_scoped_prompt_with_mode(
3675                        &available_tool_ids,
3676                        None,
3677                        self.parallel_tools.enabled,
3678                        self.runtime_config.tool_schema_prompt_mode,
3679                    );
3680                    if !tools_prompt.is_empty() {
3681                        return Ok(format!("{}\n\n{}", with_persona, tools_prompt));
3682                    }
3683                }
3684            }
3685            return Ok(with_persona);
3686        }
3687
3688        // No state machine - prepend persona to base.
3689        let with_persona = if persona_prefix.is_empty() {
3690            rendered_base
3691        } else {
3692            format!("{}\n\n{}", persona_prefix, rendered_base)
3693        };
3694
3695        if include_tool_prompt {
3696            let available_tool_ids = self.get_available_tool_ids().await?;
3697            let tools_prompt = self.tools.generate_scoped_prompt_with_mode(
3698                &available_tool_ids,
3699                None,
3700                self.parallel_tools.enabled,
3701                self.runtime_config.tool_schema_prompt_mode,
3702            );
3703            if !tools_prompt.is_empty() {
3704                return Ok(format!("{}\n\n{}", with_persona, tools_prompt));
3705            }
3706        }
3707        Ok(with_persona)
3708    }
3709
3710    fn get_state_llm(&self) -> Result<Arc<dyn LLMProvider>> {
3711        if let Some(ref sm) = self.state_machine
3712            && let Some(state_def) = sm.current_definition()
3713            && let Some(ref llm_alias) = state_def.llm
3714        {
3715            return self
3716                .llm_registry
3717                .get(llm_alias)
3718                .map_err(|e| AgentError::Config(e.to_string()));
3719        }
3720        self.llm_registry
3721            .default()
3722            .map_err(|e| AgentError::Config(e.to_string()))
3723    }
3724
3725    fn get_effective_reasoning_config(&self) -> ReasoningConfig {
3726        if let Some(ref sm) = self.state_machine
3727            && let Some(state_def) = sm.current_definition()
3728            && let Some(ref state_reasoning) = state_def.reasoning
3729        {
3730            return state_reasoning.clone();
3731        }
3732        self.reasoning_config.clone()
3733    }
3734
3735    fn get_effective_reflection_config(&self) -> ReflectionConfig {
3736        if let Some(ref sm) = self.state_machine
3737            && let Some(state_def) = sm.current_definition()
3738            && let Some(ref state_reflection) = state_def.reflection
3739        {
3740            return state_reflection.clone();
3741        }
3742        self.reflection_config.clone()
3743    }
3744
3745    fn get_skill_reasoning_config(&self, skill: &SkillDefinition) -> ReasoningConfig {
3746        skill
3747            .reasoning
3748            .clone()
3749            .unwrap_or_else(|| self.get_effective_reasoning_config())
3750    }
3751
3752    fn get_skill_reflection_config(&self, skill: &SkillDefinition) -> ReflectionConfig {
3753        skill
3754            .reflection
3755            .clone()
3756            .unwrap_or_else(|| self.get_effective_reflection_config())
3757    }
3758
3759    async fn build_disambiguation_context(&self) -> Result<DisambiguationContext> {
3760        let recent_messages: Vec<String> = self
3761            .memory
3762            .get_messages(Some(5))
3763            .await?
3764            .iter()
3765            .rev()
3766            .map(|m| format!("{:?}: {}", m.role, m.content))
3767            .collect();
3768
3769        let current_state = self.current_state().map(|s| s.to_string());
3770
3771        // Include the current state's prompt text so the detector understands
3772        // what kind of input is expected (e.g., "Ask for the order number").
3773        let state_prompt: Option<String> = self
3774            .state_machine
3775            .as_ref()
3776            .and_then(|sm| sm.current_definition())
3777            .and_then(|def| def.prompt.clone());
3778
3779        let available_tools: Vec<String> = self
3780            .get_available_tool_ids()
3781            .await
3782            .unwrap_or_else(|_| self.tools.list_ids());
3783
3784        let available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
3785
3786        let mut user_context = self.build_context_with_overlays();
3787        user_context.remove(DISAMBIGUATION_STATE_GENERATION_KEY);
3788        if let Some(state_generation) = self
3789            .state_machine
3790            .as_ref()
3791            .map(|state_machine| state_machine.generation())
3792        {
3793            user_context.insert(
3794                DISAMBIGUATION_STATE_GENERATION_KEY.to_string(),
3795                serde_json::json!(state_generation),
3796            );
3797        }
3798
3799        // Extract canonical intent labels from current state's transitions
3800        let available_intents: Vec<String> = if let Some(ref sm) = self.state_machine {
3801            sm.current_definition()
3802                .map(|def| {
3803                    def.transitions
3804                        .iter()
3805                        .filter_map(|t| t.intent.clone())
3806                        .collect()
3807                })
3808                .unwrap_or_default()
3809        } else {
3810            Vec::new()
3811        };
3812
3813        Ok(DisambiguationContext::from_agent_state(
3814            recent_messages,
3815            current_state,
3816            state_prompt,
3817            available_tools,
3818            available_skills,
3819            available_intents,
3820            user_context,
3821        ))
3822    }
3823
3824    fn get_available_skills(&self) -> Vec<&SkillDefinition> {
3825        if let Some(ref sm) = self.state_machine
3826            && let Some(state_def) = sm.current_definition()
3827        {
3828            let parent_def = sm.get_parent_definition();
3829            let effective_skills = state_def.get_effective_skills(parent_def.as_ref());
3830            if !effective_skills.is_empty() {
3831                return self
3832                    .skills
3833                    .iter()
3834                    .filter(|s| effective_skills.contains(&&s.id))
3835                    .collect();
3836            }
3837        }
3838        self.skills.iter().collect()
3839    }
3840
3841    async fn build_messages(&self) -> Result<Vec<ChatMessage>> {
3842        self.build_messages_internal(true, None, true).await
3843    }
3844
3845    async fn build_messages_for_draft(&self, user_message: &str) -> Result<Vec<ChatMessage>> {
3846        self.build_messages_internal(false, Some(user_message), true)
3847            .await
3848    }
3849
3850    async fn build_messages_internal(
3851        &self,
3852        fire_persona_hooks: bool,
3853        ephemeral_user_message: Option<&str>,
3854        include_tool_prompt: bool,
3855    ) -> Result<Vec<ChatMessage>> {
3856        let system_prompt = self
3857            .get_effective_system_prompt_with_persona_hooks(fire_persona_hooks, include_tool_prompt)
3858            .await?;
3859        let mut messages = vec![ChatMessage::system(&system_prompt)];
3860
3861        let context = self.memory.get_context().await?;
3862        let history = if let Some(ref budget) = self.memory_token_budget {
3863            context.to_llm_messages_with_allocation(&budget.allocation)
3864        } else {
3865            context.to_llm_messages()
3866        };
3867        messages.extend(history);
3868        if let Some(user_message) = ephemeral_user_message {
3869            messages.push(ChatMessage::user(user_message));
3870        }
3871
3872        let total_tokens = self.estimate_total_tokens(&messages);
3873
3874        if total_tokens > self.max_context_tokens {
3875            debug!(
3876                total = total_tokens,
3877                limit = self.max_context_tokens,
3878                "Context overflow"
3879            );
3880
3881            match &self.recovery_manager.config().llm.on_context_overflow {
3882                ContextOverflowAction::Error => {
3883                    return Err(AgentError::LLM(format!(
3884                        "Context overflow: {} tokens > {} limit",
3885                        total_tokens, self.max_context_tokens
3886                    )));
3887                }
3888                ContextOverflowAction::Truncate { keep_recent } => {
3889                    self.truncate_context(&mut messages, *keep_recent);
3890                }
3891                ContextOverflowAction::Summarize {
3892                    summarizer_llm,
3893                    max_summary_tokens,
3894                    custom_prompt,
3895                    keep_recent,
3896                    filter,
3897                } => {
3898                    self.summarize_context(
3899                        &mut messages,
3900                        summarizer_llm.as_deref(),
3901                        *max_summary_tokens,
3902                        custom_prompt.as_deref(),
3903                        *keep_recent,
3904                        filter.as_ref(),
3905                    )
3906                    .await?;
3907                }
3908            }
3909        }
3910
3911        Ok(messages)
3912    }
3913
3914    async fn main_tool_protocol(
3915        &self,
3916        llm: &dyn LLMProvider,
3917        ephemeral_new_turn: bool,
3918    ) -> Result<MainToolProtocol> {
3919        let mut choice = llm.configured_tool_choice();
3920        if matches!(choice.as_ref(), Some(ToolChoice::None)) {
3921            return Ok(MainToolProtocol {
3922                choice,
3923                tool_ids: Vec::new(),
3924                definitions: Vec::new(),
3925            });
3926        }
3927
3928        let mut tool_ids = self.get_available_tool_ids().await?;
3929        tool_ids.sort();
3930        tool_ids.dedup();
3931        if let Some(ToolChoice::Specific(expected)) = choice.as_ref() {
3932            let canonical = self.tools.canonical_id(expected).ok_or_else(|| {
3933                AgentError::Config(format!(
3934                    "specific tool choice '{expected}' is not registered"
3935                ))
3936            })?;
3937            if canonical != *expected {
3938                return Err(AgentError::Config(format!(
3939                    "specific tool choice must use canonical ID '{canonical}', not '{expected}'"
3940                )));
3941            }
3942            if !tool_ids.iter().any(|tool_id| tool_id == expected) {
3943                return Err(AgentError::Config(format!(
3944                    "specific tool choice '{expected}' is outside the effective tool grant"
3945                )));
3946            }
3947        }
3948        if matches!(
3949            choice.as_ref(),
3950            Some(ToolChoice::Required | ToolChoice::Specific(_))
3951        ) && tool_ids.is_empty()
3952        {
3953            return Err(AgentError::Config(
3954                "required tool choice has no tool inside the effective grant".to_string(),
3955            ));
3956        }
3957        if !ephemeral_new_turn
3958            && let Some(configured_choice) = choice.as_ref()
3959            && matches!(
3960                configured_choice,
3961                ToolChoice::Required | ToolChoice::Specific(_)
3962            )
3963            && self
3964                .tool_choice_satisfied_in_current_turn(configured_choice, &tool_ids)
3965                .await?
3966        {
3967            choice = Some(ToolChoice::Auto);
3968        }
3969        if let Some(ToolChoice::Specific(expected)) = choice.as_ref() {
3970            tool_ids.retain(|tool_id| tool_id == expected);
3971        }
3972
3973        let definitions = tool_ids
3974            .iter()
3975            .map(|tool_id| {
3976                let tool = self.tools.get(tool_id).ok_or_else(|| {
3977                    AgentError::Config(format!(
3978                        "effective tool '{tool_id}' disappeared before provider exposure"
3979                    ))
3980                })?;
3981                Ok(LLMToolDefinition {
3982                    name: tool_id.clone(),
3983                    description: tool.description().to_string(),
3984                    input_schema: tool.input_schema(),
3985                })
3986            })
3987            .collect::<Result<Vec<_>>>()?;
3988
3989        //
3990        // Provider-visible definitions are derived only from the current effective grant. Tool choice never expands registration, scope, or authorization.
3991        //
3992        Ok(MainToolProtocol {
3993            choice,
3994            tool_ids,
3995            definitions,
3996        })
3997    }
3998
3999    async fn tool_choice_satisfied_in_current_turn(
4000        &self,
4001        choice: &ToolChoice,
4002        effective_tool_ids: &[String],
4003    ) -> Result<bool> {
4004        let messages = self.memory.get_messages(None).await?;
4005        let mut saw_tool_result = false;
4006        for message in messages.iter().rev() {
4007            match message.role {
4008                ai_agents_core::Role::Tool | ai_agents_core::Role::Function => {
4009                    saw_tool_result = true;
4010                }
4011                ai_agents_core::Role::Assistant if saw_tool_result => {
4012                    let Some(calls) = self.parse_tool_calls(&message.content) else {
4013                        continue;
4014                    };
4015                    let calls_are_effective = !calls.is_empty()
4016                        && calls.iter().all(|call| {
4017                            self.tools
4018                                .canonical_id(&call.name)
4019                                .is_some_and(|canonical| effective_tool_ids.contains(&canonical))
4020                        });
4021                    return Ok(calls_are_effective
4022                        && match choice {
4023                            ToolChoice::Required => true,
4024                            ToolChoice::Specific(expected) => calls.iter().all(|call| {
4025                                self.tools.canonical_id(&call.name).as_deref()
4026                                    == Some(expected.as_str())
4027                            }),
4028                            _ => false,
4029                        });
4030                }
4031                ai_agents_core::Role::User => return Ok(false),
4032                _ => {}
4033            }
4034        }
4035        Ok(false)
4036    }
4037
4038    fn provider_can_use_native_tools(
4039        &self,
4040        llm: &dyn LLMProvider,
4041        protocol: &MainToolProtocol,
4042    ) -> bool {
4043        let Some(choice) = protocol.choice.as_ref() else {
4044            return false;
4045        };
4046        if matches!(choice, ToolChoice::None) || protocol.definitions.is_empty() {
4047            return false;
4048        }
4049        llm.supports_tool_choice(choice)
4050            && protocol.definitions.iter().all(|definition| {
4051                !definition.name.is_empty()
4052                    && definition.name.len() <= 64
4053                    && definition
4054                        .name
4055                        .bytes()
4056                        .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
4057            })
4058    }
4059
4060    fn prompt_messages_for_tool_protocol(
4061        &self,
4062        messages: &[ChatMessage],
4063        protocol: &MainToolProtocol,
4064        corrective: bool,
4065    ) -> Vec<ChatMessage> {
4066        let mut messages = messages.to_vec();
4067        let Some(choice) = protocol.choice.as_ref() else {
4068            return messages;
4069        };
4070        if matches!(choice, ToolChoice::None) || protocol.tool_ids.is_empty() {
4071            return messages;
4072        }
4073
4074        let mut tool_prompt = self.tools.generate_scoped_prompt_with_mode(
4075            &protocol.tool_ids,
4076            None,
4077            self.parallel_tools.enabled,
4078            self.runtime_config.tool_schema_prompt_mode,
4079        );
4080        match choice {
4081            ToolChoice::Required => tool_prompt.push_str(
4082                "\n\nYou must call at least one listed tool before giving a final answer.",
4083            ),
4084            ToolChoice::Specific(tool_id) => tool_prompt.push_str(&format!(
4085                "\n\nYou must call the '{tool_id}' tool before giving a final answer."
4086            )),
4087            ToolChoice::Auto => {}
4088            ToolChoice::None => return messages,
4089            _ => return messages,
4090        }
4091        if let Some(system) = messages
4092            .iter_mut()
4093            .find(|message| message.role == ai_agents_core::Role::System)
4094        {
4095            system.content.push_str("\n\n");
4096            system.content.push_str(&tool_prompt);
4097        } else {
4098            messages.insert(0, ChatMessage::system(tool_prompt));
4099        }
4100        if corrective {
4101            let instruction = match choice {
4102                ToolChoice::Required => {
4103                    "Your previous response did not call a required tool. Call at least one listed tool now and return only the JSON tool call."
4104                }
4105                ToolChoice::Specific(tool_id) => {
4106                    messages.push(ChatMessage::user(format!(
4107                        "Your previous response did not call the required '{tool_id}' tool. Call it now and return only the JSON tool call."
4108                    )));
4109                    return messages;
4110                }
4111                _ => return messages,
4112            };
4113            messages.push(ChatMessage::user(instruction));
4114        }
4115        messages
4116    }
4117
4118    async fn invoke_main_provider(
4119        &self,
4120        llm: Arc<dyn LLMProvider>,
4121        messages: &[ChatMessage],
4122        protocol: &MainToolProtocol,
4123        corrective: bool,
4124    ) -> std::result::Result<MainProviderResponse, LLMError> {
4125        let use_native = self.provider_can_use_native_tools(llm.as_ref(), protocol);
4126        let response = if use_native {
4127            let request = LLMToolRequest {
4128                tools: protocol.definitions.clone(),
4129                choice: protocol
4130                    .choice
4131                    .clone()
4132                    .expect("native tool requests require an explicit choice"),
4133            };
4134            self.observe_purpose(
4135                ObservationPurpose::MainResponse,
4136                llm.complete_with_tools(messages, None, &request),
4137            )
4138            .await?
4139        } else {
4140            let prompt_messages =
4141                self.prompt_messages_for_tool_protocol(messages, protocol, corrective);
4142            self.observe_purpose(
4143                ObservationPurpose::MainResponse,
4144                llm.complete(&prompt_messages, None),
4145            )
4146            .await?
4147        };
4148        Ok(MainProviderResponse {
4149            response,
4150            used_native_tools: use_native,
4151        })
4152    }
4153
4154    async fn complete_main_attempt_with_recovery(
4155        &self,
4156        llm: Arc<dyn LLMProvider>,
4157        messages: &[ChatMessage],
4158        protocol: &MainToolProtocol,
4159        corrective: bool,
4160    ) -> Result<MainProviderResponse> {
4161        let primary_result = if self.recovery_manager.config().default.max_retries > 0 {
4162            self.recovery_manager
4163                .with_retry("llm_call", None, || {
4164                    let llm = Arc::clone(&llm);
4165                    async move {
4166                        self.invoke_main_provider(llm, messages, protocol, corrective)
4167                            .await
4168                            .map_err(|error| error.classify())
4169                    }
4170                })
4171                .await
4172                .map_err(|error| AgentError::LLM(error.to_string()))
4173        } else {
4174            self.invoke_main_provider(Arc::clone(&llm), messages, protocol, corrective)
4175                .await
4176                .map_err(|error| AgentError::LLM(error.to_string()))
4177        };
4178
4179        match primary_result {
4180            Ok(response) => Ok(response),
4181            Err(primary_error) => match &self.recovery_manager.config().llm.on_failure {
4182                LLMFailureAction::FallbackLlm { fallback_llm } => {
4183                    let fallback = self.llm_registry.get(fallback_llm).map_err(|error| {
4184                        AgentError::Config(format!(
4185                            "Fallback LLM '{fallback_llm}' not found: {error}"
4186                        ))
4187                    })?;
4188                    self.invoke_main_provider(fallback, messages, protocol, corrective)
4189                        .await
4190                        .map_err(|error| AgentError::LLM(error.to_string()))
4191                }
4192                LLMFailureAction::FallbackResponse { message } => {
4193                    if matches!(
4194                        protocol.choice.as_ref(),
4195                        Some(ToolChoice::Required | ToolChoice::Specific(_))
4196                    ) {
4197                        Err(AgentError::LLM(format!(
4198                            "Required tool selection failed and cannot be satisfied by a static fallback response: {primary_error}"
4199                        )))
4200                    } else {
4201                        Ok(MainProviderResponse {
4202                            response: LLMResponse::new(message.clone(), FinishReason::Stop),
4203                            used_native_tools: false,
4204                        })
4205                    }
4206                }
4207                LLMFailureAction::Error => Err(primary_error),
4208            },
4209        }
4210    }
4211
4212    fn normalize_main_provider_response(
4213        &self,
4214        mut response: LLMResponse,
4215        protocol: &MainToolProtocol,
4216    ) -> Result<(LLMResponse, bool)> {
4217        let native_calls = response
4218            .tool_calls()
4219            .map_err(|error| AgentError::LLM(error.to_string()))?;
4220        let calls = match native_calls {
4221            Some(calls) => {
4222                let markers = calls
4223                    .iter()
4224                    .map(|call| {
4225                        serde_json::json!({
4226                            "_ai_agents_native_tool_call": true,
4227                            "id": call.id,
4228                            "tool": call.name,
4229                            "arguments": call.arguments,
4230                        })
4231                    })
4232                    .collect::<Vec<_>>();
4233                response.content = if markers.len() == 1 {
4234                    markers[0].to_string()
4235                } else {
4236                    serde_json::Value::Array(markers).to_string()
4237                };
4238                Some(calls)
4239            }
4240            None if !matches!(protocol.choice.as_ref(), Some(ToolChoice::None)) => {
4241                self.parse_tool_calls(response.content.trim())
4242            }
4243            None => None,
4244        };
4245
4246        if protocol.choice.is_some()
4247            && let Some(calls) = calls.as_ref()
4248            && calls.iter().any(|call| {
4249                self.tools
4250                    .canonical_id(&call.name)
4251                    .is_none_or(|canonical| !protocol.tool_ids.contains(&canonical))
4252            })
4253        {
4254            return Err(AgentError::LLM(
4255                "Provider returned a tool call outside the effective grant".to_string(),
4256            ));
4257        }
4258
4259        let compliant = match protocol.choice.as_ref() {
4260            Some(ToolChoice::Required) => calls.as_ref().is_some_and(|calls| !calls.is_empty()),
4261            Some(ToolChoice::Specific(expected)) => calls.as_ref().is_some_and(|calls| {
4262                !calls.is_empty()
4263                    && calls.iter().all(|call| {
4264                        self.tools.canonical_id(&call.name).as_deref() == Some(expected.as_str())
4265                    })
4266            }),
4267            _ => true,
4268        };
4269        Ok((response, compliant))
4270    }
4271
4272    async fn complete_main_llm_with_recovery(
4273        &self,
4274        llm: Arc<dyn LLMProvider>,
4275        messages: &[ChatMessage],
4276        protocol: &MainToolProtocol,
4277    ) -> Result<LLMResponse> {
4278        let first = self
4279            .complete_main_attempt_with_recovery(Arc::clone(&llm), messages, protocol, false)
4280            .await?;
4281        let (response, compliant) =
4282            self.normalize_main_provider_response(first.response, protocol)?;
4283        if compliant {
4284            return Ok(response);
4285        }
4286        if first.used_native_tools {
4287            return Err(AgentError::LLM(
4288                "Provider returned no compliant native call for required tool choice".to_string(),
4289            ));
4290        }
4291
4292        let corrected = self
4293            .complete_main_attempt_with_recovery(llm, messages, protocol, true)
4294            .await?;
4295        let (response, compliant) =
4296            self.normalize_main_provider_response(corrected.response, protocol)?;
4297        if compliant {
4298            return Ok(response);
4299        }
4300        Err(AgentError::LLM(
4301            "Provider returned no compliant tool call after one corrective retry".to_string(),
4302        ))
4303    }
4304
4305    fn is_native_tool_call_content(content: &str) -> bool {
4306        let Ok(value) = serde_json::from_str::<serde_json::Value>(content) else {
4307            return false;
4308        };
4309        match value {
4310            serde_json::Value::Array(values) => {
4311                !values.is_empty()
4312                    && values.iter().all(|value| {
4313                        value
4314                            .get("_ai_agents_native_tool_call")
4315                            .and_then(|marker| marker.as_bool())
4316                            == Some(true)
4317                    })
4318            }
4319            serde_json::Value::Object(map) => {
4320                map.get("_ai_agents_native_tool_call")
4321                    .and_then(|marker| marker.as_bool())
4322                    == Some(true)
4323            }
4324            _ => false,
4325        }
4326    }
4327
4328    fn tool_result_message(
4329        tool_call: &ToolCall,
4330        output: &str,
4331        native_tool_call: bool,
4332    ) -> ChatMessage {
4333        if !native_tool_call {
4334            return ChatMessage::function(&tool_call.name, output);
4335        }
4336        let output = serde_json::from_str::<serde_json::Value>(output)
4337            .unwrap_or_else(|_| serde_json::Value::String(output.to_string()));
4338        ChatMessage::function(
4339            &tool_call.name,
4340            serde_json::json!({
4341                "_ai_agents_native_tool_result": true,
4342                "id": tool_call.id,
4343                "tool": tool_call.name,
4344                "output": output,
4345            })
4346            .to_string(),
4347        )
4348    }
4349
4350    fn parse_main_tool_calls(
4351        &self,
4352        content: &str,
4353        protocol: &MainToolProtocol,
4354    ) -> Option<Vec<ToolCall>> {
4355        if matches!(protocol.choice.as_ref(), Some(ToolChoice::None)) {
4356            None
4357        } else {
4358            self.parse_tool_calls(content)
4359        }
4360    }
4361
4362    fn parse_tool_calls(&self, content: &str) -> Option<Vec<ToolCall>> {
4363        // Try direct JSON parse first
4364        if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(content) {
4365            // Handle JSON array of tool calls (parallel tool calling)
4366            if let Some(arr) = parsed.as_array() {
4367                let calls: Vec<ToolCall> = arr
4368                    .iter()
4369                    .filter_map(|v| self.extract_tool_call_from_value(v))
4370                    .collect();
4371                if !calls.is_empty() {
4372                    return Some(calls);
4373                }
4374            }
4375            // Handle single JSON object
4376            if let Some(tool_call) = self.extract_tool_call_from_value(&parsed) {
4377                return Some(vec![tool_call]);
4378            }
4379        }
4380
4381        // Try to extract JSON from content (handles extra text/braces from LLM)
4382        if let Some(json_str) = self.extract_json_from_content(content)
4383            && let Ok(parsed) = serde_json::from_str::<serde_json::Value>(&json_str)
4384        {
4385            // Handle JSON array of tool calls (parallel tool calling)
4386            if let Some(arr) = parsed.as_array() {
4387                let calls: Vec<ToolCall> = arr
4388                    .iter()
4389                    .filter_map(|v| self.extract_tool_call_from_value(v))
4390                    .collect();
4391                if !calls.is_empty() {
4392                    return Some(calls);
4393                }
4394            }
4395            // Handle single JSON object
4396            if let Some(tool_call) = self.extract_tool_call_from_value(&parsed) {
4397                return Some(vec![tool_call]);
4398            }
4399        }
4400
4401        None
4402    }
4403
4404    fn extract_tool_call_from_value(&self, parsed: &serde_json::Value) -> Option<ToolCall> {
4405        if let Some(tool_name) = parsed.get("tool").and_then(|v| v.as_str()) {
4406            let arguments = parsed
4407                .get("arguments")
4408                .cloned()
4409                .unwrap_or(serde_json::json!({}));
4410            return Some(ToolCall {
4411                id: parsed
4412                    .get("id")
4413                    .and_then(|value| value.as_str())
4414                    .filter(|id| !id.is_empty())
4415                    .map(str::to_string)
4416                    .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
4417                name: tool_name.to_string(),
4418                arguments,
4419            });
4420        }
4421        None
4422    }
4423
4424    // Lite models could generate unmatched braces: this function handles such cases
4425    fn extract_json_from_content(&self, content: &str) -> Option<String> {
4426        // Try array first (for parallel tool calls), then single object
4427        if let Some(result) = self.extract_json_array_from_content(content) {
4428            return Some(result);
4429        }
4430        self.extract_json_object_from_content(content)
4431    }
4432
4433    /// Extract a JSON array `[...]` containing tool calls from mixed content.
4434    fn extract_json_array_from_content(&self, content: &str) -> Option<String> {
4435        let start = content.find('[')?;
4436        let content_from_start = &content[start..];
4437
4438        let mut depth = 0;
4439        let mut end = 0;
4440        for (i, ch) in content_from_start.char_indices() {
4441            match ch {
4442                '[' => depth += 1,
4443                ']' => {
4444                    depth -= 1;
4445                    if depth == 0 {
4446                        end = i + 1;
4447                        break;
4448                    }
4449                }
4450                _ => {}
4451            }
4452        }
4453
4454        if end > 0 {
4455            let json_str = &content_from_start[..end];
4456            // Verify it looks like an array of tool calls
4457            if json_str.contains("\"tool\"") {
4458                return Some(json_str.to_string());
4459            }
4460        }
4461
4462        None
4463    }
4464
4465    /// Extract a JSON object `{...}` containing a tool call from mixed content.
4466    fn extract_json_object_from_content(&self, content: &str) -> Option<String> {
4467        let start = content.find('{')?;
4468        let content_from_start = &content[start..];
4469
4470        // Count braces to find the matching closing brace
4471        let mut depth = 0;
4472        let mut end = 0;
4473        for (i, ch) in content_from_start.char_indices() {
4474            match ch {
4475                '{' => depth += 1,
4476                '}' => {
4477                    depth -= 1;
4478                    if depth == 0 {
4479                        end = i + 1;
4480                        break;
4481                    }
4482                }
4483                _ => {}
4484            }
4485        }
4486
4487        if end > 0 {
4488            let json_str = &content_from_start[..end];
4489            // Verify it looks like a tool call
4490            if json_str.contains("\"tool\"") {
4491                return Some(json_str.to_string());
4492            }
4493        }
4494
4495        None
4496    }
4497
4498    /// Builds a structured record from executor state.
4499    ///
4500    /// The explicit fields preserve one audit boundary for execution, policy, approval, timing, and generation evidence.
4501    #[allow(clippy::too_many_arguments)]
4502    fn record_from_parts(
4503        &self,
4504        request: &ToolExecutionRequest,
4505        canonical_id: String,
4506        executed_arguments: Value,
4507        started_at: chrono::DateTime<chrono::Utc>,
4508        start: Instant,
4509        executed: bool,
4510        success: bool,
4511        output: String,
4512        metadata: HashMap<String, Value>,
4513        policy: ToolPolicyDecisionRecord,
4514        approval: Option<ToolApprovalRecord>,
4515        timed_out: bool,
4516        output_truncated: bool,
4517    ) -> ToolExecutionRecord {
4518        let versions = ToolDecisionVersions {
4519            policy: self.active_tool_security().policy_version(),
4520            registry: self.tools.version(),
4521            runtime_control: self.runtime_control.version.load(Ordering::SeqCst),
4522            state: self
4523                .state_machine
4524                .as_ref()
4525                .map(|state_machine| state_machine.generation()),
4526        };
4527        self.record_from_parts_at(
4528            request,
4529            canonical_id,
4530            executed_arguments,
4531            started_at,
4532            start,
4533            executed,
4534            success,
4535            output,
4536            metadata,
4537            policy,
4538            approval,
4539            timed_out,
4540            output_truncated,
4541            versions,
4542        )
4543    }
4544
4545    /// Builds evidence with the exact generations used by the corresponding decision.
4546    #[allow(clippy::too_many_arguments)]
4547    fn record_from_parts_at(
4548        &self,
4549        request: &ToolExecutionRequest,
4550        canonical_id: String,
4551        executed_arguments: Value,
4552        started_at: chrono::DateTime<chrono::Utc>,
4553        start: Instant,
4554        executed: bool,
4555        success: bool,
4556        output: String,
4557        metadata: HashMap<String, Value>,
4558        policy: ToolPolicyDecisionRecord,
4559        approval: Option<ToolApprovalRecord>,
4560        timed_out: bool,
4561        output_truncated: bool,
4562        versions: ToolDecisionVersions,
4563    ) -> ToolExecutionRecord {
4564        ToolExecutionRecord {
4565            call_id: request.call_id.clone(),
4566            requested_name: request.requested_name.clone(),
4567            canonical_id,
4568            source: request.source.clone(),
4569            arguments: request.arguments.clone(),
4570            executed_arguments,
4571            policy_version: versions.policy,
4572            registry_version: versions.registry,
4573            runtime_config_version: versions.runtime_control,
4574            executed,
4575            success,
4576            output,
4577            metadata,
4578            policy,
4579            approval,
4580            started_at,
4581            duration_ms: start.elapsed().as_millis() as u64,
4582            timed_out,
4583            cancelled: false,
4584            cancellation_reason: None,
4585            output_truncated,
4586        }
4587    }
4588
4589    /// Sends a finalized tool record to hooks, history, and error handling.
4590    async fn finish_tool_record(&self, record: &ToolExecutionRecord) {
4591        let result = ToolResult {
4592            success: record.success,
4593            output: record.model_output_string(),
4594            metadata: if record.metadata.is_empty() {
4595                None
4596            } else {
4597                Some(record.metadata.clone())
4598            },
4599        };
4600        self.hooks
4601            .on_tool_complete(&record.canonical_id, &result, record.duration_ms)
4602            .await;
4603        self.hooks.on_tool_execution_record(record).await;
4604        self.record_tool_call(&record.canonical_id, record.model_output_value());
4605        if !record.success {
4606            self.hooks
4607                .on_error(&AgentError::Tool(record.output.clone()))
4608                .await;
4609        }
4610    }
4611
4612    /// Releases resource locks before hooks run because hooks may invoke another tool.
4613    async fn finish_tool_record_after_resource_guards(
4614        &self,
4615        resource_guards: ToolResourceGuards,
4616        record: &ToolExecutionRecord,
4617    ) {
4618        drop(resource_guards);
4619        self.finish_tool_record(record).await;
4620    }
4621
4622    //
4623    // Validates the stable timeout range once and derives both runtime representations without narrowing or clamping.
4624    //
4625    fn validated_tool_timeout(timeout_ms: u64) -> Result<ValidatedToolTimeout> {
4626        if timeout_ms > MAX_TOOL_TIMEOUT_MS {
4627            return Err(AgentError::Config(format!(
4628                "effective tool timeout_ms must be no greater than {MAX_TOOL_TIMEOUT_MS} milliseconds"
4629            )));
4630        }
4631        let timer = Duration::from_millis(timeout_ms);
4632        let deadline_delta = chrono::Duration::from_std(timer).map_err(|_| {
4633            AgentError::Config(format!(
4634                "effective tool timeout_ms cannot be represented as a UTC deadline: {timeout_ms}"
4635            ))
4636        })?;
4637        Ok(ValidatedToolTimeout {
4638            timer,
4639            deadline_delta,
4640        })
4641    }
4642
4643    //
4644    // Selects one invocation cap while ensuring call classification and recovery can narrow but never widen security policy.
4645    //
4646    fn effective_tool_limits(
4647        security_engine: &ToolSecurityEngine,
4648        canonical_id: &str,
4649        safety: &ToolSafetyMetadata,
4650        classification: &ToolCallClassification,
4651        recovery_timeout_ms: Option<u64>,
4652    ) -> Result<(ToolExecutionLimits, ValidatedToolTimeout)> {
4653        if let Some(timeout_ms) = classification.timeout_ms {
4654            Self::validated_tool_timeout(timeout_ms)?;
4655        }
4656        if let Some(timeout_ms) = recovery_timeout_ms {
4657            Self::validated_tool_timeout(timeout_ms)?;
4658        }
4659
4660        let mut limits = security_engine.effective_limits(canonical_id, safety, classification);
4661        if let Some(recovery_timeout_ms) = recovery_timeout_ms {
4662            limits.timeout_ms = Some(limits.timeout_ms.map_or(recovery_timeout_ms, |timeout_ms| {
4663                timeout_ms.min(recovery_timeout_ms)
4664            }));
4665        }
4666        let timeout_ms = limits
4667            .timeout_ms
4668            .unwrap_or_else(|| security_engine.get_tool_timeout(canonical_id));
4669        let timeout = Self::validated_tool_timeout(timeout_ms)?;
4670        Ok((limits, timeout))
4671    }
4672
4673    /// Invokes a resolved tool once with one fresh attempt deadline, timeout, cancellation, and actor context.
4674    async fn execute_resolved_tool_once(
4675        &self,
4676        tool: Arc<dyn ai_agents_core::Tool>,
4677        args: Value,
4678        mut ctx: ToolExecutionContext,
4679        timeout: ValidatedToolTimeout,
4680    ) -> Result<(ToolResult, bool, bool, bool)> {
4681        if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4682            return Ok((
4683                ToolResult::error("Tool execution cancelled by runtime control"),
4684                false,
4685                true,
4686                false,
4687            ));
4688        }
4689        //
4690        // The tool observes the same checked timeout enforced below, starting immediately before this invocation attempt.
4691        // Each retry receives a new deadline rather than inheriting time spent in policy handling or earlier attempts.
4692        //
4693        ctx.deadline = Some(
4694            chrono::Utc::now()
4695                .checked_add_signed(timeout.deadline_delta)
4696                .ok_or_else(|| {
4697                    AgentError::Config(
4698                        "effective tool timeout_ms exceeds the current UTC deadline range"
4699                            .to_string(),
4700                    )
4701                })?,
4702        );
4703        //
4704        // Mark invocation inside the future so cancellation before the first poll remains executed false.
4705        //
4706        let invoked = Arc::new(AtomicBool::new(false));
4707        let invoked_by_future = Arc::clone(&invoked);
4708        let actor_context = current_turn_actor_context();
4709        let future = async move {
4710            invoked_by_future.store(true, Ordering::SeqCst);
4711            if let Some(actor_context) = actor_context {
4712                scope_actor_context(actor_context, tool.execute(args, ctx)).await
4713            } else {
4714                tool.execute(args, ctx).await
4715            }
4716        };
4717        tokio::pin!(future);
4718        let timer = tokio::time::sleep(timeout.timer);
4719        tokio::pin!(timer);
4720        let mut cancel_tick = tokio::time::interval(std::time::Duration::from_millis(50));
4721
4722        loop {
4723            tokio::select! {
4724                result = &mut future => return Ok((result, false, false, true)),
4725                _ = &mut timer => {
4726                    return Ok((
4727                        ToolResult::error("Tool execution timed out"),
4728                        true,
4729                        false,
4730                        invoked.load(Ordering::SeqCst),
4731                    ));
4732                }
4733                _ = cancel_tick.tick() => {
4734                    if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4735                        return Ok((
4736                            ToolResult::error("Tool execution cancelled by runtime control"),
4737                            false,
4738                            true,
4739                            invoked.load(Ordering::SeqCst),
4740                        ));
4741                    }
4742                }
4743            }
4744        }
4745    }
4746
4747    /// Truncates model-facing tool output on a character boundary.
4748    fn truncate_tool_output(output: String, max_chars: Option<usize>) -> (String, bool) {
4749        let Some(max_chars) = max_chars else {
4750            return (output, false);
4751        };
4752        let mut chars = output.chars();
4753        let truncated: String = chars.by_ref().take(max_chars).collect();
4754        if chars.next().is_some() {
4755            (truncated, true)
4756        } else {
4757            (output, false)
4758        }
4759    }
4760
4761    /// Acquires all declared resource locks in stable key order and supports emergency cancellation while waiting.
4762    async fn acquire_tool_resource_locks(&self, keys: &[String]) -> Option<ToolResourceGuards> {
4763        let locks = {
4764            let mut table = self.resource_locks.write();
4765            table.retain(|_, lock| lock.strong_count() > 0);
4766            keys.iter()
4767                .map(|key| {
4768                    if let Some(lock) = table.get(key).and_then(Weak::upgrade) {
4769                        lock
4770                    } else {
4771                        let lock = Arc::new(tokio::sync::Mutex::new(()));
4772                        table.insert(key.clone(), Arc::downgrade(&lock));
4773                        lock
4774                    }
4775                })
4776                .collect::<Vec<_>>()
4777        };
4778        let mut resource_guards = ToolResourceGuards {
4779            guards: Vec::with_capacity(locks.len()),
4780            locks: Arc::clone(&self.resource_locks),
4781        };
4782        let mut locks = locks.into_iter();
4783        while let Some(lock) = locks.next() {
4784            let mut lock = Box::pin(lock.lock_owned());
4785            loop {
4786                tokio::select! {
4787                    guard = &mut lock => {
4788                        resource_guards.guards.push(guard);
4789                        break;
4790                    }
4791                    _ = tokio::time::sleep(std::time::Duration::from_millis(10)) => {
4792                        if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4793                            drop(lock);
4794                            drop(locks);
4795                            drop(resource_guards);
4796                            return None;
4797                        }
4798                    }
4799                }
4800            }
4801        }
4802        Some(resource_guards)
4803    }
4804
4805    /// Executes retry sub-attempts inside one logical executor request while assigning every invocation a fresh deadline.
4806    ///
4807    /// Hooks and `ToolExecutionRecord` finalization remain request-level and occur once after this retry loop returns.
4808    async fn run_tool_with_retries(
4809        &self,
4810        canonical_id: &str,
4811        tool: Arc<dyn ai_agents_core::Tool>,
4812        args: Value,
4813        ctx: ToolExecutionContext,
4814        timeout: ValidatedToolTimeout,
4815        max_retries: u32,
4816    ) -> Result<(ToolResult, bool, bool, bool)> {
4817        let max_retries = if ctx.classification.safely_retryable {
4818            max_retries
4819        } else {
4820            0
4821        };
4822        let mut attempts = 0;
4823        let mut invoked = false;
4824        loop {
4825            let (result, timed_out, cancelled, attempt_invoked) = self
4826                .execute_resolved_tool_once(tool.clone(), args.clone(), ctx.clone(), timeout)
4827                .await?;
4828            invoked |= attempt_invoked;
4829            if result.success || timed_out || cancelled || attempts >= max_retries {
4830                return Ok((result, timed_out, cancelled, invoked));
4831            }
4832            attempts += 1;
4833            warn!(tool = %canonical_id, attempt = attempts, error = %result.output, "Retrying failed tool call");
4834        }
4835    }
4836
4837    /// Returns the existing model output and evidence reason when a host-backed tool cannot run.
4838    fn host_tool_unavailability(&self, canonical_id: &str) -> Option<(&'static str, &'static str)> {
4839        match canonical_id {
4840            "command" if !self.tools.command_runner_available() => Some((
4841                "Command runner is unavailable",
4842                "command runner is unavailable",
4843            )),
4844            "diagnostics" if !self.tools.diagnostics_available() => Some((
4845                "Diagnostics provider is unavailable",
4846                "diagnostics provider is unavailable",
4847            )),
4848            "web_search" if !self.tools.web_search_available() => Some((
4849                "Web search provider is unavailable",
4850                "web search provider is unavailable",
4851            )),
4852            _ => None,
4853        }
4854    }
4855
4856    /// Executes a tool request through scope, policy, HITL, timeout, bounded recovery, and evidence recording.
4857    fn execute_tool_record(
4858        &self,
4859        request: ToolExecutionRequest,
4860    ) -> Pin<Box<dyn Future<Output = Result<ToolExecutionRecord>> + Send + '_>> {
4861        Box::pin(self.execute_tool_record_inner(request, ToolFallbackState::default()))
4862    }
4863
4864    /// Implements one logical shared-executor request while preserving policy, HITL, availability, final admission, hooks, retry evidence, and bounded fallback ordering.
4865    ///
4866    /// A failed request selected for fallback releases its guards and finalizes its own record before the fallback starts as a separate shared-executor request. Canonical ancestry crosses that boundary so alias-mediated cycles and overlong acyclic chains stop before the start hook, approval, locks, or invocation while retaining terminal completion evidence.
4867    async fn execute_tool_record_inner(
4868        &self,
4869        request: ToolExecutionRequest,
4870        fallback_state: ToolFallbackState,
4871    ) -> Result<ToolExecutionRecord> {
4872        let started_at = chrono::Utc::now();
4873        let start = Instant::now();
4874        info!(tool = %request.requested_name, args = %request.arguments, "Executing tool");
4875
4876        if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4877            let record = self.record_from_parts(
4878                &request,
4879                request.requested_name.clone(),
4880                request.arguments.clone(),
4881                started_at,
4882                start,
4883                false,
4884                false,
4885                "Tool execution is disabled by runtime control".to_string(),
4886                HashMap::new(),
4887                ToolPolicyDecisionRecord::deny("runtime emergency deny is enabled"),
4888                None,
4889                false,
4890                false,
4891            );
4892            self.finish_tool_record(&record).await;
4893            return Ok(record);
4894        }
4895
4896        let Some(resolved) = self.tools.resolve(&request.requested_name) else {
4897            let record = self.record_from_parts(
4898                &request,
4899                request.requested_name.clone(),
4900                request.arguments.clone(),
4901                started_at,
4902                start,
4903                false,
4904                false,
4905                format!("Tool '{}' is unavailable", request.requested_name),
4906                HashMap::new(),
4907                ToolPolicyDecisionRecord::unavailable(format!(
4908                    "Tool '{}' is not registered",
4909                    request.requested_name
4910                )),
4911                None,
4912                false,
4913                false,
4914            );
4915            self.finish_tool_record(&record).await;
4916            return Ok(record);
4917        };
4918
4919        let canonical_id = resolved.identity.canonical_id.clone();
4920
4921        let initial_scope_snapshot = self.get_available_tool_ids_snapshot().await?;
4922        if !initial_scope_snapshot
4923            .tool_ids
4924            .iter()
4925            .any(|id| id == &canonical_id)
4926        {
4927            let record = self.record_from_parts(
4928                &request,
4929                canonical_id.clone(),
4930                request.arguments.clone(),
4931                started_at,
4932                start,
4933                false,
4934                false,
4935                format!(
4936                    "Tool '{}' is not available in the current scope",
4937                    canonical_id
4938                ),
4939                HashMap::new(),
4940                ToolPolicyDecisionRecord::deny(format!(
4941                    "Tool '{}' is not granted by the current top-level and state tool scope",
4942                    canonical_id
4943                )),
4944                None,
4945                false,
4946                false,
4947            );
4948            self.finish_tool_record(&record).await;
4949            return Ok(record);
4950        }
4951
4952        let approval_control_snapshot = self.runtime_safety_snapshot();
4953        let security_engine = approval_control_snapshot.tool_security.clone();
4954        if let Some(reason) = fallback_state.rejection_reason(&canonical_id) {
4955            let mut metadata = HashMap::new();
4956            metadata.insert(
4957                "fallback_chain".to_string(),
4958                serde_json::to_value(&fallback_state.visited_canonical_ids).unwrap_or(Value::Null),
4959            );
4960            let record = self.record_from_parts(
4961                &request,
4962                canonical_id,
4963                request.arguments.clone(),
4964                started_at,
4965                start,
4966                false,
4967                false,
4968                format!("Denied: {reason}"),
4969                metadata,
4970                ToolPolicyDecisionRecord::deny(reason),
4971                None,
4972                false,
4973                false,
4974            );
4975            self.finish_tool_record(&record).await;
4976            return Ok(record);
4977        }
4978        let admitted_canonical_id = canonical_id.clone();
4979        let fallback_state = fallback_state.with_current(canonical_id.clone());
4980        let bindings = resolved.tool.policy_bindings();
4981        let mut executed_arguments = security_engine.prepare_tool_arguments_with_bindings(
4982            &canonical_id,
4983            &request.arguments,
4984            &bindings,
4985        );
4986        let mut metadata = HashMap::new();
4987        let safety = resolved.tool.safety_metadata();
4988        let classification = resolved.tool.classify_call(&executed_arguments);
4989        let initial_recovery_timeout_ms = self.recovery_manager.get_tool_timeout(&canonical_id);
4990        let (limits, _) = Self::effective_tool_limits(
4991            &security_engine,
4992            &canonical_id,
4993            &safety,
4994            &classification,
4995            initial_recovery_timeout_ms,
4996        )?;
4997        self.hooks
4998            .on_tool_start(&canonical_id, &executed_arguments)
4999            .await;
5000        metadata.insert(
5001            "classification".to_string(),
5002            serde_json::to_value(&classification).unwrap_or(Value::Null),
5003        );
5004        metadata.insert(
5005            "effective_limits".to_string(),
5006            serde_json::to_value(&limits).unwrap_or(Value::Null),
5007        );
5008        let policy_snapshot = security_engine.policy_snapshot(&canonical_id);
5009        if !policy_snapshot.is_null() {
5010            metadata.insert("policy_snapshot".to_string(), policy_snapshot.clone());
5011        }
5012
5013        let mut approval_record = Some(ToolApprovalRecord {
5014            status: ToolApprovalStatus::NotRequired,
5015            reason: None,
5016            modified_arguments: None,
5017        });
5018
5019        let mut security_result = security_engine
5020            .validate_tool_execution_with_bindings(&canonical_id, &executed_arguments, &bindings)
5021            .await?;
5022        //
5023        // Terminal policy denials remain authoritative, while allowed or confirmation-requiring calls must fail host availability before any HITL request.
5024        // The same capability is checked again after HITL because the host can change while approval waits.
5025        //
5026        if (security_result.is_allowed()
5027            || matches!(
5028                &security_result,
5029                SecurityCheckResult::RequireConfirmation { .. }
5030            ))
5031            && let Some((output, reason)) = self.host_tool_unavailability(&canonical_id)
5032        {
5033            let record = self.record_from_parts(
5034                &request,
5035                canonical_id,
5036                executed_arguments,
5037                started_at,
5038                start,
5039                false,
5040                false,
5041                output.to_string(),
5042                metadata,
5043                ToolPolicyDecisionRecord::unavailable(reason),
5044                Some(ToolApprovalRecord {
5045                    status: ToolApprovalStatus::Unavailable,
5046                    reason: Some(reason.to_string()),
5047                    modified_arguments: None,
5048                }),
5049                false,
5050                false,
5051            );
5052            self.finish_tool_record(&record).await;
5053            return Ok(record);
5054        }
5055        match &security_result {
5056            SecurityCheckResult::Allow => {}
5057            SecurityCheckResult::Warn { message } => {
5058                warn!(tool = %canonical_id, message = %message, "Tool security warning");
5059            }
5060            SecurityCheckResult::Block { reason } => {
5061                let record = self.record_from_parts(
5062                    &request,
5063                    canonical_id,
5064                    executed_arguments,
5065                    started_at,
5066                    start,
5067                    false,
5068                    false,
5069                    format!("Denied: {}", reason),
5070                    metadata,
5071                    ToolPolicyDecisionRecord::deny(reason.clone()),
5072                    approval_record,
5073                    false,
5074                    false,
5075                );
5076                self.finish_tool_record(&record).await;
5077                return Ok(record);
5078            }
5079            SecurityCheckResult::Unavailable { reason } => {
5080                let record = self.record_from_parts(
5081                    &request,
5082                    canonical_id,
5083                    executed_arguments,
5084                    started_at,
5085                    start,
5086                    false,
5087                    false,
5088                    format!("Unavailable: {}", reason),
5089                    metadata,
5090                    ToolPolicyDecisionRecord::unavailable(reason.clone()),
5091                    approval_record,
5092                    false,
5093                    false,
5094                );
5095                self.finish_tool_record(&record).await;
5096                return Ok(record);
5097            }
5098            SecurityCheckResult::RequireConfirmation { message } => {
5099                if self.hitl_engine.is_none() {
5100                    approval_record = Some(ToolApprovalRecord {
5101                        status: ToolApprovalStatus::Unavailable,
5102                        reason: Some("No HITL engine configured".to_string()),
5103                        modified_arguments: None,
5104                    });
5105                    let record = self.record_from_parts(
5106                        &request,
5107                        canonical_id,
5108                        executed_arguments,
5109                        started_at,
5110                        start,
5111                        false,
5112                        false,
5113                        format!("Approval unavailable: {}", message),
5114                        metadata,
5115                        ToolPolicyDecisionRecord::approval(message.clone()),
5116                        approval_record,
5117                        false,
5118                        false,
5119                    );
5120                    self.finish_tool_record(&record).await;
5121                    return Ok(record);
5122                }
5123
5124                let check_result = HITLCheckResult::required(
5125                    ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
5126                    HashMap::new(),
5127                    message.clone(),
5128                    None,
5129                );
5130                match self.request_hitl_approval(check_result).await? {
5131                    ApprovalResult::Approved => {
5132                        merge_approved_record(&mut approval_record);
5133                    }
5134                    ApprovalResult::Modified { changes } => {
5135                        if let Some(obj) = executed_arguments.as_object_mut() {
5136                            for (key, value) in changes {
5137                                obj.insert(key, value);
5138                            }
5139                        }
5140                        security_result = security_engine
5141                            .validate_tool_execution_with_bindings(
5142                                &canonical_id,
5143                                &executed_arguments,
5144                                &bindings,
5145                            )
5146                            .await?;
5147                        if !matches!(
5148                            security_result,
5149                            SecurityCheckResult::Allow
5150                                | SecurityCheckResult::Warn { .. }
5151                                | SecurityCheckResult::RequireConfirmation { .. }
5152                        ) {
5153                            let reason = security_result
5154                                .reason()
5155                                .unwrap_or("modified arguments failed policy")
5156                                .to_string();
5157                            let record = self.record_from_parts(
5158                                &request,
5159                                canonical_id,
5160                                executed_arguments.clone(),
5161                                started_at,
5162                                start,
5163                                false,
5164                                false,
5165                                reason.clone(),
5166                                metadata,
5167                                ToolPolicyDecisionRecord::deny(reason),
5168                                Some(ToolApprovalRecord {
5169                                    status: ToolApprovalStatus::Modified,
5170                                    reason: None,
5171                                    modified_arguments: Some(executed_arguments),
5172                                }),
5173                                false,
5174                                false,
5175                            );
5176                            self.finish_tool_record(&record).await;
5177                            return Ok(record);
5178                        }
5179                        approval_record = Some(ToolApprovalRecord {
5180                            status: ToolApprovalStatus::Modified,
5181                            reason: None,
5182                            modified_arguments: Some(executed_arguments.clone()),
5183                        });
5184                    }
5185                    ApprovalResult::Rejected { reason } => {
5186                        let reason = reason.unwrap_or_else(|| "rejected".to_string());
5187                        approval_record = Some(ToolApprovalRecord {
5188                            status: ToolApprovalStatus::Rejected,
5189                            reason: Some(reason.clone()),
5190                            modified_arguments: None,
5191                        });
5192                        let record = self.record_from_parts(
5193                            &request,
5194                            canonical_id,
5195                            executed_arguments,
5196                            started_at,
5197                            start,
5198                            false,
5199                            false,
5200                            format!("Approval rejected: {}", reason),
5201                            metadata,
5202                            ToolPolicyDecisionRecord::approval(reason),
5203                            approval_record,
5204                            false,
5205                            false,
5206                        );
5207                        self.finish_tool_record(&record).await;
5208                        return Ok(record);
5209                    }
5210                    ApprovalResult::Timeout => {
5211                        approval_record = Some(ToolApprovalRecord {
5212                            status: ToolApprovalStatus::Timeout,
5213                            reason: Some("approval timeout".to_string()),
5214                            modified_arguments: None,
5215                        });
5216                        let record = self.record_from_parts(
5217                            &request,
5218                            canonical_id,
5219                            executed_arguments,
5220                            started_at,
5221                            start,
5222                            false,
5223                            false,
5224                            "Approval timed out".to_string(),
5225                            metadata,
5226                            ToolPolicyDecisionRecord::approval("approval timeout"),
5227                            approval_record,
5228                            false,
5229                            false,
5230                        );
5231                        self.finish_tool_record(&record).await;
5232                        return Ok(record);
5233                    }
5234                }
5235            }
5236        }
5237
5238        if approval_record
5239            .as_ref()
5240            .is_some_and(|record| matches!(record.status, ToolApprovalStatus::NotRequired))
5241            && let Some(message) =
5242                security_engine.classification_approval_message(&canonical_id, &classification)
5243        {
5244            if self.hitl_engine.is_none() {
5245                approval_record = Some(ToolApprovalRecord {
5246                    status: ToolApprovalStatus::Unavailable,
5247                    reason: Some("No HITL engine configured".to_string()),
5248                    modified_arguments: None,
5249                });
5250                let record = self.record_from_parts(
5251                    &request,
5252                    canonical_id,
5253                    executed_arguments,
5254                    started_at,
5255                    start,
5256                    false,
5257                    false,
5258                    format!("Approval unavailable: {}", message),
5259                    metadata,
5260                    ToolPolicyDecisionRecord::approval(message),
5261                    approval_record,
5262                    false,
5263                    false,
5264                );
5265                self.finish_tool_record(&record).await;
5266                return Ok(record);
5267            }
5268            let check_result = HITLCheckResult::required(
5269                ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
5270                HashMap::new(),
5271                message.clone(),
5272                None,
5273            );
5274            match self.request_hitl_approval(check_result).await? {
5275                ApprovalResult::Approved => {
5276                    merge_approved_record(&mut approval_record);
5277                }
5278                ApprovalResult::Modified { changes } => {
5279                    if let Some(obj) = executed_arguments.as_object_mut() {
5280                        for (key, value) in changes {
5281                            obj.insert(key, value);
5282                        }
5283                    }
5284                    let modified_security = security_engine
5285                        .validate_tool_execution_with_bindings(
5286                            &canonical_id,
5287                            &executed_arguments,
5288                            &bindings,
5289                        )
5290                        .await?;
5291                    if !matches!(
5292                        modified_security,
5293                        SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5294                    ) {
5295                        let reason = modified_security
5296                            .reason()
5297                            .unwrap_or("modified arguments failed policy")
5298                            .to_string();
5299                        let record = self.record_from_parts(
5300                            &request,
5301                            canonical_id,
5302                            executed_arguments.clone(),
5303                            started_at,
5304                            start,
5305                            false,
5306                            false,
5307                            reason.clone(),
5308                            metadata,
5309                            ToolPolicyDecisionRecord::deny(reason),
5310                            Some(ToolApprovalRecord {
5311                                status: ToolApprovalStatus::Modified,
5312                                reason: None,
5313                                modified_arguments: Some(executed_arguments),
5314                            }),
5315                            false,
5316                            false,
5317                        );
5318                        self.finish_tool_record(&record).await;
5319                        return Ok(record);
5320                    }
5321                    approval_record = Some(ToolApprovalRecord {
5322                        status: ToolApprovalStatus::Modified,
5323                        reason: None,
5324                        modified_arguments: Some(executed_arguments.clone()),
5325                    });
5326                }
5327                ApprovalResult::Rejected { reason } => {
5328                    let reason = reason.unwrap_or_else(|| "rejected".to_string());
5329                    let record = self.record_from_parts(
5330                        &request,
5331                        canonical_id,
5332                        executed_arguments,
5333                        started_at,
5334                        start,
5335                        false,
5336                        false,
5337                        format!("Approval rejected: {}", reason),
5338                        metadata,
5339                        ToolPolicyDecisionRecord::approval(reason.clone()),
5340                        Some(ToolApprovalRecord {
5341                            status: ToolApprovalStatus::Rejected,
5342                            reason: Some(reason),
5343                            modified_arguments: None,
5344                        }),
5345                        false,
5346                        false,
5347                    );
5348                    self.finish_tool_record(&record).await;
5349                    return Ok(record);
5350                }
5351                ApprovalResult::Timeout => {
5352                    let record = self.record_from_parts(
5353                        &request,
5354                        canonical_id,
5355                        executed_arguments,
5356                        started_at,
5357                        start,
5358                        false,
5359                        false,
5360                        "Approval timed out".to_string(),
5361                        metadata,
5362                        ToolPolicyDecisionRecord::approval("approval timeout"),
5363                        Some(ToolApprovalRecord {
5364                            status: ToolApprovalStatus::Timeout,
5365                            reason: Some("approval timeout".to_string()),
5366                            modified_arguments: None,
5367                        }),
5368                        false,
5369                        false,
5370                    );
5371                    self.finish_tool_record(&record).await;
5372                    return Ok(record);
5373                }
5374            }
5375        }
5376
5377        let hitl_lang_ctx = self.build_hitl_language_context();
5378        if let Some(ref hitl_engine) = self.hitl_engine {
5379            let check_result = self
5380                .observe_purpose(
5381                    ObservationPurpose::HitlLocalization,
5382                    hitl_engine.check_tool_with_localization(
5383                        &canonical_id,
5384                        &executed_arguments,
5385                        &hitl_lang_ctx,
5386                        self.approval_handler.as_ref(),
5387                        Some(&self.llm_registry),
5388                    ),
5389                )
5390                .await?;
5391            if check_result.is_required() {
5392                match self.request_hitl_approval(check_result).await? {
5393                    ApprovalResult::Approved => {
5394                        merge_approved_record(&mut approval_record);
5395                    }
5396                    ApprovalResult::Modified { changes } => {
5397                        if let Some(obj) = executed_arguments.as_object_mut() {
5398                            for (key, value) in changes {
5399                                obj.insert(key, value);
5400                            }
5401                        }
5402                        let modified_security = security_engine
5403                            .validate_tool_execution_with_bindings(
5404                                &canonical_id,
5405                                &executed_arguments,
5406                                &bindings,
5407                            )
5408                            .await?;
5409                        if !matches!(
5410                            modified_security,
5411                            SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5412                        ) {
5413                            let reason = modified_security
5414                                .reason()
5415                                .unwrap_or("modified arguments failed policy")
5416                                .to_string();
5417                            let record = self.record_from_parts(
5418                                &request,
5419                                canonical_id,
5420                                executed_arguments.clone(),
5421                                started_at,
5422                                start,
5423                                false,
5424                                false,
5425                                reason.clone(),
5426                                metadata,
5427                                ToolPolicyDecisionRecord::deny(reason),
5428                                Some(ToolApprovalRecord {
5429                                    status: ToolApprovalStatus::Modified,
5430                                    reason: None,
5431                                    modified_arguments: Some(executed_arguments),
5432                                }),
5433                                false,
5434                                false,
5435                            );
5436                            self.finish_tool_record(&record).await;
5437                            return Ok(record);
5438                        }
5439                        approval_record = Some(ToolApprovalRecord {
5440                            status: ToolApprovalStatus::Modified,
5441                            reason: None,
5442                            modified_arguments: Some(executed_arguments.clone()),
5443                        });
5444                    }
5445                    ApprovalResult::Rejected { reason } => {
5446                        let reason = reason.unwrap_or_else(|| "rejected".to_string());
5447                        let record = self.record_from_parts(
5448                            &request,
5449                            canonical_id,
5450                            executed_arguments,
5451                            started_at,
5452                            start,
5453                            false,
5454                            false,
5455                            format!("Approval rejected: {}", reason),
5456                            metadata,
5457                            ToolPolicyDecisionRecord::approval(reason.clone()),
5458                            Some(ToolApprovalRecord {
5459                                status: ToolApprovalStatus::Rejected,
5460                                reason: Some(reason),
5461                                modified_arguments: None,
5462                            }),
5463                            false,
5464                            false,
5465                        );
5466                        self.finish_tool_record(&record).await;
5467                        return Ok(record);
5468                    }
5469                    ApprovalResult::Timeout => {
5470                        let record = self.record_from_parts(
5471                            &request,
5472                            canonical_id,
5473                            executed_arguments,
5474                            started_at,
5475                            start,
5476                            false,
5477                            false,
5478                            "Approval timed out".to_string(),
5479                            metadata,
5480                            ToolPolicyDecisionRecord::approval("approval timeout"),
5481                            Some(ToolApprovalRecord {
5482                                status: ToolApprovalStatus::Timeout,
5483                                reason: Some("approval timeout".to_string()),
5484                                modified_arguments: None,
5485                            }),
5486                            false,
5487                            false,
5488                        );
5489                        self.finish_tool_record(&record).await;
5490                        return Ok(record);
5491                    }
5492                }
5493            }
5494
5495            let condition_check = self
5496                .observe_purpose(
5497                    ObservationPurpose::HitlLocalization,
5498                    hitl_engine.check_conditions_with_localization(
5499                        &executed_arguments,
5500                        &hitl_lang_ctx,
5501                        self.approval_handler.as_ref(),
5502                        Some(&self.llm_registry),
5503                    ),
5504                )
5505                .await?;
5506            if condition_check.is_required() {
5507                match self.request_hitl_approval(condition_check).await? {
5508                    ApprovalResult::Approved => {
5509                        merge_approved_record(&mut approval_record);
5510                    }
5511                    ApprovalResult::Modified { changes } => {
5512                        if let Some(obj) = executed_arguments.as_object_mut() {
5513                            for (key, value) in changes {
5514                                obj.insert(key, value);
5515                            }
5516                        }
5517                        let modified_security = security_engine
5518                            .validate_tool_execution_with_bindings(
5519                                &canonical_id,
5520                                &executed_arguments,
5521                                &bindings,
5522                            )
5523                            .await?;
5524                        if !matches!(
5525                            modified_security,
5526                            SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5527                        ) {
5528                            let reason = modified_security
5529                                .reason()
5530                                .unwrap_or("modified arguments failed policy")
5531                                .to_string();
5532                            let record = self.record_from_parts(
5533                                &request,
5534                                canonical_id,
5535                                executed_arguments,
5536                                started_at,
5537                                start,
5538                                false,
5539                                false,
5540                                reason.clone(),
5541                                metadata,
5542                                ToolPolicyDecisionRecord::deny(reason),
5543                                approval_record,
5544                                false,
5545                                false,
5546                            );
5547                            self.finish_tool_record(&record).await;
5548                            return Ok(record);
5549                        }
5550                        approval_record = Some(ToolApprovalRecord {
5551                            status: ToolApprovalStatus::Modified,
5552                            reason: None,
5553                            modified_arguments: Some(executed_arguments.clone()),
5554                        });
5555                    }
5556                    ApprovalResult::Rejected { reason } => {
5557                        let reason = reason.unwrap_or_else(|| "rejected".to_string());
5558                        let record = self.record_from_parts(
5559                            &request,
5560                            canonical_id,
5561                            executed_arguments,
5562                            started_at,
5563                            start,
5564                            false,
5565                            false,
5566                            format!("Approval rejected: {}", reason),
5567                            metadata,
5568                            ToolPolicyDecisionRecord::approval(reason.clone()),
5569                            Some(ToolApprovalRecord {
5570                                status: ToolApprovalStatus::Rejected,
5571                                reason: Some(reason),
5572                                modified_arguments: None,
5573                            }),
5574                            false,
5575                            false,
5576                        );
5577                        self.finish_tool_record(&record).await;
5578                        return Ok(record);
5579                    }
5580                    ApprovalResult::Timeout => {
5581                        let record = self.record_from_parts(
5582                            &request,
5583                            canonical_id,
5584                            executed_arguments,
5585                            started_at,
5586                            start,
5587                            false,
5588                            false,
5589                            "Approval timed out".to_string(),
5590                            metadata,
5591                            ToolPolicyDecisionRecord::approval("approval timeout"),
5592                            Some(ToolApprovalRecord {
5593                                status: ToolApprovalStatus::Timeout,
5594                                reason: Some("approval timeout".to_string()),
5595                                modified_arguments: None,
5596                            }),
5597                            false,
5598                            false,
5599                        );
5600                        self.finish_tool_record(&record).await;
5601                        return Ok(record);
5602                    }
5603                }
5604            }
5605        }
5606
5607        //
5608        // Freeze the action reviewed by HITL after every approved argument change and policy cap.
5609        // A later plain approval preserves any earlier Modified evidence.
5610        //
5611        executed_arguments = security_engine.prepare_tool_arguments_with_bindings(
5612            &canonical_id,
5613            &executed_arguments,
5614            &bindings,
5615        );
5616        if let Some(record) = approval_record.as_mut()
5617            && matches!(record.status, ToolApprovalStatus::Modified)
5618        {
5619            record.modified_arguments = Some(executed_arguments.clone());
5620        }
5621        let binding_security_result = security_engine
5622            .validate_tool_execution_with_bindings(&canonical_id, &executed_arguments, &bindings)
5623            .await?;
5624        let approval_confirmation_required = matches!(
5625            binding_security_result,
5626            SecurityCheckResult::RequireConfirmation { .. }
5627        ) || security_engine
5628            .classification_approval_message(
5629                &canonical_id,
5630                &resolved.tool.classify_call(&executed_arguments),
5631            )
5632            .is_some();
5633        let approval_binding = approval_record.as_ref().and_then(|record| {
5634            matches!(
5635                record.status,
5636                ToolApprovalStatus::Approved | ToolApprovalStatus::Modified
5637            )
5638            .then(|| ToolApprovalBinding {
5639                canonical_id: canonical_id.clone(),
5640                arguments: executed_arguments.clone(),
5641                confirmation_required: approval_confirmation_required,
5642                policy_version: security_engine.policy_version(),
5643                runtime_control_version: approval_control_snapshot.version,
5644                state_generation: initial_scope_snapshot.state_generation,
5645                reviewed_tool: Arc::clone(&resolved.tool),
5646            })
5647        });
5648
5649        //
5650        // Re-resolve once after HITL because registry, scope, and policy may change while approval waits.
5651        // The final Arc must still match the implementation that the approver reviewed.
5652        //
5653        let control_snapshot = self.runtime_safety_snapshot();
5654        let resolved = self.tools.resolve(&request.requested_name);
5655        let registry_version = self.tools.version();
5656        let mut versions = ToolDecisionVersions {
5657            policy: control_snapshot.tool_security.policy_version(),
5658            registry: registry_version,
5659            runtime_control: control_snapshot.version,
5660            state: None,
5661        };
5662        metadata.insert(
5663            "runtime_scope_snapshot".to_string(),
5664            serde_json::to_value(&control_snapshot.tool_scope_override).unwrap_or(Value::Null),
5665        );
5666        let resolved = match resolved {
5667            Some(resolved) => resolved,
5668            None => {
5669                let reason = format!(
5670                    "Tool '{}' became unavailable after approval",
5671                    request.requested_name
5672                );
5673                let record = self.record_from_parts_at(
5674                    &request,
5675                    request.requested_name.clone(),
5676                    executed_arguments,
5677                    started_at,
5678                    start,
5679                    false,
5680                    false,
5681                    reason.clone(),
5682                    metadata,
5683                    ToolPolicyDecisionRecord::unavailable(reason),
5684                    approval_record,
5685                    false,
5686                    false,
5687                    versions,
5688                );
5689                self.finish_tool_record(&record).await;
5690                return Ok(record);
5691            }
5692        };
5693
5694        let canonical_id = resolved.identity.canonical_id.clone();
5695        if let Some(reason) =
5696            fallback_state.final_rejection_reason(&admitted_canonical_id, &canonical_id)
5697        {
5698            //
5699            // Preserve the admitted identity used by on_tool_start for completion, record, and history. The changed final resolution remains diagnostic metadata because it never reached invocation admission.
5700            //
5701            metadata.insert(
5702                "fallback_chain".to_string(),
5703                serde_json::to_value(&fallback_state.visited_canonical_ids).unwrap_or(Value::Null),
5704            );
5705            metadata.insert(
5706                "final_resolved_canonical_id".to_string(),
5707                Value::String(canonical_id),
5708            );
5709            let record = self.record_from_parts_at(
5710                &request,
5711                admitted_canonical_id,
5712                executed_arguments,
5713                started_at,
5714                start,
5715                false,
5716                false,
5717                format!("Denied: {reason}"),
5718                metadata,
5719                ToolPolicyDecisionRecord::deny(reason),
5720                approval_record,
5721                false,
5722                false,
5723                versions,
5724            );
5725            self.finish_tool_record(&record).await;
5726            return Ok(record);
5727        }
5728        let bindings = resolved.tool.policy_bindings();
5729        let final_arguments = control_snapshot
5730            .tool_security
5731            .prepare_tool_arguments_with_bindings(&canonical_id, &executed_arguments, &bindings);
5732        if let Some(record) = approval_record.as_mut()
5733            && matches!(record.status, ToolApprovalStatus::Modified)
5734        {
5735            record.modified_arguments = Some(final_arguments.clone());
5736        }
5737        let classification = resolved.tool.classify_call(&final_arguments);
5738        let safety = resolved.tool.safety_metadata();
5739        let security_engine = control_snapshot.tool_security;
5740        let tool_config = self.recovery_manager.get_tool_config(&canonical_id).clone();
5741        let recovery_timeout_ms = self.recovery_manager.get_tool_timeout(&canonical_id);
5742        metadata.insert(
5743            "classification".to_string(),
5744            serde_json::to_value(&classification).unwrap_or(Value::Null),
5745        );
5746        //
5747        // Approved arguments can produce a new invalid call classification after the request start hook has fired. Finalize non-execution evidence instead of returning a bare configuration error that would leave the lifecycle unmatched.
5748        //
5749        let (limits, timeout) = match Self::effective_tool_limits(
5750            &security_engine,
5751            &canonical_id,
5752            &safety,
5753            &classification,
5754            recovery_timeout_ms,
5755        ) {
5756            Ok(effective) => effective,
5757            Err(error) => {
5758                let reason = error.to_string();
5759                metadata.insert(
5760                    "configuration_error".to_string(),
5761                    Value::String(reason.clone()),
5762                );
5763                let record = self.record_from_parts_at(
5764                    &request,
5765                    canonical_id,
5766                    final_arguments,
5767                    started_at,
5768                    start,
5769                    false,
5770                    false,
5771                    format!("Denied: {reason}"),
5772                    metadata,
5773                    ToolPolicyDecisionRecord::deny(reason),
5774                    approval_record,
5775                    false,
5776                    false,
5777                    versions,
5778                );
5779                self.finish_tool_record(&record).await;
5780                return Ok(record);
5781            }
5782        };
5783        let policy_snapshot = security_engine.policy_snapshot(&canonical_id);
5784        let resource_lock_keys =
5785            tool_resource_lock_keys(&canonical_id, &final_arguments, &bindings, &classification);
5786        metadata.insert(
5787            "effective_limits".to_string(),
5788            serde_json::to_value(&limits).unwrap_or(Value::Null),
5789        );
5790        metadata.insert(
5791            "resource_lock_keys".to_string(),
5792            serde_json::to_value(&resource_lock_keys).unwrap_or(Value::Null),
5793        );
5794        if policy_snapshot.is_null() {
5795            metadata.remove("policy_snapshot");
5796        } else {
5797            metadata.insert("policy_snapshot".to_string(), policy_snapshot.clone());
5798        }
5799
5800        let final_denial = |canonical_id: String,
5801                            output: String,
5802                            policy: ToolPolicyDecisionRecord,
5803                            metadata: HashMap<String, Value>,
5804                            decision_versions: ToolDecisionVersions| {
5805            self.record_from_parts_at(
5806                &request,
5807                canonical_id,
5808                final_arguments.clone(),
5809                started_at,
5810                start,
5811                false,
5812                false,
5813                output,
5814                metadata,
5815                policy,
5816                approval_record.clone(),
5817                false,
5818                false,
5819                decision_versions,
5820            )
5821        };
5822
5823        if control_snapshot.emergency_deny {
5824            let reason = "Tool execution is disabled by runtime control".to_string();
5825            let record = final_denial(
5826                canonical_id,
5827                reason.clone(),
5828                ToolPolicyDecisionRecord::deny(reason),
5829                metadata,
5830                versions,
5831            );
5832            self.finish_tool_record(&record).await;
5833            return Ok(record);
5834        }
5835
5836        //
5837        // Final scope evaluation uses the same post-HITL runtime snapshot and captures state authority before locks.
5838        // Admission must reject any later state transition, reset, or restore before consuming rate capacity or invoking the tool.
5839        //
5840        let available_snapshot = self
5841            .get_available_tool_ids_snapshot_for_scope(
5842                control_snapshot.tool_scope_override.as_deref(),
5843            )
5844            .await?;
5845        versions.state = available_snapshot.state_generation;
5846        metadata.insert(
5847            "available_tool_ids_snapshot".to_string(),
5848            serde_json::to_value(&available_snapshot.tool_ids).unwrap_or(Value::Null),
5849        );
5850        metadata.insert(
5851            "state_generation_snapshot".to_string(),
5852            serde_json::to_value(available_snapshot.state_generation).unwrap_or(Value::Null),
5853        );
5854        if !available_snapshot
5855            .tool_ids
5856            .iter()
5857            .any(|tool_id| tool_id == &canonical_id)
5858        {
5859            let reason = format!(
5860                "Tool '{}' is not available in the final runtime scope",
5861                canonical_id
5862            );
5863            let record = final_denial(
5864                canonical_id,
5865                reason.clone(),
5866                ToolPolicyDecisionRecord::deny(reason),
5867                metadata,
5868                versions,
5869            );
5870            self.finish_tool_record(&record).await;
5871            return Ok(record);
5872        }
5873
5874        //
5875        // Validation may run more than once, but it must not consume rate capacity.
5876        // Rate capacity is consumed once after resource locks are held.
5877        //
5878        let final_security_result = security_engine
5879            .validate_tool_execution_with_bindings(&canonical_id, &final_arguments, &bindings)
5880            .await?;
5881        match &final_security_result {
5882            SecurityCheckResult::Block { reason } => {
5883                let record = final_denial(
5884                    canonical_id,
5885                    format!("Denied: {}", reason),
5886                    ToolPolicyDecisionRecord::deny(reason.clone()),
5887                    metadata,
5888                    versions,
5889                );
5890                self.finish_tool_record(&record).await;
5891                return Ok(record);
5892            }
5893            SecurityCheckResult::Unavailable { reason } => {
5894                let record = final_denial(
5895                    canonical_id,
5896                    format!("Unavailable: {}", reason),
5897                    ToolPolicyDecisionRecord::unavailable(reason.clone()),
5898                    metadata,
5899                    versions,
5900                );
5901                self.finish_tool_record(&record).await;
5902                return Ok(record);
5903            }
5904            SecurityCheckResult::Warn { message } => {
5905                warn!(tool = %canonical_id, message = %message, "Tool security warning after approval");
5906            }
5907            SecurityCheckResult::Allow | SecurityCheckResult::RequireConfirmation { .. } => {}
5908        }
5909        let final_confirmation_required = matches!(
5910            final_security_result,
5911            SecurityCheckResult::RequireConfirmation { .. }
5912        ) || security_engine
5913            .classification_approval_message(&canonical_id, &classification)
5914            .is_some();
5915        let stale_approval = approval_binding.as_ref().is_some_and(|binding| {
5916            binding.is_stale(
5917                &canonical_id,
5918                &final_arguments,
5919                final_confirmation_required,
5920                versions,
5921                &resolved.tool,
5922            )
5923        });
5924        if stale_approval {
5925            let reason = "Approval became stale before final admission".to_string();
5926            let record = final_denial(
5927                canonical_id,
5928                reason.clone(),
5929                ToolPolicyDecisionRecord::deny(reason),
5930                metadata,
5931                versions,
5932            );
5933            self.finish_tool_record(&record).await;
5934            return Ok(record);
5935        }
5936        if final_confirmation_required && approval_binding.is_none() {
5937            let reason = "Final policy requires fresh approval".to_string();
5938            let record = final_denial(
5939                canonical_id,
5940                reason.clone(),
5941                ToolPolicyDecisionRecord::approval(reason),
5942                metadata,
5943                versions,
5944            );
5945            self.finish_tool_record(&record).await;
5946            return Ok(record);
5947        }
5948
5949        if let Some((_, reason)) = self.host_tool_unavailability(&canonical_id) {
5950            let record = final_denial(
5951                canonical_id,
5952                reason.to_string(),
5953                ToolPolicyDecisionRecord::unavailable(reason),
5954                metadata,
5955                versions,
5956            );
5957            self.finish_tool_record(&record).await;
5958            return Ok(record);
5959        }
5960
5961        //
5962        // Hold conflict locks across final generation admission and invocation.
5963        // Completion hooks and fallback execution run only after these guards are released.
5964        //
5965        let Some(resource_guards) = self.acquire_tool_resource_locks(&resource_lock_keys).await
5966        else {
5967            //
5968            // Cancellation while lock admission waits is terminal runtime-control evidence, even though the non-executed result retains its denial policy outcome.
5969            //
5970            let reason = "Tool execution cancelled while waiting for resource locks".to_string();
5971            let mut record = final_denial(
5972                canonical_id,
5973                reason.clone(),
5974                ToolPolicyDecisionRecord::deny(reason),
5975                metadata,
5976                versions,
5977            );
5978            record.cancelled = true;
5979            record.cancellation_reason = Some("runtime control cancellation".to_string());
5980            self.finish_tool_record(&record).await;
5981            return Ok(record);
5982        };
5983
5984        //
5985        // The lock-held admission is the last authority boundary. Runtime, policy, and state generations must still match before the one atomic rate admission.
5986        // Updates after admission apply to later calls, while emergency cancellation remains live for this call.
5987        //
5988        let admission = self.admit_tool_execution(
5989            versions.runtime_control,
5990            versions.policy,
5991            versions.state,
5992            &canonical_id,
5993        );
5994        if !matches!(admission, SecurityCheckResult::Allow) {
5995            let latest_control = self.runtime_safety_snapshot();
5996            let reason = admission
5997                .reason()
5998                .unwrap_or("tool admission was denied")
5999                .to_string();
6000            let policy = if admission.is_unavailable() {
6001                ToolPolicyDecisionRecord::unavailable(reason.clone())
6002            } else {
6003                ToolPolicyDecisionRecord::deny(reason.clone())
6004            };
6005            let record = self.record_from_parts_at(
6006                &request,
6007                canonical_id,
6008                final_arguments,
6009                started_at,
6010                start,
6011                false,
6012                false,
6013                reason,
6014                metadata,
6015                policy,
6016                approval_record,
6017                false,
6018                false,
6019                ToolDecisionVersions {
6020                    policy: latest_control.tool_security.policy_version(),
6021                    registry: versions.registry,
6022                    runtime_control: latest_control.version,
6023                    state: self
6024                        .state_machine
6025                        .as_ref()
6026                        .map(|state_machine| state_machine.generation()),
6027                },
6028            );
6029            self.finish_tool_record_after_resource_guards(resource_guards, &record)
6030                .await;
6031            return Ok(record);
6032        }
6033        let executed_arguments = final_arguments;
6034
6035        let turn_actor = current_turn_actor_context();
6036        let actor = ToolActorContext {
6037            actor_id: turn_actor
6038                .as_ref()
6039                .and_then(|context| context.effective_actor_id().map(str::to_string))
6040                .or_else(|| self.actor_id()),
6041            origin_actor_id: turn_actor
6042                .as_ref()
6043                .and_then(|context| context.origin_actor_id.clone()),
6044            sender_agent_id: turn_actor
6045                .as_ref()
6046                .and_then(|context| context.sender_agent_id.clone()),
6047        };
6048        let tool_context = ToolExecutionContext {
6049            requested_name: request.requested_name.clone(),
6050            canonical_id: canonical_id.clone(),
6051            display_name: resolved.identity.display_name.clone(),
6052            provider_id: resolved.identity.provider_id.clone(),
6053            registry_version: versions.registry,
6054            policy_version: versions.policy,
6055            runtime_control_version: versions.runtime_control,
6056            call_id: request.call_id.clone(),
6057            source: request.source.clone(),
6058            actor,
6059            cancellation: ToolCancellationToken::new(
6060                Arc::clone(&self.runtime_control.emergency_deny),
6061                Some("runtime control cancellation".to_string()),
6062            ),
6063            started_at,
6064            deadline: None,
6065            permission: ToolPolicyDecisionRecord::allow(),
6066            approval: approval_record.clone(),
6067            classification: classification.clone(),
6068            safety,
6069            limits: limits.clone(),
6070            policy_snapshot,
6071            custom_config: security_engine.custom_config(&canonical_id),
6072        };
6073        let (mut result, timed_out, cancelled, invoked) = self
6074            .run_tool_with_retries(
6075                &canonical_id,
6076                resolved.tool.clone(),
6077                executed_arguments.clone(),
6078                tool_context,
6079                timeout,
6080                tool_config.max_retries,
6081            )
6082            .await?;
6083
6084        //
6085        // Runtime cancellation is terminal for the logical request. Recovery fallback must not mask the cancelled record or start new work after the host asked active execution to stop.
6086        //
6087        let fallback_tool = if !result.success && !cancelled {
6088            match &tool_config.on_failure {
6089                ToolFailureAction::Skip => {
6090                    result = ToolResult::ok(format!(
6091                        "{{\"skipped\": true, \"reason\": \"Tool '{}' was skipped after failure\"}}",
6092                        canonical_id
6093                    ));
6094                    None
6095                }
6096                ToolFailureAction::Fallback { fallback_tool } => Some(fallback_tool.clone()),
6097                ToolFailureAction::ReportError => None,
6098            }
6099        } else {
6100            None
6101        };
6102
6103        let output_cap = limits.max_output_chars;
6104        let (output, output_truncated) =
6105            Self::truncate_tool_output(result.output.clone(), output_cap);
6106        if let Some(result_metadata) = result.metadata {
6107            metadata.extend(result_metadata);
6108        }
6109        let mut record = self.record_from_parts_at(
6110            &request,
6111            canonical_id,
6112            executed_arguments,
6113            started_at,
6114            start,
6115            invoked,
6116            result.success,
6117            output,
6118            metadata,
6119            ToolPolicyDecisionRecord::allow(),
6120            approval_record,
6121            timed_out,
6122            output_truncated,
6123            versions,
6124        );
6125        record.cancelled = cancelled;
6126        if cancelled {
6127            record.cancellation_reason = Some("runtime control cancellation".to_string());
6128        }
6129        if let Some(fallback_tool) = fallback_tool {
6130            let fallback_arguments = record.executed_arguments.clone();
6131            let original_tool = record.canonical_id.clone();
6132            //
6133            // Finish the failed logical request after releasing its resource guards so its start hook is matched and fallback hooks cannot overtake the original record.
6134            //
6135            self.finish_tool_record_after_resource_guards(resource_guards, &record)
6136                .await;
6137            let fallback_request = ToolExecutionRequest::new(
6138                request.call_id.clone(),
6139                fallback_tool,
6140                fallback_arguments,
6141                ToolCallSource::Fallback { original_tool },
6142            );
6143            return Box::pin(self.execute_tool_record_inner(fallback_request, fallback_state))
6144                .await;
6145        }
6146        self.finish_tool_record_after_resource_guards(resource_guards, &record)
6147            .await;
6148        Ok(record)
6149    }
6150
6151    #[instrument(skip(self, tool_call), fields(tool = %tool_call.name))]
6152    async fn execute_tool_smart(&self, tool_call: &ToolCall) -> Result<String> {
6153        let record = self
6154            .execute_tool_record(ToolExecutionRequest::new(
6155                tool_call.id.clone(),
6156                tool_call.name.clone(),
6157                tool_call.arguments.clone(),
6158                ToolCallSource::Model,
6159            ))
6160            .await?;
6161        if record.success {
6162            Ok(record.model_output_string())
6163        } else if matches!(record.policy.outcome, PermissionOutcome::RequiresApproval) {
6164            Err(AgentError::HITLRejected(record.model_output_string()))
6165        } else {
6166            Err(AgentError::Tool(record.model_output_string()))
6167        }
6168    }
6169
6170    //
6171    // This method is branch-safe because it only asks the router which skill matches.
6172    // Do not add disambiguation, pending-skill writes, or skill execution here.
6173    //
6174    /// Selects a skill without executing it.
6175    async fn select_skill_candidate(&self, input: &str) -> Result<Option<SkillCandidate>> {
6176        let Some(ref router) = self.skill_router else {
6177            return Ok(None);
6178        };
6179        let available_skills = self.get_available_skills();
6180        if available_skills.is_empty() {
6181            return Ok(None);
6182        }
6183        let skill_ids: Vec<&str> = available_skills.iter().map(|s| s.id.as_str()).collect();
6184        let Some(skill_id) = self
6185            .observe_purpose(
6186                ObservationPurpose::SkillRouting,
6187                router.select_skill_filtered(input, &skill_ids),
6188            )
6189            .await?
6190        else {
6191            return Ok(None);
6192        };
6193        let skill = router
6194            .get_skill(&skill_id)
6195            .cloned()
6196            .ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
6197        info!(skill_id = %skill_id, "Skill selected");
6198        Ok(Some(SkillCandidate::new(skill_id, skill)))
6199    }
6200
6201    //
6202    // This is the commit half of skill routing.
6203    // It may mutate pending skill state, run clarification, and execute skill steps.
6204    //
6205    async fn commit_skill_candidate_route_result(
6206        &self,
6207        candidate: SkillCandidate,
6208        input: &str,
6209    ) -> Result<SkillRouteResult> {
6210        let skill_id = candidate.skill_id;
6211        let skill = candidate.skill;
6212        let expected_state_generation = self
6213            .state_machine
6214            .as_ref()
6215            .map(|state_machine| state_machine.generation());
6216        let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
6217        if let Some(ref skill_disambig) = skill.disambiguation
6218            && skill_disambig.enabled.unwrap_or(false)
6219            && let Some(ref disambiguator) = self.disambiguation_manager
6220        {
6221            let context = self.build_disambiguation_context().await?;
6222            let state_override = self
6223                .state_machine
6224                .as_ref()
6225                .and_then(|sm| sm.current_definition())
6226                .and_then(|def| def.disambiguation.clone());
6227
6228            let disambiguation_result = self
6229                .observe_purpose(
6230                    ObservationPurpose::DisambiguationDetection,
6231                    disambiguator.process_input_with_override(
6232                        input,
6233                        &context,
6234                        state_override.as_ref(),
6235                        Some(skill_disambig),
6236                    ),
6237                )
6238                .await?;
6239            let current_state_generation = self
6240                .state_machine
6241                .as_ref()
6242                .map(|state_machine| state_machine.generation());
6243            if current_state_generation != expected_state_generation
6244                || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6245            {
6246                disambiguator.clear_pending().await;
6247                *self.pending_skill_id.write() = None;
6248                return Err(AgentError::Other(
6249                    "State or reset ownership changed during skill disambiguation".to_string(),
6250                ));
6251            }
6252            match disambiguation_result {
6253                DisambiguationResult::Clear => {
6254                    debug!(skill_id = %skill_id, "Skill disambiguation: clear");
6255                }
6256                DisambiguationResult::NeedsClarification {
6257                    question,
6258                    detection,
6259                } => {
6260                    let admission = self
6261                        .admit_disambiguation_redispatch(
6262                            expected_disambiguation_epoch,
6263                            expected_state_generation,
6264                        )
6265                        .await?;
6266                    let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
6267                    info!(
6268                        skill_id = %skill_id,
6269                        ambiguity_type = ?detection.ambiguity_type,
6270                        confidence = detection.confidence,
6271                        "Skill requires clarification before execution"
6272                    );
6273                    *self.pending_skill_id.write() = Some(skill_id.clone());
6274                    let response = AgentResponse::new(&question.question).with_metadata(
6275                        "disambiguation",
6276                        serde_json::json!({
6277                            "status": if awaiting_confirmation { "awaiting_confirmation" } else { "awaiting_clarification" },
6278                            "skill_id": skill_id,
6279                            "options": question.options,
6280                            "clarifying": question.clarifying,
6281                            "detection": {
6282                                "type": detection.ambiguity_type,
6283                                "confidence": detection.confidence,
6284                                "what_is_unclear": detection.what_is_unclear,
6285                            }
6286                        }),
6287                    );
6288                    drop(admission);
6289                    return Ok(SkillRouteResult::NeedsClarification {
6290                        response,
6291                        ownership: Some(DisambiguationOwnership {
6292                            epoch: expected_disambiguation_epoch,
6293                            state_generation: expected_state_generation,
6294                        }),
6295                    });
6296                }
6297                DisambiguationResult::Clarified { enriched_input, .. } => {
6298                    info!(skill_id = %skill_id, enriched = %enriched_input, "Skill disambiguation clarified");
6299                    let admission = self
6300                        .admit_disambiguation_redispatch(
6301                            expected_disambiguation_epoch,
6302                            expected_state_generation,
6303                        )
6304                        .await?;
6305                    drop(admission);
6306                    let content = self.execute_skill(&skill, &enriched_input).await?;
6307                    return Ok(SkillRouteResult::Response { skill_id, content });
6308                }
6309                DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
6310                    info!(skill_id = %skill_id, "Skill disambiguation best guess");
6311                    let admission = self
6312                        .admit_disambiguation_redispatch(
6313                            expected_disambiguation_epoch,
6314                            expected_state_generation,
6315                        )
6316                        .await?;
6317                    drop(admission);
6318                    let content = self.execute_skill(&skill, &enriched_input).await?;
6319                    return Ok(SkillRouteResult::Response { skill_id, content });
6320                }
6321                DisambiguationResult::GiveUp { reason } => {
6322                    warn!(skill_id = %skill_id, reason = %reason, "Skill disambiguation gave up");
6323                    let apology = self
6324                                .generate_localized_apology(
6325                                    "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
6326                                    &reason,
6327                                )
6328                                .await
6329                                .unwrap_or_else(|_| {
6330                                    format!("I'm sorry, I couldn't understand your request: {}", reason)
6331                                });
6332                    return Ok(SkillRouteResult::NeedsClarification {
6333                        response: AgentResponse::new(&apology),
6334                        ownership: None,
6335                    });
6336                }
6337                DisambiguationResult::Escalate { reason } => {
6338                    info!(skill_id = %skill_id, reason = %reason, "Skill disambiguation escalating");
6339                    let apology = self
6340                                .generate_localized_apology(
6341                                    "Explain briefly that you're transferring the user to a human agent for help.",
6342                                    &reason,
6343                                )
6344                                .await
6345                                .unwrap_or_else(|_| {
6346                                    format!("I need human assistance to help with your request: {}", reason)
6347                                });
6348                    return Ok(SkillRouteResult::NeedsClarification {
6349                        response: AgentResponse::new(&apology),
6350                        ownership: None,
6351                    });
6352                }
6353                DisambiguationResult::Abandoned { .. } => {
6354                    debug!(skill_id = %skill_id, "Skill disambiguation abandoned");
6355                    return Ok(SkillRouteResult::NoMatch);
6356                }
6357            }
6358        }
6359        let admission = self
6360            .admit_disambiguation_redispatch(
6361                expected_disambiguation_epoch,
6362                expected_state_generation,
6363            )
6364            .await?;
6365        drop(admission);
6366        let content = self.execute_skill(&skill, input).await?;
6367        Ok(SkillRouteResult::Response { skill_id, content })
6368    }
6369
6370    /// Result of skill routing.
6371    async fn try_skill_route(&self, input: &str) -> Result<SkillRouteResult> {
6372        if let Some(candidate) = self.select_skill_candidate(input).await? {
6373            self.commit_skill_candidate_route_result(candidate, input)
6374                .await
6375        } else {
6376            Ok(SkillRouteResult::NoMatch)
6377        }
6378    }
6379
6380    /// Execute a skill with reasoning and reflection, returning the response string.
6381    async fn execute_skill(&self, skill: &SkillDefinition, input: &str) -> Result<String> {
6382        if let Some(ref executor) = self.skill_executor {
6383            let skill_reasoning = self.get_skill_reasoning_config(skill);
6384            let skill_reflection = self.get_skill_reflection_config(skill);
6385
6386            debug!(
6387                skill_id = %skill.id,
6388                reasoning_mode = ?skill_reasoning.mode,
6389                reflection_enabled = ?skill_reflection.enabled,
6390                "Skill reasoning/reflection config"
6391            );
6392
6393            let response = self
6394                .observe_purpose(
6395                    ObservationPurpose::SkillPrompt,
6396                    executor.execute_with_invoker(skill, input, serde_json::json!({}), self),
6397                )
6398                .await?;
6399
6400            if skill_reflection.requires_evaluation() && skill_reflection.is_enabled() {
6401                let should_reflect = self
6402                    .should_reflect_with_config(input, &response, &skill_reflection)
6403                    .await?;
6404                if should_reflect {
6405                    let evaluated = self
6406                        .evaluate_and_retry_with_config(input, response, &skill_reflection)
6407                        .await?;
6408                    return Ok(evaluated);
6409                }
6410            }
6411
6412            return Ok(response);
6413        }
6414        Err(AgentError::Skill(
6415            "No skill executor configured".to_string(),
6416        ))
6417    }
6418
6419    /// Execute a skill by ID, bypassing the skill router.
6420    /// Used after skill-triggered disambiguation resolves to route directly to the matched skill.
6421    async fn execute_skill_by_id(&self, skill_id: &str, input: &str) -> Result<String> {
6422        let skill = self
6423            .skill_router
6424            .as_ref()
6425            .and_then(|r| r.get_skill(skill_id).cloned())
6426            .ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
6427        self.execute_skill(&skill, input).await
6428    }
6429
6430    async fn should_reflect_with_config(
6431        &self,
6432        input: &str,
6433        response: &str,
6434        config: &ReflectionConfig,
6435    ) -> Result<bool> {
6436        if !config.requires_evaluation() {
6437            return Ok(false);
6438        }
6439
6440        if config.is_enabled() {
6441            return Ok(true);
6442        }
6443
6444        let evaluator_llm = config
6445            .evaluator_llm
6446            .as_ref()
6447            .and_then(|alias| self.llm_registry.get(alias).ok())
6448            .or_else(|| self.llm_registry.router().ok())
6449            .or_else(|| self.llm_registry.default().ok());
6450
6451        let Some(llm) = evaluator_llm else {
6452            return Ok(false);
6453        };
6454
6455        let response_preview: String = response.chars().take(500).collect();
6456        let prompt = format!(
6457            r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
6458
6459User query: "{}"
6460Response: "{}"
6461
6462Answer YES or NO only."#,
6463            input, response_preview
6464        );
6465
6466        let messages = vec![ChatMessage::user(&prompt)];
6467        let result = self
6468            .observe_purpose(
6469                ObservationPurpose::ReflectionDecision,
6470                llm.complete(&messages, None),
6471            )
6472            .await;
6473
6474        match result {
6475            Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
6476            Err(_) => Ok(false),
6477        }
6478    }
6479
6480    async fn evaluate_and_retry_with_config(
6481        &self,
6482        input: &str,
6483        mut response: String,
6484        config: &ReflectionConfig,
6485    ) -> Result<String> {
6486        let llm = self.get_state_llm()?;
6487        let mut attempts = 0u32;
6488        let max_retries = config.max_retries;
6489
6490        loop {
6491            let evaluation = self
6492                .evaluate_response_with_config(input, &response, config)
6493                .await?;
6494
6495            if evaluation.passed || attempts >= max_retries {
6496                info!(
6497                    passed = evaluation.passed,
6498                    confidence = evaluation.confidence,
6499                    attempts = attempts + 1,
6500                    "Skill reflection evaluation complete"
6501                );
6502                return Ok(response);
6503            }
6504
6505            debug!(
6506                attempt = attempts + 1,
6507                failed_criteria = evaluation.failed_criteria().count(),
6508                "Skill response did not meet criteria, retrying"
6509            );
6510
6511            let feedback: Vec<String> = evaluation
6512                .failed_criteria()
6513                .map(|c| format!("- {}", c.criterion))
6514                .collect();
6515
6516            let retry_prompt = format!(
6517                "Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response to: {}",
6518                feedback.join("\n"),
6519                input
6520            );
6521
6522            let messages = vec![ChatMessage::user(&retry_prompt)];
6523            let retry_response = self
6524                .observe_purpose(
6525                    ObservationPurpose::ReflectionEvaluation,
6526                    llm.complete(&messages, None),
6527                )
6528                .await
6529                .map_err(|e| AgentError::LLM(e.to_string()))?;
6530
6531            response = retry_response.content.trim().to_string();
6532            attempts += 1;
6533        }
6534    }
6535
6536    async fn evaluate_response_with_config(
6537        &self,
6538        input: &str,
6539        response: &str,
6540        config: &ReflectionConfig,
6541    ) -> Result<EvaluationResult> {
6542        let evaluator_llm = config
6543            .evaluator_llm
6544            .as_ref()
6545            .and_then(|alias| self.llm_registry.get(alias).ok())
6546            .or_else(|| self.llm_registry.router().ok())
6547            .or_else(|| self.llm_registry.default().ok())
6548            .ok_or_else(|| AgentError::Config("No LLM available for evaluation".into()))?;
6549
6550        let criteria = &config.criteria;
6551        let criteria_list = criteria
6552            .iter()
6553            .enumerate()
6554            .map(|(i, c)| format!("{}. {}", i + 1, c))
6555            .collect::<Vec<_>>()
6556            .join("\n");
6557
6558        let prompt = format!(
6559            r#"Evaluate this response against the criteria.
6560
6561User query: "{}"
6562
6563Response to evaluate: "{}"
6564
6565Criteria:
6566{}
6567
6568For each criterion, respond with:
6569- criterion number
6570- PASS or FAIL
6571- brief reason
6572
6573Then provide overall confidence (0.0 to 1.0) and whether it passes overall.
6574
6575Format:
65761. PASS/FAIL - reason
65772. PASS/FAIL - reason
6578...
6579CONFIDENCE: 0.X
6580OVERALL: PASS/FAIL"#,
6581            input, response, criteria_list
6582        );
6583
6584        let messages = vec![ChatMessage::user(&prompt)];
6585        let eval_response = self
6586            .observe_purpose(
6587                ObservationPurpose::ReflectionEvaluation,
6588                evaluator_llm.complete(&messages, None),
6589            )
6590            .await
6591            .map_err(|e| AgentError::LLM(format!("Evaluation failed: {}", e)))?;
6592
6593        let content = eval_response.content.to_uppercase();
6594        let llm_pass = content.contains("OVERALL: PASS");
6595
6596        let confidence = content
6597            .lines()
6598            .find(|l| l.contains("CONFIDENCE:"))
6599            .and_then(|l| {
6600                l.split(':')
6601                    .nth(1)
6602                    .and_then(|v| v.trim().parse::<f32>().ok())
6603            })
6604            .unwrap_or(if llm_pass { 0.8 } else { 0.4 });
6605
6606        // Gate pass against confidence threshold.
6607        // LLM may say PASS but with low confidence - the threshold catches this.
6608        let overall_pass = llm_pass && confidence >= config.pass_threshold;
6609
6610        let mut criteria_results = Vec::new();
6611        for (i, criterion) in criteria.iter().enumerate() {
6612            let line_marker = format!("{}.", i + 1);
6613            let passed = eval_response
6614                .content
6615                .lines()
6616                .find(|l| l.contains(&line_marker))
6617                .map(|l| l.to_uppercase().contains("PASS"))
6618                .unwrap_or(overall_pass);
6619
6620            if passed {
6621                criteria_results.push(CriterionResult::pass(criterion));
6622            } else {
6623                criteria_results.push(CriterionResult::fail(criterion, "Did not meet criterion"));
6624            }
6625        }
6626
6627        Ok(EvaluationResult::new(overall_pass, confidence).with_criteria(criteria_results))
6628    }
6629
6630    /// Process input through the pipeline (state-level override or agent-level).
6631    async fn process_input(&self, input: &str) -> Result<ProcessData> {
6632        if let Some(processor) = self.get_state_process_processor() {
6633            let purpose = observation_purpose_for_process(processor.input_purpose_hint());
6634            return self
6635                .observe_purpose(purpose, processor.process_input(input))
6636                .await;
6637        }
6638        if let Some(ref processor) = self.process_processor {
6639            let purpose = observation_purpose_for_process(processor.input_purpose_hint());
6640            self.observe_purpose(purpose, processor.process_input(input))
6641                .await
6642        } else {
6643            Ok(ProcessData::new(input))
6644        }
6645    }
6646
6647    /// Process output through the pipeline (state-level override or agent-level).
6648    async fn process_output(
6649        &self,
6650        output: &str,
6651        input_context: &std::collections::HashMap<String, serde_json::Value>,
6652    ) -> Result<ProcessData> {
6653        if let Some(processor) = self.get_state_process_processor() {
6654            let purpose = observation_purpose_for_process(processor.output_purpose_hint());
6655            return self
6656                .observe_purpose(purpose, processor.process_output(output, input_context))
6657                .await;
6658        }
6659        if let Some(ref processor) = self.process_processor {
6660            let purpose = observation_purpose_for_process(processor.output_purpose_hint());
6661            self.observe_purpose(purpose, processor.process_output(output, input_context))
6662                .await
6663        } else {
6664            Ok(ProcessData::new(output))
6665        }
6666    }
6667
6668    /// Build a ProcessProcessor from the current state's process config, if any.
6669    fn get_state_process_processor(&self) -> Option<ProcessProcessor> {
6670        let sm = self.state_machine.as_ref()?;
6671        let def = sm.current_definition()?;
6672        let config = def.process.as_ref()?;
6673        let mut processor = ProcessProcessor::new(config.clone());
6674        if let Some(ref registry) = Some(self.llm_registry.clone()) {
6675            processor = processor.with_llm_registry(registry.clone());
6676        }
6677        processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
6678        Some(processor)
6679    }
6680
6681    // Timeout transitions reserve exit actions before committing ownership changes under the short write lock.
6682    async fn check_turn_timeout(&self) -> Result<()> {
6683        let Some(ref sm) = self.state_machine else {
6684            return Ok(());
6685        };
6686        let Some(timeout_state) = sm.check_timeout() else {
6687            return Ok(());
6688        };
6689        let claim_admission = self.disambiguation_admission.write().await;
6690        if sm.check_timeout().as_deref() != Some(timeout_state.as_str()) {
6691            return Ok(());
6692        }
6693        let Some(reservation) = self.reserve_state_transition() else {
6694            return Ok(());
6695        };
6696        let from_state = sm.current();
6697        let expected_state_generation = sm.generation();
6698        let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
6699        let history_before = sm.history();
6700        drop(claim_admission);
6701
6702        self.execute_state_exit_actions(&from_state).await;
6703
6704        let admission = self.disambiguation_admission.write().await;
6705        if sm.current() != from_state
6706            || sm.generation() != expected_state_generation
6707            || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6708            || sm.check_timeout().as_deref() != Some(timeout_state.as_str())
6709        {
6710            return Ok(());
6711        }
6712        sm.transition_to(&timeout_state, "max_turns exceeded")?;
6713        self.invalidate_pending_confirmation("state_timeout").await;
6714        let entered = sm.current();
6715        let is_reentry = Self::state_was_previously_entered(&entered, &from_state, &history_before);
6716        drop(admission);
6717
6718        self.execute_state_enter_actions(&entered, is_reentry).await;
6719        drop(reservation);
6720        info!(to = %entered, "Timeout transition");
6721        Ok(())
6722    }
6723
6724    fn increment_turn(&self) {
6725        if let Some(ref sm) = self.state_machine {
6726            sm.increment_turn();
6727        }
6728    }
6729
6730    fn transitions_available_for_commit(&self) -> Option<(Vec<Transition>, String)> {
6731        let sm = self.state_machine.as_ref()?;
6732        let current = sm.current();
6733        let transitions: Vec<_> = sm
6734            .auto_transitions()
6735            .into_iter()
6736            .filter(|t| match t.cooldown_turns {
6737                Some(cd) if cd > 0 => {
6738                    let resolved = sm.config().resolve_full_path(&current, &t.to);
6739                    !sm.is_on_cooldown(&resolved, cd)
6740                }
6741                _ => true,
6742            })
6743            .collect();
6744        Some((transitions, current))
6745    }
6746
6747    fn transition_reason(transition: &Transition) -> String {
6748        if transition.when.is_empty() {
6749            "guard condition met".to_string()
6750        } else {
6751            transition.when.clone()
6752        }
6753    }
6754
6755    /// Builds transition context with optional staged writes overlaid.
6756    fn build_transition_context(
6757        &self,
6758        user_message: &str,
6759        response: &str,
6760        current_state: &str,
6761        staged: Option<&HashMap<String, Value>>,
6762    ) -> TransitionContext {
6763        let context_map = staged
6764            .map(|writes| self.build_context_with_staged(writes))
6765            .unwrap_or_else(|| self.build_context_with_overlays());
6766        TransitionContext::new(user_message, response, current_state).with_context(context_map)
6767    }
6768
6769    /// Selects a post-response transition without committing side effects.
6770    async fn select_transition_candidate(
6771        &self,
6772        user_message: &str,
6773        response: &str,
6774    ) -> Result<Option<TransitionCandidate>> {
6775        let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6776            return Ok(None);
6777        };
6778        let transitions: Vec<Transition> = transitions
6779            .into_iter()
6780            .filter(|transition| matches!(transition.timing, TransitionTiming::PostResponse))
6781            .collect();
6782        if transitions.is_empty() {
6783            return Ok(None);
6784        }
6785        let Some(evaluator) = self.transition_evaluator.as_ref() else {
6786            return Ok(None);
6787        };
6788        let context = self.build_transition_context(user_message, response, &current_state, None);
6789        let selected = self
6790            .observe_purpose(
6791                ObservationPurpose::StateTransitionEvaluation,
6792                evaluator.select_transition(&transitions, &context),
6793            )
6794            .await?;
6795        Ok(selected.map(|index| {
6796            let transition = transitions[index].clone();
6797            TransitionCandidate::new(
6798                current_state,
6799                transition.clone(),
6800                Self::transition_reason(&transition),
6801            )
6802        }))
6803    }
6804
6805    /// Selects a guard or resolved-intent transition without an LLM call.
6806    fn select_deterministic_transition_candidate(
6807        &self,
6808        user_message: &str,
6809        current_state: &str,
6810        transitions: &[Transition],
6811        staged: &HashMap<String, Value>,
6812    ) -> Option<TransitionCandidate> {
6813        let context = self.build_transition_context(user_message, "", current_state, Some(staged));
6814
6815        for transition in transitions {
6816            if let Some(guard) = transition.guard.as_ref()
6817                && evaluate_guard(guard, &context)
6818            {
6819                return Some(TransitionCandidate::new(
6820                    current_state,
6821                    transition.clone(),
6822                    Self::transition_reason(transition),
6823                ));
6824            }
6825        }
6826
6827        let resolved_intent = context
6828            .context
6829            .get("resolved_intent")
6830            .and_then(Value::as_str)
6831            .filter(|value| !value.is_empty());
6832        if let Some(resolved_intent) = resolved_intent {
6833            for transition in transitions {
6834                if transition.intent.as_deref() == Some(resolved_intent) {
6835                    return Some(TransitionCandidate::new(
6836                        current_state,
6837                        transition.clone(),
6838                        Self::transition_reason(transition),
6839                    ));
6840                }
6841            }
6842        }
6843
6844        None
6845    }
6846
6847    /// Commits a selected transition through the shared post-response path.
6848    async fn commit_transition_candidate(&self, candidate: &TransitionCandidate) -> Result<bool> {
6849        self.commit_transition_target(&candidate.from_state, candidate.target(), &candidate.reason)
6850            .await
6851    }
6852
6853    /// Runs state transition approval before any transition side effects.
6854    async fn approve_transition_target(&self, from_state: &str, target: &str) -> Result<bool> {
6855        let approved = self.check_state_hitl(Some(from_state), target).await?;
6856        if !approved {
6857            info!(to = %target, "State transition rejected by HITL");
6858        }
6859        Ok(approved)
6860    }
6861
6862    /// Applies an approved transition after reserving exit actions and keeping async hooks outside the commit lock.
6863    async fn apply_transition_target(
6864        &self,
6865        from_state: &str,
6866        target: &str,
6867        reason: &str,
6868        staged: Option<&HashMap<String, Value>>,
6869    ) -> Result<bool> {
6870        let Some(ref sm) = self.state_machine else {
6871            return Ok(false);
6872        };
6873        let claim_admission = self.disambiguation_admission.write().await;
6874        if sm.current() != from_state {
6875            return Ok(false);
6876        }
6877        let Some(reservation) = self.reserve_state_transition() else {
6878            return Ok(false);
6879        };
6880        let expected_state_generation = sm.generation();
6881        let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
6882        let history_before = sm.history();
6883        drop(claim_admission);
6884
6885        self.execute_state_exit_actions(from_state).await;
6886
6887        let admission = self.disambiguation_admission.write().await;
6888        if sm.current() != from_state
6889            || sm.generation() != expected_state_generation
6890            || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6891        {
6892            return Ok(false);
6893        }
6894        sm.transition_to(target, reason)?;
6895        self.invalidate_pending_confirmation("state_transition")
6896            .await;
6897        sm.reset_no_transition();
6898        if let Some(staged) = staged {
6899            self.commit_staged_context_writes(staged);
6900        }
6901        let entered = sm.current();
6902        let is_reentry = Self::state_was_previously_entered(&entered, from_state, &history_before);
6903        drop(admission);
6904
6905        self.execute_state_enter_actions(&entered, is_reentry).await;
6906        drop(reservation);
6907        self.hooks
6908            .on_state_transition(Some(from_state), &entered, reason)
6909            .await;
6910        info!(from = %from_state, to = %entered, "State transition");
6911        Ok(true)
6912    }
6913
6914    /// Approves and applies a transition without staged context writes.
6915    async fn commit_transition_target(
6916        &self,
6917        from_state: &str,
6918        target: &str,
6919        reason: &str,
6920    ) -> Result<bool> {
6921        if !self.approve_transition_target(from_state, target).await? {
6922            return Ok(false);
6923        }
6924        self.apply_transition_target(from_state, target, reason, None)
6925            .await
6926    }
6927
6928    /// Applies an approved pre-response transition before redispatch.
6929    async fn apply_pre_response_transition_candidate(
6930        &self,
6931        candidate: &TransitionCandidate,
6932        staged: &HashMap<String, Value>,
6933        processed_input: &str,
6934    ) -> Result<bool> {
6935        self.commit_root_user_message(processed_input).await?;
6936        self.apply_transition_target(
6937            &candidate.from_state,
6938            candidate.target(),
6939            &candidate.reason,
6940            Some(staged),
6941        )
6942        .await
6943    }
6944
6945    /// Commits a pre-response transition after approval and before redispatch.
6946    async fn commit_pre_response_transition_candidate(
6947        &self,
6948        candidate: &TransitionCandidate,
6949        staged: &HashMap<String, Value>,
6950        processed_input: &str,
6951    ) -> Result<bool> {
6952        if !self
6953            .approve_transition_target(&candidate.from_state, candidate.target())
6954            .await?
6955        {
6956            return Ok(false);
6957        }
6958        self.apply_pre_response_transition_candidate(candidate, staged, processed_input)
6959            .await
6960    }
6961
6962    /// Handles post-response transition misses with fallback counters.
6963    async fn handle_transition_miss(&self, current_state: &str) -> Result<bool> {
6964        let Some(ref sm) = self.state_machine else {
6965            return Ok(false);
6966        };
6967        sm.increment_no_transition();
6968        let Some(fallback) = sm.check_fallback() else {
6969            return Ok(false);
6970        };
6971        self.commit_transition_target(current_state, &fallback, "fallback after no transitions")
6972            .await
6973    }
6974
6975    /// Evaluates and commits post-response transitions for the committed response path.
6976    async fn evaluate_transitions(&self, user_message: &str, response: &str) -> Result<bool> {
6977        let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6978            return Ok(false);
6979        };
6980        if transitions.is_empty() {
6981            return Ok(false);
6982        }
6983        if let Some(candidate) = self
6984            .select_transition_candidate(user_message, response)
6985            .await?
6986        {
6987            return self.commit_transition_candidate(&candidate).await;
6988        }
6989        self.handle_transition_miss(&current_state).await
6990    }
6991
6992    /// Attempts deterministic pre-response routing before old-state response generation.
6993    async fn try_pre_response_transition(
6994        &self,
6995        processed_input: &str,
6996    ) -> Result<Option<AgentResponse>> {
6997        let optimization = &self.runtime_config.optimization;
6998        if !optimization.enabled || !optimization.pre_response_deterministic_transitions {
6999            return Ok(None);
7000        }
7001        let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
7002            return Ok(None);
7003        };
7004        let eligible: Vec<Transition> = transitions
7005            .into_iter()
7006            .filter(|transition| !transition.requires_response)
7007            .filter(|transition| matches!(transition.timing, TransitionTiming::PreResponse))
7008            .collect();
7009        if eligible.is_empty() {
7010            return Ok(None);
7011        }
7012
7013        let empty_staged = HashMap::new();
7014        let mut extracted_staged: Option<HashMap<String, Value>> = None;
7015        let mut selected: Option<(TransitionCandidate, HashMap<String, Value>)> = None;
7016
7017        for transition in &eligible {
7018            let use_extractors = optimization.pre_response_extractors || transition.run_extractors;
7019            let staged_for_eval = if use_extractors {
7020                if extracted_staged.is_none() {
7021                    extracted_staged =
7022                        Some(self.run_context_extractors_staged(processed_input).await);
7023                }
7024                extracted_staged.as_ref().unwrap_or(&empty_staged)
7025            } else {
7026                &empty_staged
7027            };
7028
7029            if let Some(candidate) = self.select_deterministic_transition_candidate(
7030                processed_input,
7031                &current_state,
7032                std::slice::from_ref(transition),
7033                staged_for_eval,
7034            ) {
7035                let staged_for_commit = if use_extractors {
7036                    staged_for_eval.clone()
7037                } else {
7038                    HashMap::new()
7039                };
7040                selected = Some((candidate, staged_for_commit));
7041                break;
7042            }
7043        }
7044
7045        let Some((candidate, staged)) = selected else {
7046            return Ok(None);
7047        };
7048
7049        if !self
7050            .commit_pre_response_transition_candidate(&candidate, &staged, processed_input)
7051            .await?
7052        {
7053            return Ok(None);
7054        }
7055        self.redispatch_current_state(processed_input)
7056            .await
7057            .map(Some)
7058    }
7059
7060    //
7061    // Speculative branches overlap independent decisions but still commit exactly one path.
7062    // Losing branches must remain data only and must not write memory, run tools, or emit output.
7063    //
7064    async fn try_speculative_branches(
7065        &self,
7066        processed_input: &str,
7067        input_context: &HashMap<String, Value>,
7068    ) -> Result<Option<AgentResponse>> {
7069        let optimization = &self.runtime_config.optimization;
7070        if !optimization.enabled {
7071            return Ok(None);
7072        }
7073
7074        let effective_reasoning_mode = self.get_effective_reasoning_config().mode.clone();
7075        if !matches!(
7076            effective_reasoning_mode,
7077            ReasoningMode::None | ReasoningMode::Auto
7078        ) {
7079            return Ok(None);
7080        }
7081
7082        let mut transition_enabled =
7083            optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
7084        let mut skill_enabled = optimization.speculative_skill_routing
7085            && self.skill_router.is_some()
7086            && self.pending_skill_id.read().is_none();
7087        let mut reasoning_enabled = optimization.speculative_reasoning_auto
7088            && matches!(effective_reasoning_mode, ReasoningMode::Auto);
7089
7090        if matches!(effective_reasoning_mode, ReasoningMode::Auto)
7091            && (!reasoning_enabled || optimization.max_speculative_llm_calls_per_turn < 2)
7092        {
7093            return Ok(None);
7094        }
7095
7096        if !transition_enabled && !skill_enabled && !reasoning_enabled {
7097            return Ok(None);
7098        }
7099
7100        let mut optional_slots = optimization.max_parallel_runtime_tasks.saturating_sub(1);
7101        let mut speculative_call_slots = optimization
7102            .max_speculative_llm_calls_per_turn
7103            .saturating_sub(1);
7104        if reasoning_enabled {
7105            if optional_slots == 0 || speculative_call_slots == 0 {
7106                return Ok(None);
7107            }
7108            optional_slots -= 1;
7109            speculative_call_slots -= 1;
7110        }
7111        if transition_enabled {
7112            if optional_slots == 0 {
7113                transition_enabled = false;
7114            } else {
7115                optional_slots -= 1;
7116            }
7117        }
7118        if skill_enabled && (optional_slots == 0 || speculative_call_slots == 0) {
7119            skill_enabled = false;
7120        }
7121
7122        if !transition_enabled && !skill_enabled && !reasoning_enabled {
7123            return Ok(None);
7124        }
7125
7126        let main_kind = if transition_enabled {
7127            RuntimeOptimizationKind::ParallelStateTransition
7128        } else if skill_enabled {
7129            RuntimeOptimizationKind::SpeculativeSkillRouting
7130        } else {
7131            RuntimeOptimizationKind::SpeculativeReasoningAuto
7132        };
7133        if !self.reserve_active_speculative_llm_call(main_kind) {
7134            return Ok(None);
7135        }
7136
7137        let mut branch_set = ScheduledBranchSet::new(optimization.max_parallel_runtime_tasks)?;
7138        let main_branch = RuntimeBranch::new(
7139            RuntimeTaskPurpose::MainResponse,
7140            main_kind,
7141            RuntimeTaskPriority::Normal,
7142            RuntimeCommitBehavior::FinalResponse,
7143        );
7144        let transition_branch = RuntimeBranch::new(
7145            RuntimeTaskPurpose::StateTransition,
7146            RuntimeOptimizationKind::ParallelStateTransition,
7147            RuntimeTaskPriority::Critical,
7148            RuntimeCommitBehavior::TransitionDecision,
7149        );
7150        let skill_branch = RuntimeBranch::new(
7151            RuntimeTaskPurpose::SkillRouting,
7152            RuntimeOptimizationKind::SpeculativeSkillRouting,
7153            RuntimeTaskPriority::High,
7154            RuntimeCommitBehavior::SkillSelection,
7155        );
7156        let reasoning_branch = RuntimeBranch::new(
7157            RuntimeTaskPurpose::ReasoningJudge,
7158            RuntimeOptimizationKind::SpeculativeReasoningAuto,
7159            RuntimeTaskPriority::Normal,
7160            RuntimeCommitBehavior::ReasoningDecision,
7161        );
7162        let main_id = main_branch.branch_id();
7163        let transition_id = transition_branch.branch_id();
7164        let skill_id = skill_branch.branch_id();
7165        let reasoning_id = reasoning_branch.branch_id();
7166
7167        let main_id_for_future = main_id.clone();
7168        if !branch_set.schedule(
7169            main_branch,
7170            Box::pin(async move {
7171                match crate::optimization::observability::with_branch_observation(
7172                    &main_id_for_future,
7173                    main_kind,
7174                    RuntimeCommitBehavior::FinalResponse,
7175                    self.generate_main_response_draft(processed_input, &ReasoningMode::None),
7176                )
7177                .await
7178                {
7179                    Ok(draft) => RuntimeBranchResult::MainDraft(draft),
7180                    Err(error) => RuntimeBranchResult::Failed(error),
7181                }
7182            }),
7183        ) {
7184            return Ok(None);
7185        }
7186
7187        if transition_enabled {
7188            let transition_id_for_future = transition_id.clone();
7189            if !branch_set.schedule(
7190                transition_branch,
7191                Box::pin(async move {
7192                    match crate::optimization::observability::with_branch_observation(
7193                        &transition_id_for_future,
7194                        RuntimeOptimizationKind::ParallelStateTransition,
7195                        RuntimeCommitBehavior::TransitionDecision,
7196                        self.select_parallel_transition_candidate(processed_input),
7197                    )
7198                    .await
7199                    {
7200                        Ok(ParallelTransitionSelection::Candidate(candidate)) => {
7201                            RuntimeBranchResult::Transition(Some(candidate))
7202                        }
7203                        Ok(ParallelTransitionSelection::NoMatch) => {
7204                            RuntimeBranchResult::Transition(None)
7205                        }
7206                        Ok(ParallelTransitionSelection::ReservationExhausted) => {
7207                            RuntimeBranchResult::Cancelled
7208                        }
7209                        Err(error) => RuntimeBranchResult::Failed(error),
7210                    }
7211                }),
7212            ) {
7213                transition_enabled = false;
7214            }
7215        }
7216
7217        if skill_enabled {
7218            let skill_id_for_future = skill_id.clone();
7219            if !branch_set.schedule(
7220                skill_branch,
7221                Box::pin(async move {
7222                    if !self.reserve_active_speculative_llm_call(
7223                        RuntimeOptimizationKind::SpeculativeSkillRouting,
7224                    ) {
7225                        return RuntimeBranchResult::Cancelled;
7226                    }
7227                    match crate::optimization::observability::with_branch_observation(
7228                        &skill_id_for_future,
7229                        RuntimeOptimizationKind::SpeculativeSkillRouting,
7230                        RuntimeCommitBehavior::SkillSelection,
7231                        self.select_skill_candidate(processed_input),
7232                    )
7233                    .await
7234                    {
7235                        Ok(candidate) => RuntimeBranchResult::Skill(candidate),
7236                        Err(error) => RuntimeBranchResult::Failed(error),
7237                    }
7238                }),
7239            ) {
7240                skill_enabled = false;
7241            }
7242        }
7243
7244        if reasoning_enabled {
7245            let reasoning_id_for_future = reasoning_id.clone();
7246            if !branch_set.schedule(
7247                reasoning_branch,
7248                Box::pin(async move {
7249                    if !self.reserve_active_speculative_llm_call(
7250                        RuntimeOptimizationKind::SpeculativeReasoningAuto,
7251                    ) {
7252                        return RuntimeBranchResult::Cancelled;
7253                    }
7254                    match crate::optimization::observability::with_branch_observation(
7255                        &reasoning_id_for_future,
7256                        RuntimeOptimizationKind::SpeculativeReasoningAuto,
7257                        RuntimeCommitBehavior::ReasoningDecision,
7258                        self.determine_reasoning_mode_strict(processed_input),
7259                    )
7260                    .await
7261                    {
7262                        Ok(mode) => RuntimeBranchResult::Reasoning(mode),
7263                        Err(error) => RuntimeBranchResult::Failed(error),
7264                    }
7265                }),
7266            ) {
7267                reasoning_enabled = false;
7268            }
7269        }
7270
7271        if matches!(effective_reasoning_mode, ReasoningMode::Auto) && !reasoning_enabled {
7272            self.finalize_pending_branches(branch_set.cancel_pending());
7273            return Ok(None);
7274        }
7275
7276        if !transition_enabled && !skill_enabled && !reasoning_enabled {
7277            self.finalize_pending_branches(branch_set.cancel_pending());
7278            return Ok(None);
7279        }
7280
7281        let mut main_pending = true;
7282        let mut skill_pending = skill_enabled;
7283        let mut reasoning_pending = reasoning_enabled;
7284        let mut transition_finalized = !transition_enabled;
7285        let mut skill_finalized = !skill_enabled;
7286        let mut reasoning_finalized = !reasoning_enabled;
7287        let mut main_result: Option<Result<MainResponseDraft>> = None;
7288        let mut transition_candidate: Option<TransitionCandidate> = None;
7289        let mut skill_candidate: Option<SkillCandidate> = None;
7290        let mut reasoning_decision: Option<ReasoningMode> = None;
7291        let mut transition_fallback_required = false;
7292        let mut skill_fallback_required = false;
7293        let mut reasoning_fallback_required = false;
7294
7295        loop {
7296            if let Some(candidate) = transition_candidate.take() {
7297                if self
7298                    .approve_transition_target(&candidate.from_state, candidate.target())
7299                    .await?
7300                {
7301                    // Drop losing provider futures before transition side effects can reuse their shared resources.
7302                    self.finalize_pending_branches(branch_set.cancel_pending());
7303                    if !main_pending {
7304                        self.finalize_branch_loss(
7305                            &main_id,
7306                            main_kind,
7307                            RuntimeCommitBehavior::FinalResponse,
7308                            false,
7309                            main_result.as_ref().map(|result| result.is_err()),
7310                        );
7311                    }
7312                    if skill_enabled && !skill_pending {
7313                        self.finalize_branch_loss(
7314                            &skill_id,
7315                            RuntimeOptimizationKind::SpeculativeSkillRouting,
7316                            RuntimeCommitBehavior::SkillSelection,
7317                            false,
7318                            Some(false),
7319                        );
7320                    }
7321                    if reasoning_enabled && !reasoning_pending {
7322                        self.finalize_branch_loss(
7323                            &reasoning_id,
7324                            RuntimeOptimizationKind::SpeculativeReasoningAuto,
7325                            RuntimeCommitBehavior::ReasoningDecision,
7326                            false,
7327                            Some(false),
7328                        );
7329                    }
7330                    if !self
7331                        .apply_pre_response_transition_candidate(
7332                            &candidate,
7333                            &HashMap::new(),
7334                            processed_input,
7335                        )
7336                        .await?
7337                    {
7338                        self.finalize_optional_branch(
7339                            &transition_id,
7340                            RuntimeOptimizationKind::ParallelStateTransition,
7341                            RuntimeCommitBehavior::TransitionDecision,
7342                            "discarded",
7343                            false,
7344                        );
7345                        return Ok(None);
7346                    }
7347                    self.finalize_optional_branch(
7348                        &transition_id,
7349                        RuntimeOptimizationKind::ParallelStateTransition,
7350                        RuntimeCommitBehavior::TransitionDecision,
7351                        "committed",
7352                        true,
7353                    );
7354                    return self
7355                        .redispatch_current_state(processed_input)
7356                        .await
7357                        .map(Some);
7358                }
7359                self.finalize_optional_branch(
7360                    &transition_id,
7361                    RuntimeOptimizationKind::ParallelStateTransition,
7362                    RuntimeCommitBehavior::TransitionDecision,
7363                    "discarded",
7364                    false,
7365                );
7366                transition_finalized = true;
7367            }
7368
7369            if transition_finalized && skill_candidate.is_some() {
7370                let candidate = skill_candidate.take().unwrap();
7371                self.finalize_optional_branch(
7372                    &skill_id,
7373                    RuntimeOptimizationKind::SpeculativeSkillRouting,
7374                    RuntimeCommitBehavior::SkillSelection,
7375                    "committed",
7376                    true,
7377                );
7378                if !main_pending {
7379                    self.finalize_branch_loss(
7380                        &main_id,
7381                        main_kind,
7382                        RuntimeCommitBehavior::FinalResponse,
7383                        false,
7384                        main_result.as_ref().map(|result| result.is_err()),
7385                    );
7386                }
7387                if reasoning_enabled && !reasoning_pending {
7388                    self.finalize_branch_loss(
7389                        &reasoning_id,
7390                        RuntimeOptimizationKind::SpeculativeReasoningAuto,
7391                        RuntimeCommitBehavior::ReasoningDecision,
7392                        false,
7393                        Some(false),
7394                    );
7395                }
7396                self.finalize_pending_branches(branch_set.cancel_pending());
7397                self.commit_root_user_message(processed_input).await?;
7398                return match self
7399                    .commit_skill_candidate_route_result(candidate, processed_input)
7400                    .await?
7401                {
7402                    SkillRouteResult::Response { skill_id, content } => self
7403                        .handle_skill_response(processed_input, &skill_id, content, input_context)
7404                        .await
7405                        .map(Some),
7406                    SkillRouteResult::NeedsClarification {
7407                        response,
7408                        ownership,
7409                    } => {
7410                        let admission = self
7411                            .admit_optional_disambiguation_ownership(ownership)
7412                            .await?;
7413                        if response
7414                            .metadata
7415                            .as_ref()
7416                            .and_then(|m| m.get("disambiguation"))
7417                            .and_then(|d| d.get("status"))
7418                            .and_then(|s| s.as_str())
7419                            == Some("awaiting_clarification")
7420                        {
7421                            self.memory
7422                                .add_message(ChatMessage::assistant(&response.content))
7423                                .await?;
7424                        }
7425                        drop(admission);
7426                        self.finish_turn_if_root(&response).await?;
7427                        Ok(Some(response))
7428                    }
7429                    SkillRouteResult::NoMatch => Ok(None),
7430                };
7431            }
7432
7433            if transition_finalized
7434                && skill_finalized
7435                && let Some(reasoning_mode) = reasoning_decision.take()
7436            {
7437                if !matches!(reasoning_mode, ReasoningMode::None) {
7438                    self.finalize_optional_branch(
7439                        &reasoning_id,
7440                        RuntimeOptimizationKind::SpeculativeReasoningAuto,
7441                        RuntimeCommitBehavior::ReasoningDecision,
7442                        "committed",
7443                        true,
7444                    );
7445                    if !main_pending {
7446                        self.finalize_branch_loss(
7447                            &main_id,
7448                            main_kind,
7449                            RuntimeCommitBehavior::FinalResponse,
7450                            false,
7451                            main_result.as_ref().map(|result| result.is_err()),
7452                        );
7453                    }
7454                    self.finalize_pending_branches(branch_set.cancel_pending());
7455                    self.commit_root_user_message(processed_input).await?;
7456                    return if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
7457                        self.handle_plan_and_execute(processed_input, input_context, true)
7458                            .await
7459                            .map(Some)
7460                    } else {
7461                        self.run_committed_response_loop_with_reasoning(
7462                            processed_input,
7463                            input_context,
7464                            reasoning_mode,
7465                            true,
7466                        )
7467                        .await
7468                        .map(Some)
7469                    };
7470                }
7471                self.finalize_optional_branch(
7472                    &reasoning_id,
7473                    RuntimeOptimizationKind::SpeculativeReasoningAuto,
7474                    RuntimeCommitBehavior::ReasoningDecision,
7475                    "committed",
7476                    true,
7477                );
7478                reasoning_finalized = true;
7479            }
7480
7481            if transition_finalized && skill_finalized && reasoning_finalized {
7482                if transition_fallback_required
7483                    || skill_fallback_required
7484                    || reasoning_fallback_required
7485                {
7486                    if !main_pending {
7487                        self.finalize_branch_loss(
7488                            &main_id,
7489                            main_kind,
7490                            RuntimeCommitBehavior::FinalResponse,
7491                            false,
7492                            main_result.as_ref().map(|result| result.is_err()),
7493                        );
7494                    }
7495                    self.finalize_pending_branches(branch_set.cancel_pending());
7496                    return Ok(None);
7497                }
7498
7499                if let Some(result) = main_result.take() {
7500                    let draft = match result {
7501                        Ok(draft) => draft,
7502                        Err(error) => {
7503                            self.finalize_optional_branch(
7504                                &main_id,
7505                                main_kind,
7506                                RuntimeCommitBehavior::FinalResponse,
7507                                "failed",
7508                                false,
7509                            );
7510                            self.finalize_pending_branches(branch_set.cancel_pending());
7511                            return Err(error);
7512                        }
7513                    };
7514                    self.finalize_optional_branch(
7515                        &main_id,
7516                        main_kind,
7517                        RuntimeCommitBehavior::FinalResponse,
7518                        "committed",
7519                        true,
7520                    );
7521                    self.finalize_pending_branches(branch_set.cancel_pending());
7522                    return self
7523                        .commit_main_response_draft(
7524                            processed_input,
7525                            input_context,
7526                            draft,
7527                            ReasoningMode::None,
7528                            reasoning_enabled,
7529                        )
7530                        .await
7531                        .map(Some);
7532                }
7533            }
7534
7535            if branch_set.is_empty() {
7536                return Ok(None);
7537            }
7538
7539            let Some(outcome) = branch_set.next_completed().await else {
7540                return Ok(None);
7541            };
7542            let branch_id = outcome.branch.branch_id();
7543            match outcome.result {
7544                RuntimeBranchResult::MainDraft(draft) => {
7545                    main_pending = false;
7546                    main_result = Some(Ok(draft));
7547                }
7548                RuntimeBranchResult::Transition(candidate) => {
7549                    if let Some(candidate) = candidate {
7550                        transition_candidate = Some(candidate);
7551                    } else {
7552                        self.finalize_optional_branch(
7553                            &transition_id,
7554                            RuntimeOptimizationKind::ParallelStateTransition,
7555                            RuntimeCommitBehavior::TransitionDecision,
7556                            "discarded",
7557                            false,
7558                        );
7559                        transition_finalized = true;
7560                    }
7561                }
7562                RuntimeBranchResult::Skill(candidate) => {
7563                    skill_pending = false;
7564                    if let Some(candidate) = candidate {
7565                        skill_candidate = Some(candidate);
7566                    } else {
7567                        self.finalize_optional_branch(
7568                            &skill_id,
7569                            RuntimeOptimizationKind::SpeculativeSkillRouting,
7570                            RuntimeCommitBehavior::SkillSelection,
7571                            "discarded",
7572                            false,
7573                        );
7574                        skill_finalized = true;
7575                    }
7576                }
7577                RuntimeBranchResult::Reasoning(mode) => {
7578                    reasoning_pending = false;
7579                    reasoning_decision = Some(mode);
7580                }
7581                RuntimeBranchResult::Failed(error) => {
7582                    if branch_id == main_id {
7583                        main_pending = false;
7584                        main_result = Some(Err(error));
7585                    } else if branch_id == transition_id {
7586                        self.finalize_optional_branch(
7587                            &transition_id,
7588                            RuntimeOptimizationKind::ParallelStateTransition,
7589                            RuntimeCommitBehavior::TransitionDecision,
7590                            "failed",
7591                            false,
7592                        );
7593                        transition_finalized = true;
7594                    } else if branch_id == skill_id {
7595                        skill_pending = false;
7596                        self.finalize_optional_branch(
7597                            &skill_id,
7598                            RuntimeOptimizationKind::SpeculativeSkillRouting,
7599                            RuntimeCommitBehavior::SkillSelection,
7600                            "failed",
7601                            false,
7602                        );
7603                        skill_finalized = true;
7604                    } else if branch_id == reasoning_id {
7605                        reasoning_pending = false;
7606                        self.finalize_optional_branch(
7607                            &reasoning_id,
7608                            RuntimeOptimizationKind::SpeculativeReasoningAuto,
7609                            RuntimeCommitBehavior::ReasoningDecision,
7610                            "failed",
7611                            false,
7612                        );
7613                        reasoning_finalized = true;
7614                    }
7615                }
7616                RuntimeBranchResult::Cancelled => {
7617                    self.finalize_optional_branch(
7618                        &branch_id,
7619                        outcome.branch.optimization,
7620                        outcome.branch.commit_behavior,
7621                        "cancelled",
7622                        false,
7623                    );
7624                    if branch_id == main_id {
7625                        main_pending = false;
7626                        main_result =
7627                            Some(Err(AgentError::Other("main branch cancelled".to_string())));
7628                    } else if branch_id == transition_id {
7629                        transition_finalized = true;
7630                        transition_fallback_required = true;
7631                    } else if branch_id == skill_id {
7632                        skill_pending = false;
7633                        skill_finalized = true;
7634                        skill_fallback_required = true;
7635                    } else if branch_id == reasoning_id {
7636                        reasoning_pending = false;
7637                        reasoning_finalized = true;
7638                        reasoning_fallback_required = true;
7639                    }
7640                }
7641            }
7642        }
7643    }
7644
7645    fn finalize_pending_branches(&self, branches: Vec<RuntimeBranch>) {
7646        for branch in branches {
7647            self.finalize_optional_branch(
7648                &branch.branch_id(),
7649                branch.optimization,
7650                branch.commit_behavior,
7651                "cancelled",
7652                false,
7653            );
7654        }
7655    }
7656
7657    //
7658    // Pending losers are reported as cancelled because their futures are dropped before completion.
7659    // Completed losers keep failed or discarded status based on their recorded result.
7660    //
7661    fn finalize_branch_loss(
7662        &self,
7663        branch_id: &str,
7664        optimization: RuntimeOptimizationKind,
7665        commit_behavior: RuntimeCommitBehavior,
7666        pending: bool,
7667        completed_failed: Option<bool>,
7668    ) {
7669        let status = if pending {
7670            "cancelled"
7671        } else if completed_failed.unwrap_or(false) {
7672            "failed"
7673        } else {
7674            "discarded"
7675        };
7676        self.finalize_optional_branch(branch_id, optimization, commit_behavior, status, false);
7677    }
7678
7679    //
7680    // Finalization is separated from commit so losing branches remain observable.
7681    // This helper must not mutate runtime state other than observability.
7682    //
7683    fn finalize_optional_branch(
7684        &self,
7685        branch_id: &str,
7686        optimization: RuntimeOptimizationKind,
7687        commit_behavior: RuntimeCommitBehavior,
7688        status: &str,
7689        winner: bool,
7690    ) {
7691        crate::optimization::observability::finalize_branch(
7692            self.observability_manager.as_ref(),
7693            branch_id,
7694            status,
7695            winner,
7696            optimization,
7697            commit_behavior,
7698        );
7699    }
7700
7701    //
7702    // This is only an eligibility check.
7703    // Actual transition selection happens in select_parallel_transition_candidate.
7704    //
7705    fn has_parallel_transition_candidates(&self) -> bool {
7706        self.transitions_available_for_commit()
7707            .map(|(transitions, _)| {
7708                transitions
7709                    .iter()
7710                    .any(|transition| matches!(transition.timing, TransitionTiming::Parallel))
7711            })
7712            .unwrap_or(false)
7713    }
7714
7715    //
7716    // Parallel transition prompts must not depend on assistant response text.
7717    // Keep this branch response-independent or it can race against invalid context.
7718    //
7719    async fn select_parallel_transition_candidate(
7720        &self,
7721        processed_input: &str,
7722    ) -> Result<ParallelTransitionSelection> {
7723        let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
7724            return Ok(ParallelTransitionSelection::NoMatch);
7725        };
7726        let parallel: Vec<Transition> = transitions
7727            .into_iter()
7728            .filter(|transition| matches!(transition.timing, TransitionTiming::Parallel))
7729            .filter(|transition| !transition.requires_response)
7730            .collect();
7731        if parallel.is_empty() {
7732            return Ok(ParallelTransitionSelection::NoMatch);
7733        }
7734        let empty_staged = HashMap::new();
7735        if let Some(candidate) = self.select_deterministic_transition_candidate(
7736            processed_input,
7737            &current_state,
7738            &parallel,
7739            &empty_staged,
7740        ) {
7741            return Ok(ParallelTransitionSelection::Candidate(candidate));
7742        }
7743        let when_transitions: Vec<(usize, &Transition)> = parallel
7744            .iter()
7745            .enumerate()
7746            .filter(|(_, transition)| !transition.when.trim().is_empty())
7747            .collect();
7748        if when_transitions.is_empty() {
7749            return Ok(ParallelTransitionSelection::NoMatch);
7750        }
7751        let llm = self
7752            .llm_registry
7753            .router()
7754            .or_else(|_| self.llm_registry.default())
7755            .map_err(|e| AgentError::Config(e.to_string()))?;
7756        let conditions = when_transitions
7757            .iter()
7758            .enumerate()
7759            .map(|(display_idx, (_, transition))| {
7760                format!("{}. {}", display_idx + 1, transition.when)
7761            })
7762            .collect::<Vec<_>>()
7763            .join("\n");
7764        if !self
7765            .reserve_active_speculative_llm_call(RuntimeOptimizationKind::ParallelStateTransition)
7766        {
7767            return Ok(ParallelTransitionSelection::ReservationExhausted);
7768        }
7769        let context_preview = self.branch_context_preview();
7770        let prompt = format!(
7771            "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-{}).",
7772            current_state,
7773            processed_input,
7774            context_preview,
7775            conditions,
7776            when_transitions.len()
7777        );
7778        let response = self
7779            .observe_purpose(
7780                ObservationPurpose::StateTransitionEvaluation,
7781                llm.complete(&[ChatMessage::user(prompt)], None),
7782            )
7783            .await
7784            .map_err(|e| AgentError::LLM(e.to_string()))?;
7785        let choice = response.content.trim().parse::<usize>().unwrap_or(0);
7786        if choice == 0 || choice > when_transitions.len() {
7787            return Ok(ParallelTransitionSelection::NoMatch);
7788        }
7789        let transition = when_transitions[choice - 1].1.clone();
7790        Ok(ParallelTransitionSelection::Candidate(
7791            TransitionCandidate::new(
7792                current_state,
7793                transition.clone(),
7794                Self::transition_reason(&transition),
7795            ),
7796        ))
7797    }
7798
7799    /// Re-enters the runtime loop after an optimized transition commits.
7800    async fn redispatch_current_state(&self, processed_input: &str) -> Result<AgentResponse> {
7801        const MAX_REDISPATCH_DEPTH: u32 = 3;
7802        let current_depth = *self.redispatch_depth.read();
7803        if current_depth >= MAX_REDISPATCH_DEPTH {
7804            warn!(depth = current_depth, "Re-dispatch depth limit reached");
7805            let response = AgentResponse::new("");
7806            self.finish_turn_if_root(&response).await?;
7807            return Ok(response);
7808        }
7809        *self.redispatch_depth.write() += 1;
7810        if let Some(context) = self.active_turn_context.write().as_mut() {
7811            context.enter_redispatch();
7812        }
7813        let result = Box::pin(self.run_loop_internal(processed_input)).await;
7814        *self.redispatch_depth.write() -= 1;
7815        if let Some(context) = self.active_turn_context.write().as_mut() {
7816            context.exit_redispatch();
7817        }
7818        let response = result?;
7819        self.finish_turn_if_root(&response).await?;
7820        Ok(response)
7821    }
7822
7823    /// Runs final response hooks and maintenance only for the root dispatch.
7824    async fn finish_turn_if_root(&self, response: &AgentResponse) -> Result<()> {
7825        if *self.redispatch_depth.read() == 0 {
7826            self.post_turn_session_lifecycle().await?;
7827            if let Some(context) = self.active_turn_context.write().as_mut() {
7828                context.mark_post_turn_lifecycle_completed();
7829            }
7830            self.hooks.on_response(response).await;
7831            self.end_root_turn();
7832        }
7833        Ok(())
7834    }
7835
7836    /// Execute on_exit actions for a state being left.
7837    async fn execute_state_exit_actions(&self, state_path: &str) {
7838        if let Some(ref sm) = self.state_machine
7839            && let Some(def) = sm.get_definition(state_path)
7840            && !def.on_exit.is_empty()
7841        {
7842            debug!(state = %state_path, count = def.on_exit.len(), "Executing on_exit actions");
7843            self.execute_state_actions(&def.on_exit).await;
7844        }
7845    }
7846
7847    /// Returns whether a transition target had already been entered before this transition.
7848    fn state_was_previously_entered(
7849        state_path: &str,
7850        from_state: &str,
7851        history_before: &[StateTransitionEvent],
7852    ) -> bool {
7853        state_path == from_state
7854            || history_before
7855                .iter()
7856                .any(|event| event.from == state_path || event.to == state_path)
7857    }
7858
7859    /// Execute on_enter (or on_reenter) actions for a state being entered.
7860    async fn execute_state_enter_actions(&self, state_path: &str, is_reentry: bool) {
7861        if let Some(ref sm) = self.state_machine
7862            && let Some(def) = sm.get_definition(state_path)
7863        {
7864            if is_reentry && !def.on_reenter.is_empty() {
7865                debug!(state = %state_path, count = def.on_reenter.len(), "Executing on_reenter actions");
7866                self.execute_state_actions(&def.on_reenter).await;
7867            } else if !def.on_enter.is_empty() {
7868                debug!(state = %state_path, count = def.on_enter.len(), "Executing on_enter actions");
7869                self.execute_state_actions(&def.on_enter).await;
7870            }
7871        }
7872    }
7873
7874    /// Execute a list of state actions (tool calls, skill invocations, context updates, LLM prompts).
7875    async fn execute_state_actions(&self, actions: &[StateAction]) {
7876        for (action_index, action) in actions.iter().enumerate() {
7877            match action {
7878                StateAction::Tool { tool, args } => {
7879                    let raw_args = args.clone().unwrap_or(Value::Object(Default::default()));
7880                    let args_value = self.render_action_args(&raw_args);
7881                    let state = self.state_machine.as_ref().map(|sm| sm.current());
7882                    let request = ToolExecutionRequest::new(
7883                        uuid::Uuid::new_v4().to_string(),
7884                        tool.clone(),
7885                        args_value,
7886                        ToolCallSource::StateAction {
7887                            state,
7888                            action_index,
7889                        },
7890                    );
7891                    match self.execute_tool_record(request).await {
7892                        Ok(record) if record.success => {
7893                            debug!(tool = %record.canonical_id, "State action: tool executed");
7894                            let _ = self.context_manager.set(
7895                                "last_tool_result",
7896                                serde_json::Value::String(record.model_output_string()),
7897                            );
7898                            let _ = self.context_manager.set(
7899                                "last_tool_record",
7900                                serde_json::to_value(record).unwrap_or(Value::Null),
7901                            );
7902                        }
7903                        Ok(record) => {
7904                            warn!(tool = %record.canonical_id, error = %record.output, "State action: tool failed");
7905                        }
7906                        Err(e) => {
7907                            warn!(tool = %tool, error = %e, "State action: tool failed")
7908                        }
7909                    }
7910                }
7911                StateAction::Skill { skill } => {
7912                    if let Some(ref executor) = self.skill_executor {
7913                        if let Some(def) = self.skills.iter().find(|s| s.id == *skill) {
7914                            match executor
7915                                .execute_with_invoker(def, "", serde_json::json!({}), self)
7916                                .await
7917                            {
7918                                Ok(_) => debug!(skill = %skill, "State action: skill executed"),
7919                                Err(e) => {
7920                                    warn!(skill = %skill, error = %e, "State action: skill failed")
7921                                }
7922                            }
7923                        } else {
7924                            warn!(skill = %skill, "State action: skill not found");
7925                        }
7926                    }
7927                }
7928                StateAction::SetContext { set_context } => {
7929                    for (key, value) in set_context {
7930                        if let Err(e) = self.context_manager.set(key, value.clone()) {
7931                            warn!(key = %key, error = %e, "State action: set_context failed");
7932                        } else {
7933                            debug!(key = %key, "State action: context set");
7934                        }
7935                    }
7936                }
7937                StateAction::Prompt {
7938                    prompt,
7939                    llm,
7940                    store_as,
7941                } => {
7942                    let llm_result = if let Some(alias) = llm {
7943                        self.llm_registry.get(alias)
7944                    } else {
7945                        self.llm_registry.default()
7946                    };
7947                    match llm_result {
7948                        Ok(llm_provider) => {
7949                            // Render template variables and include conversation context
7950                            let context = self.build_context_with_overlays();
7951                            let rendered_prompt = self
7952                                .template_renderer
7953                                .render(prompt, &context)
7954                                .unwrap_or_else(|_| prompt.clone());
7955                            let recent =
7956                                self.memory.get_messages(Some(5)).await.unwrap_or_default();
7957                            let mut messages: Vec<ChatMessage> = recent;
7958                            messages.push(ChatMessage::user(&rendered_prompt));
7959                            match self
7960                                .observe_purpose(
7961                                    ObservationPurpose::StateAction,
7962                                    llm_provider.complete(&messages, None),
7963                                )
7964                                .await
7965                            {
7966                                Ok(response) => {
7967                                    if let Some(key) = store_as {
7968                                        let _ = self
7969                                            .context_manager
7970                                            .set(key, Value::String(response.content));
7971                                        debug!(key = %key, "State action: prompt result stored");
7972                                    }
7973                                }
7974                                Err(e) => {
7975                                    warn!(error = %e, "State action: prompt LLM call failed");
7976                                }
7977                            }
7978                        }
7979                        Err(e) => {
7980                            warn!(error = %e, "State action: LLM not found for prompt");
7981                        }
7982                    }
7983                }
7984            }
7985        }
7986    }
7987
7988    async fn run_context_extractors_staged(&self, user_message: &str) -> HashMap<String, Value> {
7989        let extractors = match &self.state_machine {
7990            Some(sm) => match sm.current_definition() {
7991                Some(def) if !def.extract.is_empty() => def.extract.clone(),
7992                _ => return HashMap::new(),
7993            },
7994            None => return HashMap::new(),
7995        };
7996
7997        let mut staged = HashMap::new();
7998        for extractor in &extractors {
7999            let prompt = if let Some(ref custom) = extractor.llm_extract {
8000                format!(
8001                    "User message:\n\"{}\"\n\nInstruction:\n{}",
8002                    user_message, custom
8003                )
8004            } else if let Some(ref desc) = extractor.description {
8005                format!(
8006                    "From the following message, extract: {}\n\n\
8007                     Message: \"{}\"\n\n\
8008                     If the information is present, return ONLY the extracted value.\n\
8009                     If NOT present, return exactly: __NONE__",
8010                    desc, user_message
8011                )
8012            } else {
8013                continue;
8014            };
8015
8016            let llm = match self
8017                .llm_registry
8018                .get(&extractor.llm)
8019                .or_else(|_| self.llm_registry.get("router"))
8020                .or_else(|_| self.llm_registry.get("default"))
8021            {
8022                Ok(llm) => llm,
8023                Err(e) => {
8024                    warn!(key = %extractor.key, error = %e, "Extractor LLM not found");
8025                    continue;
8026                }
8027            };
8028
8029            let messages = vec![ChatMessage::user(&prompt)];
8030            match self
8031                .observe_purpose(
8032                    ObservationPurpose::ContextExtraction,
8033                    llm.complete(&messages, None),
8034                )
8035                .await
8036            {
8037                Ok(response) => {
8038                    let value = response.content.trim().to_string();
8039                    if value != "__NONE__" && !value.is_empty() {
8040                        staged.insert(
8041                            extractor.key.clone(),
8042                            serde_json::Value::String(value.clone()),
8043                        );
8044                        debug!(key = %extractor.key, value = %value, "Context extracted");
8045                    } else if extractor.required {
8046                        warn!(key = %extractor.key, "Required extraction returned no value");
8047                    }
8048                }
8049                Err(e) => {
8050                    warn!(key = %extractor.key, error = %e, "Context extraction LLM call failed");
8051                }
8052            }
8053        }
8054        staged
8055    }
8056
8057    fn commit_staged_context_writes(&self, staged: &HashMap<String, Value>) {
8058        for (key, value) in staged {
8059            if let Err(error) = self.context_manager.update(key, value.clone()) {
8060                warn!(key = %key, error = %error, "staged context write failed");
8061            }
8062        }
8063    }
8064
8065    /// Run context extractors for the current state on the user's input.
8066    async fn run_context_extractors(&self, user_message: &str) {
8067        let staged = self.run_context_extractors_staged(user_message).await;
8068        self.commit_staged_context_writes(&staged);
8069    }
8070
8071    async fn check_memory_compression(&self) -> Result<()> {
8072        if self.memory.needs_compression() {
8073            let result = self.memory.compress(None).await?;
8074            if let CompressResult::Compressed {
8075                messages_summarized,
8076                new_summary_length,
8077                tokens_saved,
8078            } = result
8079            {
8080                let event = MemoryCompressEvent::new(
8081                    messages_summarized,
8082                    tokens_saved,
8083                    new_summary_length as u32,
8084                );
8085                self.hooks.on_memory_compress(&event).await;
8086                debug!(
8087                    messages = messages_summarized,
8088                    tokens_saved = tokens_saved,
8089                    "Memory compressed"
8090                );
8091            }
8092        }
8093
8094        // Handle overflow AFTER compression, then check warning threshold
8095        self.handle_memory_overflow().await?;
8096        self.check_memory_budget().await;
8097
8098        Ok(())
8099    }
8100
8101    async fn check_memory_budget(&self) {
8102        let Some(ref budget) = self.memory_token_budget else {
8103            return;
8104        };
8105
8106        let context = match self.memory.get_context().await {
8107            Ok(ctx) => ctx,
8108            Err(_) => return,
8109        };
8110
8111        // Overall budget warning
8112        let used_tokens = context.estimated_tokens();
8113        if budget.is_over_warn_threshold(used_tokens) {
8114            let event = MemoryBudgetEvent::new("memory", used_tokens, budget.total);
8115            self.hooks.on_memory_budget_warning(&event).await;
8116            debug!(
8117                used = used_tokens,
8118                total = budget.total,
8119                percent = event.usage_percent,
8120                "Memory budget warning"
8121            );
8122        }
8123
8124        // Per-component warning: summary
8125        if let Some(ref summary) = context.summary {
8126            let summary_tokens = ai_agents_memory::estimate_tokens(summary);
8127            let summary_budget = budget.allocation.summary;
8128            if summary_budget > 0 {
8129                let warn_threshold =
8130                    (summary_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
8131                if summary_tokens >= warn_threshold {
8132                    let event = MemoryBudgetEvent::new("summary", summary_tokens, summary_budget);
8133                    self.hooks.on_memory_budget_warning(&event).await;
8134                }
8135            }
8136        }
8137
8138        // Per-component warning: recent_messages
8139        let recent_tokens: u32 = context
8140            .messages
8141            .iter()
8142            .map(ai_agents_memory::estimate_message_tokens)
8143            .sum();
8144        let recent_budget = budget.allocation.recent_messages;
8145        if recent_budget > 0 {
8146            let warn_threshold =
8147                (recent_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
8148            if recent_tokens >= warn_threshold {
8149                let event = MemoryBudgetEvent::new("recent_messages", recent_tokens, recent_budget);
8150                self.hooks.on_memory_budget_warning(&event).await;
8151            }
8152        }
8153
8154        let relationship_budget = budget.allocation.relationships;
8155        if relationship_budget > 0 {
8156            let relationship_tokens = self
8157                .relationship_memory_text()
8158                .map(|text| ai_agents_memory::estimate_tokens(&text))
8159                .unwrap_or(0);
8160            let warn_threshold =
8161                (relationship_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
8162            if relationship_tokens >= warn_threshold {
8163                let event = MemoryBudgetEvent::new(
8164                    "relationships",
8165                    relationship_tokens,
8166                    relationship_budget,
8167                );
8168                self.hooks.on_memory_budget_warning(&event).await;
8169            }
8170        }
8171    }
8172
8173    async fn handle_memory_overflow(&self) -> Result<()> {
8174        let Some(ref budget) = self.memory_token_budget else {
8175            return Ok(());
8176        };
8177
8178        let context = self.memory.get_context().await?;
8179        let used_tokens = context.estimated_tokens();
8180
8181        if used_tokens <= budget.total {
8182            return Ok(());
8183        }
8184
8185        match budget.overflow_strategy {
8186            OverflowStrategy::TruncateOldest => {
8187                let tokens_to_free = used_tokens - budget.total;
8188                let messages_to_evict = self.calculate_eviction_count(tokens_to_free);
8189                if messages_to_evict > 0 {
8190                    self.evict_messages(messages_to_evict, EvictionReason::TokenBudgetExceeded)
8191                        .await?;
8192                }
8193            }
8194            OverflowStrategy::SummarizeMore => {
8195                let max_attempts = context.total_messages.max(1);
8196                for _ in 0..max_attempts {
8197                    match self.memory.compress(None).await? {
8198                        CompressResult::Compressed {
8199                            messages_summarized,
8200                            ..
8201                        } if messages_summarized > 0 => {
8202                            let context = self.memory.get_context().await?;
8203                            if context.estimated_tokens() <= budget.total {
8204                                return Ok(());
8205                            }
8206                        }
8207                        _ => break,
8208                    }
8209                }
8210                let context = self.memory.get_context().await?;
8211                let used_tokens = context.estimated_tokens();
8212                if used_tokens > budget.total {
8213                    return Err(AgentError::MemoryBudgetExceeded {
8214                        used: used_tokens,
8215                        budget: budget.total,
8216                    });
8217                }
8218            }
8219            OverflowStrategy::Error => {
8220                return Err(AgentError::MemoryBudgetExceeded {
8221                    used: used_tokens,
8222                    budget: budget.total,
8223                });
8224            }
8225        }
8226        Ok(())
8227    }
8228
8229    fn calculate_eviction_count(&self, tokens_to_free: u32) -> usize {
8230        // Estimate ~50 tokens per message on average
8231        ((tokens_to_free as f64 / 50.0).ceil() as usize).max(1)
8232    }
8233
8234    async fn evict_messages(&self, count: usize, reason: EvictionReason) -> Result<()> {
8235        let evicted = self.memory.evict_oldest(count).await?;
8236        if !evicted.is_empty() {
8237            let event = MemoryEvictEvent {
8238                reason,
8239                messages_evicted: evicted.len(),
8240                importance_scores: vec![],
8241            };
8242            self.hooks.on_memory_evict(&event).await;
8243            debug!(count = evicted.len(), "Messages evicted from memory");
8244        }
8245        Ok(())
8246    }
8247
8248    #[instrument(skip(self, input), fields(agent = %self.info.name))]
8249    async fn determine_reasoning_mode(&self, input: &str) -> Result<ReasoningMode> {
8250        match self.determine_reasoning_mode_strict(input).await {
8251            Ok(mode) => Ok(mode),
8252            Err(_) => Ok(ReasoningMode::None),
8253        }
8254    }
8255
8256    async fn determine_reasoning_mode_strict(&self, input: &str) -> Result<ReasoningMode> {
8257        let effective_config = self.get_effective_reasoning_config();
8258
8259        if !matches!(effective_config.mode, ReasoningMode::Auto) {
8260            return Ok(effective_config.mode.clone());
8261        }
8262
8263        let judge_llm = effective_config
8264            .judge_llm
8265            .as_ref()
8266            .and_then(|alias| self.llm_registry.get(alias).ok())
8267            .or_else(|| self.llm_registry.router().ok())
8268            .or_else(|| self.llm_registry.default().ok());
8269
8270        let Some(llm) = judge_llm else {
8271            return Ok(ReasoningMode::None);
8272        };
8273
8274        let prompt = format!(
8275            r#"Analyze this user request and determine the appropriate reasoning mode.
8276
8277User request: "{}"
8278
8279Choose ONE of these modes:
8280- none: Simple queries, greetings, direct answers (fastest)
8281- cot: Complex analysis, multi-step reasoning, math problems
8282- react: Tasks requiring multiple tool calls with observation
8283- plan_and_execute: Complex multi-step tasks requiring coordination
8284
8285Respond with ONLY the mode name (none, cot, react, or plan_and_execute)."#,
8286            input
8287        );
8288
8289        let messages = vec![ChatMessage::user(&prompt)];
8290        let response = self
8291            .observe_purpose(
8292                ObservationPurpose::ReflectionDecision,
8293                llm.complete(&messages, None),
8294            )
8295            .await
8296            .map_err(|e| AgentError::LLM(e.to_string()))?;
8297
8298        let mode_str = response.content.trim().to_lowercase();
8299        Ok(match mode_str.as_str() {
8300            "cot" => ReasoningMode::CoT,
8301            "react" => ReasoningMode::React,
8302            "plan_and_execute" => ReasoningMode::PlanAndExecute,
8303            _ => ReasoningMode::None,
8304        })
8305    }
8306
8307    async fn should_reflect(&self, input: &str, response: &str) -> Result<bool> {
8308        let effective_config = self.get_effective_reflection_config();
8309
8310        if !effective_config.requires_evaluation() {
8311            return Ok(false);
8312        }
8313
8314        if effective_config.is_enabled() {
8315            return Ok(true);
8316        }
8317
8318        let evaluator_llm = effective_config
8319            .evaluator_llm
8320            .as_ref()
8321            .and_then(|alias| self.llm_registry.get(alias).ok())
8322            .or_else(|| self.llm_registry.router().ok())
8323            .or_else(|| self.llm_registry.default().ok());
8324
8325        let Some(llm) = evaluator_llm else {
8326            return Ok(false);
8327        };
8328
8329        let response_preview: String = response.chars().take(500).collect();
8330        let prompt = format!(
8331            r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
8332
8333User query: "{}"
8334Response: "{}"
8335
8336Answer YES or NO only."#,
8337            input, response_preview
8338        );
8339
8340        let messages = vec![ChatMessage::user(&prompt)];
8341        let result = self
8342            .observe_purpose(
8343                ObservationPurpose::ReflectionDecision,
8344                llm.complete(&messages, None),
8345            )
8346            .await;
8347
8348        match result {
8349            Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
8350            Err(_) => Ok(false),
8351        }
8352    }
8353
8354    fn build_cot_system_prompt(&self, base_prompt: &str) -> String {
8355        format!(
8356            "{}\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>",
8357            base_prompt
8358        )
8359    }
8360
8361    fn build_react_system_prompt(&self, base_prompt: &str) -> String {
8362        format!(
8363            "{}\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>",
8364            base_prompt
8365        )
8366    }
8367
8368    async fn generate_plan(&self, input: &str) -> Result<Plan> {
8369        let effective = self.get_effective_reasoning_config();
8370        let planning_config = effective.get_planning();
8371
8372        let planner_llm = planning_config
8373            .and_then(|c| c.planner_llm.as_ref())
8374            .and_then(|alias| self.llm_registry.get(alias).ok())
8375            .or_else(|| self.llm_registry.router().ok())
8376            .or_else(|| self.llm_registry.default().ok())
8377            .ok_or_else(|| AgentError::Config("No LLM available for planning".into()))?;
8378
8379        let mut available_tool_ids: Vec<String> = self
8380            .get_available_tool_ids()
8381            .await
8382            .unwrap_or_else(|_| self.tools.list_ids());
8383        let mut available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
8384
8385        // Apply planning-level tool and skill filters.
8386        if let Some(config) = planning_config {
8387            if !config.available.tools.is_all() {
8388                available_tool_ids.retain(|t| config.available.tools.allows(t));
8389            }
8390            if !config.available.skills.is_all() {
8391                available_skills.retain(|s| config.available.skills.allows(s));
8392            }
8393        }
8394
8395        // Build tool descriptions with argument schemas so the planner
8396        // knows how to construct valid args for each step.
8397        let tool_descriptions: Vec<String> = available_tool_ids
8398            .iter()
8399            .filter_map(|id| {
8400                self.tools.get(id).map(|tool| {
8401                    let schema = tool.input_schema();
8402                    let args_desc = schema
8403                        .get("properties")
8404                        .and_then(|p| serde_json::to_string(p).ok())
8405                        .unwrap_or_else(|| "{}".to_string());
8406                    format!(
8407                        "- {} ({}): {}\n  Arguments: {}",
8408                        id,
8409                        tool.name(),
8410                        tool.description(),
8411                        args_desc
8412                    )
8413                })
8414            })
8415            .collect();
8416
8417        let tools_section = if tool_descriptions.is_empty() {
8418            "Available tools: none".to_string()
8419        } else {
8420            format!("Available tools:\n{}", tool_descriptions.join("\n"))
8421        };
8422
8423        let skills_section = if available_skills.is_empty() {
8424            "Available skills: none".to_string()
8425        } else {
8426            format!("Available skills: {}", available_skills.join(", "))
8427        };
8428
8429        let prompt = format!(
8430            r#"Create a step-by-step plan to accomplish this goal.
8431
8432Goal: "{}"
8433
8434{}
8435
8436{}
8437
8438Create a plan with clear steps. For each step, specify:
8439- description: What this step accomplishes
8440- action_type: "tool", "skill", "think", or "respond"
8441- action_target: The tool/skill id (if applicable)
8442- args: The arguments object matching the tool's schema (if action_type is "tool")
8443- dependencies: List of step IDs this depends on (empty if none)
8444
8445Respond in JSON format:
8446{{
8447  "steps": [
8448    {{"id": "step1", "description": "...", "action_type": "tool", "action_target": "tool_id", "args": {{"required_field": "value"}}, "dependencies": []}},
8449    {{"id": "step2", "description": "...", "action_type": "think", "action_target": "...", "dependencies": ["step1"]}}
8450  ]
8451}}"#,
8452            input, tools_section, skills_section,
8453        );
8454
8455        let messages = vec![ChatMessage::user(&prompt)];
8456        let response = self
8457            .observe_purpose(
8458                ObservationPurpose::PlanGeneration,
8459                planner_llm.complete(&messages, None),
8460            )
8461            .await
8462            .map_err(|e| AgentError::LLM(format!("Planning failed: {}", e)))?;
8463
8464        let mut plan = Plan::new(input);
8465
8466        if let Some(json_start) = response.content.find('{')
8467            && let Some(json_end) = response.content.rfind('}')
8468        {
8469            let json_str = &response.content[json_start..=json_end];
8470            if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(json_str)
8471                && let Some(steps) = parsed.get("steps").and_then(|s| s.as_array())
8472            {
8473                for step_value in steps {
8474                    let id = step_value
8475                        .get("id")
8476                        .and_then(|v| v.as_str())
8477                        .unwrap_or("step");
8478                    let desc = step_value
8479                        .get("description")
8480                        .and_then(|v| v.as_str())
8481                        .unwrap_or("");
8482                    let action_type = step_value
8483                        .get("action_type")
8484                        .and_then(|v| v.as_str())
8485                        .unwrap_or("think");
8486                    let action_target = step_value
8487                        .get("action_target")
8488                        .and_then(|v| v.as_str())
8489                        .unwrap_or("");
8490                    let args = step_value
8491                        .get("args")
8492                        .cloned()
8493                        .unwrap_or(serde_json::json!({}));
8494                    let deps: Vec<String> = step_value
8495                        .get("dependencies")
8496                        .and_then(|v| v.as_array())
8497                        .map(|arr| {
8498                            arr.iter()
8499                                .filter_map(|v| v.as_str().map(String::from))
8500                                .collect()
8501                        })
8502                        .unwrap_or_default();
8503
8504                    let action = match action_type {
8505                        "tool" => PlanAction::tool(action_target, args),
8506                        "skill" => PlanAction::skill(action_target),
8507                        "respond" => PlanAction::respond(action_target),
8508                        _ => PlanAction::think(desc),
8509                    };
8510
8511                    let step = PlanStep::new(desc, action)
8512                        .with_id(id)
8513                        .with_dependencies(deps);
8514                    plan.add_step(step);
8515                }
8516            }
8517        }
8518
8519        if plan.steps.is_empty() {
8520            plan.add_step(PlanStep::new(
8521                "Process the request",
8522                PlanAction::think(input),
8523            ));
8524            plan.add_step(PlanStep::new(
8525                "Provide response",
8526                PlanAction::respond("Answer based on analysis"),
8527            ));
8528        }
8529
8530        Ok(plan)
8531    }
8532
8533    async fn execute_plan(&self, plan: &mut Plan) -> Result<String> {
8534        let llm = self.get_state_llm()?;
8535        let mut results: HashMap<String, serde_json::Value> = HashMap::new();
8536        let effective = self.get_effective_reasoning_config();
8537        let max_steps = effective.get_planning().map(|c| c.max_steps).unwrap_or(10);
8538
8539        plan.status = PlanStatus::InProgress;
8540
8541        for step_idx in 0..plan.steps.len().min(max_steps as usize) {
8542            let step = &plan.steps[step_idx];
8543
8544            let deps_satisfied = step.dependencies.iter().all(|dep| {
8545                plan.steps
8546                    .iter()
8547                    .find(|s| &s.id == dep)
8548                    .map(|s| s.status.is_completed())
8549                    .unwrap_or(false)
8550            });
8551
8552            if !deps_satisfied {
8553                continue;
8554            }
8555
8556            plan.steps[step_idx].mark_running();
8557
8558            let result = match &plan.steps[step_idx].action {
8559                PlanAction::Tool { tool, args } => {
8560                    // When a tool step has dependency results, ask the LLM to
8561                    // produce the correct arguments given the context and tool schema.
8562                    // This avoids brittle {{stepN}} template substitution and lets
8563                    // the LLM handle type adaptation (e.g. picking the iso field
8564                    // from a datetime result for a downstream format call).
8565                    let has_dep_results = plan.steps[step_idx]
8566                        .dependencies
8567                        .iter()
8568                        .any(|dep| results.contains_key(dep));
8569
8570                    let final_args = if has_dep_results {
8571                        let dep_context: String = plan.steps[step_idx]
8572                            .dependencies
8573                            .iter()
8574                            .filter_map(|dep| results.get(dep).map(|r| format!("{}: {}", dep, r)))
8575                            .collect::<Vec<_>>()
8576                            .join("\n");
8577
8578                        let tool_schema = self
8579                            .tools
8580                            .get(tool)
8581                            .map(|t| {
8582                                let schema = t.input_schema();
8583                                let props = schema
8584                                    .get("properties")
8585                                    .and_then(|p| serde_json::to_string(p).ok())
8586                                    .unwrap_or_else(|| "{}".to_string());
8587                                format!(
8588                                    "{}: {}\nArguments schema: {}",
8589                                    t.id(),
8590                                    t.description(),
8591                                    props
8592                                )
8593                            })
8594                            .unwrap_or_default();
8595
8596                        let step_desc = &plan.steps[step_idx].description;
8597                        let arg_prompt = format!(
8598                            "Generate the JSON arguments for a tool call.\n\n\
8599                             Tool: {}\n\n\
8600                             Task: {}\n\n\
8601                             Previous step results:\n{}\n\n\
8602                             Planner's draft arguments: {}\n\n\
8603                             Produce ONLY a valid JSON object with the correct argument values.\n\
8604                             Use actual values from the previous step results, not template references.",
8605                            tool_schema,
8606                            step_desc,
8607                            dep_context,
8608                            serde_json::to_string(args).unwrap_or_default()
8609                        );
8610                        let messages = vec![ChatMessage::user(&arg_prompt)];
8611                        match self
8612                            .observe_purpose(
8613                                ObservationPurpose::PlanStep,
8614                                llm.complete(&messages, None),
8615                            )
8616                            .await
8617                        {
8618                            Ok(resp) => {
8619                                let content = resp.content.trim();
8620                                // Parse the LLM's JSON response, fall back to planner args.
8621                                let json_start = content.find('{');
8622                                let json_end = content.rfind('}');
8623                                if let (Some(start), Some(end)) = (json_start, json_end) {
8624                                    serde_json::from_str(&content[start..=end])
8625                                        .unwrap_or_else(|_| args.clone())
8626                                } else {
8627                                    args.clone()
8628                                }
8629                            }
8630                            Err(_) => args.clone(),
8631                        }
8632                    } else {
8633                        args.clone()
8634                    };
8635
8636                    let request = ToolExecutionRequest::new(
8637                        uuid::Uuid::new_v4().to_string(),
8638                        tool.clone(),
8639                        final_args,
8640                        ToolCallSource::Plan {
8641                            step_index: step_idx,
8642                        },
8643                    );
8644                    match self.execute_tool_record(request).await {
8645                        Ok(record) if record.success => {
8646                            serde_json::json!({ "output": record.model_output_string() })
8647                        }
8648                        Ok(record) => {
8649                            plan.steps[step_idx].mark_failed(record.model_output_string());
8650                            continue;
8651                        }
8652                        Err(e) => {
8653                            plan.steps[step_idx].mark_failed(e.to_string());
8654                            continue;
8655                        }
8656                    }
8657                }
8658                PlanAction::Skill { skill } => {
8659                    if let Some(skill_def) = self.skills.iter().find(|s| &s.id == skill) {
8660                        if let Some(ref executor) = self.skill_executor {
8661                            match executor
8662                                .execute_with_invoker(skill_def, "", serde_json::json!({}), self)
8663                                .await
8664                            {
8665                                Ok(output) => serde_json::json!({ "output": output }),
8666                                Err(e) => {
8667                                    plan.steps[step_idx].mark_failed(e.to_string());
8668                                    continue;
8669                                }
8670                            }
8671                        } else {
8672                            serde_json::json!({ "output": "Skill executor not available" })
8673                        }
8674                    } else {
8675                        plan.steps[step_idx].mark_failed("Skill not found");
8676                        continue;
8677                    }
8678                }
8679                PlanAction::Think { prompt } => {
8680                    let context: String = results
8681                        .iter()
8682                        .map(|(k, v)| format!("{}: {}", k, v))
8683                        .collect::<Vec<_>>()
8684                        .join("\n");
8685
8686                    let think_prompt = format!("Context:\n{}\n\nTask: {}", context, prompt);
8687                    let messages = vec![ChatMessage::user(&think_prompt)];
8688
8689                    match self
8690                        .observe_purpose(
8691                            ObservationPurpose::PlanStep,
8692                            llm.complete(&messages, None),
8693                        )
8694                        .await
8695                    {
8696                        Ok(resp) => serde_json::json!({ "output": resp.content }),
8697                        Err(e) => {
8698                            plan.steps[step_idx].mark_failed(e.to_string());
8699                            continue;
8700                        }
8701                    }
8702                }
8703                PlanAction::Respond { template } => {
8704                    let context: String = results
8705                        .iter()
8706                        .map(|(k, v)| format!("{}: {}", k, v))
8707                        .collect::<Vec<_>>()
8708                        .join("\n");
8709
8710                    let respond_prompt = format!(
8711                        "Based on this context:\n{}\n\nGenerate a response following this template/instruction: {}",
8712                        context, template
8713                    );
8714                    let messages = vec![ChatMessage::user(&respond_prompt)];
8715
8716                    match self
8717                        .observe_purpose(
8718                            ObservationPurpose::PlanStep,
8719                            llm.complete(&messages, None),
8720                        )
8721                        .await
8722                    {
8723                        Ok(resp) => serde_json::json!({ "output": resp.content }),
8724                        Err(e) => {
8725                            plan.steps[step_idx].mark_failed(e.to_string());
8726                            continue;
8727                        }
8728                    }
8729                }
8730            };
8731
8732            results.insert(plan.steps[step_idx].id.clone(), result.clone());
8733            plan.steps[step_idx].mark_completed(Some(result));
8734        }
8735
8736        // Set plan status based on whether any steps actually failed.
8737        let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
8738        if has_failures {
8739            let failed_ids: Vec<String> = plan
8740                .steps
8741                .iter()
8742                .filter(|s| s.status.is_failed())
8743                .map(|s| s.id.clone())
8744                .collect();
8745            plan.status = PlanStatus::Failed {
8746                error: format!("Steps failed: {}", failed_ids.join(", ")),
8747            };
8748        } else {
8749            plan.status = PlanStatus::Completed;
8750        }
8751
8752        // Synthesize final output from all completed step results.
8753        let all_outputs: Vec<String> = plan
8754            .steps
8755            .iter()
8756            .filter(|s| s.status.is_completed())
8757            .filter_map(|s| {
8758                s.result
8759                    .as_ref()
8760                    .and_then(|r| r.get("output"))
8761                    .and_then(|o| o.as_str())
8762                    .map(|o| format!("{}: {}", s.description, o))
8763            })
8764            .collect();
8765
8766        if all_outputs.is_empty() {
8767            return Ok("Plan execution completed but produced no results.".to_string());
8768        }
8769
8770        if all_outputs.len() == 1 {
8771            return Ok(all_outputs.into_iter().next().unwrap());
8772        }
8773
8774        // Synthesize a coherent summary from multiple step results via LLM.
8775        let context = all_outputs.join("\n\n");
8776        let prompt = format!(
8777            "You completed a multi-step plan for: \"{}\"\n\nStep results:\n{}\n\nProvide a coherent final response that synthesizes these results.",
8778            plan.goal, context
8779        );
8780        let messages = vec![ChatMessage::user(&prompt)];
8781        match self
8782            .observe_purpose(ObservationPurpose::PlanStep, llm.complete(&messages, None))
8783            .await
8784        {
8785            Ok(resp) => Ok(resp.content.trim().to_string()),
8786            Err(_) => Ok(context),
8787        }
8788    }
8789
8790    async fn evaluate_response(&self, input: &str, response: &str) -> Result<EvaluationResult> {
8791        let effective_config = self.get_effective_reflection_config();
8792        self.evaluate_response_with_config(input, response, &effective_config)
8793            .await
8794    }
8795
8796    fn extract_thinking(&self, content: &str) -> (Option<String>, String) {
8797        if let Some(start) = content.find("<thinking>")
8798            && let Some(end) = content.find("</thinking>")
8799        {
8800            let thinking = content[start + 10..end].trim().to_string();
8801            let answer = content[end + 11..].trim().to_string();
8802            return (Some(thinking), answer);
8803        }
8804        (None, content.to_string())
8805    }
8806
8807    fn format_response_with_thinking(&self, thinking: Option<&str>, answer: &str) -> String {
8808        match self.get_effective_reasoning_config().output {
8809            ReasoningOutput::Hidden => answer.to_string(),
8810            ReasoningOutput::Visible => {
8811                if let Some(t) = thinking {
8812                    format!("Thinking:\n{}\n\nAnswer:\n{}", t, answer)
8813                } else {
8814                    answer.to_string()
8815                }
8816            }
8817            ReasoningOutput::Tagged => {
8818                if let Some(t) = thinking {
8819                    format!("<thinking>{}</thinking>\n{}", t, answer)
8820                } else {
8821                    answer.to_string()
8822                }
8823            }
8824        }
8825    }
8826
8827    //
8828    // Blocking disambiguation returns every clarification or required confirmation as its own root response before redispatch.
8829    // The manager retains resolved input and pending skill ownership until a later turn explicitly confirms it.
8830    //
8831    async fn run_loop(&self, input: &str) -> Result<AgentResponse> {
8832        //
8833        // Blocking execution must fail before a turn starts when required persistence is unavailable.
8834        //
8835        self.init_storage().await?;
8836        self.begin_root_turn();
8837        let _root_cleanup = RootTurnCleanup::new(self);
8838        info!(input_len = input.len(), "Starting chat");
8839
8840        self.hooks.on_message_received(input).await;
8841
8842        // One-shot context initialization: load runtime defaults, resolve env vars,
8843        // populate builtin sources (session, agent), etc.  This must happen before
8844        // the first template render so that {{ context.* }} variables are available.
8845        if !self.context_initialized.swap(true, Ordering::SeqCst) {
8846            self.context_manager.initialize().await?;
8847            debug!("Context manager initialized (defaults, env, builtins)");
8848        }
8849
8850        self.check_turn_timeout().await?;
8851        self.context_manager.refresh_per_turn().await?;
8852
8853        // Clear stale disambiguation context from previous turns.
8854        // This prevents resolved_intent from leaking across turns and causing incorrect deterministic routing on subsequent inputs.
8855        self.clear_disambiguation_context();
8856
8857        // Disambiguation check (before input processing)
8858        if let Some(ref disambiguator) = self.disambiguation_manager {
8859            let disambiguation_context = self.build_disambiguation_context().await?;
8860
8861            // Get state-level disambiguation override
8862            let state_override = self
8863                .state_machine
8864                .as_ref()
8865                .and_then(|sm| sm.current_definition())
8866                .and_then(|def| def.disambiguation.clone());
8867
8868            let state_generation = self
8869                .state_machine
8870                .as_ref()
8871                .map(|state_machine| state_machine.generation());
8872            let disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
8873            let mut disambiguation_result = self
8874                .observe_purpose(
8875                    ObservationPurpose::DisambiguationDetection,
8876                    disambiguator.process_input_with_override(
8877                        input,
8878                        &disambiguation_context,
8879                        state_override.as_ref(),
8880                        None,
8881                    ),
8882                )
8883                .await?;
8884            let current_state_generation = self
8885                .state_machine
8886                .as_ref()
8887                .map(|state_machine| state_machine.generation());
8888            if current_state_generation != state_generation
8889                || self.disambiguation_epoch.load(Ordering::SeqCst) != disambiguation_epoch
8890            {
8891                disambiguator.clear_pending().await;
8892                *self.pending_skill_id.write() = None;
8893                disambiguation_result = DisambiguationResult::Abandoned { new_input: None };
8894                info!(
8895                    confirmation_event = "invalidated",
8896                    invalidation_reason = "state_generation_changed",
8897                    "Disambiguation result invalidated before redispatch"
8898                );
8899            }
8900            match disambiguation_result {
8901                DisambiguationResult::Clear => {
8902                    debug!("Input is clear, proceeding normally");
8903                }
8904                DisambiguationResult::NeedsClarification {
8905                    question,
8906                    detection,
8907                } => {
8908                    let admission = self
8909                        .admit_disambiguation_redispatch(disambiguation_epoch, state_generation)
8910                        .await?;
8911                    let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
8912                    info!(
8913                        ambiguity_type = ?detection.ambiguity_type,
8914                        confidence = detection.confidence,
8915                        "Input requires clarification"
8916                    );
8917
8918                    // This branch also owns post-resolution confirmation questions.
8919                    // Do not clear pending_skill_id or redispatch until the manager returns Clarified on a later turn.
8920                    self.commit_root_user_message(input).await?;
8921                    self.memory
8922                        .add_message(ChatMessage::assistant(&question.question))
8923                        .await?;
8924
8925                    let status = if awaiting_confirmation {
8926                        "awaiting_confirmation"
8927                    } else {
8928                        "awaiting_clarification"
8929                    };
8930                    let response = AgentResponse::new(&question.question).with_metadata(
8931                        "disambiguation",
8932                        serde_json::json!({
8933                            "status": status,
8934                            "options": question.options,
8935                            "clarifying": question.clarifying,
8936                            "detection": {
8937                                "type": detection.ambiguity_type,
8938                                "confidence": detection.confidence,
8939                                "what_is_unclear": detection.what_is_unclear,
8940                            }
8941                        }),
8942                    );
8943                    drop(admission);
8944                    self.finish_turn_if_root(&response).await?;
8945                    return Ok(response);
8946                }
8947                DisambiguationResult::Clarified {
8948                    enriched_input,
8949                    resolved,
8950                    ..
8951                } => {
8952                    let admission = match self
8953                        .admit_disambiguation_redispatch(disambiguation_epoch, state_generation)
8954                        .await
8955                    {
8956                        Ok(admission) => admission,
8957                        Err(error) => {
8958                            *self.pending_skill_id.write() = None;
8959                            return Err(error);
8960                        }
8961                    };
8962                    info!(
8963                        resolved_count = resolved.len(),
8964                        enriched = %enriched_input,
8965                        "Input clarified, injecting resolved intent into context"
8966                    );
8967
8968                    // Routing uses `resolved` (structured, deterministic)
8969                    // This is what makes post-disambiguation routing DETERMINISTIC
8970                    for (key, value) in &resolved {
8971                        let context_key = format!("disambiguation.{}", key);
8972                        let _ = self.context_manager.set(&context_key, value.clone());
8973                    }
8974
8975                    if let Some(intent) = resolved.get("intent") {
8976                        let _ = self.context_manager.set("resolved_intent", intent.clone());
8977                    }
8978
8979                    let _ = self
8980                        .context_manager
8981                        .set("disambiguation.resolved", serde_json::Value::Bool(true));
8982
8983                    // Check if this clarification was triggered by a skill-level override.
8984                    // If so, route directly to the matched skill instead of going through
8985                    // skill routing again (which might match a different skill).
8986                    let skill_id = self.pending_skill_id.read().clone();
8987                    if let Some(skill_id) = skill_id {
8988                        info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input");
8989                        drop(admission);
8990                        return self
8991                            .recheck_skill_disambiguation(
8992                                &skill_id,
8993                                &enriched_input,
8994                                disambiguation_epoch,
8995                                state_generation,
8996                            )
8997                            .await;
8998                    }
8999
9000                    drop(admission);
9001                    return self.run_loop_internal(&enriched_input).await;
9002                }
9003                DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
9004                    info!("Proceeding with best guess interpretation");
9005
9006                    // Same skill-id re-check for best-guess path
9007                    let skill_id = self.pending_skill_id.read().clone();
9008                    if let Some(skill_id) = skill_id {
9009                        info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input");
9010                        return self
9011                            .recheck_skill_disambiguation(
9012                                &skill_id,
9013                                &enriched_input,
9014                                disambiguation_epoch,
9015                                state_generation,
9016                            )
9017                            .await;
9018                    }
9019
9020                    return self.run_loop_internal(&enriched_input).await;
9021                }
9022                DisambiguationResult::GiveUp { reason } => {
9023                    *self.pending_skill_id.write() = None;
9024                    warn!(reason = %reason, "Disambiguation gave up");
9025                    let apology = self
9026                        .generate_localized_apology(
9027                            "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
9028                            &reason,
9029                        )
9030                        .await
9031                        .unwrap_or_else(|_| {
9032                            format!("I'm sorry, I couldn't understand your request: {}", reason)
9033                        });
9034                    let response = AgentResponse::new(&apology);
9035                    self.finish_turn_if_root(&response).await?;
9036                    return Ok(response);
9037                }
9038                DisambiguationResult::Escalate { reason } => {
9039                    *self.pending_skill_id.write() = None;
9040                    info!(reason = %reason, "Escalating to human");
9041                    if let Some(ref hitl) = self.hitl_engine {
9042                        let trigger =
9043                            ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
9044                        let mut context_map = HashMap::new();
9045                        context_map.insert("original_input".to_string(), serde_json::json!(input));
9046                        context_map.insert("reason".to_string(), serde_json::json!(&reason));
9047                        let check_result = HITLCheckResult::required(
9048                            trigger,
9049                            context_map,
9050                            format!("User request needs human assistance: {}", reason),
9051                            Some(hitl.config().default_timeout_seconds),
9052                        );
9053                        let result = self.request_hitl_approval(check_result).await?;
9054                        if matches!(
9055                            result,
9056                            ApprovalResult::Approved | ApprovalResult::Modified { .. }
9057                        ) {
9058                            return self.run_loop_internal(input).await;
9059                        }
9060                    }
9061                    let apology = self
9062                        .generate_localized_apology(
9063                            "Explain briefly that you're transferring the user to a human agent for help.",
9064                            &reason,
9065                        )
9066                        .await
9067                        .unwrap_or_else(|_| {
9068                            format!("I need human assistance to help with your request: {}", reason)
9069                        });
9070                    let response = AgentResponse::new(&apology);
9071                    self.finish_turn_if_root(&response).await?;
9072                    return Ok(response);
9073                }
9074                DisambiguationResult::Abandoned { new_input } => {
9075                    *self.pending_skill_id.write() = None;
9076
9077                    info!(
9078                        has_new_input = new_input.is_some(),
9079                        "Clarification abandoned by user"
9080                    );
9081
9082                    self.commit_root_user_message(input).await?;
9083
9084                    match new_input {
9085                        Some(fresh_input) => {
9086                            // Topic switch: process the user's new input from scratch.
9087                            // The LLM sees full conversation context including the abandoned exchange.
9088                            return self.run_loop_internal(&fresh_input).await;
9089                        }
9090                        None => {
9091                            // Pure abandonment: generate a brief acknowledgment.
9092                            let ack = self
9093                                .generate_localized_apology(
9094                                    "The user changed their mind about their previous request. \
9095                                     Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
9096                                     Do NOT apologize excessively. Be concise.",
9097                                    "User abandoned clarification",
9098                                )
9099                                .await
9100                                .unwrap_or_else(|_| {
9101                                    "OK, no problem. What else can I help with?".to_string()
9102                                });
9103
9104                            self.memory
9105                                .add_message(ChatMessage::assistant(&ack))
9106                                .await?;
9107
9108                            let response = AgentResponse::new(&ack);
9109                            self.finish_turn_if_root(&response).await?;
9110                            return Ok(response);
9111                        }
9112                    }
9113                }
9114            }
9115        }
9116
9117        self.run_loop_internal(input).await
9118    }
9119
9120    /// Generate a localized response using the router LLM
9121    async fn generate_localized_apology(&self, instruction: &str, reason: &str) -> Result<String> {
9122        let llm = self.llm_registry.router().map_err(|e| {
9123            AgentError::LLM(format!(
9124                "Router LLM not available for localized response: {}",
9125                e
9126            ))
9127        })?;
9128
9129        let recent: Vec<String> = self
9130            .memory
9131            .get_messages(Some(3))
9132            .await?
9133            .iter()
9134            .map(|m| m.content.clone())
9135            .collect();
9136
9137        let context_hint = if recent.is_empty() {
9138            String::new()
9139        } else {
9140            format!(
9141                "\nRecent conversation (detect the user's language from this):\n{}\n",
9142                recent.join("\n")
9143            )
9144        };
9145
9146        let prompt = format!(
9147            "{}\nReason: {}\n{}Respond in the same language as the user. Output ONLY the message, nothing else.",
9148            instruction, reason, context_hint
9149        );
9150
9151        let messages = vec![ChatMessage::user(&prompt)];
9152        let response = self
9153            .observe_purpose(
9154                ObservationPurpose::DisambiguationClarification,
9155                llm.complete(&messages, None),
9156            )
9157            .await
9158            .map_err(|e| AgentError::LLM(format!("Localized response generation failed: {}", e)))?;
9159
9160        Ok(response.content.trim().to_string())
9161    }
9162
9163    /// Clear disambiguation-related keys from the context manager.
9164    ///
9165    /// Render template variables in state action args using the context manager.
9166    fn render_action_args(&self, args: &Value) -> Value {
9167        let context = self.build_context_with_overlays();
9168        match args {
9169            Value::Object(map) => {
9170                let mut rendered = serde_json::Map::new();
9171                for (k, v) in map {
9172                    match v {
9173                        Value::String(s) if s.contains("{{") => {
9174                            match self.template_renderer.render(s, &context) {
9175                                Ok(rendered_str) => {
9176                                    rendered.insert(k.clone(), Value::String(rendered_str));
9177                                }
9178                                Err(_) => {
9179                                    rendered.insert(k.clone(), v.clone());
9180                                }
9181                            }
9182                        }
9183                        _ => {
9184                            rendered.insert(k.clone(), v.clone());
9185                        }
9186                    }
9187                }
9188                Value::Object(rendered)
9189            }
9190            _ => args.clone(),
9191        }
9192    }
9193
9194    /// Called at the start of each turn to prevent stale `resolved_intent` from leaking across turns.
9195    fn clear_disambiguation_context(&self) {
9196        let _ = self
9197            .context_manager
9198            .set("resolved_intent", serde_json::Value::Null);
9199
9200        let all = self.context_manager.get_all();
9201        for key in all.keys() {
9202            if key.starts_with("disambiguation.") {
9203                let _ = self.context_manager.set(key, serde_json::Value::Null);
9204            }
9205        }
9206    }
9207
9208    /// Re-run skill disambiguation on enriched input before executing the skill.
9209    /// After clarification resolves, the enriched input may still be missing required_clarity fields (e.g. "Transfer money to Jane." still lacks amount).
9210    /// This method re-runs the skill's disambiguation pass.
9211    /// If fields are still missing, it returns the new clarification question and keeps pending_skill_id set.
9212    /// If all fields are present (Clear), it executes the skill and returns the response.
9213    async fn recheck_skill_disambiguation(
9214        &self,
9215        skill_id: &str,
9216        enriched_input: &str,
9217        expected_disambiguation_epoch: u64,
9218        expected_state_generation: Option<u64>,
9219    ) -> Result<AgentResponse> {
9220        let skill = self
9221            .skill_router
9222            .as_ref()
9223            .and_then(|r| r.get_skill(skill_id).cloned());
9224
9225        // If the skill has disambiguation enabled, re-run it on the enriched input.
9226        if let Some(ref skill) = skill
9227            && let Some(ref skill_disambig) = skill.disambiguation
9228            && skill_disambig.enabled.unwrap_or(false)
9229            && let Some(ref disambiguator) = self.disambiguation_manager
9230        {
9231            let context = self.build_disambiguation_context().await?;
9232            let state_override = self
9233                .state_machine
9234                .as_ref()
9235                .and_then(|sm| sm.current_definition())
9236                .and_then(|def| def.disambiguation.clone());
9237
9238            let disambiguation_result = self
9239                .observe_purpose(
9240                    ObservationPurpose::DisambiguationDetection,
9241                    disambiguator.process_input_with_override(
9242                        enriched_input,
9243                        &context,
9244                        state_override.as_ref(),
9245                        Some(skill_disambig),
9246                    ),
9247                )
9248                .await?;
9249            let current_state_generation = self
9250                .state_machine
9251                .as_ref()
9252                .map(|state_machine| state_machine.generation());
9253            if current_state_generation != expected_state_generation
9254                || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
9255            {
9256                disambiguator.clear_pending().await;
9257                *self.pending_skill_id.write() = None;
9258                return Err(AgentError::Other(
9259                    "State or reset ownership changed during skill disambiguation recheck"
9260                        .to_string(),
9261                ));
9262            }
9263            match disambiguation_result {
9264                DisambiguationResult::Clear => {
9265                    debug!(skill_id = %skill_id, "Skill re-check: all fields present");
9266                }
9267                DisambiguationResult::NeedsClarification {
9268                    question,
9269                    detection,
9270                } => {
9271                    let admission = self
9272                        .admit_disambiguation_redispatch(
9273                            expected_disambiguation_epoch,
9274                            expected_state_generation,
9275                        )
9276                        .await?;
9277                    let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
9278                    info!(
9279                        skill_id = %skill_id,
9280                        ambiguity_type = ?detection.ambiguity_type,
9281                        what_is_unclear = ?detection.what_is_unclear,
9282                        "Skill re-check: still missing fields, asking again"
9283                    );
9284                    // Keep pending_skill_id set (do NOT clear it).
9285                    // The next turn will resolve this new clarification and
9286                    // re-enter this method until all fields are present.
9287                    self.memory
9288                        .add_message(ChatMessage::user(enriched_input))
9289                        .await?;
9290                    self.memory
9291                        .add_message(ChatMessage::assistant(&question.question))
9292                        .await?;
9293
9294                    let response = AgentResponse::new(&question.question).with_metadata(
9295                        "disambiguation",
9296                        serde_json::json!({
9297                            "status": if awaiting_confirmation { "awaiting_confirmation" } else { "awaiting_clarification" },
9298                            "skill_id": skill_id,
9299                            "options": question.options,
9300                            "clarifying": question.clarifying,
9301                            "detection": {
9302                                "type": detection.ambiguity_type,
9303                                "confidence": detection.confidence,
9304                                "what_is_unclear": detection.what_is_unclear,
9305                            }
9306                        }),
9307                    );
9308                    drop(admission);
9309                    self.finish_turn_if_root(&response).await?;
9310                    return Ok(response);
9311                }
9312                DisambiguationResult::Clarified {
9313                    enriched_input: re_enriched,
9314                    ..
9315                } => {
9316                    debug!(skill_id = %skill_id, "Skill re-check: clarified immediately, executing");
9317                    let admission = self
9318                        .admit_disambiguation_redispatch(
9319                            expected_disambiguation_epoch,
9320                            expected_state_generation,
9321                        )
9322                        .await?;
9323                    *self.pending_skill_id.write() = None;
9324                    drop(admission);
9325                    let skill_response = self.execute_skill_by_id(skill_id, &re_enriched).await?;
9326                    self.memory
9327                        .add_message(ChatMessage::user(&re_enriched))
9328                        .await?;
9329                    return self
9330                        .handle_skill_response(
9331                            &re_enriched,
9332                            skill_id,
9333                            skill_response,
9334                            &HashMap::new(),
9335                        )
9336                        .await;
9337                }
9338                DisambiguationResult::ProceedWithBestGuess {
9339                    enriched_input: re_enriched,
9340                } => {
9341                    debug!(skill_id = %skill_id, "Skill re-check: proceeding with best guess");
9342                    let admission = self
9343                        .admit_disambiguation_redispatch(
9344                            expected_disambiguation_epoch,
9345                            expected_state_generation,
9346                        )
9347                        .await?;
9348                    *self.pending_skill_id.write() = None;
9349                    drop(admission);
9350                    let skill_response = self.execute_skill_by_id(skill_id, &re_enriched).await?;
9351                    self.memory
9352                        .add_message(ChatMessage::user(&re_enriched))
9353                        .await?;
9354                    return self
9355                        .handle_skill_response(
9356                            &re_enriched,
9357                            skill_id,
9358                            skill_response,
9359                            &HashMap::new(),
9360                        )
9361                        .await;
9362                }
9363                DisambiguationResult::GiveUp { reason } => {
9364                    *self.pending_skill_id.write() = None;
9365                    let apology = self
9366                                    .generate_localized_apology(
9367                                        "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
9368                                        &reason,
9369                                    )
9370                                    .await
9371                                    .unwrap_or_else(|_| {
9372                                        format!("I'm sorry, I couldn't understand your request: {}", reason)
9373                                    });
9374                    let response = AgentResponse::new(&apology);
9375                    self.finish_turn_if_root(&response).await?;
9376                    return Ok(response);
9377                }
9378                DisambiguationResult::Escalate { reason } => {
9379                    *self.pending_skill_id.write() = None;
9380                    let apology = self
9381                                    .generate_localized_apology(
9382                                        "Explain briefly that you're transferring the user to a human agent for help.",
9383                                        &reason,
9384                                    )
9385                                    .await
9386                                    .unwrap_or_else(|_| {
9387                                        format!("I need human assistance to help with your request: {}", reason)
9388                                    });
9389                    let response = AgentResponse::new(&apology);
9390                    self.finish_turn_if_root(&response).await?;
9391                    return Ok(response);
9392                }
9393                DisambiguationResult::Abandoned { new_input } => {
9394                    // User abandoned during skill re-check.
9395                    // Clear skill routing state and fall through to normal execution.
9396                    *self.pending_skill_id.write() = None;
9397                    debug!(skill_id = %skill_id, "Skill re-check: abandoned by user");
9398                    if let Some(fresh) = new_input {
9399                        return self.run_loop_internal(&fresh).await;
9400                    }
9401                    let ack = self
9402                                    .generate_localized_apology(
9403                                        "The user changed their mind about their previous request. \
9404                                         Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
9405                                         Do NOT apologize excessively. Be concise.",
9406                                        "User abandoned clarification",
9407                                    )
9408                                    .await
9409                                    .unwrap_or_else(|_| {
9410                                        "OK, no problem. What else can I help with?".to_string()
9411                                    });
9412                    self.memory
9413                        .add_message(ChatMessage::assistant(&ack))
9414                        .await?;
9415                    let response = AgentResponse::new(&ack);
9416                    self.finish_turn_if_root(&response).await?;
9417                    return Ok(response);
9418                }
9419            }
9420        }
9421
9422        // Skill execution is admitted before the read guard is released so later invalidation cannot retroactively cancel it.
9423        let admission = self
9424            .admit_disambiguation_redispatch(
9425                expected_disambiguation_epoch,
9426                expected_state_generation,
9427            )
9428            .await?;
9429        *self.pending_skill_id.write() = None;
9430        drop(admission);
9431        let skill_response = self.execute_skill_by_id(skill_id, enriched_input).await?;
9432        self.memory
9433            .add_message(ChatMessage::user(enriched_input))
9434            .await?;
9435        self.handle_skill_response(enriched_input, skill_id, skill_response, &HashMap::new())
9436            .await
9437    }
9438
9439    /// Handle skill routing result: output processing, memory, transitions.
9440    /// Returns a fully formed AgentResponse for skill-routed requests.
9441    async fn handle_skill_response(
9442        &self,
9443        processed_input: &str,
9444        skill_id: &str,
9445        skill_response: String,
9446        input_context: &HashMap<String, Value>,
9447    ) -> Result<AgentResponse> {
9448        let output_data = self.process_output(&skill_response, input_context).await?;
9449        let final_response = output_data.content;
9450
9451        self.memory
9452            .add_message(ChatMessage::assistant(&final_response))
9453            .await?;
9454
9455        self.check_memory_compression().await?;
9456
9457        self.increment_turn();
9458        self.evaluate_transitions(processed_input, &final_response)
9459            .await?;
9460
9461        let response = AgentResponse::new(final_response)
9462            .with_metadata("skill_id", serde_json::json!(skill_id));
9463        self.finish_turn_if_root(&response).await?;
9464        Ok(response)
9465    }
9466
9467    /// Run the Plan-and-Execute flow: generate plan, execute steps, finalize.
9468    /// Supports plan-level reflection with replan loop when configured.
9469    async fn handle_plan_and_execute(
9470        &self,
9471        processed_input: &str,
9472        input_context: &HashMap<String, Value>,
9473        auto_detected: bool,
9474    ) -> Result<AgentResponse> {
9475        let effective = self.get_effective_reasoning_config();
9476        let plan_reflection = effective
9477            .get_planning()
9478            .map(|c| c.reflection.clone())
9479            .unwrap_or_default();
9480
9481        let max_attempts = if plan_reflection.enabled {
9482            1 + plan_reflection.max_replans
9483        } else {
9484            1
9485        };
9486
9487        let mut plan = self.generate_plan(processed_input).await?;
9488        info!(
9489            plan_id = %plan.id,
9490            steps = plan.steps.len(),
9491            "Plan generated"
9492        );
9493
9494        let mut plan_result = String::new();
9495
9496        for attempt in 0..max_attempts {
9497            *self.current_plan.write() = Some(plan.clone());
9498            plan_result = self.execute_plan(&mut plan).await?;
9499
9500            info!(
9501                plan_status = ?plan.status,
9502                completed_steps = plan.completed_steps().count(),
9503                attempt = attempt + 1,
9504                "Plan execution completed"
9505            );
9506
9507            if !plan_reflection.enabled {
9508                break;
9509            }
9510
9511            let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
9512            if !has_failures {
9513                break;
9514            }
9515
9516            if attempt + 1 >= max_attempts {
9517                break;
9518            }
9519
9520            match plan_reflection.on_step_failure {
9521                StepFailureAction::Replan => {
9522                    info!(attempt = attempt + 1, "Plan had failures, replanning");
9523                    plan = self.generate_plan(processed_input).await?;
9524                }
9525                StepFailureAction::Abort => {
9526                    warn!("Plan step failed, aborting");
9527                    break;
9528                }
9529                StepFailureAction::Skip | StepFailureAction::Continue => {
9530                    break;
9531                }
9532            }
9533        }
9534
9535        *self.current_plan.write() = Some(plan);
9536
9537        let output_data = self.process_output(&plan_result, input_context).await?;
9538        let final_content = output_data.content;
9539
9540        self.memory
9541            .add_message(ChatMessage::assistant(&final_content))
9542            .await?;
9543
9544        self.check_memory_compression().await?;
9545        self.increment_turn();
9546        self.evaluate_transitions(processed_input, &final_content)
9547            .await?;
9548
9549        let reasoning_metadata =
9550            ReasoningMetadata::new(ReasoningMode::PlanAndExecute).with_auto_detected(auto_detected);
9551
9552        let response = AgentResponse::new(&final_content).with_metadata(
9553            "reasoning",
9554            serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
9555        );
9556
9557        self.finish_turn_if_root(&response).await?;
9558        Ok(response)
9559    }
9560
9561    /// Inject CoT/ReAct reasoning prompt into the system message (first iteration only).
9562    fn inject_reasoning_prompt(
9563        &self,
9564        messages: &mut [ChatMessage],
9565        reasoning_mode: &ReasoningMode,
9566        is_first_iteration: bool,
9567    ) {
9568        if !is_first_iteration {
9569            return;
9570        }
9571        match reasoning_mode {
9572            ReasoningMode::CoT => {
9573                if let Some(msg) = messages.first_mut()
9574                    && matches!(msg.role, ai_agents_core::Role::System)
9575                {
9576                    msg.content = self.build_cot_system_prompt(&msg.content);
9577                    debug!("Applied Chain-of-Thought system prompt");
9578                }
9579            }
9580            ReasoningMode::React => {
9581                if let Some(msg) = messages.first_mut()
9582                    && matches!(msg.role, ai_agents_core::Role::System)
9583                {
9584                    msg.content = self.build_react_system_prompt(&msg.content);
9585                    debug!("Applied ReAct system prompt");
9586                }
9587            }
9588            _ => {}
9589        }
9590    }
9591
9592    //
9593    // Draft generation must not commit user memory or run tools.
9594    // The current user input is added only as an ephemeral message for this LLM call.
9595    //
9596    async fn generate_main_response_draft(
9597        &self,
9598        processed_input: &str,
9599        reasoning_mode: &ReasoningMode,
9600    ) -> Result<MainResponseDraft> {
9601        let llm = self.get_state_llm()?;
9602        let protocol = self.main_tool_protocol(llm.as_ref(), true).await?;
9603        let mut messages = self
9604            .build_messages_internal(false, Some(processed_input), protocol.choice.is_none())
9605            .await?;
9606        self.inject_reasoning_prompt(&mut messages, reasoning_mode, true);
9607        let response = self
9608            .complete_main_llm_with_recovery(llm, &messages, &protocol)
9609            .await?;
9610        let content = response.content.trim().to_string();
9611        let (thinking, answer) = self.extract_thinking(&content);
9612        if let Some(calls) = self.parse_main_tool_calls(&content, &protocol) {
9613            return Ok(MainResponseDraft::ToolCalls {
9614                raw_content: content,
9615                calls,
9616                thinking,
9617            });
9618        }
9619        Ok(MainResponseDraft::Text {
9620            raw_content: answer,
9621            thinking,
9622        })
9623    }
9624
9625    //
9626    // This is the only place where a winning draft is allowed to become runtime state.
9627    // Prompt and native tool calls remain inert until this method commits the draft into the shared executor path.
9628    //
9629    async fn commit_main_response_draft(
9630        &self,
9631        processed_input: &str,
9632        input_context: &HashMap<String, Value>,
9633        draft: MainResponseDraft,
9634        reasoning_mode: ReasoningMode,
9635        auto_detected: bool,
9636    ) -> Result<AgentResponse> {
9637        self.commit_root_user_message(processed_input).await?;
9638        match draft {
9639            MainResponseDraft::Text {
9640                raw_content,
9641                thinking,
9642            } => {
9643                self.finish_text_response_from_model(CommittedTextResponse {
9644                    processed_input,
9645                    input_context,
9646                    answer: raw_content,
9647                    reasoning_mode,
9648                    auto_detected,
9649                    iterations: 1,
9650                    thinking_content: thinking,
9651                    all_tool_calls: Vec::new(),
9652                })
9653                .await
9654            }
9655            MainResponseDraft::ToolCalls {
9656                raw_content,
9657                calls,
9658                thinking: _,
9659            } => {
9660                let mut all_tool_calls = Vec::new();
9661                match self
9662                    .handle_tool_calls(processed_input, &raw_content, calls, &mut all_tool_calls)
9663                    .await?
9664                {
9665                    ToolCallOutcome::Rejected(response) => {
9666                        self.finish_turn_if_root(&response).await?;
9667                        Ok(response)
9668                    }
9669                    ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => {
9670                        self.continue_after_committed_tool_draft(processed_input)
9671                            .await
9672                    }
9673                }
9674            }
9675        }
9676    }
9677
9678    //
9679    // Tool drafts need a committed continuation after function results are written.
9680    // Redispatch depth suppresses duplicate root lifecycle work during that continuation.
9681    //
9682    async fn continue_after_committed_tool_draft(
9683        &self,
9684        processed_input: &str,
9685    ) -> Result<AgentResponse> {
9686        *self.redispatch_depth.write() += 1;
9687        if let Some(context) = self.active_turn_context.write().as_mut() {
9688            context.enter_redispatch();
9689        }
9690        let result = Box::pin(self.run_loop_internal(processed_input)).await;
9691        *self.redispatch_depth.write() -= 1;
9692        if let Some(context) = self.active_turn_context.write().as_mut() {
9693            context.exit_redispatch();
9694        }
9695        let response = result?;
9696        self.finish_turn_if_root(&response).await?;
9697        Ok(response)
9698    }
9699
9700    //
9701    // Shared committed text finalization for normal responses and winning text drafts.
9702    // Keep output processing, reflection, transitions, hooks, and maintenance behind this commit boundary.
9703    //
9704    async fn finish_text_response_from_model(
9705        &self,
9706        response: CommittedTextResponse<'_>,
9707    ) -> Result<AgentResponse> {
9708        let CommittedTextResponse {
9709            processed_input,
9710            input_context,
9711            answer,
9712            reasoning_mode,
9713            auto_detected,
9714            iterations,
9715            thinking_content,
9716            all_tool_calls,
9717        } = response;
9718        let output_data = self.process_output(&answer, input_context).await?;
9719        let mut final_content = if output_data.metadata.rejected {
9720            output_data
9721                .metadata
9722                .rejection_reason
9723                .unwrap_or_else(|| answer.to_string())
9724        } else {
9725            output_data.content
9726        };
9727        let llm = self.get_state_llm()?;
9728        let reflection_metadata;
9729        (final_content, reflection_metadata) = self
9730            .run_reflection(&*llm, processed_input, final_content)
9731            .await?;
9732        final_content =
9733            self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
9734        let final_content = {
9735            let result = self
9736                .post_loop_processing(processed_input, final_content)
9737                .await?;
9738            self.apply_post_loop_result(processed_input, result).await?
9739        };
9740        let response = self.build_agent_response(AgentResponseParts {
9741            content: final_content,
9742            all_tool_calls,
9743            reasoning_mode,
9744            auto_detected,
9745            iterations,
9746            thinking: thinking_content,
9747            reflection_metadata,
9748        });
9749        self.finish_turn_if_root(&response).await?;
9750        Ok(response)
9751    }
9752
9753    //
9754    // Auto reasoning uses this path after the judge wins with a deeper mode.
9755    // It intentionally uses committed message building instead of the draft overlay.
9756    //
9757    async fn run_committed_response_loop_with_reasoning(
9758        &self,
9759        processed_input: &str,
9760        input_context: &HashMap<String, Value>,
9761        reasoning_mode: ReasoningMode,
9762        auto_detected: bool,
9763    ) -> Result<AgentResponse> {
9764        self.commit_root_user_message(processed_input).await?;
9765        let llm = self.get_state_llm()?;
9766        let mut iterations = 0u32;
9767        let mut all_tool_calls = Vec::new();
9768        let mut thinking_content = None;
9769        loop {
9770            let effective_max = if reasoning_mode != ReasoningMode::None {
9771                let rc = self.get_effective_reasoning_config();
9772                self.max_iterations.min(rc.max_iterations)
9773            } else {
9774                self.max_iterations
9775            };
9776            if iterations >= effective_max {
9777                return Err(AgentError::Other(format!(
9778                    "Max iterations ({}) exceeded",
9779                    effective_max
9780                )));
9781            }
9782            iterations += 1;
9783            *self.iteration_count.write() = iterations;
9784            let protocol = self.main_tool_protocol(llm.as_ref(), false).await?;
9785            let mut messages = self
9786                .build_messages_internal(true, None, protocol.choice.is_none())
9787                .await?;
9788            self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
9789            self.hooks.on_llm_start(&messages).await;
9790            let llm_start = Instant::now();
9791            let response = self
9792                .complete_main_llm_with_recovery(Arc::clone(&llm), &messages, &protocol)
9793                .await?;
9794            let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
9795            self.hooks.on_llm_complete(&response, llm_duration_ms).await;
9796            let content = response.content.trim();
9797            if let Some(tool_calls) = self.parse_main_tool_calls(content, &protocol) {
9798                match self
9799                    .handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
9800                    .await?
9801                {
9802                    ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
9803                    ToolCallOutcome::Rejected(resp) => {
9804                        self.finish_turn_if_root(&resp).await?;
9805                        return Ok(resp);
9806                    }
9807                }
9808            }
9809            let (extracted_thinking, answer) = self.extract_thinking(content);
9810            if extracted_thinking.is_some() {
9811                thinking_content = extracted_thinking;
9812            }
9813            return self
9814                .finish_text_response_from_model(CommittedTextResponse {
9815                    processed_input,
9816                    input_context,
9817                    answer,
9818                    reasoning_mode,
9819                    auto_detected,
9820                    iterations,
9821                    thinking_content,
9822                    all_tool_calls,
9823                })
9824                .await;
9825        }
9826    }
9827
9828    /// Handle tool calls: check transitions, execute tools in parallel, handle HITL rejection.
9829    async fn handle_tool_calls(
9830        &self,
9831        processed_input: &str,
9832        content: &str,
9833        tool_calls: Vec<ToolCall>,
9834        all_tool_calls: &mut Vec<ToolCall>,
9835    ) -> Result<ToolCallOutcome> {
9836        // Check if a transition should fire before executing the LLM's tool call.
9837        // If a transition fires, on_enter actions handle the tool call correctly
9838        // (with proper URLs from YAML), so skip the LLM's tool call.
9839        let transition_fired = self.evaluate_transitions(processed_input, content).await?;
9840        if transition_fired {
9841            self.memory
9842                .add_message(ChatMessage::assistant(
9843                    "(Transitioned to new state — tool call handled by workflow)",
9844                ))
9845                .await?;
9846            return Ok(ToolCallOutcome::TransitionFired);
9847        }
9848
9849        // 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.
9850        self.memory
9851            .add_message(ChatMessage::assistant(content))
9852            .await?;
9853        let native_tool_call = Self::is_native_tool_call_content(content);
9854
9855        let results = self.execute_tools_parallel(&tool_calls).await;
9856
9857        for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
9858            match result {
9859                Ok(output) => {
9860                    self.memory
9861                        .add_message(Self::tool_result_message(
9862                            tool_call,
9863                            &output,
9864                            native_tool_call,
9865                        ))
9866                        .await?;
9867                }
9868                Err(e) => {
9869                    // Check if this is a HITL rejection - if so, break the loop
9870                    if matches!(e, AgentError::HITLRejected(_)) {
9871                        self.memory
9872                            .add_message(ChatMessage::assistant(format!(
9873                                "The operation was rejected by the approver: {}",
9874                                e
9875                            )))
9876                            .await?;
9877                        // Return the rejection message to user, don't continue loop
9878                        return Ok(ToolCallOutcome::Rejected(AgentResponse {
9879                            content: format!("Operation cancelled: {}", e),
9880                            metadata: None,
9881                            tool_calls: Some(all_tool_calls.clone()),
9882                        }));
9883                    }
9884                    self.memory
9885                        .add_message(Self::tool_result_message(
9886                            tool_call,
9887                            &format!("Error: {}", e),
9888                            native_tool_call,
9889                        ))
9890                        .await?;
9891                }
9892            }
9893            all_tool_calls.push(tool_call.clone());
9894        }
9895        Ok(ToolCallOutcome::Continue)
9896    }
9897
9898    /// Run the reflection loop on a response, returning (improved_content, reflection_metadata).
9899    async fn run_reflection(
9900        &self,
9901        llm: &dyn LLMProvider,
9902        processed_input: &str,
9903        mut content: String,
9904    ) -> Result<(String, Option<ReflectionMetadata>)> {
9905        let should_reflect = self.should_reflect(processed_input, &content).await?;
9906        if !should_reflect {
9907            return Ok((content, None));
9908        }
9909
9910        info!("Starting response reflection evaluation");
9911        let mut attempts = 0u32;
9912        let max_retries = self.reflection_config.max_retries;
9913        let mut history: Vec<ReflectionAttempt> = Vec::new();
9914
9915        loop {
9916            let evaluation = self.evaluate_response(processed_input, &content).await?;
9917
9918            if evaluation.passed || attempts >= max_retries {
9919                info!(
9920                    passed = evaluation.passed,
9921                    confidence = evaluation.confidence,
9922                    attempts = attempts + 1,
9923                    "Reflection evaluation complete"
9924                );
9925                let reflection_metadata = Some(
9926                    ReflectionMetadata::new(evaluation)
9927                        .with_attempts(attempts + 1)
9928                        .with_history(history),
9929                );
9930                return Ok((content, reflection_metadata));
9931            }
9932
9933            debug!(
9934                attempt = attempts + 1,
9935                failed_criteria = evaluation.failed_criteria().count(),
9936                "Response did not meet criteria, retrying"
9937            );
9938
9939            history.push(
9940                ReflectionAttempt::new(&content, evaluation.clone())
9941                    .with_feedback("Response did not meet quality criteria"),
9942            );
9943
9944            let feedback: Vec<String> = evaluation
9945                .failed_criteria()
9946                .map(|c| format!("- {}", c.criterion))
9947                .collect();
9948
9949            let retry_prompt = format!(
9950                "Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response.",
9951                feedback.join("\n")
9952            );
9953
9954            self.memory
9955                .add_message(ChatMessage::user(&retry_prompt))
9956                .await?;
9957
9958            let retry_messages = self.build_messages().await?;
9959            let retry_response = self
9960                .observe_purpose(
9961                    ObservationPurpose::ReflectionEvaluation,
9962                    llm.complete(&retry_messages, None),
9963                )
9964                .await
9965                .map_err(|e| AgentError::LLM(e.to_string()))?;
9966
9967            content = retry_response.content.trim().to_string();
9968            attempts += 1;
9969        }
9970    }
9971
9972    /// Record the assistant turn, evaluate transitions, and decide what to do next.
9973    /// Returns PostLoopResult so callers can apply_post_loop_result for re-dispatch.
9974    async fn post_loop_processing(
9975        &self,
9976        processed_input: &str,
9977        content: String,
9978    ) -> Result<PostLoopResult> {
9979        // Do NOT add the assistant message to memory yet.
9980        // evaluate_transitions receives content as a direct parameter, so the message does not need to be in memory for transitions to evaluate correctly.
9981        // For NeedsRedispatch we skip adding the stale old-state response entirely, keeping memory clean for the re-dispatched handler.
9982
9983        self.increment_turn();
9984
9985        // Run context extractors so guards can check freshly-extracted values.
9986        self.run_context_extractors(processed_input).await;
9987
9988        let transitioned = self.evaluate_transitions(processed_input, &content).await?;
9989
9990        if !transitioned {
9991            self.memory
9992                .add_message(ChatMessage::assistant(&content))
9993                .await?;
9994            self.check_memory_compression().await?;
9995            return Ok(PostLoopResult::NoTransition(content));
9996        }
9997
9998        // Check if we should skip re-generation after this transition.
9999        if !self.should_regenerate_after_transition() {
10000            self.memory
10001                .add_message(ChatMessage::assistant(&content))
10002                .await?;
10003            self.check_memory_compression().await?;
10004            return Ok(PostLoopResult::Transitioned(content));
10005        }
10006
10007        // Check if the new state needs full dispatch.
10008        // Orchestration states (concurrent, group_chat, pipeline, handoff, delegate) need their dedicated handlers.
10009        // 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.
10010        if self.needs_redispatch_for_new_state() {
10011            info!("Post-transition NeedsRedispatch: new state requires full dispatch");
10012            // Stale old-state content is NOT added to memory.
10013            // apply_post_loop_result will increment redispatch_depth and call run_loop_internal, which produces the correct response and adds it.
10014            return Ok(PostLoopResult::NeedsRedispatch);
10015        }
10016
10017        // Normal post-transition re-generation (plain LLM with optional tool calls).
10018        // Add the stale response to history so the new LLM sees the conversation.
10019        self.memory
10020            .add_message(ChatMessage::assistant(&content))
10021            .await?;
10022        self.check_memory_compression().await?;
10023
10024        // If a transition fired, on_enter actions already executed (e.g., HTTP calls).
10025        // The current content was generated in the OLD state context and is stale.
10026        // Re-generate in the new state context so the LLM can reference on_enter results.
10027        // If the LLM responds with a tool call, execute it in a mini-loop so the
10028        // result is not returned as raw JSON text.
10029        let new_llm = self.get_state_llm()?;
10030        let mut final_content;
10031
10032        for post_iter in 0..self.max_iterations {
10033            let protocol = self.main_tool_protocol(new_llm.as_ref(), false).await?;
10034            let new_messages = self
10035                .build_messages_internal(true, None, protocol.choice.is_none())
10036                .await?;
10037            if post_iter == 0
10038                && let Some(system_msg) = new_messages.first()
10039                && system_msg.role == ai_agents_core::Role::System
10040            {
10041                debug!(
10042                    prompt_preview =
10043                        &system_msg.content[system_msg.content.len().saturating_sub(200)..],
10044                    "Post-transition system prompt (last 200 chars)"
10045                );
10046            }
10047
10048            let new_response = self
10049                .complete_main_llm_with_recovery(Arc::clone(&new_llm), &new_messages, &protocol)
10050                .await?;
10051            final_content = new_response.content.trim().to_string();
10052
10053            // Check if the post-transition response contains tool calls.
10054            // If so, execute them and loop so the LLM can summarize the result.
10055            if let Some(tool_calls) = self.parse_main_tool_calls(&final_content, &protocol) {
10056                let native_tool_call = Self::is_native_tool_call_content(&final_content);
10057                debug!(
10058                    post_iter = post_iter,
10059                    tools = tool_calls.len(),
10060                    "Post-transition tool call detected, executing"
10061                );
10062
10063                self.memory
10064                    .add_message(ChatMessage::assistant(&final_content))
10065                    .await?;
10066
10067                let results = self.execute_tools_parallel(&tool_calls).await;
10068                for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
10069                    match result {
10070                        Ok(output) => {
10071                            self.memory
10072                                .add_message(Self::tool_result_message(
10073                                    tool_call,
10074                                    &output,
10075                                    native_tool_call,
10076                                ))
10077                                .await?;
10078                        }
10079                        Err(e) => {
10080                            self.memory
10081                                .add_message(Self::tool_result_message(
10082                                    tool_call,
10083                                    &format!("Error: {}", e),
10084                                    native_tool_call,
10085                                ))
10086                                .await?;
10087                        }
10088                    }
10089                }
10090                // Loop to let the LLM see the tool result and produce a text response.
10091                continue;
10092            }
10093
10094            // No tool call - this is the final text response.
10095            self.memory
10096                .add_message(ChatMessage::assistant(&final_content))
10097                .await?;
10098            return Ok(PostLoopResult::Transitioned(final_content));
10099        }
10100
10101        // Exhausted post-transition iterations (unlikely). Return last content.
10102        final_content = "Post-transition processing completed.".to_string();
10103        self.memory
10104            .add_message(ChatMessage::assistant(&final_content))
10105            .await?;
10106
10107        Ok(PostLoopResult::Transitioned(final_content))
10108    }
10109
10110    /// Build the final AgentResponse with all metadata.
10111    /// Check whether to re-generate a response after a state transition.
10112    fn should_regenerate_after_transition(&self) -> bool {
10113        if let Some(ref sm) = self.state_machine {
10114            // Global setting
10115            if !sm.config().regenerate_on_transition {
10116                return false;
10117            }
10118            // Per-state override on the new (current) state
10119            if let Some(def) = sm.current_definition()
10120                && let Some(regen) = def.regenerate_on_enter
10121            {
10122                return regen;
10123            }
10124        }
10125        true
10126    }
10127
10128    /// Return true when the new state requires full dispatch via run_loop_internal.
10129    /// Covers orchestration states and any non-None effective reasoning mode.
10130    fn needs_redispatch_for_new_state(&self) -> bool {
10131        if let Some(ref sm) = self.state_machine
10132            && let Some(def) = sm.current_definition()
10133        {
10134            if def.concurrent.is_some()
10135                || def.group_chat.is_some()
10136                || def.pipeline.is_some()
10137                || def.handoff.is_some()
10138                || def.delegate.is_some()
10139            {
10140                return true;
10141            }
10142            // Any non-None effective reasoning mode requires the main dispatch loop.
10143            let effective = self.get_effective_reasoning_config();
10144            if !matches!(effective.mode, ReasoningMode::None) {
10145                return true;
10146            }
10147        }
10148        false
10149    }
10150
10151    /// Consume a PostLoopResult. NeedsRedispatch re-enters run_loop_internal.
10152    /// The user message is already in memory - redispatch_depth suppresses re-adding it.
10153    async fn apply_post_loop_result(
10154        &self,
10155        processed_input: &str,
10156        result: PostLoopResult,
10157    ) -> Result<String> {
10158        match result {
10159            PostLoopResult::NoTransition(content) | PostLoopResult::Transitioned(content) => {
10160                Ok(content)
10161            }
10162            PostLoopResult::NeedsRedispatch => {
10163                const MAX_REDISPATCH_DEPTH: u32 = 3;
10164                let current_depth = *self.redispatch_depth.read();
10165                if current_depth >= MAX_REDISPATCH_DEPTH {
10166                    warn!(
10167                        depth = current_depth,
10168                        "Post-transition re-dispatch depth limit reached, returning empty response"
10169                    );
10170                    let content = String::new();
10171                    self.memory
10172                        .add_message(ChatMessage::assistant(&content))
10173                        .await?;
10174                    return Ok(content);
10175                }
10176                *self.redispatch_depth.write() += 1;
10177                if let Some(context) = self.active_turn_context.write().as_mut() {
10178                    context.enter_redispatch();
10179                }
10180                info!(
10181                    depth = current_depth + 1,
10182                    "Re-dispatching for new state after transition"
10183                );
10184                let resp = Box::pin(self.run_loop_internal(processed_input)).await;
10185                *self.redispatch_depth.write() -= 1;
10186                if let Some(context) = self.active_turn_context.write().as_mut() {
10187                    context.exit_redispatch();
10188                }
10189                resp.map(|r| r.content)
10190            }
10191        }
10192    }
10193
10194    /// Builds the final response metadata after committed output and transition handling complete.
10195    fn build_agent_response(&self, parts: AgentResponseParts) -> AgentResponse {
10196        let AgentResponseParts {
10197            content,
10198            all_tool_calls,
10199            reasoning_mode,
10200            auto_detected,
10201            iterations,
10202            thinking,
10203            reflection_metadata,
10204        } = parts;
10205        let reasoning_metadata = ReasoningMetadata::new(reasoning_mode.clone())
10206            .with_thinking(thinking.clone().unwrap_or_default())
10207            .with_iterations(iterations)
10208            .with_auto_detected(auto_detected);
10209
10210        let mut response = AgentResponse::new(&content);
10211        if !all_tool_calls.is_empty() {
10212            response = response.with_tool_calls(all_tool_calls);
10213        }
10214
10215        if let Some(state) = self.current_state() {
10216            response = response.with_metadata("current_state", serde_json::json!(state));
10217        }
10218
10219        response = response.with_metadata(
10220            "reasoning",
10221            serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
10222        );
10223
10224        if let Some(ref refl_meta) = reflection_metadata {
10225            response = response.with_metadata(
10226                "reflection",
10227                serde_json::to_value(refl_meta).unwrap_or_default(),
10228            );
10229        }
10230
10231        response
10232    }
10233
10234    // Handle delegation: forward user input to a registry agent.
10235    async fn handle_delegated_state(
10236        &self,
10237        input: &str,
10238        delegate_id: &str,
10239        state_def: &ai_agents_state::StateDefinition,
10240    ) -> Result<AgentResponse> {
10241        use std::time::Instant;
10242
10243        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10244            AgentError::Config(format!(
10245                "State delegates to '{}' but no agent registry is configured. \
10246                 Add a spawner section with auto_spawn to your YAML.",
10247                delegate_id
10248            ))
10249        })?;
10250
10251        let state_name = self
10252            .state_machine
10253            .as_ref()
10254            .map(|sm| sm.current())
10255            .unwrap_or_else(|| "unknown".to_string());
10256
10257        self.hooks.on_delegate_start(delegate_id, &state_name).await;
10258        let start = Instant::now();
10259
10260        let delegate = registry.get(delegate_id).ok_or_else(|| {
10261            AgentError::Other(format!(
10262                "State '{}' delegates to '{}' but no agent with that ID exists in the registry.",
10263                state_name, delegate_id
10264            ))
10265        })?;
10266
10267        // Prepare input based on delegate_context mode.
10268        let context_mode = state_def.delegate_context.clone().unwrap_or_default();
10269        let effective_input = self
10270            .observe_purpose(
10271                ObservationPurpose::OrchestrationRouting,
10272                crate::orchestration::context::prepare_delegate_input(
10273                    input,
10274                    &context_mode,
10275                    &*self.memory,
10276                    self.llm_registry.get("router").ok().as_deref(),
10277                ),
10278            )
10279            .await?;
10280
10281        let response = delegate
10282            .chat_with_actor_context(&effective_input, self.outbound_actor_context())
10283            .await?;
10284
10285        let duration_ms = start.elapsed().as_millis() as u64;
10286        self.hooks
10287            .on_delegate_complete(delegate_id, &state_name, duration_ms)
10288            .await;
10289
10290        // Backward-compatible context key.
10291        let ctx_key = format!("delegation.{}.last_response", delegate_id);
10292        let _ = self.context_manager.set(
10293            &ctx_key,
10294            serde_json::Value::String(response.content.clone()),
10295        );
10296
10297        // Structured orchestration context.
10298        let _ = self.context_manager.set(
10299            "orchestration",
10300            serde_json::json!({
10301                "type": "delegate",
10302                "agent": delegate_id,
10303                "state": state_name,
10304                "response": response.content,
10305                "duration_ms": duration_ms,
10306            }),
10307        );
10308
10309        self.commit_root_user_message(input).await?;
10310
10311        // post_loop_processing records the assistant turn and evaluates transitions.
10312        // apply_post_loop_result handles NeedsRedispatch by re-entering run_loop_internal.
10313        let post_result = self
10314            .post_loop_processing(
10315                input,
10316                format!("[Delegated to {}]: {}", delegate_id, response.content),
10317            )
10318            .await?;
10319        let final_content = self.apply_post_loop_result(input, post_result).await?;
10320
10321        let mut result = AgentResponse::new(final_content);
10322
10323        let metadata = serde_json::json!({
10324            "orchestration": {
10325                "type": "delegate",
10326                "agent": delegate_id,
10327                "state": state_name,
10328                "response": response.content,
10329                "duration_ms": duration_ms,
10330            }
10331        });
10332        result.metadata = Some(
10333            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10334                metadata,
10335            )
10336            .unwrap_or_default(),
10337        );
10338
10339        self.finish_turn_if_root(&result).await?;
10340        Ok(result)
10341    }
10342
10343    // Handle concurrent execution: run multiple registry agents in parallel and aggregate.
10344    async fn handle_concurrent_state(
10345        &self,
10346        input: &str,
10347        config: &ai_agents_state::ConcurrentStateConfig,
10348    ) -> Result<AgentResponse> {
10349        use std::time::Instant;
10350
10351        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10352            AgentError::Config(
10353                "Concurrent state requires an agent registry. Add a spawner section.".into(),
10354            )
10355        })?;
10356
10357        // Render input template if provided, otherwise use the raw input.
10358        // Uses direct minijinja rendering (same approach as pipeline) so variables
10359        // are top-level: {{ user_input }}, not {{ context.user_input }}.
10360        // Enrich input with parent conversation history when context_mode is set.
10361        let context_mode = config.context_mode.clone().unwrap_or_default();
10362        let context_input = self
10363            .observe_purpose(
10364                ObservationPurpose::OrchestrationRouting,
10365                crate::orchestration::context::prepare_delegate_input(
10366                    input,
10367                    &context_mode,
10368                    &*self.memory,
10369                    self.llm_registry.get("router").ok().as_deref(),
10370                ),
10371            )
10372            .await?;
10373
10374        let effective_input = if let Some(ref tmpl) = config.input {
10375            render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10376                .unwrap_or_else(|_| context_input.clone())
10377        } else {
10378            context_input
10379        };
10380
10381        let start = Instant::now();
10382
10383        let llm_name = config
10384            .aggregation
10385            .synthesizer_llm
10386            .as_deref()
10387            .unwrap_or("router");
10388        let llm_provider = self.llm_registry.get(llm_name).ok();
10389
10390        let vote_parallelism = if self.runtime_config.optimization.enabled
10391            && self
10392                .runtime_config
10393                .optimization
10394                .parallel_orchestration_vote_extraction
10395        {
10396            Some(self.runtime_config.optimization.max_parallel_runtime_tasks)
10397        } else {
10398            None
10399        };
10400
10401        let result = self
10402            .observe_purpose(
10403                ObservationPurpose::OrchestrationAggregation,
10404                scope_actor_context(
10405                    self.outbound_actor_context(),
10406                    crate::orchestration::concurrent(
10407                        registry,
10408                        &effective_input,
10409                        &config.agents,
10410                        &config.aggregation,
10411                        llm_provider.as_deref(),
10412                        config.min_required,
10413                        config.timeout_ms,
10414                        config.on_partial_failure.clone(),
10415                        vote_parallelism,
10416                    ),
10417                ),
10418            )
10419            .await?;
10420
10421        let duration_ms = start.elapsed().as_millis() as u64;
10422        let agent_ids: Vec<String> = config.agents.iter().map(|a| a.id().to_string()).collect();
10423        let strategy = format!("{:?}", config.aggregation.strategy);
10424        self.hooks
10425            .on_concurrent_complete(&agent_ids, &strategy, duration_ms)
10426            .await;
10427
10428        // Backward-compatible context key.
10429        let _ = self.context_manager.set(
10430            "concurrent.result",
10431            serde_json::Value::String(result.response.content.clone()),
10432        );
10433
10434        // Build per-agent result data for context and metadata.
10435        let agents_json: Vec<serde_json::Value> = result
10436            .agent_results
10437            .iter()
10438            .map(|ar| {
10439                serde_json::json!({
10440                    "id": ar.agent_id,
10441                    "response": ar.response.as_ref().map(|r| r.content.as_str()),
10442                    "success": ar.success,
10443                    "error": ar.error,
10444                    "duration_ms": ar.duration_ms,
10445                })
10446            })
10447            .collect();
10448
10449        // Structured orchestration context with per-agent results.
10450        let _ = self.context_manager.set(
10451            "orchestration",
10452            serde_json::json!({
10453                "type": "concurrent",
10454                "result": result.response.content,
10455                "strategy": strategy,
10456                "agents": agents_json,
10457                "duration_ms": duration_ms,
10458            }),
10459        );
10460
10461        self.commit_root_user_message(input).await?;
10462
10463        let post_result = self
10464            .post_loop_processing(input, result.response.content.clone())
10465            .await?;
10466        let final_content = self.apply_post_loop_result(input, post_result).await?;
10467
10468        let mut response = AgentResponse::new(final_content);
10469        let metadata = serde_json::json!({
10470            "orchestration": {
10471                "type": "concurrent",
10472                "result": result.response.content,
10473                "strategy": strategy,
10474                "agents": agents_json,
10475                "duration_ms": duration_ms,
10476            }
10477        });
10478        response.metadata = Some(
10479            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10480                metadata,
10481            )
10482            .unwrap_or_default(),
10483        );
10484
10485        self.finish_turn_if_root(&response).await?;
10486        Ok(response)
10487    }
10488
10489    // Handle group chat: run a multi-turn multi-agent conversation.
10490    async fn handle_group_chat_state(
10491        &self,
10492        input: &str,
10493        config: &ai_agents_state::GroupChatStateConfig,
10494    ) -> Result<AgentResponse> {
10495        use std::time::Instant;
10496
10497        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10498            AgentError::Config(
10499                "Group chat state requires an agent registry. Add a spawner section.".into(),
10500            )
10501        })?;
10502
10503        let start = Instant::now();
10504
10505        let llm_provider = self.llm_registry.get("router").ok();
10506
10507        // Enrich input with parent conversation history when context_mode is set.
10508        let context_mode = config.context_mode.clone().unwrap_or_default();
10509        let context_input = self
10510            .observe_purpose(
10511                ObservationPurpose::OrchestrationRouting,
10512                crate::orchestration::context::prepare_delegate_input(
10513                    input,
10514                    &context_mode,
10515                    &*self.memory,
10516                    self.llm_registry.get("router").ok().as_deref(),
10517                ),
10518            )
10519            .await?;
10520
10521        // Render input template if provided, otherwise use the raw user message as topic.
10522        let effective_topic = if let Some(ref tmpl) = config.input {
10523            render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10524                .unwrap_or_else(|_| context_input.clone())
10525        } else {
10526            context_input
10527        };
10528
10529        let result = self
10530            .observe_purpose(
10531                ObservationPurpose::OrchestrationConversation,
10532                scope_actor_context(
10533                    self.outbound_actor_context(),
10534                    crate::orchestration::group_chat(
10535                        registry,
10536                        &effective_topic,
10537                        config,
10538                        llm_provider.as_deref(),
10539                        Some(&*self.hooks),
10540                    ),
10541                ),
10542            )
10543            .await?;
10544
10545        let duration_ms = start.elapsed().as_millis() as u64;
10546
10547        // Backward-compatible context key.
10548        let _ = self.context_manager.set(
10549            "group_chat.conclusion",
10550            serde_json::Value::String(result.response.content.clone()),
10551        );
10552
10553        // Build transcript data for context and metadata.
10554        let transcript_json: Vec<serde_json::Value> = result
10555            .transcript
10556            .iter()
10557            .map(|t| {
10558                serde_json::json!({
10559                    "speaker": t.speaker,
10560                    "round": t.round,
10561                    "content": t.content,
10562                })
10563            })
10564            .collect();
10565
10566        // Structured orchestration context with full transcript.
10567        let _ = self.context_manager.set(
10568            "orchestration",
10569            serde_json::json!({
10570                "type": "group_chat",
10571                "conclusion": result.response.content,
10572                "transcript": transcript_json,
10573                "rounds": result.rounds_completed,
10574                "termination": result.termination_reason,
10575                "duration_ms": duration_ms,
10576            }),
10577        );
10578
10579        self.commit_root_user_message(input).await?;
10580
10581        let post_result = self
10582            .post_loop_processing(input, result.response.content.clone())
10583            .await?;
10584        let final_content = self.apply_post_loop_result(input, post_result).await?;
10585
10586        let mut response = AgentResponse::new(final_content);
10587        let metadata = serde_json::json!({
10588            "orchestration": {
10589                "type": "group_chat",
10590                "conclusion": result.response.content,
10591                "transcript": transcript_json,
10592                "rounds": result.rounds_completed,
10593                "termination": result.termination_reason,
10594                "duration_ms": duration_ms,
10595            }
10596        });
10597        response.metadata = Some(
10598            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10599                metadata,
10600            )
10601            .unwrap_or_default(),
10602        );
10603
10604        self.finish_turn_if_root(&response).await?;
10605        Ok(response)
10606    }
10607
10608    // Handle pipeline: run agents sequentially with per-stage input templates.
10609    async fn handle_pipeline_state(
10610        &self,
10611        input: &str,
10612        config: &ai_agents_state::PipelineStateConfig,
10613    ) -> Result<AgentResponse> {
10614        use std::time::Instant;
10615
10616        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10617            AgentError::Config(
10618                "Pipeline state requires an agent registry. Add a spawner section.".into(),
10619            )
10620        })?;
10621
10622        let start = Instant::now();
10623
10624        let stages: Vec<crate::orchestration::PipelineStage> = config
10625            .stages
10626            .iter()
10627            .map(|entry| {
10628                let mut stage = crate::orchestration::PipelineStage::id(entry.id());
10629                if let Some(tmpl) = entry.input() {
10630                    stage = stage.with_input(tmpl);
10631                }
10632                stage
10633            })
10634            .collect();
10635
10636        // Enrich input with parent conversation history when context_mode is set.
10637        let context_mode = config.context_mode.clone().unwrap_or_default();
10638        let context_input = self
10639            .observe_purpose(
10640                ObservationPurpose::OrchestrationRouting,
10641                crate::orchestration::context::prepare_delegate_input(
10642                    input,
10643                    &context_mode,
10644                    &*self.memory,
10645                    self.llm_registry.get("router").ok().as_deref(),
10646                ),
10647            )
10648            .await?;
10649
10650        let context_values = self.build_context_with_overlays();
10651        let result = self
10652            .observe_purpose(
10653                ObservationPurpose::OrchestrationRouting,
10654                scope_actor_context(
10655                    self.outbound_actor_context(),
10656                    crate::orchestration::pipeline(
10657                        registry,
10658                        &context_input,
10659                        &stages,
10660                        config.timeout_ms,
10661                        Some(&*self.hooks),
10662                        Some(&context_values),
10663                    ),
10664                ),
10665            )
10666            .await?;
10667
10668        let duration_ms = start.elapsed().as_millis() as u64;
10669
10670        // Backward-compatible context key.
10671        let _ = self.context_manager.set(
10672            "pipeline.result",
10673            serde_json::Value::String(result.response.content.clone()),
10674        );
10675
10676        // Build per-stage data for context and metadata.
10677        let stages_json: Vec<serde_json::Value> = result
10678            .stage_outputs
10679            .iter()
10680            .map(|s| {
10681                serde_json::json!({
10682                    "agent_id": s.agent_id,
10683                    "output": s.output,
10684                    "duration_ms": s.duration_ms,
10685                    "skipped": s.skipped,
10686                })
10687            })
10688            .collect();
10689
10690        // Structured orchestration context.
10691        let _ = self.context_manager.set(
10692            "orchestration",
10693            serde_json::json!({
10694                "type": "pipeline",
10695                "result": result.response.content,
10696                "stages": stages_json,
10697                "duration_ms": duration_ms,
10698            }),
10699        );
10700
10701        self.commit_root_user_message(input).await?;
10702
10703        let post_result = self
10704            .post_loop_processing(input, result.response.content.clone())
10705            .await?;
10706        let final_content = self.apply_post_loop_result(input, post_result).await?;
10707
10708        let mut response = AgentResponse::new(final_content);
10709        let metadata = serde_json::json!({
10710            "orchestration": {
10711                "type": "pipeline",
10712                "result": result.response.content,
10713                "stages": stages_json,
10714                "duration_ms": duration_ms,
10715            }
10716        });
10717        response.metadata = Some(
10718            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10719                metadata,
10720            )
10721            .unwrap_or_default(),
10722        );
10723
10724        self.finish_turn_if_root(&response).await?;
10725        Ok(response)
10726    }
10727
10728    // Handle handoff: LLM-directed agent-to-agent control transfer.
10729    async fn handle_handoff_state(
10730        &self,
10731        input: &str,
10732        config: &ai_agents_state::HandoffStateConfig,
10733    ) -> Result<AgentResponse> {
10734        use std::time::Instant;
10735
10736        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10737            AgentError::Config(
10738                "Handoff state requires an agent registry. Add a spawner section.".into(),
10739            )
10740        })?;
10741
10742        let llm = self
10743            .llm_registry
10744            .get("router")
10745            .map_err(|_| AgentError::Config("Handoff state requires a router LLM.".into()))?;
10746
10747        let start = Instant::now();
10748
10749        // Enrich input with parent conversation history when context_mode is set.
10750        let context_mode = config.context_mode.clone().unwrap_or_default();
10751        let context_input = self
10752            .observe_purpose(
10753                ObservationPurpose::OrchestrationRouting,
10754                crate::orchestration::context::prepare_delegate_input(
10755                    input,
10756                    &context_mode,
10757                    &*self.memory,
10758                    self.llm_registry.get("router").ok().as_deref(),
10759                ),
10760            )
10761            .await?;
10762
10763        // Render input template if provided, otherwise forward the raw user message.
10764        let effective_input = if let Some(ref tmpl) = config.input {
10765            render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10766                .unwrap_or_else(|_| context_input.clone())
10767        } else {
10768            context_input
10769        };
10770
10771        let result = self
10772            .observe_purpose(
10773                ObservationPurpose::OrchestrationRouting,
10774                scope_actor_context(
10775                    self.outbound_actor_context(),
10776                    crate::orchestration::handoff(
10777                        registry,
10778                        &effective_input,
10779                        &config.initial_agent,
10780                        &config.available_agents,
10781                        config.max_handoffs,
10782                        llm.as_ref(),
10783                        Some(&*self.hooks),
10784                    ),
10785                ),
10786            )
10787            .await?;
10788
10789        let duration_ms = start.elapsed().as_millis() as u64;
10790
10791        // Backward-compatible context key.
10792        let _ = self.context_manager.set(
10793            "handoff.result",
10794            serde_json::Value::String(result.response.content.clone()),
10795        );
10796
10797        // Build handoff chain data for context and metadata.
10798        let chain_json: Vec<serde_json::Value> = result
10799            .handoff_chain
10800            .iter()
10801            .map(|h| {
10802                serde_json::json!({
10803                    "from": h.from_agent,
10804                    "to": h.to_agent,
10805                    "reason": h.reason,
10806                })
10807            })
10808            .collect();
10809
10810        // Structured orchestration context.
10811        let _ = self.context_manager.set(
10812            "orchestration",
10813            serde_json::json!({
10814                "type": "handoff",
10815                "result": result.response.content,
10816                "final_agent": result.final_agent,
10817                "handoff_chain": chain_json,
10818                "duration_ms": duration_ms,
10819            }),
10820        );
10821
10822        self.commit_root_user_message(input).await?;
10823
10824        let post_result = self
10825            .post_loop_processing(input, result.response.content.clone())
10826            .await?;
10827        let final_content = self.apply_post_loop_result(input, post_result).await?;
10828
10829        let mut response = AgentResponse::new(final_content);
10830        let metadata = serde_json::json!({
10831            "orchestration": {
10832                "type": "handoff",
10833                "result": result.response.content,
10834                "final_agent": result.final_agent,
10835                "handoff_chain": chain_json,
10836                "duration_ms": duration_ms,
10837            }
10838        });
10839        response.metadata = Some(
10840            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10841                metadata,
10842            )
10843            .unwrap_or_default(),
10844        );
10845
10846        self.finish_turn_if_root(&response).await?;
10847        Ok(response)
10848    }
10849
10850    // run_loop_internal: blocking (non-streaming) agent pipeline.
10851    async fn run_loop_internal(&self, input: &str) -> Result<AgentResponse> {
10852        self.begin_root_turn();
10853        // Resolve actor_id from context, reload facts if actor changed, bump counter.
10854        self.pre_turn_session_lifecycle().await;
10855
10856        let input_data = self.process_input(input).await?;
10857        self.update_active_turn_context(&input_data.content, input_data.context.clone());
10858
10859        // Inject process context (detect/extract results) into agent context
10860        // so system prompt templates can use {{ context.detected_language }} etc.
10861        for (key, value) in &input_data.context {
10862            let _ = self.context_manager.set(key, value.clone());
10863        }
10864
10865        if input_data.metadata.rejected {
10866            let reason = input_data
10867                .metadata
10868                .rejection_reason
10869                .unwrap_or_else(|| "Input rejected".to_string());
10870            warn!(reason = %reason, "Input rejected");
10871            let response = AgentResponse::new(reason);
10872            self.finish_turn_if_root(&response).await?;
10873            return Ok(response);
10874        }
10875
10876        let processed_input = &input_data.content;
10877
10878        if let Some(response) = self.try_pre_response_transition(processed_input).await? {
10879            return Ok(response);
10880        }
10881
10882        // Handle orchestration states (delegate, concurrent, group_chat, pipeline, handoff).
10883        if let Some(ref sm) = self.state_machine
10884            && let Some(def) = sm.current_definition()
10885        {
10886            if let Some(ref delegate_id) = def.delegate {
10887                return self
10888                    .handle_delegated_state(processed_input, delegate_id, &def)
10889                    .await;
10890            }
10891            if let Some(ref concurrent_config) = def.concurrent {
10892                return self
10893                    .handle_concurrent_state(processed_input, concurrent_config)
10894                    .await;
10895            }
10896            if let Some(ref group_chat_config) = def.group_chat {
10897                return self
10898                    .handle_group_chat_state(processed_input, group_chat_config)
10899                    .await;
10900            }
10901            if let Some(ref pipeline_config) = def.pipeline {
10902                return self
10903                    .handle_pipeline_state(processed_input, pipeline_config)
10904                    .await;
10905            }
10906            if let Some(ref handoff_config) = def.handoff {
10907                return self
10908                    .handle_handoff_state(processed_input, handoff_config)
10909                    .await;
10910            }
10911        }
10912
10913        //
10914        // The speculative future is boxed to keep the runtime future size manageable.
10915        // Removing the box can overflow small test stacks because this function is recursive through redispatch.
10916        //
10917        if let Some(response) =
10918            Box::pin(self.try_speculative_branches(processed_input, &input_data.context)).await?
10919        {
10920            return Ok(response);
10921        }
10922
10923        match self.try_skill_route(processed_input).await? {
10924            SkillRouteResult::Response { skill_id, content } => {
10925                self.commit_root_user_message(processed_input).await?;
10926                return self
10927                    .handle_skill_response(processed_input, &skill_id, content, &input_data.context)
10928                    .await;
10929            }
10930            SkillRouteResult::NeedsClarification {
10931                response,
10932                ownership,
10933            } => {
10934                let admission = self
10935                    .admit_optional_disambiguation_ownership(ownership)
10936                    .await?;
10937                self.commit_root_user_message(processed_input).await?;
10938                if let Some(q) = response
10939                    .metadata
10940                    .as_ref()
10941                    .and_then(|m| m.get("disambiguation"))
10942                    .and_then(|d| d.get("status"))
10943                    .and_then(|s| s.as_str())
10944                    && q == "awaiting_clarification"
10945                {
10946                    // Store the clarification question in memory so the next turn
10947                    // can be handled as a clarification response.
10948                    self.memory
10949                        .add_message(ChatMessage::assistant(&response.content))
10950                        .await?;
10951                }
10952                drop(admission);
10953                self.finish_turn_if_root(&response).await?;
10954                return Ok(response);
10955            }
10956            SkillRouteResult::NoMatch => {} // continue to normal LLM chat
10957        }
10958
10959        let effective_reasoning = self.get_effective_reasoning_config();
10960        let reasoning_mode = self.determine_reasoning_mode(processed_input).await?;
10961        let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
10962
10963        info!(
10964            reasoning_mode = ?reasoning_mode,
10965            auto_detected = auto_detected,
10966            reflection_enabled = ?self.reflection_config.enabled,
10967            "Reasoning mode determined"
10968        );
10969
10970        if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
10971            self.commit_root_user_message(processed_input).await?;
10972            return self
10973                .handle_plan_and_execute(processed_input, &input_data.context, auto_detected)
10974                .await;
10975        }
10976
10977        self.commit_root_user_message(processed_input).await?;
10978
10979        let mut iterations = 0u32;
10980        let mut all_tool_calls: Vec<ToolCall> = Vec::new();
10981        let mut thinking_content: Option<String> = None;
10982
10983        let llm = self.get_state_llm()?;
10984
10985        loop {
10986            // When reasoning is active, cap iterations at the reasoning-specific limit.
10987            let effective_max = if reasoning_mode != ReasoningMode::None {
10988                let rc = self.get_effective_reasoning_config();
10989                self.max_iterations.min(rc.max_iterations)
10990            } else {
10991                self.max_iterations
10992            };
10993
10994            if iterations >= effective_max {
10995                let err = AgentError::Other(format!("Max iterations ({}) exceeded", effective_max));
10996                self.hooks.on_error(&err).await;
10997                error!(iterations = iterations, "Max iterations exceeded");
10998                return Err(err);
10999            }
11000            iterations += 1;
11001            *self.iteration_count.write() = iterations;
11002
11003            debug!(iteration = iterations, max = effective_max, "LLM call");
11004
11005            let protocol = self.main_tool_protocol(llm.as_ref(), false).await?;
11006            let mut messages = self
11007                .build_messages_internal(true, None, protocol.choice.is_none())
11008                .await?;
11009            self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
11010
11011            self.hooks.on_llm_start(&messages).await;
11012            let llm_start = Instant::now();
11013            let response = self
11014                .complete_main_llm_with_recovery(Arc::clone(&llm), &messages, &protocol)
11015                .await?;
11016
11017            let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
11018            self.hooks.on_llm_complete(&response, llm_duration_ms).await;
11019
11020            let content = response.content.trim();
11021
11022            if let Some(tool_calls) = self.parse_main_tool_calls(content, &protocol) {
11023                match self
11024                    .handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
11025                    .await?
11026                {
11027                    ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
11028                    ToolCallOutcome::Rejected(resp) => {
11029                        self.finish_turn_if_root(&resp).await?;
11030                        return Ok(resp);
11031                    }
11032                }
11033            }
11034
11035            let (extracted_thinking, answer) = self.extract_thinking(content);
11036            if extracted_thinking.is_some() {
11037                thinking_content = extracted_thinking;
11038            }
11039
11040            let output_data = self.process_output(&answer, &input_data.context).await?;
11041
11042            let mut final_content = if output_data.metadata.rejected {
11043                output_data
11044                    .metadata
11045                    .rejection_reason
11046                    .unwrap_or_else(|| answer.to_string())
11047            } else {
11048                output_data.content
11049            };
11050
11051            // Run reflection (blocking LLM calls for retries)
11052            let reflection_metadata;
11053            (final_content, reflection_metadata) = self
11054                .run_reflection(&*llm, processed_input, final_content)
11055                .await?;
11056
11057            final_content =
11058                self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
11059
11060            // Post-loop: memory, transitions, post-transition re-generation.
11061            // apply_post_loop_result handles NeedsRedispatch by re-entering
11062            // run_loop_internal so the new state's full dispatch activates.
11063            let final_content = {
11064                let result = self
11065                    .post_loop_processing(processed_input, final_content)
11066                    .await?;
11067                self.apply_post_loop_result(processed_input, result).await?
11068            };
11069
11070            let reflected = reflection_metadata.is_some();
11071            let reasoning_mode_debug = format!("{:?}", reasoning_mode);
11072
11073            let response = self.build_agent_response(AgentResponseParts {
11074                content: final_content,
11075                all_tool_calls,
11076                reasoning_mode,
11077                auto_detected,
11078                iterations,
11079                thinking: thinking_content,
11080                reflection_metadata,
11081            });
11082
11083            self.finish_turn_if_root(&response).await?;
11084
11085            let tool_call_count = response.tool_calls.as_ref().map(|tc| tc.len()).unwrap_or(0);
11086            info!(
11087                tool_calls = tool_call_count,
11088                response_len = response.content.len(),
11089                reasoning_mode = %reasoning_mode_debug,
11090                reflected = reflected,
11091                "Chat completed"
11092            );
11093            return Ok(response);
11094        }
11095    }
11096
11097    async fn generate_buffered_streaming_draft(
11098        &self,
11099        processed_input: &str,
11100        routing_resolved: Arc<AtomicBool>,
11101    ) -> Result<StreamingDraftResult> {
11102        let llm = self.get_state_llm()?;
11103        if llm.configured_tool_choice().is_some() {
11104            let draft = self
11105                .generate_main_response_draft(processed_input, &ReasoningMode::None)
11106                .await?;
11107            return Ok(StreamingDraftResult::new(draft, Vec::new()));
11108        }
11109        let messages = self.build_messages_for_draft(processed_input).await?;
11110        let mut stream = self
11111            .observe_purpose(
11112                ObservationPurpose::MainResponse,
11113                llm.complete_stream(&messages, None),
11114            )
11115            .await
11116            .map_err(|e| AgentError::LLM(e.to_string()))?;
11117        let mut buffer = crate::optimization::StreamBranchBuffer::new(self.streaming.buffer_size)?;
11118        let mut chunks = Vec::new();
11119        let mut accumulated = String::new();
11120        while let Some(chunk_result) = stream.next().await {
11121            let chunk = chunk_result.map_err(|e| AgentError::LLM(e.to_string()))?;
11122            accumulated.push_str(&chunk.delta);
11123            let stream_chunk = StreamChunk::content(chunk.delta);
11124            if routing_resolved.load(Ordering::SeqCst) {
11125                chunks.push(stream_chunk);
11126            } else {
11127                buffer.push(stream_chunk)?;
11128            }
11129        }
11130        chunks.splice(0..0, buffer.drain());
11131        let content = accumulated.trim().to_string();
11132        let draft = if let Some(calls) = self.parse_tool_calls(&content) {
11133            MainResponseDraft::ToolCalls {
11134                raw_content: content,
11135                calls,
11136                thinking: None,
11137            }
11138        } else {
11139            MainResponseDraft::Text {
11140                raw_content: content,
11141                thinking: None,
11142            }
11143        };
11144        Ok(StreamingDraftResult::new(draft, chunks))
11145    }
11146
11147    async fn try_buffered_streaming_branches(
11148        &self,
11149        processed_input: &str,
11150        input_context: &HashMap<String, Value>,
11151    ) -> Result<Option<(AgentResponse, Vec<StreamChunk>)>> {
11152        let optimization = &self.runtime_config.optimization;
11153        if !optimization.enabled {
11154            return Ok(None);
11155        }
11156        let transition_enabled =
11157            optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
11158        if !transition_enabled {
11159            return Ok(None);
11160        }
11161        let mut branch_scheduler =
11162            TurnBranchScheduler::new(optimization.max_parallel_runtime_tasks)?;
11163        if !branch_scheduler.reserve_task() {
11164            return Ok(None);
11165        }
11166        if !self
11167            .reserve_active_speculative_llm_call(RuntimeOptimizationKind::BufferedStreamingRouting)
11168        {
11169            branch_scheduler.release_task();
11170            return Ok(None);
11171        }
11172        if !branch_scheduler.reserve_task() {
11173            branch_scheduler.release_task();
11174            return Ok(None);
11175        }
11176        let mut main_branch = RuntimeBranch::new(
11177            RuntimeTaskPurpose::MainResponse,
11178            RuntimeOptimizationKind::BufferedStreamingRouting,
11179            RuntimeTaskPriority::Normal,
11180            RuntimeCommitBehavior::FinalResponse,
11181        );
11182        let mut transition_branch = RuntimeBranch::new(
11183            RuntimeTaskPurpose::StateTransition,
11184            RuntimeOptimizationKind::ParallelStateTransition,
11185            RuntimeTaskPriority::Critical,
11186            RuntimeCommitBehavior::TransitionDecision,
11187        );
11188        let main_id = main_branch.branch_id();
11189        let transition_id = transition_branch.branch_id();
11190        let routing_resolved = Arc::new(AtomicBool::new(false));
11191        let mut main_future =
11192            Box::pin(crate::optimization::observability::with_branch_observation(
11193                &main_id,
11194                RuntimeOptimizationKind::BufferedStreamingRouting,
11195                RuntimeCommitBehavior::FinalResponse,
11196                self.generate_buffered_streaming_draft(
11197                    processed_input,
11198                    Arc::clone(&routing_resolved),
11199                ),
11200            ));
11201        let mut transition_future =
11202            Box::pin(crate::optimization::observability::with_branch_observation(
11203                &transition_id,
11204                RuntimeOptimizationKind::ParallelStateTransition,
11205                RuntimeCommitBehavior::TransitionDecision,
11206                self.select_parallel_transition_candidate(processed_input),
11207            ));
11208        let mut main_pending = true;
11209        let mut transition_pending = true;
11210        let mut main_result: Option<Result<StreamingDraftResult>> = None;
11211        let mut transition_finalized = false;
11212        let mut transition_candidate: Option<TransitionCandidate> = None;
11213        loop {
11214            if let Some(candidate) = transition_candidate.take() {
11215                if self
11216                    .approve_transition_target(&candidate.from_state, candidate.target())
11217                    .await?
11218                {
11219                    // Drop the stale stream future before transition side effects or redispatch reuse the provider.
11220                    drop(main_future);
11221                    drop(transition_future);
11222                    self.finalize_branch_loss(
11223                        &main_id,
11224                        RuntimeOptimizationKind::BufferedStreamingRouting,
11225                        RuntimeCommitBehavior::FinalResponse,
11226                        main_pending,
11227                        main_result.as_ref().map(|result| result.is_err()),
11228                    );
11229                    if !self
11230                        .apply_pre_response_transition_candidate(
11231                            &candidate,
11232                            &HashMap::new(),
11233                            processed_input,
11234                        )
11235                        .await?
11236                    {
11237                        self.finalize_optional_branch(
11238                            &transition_id,
11239                            RuntimeOptimizationKind::ParallelStateTransition,
11240                            RuntimeCommitBehavior::TransitionDecision,
11241                            "discarded",
11242                            false,
11243                        );
11244                        return Ok(None);
11245                    }
11246                    self.finalize_optional_branch(
11247                        &transition_id,
11248                        RuntimeOptimizationKind::ParallelStateTransition,
11249                        RuntimeCommitBehavior::TransitionDecision,
11250                        "committed",
11251                        true,
11252                    );
11253                    let response = self.redispatch_current_state(processed_input).await?;
11254                    return Ok(Some((
11255                        response.clone(),
11256                        vec![StreamChunk::content(response.content)],
11257                    )));
11258                }
11259                self.finalize_optional_branch(
11260                    &transition_id,
11261                    RuntimeOptimizationKind::ParallelStateTransition,
11262                    RuntimeCommitBehavior::TransitionDecision,
11263                    "discarded",
11264                    false,
11265                );
11266                routing_resolved.store(true, Ordering::SeqCst);
11267                transition_finalized = true;
11268            }
11269            if transition_finalized && let Some(result) = main_result.take() {
11270                let stream_draft = match result {
11271                    Ok(stream_draft) => stream_draft,
11272                    Err(error) => {
11273                        self.finalize_optional_branch(
11274                            &main_id,
11275                            RuntimeOptimizationKind::BufferedStreamingRouting,
11276                            RuntimeCommitBehavior::FinalResponse,
11277                            "failed",
11278                            false,
11279                        );
11280                        return Err(error);
11281                    }
11282                };
11283                let raw_draft_content = stream_draft.draft.raw_content().to_string();
11284                let buffered_chunks = stream_draft.chunks;
11285                self.finalize_optional_branch(
11286                    &main_id,
11287                    RuntimeOptimizationKind::BufferedStreamingRouting,
11288                    RuntimeCommitBehavior::FinalResponse,
11289                    "committed",
11290                    true,
11291                );
11292                let response = self
11293                    .commit_main_response_draft(
11294                        processed_input,
11295                        input_context,
11296                        stream_draft.draft,
11297                        ReasoningMode::None,
11298                        false,
11299                    )
11300                    .await?;
11301                let chunks = if response.content == raw_draft_content {
11302                    buffered_chunks
11303                } else {
11304                    vec![StreamChunk::content(response.content.clone())]
11305                };
11306                return Ok(Some((response, chunks)));
11307            }
11308            tokio::select! {
11309                result = &mut main_future, if main_pending => {
11310                    main_pending = false;
11311                    main_branch.transition_to(RuntimeBranchStatus::Completed)?;
11312                    main_result = Some(result);
11313                }
11314                result = &mut transition_future, if transition_pending => {
11315                    transition_pending = false;
11316                    transition_branch.transition_to(RuntimeBranchStatus::Completed)?;
11317                    match result {
11318                        Ok(ParallelTransitionSelection::Candidate(candidate)) => {
11319                            transition_candidate = Some(candidate)
11320                        }
11321                        Ok(ParallelTransitionSelection::NoMatch) => {
11322                            self.finalize_optional_branch(
11323                                &transition_id,
11324                                RuntimeOptimizationKind::ParallelStateTransition,
11325                                RuntimeCommitBehavior::TransitionDecision,
11326                                "discarded",
11327                                false,
11328                            );
11329                            routing_resolved.store(true, Ordering::SeqCst);
11330                            transition_finalized = true;
11331                        }
11332                        Ok(ParallelTransitionSelection::ReservationExhausted) => {
11333                            self.finalize_optional_branch(
11334                                &transition_id,
11335                                RuntimeOptimizationKind::ParallelStateTransition,
11336                                RuntimeCommitBehavior::TransitionDecision,
11337                                "cancelled",
11338                                false,
11339                            );
11340                            routing_resolved.store(true, Ordering::SeqCst);
11341                            self.finalize_branch_loss(
11342                                &main_id,
11343                                RuntimeOptimizationKind::BufferedStreamingRouting,
11344                                RuntimeCommitBehavior::FinalResponse,
11345                                main_pending,
11346                                main_result.as_ref().map(|result| result.is_err()),
11347                            );
11348                            return Ok(None);
11349                        }
11350                        Err(_) => {
11351                            self.finalize_optional_branch(
11352                                &transition_id,
11353                                RuntimeOptimizationKind::ParallelStateTransition,
11354                                RuntimeCommitBehavior::TransitionDecision,
11355                                "failed",
11356                                false,
11357                            );
11358                            routing_resolved.store(true, Ordering::SeqCst);
11359                            transition_finalized = true;
11360                        }
11361                    }
11362                }
11363            }
11364        }
11365    }
11366
11367    /// Streaming agent pipeline
11368    /// Uses all the same shared helpers as run_loop_internal.
11369    /// The ONLY difference: LLM calls use complete_stream() + yield deltas.
11370    fn run_loop_internal_stream<'a>(
11371        &'a self,
11372        input: &'a str,
11373        terminal: RuntimeStreamTerminalSlot,
11374    ) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
11375        let include_tool_events = self.streaming.include_tool_events;
11376        let include_state_events = self.streaming.include_state_events;
11377
11378        Box::pin(async_stream::stream! {
11379            self.begin_root_turn();
11380            // Parity with non-stream: resolve actor from context and load facts if changed.
11381            self.pre_turn_session_lifecycle().await;
11382
11383            let input_data = match self.process_input(input).await {
11384                Ok(data) => data,
11385                Err(e) => {
11386                    yield StreamChunk::error(e.to_string());
11387                    return;
11388                }
11389            };
11390            self.update_active_turn_context(&input_data.content, input_data.context.clone());
11391
11392            // Inject process context (detect/extract results) into agent context
11393            for (key, value) in &input_data.context {
11394                let _ = self.context_manager.set(key, value.clone());
11395            }
11396
11397            if input_data.metadata.rejected {
11398                let reason = input_data
11399                    .metadata
11400                    .rejection_reason
11401                    .unwrap_or_else(|| "Input rejected".to_string());
11402                warn!(reason = %reason, "Input rejected (stream)");
11403                yield StreamChunk::error(reason);
11404                return;
11405            }
11406
11407            let processed_input = &input_data.content;
11408
11409            if self.runtime_config.optimization.enabled
11410                && matches!(
11411                    self.runtime_config.optimization.streaming_policy,
11412                    crate::optimization::StreamingOptimizationPolicy::BufferUntilRoutingDone
11413                )
11414            {
11415                //
11416                // Buffered routing keeps stale stream output hidden until a branch winner is known.
11417                // The boxed future prevents this stream state machine from becoming too large.
11418                //
11419                match Box::pin(self.try_buffered_streaming_branches(processed_input, &input_data.context)).await {
11420                    Ok(Some((response, chunks))) => {
11421                        for chunk in chunks {
11422                            yield chunk;
11423                        }
11424                        record_runtime_stream_final(&terminal, response);
11425                        yield StreamChunk::Done {};
11426                        return;
11427                    }
11428                    Ok(None) => {}
11429                    Err(e) => {
11430                        yield StreamChunk::error(e.to_string());
11431                        return;
11432                    }
11433                }
11434            }
11435
11436            if self.runtime_config.optimization.enabled
11437                && matches!(
11438                    self.runtime_config.optimization.streaming_policy,
11439                    crate::optimization::StreamingOptimizationPolicy::PreflightOnly
11440                )
11441            {
11442                match self.try_pre_response_transition(processed_input).await {
11443                    Ok(Some(response)) => {
11444                        yield StreamChunk::content(&response.content);
11445                        record_runtime_stream_final(&terminal, response);
11446                        yield StreamChunk::Done {};
11447                        return;
11448                    }
11449                    Ok(None) => {}
11450                    Err(e) => {
11451                        yield StreamChunk::error(e.to_string());
11452                        return;
11453                    }
11454                }
11455            }
11456
11457            // Handle orchestration states in streaming mode.
11458            if let Some(ref sm) = self.state_machine
11459                && let Some(def) = sm.current_definition()
11460            {
11461                    let orchestration_result = if let Some(ref delegate_id) = def.delegate {
11462                        Some(self.handle_delegated_state(processed_input, delegate_id, &def).await)
11463                    } else if let Some(ref concurrent_config) = def.concurrent {
11464                        Some(self.handle_concurrent_state(processed_input, concurrent_config).await)
11465                    } else if let Some(ref group_chat_config) = def.group_chat {
11466                        Some(self.handle_group_chat_state(processed_input, group_chat_config).await)
11467                    } else if let Some(ref pipeline_config) = def.pipeline {
11468                        Some(self.handle_pipeline_state(processed_input, pipeline_config).await)
11469                    } else if let Some(ref handoff_config) = def.handoff {
11470                        Some(self.handle_handoff_state(processed_input, handoff_config).await)
11471                    } else {
11472                        None
11473                    };
11474
11475                    if let Some(result) = orchestration_result {
11476                        match result {
11477                            Ok(response) => {
11478                                yield StreamChunk::content(&response.content);
11479                                record_runtime_stream_final(&terminal, response);
11480                                yield StreamChunk::Done {};
11481                            }
11482                            Err(e) => {
11483                                yield StreamChunk::error(e.to_string());
11484                            }
11485                        }
11486                        return;
11487                    }
11488                }
11489
11490            // Skill routing
11491            match self.try_skill_route(processed_input).await {
11492                Ok(SkillRouteResult::Response { skill_id, content }) => {
11493                    if let Err(e) = self.commit_root_user_message(processed_input).await {
11494                        yield StreamChunk::error(e.to_string());
11495                        return;
11496                    }
11497                    match self.handle_skill_response(processed_input, &skill_id, content, &input_data.context).await {
11498                        Ok(resp) => {
11499                            yield StreamChunk::content(&resp.content);
11500                            record_runtime_stream_final(&terminal, resp);
11501                            yield StreamChunk::Done {};
11502                            return;
11503                        }
11504                        Err(e) => {
11505                            yield StreamChunk::error(e.to_string());
11506                            return;
11507                        }
11508                    }
11509                }
11510                Ok(SkillRouteResult::NeedsClarification {
11511                    response,
11512                    ownership,
11513                }) => {
11514                    let admission = match self
11515                        .admit_optional_disambiguation_ownership(ownership)
11516                        .await
11517                    {
11518                        Ok(admission) => admission,
11519                        Err(e) => {
11520                            yield StreamChunk::error(e.to_string());
11521                            return;
11522                        }
11523                    };
11524                    if let Err(e) = self.commit_root_user_message(processed_input).await {
11525                        yield StreamChunk::error(e.to_string());
11526                        return;
11527                    }
11528                    let _ = self.memory.add_message(ChatMessage::assistant(&response.content)).await;
11529                    drop(admission);
11530                    if let Err(e) = self.finish_turn_if_root(&response).await {
11531                        yield StreamChunk::error(e.to_string());
11532                        return;
11533                    }
11534                    yield StreamChunk::content(&response.content);
11535                    record_runtime_stream_final(&terminal, response);
11536                    yield StreamChunk::Done {};
11537                    return;
11538                }
11539                Ok(SkillRouteResult::NoMatch) => {} // no skill matched, continue
11540                Err(e) => {
11541                    yield StreamChunk::error(e.to_string());
11542                    return;
11543                }
11544            }
11545
11546            // Reasoning mode determination
11547            let effective_reasoning = self.get_effective_reasoning_config();
11548            let reasoning_mode = match self.determine_reasoning_mode(processed_input).await {
11549                Ok(mode) => mode,
11550                Err(e) => {
11551                    yield StreamChunk::error(e.to_string());
11552                    return;
11553                }
11554            };
11555            let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
11556
11557            info!(
11558                reasoning_mode = ?reasoning_mode,
11559                auto_detected = auto_detected,
11560                "Reasoning mode determined (stream)"
11561            );
11562
11563            // Plan-and-Execute: yield final result as single chunk (not token-by-token)
11564            if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
11565                if let Err(e) = self.commit_root_user_message(processed_input).await {
11566                    yield StreamChunk::error(e.to_string());
11567                    return;
11568                }
11569                match self.handle_plan_and_execute(processed_input, &input_data.context, auto_detected).await {
11570                    Ok(resp) => {
11571                        yield StreamChunk::content(&resp.content);
11572                        record_runtime_stream_final(&terminal, resp);
11573                        yield StreamChunk::Done {};
11574                        return;
11575                    }
11576                    Err(e) => {
11577                        yield StreamChunk::error(e.to_string());
11578                        return;
11579                    }
11580                }
11581            }
11582
11583            if let Err(e) = self.commit_root_user_message(processed_input).await {
11584                yield StreamChunk::error(e.to_string());
11585                return;
11586            }
11587
11588            let llm = match self.get_state_llm() {
11589                Ok(llm) => llm,
11590                Err(e) => {
11591                    yield StreamChunk::error(e.to_string());
11592                    return;
11593                }
11594            };
11595
11596            let mut iterations = 0u32;
11597            let mut all_tool_calls: Vec<ToolCall> = Vec::new();
11598            let mut thinking_content: Option<String> = None;
11599
11600            loop {
11601                // When reasoning is active, cap iterations at the reasoning-specific limit.
11602                let effective_max = if reasoning_mode != ReasoningMode::None {
11603                    let rc = self.get_effective_reasoning_config();
11604                    self.max_iterations.min(rc.max_iterations)
11605                } else {
11606                    self.max_iterations
11607                };
11608
11609                if iterations >= effective_max {
11610                    let err_msg = format!("Max iterations ({}) exceeded", effective_max);
11611                    let err = AgentError::Other(err_msg.clone());
11612                    self.hooks.on_error(&err).await;
11613                    error!(iterations = iterations, "Max iterations exceeded (stream)");
11614                    yield StreamChunk::error(err_msg);
11615                    return;
11616                }
11617                iterations += 1;
11618                *self.iteration_count.write() = iterations;
11619
11620                debug!(iteration = iterations, max = effective_max, "LLM call (stream)");
11621
11622                let protocol = match self.main_tool_protocol(llm.as_ref(), false).await {
11623                    Ok(protocol) => protocol,
11624                    Err(e) => {
11625                        yield StreamChunk::error(e.to_string());
11626                        return;
11627                    }
11628                };
11629                let mut messages = match self
11630                    .build_messages_internal(true, None, protocol.choice.is_none())
11631                    .await
11632                {
11633                    Ok(m) => m,
11634                    Err(e) => {
11635                        yield StreamChunk::error(e.to_string());
11636                        return;
11637                    }
11638                };
11639                self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
11640
11641                self.hooks.on_llm_start(&messages).await;
11642                let llm_start = Instant::now();
11643
11644                // Check if reflection is active — if so, suppress streaming for this LLM call
11645                // because we may need to retry and the user would see a stale first attempt.
11646                let reflection_active = self
11647                    .should_reflect(processed_input, "")
11648                    .await
11649                    .unwrap_or_default();
11650
11651                let buffered_decision = reflection_active || protocol.choice.is_some();
11652                let content = if buffered_decision {
11653                    //
11654                    // Explicit tool choice buffers the provider decision so no text or tool call is visible before the runtime validates and commits it.
11655                    //
11656                    let response = match self
11657                        .complete_main_llm_with_recovery(
11658                            Arc::clone(&llm),
11659                            &messages,
11660                            &protocol,
11661                        )
11662                        .await
11663                    {
11664                        Ok(r) => r,
11665                        Err(e) => {
11666                            yield StreamChunk::error(e.to_string());
11667                            return;
11668                        }
11669                    };
11670                    let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
11671                    self.hooks.on_llm_complete(&response, llm_duration_ms).await;
11672                    response.content.trim().to_string()
11673                } else {
11674                    // Streaming LLM call
11675                    let llm_stream = match self
11676                        .observe_purpose(
11677                            ObservationPurpose::MainResponse,
11678                            llm.complete_stream(&messages, None),
11679                        )
11680                        .await
11681                    {
11682                        Ok(s) => s,
11683                        Err(e) => {
11684                            yield StreamChunk::error(e.to_string());
11685                            return;
11686                        }
11687                    };
11688                    let mut accumulated = String::new();
11689                    let mut stream_inner = llm_stream;
11690                    while let Some(chunk_result) = stream_inner.next().await {
11691                        match chunk_result {
11692                            Ok(chunk) => {
11693                                accumulated.push_str(&chunk.delta);
11694                                yield StreamChunk::content(chunk.delta);
11695                            }
11696                            Err(e) => {
11697                                yield StreamChunk::error(e.to_string());
11698                                return;
11699                            }
11700                        }
11701                    }
11702                    let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
11703                    // Construct LLMResponse for hooks
11704                    let llm_response = ai_agents_core::LLMResponse::new(
11705                        accumulated.trim(),
11706                        ai_agents_core::FinishReason::Stop,
11707                    );
11708                    self.hooks.on_llm_complete(&llm_response, llm_duration_ms).await;
11709                    accumulated.trim().to_string()
11710                };
11711
11712                // Tool call handling
11713                if let Some(tool_calls) = self.parse_main_tool_calls(&content, &protocol) {
11714                    let native_tool_call = Self::is_native_tool_call_content(&content);
11715                    // Emit tool events for streaming
11716                    // First check transitions (same as blocking path)
11717                    let transition_fired = match self.evaluate_transitions(processed_input, &content).await {
11718                        Ok(v) => v,
11719                        Err(e) => {
11720                            yield StreamChunk::error(e.to_string());
11721                            return;
11722                        }
11723                    };
11724                    if transition_fired {
11725                        let _ = self.memory.add_message(ChatMessage::assistant(
11726                            "(Transitioned to new state — tool call handled by workflow)",
11727                        )).await;
11728
11729                        if include_state_events
11730                            && let Some(state) = self.current_state()
11731                        {
11732                                yield StreamChunk::state_transition(None, state);
11733                            }
11734                        continue;
11735                    }
11736
11737                    // Store the assistant's tool-call message (same as blocking path)
11738                    let _ = self.memory.add_message(ChatMessage::assistant(&content)).await;
11739
11740                    // Execute tools with streaming events
11741                    let results = self.execute_tools_parallel(&tool_calls).await;
11742
11743                    for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
11744                        if include_tool_events {
11745                            yield StreamChunk::tool_start(&tool_call.id, &tool_call.name);
11746                        }
11747
11748                        match result {
11749                            Ok(output) => {
11750                                if include_tool_events {
11751                                    yield StreamChunk::tool_result(
11752                                        &tool_call.id,
11753                                        &tool_call.name,
11754                                        &output,
11755                                        true,
11756                                    );
11757                                }
11758                                let _ = self.memory
11759                                    .add_message(Self::tool_result_message(
11760                                        tool_call,
11761                                        &output,
11762                                        native_tool_call,
11763                                    ))
11764                                    .await;
11765                            }
11766                            Err(e) => {
11767                                if matches!(e, AgentError::HITLRejected(_)) {
11768                                    let _ = self.memory.add_message(ChatMessage::assistant(
11769                                        format!("The operation was rejected by the approver: {}", e),
11770                                    )).await;
11771                                    let response = AgentResponse {
11772                                        content: format!("Operation cancelled: {}", e),
11773                                        metadata: None,
11774                                        tool_calls: Some(all_tool_calls.clone()),
11775                                    };
11776                                    if let Err(finalize_error) = self.finish_turn_if_root(&response).await {
11777                                        yield StreamChunk::error(finalize_error.to_string());
11778                                        return;
11779                                    }
11780                                    let legacy_error = response.content.clone();
11781                                    record_runtime_stream_final(&terminal, response);
11782                                    yield StreamChunk::error(legacy_error);
11783                                    yield StreamChunk::Done {};
11784                                    return;
11785                                }
11786                                if include_tool_events {
11787                                    yield StreamChunk::tool_result(
11788                                        &tool_call.id,
11789                                        &tool_call.name,
11790                                        e.to_string(),
11791                                        false,
11792                                    );
11793                                }
11794                                let _ = self.memory
11795                                    .add_message(Self::tool_result_message(
11796                                        tool_call,
11797                                        &format!("Error: {}", e),
11798                                        native_tool_call,
11799                                    ))
11800                                    .await;
11801                            }
11802                        }
11803                        all_tool_calls.push(tool_call.clone());
11804
11805                        if include_tool_events {
11806                            yield StreamChunk::tool_end(&tool_call.id);
11807                        }
11808                    }
11809                    continue;
11810                }
11811
11812                // Extract thinking, process output
11813                let (extracted_thinking, answer) = self.extract_thinking(&content);
11814                if extracted_thinking.is_some() {
11815                    thinking_content = extracted_thinking;
11816                }
11817
11818                let output_data = match self.process_output(&answer, &input_data.context).await {
11819                    Ok(d) => d,
11820                    Err(e) => {
11821                        yield StreamChunk::error(e.to_string());
11822                        return;
11823                    }
11824                };
11825
11826                let final_content = if output_data.metadata.rejected {
11827                    output_data
11828                        .metadata
11829                        .rejection_reason
11830                        .unwrap_or_else(|| answer.to_string())
11831                } else {
11832                    output_data.content
11833                };
11834
11835                // Reflection (uses blocking LLM calls for retries)
11836                let (final_content, reflection_metadata) = match self
11837                    .run_reflection(&*llm, processed_input, final_content)
11838                    .await
11839                {
11840                    Ok(r) => r,
11841                    Err(e) => {
11842                        yield StreamChunk::error(e.to_string());
11843                        return;
11844                    }
11845                };
11846
11847                let final_content = self.format_response_with_thinking(
11848                    thinking_content.as_deref(),
11849                    &final_content,
11850                );
11851
11852                // Buffered decisions emit only the accepted final text.
11853                if buffered_decision {
11854                    yield StreamChunk::content(&final_content);
11855                }
11856
11857                // Post-loop: memory, transitions, post-transition re-generation.
11858                // For NeedsRedispatch, run_loop_internal handles the new state's full
11859                // dispatch and its result is yielded as a single non-streamed chunk.
11860                let post_result = match self
11861                    .post_loop_processing(processed_input, final_content)
11862                    .await
11863                {
11864                    Ok(r) => r,
11865                    Err(e) => {
11866                        yield StreamChunk::error(e.to_string());
11867                        return;
11868                    }
11869                };
11870
11871                let (final_content, transitioned) = match post_result {
11872                    PostLoopResult::NoTransition(content) => (content, false),
11873                    PostLoopResult::Transitioned(content) => (content, true),
11874                    PostLoopResult::NeedsRedispatch => {
11875                        const MAX_REDISPATCH_DEPTH: u32 = 3;
11876                        let current_depth = *self.redispatch_depth.read();
11877                        let content = if current_depth >= MAX_REDISPATCH_DEPTH {
11878                            warn!(
11879                                depth = current_depth,
11880                                "Post-transition re-dispatch depth limit reached (stream)"
11881                            );
11882                            let c = String::new();
11883                            let _ = self.memory.add_message(ChatMessage::assistant(&c)).await;
11884                            c
11885                        } else {
11886                            *self.redispatch_depth.write() += 1;
11887                            if let Some(context) = self.active_turn_context.write().as_mut() {
11888                                context.enter_redispatch();
11889                            }
11890                            info!(
11891                                depth = current_depth + 1,
11892                                "Re-dispatching for new state after transition (stream)"
11893                            );
11894                            let result = self.run_loop_internal(processed_input).await;
11895                            *self.redispatch_depth.write() -= 1;
11896                            if let Some(context) = self.active_turn_context.write().as_mut() {
11897                                context.exit_redispatch();
11898                            }
11899                            match result {
11900                                Ok(resp) => resp.content,
11901                                Err(e) => {
11902                                    yield StreamChunk::error(e.to_string());
11903                                    return;
11904                                }
11905                            }
11906                        };
11907                        (content, true)
11908                    }
11909                };
11910
11911                if transitioned {
11912                    if include_state_events
11913                        && let Some(state) = self.current_state()
11914                    {
11915                            yield StreamChunk::state_transition(None, state);
11916                        }
11917                    // Yield the post-transition re-generated or re-dispatched content.
11918                    yield StreamChunk::content(&final_content);
11919                }
11920
11921                // Build and finalize the same authoritative response shape before exposing the terminal event.
11922                let final_response = self.build_agent_response(AgentResponseParts {
11923                    content: final_content,
11924                    all_tool_calls,
11925                    reasoning_mode,
11926                    auto_detected,
11927                    iterations,
11928                    thinking: thinking_content,
11929                    reflection_metadata,
11930                });
11931                if let Err(e) = self.finish_turn_if_root(&final_response).await {
11932                    yield StreamChunk::error(e.to_string());
11933                    return;
11934                }
11935
11936                record_runtime_stream_final(&terminal, final_response);
11937                yield StreamChunk::Done {};
11938                return;
11939            }
11940        })
11941    }
11942
11943    /// Streams the root turn while keeping clarification and confirmation questions as terminal responses for their turn.
11944    /// Pending manager and skill ownership must survive until explicit confirmation returns a resolved result for redispatch.
11945    fn run_loop_stream<'a>(
11946        &'a self,
11947        input: &'a str,
11948        terminal: RuntimeStreamTerminalSlot,
11949    ) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
11950        Box::pin(async_stream::stream! {
11951            self.begin_root_turn();
11952            let _root_cleanup = RootTurnCleanup::new(self);
11953            self.hooks.on_message_received(input).await;
11954
11955            // One-shot context initialization (mirrors run_loop)
11956            if !self.context_initialized.swap(true, Ordering::SeqCst) {
11957                if let Err(e) = self.context_manager.initialize().await {
11958                    yield StreamChunk::error(e.to_string());
11959                    return;
11960                }
11961                debug!("Context manager initialized (defaults, env, builtins)");
11962            }
11963
11964            if let Err(e) = self.check_turn_timeout().await {
11965                yield StreamChunk::error(e.to_string());
11966                return;
11967            }
11968            if let Err(e) = self.context_manager.refresh_per_turn().await {
11969                yield StreamChunk::error(e.to_string());
11970                return;
11971            }
11972
11973            // Clear stale disambiguation context from previous turns.
11974            self.clear_disambiguation_context();
11975
11976            // Disambiguation check (before input processing)
11977            if let Some(ref disambiguator) = self.disambiguation_manager {
11978                let disambiguation_context = match self.build_disambiguation_context().await {
11979                    Ok(ctx) => ctx,
11980                    Err(e) => {
11981                        yield StreamChunk::error(e.to_string());
11982                        return;
11983                    }
11984                };
11985
11986                let state_override = self
11987                    .state_machine
11988                    .as_ref()
11989                    .and_then(|sm| sm.current_definition())
11990                    .and_then(|def| def.disambiguation.clone());
11991
11992                let state_generation = self
11993                    .state_machine
11994                    .as_ref()
11995                    .map(|state_machine| state_machine.generation());
11996                let disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
11997                let mut result = match self
11998                    .observe_purpose(
11999                        ObservationPurpose::DisambiguationDetection,
12000                        disambiguator.process_input_with_override(
12001                            input,
12002                            &disambiguation_context,
12003                            state_override.as_ref(),
12004                            None,
12005                        ),
12006                    )
12007                    .await
12008                {
12009                    Ok(r) => r,
12010                    Err(e) => {
12011                        yield StreamChunk::error(e.to_string());
12012                        return;
12013                    }
12014                };
12015                let current_state_generation = self
12016                    .state_machine
12017                    .as_ref()
12018                    .map(|state_machine| state_machine.generation());
12019                if current_state_generation != state_generation
12020                    || self.disambiguation_epoch.load(Ordering::SeqCst) != disambiguation_epoch
12021                {
12022                    disambiguator.clear_pending().await;
12023                    *self.pending_skill_id.write() = None;
12024                    result = DisambiguationResult::Abandoned { new_input: None };
12025                    info!(
12026                        confirmation_event = "invalidated",
12027                        invalidation_reason = "state_generation_changed",
12028                        "Streaming disambiguation result invalidated before redispatch"
12029                    );
12030                }
12031                match result {
12032                    DisambiguationResult::Clear => {
12033                        debug!("Input is clear, proceeding normally (stream)");
12034                    }
12035                    DisambiguationResult::NeedsClarification {
12036                        question,
12037                        detection,
12038                    } => {
12039                        let admission = match self
12040                            .admit_disambiguation_redispatch(
12041                                disambiguation_epoch,
12042                                state_generation,
12043                            )
12044                            .await
12045                        {
12046                            Ok(admission) => admission,
12047                            Err(error) => {
12048                                *self.pending_skill_id.write() = None;
12049                                yield StreamChunk::error(error.to_string());
12050                                return;
12051                            }
12052                        };
12053                        let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
12054                        info!(
12055                            ambiguity_type = ?detection.ambiguity_type,
12056                            confidence = detection.confidence,
12057                            "Input requires clarification (stream)"
12058                        );
12059                        // Confirmation uses the same terminal branch so enriched input cannot stream before explicit agreement.
12060                        // Keep pending_skill_id intact for the later confirmed redispatch.
12061                        if let Err(e) = self.commit_root_user_message(input).await {
12062                            yield StreamChunk::error(e.to_string());
12063                            return;
12064                        }
12065                        let _ = self
12066                            .memory
12067                            .add_message(ChatMessage::assistant(&question.question))
12068                            .await;
12069                        let status = if awaiting_confirmation {
12070                            "awaiting_confirmation"
12071                        } else {
12072                            "awaiting_clarification"
12073                        };
12074                        let response = AgentResponse::new(&question.question).with_metadata(
12075                            "disambiguation",
12076                            serde_json::json!({ "status": status }),
12077                        );
12078                        drop(admission);
12079                        if let Err(e) = self.finish_turn_if_root(&response).await {
12080                            yield StreamChunk::error(e.to_string());
12081                            return;
12082                        }
12083                        yield StreamChunk::content(&question.question);
12084                        record_runtime_stream_final(&terminal, response);
12085                        yield StreamChunk::Done {};
12086                        return;
12087                    }
12088                    DisambiguationResult::Clarified {
12089                        enriched_input,
12090                        resolved,
12091                        ..
12092                    } => {
12093                        let admission = match self
12094                            .admit_disambiguation_redispatch(
12095                                disambiguation_epoch,
12096                                state_generation,
12097                            )
12098                            .await
12099                        {
12100                            Ok(admission) => admission,
12101                            Err(error) => {
12102                                *self.pending_skill_id.write() = None;
12103                                yield StreamChunk::error(error.to_string());
12104                                return;
12105                            }
12106                        };
12107                        info!(
12108                            resolved_count = resolved.len(),
12109                            enriched = %enriched_input,
12110                            "Input clarified (stream)"
12111                        );
12112                        for (key, value) in &resolved {
12113                            let context_key = format!("disambiguation.{}", key);
12114                            let _ = self.context_manager.set(&context_key, value.clone());
12115                        }
12116                        if let Some(intent) = resolved.get("intent") {
12117                            let _ = self.context_manager.set("resolved_intent", intent.clone());
12118                        }
12119                        let _ = self
12120                            .context_manager
12121                            .set("disambiguation.resolved", serde_json::Value::Bool(true));
12122
12123                        // Check if this clarification was triggered by a skill-level override.
12124                        // Re-run skill disambiguation to verify all required_clarity fields
12125                        // are present before executing.
12126                        let skill_id = self.pending_skill_id.read().clone();
12127                        if let Some(skill_id) = skill_id {
12128                            info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input (stream)");
12129                            drop(admission);
12130                            match self
12131                                .recheck_skill_disambiguation(
12132                                    &skill_id,
12133                                    &enriched_input,
12134                                    disambiguation_epoch,
12135                                    state_generation,
12136                                )
12137                                .await
12138                            {
12139                                Ok(resp) => {
12140                                    yield StreamChunk::content(&resp.content);
12141                                    record_runtime_stream_final(&terminal, resp);
12142                                    yield StreamChunk::Done {};
12143                                    return;
12144                                }
12145                                Err(e) => {
12146                                    yield StreamChunk::error(e.to_string());
12147                                    return;
12148                                }
12149                            }
12150                        }
12151
12152                        // Forward to internal stream with enriched input.
12153                        drop(admission);
12154                        let mut inner = self.run_loop_internal_stream(
12155                            &enriched_input,
12156                            Arc::clone(&terminal),
12157                        );
12158                        while let Some(chunk) = inner.next().await {
12159                            yield chunk;
12160                        }
12161                        return;
12162                    }
12163                    DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
12164                        info!("Proceeding with best guess (stream)");
12165
12166                        // Same skill-id re-check for best-guess path
12167                        let skill_id = self.pending_skill_id.read().clone();
12168                        if let Some(skill_id) = skill_id {
12169                            info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input (stream)");
12170                            match self
12171                                .recheck_skill_disambiguation(
12172                                    &skill_id,
12173                                    &enriched_input,
12174                                    disambiguation_epoch,
12175                                    state_generation,
12176                                )
12177                                .await
12178                            {
12179                                Ok(resp) => {
12180                                    yield StreamChunk::content(&resp.content);
12181                                    record_runtime_stream_final(&terminal, resp);
12182                                    yield StreamChunk::Done {};
12183                                    return;
12184                                }
12185                                Err(e) => {
12186                                    yield StreamChunk::error(e.to_string());
12187                                    return;
12188                                }
12189                            }
12190                        }
12191
12192                        let mut inner = self.run_loop_internal_stream(
12193                            &enriched_input,
12194                            Arc::clone(&terminal),
12195                        );
12196                        while let Some(chunk) = inner.next().await {
12197                            yield chunk;
12198                        }
12199                        return;
12200                    }
12201                    DisambiguationResult::GiveUp { reason } => {
12202                        *self.pending_skill_id.write() = None;
12203                        warn!(reason = %reason, "Disambiguation gave up (stream)");
12204                        let apology = self
12205                            .generate_localized_apology(
12206                                "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
12207                                &reason,
12208                            )
12209                            .await
12210                            .unwrap_or_else(|_| {
12211                                format!("I'm sorry, I couldn't understand your request: {}", reason)
12212                            });
12213                        let response = AgentResponse::new(&apology);
12214                        if let Err(e) = self.finish_turn_if_root(&response).await {
12215                            yield StreamChunk::error(e.to_string());
12216                            return;
12217                        }
12218                        yield StreamChunk::content(&apology);
12219                        record_runtime_stream_final(&terminal, response);
12220                        yield StreamChunk::Done {};
12221                        return;
12222                    }
12223                    DisambiguationResult::Escalate { reason } => {
12224                        *self.pending_skill_id.write() = None;
12225                        info!(reason = %reason, "Escalating to human (stream)");
12226                        if let Some(ref hitl) = self.hitl_engine {
12227                            let trigger =
12228                                ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
12229                            let mut context_map = HashMap::new();
12230                            context_map.insert("original_input".to_string(), serde_json::json!(input));
12231                            context_map.insert("reason".to_string(), serde_json::json!(&reason));
12232                            let check_result = HITLCheckResult::required(
12233                                trigger,
12234                                context_map,
12235                                format!("User request needs human assistance: {}", reason),
12236                                Some(hitl.config().default_timeout_seconds),
12237                            );
12238                            match self.request_hitl_approval(check_result).await {
12239                                Ok(ApprovalResult::Approved | ApprovalResult::Modified { .. }) => {
12240                                    let mut inner = self.run_loop_internal_stream(
12241                                        input,
12242                                        Arc::clone(&terminal),
12243                                    );
12244                                    while let Some(chunk) = inner.next().await {
12245                                        yield chunk;
12246                                    }
12247                                    return;
12248                                }
12249                                Ok(_) => {}
12250                                Err(e) => {
12251                                    yield StreamChunk::error(e.to_string());
12252                                    return;
12253                                }
12254                            }
12255                        }
12256                        let apology = self
12257                            .generate_localized_apology(
12258                                "Explain briefly that you're transferring the user to a human agent for help.",
12259                                &reason,
12260                            )
12261                            .await
12262                            .unwrap_or_else(|_| {
12263                                format!("I need human assistance to help with your request: {}", reason)
12264                            });
12265                        let response = AgentResponse::new(&apology);
12266                        if let Err(e) = self.finish_turn_if_root(&response).await {
12267                            yield StreamChunk::error(e.to_string());
12268                            return;
12269                        }
12270                        yield StreamChunk::content(&apology);
12271                        record_runtime_stream_final(&terminal, response);
12272                        yield StreamChunk::Done {};
12273                        return;
12274                    }
12275                    DisambiguationResult::Abandoned { new_input } => {
12276                        *self.pending_skill_id.write() = None;
12277
12278                        info!(
12279                            has_new_input = new_input.is_some(),
12280                            "Clarification abandoned by user (stream)"
12281                        );
12282
12283                        if let Err(e) = self.commit_root_user_message(input).await {
12284                            yield StreamChunk::error(e.to_string());
12285                            return;
12286                        }
12287
12288                        match new_input {
12289                            Some(fresh_input) => {
12290                                // Topic switch: forward to internal stream with fresh input.
12291                                let mut inner = self.run_loop_internal_stream(
12292                                    &fresh_input,
12293                                    Arc::clone(&terminal),
12294                                );
12295                                while let Some(chunk) = inner.next().await {
12296                                    yield chunk;
12297                                }
12298                                return;
12299                            }
12300                            None => {
12301                                // Pure abandonment: generate a brief acknowledgment.
12302                                let ack = self
12303                                    .generate_localized_apology(
12304                                        "The user changed their mind about their previous request. \
12305                                         Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
12306                                         Do NOT apologize excessively. Be concise.",
12307                                        "User abandoned clarification",
12308                                    )
12309                                    .await
12310                                    .unwrap_or_else(|_| {
12311                                        "OK, no problem. What else can I help with?".to_string()
12312                                    });
12313
12314                                let _ = self
12315                                    .memory
12316                                    .add_message(ChatMessage::assistant(&ack))
12317                                    .await;
12318
12319                                let response = AgentResponse::new(&ack);
12320                                if let Err(e) = self.finish_turn_if_root(&response).await {
12321                                    yield StreamChunk::error(e.to_string());
12322                                    return;
12323                                }
12324                                yield StreamChunk::content(&ack);
12325                                record_runtime_stream_final(&terminal, response);
12326                                yield StreamChunk::Done {};
12327                                return;
12328                            }
12329                        }
12330                    }
12331                }
12332            }
12333
12334            // No disambiguation or Clear result — proceed with internal stream
12335            let mut inner = self.run_loop_internal_stream(input, Arc::clone(&terminal));
12336            while let Some(chunk) = inner.next().await {
12337                yield chunk;
12338            }
12339        })
12340    }
12341
12342    pub fn info(&self) -> AgentInfo {
12343        self.info.clone()
12344    }
12345
12346    pub fn skills(&self) -> &[SkillDefinition] {
12347        &self.skills
12348    }
12349
12350    /// Clears conversation and pending runtime ownership through one reset contract.
12351    async fn reset_runtime_state(&self) -> Result<()> {
12352        let _admission = self.disambiguation_admission.write().await;
12353        if self.state_transition_reserved.load(Ordering::SeqCst) {
12354            return Err(AgentError::Other(
12355                "Cannot reset while a state transition is in progress".to_string(),
12356            ));
12357        }
12358        self.disambiguation_epoch.fetch_add(1, Ordering::SeqCst);
12359        *self.pending_skill_id.write() = None;
12360        if let Some(disambiguator) = self.disambiguation_manager.as_ref() {
12361            disambiguator.clear_pending().await;
12362        }
12363        self.memory.clear().await?;
12364        *self.iteration_count.write() = 0;
12365        self.tool_call_history.write().clear();
12366        if let Some(ref sm) = self.state_machine {
12367            sm.reset();
12368        }
12369        Ok(())
12370    }
12371
12372    /// Resets the runtime using the same cleanup path as the Agent trait.
12373    pub async fn reset(&self) -> Result<()> {
12374        self.reset_runtime_state().await
12375    }
12376
12377    pub fn max_context_tokens(&self) -> u32 {
12378        self.max_context_tokens
12379    }
12380
12381    pub fn llm_registry(&self) -> &Arc<LLMRegistry> {
12382        &self.llm_registry
12383    }
12384
12385    pub fn state_machine(&self) -> Option<&Arc<StateMachine>> {
12386        self.state_machine.as_ref()
12387    }
12388
12389    pub fn context_manager(&self) -> &Arc<ContextManager> {
12390        &self.context_manager
12391    }
12392
12393    pub fn tool_call_history(&self) -> Vec<ToolCallRecord> {
12394        self.tool_call_history.read().clone()
12395    }
12396
12397    pub fn memory_token_budget(&self) -> Option<&MemoryTokenBudget> {
12398        self.memory_token_budget.as_ref()
12399    }
12400
12401    pub fn parallel_tools_config(&self) -> &ParallelToolsConfig {
12402        &self.parallel_tools
12403    }
12404
12405    pub fn streaming_config(&self) -> &StreamingConfig {
12406        &self.streaming
12407    }
12408
12409    pub fn hooks(&self) -> &Arc<dyn AgentHooks> {
12410        &self.hooks
12411    }
12412
12413    pub fn hitl_engine(&self) -> Option<&HITLEngine> {
12414        self.hitl_engine.as_ref()
12415    }
12416
12417    pub fn approval_handler(&self) -> &Arc<dyn ApprovalHandler> {
12418        &self.approval_handler
12419    }
12420
12421    /// Build a context map with language hints from context_manager for HITL message localization.
12422    fn build_hitl_language_context(&self) -> HashMap<String, Value> {
12423        let mut ctx = HashMap::new();
12424        for key in &["user.language", "input.detected.language", "language"] {
12425            if let Some(val) = self.context_manager.get(key) {
12426                ctx.insert(key.to_string(), val);
12427            }
12428        }
12429        ctx
12430    }
12431
12432    /// Send a HITL check result through the approval flow and return the full ApprovalResult.
12433    async fn request_hitl_approval(&self, check_result: HITLCheckResult) -> Result<ApprovalResult> {
12434        let Some(request) = check_result.into_request() else {
12435            return Ok(ApprovalResult::Approved);
12436        };
12437
12438        self.hooks.on_approval_requested(&request).await;
12439
12440        let timeout = request.timeout;
12441
12442        let raw_result = if let Some(duration) = timeout {
12443            match tokio::time::timeout(
12444                duration,
12445                self.approval_handler.request_approval(request.clone()),
12446            )
12447            .await
12448            {
12449                Ok(result) => result,
12450                Err(_) => ApprovalResult::timeout(),
12451            }
12452        } else {
12453            self.approval_handler
12454                .request_approval(request.clone())
12455                .await
12456        };
12457
12458        self.hooks
12459            .on_approval_result(&request.id, &raw_result)
12460            .await;
12461
12462        let (outcome, effective_result): (ApprovalResolvedOutcome, Result<ApprovalResult>) =
12463            match &raw_result {
12464                ApprovalResult::Approved => (
12465                    ApprovalResolvedOutcome::Approved,
12466                    Ok(ApprovalResult::Approved),
12467                ),
12468                ApprovalResult::Rejected { reason } => (
12469                    ApprovalResolvedOutcome::Rejected {
12470                        reason: reason.clone(),
12471                    },
12472                    Ok(ApprovalResult::Rejected {
12473                        reason: reason.clone(),
12474                    }),
12475                ),
12476                ApprovalResult::Modified { changes } => (
12477                    ApprovalResolvedOutcome::Modified {
12478                        changes: changes.clone(),
12479                    },
12480                    Ok(ApprovalResult::Modified {
12481                        changes: changes.clone(),
12482                    }),
12483                ),
12484                ApprovalResult::Timeout => {
12485                    if let Some(ref engine) = self.hitl_engine {
12486                        match engine.config().on_timeout {
12487                            TimeoutAction::Approve => (
12488                                ApprovalResolvedOutcome::Approved,
12489                                Ok(ApprovalResult::Approved),
12490                            ),
12491                            TimeoutAction::Reject => {
12492                                let reason = Some("Timeout".to_string());
12493                                (
12494                                    ApprovalResolvedOutcome::Rejected {
12495                                        reason: reason.clone(),
12496                                    },
12497                                    Ok(ApprovalResult::Rejected { reason }),
12498                                )
12499                            }
12500                            TimeoutAction::Error => {
12501                                let message = "HITL approval timeout".to_string();
12502                                (
12503                                    ApprovalResolvedOutcome::Error {
12504                                        message: message.clone(),
12505                                    },
12506                                    Err(AgentError::Other(message)),
12507                                )
12508                            }
12509                        }
12510                    } else {
12511                        let reason = Some("Timeout (no engine)".to_string());
12512                        (
12513                            ApprovalResolvedOutcome::Rejected {
12514                                reason: reason.clone(),
12515                            },
12516                            Ok(ApprovalResult::Rejected { reason }),
12517                        )
12518                    }
12519                }
12520            };
12521
12522        self.hooks
12523            .on_approval_resolved(&request, &raw_result, &outcome)
12524            .await;
12525
12526        effective_result
12527    }
12528
12529    pub async fn check_state_hitl(&self, from: Option<&str>, to: &str) -> Result<bool> {
12530        if let Some(ref hitl_engine) = self.hitl_engine {
12531            let hitl_lang_ctx = self.build_hitl_language_context();
12532            let check_result = self
12533                .observe_purpose(
12534                    ObservationPurpose::HitlLocalization,
12535                    hitl_engine.check_state_transition_with_localization(
12536                        from,
12537                        to,
12538                        &hitl_lang_ctx,
12539                        self.approval_handler.as_ref(),
12540                        Some(&self.llm_registry),
12541                    ),
12542                )
12543                .await?;
12544            if check_result.is_required() {
12545                let result = self.request_hitl_approval(check_result).await?;
12546                return Ok(matches!(
12547                    result,
12548                    ApprovalResult::Approved | ApprovalResult::Modified { .. }
12549                ));
12550            }
12551        }
12552        Ok(true)
12553    }
12554
12555    /// Execute multiple tools in parallel
12556    async fn execute_tools_parallel(
12557        &self,
12558        tool_calls: &[ToolCall],
12559    ) -> Vec<(String, Result<String>)> {
12560        let can_run_parallel = tool_calls.iter().all(|tc| {
12561            self.tools
12562                .resolve(&tc.name)
12563                .map(|resolved| resolved.tool.classify_call(&tc.arguments).concurrency_safe)
12564                .unwrap_or(false)
12565        });
12566
12567        if !self.parallel_tools.enabled || tool_calls.len() <= 1 || !can_run_parallel {
12568            let mut results = Vec::new();
12569            for tc in tool_calls {
12570                let result = self
12571                    .observe_purpose(
12572                        current_observation_context()
12573                            .map(|context| context.purpose)
12574                            .unwrap_or_default(),
12575                        self.execute_tool_smart(tc),
12576                    )
12577                    .await;
12578                results.push((tc.id.clone(), result));
12579            }
12580            return results;
12581        }
12582
12583        let chunks: Vec<_> = tool_calls
12584            .chunks(self.parallel_tools.max_parallel)
12585            .collect();
12586
12587        let mut all_results = Vec::new();
12588
12589        for chunk in chunks {
12590            let futures: Vec<_> = chunk
12591                .iter()
12592                .map(|tc| {
12593                    let tc = tc.clone();
12594                    async move {
12595                        let result = self.execute_tool_smart(&tc).await;
12596                        (tc.id.clone(), result)
12597                    }
12598                })
12599                .collect();
12600
12601            let results = futures::future::join_all(futures).await;
12602            all_results.extend(results);
12603        }
12604
12605        all_results
12606    }
12607
12608    /// Streams one serialized root turn and releases its owned gate guard at Done or when the stream is dropped.
12609    ///
12610    /// The captured immutable identity stack is restored for every inner poll and export with an `Arc` clone, so nested hooks and orchestration calls detect same-runtime reentry without rebuilding ancestry. The inner stream is drained after Done before ownership is released.
12611    pub async fn chat_stream<'a>(
12612        &'a self,
12613        input: &'a str,
12614    ) -> Result<Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>>> {
12615        let RootTurnAdmission {
12616            guard: root_turn_guard,
12617            identity_stack,
12618        } = self.acquire_root_turn().await?;
12619        //
12620        // Streaming readiness runs inside the captured ownership stack while the gate is held so initialization cannot overlap or recursively enter another turn.
12621        //
12622        scope_runtime_gate_identity_stack(&identity_stack, self.init_storage()).await?;
12623        info!(input_len = input.len(), "Starting streaming chat");
12624        let terminal = new_runtime_stream_terminal_slot();
12625        let inner = self.run_loop_stream(input, terminal);
12626        let observation_context = self.build_observation_context(None);
12627        let stream: Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> =
12628            Box::pin(async_stream::stream! {
12629                let mut root_turn_guard = Some(root_turn_guard);
12630                let mut inner = inner;
12631                loop {
12632                    let next = scope_runtime_gate_identity_stack(&identity_stack, async {
12633                        if let Some(context) = observation_context.as_ref() {
12634                            with_observation_context(context.clone(), inner.next()).await
12635                        } else {
12636                            inner.next().await
12637                        }
12638                    })
12639                    .await;
12640                    match next {
12641                        Some(StreamChunk::Done {}) => {
12642                            while scope_runtime_gate_identity_stack(&identity_stack, async {
12643                                if let Some(context) = observation_context.as_ref() {
12644                                    with_observation_context(context.clone(), inner.next())
12645                                        .await
12646                                        .is_some()
12647                                } else {
12648                                    inner.next().await.is_some()
12649                                }
12650                            })
12651                            .await
12652                            {}
12653                            if observation_context.is_some() {
12654                                scope_runtime_gate_identity_stack(
12655                                    &identity_stack,
12656                                    self.export_observability_if_configured(),
12657                                )
12658                                .await;
12659                            }
12660                            drop(root_turn_guard.take());
12661                            yield StreamChunk::Done {};
12662                            return;
12663                        }
12664                        Some(chunk) => yield chunk,
12665                        None => {
12666                            if observation_context.is_some() {
12667                                scope_runtime_gate_identity_stack(
12668                                    &identity_stack,
12669                                    self.export_observability_if_configured(),
12670                                )
12671                                .await;
12672                            }
12673                            drop(root_turn_guard.take());
12674                            return;
12675                        }
12676                    }
12677                }
12678            });
12679        Ok(stream)
12680    }
12681
12682    /// Streams one serialized root turn and releases its owned gate guard at the authoritative terminal event or on drop.
12683    ///
12684    /// The captured immutable identity stack is restored for every inner poll and export with an `Arc` clone. The event API shares and drains the legacy execution stream so cleanup, nested-call detection, cancellation, and side effects cannot diverge before the gate is released.
12685    pub async fn chat_stream_events<'a>(
12686        &'a self,
12687        input: &'a str,
12688    ) -> Result<Pin<Box<dyn Stream<Item = AgentStreamEvent> + Send + 'a>>> {
12689        let RootTurnAdmission {
12690            guard: root_turn_guard,
12691            identity_stack,
12692        } = self.acquire_root_turn().await?;
12693        //
12694        // Streaming readiness runs inside the captured ownership stack while the gate is held so initialization cannot overlap or recursively enter another turn.
12695        //
12696        scope_runtime_gate_identity_stack(&identity_stack, self.init_storage()).await?;
12697        info!(input_len = input.len(), "Starting streaming chat events");
12698        let terminal = new_runtime_stream_terminal_slot();
12699        let mut inner = self.run_loop_stream(input, Arc::clone(&terminal));
12700        let observation_context = self.build_observation_context(None);
12701        let stream: Pin<Box<dyn Stream<Item = AgentStreamEvent> + Send + 'a>> =
12702            Box::pin(async_stream::stream! {
12703                let mut root_turn_guard = Some(root_turn_guard);
12704                loop {
12705                    let next = scope_runtime_gate_identity_stack(&identity_stack, async {
12706                        if let Some(context) = observation_context.as_ref() {
12707                            with_observation_context(context.clone(), inner.next()).await
12708                        } else {
12709                            inner.next().await
12710                        }
12711                    })
12712                    .await;
12713                    match next {
12714                        Some(StreamChunk::Done {}) => {
12715                            let terminal_event = { terminal.write().take() };
12716                            if let Some(response) = terminal_event {
12717                                while scope_runtime_gate_identity_stack(&identity_stack, async {
12718                                    if let Some(context) = observation_context.as_ref() {
12719                                        with_observation_context(context.clone(), inner.next())
12720                                            .await
12721                                            .is_some()
12722                                    } else {
12723                                        inner.next().await.is_some()
12724                                    }
12725                                })
12726                                .await
12727                                {}
12728                                if observation_context.is_some() {
12729                                    scope_runtime_gate_identity_stack(
12730                                        &identity_stack,
12731                                        self.export_observability_if_configured(),
12732                                    )
12733                                    .await;
12734                                }
12735                                drop(root_turn_guard.take());
12736                                yield AgentStreamEvent::Final(response);
12737                                return;
12738                            }
12739                        }
12740                        Some(StreamChunk::Error { message }) => {
12741                            let finalized = { terminal.read().is_some() };
12742                            if finalized {
12743                                continue;
12744                            }
12745                            while scope_runtime_gate_identity_stack(&identity_stack, async {
12746                                if let Some(context) = observation_context.as_ref() {
12747                                    with_observation_context(context.clone(), inner.next())
12748                                        .await
12749                                        .is_some()
12750                                } else {
12751                                    inner.next().await.is_some()
12752                                }
12753                            })
12754                            .await
12755                            {}
12756                            if observation_context.is_some() {
12757                                scope_runtime_gate_identity_stack(
12758                                    &identity_stack,
12759                                    self.export_observability_if_configured(),
12760                                )
12761                                .await;
12762                            }
12763                            drop(root_turn_guard.take());
12764                            yield AgentStreamEvent::Chunk(StreamChunk::Error { message });
12765                            return;
12766                        }
12767                        Some(chunk) => yield AgentStreamEvent::Chunk(chunk),
12768                        None => {
12769                            if observation_context.is_some() {
12770                                scope_runtime_gate_identity_stack(
12771                                    &identity_stack,
12772                                    self.export_observability_if_configured(),
12773                                )
12774                                .await;
12775                            }
12776                            drop(root_turn_guard.take());
12777                            return;
12778                        }
12779                    }
12780                }
12781            });
12782        Ok(stream)
12783    }
12784}
12785
12786#[async_trait]
12787impl ToolInvoker for RuntimeAgent {
12788    async fn invoke_tool(&self, request: ToolExecutionRequest) -> Result<ToolExecutionRecord> {
12789        self.execute_tool_record(request).await
12790    }
12791}
12792
12793#[async_trait]
12794impl Agent for RuntimeAgent {
12795    /// Runs one blocking external root turn with task-local ownership visible through finalization, hooks, orchestration, and export.
12796    async fn chat(&self, input: &str) -> Result<AgentResponse> {
12797        let RootTurnAdmission {
12798            guard,
12799            identity_stack,
12800        } = self.acquire_root_turn().await?;
12801        let result = scope_runtime_gate_identity_stack(&identity_stack, async {
12802            let result = if let Some(context) = self.build_observation_context(None) {
12803                with_observation_context(context, self.run_loop(input)).await
12804            } else {
12805                self.run_loop(input).await
12806            };
12807            self.export_observability_if_configured().await;
12808            result
12809        })
12810        .await;
12811        drop(guard);
12812        result
12813    }
12814
12815    fn info(&self) -> AgentInfo {
12816        self.info.clone()
12817    }
12818
12819    /// Resets the runtime without leaving clarification or skill ownership behind.
12820    async fn reset(&self) -> Result<()> {
12821        self.reset_runtime_state().await
12822    }
12823}
12824
12825//
12826// Render a concurrent input template using direct minijinja.
12827// Same approach as pipeline's render_stage_template so variables are top-level.
12828//
12829// Available variables:
12830//   {{ user_input }}    - the user's actual message
12831//   {{ context.<key> }} - values from the context manager
12832//
12833/// Builds safe runtime tags for background maintenance lifecycle events.
12834fn background_maintenance_tags(
12835    label: &str,
12836    stage: &str,
12837    reason: Option<&str>,
12838    policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
12839) -> HashMap<String, String> {
12840    let mut tags = HashMap::new();
12841    tags.insert("runtime.background".to_string(), "true".to_string());
12842    tags.insert("runtime.maintenance".to_string(), label.to_string());
12843    tags.insert("runtime.maintenance_stage".to_string(), stage.to_string());
12844    if let Some(policy) = policy {
12845        tags.insert(
12846            "runtime.await_before_next_turn".to_string(),
12847            await_before_next_turn_label(policy.await_before_next_turn).to_string(),
12848        );
12849        tags.insert(
12850            "runtime.maintenance_mode".to_string(),
12851            maintenance_mode_label(policy.mode).to_string(),
12852        );
12853    }
12854    if let Some(reason) = reason {
12855        tags.insert("runtime.reason".to_string(), reason.to_string());
12856    }
12857    tags
12858}
12859
12860fn await_before_next_turn_label(policy: AwaitBeforeNextTurn) -> &'static str {
12861    match policy {
12862        AwaitBeforeNextTurn::Never => "never",
12863        AwaitBeforeNextTurn::SameActor => "same_actor",
12864        AwaitBeforeNextTurn::Always => "always",
12865    }
12866}
12867
12868fn maintenance_mode_label(mode: MaintenanceMode) -> &'static str {
12869    match mode {
12870        MaintenanceMode::InlineSerial => "inline_serial",
12871        MaintenanceMode::InlineParallel => "inline_parallel",
12872        MaintenanceMode::Background => "background",
12873    }
12874}
12875
12876/// Records a background maintenance lifecycle event when observability is enabled.
12877fn record_background_maintenance_event(
12878    manager: Option<&Arc<ObservabilityManager>>,
12879    label: &str,
12880    status: EventStatus,
12881    duration_ms: u64,
12882    stage: &str,
12883    reason: Option<String>,
12884    policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
12885) {
12886    if let Some(manager) = manager {
12887        manager.record_lifecycle_event(
12888            EventType::MemoryOperation {
12889                operation: format!("{}_background_{}", label, stage),
12890            },
12891            ObservationPurpose::Other(format!("{}_maintenance", label)),
12892            status,
12893            duration_ms,
12894            background_maintenance_tags(label, stage, reason.as_deref(), policy),
12895            None,
12896        );
12897    }
12898}
12899
12900fn effective_maintenance_mode(mode: MaintenanceMode, force_parallel: bool) -> MaintenanceMode {
12901    if force_parallel && matches!(mode, MaintenanceMode::InlineSerial) {
12902        MaintenanceMode::InlineParallel
12903    } else {
12904        mode
12905    }
12906}
12907
12908fn observation_purpose_for_process(hint: ProcessPurposeHint) -> ObservationPurpose {
12909    match hint {
12910        ProcessPurposeHint::Detect => ObservationPurpose::ProcessDetect,
12911        ProcessPurposeHint::Extract => ObservationPurpose::ProcessExtract,
12912        ProcessPurposeHint::Validate => ObservationPurpose::ProcessValidate,
12913        ProcessPurposeHint::Transform | ProcessPurposeHint::Other => {
12914            ObservationPurpose::ProcessTransform
12915        }
12916    }
12917}
12918
12919fn new_tool_resource_locks() -> ToolResourceLocks {
12920    Arc::new(RwLock::new(HashMap::new()))
12921}
12922
12923//
12924// Path-bound mutations share one conservative lock so aliases and parent-child paths cannot bypass serialization.
12925// Exact domain keys remain available for non-path side effects, and unbound effects share one fallback lock.
12926//
12927fn tool_resource_lock_keys(
12928    _canonical_id: &str,
12929    args: &Value,
12930    bindings: &ai_agents_core::ToolPolicyBindings,
12931    classification: &ai_agents_core::ToolCallClassification,
12932) -> Vec<String> {
12933    if classification.concurrency_safe {
12934        return Vec::new();
12935    }
12936
12937    let mut keys = Vec::new();
12938    let mut has_path_resource = false;
12939    for binding in &bindings.path_fields {
12940        let value = value_at_argument_path(args, &binding.field)
12941            .cloned()
12942            .or_else(|| {
12943                binding
12944                    .default_path
12945                    .as_ref()
12946                    .map(|path| Value::String(path.clone()))
12947            });
12948        if let Some(value) = value {
12949            collect_resource_strings(&value, |_| {
12950                has_path_resource = true;
12951            });
12952        }
12953    }
12954    for binding in &bindings.domain_fields {
12955        if let Some(value) = value_at_argument_path(args, &binding.field) {
12956            collect_resource_strings(value, |domain| {
12957                let normalized = if binding.is_url {
12958                    normalized_url_resource_key(domain)
12959                } else {
12960                    domain.trim().trim_end_matches('.').to_ascii_lowercase()
12961                };
12962                keys.push(format!("domain:{}", normalized));
12963            });
12964        }
12965    }
12966    for binding in &bindings.command_fields {
12967        if !matches!(binding.kind, ai_agents_core::CommandBindingKind::Cwd) {
12968            continue;
12969        }
12970        if let Some(value) = value_at_argument_path(args, &binding.field) {
12971            collect_resource_strings(value, |_| {
12972                has_path_resource = true;
12973            });
12974        }
12975    }
12976    if has_path_resource {
12977        keys.push("path-mutation:global".to_string());
12978    }
12979    if keys.is_empty() {
12980        keys.push("side-effect:unbound".to_string());
12981    }
12982    keys.sort();
12983    keys.dedup();
12984    keys
12985}
12986
12987fn value_at_argument_path<'a>(value: &'a Value, field: &str) -> Option<&'a Value> {
12988    let mut current = value;
12989    for segment in field.split('.') {
12990        if segment.is_empty() {
12991            return None;
12992        }
12993        current = current.get(segment)?;
12994    }
12995    Some(current)
12996}
12997
12998fn collect_resource_strings(value: &Value, mut collect: impl FnMut(&str)) {
12999    match value {
13000        Value::String(value) => collect(value),
13001        Value::Array(values) => {
13002            for value in values {
13003                if let Some(value) = value.as_str() {
13004                    collect(value);
13005                }
13006            }
13007        }
13008        _ => {}
13009    }
13010}
13011
13012fn normalized_url_resource_key(value: &str) -> String {
13013    let value = value.trim();
13014    let Some((scheme, remainder)) = value.split_once("://") else {
13015        return value.to_ascii_lowercase();
13016    };
13017    let authority_end = remainder.find(['/', '?', '#']).unwrap_or(remainder.len());
13018    let (authority, suffix) = remainder.split_at(authority_end);
13019    format!(
13020        "{}://{}{}",
13021        scheme.to_ascii_lowercase(),
13022        authority.to_ascii_lowercase(),
13023        suffix
13024    )
13025}
13026
13027fn render_concurrent_template(
13028    template: &str,
13029    user_input: &str,
13030    context_values: &std::collections::HashMap<String, serde_json::Value>,
13031) -> Result<String> {
13032    let mut env = minijinja::Environment::new();
13033    env.add_template("concurrent", template)
13034        .map_err(|e| AgentError::Other(format!("Concurrent template parse error: {}", e)))?;
13035
13036    let mut ctx = std::collections::BTreeMap::new();
13037    ctx.insert("user_input".to_string(), minijinja::Value::from(user_input));
13038
13039    // Expose context manager values under {{ context.<key> }}.
13040    let context_obj = minijinja::Value::from_serialize(context_values);
13041    ctx.insert("context".to_string(), context_obj);
13042
13043    let tmpl = env
13044        .get_template("concurrent")
13045        .map_err(|e| AgentError::Other(format!("Concurrent template error: {}", e)))?;
13046
13047    tmpl.render(minijinja::Value::from_serialize(&ctx))
13048        .map_err(|e| AgentError::Other(format!("Concurrent template render error: {}", e)))
13049}
13050
13051#[cfg(test)]
13052mod tests {
13053    use super::*;
13054    use crate::AgentBuilder;
13055    use ai_agents_core::{LLMChunk, LLMConfig, LLMError, LLMFeature, Tool};
13056    use ai_agents_llm::mock::MockLLMProvider;
13057    use ai_agents_skills::{SkillDefinition, SkillStep};
13058    use ai_agents_tools::{
13059        CalculatorTool, CopyPathTool, DeletePathTool, FileWriteTool, MovePathTool, ToolAliases,
13060        ToolDescriptor, ToolProvider, ToolProviderError, ToolProviderType, WebFetchResolver,
13061        WebFetchTool, WebFetchTransport, WebFetchTransportRequest, WebFetchTransportResponse,
13062    };
13063
13064    fn mock_with_response(response: &str) -> MockLLMProvider {
13065        let mut mock = MockLLMProvider::new("test");
13066        mock.set_response(response);
13067        mock
13068    }
13069
13070    fn mock_with_responses(responses: Vec<&str>) -> MockLLMProvider {
13071        let mut mock = MockLLMProvider::new("test");
13072        mock.set_responses(responses.into_iter().map(String::from).collect(), true);
13073        mock
13074    }
13075
13076    /// Builds a two-state fixture so confirmation ownership can be invalidated by transition.
13077    fn disambiguation_state_machine(
13078        state_enabled: Option<bool>,
13079        require_confirmation: bool,
13080    ) -> Arc<StateMachine> {
13081        let definition = ai_agents_state::StateDefinition {
13082            prompt: Some("Handle the resolved request.".to_string()),
13083            disambiguation: Some(ai_agents_disambiguation::StateDisambiguationOverride {
13084                enabled: state_enabled,
13085                require_confirmation,
13086                ..Default::default()
13087            }),
13088            ..Default::default()
13089        };
13090        let review = ai_agents_state::StateDefinition {
13091            prompt: Some("Review a fresh request.".to_string()),
13092            ..Default::default()
13093        };
13094        Arc::new(
13095            StateMachine::new(ai_agents_state::StateConfig {
13096                initial: "active".to_string(),
13097                states: std::collections::HashMap::from([
13098                    ("active".to_string(), definition),
13099                    ("review".to_string(), review),
13100                ]),
13101                global_transitions: Vec::new(),
13102                fallback: None,
13103                max_no_transition: None,
13104                regenerate_on_transition: true,
13105            })
13106            .unwrap(),
13107        )
13108    }
13109
13110    /// Builds a state-aware disambiguation fixture without skills.
13111    fn state_disambiguation_agent(
13112        responses: Vec<&str>,
13113        manager_enabled: bool,
13114        state_enabled: Option<bool>,
13115        require_confirmation: bool,
13116    ) -> (RuntimeAgent, MockLLMProvider) {
13117        state_disambiguation_agent_with_skills(
13118            responses,
13119            manager_enabled,
13120            state_enabled,
13121            require_confirmation,
13122            Vec::new(),
13123        )
13124    }
13125
13126    /// Builds a state-aware disambiguation fixture with optional real skill routing.
13127    fn state_disambiguation_agent_with_skills(
13128        responses: Vec<&str>,
13129        manager_enabled: bool,
13130        state_enabled: Option<bool>,
13131        require_confirmation: bool,
13132        skills: Vec<SkillDefinition>,
13133    ) -> (RuntimeAgent, MockLLMProvider) {
13134        let mut mock = MockLLMProvider::new("state-confirmation");
13135        mock.set_responses(responses.into_iter().map(String::from).collect(), false);
13136        let observed = mock.clone();
13137        let agent = AgentBuilder::new()
13138            .system_prompt("Handle requests.")
13139            .llm(Arc::new(mock.clone()))
13140            .llm_alias("router", Arc::new(mock))
13141            .state_machine(disambiguation_state_machine(
13142                state_enabled,
13143                require_confirmation,
13144            ))
13145            .skills(skills)
13146            .build()
13147            .unwrap()
13148            .with_disambiguation(DisambiguationConfig {
13149                enabled: manager_enabled,
13150                ..Default::default()
13151            });
13152        (agent, observed)
13153    }
13154
13155    /// Defines a prompt skill whose provider call proves committed execution.
13156    fn confirmation_skill() -> SkillDefinition {
13157        SkillDefinition {
13158            id: "send_report".to_string(),
13159            description: "Send a report after clarification".to_string(),
13160            trigger: "When the user asks to send a report".to_string(),
13161            steps: vec![SkillStep::Prompt {
13162                prompt: "Execute confirmed report skill for: {{ input }}".to_string(),
13163                llm: None,
13164            }],
13165            reasoning: None,
13166            reflection: None,
13167            disambiguation: Some(ai_agents_disambiguation::SkillDisambiguationOverride {
13168                enabled: Some(true),
13169                ..Default::default()
13170            }),
13171        }
13172    }
13173
13174    /// Counts committed skill prompt calls without relying on final response wording.
13175    fn confirmation_skill_call_count(observed: &MockLLMProvider) -> usize {
13176        observed
13177            .call_history()
13178            .iter()
13179            .filter(|call| {
13180                call.messages
13181                    .iter()
13182                    .any(|message| message.content.contains("Execute confirmed report skill"))
13183            })
13184            .count()
13185    }
13186
13187    struct BlockingRuntimeConfirmationObserver {
13188        entered: tokio::sync::Barrier,
13189        release: tokio::sync::Notify,
13190    }
13191
13192    impl BlockingRuntimeConfirmationObserver {
13193        fn new() -> Self {
13194            Self {
13195                entered: tokio::sync::Barrier::new(2),
13196                release: tokio::sync::Notify::new(),
13197            }
13198        }
13199    }
13200
13201    struct ResetOnTransitionHooks {
13202        agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
13203        invoked: AtomicBool,
13204    }
13205
13206    #[async_trait]
13207    impl AgentHooks for ResetOnTransitionHooks {
13208        async fn on_state_transition(&self, _from: Option<&str>, _to: &str, _reason: &str) {
13209            if self.invoked.swap(true, Ordering::SeqCst) {
13210                return;
13211            }
13212            let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
13213            if let Some(agent) = agent {
13214                agent.reset().await.unwrap();
13215            }
13216        }
13217    }
13218
13219    impl ClarificationObserver for BlockingRuntimeConfirmationObserver {
13220        fn observe_question<'a>(
13221            &'a self,
13222            future: ClarificationQuestionFuture<'a>,
13223        ) -> ClarificationQuestionFuture<'a> {
13224            future
13225        }
13226
13227        fn observe_parse<'a>(
13228            &'a self,
13229            future: ClarificationParseFuture<'a>,
13230        ) -> ClarificationParseFuture<'a> {
13231            future
13232        }
13233
13234        fn observe_confirmation_parse<'a>(
13235            &'a self,
13236            future: ConfirmationParseFuture<'a>,
13237        ) -> ConfirmationParseFuture<'a> {
13238            Box::pin(async move {
13239                self.entered.wait().await;
13240                self.release.notified().await;
13241                future.await
13242            })
13243        }
13244    }
13245
13246    #[tokio::test]
13247    async fn state_confirmation_blocks_redispatch_until_explicit_agreement() {
13248        let (agent, observed) = state_disambiguation_agent(
13249            vec![
13250                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13251                r#"{"question":"What should I send?","options":null}"#,
13252                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13253                r#"{"question":"Should I send the report to Ada?"}"#,
13254                r#"{"status":"confirmed"}"#,
13255                "Request executed.",
13256            ],
13257            true,
13258            None,
13259            true,
13260        );
13261
13262        let clarification = agent.chat("Send it").await.unwrap();
13263        assert_eq!(clarification.content, "What should I send?");
13264        assert_eq!(observed.call_count(), 2);
13265
13266        let confirmation = agent.chat("The report to Ada").await.unwrap();
13267        assert_eq!(confirmation.content, "Should I send the report to Ada?");
13268        assert_eq!(
13269            confirmation
13270                .metadata
13271                .as_ref()
13272                .and_then(|metadata| metadata.get("disambiguation"))
13273                .and_then(|metadata| metadata.get("status"))
13274                .and_then(Value::as_str),
13275            Some("awaiting_confirmation")
13276        );
13277        assert_eq!(observed.call_count(), 4);
13278
13279        let completed = agent.chat("Yes").await.unwrap();
13280        assert_eq!(completed.content, "Request executed.");
13281        assert_eq!(observed.call_count(), 6);
13282    }
13283
13284    #[tokio::test]
13285    async fn streaming_state_confirmation_ends_the_turn_before_redispatch() {
13286        let (agent, observed) = state_disambiguation_agent(
13287            vec![
13288                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13289                r#"{"question":"What should I send?","options":null}"#,
13290                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13291                r#"{"question":"Should I send the report to Ada?"}"#,
13292                r#"{"status":"confirmed"}"#,
13293                "Request executed.",
13294            ],
13295            true,
13296            None,
13297            true,
13298        );
13299
13300        let mut clarification_stream = agent.chat_stream("Send it").await.unwrap();
13301        let mut clarification = String::new();
13302        while let Some(chunk) = clarification_stream.next().await {
13303            match chunk {
13304                StreamChunk::Content { text } => clarification.push_str(&text),
13305                StreamChunk::Done {} => break,
13306                StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
13307                _ => {}
13308            }
13309        }
13310        assert_eq!(clarification, "What should I send?");
13311        assert_eq!(observed.call_count(), 2);
13312
13313        let mut confirmation_stream = agent.chat_stream_events("The report to Ada").await.unwrap();
13314        let mut confirmation = None;
13315        while let Some(event) = confirmation_stream.next().await {
13316            match event {
13317                AgentStreamEvent::Final(response) => confirmation = Some(response),
13318                AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
13319                    panic!("unexpected stream error: {message}")
13320                }
13321                AgentStreamEvent::Chunk(_) => {}
13322            }
13323        }
13324        let confirmation = confirmation.expect("confirmation must finalize");
13325        assert_eq!(confirmation.content, "Should I send the report to Ada?");
13326        assert_eq!(
13327            confirmation
13328                .metadata
13329                .as_ref()
13330                .and_then(|metadata| metadata.get("disambiguation"))
13331                .and_then(|metadata| metadata.get("status"))
13332                .and_then(Value::as_str),
13333            Some("awaiting_confirmation")
13334        );
13335        assert_eq!(observed.call_count(), 4);
13336
13337        let mut completed_stream = agent.chat_stream("Yes").await.unwrap();
13338        let mut completed = String::new();
13339        while let Some(chunk) = completed_stream.next().await {
13340            match chunk {
13341                StreamChunk::Content { text } => completed.push_str(&text),
13342                StreamChunk::Done {} => break,
13343                StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
13344                _ => {}
13345            }
13346        }
13347        assert_eq!(completed, "Request executed.");
13348        assert_eq!(observed.call_count(), 6);
13349    }
13350
13351    /// Confirms a returned legacy stream blocks a blocking turn until drop and the event terminal releases the same gate.
13352    #[tokio::test]
13353    async fn root_turn_gate_serializes_blocking_and_streaming_entry_points() {
13354        let (complete_entered, mut complete_events) = tokio::sync::mpsc::unbounded_channel();
13355        let agent = Arc::new(
13356            AgentBuilder::new()
13357                .system_prompt("Serialize root turns.")
13358                .llm(Arc::new(RootTurnProbeProvider { complete_entered }))
13359                .build()
13360                .unwrap(),
13361        );
13362        let blocking_agent = Arc::clone(&agent);
13363
13364        let legacy_stream = agent.chat_stream("stream owner").await.unwrap();
13365        assert!(agent.root_turn_gate.try_lock().is_err());
13366        let blocking = tokio::spawn(async move { blocking_agent.chat("blocked").await.unwrap() });
13367        assert!(
13368            tokio::time::timeout(std::time::Duration::from_millis(50), complete_events.recv())
13369                .await
13370                .is_err(),
13371            "blocking turn reached the provider while the legacy stream owned the root gate"
13372        );
13373
13374        drop(legacy_stream);
13375        assert_eq!(
13376            tokio::time::timeout(std::time::Duration::from_secs(2), complete_events.recv())
13377                .await
13378                .expect("blocking turn did not enter after stream drop"),
13379            Some(())
13380        );
13381        let response = tokio::time::timeout(std::time::Duration::from_secs(2), blocking)
13382            .await
13383            .expect("blocking turn did not finish after stream drop")
13384            .unwrap();
13385        assert_eq!(response.content, "blocking complete");
13386
13387        let mut event_stream = agent.chat_stream_events("event terminal").await.unwrap();
13388        assert!(agent.root_turn_gate.try_lock().is_err());
13389        let mut saw_final = false;
13390        while let Some(event) = event_stream.next().await {
13391            if matches!(event, AgentStreamEvent::Final(_)) {
13392                saw_final = true;
13393                break;
13394            }
13395        }
13396        assert!(saw_final);
13397        assert!(
13398            agent.root_turn_gate.try_lock().is_ok(),
13399            "authoritative terminal event retained the root gate"
13400        );
13401    }
13402
13403    /// Confirms on_response receives a fail-fast error instead of deadlocking on the same runtime gate.
13404    #[tokio::test]
13405    async fn response_hook_rejects_same_runtime_chat_reentry() {
13406        let hooks = Arc::new(ResponseChatHooks {
13407            target: parking_lot::Mutex::new(None),
13408            invoked: AtomicBool::new(false),
13409            nested_result: parking_lot::Mutex::new(None),
13410        });
13411        let agent = Arc::new(
13412            AgentBuilder::new()
13413                .system_prompt("Reject response hook reentry.")
13414                .llm(Arc::new(mock_with_response("outer response")))
13415                .hooks(hooks.clone())
13416                .build()
13417                .unwrap(),
13418        );
13419        *hooks.target.lock() = Some(Arc::downgrade(&agent));
13420
13421        let response = tokio::time::timeout(
13422            std::time::Duration::from_secs(2),
13423            agent.chat("outer request"),
13424        )
13425        .await
13426        .expect("same-runtime response hook reentry must fail without deadlocking")
13427        .unwrap();
13428
13429        assert_eq!(response.content, "outer response");
13430        let nested_result = hooks
13431            .nested_result
13432            .lock()
13433            .clone()
13434            .expect("response hook must record its nested call");
13435        let error = nested_result.expect_err("same-runtime nested chat must be rejected");
13436        assert!(error.contains("reentrant root turn ownership"));
13437    }
13438
13439    /// Confirms a different runtime gate can nest while an A to B to A ownership cycle is rejected at A.
13440    #[tokio::test]
13441    async fn root_turn_gate_allows_nested_runtime_and_rejects_cycles() {
13442        let agent_a = AgentBuilder::new()
13443            .system_prompt("Runtime A.")
13444            .llm(Arc::new(mock_with_response("response A")))
13445            .build()
13446            .unwrap();
13447        let agent_b = AgentBuilder::new()
13448            .system_prompt("Runtime B.")
13449            .llm(Arc::new(mock_with_response("response B")))
13450            .build()
13451            .unwrap();
13452        let RootTurnAdmission {
13453            guard: guard_a,
13454            identity_stack: stack_a,
13455        } = agent_a.acquire_root_turn().await.unwrap();
13456
13457        let cycle_error = scope_runtime_gate_identity_stack(&stack_a, async {
13458            let RootTurnAdmission {
13459                guard: guard_b,
13460                identity_stack: stack_b,
13461            } = agent_b
13462                .acquire_root_turn()
13463                .await
13464                .expect("runtime B must acquire a different gate");
13465            let result =
13466                scope_runtime_gate_identity_stack(&stack_b, agent_a.acquire_root_turn()).await;
13467            drop(guard_b);
13468            match result {
13469                Err(error) => error,
13470                Ok(_) => panic!("runtime A accepted a repeated gate identity"),
13471            }
13472        })
13473        .await;
13474        drop(guard_a);
13475
13476        assert!(
13477            cycle_error
13478                .to_string()
13479                .contains("reentrant root turn ownership")
13480        );
13481    }
13482
13483    /// Confirms concurrent orchestration propagates root ancestry so an A to B to A cycle fails before waiting on A.
13484    #[tokio::test]
13485    async fn concurrent_orchestration_propagates_root_gate_ancestry() {
13486        let registry = Arc::new(crate::spawner::AgentRegistry::new());
13487        let hooks_a = Arc::new(ConcurrentResponseHooks {
13488            registry: Arc::downgrade(&registry),
13489            child_id: "runtime-b".to_string(),
13490            invoked: AtomicBool::new(false),
13491            nested_result: parking_lot::Mutex::new(None),
13492        });
13493        let hooks_b = Arc::new(ResponseChatHooks {
13494            target: parking_lot::Mutex::new(None),
13495            invoked: AtomicBool::new(false),
13496            nested_result: parking_lot::Mutex::new(None),
13497        });
13498        let agent_a = AgentBuilder::new()
13499            .system_prompt("Runtime A dispatches runtime B concurrently.")
13500            .llm(Arc::new(mock_with_response("response A")))
13501            .hooks(hooks_a.clone())
13502            .build()
13503            .unwrap();
13504        let agent_b = AgentBuilder::new()
13505            .system_prompt("Runtime B attempts to re-enter runtime A.")
13506            .llm(Arc::new(mock_with_response("response B")))
13507            .hooks(hooks_b.clone())
13508            .build()
13509            .unwrap();
13510        let spec_a = crate::spec::AgentSpec {
13511            name: "runtime-a".to_string(),
13512            system_prompt: "Runtime A dispatches runtime B concurrently.".to_string(),
13513            ..crate::spec::AgentSpec::default()
13514        };
13515        let spec_b = crate::spec::AgentSpec {
13516            name: "runtime-b".to_string(),
13517            system_prompt: "Runtime B attempts to re-enter runtime A.".to_string(),
13518            ..crate::spec::AgentSpec::default()
13519        };
13520        registry
13521            .register(crate::spawner::SpawnedAgent::from_runtime(
13522                "runtime-a".to_string(),
13523                agent_a,
13524                spec_a,
13525            ))
13526            .await
13527            .unwrap();
13528        registry
13529            .register(crate::spawner::SpawnedAgent::from_runtime(
13530                "runtime-b".to_string(),
13531                agent_b,
13532                spec_b,
13533            ))
13534            .await
13535            .unwrap();
13536        let runtime_a = registry.get("runtime-a").unwrap();
13537        *hooks_b.target.lock() = Some(Arc::downgrade(&runtime_a));
13538
13539        let response = tokio::time::timeout(
13540            std::time::Duration::from_secs(2),
13541            runtime_a.chat("outer concurrent request"),
13542        )
13543        .await
13544        .expect("concurrent orchestration cycle must fail without deadlocking")
13545        .unwrap();
13546
13547        assert_eq!(response.content, "response A");
13548        let child_result = hooks_a
13549            .nested_result
13550            .lock()
13551            .clone()
13552            .expect("runtime A hook must record runtime B completion");
13553        assert_eq!(child_result.unwrap(), "response B");
13554        let cycle_result = hooks_b
13555            .nested_result
13556            .lock()
13557            .clone()
13558            .expect("runtime B hook must record runtime A reentry");
13559        assert!(
13560            cycle_result
13561                .expect_err("runtime A accepted a repeated gate identity")
13562                .contains("reentrant root turn ownership")
13563        );
13564    }
13565
13566    /// Confirms that a pending skill remains inert until confirmation and then runs once.
13567    #[tokio::test]
13568    async fn confirmed_skill_route_executes_exactly_once() {
13569        let (agent, observed) = state_disambiguation_agent_with_skills(
13570            vec![
13571                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13572                "send_report",
13573                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13574                r#"{"question":"What should I send?","options":null}"#,
13575                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13576                r#"{"question":"Should I send the report to Ada?"}"#,
13577                r#"{"status":"confirmed"}"#,
13578                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"resolved","what_is_unclear":[],"detected_language":"en"}"#,
13579                "Report skill executed.",
13580            ],
13581            true,
13582            None,
13583            true,
13584            vec![confirmation_skill()],
13585        );
13586
13587        let clarification = agent.chat("Send it").await.unwrap();
13588        assert_eq!(clarification.content, "What should I send?");
13589        assert_eq!(confirmation_skill_call_count(&observed), 0);
13590
13591        let confirmation = agent.chat("The report to Ada").await.unwrap();
13592        assert_eq!(confirmation.content, "Should I send the report to Ada?");
13593        assert_eq!(
13594            confirmation
13595                .metadata
13596                .as_ref()
13597                .and_then(|metadata| metadata.get("disambiguation"))
13598                .and_then(|metadata| metadata.get("status"))
13599                .and_then(Value::as_str),
13600            Some("awaiting_confirmation")
13601        );
13602        assert_eq!(confirmation_skill_call_count(&observed), 0);
13603
13604        let completed = agent.chat("Yes").await.unwrap();
13605        assert_eq!(completed.content, "Report skill executed.");
13606        assert_eq!(confirmation_skill_call_count(&observed), 1);
13607        assert!(agent.pending_skill_id.read().is_none());
13608        let messages = agent.memory.get_messages(None).await.unwrap();
13609        assert!(!messages.iter().any(|message| message.content == "Yes"));
13610    }
13611
13612    /// Preserves a second clarification response after confirmation instead of overwriting its metadata.
13613    #[tokio::test]
13614    async fn confirmed_skill_recheck_preserves_new_clarification_metadata() {
13615        let (agent, observed) = state_disambiguation_agent_with_skills(
13616            vec![
13617                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13618                "send_report",
13619                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13620                r#"{"question":"What should I send?","options":null}"#,
13621                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13622                r#"{"question":"Should I send the report to Ada?"}"#,
13623                r#"{"status":"confirmed"}"#,
13624                r#"{"is_ambiguous":true,"confidence":0.3,"ambiguity_type":"missing_parameters","reasoning":"timing missing","what_is_unclear":["timing"],"detected_language":"en"}"#,
13625                r#"{"question":"When should I send it?","options":null}"#,
13626            ],
13627            true,
13628            None,
13629            true,
13630            vec![confirmation_skill()],
13631        );
13632
13633        agent.chat("Send it").await.unwrap();
13634        agent.chat("The report to Ada").await.unwrap();
13635        let follow_up = agent.chat("Yes").await.unwrap();
13636
13637        assert_eq!(follow_up.content, "When should I send it?");
13638        let metadata = follow_up
13639            .metadata
13640            .as_ref()
13641            .and_then(|metadata| metadata.get("disambiguation"))
13642            .unwrap();
13643        assert_eq!(
13644            metadata.get("status").and_then(Value::as_str),
13645            Some("awaiting_clarification")
13646        );
13647        assert_eq!(
13648            metadata.get("skill_id").and_then(Value::as_str),
13649            Some("send_report")
13650        );
13651        assert!(metadata.get("detection").is_some());
13652        assert_eq!(confirmation_skill_call_count(&observed), 0);
13653    }
13654
13655    /// Confirms rejection clears a pending skill without executing its prompt.
13656    #[tokio::test]
13657    async fn rejected_skill_confirmation_never_executes() {
13658        let (agent, observed) = state_disambiguation_agent_with_skills(
13659            vec![
13660                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13661                "send_report",
13662                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13663                r#"{"question":"What should I send?","options":null}"#,
13664                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13665                r#"{"question":"Should I send the report to Ada?"}"#,
13666                r#"{"status":"rejected"}"#,
13667                "Confirmation rejected.",
13668            ],
13669            true,
13670            None,
13671            true,
13672            vec![confirmation_skill()],
13673        );
13674
13675        agent.chat("Send it").await.unwrap();
13676        agent.chat("The report to Ada").await.unwrap();
13677        let rejected = agent.chat("No").await.unwrap();
13678
13679        assert_eq!(rejected.content, "Confirmation rejected.");
13680        assert_eq!(confirmation_skill_call_count(&observed), 0);
13681        assert!(agent.pending_skill_id.read().is_none());
13682    }
13683
13684    /// Confirms reset clears manager and skill ownership before a later streaming turn.
13685    #[tokio::test]
13686    async fn reset_invalidates_pending_skill_confirmation_before_streaming_input() {
13687        let (agent, observed) = state_disambiguation_agent_with_skills(
13688            vec![
13689                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13690                "send_report",
13691                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13692                r#"{"question":"What should I send?","options":null}"#,
13693                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13694                r#"{"question":"Should I send the report to Ada?"}"#,
13695                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"fresh input","what_is_unclear":[],"detected_language":"en"}"#,
13696                "none",
13697                "Fresh response.",
13698            ],
13699            true,
13700            None,
13701            true,
13702            vec![confirmation_skill()],
13703        );
13704
13705        agent.chat("Send it").await.unwrap();
13706        agent.chat("The report to Ada").await.unwrap();
13707        agent.reset().await.unwrap();
13708        assert!(agent.pending_skill_id.read().is_none());
13709        assert!(
13710            !agent
13711                .disambiguation_manager()
13712                .unwrap()
13713                .has_pending_clarification()
13714                .await
13715        );
13716
13717        let mut stream = agent.chat_stream("Yes").await.unwrap();
13718        let mut content = String::new();
13719        while let Some(chunk) = stream.next().await {
13720            match chunk {
13721                StreamChunk::Content { text } => content.push_str(&text),
13722                StreamChunk::Done {} => break,
13723                StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
13724                _ => {}
13725            }
13726        }
13727
13728        assert_eq!(content, "Fresh response.");
13729        assert_eq!(confirmation_skill_call_count(&observed), 0);
13730    }
13731
13732    /// Confirms the Agent trait reset uses the same pending ownership cleanup.
13733    #[tokio::test]
13734    async fn trait_reset_clears_pending_skill_confirmation() {
13735        let (agent, _) = state_disambiguation_agent_with_skills(
13736            vec![
13737                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13738                "send_report",
13739                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13740                r#"{"question":"What should I send?","options":null}"#,
13741                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13742                r#"{"question":"Should I send the report to Ada?"}"#,
13743            ],
13744            true,
13745            None,
13746            true,
13747            vec![confirmation_skill()],
13748        );
13749
13750        agent.chat("Send it").await.unwrap();
13751        agent.chat("The report to Ada").await.unwrap();
13752        <RuntimeAgent as Agent>::reset(&agent).await.unwrap();
13753
13754        assert!(agent.pending_skill_id.read().is_none());
13755        assert!(
13756            !agent
13757                .disambiguation_manager()
13758                .unwrap()
13759                .has_pending_clarification()
13760                .await
13761        );
13762    }
13763
13764    /// Confirms a state transition cancels stale skill execution before confirmation parsing.
13765    #[tokio::test]
13766    async fn state_change_invalidates_pending_skill_confirmation() {
13767        let (agent, observed) = state_disambiguation_agent_with_skills(
13768            vec![
13769                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13770                "send_report",
13771                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13772                r#"{"question":"What should I send?","options":null}"#,
13773                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13774                r#"{"question":"Should I send the report to Ada?"}"#,
13775                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"fresh input","what_is_unclear":[],"detected_language":"en"}"#,
13776                "none",
13777                "Fresh response.",
13778            ],
13779            true,
13780            None,
13781            true,
13782            vec![confirmation_skill()],
13783        );
13784
13785        agent.chat("Send it").await.unwrap();
13786        agent.chat("The report to Ada").await.unwrap();
13787        agent.transition_to("review").await.unwrap();
13788        let cancelled = agent.chat("Yes").await.unwrap();
13789
13790        assert_eq!(cancelled.content, "Fresh response.");
13791        assert_eq!(confirmation_skill_call_count(&observed), 0);
13792        assert!(agent.pending_skill_id.read().is_none());
13793    }
13794
13795    /// Confirms reset wins over a confirmation result that was already being parsed.
13796    #[tokio::test]
13797    async fn in_flight_confirmation_cannot_redispatch_after_reset() {
13798        let (mut agent, observed) = state_disambiguation_agent_with_skills(
13799            vec![
13800                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13801                "send_report",
13802                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13803                r#"{"question":"What should I send?","options":null}"#,
13804                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13805                r#"{"question":"Should I send the report to Ada?"}"#,
13806                r#"{"status":"confirmed"}"#,
13807                "Confirmation cancelled.",
13808            ],
13809            true,
13810            None,
13811            true,
13812            vec![confirmation_skill()],
13813        );
13814        let observer = Arc::new(BlockingRuntimeConfirmationObserver::new());
13815        let manager = agent
13816            .disambiguation_manager
13817            .take()
13818            .unwrap()
13819            .with_clarification_observer(observer.clone());
13820        agent.disambiguation_manager = Some(manager);
13821        let agent = Arc::new(agent);
13822
13823        agent.chat("Send it").await.unwrap();
13824        agent.chat("The report to Ada").await.unwrap();
13825
13826        let confirming_agent = Arc::clone(&agent);
13827        let confirmation = tokio::spawn(async move { confirming_agent.chat("Yes").await });
13828        observer.entered.wait().await;
13829        agent.reset().await.unwrap();
13830        observer.release.notify_one();
13831
13832        let response = confirmation.await.unwrap().unwrap();
13833        assert_eq!(response.content, "Confirmation cancelled.");
13834        assert_eq!(confirmation_skill_call_count(&observed), 0);
13835        assert!(agent.pending_skill_id.read().is_none());
13836    }
13837
13838    /// Confirms a queued reset prevents an older terminal question from being published afterward.
13839    #[tokio::test]
13840    async fn queued_reset_prevents_stale_confirmation_question_publication() {
13841        let (agent, observed) = state_disambiguation_agent(
13842            vec![
13843                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13844                r#"{"question":"What should I send?","options":null}"#,
13845                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13846                r#"{"question":"Should I send the report to Ada?"}"#,
13847            ],
13848            true,
13849            None,
13850            true,
13851        );
13852        let agent = Arc::new(agent);
13853        agent.chat("Send it").await.unwrap();
13854
13855        let admission = agent.disambiguation_admission.write().await;
13856        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
13857        let resetting_agent = Arc::clone(&agent);
13858        let reset = tokio::spawn(async move {
13859            let _ = started_tx.send(());
13860            resetting_agent.reset().await
13861        });
13862        started_rx.await.unwrap();
13863        tokio::task::yield_now().await;
13864
13865        let responding_agent = Arc::clone(&agent);
13866        let response =
13867            tokio::spawn(async move { responding_agent.chat("The report to Ada").await });
13868        tokio::time::timeout(std::time::Duration::from_secs(2), async {
13869            while observed.call_count() < 4 {
13870                tokio::task::yield_now().await;
13871            }
13872        })
13873        .await
13874        .expect("clarification processing must reach terminal publication");
13875        drop(admission);
13876
13877        reset.await.unwrap().unwrap();
13878        let error = response.await.unwrap().unwrap_err();
13879        assert!(error.to_string().contains("ownership changed"));
13880        assert!(
13881            !agent
13882                .disambiguation_manager()
13883                .unwrap()
13884                .has_pending_clarification()
13885                .await
13886        );
13887        assert!(agent.memory.get_messages(None).await.unwrap().is_empty());
13888    }
13889
13890    /// Confirms skill clarification consumers recheck ownership before publishing returned responses.
13891    #[tokio::test]
13892    async fn queued_reset_prevents_stale_skill_clarification_publication() {
13893        let (agent, observed) = state_disambiguation_agent_with_skills(
13894            vec![
13895                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13896                "send_report",
13897                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13898                r#"{"question":"What should I send?","options":null}"#,
13899            ],
13900            true,
13901            None,
13902            true,
13903            vec![confirmation_skill()],
13904        );
13905        let agent = Arc::new(agent);
13906        let admission = agent.disambiguation_admission.write().await;
13907        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
13908        let resetting_agent = Arc::clone(&agent);
13909        let reset = tokio::spawn(async move {
13910            let _ = started_tx.send(());
13911            resetting_agent.reset().await
13912        });
13913        started_rx.await.unwrap();
13914        tokio::task::yield_now().await;
13915
13916        let responding_agent = Arc::clone(&agent);
13917        let response = tokio::spawn(async move { responding_agent.chat("Send it").await });
13918        tokio::time::timeout(std::time::Duration::from_secs(2), async {
13919            while observed.call_count() < 4 {
13920                tokio::task::yield_now().await;
13921            }
13922        })
13923        .await
13924        .expect("skill clarification must reach terminal publication");
13925        drop(admission);
13926
13927        reset.await.unwrap().unwrap();
13928        let error = response.await.unwrap().unwrap_err();
13929        assert!(error.to_string().contains("ownership changed"));
13930        assert_eq!(confirmation_skill_call_count(&observed), 0);
13931        assert!(agent.pending_skill_id.read().is_none());
13932        assert!(agent.memory.get_messages(None).await.unwrap().is_empty());
13933    }
13934
13935    /// Confirms transition hooks can reenter reset after the state commit lock is released.
13936    #[tokio::test]
13937    async fn transition_hook_can_reset_without_admission_deadlock() {
13938        let hooks = Arc::new(ResetOnTransitionHooks {
13939            agent: parking_lot::Mutex::new(None),
13940            invoked: AtomicBool::new(false),
13941        });
13942        let agent = Arc::new(
13943            AgentBuilder::new()
13944                .system_prompt("Test transition hook reentrancy.")
13945                .llm(Arc::new(mock_with_response("done")))
13946                .state_machine(disambiguation_state_machine(None, false))
13947                .build()
13948                .unwrap()
13949                .with_hooks(hooks.clone()),
13950        );
13951        *hooks.agent.lock() = Some(Arc::downgrade(&agent));
13952
13953        let transitioned = tokio::time::timeout(
13954            std::time::Duration::from_secs(2),
13955            agent.apply_transition_target("active", "review", "test transition", None),
13956        )
13957        .await
13958        .expect("transition hook reset must not deadlock")
13959        .unwrap();
13960
13961        assert!(transitioned);
13962        assert!(hooks.invoked.load(Ordering::SeqCst));
13963        assert_eq!(agent.current_state().as_deref(), Some("active"));
13964    }
13965
13966    /// Confirms only the reserved transition can run source-state exit side effects.
13967    #[tokio::test]
13968    async fn concurrent_transition_cannot_duplicate_exit_actions() {
13969        let gate = PathMutationGate::new();
13970        let active = ai_agents_state::StateDefinition {
13971            on_exit: vec![StateAction::Tool {
13972                tool: "transition_exit".to_string(),
13973                args: Some(serde_json::json!({"path": "./transition-exit.txt"})),
13974            }],
13975            ..Default::default()
13976        };
13977        let state_machine = Arc::new(
13978            StateMachine::new(ai_agents_state::StateConfig {
13979                initial: "active".to_string(),
13980                states: HashMap::from([
13981                    ("active".to_string(), active),
13982                    (
13983                        "review".to_string(),
13984                        ai_agents_state::StateDefinition::default(),
13985                    ),
13986                ]),
13987                global_transitions: Vec::new(),
13988                fallback: None,
13989                max_no_transition: None,
13990                regenerate_on_transition: true,
13991            })
13992            .unwrap(),
13993        );
13994        let agent = Arc::new(
13995            AgentBuilder::new()
13996                .system_prompt("Test transition reservation.")
13997                .llm(Arc::new(mock_with_response("done")))
13998                .tool(Arc::new(BlockingPathMutationTool {
13999                    id: "transition_exit",
14000                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
14001                    gate: gate.clone(),
14002                }))
14003                .state_machine(state_machine)
14004                .build()
14005                .unwrap(),
14006        );
14007
14008        let first_agent = Arc::clone(&agent);
14009        let first = tokio::spawn(async move { first_agent.transition_to("review").await });
14010        tokio::time::timeout(std::time::Duration::from_secs(2), gate.wait_until_entered())
14011            .await
14012            .expect("reserved transition must enter its exit action");
14013
14014        let second = tokio::time::timeout(
14015            std::time::Duration::from_secs(2),
14016            agent.transition_to("review"),
14017        )
14018        .await
14019        .expect("competing transition must fail without waiting for the exit action")
14020        .unwrap_err();
14021        assert!(second.to_string().contains("already in progress"));
14022
14023        gate.release();
14024        first.await.unwrap().unwrap();
14025        assert_eq!(agent.current_state().as_deref(), Some("review"));
14026    }
14027
14028    /// Confirms a later transition cannot overtake the committed state's enter actions.
14029    #[tokio::test]
14030    async fn concurrent_transition_cannot_overtake_enter_actions() {
14031        let gate = PathMutationGate::new();
14032        let review = ai_agents_state::StateDefinition {
14033            on_enter: vec![StateAction::Tool {
14034                tool: "transition_enter".to_string(),
14035                args: Some(serde_json::json!({"path": "./transition-enter.txt"})),
14036            }],
14037            ..Default::default()
14038        };
14039        let state_machine = Arc::new(
14040            StateMachine::new(ai_agents_state::StateConfig {
14041                initial: "active".to_string(),
14042                states: HashMap::from([
14043                    (
14044                        "active".to_string(),
14045                        ai_agents_state::StateDefinition::default(),
14046                    ),
14047                    ("review".to_string(), review),
14048                ]),
14049                global_transitions: Vec::new(),
14050                fallback: None,
14051                max_no_transition: None,
14052                regenerate_on_transition: true,
14053            })
14054            .unwrap(),
14055        );
14056        let agent = Arc::new(
14057            AgentBuilder::new()
14058                .system_prompt("Test transition lifecycle reservation.")
14059                .llm(Arc::new(mock_with_response("done")))
14060                .tool(Arc::new(BlockingPathMutationTool {
14061                    id: "transition_enter",
14062                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
14063                    gate: gate.clone(),
14064                }))
14065                .state_machine(state_machine)
14066                .build()
14067                .unwrap(),
14068        );
14069
14070        let first_agent = Arc::clone(&agent);
14071        let first = tokio::spawn(async move { first_agent.transition_to("review").await });
14072        tokio::time::timeout(std::time::Duration::from_secs(2), gate.wait_until_entered())
14073            .await
14074            .expect("committed transition must enter its destination action");
14075
14076        let second = agent.transition_to("active").await.unwrap_err();
14077        assert!(second.to_string().contains("already in progress"));
14078        assert!(agent.reset().await.is_err());
14079
14080        gate.release();
14081        first.await.unwrap().unwrap();
14082        assert_eq!(agent.current_state().as_deref(), Some("review"));
14083    }
14084
14085    /// Confirms even a same-state restore invalidates manager and skill ownership.
14086    #[tokio::test]
14087    async fn same_state_restore_invalidates_pending_skill_confirmation() {
14088        let (agent, observed) = state_disambiguation_agent_with_skills(
14089            vec![
14090                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
14091                "send_report",
14092                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
14093                r#"{"question":"What should I send?","options":null}"#,
14094                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
14095                r#"{"question":"Should I send the report to Ada?"}"#,
14096            ],
14097            true,
14098            None,
14099            true,
14100            vec![confirmation_skill()],
14101        );
14102
14103        agent.chat("Send it").await.unwrap();
14104        agent.chat("The report to Ada").await.unwrap();
14105        let snapshot = agent.save_state().await.unwrap();
14106        assert_eq!(agent.current_state().as_deref(), Some("active"));
14107
14108        agent.restore_state(snapshot).await.unwrap();
14109
14110        assert_eq!(agent.current_state().as_deref(), Some("active"));
14111        assert!(agent.pending_skill_id.read().is_none());
14112        assert!(
14113            !agent
14114                .disambiguation_manager()
14115                .unwrap()
14116                .has_pending_clarification()
14117                .await
14118        );
14119        assert_eq!(confirmation_skill_call_count(&observed), 0);
14120    }
14121
14122    /// Confirms direct state mutation cannot hide behind a return to the same state path.
14123    #[tokio::test]
14124    async fn direct_state_generation_change_invalidates_confirmation() {
14125        let (agent, observed) = state_disambiguation_agent_with_skills(
14126            vec![
14127                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
14128                "send_report",
14129                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
14130                r#"{"question":"What should I send?","options":null}"#,
14131                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
14132                r#"{"question":"Should I send the report to Ada?"}"#,
14133                "Confirmation cancelled.",
14134            ],
14135            true,
14136            None,
14137            true,
14138            vec![confirmation_skill()],
14139        );
14140
14141        agent.chat("Send it").await.unwrap();
14142        agent.chat("The report to Ada").await.unwrap();
14143        let state_machine = agent.state_machine().unwrap();
14144        state_machine
14145            .transition_to("review", "external test")
14146            .unwrap();
14147        state_machine
14148            .transition_to("active", "external test")
14149            .unwrap();
14150
14151        let response = agent.chat("Yes").await.unwrap();
14152
14153        assert_eq!(response.content, "Confirmation cancelled.");
14154        assert_eq!(confirmation_skill_call_count(&observed), 0);
14155        assert!(agent.pending_skill_id.read().is_none());
14156    }
14157
14158    #[tokio::test]
14159    async fn state_confirmation_does_not_add_a_question_for_clear_input() {
14160        let (agent, observed) = state_disambiguation_agent(
14161            vec![
14162                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"clear","what_is_unclear":[],"detected_language":"en"}"#,
14163                "Request executed.",
14164            ],
14165            true,
14166            None,
14167            true,
14168        );
14169
14170        let response = agent.chat("Send the report to Ada").await.unwrap();
14171
14172        assert_eq!(response.content, "Request executed.");
14173        assert_eq!(observed.call_count(), 2);
14174    }
14175
14176    #[tokio::test]
14177    async fn state_override_cannot_activate_a_disabled_top_level_manager() {
14178        let (agent, observed) =
14179            state_disambiguation_agent(vec!["Request executed."], false, Some(true), true);
14180
14181        assert!(!agent.has_disambiguation());
14182        let response = agent.chat("Send it").await.unwrap();
14183
14184        assert_eq!(response.content, "Request executed.");
14185        assert_eq!(observed.call_count(), 1);
14186    }
14187
14188    #[tokio::test]
14189    async fn native_required_choice_executes_through_the_shared_tool_path() {
14190        let mut mock = MockLLMProvider::new("native-required");
14191        mock.set_tool_choice(Some(ToolChoice::Required));
14192        mock.add_response(
14193            LLMResponse::new("", FinishReason::ToolCall)
14194                .with_tool_calls(vec![ToolCall {
14195                    id: "provider-call-1".to_string(),
14196                    name: "calculator".to_string(),
14197                    arguments: serde_json::json!({"expression": "2 + 2"}),
14198                }])
14199                .unwrap(),
14200        );
14201        mock.add_response(LLMResponse::new("The answer is 4.", FinishReason::Stop));
14202        let observed = mock.clone();
14203        let agent = AgentBuilder::new()
14204            .system_prompt("Use the calculator when needed.")
14205            .llm(Arc::new(mock))
14206            .tool(Arc::new(CalculatorTool::new()))
14207            .build()
14208            .unwrap();
14209
14210        let response = agent.chat("What is 2 + 2?").await.unwrap();
14211
14212        assert_eq!(response.content, "The answer is 4.");
14213        assert_eq!(
14214            response.tool_calls.as_ref().unwrap()[0].id,
14215            "provider-call-1"
14216        );
14217        let calls = observed.call_history();
14218        assert_eq!(calls.len(), 2);
14219        assert!(matches!(
14220            calls[0].request.as_ref().map(|request| &request.choice),
14221            Some(ToolChoice::Required)
14222        ));
14223        assert!(matches!(
14224            calls[1].request.as_ref().map(|request| &request.choice),
14225            Some(ToolChoice::Auto)
14226        ));
14227    }
14228
14229    #[tokio::test]
14230    async fn prompt_fallback_uses_one_corrective_retry() {
14231        let mut mock = MockLLMProvider::new("prompt-required");
14232        mock.set_tool_choice(Some(ToolChoice::Required));
14233        mock.set_native_tool_support(false);
14234        mock.set_responses(
14235            vec![
14236                "I can calculate that.".to_string(),
14237                r#"{"tool":"calculator","arguments":{"expression":"2 + 2"}}"#.to_string(),
14238                "The answer is 4.".to_string(),
14239            ],
14240            false,
14241        );
14242        let observed = mock.clone();
14243        let agent = AgentBuilder::new()
14244            .system_prompt("Use tools.")
14245            .llm(Arc::new(mock))
14246            .tool(Arc::new(CalculatorTool::new()))
14247            .build()
14248            .unwrap();
14249
14250        let response = agent.chat("What is 2 + 2?").await.unwrap();
14251
14252        assert_eq!(response.content, "The answer is 4.");
14253        assert_eq!(observed.call_count(), 3);
14254        let corrective = &observed.call_history()[1].messages;
14255        assert!(
14256            corrective
14257                .last()
14258                .unwrap()
14259                .content
14260                .contains("previous response")
14261        );
14262    }
14263
14264    #[tokio::test]
14265    async fn prompt_fallback_fails_after_one_noncompliant_retry() {
14266        let mut mock = MockLLMProvider::new("prompt-required-failure");
14267        mock.set_tool_choice(Some(ToolChoice::Required));
14268        mock.set_native_tool_support(false);
14269        mock.set_responses(
14270            vec!["No tool.".to_string(), "Still no tool.".to_string()],
14271            false,
14272        );
14273        let observed = mock.clone();
14274        let agent = AgentBuilder::new()
14275            .system_prompt("Use tools.")
14276            .llm(Arc::new(mock))
14277            .tool(Arc::new(CalculatorTool::new()))
14278            .build()
14279            .unwrap();
14280
14281        let error = agent.chat("What is 2 + 2?").await.unwrap_err();
14282
14283        assert!(error.to_string().contains("one corrective retry"));
14284        assert_eq!(observed.call_count(), 2);
14285    }
14286
14287    #[tokio::test]
14288    async fn specific_choice_cannot_widen_the_effective_grant() {
14289        let mut mock = MockLLMProvider::new("specific-outside-grant");
14290        mock.set_tool_choice(Some(ToolChoice::Specific("random".to_string())));
14291        let observed = mock.clone();
14292        let agent = AgentBuilder::new()
14293            .system_prompt("Use tools.")
14294            .llm(Arc::new(mock))
14295            .tool(Arc::new(CalculatorTool::new()))
14296            .build()
14297            .unwrap();
14298
14299        let error = agent.chat("Generate a value.").await.unwrap_err();
14300
14301        assert!(error.to_string().contains("is not registered"));
14302        assert_eq!(observed.call_count(), 0);
14303    }
14304
14305    #[tokio::test]
14306    async fn none_choice_exposes_no_tool_protocol() {
14307        let mut mock = MockLLMProvider::new("no-tools");
14308        mock.set_tool_choice(Some(ToolChoice::None));
14309        mock.set_response(r#"{"tool":"calculator","arguments":{"expression":"2 + 2"}}"#);
14310        let observed = mock.clone();
14311        let agent = AgentBuilder::new()
14312            .system_prompt("Answer directly.")
14313            .llm(Arc::new(mock))
14314            .tool(Arc::new(CalculatorTool::new()))
14315            .build()
14316            .unwrap();
14317
14318        let response = agent.chat("Hello").await.unwrap();
14319
14320        assert!(response.tool_calls.is_none());
14321        assert_eq!(observed.call_count(), 1);
14322        let call = observed.last_call().unwrap();
14323        assert!(call.request.is_none());
14324        assert!(
14325            call.messages
14326                .iter()
14327                .all(|message| !message.content.contains("Available tools:"))
14328        );
14329    }
14330
14331    struct RuntimeStorage {
14332        capabilities: Box<[StorageCapability]>,
14333        snapshots: RwLock<HashMap<String, AgentSnapshot>>,
14334        metadata: RwLock<HashMap<String, ai_agents_core::SessionMetadata>>,
14335        metadata_save_calls: AtomicU64,
14336        metadata_load_calls: AtomicU64,
14337        fail_metadata_save: AtomicBool,
14338        fail_metadata_load: AtomicBool,
14339    }
14340
14341    impl RuntimeStorage {
14342        fn new(capabilities: impl IntoIterator<Item = StorageCapability>) -> Self {
14343            Self {
14344                capabilities: capabilities.into_iter().collect(),
14345                snapshots: RwLock::new(HashMap::new()),
14346                metadata: RwLock::new(HashMap::new()),
14347                metadata_save_calls: AtomicU64::new(0),
14348                metadata_load_calls: AtomicU64::new(0),
14349                fail_metadata_save: AtomicBool::new(false),
14350                fail_metadata_load: AtomicBool::new(false),
14351            }
14352        }
14353    }
14354
14355    #[async_trait]
14356    impl AgentStorage for RuntimeStorage {
14357        fn supports(&self, capability: StorageCapability) -> bool {
14358            self.capabilities.contains(&capability)
14359        }
14360
14361        async fn save(&self, session_id: &str, snapshot: &AgentSnapshot) -> Result<()> {
14362            self.snapshots
14363                .write()
14364                .insert(session_id.to_string(), snapshot.clone());
14365            Ok(())
14366        }
14367
14368        async fn load(&self, session_id: &str) -> Result<Option<AgentSnapshot>> {
14369            Ok(self.snapshots.read().get(session_id).cloned())
14370        }
14371
14372        async fn delete(&self, session_id: &str) -> Result<()> {
14373            self.snapshots.write().remove(session_id);
14374            Ok(())
14375        }
14376
14377        async fn list_sessions(&self) -> Result<Vec<String>> {
14378            Ok(self.snapshots.read().keys().cloned().collect())
14379        }
14380
14381        async fn save_snapshot_with_metadata(
14382            &self,
14383            session_id: &str,
14384            snapshot: &AgentSnapshot,
14385            metadata: &ai_agents_core::SessionMetadata,
14386        ) -> Result<()> {
14387            self.metadata_save_calls.fetch_add(1, Ordering::SeqCst);
14388            if self.fail_metadata_save.load(Ordering::SeqCst) {
14389                return Err(AgentError::Persistence("metadata save failed".into()));
14390            }
14391            self.snapshots
14392                .write()
14393                .insert(session_id.to_string(), snapshot.clone());
14394            self.metadata
14395                .write()
14396                .insert(session_id.to_string(), metadata.clone());
14397            Ok(())
14398        }
14399
14400        async fn save_metadata(
14401            &self,
14402            session_id: &str,
14403            metadata: &ai_agents_core::SessionMetadata,
14404        ) -> Result<()> {
14405            self.metadata_save_calls.fetch_add(1, Ordering::SeqCst);
14406            if self.fail_metadata_save.load(Ordering::SeqCst) {
14407                return Err(AgentError::Persistence("metadata save failed".into()));
14408            }
14409            self.metadata
14410                .write()
14411                .insert(session_id.to_string(), metadata.clone());
14412            Ok(())
14413        }
14414
14415        async fn load_metadata(
14416            &self,
14417            session_id: &str,
14418        ) -> Result<Option<ai_agents_core::SessionMetadata>> {
14419            self.metadata_load_calls.fetch_add(1, Ordering::SeqCst);
14420            if self.fail_metadata_load.load(Ordering::SeqCst) {
14421                return Err(AgentError::Persistence("metadata load failed".into()));
14422            }
14423            Ok(self.metadata.read().get(session_id).cloned())
14424        }
14425    }
14426
14427    fn runtime_storage_agent() -> RuntimeAgent {
14428        AgentBuilder::new()
14429            .system_prompt("Test runtime storage integration.")
14430            .llm(Arc::new(mock_with_response("done")))
14431            .build()
14432            .unwrap()
14433    }
14434
14435    fn restore_spec(id: &str) -> crate::spec::AgentSpec {
14436        crate::spec::AgentSpec {
14437            name: id.to_string(),
14438            system_prompt: format!("Restore child {id}."),
14439            ..crate::spec::AgentSpec::default()
14440        }
14441    }
14442
14443    fn restore_entry(id: &str) -> ai_agents_core::SpawnedAgentEntry {
14444        ai_agents_core::SpawnedAgentEntry {
14445            id: id.to_string(),
14446            name: id.to_string(),
14447            spec_yaml: serde_yaml::to_string(&restore_spec(id)).unwrap(),
14448        }
14449    }
14450
14451    fn restore_spawner(
14452        storage: Arc<RuntimeStorage>,
14453        max_agents: usize,
14454    ) -> (
14455        Arc<crate::spawner::AgentSpawner>,
14456        Arc<crate::spawner::AgentRegistry>,
14457    ) {
14458        let mut llms = LLMRegistry::new();
14459        llms.register("default", Arc::new(mock_with_response("done")));
14460        (
14461            Arc::new(
14462                crate::spawner::AgentSpawner::new()
14463                    .with_shared_llms(llms)
14464                    .with_shared_storage(storage)
14465                    .with_max_agents(max_agents),
14466            ),
14467            Arc::new(crate::spawner::AgentRegistry::new()),
14468        )
14469    }
14470
14471    async fn save_restore_target(
14472        parent: &RuntimeAgent,
14473        storage: &RuntimeStorage,
14474        session_id: &str,
14475        entries: Vec<ai_agents_core::SpawnedAgentEntry>,
14476    ) {
14477        let mut snapshot = parent.save_state().await.unwrap();
14478        snapshot.spawned_agents = Some(entries);
14479        storage.save(session_id, &snapshot).await.unwrap();
14480        storage
14481            .save_metadata(session_id, &ai_agents_core::SessionMetadata::default())
14482            .await
14483            .unwrap();
14484    }
14485
14486    #[tokio::test]
14487    async fn storage_init_requires_storage_for_actor_facts() {
14488        let facts = ai_agents_facts::FactsConfig {
14489            enabled: true,
14490            ..Default::default()
14491        };
14492        let agent = runtime_storage_agent().with_facts_config(None, Some(facts));
14493
14494        let error = agent.init_storage().await.unwrap_err();
14495        assert!(matches!(
14496            error,
14497            AgentError::Config(message)
14498                if message.contains("actor facts or actor memory")
14499                    && message.contains("none is configured or injected")
14500        ));
14501    }
14502
14503    #[tokio::test]
14504    async fn storage_init_validates_actor_facts_capability() {
14505        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14506        let actor_memory = ai_agents_facts::ActorMemoryConfig {
14507            enabled: true,
14508            ..Default::default()
14509        };
14510        let agent = runtime_storage_agent()
14511            .with_storage(storage)
14512            .with_facts_config(Some(actor_memory), None);
14513
14514        assert!(matches!(
14515            agent.init_storage().await,
14516            Err(AgentError::UnsupportedStorageCapability(
14517                StorageCapability::ActorFacts
14518            ))
14519        ));
14520    }
14521
14522    #[tokio::test]
14523    async fn blocking_chat_rejects_unsupported_required_storage() {
14524        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14525        let facts = ai_agents_facts::FactsConfig {
14526            enabled: true,
14527            ..Default::default()
14528        };
14529        let agent = runtime_storage_agent()
14530            .with_storage(storage)
14531            .with_facts_config(None, Some(facts));
14532
14533        assert!(matches!(
14534            agent.chat("hello").await,
14535            Err(AgentError::UnsupportedStorageCapability(
14536                StorageCapability::ActorFacts
14537            ))
14538        ));
14539    }
14540
14541    #[tokio::test]
14542    async fn streaming_chat_rejects_unsupported_required_storage_before_stream_creation() {
14543        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14544        let config = ai_agents_relationships::RelationshipConfig {
14545            enabled: true,
14546            ..Default::default()
14547        };
14548        let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
14549        let agent = runtime_storage_agent()
14550            .with_storage(storage)
14551            .with_relationships(manager);
14552
14553        assert!(matches!(
14554            agent.chat_stream("hello").await,
14555            Err(AgentError::UnsupportedStorageCapability(
14556                StorageCapability::ActorRelationships
14557            ))
14558        ));
14559    }
14560
14561    #[tokio::test]
14562    async fn storage_init_completes_facts_for_injected_storage() {
14563        let storage = Arc::new(RuntimeStorage::new([
14564            StorageCapability::Snapshot,
14565            StorageCapability::ActorFacts,
14566        ]));
14567        let facts = ai_agents_facts::FactsConfig {
14568            enabled: true,
14569            ..Default::default()
14570        };
14571        let agent = runtime_storage_agent()
14572            .with_storage(storage)
14573            .with_facts_config(None, Some(facts));
14574
14575        agent.init_storage().await.unwrap();
14576        assert!(agent.fact_store().is_some());
14577    }
14578
14579    #[tokio::test]
14580    async fn storage_init_requires_storage_for_persistent_relationships() {
14581        let config = ai_agents_relationships::RelationshipConfig {
14582            enabled: true,
14583            ..Default::default()
14584        };
14585        let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
14586        let agent = runtime_storage_agent().with_relationships(manager);
14587
14588        let error = agent.init_storage().await.unwrap_err();
14589        assert!(matches!(
14590            error,
14591            AgentError::Config(message)
14592                if message.contains("persistent relationships")
14593                    && message.contains("none is configured or injected")
14594        ));
14595    }
14596
14597    #[tokio::test]
14598    async fn storage_init_validates_persistent_relationships_capability() {
14599        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14600        let config = ai_agents_relationships::RelationshipConfig {
14601            enabled: true,
14602            ..Default::default()
14603        };
14604        let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
14605        let agent = runtime_storage_agent()
14606            .with_storage(storage)
14607            .with_relationships(manager);
14608
14609        assert!(matches!(
14610            agent.init_storage().await,
14611            Err(AgentError::UnsupportedStorageCapability(
14612                StorageCapability::ActorRelationships
14613            ))
14614        ));
14615    }
14616
14617    #[tokio::test]
14618    async fn session_restore_updates_identity_and_clears_stale_actor_binding() {
14619        let storage = Arc::new(RuntimeStorage::new([
14620            StorageCapability::Snapshot,
14621            StorageCapability::SessionMetadata,
14622        ]));
14623        let agent = runtime_storage_agent().with_storage(storage.clone());
14624        agent.set_actor_id("old-actor").unwrap();
14625        agent.save_session("old").await.unwrap();
14626        storage
14627            .save("target", &agent.save_state().await.unwrap())
14628            .await
14629            .unwrap();
14630        storage
14631            .save_metadata("target", &ai_agents_core::SessionMetadata::default())
14632            .await
14633            .unwrap();
14634
14635        assert!(agent.load_session("target").await.unwrap());
14636
14637        assert_eq!(agent.current_session_id.read().as_deref(), Some("target"));
14638        assert_eq!(agent.actor_id(), None);
14639    }
14640
14641    #[tokio::test]
14642    async fn complete_restore_reconciles_growth_shrink_and_empty_topologies() {
14643        let storage = Arc::new(RuntimeStorage::new([
14644            StorageCapability::Snapshot,
14645            StorageCapability::SessionMetadata,
14646        ]));
14647        let (spawner, registry) = restore_spawner(storage.clone(), 3);
14648        let parent = runtime_storage_agent()
14649            .with_storage(storage.clone())
14650            .with_spawner_handles(Arc::clone(&spawner), Arc::clone(&registry));
14651
14652        for id in ["a", "b"] {
14653            let spawned = spawner
14654                .spawn_with_id(id.to_string(), restore_spec(id))
14655                .await
14656                .unwrap();
14657            spawned.agent.save_session("grow").await.unwrap();
14658            registry.register(spawned).await.unwrap();
14659        }
14660        let staged_c = crate::spawner::storage::NamespacedStorage::new(storage.clone(), "c");
14661        staged_c
14662            .save("grow", &AgentSnapshot::new("c".into()))
14663            .await
14664            .unwrap();
14665        staged_c
14666            .save_metadata("grow", &ai_agents_core::SessionMetadata::default())
14667            .await
14668            .unwrap();
14669        save_restore_target(
14670            &parent,
14671            storage.as_ref(),
14672            "grow",
14673            vec![restore_entry("a"), restore_entry("b"), restore_entry("c")],
14674        )
14675        .await;
14676
14677        assert_eq!(parent.restore_session_full("grow").await.unwrap(), 3);
14678        assert_eq!(registry.count(), 3);
14679        assert!(registry.contains("c"));
14680        assert_eq!(spawner.spawned_count(), 3);
14681
14682        for id in ["a", "b"] {
14683            registry
14684                .get(id)
14685                .unwrap()
14686                .save_session("shrink")
14687                .await
14688                .unwrap();
14689        }
14690        save_restore_target(
14691            &parent,
14692            storage.as_ref(),
14693            "shrink",
14694            vec![restore_entry("a"), restore_entry("b")],
14695        )
14696        .await;
14697
14698        assert_eq!(parent.restore_session_full("shrink").await.unwrap(), 2);
14699        assert_eq!(registry.count(), 2);
14700        assert!(!registry.contains("c"));
14701        assert_eq!(spawner.spawned_count(), 2);
14702
14703        save_restore_target(&parent, storage.as_ref(), "empty", Vec::new()).await;
14704
14705        assert_eq!(parent.restore_session_full("empty").await.unwrap(), 0);
14706        assert_eq!(registry.count(), 0);
14707        assert_eq!(spawner.spawned_count(), 0);
14708        assert_eq!(parent.current_session_id.read().as_deref(), Some("empty"));
14709    }
14710
14711    #[tokio::test]
14712    async fn storage_session_metadata_is_called_only_when_advertised() {
14713        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14714        storage.fail_metadata_save.store(true, Ordering::SeqCst);
14715        storage.fail_metadata_load.store(true, Ordering::SeqCst);
14716        let agent = runtime_storage_agent().with_storage(storage.clone());
14717
14718        agent.save_session("session").await.unwrap();
14719        assert!(agent.load_session("session").await.unwrap());
14720        assert_eq!(storage.metadata_save_calls.load(Ordering::SeqCst), 0);
14721        assert_eq!(storage.metadata_load_calls.load(Ordering::SeqCst), 0);
14722    }
14723
14724    #[cfg(feature = "sqlite")]
14725    #[tokio::test]
14726    async fn sqlite_runtime_save_filter_reopen_and_reload_stay_consistent() {
14727        let directory =
14728            std::env::temp_dir().join(format!("ai-agents-runtime-sqlite-{}", uuid::Uuid::new_v4()));
14729        let path = directory.join("sessions.sqlite");
14730        let path_string = path.to_string_lossy().into_owned();
14731        let storage = Arc::new(
14732            ai_agents_storage::SqliteStorage::new(&path_string)
14733                .await
14734                .unwrap(),
14735        );
14736        let agent = runtime_storage_agent().with_storage(storage.clone());
14737        agent.set_session_metadata(ai_agents_core::SessionMetadata {
14738            tags: vec!["initial".into()],
14739            ..Default::default()
14740        });
14741        agent.chat("persist this turn").await.unwrap();
14742        agent.save_session("session").await.unwrap();
14743
14744        agent.set_session_metadata(ai_agents_core::SessionMetadata {
14745            tags: vec!["updated".into()],
14746            ..Default::default()
14747        });
14748        agent.save_session("session").await.unwrap();
14749        assert!(
14750            agent
14751                .list_sessions_filtered(&ai_agents_core::SessionFilter {
14752                    tags: Some(vec!["initial".into()]),
14753                    ..Default::default()
14754                })
14755                .await
14756                .unwrap()
14757                .is_empty()
14758        );
14759        assert_eq!(
14760            agent
14761                .list_sessions_filtered(&ai_agents_core::SessionFilter {
14762                    tags: Some(vec!["updated".into()]),
14763                    ..Default::default()
14764                })
14765                .await
14766                .unwrap()
14767                .len(),
14768            1
14769        );
14770        drop(agent);
14771        storage.close().await;
14772        drop(storage);
14773
14774        let reopened_storage = Arc::new(
14775            ai_agents_storage::SqliteStorage::new(&path_string)
14776                .await
14777                .unwrap(),
14778        );
14779        let restored = runtime_storage_agent().with_storage(reopened_storage.clone());
14780        assert!(restored.load_session("session").await.unwrap());
14781        assert_eq!(restored.session_metadata().tags, vec!["updated"]);
14782        assert_eq!(
14783            restored.current_session_id.read().as_deref(),
14784            Some("session")
14785        );
14786        assert!(restored.save_state().await.unwrap().memory.messages.len() >= 2);
14787        assert_eq!(
14788            restored
14789                .list_sessions_filtered(&ai_agents_core::SessionFilter {
14790                    tags: Some(vec!["updated".into()]),
14791                    ..Default::default()
14792                })
14793                .await
14794                .unwrap()
14795                .len(),
14796            1
14797        );
14798
14799        drop(restored);
14800        reopened_storage.close().await;
14801        drop(reopened_storage);
14802        crate::remove_sqlite_test_directory(&directory)
14803            .await
14804            .unwrap();
14805    }
14806
14807    #[tokio::test]
14808    async fn storage_session_metadata_backend_failures_propagate() {
14809        let storage = Arc::new(RuntimeStorage::new([
14810            StorageCapability::Snapshot,
14811            StorageCapability::SessionMetadata,
14812        ]));
14813        let agent = runtime_storage_agent().with_storage(storage.clone());
14814
14815        agent.save_session("session").await.unwrap();
14816        storage
14817            .save("target", &agent.save_state().await.unwrap())
14818            .await
14819            .unwrap();
14820        storage.fail_metadata_load.store(true, Ordering::SeqCst);
14821        assert!(matches!(
14822            agent.load_session("target").await,
14823            Err(AgentError::Persistence(message)) if message == "metadata load failed"
14824        ));
14825        assert_eq!(agent.current_session_id.read().as_deref(), Some("session"));
14826
14827        storage.fail_metadata_save.store(true, Ordering::SeqCst);
14828        assert!(matches!(
14829            agent.save_session("session").await,
14830            Err(AgentError::Persistence(message)) if message == "metadata save failed"
14831        ));
14832    }
14833
14834    struct ProviderFutureDropSignal {
14835        dropped: Arc<AtomicBool>,
14836    }
14837
14838    impl Drop for ProviderFutureDropSignal {
14839        fn drop(&mut self) {
14840            self.dropped.store(true, Ordering::SeqCst);
14841        }
14842    }
14843
14844    struct BufferedLockingProvider {
14845        lock: Arc<tokio::sync::Mutex<()>>,
14846        stream_started: Arc<tokio::sync::Notify>,
14847        stream_dropped: Arc<AtomicBool>,
14848        committed_after_drop: Arc<AtomicBool>,
14849    }
14850
14851    #[async_trait]
14852    impl LLMProvider for BufferedLockingProvider {
14853        async fn complete(
14854            &self,
14855            _messages: &[ChatMessage],
14856            _config: Option<&LLMConfig>,
14857        ) -> std::result::Result<LLMResponse, LLMError> {
14858            let _guard = self.lock.lock().await;
14859            self.committed_after_drop
14860                .store(self.stream_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14861            Ok(LLMResponse::new(
14862                "Committed technical response.",
14863                FinishReason::Stop,
14864            ))
14865        }
14866
14867        async fn complete_stream(
14868            &self,
14869            _messages: &[ChatMessage],
14870            _config: Option<&LLMConfig>,
14871        ) -> std::result::Result<
14872            Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14873            LLMError,
14874        > {
14875            let _guard = self.lock.lock().await;
14876            let _drop_signal = ProviderFutureDropSignal {
14877                dropped: Arc::clone(&self.stream_dropped),
14878            };
14879            self.stream_started.notify_one();
14880            std::future::pending().await
14881        }
14882
14883        fn provider_name(&self) -> &str {
14884            "buffered-locking"
14885        }
14886
14887        fn supports(&self, _feature: LLMFeature) -> bool {
14888            false
14889        }
14890    }
14891
14892    struct PendingDropStream {
14893        dropped: Arc<AtomicBool>,
14894        dropped_notify: Arc<tokio::sync::Notify>,
14895    }
14896
14897    impl Stream for PendingDropStream {
14898        type Item = std::result::Result<LLMChunk, LLMError>;
14899
14900        fn poll_next(
14901            self: Pin<&mut Self>,
14902            _cx: &mut std::task::Context<'_>,
14903        ) -> std::task::Poll<Option<Self::Item>> {
14904            std::task::Poll::Pending
14905        }
14906    }
14907
14908    impl Drop for PendingDropStream {
14909        fn drop(&mut self) {
14910            self.dropped.store(true, Ordering::SeqCst);
14911            self.dropped_notify.notify_one();
14912        }
14913    }
14914
14915    struct EstablishedStreamProvider {
14916        stream_started: Arc<tokio::sync::Notify>,
14917        stream_dropped: Arc<AtomicBool>,
14918        stream_dropped_notify: Arc<tokio::sync::Notify>,
14919        committed_after_drop: Arc<AtomicBool>,
14920    }
14921
14922    #[async_trait]
14923    impl LLMProvider for EstablishedStreamProvider {
14924        async fn complete(
14925            &self,
14926            _messages: &[ChatMessage],
14927            _config: Option<&LLMConfig>,
14928        ) -> std::result::Result<LLMResponse, LLMError> {
14929            if !self.stream_dropped.load(Ordering::SeqCst) {
14930                self.stream_dropped_notify.notified().await;
14931            }
14932            self.committed_after_drop
14933                .store(self.stream_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14934            Ok(LLMResponse::new(
14935                "Committed technical response.",
14936                FinishReason::Stop,
14937            ))
14938        }
14939
14940        async fn complete_stream(
14941            &self,
14942            _messages: &[ChatMessage],
14943            _config: Option<&LLMConfig>,
14944        ) -> std::result::Result<
14945            Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14946            LLMError,
14947        > {
14948            self.stream_started.notify_one();
14949            Ok(Box::new(PendingDropStream {
14950                dropped: Arc::clone(&self.stream_dropped),
14951                dropped_notify: Arc::clone(&self.stream_dropped_notify),
14952            }))
14953        }
14954
14955        fn provider_name(&self) -> &str {
14956            "established-stream"
14957        }
14958
14959        fn supports(&self, _feature: LLMFeature) -> bool {
14960            false
14961        }
14962    }
14963
14964    struct FirstCallLockingProvider {
14965        lock: Arc<tokio::sync::Mutex<()>>,
14966        first_started: Arc<tokio::sync::Notify>,
14967        first_dropped: Arc<AtomicBool>,
14968        committed_after_drop: Arc<AtomicBool>,
14969        calls: AtomicU64,
14970    }
14971
14972    #[async_trait]
14973    impl LLMProvider for FirstCallLockingProvider {
14974        async fn complete(
14975            &self,
14976            _messages: &[ChatMessage],
14977            _config: Option<&LLMConfig>,
14978        ) -> std::result::Result<LLMResponse, LLMError> {
14979            let _guard = self.lock.lock().await;
14980            let call = self.calls.fetch_add(1, Ordering::SeqCst);
14981            if call == 0 {
14982                let _drop_signal = ProviderFutureDropSignal {
14983                    dropped: Arc::clone(&self.first_dropped),
14984                };
14985                self.first_started.notify_one();
14986                return std::future::pending().await;
14987            }
14988            self.committed_after_drop
14989                .store(self.first_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14990            Ok(LLMResponse::new(
14991                "Committed technical response.",
14992                FinishReason::Stop,
14993            ))
14994        }
14995
14996        async fn complete_stream(
14997            &self,
14998            _messages: &[ChatMessage],
14999            _config: Option<&LLMConfig>,
15000        ) -> std::result::Result<
15001            Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
15002            LLMError,
15003        > {
15004            Err(LLMError::Other(
15005                "streaming is not used in this test".to_string(),
15006            ))
15007        }
15008
15009        fn provider_name(&self) -> &str {
15010            "first-call-locking"
15011        }
15012
15013        fn supports(&self, _feature: LLMFeature) -> bool {
15014            false
15015        }
15016    }
15017
15018    struct RoutingAfterProviderStart {
15019        provider_started: Arc<tokio::sync::Notify>,
15020    }
15021
15022    #[async_trait]
15023    impl LLMProvider for RoutingAfterProviderStart {
15024        async fn complete(
15025            &self,
15026            _messages: &[ChatMessage],
15027            _config: Option<&LLMConfig>,
15028        ) -> std::result::Result<LLMResponse, LLMError> {
15029            self.provider_started.notified().await;
15030            Ok(LLMResponse::new("1", FinishReason::Stop))
15031        }
15032
15033        async fn complete_stream(
15034            &self,
15035            _messages: &[ChatMessage],
15036            _config: Option<&LLMConfig>,
15037        ) -> std::result::Result<
15038            Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
15039            LLMError,
15040        > {
15041            Err(LLMError::Other(
15042                "streaming is not used in this test".to_string(),
15043            ))
15044        }
15045
15046        fn provider_name(&self) -> &str {
15047            "routing-after-start"
15048        }
15049
15050        fn supports(&self, _feature: LLMFeature) -> bool {
15051            false
15052        }
15053    }
15054
15055    /// Test hook that counts completed responses.
15056    struct ResponseCountingHooks {
15057        responses: Arc<std::sync::atomic::AtomicUsize>,
15058    }
15059
15060    /// Test provider that reports blocking entry and emits a deterministic event stream.
15061    struct RootTurnProbeProvider {
15062        complete_entered: tokio::sync::mpsc::UnboundedSender<()>,
15063    }
15064
15065    /// Calls a configured runtime from on_response and records whether nested root admission succeeded.
15066    struct ResponseChatHooks {
15067        target: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
15068        invoked: AtomicBool,
15069        nested_result: parking_lot::Mutex<Option<std::result::Result<String, String>>>,
15070    }
15071
15072    /// Dispatches one concurrent child from on_response and records the spawned orchestration result.
15073    struct ConcurrentResponseHooks {
15074        registry: Weak<crate::spawner::AgentRegistry>,
15075        child_id: String,
15076        invoked: AtomicBool,
15077        nested_result: parking_lot::Mutex<Option<std::result::Result<String, String>>>,
15078    }
15079
15080    /// Test tool that fails once and records the deadline observed by each invocation attempt.
15081    struct RetryDeadlineTool {
15082        calls: Arc<std::sync::atomic::AtomicUsize>,
15083        deadlines: Arc<parking_lot::Mutex<Vec<chrono::DateTime<chrono::Utc>>>>,
15084        remaining_ms: Arc<parking_lot::Mutex<Vec<i64>>>,
15085    }
15086
15087    /// Test hook that records shared-executor lifecycle order and authoritative records.
15088    struct ToolLifecycleRecordingHooks {
15089        events: parking_lot::Mutex<Vec<String>>,
15090        records: parking_lot::Mutex<Vec<ToolExecutionRecord>>,
15091    }
15092
15093    impl ToolLifecycleRecordingHooks {
15094        /// Creates an empty lifecycle recorder.
15095        fn new() -> Self {
15096            Self {
15097                events: parking_lot::Mutex::new(Vec::new()),
15098                records: parking_lot::Mutex::new(Vec::new()),
15099            }
15100        }
15101
15102        /// Returns a stable snapshot of recorded hook order.
15103        fn events(&self) -> Vec<String> {
15104            self.events.lock().clone()
15105        }
15106
15107        /// Returns a stable snapshot of authoritative execution records.
15108        fn records(&self) -> Vec<ToolExecutionRecord> {
15109            self.records.lock().clone()
15110        }
15111    }
15112
15113    /// Test tool that returns the execution context it received.
15114    struct ContextEchoTool;
15115
15116    #[async_trait]
15117    impl LLMProvider for RootTurnProbeProvider {
15118        async fn complete(
15119            &self,
15120            _messages: &[ChatMessage],
15121            _config: Option<&LLMConfig>,
15122        ) -> std::result::Result<LLMResponse, LLMError> {
15123            let _ = self.complete_entered.send(());
15124            Ok(LLMResponse::new("blocking complete", FinishReason::Stop))
15125        }
15126
15127        async fn complete_stream(
15128            &self,
15129            _messages: &[ChatMessage],
15130            _config: Option<&LLMConfig>,
15131        ) -> std::result::Result<
15132            Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
15133            LLMError,
15134        > {
15135            Ok(Box::new(futures::stream::iter(vec![Ok(
15136                LLMChunk::final_chunk("stream complete", FinishReason::Stop, None),
15137            )])))
15138        }
15139
15140        fn provider_name(&self) -> &str {
15141            "root-turn-probe"
15142        }
15143
15144        fn supports(&self, feature: LLMFeature) -> bool {
15145            matches!(feature, LLMFeature::Streaming)
15146        }
15147    }
15148
15149    #[async_trait]
15150    impl ai_agents_core::Tool for ContextEchoTool {
15151        fn id(&self) -> &str {
15152            "context_echo"
15153        }
15154
15155        fn name(&self) -> &str {
15156            "Context Echo"
15157        }
15158
15159        fn description(&self) -> &str {
15160            "Returns selected execution context fields."
15161        }
15162
15163        fn input_schema(&self) -> Value {
15164            serde_json::json!({"type": "object"})
15165        }
15166
15167        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15168            ai_agents_core::ToolPolicyBindings {
15169                path_fields: vec![ai_agents_core::PathPolicyBinding::read("path")],
15170                result_limit_fields: vec![ai_agents_core::ResultLimitBinding::new(
15171                    "max_results",
15172                    ai_agents_core::ResultLimitKind::MaxResults,
15173                )],
15174                ..Default::default()
15175            }
15176        }
15177
15178        async fn execute(
15179            &self,
15180            _args: Value,
15181            ctx: ai_agents_core::ToolExecutionContext,
15182        ) -> ToolResult {
15183            ToolResult::ok(
15184                serde_json::json!({
15185                    "requested_name": ctx.requested_name,
15186                    "canonical_id": ctx.canonical_id,
15187                    "display_name": ctx.display_name,
15188                    "max_results": ctx.limits.max_results,
15189                    "custom_config": ctx.custom_config,
15190                })
15191                .to_string(),
15192            )
15193        }
15194    }
15195
15196    #[async_trait]
15197    impl ai_agents_core::Tool for RetryDeadlineTool {
15198        fn id(&self) -> &str {
15199            "retry_deadline"
15200        }
15201
15202        fn name(&self) -> &str {
15203            "Retry Deadline"
15204        }
15205
15206        fn description(&self) -> &str {
15207            "Records one deadline per retry invocation."
15208        }
15209
15210        fn input_schema(&self) -> Value {
15211            serde_json::json!({"type": "object"})
15212        }
15213
15214        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15215            ai_agents_core::ToolSafetyMetadata::compute()
15216        }
15217
15218        fn classify_call(&self, _args: &Value) -> ai_agents_core::ToolCallClassification {
15219            let mut classification =
15220                ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
15221            classification.timeout_ms = Some(1_000);
15222            classification.safely_retryable = true;
15223            classification
15224        }
15225
15226        // Records each fresh retry deadline and its initial remaining budget before deciding whether to retry.
15227        async fn execute(
15228            &self,
15229            _args: Value,
15230            ctx: ai_agents_core::ToolExecutionContext,
15231        ) -> ToolResult {
15232            let deadline = ctx
15233                .deadline
15234                .expect("each invocation must receive a deadline");
15235            self.remaining_ms.lock().push(
15236                deadline
15237                    .signed_duration_since(chrono::Utc::now())
15238                    .num_milliseconds(),
15239            );
15240            self.deadlines.lock().push(deadline);
15241            let call = self.calls.fetch_add(1, Ordering::SeqCst);
15242            if call == 0 {
15243                tokio::time::sleep(std::time::Duration::from_millis(20)).await;
15244                ToolResult::error("retry")
15245            } else {
15246                ToolResult::ok("done")
15247            }
15248        }
15249    }
15250
15251    /// Test tool that exposes one call-level timeout and waits long enough to observe the selected timer.
15252    struct ClassifiedTimeoutTool {
15253        id: &'static str,
15254        calls: Arc<std::sync::atomic::AtomicUsize>,
15255        timeout_ms: u64,
15256        sleep_ms: u64,
15257        requires_approval: bool,
15258        remaining_ms: Arc<parking_lot::Mutex<Vec<i64>>>,
15259    }
15260
15261    /// Test tool that changes timeout classification after approval and can hold a shared path lock.
15262    struct ApprovalModifiedTimeoutTool {
15263        calls: Arc<std::sync::atomic::AtomicUsize>,
15264    }
15265
15266    /// Test tool that stays active long enough for runtime cancellation.
15267    struct SlowTool;
15268
15269    /// Test tool that fails once and must not be retried for writes.
15270    struct FlakyWriteTool {
15271        calls: Arc<std::sync::atomic::AtomicUsize>,
15272    }
15273
15274    /// Test tool that tracks concurrent execution on one path.
15275    struct LockedWriteTool {
15276        active: Arc<std::sync::atomic::AtomicUsize>,
15277        max_active: Arc<std::sync::atomic::AtomicUsize>,
15278    }
15279
15280    struct MultiResourceWriteTool {
15281        active: Arc<std::sync::atomic::AtomicUsize>,
15282        max_active: Arc<std::sync::atomic::AtomicUsize>,
15283    }
15284
15285    #[derive(Clone)]
15286    struct PathMutationGate {
15287        entered: Arc<AtomicBool>,
15288        entered_notify: Arc<tokio::sync::Notify>,
15289        release: Arc<tokio::sync::Notify>,
15290    }
15291
15292    impl PathMutationGate {
15293        fn new() -> Self {
15294            Self {
15295                entered: Arc::new(AtomicBool::new(false)),
15296                entered_notify: Arc::new(tokio::sync::Notify::new()),
15297                release: Arc::new(tokio::sync::Notify::new()),
15298            }
15299        }
15300
15301        async fn wait_until_entered(&self) {
15302            if !self.entered.load(Ordering::SeqCst) {
15303                self.entered_notify.notified().await;
15304            }
15305        }
15306
15307        fn release(&self) {
15308            self.release.notify_one();
15309        }
15310    }
15311
15312    struct BlockingPathMutationTool {
15313        id: &'static str,
15314        path_fields: Vec<ai_agents_core::PathPolicyBinding>,
15315        gate: PathMutationGate,
15316    }
15317
15318    struct NoBindingWriteTool {
15319        active: Arc<std::sync::atomic::AtomicUsize>,
15320        max_active: Arc<std::sync::atomic::AtomicUsize>,
15321    }
15322
15323    struct RecoveryTestTool {
15324        id: String,
15325        succeeds: bool,
15326        calls: Arc<std::sync::atomic::AtomicUsize>,
15327        max_output_chars: Option<usize>,
15328    }
15329
15330    struct BlockingApprovalHandler {
15331        entered: Arc<tokio::sync::Barrier>,
15332        release: Arc<tokio::sync::Notify>,
15333        result: ApprovalResult,
15334    }
15335
15336    struct CountingApprovalHandler {
15337        calls: Arc<std::sync::atomic::AtomicUsize>,
15338    }
15339
15340    /// Test provider that moves one fallback alias to an ancestor canonical ID during refresh.
15341    struct DriftingFallbackProvider {
15342        refreshed: AtomicBool,
15343        primary_calls: Arc<std::sync::atomic::AtomicUsize>,
15344        secondary_calls: Arc<std::sync::atomic::AtomicUsize>,
15345    }
15346
15347    /// Test hook that records fallback lifecycle evidence and refreshes the provider after initial canonical admission.
15348    struct RefreshFallbackProviderHooks {
15349        agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
15350        lifecycle: Arc<ToolLifecycleRecordingHooks>,
15351    }
15352
15353    struct RuntimeWebFetchTransport {
15354        calls: Arc<std::sync::atomic::AtomicUsize>,
15355    }
15356
15357    struct RuntimeWebFetchResolver;
15358
15359    struct ReentrantToolHooks {
15360        agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
15361        invoked: AtomicBool,
15362        nested_success: AtomicBool,
15363    }
15364
15365    #[async_trait]
15366    impl ai_agents_core::Tool for ClassifiedTimeoutTool {
15367        // Returns the configurable ID used to select timeout policy in each test.
15368        fn id(&self) -> &str {
15369            self.id
15370        }
15371
15372        // Provides a stable display name for the test registry.
15373        fn name(&self) -> &str {
15374            "Classified Timeout"
15375        }
15376
15377        // Describes the test tool's deadline observation behavior.
15378        fn description(&self) -> &str {
15379            "Records and waits under one call-level timeout."
15380        }
15381
15382        // Accepts empty object arguments so timeout tests isolate execution limits.
15383        fn input_schema(&self) -> Value {
15384            serde_json::json!({"type": "object"})
15385        }
15386
15387        // Exposes the configured call-level cap and approval requirement to the runtime.
15388        fn classify_call(&self, _args: &Value) -> ai_agents_core::ToolCallClassification {
15389            let mut classification =
15390                ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
15391            classification.timeout_ms = Some(self.timeout_ms);
15392            classification.requires_approval = self.requires_approval;
15393            classification
15394        }
15395
15396        // Records the visible deadline and sleeps past the selected cap so the runtime timer must terminate execution.
15397        async fn execute(
15398            &self,
15399            _args: Value,
15400            ctx: ai_agents_core::ToolExecutionContext,
15401        ) -> ToolResult {
15402            self.calls.fetch_add(1, Ordering::SeqCst);
15403            let deadline = ctx
15404                .deadline
15405                .expect("each invocation must receive a deadline");
15406            self.remaining_ms.lock().push(
15407                deadline
15408                    .signed_duration_since(chrono::Utc::now())
15409                    .num_milliseconds(),
15410            );
15411            tokio::time::sleep(Duration::from_millis(self.sleep_ms)).await;
15412            ToolResult::ok("done")
15413        }
15414    }
15415
15416    #[async_trait]
15417    impl ai_agents_core::Tool for ApprovalModifiedTimeoutTool {
15418        // Returns the canonical ID used by the approval-modification timeout test.
15419        fn id(&self) -> &str {
15420            "approval_modified_timeout"
15421        }
15422
15423        // Provides a stable display name for approval requests.
15424        fn name(&self) -> &str {
15425            "Approval Modified Timeout"
15426        }
15427
15428        // Describes the argument-dependent timeout behavior under test.
15429        fn description(&self) -> &str {
15430            "Becomes invalid only after approval modifies its arguments."
15431        }
15432
15433        // Accepts the path and approval-controlled invalid-timeout switch.
15434        fn input_schema(&self) -> Value {
15435            serde_json::json!({"type": "object"})
15436        }
15437
15438        // Binds one path so a misplaced final validation would wait for a resource lock.
15439        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15440            ai_agents_core::ToolPolicyBindings {
15441                path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15442                ..Default::default()
15443            }
15444        }
15445
15446        // Keeps the test call non-concurrent so final timeout validation must precede its path lock.
15447        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15448            ai_agents_core::ToolSafetyMetadata {
15449                read_only: false,
15450                concurrency_safe: false,
15451                operation: ai_agents_core::ToolOperationKind::Write,
15452                side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
15453                requires_network: false,
15454                destructive: false,
15455                open_world: false,
15456                host_dependent: false,
15457                requires_user_interaction: false,
15458                supports_cancellation: true,
15459                default_requires_approval: true,
15460                should_defer_schema: false,
15461                max_output_chars: Some(1024),
15462                max_result_size_chars: Some(1024),
15463            }
15464        }
15465
15466        // Produces an invalid cap only from the arguments returned by the approval handler.
15467        fn classify_call(&self, args: &Value) -> ai_agents_core::ToolCallClassification {
15468            let mut classification =
15469                ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
15470            classification.timeout_ms = Some(if args["invalid_timeout"].as_bool() == Some(true) {
15471                u64::MAX
15472            } else {
15473                1_000
15474            });
15475            classification
15476        }
15477
15478        // Records any incorrect invocation after final classification should have failed.
15479        async fn execute(
15480            &self,
15481            _args: Value,
15482            _ctx: ai_agents_core::ToolExecutionContext,
15483        ) -> ToolResult {
15484            self.calls.fetch_add(1, Ordering::SeqCst);
15485            ToolResult::ok("unexpected")
15486        }
15487    }
15488
15489    #[async_trait]
15490    impl ai_agents_core::Tool for SlowTool {
15491        fn id(&self) -> &str {
15492            "slow"
15493        }
15494
15495        fn name(&self) -> &str {
15496            "Slow"
15497        }
15498
15499        fn description(&self) -> &str {
15500            "Waits until cancelled or timed out."
15501        }
15502
15503        fn input_schema(&self) -> Value {
15504            serde_json::json!({"type": "object"})
15505        }
15506
15507        async fn execute(
15508            &self,
15509            _args: Value,
15510            _ctx: ai_agents_core::ToolExecutionContext,
15511        ) -> ToolResult {
15512            tokio::time::sleep(std::time::Duration::from_secs(5)).await;
15513            ToolResult::ok("done")
15514        }
15515    }
15516
15517    #[async_trait]
15518    impl ai_agents_core::Tool for FlakyWriteTool {
15519        fn id(&self) -> &str {
15520            "flaky_write"
15521        }
15522
15523        fn name(&self) -> &str {
15524            "Flaky Write"
15525        }
15526
15527        fn description(&self) -> &str {
15528            "Fails on the first write attempt."
15529        }
15530
15531        fn input_schema(&self) -> Value {
15532            serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
15533        }
15534
15535        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15536            ai_agents_core::ToolPolicyBindings {
15537                path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15538                ..Default::default()
15539            }
15540        }
15541
15542        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15543            ai_agents_core::ToolSafetyMetadata {
15544                read_only: false,
15545                concurrency_safe: false,
15546                operation: ai_agents_core::ToolOperationKind::Write,
15547                side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
15548                requires_network: false,
15549                destructive: false,
15550                open_world: false,
15551                host_dependent: false,
15552                requires_user_interaction: false,
15553                supports_cancellation: true,
15554                default_requires_approval: false,
15555                should_defer_schema: false,
15556                max_output_chars: Some(1024),
15557                max_result_size_chars: Some(1024),
15558            }
15559        }
15560
15561        fn classify_call(&self, _args: &Value) -> ai_agents_core::ToolCallClassification {
15562            let mut classification =
15563                ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
15564            classification.safely_retryable = false;
15565            classification
15566        }
15567
15568        async fn execute(
15569            &self,
15570            _args: Value,
15571            _ctx: ai_agents_core::ToolExecutionContext,
15572        ) -> ToolResult {
15573            let call = self.calls.fetch_add(1, Ordering::SeqCst);
15574            if call == 0 {
15575                ToolResult::error("first failure")
15576            } else {
15577                ToolResult::ok("second success")
15578            }
15579        }
15580    }
15581
15582    #[async_trait]
15583    impl ai_agents_core::Tool for LockedWriteTool {
15584        fn id(&self) -> &str {
15585            "locked_write"
15586        }
15587
15588        fn name(&self) -> &str {
15589            "Locked Write"
15590        }
15591
15592        fn description(&self) -> &str {
15593            "Tracks concurrent execution on one resource."
15594        }
15595
15596        fn input_schema(&self) -> Value {
15597            serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
15598        }
15599
15600        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15601            ai_agents_core::ToolPolicyBindings {
15602                path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15603                ..Default::default()
15604            }
15605        }
15606
15607        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15608            ai_agents_core::ToolSafetyMetadata {
15609                read_only: false,
15610                concurrency_safe: false,
15611                operation: ai_agents_core::ToolOperationKind::Write,
15612                side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
15613                requires_network: false,
15614                destructive: false,
15615                open_world: false,
15616                host_dependent: false,
15617                requires_user_interaction: false,
15618                supports_cancellation: true,
15619                default_requires_approval: false,
15620                should_defer_schema: false,
15621                max_output_chars: Some(1024),
15622                max_result_size_chars: Some(1024),
15623            }
15624        }
15625
15626        async fn execute(
15627            &self,
15628            _args: Value,
15629            _ctx: ai_agents_core::ToolExecutionContext,
15630        ) -> ToolResult {
15631            let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
15632            loop {
15633                let current_max = self.max_active.load(Ordering::SeqCst);
15634                if active <= current_max {
15635                    break;
15636                }
15637                if self
15638                    .max_active
15639                    .compare_exchange(current_max, active, Ordering::SeqCst, Ordering::SeqCst)
15640                    .is_ok()
15641                {
15642                    break;
15643                }
15644            }
15645            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
15646            self.active.fetch_sub(1, Ordering::SeqCst);
15647            ToolResult::ok("done")
15648        }
15649    }
15650
15651    #[async_trait]
15652    impl ai_agents_core::Tool for MultiResourceWriteTool {
15653        fn id(&self) -> &str {
15654            "multi_resource_write"
15655        }
15656
15657        fn name(&self) -> &str {
15658            "Multi Resource Write"
15659        }
15660
15661        fn description(&self) -> &str {
15662            "Tracks concurrent execution across source and destination resources."
15663        }
15664
15665        fn input_schema(&self) -> Value {
15666            serde_json::json!({"type": "object"})
15667        }
15668
15669        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15670            ai_agents_core::ToolPolicyBindings {
15671                path_fields: vec![
15672                    ai_agents_core::PathPolicyBinding::read_write("source_path"),
15673                    ai_agents_core::PathPolicyBinding::write("destination_path"),
15674                ],
15675                ..Default::default()
15676            }
15677        }
15678
15679        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15680            LockedWriteTool {
15681                active: Arc::clone(&self.active),
15682                max_active: Arc::clone(&self.max_active),
15683            }
15684            .safety_metadata()
15685        }
15686
15687        async fn execute(
15688            &self,
15689            _args: Value,
15690            _ctx: ai_agents_core::ToolExecutionContext,
15691        ) -> ToolResult {
15692            let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
15693            self.max_active.fetch_max(active, Ordering::SeqCst);
15694            tokio::time::sleep(std::time::Duration::from_millis(75)).await;
15695            self.active.fetch_sub(1, Ordering::SeqCst);
15696            ToolResult::ok("done")
15697        }
15698    }
15699
15700    #[async_trait]
15701    impl ai_agents_core::Tool for BlockingPathMutationTool {
15702        fn id(&self) -> &str {
15703            self.id
15704        }
15705
15706        fn name(&self) -> &str {
15707            self.id
15708        }
15709
15710        fn description(&self) -> &str {
15711            "Blocks a path mutation until the test releases it."
15712        }
15713
15714        fn input_schema(&self) -> Value {
15715            serde_json::json!({"type": "object"})
15716        }
15717
15718        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15719            ai_agents_core::ToolPolicyBindings {
15720                path_fields: self.path_fields.clone(),
15721                ..Default::default()
15722            }
15723        }
15724
15725        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15726            ai_agents_core::ToolSafetyMetadata {
15727                read_only: false,
15728                concurrency_safe: false,
15729                operation: ai_agents_core::ToolOperationKind::Write,
15730                side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
15731                requires_network: false,
15732                destructive: false,
15733                open_world: false,
15734                host_dependent: false,
15735                requires_user_interaction: false,
15736                supports_cancellation: true,
15737                default_requires_approval: false,
15738                should_defer_schema: false,
15739                max_output_chars: Some(1024),
15740                max_result_size_chars: Some(1024),
15741            }
15742        }
15743
15744        async fn execute(
15745            &self,
15746            _args: Value,
15747            _ctx: ai_agents_core::ToolExecutionContext,
15748        ) -> ToolResult {
15749            self.gate.entered.store(true, Ordering::SeqCst);
15750            self.gate.entered_notify.notify_one();
15751            self.gate.release.notified().await;
15752            ToolResult::ok("done")
15753        }
15754    }
15755
15756    #[async_trait]
15757    impl ai_agents_core::Tool for NoBindingWriteTool {
15758        fn id(&self) -> &str {
15759            "no_binding_write"
15760        }
15761
15762        fn name(&self) -> &str {
15763            "No Binding Write"
15764        }
15765
15766        fn description(&self) -> &str {
15767            "Tracks concurrent execution without resource bindings."
15768        }
15769
15770        fn input_schema(&self) -> Value {
15771            serde_json::json!({"type": "object"})
15772        }
15773
15774        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15775            LockedWriteTool {
15776                active: Arc::clone(&self.active),
15777                max_active: Arc::clone(&self.max_active),
15778            }
15779            .safety_metadata()
15780        }
15781
15782        async fn execute(
15783            &self,
15784            _args: Value,
15785            _ctx: ai_agents_core::ToolExecutionContext,
15786        ) -> ToolResult {
15787            let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
15788            self.max_active.fetch_max(active, Ordering::SeqCst);
15789            tokio::time::sleep(std::time::Duration::from_millis(75)).await;
15790            self.active.fetch_sub(1, Ordering::SeqCst);
15791            ToolResult::ok("done")
15792        }
15793    }
15794
15795    #[async_trait]
15796    impl ai_agents_core::Tool for RecoveryTestTool {
15797        fn id(&self) -> &str {
15798            &self.id
15799        }
15800
15801        fn name(&self) -> &str {
15802            &self.id
15803        }
15804
15805        fn description(&self) -> &str {
15806            "Records recovery execution and returns a configured result."
15807        }
15808
15809        fn input_schema(&self) -> Value {
15810            serde_json::json!({"type": "object"})
15811        }
15812
15813        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15814            ai_agents_core::ToolPolicyBindings {
15815                path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15816                ..Default::default()
15817            }
15818        }
15819
15820        /// Supplies a configurable output cap while preserving non-concurrent path mutation behavior.
15821        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15822            ai_agents_core::ToolSafetyMetadata {
15823                read_only: false,
15824                concurrency_safe: false,
15825                operation: ai_agents_core::ToolOperationKind::Write,
15826                side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
15827                requires_network: false,
15828                destructive: false,
15829                open_world: false,
15830                host_dependent: false,
15831                requires_user_interaction: false,
15832                supports_cancellation: true,
15833                default_requires_approval: false,
15834                should_defer_schema: false,
15835                max_output_chars: Some(self.max_output_chars.unwrap_or(1024)),
15836                max_result_size_chars: Some(1024),
15837            }
15838        }
15839
15840        /// Returns deterministic output and metadata for recovery lifecycle assertions.
15841        async fn execute(
15842            &self,
15843            _args: Value,
15844            _ctx: ai_agents_core::ToolExecutionContext,
15845        ) -> ToolResult {
15846            self.calls.fetch_add(1, Ordering::SeqCst);
15847            let mut result = if self.succeeds {
15848                ToolResult::ok(format!("{} succeeded", self.id))
15849            } else {
15850                ToolResult::error(format!("{} failed", self.id))
15851            };
15852            result.metadata = Some(HashMap::from([(
15853                "recovery_test_tool".to_string(),
15854                Value::String(self.id.clone()),
15855            )]));
15856            result
15857        }
15858    }
15859
15860    #[async_trait]
15861    impl WebFetchTransport for RuntimeWebFetchTransport {
15862        /// Rejects unvalidated transport calls in the shared-executor fixture.
15863        async fn send(
15864            &self,
15865            _request: WebFetchTransportRequest,
15866        ) -> std::result::Result<WebFetchTransportResponse, String> {
15867            Err("validated addresses are required".to_string())
15868        }
15869
15870        /// Records transport only after runtime and tool policy validation complete.
15871        async fn send_validated(
15872            &self,
15873            _request: WebFetchTransportRequest,
15874            _addresses: &[std::net::SocketAddr],
15875        ) -> std::result::Result<WebFetchTransportResponse, String> {
15876            self.calls.fetch_add(1, Ordering::SeqCst);
15877            Ok(WebFetchTransportResponse {
15878                status: 200,
15879                content_type: Some("text/plain".to_string()),
15880                location: None,
15881                body: b"approved".to_vec(),
15882            })
15883        }
15884    }
15885
15886    #[async_trait]
15887    impl WebFetchResolver for RuntimeWebFetchResolver {
15888        /// Returns one public fixture address without external DNS.
15889        async fn resolve(
15890            &self,
15891            _host: &str,
15892            _port: u16,
15893        ) -> std::result::Result<Vec<std::net::IpAddr>, String> {
15894            Ok(vec![std::net::IpAddr::V4(std::net::Ipv4Addr::new(
15895                93, 184, 216, 34,
15896            ))])
15897        }
15898    }
15899
15900    #[async_trait]
15901    impl ToolProvider for DriftingFallbackProvider {
15902        // Returns the provider ID used by the registry refresh test.
15903        fn id(&self) -> &str {
15904            "drifting_fallback"
15905        }
15906
15907        // Provides a stable test provider name.
15908        fn name(&self) -> &str {
15909            "Drifting Fallback"
15910        }
15911
15912        // Marks the provider as a custom in-process fixture.
15913        fn provider_type(&self) -> ToolProviderType {
15914            ToolProviderType::Custom
15915        }
15916
15917        // Moves the fallback alias from secondary to primary after refresh.
15918        async fn list_tools(&self) -> Vec<ToolDescriptor> {
15919            let alias = ToolAliases::new().with_name("en", "fallback alias");
15920            let mut primary = ToolDescriptor::new(
15921                "primary",
15922                "Primary",
15923                "Fails before fallback.",
15924                serde_json::json!({"type": "object"}),
15925            );
15926            let mut secondary = ToolDescriptor::new(
15927                "secondary",
15928                "Secondary",
15929                "Must not execute after final canonical drift.",
15930                serde_json::json!({"type": "object"}),
15931            );
15932            if self.refreshed.load(Ordering::SeqCst) {
15933                primary = primary.with_aliases(alias);
15934            } else {
15935                secondary = secondary.with_aliases(alias);
15936            }
15937            vec![primary, secondary]
15938        }
15939
15940        // Returns deterministic failing tools so fallback ancestry is the only terminal control.
15941        async fn get_tool(&self, tool_id: &str) -> Option<Arc<dyn Tool>> {
15942            let calls = match tool_id {
15943                "primary" => Arc::clone(&self.primary_calls),
15944                "secondary" => Arc::clone(&self.secondary_calls),
15945                _ => return None,
15946            };
15947            Some(Arc::new(RecoveryTestTool {
15948                id: tool_id.to_string(),
15949                succeeds: false,
15950                calls,
15951                max_output_chars: None,
15952            }))
15953        }
15954
15955        // Allows the registry to rebuild canonical and alias indexes during the test hook.
15956        fn supports_refresh(&self) -> bool {
15957            true
15958        }
15959
15960        // Switches the alias mapping before the registry recreates its provider snapshot.
15961        async fn refresh(&self) -> std::result::Result<(), ToolProviderError> {
15962            self.refreshed.store(true, Ordering::SeqCst);
15963            Ok(())
15964        }
15965    }
15966
15967    #[async_trait]
15968    impl AgentHooks for RefreshFallbackProviderHooks {
15969        // Records the admitted lifecycle identity before refreshing the fallback provider at the intended async boundary.
15970        async fn on_tool_start(&self, tool: &str, args: &Value) {
15971            self.lifecycle.on_tool_start(tool, args).await;
15972            if tool != "secondary" {
15973                return;
15974            }
15975            let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
15976            if let Some(agent) = agent {
15977                agent
15978                    .tools
15979                    .refresh_provider("drifting_fallback")
15980                    .await
15981                    .unwrap();
15982            }
15983        }
15984
15985        async fn on_tool_complete(&self, tool: &str, result: &ToolResult, duration_ms: u64) {
15986            self.lifecycle
15987                .on_tool_complete(tool, result, duration_ms)
15988                .await;
15989        }
15990
15991        async fn on_tool_execution_record(&self, record: &ToolExecutionRecord) {
15992            self.lifecycle.on_tool_execution_record(record).await;
15993        }
15994
15995        async fn on_error(&self, error: &AgentError) {
15996            self.lifecycle.on_error(error).await;
15997        }
15998    }
15999
16000    #[async_trait]
16001    impl ApprovalHandler for BlockingApprovalHandler {
16002        async fn request_approval(
16003            &self,
16004            _request: ai_agents_hitl::ApprovalRequest,
16005        ) -> ApprovalResult {
16006            self.entered.wait().await;
16007            self.release.notified().await;
16008            self.result.clone()
16009        }
16010    }
16011
16012    #[async_trait]
16013    impl ApprovalHandler for CountingApprovalHandler {
16014        async fn request_approval(
16015            &self,
16016            _request: ai_agents_hitl::ApprovalRequest,
16017        ) -> ApprovalResult {
16018            self.calls.fetch_add(1, Ordering::SeqCst);
16019            ApprovalResult::Approved
16020        }
16021    }
16022
16023    #[async_trait]
16024    impl AgentHooks for ReentrantToolHooks {
16025        async fn on_tool_complete(&self, tool: &str, _result: &ToolResult, _duration_ms: u64) {
16026            if tool != "reentrant_write" || self.invoked.swap(true, Ordering::SeqCst) {
16027                return;
16028            }
16029            let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
16030            if let Some(agent) = agent {
16031                let result = agent
16032                    .invoke_tool(ToolExecutionRequest::new(
16033                        "nested-hook-call",
16034                        "reentrant_write",
16035                        serde_json::json!({"path": "./hook.txt"}),
16036                        ToolCallSource::Manual,
16037                    ))
16038                    .await;
16039                self.nested_success
16040                    .store(result.is_ok_and(|record| record.success), Ordering::SeqCst);
16041            }
16042        }
16043    }
16044
16045    #[async_trait]
16046    impl AgentHooks for ResponseCountingHooks {
16047        async fn on_response(&self, _response: &AgentResponse) {
16048            self.responses.fetch_add(1, Ordering::SeqCst);
16049        }
16050    }
16051
16052    #[async_trait]
16053    impl AgentHooks for ResponseChatHooks {
16054        /// Attempts one nested blocking turn without retaining the target mutex across the await.
16055        async fn on_response(&self, _response: &AgentResponse) {
16056            if self.invoked.swap(true, Ordering::SeqCst) {
16057                return;
16058            }
16059            let target = self.target.lock().as_ref().and_then(Weak::upgrade);
16060            let result = if let Some(target) = target {
16061                target
16062                    .chat("nested response hook call")
16063                    .await
16064                    .map(|response| response.content)
16065                    .map_err(|error| error.to_string())
16066            } else {
16067                Err("response hook target is unavailable".to_string())
16068            };
16069            *self.nested_result.lock() = Some(result);
16070        }
16071    }
16072
16073    #[async_trait]
16074    impl AgentHooks for ConcurrentResponseHooks {
16075        /// Runs one child through the real concurrent JoinSet boundary without retaining registry state across the await.
16076        async fn on_response(&self, _response: &AgentResponse) {
16077            if self.invoked.swap(true, Ordering::SeqCst) {
16078                return;
16079            }
16080            let Some(registry) = self.registry.upgrade() else {
16081                *self.nested_result.lock() =
16082                    Some(Err("concurrent registry is unavailable".to_string()));
16083                return;
16084            };
16085            let agents = [ai_agents_state::ConcurrentAgentRef::Id(
16086                self.child_id.clone(),
16087            )];
16088            let aggregation = ai_agents_state::AggregationConfig {
16089                strategy: ai_agents_state::AggregationStrategy::FirstWins,
16090                synthesizer_llm: None,
16091                synthesizer_prompt: None,
16092                vote: None,
16093            };
16094            let result = crate::orchestration::concurrent(
16095                &registry,
16096                "nested concurrent response hook call",
16097                &agents,
16098                &aggregation,
16099                None,
16100                Some(1),
16101                None,
16102                ai_agents_state::PartialFailureAction::Abort,
16103                None,
16104            )
16105            .await
16106            .map(|result| result.response.content)
16107            .map_err(|error| error.to_string());
16108            *self.nested_result.lock() = Some(result);
16109        }
16110    }
16111
16112    #[async_trait]
16113    impl AgentHooks for ToolLifecycleRecordingHooks {
16114        async fn on_tool_start(&self, tool: &str, _args: &Value) {
16115            self.events.lock().push(format!("start:{tool}"));
16116        }
16117
16118        async fn on_tool_complete(&self, tool: &str, result: &ToolResult, _duration_ms: u64) {
16119            self.events
16120                .lock()
16121                .push(format!("complete:{tool}:{}", result.success));
16122        }
16123
16124        async fn on_tool_execution_record(&self, record: &ToolExecutionRecord) {
16125            self.events.lock().push(format!(
16126                "record:{}:{}",
16127                record.canonical_id, record.executed
16128            ));
16129            self.records.lock().push(record.clone());
16130        }
16131
16132        /// Records error reporting so fallback cannot overtake the failed original lifecycle.
16133        async fn on_error(&self, _error: &AgentError) {
16134            self.events.lock().push("error".to_string());
16135        }
16136    }
16137
16138    struct ApprovalRecordingHooks {
16139        events: parking_lot::Mutex<Vec<String>>,
16140    }
16141
16142    impl ApprovalRecordingHooks {
16143        fn new() -> Self {
16144            Self {
16145                events: parking_lot::Mutex::new(Vec::new()),
16146            }
16147        }
16148
16149        fn events(&self) -> Vec<String> {
16150            self.events.lock().clone()
16151        }
16152    }
16153
16154    #[async_trait]
16155    impl AgentHooks for ApprovalRecordingHooks {
16156        async fn on_approval_result(&self, request_id: &str, result: &ApprovalResult) {
16157            self.events.lock().push(format!(
16158                "raw:{}:{}",
16159                request_id,
16160                approval_result_name(result)
16161            ));
16162        }
16163
16164        async fn on_approval_resolved(
16165            &self,
16166            request: &ai_agents_hitl::ApprovalRequest,
16167            raw_result: &ApprovalResult,
16168            outcome: &ApprovalResolvedOutcome,
16169        ) {
16170            self.events.lock().push(format!(
16171                "resolved:{}:{}:{}",
16172                request.id,
16173                approval_result_name(raw_result),
16174                approval_outcome_name(outcome)
16175            ));
16176        }
16177    }
16178
16179    fn approval_result_name(result: &ApprovalResult) -> &'static str {
16180        match result {
16181            ApprovalResult::Approved => "approved",
16182            ApprovalResult::Rejected { .. } => "rejected",
16183            ApprovalResult::Modified { .. } => "modified",
16184            ApprovalResult::Timeout => "timeout",
16185        }
16186    }
16187
16188    fn approval_outcome_name(outcome: &ApprovalResolvedOutcome) -> &'static str {
16189        match outcome {
16190            ApprovalResolvedOutcome::Approved => "approved",
16191            ApprovalResolvedOutcome::Rejected { .. } => "rejected",
16192            ApprovalResolvedOutcome::Modified { .. } => "modified",
16193            ApprovalResolvedOutcome::Error { .. } => "error",
16194        }
16195    }
16196
16197    fn assert_correlated_approval_events(
16198        events: &[String],
16199        raw_status: &str,
16200        outcome_status: &str,
16201    ) {
16202        assert_eq!(events.len(), 2);
16203        let raw: Vec<_> = events[0].split(':').collect();
16204        let resolved: Vec<_> = events[1].split(':').collect();
16205        assert_eq!(raw[0], "raw");
16206        assert_eq!(resolved[0], "resolved");
16207        assert_eq!(raw[1], resolved[1]);
16208        assert_eq!(raw[2], raw_status);
16209        assert_eq!(resolved[2], raw_status);
16210        assert_eq!(resolved[3], outcome_status);
16211    }
16212
16213    fn approval_security_config(policy_enabled: bool) -> ToolSecurityConfig {
16214        let mut security = ToolSecurityConfig {
16215            enabled: true,
16216            fail_closed: true,
16217            ..Default::default()
16218        };
16219        let policy = ai_agents_tools::ToolPolicyConfig {
16220            enabled: policy_enabled,
16221            write_paths: vec![".".to_string()],
16222            require_confirmation: true,
16223            ..Default::default()
16224        };
16225        security.tools.insert("locked_write".to_string(), policy);
16226        security
16227    }
16228
16229    struct MutationTestWorkspace {
16230        root: std::path::PathBuf,
16231    }
16232
16233    impl MutationTestWorkspace {
16234        fn new() -> Self {
16235            let root = std::env::temp_dir().join(format!(
16236                "ai-agents-runtime-mutation-{}",
16237                uuid::Uuid::new_v4()
16238            ));
16239            std::fs::create_dir_all(&root).unwrap();
16240            Self { root }
16241        }
16242    }
16243
16244    impl Drop for MutationTestWorkspace {
16245        fn drop(&mut self) {
16246            let _ = std::fs::remove_dir_all(&self.root);
16247        }
16248    }
16249
16250    async fn wait_for_resource_lock_strong_count(locks: &ToolResourceLocks, minimum: usize) {
16251        tokio::time::timeout(std::time::Duration::from_secs(2), async {
16252            loop {
16253                let strong_count = locks
16254                    .read()
16255                    .get("path-mutation:global")
16256                    .map_or(0, |lock| lock.strong_count());
16257                if strong_count >= minimum {
16258                    break;
16259                }
16260                tokio::task::yield_now().await;
16261            }
16262        })
16263        .await
16264        .expect("path mutation call did not reach the shared lock");
16265    }
16266
16267    async fn assert_path_mutation_pair_serialized(
16268        first_id: &'static str,
16269        first_fields: Vec<ai_agents_core::PathPolicyBinding>,
16270        first_args: Value,
16271        second_id: &'static str,
16272        second_fields: Vec<ai_agents_core::PathPolicyBinding>,
16273        second_args: Value,
16274    ) {
16275        let locks = new_tool_resource_locks();
16276        let first_gate = PathMutationGate::new();
16277        let second_gate = PathMutationGate::new();
16278        second_gate.release();
16279        let agent = Arc::new(
16280            AgentBuilder::new()
16281                .system_prompt("Test global path mutation locking.")
16282                .llm(Arc::new(mock_with_response("done")))
16283                .tool(Arc::new(BlockingPathMutationTool {
16284                    id: first_id,
16285                    path_fields: first_fields,
16286                    gate: first_gate.clone(),
16287                }))
16288                .tool(Arc::new(BlockingPathMutationTool {
16289                    id: second_id,
16290                    path_fields: second_fields,
16291                    gate: second_gate.clone(),
16292                }))
16293                .build()
16294                .unwrap()
16295                .with_shared_resource_locks(Arc::clone(&locks)),
16296        );
16297
16298        let first = {
16299            let agent = Arc::clone(&agent);
16300            tokio::spawn(async move {
16301                agent
16302                    .invoke_tool(ToolExecutionRequest::new(
16303                        format!("{}-first", first_id),
16304                        first_id,
16305                        first_args,
16306                        ToolCallSource::Manual,
16307                    ))
16308                    .await
16309                    .unwrap()
16310            })
16311        };
16312        first_gate.wait_until_entered().await;
16313
16314        let second = {
16315            let agent = Arc::clone(&agent);
16316            tokio::spawn(async move {
16317                agent
16318                    .invoke_tool(ToolExecutionRequest::new(
16319                        format!("{}-second", second_id),
16320                        second_id,
16321                        second_args,
16322                        ToolCallSource::Manual,
16323                    ))
16324                    .await
16325                    .unwrap()
16326            })
16327        };
16328        wait_for_resource_lock_strong_count(&locks, 2).await;
16329        assert!(!second_gate.entered.load(Ordering::SeqCst));
16330        assert!(!second.is_finished());
16331
16332        first_gate.release();
16333        let (first, second) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
16334            tokio::join!(first, second)
16335        })
16336        .await
16337        .expect("serialized path mutation calls did not finish");
16338        assert!(first.unwrap().success);
16339        assert!(second.unwrap().success);
16340        assert!(second_gate.entered.load(Ordering::SeqCst));
16341        assert!(locks.read().is_empty());
16342    }
16343
16344    #[derive(Clone, Copy)]
16345    enum MutationDenial {
16346        Policy,
16347        Approval,
16348    }
16349
16350    fn mutation_denial_security_config(
16351        tool_id: &str,
16352        workspace: &std::path::Path,
16353        denial: MutationDenial,
16354    ) -> ToolSecurityConfig {
16355        let workspace = workspace.to_string_lossy().into_owned();
16356        let mut policy = ai_agents_tools::ToolPolicyConfig {
16357            read_paths: vec![workspace.clone()],
16358            write_paths: vec![workspace.clone()],
16359            ..Default::default()
16360        };
16361        match denial {
16362            MutationDenial::Policy => policy.blocked_paths = vec![workspace],
16363            MutationDenial::Approval => policy.require_confirmation = true,
16364        }
16365
16366        let mut security = ToolSecurityConfig {
16367            enabled: true,
16368            fail_closed: true,
16369            ..Default::default()
16370        };
16371        security.tools.insert(tool_id.to_string(), policy);
16372        security
16373    }
16374
16375    async fn assert_path_mutation_denied(tool: Arc<dyn Tool>, denial: MutationDenial) {
16376        let workspace = MutationTestWorkspace::new();
16377        let tool_id = tool.id().to_string();
16378        let preserved = workspace.root.join(format!("{}-preserved.txt", tool_id));
16379        let destination = workspace.root.join(format!("{}-destination.txt", tool_id));
16380        std::fs::write(&preserved, "preserved").unwrap();
16381        let arguments = match tool_id.as_str() {
16382            "copy_path" | "move_path" => serde_json::json!({
16383                "source_path": preserved.to_string_lossy(),
16384                "destination_path": destination.to_string_lossy(),
16385                "dry_run": false
16386            }),
16387            "delete_path" => serde_json::json!({
16388                "path": preserved.to_string_lossy(),
16389                "recursive": false,
16390                "dry_run": false
16391            }),
16392            _ => panic!("unsupported mutation tool: {}", tool_id),
16393        };
16394        let security = mutation_denial_security_config(&tool_id, &workspace.root, denial);
16395        let builder = AgentBuilder::new()
16396            .system_prompt("Test mutation denial.")
16397            .llm(Arc::new(mock_with_response("done")))
16398            .tool(tool)
16399            .tool_security(ToolSecurityEngine::new(security));
16400        let builder = match denial {
16401            MutationDenial::Policy => builder,
16402            MutationDenial::Approval => builder
16403                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16404                .approval_handler(Arc::new(RejectAllHandler::new())),
16405        };
16406        let agent = builder.build().unwrap();
16407
16408        let record = agent
16409            .invoke_tool(ToolExecutionRequest::new(
16410                format!("{}-denied", tool_id),
16411                tool_id.clone(),
16412                arguments,
16413                ToolCallSource::Manual,
16414            ))
16415            .await
16416            .unwrap();
16417
16418        assert!(!record.executed, "{} must not be invoked", tool_id);
16419        assert!(!record.success);
16420        match denial {
16421            MutationDenial::Policy => {
16422                assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
16423                assert!(record.approval.as_ref().is_some_and(|approval| matches!(
16424                    &approval.status,
16425                    ToolApprovalStatus::NotRequired
16426                )));
16427            }
16428            MutationDenial::Approval => {
16429                assert_eq!(record.policy.outcome, PermissionOutcome::RequiresApproval);
16430                assert!(record.approval.as_ref().is_some_and(|approval| matches!(
16431                    &approval.status,
16432                    ToolApprovalStatus::Rejected
16433                )));
16434            }
16435        }
16436        assert_eq!(std::fs::read_to_string(&preserved).unwrap(), "preserved");
16437        assert!(!destination.exists());
16438    }
16439
16440    fn recovery_manager_with_fallbacks(
16441        fallbacks: impl IntoIterator<Item = (String, String)>,
16442    ) -> RecoveryManager {
16443        use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
16444
16445        let per_tool = fallbacks
16446            .into_iter()
16447            .map(|(tool, fallback_tool)| {
16448                (
16449                    tool,
16450                    ToolRetryConfig {
16451                        max_retries: 0,
16452                        timeout_ms: Some(1_000),
16453                        on_failure: ToolFailureAction::Fallback { fallback_tool },
16454                    },
16455                )
16456            })
16457            .collect();
16458        RecoveryManager::new(ErrorRecoveryConfig {
16459            tools: ToolRecoveryConfig {
16460                per_tool,
16461                ..Default::default()
16462            },
16463            ..Default::default()
16464        })
16465    }
16466
16467    fn approval_check() -> HITLCheckResult {
16468        HITLCheckResult::required(
16469            ApprovalTrigger::tool("test", serde_json::json!({})),
16470            HashMap::new(),
16471            "Approve?",
16472            None,
16473        )
16474    }
16475
16476    fn agent_with_approval_result(
16477        raw_result: ApprovalResult,
16478        timeout_action: TimeoutAction,
16479        hooks: Arc<ApprovalRecordingHooks>,
16480    ) -> RuntimeAgent {
16481        use ai_agents_hitl::{CallbackHandler, HITLConfig};
16482
16483        let config = HITLConfig {
16484            on_timeout: timeout_action,
16485            ..Default::default()
16486        };
16487        let handler = CallbackHandler::new(move |_| raw_result.clone());
16488        AgentBuilder::new()
16489            .system_prompt("Test HITL hooks.")
16490            .llm(Arc::new(mock_with_response("done")))
16491            .build()
16492            .unwrap()
16493            .with_hooks(hooks)
16494            .with_hitl(HITLEngine::new(config), Arc::new(handler))
16495    }
16496
16497    #[tokio::test]
16498    async fn approval_hooks_expose_direct_effective_decisions_after_raw_results() {
16499        let cases = vec![
16500            (ApprovalResult::Approved, "approved"),
16501            (
16502                ApprovalResult::Rejected {
16503                    reason: Some("denied".to_string()),
16504                },
16505                "rejected",
16506            ),
16507            (
16508                ApprovalResult::Modified {
16509                    changes: HashMap::from([("value".to_string(), serde_json::json!(2))]),
16510                },
16511                "modified",
16512            ),
16513        ];
16514
16515        for (raw_result, expected) in cases {
16516            let hooks = Arc::new(ApprovalRecordingHooks::new());
16517            let agent =
16518                agent_with_approval_result(raw_result, TimeoutAction::Reject, hooks.clone());
16519
16520            let result = agent.request_hitl_approval(approval_check()).await.unwrap();
16521
16522            assert_eq!(approval_result_name(&result), expected);
16523            assert_correlated_approval_events(&hooks.events(), expected, expected);
16524        }
16525    }
16526
16527    #[tokio::test]
16528    async fn approval_hooks_expose_timeout_policy_decisions() {
16529        for (timeout_action, expected) in [
16530            (TimeoutAction::Approve, "approved"),
16531            (TimeoutAction::Reject, "rejected"),
16532        ] {
16533            let hooks = Arc::new(ApprovalRecordingHooks::new());
16534            let agent =
16535                agent_with_approval_result(ApprovalResult::Timeout, timeout_action, hooks.clone());
16536
16537            let result = agent.request_hitl_approval(approval_check()).await.unwrap();
16538
16539            assert_eq!(approval_result_name(&result), expected);
16540            assert_correlated_approval_events(&hooks.events(), "timeout", expected);
16541        }
16542    }
16543
16544    #[tokio::test]
16545    async fn timeout_error_fires_correlated_resolved_error_before_returning() {
16546        let hooks = Arc::new(ApprovalRecordingHooks::new());
16547        let agent = agent_with_approval_result(
16548            ApprovalResult::Timeout,
16549            TimeoutAction::Error,
16550            hooks.clone(),
16551        );
16552
16553        let error = agent
16554            .request_hitl_approval(approval_check())
16555            .await
16556            .unwrap_err();
16557
16558        assert!(error.to_string().contains("HITL approval timeout"));
16559        assert_correlated_approval_events(&hooks.events(), "timeout", "error");
16560    }
16561
16562    // Basic YAML → Build → Chat flow
16563    #[tokio::test]
16564    async fn test_integration_yaml_to_chat_basic() {
16565        let mock = mock_with_response("Hello! How can I help you?");
16566        let agent = AgentBuilder::new()
16567            .system_prompt("You are a test assistant.")
16568            .llm(Arc::new(mock))
16569            .build()
16570            .unwrap();
16571
16572        let response = agent.chat("Hi").await.unwrap();
16573        assert!(!response.content.is_empty());
16574        assert_eq!(response.content, "Hello! How can I help you?");
16575    }
16576
16577    #[tokio::test]
16578    async fn stream_events_emit_one_authoritative_final_without_legacy_done() {
16579        let agent = AgentBuilder::new()
16580            .system_prompt("You are a test assistant.")
16581            .llm(Arc::new(mock_with_response(
16582                "Hello from the final response.",
16583            )))
16584            .build()
16585            .unwrap();
16586
16587        let mut stream = agent.chat_stream_events("Hi").await.unwrap();
16588        let mut final_responses = Vec::new();
16589        let mut legacy_done = 0;
16590        while let Some(event) = stream.next().await {
16591            match event {
16592                AgentStreamEvent::Chunk(StreamChunk::Done {}) => legacy_done += 1,
16593                AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
16594                    panic!("unexpected stream error: {message}")
16595                }
16596                AgentStreamEvent::Final(response) => final_responses.push(response),
16597                AgentStreamEvent::Chunk(_) => {}
16598            }
16599        }
16600
16601        assert_eq!(legacy_done, 0);
16602        assert_eq!(final_responses.len(), 1);
16603        let response = final_responses.pop().unwrap();
16604        assert_eq!(response.content, "Hello from the final response.");
16605        assert!(
16606            response
16607                .metadata
16608                .as_ref()
16609                .is_some_and(|metadata| { metadata.contains_key("reasoning") })
16610        );
16611    }
16612
16613    #[tokio::test]
16614    async fn stream_final_content_includes_output_processing_after_provisional_chunks() {
16615        let yaml = r#"
16616name: ProcessedStreamAgent
16617system_prompt: "Answer directly."
16618process:
16619  output:
16620    - type: format
16621      config:
16622        template: "{{ response }} [finalized]"
16623streaming:
16624  enabled: true
16625"#;
16626        let agent = AgentBuilder::from_yaml(yaml)
16627            .unwrap()
16628            .llm(Arc::new(mock_with_response("provisional answer")))
16629            .auto_configure_features()
16630            .unwrap()
16631            .build()
16632            .unwrap();
16633
16634        let mut stream = agent.chat_stream_events("Hi").await.unwrap();
16635        let mut provisional = String::new();
16636        let mut final_content = None;
16637        while let Some(event) = stream.next().await {
16638            match event {
16639                AgentStreamEvent::Chunk(StreamChunk::Content { text }) => {
16640                    provisional.push_str(&text)
16641                }
16642                AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
16643                    panic!("unexpected stream error: {message}")
16644                }
16645                AgentStreamEvent::Final(response) => final_content = Some(response.content),
16646                AgentStreamEvent::Chunk(_) => {}
16647            }
16648        }
16649
16650        assert_eq!(provisional, "provisional answer");
16651        assert_eq!(
16652            final_content.as_deref(),
16653            Some("provisional answer [finalized]")
16654        );
16655    }
16656
16657    #[tokio::test]
16658    async fn stream_events_preserve_tool_progress_and_final_tool_calls() {
16659        let agent = AgentBuilder::new()
16660            .system_prompt("Use the echo tool once, then answer.")
16661            .llm(Arc::new(mock_with_responses(vec![
16662                r#"{"tool":"echo","arguments":{"message":"hello"}}"#,
16663                "Echo completed.",
16664            ])))
16665            .tool(Arc::new(ai_agents_tools::EchoTool::new()))
16666            .build()
16667            .unwrap();
16668
16669        let mut stream = agent.chat_stream_events("echo hello").await.unwrap();
16670        let mut starts = 0;
16671        let mut results = 0;
16672        let mut ends = 0;
16673        let mut final_response = None;
16674        while let Some(event) = stream.next().await {
16675            match event {
16676                AgentStreamEvent::Chunk(StreamChunk::ToolCallStart { name, .. }) => {
16677                    assert_eq!(name, "echo");
16678                    starts += 1;
16679                }
16680                AgentStreamEvent::Chunk(StreamChunk::ToolResult { name, success, .. }) => {
16681                    assert_eq!(name, "echo");
16682                    assert!(success);
16683                    results += 1;
16684                }
16685                AgentStreamEvent::Chunk(StreamChunk::ToolCallEnd { .. }) => ends += 1,
16686                AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
16687                    panic!("unexpected stream error: {message}")
16688                }
16689                AgentStreamEvent::Final(response) => final_response = Some(response),
16690                AgentStreamEvent::Chunk(_) => {}
16691            }
16692        }
16693
16694        assert_eq!((starts, results, ends), (1, 1, 1));
16695        let response = final_response.expect("tool stream must finalize");
16696        assert_eq!(response.content, "Echo completed.");
16697        assert_eq!(
16698            response.tool_calls.as_ref().map(|calls| calls
16699                .iter()
16700                .map(|call| call.name.as_str())
16701                .collect::<Vec<_>>()),
16702            Some(vec!["echo"])
16703        );
16704    }
16705
16706    #[tokio::test]
16707    async fn legacy_stream_still_emits_one_done_chunk() {
16708        let agent = AgentBuilder::new()
16709            .system_prompt("You are a test assistant.")
16710            .llm(Arc::new(mock_with_response(
16711                "Hello from the legacy stream.",
16712            )))
16713            .build()
16714            .unwrap();
16715
16716        let mut stream = agent.chat_stream("Hi").await.unwrap();
16717        let mut done = 0;
16718        while let Some(chunk) = stream.next().await {
16719            match chunk {
16720                StreamChunk::Done {} => done += 1,
16721                StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
16722                _ => {}
16723            }
16724        }
16725
16726        assert_eq!(done, 1);
16727    }
16728
16729    // Multi-turn conversation
16730    #[tokio::test]
16731    async fn test_integration_multi_turn_conversation() {
16732        let mock = mock_with_responses(vec![
16733            "Hello! I'm your assistant.",
16734            "The weather is sunny today.",
16735            "Goodbye!",
16736        ]);
16737        let agent = AgentBuilder::new()
16738            .system_prompt("You are helpful.")
16739            .llm(Arc::new(mock))
16740            .build()
16741            .unwrap();
16742
16743        let r1 = agent.chat("Hi").await.unwrap();
16744        assert_eq!(r1.content, "Hello! I'm your assistant.");
16745
16746        let r2 = agent.chat("What's the weather?").await.unwrap();
16747        assert_eq!(r2.content, "The weather is sunny today.");
16748
16749        let r3 = agent.chat("Bye").await.unwrap();
16750        assert_eq!(r3.content, "Goodbye!");
16751
16752        // Verify memory accumulated messages
16753        let messages = agent.memory.get_messages(None).await.unwrap();
16754        // 3 user + 3 assistant = 6 messages
16755        assert_eq!(messages.len(), 6);
16756    }
16757
16758    #[test]
16759    fn later_approval_preserves_modified_evidence() {
16760        let arguments = serde_json::json!({"dry_run": true});
16761        let mut record = Some(ToolApprovalRecord {
16762            status: ToolApprovalStatus::Modified,
16763            reason: None,
16764            modified_arguments: Some(arguments.clone()),
16765        });
16766
16767        merge_approved_record(&mut record);
16768
16769        let record = record.unwrap();
16770        assert!(matches!(record.status, ToolApprovalStatus::Modified));
16771        assert_eq!(record.modified_arguments, Some(arguments));
16772    }
16773
16774    #[test]
16775    fn approval_binding_rejects_replaced_tool_implementation() {
16776        let reviewed_tool: Arc<dyn ai_agents_core::Tool> = Arc::new(ContextEchoTool);
16777        let same_tool = Arc::clone(&reviewed_tool);
16778        let replacement_tool: Arc<dyn ai_agents_core::Tool> = Arc::new(ContextEchoTool);
16779        let arguments = serde_json::json!({"path": "."});
16780        let versions = ToolDecisionVersions {
16781            policy: 2,
16782            registry: 3,
16783            runtime_control: 4,
16784            state: Some(5),
16785        };
16786        let binding = ToolApprovalBinding {
16787            canonical_id: "context_echo".to_string(),
16788            arguments: arguments.clone(),
16789            confirmation_required: true,
16790            policy_version: versions.policy,
16791            runtime_control_version: versions.runtime_control,
16792            state_generation: versions.state,
16793            reviewed_tool,
16794        };
16795
16796        assert!(!binding.is_stale("context_echo", &arguments, true, versions, &same_tool,));
16797        assert!(binding.is_stale(
16798            "context_echo",
16799            &arguments,
16800            true,
16801            versions,
16802            &replacement_tool,
16803        ));
16804    }
16805
16806    #[tokio::test]
16807    async fn approved_mutation_to_dry_run_remains_executable() {
16808        use ai_agents_hitl::CallbackHandler;
16809
16810        let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
16811            changes: HashMap::from([("dry_run".to_string(), serde_json::json!(true))]),
16812        });
16813        let agent = AgentBuilder::new()
16814            .system_prompt("Test safer approval modifications.")
16815            .llm(Arc::new(mock_with_response("done")))
16816            .tool(Arc::new(ai_agents_tools::FileWriteTool::new()))
16817            .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16818            .approval_handler(Arc::new(handler))
16819            .build()
16820            .unwrap();
16821
16822        let record = agent
16823            .invoke_tool(ToolExecutionRequest::new(
16824                "approved-dry-run",
16825                "file_write",
16826                serde_json::json!({
16827                    "path": "./approval-dry-run.txt",
16828                    "content": "not written"
16829                }),
16830                ToolCallSource::Manual,
16831            ))
16832            .await
16833            .unwrap();
16834
16835        assert!(record.executed);
16836        assert!(record.success);
16837        assert_eq!(record.executed_arguments["dry_run"], true);
16838        assert!(matches!(
16839            record.approval.as_ref().map(|approval| &approval.status),
16840            Some(ToolApprovalStatus::Modified)
16841        ));
16842        let output: Value = serde_json::from_str(&record.output).unwrap();
16843        assert_eq!(output["mutation_performed"], false);
16844    }
16845
16846    /// Proves shared HITL approval evidence reaches the web fetch implementation unchanged.
16847    #[tokio::test]
16848    async fn shared_executor_approval_reaches_web_fetch_transport() {
16849        use ai_agents_hitl::{CallbackHandler, HITLConfig};
16850        use ai_agents_tools::{DomainPolicyConfig, ToolPolicyConfig};
16851
16852        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16853        let tool = WebFetchTool::with_transport_and_resolver(
16854            Arc::new(RuntimeWebFetchTransport {
16855                calls: Arc::clone(&calls),
16856            }),
16857            Arc::new(RuntimeWebFetchResolver),
16858        );
16859        let mut security = ToolSecurityConfig {
16860            enabled: true,
16861            fail_closed: true,
16862            ..Default::default()
16863        };
16864        security.tools.insert(
16865            "web_fetch".to_string(),
16866            ToolPolicyConfig {
16867                domains: DomainPolicyConfig {
16868                    requires_approval: vec!["approval.test".to_string()],
16869                    ..Default::default()
16870                },
16871                allowed_schemes: vec!["https".to_string()],
16872                allowed_ports: vec![443],
16873                ..Default::default()
16874            },
16875        );
16876        let handler = CallbackHandler::new(|_| ApprovalResult::Approved);
16877        let agent = AgentBuilder::new()
16878            .system_prompt("Test approved web fetch execution.")
16879            .llm(Arc::new(mock_with_response("done")))
16880            .tool(Arc::new(tool))
16881            .tool_security(ToolSecurityEngine::new(security))
16882            .build()
16883            .unwrap()
16884            .with_hitl(HITLEngine::new(HITLConfig::default()), Arc::new(handler));
16885
16886        let record = agent
16887            .invoke_tool(ToolExecutionRequest::new(
16888                "approved-web-fetch",
16889                "web_fetch",
16890                serde_json::json!({
16891                    "url": "https://approval.test/page",
16892                    "cache_ttl_seconds": 0
16893                }),
16894                ToolCallSource::Manual,
16895            ))
16896            .await
16897            .unwrap();
16898
16899        assert!(record.success);
16900        assert!(
16901            record
16902                .approval
16903                .as_ref()
16904                .is_some_and(|approval| matches!(approval.status, ToolApprovalStatus::Approved))
16905        );
16906        assert_eq!(calls.load(Ordering::SeqCst), 1);
16907    }
16908
16909    #[tokio::test]
16910    async fn context_preserves_requested_and_canonical_identity() {
16911        let mock = mock_with_response("hello");
16912        let mut tools = ai_agents_tools::ToolRegistry::new();
16913        tools.register(Arc::new(ContextEchoTool)).unwrap();
16914
16915        let mut security = ToolSecurityConfig {
16916            enabled: true,
16917            fail_closed: true,
16918            ..Default::default()
16919        };
16920        let mut policy = ai_agents_tools::ToolPolicyConfig {
16921            read_paths: vec![".".to_string()],
16922            max_results: Some(7),
16923            ..Default::default()
16924        };
16925        policy
16926            .config
16927            .insert("backend".to_string(), serde_json::json!("memory"));
16928        security.tools.insert("context_echo".to_string(), policy);
16929
16930        let agent = AgentBuilder::new()
16931            .system_prompt("You are helpful.")
16932            .llm(Arc::new(mock))
16933            .tools(tools)
16934            .tool_security(ToolSecurityEngine::new(security))
16935            .build()
16936            .unwrap();
16937
16938        let record = agent
16939            .invoke_tool(ToolExecutionRequest::new(
16940                "ctx-call",
16941                "Context Echo",
16942                serde_json::json!({"path": ".", "max_results": 99}),
16943                ToolCallSource::Manual,
16944            ))
16945            .await
16946            .unwrap();
16947
16948        assert!(record.success);
16949        assert!(matches!(&record.source, ToolCallSource::Manual));
16950        assert_eq!(record.requested_name, "Context Echo");
16951        assert_eq!(record.canonical_id, "context_echo");
16952        assert_eq!(record.policy.outcome, PermissionOutcome::Allow);
16953        assert_eq!(record.executed_arguments["max_results"], 7);
16954        let output: Value = serde_json::from_str(&record.output).unwrap();
16955        assert_eq!(output["requested_name"], "Context Echo");
16956        assert_eq!(output["canonical_id"], "context_echo");
16957        assert_eq!(output["max_results"], 7);
16958        assert_eq!(output["custom_config"]["backend"], "memory");
16959        assert!(record.metadata.contains_key("effective_limits"));
16960        assert!(record.metadata.contains_key("policy_snapshot"));
16961    }
16962
16963    #[tokio::test]
16964    async fn test_runtime_control_cancels_active_tool_call() {
16965        let mock = mock_with_response("hello");
16966        let agent = Arc::new(
16967            AgentBuilder::new()
16968                .system_prompt("You are helpful.")
16969                .llm(Arc::new(mock))
16970                .tool(Arc::new(SlowTool))
16971                .build()
16972                .unwrap(),
16973        );
16974        let control = agent.runtime_control();
16975        let running_agent = Arc::clone(&agent);
16976        let handle = tokio::spawn(async move {
16977            running_agent
16978                .invoke_tool(ToolExecutionRequest::new(
16979                    "slow-call",
16980                    "slow",
16981                    serde_json::json!({}),
16982                    ToolCallSource::Manual,
16983                ))
16984                .await
16985                .unwrap()
16986        });
16987
16988        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
16989        control.cancel_all();
16990        let record = handle.await.unwrap();
16991
16992        assert!(record.executed);
16993        assert!(record.cancelled);
16994        assert!(!record.success);
16995        assert!(record.cancellation_reason.is_some());
16996    }
16997
16998    // Verifies runtime cancellation returns the original terminal record without starting configured fallback work.
16999    #[tokio::test]
17000    async fn cancelled_tool_does_not_enter_fallback() {
17001        let fallback_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17002        let agent = Arc::new(
17003            AgentBuilder::new()
17004                .system_prompt("Test cancellation before fallback.")
17005                .llm(Arc::new(mock_with_response("done")))
17006                .tool(Arc::new(SlowTool))
17007                .tool(Arc::new(RecoveryTestTool {
17008                    id: "fallback".to_string(),
17009                    succeeds: true,
17010                    calls: Arc::clone(&fallback_calls),
17011                    max_output_chars: None,
17012                }))
17013                .recovery_manager(recovery_manager_with_fallbacks([(
17014                    "slow".to_string(),
17015                    "fallback".to_string(),
17016                )]))
17017                .build()
17018                .unwrap(),
17019        );
17020        let control = agent.runtime_control();
17021        let running_agent = Arc::clone(&agent);
17022        let handle = tokio::spawn(async move {
17023            running_agent
17024                .invoke_tool(ToolExecutionRequest::new(
17025                    "cancelled-fallback-call",
17026                    "slow",
17027                    serde_json::json!({}),
17028                    ToolCallSource::Manual,
17029                ))
17030                .await
17031                .unwrap()
17032        });
17033
17034        tokio::time::sleep(Duration::from_millis(100)).await;
17035        control.cancel_all();
17036        let record = handle.await.unwrap();
17037
17038        assert!(record.executed);
17039        assert!(record.cancelled);
17040        assert!(!record.success);
17041        assert_eq!(record.canonical_id, "slow");
17042        assert_eq!(fallback_calls.load(Ordering::SeqCst), 0);
17043        assert_eq!(agent.tool_call_history().len(), 1);
17044    }
17045
17046    #[tokio::test]
17047    async fn non_idempotent_tool_calls_are_not_retried() {
17048        use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
17049
17050        let mock = mock_with_response("hello");
17051        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17052        let agent = AgentBuilder::new()
17053            .system_prompt("You are helpful.")
17054            .llm(Arc::new(mock))
17055            .tool(Arc::new(FlakyWriteTool {
17056                calls: Arc::clone(&calls),
17057            }))
17058            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
17059                tools: ToolRecoveryConfig {
17060                    default: ToolRetryConfig {
17061                        max_retries: 2,
17062                        ..Default::default()
17063                    },
17064                    ..Default::default()
17065                },
17066                ..Default::default()
17067            }))
17068            .build()
17069            .unwrap();
17070
17071        let record = agent
17072            .invoke_tool(ToolExecutionRequest::new(
17073                "flaky-call",
17074                "flaky_write",
17075                serde_json::json!({"path": "./tmp.txt"}),
17076                ToolCallSource::Manual,
17077            ))
17078            .await
17079            .unwrap();
17080
17081        assert!(!record.success);
17082        assert_eq!(calls.load(Ordering::SeqCst), 1);
17083    }
17084
17085    #[tokio::test]
17086    async fn safely_retryable_tool_receives_a_fresh_deadline_per_attempt() {
17087        use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
17088
17089        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17090        let deadlines = Arc::new(parking_lot::Mutex::new(Vec::new()));
17091        let remaining_ms = Arc::new(parking_lot::Mutex::new(Vec::new()));
17092        let agent = AgentBuilder::new()
17093            .system_prompt("Test retry deadlines.")
17094            .llm(Arc::new(mock_with_response("done")))
17095            .tool(Arc::new(RetryDeadlineTool {
17096                calls: Arc::clone(&calls),
17097                deadlines: Arc::clone(&deadlines),
17098                remaining_ms: Arc::clone(&remaining_ms),
17099            }))
17100            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
17101                tools: ToolRecoveryConfig {
17102                    per_tool: HashMap::from([(
17103                        "retry_deadline".to_string(),
17104                        ToolRetryConfig {
17105                            max_retries: 1,
17106                            ..Default::default()
17107                        },
17108                    )]),
17109                    ..Default::default()
17110                },
17111                ..Default::default()
17112            }))
17113            .build()
17114            .unwrap();
17115
17116        let record = agent
17117            .invoke_tool(ToolExecutionRequest::new(
17118                "retry-deadline-call",
17119                "retry_deadline",
17120                serde_json::json!({}),
17121                ToolCallSource::Manual,
17122            ))
17123            .await
17124            .unwrap();
17125
17126        assert!(record.executed);
17127        assert!(record.success);
17128        assert_eq!(calls.load(Ordering::SeqCst), 2);
17129        let deadlines = deadlines.lock();
17130        assert_eq!(deadlines.len(), 2);
17131        assert!(
17132            deadlines[1] > deadlines[0],
17133            "retry inherited the first invocation deadline"
17134        );
17135        let remaining_ms = remaining_ms.lock();
17136        assert_eq!(remaining_ms.len(), 2);
17137        assert!(
17138            remaining_ms
17139                .iter()
17140                .all(|remaining| (800..=1_000).contains(remaining))
17141        );
17142    }
17143
17144    // Verifies a call-level cap drives both the visible deadline and the executor timer.
17145    #[tokio::test]
17146    async fn call_classification_timeout_controls_deadline_and_timer() {
17147        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17148        let remaining_ms = Arc::new(parking_lot::Mutex::new(Vec::new()));
17149        let agent = AgentBuilder::new()
17150            .system_prompt("Test call-level timeout.")
17151            .llm(Arc::new(mock_with_response("done")))
17152            .tool(Arc::new(ClassifiedTimeoutTool {
17153                id: "classified_timeout",
17154                calls: Arc::clone(&calls),
17155                timeout_ms: 100,
17156                sleep_ms: 150,
17157                requires_approval: false,
17158                remaining_ms: Arc::clone(&remaining_ms),
17159            }))
17160            .build()
17161            .unwrap();
17162
17163        let started = Instant::now();
17164        let record = agent
17165            .invoke_tool(ToolExecutionRequest::new(
17166                "classified-timeout-call",
17167                "classified_timeout",
17168                serde_json::json!({}),
17169                ToolCallSource::Manual,
17170            ))
17171            .await
17172            .unwrap();
17173
17174        assert!(record.executed);
17175        assert!(record.timed_out);
17176        assert!(!record.success);
17177        assert_eq!(calls.load(Ordering::SeqCst), 1);
17178        assert!(started.elapsed() < Duration::from_secs(1));
17179        let remaining_ms = remaining_ms.lock();
17180        assert_eq!(remaining_ms.len(), 1);
17181        assert!((1..=100).contains(&remaining_ms[0]));
17182    }
17183
17184    // Verifies a per-tool recovery cap lowers broader call and security timeout values.
17185    #[tokio::test]
17186    async fn recovery_timeout_only_lowers_call_and_policy_timeouts() {
17187        use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
17188
17189        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17190        let remaining_ms = Arc::new(parking_lot::Mutex::new(Vec::new()));
17191        let agent = AgentBuilder::new()
17192            .system_prompt("Test recovery timeout.")
17193            .llm(Arc::new(mock_with_response("done")))
17194            .tool(Arc::new(ClassifiedTimeoutTool {
17195                id: "recovery_timeout",
17196                calls: Arc::clone(&calls),
17197                timeout_ms: 1_000,
17198                sleep_ms: 150,
17199                requires_approval: false,
17200                remaining_ms: Arc::clone(&remaining_ms),
17201            }))
17202            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
17203                tools: ToolRecoveryConfig {
17204                    per_tool: HashMap::from([(
17205                        "recovery_timeout".to_string(),
17206                        ToolRetryConfig {
17207                            timeout_ms: Some(100),
17208                            ..Default::default()
17209                        },
17210                    )]),
17211                    ..Default::default()
17212                },
17213                ..Default::default()
17214            }))
17215            .build()
17216            .unwrap();
17217
17218        let started = Instant::now();
17219        let record = agent
17220            .invoke_tool(ToolExecutionRequest::new(
17221                "recovery-timeout-call",
17222                "recovery_timeout",
17223                serde_json::json!({}),
17224                ToolCallSource::Manual,
17225            ))
17226            .await
17227            .unwrap();
17228
17229        assert!(record.executed);
17230        assert!(record.timed_out);
17231        assert!(!record.success);
17232        assert_eq!(calls.load(Ordering::SeqCst), 1);
17233        assert!(started.elapsed() < Duration::from_secs(1));
17234        assert_eq!(record.metadata["effective_limits"]["timeout_ms"], 100);
17235        let remaining_ms = remaining_ms.lock();
17236        assert_eq!(remaining_ms.len(), 1);
17237        assert!((1..=100).contains(&remaining_ms[0]));
17238    }
17239
17240    // Verifies the complete default recovery policy supplies the timeout when no per-tool policy exists.
17241    #[tokio::test]
17242    async fn recovery_default_timeout_controls_deadline_and_timer() {
17243        use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
17244
17245        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17246        let remaining_ms = Arc::new(parking_lot::Mutex::new(Vec::new()));
17247        let agent = AgentBuilder::new()
17248            .system_prompt("Test default recovery timeout.")
17249            .llm(Arc::new(mock_with_response("done")))
17250            .tool(Arc::new(ClassifiedTimeoutTool {
17251                id: "default_recovery_timeout",
17252                calls: Arc::clone(&calls),
17253                timeout_ms: 1_000,
17254                sleep_ms: 150,
17255                requires_approval: false,
17256                remaining_ms: Arc::clone(&remaining_ms),
17257            }))
17258            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
17259                tools: ToolRecoveryConfig {
17260                    default: ToolRetryConfig {
17261                        timeout_ms: Some(100),
17262                        ..Default::default()
17263                    },
17264                    ..Default::default()
17265                },
17266                ..Default::default()
17267            }))
17268            .build()
17269            .unwrap();
17270
17271        let started = Instant::now();
17272        let record = agent
17273            .invoke_tool(ToolExecutionRequest::new(
17274                "default-recovery-timeout-call",
17275                "default_recovery_timeout",
17276                serde_json::json!({}),
17277                ToolCallSource::Manual,
17278            ))
17279            .await
17280            .unwrap();
17281
17282        assert!(record.executed);
17283        assert!(record.timed_out);
17284        assert!(!record.success);
17285        assert_eq!(calls.load(Ordering::SeqCst), 1);
17286        assert!(started.elapsed() < Duration::from_secs(1));
17287        assert_eq!(record.metadata["effective_limits"]["timeout_ms"], 100);
17288        let remaining_ms = remaining_ms.lock();
17289        assert_eq!(remaining_ms.len(), 1);
17290        assert!((1..=100).contains(&remaining_ms[0]));
17291    }
17292
17293    // Verifies classification and recovery values cannot widen the security baseline.
17294    #[test]
17295    fn recovery_timeout_cannot_widen_security_baseline() {
17296        let security_engine = ToolSecurityEngine::new(ToolSecurityConfig {
17297            default_timeout_ms: 100,
17298            ..Default::default()
17299        });
17300        let safety = ToolSafetyMetadata::compute();
17301        let mut classification = ToolCallClassification::from_metadata(&safety);
17302        classification.timeout_ms = Some(500);
17303
17304        let (limits, timeout) = RuntimeAgent::effective_tool_limits(
17305            &security_engine,
17306            "recovery_cannot_widen",
17307            &safety,
17308            &classification,
17309            Some(1_000),
17310        )
17311        .unwrap();
17312
17313        assert_eq!(limits.timeout_ms, Some(100));
17314        assert_eq!(timeout.timer, Duration::from_millis(100));
17315    }
17316
17317    // Verifies an invalid call-level cap fails before approval or invocation side effects.
17318    #[tokio::test]
17319    async fn invalid_call_timeout_stops_before_approval_or_tool_invocation() {
17320        let tool_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17321        let approval_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17322        let remaining_ms = Arc::new(parking_lot::Mutex::new(Vec::new()));
17323        let mut security = ToolSecurityConfig {
17324            enabled: true,
17325            ..Default::default()
17326        };
17327        security.tools.insert(
17328            "invalid_call_timeout".to_string(),
17329            ai_agents_tools::ToolPolicyConfig {
17330                require_confirmation: true,
17331                ..Default::default()
17332            },
17333        );
17334        let agent = AgentBuilder::new()
17335            .system_prompt("Test invalid call timeout.")
17336            .llm(Arc::new(mock_with_response("done")))
17337            .tool(Arc::new(ClassifiedTimeoutTool {
17338                id: "invalid_call_timeout",
17339                calls: Arc::clone(&tool_calls),
17340                timeout_ms: u64::MAX,
17341                sleep_ms: 0,
17342                requires_approval: false,
17343                remaining_ms,
17344            }))
17345            .tool_security(ToolSecurityEngine::new(security))
17346            .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17347            .approval_handler(Arc::new(CountingApprovalHandler {
17348                calls: Arc::clone(&approval_calls),
17349            }))
17350            .build()
17351            .unwrap();
17352
17353        let error = agent
17354            .invoke_tool(ToolExecutionRequest::new(
17355                "invalid-call-timeout",
17356                "invalid_call_timeout",
17357                serde_json::json!({}),
17358                ToolCallSource::Manual,
17359            ))
17360            .await
17361            .unwrap_err();
17362
17363        assert!(error.to_string().contains(
17364            "effective tool timeout_ms must be no greater than 3153600000000000 milliseconds"
17365        ));
17366        assert_eq!(approval_calls.load(Ordering::SeqCst), 0);
17367        assert_eq!(tool_calls.load(Ordering::SeqCst), 0);
17368    }
17369
17370    // Verifies approval-modified arguments are reclassified before resource-lock waiting or implementation invocation.
17371    #[tokio::test]
17372    async fn invalid_modified_call_timeout_stops_before_lock_or_invocation() {
17373        use ai_agents_hitl::CallbackHandler;
17374
17375        let blocker_gate = PathMutationGate::new();
17376        let tool_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17377        let hooks = Arc::new(ToolLifecycleRecordingHooks::new());
17378        let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
17379            changes: HashMap::from([("invalid_timeout".to_string(), Value::Bool(true))]),
17380        });
17381        let agent = Arc::new(
17382            AgentBuilder::new()
17383                .system_prompt("Test final call timeout validation.")
17384                .llm(Arc::new(mock_with_response("done")))
17385                .tool(Arc::new(BlockingPathMutationTool {
17386                    id: "timeout_lock_blocker",
17387                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
17388                    gate: blocker_gate.clone(),
17389                }))
17390                .tool(Arc::new(ApprovalModifiedTimeoutTool {
17391                    calls: Arc::clone(&tool_calls),
17392                }))
17393                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17394                .approval_handler(Arc::new(handler))
17395                .hooks(hooks.clone())
17396                .build()
17397                .unwrap(),
17398        );
17399        let blocking_agent = Arc::clone(&agent);
17400        let blocker = tokio::spawn(async move {
17401            blocking_agent
17402                .invoke_tool(ToolExecutionRequest::new(
17403                    "timeout-lock-blocker",
17404                    "timeout_lock_blocker",
17405                    serde_json::json!({"path": "./shared-timeout.txt"}),
17406                    ToolCallSource::Manual,
17407                ))
17408                .await
17409                .unwrap()
17410        });
17411        blocker_gate.wait_until_entered().await;
17412
17413        let record = tokio::time::timeout(
17414            Duration::from_millis(500),
17415            agent.invoke_tool(ToolExecutionRequest::new(
17416                "invalid-modified-timeout",
17417                "approval_modified_timeout",
17418                serde_json::json!({
17419                    "path": "./shared-timeout.txt",
17420                    "invalid_timeout": false
17421                }),
17422                ToolCallSource::Manual,
17423            )),
17424        )
17425        .await
17426        .expect("final timeout validation must not wait for the held path lock")
17427        .unwrap();
17428
17429        blocker_gate.release();
17430        assert!(blocker.await.unwrap().success);
17431        assert!(!record.executed);
17432        assert!(!record.success);
17433        assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
17434        assert!(record.output.contains(
17435            "effective tool timeout_ms must be no greater than 3153600000000000 milliseconds"
17436        ));
17437        assert_eq!(tool_calls.load(Ordering::SeqCst), 0);
17438        let invalid_request_events = hooks
17439            .events()
17440            .into_iter()
17441            .filter(|event| event.contains("approval_modified_timeout") || event == "error")
17442            .collect::<Vec<_>>();
17443        assert_eq!(
17444            invalid_request_events,
17445            vec![
17446                "start:approval_modified_timeout",
17447                "complete:approval_modified_timeout:false",
17448                "record:approval_modified_timeout:false",
17449                "error"
17450            ]
17451        );
17452    }
17453
17454    #[tokio::test]
17455    async fn side_effecting_tools_are_serialized_per_resource() {
17456        let mock = mock_with_response("hello");
17457        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17458        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17459        let agent = Arc::new(
17460            AgentBuilder::new()
17461                .system_prompt("You are helpful.")
17462                .llm(Arc::new(mock))
17463                .tool(Arc::new(LockedWriteTool {
17464                    active: Arc::clone(&active),
17465                    max_active: Arc::clone(&max_active),
17466                }))
17467                .build()
17468                .unwrap(),
17469        );
17470
17471        let left = {
17472            let agent = Arc::clone(&agent);
17473            tokio::spawn(async move {
17474                agent
17475                    .invoke_tool(ToolExecutionRequest::new(
17476                        "lock-1",
17477                        "locked_write",
17478                        serde_json::json!({"path": "./same.txt"}),
17479                        ToolCallSource::Manual,
17480                    ))
17481                    .await
17482                    .unwrap()
17483            })
17484        };
17485        let right = {
17486            let agent = Arc::clone(&agent);
17487            tokio::spawn(async move {
17488                agent
17489                    .invoke_tool(ToolExecutionRequest::new(
17490                        "lock-2",
17491                        "locked_write",
17492                        serde_json::json!({"path": "./same.txt"}),
17493                        ToolCallSource::Manual,
17494                    ))
17495                    .await
17496                    .unwrap()
17497            })
17498        };
17499
17500        let left = left.await.unwrap();
17501        let right = right.await.unwrap();
17502        assert!(left.success);
17503        assert!(right.success);
17504        assert_eq!(max_active.load(Ordering::SeqCst), 1);
17505    }
17506
17507    #[tokio::test]
17508    async fn path_resources_use_shared_global_lock_and_cleanup() {
17509        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17510        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17511        let bindings = ai_agents_core::ToolPolicyBindings {
17512            path_fields: vec![
17513                ai_agents_core::PathPolicyBinding::read_write("source_path"),
17514                ai_agents_core::PathPolicyBinding::write("destination_path"),
17515            ],
17516            ..Default::default()
17517        };
17518        let classification = ai_agents_core::ToolCallClassification::from_metadata(
17519            &MultiResourceWriteTool {
17520                active: Arc::clone(&active),
17521                max_active: Arc::clone(&max_active),
17522            }
17523            .safety_metadata(),
17524        );
17525        let left_args = serde_json::json!({
17526            "source_path": "./a/../first.txt",
17527            "destination_path": "./second.txt"
17528        });
17529        let right_args = serde_json::json!({
17530            "source_path": "./second.txt",
17531            "destination_path": "./first.txt"
17532        });
17533        let left_keys = tool_resource_lock_keys(
17534            "multi_resource_write",
17535            &left_args,
17536            &bindings,
17537            &classification,
17538        );
17539        let right_keys = tool_resource_lock_keys(
17540            "multi_resource_write",
17541            &right_args,
17542            &bindings,
17543            &classification,
17544        );
17545        assert_eq!(left_keys, right_keys);
17546        assert_eq!(left_keys, vec!["path-mutation:global".to_string()]);
17547
17548        let locks = new_tool_resource_locks();
17549        let build_agent = || {
17550            AgentBuilder::new()
17551                .system_prompt("Test shared resource locks.")
17552                .llm(Arc::new(mock_with_response("done")))
17553                .tool(Arc::new(MultiResourceWriteTool {
17554                    active: Arc::clone(&active),
17555                    max_active: Arc::clone(&max_active),
17556                }))
17557                .build()
17558                .unwrap()
17559                .with_shared_resource_locks(Arc::clone(&locks))
17560        };
17561        let left_agent = Arc::new(build_agent());
17562        let right_agent = Arc::new(build_agent());
17563        let left = tokio::spawn(async move {
17564            left_agent
17565                .invoke_tool(ToolExecutionRequest::new(
17566                    "multi-left",
17567                    "multi_resource_write",
17568                    left_args,
17569                    ToolCallSource::Manual,
17570                ))
17571                .await
17572                .unwrap()
17573        });
17574        let right = tokio::spawn(async move {
17575            right_agent
17576                .invoke_tool(ToolExecutionRequest::new(
17577                    "multi-right",
17578                    "multi_resource_write",
17579                    right_args,
17580                    ToolCallSource::Manual,
17581                ))
17582                .await
17583                .unwrap()
17584        });
17585        let (left, right) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
17586            tokio::join!(left, right)
17587        })
17588        .await
17589        .expect("reversed resource acquisition must not deadlock");
17590
17591        assert!(left.unwrap().success);
17592        assert!(right.unwrap().success);
17593        assert_eq!(max_active.load(Ordering::SeqCst), 1);
17594        assert!(locks.read().is_empty());
17595    }
17596
17597    #[tokio::test]
17598    async fn global_path_lock_serializes_copy_destination_with_file_write() {
17599        assert_path_mutation_pair_serialized(
17600            "copy_path",
17601            CopyPathTool::new().policy_bindings().path_fields,
17602            serde_json::json!({
17603                "source_path": "./source.txt",
17604                "destination_path": "./shared.txt"
17605            }),
17606            "file_write",
17607            FileWriteTool::new().policy_bindings().path_fields,
17608            serde_json::json!({"path": "./shared.txt"}),
17609        )
17610        .await;
17611    }
17612
17613    #[tokio::test]
17614    async fn parent_and_spawned_runtime_share_global_path_lock() {
17615        let workspace = MutationTestWorkspace::new();
17616        let destination = workspace.root.join("spawned.txt");
17617        let parent_gate = PathMutationGate::new();
17618        let parent = Arc::new(
17619            AgentBuilder::from_yaml(
17620                r#"
17621name: LockParent
17622system_prompt: parent
17623llm:
17624  default: default
17625tools:
17626  - parent_path_write
17627spawner:
17628  shared_llms: true
17629"#,
17630            )
17631            .unwrap()
17632            .llm(Arc::new(mock_with_response("done")))
17633            .auto_configure_spawner()
17634            .await
17635            .unwrap()
17636            .tool(Arc::new(BlockingPathMutationTool {
17637                id: "parent_path_write",
17638                path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
17639                gate: parent_gate.clone(),
17640            }))
17641            .build()
17642            .unwrap(),
17643        );
17644
17645        let mut child_spec = crate::spec::AgentSpec {
17646            name: "LockChild".to_string(),
17647            system_prompt: "child".to_string(),
17648            tools: Some(vec![crate::spec::ToolEntry::Simple(
17649                "file_write".to_string(),
17650            )]),
17651            ..Default::default()
17652        };
17653        child_spec.tool_security.enabled = true;
17654        child_spec.tool_security.fail_closed = true;
17655        let file_write_policy = ai_agents_tools::ToolPolicyConfig {
17656            write_paths: vec![workspace.root.to_string_lossy().into_owned()],
17657            allow_without_confirmation: true,
17658            ..Default::default()
17659        };
17660        child_spec
17661            .tool_security
17662            .tools
17663            .insert("file_write".to_string(), file_write_policy);
17664        let spawned = parent
17665            .spawner()
17666            .unwrap()
17667            .spawn_from_spec(child_spec)
17668            .await
17669            .unwrap();
17670        assert!(Arc::ptr_eq(
17671            &parent.resource_locks,
17672            &spawned.agent.resource_locks
17673        ));
17674        assert!(!Arc::ptr_eq(
17675            &parent.runtime_control,
17676            &spawned.agent.runtime_control
17677        ));
17678
17679        let parent_call = {
17680            let parent = Arc::clone(&parent);
17681            let destination = destination.clone();
17682            tokio::spawn(async move {
17683                parent
17684                    .invoke_tool(ToolExecutionRequest::new(
17685                        "parent-lock-holder",
17686                        "parent_path_write",
17687                        serde_json::json!({"path": destination}),
17688                        ToolCallSource::Manual,
17689                    ))
17690                    .await
17691                    .unwrap()
17692            })
17693        };
17694        parent_gate.wait_until_entered().await;
17695
17696        let child_call = {
17697            let child = Arc::clone(&spawned.agent);
17698            let destination = destination.clone();
17699            tokio::spawn(async move {
17700                child
17701                    .invoke_tool(ToolExecutionRequest::new(
17702                        "spawned-file-write",
17703                        "file_write",
17704                        serde_json::json!({
17705                            "path": destination,
17706                            "content": "spawned",
17707                            "dry_run": false
17708                        }),
17709                        ToolCallSource::Manual,
17710                    ))
17711                    .await
17712                    .unwrap()
17713            })
17714        };
17715        wait_for_resource_lock_strong_count(&parent.resource_locks, 2).await;
17716        assert!(!child_call.is_finished());
17717
17718        parent_gate.release();
17719        let (parent_record, child_record) =
17720            tokio::time::timeout(std::time::Duration::from_secs(2), async {
17721                tokio::join!(parent_call, child_call)
17722            })
17723            .await
17724            .expect("parent and spawned path mutations did not finish");
17725        assert!(parent_record.unwrap().success);
17726        assert!(child_record.unwrap().success);
17727        assert_eq!(std::fs::read_to_string(destination).unwrap(), "spawned");
17728        assert!(parent.resource_locks.read().is_empty());
17729    }
17730
17731    #[tokio::test]
17732    async fn cancelled_global_path_lock_waiter_does_not_retain_weak_entry() {
17733        let locks = new_tool_resource_locks();
17734        let holder_gate = PathMutationGate::new();
17735        let waiter_gate = PathMutationGate::new();
17736        waiter_gate.release();
17737        let holder = Arc::new(
17738            AgentBuilder::new()
17739                .system_prompt("Hold the global path lock.")
17740                .llm(Arc::new(mock_with_response("done")))
17741                .tool(Arc::new(BlockingPathMutationTool {
17742                    id: "holder_write",
17743                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
17744                    gate: holder_gate.clone(),
17745                }))
17746                .build()
17747                .unwrap()
17748                .with_shared_resource_locks(Arc::clone(&locks)),
17749        );
17750        let waiter = Arc::new(
17751            AgentBuilder::new()
17752                .system_prompt("Wait for the global path lock.")
17753                .llm(Arc::new(mock_with_response("done")))
17754                .tool(Arc::new(BlockingPathMutationTool {
17755                    id: "waiter_write",
17756                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
17757                    gate: waiter_gate.clone(),
17758                }))
17759                .build()
17760                .unwrap()
17761                .with_shared_resource_locks(Arc::clone(&locks)),
17762        );
17763
17764        let holder_call = {
17765            let holder = Arc::clone(&holder);
17766            tokio::spawn(async move {
17767                holder
17768                    .invoke_tool(ToolExecutionRequest::new(
17769                        "holder-call",
17770                        "holder_write",
17771                        serde_json::json!({"path": "./shared.txt"}),
17772                        ToolCallSource::Manual,
17773                    ))
17774                    .await
17775                    .unwrap()
17776            })
17777        };
17778        holder_gate.wait_until_entered().await;
17779
17780        let waiter_call = {
17781            let waiter = Arc::clone(&waiter);
17782            tokio::spawn(async move {
17783                waiter
17784                    .invoke_tool(ToolExecutionRequest::new(
17785                        "waiter-call",
17786                        "waiter_write",
17787                        serde_json::json!({"path": "./shared.txt"}),
17788                        ToolCallSource::Manual,
17789                    ))
17790                    .await
17791                    .unwrap()
17792            })
17793        };
17794        wait_for_resource_lock_strong_count(&locks, 2).await;
17795        waiter.runtime_control().cancel_all();
17796
17797        let waiter_record = tokio::time::timeout(std::time::Duration::from_secs(2), waiter_call)
17798            .await
17799            .expect("cancelled lock waiter did not finish")
17800            .unwrap();
17801        assert!(!waiter_record.success);
17802        assert!(!waiter_record.executed);
17803        assert!(waiter_record.cancelled);
17804        assert_eq!(
17805            waiter_record.cancellation_reason.as_deref(),
17806            Some("runtime control cancellation")
17807        );
17808        assert!(!waiter_gate.entered.load(Ordering::SeqCst));
17809        assert_eq!(
17810            locks
17811                .read()
17812                .get("path-mutation:global")
17813                .map_or(0, |lock| lock.strong_count()),
17814            1
17815        );
17816
17817        holder_gate.release();
17818        let holder_record = tokio::time::timeout(std::time::Duration::from_secs(2), holder_call)
17819            .await
17820            .expect("lock holder did not finish")
17821            .unwrap();
17822        assert!(holder_record.success);
17823        assert!(locks.read().is_empty());
17824    }
17825
17826    #[tokio::test]
17827    async fn path_mutation_policy_and_approval_denials_do_not_invoke_tools() {
17828        for denial in [MutationDenial::Policy, MutationDenial::Approval] {
17829            let tools: [Arc<dyn Tool>; 3] = [
17830                Arc::new(CopyPathTool::new()),
17831                Arc::new(MovePathTool::new()),
17832                Arc::new(DeletePathTool::new()),
17833            ];
17834            for tool in tools {
17835                assert_path_mutation_denied(tool, denial).await;
17836            }
17837        }
17838    }
17839
17840    #[tokio::test]
17841    async fn policy_denial_keeps_executor_hook_lifecycle_and_record_authority() {
17842        let workspace = MutationTestWorkspace::new();
17843        let target = workspace.root.join("denied.txt");
17844        let hooks = Arc::new(ToolLifecycleRecordingHooks::new());
17845        let agent = AgentBuilder::new()
17846            .system_prompt("Test denied tool hooks.")
17847            .llm(Arc::new(mock_with_response("done")))
17848            .tool(Arc::new(FileWriteTool::new()))
17849            .tool_security(ToolSecurityEngine::new(mutation_denial_security_config(
17850                "file_write",
17851                &workspace.root,
17852                MutationDenial::Policy,
17853            )))
17854            .hooks(hooks.clone())
17855            .build()
17856            .unwrap();
17857
17858        let record = agent
17859            .invoke_tool(ToolExecutionRequest::new(
17860                "denied-hook-call",
17861                "file_write",
17862                serde_json::json!({
17863                    "path": target.to_string_lossy(),
17864                    "content": "blocked"
17865                }),
17866                ToolCallSource::Manual,
17867            ))
17868            .await
17869            .unwrap();
17870
17871        assert!(!record.executed);
17872        assert!(!record.success);
17873        assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
17874        assert_eq!(
17875            hooks.events(),
17876            vec![
17877                "start:file_write",
17878                "complete:file_write:false",
17879                "record:file_write:false",
17880                "error"
17881            ]
17882        );
17883        assert!(!target.exists());
17884    }
17885
17886    #[tokio::test]
17887    async fn approval_argument_changes_are_rechecked_against_final_scope() {
17888        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17889        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17890        let entered = Arc::new(tokio::sync::Barrier::new(2));
17891        let release = Arc::new(tokio::sync::Notify::new());
17892        let handler = Arc::new(BlockingApprovalHandler {
17893            entered: Arc::clone(&entered),
17894            release: Arc::clone(&release),
17895            result: ApprovalResult::Modified {
17896                changes: HashMap::from([(
17897                    "path".to_string(),
17898                    Value::String("./after-approval.txt".to_string()),
17899                )]),
17900            },
17901        });
17902        let agent = Arc::new(
17903            AgentBuilder::new()
17904                .system_prompt("Test final scope validation.")
17905                .llm(Arc::new(mock_with_response("done")))
17906                .tool(Arc::new(LockedWriteTool {
17907                    active: Arc::clone(&active),
17908                    max_active: Arc::clone(&max_active),
17909                }))
17910                .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
17911                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17912                .approval_handler(handler)
17913                .build()
17914                .unwrap(),
17915        );
17916        let control = agent.runtime_control();
17917        let running = Arc::clone(&agent);
17918        let call = tokio::spawn(async move {
17919            running
17920                .invoke_tool(ToolExecutionRequest::new(
17921                    "approval-scope",
17922                    "locked_write",
17923                    serde_json::json!({"path": "./before-approval.txt"}),
17924                    ToolCallSource::Manual,
17925                ))
17926                .await
17927                .unwrap()
17928        });
17929        entered.wait().await;
17930        let expected_version = control.set_tool_scope(Vec::new());
17931        release.notify_one();
17932        let record = call.await.unwrap();
17933
17934        assert!(!record.executed);
17935        assert!(!record.success);
17936        assert_eq!(record.runtime_config_version, expected_version);
17937        assert_eq!(record.executed_arguments["path"], "./after-approval.txt");
17938        assert_eq!(max_active.load(Ordering::SeqCst), 0);
17939        assert_eq!(
17940            record.metadata["runtime_scope_snapshot"],
17941            serde_json::json!([])
17942        );
17943    }
17944
17945    #[tokio::test]
17946    async fn approval_is_rechecked_against_final_policy_snapshot() {
17947        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17948        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17949        let entered = Arc::new(tokio::sync::Barrier::new(2));
17950        let release = Arc::new(tokio::sync::Notify::new());
17951        let handler = Arc::new(BlockingApprovalHandler {
17952            entered: Arc::clone(&entered),
17953            release: Arc::clone(&release),
17954            result: ApprovalResult::Approved,
17955        });
17956        let agent = Arc::new(
17957            AgentBuilder::new()
17958                .system_prompt("Test final policy validation.")
17959                .llm(Arc::new(mock_with_response("done")))
17960                .tool(Arc::new(LockedWriteTool {
17961                    active: Arc::clone(&active),
17962                    max_active: Arc::clone(&max_active),
17963                }))
17964                .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
17965                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17966                .approval_handler(handler)
17967                .build()
17968                .unwrap(),
17969        );
17970        let control = agent.runtime_control();
17971        let running = Arc::clone(&agent);
17972        let call = tokio::spawn(async move {
17973            running
17974                .invoke_tool(ToolExecutionRequest::new(
17975                    "approval-policy",
17976                    "locked_write",
17977                    serde_json::json!({"path": "./policy.txt"}),
17978                    ToolCallSource::Manual,
17979                ))
17980                .await
17981                .unwrap()
17982        });
17983        entered.wait().await;
17984        let expected_version = control.set_tool_security(approval_security_config(false));
17985        release.notify_one();
17986        let record = call.await.unwrap();
17987
17988        assert!(!record.executed);
17989        assert!(!record.success);
17990        assert_eq!(record.runtime_config_version, expected_version);
17991        assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
17992        assert_eq!(max_active.load(Ordering::SeqCst), 0);
17993        assert!(record.metadata.contains_key("policy_snapshot"));
17994    }
17995
17996    #[test]
17997    fn invalid_live_policy_does_not_replace_snapshot_or_generation() {
17998        let agent = AgentBuilder::new()
17999            .system_prompt("Test runtime policy validation.")
18000            .llm(Arc::new(mock_with_response("done")))
18001            .build()
18002            .unwrap();
18003        let control = agent.runtime_control();
18004        let mut valid = ToolSecurityConfig::default();
18005        valid.tools.insert(
18006            "web_search".to_string(),
18007            ai_agents_tools::ToolPolicyConfig {
18008                max_results: Some(5),
18009                ..Default::default()
18010            },
18011        );
18012        let generation = control.try_set_tool_security(valid).unwrap();
18013
18014        let mut invalid = ToolSecurityConfig::default();
18015        invalid.tools.insert(
18016            "web_search".to_string(),
18017            ai_agents_tools::ToolPolicyConfig {
18018                max_results: Some(0),
18019                ..Default::default()
18020            },
18021        );
18022        let error = control.try_set_tool_security(invalid).unwrap_err();
18023
18024        assert!(
18025            error
18026                .to_string()
18027                .contains("max_results must be greater than 0")
18028        );
18029        assert_eq!(control.version(), generation);
18030        assert_eq!(
18031            control
18032                .state
18033                .tool_security_override
18034                .read()
18035                .as_ref()
18036                .unwrap()
18037                .config()
18038                .tools["web_search"]
18039                .max_results,
18040            Some(5)
18041        );
18042    }
18043
18044    // Verifies invalid security timeout configuration prevents agent construction and later side effects.
18045    #[test]
18046    fn invalid_timeout_config_stops_before_approval_or_tool_invocation() {
18047        let tool_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18048        let approval_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18049        let spec = crate::spec::AgentSpec {
18050            tool_security: ToolSecurityConfig {
18051                enabled: true,
18052                default_timeout_ms: u64::MAX,
18053                ..Default::default()
18054            },
18055            ..Default::default()
18056        };
18057
18058        let result = AgentBuilder::from_spec(spec)
18059            .llm(Arc::new(mock_with_response("done")))
18060            .tool(Arc::new(FlakyWriteTool {
18061                calls: Arc::clone(&tool_calls),
18062            }))
18063            .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
18064            .approval_handler(Arc::new(CountingApprovalHandler {
18065                calls: Arc::clone(&approval_calls),
18066            }))
18067            .build();
18068
18069        assert!(result.is_err());
18070        assert_eq!(approval_calls.load(Ordering::SeqCst), 0);
18071        assert_eq!(tool_calls.load(Ordering::SeqCst), 0);
18072    }
18073
18074    // Verifies invalid recovery timeout configuration prevents agent construction and later side effects.
18075    #[test]
18076    fn invalid_recovery_timeout_config_stops_before_approval_or_tool_invocation() {
18077        use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
18078
18079        let tool_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18080        let approval_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18081        let spec = crate::spec::AgentSpec {
18082            error_recovery: ErrorRecoveryConfig {
18083                tools: ToolRecoveryConfig {
18084                    default: ToolRetryConfig {
18085                        timeout_ms: Some(u64::MAX),
18086                        ..Default::default()
18087                    },
18088                    ..Default::default()
18089                },
18090                ..Default::default()
18091            },
18092            ..Default::default()
18093        };
18094
18095        let result = AgentBuilder::from_spec(spec)
18096            .llm(Arc::new(mock_with_response("done")))
18097            .tool(Arc::new(FlakyWriteTool {
18098                calls: Arc::clone(&tool_calls),
18099            }))
18100            .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
18101            .approval_handler(Arc::new(CountingApprovalHandler {
18102                calls: Arc::clone(&approval_calls),
18103            }))
18104            .build();
18105
18106        assert!(result.is_err());
18107        assert_eq!(approval_calls.load(Ordering::SeqCst), 0);
18108        assert_eq!(tool_calls.load(Ordering::SeqCst), 0);
18109    }
18110
18111    // Verifies rejected runtime policy replacement preserves both the valid snapshot and generation.
18112    #[test]
18113    fn invalid_timeout_policy_does_not_replace_snapshot_or_generation() {
18114        let agent = AgentBuilder::new()
18115            .system_prompt("Test runtime timeout policy validation.")
18116            .llm(Arc::new(mock_with_response("done")))
18117            .build()
18118            .unwrap();
18119        let control = agent.runtime_control();
18120        let valid = ToolSecurityConfig {
18121            default_timeout_ms: 5_000,
18122            ..Default::default()
18123        };
18124        let generation = control.try_set_tool_security(valid).unwrap();
18125
18126        let invalid = ToolSecurityConfig {
18127            default_timeout_ms: MAX_TOOL_TIMEOUT_MS + 1,
18128            ..Default::default()
18129        };
18130        let error = control.try_set_tool_security(invalid).unwrap_err();
18131
18132        assert!(error.to_string().contains(&format!(
18133            "tool_security.default_timeout_ms must be no greater than {MAX_TOOL_TIMEOUT_MS} milliseconds"
18134        )));
18135        assert_eq!(control.version(), generation);
18136        assert_eq!(
18137            control
18138                .state
18139                .tool_security_override
18140                .read()
18141                .as_ref()
18142                .unwrap()
18143                .config()
18144                .default_timeout_ms,
18145            5_000
18146        );
18147    }
18148
18149    // Verifies the shared maximum converts exactly while larger public u64 values fail closed.
18150    #[test]
18151    fn runtime_tool_timeout_conversion_enforces_the_stable_boundary() {
18152        let timeout = RuntimeAgent::validated_tool_timeout(MAX_TOOL_TIMEOUT_MS).unwrap();
18153        assert_eq!(timeout.timer, Duration::from_millis(MAX_TOOL_TIMEOUT_MS));
18154        assert_eq!(
18155            timeout.deadline_delta,
18156            chrono::Duration::milliseconds(MAX_TOOL_TIMEOUT_MS as i64)
18157        );
18158
18159        for timeout_ms in [MAX_TOOL_TIMEOUT_MS + 1, u64::MAX] {
18160            let error = RuntimeAgent::validated_tool_timeout(timeout_ms).unwrap_err();
18161            assert!(error.to_string().contains(&format!(
18162                "effective tool timeout_ms must be no greater than {MAX_TOOL_TIMEOUT_MS} milliseconds"
18163            )));
18164        }
18165    }
18166
18167    #[tokio::test]
18168    async fn persistent_override_preserves_rate_history_within_generation() {
18169        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18170        let agent = AgentBuilder::new()
18171            .system_prompt("Test persistent policy overrides.")
18172            .llm(Arc::new(mock_with_response("done")))
18173            .tool(Arc::new(RecoveryTestTool {
18174                id: "limited_override".to_string(),
18175                succeeds: true,
18176                calls: Arc::clone(&calls),
18177                max_output_chars: None,
18178            }))
18179            .build()
18180            .unwrap();
18181        let mut security = ToolSecurityConfig {
18182            enabled: true,
18183            fail_closed: true,
18184            ..Default::default()
18185        };
18186        let policy = ai_agents_tools::ToolPolicyConfig {
18187            write_paths: vec![".".to_string()],
18188            rate_limit: Some(1),
18189            ..Default::default()
18190        };
18191        security
18192            .tools
18193            .insert("limited_override".to_string(), policy);
18194        let generation = agent.runtime_control().set_tool_security(security);
18195
18196        let first = agent
18197            .invoke_tool(ToolExecutionRequest::new(
18198                "limited-first",
18199                "limited_override",
18200                serde_json::json!({"path": "./limited.txt"}),
18201                ToolCallSource::Manual,
18202            ))
18203            .await
18204            .unwrap();
18205        let second = agent
18206            .invoke_tool(ToolExecutionRequest::new(
18207                "limited-second",
18208                "limited_override",
18209                serde_json::json!({"path": "./limited.txt"}),
18210                ToolCallSource::Manual,
18211            ))
18212            .await
18213            .unwrap();
18214
18215        assert!(first.success);
18216        assert_eq!(first.policy_version, generation);
18217        assert!(!second.executed);
18218        assert!(second.output.contains("Rate limit exceeded"));
18219        assert_eq!(second.policy_version, generation);
18220        assert_eq!(calls.load(Ordering::SeqCst), 1);
18221    }
18222
18223    #[tokio::test]
18224    async fn concurrent_rate_admission_consumes_capacity_atomically() {
18225        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18226        let tool = Arc::new(RecoveryTestTool {
18227            id: "atomic_rate".to_string(),
18228            succeeds: true,
18229            calls: Arc::clone(&calls),
18230            max_output_chars: None,
18231        });
18232        let arguments = serde_json::json!({"path": "./atomic-rate.txt"});
18233        let bindings = tool.policy_bindings();
18234        let classification = tool.classify_call(&arguments);
18235        let resource_keys =
18236            tool_resource_lock_keys(tool.id(), &arguments, &bindings, &classification);
18237        let mut security = ToolSecurityConfig {
18238            enabled: true,
18239            fail_closed: true,
18240            ..Default::default()
18241        };
18242        let policy = ai_agents_tools::ToolPolicyConfig {
18243            write_paths: vec![".".to_string()],
18244            rate_limit: Some(1),
18245            ..Default::default()
18246        };
18247        security.tools.insert(tool.id().to_string(), policy);
18248        let agent = Arc::new(
18249            AgentBuilder::new()
18250                .system_prompt("Test atomic rate admission.")
18251                .llm(Arc::new(mock_with_response("done")))
18252                .tool(tool)
18253                .tool_security(ToolSecurityEngine::new(security))
18254                .build()
18255                .unwrap(),
18256        );
18257        let held = agent
18258            .acquire_tool_resource_locks(&resource_keys)
18259            .await
18260            .unwrap();
18261        let left = {
18262            let agent = Arc::clone(&agent);
18263            let arguments = arguments.clone();
18264            tokio::spawn(async move {
18265                agent
18266                    .invoke_tool(ToolExecutionRequest::new(
18267                        "atomic-rate-left",
18268                        "atomic_rate",
18269                        arguments,
18270                        ToolCallSource::Manual,
18271                    ))
18272                    .await
18273                    .unwrap()
18274            })
18275        };
18276        let right = {
18277            let agent = Arc::clone(&agent);
18278            tokio::spawn(async move {
18279                agent
18280                    .invoke_tool(ToolExecutionRequest::new(
18281                        "atomic-rate-right",
18282                        "atomic_rate",
18283                        arguments,
18284                        ToolCallSource::Manual,
18285                    ))
18286                    .await
18287                    .unwrap()
18288            })
18289        };
18290        tokio::time::sleep(std::time::Duration::from_millis(25)).await;
18291        drop(held);
18292        let (left, right) = tokio::join!(left, right);
18293        let records = [left.unwrap(), right.unwrap()];
18294
18295        assert_eq!(records.iter().filter(|record| record.success).count(), 1);
18296        assert_eq!(records.iter().filter(|record| record.executed).count(), 1);
18297        assert!(
18298            records.iter().any(|record| {
18299                !record.executed && record.output.contains("Rate limit exceeded")
18300            })
18301        );
18302        assert_eq!(calls.load(Ordering::SeqCst), 1);
18303    }
18304
18305    #[tokio::test]
18306    async fn changed_policy_generation_invalidates_pending_approval() {
18307        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18308        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18309        let entered = Arc::new(tokio::sync::Barrier::new(2));
18310        let release = Arc::new(tokio::sync::Notify::new());
18311        let handler = Arc::new(BlockingApprovalHandler {
18312            entered: Arc::clone(&entered),
18313            release: Arc::clone(&release),
18314            result: ApprovalResult::Approved,
18315        });
18316        let agent = Arc::new(
18317            AgentBuilder::new()
18318                .system_prompt("Test stale approval denial.")
18319                .llm(Arc::new(mock_with_response("done")))
18320                .tool(Arc::new(LockedWriteTool {
18321                    active: Arc::clone(&active),
18322                    max_active: Arc::clone(&max_active),
18323                }))
18324                .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
18325                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
18326                .approval_handler(handler)
18327                .build()
18328                .unwrap(),
18329        );
18330        let running = Arc::clone(&agent);
18331        let call = tokio::spawn(async move {
18332            running
18333                .invoke_tool(ToolExecutionRequest::new(
18334                    "stale-approval",
18335                    "locked_write",
18336                    serde_json::json!({"path": "./stale.txt"}),
18337                    ToolCallSource::Manual,
18338                ))
18339                .await
18340                .unwrap()
18341        });
18342        entered.wait().await;
18343        let generation = agent
18344            .runtime_control()
18345            .set_tool_security(approval_security_config(true));
18346        release.notify_one();
18347        let record = call.await.unwrap();
18348
18349        assert!(!record.executed);
18350        assert!(record.output.contains("Approval became stale"));
18351        assert_eq!(record.policy_version, generation);
18352        assert_eq!(max_active.load(Ordering::SeqCst), 0);
18353    }
18354
18355    #[tokio::test]
18356    async fn final_policy_reapplies_argument_caps_after_approval_changes() {
18357        use ai_agents_hitl::CallbackHandler;
18358
18359        let mut security = ToolSecurityConfig {
18360            enabled: true,
18361            fail_closed: true,
18362            ..Default::default()
18363        };
18364        let policy = ai_agents_tools::ToolPolicyConfig {
18365            read_paths: vec![".".to_string()],
18366            max_results: Some(5),
18367            require_confirmation: true,
18368            ..Default::default()
18369        };
18370        security.tools.insert("context_echo".to_string(), policy);
18371        let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
18372            changes: HashMap::from([("max_results".to_string(), serde_json::json!(99))]),
18373        });
18374        let agent = AgentBuilder::new()
18375            .system_prompt("Test final argument caps.")
18376            .llm(Arc::new(mock_with_response("done")))
18377            .tool(Arc::new(ContextEchoTool))
18378            .tool_security(ToolSecurityEngine::new(security))
18379            .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
18380            .approval_handler(Arc::new(handler))
18381            .build()
18382            .unwrap();
18383
18384        let record = agent
18385            .invoke_tool(ToolExecutionRequest::new(
18386                "final-cap",
18387                "context_echo",
18388                serde_json::json!({"path": ".", "max_results": 1}),
18389                ToolCallSource::Manual,
18390            ))
18391            .await
18392            .unwrap();
18393
18394        assert!(record.success);
18395        assert_eq!(record.executed_arguments["max_results"], 5);
18396        assert_eq!(
18397            record.approval.unwrap().modified_arguments.unwrap()["max_results"],
18398            5
18399        );
18400    }
18401
18402    #[tokio::test]
18403    async fn no_binding_writes_use_canonical_fallback_lock() {
18404        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18405        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18406        let agent = Arc::new(
18407            AgentBuilder::new()
18408                .system_prompt("Test fallback resource locks.")
18409                .llm(Arc::new(mock_with_response("done")))
18410                .tool(Arc::new(NoBindingWriteTool {
18411                    active: Arc::clone(&active),
18412                    max_active: Arc::clone(&max_active),
18413                }))
18414                .build()
18415                .unwrap(),
18416        );
18417        let left = {
18418            let agent = Arc::clone(&agent);
18419            tokio::spawn(async move {
18420                agent
18421                    .invoke_tool(ToolExecutionRequest::new(
18422                        "no-binding-left",
18423                        "no_binding_write",
18424                        serde_json::json!({}),
18425                        ToolCallSource::Manual,
18426                    ))
18427                    .await
18428                    .unwrap()
18429            })
18430        };
18431        let right = {
18432            let agent = Arc::clone(&agent);
18433            tokio::spawn(async move {
18434                agent
18435                    .invoke_tool(ToolExecutionRequest::new(
18436                        "no-binding-right",
18437                        "no_binding_write",
18438                        serde_json::json!({}),
18439                        ToolCallSource::Manual,
18440                    ))
18441                    .await
18442                    .unwrap()
18443            })
18444        };
18445        let (left, right) = tokio::join!(left, right);
18446
18447        assert!(left.unwrap().success);
18448        assert!(right.unwrap().success);
18449        assert_eq!(max_active.load(Ordering::SeqCst), 1);
18450    }
18451
18452    #[tokio::test]
18453    async fn parent_and_child_paths_share_a_resource_lock() {
18454        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18455        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18456        let agent = Arc::new(
18457            AgentBuilder::new()
18458                .system_prompt("Test parent child resource locks.")
18459                .llm(Arc::new(mock_with_response("done")))
18460                .tool(Arc::new(LockedWriteTool {
18461                    active: Arc::clone(&active),
18462                    max_active: Arc::clone(&max_active),
18463                }))
18464                .build()
18465                .unwrap(),
18466        );
18467        let parent = format!("./lock-parent-{}", uuid::Uuid::new_v4());
18468        let child = format!("{}/child.txt", parent);
18469        let left = {
18470            let agent = Arc::clone(&agent);
18471            tokio::spawn(async move {
18472                agent
18473                    .invoke_tool(ToolExecutionRequest::new(
18474                        "parent-lock",
18475                        "locked_write",
18476                        serde_json::json!({"path": parent}),
18477                        ToolCallSource::Manual,
18478                    ))
18479                    .await
18480                    .unwrap()
18481            })
18482        };
18483        let right = {
18484            let agent = Arc::clone(&agent);
18485            tokio::spawn(async move {
18486                agent
18487                    .invoke_tool(ToolExecutionRequest::new(
18488                        "child-lock",
18489                        "locked_write",
18490                        serde_json::json!({"path": child}),
18491                        ToolCallSource::Manual,
18492                    ))
18493                    .await
18494                    .unwrap()
18495            })
18496        };
18497        let (left, right) = tokio::join!(left, right);
18498
18499        assert!(left.unwrap().success);
18500        assert!(right.unwrap().success);
18501        assert_eq!(max_active.load(Ordering::SeqCst), 1);
18502    }
18503
18504    #[tokio::test]
18505    async fn tool_hooks_can_reenter_after_resource_guards_are_dropped() {
18506        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18507        let hooks = Arc::new(ReentrantToolHooks {
18508            agent: parking_lot::Mutex::new(None),
18509            invoked: AtomicBool::new(false),
18510            nested_success: AtomicBool::new(false),
18511        });
18512        let agent = Arc::new(
18513            AgentBuilder::new()
18514                .system_prompt("Test hook reentrancy.")
18515                .llm(Arc::new(mock_with_response("done")))
18516                .tool(Arc::new(RecoveryTestTool {
18517                    id: "reentrant_write".to_string(),
18518                    succeeds: true,
18519                    calls: Arc::clone(&calls),
18520                    max_output_chars: None,
18521                }))
18522                .hooks(hooks.clone())
18523                .build()
18524                .unwrap(),
18525        );
18526        *hooks.agent.lock() = Some(Arc::downgrade(&agent));
18527        let record = tokio::time::timeout(
18528            std::time::Duration::from_secs(2),
18529            agent.invoke_tool(ToolExecutionRequest::new(
18530                "outer-hook-call",
18531                "reentrant_write",
18532                serde_json::json!({"path": "./hook.txt"}),
18533                ToolCallSource::Manual,
18534            )),
18535        )
18536        .await
18537        .expect("tool completion hook must not retain resource guards")
18538        .unwrap();
18539
18540        assert!(record.success);
18541        assert!(hooks.nested_success.load(Ordering::SeqCst));
18542        assert_eq!(calls.load(Ordering::SeqCst), 2);
18543    }
18544
18545    /// Confirms fallback starts only after the failed original lifecycle and resource ownership complete.
18546    #[tokio::test]
18547    async fn fallback_finalizes_original_record_before_shared_execution() {
18548        let primary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18549        let fallback_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18550        let hooks = Arc::new(ToolLifecycleRecordingHooks::new());
18551        let agent = AgentBuilder::new()
18552            .system_prompt("Test fallback execution.")
18553            .llm(Arc::new(mock_with_response("done")))
18554            .tool(Arc::new(RecoveryTestTool {
18555                id: "primary".to_string(),
18556                succeeds: false,
18557                calls: Arc::clone(&primary_calls),
18558                max_output_chars: None,
18559            }))
18560            .tool(Arc::new(RecoveryTestTool {
18561                id: "fallback".to_string(),
18562                succeeds: true,
18563                calls: Arc::clone(&fallback_calls),
18564                max_output_chars: None,
18565            }))
18566            .recovery_manager(recovery_manager_with_fallbacks([(
18567                "primary".to_string(),
18568                "fallback".to_string(),
18569            )]))
18570            .hooks(hooks.clone())
18571            .build()
18572            .unwrap();
18573        let record = tokio::time::timeout(
18574            std::time::Duration::from_secs(2),
18575            agent.invoke_tool(ToolExecutionRequest::new(
18576                "fallback-call",
18577                "primary",
18578                serde_json::json!({"path": "./shared.txt"}),
18579                ToolCallSource::Manual,
18580            )),
18581        )
18582        .await
18583        .expect("fallback must not retain the primary resource guard")
18584        .unwrap();
18585
18586        assert_eq!(
18587            hooks.events(),
18588            vec![
18589                "start:primary",
18590                "complete:primary:false",
18591                "record:primary:true",
18592                "error",
18593                "start:fallback",
18594                "complete:fallback:true",
18595                "record:fallback:true",
18596            ]
18597        );
18598        let records = hooks.records();
18599        assert_eq!(records.len(), 2);
18600        let original = &records[0];
18601        assert_eq!(original.canonical_id, "primary");
18602        assert!(matches!(original.source, ToolCallSource::Manual));
18603        assert!(original.executed);
18604        assert!(!original.success);
18605
18606        let fallback = &records[1];
18607        assert_eq!(fallback.canonical_id, "fallback");
18608        assert_eq!(fallback.call_id, "fallback-call");
18609        assert!(matches!(
18610            &fallback.source,
18611            ToolCallSource::Fallback { original_tool } if original_tool == "primary"
18612        ));
18613        assert!(fallback.executed);
18614        assert!(fallback.success);
18615        assert_eq!(record.canonical_id, fallback.canonical_id);
18616        assert_eq!(record.output, fallback.output);
18617
18618        let history = agent.tool_call_history();
18619        assert_eq!(
18620            history
18621                .iter()
18622                .map(|entry| entry.tool_id.as_str())
18623                .collect::<Vec<_>>(),
18624            vec!["primary", "fallback"]
18625        );
18626        assert_eq!(history[0].result.get("success"), Some(&Value::Bool(false)));
18627        assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
18628        assert_eq!(fallback_calls.load(Ordering::SeqCst), 1);
18629    }
18630
18631    /// Confirms a self-referential fallback produces bounded non-execution evidence after the failed original lifecycle.
18632    #[tokio::test]
18633    async fn self_fallback_cycle_is_denied_before_reinvocation() {
18634        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18635        let hooks = Arc::new(ToolLifecycleRecordingHooks::new());
18636        let agent = AgentBuilder::new()
18637            .system_prompt("Test self-fallback cycle admission.")
18638            .llm(Arc::new(mock_with_response("done")))
18639            .tool(Arc::new(RecoveryTestTool {
18640                id: "primary".to_string(),
18641                succeeds: false,
18642                calls: Arc::clone(&calls),
18643                max_output_chars: None,
18644            }))
18645            .recovery_manager(recovery_manager_with_fallbacks([(
18646                "primary".to_string(),
18647                "primary".to_string(),
18648            )]))
18649            .hooks(hooks.clone())
18650            .build()
18651            .unwrap();
18652
18653        let record = tokio::time::timeout(
18654            std::time::Duration::from_secs(2),
18655            agent.invoke_tool(ToolExecutionRequest::new(
18656                "self-fallback-call",
18657                "primary",
18658                serde_json::json!({"path": "./shared.txt"}),
18659                ToolCallSource::Manual,
18660            )),
18661        )
18662        .await
18663        .expect("self fallback must terminate without recursive execution")
18664        .unwrap();
18665
18666        assert_eq!(calls.load(Ordering::SeqCst), 1);
18667        assert_eq!(record.canonical_id, "primary");
18668        assert!(!record.executed);
18669        assert!(!record.success);
18670        assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
18671        assert!(record.output.contains("fallback cycle"));
18672        assert!(matches!(
18673            record.source,
18674            ToolCallSource::Fallback { ref original_tool } if original_tool == "primary"
18675        ));
18676        assert_eq!(
18677            record.metadata.get("fallback_chain"),
18678            Some(&serde_json::json!(["primary"]))
18679        );
18680        assert_eq!(
18681            hooks.events(),
18682            vec![
18683                "start:primary",
18684                "complete:primary:false",
18685                "record:primary:true",
18686                "error",
18687                "complete:primary:false",
18688                "record:primary:false",
18689                "error",
18690            ]
18691        );
18692        assert_eq!(agent.tool_call_history().len(), 2);
18693    }
18694
18695    /// Confirms fallback ancestry compares resolved canonical IDs so display aliases cannot conceal a multi-tool cycle.
18696    #[tokio::test]
18697    async fn alias_mediated_fallback_cycle_is_denied_canonically() {
18698        let primary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18699        let secondary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18700        let hooks = Arc::new(ToolLifecycleRecordingHooks::new());
18701        let agent = AgentBuilder::new()
18702            .system_prompt("Test canonical fallback cycle admission.")
18703            .llm(Arc::new(mock_with_response("done")))
18704            .tool(Arc::new(RecoveryTestTool {
18705                id: "primary".to_string(),
18706                succeeds: false,
18707                calls: Arc::clone(&primary_calls),
18708                max_output_chars: None,
18709            }))
18710            .tool(Arc::new(RecoveryTestTool {
18711                id: "secondary".to_string(),
18712                succeeds: false,
18713                calls: Arc::clone(&secondary_calls),
18714                max_output_chars: None,
18715            }))
18716            .recovery_manager(recovery_manager_with_fallbacks([
18717                ("primary".to_string(), "secondary".to_string()),
18718                ("secondary".to_string(), "primary alias".to_string()),
18719            ]))
18720            .hooks(hooks.clone())
18721            .build()
18722            .unwrap();
18723        agent.tools.set_tool_aliases(
18724            "primary",
18725            ToolAliases::new().with_name("en", "primary alias"),
18726        );
18727
18728        let record = tokio::time::timeout(
18729            std::time::Duration::from_secs(2),
18730            agent.invoke_tool(ToolExecutionRequest::new(
18731                "alias-fallback-call",
18732                "primary",
18733                serde_json::json!({"path": "./shared.txt"}),
18734                ToolCallSource::Manual,
18735            )),
18736        )
18737        .await
18738        .expect("alias-mediated fallback cycle must terminate")
18739        .unwrap();
18740
18741        assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
18742        assert_eq!(secondary_calls.load(Ordering::SeqCst), 1);
18743        assert_eq!(record.requested_name, "primary alias");
18744        assert_eq!(record.canonical_id, "primary");
18745        assert!(!record.executed);
18746        assert!(record.output.contains("fallback cycle"));
18747        assert_eq!(
18748            record.metadata.get("fallback_chain"),
18749            Some(&serde_json::json!(["primary", "secondary"]))
18750        );
18751        assert_eq!(hooks.records().len(), 3);
18752        assert_eq!(agent.tool_call_history().len(), 3);
18753    }
18754
18755    /// Confirms final provider refresh cannot move a fallback alias onto an already visited canonical ancestor.
18756    #[tokio::test]
18757    async fn final_canonical_drift_cannot_bypass_fallback_ancestry() {
18758        let primary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18759        let secondary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18760        let provider = Arc::new(DriftingFallbackProvider {
18761            refreshed: AtomicBool::new(false),
18762            primary_calls: Arc::clone(&primary_calls),
18763            secondary_calls: Arc::clone(&secondary_calls),
18764        });
18765        let registry = ToolRegistry::new();
18766        registry.register_provider(provider).await.unwrap();
18767        let lifecycle = Arc::new(ToolLifecycleRecordingHooks::new());
18768        let hooks = Arc::new(RefreshFallbackProviderHooks {
18769            agent: parking_lot::Mutex::new(None),
18770            lifecycle: Arc::clone(&lifecycle),
18771        });
18772        let agent = Arc::new(
18773            AgentBuilder::new()
18774                .system_prompt("Test final canonical fallback admission.")
18775                .llm(Arc::new(mock_with_response("done")))
18776                .tools(registry)
18777                .recovery_manager(recovery_manager_with_fallbacks([(
18778                    "primary".to_string(),
18779                    "fallback alias".to_string(),
18780                )]))
18781                .hooks(hooks.clone())
18782                .build()
18783                .unwrap(),
18784        );
18785        *hooks.agent.lock() = Some(Arc::downgrade(&agent));
18786
18787        let record = agent
18788            .invoke_tool(ToolExecutionRequest::new(
18789                "drifting-fallback-call",
18790                "primary",
18791                serde_json::json!({"path": "./shared.txt"}),
18792                ToolCallSource::Manual,
18793            ))
18794            .await
18795            .unwrap();
18796
18797        assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
18798        assert_eq!(secondary_calls.load(Ordering::SeqCst), 0);
18799        assert_eq!(record.canonical_id, "secondary");
18800        assert!(!record.executed);
18801        assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
18802        assert!(record.output.contains("fallback cycle"));
18803        assert_eq!(
18804            record.metadata.get("fallback_chain"),
18805            Some(&serde_json::json!(["primary", "secondary"]))
18806        );
18807        assert_eq!(
18808            record.metadata.get("final_resolved_canonical_id"),
18809            Some(&serde_json::json!("primary"))
18810        );
18811        assert_eq!(
18812            lifecycle.events(),
18813            vec![
18814                "start:primary",
18815                "complete:primary:false",
18816                "record:primary:true",
18817                "error",
18818                "start:secondary",
18819                "complete:secondary:false",
18820                "record:secondary:false",
18821                "error",
18822            ]
18823        );
18824        let records = lifecycle.records();
18825        assert_eq!(records.len(), 2);
18826        assert_eq!(records[1].canonical_id, "secondary");
18827        assert_eq!(
18828            records[1].metadata.get("final_resolved_canonical_id"),
18829            Some(&serde_json::json!("primary"))
18830        );
18831        let history = agent.tool_call_history();
18832        assert_eq!(
18833            history
18834                .iter()
18835                .map(|entry| entry.tool_id.as_str())
18836                .collect::<Vec<_>>(),
18837            vec!["primary", "secondary"]
18838        );
18839    }
18840
18841    /// Confirms a unique acyclic fallback chain is bounded before the first request beyond the fixed hop limit invokes its tool.
18842    #[tokio::test]
18843    async fn acyclic_fallback_chain_is_denied_after_the_hop_limit() {
18844        let tool_count = MAX_TOOL_FALLBACK_HOPS + 2;
18845        let calls = (0..tool_count)
18846            .map(|_| Arc::new(std::sync::atomic::AtomicUsize::new(0)))
18847            .collect::<Vec<_>>();
18848        let mut builder = AgentBuilder::new()
18849            .system_prompt("Test bounded acyclic fallback admission.")
18850            .llm(Arc::new(mock_with_response("done")));
18851        for (index, counter) in calls.iter().enumerate() {
18852            builder = builder.tool(Arc::new(RecoveryTestTool {
18853                id: format!("fallback_{index}"),
18854                succeeds: false,
18855                calls: Arc::clone(counter),
18856                max_output_chars: None,
18857            }));
18858        }
18859        let fallbacks = (0..tool_count - 1).map(|index| {
18860            (
18861                format!("fallback_{index}"),
18862                format!("fallback_{}", index + 1),
18863            )
18864        });
18865        let agent = builder
18866            .recovery_manager(recovery_manager_with_fallbacks(fallbacks))
18867            .build()
18868            .unwrap();
18869
18870        let record = tokio::time::timeout(
18871            std::time::Duration::from_secs(2),
18872            agent.invoke_tool(ToolExecutionRequest::new(
18873                "bounded-fallback-call",
18874                "fallback_0",
18875                serde_json::json!({"path": "./shared.txt"}),
18876                ToolCallSource::Manual,
18877            )),
18878        )
18879        .await
18880        .expect("bounded fallback chain must terminate")
18881        .unwrap();
18882
18883        for counter in calls.iter().take(MAX_TOOL_FALLBACK_HOPS + 1) {
18884            assert_eq!(counter.load(Ordering::SeqCst), 1);
18885        }
18886        assert_eq!(calls[MAX_TOOL_FALLBACK_HOPS + 1].load(Ordering::SeqCst), 0);
18887        assert_eq!(
18888            record.canonical_id,
18889            format!("fallback_{}", MAX_TOOL_FALLBACK_HOPS + 1)
18890        );
18891        assert!(!record.executed);
18892        assert!(record.output.contains("maximum of 16 hops"));
18893        assert_eq!(agent.tool_call_history().len(), tool_count);
18894    }
18895
18896    #[tokio::test]
18897    async fn diagnostics_without_provider_records_unavailable_without_execution() {
18898        let mock = mock_with_response("hello");
18899        let yaml = r#"
18900name: DiagnosticsNoProviderAgent
18901system_prompt: "Review diagnostics."
18902tools: [diagnostics]
18903"#;
18904        let agent = AgentBuilder::from_yaml(yaml)
18905            .unwrap()
18906            .llm(Arc::new(mock))
18907            .auto_configure_features()
18908            .unwrap()
18909            .build()
18910            .unwrap();
18911
18912        let record = agent
18913            .invoke_tool(ToolExecutionRequest::new(
18914                "diagnostics-call",
18915                "diagnostics",
18916                serde_json::json!({}),
18917                ToolCallSource::Manual,
18918            ))
18919            .await
18920            .unwrap();
18921
18922        assert!(!record.executed);
18923        assert!(!record.success);
18924        assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
18925    }
18926
18927    #[tokio::test]
18928    async fn web_search_without_provider_records_unavailable_without_execution() {
18929        let mock = mock_with_response("hello");
18930        let yaml = r#"
18931name: WebSearchNoProviderAgent
18932system_prompt: "You search the web."
18933tools: [web_search]
18934"#;
18935        let agent = AgentBuilder::from_yaml(yaml)
18936            .unwrap()
18937            .llm(Arc::new(mock))
18938            .auto_configure_features()
18939            .unwrap()
18940            .build()
18941            .unwrap();
18942
18943        let record = agent
18944            .invoke_tool(ToolExecutionRequest::new(
18945                "web-search-call",
18946                "web_search",
18947                serde_json::json!({"query": "rust async"}),
18948                ToolCallSource::Manual,
18949            ))
18950            .await
18951            .unwrap();
18952
18953        assert!(!record.executed);
18954        assert!(!record.success);
18955        assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
18956    }
18957
18958    #[tokio::test]
18959    async fn unavailable_host_tool_does_not_request_approval() {
18960        let approvals = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18961        let handler = Arc::new(CountingApprovalHandler {
18962            calls: Arc::clone(&approvals),
18963        });
18964        let mut security = ToolSecurityConfig {
18965            enabled: true,
18966            fail_closed: true,
18967            ..Default::default()
18968        };
18969        security.tools.insert(
18970            "web_search".to_string(),
18971            ai_agents_tools::ToolPolicyConfig {
18972                enabled: true,
18973                require_confirmation: true,
18974                ..Default::default()
18975            },
18976        );
18977        let yaml = r#"
18978name: UnavailableApprovalAgent
18979system_prompt: "Search only with approval."
18980tools: [web_search]
18981"#;
18982        let agent = AgentBuilder::from_yaml(yaml)
18983            .unwrap()
18984            .llm(Arc::new(mock_with_response("done")))
18985            .auto_configure_features()
18986            .unwrap()
18987            .tool_security(ToolSecurityEngine::new(security))
18988            .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
18989            .approval_handler(handler)
18990            .build()
18991            .unwrap();
18992
18993        let record = agent
18994            .invoke_tool(ToolExecutionRequest::new(
18995                "unavailable-before-approval",
18996                "web_search",
18997                serde_json::json!({"query": "rust async"}),
18998                ToolCallSource::Manual,
18999            ))
19000            .await
19001            .unwrap();
19002
19003        assert_eq!(approvals.load(Ordering::SeqCst), 0);
19004        assert!(!record.executed);
19005        assert!(!record.success);
19006        assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
19007        assert!(
19008            record
19009                .approval
19010                .as_ref()
19011                .is_some_and(|approval| matches!(approval.status, ToolApprovalStatus::Unavailable))
19012        );
19013    }
19014
19015    #[tokio::test]
19016    async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_omitted() {
19017        let mock = mock_with_response("hello");
19018        let yaml = r#"
19019name: SpawnerNoGrantAgent
19020system_prompt: "You manage agents."
19021spawner:
19022  max_agents: 2
19023"#;
19024        let agent = AgentBuilder::from_yaml(yaml)
19025            .unwrap()
19026            .llm(Arc::new(mock))
19027            .auto_configure_features()
19028            .unwrap()
19029            .auto_configure_spawner()
19030            .await
19031            .unwrap()
19032            .build()
19033            .unwrap();
19034
19035        let available = agent.get_available_tool_ids().await.unwrap();
19036        assert!(available.is_empty());
19037    }
19038
19039    #[tokio::test]
19040    async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_empty() {
19041        let mock = mock_with_response("hello");
19042        let yaml = r#"
19043name: EmptySpawnerNoGrantAgent
19044system_prompt: "You manage agents."
19045tools: []
19046spawner:
19047  max_agents: 2
19048"#;
19049        let agent = AgentBuilder::from_yaml(yaml)
19050            .unwrap()
19051            .llm(Arc::new(mock))
19052            .auto_configure_features()
19053            .unwrap()
19054            .auto_configure_spawner()
19055            .await
19056            .unwrap()
19057            .build()
19058            .unwrap();
19059
19060        let available = agent.get_available_tool_ids().await.unwrap();
19061        assert!(available.is_empty());
19062    }
19063
19064    #[tokio::test]
19065    async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_empty() {
19066        let mock = mock_with_response("hello");
19067        let yaml = r#"
19068name: ManagementGrantAgent
19069system_prompt: "You manage agents."
19070tools: []
19071spawner:
19072  management_tools: true
19073"#;
19074        let agent = AgentBuilder::from_yaml(yaml)
19075            .unwrap()
19076            .llm(Arc::new(mock))
19077            .auto_configure_features()
19078            .unwrap()
19079            .auto_configure_spawner()
19080            .await
19081            .unwrap()
19082            .build()
19083            .unwrap();
19084
19085        let available = agent.get_available_tool_ids().await.unwrap();
19086        assert_eq!(available.len(), 4);
19087        assert!(available.contains(&"spawn_agent".to_string()));
19088        assert!(available.contains(&"send_agent_message".to_string()));
19089        assert!(available.contains(&"list_agents".to_string()));
19090        assert!(available.contains(&"remove_agent".to_string()));
19091    }
19092
19093    #[tokio::test]
19094    async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_omitted() {
19095        let mock = mock_with_response("hello");
19096        let yaml = r#"
19097name: ManagementOmittedToolsGrantAgent
19098system_prompt: "You manage agents."
19099spawner:
19100  management_tools: true
19101"#;
19102        let agent = AgentBuilder::from_yaml(yaml)
19103            .unwrap()
19104            .llm(Arc::new(mock))
19105            .auto_configure_features()
19106            .unwrap()
19107            .auto_configure_spawner()
19108            .await
19109            .unwrap()
19110            .build()
19111            .unwrap();
19112
19113        let available = agent.get_available_tool_ids().await.unwrap();
19114        assert_eq!(available.len(), 4);
19115        assert!(available.contains(&"spawn_agent".to_string()));
19116        assert!(available.contains(&"send_agent_message".to_string()));
19117        assert!(available.contains(&"list_agents".to_string()));
19118        assert!(available.contains(&"remove_agent".to_string()));
19119    }
19120
19121    #[tokio::test]
19122    async fn test_management_tools_selected_grants_only_selected_tools() {
19123        let mock = mock_with_response("hello");
19124        let yaml = r#"
19125name: ManagementSelectedGrantAgent
19126system_prompt: "You manage agents."
19127tools: []
19128spawner:
19129  management_tools:
19130    - spawn_agent
19131    - send_agent_message
19132    - list_agents
19133"#;
19134        let agent = AgentBuilder::from_yaml(yaml)
19135            .unwrap()
19136            .llm(Arc::new(mock))
19137            .auto_configure_features()
19138            .unwrap()
19139            .auto_configure_spawner()
19140            .await
19141            .unwrap()
19142            .build()
19143            .unwrap();
19144
19145        let available = agent.get_available_tool_ids().await.unwrap();
19146        assert_eq!(available.len(), 3);
19147        assert!(available.contains(&"spawn_agent".to_string()));
19148        assert!(available.contains(&"send_agent_message".to_string()));
19149        assert!(available.contains(&"list_agents".to_string()));
19150        assert!(!available.contains(&"remove_agent".to_string()));
19151    }
19152
19153    #[tokio::test]
19154    async fn test_orchestration_tools_flag_grants_tools_when_top_level_tools_empty() {
19155        let mock = mock_with_response("hello");
19156        let yaml = r#"
19157name: OrchestrationGrantAgent
19158system_prompt: "You coordinate agents."
19159llms:
19160  default:
19161    provider: openai
19162    model: gpt-4
19163  router:
19164    provider: openai
19165    model: gpt-4
19166llm:
19167  default: default
19168  router: router
19169tools: []
19170spawner:
19171  orchestration_tools: true
19172"#;
19173        let agent = AgentBuilder::from_yaml(yaml)
19174            .unwrap()
19175            .llm(Arc::new(mock))
19176            .auto_configure_features()
19177            .unwrap()
19178            .auto_configure_spawner()
19179            .await
19180            .unwrap()
19181            .build()
19182            .unwrap();
19183
19184        let available = agent.get_available_tool_ids().await.unwrap();
19185        assert_eq!(available.len(), 5);
19186        assert!(available.contains(&"route_to_agent".to_string()));
19187        assert!(available.contains(&"pipeline_process".to_string()));
19188        assert!(available.contains(&"concurrent_ask".to_string()));
19189        assert!(available.contains(&"group_discussion".to_string()));
19190        assert!(available.contains(&"handoff_conversation".to_string()));
19191    }
19192
19193    #[tokio::test]
19194    async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_empty() {
19195        let mock = mock_with_response("hello");
19196        let yaml = r#"
19197name: PersonaGrantAgent
19198system_prompt: "You can evolve persona."
19199llm:
19200  provider: openai
19201  model: gpt-4
19202tools: []
19203persona:
19204  identity:
19205    name: "Guide"
19206    role: "Helper"
19207  evolution:
19208    enabled: true
19209    allow_llm_evolve: true
19210    mutable_fields:
19211      - traits.personality
19212"#;
19213        let agent = AgentBuilder::from_yaml(yaml)
19214            .unwrap()
19215            .llm(Arc::new(mock))
19216            .build()
19217            .unwrap();
19218
19219        let available = agent.get_available_tool_ids().await.unwrap();
19220        assert_eq!(available, vec!["persona_evolve".to_string()]);
19221    }
19222
19223    #[tokio::test]
19224    async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_omitted() {
19225        let mock = mock_with_response("hello");
19226        let yaml = r#"
19227name: PersonaOmittedToolsGrantAgent
19228system_prompt: "You can evolve persona."
19229llm:
19230  provider: openai
19231  model: gpt-4
19232persona:
19233  identity:
19234    name: "Guide"
19235    role: "Helper"
19236  evolution:
19237    enabled: true
19238    allow_llm_evolve: true
19239    mutable_fields:
19240      - traits.personality
19241"#;
19242        let agent = AgentBuilder::from_yaml(yaml)
19243            .unwrap()
19244            .llm(Arc::new(mock))
19245            .build()
19246            .unwrap();
19247
19248        let available = agent.get_available_tool_ids().await.unwrap();
19249        assert_eq!(available, vec!["persona_evolve".to_string()]);
19250    }
19251
19252    #[tokio::test]
19253    async fn test_omitted_yaml_tools_exposes_no_tools() {
19254        let mock = mock_with_response("hello");
19255        let yaml = r#"
19256name: NoToolsAgent
19257system_prompt: "You are helpful."
19258"#;
19259        let agent = AgentBuilder::from_yaml(yaml)
19260            .unwrap()
19261            .llm(Arc::new(mock))
19262            .auto_configure_features()
19263            .unwrap()
19264            .build()
19265            .unwrap();
19266
19267        let available = agent.get_available_tool_ids().await.unwrap();
19268        assert!(available.is_empty());
19269    }
19270
19271    #[tokio::test]
19272    async fn runtime_scope_cannot_widen_omitted_or_empty_yaml_grants() {
19273        for tools in ["", "tools: []"] {
19274            let yaml = format!(
19275                r#"
19276name: RuntimeScopeNoGrantAgent
19277system_prompt: "No ordinary tools are granted."
19278{tools}
19279"#
19280            );
19281            let agent = AgentBuilder::from_yaml(&yaml)
19282                .unwrap()
19283                .llm(Arc::new(mock_with_response("done")))
19284                .auto_configure_features()
19285                .unwrap()
19286                .build()
19287                .unwrap();
19288
19289            agent
19290                .runtime_control()
19291                .set_tool_scope(vec!["calculator".to_string()]);
19292
19293            assert!(agent.get_available_tool_ids().await.unwrap().is_empty());
19294        }
19295    }
19296
19297    #[tokio::test]
19298    async fn runtime_scope_widening_attempt_keeps_only_declared_tools() {
19299        let yaml = r#"
19300name: RuntimeScopeWideningAgent
19301system_prompt: "Runtime scope cannot add authority."
19302tools: [calculator]
19303"#;
19304        let agent = AgentBuilder::from_yaml(yaml)
19305            .unwrap()
19306            .llm(Arc::new(mock_with_response("done")))
19307            .auto_configure_features()
19308            .unwrap()
19309            .build()
19310            .unwrap();
19311
19312        agent
19313            .runtime_control()
19314            .set_tool_scope(vec!["calculator".to_string(), "datetime".to_string()]);
19315
19316        assert_eq!(
19317            agent.get_available_tool_ids().await.unwrap(),
19318            vec!["calculator".to_string()]
19319        );
19320    }
19321
19322    #[tokio::test]
19323    async fn runtime_scope_is_canonical_unique_ordered_and_clear_restores_declared_grant() {
19324        let yaml = r#"
19325name: RuntimeScopeIntersectionAgent
19326system_prompt: "Use only declared tools."
19327tools: [calculator, datetime]
19328"#;
19329        let agent = AgentBuilder::from_yaml(yaml)
19330            .unwrap()
19331            .llm(Arc::new(mock_with_response("done")))
19332            .auto_configure_features()
19333            .unwrap()
19334            .build()
19335            .unwrap();
19336        let mut aliases = ai_agents_tools::ToolAliases::default();
19337        aliases
19338            .names
19339            .insert("en".to_string(), "calculate_alias".to_string());
19340        agent.tools.set_tool_aliases("calculator", aliases);
19341        let control = agent.runtime_control();
19342
19343        control.set_tool_scope(vec![
19344            "datetime".to_string(),
19345            "calculate_alias".to_string(),
19346            "calculator".to_string(),
19347            "unknown".to_string(),
19348            "datetime".to_string(),
19349        ]);
19350        assert_eq!(
19351            agent.get_available_tool_ids().await.unwrap(),
19352            vec!["calculator".to_string(), "datetime".to_string()]
19353        );
19354
19355        control.set_tool_scope(vec!["datetime".to_string()]);
19356        assert_eq!(
19357            agent.get_available_tool_ids().await.unwrap(),
19358            vec!["datetime".to_string()]
19359        );
19360
19361        control.clear_tool_scope_override();
19362        assert_eq!(
19363            agent.get_available_tool_ids().await.unwrap(),
19364            vec!["calculator".to_string(), "datetime".to_string()]
19365        );
19366    }
19367
19368    #[tokio::test]
19369    async fn runtime_scope_preserves_programmatic_registration_as_declared_grant() {
19370        let agent = AgentBuilder::new()
19371            .system_prompt("Use registered tools.")
19372            .llm(Arc::new(mock_with_response("done")))
19373            .tool(Arc::new(ContextEchoTool))
19374            .tool(Arc::new(SlowTool))
19375            .build()
19376            .unwrap();
19377
19378        agent.runtime_control().set_tool_scope(vec![
19379            "Context Echo".to_string(),
19380            "context_echo".to_string(),
19381            "unknown".to_string(),
19382        ]);
19383
19384        assert_eq!(
19385            agent.get_available_tool_ids().await.unwrap(),
19386            vec!["context_echo".to_string()]
19387        );
19388    }
19389
19390    #[tokio::test]
19391    async fn nested_state_scopes_intersect_every_ancestor_with_aliases() {
19392        let yaml = r#"
19393name: NestedStateScopeAgent
19394system_prompt: "Honor every state scope."
19395tools: [calculator, datetime, echo]
19396states:
19397  initial: root
19398  states:
19399    root:
19400      tools: [calculate_alias, datetime]
19401      initial: middle
19402      states:
19403        middle:
19404          initial: leaf
19405          states:
19406            leaf:
19407              tools: [datetime_alias, echo]
19408"#;
19409        let agent = AgentBuilder::from_yaml(yaml)
19410            .unwrap()
19411            .llm(Arc::new(mock_with_response("done")))
19412            .auto_configure_features()
19413            .unwrap()
19414            .build()
19415            .unwrap();
19416        let mut calculator_aliases = ai_agents_tools::ToolAliases::default();
19417        calculator_aliases
19418            .names
19419            .insert("en".to_string(), "calculate_alias".to_string());
19420        agent
19421            .tools
19422            .set_tool_aliases("calculator", calculator_aliases);
19423        let mut datetime_aliases = ai_agents_tools::ToolAliases::default();
19424        datetime_aliases
19425            .names
19426            .insert("en".to_string(), "datetime_alias".to_string());
19427        agent.tools.set_tool_aliases("datetime", datetime_aliases);
19428        agent.runtime_control().set_tool_scope(vec![
19429            "unknown".to_string(),
19430            "datetime_alias".to_string(),
19431            "calculate_alias".to_string(),
19432            "datetime".to_string(),
19433        ]);
19434
19435        assert_eq!(agent.current_state().as_deref(), Some("root.middle.leaf"));
19436        assert_eq!(
19437            agent.get_available_tool_ids().await.unwrap(),
19438            vec!["datetime".to_string()]
19439        );
19440    }
19441
19442    #[tokio::test]
19443    async fn ancestor_empty_state_scope_denies_omitted_descendants() {
19444        let yaml = r#"
19445name: NestedEmptyStateScopeAgent
19446system_prompt: "An empty ancestor scope denies all tools."
19447tools: [calculator]
19448states:
19449  initial: root
19450  states:
19451    root:
19452      tools: []
19453      initial: middle
19454      states:
19455        middle:
19456          initial: leaf
19457          states:
19458            leaf: {}
19459"#;
19460        let agent = AgentBuilder::from_yaml(yaml)
19461            .unwrap()
19462            .llm(Arc::new(mock_with_response("done")))
19463            .auto_configure_features()
19464            .unwrap()
19465            .build()
19466            .unwrap();
19467
19468        assert!(agent.get_available_tool_ids().await.unwrap().is_empty());
19469    }
19470
19471    #[tokio::test]
19472    async fn state_change_during_approval_invalidates_the_reviewed_authority() {
19473        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
19474        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
19475        let entered = Arc::new(tokio::sync::Barrier::new(2));
19476        let release = Arc::new(tokio::sync::Notify::new());
19477        let handler = Arc::new(BlockingApprovalHandler {
19478            entered: Arc::clone(&entered),
19479            release: Arc::clone(&release),
19480            result: ApprovalResult::Approved,
19481        });
19482        let yaml = r#"
19483name: ApprovalStateGenerationAgent
19484system_prompt: "State authority may change during approval."
19485tools: [locked_write]
19486states:
19487  initial: first
19488  states:
19489    first:
19490      tools: [locked_write]
19491    second:
19492      tools: [locked_write]
19493"#;
19494        let agent = Arc::new(
19495            AgentBuilder::from_yaml(yaml)
19496                .unwrap()
19497                .llm(Arc::new(mock_with_response("done")))
19498                .tool(Arc::new(LockedWriteTool {
19499                    active: Arc::clone(&active),
19500                    max_active: Arc::clone(&max_active),
19501                }))
19502                .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
19503                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
19504                .approval_handler(handler)
19505                .build()
19506                .unwrap(),
19507        );
19508        let running = Arc::clone(&agent);
19509        let call = tokio::spawn(async move {
19510            running
19511                .invoke_tool(ToolExecutionRequest::new(
19512                    "approval-state-generation",
19513                    "locked_write",
19514                    serde_json::json!({"path": "./state-generation.txt"}),
19515                    ToolCallSource::Manual,
19516                ))
19517                .await
19518                .unwrap()
19519        });
19520
19521        entered.wait().await;
19522        agent.transition_to("second").await.unwrap();
19523        release.notify_one();
19524        let record = call.await.unwrap();
19525
19526        assert!(!record.executed);
19527        assert!(record.output.contains("Approval became stale"));
19528        assert_eq!(max_active.load(Ordering::SeqCst), 0);
19529    }
19530
19531    #[tokio::test]
19532    async fn state_change_while_waiting_for_resource_lock_fails_final_admission() {
19533        let holder_gate = PathMutationGate::new();
19534        let waiter_gate = PathMutationGate::new();
19535        let yaml = r#"
19536name: LockedStateGenerationAgent
19537system_prompt: "State authority must remain stable through admission."
19538tools: [state_lock_holder, state_lock_waiter]
19539states:
19540  initial: first
19541  states:
19542    first:
19543      tools: [state_lock_holder, state_lock_waiter]
19544    second:
19545      tools: [state_lock_holder, state_lock_waiter]
19546"#;
19547        let agent = Arc::new(
19548            AgentBuilder::from_yaml(yaml)
19549                .unwrap()
19550                .llm(Arc::new(mock_with_response("done")))
19551                .tool(Arc::new(BlockingPathMutationTool {
19552                    id: "state_lock_holder",
19553                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
19554                    gate: holder_gate.clone(),
19555                }))
19556                .tool(Arc::new(BlockingPathMutationTool {
19557                    id: "state_lock_waiter",
19558                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
19559                    gate: waiter_gate.clone(),
19560                }))
19561                .build()
19562                .unwrap(),
19563        );
19564        let holder_call = {
19565            let agent = Arc::clone(&agent);
19566            tokio::spawn(async move {
19567                agent
19568                    .invoke_tool(ToolExecutionRequest::new(
19569                        "state-lock-holder",
19570                        "state_lock_holder",
19571                        serde_json::json!({"path": "./shared-state-path.txt"}),
19572                        ToolCallSource::Manual,
19573                    ))
19574                    .await
19575                    .unwrap()
19576            })
19577        };
19578        holder_gate.wait_until_entered().await;
19579        let waiter_call = {
19580            let agent = Arc::clone(&agent);
19581            tokio::spawn(async move {
19582                agent
19583                    .invoke_tool(ToolExecutionRequest::new(
19584                        "state-lock-waiter",
19585                        "state_lock_waiter",
19586                        serde_json::json!({"path": "./shared-state-path.txt"}),
19587                        ToolCallSource::Manual,
19588                    ))
19589                    .await
19590                    .unwrap()
19591            })
19592        };
19593
19594        wait_for_resource_lock_strong_count(&agent.resource_locks, 2).await;
19595        agent.transition_to("second").await.unwrap();
19596        holder_gate.release();
19597        let holder_record = holder_call.await.unwrap();
19598        let waiter_record = waiter_call.await.unwrap();
19599
19600        assert!(holder_record.success);
19601        assert!(!waiter_record.executed);
19602        assert!(
19603            waiter_record
19604                .output
19605                .contains("state scope changed before admission")
19606        );
19607        assert!(!waiter_gate.entered.load(Ordering::SeqCst));
19608    }
19609
19610    #[tokio::test]
19611    async fn test_state_tools_cannot_widen_top_level_grant() {
19612        let mock = mock_with_response("hello");
19613        let yaml = r#"
19614name: NarrowToolsAgent
19615system_prompt: "You are helpful."
19616tools:
19617  - calculator
19618states:
19619  initial: current
19620  states:
19621    current:
19622      tools: [datetime]
19623"#;
19624        let agent = AgentBuilder::from_yaml(yaml)
19625            .unwrap()
19626            .llm(Arc::new(mock))
19627            .auto_configure_features()
19628            .unwrap()
19629            .build()
19630            .unwrap();
19631
19632        let available = agent.get_available_tool_ids().await.unwrap();
19633        assert!(available.is_empty());
19634    }
19635
19636    // Tool execution in chat flow
19637    #[tokio::test]
19638    async fn test_integration_tool_execution() {
19639        // Mock LLM that returns a tool call then a final answer
19640        let mock = mock_with_responses(vec![
19641            // First response: tool call
19642            r#"I'll calculate that for you.
19643[TOOL_CALL: {"name": "calculator", "arguments": {"expression": "2+2"}}]"#,
19644            // After tool result: final answer
19645            "The answer is 4.",
19646        ]);
19647        let mut tools = ai_agents_tools::ToolRegistry::new();
19648        tools
19649            .register(Arc::new(ai_agents_tools::CalculatorTool))
19650            .unwrap();
19651
19652        let agent = AgentBuilder::new()
19653            .system_prompt("You are a calculator assistant.")
19654            .llm(Arc::new(mock))
19655            .tools(tools)
19656            .build()
19657            .unwrap();
19658
19659        let response = agent.chat("What is 2+2?").await.unwrap();
19660        // The agent should eventually produce a response
19661        assert!(!response.content.is_empty());
19662    }
19663
19664    #[tokio::test]
19665    async fn test_tool_hitl_rejection_finalizes_blocking_turn() {
19666        let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
19667        let hooks = Arc::new(ResponseCountingHooks {
19668            responses: Arc::clone(&responses),
19669        });
19670        let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
19671        let yaml = r#"
19672name: ToolRejectAgent
19673system_prompt: "You use tools when requested."
19674tools:
19675  - echo
19676hitl:
19677  tools:
19678    echo:
19679      require_approval: true
19680      approval_message: "Approve echo?"
19681"#;
19682        let agent = AgentBuilder::from_yaml(yaml)
19683            .unwrap()
19684            .llm(Arc::new(mock))
19685            .auto_configure_features()
19686            .unwrap()
19687            .hooks(hooks)
19688            .build()
19689            .unwrap();
19690
19691        let response = agent.chat("echo hello").await.unwrap();
19692
19693        assert!(
19694            response.content.contains("Operation cancelled"),
19695            "unexpected response: {}",
19696            response.content
19697        );
19698        assert_eq!(responses.load(Ordering::SeqCst), 1);
19699        let messages = agent.memory.get_messages(None).await.unwrap();
19700        assert_eq!(messages.len(), 3);
19701        assert_eq!(messages[0].content, "echo hello");
19702        assert!(messages[1].content.contains("\"tool\":\"echo\""));
19703        assert!(messages[2].content.contains("rejected by the approver"));
19704    }
19705
19706    #[tokio::test]
19707    async fn test_tool_hitl_rejection_finalizes_streaming_turn() {
19708        use futures::StreamExt;
19709
19710        let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
19711        let hooks = Arc::new(ResponseCountingHooks {
19712            responses: Arc::clone(&responses),
19713        });
19714        let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
19715        let yaml = r#"
19716name: ToolRejectStreamingAgent
19717system_prompt: "You use tools when requested."
19718tools:
19719  - echo
19720streaming:
19721  enabled: true
19722hitl:
19723  tools:
19724    echo:
19725      require_approval: true
19726      approval_message: "Approve echo?"
19727"#;
19728        let agent = AgentBuilder::from_yaml(yaml)
19729            .unwrap()
19730            .llm(Arc::new(mock))
19731            .auto_configure_features()
19732            .unwrap()
19733            .hooks(hooks)
19734            .build()
19735            .unwrap();
19736
19737        let mut stream = agent.chat_stream("echo hello").await.unwrap();
19738        let mut terminal_error = String::new();
19739        let mut done = false;
19740        while let Some(chunk) = stream.next().await {
19741            match chunk {
19742                StreamChunk::Error { message } => terminal_error = message,
19743                StreamChunk::Done {} => {
19744                    done = true;
19745                    break;
19746                }
19747                _ => {}
19748            }
19749        }
19750
19751        assert!(done);
19752        assert!(
19753            terminal_error.contains("Operation cancelled"),
19754            "unexpected terminal error: {}",
19755            terminal_error
19756        );
19757        assert_eq!(responses.load(Ordering::SeqCst), 1);
19758        let messages = agent.memory.get_messages(None).await.unwrap();
19759        assert_eq!(messages.len(), 3);
19760        assert_eq!(messages[0].content, "echo hello");
19761        assert!(messages[1].content.contains("\"tool\":\"echo\""));
19762        assert!(messages[2].content.contains("rejected by the approver"));
19763    }
19764
19765    #[tokio::test]
19766    async fn tool_hitl_rejection_preserves_legacy_error_but_finalizes_event_stream() {
19767        let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
19768        let yaml = r#"
19769name: ToolRejectEventAgent
19770system_prompt: "You use tools when requested."
19771tools:
19772  - echo
19773streaming:
19774  enabled: true
19775hitl:
19776  tools:
19777    echo:
19778      require_approval: true
19779      approval_message: "Approve echo?"
19780"#;
19781        let agent = AgentBuilder::from_yaml(yaml)
19782            .unwrap()
19783            .llm(Arc::new(mock))
19784            .auto_configure_features()
19785            .unwrap()
19786            .build()
19787            .unwrap();
19788
19789        let mut stream = agent.chat_stream_events("echo hello").await.unwrap();
19790        let mut error_seen = false;
19791        let mut final_response = None;
19792        while let Some(event) = stream.next().await {
19793            match event {
19794                AgentStreamEvent::Chunk(StreamChunk::Error { .. }) => error_seen = true,
19795                AgentStreamEvent::Final(response) => final_response = Some(response),
19796                AgentStreamEvent::Chunk(_) => {}
19797            }
19798        }
19799
19800        assert!(!error_seen);
19801        assert!(
19802            final_response
19803                .is_some_and(|response| { response.content.contains("Operation cancelled") })
19804        );
19805    }
19806
19807    #[tokio::test]
19808    async fn test_pre_response_guard_transition_skips_old_state_llm() {
19809        let mock = mock_with_response("Billing state response");
19810        let call_counter = mock.clone();
19811        let yaml = r#"
19812name: OptimizedStateAgent
19813system_prompt: "You route before answering."
19814runtime:
19815  optimization:
19816    enabled: true
19817    pre_response_deterministic_transitions: true
19818states:
19819  initial: greeting
19820  states:
19821    greeting:
19822      prompt: "Old state prompt that should be skipped."
19823      transitions:
19824        - to: billing
19825          guard:
19826            context:
19827              topic:
19828                eq: billing
19829          timing: pre_response
19830    billing:
19831      prompt: "Answer from the billing state."
19832"#;
19833        let agent = AgentBuilder::from_yaml(yaml)
19834            .unwrap()
19835            .llm(Arc::new(mock))
19836            .build()
19837            .unwrap();
19838        agent
19839            .set_context("topic", serde_json::json!("billing"))
19840            .unwrap();
19841
19842        let response = agent.chat("I need billing help").await.unwrap();
19843
19844        assert_eq!(agent.current_state().as_deref(), Some("billing"));
19845        assert_eq!(response.content, "Billing state response");
19846        assert_eq!(call_counter.call_count(), 1);
19847        assert_eq!(agent.actor_facts().len(), 0);
19848    }
19849
19850    #[tokio::test]
19851    async fn test_set_context_supports_dotted_paths_for_pre_response_guards() {
19852        let mock = mock_with_response("Billing state response");
19853        let call_counter = mock.clone();
19854        let yaml = r#"
19855name: OptimizedStateAgent
19856system_prompt: "You route before answering."
19857runtime:
19858  optimization:
19859    enabled: true
19860    pre_response_deterministic_transitions: true
19861context:
19862  request:
19863    type: runtime
19864    default:
19865      topic: general
19866states:
19867  initial: greeting
19868  states:
19869    greeting:
19870      prompt: "Old state prompt that should be skipped."
19871      transitions:
19872        - to: billing
19873          guard:
19874            context:
19875              request.topic:
19876                eq: billing
19877          timing: pre_response
19878    billing:
19879      prompt: "Answer from the billing state."
19880"#;
19881        let agent = AgentBuilder::from_yaml(yaml)
19882            .unwrap()
19883            .llm(Arc::new(mock))
19884            .build()
19885            .unwrap();
19886        agent
19887            .set_context("request.topic", serde_json::json!("billing"))
19888            .unwrap();
19889
19890        let response = agent.chat("I need billing help").await.unwrap();
19891
19892        assert_eq!(agent.current_state().as_deref(), Some("billing"));
19893        assert_eq!(response.content, "Billing state response");
19894        assert_eq!(call_counter.call_count(), 1);
19895        assert_eq!(
19896            agent.get_context().get("request"),
19897            Some(&serde_json::json!({"topic": "billing"}))
19898        );
19899    }
19900
19901    #[tokio::test]
19902    async fn test_pre_response_rejection_does_not_commit_staged_context_or_user() {
19903        let mock = mock_with_response("billing");
19904        let yaml = r#"
19905name: OptimizedStateAgent
19906system_prompt: "You route before answering."
19907runtime:
19908  optimization:
19909    enabled: true
19910    pre_response_deterministic_transitions: true
19911hitl:
19912  states:
19913    billing:
19914      on_enter: require_approval
19915      approval_message: "Approve billing route?"
19916states:
19917  initial: greeting
19918  states:
19919    greeting:
19920      prompt: "Old state prompt."
19921      extract:
19922        - key: topic
19923          description: "Support topic"
19924      transitions:
19925        - to: billing
19926          guard:
19927            context:
19928              topic:
19929                eq: billing
19930          timing: pre_response
19931          run_extractors: true
19932    billing:
19933      prompt: "Billing state."
19934"#;
19935        let agent = AgentBuilder::from_yaml(yaml)
19936            .unwrap()
19937            .llm(Arc::new(mock))
19938            .build()
19939            .unwrap();
19940
19941        let response = agent
19942            .try_pre_response_transition("billing please")
19943            .await
19944            .unwrap();
19945
19946        assert!(response.is_none());
19947        assert_eq!(agent.current_state().as_deref(), Some("greeting"));
19948        assert!(!agent.get_context().contains_key("topic"));
19949        assert_eq!(agent.memory.get_messages(None).await.unwrap().len(), 0);
19950    }
19951
19952    #[tokio::test]
19953    async fn test_pre_response_extractor_commits_context_on_winning_path() {
19954        let mock = mock_with_responses(vec!["billing", "Billing response"]);
19955        let yaml = r#"
19956name: OptimizedStateAgent
19957system_prompt: "You route before answering."
19958runtime:
19959  optimization:
19960    enabled: true
19961    pre_response_deterministic_transitions: true
19962states:
19963  initial: greeting
19964  states:
19965    greeting:
19966      prompt: "Old state prompt."
19967      extract:
19968        - key: topic
19969          description: "Support topic"
19970      transitions:
19971        - to: billing
19972          guard:
19973            context:
19974              topic:
19975                eq: billing
19976          timing: pre_response
19977          run_extractors: true
19978    billing:
19979      prompt: "Billing state."
19980"#;
19981        let agent = AgentBuilder::from_yaml(yaml)
19982            .unwrap()
19983            .llm(Arc::new(mock))
19984            .build()
19985            .unwrap();
19986
19987        let response = agent.chat("billing please").await.unwrap();
19988
19989        assert_eq!(agent.current_state().as_deref(), Some("billing"));
19990        assert_eq!(response.content, "Billing response");
19991        assert_eq!(
19992            agent.get_context().get("topic"),
19993            Some(&serde_json::json!("billing"))
19994        );
19995    }
19996
19997    #[tokio::test]
19998    async fn test_pre_response_extractor_miss_does_not_mutate_context() {
19999        let mock = mock_with_response("__NONE__");
20000        let yaml = r#"
20001name: OptimizedStateAgent
20002system_prompt: "You route before answering."
20003runtime:
20004  optimization:
20005    enabled: true
20006    pre_response_deterministic_transitions: true
20007states:
20008  initial: greeting
20009  states:
20010    greeting:
20011      prompt: "Old state prompt."
20012      extract:
20013        - key: topic
20014          description: "Support topic"
20015      transitions:
20016        - to: billing
20017          guard:
20018            context:
20019              topic:
20020                eq: billing
20021          timing: pre_response
20022          run_extractors: true
20023    billing:
20024      prompt: "Billing state."
20025"#;
20026        let agent = AgentBuilder::from_yaml(yaml)
20027            .unwrap()
20028            .llm(Arc::new(mock))
20029            .build()
20030            .unwrap();
20031
20032        let response = agent.try_pre_response_transition("hello").await.unwrap();
20033
20034        assert!(response.is_none());
20035        assert_eq!(agent.current_state().as_deref(), Some("greeting"));
20036        assert!(!agent.get_context().contains_key("topic"));
20037    }
20038
20039    #[tokio::test]
20040    async fn test_default_guard_transition_stays_post_response() {
20041        let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
20042        let call_counter = mock.clone();
20043        let yaml = r#"
20044name: TimingAgent
20045system_prompt: "You route carefully."
20046runtime:
20047  optimization:
20048    enabled: true
20049    pre_response_deterministic_transitions: true
20050states:
20051  initial: greeting
20052  states:
20053    greeting:
20054      prompt: "Old state prompt."
20055      transitions:
20056        - to: billing
20057          guard:
20058            context:
20059              topic:
20060                eq: billing
20061    billing:
20062      prompt: "Billing state."
20063"#;
20064        let agent = AgentBuilder::from_yaml(yaml)
20065            .unwrap()
20066            .llm(Arc::new(mock))
20067            .build()
20068            .unwrap();
20069        agent
20070            .set_context("topic", serde_json::json!("billing"))
20071            .unwrap();
20072
20073        let response = agent.chat("billing please").await.unwrap();
20074
20075        assert_eq!(agent.current_state().as_deref(), Some("billing"));
20076        assert_eq!(response.content, "Billing response");
20077        assert_eq!(call_counter.call_count(), 2);
20078    }
20079
20080    #[tokio::test]
20081    async fn test_explicit_post_response_guard_transition_stays_post_response() {
20082        let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
20083        let call_counter = mock.clone();
20084        let yaml = r#"
20085name: TimingAgent
20086system_prompt: "You route carefully."
20087runtime:
20088  optimization:
20089    enabled: true
20090    pre_response_deterministic_transitions: true
20091states:
20092  initial: greeting
20093  states:
20094    greeting:
20095      prompt: "Old state prompt."
20096      transitions:
20097        - to: billing
20098          guard:
20099            context:
20100              topic:
20101                eq: billing
20102          timing: post_response
20103    billing:
20104      prompt: "Billing state."
20105"#;
20106        let agent = AgentBuilder::from_yaml(yaml)
20107            .unwrap()
20108            .llm(Arc::new(mock))
20109            .build()
20110            .unwrap();
20111        agent
20112            .set_context("topic", serde_json::json!("billing"))
20113            .unwrap();
20114
20115        let response = agent.chat("billing please").await.unwrap();
20116
20117        assert_eq!(agent.current_state().as_deref(), Some("billing"));
20118        assert_eq!(response.content, "Billing response");
20119        assert_eq!(call_counter.call_count(), 2);
20120    }
20121
20122    #[tokio::test]
20123    async fn test_pre_response_extractors_are_transition_scoped() {
20124        let mock = mock_with_responses(vec!["billing", "Billing response"]);
20125        let yaml = r#"
20126name: ScopedExtractorAgent
20127system_prompt: "You route carefully."
20128runtime:
20129  optimization:
20130    enabled: true
20131    pre_response_deterministic_transitions: true
20132states:
20133  initial: greeting
20134  states:
20135    greeting:
20136      prompt: "Old state prompt."
20137      extract:
20138        - key: topic
20139          description: "Support topic"
20140      transitions:
20141        - to: wrong
20142          guard:
20143            context:
20144              topic:
20145                eq: billing
20146          timing: pre_response
20147        - to: billing
20148          guard:
20149            context:
20150              topic:
20151                eq: billing
20152          timing: pre_response
20153          run_extractors: true
20154    wrong:
20155      prompt: "Wrong state."
20156    billing:
20157      prompt: "Billing state."
20158"#;
20159        let agent = AgentBuilder::from_yaml(yaml)
20160            .unwrap()
20161            .llm(Arc::new(mock))
20162            .build()
20163            .unwrap();
20164
20165        let response = agent.chat("billing please").await.unwrap();
20166
20167        assert_eq!(agent.current_state().as_deref(), Some("billing"));
20168        assert_eq!(response.content, "Billing response");
20169    }
20170
20171    #[tokio::test]
20172    async fn test_pre_response_resolved_intent_routes_early() {
20173        let mock = mock_with_response("Billing response");
20174        let yaml = r#"
20175name: IntentAgent
20176system_prompt: "You route carefully."
20177runtime:
20178  optimization:
20179    enabled: true
20180    pre_response_deterministic_transitions: true
20181states:
20182  initial: greeting
20183  states:
20184    greeting:
20185      prompt: "Old state prompt."
20186      transitions:
20187        - to: billing
20188          intent: billing
20189          timing: pre_response
20190    billing:
20191      prompt: "Billing state."
20192"#;
20193        let agent = AgentBuilder::from_yaml(yaml)
20194            .unwrap()
20195            .llm(Arc::new(mock))
20196            .build()
20197            .unwrap();
20198        agent
20199            .set_context("resolved_intent", serde_json::json!("billing"))
20200            .unwrap();
20201
20202        let response = agent
20203            .try_pre_response_transition("I need billing help")
20204            .await
20205            .unwrap()
20206            .unwrap();
20207
20208        assert_eq!(agent.current_state().as_deref(), Some("billing"));
20209        assert_eq!(response.content, "Billing response");
20210    }
20211
20212    #[tokio::test]
20213    async fn test_background_overflow_error_surfaces() {
20214        let mut config = RuntimeConfig::default();
20215        config.optimization.enabled = true;
20216        config.optimization.post_turn.max_background_tasks = 1;
20217        config.optimization.post_turn.on_background_overflow = BackgroundOverflowPolicy::Error;
20218        let policy = crate::optimization::MaintenanceTaskPolicy {
20219            mode: MaintenanceMode::Background,
20220            await_before_next_turn: AwaitBeforeNextTurn::Always,
20221        };
20222        let agent = AgentBuilder::new()
20223            .system_prompt("You are helpful.")
20224            .llm(Arc::new(mock_with_response("ok")))
20225            .build()
20226            .unwrap()
20227            .with_runtime_config(config);
20228        agent
20229            .background_maintenance
20230            .spawn(None, async { std::future::pending::<Result<()>>().await })
20231            .unwrap();
20232
20233        let result = agent
20234            .spawn_or_handle_background(None, async { Ok(()) }, "facts", &policy)
20235            .await;
20236
20237        assert!(result.is_err());
20238    }
20239
20240    #[tokio::test]
20241    async fn test_speculative_reasoning_low_cap_uses_serial_reasoning() {
20242        let default_mock = mock_with_response("Plain draft response");
20243        let router_mock = mock_with_response("cot");
20244        let router_counter = router_mock.clone();
20245        let yaml = r#"
20246name: ReasoningReservationAgent
20247system_prompt: "You answer plainly unless reasoning wins."
20248llm:
20249  default: default
20250  router: router
20251observability:
20252  enabled: true
20253  export:
20254    write_raw_events: true
20255reasoning:
20256  mode: auto
20257  judge_llm: router
20258runtime:
20259  optimization:
20260    enabled: true
20261    max_speculative_llm_calls_per_turn: 1
20262    speculative_reasoning_auto: true
20263    max_parallel_runtime_tasks: 2
20264"#;
20265        let agent = AgentBuilder::from_yaml(yaml)
20266            .unwrap()
20267            .llm_alias("default", Arc::new(default_mock))
20268            .llm_alias("router", Arc::new(router_mock))
20269            .build()
20270            .unwrap();
20271
20272        let response = agent.chat("hello").await.unwrap();
20273
20274        assert_eq!(response.content, "Plain draft response");
20275        assert_eq!(router_counter.call_count(), 1);
20276        let events = agent.observability().unwrap().raw_events();
20277        assert!(!events.iter().any(|event| {
20278            event.dimensions.get("commit_behavior") == Some(&"reasoning_decision".to_string())
20279        }));
20280    }
20281
20282    #[tokio::test]
20283    async fn test_forced_reasoning_skips_plain_speculative_draft() {
20284        let mock = mock_with_response("Reasoned response");
20285        let yaml = r#"
20286name: ForcedReasoningAgent
20287system_prompt: "You reason before answering."
20288observability:
20289  enabled: true
20290  export:
20291    write_raw_events: true
20292reasoning:
20293  mode: cot
20294runtime:
20295  optimization:
20296    enabled: true
20297    max_speculative_llm_calls_per_turn: 2
20298    speculative_state_transitions: true
20299    max_parallel_runtime_tasks: 2
20300states:
20301  initial: triage
20302  states:
20303    triage:
20304      prompt: "Answer from triage."
20305      transitions:
20306        - to: billing
20307          guard:
20308            context:
20309              route:
20310                eq: billing
20311          timing: parallel
20312    billing:
20313      prompt: "Billing state."
20314"#;
20315        let agent = AgentBuilder::from_yaml(yaml)
20316            .unwrap()
20317            .llm(Arc::new(mock))
20318            .build()
20319            .unwrap();
20320
20321        let response = agent.chat("hello").await.unwrap();
20322
20323        assert_eq!(response.content, "Reasoned response");
20324        let events = agent.observability().unwrap().raw_events();
20325        assert!(
20326            !events
20327                .iter()
20328                .any(|event| event.dimensions.contains_key("branch_status"))
20329        );
20330    }
20331
20332    #[tokio::test]
20333    async fn test_speculative_skill_low_cap_uses_serial_skill_route() {
20334        let default_mock = mock_with_response("Skill committed response");
20335        let router_mock = mock_with_response("helper");
20336        let router_counter = router_mock.clone();
20337        let yaml = r#"
20338name: SkillReservationAgent
20339system_prompt: "Use skills when they match."
20340llm:
20341  default: default
20342  router: router
20343observability:
20344  enabled: true
20345  export:
20346    write_raw_events: true
20347runtime:
20348  optimization:
20349    enabled: true
20350    max_speculative_llm_calls_per_turn: 1
20351    speculative_skill_routing: true
20352    max_parallel_runtime_tasks: 2
20353skills:
20354  - id: helper
20355    description: "Answer helper requests"
20356    trigger: "User asks for helper"
20357    steps:
20358      - prompt: "Answer the helper request: {{ user_input }}"
20359"#;
20360        let agent = AgentBuilder::from_yaml(yaml)
20361            .unwrap()
20362            .llm_alias("default", Arc::new(default_mock))
20363            .llm_alias("router", Arc::new(router_mock))
20364            .build()
20365            .unwrap();
20366
20367        let response = agent.chat("please use helper").await.unwrap();
20368
20369        assert_eq!(response.content, "Skill committed response");
20370        assert_eq!(router_counter.call_count(), 1);
20371        let events = agent.observability().unwrap().raw_events();
20372        assert!(
20373            !events
20374                .iter()
20375                .any(|event| event.dimensions.contains_key("branch_status"))
20376        );
20377    }
20378
20379    #[tokio::test]
20380    async fn test_parallel_transition_low_cap_allows_deterministic_route() {
20381        let mock = mock_with_response("unused");
20382        let call_counter = mock.clone();
20383        let yaml = r#"
20384name: ParallelTransitionLowCapAgent
20385system_prompt: "Route before stale responses when safe."
20386runtime:
20387  optimization:
20388    enabled: true
20389    max_speculative_llm_calls_per_turn: 1
20390    speculative_state_transitions: true
20391    max_parallel_runtime_tasks: 2
20392states:
20393  initial: triage
20394  states:
20395    triage:
20396      prompt: "Triage state."
20397      transitions:
20398        - to: billing
20399          guard:
20400            context:
20401              route:
20402                eq: billing
20403          timing: parallel
20404    billing:
20405      prompt: "Billing state."
20406"#;
20407        let agent = AgentBuilder::from_yaml(yaml)
20408            .unwrap()
20409            .llm(Arc::new(mock))
20410            .build()
20411            .unwrap();
20412        agent
20413            .set_context("route", serde_json::json!("billing"))
20414            .unwrap();
20415        agent.update_active_turn_context("billing help", HashMap::new());
20416        assert!(
20417            agent.reserve_active_speculative_llm_call(
20418                RuntimeOptimizationKind::ParallelStateTransition
20419            )
20420        );
20421
20422        let selection = agent
20423            .select_parallel_transition_candidate("billing help")
20424            .await
20425            .unwrap();
20426        agent.end_root_turn();
20427
20428        match selection {
20429            ParallelTransitionSelection::Candidate(candidate) => {
20430                assert_eq!(candidate.target(), "billing");
20431            }
20432            ParallelTransitionSelection::NoMatch => panic!("deterministic route did not match"),
20433            ParallelTransitionSelection::ReservationExhausted => {
20434                panic!("deterministic route consumed LLM budget")
20435            }
20436        }
20437        assert_eq!(call_counter.call_count(), 0);
20438    }
20439
20440    #[tokio::test]
20441    async fn speculative_transition_drops_loser_before_state_actions() {
20442        let lock = Arc::new(tokio::sync::Mutex::new(()));
20443        let first_started = Arc::new(tokio::sync::Notify::new());
20444        let first_dropped = Arc::new(AtomicBool::new(false));
20445        let committed_after_drop = Arc::new(AtomicBool::new(false));
20446        let default = Arc::new(FirstCallLockingProvider {
20447            lock,
20448            first_started: Arc::clone(&first_started),
20449            first_dropped: Arc::clone(&first_dropped),
20450            committed_after_drop: Arc::clone(&committed_after_drop),
20451            calls: AtomicU64::new(0),
20452        });
20453        let router = Arc::new(RoutingAfterProviderStart {
20454            provider_started: first_started,
20455        });
20456        let yaml = r#"
20457name: SpeculativeCancellationAgent
20458system_prompt: "Route before committed work."
20459llm:
20460  default: default
20461  router: router
20462runtime:
20463  optimization:
20464    enabled: true
20465    max_speculative_llm_calls_per_turn: 2
20466    speculative_state_transitions: true
20467    max_parallel_runtime_tasks: 2
20468states:
20469  initial: triage
20470  states:
20471    triage:
20472      prompt: "Triage state."
20473      transitions:
20474        - to: technical
20475          when: "The request needs technical support"
20476          timing: parallel
20477    technical:
20478      prompt: "Technical state."
20479      on_enter:
20480        - prompt: "Prepare technical context."
20481          llm: default
20482          store_as: preparation
20483"#;
20484        let agent = AgentBuilder::from_yaml(yaml)
20485            .unwrap()
20486            .llm_alias("default", default)
20487            .llm_alias("router", router)
20488            .build()
20489            .unwrap();
20490
20491        let response = tokio::time::timeout(
20492            std::time::Duration::from_secs(2),
20493            agent.chat("I cannot log in because of AUTH-17."),
20494        )
20495        .await
20496        .expect("committed work must not wait on the losing provider future")
20497        .unwrap();
20498
20499        assert_eq!(response.content, "Committed technical response.");
20500        assert_eq!(agent.current_state().as_deref(), Some("technical"));
20501        assert!(first_dropped.load(Ordering::SeqCst));
20502        assert!(committed_after_drop.load(Ordering::SeqCst));
20503    }
20504
20505    #[tokio::test]
20506    async fn buffered_transition_drops_stale_stream_before_redispatch() {
20507        use futures::StreamExt;
20508
20509        let lock = Arc::new(tokio::sync::Mutex::new(()));
20510        let stream_started = Arc::new(tokio::sync::Notify::new());
20511        let stream_dropped = Arc::new(AtomicBool::new(false));
20512        let committed_after_drop = Arc::new(AtomicBool::new(false));
20513        let default = Arc::new(BufferedLockingProvider {
20514            lock,
20515            stream_started: Arc::clone(&stream_started),
20516            stream_dropped: Arc::clone(&stream_dropped),
20517            committed_after_drop: Arc::clone(&committed_after_drop),
20518        });
20519        let router = Arc::new(RoutingAfterProviderStart {
20520            provider_started: stream_started,
20521        });
20522        let yaml = r#"
20523name: BufferedCancellationAgent
20524system_prompt: "Hide stale streamed output."
20525llm:
20526  default: default
20527  router: router
20528streaming:
20529  enabled: true
20530  buffer_size: 8
20531runtime:
20532  optimization:
20533    enabled: true
20534    max_speculative_llm_calls_per_turn: 2
20535    speculative_state_transitions: true
20536    streaming_policy: buffer_until_routing_done
20537    max_parallel_runtime_tasks: 2
20538states:
20539  initial: triage
20540  states:
20541    triage:
20542      prompt: "Triage state."
20543      transitions:
20544        - to: technical
20545          when: "The request needs technical support"
20546          timing: parallel
20547    technical:
20548      prompt: "Technical state."
20549"#;
20550        let agent = AgentBuilder::from_yaml(yaml)
20551            .unwrap()
20552            .llm_alias("default", default)
20553            .llm_alias("router", router)
20554            .build()
20555            .unwrap();
20556
20557        let content = tokio::time::timeout(std::time::Duration::from_secs(2), async {
20558            let mut stream = agent
20559                .chat_stream("AUTH-17 needs technical help.")
20560                .await
20561                .unwrap();
20562            let mut content = String::new();
20563            while let Some(chunk) = stream.next().await {
20564                match chunk {
20565                    StreamChunk::Content { text } => content.push_str(&text),
20566                    StreamChunk::Done {} => break,
20567                    StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
20568                    _ => {}
20569                }
20570            }
20571            content
20572        })
20573        .await
20574        .expect("redispatch must not wait on the stale streaming future");
20575
20576        assert_eq!(content, "Committed technical response.");
20577        assert_eq!(agent.current_state().as_deref(), Some("technical"));
20578        assert!(stream_dropped.load(Ordering::SeqCst));
20579        assert!(committed_after_drop.load(Ordering::SeqCst));
20580    }
20581
20582    #[tokio::test]
20583    async fn buffered_transition_drops_established_stream_before_redispatch() {
20584        use futures::StreamExt;
20585
20586        let stream_started = Arc::new(tokio::sync::Notify::new());
20587        let stream_dropped = Arc::new(AtomicBool::new(false));
20588        let stream_dropped_notify = Arc::new(tokio::sync::Notify::new());
20589        let committed_after_drop = Arc::new(AtomicBool::new(false));
20590        let default = Arc::new(EstablishedStreamProvider {
20591            stream_started: Arc::clone(&stream_started),
20592            stream_dropped: Arc::clone(&stream_dropped),
20593            stream_dropped_notify,
20594            committed_after_drop: Arc::clone(&committed_after_drop),
20595        });
20596        let router = Arc::new(RoutingAfterProviderStart {
20597            provider_started: stream_started,
20598        });
20599        let yaml = r#"
20600name: EstablishedStreamCancellationAgent
20601system_prompt: "Hide stale streamed output."
20602llm:
20603  default: default
20604  router: router
20605streaming:
20606  enabled: true
20607  buffer_size: 8
20608runtime:
20609  optimization:
20610    enabled: true
20611    max_speculative_llm_calls_per_turn: 2
20612    speculative_state_transitions: true
20613    streaming_policy: buffer_until_routing_done
20614    max_parallel_runtime_tasks: 2
20615states:
20616  initial: triage
20617  states:
20618    triage:
20619      prompt: "Triage state."
20620      transitions:
20621        - to: technical
20622          when: "The request needs technical support"
20623          timing: parallel
20624    technical:
20625      prompt: "Technical state."
20626"#;
20627        let agent = AgentBuilder::from_yaml(yaml)
20628            .unwrap()
20629            .llm_alias("default", default)
20630            .llm_alias("router", router)
20631            .build()
20632            .unwrap();
20633
20634        let content = tokio::time::timeout(std::time::Duration::from_secs(2), async {
20635            let mut stream = agent
20636                .chat_stream("AUTH-17 needs technical help.")
20637                .await
20638                .unwrap();
20639            let mut content = String::new();
20640            while let Some(chunk) = stream.next().await {
20641                match chunk {
20642                    StreamChunk::Content { text } => content.push_str(&text),
20643                    StreamChunk::Done {} => break,
20644                    StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
20645                    _ => {}
20646                }
20647            }
20648            content
20649        })
20650        .await
20651        .expect("redispatch must wait for the established stale stream to be dropped");
20652
20653        assert_eq!(content, "Committed technical response.");
20654        assert_eq!(agent.current_state().as_deref(), Some("technical"));
20655        assert!(stream_dropped.load(Ordering::SeqCst));
20656        assert!(committed_after_drop.load(Ordering::SeqCst));
20657    }
20658
20659    #[tokio::test]
20660    async fn test_buffered_streaming_transition_reservation_falls_back() {
20661        use futures::StreamExt;
20662
20663        let mock = mock_with_responses(vec![
20664            "Serial streaming response",
20665            "Serial streaming response",
20666        ]);
20667        let router_mock = mock_with_response("1");
20668        let router_counter = router_mock.clone();
20669        let yaml = r#"
20670name: BufferedReservationFallbackAgent
20671system_prompt: "Stream normally if speculative routing cannot be evaluated."
20672llm:
20673  default: default
20674  router: router
20675observability:
20676  enabled: true
20677  export:
20678    write_raw_events: true
20679streaming:
20680  enabled: true
20681  buffer_size: 8
20682runtime:
20683  optimization:
20684    enabled: true
20685    max_speculative_llm_calls_per_turn: 1
20686    speculative_state_transitions: true
20687    streaming_policy: buffer_until_routing_done
20688    max_parallel_runtime_tasks: 2
20689states:
20690  initial: triage
20691  states:
20692    triage:
20693      prompt: "Triage state."
20694      transitions:
20695        - to: billing
20696          guard:
20697            context:
20698              route:
20699                eq: billing
20700          when: "User asks about billing"
20701          timing: parallel
20702    billing:
20703      prompt: "Billing state."
20704"#;
20705        let agent = AgentBuilder::from_yaml(yaml)
20706            .unwrap()
20707            .llm_alias("default", Arc::new(mock))
20708            .llm_alias("router", Arc::new(router_mock))
20709            .build()
20710            .unwrap();
20711
20712        let mut stream = agent.chat_stream("hello").await.unwrap();
20713        let mut content = String::new();
20714        let mut error = None;
20715        while let Some(chunk) = stream.next().await {
20716            match chunk {
20717                StreamChunk::Content { text } => content.push_str(&text),
20718                StreamChunk::Error { message } => error = Some(message),
20719                StreamChunk::Done {} => break,
20720                _ => {}
20721            }
20722        }
20723
20724        assert_eq!(error, None);
20725        assert_eq!(content, "Serial streaming response");
20726        assert_eq!(router_counter.call_count(), 0);
20727        let events = agent.observability().unwrap().raw_events();
20728        assert!(events.iter().any(|event| {
20729            event.dimensions.get("branch_status") == Some(&"cancelled".to_string())
20730                && event.dimensions.get("commit_behavior")
20731                    == Some(&"transition_decision".to_string())
20732        }));
20733    }
20734
20735    #[tokio::test]
20736    async fn test_blocking_error_cleanup_resets_root_turn_for_next_chat() {
20737        let mut mock = mock_with_response("Recovered response");
20738        mock.set_error("boom");
20739        let mut handle = mock.clone();
20740        let agent = AgentBuilder::new()
20741            .system_prompt("You are helpful.")
20742            .llm(Arc::new(mock))
20743            .build()
20744            .unwrap();
20745
20746        assert!(agent.chat("first").await.is_err());
20747        handle.clear_error();
20748        let response = agent.chat("second").await.unwrap();
20749
20750        assert_eq!(response.content, "Recovered response");
20751        let messages = agent.memory.get_messages(None).await.unwrap();
20752        let user_count = messages
20753            .iter()
20754            .filter(|message| message.role == ai_agents_core::Role::User)
20755            .count();
20756        assert_eq!(user_count, 2);
20757    }
20758
20759    #[tokio::test]
20760    async fn test_streaming_error_cleanup_resets_root_turn_for_next_chat() {
20761        use futures::StreamExt;
20762
20763        let mut mock = mock_with_response("Recovered response");
20764        mock.set_error("stream boom");
20765        let mut handle = mock.clone();
20766        let agent = AgentBuilder::new()
20767            .system_prompt("You are helpful.")
20768            .llm(Arc::new(mock))
20769            .build()
20770            .unwrap();
20771
20772        let mut stream = agent.chat_stream("first").await.unwrap();
20773        let mut saw_error = false;
20774        while let Some(chunk) = stream.next().await {
20775            if matches!(chunk, StreamChunk::Error { .. }) {
20776                saw_error = true;
20777            }
20778        }
20779        assert!(saw_error);
20780
20781        handle.clear_error();
20782        let response = agent.chat("second").await.unwrap();
20783
20784        assert_eq!(response.content, "Recovered response");
20785        let messages = agent.memory.get_messages(None).await.unwrap();
20786        let user_count = messages
20787            .iter()
20788            .filter(|message| message.role == ai_agents_core::Role::User)
20789            .count();
20790        assert_eq!(user_count, 2);
20791    }
20792
20793    #[tokio::test]
20794    async fn test_buffered_streaming_route_miss_releases_buffer_limit() {
20795        use futures::StreamExt;
20796
20797        let mut mock = mock_with_response("one two three");
20798        mock.set_latency(10);
20799        let yaml = r#"
20800name: BufferedMissAgent
20801system_prompt: "You stream safely."
20802llm:
20803  default: default
20804streaming:
20805  enabled: true
20806  buffer_size: 1
20807runtime:
20808  optimization:
20809    enabled: true
20810    max_speculative_llm_calls_per_turn: 2
20811    speculative_state_transitions: true
20812    streaming_policy: buffer_until_routing_done
20813    max_parallel_runtime_tasks: 2
20814states:
20815  initial: triage
20816  states:
20817    triage:
20818      prompt: "Answer from triage."
20819      transitions:
20820        - to: billing
20821          guard:
20822            context:
20823              route:
20824                eq: billing
20825          timing: parallel
20826    billing:
20827      prompt: "Billing state."
20828"#;
20829        let agent = AgentBuilder::from_yaml(yaml)
20830            .unwrap()
20831            .llm_alias("default", Arc::new(mock))
20832            .build()
20833            .unwrap();
20834
20835        let mut stream = agent.chat_stream("hello").await.unwrap();
20836        let mut content = String::new();
20837        let mut error = None;
20838        while let Some(chunk) = stream.next().await {
20839            match chunk {
20840                StreamChunk::Content { text } => content.push_str(&text),
20841                StreamChunk::Error { message } => error = Some(message),
20842                StreamChunk::Done {} => break,
20843                _ => {}
20844            }
20845        }
20846
20847        assert_eq!(error, None);
20848        assert_eq!(content, "one two three");
20849    }
20850
20851    #[tokio::test]
20852    async fn test_buffered_streaming_main_failure_finalizes_branch() {
20853        use futures::StreamExt;
20854
20855        let mock = mock_with_response("one two");
20856        let mut router_mock = mock_with_response("0");
20857        router_mock.set_latency(50);
20858        let yaml = r#"
20859name: BufferedFailureAgent
20860system_prompt: "You stream safely."
20861llm:
20862  default: default
20863  router: router
20864observability:
20865  enabled: true
20866  export:
20867    write_raw_events: true
20868streaming:
20869  enabled: true
20870  buffer_size: 1
20871runtime:
20872  optimization:
20873    enabled: true
20874    max_speculative_llm_calls_per_turn: 2
20875    speculative_state_transitions: true
20876    streaming_policy: buffer_until_routing_done
20877    max_parallel_runtime_tasks: 2
20878states:
20879  initial: triage
20880  states:
20881    triage:
20882      prompt: "Ask for the category."
20883      transitions:
20884        - to: billing
20885          when: "User asks about billing"
20886          timing: parallel
20887    billing:
20888      prompt: "Billing state."
20889"#;
20890        let agent = AgentBuilder::from_yaml(yaml)
20891            .unwrap()
20892            .llm_alias("default", Arc::new(mock))
20893            .llm_alias("router", Arc::new(router_mock))
20894            .build()
20895            .unwrap();
20896
20897        let mut stream = agent.chat_stream("hello").await.unwrap();
20898        let mut error = String::new();
20899        while let Some(chunk) = stream.next().await {
20900            if let StreamChunk::Error { message } = chunk {
20901                error = message;
20902            }
20903        }
20904
20905        assert!(
20906            error.contains("stream buffer filled"),
20907            "unexpected stream error: {}",
20908            error
20909        );
20910        let events = agent.observability().unwrap().raw_events();
20911        assert!(events.iter().any(|event| {
20912            event.dimensions.get("branch_status") == Some(&"failed".to_string())
20913                && event.dimensions.get("commit_behavior") == Some(&"final_response".to_string())
20914                && event.dimensions.get("optimization")
20915                    == Some(&"buffered_streaming_routing".to_string())
20916        }));
20917    }
20918
20919    #[tokio::test]
20920    async fn test_streaming_preflight_does_not_emit_old_state_content() {
20921        use futures::StreamExt;
20922
20923        let mock = mock_with_response("Billing streamed response");
20924        let yaml = r#"
20925name: StreamingOptimizedAgent
20926system_prompt: "You route before streaming."
20927runtime:
20928  optimization:
20929    enabled: true
20930    pre_response_deterministic_transitions: true
20931streaming:
20932  enabled: true
20933states:
20934  initial: greeting
20935  states:
20936    greeting:
20937      prompt: "OLD_STATE_SENTINEL"
20938      transitions:
20939        - to: billing
20940          guard:
20941            context:
20942              topic:
20943                eq: billing
20944          timing: pre_response
20945    billing:
20946      prompt: "Billing state."
20947"#;
20948        let agent = AgentBuilder::from_yaml(yaml)
20949            .unwrap()
20950            .llm(Arc::new(mock))
20951            .build()
20952            .unwrap();
20953        agent
20954            .set_context("topic", serde_json::json!("billing"))
20955            .unwrap();
20956
20957        let mut stream = agent.chat_stream("billing please").await.unwrap();
20958        let mut content = String::new();
20959        while let Some(chunk) = stream.next().await {
20960            match chunk {
20961                StreamChunk::Content { text } => content.push_str(&text),
20962                StreamChunk::Error { message } => panic!("stream error: {}", message),
20963                StreamChunk::Done {} => break,
20964                _ => {}
20965            }
20966        }
20967
20968        assert_eq!(agent.current_state().as_deref(), Some("billing"));
20969        assert!(content.contains("Billing streamed response"));
20970        assert!(!content.contains("OLD_STATE_SENTINEL"));
20971    }
20972
20973    // State machine transitions
20974    #[tokio::test]
20975    async fn test_integration_state_machine_basic() {
20976        let yaml = r#"
20977name: StateAgent
20978system_prompt: "You are a support agent."
20979states:
20980  initial: greeting
20981  states:
20982    greeting:
20983      prompt: "Welcome the user warmly."
20984      transitions:
20985        - to: support
20986          when: "User needs help"
20987          auto: true
20988    support:
20989      prompt: "Help solve the user's problem."
20990"#;
20991        let mock = mock_with_responses(vec![
20992            "Welcome! How can I help?", // greeting response
20993            "1",                        // transition evaluator picks first (index 0)
20994            "I'll help you with that.", // support response
20995        ]);
20996        let builder = AgentBuilder::from_yaml(yaml).unwrap();
20997        let agent = builder.llm(Arc::new(mock)).build().unwrap();
20998
20999        assert_eq!(agent.current_state(), Some("greeting".to_string()));
21000        let _ = agent.chat("I need help").await.unwrap();
21001        // After transition evaluation, state may or may not have changed
21002        // depending on mock evaluator response - the key is that it doesn't crash
21003    }
21004
21005    // State on_enter/on_exit actions
21006    #[tokio::test]
21007    async fn test_integration_state_on_enter_set_context() {
21008        let yaml = r#"
21009name: ActionAgent
21010system_prompt: "You are helpful."
21011states:
21012  initial: step1
21013  states:
21014    step1:
21015      prompt: "Step 1"
21016      on_exit:
21017        - set_context:
21018            step1_exited: true
21019      transitions:
21020        - to: step2
21021          when: "always"
21022          auto: true
21023    step2:
21024      prompt: "Step 2"
21025      on_enter:
21026        - set_context:
21027            step2_entered: true
21028"#;
21029        // The transition evaluator will pick the first transition (index 0)
21030        let mock = mock_with_responses(vec![
21031            "Processing step 1.",
21032            "0", // transition evaluator response: select first transition
21033        ]);
21034        let builder = AgentBuilder::from_yaml(yaml).unwrap();
21035        let agent = builder.llm(Arc::new(mock)).build().unwrap();
21036
21037        assert_eq!(agent.current_state(), Some("step1".to_string()));
21038
21039        // Manually transition to test on_enter/on_exit
21040        agent.transition_to("step2").await.unwrap();
21041
21042        assert_eq!(agent.current_state(), Some("step2".to_string()));
21043
21044        // Verify context was set by on_exit and on_enter actions
21045        let ctx = agent.get_context();
21046        assert_eq!(ctx.get("step1_exited"), Some(&serde_json::json!(true)));
21047        assert_eq!(ctx.get("step2_entered"), Some(&serde_json::json!(true)));
21048    }
21049
21050    #[tokio::test]
21051    async fn state_action_tool_preserves_source_in_stored_record() {
21052        let yaml = r#"
21053name: StateActionToolAgent
21054system_prompt: "You are helpful."
21055tools:
21056  - context_echo
21057states:
21058  initial: idle
21059  states:
21060    idle:
21061      prompt: "Idle"
21062    active:
21063      prompt: "Active"
21064      on_enter:
21065        - set_context:
21066            action_started: true
21067        - tool: context_echo
21068          args: {}
21069"#;
21070        let agent = AgentBuilder::from_yaml(yaml)
21071            .unwrap()
21072            .llm(Arc::new(mock_with_response("unused")))
21073            .tool(Arc::new(ContextEchoTool))
21074            .build()
21075            .unwrap();
21076
21077        agent.transition_to("active").await.unwrap();
21078
21079        let record: ToolExecutionRecord = serde_json::from_value(
21080            agent
21081                .get_context()
21082                .get("last_tool_record")
21083                .cloned()
21084                .expect("successful state action must store its execution record"),
21085        )
21086        .unwrap();
21087        assert!(record.executed);
21088        assert!(record.success);
21089        assert_eq!(record.canonical_id, "context_echo");
21090        assert!(matches!(
21091            &record.source,
21092            ToolCallSource::StateAction {
21093                state: Some(state),
21094                action_index: 1,
21095            } if state == "active"
21096        ));
21097    }
21098
21099    #[tokio::test]
21100    async fn test_ordinary_transition_uses_on_enter_then_on_reenter() {
21101        let yaml = r#"
21102name: OrdinaryLifecycleAgent
21103system_prompt: "You are helpful."
21104states:
21105  initial: intake
21106  regenerate_on_transition: false
21107  states:
21108    intake:
21109      prompt: "Intake"
21110      transitions:
21111        - to: drafting
21112          guard:
21113            context:
21114              route:
21115                eq: drafting
21116    drafting:
21117      prompt: "Drafting"
21118      on_enter:
21119        - set_context:
21120            draft_version: 1
21121      on_reenter:
21122        - set_context:
21123            draft_version: 2
21124      transitions:
21125        - to: review
21126          guard:
21127            context:
21128              route:
21129                eq: review
21130    review:
21131      prompt: "Review"
21132      on_enter:
21133        - set_context:
21134            review_entry: first
21135      transitions:
21136        - to: drafting
21137          guard:
21138            context:
21139              route:
21140                eq: drafting
21141"#;
21142        let agent = AgentBuilder::from_yaml(yaml)
21143            .unwrap()
21144            .llm(Arc::new(mock_with_responses(vec![
21145                "Intake response",
21146                "Draft response",
21147                "Review response",
21148            ])))
21149            .build()
21150            .unwrap();
21151
21152        agent
21153            .set_context("route", serde_json::json!("drafting"))
21154            .unwrap();
21155        agent.chat("Start a draft").await.unwrap();
21156        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
21157        assert_eq!(
21158            agent.get_context().get("draft_version"),
21159            Some(&serde_json::json!(1))
21160        );
21161
21162        agent
21163            .set_context("route", serde_json::json!("review"))
21164            .unwrap();
21165        agent.chat("Review this").await.unwrap();
21166        assert_eq!(agent.current_state().as_deref(), Some("review"));
21167        assert_eq!(
21168            agent.get_context().get("review_entry"),
21169            Some(&serde_json::json!("first"))
21170        );
21171
21172        agent
21173            .set_context("route", serde_json::json!("drafting"))
21174            .unwrap();
21175        agent.chat("Revise this").await.unwrap();
21176        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
21177        assert_eq!(
21178            agent.get_context().get("draft_version"),
21179            Some(&serde_json::json!(2))
21180        );
21181    }
21182
21183    #[tokio::test]
21184    async fn test_manual_transition_uses_on_enter_then_on_reenter() {
21185        let yaml = r#"
21186name: ManualLifecycleAgent
21187system_prompt: "You are helpful."
21188states:
21189  initial: intake
21190  states:
21191    intake:
21192      prompt: "Intake"
21193    drafting:
21194      prompt: "Drafting"
21195      on_enter:
21196        - set_context:
21197            draft_version: 1
21198      on_reenter:
21199        - set_context:
21200            draft_version: 2
21201    review:
21202      prompt: "Review"
21203"#;
21204        let agent = AgentBuilder::from_yaml(yaml)
21205            .unwrap()
21206            .llm(Arc::new(mock_with_response("unused")))
21207            .build()
21208            .unwrap();
21209
21210        assert!(!agent.get_context().contains_key("draft_version"));
21211        agent.transition_to("drafting").await.unwrap();
21212        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
21213        assert_eq!(
21214            agent.get_context().get("draft_version"),
21215            Some(&serde_json::json!(1))
21216        );
21217
21218        agent.transition_to("review").await.unwrap();
21219        agent.transition_to("drafting").await.unwrap();
21220        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
21221        assert_eq!(
21222            agent.get_context().get("draft_version"),
21223            Some(&serde_json::json!(2))
21224        );
21225    }
21226
21227    #[tokio::test]
21228    async fn test_timeout_transition_uses_on_enter_then_on_reenter() {
21229        let yaml = r#"
21230name: TimeoutLifecycleAgent
21231system_prompt: "You are helpful."
21232states:
21233  initial: intake
21234  regenerate_on_transition: false
21235  states:
21236    intake:
21237      prompt: "Intake"
21238      max_turns: 1
21239      timeout_to: drafting
21240    drafting:
21241      prompt: "Drafting"
21242      max_turns: 1
21243      timeout_to: review
21244      on_enter:
21245        - set_context:
21246            draft_version: 1
21247      on_reenter:
21248        - set_context:
21249            draft_version: 2
21250    review:
21251      prompt: "Review"
21252      max_turns: 1
21253      timeout_to: drafting
21254      on_enter:
21255        - set_context:
21256            review_entry: first
21257"#;
21258        let agent = AgentBuilder::from_yaml(yaml)
21259            .unwrap()
21260            .llm(Arc::new(mock_with_responses(vec![
21261                "Intake",
21262                "First draft",
21263                "Review",
21264                "Revised draft",
21265            ])))
21266            .build()
21267            .unwrap();
21268
21269        agent.chat("First turn").await.unwrap();
21270        assert_eq!(agent.current_state().as_deref(), Some("intake"));
21271        assert!(!agent.get_context().contains_key("draft_version"));
21272
21273        agent.chat("Second turn").await.unwrap();
21274        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
21275        assert_eq!(
21276            agent.get_context().get("draft_version"),
21277            Some(&serde_json::json!(1))
21278        );
21279
21280        agent.chat("Third turn").await.unwrap();
21281        assert_eq!(agent.current_state().as_deref(), Some("review"));
21282        assert_eq!(
21283            agent.get_context().get("review_entry"),
21284            Some(&serde_json::json!("first"))
21285        );
21286
21287        agent.chat("Fourth turn").await.unwrap();
21288        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
21289        assert_eq!(
21290            agent.get_context().get("draft_version"),
21291            Some(&serde_json::json!(2))
21292        );
21293    }
21294
21295    // Process pipeline transforms input
21296    #[tokio::test]
21297    async fn test_integration_process_normalize() {
21298        let yaml = r#"
21299name: ProcessAgent
21300system_prompt: "You are helpful."
21301process:
21302  input:
21303    - type: normalize
21304      config:
21305        trim: true
21306        collapse_whitespace: true
21307"#;
21308        let mock = mock_with_response("Got your message.");
21309        let builder = AgentBuilder::from_yaml(yaml).unwrap();
21310        let agent = builder.llm(Arc::new(mock.clone())).build().unwrap();
21311
21312        let _ = agent.chat("  hello   world  ").await.unwrap();
21313
21314        // Verify the LLM received the normalized input (trimmed + collapsed whitespace)
21315        let history = mock.call_history();
21316        assert!(!history.is_empty());
21317        // The user message in LLM call should be normalized
21318        let last_call = history.last().unwrap();
21319        let user_msg = last_call
21320            .messages
21321            .iter()
21322            .find(|m| m.role == ai_agents_core::Role::User)
21323            .unwrap();
21324        assert_eq!(user_msg.content, "hello world");
21325    }
21326
21327    // ═══════════════════════════════════════════════════════════
21328    // Integration Test 2.1.7: Memory compression triggers
21329    // ═══════════════════════════════════════════════════════════
21330    #[tokio::test]
21331    async fn test_integration_memory_compression() {
21332        let yaml = r#"
21333name: MemoryAgent
21334system_prompt: "You are helpful."
21335memory:
21336  type: compacting
21337  max_messages: 100
21338  compress_threshold: 5
21339  max_recent_messages: 3
21340  summarize_batch_size: 2
21341"#;
21342        // Provide enough responses for compression to trigger
21343        let responses: Vec<&str> = (0..8).map(|_| "Response from assistant.").collect();
21344        let mock = mock_with_responses(responses);
21345        let builder = AgentBuilder::from_yaml(yaml).unwrap();
21346        let agent = builder.llm(Arc::new(mock)).build().unwrap();
21347
21348        // Send enough messages to trigger compression
21349        for i in 0..6 {
21350            let _ = agent.chat(&format!("Message {}", i)).await.unwrap();
21351        }
21352
21353        // Memory should have compressed - verify it didn't crash
21354        // and that messages are bounded
21355        let messages = agent.memory.get_messages(None).await.unwrap();
21356        // With compress_threshold=5 and max_recent_messages=3,
21357        // after 6 turns (12 messages), compression should have run
21358        assert!(messages.len() <= 12); // At most all messages if no compression, fewer if compressed
21359    }
21360
21361    // YAML with multiple LLMs
21362    #[tokio::test]
21363    async fn test_integration_multi_llm_registry() {
21364        let mut mock_default = MockLLMProvider::new("default");
21365        mock_default.set_response("Default LLM response.");
21366        let mut mock_router = MockLLMProvider::new("router");
21367        mock_router.set_response("Router response.");
21368
21369        let agent = AgentBuilder::new()
21370            .system_prompt("You are helpful.")
21371            .llm_alias("default", Arc::new(mock_default))
21372            .llm_alias("router", Arc::new(mock_router))
21373            .build()
21374            .unwrap();
21375
21376        let response = agent.chat("Hello").await.unwrap();
21377        assert_eq!(response.content, "Default LLM response.");
21378    }
21379
21380    // Agent reset clears state
21381    #[tokio::test]
21382    async fn test_integration_agent_reset() {
21383        let mock = mock_with_responses(vec!["Hello!", "Hello again!"]);
21384        let agent = AgentBuilder::new()
21385            .system_prompt("You are helpful.")
21386            .llm(Arc::new(mock))
21387            .build()
21388            .unwrap();
21389
21390        let _ = agent.chat("Hi").await.unwrap();
21391        let messages = agent.memory.get_messages(None).await.unwrap();
21392        assert_eq!(messages.len(), 2); // user + assistant
21393
21394        agent.reset().await.unwrap();
21395        let messages = agent.memory.get_messages(None).await.unwrap();
21396        assert_eq!(messages.len(), 0);
21397    }
21398
21399    // Process pipeline rejects input
21400    #[tokio::test]
21401    async fn test_integration_process_validate_reject() {
21402        use ai_agents_process::{ProcessConfig, ProcessProcessor};
21403
21404        let validate_config = ai_agents_process::ValidateStage {
21405            id: Some("length_check".to_string()),
21406            condition: None,
21407            config: ai_agents_process::ValidateConfig {
21408                rules: vec![ai_agents_process::ValidationRule::MinLength {
21409                    min_length: 10,
21410                    on_fail: ai_agents_process::ValidationAction {
21411                        action: ai_agents_process::ValidationActionType::Reject,
21412                        message: None,
21413                    },
21414                }],
21415                ..Default::default()
21416            },
21417        };
21418        let process_config = ProcessConfig {
21419            input: vec![ai_agents_process::ProcessStage::Validate(validate_config)],
21420            ..Default::default()
21421        };
21422        let processor = ProcessProcessor::new(process_config);
21423
21424        let mock = mock_with_response("Should not reach here.");
21425        let agent = AgentBuilder::new()
21426            .system_prompt("You are helpful.")
21427            .llm(Arc::new(mock))
21428            .process_processor(processor)
21429            .build()
21430            .unwrap();
21431
21432        let response = agent.chat("Hi").await.unwrap();
21433        // Rejected input should produce a rejection response, not call LLM
21434        assert!(
21435            response.content.contains("rejected")
21436                || response.content.contains("Input rejected")
21437                || response.content.contains("too short")
21438                || response.content.contains("Too short")
21439                || response.content.len() < 50, // rejection message is typically short
21440            "Expected rejection response, got: {}",
21441            response.content
21442        );
21443    }
21444
21445    // LLM fallback: primary fails, fallback LLM responds
21446    #[tokio::test]
21447    async fn test_llm_fallback_on_failure() {
21448        use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
21449
21450        let mut primary = MockLLMProvider::new("primary");
21451        primary.set_error("Primary LLM is unavailable");
21452
21453        let mut fallback = MockLLMProvider::new("fallback");
21454        fallback.set_response("Fallback response works!");
21455
21456        let agent = AgentBuilder::new()
21457            .system_prompt("You are helpful.")
21458            .llm_alias("default", Arc::new(primary))
21459            .llm_alias("backup", Arc::new(fallback))
21460            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
21461                llm: LLMRecoveryConfig {
21462                    on_failure: LLMFailureAction::FallbackLlm {
21463                        fallback_llm: "backup".to_string(),
21464                    },
21465                    ..Default::default()
21466                },
21467                ..Default::default()
21468            }))
21469            .build()
21470            .unwrap();
21471
21472        let response = agent.chat("Hello").await.unwrap();
21473        assert!(
21474            response.content.contains("Fallback response"),
21475            "Expected fallback response, got: {}",
21476            response.content
21477        );
21478    }
21479
21480    // LLM fallback: primary fails, static message returned
21481    #[tokio::test]
21482    async fn test_llm_fallback_response_static_message() {
21483        use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
21484
21485        let mut primary = MockLLMProvider::new("primary");
21486        primary.set_error("Primary LLM is unavailable");
21487
21488        let agent = AgentBuilder::new()
21489            .system_prompt("You are helpful.")
21490            .llm(Arc::new(primary))
21491            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
21492                llm: LLMRecoveryConfig {
21493                    on_failure: LLMFailureAction::FallbackResponse {
21494                        message: "I am temporarily unavailable. Please try again later."
21495                            .to_string(),
21496                    },
21497                    ..Default::default()
21498                },
21499                ..Default::default()
21500            }))
21501            .build()
21502            .unwrap();
21503
21504        let response = agent.chat("Hello").await.unwrap();
21505        assert!(
21506            response.content.contains("temporarily unavailable"),
21507            "Expected static fallback message, got: {}",
21508            response.content
21509        );
21510    }
21511
21512    // Tool skip: tool fails, on_failure: skip absorbs the error
21513    #[tokio::test]
21514    async fn test_tool_failure_skip() {
21515        use ai_agents_recovery::{
21516            ErrorRecoveryConfig, ToolFailureAction, ToolRecoveryConfig, ToolRetryConfig,
21517        };
21518
21519        // LLM requests a nonexistent tool, then responds after seeing the skip result
21520        let mock = mock_with_responses(vec![
21521            r#"I'll use the nonexistent tool.
21522[TOOL_CALL: {"name": "nonexistent_tool", "arguments": {}}]"#,
21523            "The tool was unavailable, but I can still help you.",
21524        ]);
21525
21526        let agent = AgentBuilder::new()
21527            .system_prompt("You are helpful.")
21528            .llm(Arc::new(mock))
21529            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
21530                tools: ToolRecoveryConfig {
21531                    default: ToolRetryConfig {
21532                        max_retries: 0,
21533                        timeout_ms: None,
21534                        on_failure: ToolFailureAction::Skip,
21535                    },
21536                    ..Default::default()
21537                },
21538                ..Default::default()
21539            }))
21540            .build()
21541            .unwrap();
21542
21543        // The tool will fail (not found), but on_failure: skip absorbs the error
21544        let response = agent.chat("Use the nonexistent tool").await;
21545        assert!(
21546            response.is_ok(),
21547            "Expected Ok with skip policy, got: {:?}",
21548            response
21549        );
21550    }
21551}