Skip to main content

ai_agents_runtime/
runtime.rs

1use async_trait::async_trait;
2use futures::stream::{Stream, StreamExt};
3use parking_lot::RwLock;
4use serde_json::Value;
5use std::collections::{HashMap, HashSet};
6use std::future::Future;
7use std::pin::Pin;
8use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
9use std::sync::{Arc, Weak};
10use std::time::Instant;
11use tracing::{debug, error, info, instrument, warn};
12
13const DISAMBIGUATION_STATE_GENERATION_KEY: &str = "_runtime.disambiguation_state_generation";
14
15/// Keeps the strong identity of one non-reentrant root-turn gate.
16pub(crate) type RootTurnGate = Arc<tokio::sync::Mutex<()>>;
17
18/// Owns one immutable root-gate ancestry snapshot that can cross polls and spawned tasks by cloning one `Arc`.
19pub(crate) type RootTurnGateIdentityStack = Arc<[RootTurnGate]>;
20
21tokio::task_local! {
22    static RUNTIME_GATE_IDENTITY_STACK: RootTurnGateIdentityStack;
23}
24
25/// Captures the immutable gate ownership chain for propagation and returns an empty chain outside any root-turn scope.
26pub(crate) fn current_runtime_gate_identity_stack() -> RootTurnGateIdentityStack {
27    RUNTIME_GATE_IDENTITY_STACK
28        .try_with(Arc::clone)
29        .unwrap_or_default()
30}
31
32/// Polls a future with an explicit gate ownership chain while cloning only its immutable owner instead of rebuilding the ancestry.
33pub(crate) async fn scope_runtime_gate_identity_stack<F, T>(
34    identity_stack: &RootTurnGateIdentityStack,
35    future: F,
36) -> T
37where
38    F: Future<Output = T>,
39{
40    RUNTIME_GATE_IDENTITY_STACK
41        .scope(Arc::clone(identity_stack), future)
42        .await
43}
44
45/// Shared lock table used to serialize side-effecting tool calls by canonical resource.
46pub(crate) type ToolResourceLocks = Arc<RwLock<HashMap<String, Weak<tokio::sync::Mutex<()>>>>>;
47
48//
49// Owns every lock acquired for one tool call and removes dead weak entries after release.
50//
51struct ToolResourceGuards {
52    guards: Vec<tokio::sync::OwnedMutexGuard<()>>,
53    locks: ToolResourceLocks,
54}
55
56//
57// Couples one acquired root-turn gate with the immutable ancestry owner that must remain active through hooks, orchestration, and stream polling.
58//
59struct RootTurnAdmission {
60    guard: tokio::sync::OwnedMutexGuard<()>,
61    identity_stack: RootTurnGateIdentityStack,
62}
63
64#[derive(Clone)]
65struct StoredSessionRestore {
66    snapshot: AgentSnapshot,
67    metadata: Option<ai_agents_core::SessionMetadata>,
68}
69
70struct RuntimeSessionRestorePoint {
71    snapshot: AgentSnapshot,
72    metadata: ai_agents_core::SessionMetadata,
73    actor_id: Option<String>,
74    session_id: Option<String>,
75}
76
77impl Drop for ToolResourceGuards {
78    fn drop(&mut self) {
79        self.guards.clear();
80        self.locks.write().retain(|_, lock| lock.strong_count() > 0);
81    }
82}
83
84//
85// Captures runtime controls under one read guard so policy and scope belong to the same generation.
86//
87#[derive(Clone)]
88struct RuntimeSafetySnapshot {
89    version: u64,
90    emergency_deny: bool,
91    tool_security: ToolSecurityEngine,
92    tool_scope_override: Option<Vec<String>>,
93}
94
95//
96// Records the generations used by the final authorization decision.
97//
98#[derive(Clone, Copy)]
99struct ToolDecisionVersions {
100    policy: u64,
101    registry: u64,
102    runtime_control: u64,
103    state: Option<u64>,
104}
105
106//
107// Carries deterministic effective tools with the state generation that authorized their scopes.
108//
109struct AvailableToolIdsSnapshot {
110    tool_ids: Vec<String>,
111    state_generation: Option<u64>,
112}
113
114//
115// Binds human approval to the reviewed action and exact tool implementation.
116//
117#[derive(Clone)]
118struct ToolApprovalBinding {
119    canonical_id: String,
120    arguments: Value,
121    confirmation_required: bool,
122    policy_version: u64,
123    runtime_control_version: u64,
124    state_generation: Option<u64>,
125    reviewed_tool: Arc<dyn ai_agents_core::Tool>,
126}
127
128//
129// A later plain approval must not erase arguments modified by an earlier approval stage.
130//
131fn merge_approved_record(record: &mut Option<ToolApprovalRecord>) {
132    if record
133        .as_ref()
134        .is_some_and(|record| matches!(record.status, ToolApprovalStatus::Modified))
135    {
136        return;
137    }
138    *record = Some(ToolApprovalRecord {
139        status: ToolApprovalStatus::Approved,
140        reason: None,
141        modified_arguments: None,
142    });
143}
144
145impl ToolApprovalBinding {
146    /// Returns true when final authorization no longer matches the reviewed action.
147    fn is_stale(
148        &self,
149        canonical_id: &str,
150        arguments: &Value,
151        confirmation_required: bool,
152        versions: ToolDecisionVersions,
153        resolved_tool: &Arc<dyn ai_agents_core::Tool>,
154    ) -> bool {
155        self.canonical_id != canonical_id
156            || self.arguments != *arguments
157            || self.confirmation_required != confirmation_required
158            || self.policy_version != versions.policy
159            || self.runtime_control_version != versions.runtime_control
160            || self.state_generation != versions.state
161            || !Arc::ptr_eq(&self.reviewed_tool, resolved_tool)
162    }
163}
164
165use crate::turn_context::{current_turn_actor_context, scope_actor_context};
166
167use ai_agents_context::{ContextManager, ContextProvider, TemplateRenderer};
168use ai_agents_core::traits::storage::StorageCapability;
169use ai_agents_core::{
170    AgentError, AgentSnapshot, AgentStorage, ChatMessage, FinishReason, LLMError, LLMProvider,
171    LLMResponse, LLMToolDefinition, LLMToolRequest, PermissionOutcome, Result, ToolActorContext,
172    ToolApprovalRecord, ToolApprovalStatus, ToolCallSource, ToolCancellationToken, ToolChoice,
173    ToolExecutionContext, ToolExecutionRecord, ToolExecutionRequest, ToolInvoker,
174    ToolPolicyDecisionRecord, ToolResult,
175};
176use ai_agents_disambiguation::{
177    ClarificationObserver, ClarificationParseFuture, ClarificationQuestionFuture,
178    ConfirmationParseFuture, DisambiguationConfig, DisambiguationContext, DisambiguationManager,
179    DisambiguationResult,
180};
181use ai_agents_hitl::{
182    ApprovalHandler, ApprovalResolvedOutcome, ApprovalResult, ApprovalTrigger, HITLCheckResult,
183    HITLEngine, RejectAllHandler, TimeoutAction,
184};
185use ai_agents_hooks::{AgentHooks, NoopHooks};
186use ai_agents_llm::LLMRegistry;
187use ai_agents_memory::{
188    CompressResult, EvictionReason, Memory, MemoryBudgetEvent, MemoryCompressEvent,
189    MemoryEvictEvent, MemoryTokenBudget, OverflowStrategy,
190};
191use ai_agents_observability::{
192    EventStatus, EventType, ObservabilityManager, ObservationPurpose, SpanContext,
193    current_observation_context, new_session_id as new_observation_session_id,
194    resolve_language_from_context, with_observation_context, with_observation_purpose,
195};
196use ai_agents_process::{
197    ProcessData, ProcessProcessor, ProcessPurposeHint, ProcessStageFuture, ProcessStageObserver,
198};
199use ai_agents_reasoning::{
200    CriterionResult, EvaluationResult, Plan, PlanAction, PlanStatus, PlanStep, ReasoningConfig,
201    ReasoningMetadata, ReasoningMode, ReasoningOutput, ReflectionAttempt, ReflectionConfig,
202    ReflectionMetadata, StepFailureAction,
203};
204use ai_agents_recovery::{
205    ByRoleFilter, ContextOverflowAction, FilterConfig, IntoClassifiedError, KeepRecentFilter,
206    LLMFailureAction, MessageFilter, RecoveryManager, SkipPatternFilter, ToolFailureAction,
207};
208use ai_agents_relationships::RelationshipManager;
209use ai_agents_skills::{SkillDefinition, SkillExecutor, SkillRouter};
210use ai_agents_state::{
211    PromptMode, StateAction, StateMachine, StateMachineSnapshot, StateTransitionEvent, Transition,
212    TransitionContext, TransitionEvaluator, TransitionTiming, evaluate_guard,
213};
214use ai_agents_storage::{StorageConfig as StorageStorageConfig, create_storage};
215use ai_agents_tools::{
216    CommandRunner, ConditionEvaluator, DiagnosticsProvider, EvaluationContext, LLMGetter,
217    QuestionHandler, SecurityCheckResult, TodoItem, ToolCallRecord, ToolRegistry,
218    ToolSecurityConfig, ToolSecurityEngine,
219};
220
221use super::{
222    Agent, AgentInfo, AgentResponse, AgentStreamEvent, ParallelToolsConfig, StreamChunk,
223    StreamingConfig, ToolCall,
224};
225use crate::optimization::{
226    AwaitBeforeNextTurn, BackgroundMaintenanceQueue, BackgroundOverflowPolicy, MainResponseDraft,
227    MaintenanceMode, MaintenanceSequenceKey, RuntimeBranch, RuntimeBranchResult,
228    RuntimeBranchStatus, RuntimeCommitBehavior, RuntimeConfig, RuntimeOptimizationKind,
229    RuntimeTaskPriority, RuntimeTaskPurpose, ScheduledBranchSet, SkillCandidate,
230    StreamingDraftResult, TransitionCandidate, TurnBranchScheduler, TurnOptimizationContext,
231};
232use crate::spec::StorageConfig;
233
234/// Outcome of processing tool calls within the agent loop.
235enum ToolCallOutcome {
236    /// Tools executed successfully, continue the LLM loop for the next iteration.
237    Continue,
238    /// A state transition fired during tool call handling, continue the loop.
239    TransitionFired,
240    /// HITL rejected a tool call, return this response immediately.
241    Rejected(AgentResponse),
242}
243
244#[derive(Clone)]
245struct MainToolProtocol {
246    choice: Option<ToolChoice>,
247    tool_ids: Vec<String>,
248    definitions: Vec<LLMToolDefinition>,
249}
250
251struct MainProviderResponse {
252    response: LLMResponse,
253    used_native_tools: bool,
254}
255
256//
257// Carries one committed model answer through shared output processing and root-turn finalization.
258//
259struct CommittedTextResponse<'a> {
260    processed_input: &'a str,
261    input_context: &'a HashMap<String, Value>,
262    answer: String,
263    reasoning_mode: ReasoningMode,
264    auto_detected: bool,
265    iterations: u32,
266    thinking_content: Option<String>,
267    all_tool_calls: Vec<ToolCall>,
268}
269
270//
271// Carries the finalized response fields into metadata assembly without changing their ownership.
272//
273struct AgentResponseParts {
274    content: String,
275    all_tool_calls: Vec<ToolCall>,
276    reasoning_mode: ReasoningMode,
277    auto_detected: bool,
278    iterations: u32,
279    thinking: Option<String>,
280    reflection_metadata: Option<ReflectionMetadata>,
281}
282
283//
284// Owns the finalized response paired with a legacy Done chunk without changing provisional chunk timing.
285//
286type RuntimeStreamTerminalSlot = Arc<RwLock<Option<AgentResponse>>>;
287
288//
289// Creates isolated terminal ownership for one public stream.
290//
291fn new_runtime_stream_terminal_slot() -> RuntimeStreamTerminalSlot {
292    Arc::new(RwLock::new(None))
293}
294
295//
296// Records the exact response that already completed root-turn finalization before Done is emitted.
297//
298fn record_runtime_stream_final(slot: &RuntimeStreamTerminalSlot, response: AgentResponse) {
299    *slot.write() = Some(response);
300}
301
302#[derive(Clone, Copy)]
303struct DisambiguationOwnership {
304    epoch: u64,
305    state_generation: Option<u64>,
306}
307
308/// Outcome of skill routing — used by `try_skill_route`.
309enum SkillRouteResult {
310    /// No skill matched, continue to normal LLM chat.
311    NoMatch,
312    /// Skill executed successfully.
313    Response { skill_id: String, content: String },
314    /// Skill matched but needs disambiguation first or returned a terminal disambiguation response.
315    NeedsClarification {
316        response: AgentResponse,
317        ownership: Option<DisambiguationOwnership>,
318    },
319}
320
321/// Result of response-independent parallel transition selection.
322enum ParallelTransitionSelection {
323    /// A transition matched and can be committed before the old-state response.
324    Candidate(TransitionCandidate),
325    /// All eligible transition checks ran and none matched.
326    NoMatch,
327    /// LLM-based route evaluation needed speculative capacity that was unavailable.
328    ReservationExhausted,
329}
330
331/// Outcome of post_loop_processing - drives the caller's next step.
332enum PostLoopResult {
333    /// No transition fired. Content is the LLM response for this turn.
334    NoTransition(String),
335    /// Transition fired. Content is from plain post-transition re-generation.
336    Transitioned(String),
337    /// Transition fired into a state that requires full dispatch.
338    /// Caller re-enters run_loop_internal to apply the correct handler.
339    NeedsRedispatch,
340}
341
342struct StateTransitionReservation<'a> {
343    reserved: &'a AtomicBool,
344}
345
346impl Drop for StateTransitionReservation<'_> {
347    fn drop(&mut self) {
348        self.reserved.store(false, Ordering::SeqCst);
349    }
350}
351
352struct RootTurnCleanup<'a> {
353    agent: &'a RuntimeAgent,
354}
355
356impl<'a> RootTurnCleanup<'a> {
357    fn new(agent: &'a RuntimeAgent) -> Self {
358        Self { agent }
359    }
360}
361
362impl Drop for RootTurnCleanup<'_> {
363    fn drop(&mut self) {
364        self.agent.end_root_turn();
365    }
366}
367
368/// Host-owned runtime control state shared with active agents.
369#[derive(Debug)]
370struct RuntimeControlState {
371    /// Serializes control mutations with exact runtime safety snapshots.
372    snapshot_guard: RwLock<()>,
373    /// Monotonic version for runtime-control snapshots.
374    version: AtomicU64,
375    /// Emergency switch that denies future calls and is shared with active tool contexts.
376    emergency_deny: Arc<AtomicBool>,
377    /// Optional live replacement for tool security policy.
378    tool_security_override: RwLock<Option<ToolSecurityEngine>>,
379    /// Optional live narrowing scope applied after the runtime's declared grant.
380    tool_scope_override: RwLock<Option<Vec<String>>>,
381}
382
383impl Default for RuntimeControlState {
384    fn default() -> Self {
385        Self {
386            snapshot_guard: RwLock::new(()),
387            version: AtomicU64::new(1),
388            emergency_deny: Arc::new(AtomicBool::new(false)),
389            tool_security_override: RwLock::new(None),
390            tool_scope_override: RwLock::new(None),
391        }
392    }
393}
394
395/// Host-only handle for live runtime safety controls.
396#[derive(Clone)]
397pub struct RuntimeControlHandle {
398    state: Arc<RuntimeControlState>,
399}
400
401impl RuntimeControlHandle {
402    /// Returns the current runtime-control version.
403    pub fn version(&self) -> u64 {
404        self.state.version.load(Ordering::SeqCst)
405    }
406
407    fn bump(&self) -> u64 {
408        self.state.version.fetch_add(1, Ordering::SeqCst) + 1
409    }
410
411    /// Overrides tool security for later tool calls and panics if host policy is invalid.
412    pub fn set_tool_security(&self, config: ToolSecurityConfig) -> u64 {
413        self.try_set_tool_security(config)
414            .expect("invalid tool security configuration")
415    }
416
417    /// Validates replacement policy before changing the runtime-control generation or active snapshot.
418    pub fn try_set_tool_security(&self, config: ToolSecurityConfig) -> Result<u64> {
419        config.validate()?;
420        let _guard = self.state.snapshot_guard.write();
421        let generation = self.bump();
422        *self.state.tool_security_override.write() = Some(
423            ToolSecurityEngine::new_with_policy_version(config, generation),
424        );
425        Ok(generation)
426    }
427
428    /// Clears the live tool security override.
429    pub fn clear_tool_security_override(&self) -> u64 {
430        let _guard = self.state.snapshot_guard.write();
431        *self.state.tool_security_override.write() = None;
432        self.bump()
433    }
434
435    /// Narrows the runtime's declared tool grant for later calls without adding authority.
436    pub fn set_tool_scope(&self, tool_ids: Vec<String>) -> u64 {
437        let _guard = self.state.snapshot_guard.write();
438        *self.state.tool_scope_override.write() = Some(tool_ids);
439        self.bump()
440    }
441
442    /// Clears live narrowing so later calls return to the runtime's declared grant.
443    pub fn clear_tool_scope_override(&self) -> u64 {
444        let _guard = self.state.snapshot_guard.write();
445        *self.state.tool_scope_override.write() = None;
446        self.bump()
447    }
448
449    /// Enables or disables emergency denial for future tool calls.
450    pub fn set_emergency_deny(&self, enabled: bool) -> u64 {
451        let _guard = self.state.snapshot_guard.write();
452        self.state.emergency_deny.store(enabled, Ordering::SeqCst);
453        self.bump()
454    }
455
456    /// Denies future calls and asks active cancellable work to stop.
457    pub fn cancel_all(&self) -> u64 {
458        self.set_emergency_deny(true)
459    }
460}
461
462pub struct RuntimeAgent {
463    info: AgentInfo,
464    llm_registry: Arc<LLMRegistry>,
465    memory: Arc<dyn Memory>,
466    tools: Arc<ToolRegistry>,
467    skills: Vec<SkillDefinition>,
468    skill_router: Option<SkillRouter>,
469    skill_executor: Option<SkillExecutor>,
470    base_system_prompt: String,
471    max_iterations: u32,
472    iteration_count: RwLock<u32>,
473    max_context_tokens: u32,
474    memory_token_budget: Option<MemoryTokenBudget>,
475    recovery_manager: RecoveryManager,
476    tool_security: ToolSecurityEngine,
477    process_processor: Option<ProcessProcessor>,
478    message_filters: RwLock<HashMap<String, Arc<dyn MessageFilter>>>,
479    state_machine: Option<Arc<StateMachine>>,
480    transition_evaluator: Option<Arc<dyn TransitionEvaluator>>,
481    context_manager: Arc<ContextManager>,
482    template_renderer: TemplateRenderer,
483    tool_call_history: RwLock<Vec<ToolCallRecord>>,
484    parallel_tools: ParallelToolsConfig,
485    streaming: StreamingConfig,
486    hooks: Arc<dyn AgentHooks>,
487    hitl_engine: Option<HITLEngine>,
488    approval_handler: Arc<dyn ApprovalHandler>,
489    storage_config: StorageConfig,
490    storage: RwLock<Option<Arc<dyn AgentStorage>>>,
491    storage_init: tokio::sync::Mutex<()>,
492    reasoning_config: ReasoningConfig,
493    reflection_config: ReflectionConfig,
494    disambiguation_manager: Option<DisambiguationManager>,
495    /// Generation invalidating confirmation work across reset and state changes.
496    disambiguation_epoch: AtomicU64,
497    /// Serializes confirmation redispatch admission with reset and state mutation.
498    disambiguation_admission: tokio::sync::RwLock<()>,
499    /// Reserves one runtime-owned transition before its exit actions can produce side effects.
500    state_transition_reserved: AtomicBool,
501    /// Structured persona manager for identity, evolution, and secrets.
502    persona_manager: Option<Arc<ai_agents_persona::PersonaManager>>,
503    /// Skill ID that triggered the current pending disambiguation.
504    /// Set by try_skill_route() when skill-level disambiguation triggers clarification.
505    /// Read by run_loop() when clarification resolves to route directly to the skill.
506    pending_skill_id: RwLock<Option<String>>,
507    current_plan: RwLock<Option<Plan>>,
508    /// Tool IDs declared in the top-level `tools:` spec.
509    declared_tool_ids: Option<Vec<String>>,
510    /// Whether the context manager has been initialized (defaults loaded, env resolved, etc.)
511    context_initialized: AtomicBool,
512    /// Spawner for dynamic agent creation (set when YAML has a spawner: section).
513    spawner: Option<Arc<crate::spawner::AgentSpawner>>,
514    /// Registry tracking spawned agents (set when YAML has a spawner: section).
515    spawner_registry: Option<Arc<crate::spawner::AgentRegistry>>,
516    /// Re-dispatch depth for post-transition full dispatch.
517    /// 0 = not re-dispatching. > 0 = user message already in memory, skip re-adding.
518    redispatch_depth: RwLock<u32>,
519    /// Active optimized turn context used to keep root lifecycle state in one place.
520    active_turn_context: RwLock<Option<TurnOptimizationContext>>,
521    /// Tracks whether the root turn already wrote the processed user message.
522    root_user_message_committed: AtomicBool,
523    /// Current actor ID for cross-session memory.
524    actor_id: RwLock<Option<String>>,
525    /// Fact store for managing per-actor extracted facts.
526    fact_store: RwLock<Option<Arc<ai_agents_facts::FactStore>>>,
527    /// Fact extractor for LLM-based fact extraction.
528    /// None when actor_memory is enabled without facts.enabled.
529    fact_extractor: RwLock<Option<Arc<dyn ai_agents_facts::FactExtractor>>>,
530    /// Cached actor facts keyed by actor ID so concurrent or alternating turns do not overwrite one another.
531    actor_facts_cache: Arc<RwLock<HashMap<String, Vec<ai_agents_core::KeyFact>>>>,
532    /// Number of messages since last fact extraction.
533    messages_since_extraction: Arc<RwLock<usize>>,
534    /// Actor memory configuration.
535    actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
536    /// Facts configuration.
537    facts_config: Option<ai_agents_facts::FactsConfig>,
538    /// Session-scoped metadata (tags, ttl, actor roster).
539    session_metadata: RwLock<ai_agents_core::SessionMetadata>,
540    /// Session id currently bound to this runtime instance.
541    current_session_id: RwLock<Option<String>>,
542    /// Relationship manager for actor-scoped social memory.
543    relationship_manager: Option<Arc<RelationshipManager>>,
544    /// Observability manager for traces, metrics, reports, and exports.
545    observability_manager: Option<Arc<ObservabilityManager>>,
546    /// Runtime optimization and maintenance policy.
547    runtime_config: RuntimeConfig,
548    /// Queue for background maintenance tasks.
549    background_maintenance: Arc<BackgroundMaintenanceQueue>,
550    /// Cross-tool locks for side-effecting calls that target the same resource.
551    resource_locks: ToolResourceLocks,
552    /// Host-only runtime control state.
553    runtime_control: Arc<RuntimeControlState>,
554    /// Serializes independent externally initiated root turns across blocking and streaming APIs.
555    root_turn_gate: RootTurnGate,
556}
557
558impl std::fmt::Debug for RuntimeAgent {
559    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
560        f.debug_struct("RuntimeAgent")
561            .field("info", &self.info)
562            .field("base_system_prompt", &self.base_system_prompt)
563            .field("max_iterations", &self.max_iterations)
564            .field("skills_count", &self.skills.len())
565            .field("max_context_tokens", &self.max_context_tokens)
566            .field("has_state_machine", &self.state_machine.is_some())
567            .field("parallel_tools", &self.parallel_tools)
568            .field("streaming", &self.streaming)
569            .field("has_hooks", &true)
570            .field("has_hitl", &self.hitl_engine.is_some())
571            .field("storage_type", &self.storage_config.storage_type())
572            .field("reasoning_mode", &self.reasoning_config.mode)
573            .field("reflection_enabled", &self.reflection_config.enabled)
574            .field("declared_tool_ids", &self.declared_tool_ids)
575            .field("has_persona", &self.persona_manager.is_some())
576            .field("has_observability", &self.observability_manager.is_some())
577            .finish_non_exhaustive()
578    }
579}
580
581struct ObservabilityClarificationObserver;
582
583impl ClarificationObserver for ObservabilityClarificationObserver {
584    /// Scopes clarification question generation as disambiguation_clarification.
585    fn observe_question<'a>(
586        &'a self,
587        future: ClarificationQuestionFuture<'a>,
588    ) -> ClarificationQuestionFuture<'a> {
589        Box::pin(async move {
590            with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
591        })
592    }
593
594    /// Scopes clarification response parsing as disambiguation_clarification.
595    fn observe_parse<'a>(
596        &'a self,
597        future: ClarificationParseFuture<'a>,
598    ) -> ClarificationParseFuture<'a> {
599        Box::pin(async move {
600            with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
601        })
602    }
603
604    /// Scopes semantic confirmation parsing as disambiguation_clarification.
605    fn observe_confirmation_parse<'a>(
606        &'a self,
607        future: ConfirmationParseFuture<'a>,
608    ) -> ConfirmationParseFuture<'a> {
609        Box::pin(async move {
610            with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
611        })
612    }
613}
614
615struct ObservabilityProcessStageObserver;
616
617impl ProcessStageObserver for ObservabilityProcessStageObserver {
618    /// Scopes one process stage using the purpose implied by its stage type.
619    fn observe<'a>(
620        &'a self,
621        hint: ProcessPurposeHint,
622        future: ProcessStageFuture<'a>,
623    ) -> ProcessStageFuture<'a> {
624        Box::pin(async move {
625            with_observation_purpose(observation_purpose_for_process(hint), future).await
626        })
627    }
628}
629
630struct RegistryLLMGetter {
631    registry: Arc<LLMRegistry>,
632}
633
634impl LLMGetter for RegistryLLMGetter {
635    fn get_llm(&self, alias: &str) -> Option<Arc<dyn LLMProvider>> {
636        self.registry.get(alias).ok()
637    }
638}
639
640impl RuntimeAgent {
641    /// Constructs a runtime with one gate shared by every external root-turn entry point.
642    #[allow(clippy::too_many_arguments)]
643    pub fn new(
644        info: AgentInfo,
645        llm_registry: Arc<LLMRegistry>,
646        memory: Arc<dyn Memory>,
647        tools: Arc<ToolRegistry>,
648        skills: Vec<SkillDefinition>,
649        system_prompt: String,
650        max_iterations: u32,
651    ) -> Self {
652        let (skill_router, skill_executor) = if !skills.is_empty() {
653            let router_llm = llm_registry.router().ok();
654            let router = router_llm.map(|llm| SkillRouter::new(llm, skills.clone()));
655            let executor = SkillExecutor::new(llm_registry.clone(), tools.clone());
656            (router, Some(executor))
657        } else {
658            (None, None)
659        };
660
661        let context_manager =
662            ContextManager::new(HashMap::new(), info.name.clone(), info.version.clone());
663
664        Self {
665            info,
666            llm_registry,
667            memory,
668            tools,
669            skills,
670            skill_router,
671            skill_executor,
672            base_system_prompt: system_prompt,
673            max_iterations,
674            iteration_count: RwLock::new(0),
675            max_context_tokens: 128000,
676            memory_token_budget: None,
677            recovery_manager: RecoveryManager::default(),
678            tool_security: ToolSecurityEngine::default(),
679            process_processor: None,
680            message_filters: RwLock::new(HashMap::new()),
681            state_machine: None,
682            transition_evaluator: None,
683            context_manager: Arc::new(context_manager),
684            template_renderer: TemplateRenderer::new(),
685            tool_call_history: RwLock::new(Vec::new()),
686            parallel_tools: ParallelToolsConfig::default(),
687            streaming: StreamingConfig::default(),
688            hooks: Arc::new(NoopHooks),
689            hitl_engine: None,
690            approval_handler: Arc::new(RejectAllHandler::new()),
691            storage_config: StorageConfig::default(),
692            storage: RwLock::new(None),
693            storage_init: tokio::sync::Mutex::new(()),
694            reasoning_config: ReasoningConfig::default(),
695            reflection_config: ReflectionConfig::default(),
696            disambiguation_manager: None,
697            disambiguation_epoch: AtomicU64::new(0),
698            disambiguation_admission: tokio::sync::RwLock::new(()),
699            state_transition_reserved: AtomicBool::new(false),
700            persona_manager: None,
701            pending_skill_id: RwLock::new(None),
702            current_plan: RwLock::new(None),
703            declared_tool_ids: None,
704            context_initialized: AtomicBool::new(false),
705            spawner: None,
706            spawner_registry: None,
707            redispatch_depth: RwLock::new(0),
708            active_turn_context: RwLock::new(None),
709            root_user_message_committed: AtomicBool::new(false),
710            actor_id: RwLock::new(None),
711            fact_store: RwLock::new(None),
712            fact_extractor: RwLock::new(None),
713            actor_facts_cache: Arc::new(RwLock::new(HashMap::new())),
714            messages_since_extraction: Arc::new(RwLock::new(0)),
715            actor_memory_config: None,
716            facts_config: None,
717            session_metadata: RwLock::new(ai_agents_core::SessionMetadata::default()),
718            current_session_id: RwLock::new(None),
719            relationship_manager: None,
720            observability_manager: None,
721            runtime_config: RuntimeConfig::default(),
722            background_maintenance: Arc::new(BackgroundMaintenanceQueue::default()),
723            resource_locks: new_tool_resource_locks(),
724            runtime_control: Arc::new(RuntimeControlState::default()),
725            root_turn_gate: Arc::new(tokio::sync::Mutex::new(())),
726        }
727    }
728
729    pub fn with_declared_tool_ids(mut self, ids: Option<Vec<String>>) -> Self {
730        self.declared_tool_ids = ids;
731        self
732    }
733
734    pub fn with_storage_config(mut self, config: StorageConfig) -> Self {
735        self.storage_config = config;
736        self
737    }
738
739    pub fn with_storage(self, storage: Arc<dyn AgentStorage>) -> Self {
740        *self.storage.write() = Some(storage);
741        self
742    }
743
744    pub(crate) fn with_shared_resource_locks(mut self, locks: ToolResourceLocks) -> Self {
745        self.resource_locks = locks;
746        self
747    }
748
749    pub fn with_reasoning(mut self, config: ReasoningConfig) -> Self {
750        self.reasoning_config = config;
751        self
752    }
753
754    pub fn with_reflection(mut self, config: ReflectionConfig) -> Self {
755        self.reflection_config = config;
756        self
757    }
758
759    /// Attach a relationship manager configured by the builder or host application.
760    pub fn with_relationships(mut self, manager: Arc<RelationshipManager>) -> Self {
761        self.relationship_manager = Some(manager);
762        self
763    }
764
765    /// Attach a shared observability manager for traces, metrics, reports, and exports.
766    pub fn with_observability(mut self, manager: Arc<ObservabilityManager>) -> Self {
767        self.observability_manager = Some(manager);
768        self
769    }
770
771    /// Attach runtime optimization policy and resize the background queue.
772    pub fn with_runtime_config(mut self, config: RuntimeConfig) -> Self {
773        let max_tasks = config.optimization.post_turn.max_background_tasks;
774        self.background_maintenance = Arc::new(BackgroundMaintenanceQueue::new(max_tasks));
775        self.runtime_config = config;
776        self
777    }
778
779    /// Returns the runtime optimization policy.
780    pub fn runtime_config(&self) -> &RuntimeConfig {
781        &self.runtime_config
782    }
783
784    /// Wait for all background maintenance tasks to finish.
785    pub async fn flush_background_tasks(&self) -> Result<()> {
786        self.background_maintenance.flush_all().await
787    }
788
789    /// Wait for background maintenance associated with one actor to finish.
790    pub async fn flush_background_tasks_for_actor(&self, actor_id: &str) -> Result<()> {
791        self.background_maintenance.flush_scope(actor_id).await
792    }
793
794    /// Wait for background maintenance associated with one task kind to finish.
795    pub async fn flush_background_tasks_for_purpose(
796        &self,
797        purpose: RuntimeTaskPurpose,
798    ) -> Result<()> {
799        self.background_maintenance.flush_purpose(purpose).await
800    }
801
802    /// Wait for background maintenance associated with one actor and task kind to finish.
803    pub async fn flush_background_tasks_for_actor_purpose(
804        &self,
805        actor_id: &str,
806        purpose: RuntimeTaskPurpose,
807    ) -> Result<()> {
808        self.background_maintenance
809            .flush_scope_purpose(actor_id, purpose)
810            .await
811    }
812
813    /// Flush background maintenance before a host shuts down the runtime.
814    pub async fn shutdown_background_tasks(&self) -> Result<()> {
815        self.flush_background_tasks().await
816    }
817
818    /// Returns the configured observability manager for report and export access.
819    pub fn observability(&self) -> Option<Arc<ObservabilityManager>> {
820        self.observability_manager.clone()
821    }
822
823    /// Exports observability files after a turn when export settings request it.
824    async fn export_observability_if_configured(&self) {
825        let Some(manager) = self.observability_manager.as_ref() else {
826            return;
827        };
828        let export = &manager.config().export;
829        if !export.write_report && !export.write_raw_events {
830            return;
831        }
832        if let Err(error) = manager.export().await {
833            warn!(error = %error, "Observability export failed");
834        }
835    }
836
837    /// Returns the configured relationship manager, if relationship memory is enabled.
838    pub fn relationship_manager(&self) -> Option<Arc<RelationshipManager>> {
839        self.relationship_manager.clone()
840    }
841
842    fn current_turn_actor_context(&self) -> Option<crate::TurnActorContext> {
843        current_turn_actor_context()
844    }
845
846    fn effective_actor_id(&self) -> Option<String> {
847        self.current_turn_actor_context()
848            .and_then(|ctx| ctx.effective_actor_id().map(|id| id.to_string()))
849            .or_else(|| self.actor_id.read().clone())
850    }
851
852    fn effective_origin_actor_id(&self) -> Option<String> {
853        self.current_turn_actor_context()
854            .and_then(|ctx| ctx.origin_actor_id.clone())
855            .or_else(|| self.actor_id.read().clone())
856    }
857
858    fn record_session_actor_if_needed(&self) {
859        if let Some(actor_id) = self.effective_origin_actor_id() {
860            let mut meta = self.session_metadata.write();
861            meta.actor_id = Some(actor_id.clone());
862            if !meta.actors.iter().any(|a| a == &actor_id) {
863                meta.actors.push(actor_id);
864            }
865        }
866    }
867
868    fn outbound_actor_context(&self) -> crate::TurnActorContext {
869        let mut context = self.current_turn_actor_context().unwrap_or_default();
870        if context.origin_actor_id.is_none() {
871            context.origin_actor_id = self.effective_origin_actor_id();
872        }
873        context.sender_agent_id = Some(self.info.id.clone());
874        context
875    }
876
877    /// Returns the current session ID or creates one for unsaved observed turns.
878    fn observation_session_id(&self) -> Option<String> {
879        let mut current = self.current_session_id.write();
880        if current.is_none() {
881            *current = Some(new_observation_session_id());
882        }
883        current.clone()
884    }
885
886    /// Builds the root or child observation context for a chat entry point.
887    fn build_observation_context(&self, actor_id: Option<String>) -> Option<SpanContext> {
888        let manager = self.observability_manager.as_ref()?;
889        let context = self.build_context_with_overlays();
890        let language = resolve_language_from_context(manager.config(), &context);
891        let context = current_observation_context()
892            .map(|parent| parent.child_for_agent(self.info.id.clone()).with_new_turn())
893            .unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
894        Some(
895            context
896                .with_actor(actor_id.or_else(|| self.effective_actor_id()))
897                .with_session(self.observation_session_id())
898                .with_state(self.current_state())
899                .with_language(Some(language)),
900        )
901    }
902
903    /// Refreshes task-local context with current runtime labels and a purpose.
904    fn current_runtime_observation_context(
905        &self,
906        purpose: ObservationPurpose,
907    ) -> Option<SpanContext> {
908        let manager = self.observability_manager.as_ref()?;
909        let context = self.build_context_with_overlays();
910        let language = resolve_language_from_context(manager.config(), &context);
911        let mut observation = current_observation_context()
912            .unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
913        observation.agent_id = self.info.id.clone();
914        observation.actor_id = self.effective_actor_id();
915        observation.session_id = self.observation_session_id();
916        observation.state = self.current_state();
917        observation.language = Some(language);
918        observation.purpose = purpose;
919        Some(observation)
920    }
921
922    /// Runs a future under a purpose while preserving current trace context.
923    async fn observe_purpose<F, T>(&self, purpose: ObservationPurpose, future: F) -> T
924    where
925        F: Future<Output = T>,
926    {
927        if let Some(context) = self.current_runtime_observation_context(purpose) {
928            with_observation_context(context, future).await
929        } else {
930            future.await
931        }
932    }
933
934    /// Runs one actor-scoped external root turn with gate ownership visible through finalization, hooks, orchestration, and export.
935    ///
936    /// Internal redispatch calls `run_loop` directly and therefore keeps the original guard and identity stack instead of acquiring this non-reentrant gate again.
937    fn chat_with_actor_context_boxed<'a>(
938        &'a self,
939        input: &'a str,
940        actor_context: crate::TurnActorContext,
941    ) -> Pin<Box<dyn Future<Output = Result<AgentResponse>> + Send + 'a>> {
942        Box::pin(async move {
943            let RootTurnAdmission {
944                guard,
945                identity_stack,
946            } = self.acquire_root_turn().await?;
947            let result = scope_runtime_gate_identity_stack(&identity_stack, async move {
948                let actor_id = actor_context.effective_actor_id().map(str::to_string);
949                let run = async move {
950                    scope_actor_context(
951                        actor_context,
952                        Box::pin(async move { self.run_loop(input).await }),
953                    )
954                    .await
955                };
956                let result = if let Some(context) = self.build_observation_context(actor_id) {
957                    with_observation_context(context, run).await
958                } else {
959                    run.await
960                };
961                self.export_observability_if_configured().await;
962                result
963            })
964            .await;
965            drop(guard);
966            result
967        })
968    }
969
970    /// Rejects recursive ownership before waiting, then acquires the external root-turn gate and builds one immutable extended ancestry snapshot.
971    async fn acquire_root_turn(&self) -> Result<RootTurnAdmission> {
972        let gate_identity = Arc::clone(&self.root_turn_gate);
973        let current_identity_stack = current_runtime_gate_identity_stack();
974        if current_identity_stack
975            .iter()
976            .any(|owned_gate| Arc::ptr_eq(owned_gate, &gate_identity))
977        {
978            return Err(AgentError::Other(format!(
979                "RuntimeAgent '{}' rejected reentrant root turn ownership",
980                self.info.id
981            )));
982        }
983        let guard = Arc::clone(&gate_identity).lock_owned().await;
984        //
985        // Root admission is the only place that extends ancestry, so later scopes can preserve the complete chain with an `Arc` clone.
986        //
987        let mut identity_stack = Vec::with_capacity(current_identity_stack.len() + 1);
988        identity_stack.extend(current_identity_stack.iter().cloned());
989        identity_stack.push(gate_identity);
990        Ok(RootTurnAdmission {
991            guard,
992            identity_stack: identity_stack.into(),
993        })
994    }
995
996    /// Run one serialized turn with turn-scoped actor context without mutating the runtime's global actor ID.
997    ///
998    /// The supplied context is available to actor-scoped facts, relationship memory, orchestration, and prompt templates only for the lifetime of this call.
999    pub async fn chat_with_actor_context(
1000        &self,
1001        input: &str,
1002        actor_context: crate::TurnActorContext,
1003    ) -> Result<AgentResponse> {
1004        self.chat_with_actor_context_boxed(input, actor_context)
1005            .await
1006    }
1007
1008    /// Convenience wrapper around [`Self::chat_with_actor_context`] for a turn whose original actor is known up front.
1009    pub async fn chat_as_actor(&self, actor_id: &str, input: &str) -> Result<AgentResponse> {
1010        let actor_context = crate::TurnActorContext::new().with_origin_actor(actor_id);
1011        self.chat_with_actor_context(input, actor_context).await
1012    }
1013
1014    /// Ensure the effective actor's relationship is loaded from storage into the relationship manager.
1015    pub async fn load_actor_relationship(&self) -> Result<()> {
1016        self.maybe_load_actor_relationship().await;
1017        Ok(())
1018    }
1019
1020    /// Manually apply a delta to the effective actor's `agent_to_actor` relationship perspective and persist the updated relationship when storage is configured.
1021    pub async fn update_relationship_dimension(
1022        &self,
1023        dimension: &str,
1024        delta: f64,
1025        reason: Option<&str>,
1026    ) -> Result<ai_agents_relationships::DimensionChange> {
1027        self.update_relationship_dimension_for_perspective(
1028            ai_agents_relationships::RelationshipPerspective::AgentToActor,
1029            dimension,
1030            delta,
1031            reason,
1032        )
1033        .await
1034    }
1035
1036    /// Manually apply a delta to a specific relationship perspective for the effective actor.
1037    ///
1038    /// Use this for two-sided configurations when you need to update `agent_to_actor`, `perceived_actor_to_agent`, or `mutual` explicitly from application logic.
1039    pub async fn update_relationship_dimension_for_perspective(
1040        &self,
1041        perspective: ai_agents_relationships::RelationshipPerspective,
1042        dimension: &str,
1043        delta: f64,
1044        reason: Option<&str>,
1045    ) -> Result<ai_agents_relationships::DimensionChange> {
1046        let manager = self
1047            .relationship_manager
1048            .as_ref()
1049            .ok_or_else(|| AgentError::Config("Relationship memory is not configured".into()))?;
1050        let actor_id = self.effective_actor_id().ok_or_else(|| {
1051            AgentError::Config("No actor ID set. Use set_actor_id() first".into())
1052        })?;
1053        let change = manager.update_dimension_for_perspective(
1054            &actor_id,
1055            perspective,
1056            dimension,
1057            delta,
1058            1.0,
1059            reason.unwrap_or("manual relationship update"),
1060        )?;
1061        self.persist_actor_relationship(&actor_id).await?;
1062        info!(
1063            actor_id = %actor_id,
1064            perspective = %change.perspective,
1065            dimension = %change.dimension,
1066            delta = change.delta,
1067            current = change.current,
1068            "relationship updated manually"
1069        );
1070        self.hooks
1071            .on_relationship_change(&actor_id, std::slice::from_ref(&change))
1072            .await;
1073        Ok(change)
1074    }
1075
1076    pub fn reasoning_config(&self) -> &ReasoningConfig {
1077        &self.reasoning_config
1078    }
1079
1080    pub fn reflection_config(&self) -> &ReflectionConfig {
1081        &self.reflection_config
1082    }
1083
1084    /// Set only the actor memory and facts configs without creating the store.
1085    /// The store and extractor are created lazily in init_storage().
1086    pub fn with_facts_config(
1087        mut self,
1088        actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
1089        facts_config: Option<ai_agents_facts::FactsConfig>,
1090    ) -> Self {
1091        self.actor_memory_config = actor_memory_config;
1092        self.facts_config = facts_config;
1093        self
1094    }
1095
1096    /// Configure fact store and optional extractor for actor memory.
1097    /// Pass `None` for `extractor` to load existing facts without running extraction.
1098    pub fn with_facts(
1099        mut self,
1100        store: Arc<ai_agents_facts::FactStore>,
1101        extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>>,
1102        actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
1103        facts_config: Option<ai_agents_facts::FactsConfig>,
1104    ) -> Self {
1105        *self.fact_store.write() = Some(store);
1106        *self.fact_extractor.write() = extractor;
1107        self.actor_memory_config = actor_memory_config;
1108        self.facts_config = facts_config;
1109        self
1110    }
1111
1112    /// Get the fact store for direct fact manipulation.
1113    pub fn fact_store(&self) -> Option<Arc<ai_agents_facts::FactStore>> {
1114        self.fact_store.read().clone()
1115    }
1116
1117    /// Get the current actor ID.
1118    pub fn actor_id(&self) -> Option<String> {
1119        self.actor_id.read().clone()
1120    }
1121
1122    /// Set the current actor ID (player, user, another agent, etc.).
1123    pub fn set_actor_id(&self, actor_id: &str) -> ai_agents_core::Result<()> {
1124        *self.actor_id.write() = Some(actor_id.to_string());
1125        {
1126            let mut meta = self.session_metadata.write();
1127            meta.actor_id = Some(actor_id.to_string());
1128            if !meta.actors.iter().any(|a| a == actor_id) {
1129                meta.actors.push(actor_id.to_string());
1130            }
1131        }
1132        Ok(())
1133    }
1134
1135    /// Clear the current actor binding while retaining the session actor roster.
1136    pub fn clear_actor_id(&self) {
1137        *self.actor_id.write() = None;
1138        self.session_metadata.write().actor_id = None;
1139    }
1140
1141    /// Set the current actor ID. Convenience wrapper around set_actor_id.
1142    pub fn set_user_id(&self, user_id: &str) -> ai_agents_core::Result<()> {
1143        self.set_actor_id(user_id)
1144    }
1145
1146    /// Load facts for the current actor from storage and cache them for prompt injection.
1147    pub async fn load_actor_memory(&self) -> ai_agents_core::Result<()> {
1148        let actor_id = match self.effective_actor_id() {
1149            Some(id) => id,
1150            None => return Ok(()),
1151        };
1152
1153        let store_opt = self.fact_store.read().clone();
1154        if let Some(store) = store_opt {
1155            let facts = store.get_facts(&actor_id).await?;
1156            let count = facts.len();
1157            self.actor_facts_cache
1158                .write()
1159                .insert(actor_id.clone(), facts);
1160            self.hooks.on_actor_memory_loaded(&actor_id, count).await;
1161            tracing::debug!("loaded {} facts for actor {}", count, actor_id);
1162        }
1163
1164        Ok(())
1165    }
1166
1167    /// Load actor memory only when the effective actor has no cached facts yet.
1168    async fn maybe_load_actor_memory(&self) {
1169        let Some(actor_id) = self.effective_actor_id() else {
1170            return;
1171        };
1172        if self.actor_facts_cache.read().contains_key(&actor_id) {
1173            return;
1174        }
1175        let _ = self.load_actor_memory().await;
1176    }
1177
1178    /// Pre-turn lifecycle shared by streaming and non-streaming paths.
1179    async fn pre_turn_session_lifecycle(&self) {
1180        if *self.redispatch_depth.read() > 0 {
1181            return;
1182        }
1183        self.resolve_actor_id_from_context();
1184        self.await_background_before_next_turn().await;
1185        self.record_session_actor_if_needed();
1186        self.maybe_load_actor_memory().await;
1187        self.maybe_load_actor_relationship().await;
1188        *self.messages_since_extraction.write() += 1;
1189    }
1190
1191    /// Post-turn lifecycle shared by streaming and non-streaming paths.
1192    async fn post_turn_session_lifecycle(&self) -> Result<()> {
1193        if *self.redispatch_depth.read() > 0 {
1194            return Ok(());
1195        }
1196        *self.messages_since_extraction.write() += 1;
1197        self.run_post_turn_maintenance().await
1198    }
1199
1200    /// Starts root-turn bookkeeping for user-message commit tracking.
1201    fn begin_root_turn(&self) {
1202        if *self.redispatch_depth.read() == 0 {
1203            let mut guard = self.active_turn_context.write();
1204            if guard.is_none() {
1205                self.root_user_message_committed
1206                    .store(false, Ordering::SeqCst);
1207                let max_calls = self
1208                    .runtime_config
1209                    .optimization
1210                    .max_speculative_llm_calls_per_turn;
1211                *guard = Some(TurnOptimizationContext::new(
1212                    String::new(),
1213                    HashMap::new(),
1214                    max_calls,
1215                ));
1216            }
1217        }
1218    }
1219
1220    fn update_active_turn_context(
1221        &self,
1222        processed_input: &str,
1223        input_context: HashMap<String, Value>,
1224    ) {
1225        if *self.redispatch_depth.read() > 0 {
1226            return;
1227        }
1228        let max_calls = self
1229            .runtime_config
1230            .optimization
1231            .max_speculative_llm_calls_per_turn;
1232        let mut guard = self.active_turn_context.write();
1233        match guard.as_mut() {
1234            Some(context) => {
1235                context.processed_input = processed_input.to_string();
1236                context.input_context = input_context;
1237                context.max_speculative_llm_calls = max_calls;
1238            }
1239            None => {
1240                *guard = Some(TurnOptimizationContext::new(
1241                    processed_input,
1242                    input_context,
1243                    max_calls,
1244                ));
1245            }
1246        }
1247    }
1248
1249    /// Writes the processed user message once for the root turn.
1250    async fn commit_root_user_message(&self, processed_input: &str) -> Result<()> {
1251        if *self.redispatch_depth.read() > 0 {
1252            return Ok(());
1253        }
1254        if !self
1255            .root_user_message_committed
1256            .swap(true, Ordering::SeqCst)
1257        {
1258            self.memory
1259                .add_message(ChatMessage::user(processed_input))
1260                .await?;
1261            if let Some(context) = self.active_turn_context.write().as_mut() {
1262                context.mark_user_message_committed();
1263            }
1264        }
1265        Ok(())
1266    }
1267
1268    /// Clears root-turn bookkeeping after final response handling.
1269    fn end_root_turn(&self) {
1270        if *self.redispatch_depth.read() == 0 {
1271            self.root_user_message_committed
1272                .store(false, Ordering::SeqCst);
1273            *self.active_turn_context.write() = None;
1274        }
1275    }
1276
1277    fn reserve_active_speculative_llm_call(&self, kind: RuntimeOptimizationKind) -> bool {
1278        self.begin_root_turn();
1279        let mut guard = self.active_turn_context.write();
1280        let Some(context) = guard.as_mut() else {
1281            return false;
1282        };
1283        context.reserve_speculative_llm_call_for(kind)
1284    }
1285
1286    fn branch_context_preview(&self) -> String {
1287        let context = self.build_context_with_overlays();
1288        let mut value = serde_json::to_string_pretty(&context).unwrap_or_else(|_| "{}".to_string());
1289        const MAX_CONTEXT_PREVIEW_CHARS: usize = 2048;
1290        if value.chars().count() > MAX_CONTEXT_PREVIEW_CHARS {
1291            value = value
1292                .chars()
1293                .take(MAX_CONTEXT_PREVIEW_CHARS)
1294                .collect::<String>();
1295            value.push_str("...");
1296        }
1297        value
1298    }
1299
1300    /// Applies freshness policy before rendering the next prompt.
1301    async fn await_background_before_next_turn(&self) {
1302        let optimization = &self.runtime_config.optimization;
1303        if !optimization.enabled {
1304            return;
1305        }
1306        let actor_id = self.effective_actor_id();
1307        let post = &optimization.post_turn;
1308        self.await_background_task(
1309            post.facts.await_before_next_turn,
1310            RuntimeTaskPurpose::PostTurnFacts,
1311            actor_id.as_deref(),
1312            "facts",
1313        )
1314        .await;
1315        self.await_background_task(
1316            post.relationships.await_before_next_turn,
1317            RuntimeTaskPurpose::PostTurnRelationship,
1318            actor_id.as_deref(),
1319            "relationships",
1320        )
1321        .await;
1322    }
1323
1324    async fn await_background_task(
1325        &self,
1326        policy: AwaitBeforeNextTurn,
1327        purpose: RuntimeTaskPurpose,
1328        actor_id: Option<&str>,
1329        label: &str,
1330    ) {
1331        match policy {
1332            AwaitBeforeNextTurn::Never => {}
1333            AwaitBeforeNextTurn::Always => {
1334                if let Err(error) = self.flush_background_tasks_for_purpose(purpose).await {
1335                    warn!(label = label, error = %error, "background maintenance flush failed");
1336                }
1337            }
1338            AwaitBeforeNextTurn::SameActor => {
1339                if let Some(actor_id) = actor_id
1340                    && let Err(error) = self
1341                        .flush_background_tasks_for_actor_purpose(actor_id, purpose)
1342                        .await
1343                {
1344                    warn!(label = label, actor_id = %actor_id, error = %error, "actor background maintenance flush failed");
1345                }
1346            }
1347        }
1348    }
1349
1350    /// Runs post-turn facts and relationship maintenance according to runtime policy.
1351    async fn run_post_turn_maintenance(&self) -> Result<()> {
1352        let optimization = &self.runtime_config.optimization;
1353        if !optimization.enabled {
1354            self.auto_extract_facts().await;
1355            self.auto_update_relationship().await;
1356            return Ok(());
1357        }
1358
1359        let facts_mode = effective_maintenance_mode(
1360            optimization.post_turn.facts.mode,
1361            optimization.parallel_post_turn_memory,
1362        );
1363        let relationships_mode = effective_maintenance_mode(
1364            optimization.post_turn.relationships.mode,
1365            optimization.parallel_post_turn_memory,
1366        );
1367
1368        match (facts_mode, relationships_mode) {
1369            (MaintenanceMode::InlineSerial, MaintenanceMode::InlineSerial) => {
1370                self.auto_extract_facts().await;
1371                self.auto_update_relationship().await;
1372            }
1373            (MaintenanceMode::InlineParallel, MaintenanceMode::InlineParallel) => {
1374                let facts = self.auto_extract_facts();
1375                let relationships = self.auto_update_relationship();
1376                tokio::join!(facts, relationships);
1377            }
1378            (MaintenanceMode::Background, MaintenanceMode::Background) => {
1379                self.schedule_facts_background().await?;
1380                self.schedule_relationship_background().await?;
1381            }
1382            (MaintenanceMode::Background, MaintenanceMode::InlineParallel)
1383            | (MaintenanceMode::Background, MaintenanceMode::InlineSerial) => {
1384                self.schedule_facts_background().await?;
1385                self.auto_update_relationship().await;
1386            }
1387            (MaintenanceMode::InlineParallel, MaintenanceMode::Background)
1388            | (MaintenanceMode::InlineSerial, MaintenanceMode::Background) => {
1389                self.auto_extract_facts().await;
1390                self.schedule_relationship_background().await?;
1391            }
1392            _ => {
1393                self.auto_extract_facts().await;
1394                self.auto_update_relationship().await;
1395            }
1396        }
1397        Ok(())
1398    }
1399
1400    async fn schedule_facts_background(&self) -> Result<()> {
1401        let policy = self.runtime_config.optimization.post_turn.facts.clone();
1402        let should_extract = self
1403            .facts_config
1404            .as_ref()
1405            .map(|c| c.enabled && c.auto_extract)
1406            .unwrap_or(false);
1407        if !should_extract {
1408            return Ok(());
1409        }
1410        let msgs_since = *self.messages_since_extraction.read();
1411        if msgs_since < 2 {
1412            return Ok(());
1413        }
1414        let Some(actor_id) = self.effective_actor_id() else {
1415            self.record_skipped_maintenance(
1416                "facts",
1417                ObservationPurpose::FactsExtraction,
1418                "missing_actor",
1419                Some(&policy),
1420            );
1421            return Ok(());
1422        };
1423        let Some(extractor) = self.fact_extractor.read().clone() else {
1424            return Ok(());
1425        };
1426        let messages = match self.memory.get_messages(None).await {
1427            Ok(messages) => messages,
1428            Err(error) => {
1429                warn!(error = %error, "failed to snapshot messages for fact extraction");
1430                return Ok(());
1431            }
1432        };
1433        let recent: Vec<_> = messages
1434            .iter()
1435            .rev()
1436            .take(msgs_since)
1437            .rev()
1438            .cloned()
1439            .collect();
1440        if recent.is_empty() {
1441            return Ok(());
1442        }
1443        let existing = self
1444            .actor_facts_cache
1445            .read()
1446            .get(&actor_id)
1447            .cloned()
1448            .unwrap_or_default();
1449        let categories = self
1450            .facts_config
1451            .as_ref()
1452            .map(|c| c.custom_categories.clone())
1453            .unwrap_or_default();
1454        let store = self.fact_store.read().clone();
1455        let cache = Arc::clone(&self.actor_facts_cache);
1456        let counter = Arc::clone(&self.messages_since_extraction);
1457        let hooks = Arc::clone(&self.hooks);
1458        let agent_id = self.info.id.clone();
1459        let observation = current_observation_context();
1460        let key = MaintenanceSequenceKey::actor(
1461            agent_id,
1462            actor_id.clone(),
1463            RuntimeTaskPurpose::PostTurnFacts,
1464        );
1465        let actor_for_task = actor_id.clone();
1466        let task = async move {
1467            let run = async move {
1468                let facts = extractor
1469                    .extract(&recent, &existing, Some(&actor_for_task), &categories)
1470                    .await?;
1471                if !facts.is_empty() {
1472                    if let Some(store) = store {
1473                        let authoritative = store.add_facts(&actor_for_task, facts.clone()).await?;
1474                        cache.write().insert(actor_for_task.clone(), authoritative);
1475                    } else {
1476                        cache
1477                            .write()
1478                            .entry(actor_for_task.clone())
1479                            .or_default()
1480                            .extend(facts.clone());
1481                    }
1482                    {
1483                        let mut count = counter.write();
1484                        if *count <= msgs_since {
1485                            *count = 0;
1486                        } else {
1487                            *count -= msgs_since;
1488                        }
1489                    }
1490                    hooks.on_facts_extracted(&actor_for_task, &facts).await;
1491                }
1492                Ok(())
1493            };
1494            if let Some(context) = observation {
1495                with_observation_context(
1496                    context.with_purpose(ObservationPurpose::FactsExtraction),
1497                    run,
1498                )
1499                .await
1500            } else {
1501                run.await
1502            }
1503        };
1504        self.spawn_or_handle_background(Some(key), task, "facts", &policy)
1505            .await
1506    }
1507
1508    async fn schedule_relationship_background(&self) -> Result<()> {
1509        let policy = self
1510            .runtime_config
1511            .optimization
1512            .post_turn
1513            .relationships
1514            .clone();
1515        let Some(manager) = self.relationship_manager.as_ref().cloned() else {
1516            return Ok(());
1517        };
1518        let Some(actor_id) = self.effective_actor_id() else {
1519            self.record_skipped_maintenance(
1520                "relationships",
1521                ObservationPurpose::RelationshipUpdate,
1522                "missing_actor",
1523                Some(&policy),
1524            );
1525            return Ok(());
1526        };
1527        let recent_messages = manager.config().auto_update.recent_messages;
1528        let messages = match self.memory.get_messages(Some(recent_messages)).await {
1529            Ok(messages) => messages,
1530            Err(error) => {
1531                warn!(actor = %actor_id, error = %error, "failed to snapshot messages for relationship update");
1532                return Ok(());
1533            }
1534        };
1535        let storage = self.storage.read().clone();
1536        let hooks = Arc::clone(&self.hooks);
1537        let agent_id = self.info.id.clone();
1538        let observation = current_observation_context();
1539        let key = MaintenanceSequenceKey::actor(
1540            agent_id.clone(),
1541            actor_id.clone(),
1542            RuntimeTaskPurpose::PostTurnRelationship,
1543        );
1544        let actor_for_task = actor_id.clone();
1545        let task = async move {
1546            let run = async move {
1547                if manager.config().auto_update.enabled {
1548                    let update = manager.auto_update(&actor_for_task, &messages).await?;
1549                    if !update.changes.is_empty() {
1550                        hooks
1551                            .on_relationship_change(&actor_for_task, &update.changes)
1552                            .await;
1553                    }
1554                    if let Some(ref event) = update.event {
1555                        hooks.on_notable_event(&actor_for_task, event).await;
1556                    }
1557                }
1558                if manager.config().persistence.enabled
1559                    && let (Some(storage), Some(value)) =
1560                        (storage, manager.relationship_as_value(&actor_for_task)?)
1561                {
1562                    storage
1563                        .save_relationship(&agent_id, &actor_for_task, &value)
1564                        .await?;
1565                }
1566                Ok(())
1567            };
1568            if let Some(context) = observation {
1569                with_observation_context(
1570                    context.with_purpose(ObservationPurpose::RelationshipUpdate),
1571                    run,
1572                )
1573                .await
1574            } else {
1575                run.await
1576            }
1577        };
1578        self.spawn_or_handle_background(Some(key), task, "relationships", &policy)
1579            .await
1580    }
1581
1582    /// Queues background maintenance or applies the configured overflow behavior.
1583    async fn spawn_or_handle_background<F>(
1584        &self,
1585        key: Option<MaintenanceSequenceKey>,
1586        task: F,
1587        label: &'static str,
1588        policy: &crate::optimization::config::MaintenanceTaskPolicy,
1589    ) -> Result<()>
1590    where
1591        F: Future<Output = Result<()>> + Send + 'static,
1592    {
1593        if self.background_maintenance.is_full() {
1594            match self
1595                .runtime_config
1596                .optimization
1597                .post_turn
1598                .on_background_overflow
1599            {
1600                BackgroundOverflowPolicy::RunInline => {
1601                    record_background_maintenance_event(
1602                        self.observability_manager.as_ref(),
1603                        label,
1604                        EventStatus::Success,
1605                        0,
1606                        "inline_overflow",
1607                        None,
1608                        Some(policy),
1609                    );
1610                    let start = Instant::now();
1611                    match task.await {
1612                        Ok(()) => record_background_maintenance_event(
1613                            self.observability_manager.as_ref(),
1614                            label,
1615                            EventStatus::Success,
1616                            start.elapsed().as_millis() as u64,
1617                            "inline_completed",
1618                            None,
1619                            Some(policy),
1620                        ),
1621                        Err(error) => {
1622                            warn!(label = label, error = %error, "inline maintenance fallback failed");
1623                            record_background_maintenance_event(
1624                                self.observability_manager.as_ref(),
1625                                label,
1626                                EventStatus::Error,
1627                                start.elapsed().as_millis() as u64,
1628                                "inline_failed",
1629                                Some(error.to_string()),
1630                                Some(policy),
1631                            );
1632                            return Err(error);
1633                        }
1634                    }
1635                }
1636                BackgroundOverflowPolicy::Drop => {
1637                    self.record_skipped_maintenance(
1638                        label,
1639                        ObservationPurpose::Other(label.to_string()),
1640                        "queue_full",
1641                        Some(policy),
1642                    );
1643                }
1644                BackgroundOverflowPolicy::Error => {
1645                    record_background_maintenance_event(
1646                        self.observability_manager.as_ref(),
1647                        label,
1648                        EventStatus::Error,
1649                        0,
1650                        "queue_full",
1651                        None,
1652                        Some(policy),
1653                    );
1654                    warn!(label = label, "background maintenance queue full");
1655                    return Err(AgentError::Other(format!(
1656                        "background maintenance queue is full for {}",
1657                        label
1658                    )));
1659                }
1660            }
1661            return Ok(());
1662        }
1663
1664        record_background_maintenance_event(
1665            self.observability_manager.as_ref(),
1666            label,
1667            EventStatus::Success,
1668            0,
1669            "scheduled",
1670            None,
1671            Some(policy),
1672        );
1673        let manager = self.observability_manager.clone();
1674        let policy_for_task = policy.clone();
1675        let observed_task = async move {
1676            let start = Instant::now();
1677            let result = task.await;
1678            match &result {
1679                Ok(()) => record_background_maintenance_event(
1680                    manager.as_ref(),
1681                    label,
1682                    EventStatus::Success,
1683                    start.elapsed().as_millis() as u64,
1684                    "completed",
1685                    None,
1686                    Some(&policy_for_task),
1687                ),
1688                Err(error) => record_background_maintenance_event(
1689                    manager.as_ref(),
1690                    label,
1691                    EventStatus::Error,
1692                    start.elapsed().as_millis() as u64,
1693                    "failed",
1694                    Some(error.to_string()),
1695                    Some(&policy_for_task),
1696                ),
1697            }
1698            result
1699        };
1700
1701        if let Err(error) = self.background_maintenance.spawn(key, observed_task) {
1702            record_background_maintenance_event(
1703                self.observability_manager.as_ref(),
1704                label,
1705                EventStatus::Error,
1706                0,
1707                "spawn_failed",
1708                Some(error.to_string()),
1709                Some(policy),
1710            );
1711            warn!(label = label, error = %error, "background maintenance spawn failed");
1712            return Err(error);
1713        }
1714        Ok(())
1715    }
1716
1717    /// Records a skipped background maintenance event when work cannot run.
1718    fn record_skipped_maintenance(
1719        &self,
1720        label: &str,
1721        purpose: ObservationPurpose,
1722        reason: &str,
1723        policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
1724    ) {
1725        if let Some(manager) = self.observability_manager.as_ref() {
1726            let mut tags = background_maintenance_tags(label, "skipped", Some(reason), policy);
1727            tags.insert("runtime.skip_reason".to_string(), reason.to_string());
1728            manager.record_lifecycle_event(
1729                EventType::MemoryOperation {
1730                    operation: format!("{}_maintenance", label),
1731                },
1732                purpose,
1733                EventStatus::Skipped,
1734                0,
1735                tags,
1736                None,
1737            );
1738        }
1739    }
1740
1741    /// Get cached actor facts for the effective actor.
1742    pub fn actor_facts(&self) -> Vec<ai_agents_core::KeyFact> {
1743        let Some(actor_id) = self.effective_actor_id() else {
1744            return Vec::new();
1745        };
1746        self.actor_facts_cache
1747            .read()
1748            .get(&actor_id)
1749            .cloned()
1750            .unwrap_or_default()
1751    }
1752
1753    /// Returns the formatted relationship prompt text for the effective actor, if relationship injection produced any text for this turn.
1754    pub fn relationship_memory_text(&self) -> Option<String> {
1755        self.format_relationship_for_context().map(|(_, text)| text)
1756    }
1757
1758    /// Manually extract facts from the last N messages.
1759    pub async fn extract_facts(
1760        &self,
1761        last_n: usize,
1762    ) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
1763        self.extract_facts_with_source(last_n, "manual").await
1764    }
1765
1766    async fn extract_facts_with_source(
1767        &self,
1768        last_n: usize,
1769        source: &'static str,
1770    ) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
1771        let extractor = match self.fact_extractor.read().clone() {
1772            Some(e) => e,
1773            None => return Ok(vec![]),
1774        };
1775
1776        let messages = self.memory.get_messages(None).await?;
1777        let recent: Vec<_> = messages.iter().rev().take(last_n).rev().cloned().collect();
1778
1779        if recent.is_empty() {
1780            return Ok(vec![]);
1781        }
1782
1783        let actor_id = self.effective_actor_id();
1784        let existing = actor_id
1785            .as_ref()
1786            .and_then(|aid| self.actor_facts_cache.read().get(aid).cloned())
1787            .unwrap_or_default();
1788
1789        let categories = self
1790            .facts_config
1791            .as_ref()
1792            .map(|c| c.custom_categories.clone())
1793            .unwrap_or_default();
1794
1795        let facts = self
1796            .observe_purpose(
1797                ObservationPurpose::FactsExtraction,
1798                extractor.extract(&recent, &existing, actor_id.as_deref(), &categories),
1799            )
1800            .await?;
1801
1802        // Save to storage and update the actor-scoped cache when an actor is known.
1803        if !facts.is_empty() {
1804            let fact_store_opt = self.fact_store.read().clone();
1805            let mut stored_total = 0usize;
1806            let mut cache_updated = false;
1807            if let (Some(store), Some(aid)) = (fact_store_opt, &actor_id) {
1808                // add_facts now returns the authoritative post-write set.
1809                let authoritative = store.add_facts(aid, facts.clone()).await?;
1810                stored_total = authoritative.len();
1811                self.actor_facts_cache
1812                    .write()
1813                    .insert(aid.clone(), authoritative);
1814                cache_updated = true;
1815            } else if let Some(aid) = &actor_id {
1816                let mut cache = self.actor_facts_cache.write();
1817                let entry = cache.entry(aid.clone()).or_default();
1818                entry.extend(facts.clone());
1819                stored_total = entry.len();
1820                cache_updated = true;
1821            }
1822
1823            info!(
1824                actor_id = %actor_id.as_deref().unwrap_or("<none>"),
1825                source = source,
1826                requested_messages = last_n,
1827                message_count = recent.len(),
1828                extracted_count = facts.len(),
1829                cache_updated = cache_updated,
1830                stored_total = stored_total,
1831                "facts extracted"
1832            );
1833
1834            if let Some(ref aid) = actor_id {
1835                self.hooks.on_facts_extracted(aid, &facts).await;
1836            }
1837        }
1838
1839        Ok(facts)
1840    }
1841
1842    /// Resolve actor_id from context if method is from_context.
1843    /// Supports dotted paths (e.g. "player.id", "user.profile.id").
1844    fn resolve_actor_id_from_context(&self) {
1845        if self
1846            .current_turn_actor_context()
1847            .and_then(|ctx| ctx.effective_actor_id().map(str::to_string))
1848            .is_some()
1849        {
1850            return;
1851        }
1852
1853        if let Some(ref am_config) = self.actor_memory_config
1854            && am_config.identification.method == ai_agents_facts::IdentificationMethod::FromContext
1855            && let Some(ref path) = am_config.identification.context_path
1856        {
1857            // get_path resolves dotted paths; get only handles top-level keys.
1858            let val = self
1859                .context_manager
1860                .get_path(path)
1861                .or_else(|| self.context_manager.get(path));
1862            if let Some(val) = val
1863                && let Some(id_str) = val.as_str()
1864            {
1865                let current = self.actor_id.read().clone();
1866                if current.as_deref() != Some(id_str) {
1867                    *self.actor_id.write() = Some(id_str.to_string());
1868                    let mut meta = self.session_metadata.write();
1869                    meta.actor_id = Some(id_str.to_string());
1870                    if !meta.actors.iter().any(|a| a == id_str) {
1871                        meta.actors.push(id_str.to_string());
1872                    }
1873                }
1874            }
1875        }
1876    }
1877
1878    /// Format actor facts for template injection.
1879    fn format_actor_facts_for_context(&self) -> String {
1880        // Respect inject_in_context: false.
1881        let should_inject = self
1882            .facts_config
1883            .as_ref()
1884            .map(|c| c.inject_in_context)
1885            .unwrap_or(true);
1886        if !should_inject {
1887            return String::new();
1888        }
1889
1890        let Some(actor_id) = self.effective_actor_id() else {
1891            return String::new();
1892        };
1893
1894        let facts = self
1895            .actor_facts_cache
1896            .read()
1897            .get(&actor_id)
1898            .cloned()
1899            .unwrap_or_default();
1900        if facts.is_empty() {
1901            return String::new();
1902        }
1903
1904        let am_config = self.actor_memory_config.as_ref();
1905        // Effective token cap: prefer memory.token_budget.allocation.facts when present,
1906        // otherwise fall back to actor_memory.injection.max_tokens.
1907        let facts_budget = self
1908            .memory_token_budget
1909            .as_ref()
1910            .map(|b| b.allocation.facts as usize)
1911            .filter(|n| *n > 0);
1912        let default_max = am_config.map(|c| c.injection.max_tokens).unwrap_or(800);
1913        let max_tokens = facts_budget.unwrap_or(default_max);
1914
1915        // Filter by category when injection.mode = category.
1916        let filtered: Vec<ai_agents_core::KeyFact> = if let Some(cfg) = am_config {
1917            if cfg.injection.mode == ai_agents_facts::InjectionMode::OnDemand {
1918                return String::new();
1919            }
1920            if cfg.injection.mode == ai_agents_facts::InjectionMode::Category
1921                && !cfg.injection.categories.is_empty()
1922            {
1923                facts
1924                    .iter()
1925                    .filter(|f| {
1926                        cfg.injection
1927                            .categories
1928                            .iter()
1929                            .any(|c| f.category.to_string() == *c)
1930                    })
1931                    .cloned()
1932                    .collect()
1933            } else {
1934                facts.clone()
1935            }
1936        } else {
1937            facts.clone()
1938        };
1939
1940        if filtered.is_empty() {
1941            return String::new();
1942        }
1943
1944        if let Some(store) = self.fact_store.read().clone() {
1945            store.format_for_context(&filtered, max_tokens)
1946        } else {
1947            String::new()
1948        }
1949    }
1950
1951    fn build_context_with_staged(&self, staged: &HashMap<String, Value>) -> HashMap<String, Value> {
1952        let context = self.build_context_with_overlays();
1953        let mut root = Value::Object(context.into_iter().collect());
1954        for (path, value) in staged {
1955            if let Ok(updated) = ai_agents_core::set_dot_path(root.clone(), path, value.clone()) {
1956                root = updated;
1957            }
1958        }
1959        match root {
1960            Value::Object(obj) => obj.into_iter().collect(),
1961            _ => HashMap::new(),
1962        }
1963    }
1964
1965    fn build_context_with_overlays(&self) -> HashMap<String, Value> {
1966        let mut context = self.context_manager.get_all();
1967        let mut root = Value::Object(context.clone().into_iter().collect());
1968
1969        if let Some(turn_ctx) = self.current_turn_actor_context() {
1970            if let Some(ref origin_actor_id) = turn_ctx.origin_actor_id
1971                && let Ok(updated) = ai_agents_core::set_dot_path(
1972                    root.clone(),
1973                    "interaction.origin_actor_id",
1974                    serde_json::json!(origin_actor_id),
1975                )
1976            {
1977                root = updated;
1978            }
1979            if let Some(ref sender_agent_id) = turn_ctx.sender_agent_id
1980                && let Ok(updated) = ai_agents_core::set_dot_path(
1981                    root.clone(),
1982                    "interaction.sender_agent_id",
1983                    serde_json::json!(sender_agent_id),
1984                )
1985            {
1986                root = updated;
1987            }
1988        }
1989
1990        if let Some(ref actor_id) = self.effective_actor_id()
1991            && let Ok(updated) = ai_agents_core::set_dot_path(
1992                root.clone(),
1993                "interaction.actor_id",
1994                serde_json::json!(actor_id),
1995            )
1996        {
1997            root = updated;
1998        }
1999
2000        if let Some(manager) = self.relationship_manager.as_ref()
2001            && let Some(actor_id) = self.effective_actor_id()
2002            && let Some(value) = manager.to_context_value(&actor_id)
2003            && let Ok(updated) = ai_agents_core::set_dot_path(
2004                root.clone(),
2005                &manager.config().injection.context_path,
2006                value,
2007            )
2008        {
2009            root = updated;
2010        }
2011
2012        if let Value::Object(obj) = root {
2013            context = obj.into_iter().collect();
2014        }
2015
2016        context
2017    }
2018
2019    fn resolve_actor_name_from_context(&self) -> Option<String> {
2020        for path in ["actor.name", "user.name", "player.name", "customer.name"] {
2021            if let Some(value) = self.context_manager.get_path(path)
2022                && let Some(name) = value.as_str()
2023            {
2024                return Some(name.to_string());
2025            }
2026        }
2027        None
2028    }
2029
2030    async fn maybe_load_actor_relationship(&self) {
2031        let Some(manager) = self.relationship_manager.as_ref() else {
2032            return;
2033        };
2034        let Some(actor_id) = self.effective_actor_id() else {
2035            return;
2036        };
2037
2038        let mut should_fire_loaded = false;
2039        if manager.get(&actor_id).is_none() {
2040            let mut loaded = false;
2041            if manager.config().persistence.enabled {
2042                let storage = self.storage.read().clone();
2043                if let Some(storage) = storage {
2044                    match storage.load_relationship(&self.info.id, &actor_id).await {
2045                        Ok(Some(value)) => match manager.insert_from_value(value) {
2046                            Ok(_) => loaded = true,
2047                            Err(e) => {
2048                                warn!(actor = %actor_id, error = %e, "failed to restore relationship")
2049                            }
2050                        },
2051                        Ok(None) => {}
2052                        Err(e) => {
2053                            warn!(actor = %actor_id, error = %e, "failed to load relationship")
2054                        }
2055                    }
2056                }
2057            }
2058
2059            if !loaded {
2060                manager.get_or_create(&actor_id, self.resolve_actor_name_from_context().as_deref());
2061            }
2062            should_fire_loaded = true;
2063        }
2064
2065        let actor_name = self.resolve_actor_name_from_context();
2066        let relationship = manager.touch_interaction(&actor_id, actor_name.as_deref());
2067        if should_fire_loaded {
2068            self.hooks
2069                .on_relationship_loaded(&actor_id, &relationship)
2070                .await;
2071        }
2072    }
2073
2074    fn format_relationship_for_context(&self) -> Option<(String, String)> {
2075        let manager = self.relationship_manager.as_ref()?;
2076        if !manager.config().injection.enabled {
2077            return None;
2078        }
2079        let actor_id = self.effective_actor_id()?;
2080        let relationship = manager.get(&actor_id)?;
2081        let local_cap = manager.config().injection.max_tokens;
2082        let global_cap = self
2083            .memory_token_budget
2084            .as_ref()
2085            .map(|b| b.allocation.relationships as usize)
2086            .filter(|n| *n > 0);
2087        let max_tokens = global_cap.map(|g| g.min(local_cap)).unwrap_or(local_cap);
2088        let text = ai_agents_relationships::format_relationship(
2089            &relationship,
2090            &manager.config().injection.format,
2091            max_tokens,
2092        );
2093        if text.is_empty() {
2094            None
2095        } else {
2096            Some((manager.config().injection.prompt_variable.clone(), text))
2097        }
2098    }
2099
2100    async fn persist_actor_relationship(&self, actor_id: &str) -> Result<()> {
2101        let Some(manager) = self.relationship_manager.as_ref() else {
2102            return Ok(());
2103        };
2104        if !manager.config().persistence.enabled {
2105            return Ok(());
2106        }
2107        let storage = self.storage.read().clone();
2108        let Some(storage) = storage else {
2109            return Ok(());
2110        };
2111        if let Some(value) = manager.relationship_as_value(actor_id)? {
2112            storage
2113                .save_relationship(&self.info.id, actor_id, &value)
2114                .await?;
2115        }
2116        Ok(())
2117    }
2118
2119    async fn auto_update_relationship(&self) {
2120        let Some(manager) = self.relationship_manager.as_ref() else {
2121            return;
2122        };
2123        let Some(actor_id) = self.effective_actor_id() else {
2124            return;
2125        };
2126        if !manager.config().auto_update.enabled {
2127            let _ = self.persist_actor_relationship(&actor_id).await;
2128            return;
2129        }
2130
2131        let recent_messages = manager.config().auto_update.recent_messages;
2132        let messages = match self.memory.get_messages(Some(recent_messages)).await {
2133            Ok(messages) => messages,
2134            Err(e) => {
2135                warn!(actor = %actor_id, error = %e, "failed to read messages for relationship update");
2136                return;
2137            }
2138        };
2139
2140        match self
2141            .observe_purpose(
2142                ObservationPurpose::RelationshipUpdate,
2143                manager.auto_update(&actor_id, &messages),
2144            )
2145            .await
2146        {
2147            Ok(update) => {
2148                if !update.changes.is_empty() {
2149                    self.hooks
2150                        .on_relationship_change(&actor_id, &update.changes)
2151                        .await;
2152                }
2153                if let Some(ref event) = update.event {
2154                    self.hooks.on_notable_event(&actor_id, event).await;
2155                }
2156                let persisted = match self.persist_actor_relationship(&actor_id).await {
2157                    Ok(()) => true,
2158                    Err(e) => {
2159                        warn!(actor = %actor_id, error = %e, "failed to persist relationship");
2160                        false
2161                    }
2162                };
2163                if !update.changes.is_empty() || update.event.is_some() {
2164                    let changed_dimensions: Vec<String> = update
2165                        .changes
2166                        .iter()
2167                        .map(|change| format!("{}:{}", change.perspective, change.dimension))
2168                        .collect();
2169                    info!(
2170                        actor_id = %actor_id,
2171                        change_count = update.changes.len(),
2172                        changed_dimensions = ?changed_dimensions,
2173                        event_present = update.event.is_some(),
2174                        persisted = persisted,
2175                        "relationship updated"
2176                    );
2177                } else {
2178                    debug!(actor_id = %actor_id, persisted = persisted, "relationship evaluation ran but found no changes");
2179                }
2180            }
2181            Err(e) => warn!(actor = %actor_id, error = %e, "relationship update failed"),
2182        }
2183    }
2184
2185    /// Run auto-extraction after a chat turn.
2186    async fn auto_extract_facts(&self) {
2187        let should_extract = self
2188            .facts_config
2189            .as_ref()
2190            .map(|c| c.enabled && c.auto_extract)
2191            .unwrap_or(false);
2192
2193        if !should_extract {
2194            debug!("fact extraction skipped because auto extraction is disabled");
2195            return;
2196        }
2197
2198        let msgs_since = *self.messages_since_extraction.read();
2199        if msgs_since < 2 {
2200            debug!(
2201                messages_since_extraction = msgs_since,
2202                "fact extraction skipped until threshold is reached"
2203            );
2204            return;
2205        }
2206
2207        match self.extract_facts_with_source(msgs_since, "auto").await {
2208            Ok(facts) => {
2209                if !facts.is_empty() {
2210                    *self.messages_since_extraction.write() = 0;
2211                } else {
2212                    debug!("fact extraction ran but found no new facts");
2213                }
2214            }
2215            Err(e) => {
2216                warn!("fact extraction failed: {}", e);
2217            }
2218        }
2219    }
2220
2221    pub fn with_persona(mut self, manager: Arc<ai_agents_persona::PersonaManager>) -> Self {
2222        self.persona_manager = Some(manager);
2223        self
2224    }
2225
2226    pub fn persona_manager(&self) -> Option<&Arc<ai_agents_persona::PersonaManager>> {
2227        self.persona_manager.as_ref()
2228    }
2229
2230    pub fn with_disambiguation(mut self, config: DisambiguationConfig) -> Self {
2231        if config.is_enabled() {
2232            let manager = DisambiguationManager::new(config, Arc::clone(&self.llm_registry))
2233                .with_clarification_observer(Arc::new(ObservabilityClarificationObserver));
2234            self.disambiguation_manager = Some(manager);
2235        }
2236        self
2237    }
2238
2239    pub fn disambiguation_manager(&self) -> Option<&DisambiguationManager> {
2240        self.disambiguation_manager.as_ref()
2241    }
2242
2243    pub fn has_disambiguation(&self) -> bool {
2244        self.disambiguation_manager
2245            .as_ref()
2246            .is_some_and(|m| m.is_enabled())
2247    }
2248
2249    pub async fn init_storage(&self) -> Result<()> {
2250        //
2251        // One readiness guard prevents concurrent entry points from constructing storage or fact state more than once.
2252        //
2253        let _guard = self.storage_init.lock().await;
2254        let mut storage = self.storage.read().clone();
2255        if storage.is_none() && !self.storage_config.is_none() {
2256            let storage_config = self.convert_storage_config();
2257            storage = create_storage(&storage_config).await?;
2258            *self.storage.write() = storage.clone();
2259        }
2260
2261        self.validate_storage_requirements(storage.as_deref())?;
2262        self.complete_facts_init().await;
2263        Ok(())
2264    }
2265
2266    fn validate_storage_requirements(&self, storage: Option<&dyn AgentStorage>) -> Result<()> {
2267        let facts_required = self
2268            .facts_config
2269            .as_ref()
2270            .is_some_and(|config| config.enabled)
2271            || self
2272                .actor_memory_config
2273                .as_ref()
2274                .is_some_and(|config| config.enabled);
2275        let relationships_required = self
2276            .relationship_manager
2277            .as_ref()
2278            .is_some_and(|manager| manager.config().persistence.enabled);
2279
2280        let Some(storage) = storage else {
2281            let mut requirements = Vec::new();
2282            if facts_required {
2283                requirements.push("actor facts or actor memory");
2284            }
2285            if relationships_required {
2286                requirements.push("persistent relationships");
2287            }
2288            if requirements.is_empty() {
2289                return Ok(());
2290            }
2291            return Err(AgentError::Config(format!(
2292                "Storage is required for enabled {} but none is configured or injected",
2293                requirements.join(" and ")
2294            )));
2295        };
2296
2297        //
2298        // Durable features must fail before execution instead of degrading to volatile state.
2299        //
2300        if facts_required && !storage.supports(StorageCapability::ActorFacts) {
2301            return Err(AgentError::UnsupportedStorageCapability(
2302                StorageCapability::ActorFacts,
2303            ));
2304        }
2305        if relationships_required && !storage.supports(StorageCapability::ActorRelationships) {
2306            return Err(AgentError::UnsupportedStorageCapability(
2307                StorageCapability::ActorRelationships,
2308            ));
2309        }
2310        Ok(())
2311    }
2312
2313    /// Initialize fact store and extractor from stored config and current storage.
2314    /// Called from init_storage() so facts are ready before the first turn.
2315    async fn complete_facts_init(&self) {
2316        if self.fact_store.read().is_some() {
2317            return;
2318        }
2319        let storage = match self.storage.read().clone() {
2320            Some(s) => s,
2321            None => return,
2322        };
2323
2324        let facts_enabled = self
2325            .facts_config
2326            .as_ref()
2327            .map(|f| f.enabled)
2328            .unwrap_or(false);
2329        let actor_memory_enabled = self
2330            .actor_memory_config
2331            .as_ref()
2332            .map(|a| a.enabled)
2333            .unwrap_or(false);
2334
2335        if !facts_enabled && !actor_memory_enabled {
2336            return;
2337        }
2338
2339        let fc = self.facts_config.clone().unwrap_or_default();
2340        let store = Arc::new(ai_agents_facts::FactStore::new(
2341            storage,
2342            self.info.id.clone(),
2343            fc.clone(),
2344        ));
2345
2346        let extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>> = if facts_enabled {
2347            let extractor_llm = fc
2348                .extractor_llm
2349                .as_ref()
2350                .and_then(|alias| self.llm_registry.get(alias).ok())
2351                .or_else(|| self.llm_registry.router().ok())
2352                .or_else(|| self.llm_registry.default().ok());
2353            extractor_llm.map(|llm| {
2354                Arc::new(ai_agents_facts::LLMFactExtractor::new(llm, fc.clone()))
2355                    as Arc<dyn ai_agents_facts::FactExtractor>
2356            })
2357        } else {
2358            None
2359        };
2360
2361        *self.fact_store.write() = Some(store);
2362        *self.fact_extractor.write() = extractor;
2363        debug!(
2364            agent = %self.info.id,
2365            facts_enabled,
2366            actor_memory_enabled,
2367            "facts storage initialized"
2368        );
2369    }
2370
2371    fn convert_storage_config(&self) -> StorageStorageConfig {
2372        crate::spec::storage::to_storage_config(&self.storage_config)
2373    }
2374
2375    pub fn storage(&self) -> Option<Arc<dyn AgentStorage>> {
2376        self.storage.read().clone()
2377    }
2378
2379    pub fn storage_config(&self) -> &StorageConfig {
2380        &self.storage_config
2381    }
2382
2383    /// Returns the spawner if configured via a spawner: YAML section.
2384    pub fn spawner(&self) -> Option<&Arc<crate::spawner::AgentSpawner>> {
2385        self.spawner.as_ref()
2386    }
2387
2388    /// Returns the agent registry if configured via a spawner: YAML section.
2389    pub fn spawner_registry(&self) -> Option<&Arc<crate::spawner::AgentRegistry>> {
2390        self.spawner_registry.as_ref()
2391    }
2392
2393    pub fn has_spawner(&self) -> bool {
2394        self.spawner_registry.is_some()
2395    }
2396
2397    pub fn with_spawner_handles(
2398        mut self,
2399        spawner: Arc<crate::spawner::AgentSpawner>,
2400        registry: Arc<crate::spawner::AgentRegistry>,
2401    ) -> Self {
2402        self.spawner = Some(spawner);
2403        self.spawner_registry = Some(registry);
2404        self
2405    }
2406
2407    pub fn with_hooks(mut self, hooks: Arc<dyn AgentHooks>) -> Self {
2408        self.hooks = hooks;
2409        self
2410    }
2411
2412    pub fn with_parallel_tools(mut self, config: ParallelToolsConfig) -> Self {
2413        self.parallel_tools = config;
2414        self
2415    }
2416
2417    pub fn with_streaming(mut self, config: StreamingConfig) -> Self {
2418        self.streaming = config;
2419        self
2420    }
2421
2422    pub fn with_hitl(mut self, engine: HITLEngine, handler: Arc<dyn ApprovalHandler>) -> Self {
2423        self.hitl_engine = Some(engine);
2424        self.approval_handler = handler;
2425        self
2426    }
2427
2428    pub fn with_max_context_tokens(mut self, tokens: u32) -> Self {
2429        self.max_context_tokens = tokens;
2430        self
2431    }
2432
2433    pub fn with_memory_token_budget(mut self, budget: MemoryTokenBudget) -> Self {
2434        self.memory_token_budget = Some(budget);
2435        self
2436    }
2437
2438    pub fn with_recovery_manager(mut self, manager: RecoveryManager) -> Self {
2439        self.recovery_manager = manager;
2440        self
2441    }
2442
2443    pub fn with_tool_security(mut self, engine: ToolSecurityEngine) -> Self {
2444        self.tool_security = engine;
2445        self
2446    }
2447
2448    /// Returns the host-only runtime control handle.
2449    pub fn runtime_control(&self) -> RuntimeControlHandle {
2450        RuntimeControlHandle {
2451            state: Arc::clone(&self.runtime_control),
2452        }
2453    }
2454
2455    /// Installs or clears the host question handler used by `ask_user`.
2456    pub fn set_question_handler(&self, handler: Option<Arc<dyn QuestionHandler>>) {
2457        self.tools.set_question_handler(handler);
2458    }
2459
2460    /// Installs the host diagnostics provider used by `diagnostics`.
2461    pub fn set_diagnostics_provider(&self, provider: Arc<dyn DiagnosticsProvider>) {
2462        self.tools.set_diagnostics_provider(provider);
2463    }
2464
2465    /// Installs the host command runner used by `command`.
2466    pub fn set_command_runner(&self, runner: Arc<dyn CommandRunner>) {
2467        self.tools.set_command_runner(runner);
2468    }
2469
2470    /// Installs the host web-search provider used by `web_search`.
2471    pub fn set_web_search_provider(&self, provider: Arc<dyn ai_agents_tools::WebSearchProvider>) {
2472        self.tools.set_web_search_provider(provider);
2473    }
2474
2475    /// Returns the current session-local todo list.
2476    pub fn todos(&self) -> Vec<TodoItem> {
2477        self.tools.todos()
2478    }
2479
2480    /// Returns the effective tool security snapshot for the next decision.
2481    fn active_tool_security(&self) -> ToolSecurityEngine {
2482        self.runtime_control
2483            .tool_security_override
2484            .read()
2485            .clone()
2486            .unwrap_or_else(|| self.tool_security.clone())
2487    }
2488
2489    /// Reads policy, scope, emergency state, and generation as one safety snapshot.
2490    fn runtime_safety_snapshot(&self) -> RuntimeSafetySnapshot {
2491        let _guard = self.runtime_control.snapshot_guard.read();
2492        RuntimeSafetySnapshot {
2493            version: self.runtime_control.version.load(Ordering::SeqCst),
2494            emergency_deny: self.runtime_control.emergency_deny.load(Ordering::SeqCst),
2495            tool_security: self
2496                .runtime_control
2497                .tool_security_override
2498                .read()
2499                .clone()
2500                .unwrap_or_else(|| self.tool_security.clone()),
2501            tool_scope_override: self.runtime_control.tool_scope_override.read().clone(),
2502        }
2503    }
2504
2505    /// Performs the single rate admission after locks are held and rejects changed runtime, policy, or state authority.
2506    fn admit_tool_execution(
2507        &self,
2508        expected_runtime_version: u64,
2509        expected_policy_version: u64,
2510        expected_state_generation: Option<u64>,
2511        canonical_id: &str,
2512    ) -> SecurityCheckResult {
2513        let _guard = self.runtime_control.snapshot_guard.read();
2514        if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
2515            return SecurityCheckResult::Block {
2516                reason: "runtime emergency deny is enabled".to_string(),
2517            };
2518        }
2519        let runtime_version = self.runtime_control.version.load(Ordering::SeqCst);
2520        let security_engine = self
2521            .runtime_control
2522            .tool_security_override
2523            .read()
2524            .clone()
2525            .unwrap_or_else(|| self.tool_security.clone());
2526        if runtime_version != expected_runtime_version
2527            || security_engine.policy_version() != expected_policy_version
2528        {
2529            return SecurityCheckResult::Block {
2530                reason: "runtime safety controls changed before admission".to_string(),
2531            };
2532        }
2533        let current_state_generation = self
2534            .state_machine
2535            .as_ref()
2536            .map(|state_machine| state_machine.generation());
2537        if current_state_generation != expected_state_generation {
2538            return SecurityCheckResult::Block {
2539                reason: "state scope changed before admission".to_string(),
2540            };
2541        }
2542        security_engine.admit_tool_execution(canonical_id)
2543    }
2544
2545    pub fn with_process_processor(mut self, processor: ProcessProcessor) -> Self {
2546        let processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
2547        self.process_processor = Some(processor);
2548        self
2549    }
2550
2551    pub fn with_state_machine(
2552        mut self,
2553        state_machine: Arc<StateMachine>,
2554        evaluator: Arc<dyn TransitionEvaluator>,
2555    ) -> Self {
2556        self.state_machine = Some(state_machine);
2557        self.transition_evaluator = Some(evaluator);
2558        self
2559    }
2560
2561    pub fn with_context_manager(mut self, manager: Arc<ContextManager>) -> Self {
2562        self.context_manager = manager;
2563        self
2564    }
2565
2566    pub fn register_message_filter(&self, name: impl Into<String>, filter: Arc<dyn MessageFilter>) {
2567        self.message_filters.write().insert(name.into(), filter);
2568    }
2569
2570    pub fn set_context(&self, key: &str, value: Value) -> Result<()> {
2571        self.context_manager.update(key, value)
2572    }
2573
2574    pub fn update_context(&self, path: &str, value: Value) -> Result<()> {
2575        self.context_manager.update(path, value)
2576    }
2577
2578    pub fn get_context(&self) -> HashMap<String, Value> {
2579        self.build_context_with_overlays()
2580    }
2581
2582    pub fn remove_context(&self, key: &str) -> Option<Value> {
2583        self.context_manager.remove(key)
2584    }
2585
2586    pub async fn refresh_context(&self, key: &str) -> Result<()> {
2587        self.context_manager.refresh(key).await
2588    }
2589
2590    pub fn register_context_provider(&self, name: &str, provider: Arc<dyn ContextProvider>) {
2591        self.context_manager.register_provider(name, provider);
2592    }
2593
2594    pub fn current_state(&self) -> Option<String> {
2595        self.state_machine.as_ref().map(|sm| sm.current())
2596    }
2597
2598    // State changes invalidate only resolved confirmation ownership; ordinary clarification remains manager-owned across layered enablement changes.
2599    async fn invalidate_pending_confirmation(&self, reason: &'static str) {
2600        self.disambiguation_epoch.fetch_add(1, Ordering::SeqCst);
2601        let Some(disambiguator) = self.disambiguation_manager.as_ref() else {
2602            return;
2603        };
2604        if disambiguator.has_pending_confirmation().await {
2605            disambiguator.clear_pending().await;
2606            *self.pending_skill_id.write() = None;
2607            info!(
2608                confirmation_event = "invalidated",
2609                invalidation_reason = reason,
2610                "Runtime invalidated pending confirmation"
2611            );
2612        }
2613    }
2614
2615    // Admission linearizes pending publication or redispatch before reset and state mutation, then rechecks both ownership generations.
2616    async fn admit_disambiguation_redispatch(
2617        &self,
2618        expected_epoch: u64,
2619        expected_state_generation: Option<u64>,
2620    ) -> Result<tokio::sync::RwLockReadGuard<'_, ()>> {
2621        let admission = self.disambiguation_admission.read().await;
2622        let state_generation = self
2623            .state_machine
2624            .as_ref()
2625            .map(|state_machine| state_machine.generation());
2626        if self.disambiguation_epoch.load(Ordering::SeqCst) != expected_epoch
2627            || state_generation != expected_state_generation
2628        {
2629            return Err(AgentError::Other(
2630                "Disambiguation ownership changed before redispatch admission".to_string(),
2631            ));
2632        }
2633        Ok(admission)
2634    }
2635
2636    // A transition reservation prevents losing runtime transitions from running duplicate exit side effects.
2637    fn reserve_state_transition(&self) -> Option<StateTransitionReservation<'_>> {
2638        self.state_transition_reserved
2639            .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
2640            .ok()
2641            .map(|_| StateTransitionReservation {
2642                reserved: &self.state_transition_reserved,
2643            })
2644    }
2645
2646    // Optional ownership is used only for terminal skill responses that came from manager-owned pending state.
2647    async fn admit_optional_disambiguation_ownership(
2648        &self,
2649        ownership: Option<DisambiguationOwnership>,
2650    ) -> Result<Option<tokio::sync::RwLockReadGuard<'_, ()>>> {
2651        match ownership {
2652            Some(ownership) => self
2653                .admit_disambiguation_redispatch(ownership.epoch, ownership.state_generation)
2654                .await
2655                .map(Some),
2656            None => Ok(None),
2657        }
2658    }
2659
2660    /// Applies a manual transition after reserving its exit actions and keeping async lifecycle work outside the commit lock.
2661    pub async fn transition_to(&self, state: &str) -> Result<()> {
2662        let Some(ref sm) = self.state_machine else {
2663            return Ok(());
2664        };
2665        let claim_admission = self.disambiguation_admission.write().await;
2666        let reservation = self.reserve_state_transition().ok_or_else(|| {
2667            AgentError::Other("Another state transition is already in progress".to_string())
2668        })?;
2669        let from_state = sm.current();
2670        let expected_state_generation = sm.generation();
2671        let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
2672        let history_before = sm.history();
2673        drop(claim_admission);
2674
2675        self.execute_state_exit_actions(&from_state).await;
2676
2677        let admission = self.disambiguation_admission.write().await;
2678        if sm.current() != from_state
2679            || sm.generation() != expected_state_generation
2680            || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
2681        {
2682            return Err(AgentError::Other(
2683                "State ownership changed during manual transition preparation".to_string(),
2684            ));
2685        }
2686        sm.transition_to(state, "manual transition")?;
2687        self.invalidate_pending_confirmation("state_transition")
2688            .await;
2689        let entered = sm.current();
2690        let is_reentry = Self::state_was_previously_entered(&entered, &from_state, &history_before);
2691        drop(admission);
2692
2693        self.execute_state_enter_actions(&entered, is_reentry).await;
2694        drop(reservation);
2695        info!(to = %entered, "Manual state transition");
2696        Ok(())
2697    }
2698
2699    pub fn state_history(&self) -> Vec<StateTransitionEvent> {
2700        self.state_machine
2701            .as_ref()
2702            .map(|sm| sm.history())
2703            .unwrap_or_default()
2704    }
2705
2706    /// Get a copy of current session metadata.
2707    pub fn session_metadata(&self) -> ai_agents_core::SessionMetadata {
2708        self.session_metadata.read().clone()
2709    }
2710
2711    /// Delete all facts and sessions for an actor, gated by privacy.allow_deletion.
2712    /// Returns Err when actor_memory.privacy.allow_deletion is false.
2713    pub async fn delete_actor_data(&self, actor_id: &str) -> Result<()> {
2714        let allowed = self
2715            .actor_memory_config
2716            .as_ref()
2717            .map(|c| c.privacy.allow_deletion)
2718            .unwrap_or(true);
2719        if !allowed {
2720            return Err(AgentError::Config(
2721                "privacy.allow_deletion is false; actor data deletion is not permitted".into(),
2722            ));
2723        }
2724        let storage = self.storage.read().clone();
2725        if let Some(storage) = storage {
2726            //
2727            // Composite support is checked before mutation so privacy deletion cannot become partial.
2728            //
2729            if !storage.supports(StorageCapability::ActorDataDeletion) {
2730                return Err(AgentError::UnsupportedStorageCapability(
2731                    StorageCapability::ActorDataDeletion,
2732                ));
2733            }
2734            storage.delete_actor_data(&self.info.id, actor_id).await?;
2735        } else {
2736            //
2737            // Clone the fallback store before awaiting so backend I/O never holds the runtime read guard.
2738            //
2739            let store = { self.fact_store.read().clone() };
2740            if let Some(store) = store {
2741                store.delete_actor_data(actor_id).await?;
2742            }
2743        }
2744        if let Some(manager) = self.relationship_manager.as_ref() {
2745            manager.remove(actor_id);
2746        }
2747        self.actor_facts_cache.write().remove(actor_id);
2748        Ok(())
2749    }
2750
2751    /// Overwrite session metadata (tags, ttl, custom fields).
2752    pub fn set_session_metadata(&self, meta: ai_agents_core::SessionMetadata) {
2753        *self.session_metadata.write() = meta;
2754    }
2755
2756    /// Delete sessions whose TTL has expired. Returns number of sessions removed.
2757    pub async fn cleanup_expired_sessions(&self) -> Result<usize> {
2758        let storage = self.storage.read().clone();
2759        match storage {
2760            Some(s) => {
2761                let count = s.cleanup_expired().await?;
2762                if count > 0 {
2763                    self.hooks.on_sessions_expired(count).await;
2764                }
2765                Ok(count)
2766            }
2767            None => Err(AgentError::Config(
2768                "No storage configured. Use with_storage_config() or with_storage() first".into(),
2769            )),
2770        }
2771    }
2772
2773    /// List sessions matching a filter. Supports actor, tag, and date filters.
2774    pub async fn list_sessions_filtered(
2775        &self,
2776        filter: &ai_agents_core::SessionFilter,
2777    ) -> Result<Vec<ai_agents_core::SessionSummary>> {
2778        let storage = self.storage.read().clone();
2779        match storage {
2780            Some(s) => s.list_sessions_filtered(filter).await,
2781            None => Err(AgentError::Config(
2782                "No storage configured. Use with_storage_config() or with_storage() first".into(),
2783            )),
2784        }
2785    }
2786
2787    pub async fn save_state(&self) -> Result<AgentSnapshot> {
2788        let memory_snapshot = self.memory.snapshot().await?;
2789        let state_machine_snapshot = self.state_machine.as_ref().map(|sm| sm.snapshot());
2790        let context_snapshot = self.context_manager.snapshot();
2791
2792        let mut snapshot = AgentSnapshot::new(self.info.id.clone())
2793            .with_memory(memory_snapshot)
2794            .with_context(context_snapshot)
2795            .with_state_machine(
2796                state_machine_snapshot.unwrap_or_else(|| StateMachineSnapshot {
2797                    current_state: String::new(),
2798                    previous_state: None,
2799                    turn_count: 0,
2800                    no_transition_count: 0,
2801                    history: vec![],
2802                }),
2803            );
2804
2805        if let Some(ref persona) = self.persona_manager {
2806            snapshot.persona = Some(persona.snapshot_as_value()?);
2807        }
2808
2809        if let Some(ref relationships) = self.relationship_manager {
2810            snapshot.relationships = Some(relationships.snapshot_as_value()?);
2811        }
2812
2813        Ok(snapshot)
2814    }
2815
2816    /// Save state including spawned agents manifest for session persistence.
2817    pub async fn save_state_full(&self) -> Result<AgentSnapshot> {
2818        let mut snapshot = self.save_state().await?;
2819        if let Some(ref registry) = self.spawner_registry {
2820            let entries = registry.list_with_specs();
2821            if !entries.is_empty() {
2822                snapshot = snapshot.with_spawned_agents(entries);
2823            }
2824        }
2825        Ok(snapshot)
2826    }
2827
2828    /// Restores persisted state after invalidating any pending confirmation ownership.
2829    pub async fn restore_state(&self, snapshot: AgentSnapshot) -> Result<()> {
2830        let _admission = self.disambiguation_admission.write().await;
2831        if self.state_transition_reserved.load(Ordering::SeqCst) {
2832            return Err(AgentError::Other(
2833                "Cannot restore state while a state transition is in progress".to_string(),
2834            ));
2835        }
2836        self.invalidate_pending_confirmation("state_restore").await;
2837        *self.pending_skill_id.write() = None;
2838        if let Some(disambiguator) = self.disambiguation_manager.as_ref() {
2839            disambiguator.clear_pending().await;
2840        }
2841        self.memory.restore(snapshot.memory).await?;
2842
2843        if let (Some(sm), Some(sm_snapshot)) = (&self.state_machine, snapshot.state_machine)
2844            && !sm_snapshot.current_state.is_empty()
2845        {
2846            sm.restore(sm_snapshot)?;
2847        }
2848
2849        self.context_manager.restore(snapshot.context);
2850
2851        if let (Some(persona_value), Some(persona_manager)) =
2852            (snapshot.persona, &self.persona_manager)
2853        {
2854            persona_manager.restore_from_value(persona_value)?;
2855        }
2856
2857        if let (Some(relationship_value), Some(relationship_manager)) =
2858            (snapshot.relationships, &self.relationship_manager)
2859        {
2860            relationship_manager.restore_from_value(relationship_value)?;
2861        }
2862
2863        info!(agent_id = %snapshot.agent_id, "State restored");
2864        Ok(())
2865    }
2866
2867    pub async fn save_to(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<()> {
2868        let snapshot = self.save_state().await?;
2869        storage.save(session_id, &snapshot).await
2870    }
2871
2872    async fn load_session_restore(
2873        storage: &dyn AgentStorage,
2874        session_id: &str,
2875    ) -> Result<Option<StoredSessionRestore>> {
2876        let Some(snapshot) = storage.load(session_id).await? else {
2877            return Ok(None);
2878        };
2879        //
2880        // Metadata is restored only when the backend explicitly owns that durable contract.
2881        //
2882        let metadata = if storage.supports(StorageCapability::SessionMetadata) {
2883            storage.load_metadata(session_id).await?
2884        } else {
2885            None
2886        };
2887        Ok(Some(StoredSessionRestore { snapshot, metadata }))
2888    }
2889
2890    async fn capture_session_restore_point(&self) -> Result<RuntimeSessionRestorePoint> {
2891        Ok(RuntimeSessionRestorePoint {
2892            snapshot: self.save_state().await?,
2893            metadata: self.session_metadata(),
2894            actor_id: self.actor_id(),
2895            session_id: self.current_session_id.read().clone(),
2896        })
2897    }
2898
2899    async fn apply_session_restore_unchecked(
2900        &self,
2901        session_id: &str,
2902        stored: StoredSessionRestore,
2903    ) -> Result<()> {
2904        self.restore_state(stored.snapshot).await?;
2905        let metadata = stored.metadata.unwrap_or_default();
2906        if let Some(actor_id) = metadata.actor_id.as_deref() {
2907            self.set_actor_id(actor_id)?;
2908        } else {
2909            self.clear_actor_id();
2910        }
2911        self.set_session_metadata(metadata);
2912        *self.current_session_id.write() = Some(session_id.to_string());
2913        Ok(())
2914    }
2915
2916    async fn restore_session_restore_point(
2917        &self,
2918        restore_point: &RuntimeSessionRestorePoint,
2919    ) -> Result<()> {
2920        self.restore_state(restore_point.snapshot.clone()).await?;
2921        if let Some(actor_id) = restore_point.actor_id.as_deref() {
2922            self.set_actor_id(actor_id)?;
2923        } else {
2924            self.clear_actor_id();
2925        }
2926        self.set_session_metadata(restore_point.metadata.clone());
2927        *self.current_session_id.write() = restore_point.session_id.clone();
2928        Ok(())
2929    }
2930
2931    async fn apply_session_restore(
2932        &self,
2933        session_id: &str,
2934        stored: StoredSessionRestore,
2935    ) -> Result<()> {
2936        let before = self.capture_session_restore_point().await?;
2937        if let Err(error) = self
2938            .apply_session_restore_unchecked(session_id, stored)
2939            .await
2940        {
2941            return match self.restore_session_restore_point(&before).await {
2942                Ok(()) => Err(error),
2943                Err(rollback_error) => Err(AgentError::Other(format!(
2944                    "Session restore failed: {error}; rollback failed: {rollback_error}"
2945                ))),
2946            };
2947        }
2948        Ok(())
2949    }
2950
2951    async fn rollback_session_restore_set(
2952        parent: Option<(&RuntimeAgent, &RuntimeSessionRestorePoint)>,
2953        children: &[(String, Arc<RuntimeAgent>, RuntimeSessionRestorePoint)],
2954    ) -> Vec<String> {
2955        let mut errors = Vec::new();
2956        if let Some((agent, restore_point)) = parent
2957            && let Err(error) = agent.restore_session_restore_point(restore_point).await
2958        {
2959            errors.push(format!("parent: {error}"));
2960        }
2961        for (id, agent, restore_point) in children {
2962            if let Err(error) = agent.restore_session_restore_point(restore_point).await {
2963                errors.push(format!("child '{id}': {error}"));
2964            }
2965        }
2966        errors
2967    }
2968
2969    fn restore_failure(error: impl std::fmt::Display, rollback_errors: Vec<String>) -> AgentError {
2970        if rollback_errors.is_empty() {
2971            AgentError::Other(format!(
2972                "Session restore failed: {error}; runtime state was rolled back"
2973            ))
2974        } else {
2975            AgentError::Other(format!(
2976                "Session restore failed: {error}; rollback also failed for {}",
2977                rollback_errors.join(", ")
2978            ))
2979        }
2980    }
2981
2982    pub async fn load_from(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<bool> {
2983        let Some(stored) = Self::load_session_restore(storage, session_id).await? else {
2984            return Ok(false);
2985        };
2986        self.apply_session_restore(session_id, stored).await?;
2987        Ok(true)
2988    }
2989
2990    pub async fn save_session(&self, session_id: &str) -> Result<()> {
2991        let storage = self.storage.read().clone();
2992        match storage {
2993            Some(s) => {
2994                // Fire on_session_created when this session id is first seen on this runtime.
2995                let is_new = {
2996                    let cur = self.current_session_id.read().clone();
2997                    cur.as_deref() != Some(session_id)
2998                };
2999                if is_new {
3000                    *self.current_session_id.write() = Some(session_id.to_string());
3001                    self.hooks.on_session_created(session_id).await;
3002                }
3003
3004                // Update metadata before persisting.
3005                {
3006                    let now = chrono::Utc::now();
3007                    let msg_count = self
3008                        .memory
3009                        .get_messages(None)
3010                        .await
3011                        .map(|v| v.len())
3012                        .unwrap_or(0);
3013                    let mut meta = self.session_metadata.write();
3014                    meta.last_active = now;
3015                    meta.message_count = msg_count;
3016                    if meta.actor_id.is_none() {
3017                        meta.actor_id = self.actor_id.read().clone();
3018                    }
3019                }
3020
3021                let snapshot = self.save_state().await?;
3022                //
3023                // Metadata-capable backends own atomic snapshot and metadata persistence.
3024                //
3025                if s.supports(StorageCapability::SessionMetadata) {
3026                    let metadata = self.session_metadata.read().clone();
3027                    s.save_snapshot_with_metadata(session_id, &snapshot, &metadata)
3028                        .await
3029                } else {
3030                    s.save(session_id, &snapshot).await
3031                }
3032            }
3033            None => Err(AgentError::Config(
3034                "No storage configured. Use with_storage_config() or with_storage() first".into(),
3035            )),
3036        }
3037    }
3038
3039    pub async fn load_session(&self, session_id: &str) -> Result<bool> {
3040        let storage = self.storage.read().clone();
3041        match storage {
3042            Some(storage) => self.load_from(storage.as_ref(), session_id).await,
3043            None => Err(AgentError::Config(
3044                "No storage configured. Use with_storage_config() or with_storage() first".into(),
3045            )),
3046        }
3047    }
3048
3049    /// Restore this runtime and its complete saved child topology from one named session.
3050    pub async fn restore_session_full(&self, session_id: &str) -> Result<usize> {
3051        self.init_storage().await?;
3052        let storage = self.storage.read().clone().ok_or_else(|| {
3053            AgentError::Config(
3054                "No storage configured. Use with_storage_config() or with_storage() first".into(),
3055            )
3056        })?;
3057        let target_parent = Self::load_session_restore(storage.as_ref(), session_id)
3058            .await?
3059            .ok_or_else(|| AgentError::Persistence(format!("Session not found: {session_id}")))?;
3060        let manifest = target_parent
3061            .snapshot
3062            .spawned_agents
3063            .clone()
3064            .unwrap_or_default();
3065
3066        let registry = self.spawner_registry.as_ref().cloned();
3067        let spawner = if manifest.is_empty() {
3068            self.spawner.as_ref().cloned()
3069        } else {
3070            Some(self.spawner.as_ref().cloned().ok_or_else(|| {
3071                AgentError::Config(
3072                    "Saved session contains child agents but this runtime has no spawner".into(),
3073                )
3074            })?)
3075        };
3076        let registry = if manifest.is_empty() {
3077            registry
3078        } else {
3079            Some(registry.ok_or_else(|| {
3080                AgentError::Config(
3081                    "Saved session contains child agents but this runtime has no registry".into(),
3082                )
3083            })?)
3084        };
3085
3086        let mut target_ids = HashSet::with_capacity(manifest.len());
3087        let mut prepared = Vec::with_capacity(manifest.len());
3088        for entry in manifest {
3089            if !target_ids.insert(entry.id.clone()) {
3090                return Err(AgentError::InvalidSpec(format!(
3091                    "Saved child manifest contains duplicate ID: {}",
3092                    entry.id
3093                )));
3094            }
3095            let spec = crate::spec::AgentSpec::from_yaml_strict(&entry.spec_yaml)?;
3096            spawner
3097                .as_ref()
3098                .expect("non-empty manifests require a spawner")
3099                .validate_explicit_child(&entry.id, &spec)?;
3100            prepared.push((entry.id, spec));
3101        }
3102
3103        let current_ids = registry
3104            .as_ref()
3105            .map(|registry| {
3106                registry
3107                    .list()
3108                    .into_iter()
3109                    .map(|info| info.id)
3110                    .collect::<HashSet<_>>()
3111            })
3112            .unwrap_or_default();
3113        let removal_count = current_ids.difference(&target_ids).count();
3114        let additions = prepared
3115            .iter()
3116            .filter(|(id, _)| !current_ids.contains(id))
3117            .cloned()
3118            .collect::<Vec<_>>();
3119
3120        let mut existing = Vec::new();
3121        if let Some(registry) = registry.as_ref() {
3122            for (id, _) in prepared.iter().filter(|(id, _)| current_ids.contains(id)) {
3123                let agent = registry.get(id).ok_or_else(|| {
3124                    AgentError::Config(format!("Retained child disappeared during restore: {id}"))
3125                })?;
3126                let child_storage = agent.storage().ok_or_else(|| {
3127                    AgentError::Config(format!("Child '{id}' has no storage for session restore"))
3128                })?;
3129                let stored = Self::load_session_restore(child_storage.as_ref(), session_id)
3130                    .await?
3131                    .ok_or_else(|| {
3132                        AgentError::Persistence(format!(
3133                            "Child '{id}' has no saved session '{session_id}'"
3134                        ))
3135                    })?;
3136                existing.push((id.clone(), agent, stored));
3137            }
3138        }
3139
3140        let mut staged = Vec::with_capacity(additions.len());
3141        if !additions.is_empty() {
3142            let spawner = spawner
3143                .as_ref()
3144                .expect("restored additions require a spawner");
3145            let reservations = spawner.reserve_restore_capacity(additions.len(), removal_count)?;
3146            for ((id, spec), reservation) in additions.into_iter().zip(reservations) {
3147                let spawned = spawner
3148                    .spawn_with_reserved_capacity(id.clone(), spec, reservation)
3149                    .await?;
3150                let child_storage = spawned.agent.storage().ok_or_else(|| {
3151                    AgentError::Config(format!("Child '{id}' has no storage for session restore"))
3152                })?;
3153                let stored = Self::load_session_restore(child_storage.as_ref(), session_id)
3154                    .await?
3155                    .ok_or_else(|| {
3156                        AgentError::Persistence(format!(
3157                            "Child '{id}' has no saved session '{session_id}'"
3158                        ))
3159                    })?;
3160                staged.push((spawned, stored));
3161            }
3162        } else if let Some(spawner) = spawner.as_ref() {
3163            spawner.reserve_restore_capacity(0, removal_count)?;
3164        }
3165
3166        let parent_before = self.capture_session_restore_point().await?;
3167        let mut existing_before = Vec::with_capacity(existing.len());
3168        for (id, agent, _) in &existing {
3169            existing_before.push((
3170                id.clone(),
3171                Arc::clone(agent),
3172                agent.capture_session_restore_point().await?,
3173            ));
3174        }
3175
3176        //
3177        // Every record is loaded before runtime state changes. Registry replacement is the final topology commit point.
3178        //
3179        for (_, agent, stored) in &existing {
3180            if let Err(error) = agent
3181                .apply_session_restore_unchecked(session_id, stored.clone())
3182                .await
3183            {
3184                drop(staged);
3185                let rollback_errors =
3186                    Self::rollback_session_restore_set(None, &existing_before).await;
3187                return Err(Self::restore_failure(error, rollback_errors));
3188            }
3189        }
3190        for (spawned, stored) in &staged {
3191            if let Err(error) = spawned
3192                .agent
3193                .apply_session_restore_unchecked(session_id, stored.clone())
3194                .await
3195            {
3196                drop(staged);
3197                let rollback_errors =
3198                    Self::rollback_session_restore_set(None, &existing_before).await;
3199                return Err(Self::restore_failure(error, rollback_errors));
3200            }
3201        }
3202        if let Err(error) = self
3203            .apply_session_restore_unchecked(session_id, target_parent)
3204            .await
3205        {
3206            drop(staged);
3207            let rollback_errors =
3208                Self::rollback_session_restore_set(Some((self, &parent_before)), &existing_before)
3209                    .await;
3210            return Err(Self::restore_failure(error, rollback_errors));
3211        }
3212
3213        if let Some(registry) = registry.as_ref()
3214            && let Err(error) = registry
3215                .reconcile(
3216                    &target_ids,
3217                    staged.into_iter().map(|(spawned, _)| spawned).collect(),
3218                )
3219                .await
3220        {
3221            let rollback_errors =
3222                Self::rollback_session_restore_set(Some((self, &parent_before)), &existing_before)
3223                    .await;
3224            return Err(Self::restore_failure(error, rollback_errors));
3225        }
3226
3227        Ok(target_ids.len())
3228    }
3229
3230    pub async fn delete_session(&self, session_id: &str) -> Result<()> {
3231        let storage = self.storage.read().clone();
3232        match storage {
3233            Some(s) => s.delete(session_id).await,
3234            None => Err(AgentError::Config(
3235                "No storage configured. Use with_storage_config() or with_storage() first".into(),
3236            )),
3237        }
3238    }
3239
3240    pub async fn list_sessions(&self) -> Result<Vec<String>> {
3241        let storage = self.storage.read().clone();
3242        match storage {
3243            Some(s) => s.list_sessions().await,
3244            None => Err(AgentError::Config(
3245                "No storage configured. Use with_storage_config() or with_storage() first".into(),
3246            )),
3247        }
3248    }
3249
3250    fn estimate_tokens(&self, text: &str) -> u32 {
3251        (text.len() as f32 / 4.0).ceil() as u32
3252    }
3253
3254    fn estimate_total_tokens(&self, messages: &[ChatMessage]) -> u32 {
3255        messages
3256            .iter()
3257            .map(|m| self.estimate_tokens(&m.content))
3258            .sum()
3259    }
3260
3261    fn truncate_context(&self, messages: &mut Vec<ChatMessage>, keep_recent: usize) {
3262        if messages.len() <= keep_recent + 1 {
3263            return;
3264        }
3265        let system_msg = messages.remove(0);
3266        let to_remove = messages.len().saturating_sub(keep_recent);
3267        messages.drain(..to_remove);
3268        messages.insert(0, system_msg);
3269    }
3270
3271    fn get_filter(&self, config: &FilterConfig) -> Arc<dyn MessageFilter> {
3272        match config {
3273            FilterConfig::KeepRecent(n) => Arc::new(KeepRecentFilter::new(*n)),
3274            FilterConfig::ByRole { keep_roles } => Arc::new(ByRoleFilter::new(keep_roles.clone())),
3275            FilterConfig::SkipPattern { skip_if_contains } => {
3276                Arc::new(SkipPatternFilter::new(skip_if_contains.clone()))
3277            }
3278            FilterConfig::Custom { name } => {
3279                let filters = self.message_filters.read();
3280                filters
3281                    .get(name)
3282                    .cloned()
3283                    .unwrap_or_else(|| Arc::new(KeepRecentFilter::new(10)))
3284            }
3285        }
3286    }
3287
3288    async fn summarize_context(
3289        &self,
3290        messages: &mut Vec<ChatMessage>,
3291        summarizer_llm: Option<&str>,
3292        max_summary_tokens: u32,
3293        custom_prompt: Option<&str>,
3294        keep_recent: usize,
3295        filter: Option<&FilterConfig>,
3296    ) -> Result<()> {
3297        let system_msg = messages.remove(0);
3298
3299        let to_summarize_count = messages.len().saturating_sub(keep_recent);
3300        if to_summarize_count == 0 {
3301            messages.insert(0, system_msg);
3302            return Ok(());
3303        }
3304
3305        let recent_msgs: Vec<ChatMessage> = messages.drain(to_summarize_count..).collect();
3306        let mut to_summarize = std::mem::take(messages);
3307
3308        if let Some(filter_config) = filter {
3309            let filter = self.get_filter(filter_config);
3310            to_summarize = filter.filter(to_summarize);
3311        }
3312
3313        if to_summarize.is_empty() {
3314            *messages = recent_msgs;
3315            messages.insert(0, system_msg);
3316            return Ok(());
3317        }
3318
3319        let conversation_text = to_summarize
3320            .iter()
3321            .map(|m| format!("{:?}: {}", m.role, m.content))
3322            .collect::<Vec<_>>()
3323            .join("\n");
3324
3325        let default_prompt = format!(
3326            "Summarize the following conversation in under {} tokens, preserving key information:\n\n{}",
3327            max_summary_tokens, conversation_text
3328        );
3329
3330        let summary_prompt = custom_prompt
3331            .map(|p| format!("{}\n\n{}", p, conversation_text))
3332            .unwrap_or(default_prompt);
3333
3334        let summarizer = if let Some(alias) = summarizer_llm {
3335            self.llm_registry
3336                .get(alias)
3337                .map_err(|e| AgentError::Config(e.to_string()))?
3338        } else {
3339            self.llm_registry
3340                .router()
3341                .or_else(|_| self.llm_registry.default())
3342                .map_err(|e| AgentError::Config(e.to_string()))?
3343        };
3344
3345        let summary_msgs = vec![ChatMessage::user(&summary_prompt)];
3346        let response = self
3347            .observe_purpose(
3348                ObservationPurpose::Summarization,
3349                summarizer.complete(&summary_msgs, None),
3350            )
3351            .await?;
3352
3353        let summary_message = ChatMessage::system(format!(
3354            "[Previous conversation summary]\n{}",
3355            response.content
3356        ));
3357
3358        *messages = vec![system_msg, summary_message];
3359        messages.extend(recent_msgs);
3360
3361        debug!(
3362            summarized_count = to_summarize_count,
3363            kept_recent = keep_recent,
3364            "Context summarized"
3365        );
3366
3367        Ok(())
3368    }
3369
3370    fn render_system_prompt(&self) -> Result<String> {
3371        let mut context = self.build_context_with_overlays();
3372
3373        // Inject actor_facts for {{ actor_facts }} template variable.
3374        let facts_text = self.format_actor_facts_for_context();
3375        if !facts_text.is_empty() {
3376            context.insert(
3377                "actor_facts".to_string(),
3378                serde_json::Value::String(facts_text),
3379            );
3380        }
3381
3382        if let Some((key, text)) = self.format_relationship_for_context() {
3383            context.insert(key, serde_json::Value::String(text));
3384        }
3385
3386        self.template_renderer
3387            .render(&self.base_system_prompt, &context)
3388    }
3389
3390    /// Canonicalizes IDs once while preserving their first declaration order and ignoring unknown entries.
3391    fn canonical_unique_tool_ids(&self, ids: &[String]) -> Vec<String> {
3392        let mut seen = HashSet::new();
3393        ids.iter()
3394            .filter_map(|id| self.tools.canonical_id(id))
3395            .filter(|canonical_id| seen.insert(canonical_id.clone()))
3396            .collect()
3397    }
3398
3399    /// Applies runtime narrowing to the canonical declared grant without permitting scope-based expansion.
3400    fn get_top_level_tool_ids_for_scope(&self, scope_override: Option<&[String]>) -> Vec<String> {
3401        let Some(declared) = self.declared_tool_ids.as_deref() else {
3402            return Vec::new();
3403        };
3404        let mut effective = self.canonical_unique_tool_ids(declared);
3405        if let Some(scope) = scope_override {
3406            let scope: HashSet<String> =
3407                self.canonical_unique_tool_ids(scope).into_iter().collect();
3408            effective.retain(|canonical_id| scope.contains(canonical_id));
3409        }
3410        effective
3411    }
3412
3413    /// Reads the live runtime narrowing and returns the current deterministic effective tool IDs.
3414    async fn get_available_tool_ids(&self) -> Result<Vec<String>> {
3415        Ok(self.get_available_tool_ids_snapshot().await?.tool_ids)
3416    }
3417
3418    /// Captures the runtime scope used for ordinary availability checks before applying state narrowing.
3419    async fn get_available_tool_ids_snapshot(&self) -> Result<AvailableToolIdsSnapshot> {
3420        let scope_override = self.runtime_control.tool_scope_override.read().clone();
3421        self.get_available_tool_ids_snapshot_for_scope(scope_override.as_deref())
3422            .await
3423    }
3424
3425    /// Applies every explicit ancestor and current-state scope to one runtime snapshot and retains its state generation.
3426    async fn get_available_tool_ids_snapshot_for_scope(
3427        &self,
3428        scope_override: Option<&[String]>,
3429    ) -> Result<AvailableToolIdsSnapshot> {
3430        let mut available = self.get_top_level_tool_ids_for_scope(scope_override);
3431        let (state_generation, state_scopes) = self
3432            .state_machine
3433            .as_ref()
3434            .map(|state_machine| {
3435                let (generation, scopes) = state_machine.current_tool_scope_snapshot();
3436                (Some(generation), scopes)
3437            })
3438            .unwrap_or((None, Vec::new()));
3439
3440        if available.is_empty() || state_scopes.is_empty() {
3441            return Ok(AvailableToolIdsSnapshot {
3442                tool_ids: available,
3443                state_generation,
3444            });
3445        }
3446
3447        let eval_ctx = self.build_evaluation_context().await?;
3448        let llm_getter = RegistryLLMGetter {
3449            registry: self.llm_registry.clone(),
3450        };
3451        let evaluator = ConditionEvaluator::new(llm_getter);
3452
3453        for state_scope in state_scopes {
3454            if state_scope.is_empty() {
3455                available.clear();
3456                break;
3457            }
3458
3459            let mut allowed = HashSet::new();
3460            for tool_ref in &state_scope {
3461                let tool_id = tool_ref.id();
3462                let Some(canonical_id) = self.tools.canonical_id(tool_id) else {
3463                    continue;
3464                };
3465                let condition_matches = if let Some(condition) = tool_ref.condition() {
3466                    match evaluator.evaluate(condition, &eval_ctx).await {
3467                        Ok(matches) => matches,
3468                        Err(error) => {
3469                            warn!(tool = tool_id, error = %error, "Error evaluating tool condition");
3470                            false
3471                        }
3472                    }
3473                } else {
3474                    true
3475                };
3476                if condition_matches {
3477                    allowed.insert(canonical_id);
3478                } else {
3479                    debug!(tool = tool_id, "Tool condition not met, skipping");
3480                }
3481            }
3482            available.retain(|canonical_id| allowed.contains(canonical_id));
3483            if available.is_empty() {
3484                break;
3485            }
3486        }
3487
3488        Ok(AvailableToolIdsSnapshot {
3489            tool_ids: available,
3490            state_generation,
3491        })
3492    }
3493
3494    async fn build_evaluation_context(&self) -> Result<EvaluationContext> {
3495        let context = self.build_context_with_overlays();
3496        let messages = self.memory.get_messages(Some(10)).await?;
3497        let tool_history = self.tool_call_history.read().clone();
3498
3499        let (state_name, turn_count, previous_state) = if let Some(ref sm) = self.state_machine {
3500            (Some(sm.current()), sm.turn_count(), sm.previous())
3501        } else {
3502            (None, 0, None)
3503        };
3504
3505        Ok(EvaluationContext::default()
3506            .with_context(context)
3507            .with_state(state_name, turn_count, previous_state)
3508            .with_called_tools(tool_history)
3509            .with_messages(messages))
3510    }
3511
3512    fn record_tool_call(&self, tool_id: &str, result: Value) {
3513        self.tool_call_history.write().push(ToolCallRecord {
3514            tool_id: tool_id.to_string(),
3515            result,
3516            timestamp: chrono::Utc::now(),
3517        });
3518    }
3519
3520    async fn get_effective_system_prompt_with_persona_hooks(
3521        &self,
3522        fire_persona_hooks: bool,
3523        include_tool_prompt: bool,
3524    ) -> Result<String> {
3525        let rendered_base = self.render_system_prompt()?;
3526
3527        let persona_prefix = if let Some(ref persona) = self.persona_manager {
3528            let context = self.build_context_with_overlays();
3529            if fire_persona_hooks {
3530                let render_result = persona.render_prompt(&context)?;
3531                for content in &render_result.newly_revealed {
3532                    self.hooks.on_secret_revealed(content).await;
3533                }
3534                render_result.prompt
3535            } else {
3536                persona.render_prompt_preview(&context)?
3537            }
3538        } else {
3539            String::new()
3540        };
3541
3542        if let Some(ref sm) = self.state_machine
3543            && let Some(state_def) = sm.current_definition()
3544        {
3545            let state_prompt = if let Some(ref prompt) = state_def.prompt {
3546                let context = self.build_context_with_overlays();
3547                self.template_renderer.render_with_state(
3548                    prompt,
3549                    &context,
3550                    &sm.current(),
3551                    sm.previous().as_deref(),
3552                    sm.turn_count(),
3553                    state_def.max_turns,
3554                )?
3555            } else {
3556                String::new()
3557            };
3558
3559            let combined = match state_def.prompt_mode {
3560                PromptMode::Append => {
3561                    if state_prompt.is_empty() {
3562                        rendered_base
3563                    } else {
3564                        format!(
3565                            "{}\n\n[Current State: {}]\n{}",
3566                            rendered_base,
3567                            sm.current(),
3568                            state_prompt
3569                        )
3570                    }
3571                }
3572                PromptMode::Replace => {
3573                    if state_prompt.is_empty() {
3574                        rendered_base
3575                    } else {
3576                        state_prompt
3577                    }
3578                }
3579                PromptMode::Prepend => {
3580                    if state_prompt.is_empty() {
3581                        rendered_base
3582                    } else {
3583                        format!("{}\n\n{}", state_prompt, rendered_base)
3584                    }
3585                }
3586            };
3587
3588            // Persona always prepended regardless of prompt_mode.
3589            let with_persona = if persona_prefix.is_empty() {
3590                combined
3591            } else {
3592                format!("{}\n\n{}", persona_prefix, combined)
3593            };
3594
3595            if include_tool_prompt {
3596                let available_tool_ids = self.get_available_tool_ids().await?;
3597                if !available_tool_ids.is_empty() {
3598                    let tools_prompt = self.tools.generate_scoped_prompt_with_mode(
3599                        &available_tool_ids,
3600                        None,
3601                        self.parallel_tools.enabled,
3602                        self.runtime_config.tool_schema_prompt_mode,
3603                    );
3604                    if !tools_prompt.is_empty() {
3605                        return Ok(format!("{}\n\n{}", with_persona, tools_prompt));
3606                    }
3607                }
3608            }
3609            return Ok(with_persona);
3610        }
3611
3612        // No state machine - prepend persona to base.
3613        let with_persona = if persona_prefix.is_empty() {
3614            rendered_base
3615        } else {
3616            format!("{}\n\n{}", persona_prefix, rendered_base)
3617        };
3618
3619        if include_tool_prompt {
3620            let available_tool_ids = self.get_available_tool_ids().await?;
3621            let tools_prompt = self.tools.generate_scoped_prompt_with_mode(
3622                &available_tool_ids,
3623                None,
3624                self.parallel_tools.enabled,
3625                self.runtime_config.tool_schema_prompt_mode,
3626            );
3627            if !tools_prompt.is_empty() {
3628                return Ok(format!("{}\n\n{}", with_persona, tools_prompt));
3629            }
3630        }
3631        Ok(with_persona)
3632    }
3633
3634    fn get_state_llm(&self) -> Result<Arc<dyn LLMProvider>> {
3635        if let Some(ref sm) = self.state_machine
3636            && let Some(state_def) = sm.current_definition()
3637            && let Some(ref llm_alias) = state_def.llm
3638        {
3639            return self
3640                .llm_registry
3641                .get(llm_alias)
3642                .map_err(|e| AgentError::Config(e.to_string()));
3643        }
3644        self.llm_registry
3645            .default()
3646            .map_err(|e| AgentError::Config(e.to_string()))
3647    }
3648
3649    fn get_effective_reasoning_config(&self) -> ReasoningConfig {
3650        if let Some(ref sm) = self.state_machine
3651            && let Some(state_def) = sm.current_definition()
3652            && let Some(ref state_reasoning) = state_def.reasoning
3653        {
3654            return state_reasoning.clone();
3655        }
3656        self.reasoning_config.clone()
3657    }
3658
3659    fn get_effective_reflection_config(&self) -> ReflectionConfig {
3660        if let Some(ref sm) = self.state_machine
3661            && let Some(state_def) = sm.current_definition()
3662            && let Some(ref state_reflection) = state_def.reflection
3663        {
3664            return state_reflection.clone();
3665        }
3666        self.reflection_config.clone()
3667    }
3668
3669    fn get_skill_reasoning_config(&self, skill: &SkillDefinition) -> ReasoningConfig {
3670        skill
3671            .reasoning
3672            .clone()
3673            .unwrap_or_else(|| self.get_effective_reasoning_config())
3674    }
3675
3676    fn get_skill_reflection_config(&self, skill: &SkillDefinition) -> ReflectionConfig {
3677        skill
3678            .reflection
3679            .clone()
3680            .unwrap_or_else(|| self.get_effective_reflection_config())
3681    }
3682
3683    async fn build_disambiguation_context(&self) -> Result<DisambiguationContext> {
3684        let recent_messages: Vec<String> = self
3685            .memory
3686            .get_messages(Some(5))
3687            .await?
3688            .iter()
3689            .rev()
3690            .map(|m| format!("{:?}: {}", m.role, m.content))
3691            .collect();
3692
3693        let current_state = self.current_state().map(|s| s.to_string());
3694
3695        // Include the current state's prompt text so the detector understands
3696        // what kind of input is expected (e.g., "Ask for the order number").
3697        let state_prompt: Option<String> = self
3698            .state_machine
3699            .as_ref()
3700            .and_then(|sm| sm.current_definition())
3701            .and_then(|def| def.prompt.clone());
3702
3703        let available_tools: Vec<String> = self
3704            .get_available_tool_ids()
3705            .await
3706            .unwrap_or_else(|_| self.tools.list_ids());
3707
3708        let available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
3709
3710        let mut user_context = self.build_context_with_overlays();
3711        user_context.remove(DISAMBIGUATION_STATE_GENERATION_KEY);
3712        if let Some(state_generation) = self
3713            .state_machine
3714            .as_ref()
3715            .map(|state_machine| state_machine.generation())
3716        {
3717            user_context.insert(
3718                DISAMBIGUATION_STATE_GENERATION_KEY.to_string(),
3719                serde_json::json!(state_generation),
3720            );
3721        }
3722
3723        // Extract canonical intent labels from current state's transitions
3724        let available_intents: Vec<String> = if let Some(ref sm) = self.state_machine {
3725            sm.current_definition()
3726                .map(|def| {
3727                    def.transitions
3728                        .iter()
3729                        .filter_map(|t| t.intent.clone())
3730                        .collect()
3731                })
3732                .unwrap_or_default()
3733        } else {
3734            Vec::new()
3735        };
3736
3737        Ok(DisambiguationContext::from_agent_state(
3738            recent_messages,
3739            current_state,
3740            state_prompt,
3741            available_tools,
3742            available_skills,
3743            available_intents,
3744            user_context,
3745        ))
3746    }
3747
3748    fn get_available_skills(&self) -> Vec<&SkillDefinition> {
3749        if let Some(ref sm) = self.state_machine
3750            && let Some(state_def) = sm.current_definition()
3751        {
3752            let parent_def = sm.get_parent_definition();
3753            let effective_skills = state_def.get_effective_skills(parent_def.as_ref());
3754            if !effective_skills.is_empty() {
3755                return self
3756                    .skills
3757                    .iter()
3758                    .filter(|s| effective_skills.contains(&&s.id))
3759                    .collect();
3760            }
3761        }
3762        self.skills.iter().collect()
3763    }
3764
3765    async fn build_messages(&self) -> Result<Vec<ChatMessage>> {
3766        self.build_messages_internal(true, None, true).await
3767    }
3768
3769    async fn build_messages_for_draft(&self, user_message: &str) -> Result<Vec<ChatMessage>> {
3770        self.build_messages_internal(false, Some(user_message), true)
3771            .await
3772    }
3773
3774    async fn build_messages_internal(
3775        &self,
3776        fire_persona_hooks: bool,
3777        ephemeral_user_message: Option<&str>,
3778        include_tool_prompt: bool,
3779    ) -> Result<Vec<ChatMessage>> {
3780        let system_prompt = self
3781            .get_effective_system_prompt_with_persona_hooks(fire_persona_hooks, include_tool_prompt)
3782            .await?;
3783        let mut messages = vec![ChatMessage::system(&system_prompt)];
3784
3785        let context = self.memory.get_context().await?;
3786        let history = if let Some(ref budget) = self.memory_token_budget {
3787            context.to_llm_messages_with_allocation(&budget.allocation)
3788        } else {
3789            context.to_llm_messages()
3790        };
3791        messages.extend(history);
3792        if let Some(user_message) = ephemeral_user_message {
3793            messages.push(ChatMessage::user(user_message));
3794        }
3795
3796        let total_tokens = self.estimate_total_tokens(&messages);
3797
3798        if total_tokens > self.max_context_tokens {
3799            debug!(
3800                total = total_tokens,
3801                limit = self.max_context_tokens,
3802                "Context overflow"
3803            );
3804
3805            match &self.recovery_manager.config().llm.on_context_overflow {
3806                ContextOverflowAction::Error => {
3807                    return Err(AgentError::LLM(format!(
3808                        "Context overflow: {} tokens > {} limit",
3809                        total_tokens, self.max_context_tokens
3810                    )));
3811                }
3812                ContextOverflowAction::Truncate { keep_recent } => {
3813                    self.truncate_context(&mut messages, *keep_recent);
3814                }
3815                ContextOverflowAction::Summarize {
3816                    summarizer_llm,
3817                    max_summary_tokens,
3818                    custom_prompt,
3819                    keep_recent,
3820                    filter,
3821                } => {
3822                    self.summarize_context(
3823                        &mut messages,
3824                        summarizer_llm.as_deref(),
3825                        *max_summary_tokens,
3826                        custom_prompt.as_deref(),
3827                        *keep_recent,
3828                        filter.as_ref(),
3829                    )
3830                    .await?;
3831                }
3832            }
3833        }
3834
3835        Ok(messages)
3836    }
3837
3838    async fn main_tool_protocol(
3839        &self,
3840        llm: &dyn LLMProvider,
3841        ephemeral_new_turn: bool,
3842    ) -> Result<MainToolProtocol> {
3843        let mut choice = llm.configured_tool_choice();
3844        if matches!(choice.as_ref(), Some(ToolChoice::None)) {
3845            return Ok(MainToolProtocol {
3846                choice,
3847                tool_ids: Vec::new(),
3848                definitions: Vec::new(),
3849            });
3850        }
3851
3852        let mut tool_ids = self.get_available_tool_ids().await?;
3853        tool_ids.sort();
3854        tool_ids.dedup();
3855        if let Some(ToolChoice::Specific(expected)) = choice.as_ref() {
3856            let canonical = self.tools.canonical_id(expected).ok_or_else(|| {
3857                AgentError::Config(format!(
3858                    "specific tool choice '{expected}' is not registered"
3859                ))
3860            })?;
3861            if canonical != *expected {
3862                return Err(AgentError::Config(format!(
3863                    "specific tool choice must use canonical ID '{canonical}', not '{expected}'"
3864                )));
3865            }
3866            if !tool_ids.iter().any(|tool_id| tool_id == expected) {
3867                return Err(AgentError::Config(format!(
3868                    "specific tool choice '{expected}' is outside the effective tool grant"
3869                )));
3870            }
3871        }
3872        if matches!(
3873            choice.as_ref(),
3874            Some(ToolChoice::Required | ToolChoice::Specific(_))
3875        ) && tool_ids.is_empty()
3876        {
3877            return Err(AgentError::Config(
3878                "required tool choice has no tool inside the effective grant".to_string(),
3879            ));
3880        }
3881        if !ephemeral_new_turn
3882            && let Some(configured_choice) = choice.as_ref()
3883            && matches!(
3884                configured_choice,
3885                ToolChoice::Required | ToolChoice::Specific(_)
3886            )
3887            && self
3888                .tool_choice_satisfied_in_current_turn(configured_choice, &tool_ids)
3889                .await?
3890        {
3891            choice = Some(ToolChoice::Auto);
3892        }
3893        if let Some(ToolChoice::Specific(expected)) = choice.as_ref() {
3894            tool_ids.retain(|tool_id| tool_id == expected);
3895        }
3896
3897        let definitions = tool_ids
3898            .iter()
3899            .map(|tool_id| {
3900                let tool = self.tools.get(tool_id).ok_or_else(|| {
3901                    AgentError::Config(format!(
3902                        "effective tool '{tool_id}' disappeared before provider exposure"
3903                    ))
3904                })?;
3905                Ok(LLMToolDefinition {
3906                    name: tool_id.clone(),
3907                    description: tool.description().to_string(),
3908                    input_schema: tool.input_schema(),
3909                })
3910            })
3911            .collect::<Result<Vec<_>>>()?;
3912
3913        //
3914        // Provider-visible definitions are derived only from the current effective grant. Tool choice never expands registration, scope, or authorization.
3915        //
3916        Ok(MainToolProtocol {
3917            choice,
3918            tool_ids,
3919            definitions,
3920        })
3921    }
3922
3923    async fn tool_choice_satisfied_in_current_turn(
3924        &self,
3925        choice: &ToolChoice,
3926        effective_tool_ids: &[String],
3927    ) -> Result<bool> {
3928        let messages = self.memory.get_messages(None).await?;
3929        let mut saw_tool_result = false;
3930        for message in messages.iter().rev() {
3931            match message.role {
3932                ai_agents_core::Role::Tool | ai_agents_core::Role::Function => {
3933                    saw_tool_result = true;
3934                }
3935                ai_agents_core::Role::Assistant if saw_tool_result => {
3936                    let Some(calls) = self.parse_tool_calls(&message.content) else {
3937                        continue;
3938                    };
3939                    let calls_are_effective = !calls.is_empty()
3940                        && calls.iter().all(|call| {
3941                            self.tools
3942                                .canonical_id(&call.name)
3943                                .is_some_and(|canonical| effective_tool_ids.contains(&canonical))
3944                        });
3945                    return Ok(calls_are_effective
3946                        && match choice {
3947                            ToolChoice::Required => true,
3948                            ToolChoice::Specific(expected) => calls.iter().all(|call| {
3949                                self.tools.canonical_id(&call.name).as_deref()
3950                                    == Some(expected.as_str())
3951                            }),
3952                            _ => false,
3953                        });
3954                }
3955                ai_agents_core::Role::User => return Ok(false),
3956                _ => {}
3957            }
3958        }
3959        Ok(false)
3960    }
3961
3962    fn provider_can_use_native_tools(
3963        &self,
3964        llm: &dyn LLMProvider,
3965        protocol: &MainToolProtocol,
3966    ) -> bool {
3967        let Some(choice) = protocol.choice.as_ref() else {
3968            return false;
3969        };
3970        if matches!(choice, ToolChoice::None) || protocol.definitions.is_empty() {
3971            return false;
3972        }
3973        llm.supports_tool_choice(choice)
3974            && protocol.definitions.iter().all(|definition| {
3975                !definition.name.is_empty()
3976                    && definition.name.len() <= 64
3977                    && definition
3978                        .name
3979                        .bytes()
3980                        .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
3981            })
3982    }
3983
3984    fn prompt_messages_for_tool_protocol(
3985        &self,
3986        messages: &[ChatMessage],
3987        protocol: &MainToolProtocol,
3988        corrective: bool,
3989    ) -> Vec<ChatMessage> {
3990        let mut messages = messages.to_vec();
3991        let Some(choice) = protocol.choice.as_ref() else {
3992            return messages;
3993        };
3994        if matches!(choice, ToolChoice::None) || protocol.tool_ids.is_empty() {
3995            return messages;
3996        }
3997
3998        let mut tool_prompt = self.tools.generate_scoped_prompt_with_mode(
3999            &protocol.tool_ids,
4000            None,
4001            self.parallel_tools.enabled,
4002            self.runtime_config.tool_schema_prompt_mode,
4003        );
4004        match choice {
4005            ToolChoice::Required => tool_prompt.push_str(
4006                "\n\nYou must call at least one listed tool before giving a final answer.",
4007            ),
4008            ToolChoice::Specific(tool_id) => tool_prompt.push_str(&format!(
4009                "\n\nYou must call the '{tool_id}' tool before giving a final answer."
4010            )),
4011            ToolChoice::Auto => {}
4012            ToolChoice::None => return messages,
4013            _ => return messages,
4014        }
4015        if let Some(system) = messages
4016            .iter_mut()
4017            .find(|message| message.role == ai_agents_core::Role::System)
4018        {
4019            system.content.push_str("\n\n");
4020            system.content.push_str(&tool_prompt);
4021        } else {
4022            messages.insert(0, ChatMessage::system(tool_prompt));
4023        }
4024        if corrective {
4025            let instruction = match choice {
4026                ToolChoice::Required => {
4027                    "Your previous response did not call a required tool. Call at least one listed tool now and return only the JSON tool call."
4028                }
4029                ToolChoice::Specific(tool_id) => {
4030                    messages.push(ChatMessage::user(format!(
4031                        "Your previous response did not call the required '{tool_id}' tool. Call it now and return only the JSON tool call."
4032                    )));
4033                    return messages;
4034                }
4035                _ => return messages,
4036            };
4037            messages.push(ChatMessage::user(instruction));
4038        }
4039        messages
4040    }
4041
4042    async fn invoke_main_provider(
4043        &self,
4044        llm: Arc<dyn LLMProvider>,
4045        messages: &[ChatMessage],
4046        protocol: &MainToolProtocol,
4047        corrective: bool,
4048    ) -> std::result::Result<MainProviderResponse, LLMError> {
4049        let use_native = self.provider_can_use_native_tools(llm.as_ref(), protocol);
4050        let response = if use_native {
4051            let request = LLMToolRequest {
4052                tools: protocol.definitions.clone(),
4053                choice: protocol
4054                    .choice
4055                    .clone()
4056                    .expect("native tool requests require an explicit choice"),
4057            };
4058            self.observe_purpose(
4059                ObservationPurpose::MainResponse,
4060                llm.complete_with_tools(messages, None, &request),
4061            )
4062            .await?
4063        } else {
4064            let prompt_messages =
4065                self.prompt_messages_for_tool_protocol(messages, protocol, corrective);
4066            self.observe_purpose(
4067                ObservationPurpose::MainResponse,
4068                llm.complete(&prompt_messages, None),
4069            )
4070            .await?
4071        };
4072        Ok(MainProviderResponse {
4073            response,
4074            used_native_tools: use_native,
4075        })
4076    }
4077
4078    async fn complete_main_attempt_with_recovery(
4079        &self,
4080        llm: Arc<dyn LLMProvider>,
4081        messages: &[ChatMessage],
4082        protocol: &MainToolProtocol,
4083        corrective: bool,
4084    ) -> Result<MainProviderResponse> {
4085        let primary_result = if self.recovery_manager.config().default.max_retries > 0 {
4086            self.recovery_manager
4087                .with_retry("llm_call", None, || {
4088                    let llm = Arc::clone(&llm);
4089                    async move {
4090                        self.invoke_main_provider(llm, messages, protocol, corrective)
4091                            .await
4092                            .map_err(|error| error.classify())
4093                    }
4094                })
4095                .await
4096                .map_err(|error| AgentError::LLM(error.to_string()))
4097        } else {
4098            self.invoke_main_provider(Arc::clone(&llm), messages, protocol, corrective)
4099                .await
4100                .map_err(|error| AgentError::LLM(error.to_string()))
4101        };
4102
4103        match primary_result {
4104            Ok(response) => Ok(response),
4105            Err(primary_error) => match &self.recovery_manager.config().llm.on_failure {
4106                LLMFailureAction::FallbackLlm { fallback_llm } => {
4107                    let fallback = self.llm_registry.get(fallback_llm).map_err(|error| {
4108                        AgentError::Config(format!(
4109                            "Fallback LLM '{fallback_llm}' not found: {error}"
4110                        ))
4111                    })?;
4112                    self.invoke_main_provider(fallback, messages, protocol, corrective)
4113                        .await
4114                        .map_err(|error| AgentError::LLM(error.to_string()))
4115                }
4116                LLMFailureAction::FallbackResponse { message } => {
4117                    if matches!(
4118                        protocol.choice.as_ref(),
4119                        Some(ToolChoice::Required | ToolChoice::Specific(_))
4120                    ) {
4121                        Err(AgentError::LLM(format!(
4122                            "Required tool selection failed and cannot be satisfied by a static fallback response: {primary_error}"
4123                        )))
4124                    } else {
4125                        Ok(MainProviderResponse {
4126                            response: LLMResponse::new(message.clone(), FinishReason::Stop),
4127                            used_native_tools: false,
4128                        })
4129                    }
4130                }
4131                LLMFailureAction::Error => Err(primary_error),
4132            },
4133        }
4134    }
4135
4136    fn normalize_main_provider_response(
4137        &self,
4138        mut response: LLMResponse,
4139        protocol: &MainToolProtocol,
4140    ) -> Result<(LLMResponse, bool)> {
4141        let native_calls = response
4142            .tool_calls()
4143            .map_err(|error| AgentError::LLM(error.to_string()))?;
4144        let calls = match native_calls {
4145            Some(calls) => {
4146                let markers = calls
4147                    .iter()
4148                    .map(|call| {
4149                        serde_json::json!({
4150                            "_ai_agents_native_tool_call": true,
4151                            "id": call.id,
4152                            "tool": call.name,
4153                            "arguments": call.arguments,
4154                        })
4155                    })
4156                    .collect::<Vec<_>>();
4157                response.content = if markers.len() == 1 {
4158                    markers[0].to_string()
4159                } else {
4160                    serde_json::Value::Array(markers).to_string()
4161                };
4162                Some(calls)
4163            }
4164            None if !matches!(protocol.choice.as_ref(), Some(ToolChoice::None)) => {
4165                self.parse_tool_calls(response.content.trim())
4166            }
4167            None => None,
4168        };
4169
4170        if protocol.choice.is_some()
4171            && let Some(calls) = calls.as_ref()
4172            && calls.iter().any(|call| {
4173                self.tools
4174                    .canonical_id(&call.name)
4175                    .is_none_or(|canonical| !protocol.tool_ids.contains(&canonical))
4176            })
4177        {
4178            return Err(AgentError::LLM(
4179                "Provider returned a tool call outside the effective grant".to_string(),
4180            ));
4181        }
4182
4183        let compliant = match protocol.choice.as_ref() {
4184            Some(ToolChoice::Required) => calls.as_ref().is_some_and(|calls| !calls.is_empty()),
4185            Some(ToolChoice::Specific(expected)) => calls.as_ref().is_some_and(|calls| {
4186                !calls.is_empty()
4187                    && calls.iter().all(|call| {
4188                        self.tools.canonical_id(&call.name).as_deref() == Some(expected.as_str())
4189                    })
4190            }),
4191            _ => true,
4192        };
4193        Ok((response, compliant))
4194    }
4195
4196    async fn complete_main_llm_with_recovery(
4197        &self,
4198        llm: Arc<dyn LLMProvider>,
4199        messages: &[ChatMessage],
4200        protocol: &MainToolProtocol,
4201    ) -> Result<LLMResponse> {
4202        let first = self
4203            .complete_main_attempt_with_recovery(Arc::clone(&llm), messages, protocol, false)
4204            .await?;
4205        let (response, compliant) =
4206            self.normalize_main_provider_response(first.response, protocol)?;
4207        if compliant {
4208            return Ok(response);
4209        }
4210        if first.used_native_tools {
4211            return Err(AgentError::LLM(
4212                "Provider returned no compliant native call for required tool choice".to_string(),
4213            ));
4214        }
4215
4216        let corrected = self
4217            .complete_main_attempt_with_recovery(llm, messages, protocol, true)
4218            .await?;
4219        let (response, compliant) =
4220            self.normalize_main_provider_response(corrected.response, protocol)?;
4221        if compliant {
4222            return Ok(response);
4223        }
4224        Err(AgentError::LLM(
4225            "Provider returned no compliant tool call after one corrective retry".to_string(),
4226        ))
4227    }
4228
4229    fn is_native_tool_call_content(content: &str) -> bool {
4230        let Ok(value) = serde_json::from_str::<serde_json::Value>(content) else {
4231            return false;
4232        };
4233        match value {
4234            serde_json::Value::Array(values) => {
4235                !values.is_empty()
4236                    && values.iter().all(|value| {
4237                        value
4238                            .get("_ai_agents_native_tool_call")
4239                            .and_then(|marker| marker.as_bool())
4240                            == Some(true)
4241                    })
4242            }
4243            serde_json::Value::Object(map) => {
4244                map.get("_ai_agents_native_tool_call")
4245                    .and_then(|marker| marker.as_bool())
4246                    == Some(true)
4247            }
4248            _ => false,
4249        }
4250    }
4251
4252    fn tool_result_message(
4253        tool_call: &ToolCall,
4254        output: &str,
4255        native_tool_call: bool,
4256    ) -> ChatMessage {
4257        if !native_tool_call {
4258            return ChatMessage::function(&tool_call.name, output);
4259        }
4260        let output = serde_json::from_str::<serde_json::Value>(output)
4261            .unwrap_or_else(|_| serde_json::Value::String(output.to_string()));
4262        ChatMessage::function(
4263            &tool_call.name,
4264            serde_json::json!({
4265                "_ai_agents_native_tool_result": true,
4266                "id": tool_call.id,
4267                "tool": tool_call.name,
4268                "output": output,
4269            })
4270            .to_string(),
4271        )
4272    }
4273
4274    fn parse_main_tool_calls(
4275        &self,
4276        content: &str,
4277        protocol: &MainToolProtocol,
4278    ) -> Option<Vec<ToolCall>> {
4279        if matches!(protocol.choice.as_ref(), Some(ToolChoice::None)) {
4280            None
4281        } else {
4282            self.parse_tool_calls(content)
4283        }
4284    }
4285
4286    fn parse_tool_calls(&self, content: &str) -> Option<Vec<ToolCall>> {
4287        // Try direct JSON parse first
4288        if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(content) {
4289            // Handle JSON array of tool calls (parallel tool calling)
4290            if let Some(arr) = parsed.as_array() {
4291                let calls: Vec<ToolCall> = arr
4292                    .iter()
4293                    .filter_map(|v| self.extract_tool_call_from_value(v))
4294                    .collect();
4295                if !calls.is_empty() {
4296                    return Some(calls);
4297                }
4298            }
4299            // Handle single JSON object
4300            if let Some(tool_call) = self.extract_tool_call_from_value(&parsed) {
4301                return Some(vec![tool_call]);
4302            }
4303        }
4304
4305        // Try to extract JSON from content (handles extra text/braces from LLM)
4306        if let Some(json_str) = self.extract_json_from_content(content)
4307            && let Ok(parsed) = serde_json::from_str::<serde_json::Value>(&json_str)
4308        {
4309            // Handle JSON array of tool calls (parallel tool calling)
4310            if let Some(arr) = parsed.as_array() {
4311                let calls: Vec<ToolCall> = arr
4312                    .iter()
4313                    .filter_map(|v| self.extract_tool_call_from_value(v))
4314                    .collect();
4315                if !calls.is_empty() {
4316                    return Some(calls);
4317                }
4318            }
4319            // Handle single JSON object
4320            if let Some(tool_call) = self.extract_tool_call_from_value(&parsed) {
4321                return Some(vec![tool_call]);
4322            }
4323        }
4324
4325        None
4326    }
4327
4328    fn extract_tool_call_from_value(&self, parsed: &serde_json::Value) -> Option<ToolCall> {
4329        if let Some(tool_name) = parsed.get("tool").and_then(|v| v.as_str()) {
4330            let arguments = parsed
4331                .get("arguments")
4332                .cloned()
4333                .unwrap_or(serde_json::json!({}));
4334            return Some(ToolCall {
4335                id: parsed
4336                    .get("id")
4337                    .and_then(|value| value.as_str())
4338                    .filter(|id| !id.is_empty())
4339                    .map(str::to_string)
4340                    .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
4341                name: tool_name.to_string(),
4342                arguments,
4343            });
4344        }
4345        None
4346    }
4347
4348    // Lite models could generate unmatched braces: this function handles such cases
4349    fn extract_json_from_content(&self, content: &str) -> Option<String> {
4350        // Try array first (for parallel tool calls), then single object
4351        if let Some(result) = self.extract_json_array_from_content(content) {
4352            return Some(result);
4353        }
4354        self.extract_json_object_from_content(content)
4355    }
4356
4357    /// Extract a JSON array `[...]` containing tool calls from mixed content.
4358    fn extract_json_array_from_content(&self, content: &str) -> Option<String> {
4359        let start = content.find('[')?;
4360        let content_from_start = &content[start..];
4361
4362        let mut depth = 0;
4363        let mut end = 0;
4364        for (i, ch) in content_from_start.char_indices() {
4365            match ch {
4366                '[' => depth += 1,
4367                ']' => {
4368                    depth -= 1;
4369                    if depth == 0 {
4370                        end = i + 1;
4371                        break;
4372                    }
4373                }
4374                _ => {}
4375            }
4376        }
4377
4378        if end > 0 {
4379            let json_str = &content_from_start[..end];
4380            // Verify it looks like an array of tool calls
4381            if json_str.contains("\"tool\"") {
4382                return Some(json_str.to_string());
4383            }
4384        }
4385
4386        None
4387    }
4388
4389    /// Extract a JSON object `{...}` containing a tool call from mixed content.
4390    fn extract_json_object_from_content(&self, content: &str) -> Option<String> {
4391        let start = content.find('{')?;
4392        let content_from_start = &content[start..];
4393
4394        // Count braces to find the matching closing brace
4395        let mut depth = 0;
4396        let mut end = 0;
4397        for (i, ch) in content_from_start.char_indices() {
4398            match ch {
4399                '{' => depth += 1,
4400                '}' => {
4401                    depth -= 1;
4402                    if depth == 0 {
4403                        end = i + 1;
4404                        break;
4405                    }
4406                }
4407                _ => {}
4408            }
4409        }
4410
4411        if end > 0 {
4412            let json_str = &content_from_start[..end];
4413            // Verify it looks like a tool call
4414            if json_str.contains("\"tool\"") {
4415                return Some(json_str.to_string());
4416            }
4417        }
4418
4419        None
4420    }
4421
4422    /// Builds a structured record from executor state.
4423    ///
4424    /// The explicit fields preserve one audit boundary for execution, policy, approval, timing, and generation evidence.
4425    #[allow(clippy::too_many_arguments)]
4426    fn record_from_parts(
4427        &self,
4428        request: &ToolExecutionRequest,
4429        canonical_id: String,
4430        executed_arguments: Value,
4431        started_at: chrono::DateTime<chrono::Utc>,
4432        start: Instant,
4433        executed: bool,
4434        success: bool,
4435        output: String,
4436        metadata: HashMap<String, Value>,
4437        policy: ToolPolicyDecisionRecord,
4438        approval: Option<ToolApprovalRecord>,
4439        timed_out: bool,
4440        output_truncated: bool,
4441    ) -> ToolExecutionRecord {
4442        let versions = ToolDecisionVersions {
4443            policy: self.active_tool_security().policy_version(),
4444            registry: self.tools.version(),
4445            runtime_control: self.runtime_control.version.load(Ordering::SeqCst),
4446            state: self
4447                .state_machine
4448                .as_ref()
4449                .map(|state_machine| state_machine.generation()),
4450        };
4451        self.record_from_parts_at(
4452            request,
4453            canonical_id,
4454            executed_arguments,
4455            started_at,
4456            start,
4457            executed,
4458            success,
4459            output,
4460            metadata,
4461            policy,
4462            approval,
4463            timed_out,
4464            output_truncated,
4465            versions,
4466        )
4467    }
4468
4469    /// Builds evidence with the exact generations used by the corresponding decision.
4470    #[allow(clippy::too_many_arguments)]
4471    fn record_from_parts_at(
4472        &self,
4473        request: &ToolExecutionRequest,
4474        canonical_id: String,
4475        executed_arguments: Value,
4476        started_at: chrono::DateTime<chrono::Utc>,
4477        start: Instant,
4478        executed: bool,
4479        success: bool,
4480        output: String,
4481        metadata: HashMap<String, Value>,
4482        policy: ToolPolicyDecisionRecord,
4483        approval: Option<ToolApprovalRecord>,
4484        timed_out: bool,
4485        output_truncated: bool,
4486        versions: ToolDecisionVersions,
4487    ) -> ToolExecutionRecord {
4488        ToolExecutionRecord {
4489            call_id: request.call_id.clone(),
4490            requested_name: request.requested_name.clone(),
4491            canonical_id,
4492            source: request.source.clone(),
4493            arguments: request.arguments.clone(),
4494            executed_arguments,
4495            policy_version: versions.policy,
4496            registry_version: versions.registry,
4497            runtime_config_version: versions.runtime_control,
4498            executed,
4499            success,
4500            output,
4501            metadata,
4502            policy,
4503            approval,
4504            started_at,
4505            duration_ms: start.elapsed().as_millis() as u64,
4506            timed_out,
4507            cancelled: false,
4508            cancellation_reason: None,
4509            output_truncated,
4510        }
4511    }
4512
4513    /// Sends a finalized tool record to hooks, history, and error handling.
4514    async fn finish_tool_record(&self, record: &ToolExecutionRecord) {
4515        let result = ToolResult {
4516            success: record.success,
4517            output: record.model_output_string(),
4518            metadata: if record.metadata.is_empty() {
4519                None
4520            } else {
4521                Some(record.metadata.clone())
4522            },
4523        };
4524        self.hooks
4525            .on_tool_complete(&record.canonical_id, &result, record.duration_ms)
4526            .await;
4527        self.hooks.on_tool_execution_record(record).await;
4528        self.record_tool_call(&record.canonical_id, record.model_output_value());
4529        if !record.success {
4530            self.hooks
4531                .on_error(&AgentError::Tool(record.output.clone()))
4532                .await;
4533        }
4534    }
4535
4536    /// Releases resource locks before hooks run because hooks may invoke another tool.
4537    async fn finish_tool_record_after_resource_guards(
4538        &self,
4539        resource_guards: ToolResourceGuards,
4540        record: &ToolExecutionRecord,
4541    ) {
4542        drop(resource_guards);
4543        self.finish_tool_record(record).await;
4544    }
4545
4546    /// Invokes a resolved tool once with one fresh attempt deadline, timeout, cancellation, and actor context.
4547    async fn execute_resolved_tool_once(
4548        &self,
4549        tool: Arc<dyn ai_agents_core::Tool>,
4550        args: Value,
4551        mut ctx: ToolExecutionContext,
4552        timeout_ms: u64,
4553    ) -> Result<(ToolResult, bool, bool, bool)> {
4554        if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4555            return Ok((
4556                ToolResult::error("Tool execution cancelled by runtime control"),
4557                false,
4558                true,
4559                false,
4560            ));
4561        }
4562        //
4563        // The tool observes the same timeout budget enforced below, starting immediately before this invocation attempt.
4564        // Each retry receives a new deadline rather than inheriting time spent in policy handling or earlier attempts.
4565        //
4566        ctx.deadline = Some(chrono::Utc::now() + chrono::Duration::milliseconds(timeout_ms as i64));
4567        //
4568        // Mark invocation inside the future so cancellation before the first poll remains executed false.
4569        //
4570        let invoked = Arc::new(AtomicBool::new(false));
4571        let invoked_by_future = Arc::clone(&invoked);
4572        let actor_context = current_turn_actor_context();
4573        let future = async move {
4574            invoked_by_future.store(true, Ordering::SeqCst);
4575            if let Some(actor_context) = actor_context {
4576                scope_actor_context(actor_context, tool.execute(args, ctx)).await
4577            } else {
4578                tool.execute(args, ctx).await
4579            }
4580        };
4581        tokio::pin!(future);
4582        let timeout = tokio::time::sleep(std::time::Duration::from_millis(timeout_ms));
4583        tokio::pin!(timeout);
4584        let mut cancel_tick = tokio::time::interval(std::time::Duration::from_millis(50));
4585
4586        loop {
4587            tokio::select! {
4588                result = &mut future => return Ok((result, false, false, true)),
4589                _ = &mut timeout => {
4590                    return Ok((
4591                        ToolResult::error("Tool execution timed out"),
4592                        true,
4593                        false,
4594                        invoked.load(Ordering::SeqCst),
4595                    ));
4596                }
4597                _ = cancel_tick.tick() => {
4598                    if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4599                        return Ok((
4600                            ToolResult::error("Tool execution cancelled by runtime control"),
4601                            false,
4602                            true,
4603                            invoked.load(Ordering::SeqCst),
4604                        ));
4605                    }
4606                }
4607            }
4608        }
4609    }
4610
4611    /// Truncates model-facing tool output on a character boundary.
4612    fn truncate_tool_output(output: String, max_chars: Option<usize>) -> (String, bool) {
4613        let Some(max_chars) = max_chars else {
4614            return (output, false);
4615        };
4616        let mut chars = output.chars();
4617        let truncated: String = chars.by_ref().take(max_chars).collect();
4618        if chars.next().is_some() {
4619            (truncated, true)
4620        } else {
4621            (output, false)
4622        }
4623    }
4624
4625    /// Acquires all declared resource locks in stable key order and supports emergency cancellation while waiting.
4626    async fn acquire_tool_resource_locks(&self, keys: &[String]) -> Option<ToolResourceGuards> {
4627        let locks = {
4628            let mut table = self.resource_locks.write();
4629            table.retain(|_, lock| lock.strong_count() > 0);
4630            keys.iter()
4631                .map(|key| {
4632                    if let Some(lock) = table.get(key).and_then(Weak::upgrade) {
4633                        lock
4634                    } else {
4635                        let lock = Arc::new(tokio::sync::Mutex::new(()));
4636                        table.insert(key.clone(), Arc::downgrade(&lock));
4637                        lock
4638                    }
4639                })
4640                .collect::<Vec<_>>()
4641        };
4642        let mut resource_guards = ToolResourceGuards {
4643            guards: Vec::with_capacity(locks.len()),
4644            locks: Arc::clone(&self.resource_locks),
4645        };
4646        let mut locks = locks.into_iter();
4647        while let Some(lock) = locks.next() {
4648            let mut lock = Box::pin(lock.lock_owned());
4649            loop {
4650                tokio::select! {
4651                    guard = &mut lock => {
4652                        resource_guards.guards.push(guard);
4653                        break;
4654                    }
4655                    _ = tokio::time::sleep(std::time::Duration::from_millis(10)) => {
4656                        if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4657                            drop(lock);
4658                            drop(locks);
4659                            drop(resource_guards);
4660                            return None;
4661                        }
4662                    }
4663                }
4664            }
4665        }
4666        Some(resource_guards)
4667    }
4668
4669    /// Executes retry sub-attempts inside one logical executor request while assigning every invocation a fresh deadline.
4670    ///
4671    /// Hooks and `ToolExecutionRecord` finalization remain request-level and occur once after this retry loop returns.
4672    async fn run_tool_with_retries(
4673        &self,
4674        canonical_id: &str,
4675        tool: Arc<dyn ai_agents_core::Tool>,
4676        args: Value,
4677        ctx: ToolExecutionContext,
4678        timeout_ms: u64,
4679        max_retries: u32,
4680    ) -> Result<(ToolResult, bool, bool, bool)> {
4681        let max_retries = if ctx.classification.safely_retryable {
4682            max_retries
4683        } else {
4684            0
4685        };
4686        let mut attempts = 0;
4687        let mut invoked = false;
4688        loop {
4689            let (result, timed_out, cancelled, attempt_invoked) = self
4690                .execute_resolved_tool_once(tool.clone(), args.clone(), ctx.clone(), timeout_ms)
4691                .await?;
4692            invoked |= attempt_invoked;
4693            if result.success || timed_out || cancelled || attempts >= max_retries {
4694                return Ok((result, timed_out, cancelled, invoked));
4695            }
4696            attempts += 1;
4697            warn!(tool = %canonical_id, attempt = attempts, error = %result.output, "Retrying failed tool call");
4698        }
4699    }
4700
4701    /// Returns the existing model output and evidence reason when a host-backed tool cannot run.
4702    fn host_tool_unavailability(&self, canonical_id: &str) -> Option<(&'static str, &'static str)> {
4703        match canonical_id {
4704            "command" if !self.tools.command_runner_available() => Some((
4705                "Command runner is unavailable",
4706                "command runner is unavailable",
4707            )),
4708            "diagnostics" if !self.tools.diagnostics_available() => Some((
4709                "Diagnostics provider is unavailable",
4710                "diagnostics provider is unavailable",
4711            )),
4712            "web_search" if !self.tools.web_search_available() => Some((
4713                "Web search provider is unavailable",
4714                "web search provider is unavailable",
4715            )),
4716            _ => None,
4717        }
4718    }
4719
4720    /// Executes a tool request through scope, policy, HITL, timeout, recovery, and evidence recording.
4721    fn execute_tool_record(
4722        &self,
4723        request: ToolExecutionRequest,
4724    ) -> Pin<Box<dyn Future<Output = Result<ToolExecutionRecord>> + Send + '_>> {
4725        Box::pin(self.execute_tool_record_inner(request))
4726    }
4727
4728    /// Implements one logical shared-executor request while preserving policy, HITL, availability, final admission, hooks, retry evidence, and fallback ordering.
4729    ///
4730    /// A failed request selected for fallback releases its guards and finalizes its own record before the fallback starts as a separate shared-executor request.
4731    async fn execute_tool_record_inner(
4732        &self,
4733        request: ToolExecutionRequest,
4734    ) -> Result<ToolExecutionRecord> {
4735        let started_at = chrono::Utc::now();
4736        let start = Instant::now();
4737        info!(tool = %request.requested_name, args = %request.arguments, "Executing tool");
4738
4739        if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4740            let record = self.record_from_parts(
4741                &request,
4742                request.requested_name.clone(),
4743                request.arguments.clone(),
4744                started_at,
4745                start,
4746                false,
4747                false,
4748                "Tool execution is disabled by runtime control".to_string(),
4749                HashMap::new(),
4750                ToolPolicyDecisionRecord::deny("runtime emergency deny is enabled"),
4751                None,
4752                false,
4753                false,
4754            );
4755            self.finish_tool_record(&record).await;
4756            return Ok(record);
4757        }
4758
4759        let Some(resolved) = self.tools.resolve(&request.requested_name) else {
4760            let record = self.record_from_parts(
4761                &request,
4762                request.requested_name.clone(),
4763                request.arguments.clone(),
4764                started_at,
4765                start,
4766                false,
4767                false,
4768                format!("Tool '{}' is unavailable", request.requested_name),
4769                HashMap::new(),
4770                ToolPolicyDecisionRecord::unavailable(format!(
4771                    "Tool '{}' is not registered",
4772                    request.requested_name
4773                )),
4774                None,
4775                false,
4776                false,
4777            );
4778            self.finish_tool_record(&record).await;
4779            return Ok(record);
4780        };
4781
4782        let canonical_id = resolved.identity.canonical_id.clone();
4783
4784        let initial_scope_snapshot = self.get_available_tool_ids_snapshot().await?;
4785        if !initial_scope_snapshot
4786            .tool_ids
4787            .iter()
4788            .any(|id| id == &canonical_id)
4789        {
4790            let record = self.record_from_parts(
4791                &request,
4792                canonical_id.clone(),
4793                request.arguments.clone(),
4794                started_at,
4795                start,
4796                false,
4797                false,
4798                format!(
4799                    "Tool '{}' is not available in the current scope",
4800                    canonical_id
4801                ),
4802                HashMap::new(),
4803                ToolPolicyDecisionRecord::deny(format!(
4804                    "Tool '{}' is not granted by the current top-level and state tool scope",
4805                    canonical_id
4806                )),
4807                None,
4808                false,
4809                false,
4810            );
4811            self.finish_tool_record(&record).await;
4812            return Ok(record);
4813        }
4814
4815        let approval_control_snapshot = self.runtime_safety_snapshot();
4816        let security_engine = approval_control_snapshot.tool_security.clone();
4817        let bindings = resolved.tool.policy_bindings();
4818        let mut executed_arguments = security_engine.prepare_tool_arguments_with_bindings(
4819            &canonical_id,
4820            &request.arguments,
4821            &bindings,
4822        );
4823        self.hooks
4824            .on_tool_start(&canonical_id, &executed_arguments)
4825            .await;
4826
4827        let mut metadata = HashMap::new();
4828        let safety = resolved.tool.safety_metadata();
4829        let classification = resolved.tool.classify_call(&executed_arguments);
4830        let limits = security_engine.effective_limits(&canonical_id, &safety, &classification);
4831        metadata.insert(
4832            "classification".to_string(),
4833            serde_json::to_value(&classification).unwrap_or(Value::Null),
4834        );
4835        metadata.insert(
4836            "effective_limits".to_string(),
4837            serde_json::to_value(&limits).unwrap_or(Value::Null),
4838        );
4839        let policy_snapshot = security_engine.policy_snapshot(&canonical_id);
4840        if !policy_snapshot.is_null() {
4841            metadata.insert("policy_snapshot".to_string(), policy_snapshot.clone());
4842        }
4843
4844        let mut approval_record = Some(ToolApprovalRecord {
4845            status: ToolApprovalStatus::NotRequired,
4846            reason: None,
4847            modified_arguments: None,
4848        });
4849
4850        let mut security_result = security_engine
4851            .validate_tool_execution_with_bindings(&canonical_id, &executed_arguments, &bindings)
4852            .await?;
4853        //
4854        // Terminal policy denials remain authoritative, while allowed or confirmation-requiring calls must fail host availability before any HITL request.
4855        // The same capability is checked again after HITL because the host can change while approval waits.
4856        //
4857        if (security_result.is_allowed()
4858            || matches!(
4859                &security_result,
4860                SecurityCheckResult::RequireConfirmation { .. }
4861            ))
4862            && let Some((output, reason)) = self.host_tool_unavailability(&canonical_id)
4863        {
4864            let record = self.record_from_parts(
4865                &request,
4866                canonical_id,
4867                executed_arguments,
4868                started_at,
4869                start,
4870                false,
4871                false,
4872                output.to_string(),
4873                metadata,
4874                ToolPolicyDecisionRecord::unavailable(reason),
4875                Some(ToolApprovalRecord {
4876                    status: ToolApprovalStatus::Unavailable,
4877                    reason: Some(reason.to_string()),
4878                    modified_arguments: None,
4879                }),
4880                false,
4881                false,
4882            );
4883            self.finish_tool_record(&record).await;
4884            return Ok(record);
4885        }
4886        match &security_result {
4887            SecurityCheckResult::Allow => {}
4888            SecurityCheckResult::Warn { message } => {
4889                warn!(tool = %canonical_id, message = %message, "Tool security warning");
4890            }
4891            SecurityCheckResult::Block { reason } => {
4892                let record = self.record_from_parts(
4893                    &request,
4894                    canonical_id,
4895                    executed_arguments,
4896                    started_at,
4897                    start,
4898                    false,
4899                    false,
4900                    format!("Denied: {}", reason),
4901                    metadata,
4902                    ToolPolicyDecisionRecord::deny(reason.clone()),
4903                    approval_record,
4904                    false,
4905                    false,
4906                );
4907                self.finish_tool_record(&record).await;
4908                return Ok(record);
4909            }
4910            SecurityCheckResult::Unavailable { reason } => {
4911                let record = self.record_from_parts(
4912                    &request,
4913                    canonical_id,
4914                    executed_arguments,
4915                    started_at,
4916                    start,
4917                    false,
4918                    false,
4919                    format!("Unavailable: {}", reason),
4920                    metadata,
4921                    ToolPolicyDecisionRecord::unavailable(reason.clone()),
4922                    approval_record,
4923                    false,
4924                    false,
4925                );
4926                self.finish_tool_record(&record).await;
4927                return Ok(record);
4928            }
4929            SecurityCheckResult::RequireConfirmation { message } => {
4930                if self.hitl_engine.is_none() {
4931                    approval_record = Some(ToolApprovalRecord {
4932                        status: ToolApprovalStatus::Unavailable,
4933                        reason: Some("No HITL engine configured".to_string()),
4934                        modified_arguments: None,
4935                    });
4936                    let record = self.record_from_parts(
4937                        &request,
4938                        canonical_id,
4939                        executed_arguments,
4940                        started_at,
4941                        start,
4942                        false,
4943                        false,
4944                        format!("Approval unavailable: {}", message),
4945                        metadata,
4946                        ToolPolicyDecisionRecord::approval(message.clone()),
4947                        approval_record,
4948                        false,
4949                        false,
4950                    );
4951                    self.finish_tool_record(&record).await;
4952                    return Ok(record);
4953                }
4954
4955                let check_result = HITLCheckResult::required(
4956                    ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
4957                    HashMap::new(),
4958                    message.clone(),
4959                    None,
4960                );
4961                match self.request_hitl_approval(check_result).await? {
4962                    ApprovalResult::Approved => {
4963                        merge_approved_record(&mut approval_record);
4964                    }
4965                    ApprovalResult::Modified { changes } => {
4966                        if let Some(obj) = executed_arguments.as_object_mut() {
4967                            for (key, value) in changes {
4968                                obj.insert(key, value);
4969                            }
4970                        }
4971                        security_result = security_engine
4972                            .validate_tool_execution_with_bindings(
4973                                &canonical_id,
4974                                &executed_arguments,
4975                                &bindings,
4976                            )
4977                            .await?;
4978                        if !matches!(
4979                            security_result,
4980                            SecurityCheckResult::Allow
4981                                | SecurityCheckResult::Warn { .. }
4982                                | SecurityCheckResult::RequireConfirmation { .. }
4983                        ) {
4984                            let reason = security_result
4985                                .reason()
4986                                .unwrap_or("modified arguments failed policy")
4987                                .to_string();
4988                            let record = self.record_from_parts(
4989                                &request,
4990                                canonical_id,
4991                                executed_arguments.clone(),
4992                                started_at,
4993                                start,
4994                                false,
4995                                false,
4996                                reason.clone(),
4997                                metadata,
4998                                ToolPolicyDecisionRecord::deny(reason),
4999                                Some(ToolApprovalRecord {
5000                                    status: ToolApprovalStatus::Modified,
5001                                    reason: None,
5002                                    modified_arguments: Some(executed_arguments),
5003                                }),
5004                                false,
5005                                false,
5006                            );
5007                            self.finish_tool_record(&record).await;
5008                            return Ok(record);
5009                        }
5010                        approval_record = Some(ToolApprovalRecord {
5011                            status: ToolApprovalStatus::Modified,
5012                            reason: None,
5013                            modified_arguments: Some(executed_arguments.clone()),
5014                        });
5015                    }
5016                    ApprovalResult::Rejected { reason } => {
5017                        let reason = reason.unwrap_or_else(|| "rejected".to_string());
5018                        approval_record = Some(ToolApprovalRecord {
5019                            status: ToolApprovalStatus::Rejected,
5020                            reason: Some(reason.clone()),
5021                            modified_arguments: None,
5022                        });
5023                        let record = self.record_from_parts(
5024                            &request,
5025                            canonical_id,
5026                            executed_arguments,
5027                            started_at,
5028                            start,
5029                            false,
5030                            false,
5031                            format!("Approval rejected: {}", reason),
5032                            metadata,
5033                            ToolPolicyDecisionRecord::approval(reason),
5034                            approval_record,
5035                            false,
5036                            false,
5037                        );
5038                        self.finish_tool_record(&record).await;
5039                        return Ok(record);
5040                    }
5041                    ApprovalResult::Timeout => {
5042                        approval_record = Some(ToolApprovalRecord {
5043                            status: ToolApprovalStatus::Timeout,
5044                            reason: Some("approval timeout".to_string()),
5045                            modified_arguments: None,
5046                        });
5047                        let record = self.record_from_parts(
5048                            &request,
5049                            canonical_id,
5050                            executed_arguments,
5051                            started_at,
5052                            start,
5053                            false,
5054                            false,
5055                            "Approval timed out".to_string(),
5056                            metadata,
5057                            ToolPolicyDecisionRecord::approval("approval timeout"),
5058                            approval_record,
5059                            false,
5060                            false,
5061                        );
5062                        self.finish_tool_record(&record).await;
5063                        return Ok(record);
5064                    }
5065                }
5066            }
5067        }
5068
5069        if approval_record
5070            .as_ref()
5071            .is_some_and(|record| matches!(record.status, ToolApprovalStatus::NotRequired))
5072            && let Some(message) =
5073                security_engine.classification_approval_message(&canonical_id, &classification)
5074        {
5075            if self.hitl_engine.is_none() {
5076                approval_record = Some(ToolApprovalRecord {
5077                    status: ToolApprovalStatus::Unavailable,
5078                    reason: Some("No HITL engine configured".to_string()),
5079                    modified_arguments: None,
5080                });
5081                let record = self.record_from_parts(
5082                    &request,
5083                    canonical_id,
5084                    executed_arguments,
5085                    started_at,
5086                    start,
5087                    false,
5088                    false,
5089                    format!("Approval unavailable: {}", message),
5090                    metadata,
5091                    ToolPolicyDecisionRecord::approval(message),
5092                    approval_record,
5093                    false,
5094                    false,
5095                );
5096                self.finish_tool_record(&record).await;
5097                return Ok(record);
5098            }
5099            let check_result = HITLCheckResult::required(
5100                ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
5101                HashMap::new(),
5102                message.clone(),
5103                None,
5104            );
5105            match self.request_hitl_approval(check_result).await? {
5106                ApprovalResult::Approved => {
5107                    merge_approved_record(&mut approval_record);
5108                }
5109                ApprovalResult::Modified { changes } => {
5110                    if let Some(obj) = executed_arguments.as_object_mut() {
5111                        for (key, value) in changes {
5112                            obj.insert(key, value);
5113                        }
5114                    }
5115                    let modified_security = security_engine
5116                        .validate_tool_execution_with_bindings(
5117                            &canonical_id,
5118                            &executed_arguments,
5119                            &bindings,
5120                        )
5121                        .await?;
5122                    if !matches!(
5123                        modified_security,
5124                        SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5125                    ) {
5126                        let reason = modified_security
5127                            .reason()
5128                            .unwrap_or("modified arguments failed policy")
5129                            .to_string();
5130                        let record = self.record_from_parts(
5131                            &request,
5132                            canonical_id,
5133                            executed_arguments.clone(),
5134                            started_at,
5135                            start,
5136                            false,
5137                            false,
5138                            reason.clone(),
5139                            metadata,
5140                            ToolPolicyDecisionRecord::deny(reason),
5141                            Some(ToolApprovalRecord {
5142                                status: ToolApprovalStatus::Modified,
5143                                reason: None,
5144                                modified_arguments: Some(executed_arguments),
5145                            }),
5146                            false,
5147                            false,
5148                        );
5149                        self.finish_tool_record(&record).await;
5150                        return Ok(record);
5151                    }
5152                    approval_record = Some(ToolApprovalRecord {
5153                        status: ToolApprovalStatus::Modified,
5154                        reason: None,
5155                        modified_arguments: Some(executed_arguments.clone()),
5156                    });
5157                }
5158                ApprovalResult::Rejected { reason } => {
5159                    let reason = reason.unwrap_or_else(|| "rejected".to_string());
5160                    let record = self.record_from_parts(
5161                        &request,
5162                        canonical_id,
5163                        executed_arguments,
5164                        started_at,
5165                        start,
5166                        false,
5167                        false,
5168                        format!("Approval rejected: {}", reason),
5169                        metadata,
5170                        ToolPolicyDecisionRecord::approval(reason.clone()),
5171                        Some(ToolApprovalRecord {
5172                            status: ToolApprovalStatus::Rejected,
5173                            reason: Some(reason),
5174                            modified_arguments: None,
5175                        }),
5176                        false,
5177                        false,
5178                    );
5179                    self.finish_tool_record(&record).await;
5180                    return Ok(record);
5181                }
5182                ApprovalResult::Timeout => {
5183                    let record = self.record_from_parts(
5184                        &request,
5185                        canonical_id,
5186                        executed_arguments,
5187                        started_at,
5188                        start,
5189                        false,
5190                        false,
5191                        "Approval timed out".to_string(),
5192                        metadata,
5193                        ToolPolicyDecisionRecord::approval("approval timeout"),
5194                        Some(ToolApprovalRecord {
5195                            status: ToolApprovalStatus::Timeout,
5196                            reason: Some("approval timeout".to_string()),
5197                            modified_arguments: None,
5198                        }),
5199                        false,
5200                        false,
5201                    );
5202                    self.finish_tool_record(&record).await;
5203                    return Ok(record);
5204                }
5205            }
5206        }
5207
5208        let hitl_lang_ctx = self.build_hitl_language_context();
5209        if let Some(ref hitl_engine) = self.hitl_engine {
5210            let check_result = self
5211                .observe_purpose(
5212                    ObservationPurpose::HitlLocalization,
5213                    hitl_engine.check_tool_with_localization(
5214                        &canonical_id,
5215                        &executed_arguments,
5216                        &hitl_lang_ctx,
5217                        self.approval_handler.as_ref(),
5218                        Some(&self.llm_registry),
5219                    ),
5220                )
5221                .await?;
5222            if check_result.is_required() {
5223                match self.request_hitl_approval(check_result).await? {
5224                    ApprovalResult::Approved => {
5225                        merge_approved_record(&mut approval_record);
5226                    }
5227                    ApprovalResult::Modified { changes } => {
5228                        if let Some(obj) = executed_arguments.as_object_mut() {
5229                            for (key, value) in changes {
5230                                obj.insert(key, value);
5231                            }
5232                        }
5233                        let modified_security = security_engine
5234                            .validate_tool_execution_with_bindings(
5235                                &canonical_id,
5236                                &executed_arguments,
5237                                &bindings,
5238                            )
5239                            .await?;
5240                        if !matches!(
5241                            modified_security,
5242                            SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5243                        ) {
5244                            let reason = modified_security
5245                                .reason()
5246                                .unwrap_or("modified arguments failed policy")
5247                                .to_string();
5248                            let record = self.record_from_parts(
5249                                &request,
5250                                canonical_id,
5251                                executed_arguments.clone(),
5252                                started_at,
5253                                start,
5254                                false,
5255                                false,
5256                                reason.clone(),
5257                                metadata,
5258                                ToolPolicyDecisionRecord::deny(reason),
5259                                Some(ToolApprovalRecord {
5260                                    status: ToolApprovalStatus::Modified,
5261                                    reason: None,
5262                                    modified_arguments: Some(executed_arguments),
5263                                }),
5264                                false,
5265                                false,
5266                            );
5267                            self.finish_tool_record(&record).await;
5268                            return Ok(record);
5269                        }
5270                        approval_record = Some(ToolApprovalRecord {
5271                            status: ToolApprovalStatus::Modified,
5272                            reason: None,
5273                            modified_arguments: Some(executed_arguments.clone()),
5274                        });
5275                    }
5276                    ApprovalResult::Rejected { reason } => {
5277                        let reason = reason.unwrap_or_else(|| "rejected".to_string());
5278                        let record = self.record_from_parts(
5279                            &request,
5280                            canonical_id,
5281                            executed_arguments,
5282                            started_at,
5283                            start,
5284                            false,
5285                            false,
5286                            format!("Approval rejected: {}", reason),
5287                            metadata,
5288                            ToolPolicyDecisionRecord::approval(reason.clone()),
5289                            Some(ToolApprovalRecord {
5290                                status: ToolApprovalStatus::Rejected,
5291                                reason: Some(reason),
5292                                modified_arguments: None,
5293                            }),
5294                            false,
5295                            false,
5296                        );
5297                        self.finish_tool_record(&record).await;
5298                        return Ok(record);
5299                    }
5300                    ApprovalResult::Timeout => {
5301                        let record = self.record_from_parts(
5302                            &request,
5303                            canonical_id,
5304                            executed_arguments,
5305                            started_at,
5306                            start,
5307                            false,
5308                            false,
5309                            "Approval timed out".to_string(),
5310                            metadata,
5311                            ToolPolicyDecisionRecord::approval("approval timeout"),
5312                            Some(ToolApprovalRecord {
5313                                status: ToolApprovalStatus::Timeout,
5314                                reason: Some("approval timeout".to_string()),
5315                                modified_arguments: None,
5316                            }),
5317                            false,
5318                            false,
5319                        );
5320                        self.finish_tool_record(&record).await;
5321                        return Ok(record);
5322                    }
5323                }
5324            }
5325
5326            let condition_check = self
5327                .observe_purpose(
5328                    ObservationPurpose::HitlLocalization,
5329                    hitl_engine.check_conditions_with_localization(
5330                        &executed_arguments,
5331                        &hitl_lang_ctx,
5332                        self.approval_handler.as_ref(),
5333                        Some(&self.llm_registry),
5334                    ),
5335                )
5336                .await?;
5337            if condition_check.is_required() {
5338                match self.request_hitl_approval(condition_check).await? {
5339                    ApprovalResult::Approved => {
5340                        merge_approved_record(&mut approval_record);
5341                    }
5342                    ApprovalResult::Modified { changes } => {
5343                        if let Some(obj) = executed_arguments.as_object_mut() {
5344                            for (key, value) in changes {
5345                                obj.insert(key, value);
5346                            }
5347                        }
5348                        let modified_security = security_engine
5349                            .validate_tool_execution_with_bindings(
5350                                &canonical_id,
5351                                &executed_arguments,
5352                                &bindings,
5353                            )
5354                            .await?;
5355                        if !matches!(
5356                            modified_security,
5357                            SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5358                        ) {
5359                            let reason = modified_security
5360                                .reason()
5361                                .unwrap_or("modified arguments failed policy")
5362                                .to_string();
5363                            let record = self.record_from_parts(
5364                                &request,
5365                                canonical_id,
5366                                executed_arguments,
5367                                started_at,
5368                                start,
5369                                false,
5370                                false,
5371                                reason.clone(),
5372                                metadata,
5373                                ToolPolicyDecisionRecord::deny(reason),
5374                                approval_record,
5375                                false,
5376                                false,
5377                            );
5378                            self.finish_tool_record(&record).await;
5379                            return Ok(record);
5380                        }
5381                        approval_record = Some(ToolApprovalRecord {
5382                            status: ToolApprovalStatus::Modified,
5383                            reason: None,
5384                            modified_arguments: Some(executed_arguments.clone()),
5385                        });
5386                    }
5387                    ApprovalResult::Rejected { reason } => {
5388                        let reason = reason.unwrap_or_else(|| "rejected".to_string());
5389                        let record = self.record_from_parts(
5390                            &request,
5391                            canonical_id,
5392                            executed_arguments,
5393                            started_at,
5394                            start,
5395                            false,
5396                            false,
5397                            format!("Approval rejected: {}", reason),
5398                            metadata,
5399                            ToolPolicyDecisionRecord::approval(reason.clone()),
5400                            Some(ToolApprovalRecord {
5401                                status: ToolApprovalStatus::Rejected,
5402                                reason: Some(reason),
5403                                modified_arguments: None,
5404                            }),
5405                            false,
5406                            false,
5407                        );
5408                        self.finish_tool_record(&record).await;
5409                        return Ok(record);
5410                    }
5411                    ApprovalResult::Timeout => {
5412                        let record = self.record_from_parts(
5413                            &request,
5414                            canonical_id,
5415                            executed_arguments,
5416                            started_at,
5417                            start,
5418                            false,
5419                            false,
5420                            "Approval timed out".to_string(),
5421                            metadata,
5422                            ToolPolicyDecisionRecord::approval("approval timeout"),
5423                            Some(ToolApprovalRecord {
5424                                status: ToolApprovalStatus::Timeout,
5425                                reason: Some("approval timeout".to_string()),
5426                                modified_arguments: None,
5427                            }),
5428                            false,
5429                            false,
5430                        );
5431                        self.finish_tool_record(&record).await;
5432                        return Ok(record);
5433                    }
5434                }
5435            }
5436        }
5437
5438        //
5439        // Freeze the action reviewed by HITL after every approved argument change and policy cap.
5440        // A later plain approval preserves any earlier Modified evidence.
5441        //
5442        executed_arguments = security_engine.prepare_tool_arguments_with_bindings(
5443            &canonical_id,
5444            &executed_arguments,
5445            &bindings,
5446        );
5447        if let Some(record) = approval_record.as_mut()
5448            && matches!(record.status, ToolApprovalStatus::Modified)
5449        {
5450            record.modified_arguments = Some(executed_arguments.clone());
5451        }
5452        let binding_security_result = security_engine
5453            .validate_tool_execution_with_bindings(&canonical_id, &executed_arguments, &bindings)
5454            .await?;
5455        let approval_confirmation_required = matches!(
5456            binding_security_result,
5457            SecurityCheckResult::RequireConfirmation { .. }
5458        ) || security_engine
5459            .classification_approval_message(
5460                &canonical_id,
5461                &resolved.tool.classify_call(&executed_arguments),
5462            )
5463            .is_some();
5464        let approval_binding = approval_record.as_ref().and_then(|record| {
5465            matches!(
5466                record.status,
5467                ToolApprovalStatus::Approved | ToolApprovalStatus::Modified
5468            )
5469            .then(|| ToolApprovalBinding {
5470                canonical_id: canonical_id.clone(),
5471                arguments: executed_arguments.clone(),
5472                confirmation_required: approval_confirmation_required,
5473                policy_version: security_engine.policy_version(),
5474                runtime_control_version: approval_control_snapshot.version,
5475                state_generation: initial_scope_snapshot.state_generation,
5476                reviewed_tool: Arc::clone(&resolved.tool),
5477            })
5478        });
5479
5480        //
5481        // Re-resolve once after HITL because registry, scope, and policy may change while approval waits.
5482        // The final Arc must still match the implementation that the approver reviewed.
5483        //
5484        let control_snapshot = self.runtime_safety_snapshot();
5485        let resolved = self.tools.resolve(&request.requested_name);
5486        let registry_version = self.tools.version();
5487        let mut versions = ToolDecisionVersions {
5488            policy: control_snapshot.tool_security.policy_version(),
5489            registry: registry_version,
5490            runtime_control: control_snapshot.version,
5491            state: None,
5492        };
5493        metadata.insert(
5494            "runtime_scope_snapshot".to_string(),
5495            serde_json::to_value(&control_snapshot.tool_scope_override).unwrap_or(Value::Null),
5496        );
5497        let resolved = match resolved {
5498            Some(resolved) => resolved,
5499            None => {
5500                let reason = format!(
5501                    "Tool '{}' became unavailable after approval",
5502                    request.requested_name
5503                );
5504                let record = self.record_from_parts_at(
5505                    &request,
5506                    request.requested_name.clone(),
5507                    executed_arguments,
5508                    started_at,
5509                    start,
5510                    false,
5511                    false,
5512                    reason.clone(),
5513                    metadata,
5514                    ToolPolicyDecisionRecord::unavailable(reason),
5515                    approval_record,
5516                    false,
5517                    false,
5518                    versions,
5519                );
5520                self.finish_tool_record(&record).await;
5521                return Ok(record);
5522            }
5523        };
5524
5525        let canonical_id = resolved.identity.canonical_id.clone();
5526        let bindings = resolved.tool.policy_bindings();
5527        let final_arguments = control_snapshot
5528            .tool_security
5529            .prepare_tool_arguments_with_bindings(&canonical_id, &executed_arguments, &bindings);
5530        if let Some(record) = approval_record.as_mut()
5531            && matches!(record.status, ToolApprovalStatus::Modified)
5532        {
5533            record.modified_arguments = Some(final_arguments.clone());
5534        }
5535        let classification = resolved.tool.classify_call(&final_arguments);
5536        let safety = resolved.tool.safety_metadata();
5537        let security_engine = control_snapshot.tool_security;
5538        let limits = security_engine.effective_limits(&canonical_id, &safety, &classification);
5539        let policy_snapshot = security_engine.policy_snapshot(&canonical_id);
5540        let resource_lock_keys =
5541            tool_resource_lock_keys(&canonical_id, &final_arguments, &bindings, &classification);
5542        metadata.insert(
5543            "classification".to_string(),
5544            serde_json::to_value(&classification).unwrap_or(Value::Null),
5545        );
5546        metadata.insert(
5547            "effective_limits".to_string(),
5548            serde_json::to_value(&limits).unwrap_or(Value::Null),
5549        );
5550        metadata.insert(
5551            "resource_lock_keys".to_string(),
5552            serde_json::to_value(&resource_lock_keys).unwrap_or(Value::Null),
5553        );
5554        if policy_snapshot.is_null() {
5555            metadata.remove("policy_snapshot");
5556        } else {
5557            metadata.insert("policy_snapshot".to_string(), policy_snapshot.clone());
5558        }
5559
5560        let final_denial = |canonical_id: String,
5561                            output: String,
5562                            policy: ToolPolicyDecisionRecord,
5563                            metadata: HashMap<String, Value>,
5564                            decision_versions: ToolDecisionVersions| {
5565            self.record_from_parts_at(
5566                &request,
5567                canonical_id,
5568                final_arguments.clone(),
5569                started_at,
5570                start,
5571                false,
5572                false,
5573                output,
5574                metadata,
5575                policy,
5576                approval_record.clone(),
5577                false,
5578                false,
5579                decision_versions,
5580            )
5581        };
5582
5583        if control_snapshot.emergency_deny {
5584            let reason = "Tool execution is disabled by runtime control".to_string();
5585            let record = final_denial(
5586                canonical_id,
5587                reason.clone(),
5588                ToolPolicyDecisionRecord::deny(reason),
5589                metadata,
5590                versions,
5591            );
5592            self.finish_tool_record(&record).await;
5593            return Ok(record);
5594        }
5595
5596        //
5597        // Final scope evaluation uses the same post-HITL runtime snapshot and captures state authority before locks.
5598        // Admission must reject any later state transition, reset, or restore before consuming rate capacity or invoking the tool.
5599        //
5600        let available_snapshot = self
5601            .get_available_tool_ids_snapshot_for_scope(
5602                control_snapshot.tool_scope_override.as_deref(),
5603            )
5604            .await?;
5605        versions.state = available_snapshot.state_generation;
5606        metadata.insert(
5607            "available_tool_ids_snapshot".to_string(),
5608            serde_json::to_value(&available_snapshot.tool_ids).unwrap_or(Value::Null),
5609        );
5610        metadata.insert(
5611            "state_generation_snapshot".to_string(),
5612            serde_json::to_value(available_snapshot.state_generation).unwrap_or(Value::Null),
5613        );
5614        if !available_snapshot
5615            .tool_ids
5616            .iter()
5617            .any(|tool_id| tool_id == &canonical_id)
5618        {
5619            let reason = format!(
5620                "Tool '{}' is not available in the final runtime scope",
5621                canonical_id
5622            );
5623            let record = final_denial(
5624                canonical_id,
5625                reason.clone(),
5626                ToolPolicyDecisionRecord::deny(reason),
5627                metadata,
5628                versions,
5629            );
5630            self.finish_tool_record(&record).await;
5631            return Ok(record);
5632        }
5633
5634        //
5635        // Validation may run more than once, but it must not consume rate capacity.
5636        // Rate capacity is consumed once after resource locks are held.
5637        //
5638        let final_security_result = security_engine
5639            .validate_tool_execution_with_bindings(&canonical_id, &final_arguments, &bindings)
5640            .await?;
5641        match &final_security_result {
5642            SecurityCheckResult::Block { reason } => {
5643                let record = final_denial(
5644                    canonical_id,
5645                    format!("Denied: {}", reason),
5646                    ToolPolicyDecisionRecord::deny(reason.clone()),
5647                    metadata,
5648                    versions,
5649                );
5650                self.finish_tool_record(&record).await;
5651                return Ok(record);
5652            }
5653            SecurityCheckResult::Unavailable { reason } => {
5654                let record = final_denial(
5655                    canonical_id,
5656                    format!("Unavailable: {}", reason),
5657                    ToolPolicyDecisionRecord::unavailable(reason.clone()),
5658                    metadata,
5659                    versions,
5660                );
5661                self.finish_tool_record(&record).await;
5662                return Ok(record);
5663            }
5664            SecurityCheckResult::Warn { message } => {
5665                warn!(tool = %canonical_id, message = %message, "Tool security warning after approval");
5666            }
5667            SecurityCheckResult::Allow | SecurityCheckResult::RequireConfirmation { .. } => {}
5668        }
5669        let final_confirmation_required = matches!(
5670            final_security_result,
5671            SecurityCheckResult::RequireConfirmation { .. }
5672        ) || security_engine
5673            .classification_approval_message(&canonical_id, &classification)
5674            .is_some();
5675        let stale_approval = approval_binding.as_ref().is_some_and(|binding| {
5676            binding.is_stale(
5677                &canonical_id,
5678                &final_arguments,
5679                final_confirmation_required,
5680                versions,
5681                &resolved.tool,
5682            )
5683        });
5684        if stale_approval {
5685            let reason = "Approval became stale before final admission".to_string();
5686            let record = final_denial(
5687                canonical_id,
5688                reason.clone(),
5689                ToolPolicyDecisionRecord::deny(reason),
5690                metadata,
5691                versions,
5692            );
5693            self.finish_tool_record(&record).await;
5694            return Ok(record);
5695        }
5696        if final_confirmation_required && approval_binding.is_none() {
5697            let reason = "Final policy requires fresh approval".to_string();
5698            let record = final_denial(
5699                canonical_id,
5700                reason.clone(),
5701                ToolPolicyDecisionRecord::approval(reason),
5702                metadata,
5703                versions,
5704            );
5705            self.finish_tool_record(&record).await;
5706            return Ok(record);
5707        }
5708
5709        if let Some((_, reason)) = self.host_tool_unavailability(&canonical_id) {
5710            let record = final_denial(
5711                canonical_id,
5712                reason.to_string(),
5713                ToolPolicyDecisionRecord::unavailable(reason),
5714                metadata,
5715                versions,
5716            );
5717            self.finish_tool_record(&record).await;
5718            return Ok(record);
5719        }
5720
5721        //
5722        // Hold conflict locks across final generation admission and invocation.
5723        // Completion hooks and fallback execution run only after these guards are released.
5724        //
5725        let Some(resource_guards) = self.acquire_tool_resource_locks(&resource_lock_keys).await
5726        else {
5727            let reason = "Tool execution cancelled while waiting for resource locks".to_string();
5728            let record = final_denial(
5729                canonical_id,
5730                reason.clone(),
5731                ToolPolicyDecisionRecord::deny(reason),
5732                metadata,
5733                versions,
5734            );
5735            self.finish_tool_record(&record).await;
5736            return Ok(record);
5737        };
5738
5739        //
5740        // The lock-held admission is the last authority boundary. Runtime, policy, and state generations must still match before the one atomic rate admission.
5741        // Updates after admission apply to later calls, while emergency cancellation remains live for this call.
5742        //
5743        let admission = self.admit_tool_execution(
5744            versions.runtime_control,
5745            versions.policy,
5746            versions.state,
5747            &canonical_id,
5748        );
5749        if !matches!(admission, SecurityCheckResult::Allow) {
5750            let latest_control = self.runtime_safety_snapshot();
5751            let reason = admission
5752                .reason()
5753                .unwrap_or("tool admission was denied")
5754                .to_string();
5755            let policy = if admission.is_unavailable() {
5756                ToolPolicyDecisionRecord::unavailable(reason.clone())
5757            } else {
5758                ToolPolicyDecisionRecord::deny(reason.clone())
5759            };
5760            let record = self.record_from_parts_at(
5761                &request,
5762                canonical_id,
5763                final_arguments,
5764                started_at,
5765                start,
5766                false,
5767                false,
5768                reason,
5769                metadata,
5770                policy,
5771                approval_record,
5772                false,
5773                false,
5774                ToolDecisionVersions {
5775                    policy: latest_control.tool_security.policy_version(),
5776                    registry: versions.registry,
5777                    runtime_control: latest_control.version,
5778                    state: self
5779                        .state_machine
5780                        .as_ref()
5781                        .map(|state_machine| state_machine.generation()),
5782                },
5783            );
5784            self.finish_tool_record_after_resource_guards(resource_guards, &record)
5785                .await;
5786            return Ok(record);
5787        }
5788        let executed_arguments = final_arguments;
5789
5790        let tool_config = self.recovery_manager.get_tool_config(&canonical_id);
5791        let timeout_ms = limits
5792            .timeout_ms
5793            .unwrap_or_else(|| security_engine.get_tool_timeout(&canonical_id));
5794        let turn_actor = current_turn_actor_context();
5795        let actor = ToolActorContext {
5796            actor_id: turn_actor
5797                .as_ref()
5798                .and_then(|context| context.effective_actor_id().map(str::to_string))
5799                .or_else(|| self.actor_id()),
5800            origin_actor_id: turn_actor
5801                .as_ref()
5802                .and_then(|context| context.origin_actor_id.clone()),
5803            sender_agent_id: turn_actor
5804                .as_ref()
5805                .and_then(|context| context.sender_agent_id.clone()),
5806        };
5807        let tool_context = ToolExecutionContext {
5808            requested_name: request.requested_name.clone(),
5809            canonical_id: canonical_id.clone(),
5810            display_name: resolved.identity.display_name.clone(),
5811            provider_id: resolved.identity.provider_id.clone(),
5812            registry_version: versions.registry,
5813            policy_version: versions.policy,
5814            runtime_control_version: versions.runtime_control,
5815            call_id: request.call_id.clone(),
5816            source: request.source.clone(),
5817            actor,
5818            cancellation: ToolCancellationToken::new(
5819                Arc::clone(&self.runtime_control.emergency_deny),
5820                Some("runtime control cancellation".to_string()),
5821            ),
5822            started_at,
5823            deadline: None,
5824            permission: ToolPolicyDecisionRecord::allow(),
5825            approval: approval_record.clone(),
5826            classification: classification.clone(),
5827            safety,
5828            limits: limits.clone(),
5829            policy_snapshot,
5830            custom_config: security_engine.custom_config(&canonical_id),
5831        };
5832        let (mut result, timed_out, cancelled, invoked) = self
5833            .run_tool_with_retries(
5834                &canonical_id,
5835                resolved.tool.clone(),
5836                executed_arguments.clone(),
5837                tool_context,
5838                timeout_ms,
5839                tool_config.max_retries,
5840            )
5841            .await?;
5842
5843        let fallback_tool = if !result.success {
5844            match &tool_config.on_failure {
5845                ToolFailureAction::Skip => {
5846                    result = ToolResult::ok(format!(
5847                        "{{\"skipped\": true, \"reason\": \"Tool '{}' was skipped after failure\"}}",
5848                        canonical_id
5849                    ));
5850                    None
5851                }
5852                ToolFailureAction::Fallback { fallback_tool } => Some(fallback_tool.clone()),
5853                ToolFailureAction::ReportError => None,
5854            }
5855        } else {
5856            None
5857        };
5858
5859        let output_cap = limits.max_output_chars;
5860        let (output, output_truncated) =
5861            Self::truncate_tool_output(result.output.clone(), output_cap);
5862        if let Some(result_metadata) = result.metadata {
5863            metadata.extend(result_metadata);
5864        }
5865        let mut record = self.record_from_parts_at(
5866            &request,
5867            canonical_id,
5868            executed_arguments,
5869            started_at,
5870            start,
5871            invoked,
5872            result.success,
5873            output,
5874            metadata,
5875            ToolPolicyDecisionRecord::allow(),
5876            approval_record,
5877            timed_out,
5878            output_truncated,
5879            versions,
5880        );
5881        record.cancelled = cancelled;
5882        if cancelled {
5883            record.cancellation_reason = Some("runtime control cancellation".to_string());
5884        }
5885        if let Some(fallback_tool) = fallback_tool {
5886            let fallback_arguments = record.executed_arguments.clone();
5887            let original_tool = record.canonical_id.clone();
5888            //
5889            // Finish the failed logical request after releasing its resource guards so its start hook is matched and fallback hooks cannot overtake the original record.
5890            //
5891            self.finish_tool_record_after_resource_guards(resource_guards, &record)
5892                .await;
5893            let fallback_request = ToolExecutionRequest::new(
5894                request.call_id.clone(),
5895                fallback_tool,
5896                fallback_arguments,
5897                ToolCallSource::Fallback { original_tool },
5898            );
5899            return Box::pin(self.execute_tool_record(fallback_request)).await;
5900        }
5901        self.finish_tool_record_after_resource_guards(resource_guards, &record)
5902            .await;
5903        Ok(record)
5904    }
5905
5906    #[instrument(skip(self, tool_call), fields(tool = %tool_call.name))]
5907    async fn execute_tool_smart(&self, tool_call: &ToolCall) -> Result<String> {
5908        let record = self
5909            .execute_tool_record(ToolExecutionRequest::new(
5910                tool_call.id.clone(),
5911                tool_call.name.clone(),
5912                tool_call.arguments.clone(),
5913                ToolCallSource::Model,
5914            ))
5915            .await?;
5916        if record.success {
5917            Ok(record.model_output_string())
5918        } else if matches!(record.policy.outcome, PermissionOutcome::RequiresApproval) {
5919            Err(AgentError::HITLRejected(record.model_output_string()))
5920        } else {
5921            Err(AgentError::Tool(record.model_output_string()))
5922        }
5923    }
5924
5925    //
5926    // This method is branch-safe because it only asks the router which skill matches.
5927    // Do not add disambiguation, pending-skill writes, or skill execution here.
5928    //
5929    /// Selects a skill without executing it.
5930    async fn select_skill_candidate(&self, input: &str) -> Result<Option<SkillCandidate>> {
5931        let Some(ref router) = self.skill_router else {
5932            return Ok(None);
5933        };
5934        let available_skills = self.get_available_skills();
5935        if available_skills.is_empty() {
5936            return Ok(None);
5937        }
5938        let skill_ids: Vec<&str> = available_skills.iter().map(|s| s.id.as_str()).collect();
5939        let Some(skill_id) = self
5940            .observe_purpose(
5941                ObservationPurpose::SkillRouting,
5942                router.select_skill_filtered(input, &skill_ids),
5943            )
5944            .await?
5945        else {
5946            return Ok(None);
5947        };
5948        let skill = router
5949            .get_skill(&skill_id)
5950            .cloned()
5951            .ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
5952        info!(skill_id = %skill_id, "Skill selected");
5953        Ok(Some(SkillCandidate::new(skill_id, skill)))
5954    }
5955
5956    //
5957    // This is the commit half of skill routing.
5958    // It may mutate pending skill state, run clarification, and execute skill steps.
5959    //
5960    async fn commit_skill_candidate_route_result(
5961        &self,
5962        candidate: SkillCandidate,
5963        input: &str,
5964    ) -> Result<SkillRouteResult> {
5965        let skill_id = candidate.skill_id;
5966        let skill = candidate.skill;
5967        let expected_state_generation = self
5968            .state_machine
5969            .as_ref()
5970            .map(|state_machine| state_machine.generation());
5971        let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
5972        if let Some(ref skill_disambig) = skill.disambiguation
5973            && skill_disambig.enabled.unwrap_or(false)
5974            && let Some(ref disambiguator) = self.disambiguation_manager
5975        {
5976            let context = self.build_disambiguation_context().await?;
5977            let state_override = self
5978                .state_machine
5979                .as_ref()
5980                .and_then(|sm| sm.current_definition())
5981                .and_then(|def| def.disambiguation.clone());
5982
5983            let disambiguation_result = self
5984                .observe_purpose(
5985                    ObservationPurpose::DisambiguationDetection,
5986                    disambiguator.process_input_with_override(
5987                        input,
5988                        &context,
5989                        state_override.as_ref(),
5990                        Some(skill_disambig),
5991                    ),
5992                )
5993                .await?;
5994            let current_state_generation = self
5995                .state_machine
5996                .as_ref()
5997                .map(|state_machine| state_machine.generation());
5998            if current_state_generation != expected_state_generation
5999                || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6000            {
6001                disambiguator.clear_pending().await;
6002                *self.pending_skill_id.write() = None;
6003                return Err(AgentError::Other(
6004                    "State or reset ownership changed during skill disambiguation".to_string(),
6005                ));
6006            }
6007            match disambiguation_result {
6008                DisambiguationResult::Clear => {
6009                    debug!(skill_id = %skill_id, "Skill disambiguation: clear");
6010                }
6011                DisambiguationResult::NeedsClarification {
6012                    question,
6013                    detection,
6014                } => {
6015                    let admission = self
6016                        .admit_disambiguation_redispatch(
6017                            expected_disambiguation_epoch,
6018                            expected_state_generation,
6019                        )
6020                        .await?;
6021                    let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
6022                    info!(
6023                        skill_id = %skill_id,
6024                        ambiguity_type = ?detection.ambiguity_type,
6025                        confidence = detection.confidence,
6026                        "Skill requires clarification before execution"
6027                    );
6028                    *self.pending_skill_id.write() = Some(skill_id.clone());
6029                    let response = AgentResponse::new(&question.question).with_metadata(
6030                        "disambiguation",
6031                        serde_json::json!({
6032                            "status": if awaiting_confirmation { "awaiting_confirmation" } else { "awaiting_clarification" },
6033                            "skill_id": skill_id,
6034                            "options": question.options,
6035                            "clarifying": question.clarifying,
6036                            "detection": {
6037                                "type": detection.ambiguity_type,
6038                                "confidence": detection.confidence,
6039                                "what_is_unclear": detection.what_is_unclear,
6040                            }
6041                        }),
6042                    );
6043                    drop(admission);
6044                    return Ok(SkillRouteResult::NeedsClarification {
6045                        response,
6046                        ownership: Some(DisambiguationOwnership {
6047                            epoch: expected_disambiguation_epoch,
6048                            state_generation: expected_state_generation,
6049                        }),
6050                    });
6051                }
6052                DisambiguationResult::Clarified { enriched_input, .. } => {
6053                    info!(skill_id = %skill_id, enriched = %enriched_input, "Skill disambiguation clarified");
6054                    let admission = self
6055                        .admit_disambiguation_redispatch(
6056                            expected_disambiguation_epoch,
6057                            expected_state_generation,
6058                        )
6059                        .await?;
6060                    drop(admission);
6061                    let content = self.execute_skill(&skill, &enriched_input).await?;
6062                    return Ok(SkillRouteResult::Response { skill_id, content });
6063                }
6064                DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
6065                    info!(skill_id = %skill_id, "Skill disambiguation best guess");
6066                    let admission = self
6067                        .admit_disambiguation_redispatch(
6068                            expected_disambiguation_epoch,
6069                            expected_state_generation,
6070                        )
6071                        .await?;
6072                    drop(admission);
6073                    let content = self.execute_skill(&skill, &enriched_input).await?;
6074                    return Ok(SkillRouteResult::Response { skill_id, content });
6075                }
6076                DisambiguationResult::GiveUp { reason } => {
6077                    warn!(skill_id = %skill_id, reason = %reason, "Skill disambiguation gave up");
6078                    let apology = self
6079                                .generate_localized_apology(
6080                                    "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
6081                                    &reason,
6082                                )
6083                                .await
6084                                .unwrap_or_else(|_| {
6085                                    format!("I'm sorry, I couldn't understand your request: {}", reason)
6086                                });
6087                    return Ok(SkillRouteResult::NeedsClarification {
6088                        response: AgentResponse::new(&apology),
6089                        ownership: None,
6090                    });
6091                }
6092                DisambiguationResult::Escalate { reason } => {
6093                    info!(skill_id = %skill_id, reason = %reason, "Skill disambiguation escalating");
6094                    let apology = self
6095                                .generate_localized_apology(
6096                                    "Explain briefly that you're transferring the user to a human agent for help.",
6097                                    &reason,
6098                                )
6099                                .await
6100                                .unwrap_or_else(|_| {
6101                                    format!("I need human assistance to help with your request: {}", reason)
6102                                });
6103                    return Ok(SkillRouteResult::NeedsClarification {
6104                        response: AgentResponse::new(&apology),
6105                        ownership: None,
6106                    });
6107                }
6108                DisambiguationResult::Abandoned { .. } => {
6109                    debug!(skill_id = %skill_id, "Skill disambiguation abandoned");
6110                    return Ok(SkillRouteResult::NoMatch);
6111                }
6112            }
6113        }
6114        let admission = self
6115            .admit_disambiguation_redispatch(
6116                expected_disambiguation_epoch,
6117                expected_state_generation,
6118            )
6119            .await?;
6120        drop(admission);
6121        let content = self.execute_skill(&skill, input).await?;
6122        Ok(SkillRouteResult::Response { skill_id, content })
6123    }
6124
6125    /// Result of skill routing.
6126    async fn try_skill_route(&self, input: &str) -> Result<SkillRouteResult> {
6127        if let Some(candidate) = self.select_skill_candidate(input).await? {
6128            self.commit_skill_candidate_route_result(candidate, input)
6129                .await
6130        } else {
6131            Ok(SkillRouteResult::NoMatch)
6132        }
6133    }
6134
6135    /// Execute a skill with reasoning and reflection, returning the response string.
6136    async fn execute_skill(&self, skill: &SkillDefinition, input: &str) -> Result<String> {
6137        if let Some(ref executor) = self.skill_executor {
6138            let skill_reasoning = self.get_skill_reasoning_config(skill);
6139            let skill_reflection = self.get_skill_reflection_config(skill);
6140
6141            debug!(
6142                skill_id = %skill.id,
6143                reasoning_mode = ?skill_reasoning.mode,
6144                reflection_enabled = ?skill_reflection.enabled,
6145                "Skill reasoning/reflection config"
6146            );
6147
6148            let response = self
6149                .observe_purpose(
6150                    ObservationPurpose::SkillPrompt,
6151                    executor.execute_with_invoker(skill, input, serde_json::json!({}), self),
6152                )
6153                .await?;
6154
6155            if skill_reflection.requires_evaluation() && skill_reflection.is_enabled() {
6156                let should_reflect = self
6157                    .should_reflect_with_config(input, &response, &skill_reflection)
6158                    .await?;
6159                if should_reflect {
6160                    let evaluated = self
6161                        .evaluate_and_retry_with_config(input, response, &skill_reflection)
6162                        .await?;
6163                    return Ok(evaluated);
6164                }
6165            }
6166
6167            return Ok(response);
6168        }
6169        Err(AgentError::Skill(
6170            "No skill executor configured".to_string(),
6171        ))
6172    }
6173
6174    /// Execute a skill by ID, bypassing the skill router.
6175    /// Used after skill-triggered disambiguation resolves to route directly to the matched skill.
6176    async fn execute_skill_by_id(&self, skill_id: &str, input: &str) -> Result<String> {
6177        let skill = self
6178            .skill_router
6179            .as_ref()
6180            .and_then(|r| r.get_skill(skill_id).cloned())
6181            .ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
6182        self.execute_skill(&skill, input).await
6183    }
6184
6185    async fn should_reflect_with_config(
6186        &self,
6187        input: &str,
6188        response: &str,
6189        config: &ReflectionConfig,
6190    ) -> Result<bool> {
6191        if !config.requires_evaluation() {
6192            return Ok(false);
6193        }
6194
6195        if config.is_enabled() {
6196            return Ok(true);
6197        }
6198
6199        let evaluator_llm = config
6200            .evaluator_llm
6201            .as_ref()
6202            .and_then(|alias| self.llm_registry.get(alias).ok())
6203            .or_else(|| self.llm_registry.router().ok())
6204            .or_else(|| self.llm_registry.default().ok());
6205
6206        let Some(llm) = evaluator_llm else {
6207            return Ok(false);
6208        };
6209
6210        let response_preview: String = response.chars().take(500).collect();
6211        let prompt = format!(
6212            r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
6213
6214User query: "{}"
6215Response: "{}"
6216
6217Answer YES or NO only."#,
6218            input, response_preview
6219        );
6220
6221        let messages = vec![ChatMessage::user(&prompt)];
6222        let result = self
6223            .observe_purpose(
6224                ObservationPurpose::ReflectionDecision,
6225                llm.complete(&messages, None),
6226            )
6227            .await;
6228
6229        match result {
6230            Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
6231            Err(_) => Ok(false),
6232        }
6233    }
6234
6235    async fn evaluate_and_retry_with_config(
6236        &self,
6237        input: &str,
6238        mut response: String,
6239        config: &ReflectionConfig,
6240    ) -> Result<String> {
6241        let llm = self.get_state_llm()?;
6242        let mut attempts = 0u32;
6243        let max_retries = config.max_retries;
6244
6245        loop {
6246            let evaluation = self
6247                .evaluate_response_with_config(input, &response, config)
6248                .await?;
6249
6250            if evaluation.passed || attempts >= max_retries {
6251                info!(
6252                    passed = evaluation.passed,
6253                    confidence = evaluation.confidence,
6254                    attempts = attempts + 1,
6255                    "Skill reflection evaluation complete"
6256                );
6257                return Ok(response);
6258            }
6259
6260            debug!(
6261                attempt = attempts + 1,
6262                failed_criteria = evaluation.failed_criteria().count(),
6263                "Skill response did not meet criteria, retrying"
6264            );
6265
6266            let feedback: Vec<String> = evaluation
6267                .failed_criteria()
6268                .map(|c| format!("- {}", c.criterion))
6269                .collect();
6270
6271            let retry_prompt = format!(
6272                "Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response to: {}",
6273                feedback.join("\n"),
6274                input
6275            );
6276
6277            let messages = vec![ChatMessage::user(&retry_prompt)];
6278            let retry_response = self
6279                .observe_purpose(
6280                    ObservationPurpose::ReflectionEvaluation,
6281                    llm.complete(&messages, None),
6282                )
6283                .await
6284                .map_err(|e| AgentError::LLM(e.to_string()))?;
6285
6286            response = retry_response.content.trim().to_string();
6287            attempts += 1;
6288        }
6289    }
6290
6291    async fn evaluate_response_with_config(
6292        &self,
6293        input: &str,
6294        response: &str,
6295        config: &ReflectionConfig,
6296    ) -> Result<EvaluationResult> {
6297        let evaluator_llm = config
6298            .evaluator_llm
6299            .as_ref()
6300            .and_then(|alias| self.llm_registry.get(alias).ok())
6301            .or_else(|| self.llm_registry.router().ok())
6302            .or_else(|| self.llm_registry.default().ok())
6303            .ok_or_else(|| AgentError::Config("No LLM available for evaluation".into()))?;
6304
6305        let criteria = &config.criteria;
6306        let criteria_list = criteria
6307            .iter()
6308            .enumerate()
6309            .map(|(i, c)| format!("{}. {}", i + 1, c))
6310            .collect::<Vec<_>>()
6311            .join("\n");
6312
6313        let prompt = format!(
6314            r#"Evaluate this response against the criteria.
6315
6316User query: "{}"
6317
6318Response to evaluate: "{}"
6319
6320Criteria:
6321{}
6322
6323For each criterion, respond with:
6324- criterion number
6325- PASS or FAIL
6326- brief reason
6327
6328Then provide overall confidence (0.0 to 1.0) and whether it passes overall.
6329
6330Format:
63311. PASS/FAIL - reason
63322. PASS/FAIL - reason
6333...
6334CONFIDENCE: 0.X
6335OVERALL: PASS/FAIL"#,
6336            input, response, criteria_list
6337        );
6338
6339        let messages = vec![ChatMessage::user(&prompt)];
6340        let eval_response = self
6341            .observe_purpose(
6342                ObservationPurpose::ReflectionEvaluation,
6343                evaluator_llm.complete(&messages, None),
6344            )
6345            .await
6346            .map_err(|e| AgentError::LLM(format!("Evaluation failed: {}", e)))?;
6347
6348        let content = eval_response.content.to_uppercase();
6349        let llm_pass = content.contains("OVERALL: PASS");
6350
6351        let confidence = content
6352            .lines()
6353            .find(|l| l.contains("CONFIDENCE:"))
6354            .and_then(|l| {
6355                l.split(':')
6356                    .nth(1)
6357                    .and_then(|v| v.trim().parse::<f32>().ok())
6358            })
6359            .unwrap_or(if llm_pass { 0.8 } else { 0.4 });
6360
6361        // Gate pass against confidence threshold.
6362        // LLM may say PASS but with low confidence - the threshold catches this.
6363        let overall_pass = llm_pass && confidence >= config.pass_threshold;
6364
6365        let mut criteria_results = Vec::new();
6366        for (i, criterion) in criteria.iter().enumerate() {
6367            let line_marker = format!("{}.", i + 1);
6368            let passed = eval_response
6369                .content
6370                .lines()
6371                .find(|l| l.contains(&line_marker))
6372                .map(|l| l.to_uppercase().contains("PASS"))
6373                .unwrap_or(overall_pass);
6374
6375            if passed {
6376                criteria_results.push(CriterionResult::pass(criterion));
6377            } else {
6378                criteria_results.push(CriterionResult::fail(criterion, "Did not meet criterion"));
6379            }
6380        }
6381
6382        Ok(EvaluationResult::new(overall_pass, confidence).with_criteria(criteria_results))
6383    }
6384
6385    /// Process input through the pipeline (state-level override or agent-level).
6386    async fn process_input(&self, input: &str) -> Result<ProcessData> {
6387        if let Some(processor) = self.get_state_process_processor() {
6388            let purpose = observation_purpose_for_process(processor.input_purpose_hint());
6389            return self
6390                .observe_purpose(purpose, processor.process_input(input))
6391                .await;
6392        }
6393        if let Some(ref processor) = self.process_processor {
6394            let purpose = observation_purpose_for_process(processor.input_purpose_hint());
6395            self.observe_purpose(purpose, processor.process_input(input))
6396                .await
6397        } else {
6398            Ok(ProcessData::new(input))
6399        }
6400    }
6401
6402    /// Process output through the pipeline (state-level override or agent-level).
6403    async fn process_output(
6404        &self,
6405        output: &str,
6406        input_context: &std::collections::HashMap<String, serde_json::Value>,
6407    ) -> Result<ProcessData> {
6408        if let Some(processor) = self.get_state_process_processor() {
6409            let purpose = observation_purpose_for_process(processor.output_purpose_hint());
6410            return self
6411                .observe_purpose(purpose, processor.process_output(output, input_context))
6412                .await;
6413        }
6414        if let Some(ref processor) = self.process_processor {
6415            let purpose = observation_purpose_for_process(processor.output_purpose_hint());
6416            self.observe_purpose(purpose, processor.process_output(output, input_context))
6417                .await
6418        } else {
6419            Ok(ProcessData::new(output))
6420        }
6421    }
6422
6423    /// Build a ProcessProcessor from the current state's process config, if any.
6424    fn get_state_process_processor(&self) -> Option<ProcessProcessor> {
6425        let sm = self.state_machine.as_ref()?;
6426        let def = sm.current_definition()?;
6427        let config = def.process.as_ref()?;
6428        let mut processor = ProcessProcessor::new(config.clone());
6429        if let Some(ref registry) = Some(self.llm_registry.clone()) {
6430            processor = processor.with_llm_registry(registry.clone());
6431        }
6432        processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
6433        Some(processor)
6434    }
6435
6436    // Timeout transitions reserve exit actions before committing ownership changes under the short write lock.
6437    async fn check_turn_timeout(&self) -> Result<()> {
6438        let Some(ref sm) = self.state_machine else {
6439            return Ok(());
6440        };
6441        let Some(timeout_state) = sm.check_timeout() else {
6442            return Ok(());
6443        };
6444        let claim_admission = self.disambiguation_admission.write().await;
6445        if sm.check_timeout().as_deref() != Some(timeout_state.as_str()) {
6446            return Ok(());
6447        }
6448        let Some(reservation) = self.reserve_state_transition() else {
6449            return Ok(());
6450        };
6451        let from_state = sm.current();
6452        let expected_state_generation = sm.generation();
6453        let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
6454        let history_before = sm.history();
6455        drop(claim_admission);
6456
6457        self.execute_state_exit_actions(&from_state).await;
6458
6459        let admission = self.disambiguation_admission.write().await;
6460        if sm.current() != from_state
6461            || sm.generation() != expected_state_generation
6462            || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6463            || sm.check_timeout().as_deref() != Some(timeout_state.as_str())
6464        {
6465            return Ok(());
6466        }
6467        sm.transition_to(&timeout_state, "max_turns exceeded")?;
6468        self.invalidate_pending_confirmation("state_timeout").await;
6469        let entered = sm.current();
6470        let is_reentry = Self::state_was_previously_entered(&entered, &from_state, &history_before);
6471        drop(admission);
6472
6473        self.execute_state_enter_actions(&entered, is_reentry).await;
6474        drop(reservation);
6475        info!(to = %entered, "Timeout transition");
6476        Ok(())
6477    }
6478
6479    fn increment_turn(&self) {
6480        if let Some(ref sm) = self.state_machine {
6481            sm.increment_turn();
6482        }
6483    }
6484
6485    fn transitions_available_for_commit(&self) -> Option<(Vec<Transition>, String)> {
6486        let sm = self.state_machine.as_ref()?;
6487        let current = sm.current();
6488        let transitions: Vec<_> = sm
6489            .auto_transitions()
6490            .into_iter()
6491            .filter(|t| match t.cooldown_turns {
6492                Some(cd) if cd > 0 => {
6493                    let resolved = sm.config().resolve_full_path(&current, &t.to);
6494                    !sm.is_on_cooldown(&resolved, cd)
6495                }
6496                _ => true,
6497            })
6498            .collect();
6499        Some((transitions, current))
6500    }
6501
6502    fn transition_reason(transition: &Transition) -> String {
6503        if transition.when.is_empty() {
6504            "guard condition met".to_string()
6505        } else {
6506            transition.when.clone()
6507        }
6508    }
6509
6510    /// Builds transition context with optional staged writes overlaid.
6511    fn build_transition_context(
6512        &self,
6513        user_message: &str,
6514        response: &str,
6515        current_state: &str,
6516        staged: Option<&HashMap<String, Value>>,
6517    ) -> TransitionContext {
6518        let context_map = staged
6519            .map(|writes| self.build_context_with_staged(writes))
6520            .unwrap_or_else(|| self.build_context_with_overlays());
6521        TransitionContext::new(user_message, response, current_state).with_context(context_map)
6522    }
6523
6524    /// Selects a post-response transition without committing side effects.
6525    async fn select_transition_candidate(
6526        &self,
6527        user_message: &str,
6528        response: &str,
6529    ) -> Result<Option<TransitionCandidate>> {
6530        let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6531            return Ok(None);
6532        };
6533        let transitions: Vec<Transition> = transitions
6534            .into_iter()
6535            .filter(|transition| matches!(transition.timing, TransitionTiming::PostResponse))
6536            .collect();
6537        if transitions.is_empty() {
6538            return Ok(None);
6539        }
6540        let Some(evaluator) = self.transition_evaluator.as_ref() else {
6541            return Ok(None);
6542        };
6543        let context = self.build_transition_context(user_message, response, &current_state, None);
6544        let selected = self
6545            .observe_purpose(
6546                ObservationPurpose::StateTransitionEvaluation,
6547                evaluator.select_transition(&transitions, &context),
6548            )
6549            .await?;
6550        Ok(selected.map(|index| {
6551            let transition = transitions[index].clone();
6552            TransitionCandidate::new(
6553                current_state,
6554                transition.clone(),
6555                Self::transition_reason(&transition),
6556            )
6557        }))
6558    }
6559
6560    /// Selects a guard or resolved-intent transition without an LLM call.
6561    fn select_deterministic_transition_candidate(
6562        &self,
6563        user_message: &str,
6564        current_state: &str,
6565        transitions: &[Transition],
6566        staged: &HashMap<String, Value>,
6567    ) -> Option<TransitionCandidate> {
6568        let context = self.build_transition_context(user_message, "", current_state, Some(staged));
6569
6570        for transition in transitions {
6571            if let Some(guard) = transition.guard.as_ref()
6572                && evaluate_guard(guard, &context)
6573            {
6574                return Some(TransitionCandidate::new(
6575                    current_state,
6576                    transition.clone(),
6577                    Self::transition_reason(transition),
6578                ));
6579            }
6580        }
6581
6582        let resolved_intent = context
6583            .context
6584            .get("resolved_intent")
6585            .and_then(Value::as_str)
6586            .filter(|value| !value.is_empty());
6587        if let Some(resolved_intent) = resolved_intent {
6588            for transition in transitions {
6589                if transition.intent.as_deref() == Some(resolved_intent) {
6590                    return Some(TransitionCandidate::new(
6591                        current_state,
6592                        transition.clone(),
6593                        Self::transition_reason(transition),
6594                    ));
6595                }
6596            }
6597        }
6598
6599        None
6600    }
6601
6602    /// Commits a selected transition through the shared post-response path.
6603    async fn commit_transition_candidate(&self, candidate: &TransitionCandidate) -> Result<bool> {
6604        self.commit_transition_target(&candidate.from_state, candidate.target(), &candidate.reason)
6605            .await
6606    }
6607
6608    /// Runs state transition approval before any transition side effects.
6609    async fn approve_transition_target(&self, from_state: &str, target: &str) -> Result<bool> {
6610        let approved = self.check_state_hitl(Some(from_state), target).await?;
6611        if !approved {
6612            info!(to = %target, "State transition rejected by HITL");
6613        }
6614        Ok(approved)
6615    }
6616
6617    /// Applies an approved transition after reserving exit actions and keeping async hooks outside the commit lock.
6618    async fn apply_transition_target(
6619        &self,
6620        from_state: &str,
6621        target: &str,
6622        reason: &str,
6623        staged: Option<&HashMap<String, Value>>,
6624    ) -> Result<bool> {
6625        let Some(ref sm) = self.state_machine else {
6626            return Ok(false);
6627        };
6628        let claim_admission = self.disambiguation_admission.write().await;
6629        if sm.current() != from_state {
6630            return Ok(false);
6631        }
6632        let Some(reservation) = self.reserve_state_transition() else {
6633            return Ok(false);
6634        };
6635        let expected_state_generation = sm.generation();
6636        let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
6637        let history_before = sm.history();
6638        drop(claim_admission);
6639
6640        self.execute_state_exit_actions(from_state).await;
6641
6642        let admission = self.disambiguation_admission.write().await;
6643        if sm.current() != from_state
6644            || sm.generation() != expected_state_generation
6645            || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6646        {
6647            return Ok(false);
6648        }
6649        sm.transition_to(target, reason)?;
6650        self.invalidate_pending_confirmation("state_transition")
6651            .await;
6652        sm.reset_no_transition();
6653        if let Some(staged) = staged {
6654            self.commit_staged_context_writes(staged);
6655        }
6656        let entered = sm.current();
6657        let is_reentry = Self::state_was_previously_entered(&entered, from_state, &history_before);
6658        drop(admission);
6659
6660        self.execute_state_enter_actions(&entered, is_reentry).await;
6661        drop(reservation);
6662        self.hooks
6663            .on_state_transition(Some(from_state), &entered, reason)
6664            .await;
6665        info!(from = %from_state, to = %entered, "State transition");
6666        Ok(true)
6667    }
6668
6669    /// Approves and applies a transition without staged context writes.
6670    async fn commit_transition_target(
6671        &self,
6672        from_state: &str,
6673        target: &str,
6674        reason: &str,
6675    ) -> Result<bool> {
6676        if !self.approve_transition_target(from_state, target).await? {
6677            return Ok(false);
6678        }
6679        self.apply_transition_target(from_state, target, reason, None)
6680            .await
6681    }
6682
6683    /// Applies an approved pre-response transition before redispatch.
6684    async fn apply_pre_response_transition_candidate(
6685        &self,
6686        candidate: &TransitionCandidate,
6687        staged: &HashMap<String, Value>,
6688        processed_input: &str,
6689    ) -> Result<bool> {
6690        self.commit_root_user_message(processed_input).await?;
6691        self.apply_transition_target(
6692            &candidate.from_state,
6693            candidate.target(),
6694            &candidate.reason,
6695            Some(staged),
6696        )
6697        .await
6698    }
6699
6700    /// Commits a pre-response transition after approval and before redispatch.
6701    async fn commit_pre_response_transition_candidate(
6702        &self,
6703        candidate: &TransitionCandidate,
6704        staged: &HashMap<String, Value>,
6705        processed_input: &str,
6706    ) -> Result<bool> {
6707        if !self
6708            .approve_transition_target(&candidate.from_state, candidate.target())
6709            .await?
6710        {
6711            return Ok(false);
6712        }
6713        self.apply_pre_response_transition_candidate(candidate, staged, processed_input)
6714            .await
6715    }
6716
6717    /// Handles post-response transition misses with fallback counters.
6718    async fn handle_transition_miss(&self, current_state: &str) -> Result<bool> {
6719        let Some(ref sm) = self.state_machine else {
6720            return Ok(false);
6721        };
6722        sm.increment_no_transition();
6723        let Some(fallback) = sm.check_fallback() else {
6724            return Ok(false);
6725        };
6726        self.commit_transition_target(current_state, &fallback, "fallback after no transitions")
6727            .await
6728    }
6729
6730    /// Evaluates and commits post-response transitions for the committed response path.
6731    async fn evaluate_transitions(&self, user_message: &str, response: &str) -> Result<bool> {
6732        let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6733            return Ok(false);
6734        };
6735        if transitions.is_empty() {
6736            return Ok(false);
6737        }
6738        if let Some(candidate) = self
6739            .select_transition_candidate(user_message, response)
6740            .await?
6741        {
6742            return self.commit_transition_candidate(&candidate).await;
6743        }
6744        self.handle_transition_miss(&current_state).await
6745    }
6746
6747    /// Attempts deterministic pre-response routing before old-state response generation.
6748    async fn try_pre_response_transition(
6749        &self,
6750        processed_input: &str,
6751    ) -> Result<Option<AgentResponse>> {
6752        let optimization = &self.runtime_config.optimization;
6753        if !optimization.enabled || !optimization.pre_response_deterministic_transitions {
6754            return Ok(None);
6755        }
6756        let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6757            return Ok(None);
6758        };
6759        let eligible: Vec<Transition> = transitions
6760            .into_iter()
6761            .filter(|transition| !transition.requires_response)
6762            .filter(|transition| matches!(transition.timing, TransitionTiming::PreResponse))
6763            .collect();
6764        if eligible.is_empty() {
6765            return Ok(None);
6766        }
6767
6768        let empty_staged = HashMap::new();
6769        let mut extracted_staged: Option<HashMap<String, Value>> = None;
6770        let mut selected: Option<(TransitionCandidate, HashMap<String, Value>)> = None;
6771
6772        for transition in &eligible {
6773            let use_extractors = optimization.pre_response_extractors || transition.run_extractors;
6774            let staged_for_eval = if use_extractors {
6775                if extracted_staged.is_none() {
6776                    extracted_staged =
6777                        Some(self.run_context_extractors_staged(processed_input).await);
6778                }
6779                extracted_staged.as_ref().unwrap_or(&empty_staged)
6780            } else {
6781                &empty_staged
6782            };
6783
6784            if let Some(candidate) = self.select_deterministic_transition_candidate(
6785                processed_input,
6786                &current_state,
6787                std::slice::from_ref(transition),
6788                staged_for_eval,
6789            ) {
6790                let staged_for_commit = if use_extractors {
6791                    staged_for_eval.clone()
6792                } else {
6793                    HashMap::new()
6794                };
6795                selected = Some((candidate, staged_for_commit));
6796                break;
6797            }
6798        }
6799
6800        let Some((candidate, staged)) = selected else {
6801            return Ok(None);
6802        };
6803
6804        if !self
6805            .commit_pre_response_transition_candidate(&candidate, &staged, processed_input)
6806            .await?
6807        {
6808            return Ok(None);
6809        }
6810        self.redispatch_current_state(processed_input)
6811            .await
6812            .map(Some)
6813    }
6814
6815    //
6816    // Speculative branches overlap independent decisions but still commit exactly one path.
6817    // Losing branches must remain data only and must not write memory, run tools, or emit output.
6818    //
6819    async fn try_speculative_branches(
6820        &self,
6821        processed_input: &str,
6822        input_context: &HashMap<String, Value>,
6823    ) -> Result<Option<AgentResponse>> {
6824        let optimization = &self.runtime_config.optimization;
6825        if !optimization.enabled {
6826            return Ok(None);
6827        }
6828
6829        let effective_reasoning_mode = self.get_effective_reasoning_config().mode.clone();
6830        if !matches!(
6831            effective_reasoning_mode,
6832            ReasoningMode::None | ReasoningMode::Auto
6833        ) {
6834            return Ok(None);
6835        }
6836
6837        let mut transition_enabled =
6838            optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
6839        let mut skill_enabled = optimization.speculative_skill_routing
6840            && self.skill_router.is_some()
6841            && self.pending_skill_id.read().is_none();
6842        let mut reasoning_enabled = optimization.speculative_reasoning_auto
6843            && matches!(effective_reasoning_mode, ReasoningMode::Auto);
6844
6845        if matches!(effective_reasoning_mode, ReasoningMode::Auto)
6846            && (!reasoning_enabled || optimization.max_speculative_llm_calls_per_turn < 2)
6847        {
6848            return Ok(None);
6849        }
6850
6851        if !transition_enabled && !skill_enabled && !reasoning_enabled {
6852            return Ok(None);
6853        }
6854
6855        let mut optional_slots = optimization.max_parallel_runtime_tasks.saturating_sub(1);
6856        let mut speculative_call_slots = optimization
6857            .max_speculative_llm_calls_per_turn
6858            .saturating_sub(1);
6859        if reasoning_enabled {
6860            if optional_slots == 0 || speculative_call_slots == 0 {
6861                return Ok(None);
6862            }
6863            optional_slots -= 1;
6864            speculative_call_slots -= 1;
6865        }
6866        if transition_enabled {
6867            if optional_slots == 0 {
6868                transition_enabled = false;
6869            } else {
6870                optional_slots -= 1;
6871            }
6872        }
6873        if skill_enabled && (optional_slots == 0 || speculative_call_slots == 0) {
6874            skill_enabled = false;
6875        }
6876
6877        if !transition_enabled && !skill_enabled && !reasoning_enabled {
6878            return Ok(None);
6879        }
6880
6881        let main_kind = if transition_enabled {
6882            RuntimeOptimizationKind::ParallelStateTransition
6883        } else if skill_enabled {
6884            RuntimeOptimizationKind::SpeculativeSkillRouting
6885        } else {
6886            RuntimeOptimizationKind::SpeculativeReasoningAuto
6887        };
6888        if !self.reserve_active_speculative_llm_call(main_kind) {
6889            return Ok(None);
6890        }
6891
6892        let mut branch_set = ScheduledBranchSet::new(optimization.max_parallel_runtime_tasks)?;
6893        let main_branch = RuntimeBranch::new(
6894            RuntimeTaskPurpose::MainResponse,
6895            main_kind,
6896            RuntimeTaskPriority::Normal,
6897            RuntimeCommitBehavior::FinalResponse,
6898        );
6899        let transition_branch = RuntimeBranch::new(
6900            RuntimeTaskPurpose::StateTransition,
6901            RuntimeOptimizationKind::ParallelStateTransition,
6902            RuntimeTaskPriority::Critical,
6903            RuntimeCommitBehavior::TransitionDecision,
6904        );
6905        let skill_branch = RuntimeBranch::new(
6906            RuntimeTaskPurpose::SkillRouting,
6907            RuntimeOptimizationKind::SpeculativeSkillRouting,
6908            RuntimeTaskPriority::High,
6909            RuntimeCommitBehavior::SkillSelection,
6910        );
6911        let reasoning_branch = RuntimeBranch::new(
6912            RuntimeTaskPurpose::ReasoningJudge,
6913            RuntimeOptimizationKind::SpeculativeReasoningAuto,
6914            RuntimeTaskPriority::Normal,
6915            RuntimeCommitBehavior::ReasoningDecision,
6916        );
6917        let main_id = main_branch.branch_id();
6918        let transition_id = transition_branch.branch_id();
6919        let skill_id = skill_branch.branch_id();
6920        let reasoning_id = reasoning_branch.branch_id();
6921
6922        let main_id_for_future = main_id.clone();
6923        if !branch_set.schedule(
6924            main_branch,
6925            Box::pin(async move {
6926                match crate::optimization::observability::with_branch_observation(
6927                    &main_id_for_future,
6928                    main_kind,
6929                    RuntimeCommitBehavior::FinalResponse,
6930                    self.generate_main_response_draft(processed_input, &ReasoningMode::None),
6931                )
6932                .await
6933                {
6934                    Ok(draft) => RuntimeBranchResult::MainDraft(draft),
6935                    Err(error) => RuntimeBranchResult::Failed(error),
6936                }
6937            }),
6938        ) {
6939            return Ok(None);
6940        }
6941
6942        if transition_enabled {
6943            let transition_id_for_future = transition_id.clone();
6944            if !branch_set.schedule(
6945                transition_branch,
6946                Box::pin(async move {
6947                    match crate::optimization::observability::with_branch_observation(
6948                        &transition_id_for_future,
6949                        RuntimeOptimizationKind::ParallelStateTransition,
6950                        RuntimeCommitBehavior::TransitionDecision,
6951                        self.select_parallel_transition_candidate(processed_input),
6952                    )
6953                    .await
6954                    {
6955                        Ok(ParallelTransitionSelection::Candidate(candidate)) => {
6956                            RuntimeBranchResult::Transition(Some(candidate))
6957                        }
6958                        Ok(ParallelTransitionSelection::NoMatch) => {
6959                            RuntimeBranchResult::Transition(None)
6960                        }
6961                        Ok(ParallelTransitionSelection::ReservationExhausted) => {
6962                            RuntimeBranchResult::Cancelled
6963                        }
6964                        Err(error) => RuntimeBranchResult::Failed(error),
6965                    }
6966                }),
6967            ) {
6968                transition_enabled = false;
6969            }
6970        }
6971
6972        if skill_enabled {
6973            let skill_id_for_future = skill_id.clone();
6974            if !branch_set.schedule(
6975                skill_branch,
6976                Box::pin(async move {
6977                    if !self.reserve_active_speculative_llm_call(
6978                        RuntimeOptimizationKind::SpeculativeSkillRouting,
6979                    ) {
6980                        return RuntimeBranchResult::Cancelled;
6981                    }
6982                    match crate::optimization::observability::with_branch_observation(
6983                        &skill_id_for_future,
6984                        RuntimeOptimizationKind::SpeculativeSkillRouting,
6985                        RuntimeCommitBehavior::SkillSelection,
6986                        self.select_skill_candidate(processed_input),
6987                    )
6988                    .await
6989                    {
6990                        Ok(candidate) => RuntimeBranchResult::Skill(candidate),
6991                        Err(error) => RuntimeBranchResult::Failed(error),
6992                    }
6993                }),
6994            ) {
6995                skill_enabled = false;
6996            }
6997        }
6998
6999        if reasoning_enabled {
7000            let reasoning_id_for_future = reasoning_id.clone();
7001            if !branch_set.schedule(
7002                reasoning_branch,
7003                Box::pin(async move {
7004                    if !self.reserve_active_speculative_llm_call(
7005                        RuntimeOptimizationKind::SpeculativeReasoningAuto,
7006                    ) {
7007                        return RuntimeBranchResult::Cancelled;
7008                    }
7009                    match crate::optimization::observability::with_branch_observation(
7010                        &reasoning_id_for_future,
7011                        RuntimeOptimizationKind::SpeculativeReasoningAuto,
7012                        RuntimeCommitBehavior::ReasoningDecision,
7013                        self.determine_reasoning_mode_strict(processed_input),
7014                    )
7015                    .await
7016                    {
7017                        Ok(mode) => RuntimeBranchResult::Reasoning(mode),
7018                        Err(error) => RuntimeBranchResult::Failed(error),
7019                    }
7020                }),
7021            ) {
7022                reasoning_enabled = false;
7023            }
7024        }
7025
7026        if matches!(effective_reasoning_mode, ReasoningMode::Auto) && !reasoning_enabled {
7027            self.finalize_pending_branches(branch_set.cancel_pending());
7028            return Ok(None);
7029        }
7030
7031        if !transition_enabled && !skill_enabled && !reasoning_enabled {
7032            self.finalize_pending_branches(branch_set.cancel_pending());
7033            return Ok(None);
7034        }
7035
7036        let mut main_pending = true;
7037        let mut skill_pending = skill_enabled;
7038        let mut reasoning_pending = reasoning_enabled;
7039        let mut transition_finalized = !transition_enabled;
7040        let mut skill_finalized = !skill_enabled;
7041        let mut reasoning_finalized = !reasoning_enabled;
7042        let mut main_result: Option<Result<MainResponseDraft>> = None;
7043        let mut transition_candidate: Option<TransitionCandidate> = None;
7044        let mut skill_candidate: Option<SkillCandidate> = None;
7045        let mut reasoning_decision: Option<ReasoningMode> = None;
7046        let mut transition_fallback_required = false;
7047        let mut skill_fallback_required = false;
7048        let mut reasoning_fallback_required = false;
7049
7050        loop {
7051            if let Some(candidate) = transition_candidate.take() {
7052                if self
7053                    .approve_transition_target(&candidate.from_state, candidate.target())
7054                    .await?
7055                {
7056                    // Drop losing provider futures before transition side effects can reuse their shared resources.
7057                    self.finalize_pending_branches(branch_set.cancel_pending());
7058                    if !main_pending {
7059                        self.finalize_branch_loss(
7060                            &main_id,
7061                            main_kind,
7062                            RuntimeCommitBehavior::FinalResponse,
7063                            false,
7064                            main_result.as_ref().map(|result| result.is_err()),
7065                        );
7066                    }
7067                    if skill_enabled && !skill_pending {
7068                        self.finalize_branch_loss(
7069                            &skill_id,
7070                            RuntimeOptimizationKind::SpeculativeSkillRouting,
7071                            RuntimeCommitBehavior::SkillSelection,
7072                            false,
7073                            Some(false),
7074                        );
7075                    }
7076                    if reasoning_enabled && !reasoning_pending {
7077                        self.finalize_branch_loss(
7078                            &reasoning_id,
7079                            RuntimeOptimizationKind::SpeculativeReasoningAuto,
7080                            RuntimeCommitBehavior::ReasoningDecision,
7081                            false,
7082                            Some(false),
7083                        );
7084                    }
7085                    if !self
7086                        .apply_pre_response_transition_candidate(
7087                            &candidate,
7088                            &HashMap::new(),
7089                            processed_input,
7090                        )
7091                        .await?
7092                    {
7093                        self.finalize_optional_branch(
7094                            &transition_id,
7095                            RuntimeOptimizationKind::ParallelStateTransition,
7096                            RuntimeCommitBehavior::TransitionDecision,
7097                            "discarded",
7098                            false,
7099                        );
7100                        return Ok(None);
7101                    }
7102                    self.finalize_optional_branch(
7103                        &transition_id,
7104                        RuntimeOptimizationKind::ParallelStateTransition,
7105                        RuntimeCommitBehavior::TransitionDecision,
7106                        "committed",
7107                        true,
7108                    );
7109                    return self
7110                        .redispatch_current_state(processed_input)
7111                        .await
7112                        .map(Some);
7113                }
7114                self.finalize_optional_branch(
7115                    &transition_id,
7116                    RuntimeOptimizationKind::ParallelStateTransition,
7117                    RuntimeCommitBehavior::TransitionDecision,
7118                    "discarded",
7119                    false,
7120                );
7121                transition_finalized = true;
7122            }
7123
7124            if transition_finalized && skill_candidate.is_some() {
7125                let candidate = skill_candidate.take().unwrap();
7126                self.finalize_optional_branch(
7127                    &skill_id,
7128                    RuntimeOptimizationKind::SpeculativeSkillRouting,
7129                    RuntimeCommitBehavior::SkillSelection,
7130                    "committed",
7131                    true,
7132                );
7133                if !main_pending {
7134                    self.finalize_branch_loss(
7135                        &main_id,
7136                        main_kind,
7137                        RuntimeCommitBehavior::FinalResponse,
7138                        false,
7139                        main_result.as_ref().map(|result| result.is_err()),
7140                    );
7141                }
7142                if reasoning_enabled && !reasoning_pending {
7143                    self.finalize_branch_loss(
7144                        &reasoning_id,
7145                        RuntimeOptimizationKind::SpeculativeReasoningAuto,
7146                        RuntimeCommitBehavior::ReasoningDecision,
7147                        false,
7148                        Some(false),
7149                    );
7150                }
7151                self.finalize_pending_branches(branch_set.cancel_pending());
7152                self.commit_root_user_message(processed_input).await?;
7153                return match self
7154                    .commit_skill_candidate_route_result(candidate, processed_input)
7155                    .await?
7156                {
7157                    SkillRouteResult::Response { skill_id, content } => self
7158                        .handle_skill_response(processed_input, &skill_id, content, input_context)
7159                        .await
7160                        .map(Some),
7161                    SkillRouteResult::NeedsClarification {
7162                        response,
7163                        ownership,
7164                    } => {
7165                        let admission = self
7166                            .admit_optional_disambiguation_ownership(ownership)
7167                            .await?;
7168                        if response
7169                            .metadata
7170                            .as_ref()
7171                            .and_then(|m| m.get("disambiguation"))
7172                            .and_then(|d| d.get("status"))
7173                            .and_then(|s| s.as_str())
7174                            == Some("awaiting_clarification")
7175                        {
7176                            self.memory
7177                                .add_message(ChatMessage::assistant(&response.content))
7178                                .await?;
7179                        }
7180                        drop(admission);
7181                        self.finish_turn_if_root(&response).await?;
7182                        Ok(Some(response))
7183                    }
7184                    SkillRouteResult::NoMatch => Ok(None),
7185                };
7186            }
7187
7188            if transition_finalized
7189                && skill_finalized
7190                && let Some(reasoning_mode) = reasoning_decision.take()
7191            {
7192                if !matches!(reasoning_mode, ReasoningMode::None) {
7193                    self.finalize_optional_branch(
7194                        &reasoning_id,
7195                        RuntimeOptimizationKind::SpeculativeReasoningAuto,
7196                        RuntimeCommitBehavior::ReasoningDecision,
7197                        "committed",
7198                        true,
7199                    );
7200                    if !main_pending {
7201                        self.finalize_branch_loss(
7202                            &main_id,
7203                            main_kind,
7204                            RuntimeCommitBehavior::FinalResponse,
7205                            false,
7206                            main_result.as_ref().map(|result| result.is_err()),
7207                        );
7208                    }
7209                    self.finalize_pending_branches(branch_set.cancel_pending());
7210                    self.commit_root_user_message(processed_input).await?;
7211                    return if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
7212                        self.handle_plan_and_execute(processed_input, input_context, true)
7213                            .await
7214                            .map(Some)
7215                    } else {
7216                        self.run_committed_response_loop_with_reasoning(
7217                            processed_input,
7218                            input_context,
7219                            reasoning_mode,
7220                            true,
7221                        )
7222                        .await
7223                        .map(Some)
7224                    };
7225                }
7226                self.finalize_optional_branch(
7227                    &reasoning_id,
7228                    RuntimeOptimizationKind::SpeculativeReasoningAuto,
7229                    RuntimeCommitBehavior::ReasoningDecision,
7230                    "committed",
7231                    true,
7232                );
7233                reasoning_finalized = true;
7234            }
7235
7236            if transition_finalized && skill_finalized && reasoning_finalized {
7237                if transition_fallback_required
7238                    || skill_fallback_required
7239                    || reasoning_fallback_required
7240                {
7241                    if !main_pending {
7242                        self.finalize_branch_loss(
7243                            &main_id,
7244                            main_kind,
7245                            RuntimeCommitBehavior::FinalResponse,
7246                            false,
7247                            main_result.as_ref().map(|result| result.is_err()),
7248                        );
7249                    }
7250                    self.finalize_pending_branches(branch_set.cancel_pending());
7251                    return Ok(None);
7252                }
7253
7254                if let Some(result) = main_result.take() {
7255                    let draft = match result {
7256                        Ok(draft) => draft,
7257                        Err(error) => {
7258                            self.finalize_optional_branch(
7259                                &main_id,
7260                                main_kind,
7261                                RuntimeCommitBehavior::FinalResponse,
7262                                "failed",
7263                                false,
7264                            );
7265                            self.finalize_pending_branches(branch_set.cancel_pending());
7266                            return Err(error);
7267                        }
7268                    };
7269                    self.finalize_optional_branch(
7270                        &main_id,
7271                        main_kind,
7272                        RuntimeCommitBehavior::FinalResponse,
7273                        "committed",
7274                        true,
7275                    );
7276                    self.finalize_pending_branches(branch_set.cancel_pending());
7277                    return self
7278                        .commit_main_response_draft(
7279                            processed_input,
7280                            input_context,
7281                            draft,
7282                            ReasoningMode::None,
7283                            reasoning_enabled,
7284                        )
7285                        .await
7286                        .map(Some);
7287                }
7288            }
7289
7290            if branch_set.is_empty() {
7291                return Ok(None);
7292            }
7293
7294            let Some(outcome) = branch_set.next_completed().await else {
7295                return Ok(None);
7296            };
7297            let branch_id = outcome.branch.branch_id();
7298            match outcome.result {
7299                RuntimeBranchResult::MainDraft(draft) => {
7300                    main_pending = false;
7301                    main_result = Some(Ok(draft));
7302                }
7303                RuntimeBranchResult::Transition(candidate) => {
7304                    if let Some(candidate) = candidate {
7305                        transition_candidate = Some(candidate);
7306                    } else {
7307                        self.finalize_optional_branch(
7308                            &transition_id,
7309                            RuntimeOptimizationKind::ParallelStateTransition,
7310                            RuntimeCommitBehavior::TransitionDecision,
7311                            "discarded",
7312                            false,
7313                        );
7314                        transition_finalized = true;
7315                    }
7316                }
7317                RuntimeBranchResult::Skill(candidate) => {
7318                    skill_pending = false;
7319                    if let Some(candidate) = candidate {
7320                        skill_candidate = Some(candidate);
7321                    } else {
7322                        self.finalize_optional_branch(
7323                            &skill_id,
7324                            RuntimeOptimizationKind::SpeculativeSkillRouting,
7325                            RuntimeCommitBehavior::SkillSelection,
7326                            "discarded",
7327                            false,
7328                        );
7329                        skill_finalized = true;
7330                    }
7331                }
7332                RuntimeBranchResult::Reasoning(mode) => {
7333                    reasoning_pending = false;
7334                    reasoning_decision = Some(mode);
7335                }
7336                RuntimeBranchResult::Failed(error) => {
7337                    if branch_id == main_id {
7338                        main_pending = false;
7339                        main_result = Some(Err(error));
7340                    } else if branch_id == transition_id {
7341                        self.finalize_optional_branch(
7342                            &transition_id,
7343                            RuntimeOptimizationKind::ParallelStateTransition,
7344                            RuntimeCommitBehavior::TransitionDecision,
7345                            "failed",
7346                            false,
7347                        );
7348                        transition_finalized = true;
7349                    } else if branch_id == skill_id {
7350                        skill_pending = false;
7351                        self.finalize_optional_branch(
7352                            &skill_id,
7353                            RuntimeOptimizationKind::SpeculativeSkillRouting,
7354                            RuntimeCommitBehavior::SkillSelection,
7355                            "failed",
7356                            false,
7357                        );
7358                        skill_finalized = true;
7359                    } else if branch_id == reasoning_id {
7360                        reasoning_pending = false;
7361                        self.finalize_optional_branch(
7362                            &reasoning_id,
7363                            RuntimeOptimizationKind::SpeculativeReasoningAuto,
7364                            RuntimeCommitBehavior::ReasoningDecision,
7365                            "failed",
7366                            false,
7367                        );
7368                        reasoning_finalized = true;
7369                    }
7370                }
7371                RuntimeBranchResult::Cancelled => {
7372                    self.finalize_optional_branch(
7373                        &branch_id,
7374                        outcome.branch.optimization,
7375                        outcome.branch.commit_behavior,
7376                        "cancelled",
7377                        false,
7378                    );
7379                    if branch_id == main_id {
7380                        main_pending = false;
7381                        main_result =
7382                            Some(Err(AgentError::Other("main branch cancelled".to_string())));
7383                    } else if branch_id == transition_id {
7384                        transition_finalized = true;
7385                        transition_fallback_required = true;
7386                    } else if branch_id == skill_id {
7387                        skill_pending = false;
7388                        skill_finalized = true;
7389                        skill_fallback_required = true;
7390                    } else if branch_id == reasoning_id {
7391                        reasoning_pending = false;
7392                        reasoning_finalized = true;
7393                        reasoning_fallback_required = true;
7394                    }
7395                }
7396            }
7397        }
7398    }
7399
7400    fn finalize_pending_branches(&self, branches: Vec<RuntimeBranch>) {
7401        for branch in branches {
7402            self.finalize_optional_branch(
7403                &branch.branch_id(),
7404                branch.optimization,
7405                branch.commit_behavior,
7406                "cancelled",
7407                false,
7408            );
7409        }
7410    }
7411
7412    //
7413    // Pending losers are reported as cancelled because their futures are dropped before completion.
7414    // Completed losers keep failed or discarded status based on their recorded result.
7415    //
7416    fn finalize_branch_loss(
7417        &self,
7418        branch_id: &str,
7419        optimization: RuntimeOptimizationKind,
7420        commit_behavior: RuntimeCommitBehavior,
7421        pending: bool,
7422        completed_failed: Option<bool>,
7423    ) {
7424        let status = if pending {
7425            "cancelled"
7426        } else if completed_failed.unwrap_or(false) {
7427            "failed"
7428        } else {
7429            "discarded"
7430        };
7431        self.finalize_optional_branch(branch_id, optimization, commit_behavior, status, false);
7432    }
7433
7434    //
7435    // Finalization is separated from commit so losing branches remain observable.
7436    // This helper must not mutate runtime state other than observability.
7437    //
7438    fn finalize_optional_branch(
7439        &self,
7440        branch_id: &str,
7441        optimization: RuntimeOptimizationKind,
7442        commit_behavior: RuntimeCommitBehavior,
7443        status: &str,
7444        winner: bool,
7445    ) {
7446        crate::optimization::observability::finalize_branch(
7447            self.observability_manager.as_ref(),
7448            branch_id,
7449            status,
7450            winner,
7451            optimization,
7452            commit_behavior,
7453        );
7454    }
7455
7456    //
7457    // This is only an eligibility check.
7458    // Actual transition selection happens in select_parallel_transition_candidate.
7459    //
7460    fn has_parallel_transition_candidates(&self) -> bool {
7461        self.transitions_available_for_commit()
7462            .map(|(transitions, _)| {
7463                transitions
7464                    .iter()
7465                    .any(|transition| matches!(transition.timing, TransitionTiming::Parallel))
7466            })
7467            .unwrap_or(false)
7468    }
7469
7470    //
7471    // Parallel transition prompts must not depend on assistant response text.
7472    // Keep this branch response-independent or it can race against invalid context.
7473    //
7474    async fn select_parallel_transition_candidate(
7475        &self,
7476        processed_input: &str,
7477    ) -> Result<ParallelTransitionSelection> {
7478        let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
7479            return Ok(ParallelTransitionSelection::NoMatch);
7480        };
7481        let parallel: Vec<Transition> = transitions
7482            .into_iter()
7483            .filter(|transition| matches!(transition.timing, TransitionTiming::Parallel))
7484            .filter(|transition| !transition.requires_response)
7485            .collect();
7486        if parallel.is_empty() {
7487            return Ok(ParallelTransitionSelection::NoMatch);
7488        }
7489        let empty_staged = HashMap::new();
7490        if let Some(candidate) = self.select_deterministic_transition_candidate(
7491            processed_input,
7492            &current_state,
7493            &parallel,
7494            &empty_staged,
7495        ) {
7496            return Ok(ParallelTransitionSelection::Candidate(candidate));
7497        }
7498        let when_transitions: Vec<(usize, &Transition)> = parallel
7499            .iter()
7500            .enumerate()
7501            .filter(|(_, transition)| !transition.when.trim().is_empty())
7502            .collect();
7503        if when_transitions.is_empty() {
7504            return Ok(ParallelTransitionSelection::NoMatch);
7505        }
7506        let llm = self
7507            .llm_registry
7508            .router()
7509            .or_else(|_| self.llm_registry.default())
7510            .map_err(|e| AgentError::Config(e.to_string()))?;
7511        let conditions = when_transitions
7512            .iter()
7513            .enumerate()
7514            .map(|(display_idx, (_, transition))| {
7515                format!("{}. {}", display_idx + 1, transition.when)
7516            })
7517            .collect::<Vec<_>>()
7518            .join("\n");
7519        if !self
7520            .reserve_active_speculative_llm_call(RuntimeOptimizationKind::ParallelStateTransition)
7521        {
7522            return Ok(ParallelTransitionSelection::ReservationExhausted);
7523        }
7524        let context_preview = self.branch_context_preview();
7525        let prompt = format!(
7526            "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-{}).",
7527            current_state,
7528            processed_input,
7529            context_preview,
7530            conditions,
7531            when_transitions.len()
7532        );
7533        let response = self
7534            .observe_purpose(
7535                ObservationPurpose::StateTransitionEvaluation,
7536                llm.complete(&[ChatMessage::user(prompt)], None),
7537            )
7538            .await
7539            .map_err(|e| AgentError::LLM(e.to_string()))?;
7540        let choice = response.content.trim().parse::<usize>().unwrap_or(0);
7541        if choice == 0 || choice > when_transitions.len() {
7542            return Ok(ParallelTransitionSelection::NoMatch);
7543        }
7544        let transition = when_transitions[choice - 1].1.clone();
7545        Ok(ParallelTransitionSelection::Candidate(
7546            TransitionCandidate::new(
7547                current_state,
7548                transition.clone(),
7549                Self::transition_reason(&transition),
7550            ),
7551        ))
7552    }
7553
7554    /// Re-enters the runtime loop after an optimized transition commits.
7555    async fn redispatch_current_state(&self, processed_input: &str) -> Result<AgentResponse> {
7556        const MAX_REDISPATCH_DEPTH: u32 = 3;
7557        let current_depth = *self.redispatch_depth.read();
7558        if current_depth >= MAX_REDISPATCH_DEPTH {
7559            warn!(depth = current_depth, "Re-dispatch depth limit reached");
7560            let response = AgentResponse::new("");
7561            self.finish_turn_if_root(&response).await?;
7562            return Ok(response);
7563        }
7564        *self.redispatch_depth.write() += 1;
7565        if let Some(context) = self.active_turn_context.write().as_mut() {
7566            context.enter_redispatch();
7567        }
7568        let result = Box::pin(self.run_loop_internal(processed_input)).await;
7569        *self.redispatch_depth.write() -= 1;
7570        if let Some(context) = self.active_turn_context.write().as_mut() {
7571            context.exit_redispatch();
7572        }
7573        let response = result?;
7574        self.finish_turn_if_root(&response).await?;
7575        Ok(response)
7576    }
7577
7578    /// Runs final response hooks and maintenance only for the root dispatch.
7579    async fn finish_turn_if_root(&self, response: &AgentResponse) -> Result<()> {
7580        if *self.redispatch_depth.read() == 0 {
7581            self.post_turn_session_lifecycle().await?;
7582            if let Some(context) = self.active_turn_context.write().as_mut() {
7583                context.mark_post_turn_lifecycle_completed();
7584            }
7585            self.hooks.on_response(response).await;
7586            self.end_root_turn();
7587        }
7588        Ok(())
7589    }
7590
7591    /// Execute on_exit actions for a state being left.
7592    async fn execute_state_exit_actions(&self, state_path: &str) {
7593        if let Some(ref sm) = self.state_machine
7594            && let Some(def) = sm.get_definition(state_path)
7595            && !def.on_exit.is_empty()
7596        {
7597            debug!(state = %state_path, count = def.on_exit.len(), "Executing on_exit actions");
7598            self.execute_state_actions(&def.on_exit).await;
7599        }
7600    }
7601
7602    /// Returns whether a transition target had already been entered before this transition.
7603    fn state_was_previously_entered(
7604        state_path: &str,
7605        from_state: &str,
7606        history_before: &[StateTransitionEvent],
7607    ) -> bool {
7608        state_path == from_state
7609            || history_before
7610                .iter()
7611                .any(|event| event.from == state_path || event.to == state_path)
7612    }
7613
7614    /// Execute on_enter (or on_reenter) actions for a state being entered.
7615    async fn execute_state_enter_actions(&self, state_path: &str, is_reentry: bool) {
7616        if let Some(ref sm) = self.state_machine
7617            && let Some(def) = sm.get_definition(state_path)
7618        {
7619            if is_reentry && !def.on_reenter.is_empty() {
7620                debug!(state = %state_path, count = def.on_reenter.len(), "Executing on_reenter actions");
7621                self.execute_state_actions(&def.on_reenter).await;
7622            } else if !def.on_enter.is_empty() {
7623                debug!(state = %state_path, count = def.on_enter.len(), "Executing on_enter actions");
7624                self.execute_state_actions(&def.on_enter).await;
7625            }
7626        }
7627    }
7628
7629    /// Execute a list of state actions (tool calls, skill invocations, context updates, LLM prompts).
7630    async fn execute_state_actions(&self, actions: &[StateAction]) {
7631        for (action_index, action) in actions.iter().enumerate() {
7632            match action {
7633                StateAction::Tool { tool, args } => {
7634                    let raw_args = args.clone().unwrap_or(Value::Object(Default::default()));
7635                    let args_value = self.render_action_args(&raw_args);
7636                    let state = self.state_machine.as_ref().map(|sm| sm.current());
7637                    let request = ToolExecutionRequest::new(
7638                        uuid::Uuid::new_v4().to_string(),
7639                        tool.clone(),
7640                        args_value,
7641                        ToolCallSource::StateAction {
7642                            state,
7643                            action_index,
7644                        },
7645                    );
7646                    match self.execute_tool_record(request).await {
7647                        Ok(record) if record.success => {
7648                            debug!(tool = %record.canonical_id, "State action: tool executed");
7649                            let _ = self.context_manager.set(
7650                                "last_tool_result",
7651                                serde_json::Value::String(record.model_output_string()),
7652                            );
7653                            let _ = self.context_manager.set(
7654                                "last_tool_record",
7655                                serde_json::to_value(record).unwrap_or(Value::Null),
7656                            );
7657                        }
7658                        Ok(record) => {
7659                            warn!(tool = %record.canonical_id, error = %record.output, "State action: tool failed");
7660                        }
7661                        Err(e) => {
7662                            warn!(tool = %tool, error = %e, "State action: tool failed")
7663                        }
7664                    }
7665                }
7666                StateAction::Skill { skill } => {
7667                    if let Some(ref executor) = self.skill_executor {
7668                        if let Some(def) = self.skills.iter().find(|s| s.id == *skill) {
7669                            match executor
7670                                .execute_with_invoker(def, "", serde_json::json!({}), self)
7671                                .await
7672                            {
7673                                Ok(_) => debug!(skill = %skill, "State action: skill executed"),
7674                                Err(e) => {
7675                                    warn!(skill = %skill, error = %e, "State action: skill failed")
7676                                }
7677                            }
7678                        } else {
7679                            warn!(skill = %skill, "State action: skill not found");
7680                        }
7681                    }
7682                }
7683                StateAction::SetContext { set_context } => {
7684                    for (key, value) in set_context {
7685                        if let Err(e) = self.context_manager.set(key, value.clone()) {
7686                            warn!(key = %key, error = %e, "State action: set_context failed");
7687                        } else {
7688                            debug!(key = %key, "State action: context set");
7689                        }
7690                    }
7691                }
7692                StateAction::Prompt {
7693                    prompt,
7694                    llm,
7695                    store_as,
7696                } => {
7697                    let llm_result = if let Some(alias) = llm {
7698                        self.llm_registry.get(alias)
7699                    } else {
7700                        self.llm_registry.default()
7701                    };
7702                    match llm_result {
7703                        Ok(llm_provider) => {
7704                            // Render template variables and include conversation context
7705                            let context = self.build_context_with_overlays();
7706                            let rendered_prompt = self
7707                                .template_renderer
7708                                .render(prompt, &context)
7709                                .unwrap_or_else(|_| prompt.clone());
7710                            let recent =
7711                                self.memory.get_messages(Some(5)).await.unwrap_or_default();
7712                            let mut messages: Vec<ChatMessage> = recent;
7713                            messages.push(ChatMessage::user(&rendered_prompt));
7714                            match self
7715                                .observe_purpose(
7716                                    ObservationPurpose::StateAction,
7717                                    llm_provider.complete(&messages, None),
7718                                )
7719                                .await
7720                            {
7721                                Ok(response) => {
7722                                    if let Some(key) = store_as {
7723                                        let _ = self
7724                                            .context_manager
7725                                            .set(key, Value::String(response.content));
7726                                        debug!(key = %key, "State action: prompt result stored");
7727                                    }
7728                                }
7729                                Err(e) => {
7730                                    warn!(error = %e, "State action: prompt LLM call failed");
7731                                }
7732                            }
7733                        }
7734                        Err(e) => {
7735                            warn!(error = %e, "State action: LLM not found for prompt");
7736                        }
7737                    }
7738                }
7739            }
7740        }
7741    }
7742
7743    async fn run_context_extractors_staged(&self, user_message: &str) -> HashMap<String, Value> {
7744        let extractors = match &self.state_machine {
7745            Some(sm) => match sm.current_definition() {
7746                Some(def) if !def.extract.is_empty() => def.extract.clone(),
7747                _ => return HashMap::new(),
7748            },
7749            None => return HashMap::new(),
7750        };
7751
7752        let mut staged = HashMap::new();
7753        for extractor in &extractors {
7754            let prompt = if let Some(ref custom) = extractor.llm_extract {
7755                format!(
7756                    "User message:\n\"{}\"\n\nInstruction:\n{}",
7757                    user_message, custom
7758                )
7759            } else if let Some(ref desc) = extractor.description {
7760                format!(
7761                    "From the following message, extract: {}\n\n\
7762                     Message: \"{}\"\n\n\
7763                     If the information is present, return ONLY the extracted value.\n\
7764                     If NOT present, return exactly: __NONE__",
7765                    desc, user_message
7766                )
7767            } else {
7768                continue;
7769            };
7770
7771            let llm = match self
7772                .llm_registry
7773                .get(&extractor.llm)
7774                .or_else(|_| self.llm_registry.get("router"))
7775                .or_else(|_| self.llm_registry.get("default"))
7776            {
7777                Ok(llm) => llm,
7778                Err(e) => {
7779                    warn!(key = %extractor.key, error = %e, "Extractor LLM not found");
7780                    continue;
7781                }
7782            };
7783
7784            let messages = vec![ChatMessage::user(&prompt)];
7785            match self
7786                .observe_purpose(
7787                    ObservationPurpose::ContextExtraction,
7788                    llm.complete(&messages, None),
7789                )
7790                .await
7791            {
7792                Ok(response) => {
7793                    let value = response.content.trim().to_string();
7794                    if value != "__NONE__" && !value.is_empty() {
7795                        staged.insert(
7796                            extractor.key.clone(),
7797                            serde_json::Value::String(value.clone()),
7798                        );
7799                        debug!(key = %extractor.key, value = %value, "Context extracted");
7800                    } else if extractor.required {
7801                        warn!(key = %extractor.key, "Required extraction returned no value");
7802                    }
7803                }
7804                Err(e) => {
7805                    warn!(key = %extractor.key, error = %e, "Context extraction LLM call failed");
7806                }
7807            }
7808        }
7809        staged
7810    }
7811
7812    fn commit_staged_context_writes(&self, staged: &HashMap<String, Value>) {
7813        for (key, value) in staged {
7814            if let Err(error) = self.context_manager.update(key, value.clone()) {
7815                warn!(key = %key, error = %error, "staged context write failed");
7816            }
7817        }
7818    }
7819
7820    /// Run context extractors for the current state on the user's input.
7821    async fn run_context_extractors(&self, user_message: &str) {
7822        let staged = self.run_context_extractors_staged(user_message).await;
7823        self.commit_staged_context_writes(&staged);
7824    }
7825
7826    async fn check_memory_compression(&self) -> Result<()> {
7827        if self.memory.needs_compression() {
7828            let result = self.memory.compress(None).await?;
7829            if let CompressResult::Compressed {
7830                messages_summarized,
7831                new_summary_length,
7832                tokens_saved,
7833            } = result
7834            {
7835                let event = MemoryCompressEvent::new(
7836                    messages_summarized,
7837                    tokens_saved,
7838                    new_summary_length as u32,
7839                );
7840                self.hooks.on_memory_compress(&event).await;
7841                debug!(
7842                    messages = messages_summarized,
7843                    tokens_saved = tokens_saved,
7844                    "Memory compressed"
7845                );
7846            }
7847        }
7848
7849        // Handle overflow AFTER compression, then check warning threshold
7850        self.handle_memory_overflow().await?;
7851        self.check_memory_budget().await;
7852
7853        Ok(())
7854    }
7855
7856    async fn check_memory_budget(&self) {
7857        let Some(ref budget) = self.memory_token_budget else {
7858            return;
7859        };
7860
7861        let context = match self.memory.get_context().await {
7862            Ok(ctx) => ctx,
7863            Err(_) => return,
7864        };
7865
7866        // Overall budget warning
7867        let used_tokens = context.estimated_tokens();
7868        if budget.is_over_warn_threshold(used_tokens) {
7869            let event = MemoryBudgetEvent::new("memory", used_tokens, budget.total);
7870            self.hooks.on_memory_budget_warning(&event).await;
7871            debug!(
7872                used = used_tokens,
7873                total = budget.total,
7874                percent = event.usage_percent,
7875                "Memory budget warning"
7876            );
7877        }
7878
7879        // Per-component warning: summary
7880        if let Some(ref summary) = context.summary {
7881            let summary_tokens = ai_agents_memory::estimate_tokens(summary);
7882            let summary_budget = budget.allocation.summary;
7883            if summary_budget > 0 {
7884                let warn_threshold =
7885                    (summary_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
7886                if summary_tokens >= warn_threshold {
7887                    let event = MemoryBudgetEvent::new("summary", summary_tokens, summary_budget);
7888                    self.hooks.on_memory_budget_warning(&event).await;
7889                }
7890            }
7891        }
7892
7893        // Per-component warning: recent_messages
7894        let recent_tokens: u32 = context
7895            .messages
7896            .iter()
7897            .map(ai_agents_memory::estimate_message_tokens)
7898            .sum();
7899        let recent_budget = budget.allocation.recent_messages;
7900        if recent_budget > 0 {
7901            let warn_threshold =
7902                (recent_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
7903            if recent_tokens >= warn_threshold {
7904                let event = MemoryBudgetEvent::new("recent_messages", recent_tokens, recent_budget);
7905                self.hooks.on_memory_budget_warning(&event).await;
7906            }
7907        }
7908
7909        let relationship_budget = budget.allocation.relationships;
7910        if relationship_budget > 0 {
7911            let relationship_tokens = self
7912                .relationship_memory_text()
7913                .map(|text| ai_agents_memory::estimate_tokens(&text))
7914                .unwrap_or(0);
7915            let warn_threshold =
7916                (relationship_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
7917            if relationship_tokens >= warn_threshold {
7918                let event = MemoryBudgetEvent::new(
7919                    "relationships",
7920                    relationship_tokens,
7921                    relationship_budget,
7922                );
7923                self.hooks.on_memory_budget_warning(&event).await;
7924            }
7925        }
7926    }
7927
7928    async fn handle_memory_overflow(&self) -> Result<()> {
7929        let Some(ref budget) = self.memory_token_budget else {
7930            return Ok(());
7931        };
7932
7933        let context = self.memory.get_context().await?;
7934        let used_tokens = context.estimated_tokens();
7935
7936        if used_tokens <= budget.total {
7937            return Ok(());
7938        }
7939
7940        match budget.overflow_strategy {
7941            OverflowStrategy::TruncateOldest => {
7942                let tokens_to_free = used_tokens - budget.total;
7943                let messages_to_evict = self.calculate_eviction_count(tokens_to_free);
7944                if messages_to_evict > 0 {
7945                    self.evict_messages(messages_to_evict, EvictionReason::TokenBudgetExceeded)
7946                        .await?;
7947                }
7948            }
7949            OverflowStrategy::SummarizeMore => {
7950                let max_attempts = context.total_messages.max(1);
7951                for _ in 0..max_attempts {
7952                    match self.memory.compress(None).await? {
7953                        CompressResult::Compressed {
7954                            messages_summarized,
7955                            ..
7956                        } if messages_summarized > 0 => {
7957                            let context = self.memory.get_context().await?;
7958                            if context.estimated_tokens() <= budget.total {
7959                                return Ok(());
7960                            }
7961                        }
7962                        _ => break,
7963                    }
7964                }
7965                let context = self.memory.get_context().await?;
7966                let used_tokens = context.estimated_tokens();
7967                if used_tokens > budget.total {
7968                    return Err(AgentError::MemoryBudgetExceeded {
7969                        used: used_tokens,
7970                        budget: budget.total,
7971                    });
7972                }
7973            }
7974            OverflowStrategy::Error => {
7975                return Err(AgentError::MemoryBudgetExceeded {
7976                    used: used_tokens,
7977                    budget: budget.total,
7978                });
7979            }
7980        }
7981        Ok(())
7982    }
7983
7984    fn calculate_eviction_count(&self, tokens_to_free: u32) -> usize {
7985        // Estimate ~50 tokens per message on average
7986        ((tokens_to_free as f64 / 50.0).ceil() as usize).max(1)
7987    }
7988
7989    async fn evict_messages(&self, count: usize, reason: EvictionReason) -> Result<()> {
7990        let evicted = self.memory.evict_oldest(count).await?;
7991        if !evicted.is_empty() {
7992            let event = MemoryEvictEvent {
7993                reason,
7994                messages_evicted: evicted.len(),
7995                importance_scores: vec![],
7996            };
7997            self.hooks.on_memory_evict(&event).await;
7998            debug!(count = evicted.len(), "Messages evicted from memory");
7999        }
8000        Ok(())
8001    }
8002
8003    #[instrument(skip(self, input), fields(agent = %self.info.name))]
8004    async fn determine_reasoning_mode(&self, input: &str) -> Result<ReasoningMode> {
8005        match self.determine_reasoning_mode_strict(input).await {
8006            Ok(mode) => Ok(mode),
8007            Err(_) => Ok(ReasoningMode::None),
8008        }
8009    }
8010
8011    async fn determine_reasoning_mode_strict(&self, input: &str) -> Result<ReasoningMode> {
8012        let effective_config = self.get_effective_reasoning_config();
8013
8014        if !matches!(effective_config.mode, ReasoningMode::Auto) {
8015            return Ok(effective_config.mode.clone());
8016        }
8017
8018        let judge_llm = effective_config
8019            .judge_llm
8020            .as_ref()
8021            .and_then(|alias| self.llm_registry.get(alias).ok())
8022            .or_else(|| self.llm_registry.router().ok())
8023            .or_else(|| self.llm_registry.default().ok());
8024
8025        let Some(llm) = judge_llm else {
8026            return Ok(ReasoningMode::None);
8027        };
8028
8029        let prompt = format!(
8030            r#"Analyze this user request and determine the appropriate reasoning mode.
8031
8032User request: "{}"
8033
8034Choose ONE of these modes:
8035- none: Simple queries, greetings, direct answers (fastest)
8036- cot: Complex analysis, multi-step reasoning, math problems
8037- react: Tasks requiring multiple tool calls with observation
8038- plan_and_execute: Complex multi-step tasks requiring coordination
8039
8040Respond with ONLY the mode name (none, cot, react, or plan_and_execute)."#,
8041            input
8042        );
8043
8044        let messages = vec![ChatMessage::user(&prompt)];
8045        let response = self
8046            .observe_purpose(
8047                ObservationPurpose::ReflectionDecision,
8048                llm.complete(&messages, None),
8049            )
8050            .await
8051            .map_err(|e| AgentError::LLM(e.to_string()))?;
8052
8053        let mode_str = response.content.trim().to_lowercase();
8054        Ok(match mode_str.as_str() {
8055            "cot" => ReasoningMode::CoT,
8056            "react" => ReasoningMode::React,
8057            "plan_and_execute" => ReasoningMode::PlanAndExecute,
8058            _ => ReasoningMode::None,
8059        })
8060    }
8061
8062    async fn should_reflect(&self, input: &str, response: &str) -> Result<bool> {
8063        let effective_config = self.get_effective_reflection_config();
8064
8065        if !effective_config.requires_evaluation() {
8066            return Ok(false);
8067        }
8068
8069        if effective_config.is_enabled() {
8070            return Ok(true);
8071        }
8072
8073        let evaluator_llm = effective_config
8074            .evaluator_llm
8075            .as_ref()
8076            .and_then(|alias| self.llm_registry.get(alias).ok())
8077            .or_else(|| self.llm_registry.router().ok())
8078            .or_else(|| self.llm_registry.default().ok());
8079
8080        let Some(llm) = evaluator_llm else {
8081            return Ok(false);
8082        };
8083
8084        let response_preview: String = response.chars().take(500).collect();
8085        let prompt = format!(
8086            r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
8087
8088User query: "{}"
8089Response: "{}"
8090
8091Answer YES or NO only."#,
8092            input, response_preview
8093        );
8094
8095        let messages = vec![ChatMessage::user(&prompt)];
8096        let result = self
8097            .observe_purpose(
8098                ObservationPurpose::ReflectionDecision,
8099                llm.complete(&messages, None),
8100            )
8101            .await;
8102
8103        match result {
8104            Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
8105            Err(_) => Ok(false),
8106        }
8107    }
8108
8109    fn build_cot_system_prompt(&self, base_prompt: &str) -> String {
8110        format!(
8111            "{}\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>",
8112            base_prompt
8113        )
8114    }
8115
8116    fn build_react_system_prompt(&self, base_prompt: &str) -> String {
8117        format!(
8118            "{}\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>",
8119            base_prompt
8120        )
8121    }
8122
8123    async fn generate_plan(&self, input: &str) -> Result<Plan> {
8124        let effective = self.get_effective_reasoning_config();
8125        let planning_config = effective.get_planning();
8126
8127        let planner_llm = planning_config
8128            .and_then(|c| c.planner_llm.as_ref())
8129            .and_then(|alias| self.llm_registry.get(alias).ok())
8130            .or_else(|| self.llm_registry.router().ok())
8131            .or_else(|| self.llm_registry.default().ok())
8132            .ok_or_else(|| AgentError::Config("No LLM available for planning".into()))?;
8133
8134        let mut available_tool_ids: Vec<String> = self
8135            .get_available_tool_ids()
8136            .await
8137            .unwrap_or_else(|_| self.tools.list_ids());
8138        let mut available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
8139
8140        // Apply planning-level tool and skill filters.
8141        if let Some(config) = planning_config {
8142            if !config.available.tools.is_all() {
8143                available_tool_ids.retain(|t| config.available.tools.allows(t));
8144            }
8145            if !config.available.skills.is_all() {
8146                available_skills.retain(|s| config.available.skills.allows(s));
8147            }
8148        }
8149
8150        // Build tool descriptions with argument schemas so the planner
8151        // knows how to construct valid args for each step.
8152        let tool_descriptions: Vec<String> = available_tool_ids
8153            .iter()
8154            .filter_map(|id| {
8155                self.tools.get(id).map(|tool| {
8156                    let schema = tool.input_schema();
8157                    let args_desc = schema
8158                        .get("properties")
8159                        .and_then(|p| serde_json::to_string(p).ok())
8160                        .unwrap_or_else(|| "{}".to_string());
8161                    format!(
8162                        "- {} ({}): {}\n  Arguments: {}",
8163                        id,
8164                        tool.name(),
8165                        tool.description(),
8166                        args_desc
8167                    )
8168                })
8169            })
8170            .collect();
8171
8172        let tools_section = if tool_descriptions.is_empty() {
8173            "Available tools: none".to_string()
8174        } else {
8175            format!("Available tools:\n{}", tool_descriptions.join("\n"))
8176        };
8177
8178        let skills_section = if available_skills.is_empty() {
8179            "Available skills: none".to_string()
8180        } else {
8181            format!("Available skills: {}", available_skills.join(", "))
8182        };
8183
8184        let prompt = format!(
8185            r#"Create a step-by-step plan to accomplish this goal.
8186
8187Goal: "{}"
8188
8189{}
8190
8191{}
8192
8193Create a plan with clear steps. For each step, specify:
8194- description: What this step accomplishes
8195- action_type: "tool", "skill", "think", or "respond"
8196- action_target: The tool/skill id (if applicable)
8197- args: The arguments object matching the tool's schema (if action_type is "tool")
8198- dependencies: List of step IDs this depends on (empty if none)
8199
8200Respond in JSON format:
8201{{
8202  "steps": [
8203    {{"id": "step1", "description": "...", "action_type": "tool", "action_target": "tool_id", "args": {{"required_field": "value"}}, "dependencies": []}},
8204    {{"id": "step2", "description": "...", "action_type": "think", "action_target": "...", "dependencies": ["step1"]}}
8205  ]
8206}}"#,
8207            input, tools_section, skills_section,
8208        );
8209
8210        let messages = vec![ChatMessage::user(&prompt)];
8211        let response = self
8212            .observe_purpose(
8213                ObservationPurpose::PlanGeneration,
8214                planner_llm.complete(&messages, None),
8215            )
8216            .await
8217            .map_err(|e| AgentError::LLM(format!("Planning failed: {}", e)))?;
8218
8219        let mut plan = Plan::new(input);
8220
8221        if let Some(json_start) = response.content.find('{')
8222            && let Some(json_end) = response.content.rfind('}')
8223        {
8224            let json_str = &response.content[json_start..=json_end];
8225            if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(json_str)
8226                && let Some(steps) = parsed.get("steps").and_then(|s| s.as_array())
8227            {
8228                for step_value in steps {
8229                    let id = step_value
8230                        .get("id")
8231                        .and_then(|v| v.as_str())
8232                        .unwrap_or("step");
8233                    let desc = step_value
8234                        .get("description")
8235                        .and_then(|v| v.as_str())
8236                        .unwrap_or("");
8237                    let action_type = step_value
8238                        .get("action_type")
8239                        .and_then(|v| v.as_str())
8240                        .unwrap_or("think");
8241                    let action_target = step_value
8242                        .get("action_target")
8243                        .and_then(|v| v.as_str())
8244                        .unwrap_or("");
8245                    let args = step_value
8246                        .get("args")
8247                        .cloned()
8248                        .unwrap_or(serde_json::json!({}));
8249                    let deps: Vec<String> = step_value
8250                        .get("dependencies")
8251                        .and_then(|v| v.as_array())
8252                        .map(|arr| {
8253                            arr.iter()
8254                                .filter_map(|v| v.as_str().map(String::from))
8255                                .collect()
8256                        })
8257                        .unwrap_or_default();
8258
8259                    let action = match action_type {
8260                        "tool" => PlanAction::tool(action_target, args),
8261                        "skill" => PlanAction::skill(action_target),
8262                        "respond" => PlanAction::respond(action_target),
8263                        _ => PlanAction::think(desc),
8264                    };
8265
8266                    let step = PlanStep::new(desc, action)
8267                        .with_id(id)
8268                        .with_dependencies(deps);
8269                    plan.add_step(step);
8270                }
8271            }
8272        }
8273
8274        if plan.steps.is_empty() {
8275            plan.add_step(PlanStep::new(
8276                "Process the request",
8277                PlanAction::think(input),
8278            ));
8279            plan.add_step(PlanStep::new(
8280                "Provide response",
8281                PlanAction::respond("Answer based on analysis"),
8282            ));
8283        }
8284
8285        Ok(plan)
8286    }
8287
8288    async fn execute_plan(&self, plan: &mut Plan) -> Result<String> {
8289        let llm = self.get_state_llm()?;
8290        let mut results: HashMap<String, serde_json::Value> = HashMap::new();
8291        let effective = self.get_effective_reasoning_config();
8292        let max_steps = effective.get_planning().map(|c| c.max_steps).unwrap_or(10);
8293
8294        plan.status = PlanStatus::InProgress;
8295
8296        for step_idx in 0..plan.steps.len().min(max_steps as usize) {
8297            let step = &plan.steps[step_idx];
8298
8299            let deps_satisfied = step.dependencies.iter().all(|dep| {
8300                plan.steps
8301                    .iter()
8302                    .find(|s| &s.id == dep)
8303                    .map(|s| s.status.is_completed())
8304                    .unwrap_or(false)
8305            });
8306
8307            if !deps_satisfied {
8308                continue;
8309            }
8310
8311            plan.steps[step_idx].mark_running();
8312
8313            let result = match &plan.steps[step_idx].action {
8314                PlanAction::Tool { tool, args } => {
8315                    // When a tool step has dependency results, ask the LLM to
8316                    // produce the correct arguments given the context and tool schema.
8317                    // This avoids brittle {{stepN}} template substitution and lets
8318                    // the LLM handle type adaptation (e.g. picking the iso field
8319                    // from a datetime result for a downstream format call).
8320                    let has_dep_results = plan.steps[step_idx]
8321                        .dependencies
8322                        .iter()
8323                        .any(|dep| results.contains_key(dep));
8324
8325                    let final_args = if has_dep_results {
8326                        let dep_context: String = plan.steps[step_idx]
8327                            .dependencies
8328                            .iter()
8329                            .filter_map(|dep| results.get(dep).map(|r| format!("{}: {}", dep, r)))
8330                            .collect::<Vec<_>>()
8331                            .join("\n");
8332
8333                        let tool_schema = self
8334                            .tools
8335                            .get(tool)
8336                            .map(|t| {
8337                                let schema = t.input_schema();
8338                                let props = schema
8339                                    .get("properties")
8340                                    .and_then(|p| serde_json::to_string(p).ok())
8341                                    .unwrap_or_else(|| "{}".to_string());
8342                                format!(
8343                                    "{}: {}\nArguments schema: {}",
8344                                    t.id(),
8345                                    t.description(),
8346                                    props
8347                                )
8348                            })
8349                            .unwrap_or_default();
8350
8351                        let step_desc = &plan.steps[step_idx].description;
8352                        let arg_prompt = format!(
8353                            "Generate the JSON arguments for a tool call.\n\n\
8354                             Tool: {}\n\n\
8355                             Task: {}\n\n\
8356                             Previous step results:\n{}\n\n\
8357                             Planner's draft arguments: {}\n\n\
8358                             Produce ONLY a valid JSON object with the correct argument values.\n\
8359                             Use actual values from the previous step results, not template references.",
8360                            tool_schema,
8361                            step_desc,
8362                            dep_context,
8363                            serde_json::to_string(args).unwrap_or_default()
8364                        );
8365                        let messages = vec![ChatMessage::user(&arg_prompt)];
8366                        match self
8367                            .observe_purpose(
8368                                ObservationPurpose::PlanStep,
8369                                llm.complete(&messages, None),
8370                            )
8371                            .await
8372                        {
8373                            Ok(resp) => {
8374                                let content = resp.content.trim();
8375                                // Parse the LLM's JSON response, fall back to planner args.
8376                                let json_start = content.find('{');
8377                                let json_end = content.rfind('}');
8378                                if let (Some(start), Some(end)) = (json_start, json_end) {
8379                                    serde_json::from_str(&content[start..=end])
8380                                        .unwrap_or_else(|_| args.clone())
8381                                } else {
8382                                    args.clone()
8383                                }
8384                            }
8385                            Err(_) => args.clone(),
8386                        }
8387                    } else {
8388                        args.clone()
8389                    };
8390
8391                    let request = ToolExecutionRequest::new(
8392                        uuid::Uuid::new_v4().to_string(),
8393                        tool.clone(),
8394                        final_args,
8395                        ToolCallSource::Plan {
8396                            step_index: step_idx,
8397                        },
8398                    );
8399                    match self.execute_tool_record(request).await {
8400                        Ok(record) if record.success => {
8401                            serde_json::json!({ "output": record.model_output_string() })
8402                        }
8403                        Ok(record) => {
8404                            plan.steps[step_idx].mark_failed(record.model_output_string());
8405                            continue;
8406                        }
8407                        Err(e) => {
8408                            plan.steps[step_idx].mark_failed(e.to_string());
8409                            continue;
8410                        }
8411                    }
8412                }
8413                PlanAction::Skill { skill } => {
8414                    if let Some(skill_def) = self.skills.iter().find(|s| &s.id == skill) {
8415                        if let Some(ref executor) = self.skill_executor {
8416                            match executor
8417                                .execute_with_invoker(skill_def, "", serde_json::json!({}), self)
8418                                .await
8419                            {
8420                                Ok(output) => serde_json::json!({ "output": output }),
8421                                Err(e) => {
8422                                    plan.steps[step_idx].mark_failed(e.to_string());
8423                                    continue;
8424                                }
8425                            }
8426                        } else {
8427                            serde_json::json!({ "output": "Skill executor not available" })
8428                        }
8429                    } else {
8430                        plan.steps[step_idx].mark_failed("Skill not found");
8431                        continue;
8432                    }
8433                }
8434                PlanAction::Think { prompt } => {
8435                    let context: String = results
8436                        .iter()
8437                        .map(|(k, v)| format!("{}: {}", k, v))
8438                        .collect::<Vec<_>>()
8439                        .join("\n");
8440
8441                    let think_prompt = format!("Context:\n{}\n\nTask: {}", context, prompt);
8442                    let messages = vec![ChatMessage::user(&think_prompt)];
8443
8444                    match self
8445                        .observe_purpose(
8446                            ObservationPurpose::PlanStep,
8447                            llm.complete(&messages, None),
8448                        )
8449                        .await
8450                    {
8451                        Ok(resp) => serde_json::json!({ "output": resp.content }),
8452                        Err(e) => {
8453                            plan.steps[step_idx].mark_failed(e.to_string());
8454                            continue;
8455                        }
8456                    }
8457                }
8458                PlanAction::Respond { template } => {
8459                    let context: String = results
8460                        .iter()
8461                        .map(|(k, v)| format!("{}: {}", k, v))
8462                        .collect::<Vec<_>>()
8463                        .join("\n");
8464
8465                    let respond_prompt = format!(
8466                        "Based on this context:\n{}\n\nGenerate a response following this template/instruction: {}",
8467                        context, template
8468                    );
8469                    let messages = vec![ChatMessage::user(&respond_prompt)];
8470
8471                    match self
8472                        .observe_purpose(
8473                            ObservationPurpose::PlanStep,
8474                            llm.complete(&messages, None),
8475                        )
8476                        .await
8477                    {
8478                        Ok(resp) => serde_json::json!({ "output": resp.content }),
8479                        Err(e) => {
8480                            plan.steps[step_idx].mark_failed(e.to_string());
8481                            continue;
8482                        }
8483                    }
8484                }
8485            };
8486
8487            results.insert(plan.steps[step_idx].id.clone(), result.clone());
8488            plan.steps[step_idx].mark_completed(Some(result));
8489        }
8490
8491        // Set plan status based on whether any steps actually failed.
8492        let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
8493        if has_failures {
8494            let failed_ids: Vec<String> = plan
8495                .steps
8496                .iter()
8497                .filter(|s| s.status.is_failed())
8498                .map(|s| s.id.clone())
8499                .collect();
8500            plan.status = PlanStatus::Failed {
8501                error: format!("Steps failed: {}", failed_ids.join(", ")),
8502            };
8503        } else {
8504            plan.status = PlanStatus::Completed;
8505        }
8506
8507        // Synthesize final output from all completed step results.
8508        let all_outputs: Vec<String> = plan
8509            .steps
8510            .iter()
8511            .filter(|s| s.status.is_completed())
8512            .filter_map(|s| {
8513                s.result
8514                    .as_ref()
8515                    .and_then(|r| r.get("output"))
8516                    .and_then(|o| o.as_str())
8517                    .map(|o| format!("{}: {}", s.description, o))
8518            })
8519            .collect();
8520
8521        if all_outputs.is_empty() {
8522            return Ok("Plan execution completed but produced no results.".to_string());
8523        }
8524
8525        if all_outputs.len() == 1 {
8526            return Ok(all_outputs.into_iter().next().unwrap());
8527        }
8528
8529        // Synthesize a coherent summary from multiple step results via LLM.
8530        let context = all_outputs.join("\n\n");
8531        let prompt = format!(
8532            "You completed a multi-step plan for: \"{}\"\n\nStep results:\n{}\n\nProvide a coherent final response that synthesizes these results.",
8533            plan.goal, context
8534        );
8535        let messages = vec![ChatMessage::user(&prompt)];
8536        match self
8537            .observe_purpose(ObservationPurpose::PlanStep, llm.complete(&messages, None))
8538            .await
8539        {
8540            Ok(resp) => Ok(resp.content.trim().to_string()),
8541            Err(_) => Ok(context),
8542        }
8543    }
8544
8545    async fn evaluate_response(&self, input: &str, response: &str) -> Result<EvaluationResult> {
8546        let effective_config = self.get_effective_reflection_config();
8547        self.evaluate_response_with_config(input, response, &effective_config)
8548            .await
8549    }
8550
8551    fn extract_thinking(&self, content: &str) -> (Option<String>, String) {
8552        if let Some(start) = content.find("<thinking>")
8553            && let Some(end) = content.find("</thinking>")
8554        {
8555            let thinking = content[start + 10..end].trim().to_string();
8556            let answer = content[end + 11..].trim().to_string();
8557            return (Some(thinking), answer);
8558        }
8559        (None, content.to_string())
8560    }
8561
8562    fn format_response_with_thinking(&self, thinking: Option<&str>, answer: &str) -> String {
8563        match self.get_effective_reasoning_config().output {
8564            ReasoningOutput::Hidden => answer.to_string(),
8565            ReasoningOutput::Visible => {
8566                if let Some(t) = thinking {
8567                    format!("Thinking:\n{}\n\nAnswer:\n{}", t, answer)
8568                } else {
8569                    answer.to_string()
8570                }
8571            }
8572            ReasoningOutput::Tagged => {
8573                if let Some(t) = thinking {
8574                    format!("<thinking>{}</thinking>\n{}", t, answer)
8575                } else {
8576                    answer.to_string()
8577                }
8578            }
8579        }
8580    }
8581
8582    //
8583    // Blocking disambiguation returns every clarification or required confirmation as its own root response before redispatch.
8584    // The manager retains resolved input and pending skill ownership until a later turn explicitly confirms it.
8585    //
8586    async fn run_loop(&self, input: &str) -> Result<AgentResponse> {
8587        //
8588        // Blocking execution must fail before a turn starts when required persistence is unavailable.
8589        //
8590        self.init_storage().await?;
8591        self.begin_root_turn();
8592        let _root_cleanup = RootTurnCleanup::new(self);
8593        info!(input_len = input.len(), "Starting chat");
8594
8595        self.hooks.on_message_received(input).await;
8596
8597        // One-shot context initialization: load runtime defaults, resolve env vars,
8598        // populate builtin sources (session, agent), etc.  This must happen before
8599        // the first template render so that {{ context.* }} variables are available.
8600        if !self.context_initialized.swap(true, Ordering::SeqCst) {
8601            self.context_manager.initialize().await?;
8602            debug!("Context manager initialized (defaults, env, builtins)");
8603        }
8604
8605        self.check_turn_timeout().await?;
8606        self.context_manager.refresh_per_turn().await?;
8607
8608        // Clear stale disambiguation context from previous turns.
8609        // This prevents resolved_intent from leaking across turns and causing incorrect deterministic routing on subsequent inputs.
8610        self.clear_disambiguation_context();
8611
8612        // Disambiguation check (before input processing)
8613        if let Some(ref disambiguator) = self.disambiguation_manager {
8614            let disambiguation_context = self.build_disambiguation_context().await?;
8615
8616            // Get state-level disambiguation override
8617            let state_override = self
8618                .state_machine
8619                .as_ref()
8620                .and_then(|sm| sm.current_definition())
8621                .and_then(|def| def.disambiguation.clone());
8622
8623            let state_generation = self
8624                .state_machine
8625                .as_ref()
8626                .map(|state_machine| state_machine.generation());
8627            let disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
8628            let mut disambiguation_result = self
8629                .observe_purpose(
8630                    ObservationPurpose::DisambiguationDetection,
8631                    disambiguator.process_input_with_override(
8632                        input,
8633                        &disambiguation_context,
8634                        state_override.as_ref(),
8635                        None,
8636                    ),
8637                )
8638                .await?;
8639            let current_state_generation = self
8640                .state_machine
8641                .as_ref()
8642                .map(|state_machine| state_machine.generation());
8643            if current_state_generation != state_generation
8644                || self.disambiguation_epoch.load(Ordering::SeqCst) != disambiguation_epoch
8645            {
8646                disambiguator.clear_pending().await;
8647                *self.pending_skill_id.write() = None;
8648                disambiguation_result = DisambiguationResult::Abandoned { new_input: None };
8649                info!(
8650                    confirmation_event = "invalidated",
8651                    invalidation_reason = "state_generation_changed",
8652                    "Disambiguation result invalidated before redispatch"
8653                );
8654            }
8655            match disambiguation_result {
8656                DisambiguationResult::Clear => {
8657                    debug!("Input is clear, proceeding normally");
8658                }
8659                DisambiguationResult::NeedsClarification {
8660                    question,
8661                    detection,
8662                } => {
8663                    let admission = self
8664                        .admit_disambiguation_redispatch(disambiguation_epoch, state_generation)
8665                        .await?;
8666                    let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
8667                    info!(
8668                        ambiguity_type = ?detection.ambiguity_type,
8669                        confidence = detection.confidence,
8670                        "Input requires clarification"
8671                    );
8672
8673                    // This branch also owns post-resolution confirmation questions.
8674                    // Do not clear pending_skill_id or redispatch until the manager returns Clarified on a later turn.
8675                    self.commit_root_user_message(input).await?;
8676                    self.memory
8677                        .add_message(ChatMessage::assistant(&question.question))
8678                        .await?;
8679
8680                    let status = if awaiting_confirmation {
8681                        "awaiting_confirmation"
8682                    } else {
8683                        "awaiting_clarification"
8684                    };
8685                    let response = AgentResponse::new(&question.question).with_metadata(
8686                        "disambiguation",
8687                        serde_json::json!({
8688                            "status": status,
8689                            "options": question.options,
8690                            "clarifying": question.clarifying,
8691                            "detection": {
8692                                "type": detection.ambiguity_type,
8693                                "confidence": detection.confidence,
8694                                "what_is_unclear": detection.what_is_unclear,
8695                            }
8696                        }),
8697                    );
8698                    drop(admission);
8699                    self.finish_turn_if_root(&response).await?;
8700                    return Ok(response);
8701                }
8702                DisambiguationResult::Clarified {
8703                    enriched_input,
8704                    resolved,
8705                    ..
8706                } => {
8707                    let admission = match self
8708                        .admit_disambiguation_redispatch(disambiguation_epoch, state_generation)
8709                        .await
8710                    {
8711                        Ok(admission) => admission,
8712                        Err(error) => {
8713                            *self.pending_skill_id.write() = None;
8714                            return Err(error);
8715                        }
8716                    };
8717                    info!(
8718                        resolved_count = resolved.len(),
8719                        enriched = %enriched_input,
8720                        "Input clarified, injecting resolved intent into context"
8721                    );
8722
8723                    // Routing uses `resolved` (structured, deterministic)
8724                    // This is what makes post-disambiguation routing DETERMINISTIC
8725                    for (key, value) in &resolved {
8726                        let context_key = format!("disambiguation.{}", key);
8727                        let _ = self.context_manager.set(&context_key, value.clone());
8728                    }
8729
8730                    if let Some(intent) = resolved.get("intent") {
8731                        let _ = self.context_manager.set("resolved_intent", intent.clone());
8732                    }
8733
8734                    let _ = self
8735                        .context_manager
8736                        .set("disambiguation.resolved", serde_json::Value::Bool(true));
8737
8738                    // Check if this clarification was triggered by a skill-level override.
8739                    // If so, route directly to the matched skill instead of going through
8740                    // skill routing again (which might match a different skill).
8741                    let skill_id = self.pending_skill_id.read().clone();
8742                    if let Some(skill_id) = skill_id {
8743                        info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input");
8744                        drop(admission);
8745                        return self
8746                            .recheck_skill_disambiguation(
8747                                &skill_id,
8748                                &enriched_input,
8749                                disambiguation_epoch,
8750                                state_generation,
8751                            )
8752                            .await;
8753                    }
8754
8755                    drop(admission);
8756                    return self.run_loop_internal(&enriched_input).await;
8757                }
8758                DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
8759                    info!("Proceeding with best guess interpretation");
8760
8761                    // Same skill-id re-check for best-guess path
8762                    let skill_id = self.pending_skill_id.read().clone();
8763                    if let Some(skill_id) = skill_id {
8764                        info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input");
8765                        return self
8766                            .recheck_skill_disambiguation(
8767                                &skill_id,
8768                                &enriched_input,
8769                                disambiguation_epoch,
8770                                state_generation,
8771                            )
8772                            .await;
8773                    }
8774
8775                    return self.run_loop_internal(&enriched_input).await;
8776                }
8777                DisambiguationResult::GiveUp { reason } => {
8778                    *self.pending_skill_id.write() = None;
8779                    warn!(reason = %reason, "Disambiguation gave up");
8780                    let apology = self
8781                        .generate_localized_apology(
8782                            "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
8783                            &reason,
8784                        )
8785                        .await
8786                        .unwrap_or_else(|_| {
8787                            format!("I'm sorry, I couldn't understand your request: {}", reason)
8788                        });
8789                    let response = AgentResponse::new(&apology);
8790                    self.finish_turn_if_root(&response).await?;
8791                    return Ok(response);
8792                }
8793                DisambiguationResult::Escalate { reason } => {
8794                    *self.pending_skill_id.write() = None;
8795                    info!(reason = %reason, "Escalating to human");
8796                    if let Some(ref hitl) = self.hitl_engine {
8797                        let trigger =
8798                            ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
8799                        let mut context_map = HashMap::new();
8800                        context_map.insert("original_input".to_string(), serde_json::json!(input));
8801                        context_map.insert("reason".to_string(), serde_json::json!(&reason));
8802                        let check_result = HITLCheckResult::required(
8803                            trigger,
8804                            context_map,
8805                            format!("User request needs human assistance: {}", reason),
8806                            Some(hitl.config().default_timeout_seconds),
8807                        );
8808                        let result = self.request_hitl_approval(check_result).await?;
8809                        if matches!(
8810                            result,
8811                            ApprovalResult::Approved | ApprovalResult::Modified { .. }
8812                        ) {
8813                            return self.run_loop_internal(input).await;
8814                        }
8815                    }
8816                    let apology = self
8817                        .generate_localized_apology(
8818                            "Explain briefly that you're transferring the user to a human agent for help.",
8819                            &reason,
8820                        )
8821                        .await
8822                        .unwrap_or_else(|_| {
8823                            format!("I need human assistance to help with your request: {}", reason)
8824                        });
8825                    let response = AgentResponse::new(&apology);
8826                    self.finish_turn_if_root(&response).await?;
8827                    return Ok(response);
8828                }
8829                DisambiguationResult::Abandoned { new_input } => {
8830                    *self.pending_skill_id.write() = None;
8831
8832                    info!(
8833                        has_new_input = new_input.is_some(),
8834                        "Clarification abandoned by user"
8835                    );
8836
8837                    self.commit_root_user_message(input).await?;
8838
8839                    match new_input {
8840                        Some(fresh_input) => {
8841                            // Topic switch: process the user's new input from scratch.
8842                            // The LLM sees full conversation context including the abandoned exchange.
8843                            return self.run_loop_internal(&fresh_input).await;
8844                        }
8845                        None => {
8846                            // Pure abandonment: generate a brief acknowledgment.
8847                            let ack = self
8848                                .generate_localized_apology(
8849                                    "The user changed their mind about their previous request. \
8850                                     Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
8851                                     Do NOT apologize excessively. Be concise.",
8852                                    "User abandoned clarification",
8853                                )
8854                                .await
8855                                .unwrap_or_else(|_| {
8856                                    "OK, no problem. What else can I help with?".to_string()
8857                                });
8858
8859                            self.memory
8860                                .add_message(ChatMessage::assistant(&ack))
8861                                .await?;
8862
8863                            let response = AgentResponse::new(&ack);
8864                            self.finish_turn_if_root(&response).await?;
8865                            return Ok(response);
8866                        }
8867                    }
8868                }
8869            }
8870        }
8871
8872        self.run_loop_internal(input).await
8873    }
8874
8875    /// Generate a localized response using the router LLM
8876    async fn generate_localized_apology(&self, instruction: &str, reason: &str) -> Result<String> {
8877        let llm = self.llm_registry.router().map_err(|e| {
8878            AgentError::LLM(format!(
8879                "Router LLM not available for localized response: {}",
8880                e
8881            ))
8882        })?;
8883
8884        let recent: Vec<String> = self
8885            .memory
8886            .get_messages(Some(3))
8887            .await?
8888            .iter()
8889            .map(|m| m.content.clone())
8890            .collect();
8891
8892        let context_hint = if recent.is_empty() {
8893            String::new()
8894        } else {
8895            format!(
8896                "\nRecent conversation (detect the user's language from this):\n{}\n",
8897                recent.join("\n")
8898            )
8899        };
8900
8901        let prompt = format!(
8902            "{}\nReason: {}\n{}Respond in the same language as the user. Output ONLY the message, nothing else.",
8903            instruction, reason, context_hint
8904        );
8905
8906        let messages = vec![ChatMessage::user(&prompt)];
8907        let response = self
8908            .observe_purpose(
8909                ObservationPurpose::DisambiguationClarification,
8910                llm.complete(&messages, None),
8911            )
8912            .await
8913            .map_err(|e| AgentError::LLM(format!("Localized response generation failed: {}", e)))?;
8914
8915        Ok(response.content.trim().to_string())
8916    }
8917
8918    /// Clear disambiguation-related keys from the context manager.
8919    ///
8920    /// Render template variables in state action args using the context manager.
8921    fn render_action_args(&self, args: &Value) -> Value {
8922        let context = self.build_context_with_overlays();
8923        match args {
8924            Value::Object(map) => {
8925                let mut rendered = serde_json::Map::new();
8926                for (k, v) in map {
8927                    match v {
8928                        Value::String(s) if s.contains("{{") => {
8929                            match self.template_renderer.render(s, &context) {
8930                                Ok(rendered_str) => {
8931                                    rendered.insert(k.clone(), Value::String(rendered_str));
8932                                }
8933                                Err(_) => {
8934                                    rendered.insert(k.clone(), v.clone());
8935                                }
8936                            }
8937                        }
8938                        _ => {
8939                            rendered.insert(k.clone(), v.clone());
8940                        }
8941                    }
8942                }
8943                Value::Object(rendered)
8944            }
8945            _ => args.clone(),
8946        }
8947    }
8948
8949    /// Called at the start of each turn to prevent stale `resolved_intent` from leaking across turns.
8950    fn clear_disambiguation_context(&self) {
8951        let _ = self
8952            .context_manager
8953            .set("resolved_intent", serde_json::Value::Null);
8954
8955        let all = self.context_manager.get_all();
8956        for key in all.keys() {
8957            if key.starts_with("disambiguation.") {
8958                let _ = self.context_manager.set(key, serde_json::Value::Null);
8959            }
8960        }
8961    }
8962
8963    /// Re-run skill disambiguation on enriched input before executing the skill.
8964    /// After clarification resolves, the enriched input may still be missing required_clarity fields (e.g. "Transfer money to Jane." still lacks amount).
8965    /// This method re-runs the skill's disambiguation pass.
8966    /// If fields are still missing, it returns the new clarification question and keeps pending_skill_id set.
8967    /// If all fields are present (Clear), it executes the skill and returns the response.
8968    async fn recheck_skill_disambiguation(
8969        &self,
8970        skill_id: &str,
8971        enriched_input: &str,
8972        expected_disambiguation_epoch: u64,
8973        expected_state_generation: Option<u64>,
8974    ) -> Result<AgentResponse> {
8975        let skill = self
8976            .skill_router
8977            .as_ref()
8978            .and_then(|r| r.get_skill(skill_id).cloned());
8979
8980        // If the skill has disambiguation enabled, re-run it on the enriched input.
8981        if let Some(ref skill) = skill
8982            && let Some(ref skill_disambig) = skill.disambiguation
8983            && skill_disambig.enabled.unwrap_or(false)
8984            && let Some(ref disambiguator) = self.disambiguation_manager
8985        {
8986            let context = self.build_disambiguation_context().await?;
8987            let state_override = self
8988                .state_machine
8989                .as_ref()
8990                .and_then(|sm| sm.current_definition())
8991                .and_then(|def| def.disambiguation.clone());
8992
8993            let disambiguation_result = self
8994                .observe_purpose(
8995                    ObservationPurpose::DisambiguationDetection,
8996                    disambiguator.process_input_with_override(
8997                        enriched_input,
8998                        &context,
8999                        state_override.as_ref(),
9000                        Some(skill_disambig),
9001                    ),
9002                )
9003                .await?;
9004            let current_state_generation = self
9005                .state_machine
9006                .as_ref()
9007                .map(|state_machine| state_machine.generation());
9008            if current_state_generation != expected_state_generation
9009                || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
9010            {
9011                disambiguator.clear_pending().await;
9012                *self.pending_skill_id.write() = None;
9013                return Err(AgentError::Other(
9014                    "State or reset ownership changed during skill disambiguation recheck"
9015                        .to_string(),
9016                ));
9017            }
9018            match disambiguation_result {
9019                DisambiguationResult::Clear => {
9020                    debug!(skill_id = %skill_id, "Skill re-check: all fields present");
9021                }
9022                DisambiguationResult::NeedsClarification {
9023                    question,
9024                    detection,
9025                } => {
9026                    let admission = self
9027                        .admit_disambiguation_redispatch(
9028                            expected_disambiguation_epoch,
9029                            expected_state_generation,
9030                        )
9031                        .await?;
9032                    let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
9033                    info!(
9034                        skill_id = %skill_id,
9035                        ambiguity_type = ?detection.ambiguity_type,
9036                        what_is_unclear = ?detection.what_is_unclear,
9037                        "Skill re-check: still missing fields, asking again"
9038                    );
9039                    // Keep pending_skill_id set (do NOT clear it).
9040                    // The next turn will resolve this new clarification and
9041                    // re-enter this method until all fields are present.
9042                    self.memory
9043                        .add_message(ChatMessage::user(enriched_input))
9044                        .await?;
9045                    self.memory
9046                        .add_message(ChatMessage::assistant(&question.question))
9047                        .await?;
9048
9049                    let response = AgentResponse::new(&question.question).with_metadata(
9050                        "disambiguation",
9051                        serde_json::json!({
9052                            "status": if awaiting_confirmation { "awaiting_confirmation" } else { "awaiting_clarification" },
9053                            "skill_id": skill_id,
9054                            "options": question.options,
9055                            "clarifying": question.clarifying,
9056                            "detection": {
9057                                "type": detection.ambiguity_type,
9058                                "confidence": detection.confidence,
9059                                "what_is_unclear": detection.what_is_unclear,
9060                            }
9061                        }),
9062                    );
9063                    drop(admission);
9064                    self.finish_turn_if_root(&response).await?;
9065                    return Ok(response);
9066                }
9067                DisambiguationResult::Clarified {
9068                    enriched_input: re_enriched,
9069                    ..
9070                } => {
9071                    debug!(skill_id = %skill_id, "Skill re-check: clarified immediately, executing");
9072                    let admission = self
9073                        .admit_disambiguation_redispatch(
9074                            expected_disambiguation_epoch,
9075                            expected_state_generation,
9076                        )
9077                        .await?;
9078                    *self.pending_skill_id.write() = None;
9079                    drop(admission);
9080                    let skill_response = self.execute_skill_by_id(skill_id, &re_enriched).await?;
9081                    self.memory
9082                        .add_message(ChatMessage::user(&re_enriched))
9083                        .await?;
9084                    return self
9085                        .handle_skill_response(
9086                            &re_enriched,
9087                            skill_id,
9088                            skill_response,
9089                            &HashMap::new(),
9090                        )
9091                        .await;
9092                }
9093                DisambiguationResult::ProceedWithBestGuess {
9094                    enriched_input: re_enriched,
9095                } => {
9096                    debug!(skill_id = %skill_id, "Skill re-check: proceeding with best guess");
9097                    let admission = self
9098                        .admit_disambiguation_redispatch(
9099                            expected_disambiguation_epoch,
9100                            expected_state_generation,
9101                        )
9102                        .await?;
9103                    *self.pending_skill_id.write() = None;
9104                    drop(admission);
9105                    let skill_response = self.execute_skill_by_id(skill_id, &re_enriched).await?;
9106                    self.memory
9107                        .add_message(ChatMessage::user(&re_enriched))
9108                        .await?;
9109                    return self
9110                        .handle_skill_response(
9111                            &re_enriched,
9112                            skill_id,
9113                            skill_response,
9114                            &HashMap::new(),
9115                        )
9116                        .await;
9117                }
9118                DisambiguationResult::GiveUp { reason } => {
9119                    *self.pending_skill_id.write() = None;
9120                    let apology = self
9121                                    .generate_localized_apology(
9122                                        "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
9123                                        &reason,
9124                                    )
9125                                    .await
9126                                    .unwrap_or_else(|_| {
9127                                        format!("I'm sorry, I couldn't understand your request: {}", reason)
9128                                    });
9129                    let response = AgentResponse::new(&apology);
9130                    self.finish_turn_if_root(&response).await?;
9131                    return Ok(response);
9132                }
9133                DisambiguationResult::Escalate { reason } => {
9134                    *self.pending_skill_id.write() = None;
9135                    let apology = self
9136                                    .generate_localized_apology(
9137                                        "Explain briefly that you're transferring the user to a human agent for help.",
9138                                        &reason,
9139                                    )
9140                                    .await
9141                                    .unwrap_or_else(|_| {
9142                                        format!("I need human assistance to help with your request: {}", reason)
9143                                    });
9144                    let response = AgentResponse::new(&apology);
9145                    self.finish_turn_if_root(&response).await?;
9146                    return Ok(response);
9147                }
9148                DisambiguationResult::Abandoned { new_input } => {
9149                    // User abandoned during skill re-check.
9150                    // Clear skill routing state and fall through to normal execution.
9151                    *self.pending_skill_id.write() = None;
9152                    debug!(skill_id = %skill_id, "Skill re-check: abandoned by user");
9153                    if let Some(fresh) = new_input {
9154                        return self.run_loop_internal(&fresh).await;
9155                    }
9156                    let ack = self
9157                                    .generate_localized_apology(
9158                                        "The user changed their mind about their previous request. \
9159                                         Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
9160                                         Do NOT apologize excessively. Be concise.",
9161                                        "User abandoned clarification",
9162                                    )
9163                                    .await
9164                                    .unwrap_or_else(|_| {
9165                                        "OK, no problem. What else can I help with?".to_string()
9166                                    });
9167                    self.memory
9168                        .add_message(ChatMessage::assistant(&ack))
9169                        .await?;
9170                    let response = AgentResponse::new(&ack);
9171                    self.finish_turn_if_root(&response).await?;
9172                    return Ok(response);
9173                }
9174            }
9175        }
9176
9177        // Skill execution is admitted before the read guard is released so later invalidation cannot retroactively cancel it.
9178        let admission = self
9179            .admit_disambiguation_redispatch(
9180                expected_disambiguation_epoch,
9181                expected_state_generation,
9182            )
9183            .await?;
9184        *self.pending_skill_id.write() = None;
9185        drop(admission);
9186        let skill_response = self.execute_skill_by_id(skill_id, enriched_input).await?;
9187        self.memory
9188            .add_message(ChatMessage::user(enriched_input))
9189            .await?;
9190        self.handle_skill_response(enriched_input, skill_id, skill_response, &HashMap::new())
9191            .await
9192    }
9193
9194    /// Handle skill routing result: output processing, memory, transitions.
9195    /// Returns a fully formed AgentResponse for skill-routed requests.
9196    async fn handle_skill_response(
9197        &self,
9198        processed_input: &str,
9199        skill_id: &str,
9200        skill_response: String,
9201        input_context: &HashMap<String, Value>,
9202    ) -> Result<AgentResponse> {
9203        let output_data = self.process_output(&skill_response, input_context).await?;
9204        let final_response = output_data.content;
9205
9206        self.memory
9207            .add_message(ChatMessage::assistant(&final_response))
9208            .await?;
9209
9210        self.check_memory_compression().await?;
9211
9212        self.increment_turn();
9213        self.evaluate_transitions(processed_input, &final_response)
9214            .await?;
9215
9216        let response = AgentResponse::new(final_response)
9217            .with_metadata("skill_id", serde_json::json!(skill_id));
9218        self.finish_turn_if_root(&response).await?;
9219        Ok(response)
9220    }
9221
9222    /// Run the Plan-and-Execute flow: generate plan, execute steps, finalize.
9223    /// Supports plan-level reflection with replan loop when configured.
9224    async fn handle_plan_and_execute(
9225        &self,
9226        processed_input: &str,
9227        input_context: &HashMap<String, Value>,
9228        auto_detected: bool,
9229    ) -> Result<AgentResponse> {
9230        let effective = self.get_effective_reasoning_config();
9231        let plan_reflection = effective
9232            .get_planning()
9233            .map(|c| c.reflection.clone())
9234            .unwrap_or_default();
9235
9236        let max_attempts = if plan_reflection.enabled {
9237            1 + plan_reflection.max_replans
9238        } else {
9239            1
9240        };
9241
9242        let mut plan = self.generate_plan(processed_input).await?;
9243        info!(
9244            plan_id = %plan.id,
9245            steps = plan.steps.len(),
9246            "Plan generated"
9247        );
9248
9249        let mut plan_result = String::new();
9250
9251        for attempt in 0..max_attempts {
9252            *self.current_plan.write() = Some(plan.clone());
9253            plan_result = self.execute_plan(&mut plan).await?;
9254
9255            info!(
9256                plan_status = ?plan.status,
9257                completed_steps = plan.completed_steps().count(),
9258                attempt = attempt + 1,
9259                "Plan execution completed"
9260            );
9261
9262            if !plan_reflection.enabled {
9263                break;
9264            }
9265
9266            let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
9267            if !has_failures {
9268                break;
9269            }
9270
9271            if attempt + 1 >= max_attempts {
9272                break;
9273            }
9274
9275            match plan_reflection.on_step_failure {
9276                StepFailureAction::Replan => {
9277                    info!(attempt = attempt + 1, "Plan had failures, replanning");
9278                    plan = self.generate_plan(processed_input).await?;
9279                }
9280                StepFailureAction::Abort => {
9281                    warn!("Plan step failed, aborting");
9282                    break;
9283                }
9284                StepFailureAction::Skip | StepFailureAction::Continue => {
9285                    break;
9286                }
9287            }
9288        }
9289
9290        *self.current_plan.write() = Some(plan);
9291
9292        let output_data = self.process_output(&plan_result, input_context).await?;
9293        let final_content = output_data.content;
9294
9295        self.memory
9296            .add_message(ChatMessage::assistant(&final_content))
9297            .await?;
9298
9299        self.check_memory_compression().await?;
9300        self.increment_turn();
9301        self.evaluate_transitions(processed_input, &final_content)
9302            .await?;
9303
9304        let reasoning_metadata =
9305            ReasoningMetadata::new(ReasoningMode::PlanAndExecute).with_auto_detected(auto_detected);
9306
9307        let response = AgentResponse::new(&final_content).with_metadata(
9308            "reasoning",
9309            serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
9310        );
9311
9312        self.finish_turn_if_root(&response).await?;
9313        Ok(response)
9314    }
9315
9316    /// Inject CoT/ReAct reasoning prompt into the system message (first iteration only).
9317    fn inject_reasoning_prompt(
9318        &self,
9319        messages: &mut [ChatMessage],
9320        reasoning_mode: &ReasoningMode,
9321        is_first_iteration: bool,
9322    ) {
9323        if !is_first_iteration {
9324            return;
9325        }
9326        match reasoning_mode {
9327            ReasoningMode::CoT => {
9328                if let Some(msg) = messages.first_mut()
9329                    && matches!(msg.role, ai_agents_core::Role::System)
9330                {
9331                    msg.content = self.build_cot_system_prompt(&msg.content);
9332                    debug!("Applied Chain-of-Thought system prompt");
9333                }
9334            }
9335            ReasoningMode::React => {
9336                if let Some(msg) = messages.first_mut()
9337                    && matches!(msg.role, ai_agents_core::Role::System)
9338                {
9339                    msg.content = self.build_react_system_prompt(&msg.content);
9340                    debug!("Applied ReAct system prompt");
9341                }
9342            }
9343            _ => {}
9344        }
9345    }
9346
9347    //
9348    // Draft generation must not commit user memory or run tools.
9349    // The current user input is added only as an ephemeral message for this LLM call.
9350    //
9351    async fn generate_main_response_draft(
9352        &self,
9353        processed_input: &str,
9354        reasoning_mode: &ReasoningMode,
9355    ) -> Result<MainResponseDraft> {
9356        let llm = self.get_state_llm()?;
9357        let protocol = self.main_tool_protocol(llm.as_ref(), true).await?;
9358        let mut messages = self
9359            .build_messages_internal(false, Some(processed_input), protocol.choice.is_none())
9360            .await?;
9361        self.inject_reasoning_prompt(&mut messages, reasoning_mode, true);
9362        let response = self
9363            .complete_main_llm_with_recovery(llm, &messages, &protocol)
9364            .await?;
9365        let content = response.content.trim().to_string();
9366        let (thinking, answer) = self.extract_thinking(&content);
9367        if let Some(calls) = self.parse_main_tool_calls(&content, &protocol) {
9368            return Ok(MainResponseDraft::ToolCalls {
9369                raw_content: content,
9370                calls,
9371                thinking,
9372            });
9373        }
9374        Ok(MainResponseDraft::Text {
9375            raw_content: answer,
9376            thinking,
9377        })
9378    }
9379
9380    //
9381    // This is the only place where a winning draft is allowed to become runtime state.
9382    // Prompt and native tool calls remain inert until this method commits the draft into the shared executor path.
9383    //
9384    async fn commit_main_response_draft(
9385        &self,
9386        processed_input: &str,
9387        input_context: &HashMap<String, Value>,
9388        draft: MainResponseDraft,
9389        reasoning_mode: ReasoningMode,
9390        auto_detected: bool,
9391    ) -> Result<AgentResponse> {
9392        self.commit_root_user_message(processed_input).await?;
9393        match draft {
9394            MainResponseDraft::Text {
9395                raw_content,
9396                thinking,
9397            } => {
9398                self.finish_text_response_from_model(CommittedTextResponse {
9399                    processed_input,
9400                    input_context,
9401                    answer: raw_content,
9402                    reasoning_mode,
9403                    auto_detected,
9404                    iterations: 1,
9405                    thinking_content: thinking,
9406                    all_tool_calls: Vec::new(),
9407                })
9408                .await
9409            }
9410            MainResponseDraft::ToolCalls {
9411                raw_content,
9412                calls,
9413                thinking: _,
9414            } => {
9415                let mut all_tool_calls = Vec::new();
9416                match self
9417                    .handle_tool_calls(processed_input, &raw_content, calls, &mut all_tool_calls)
9418                    .await?
9419                {
9420                    ToolCallOutcome::Rejected(response) => {
9421                        self.finish_turn_if_root(&response).await?;
9422                        Ok(response)
9423                    }
9424                    ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => {
9425                        self.continue_after_committed_tool_draft(processed_input)
9426                            .await
9427                    }
9428                }
9429            }
9430        }
9431    }
9432
9433    //
9434    // Tool drafts need a committed continuation after function results are written.
9435    // Redispatch depth suppresses duplicate root lifecycle work during that continuation.
9436    //
9437    async fn continue_after_committed_tool_draft(
9438        &self,
9439        processed_input: &str,
9440    ) -> Result<AgentResponse> {
9441        *self.redispatch_depth.write() += 1;
9442        if let Some(context) = self.active_turn_context.write().as_mut() {
9443            context.enter_redispatch();
9444        }
9445        let result = Box::pin(self.run_loop_internal(processed_input)).await;
9446        *self.redispatch_depth.write() -= 1;
9447        if let Some(context) = self.active_turn_context.write().as_mut() {
9448            context.exit_redispatch();
9449        }
9450        let response = result?;
9451        self.finish_turn_if_root(&response).await?;
9452        Ok(response)
9453    }
9454
9455    //
9456    // Shared committed text finalization for normal responses and winning text drafts.
9457    // Keep output processing, reflection, transitions, hooks, and maintenance behind this commit boundary.
9458    //
9459    async fn finish_text_response_from_model(
9460        &self,
9461        response: CommittedTextResponse<'_>,
9462    ) -> Result<AgentResponse> {
9463        let CommittedTextResponse {
9464            processed_input,
9465            input_context,
9466            answer,
9467            reasoning_mode,
9468            auto_detected,
9469            iterations,
9470            thinking_content,
9471            all_tool_calls,
9472        } = response;
9473        let output_data = self.process_output(&answer, input_context).await?;
9474        let mut final_content = if output_data.metadata.rejected {
9475            output_data
9476                .metadata
9477                .rejection_reason
9478                .unwrap_or_else(|| answer.to_string())
9479        } else {
9480            output_data.content
9481        };
9482        let llm = self.get_state_llm()?;
9483        let reflection_metadata;
9484        (final_content, reflection_metadata) = self
9485            .run_reflection(&*llm, processed_input, final_content)
9486            .await?;
9487        final_content =
9488            self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
9489        let final_content = {
9490            let result = self
9491                .post_loop_processing(processed_input, final_content)
9492                .await?;
9493            self.apply_post_loop_result(processed_input, result).await?
9494        };
9495        let response = self.build_agent_response(AgentResponseParts {
9496            content: final_content,
9497            all_tool_calls,
9498            reasoning_mode,
9499            auto_detected,
9500            iterations,
9501            thinking: thinking_content,
9502            reflection_metadata,
9503        });
9504        self.finish_turn_if_root(&response).await?;
9505        Ok(response)
9506    }
9507
9508    //
9509    // Auto reasoning uses this path after the judge wins with a deeper mode.
9510    // It intentionally uses committed message building instead of the draft overlay.
9511    //
9512    async fn run_committed_response_loop_with_reasoning(
9513        &self,
9514        processed_input: &str,
9515        input_context: &HashMap<String, Value>,
9516        reasoning_mode: ReasoningMode,
9517        auto_detected: bool,
9518    ) -> Result<AgentResponse> {
9519        self.commit_root_user_message(processed_input).await?;
9520        let llm = self.get_state_llm()?;
9521        let mut iterations = 0u32;
9522        let mut all_tool_calls = Vec::new();
9523        let mut thinking_content = None;
9524        loop {
9525            let effective_max = if reasoning_mode != ReasoningMode::None {
9526                let rc = self.get_effective_reasoning_config();
9527                self.max_iterations.min(rc.max_iterations)
9528            } else {
9529                self.max_iterations
9530            };
9531            if iterations >= effective_max {
9532                return Err(AgentError::Other(format!(
9533                    "Max iterations ({}) exceeded",
9534                    effective_max
9535                )));
9536            }
9537            iterations += 1;
9538            *self.iteration_count.write() = iterations;
9539            let protocol = self.main_tool_protocol(llm.as_ref(), false).await?;
9540            let mut messages = self
9541                .build_messages_internal(true, None, protocol.choice.is_none())
9542                .await?;
9543            self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
9544            self.hooks.on_llm_start(&messages).await;
9545            let llm_start = Instant::now();
9546            let response = self
9547                .complete_main_llm_with_recovery(Arc::clone(&llm), &messages, &protocol)
9548                .await?;
9549            let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
9550            self.hooks.on_llm_complete(&response, llm_duration_ms).await;
9551            let content = response.content.trim();
9552            if let Some(tool_calls) = self.parse_main_tool_calls(content, &protocol) {
9553                match self
9554                    .handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
9555                    .await?
9556                {
9557                    ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
9558                    ToolCallOutcome::Rejected(resp) => {
9559                        self.finish_turn_if_root(&resp).await?;
9560                        return Ok(resp);
9561                    }
9562                }
9563            }
9564            let (extracted_thinking, answer) = self.extract_thinking(content);
9565            if extracted_thinking.is_some() {
9566                thinking_content = extracted_thinking;
9567            }
9568            return self
9569                .finish_text_response_from_model(CommittedTextResponse {
9570                    processed_input,
9571                    input_context,
9572                    answer,
9573                    reasoning_mode,
9574                    auto_detected,
9575                    iterations,
9576                    thinking_content,
9577                    all_tool_calls,
9578                })
9579                .await;
9580        }
9581    }
9582
9583    /// Handle tool calls: check transitions, execute tools in parallel, handle HITL rejection.
9584    async fn handle_tool_calls(
9585        &self,
9586        processed_input: &str,
9587        content: &str,
9588        tool_calls: Vec<ToolCall>,
9589        all_tool_calls: &mut Vec<ToolCall>,
9590    ) -> Result<ToolCallOutcome> {
9591        // Check if a transition should fire before executing the LLM's tool call.
9592        // If a transition fires, on_enter actions handle the tool call correctly
9593        // (with proper URLs from YAML), so skip the LLM's tool call.
9594        let transition_fired = self.evaluate_transitions(processed_input, content).await?;
9595        if transition_fired {
9596            self.memory
9597                .add_message(ChatMessage::assistant(
9598                    "(Transitioned to new state — tool call handled by workflow)",
9599                ))
9600                .await?;
9601            return Ok(ToolCallOutcome::TransitionFired);
9602        }
9603
9604        // 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.
9605        self.memory
9606            .add_message(ChatMessage::assistant(content))
9607            .await?;
9608        let native_tool_call = Self::is_native_tool_call_content(content);
9609
9610        let results = self.execute_tools_parallel(&tool_calls).await;
9611
9612        for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
9613            match result {
9614                Ok(output) => {
9615                    self.memory
9616                        .add_message(Self::tool_result_message(
9617                            tool_call,
9618                            &output,
9619                            native_tool_call,
9620                        ))
9621                        .await?;
9622                }
9623                Err(e) => {
9624                    // Check if this is a HITL rejection - if so, break the loop
9625                    if matches!(e, AgentError::HITLRejected(_)) {
9626                        self.memory
9627                            .add_message(ChatMessage::assistant(format!(
9628                                "The operation was rejected by the approver: {}",
9629                                e
9630                            )))
9631                            .await?;
9632                        // Return the rejection message to user, don't continue loop
9633                        return Ok(ToolCallOutcome::Rejected(AgentResponse {
9634                            content: format!("Operation cancelled: {}", e),
9635                            metadata: None,
9636                            tool_calls: Some(all_tool_calls.clone()),
9637                        }));
9638                    }
9639                    self.memory
9640                        .add_message(Self::tool_result_message(
9641                            tool_call,
9642                            &format!("Error: {}", e),
9643                            native_tool_call,
9644                        ))
9645                        .await?;
9646                }
9647            }
9648            all_tool_calls.push(tool_call.clone());
9649        }
9650        Ok(ToolCallOutcome::Continue)
9651    }
9652
9653    /// Run the reflection loop on a response, returning (improved_content, reflection_metadata).
9654    async fn run_reflection(
9655        &self,
9656        llm: &dyn LLMProvider,
9657        processed_input: &str,
9658        mut content: String,
9659    ) -> Result<(String, Option<ReflectionMetadata>)> {
9660        let should_reflect = self.should_reflect(processed_input, &content).await?;
9661        if !should_reflect {
9662            return Ok((content, None));
9663        }
9664
9665        info!("Starting response reflection evaluation");
9666        let mut attempts = 0u32;
9667        let max_retries = self.reflection_config.max_retries;
9668        let mut history: Vec<ReflectionAttempt> = Vec::new();
9669
9670        loop {
9671            let evaluation = self.evaluate_response(processed_input, &content).await?;
9672
9673            if evaluation.passed || attempts >= max_retries {
9674                info!(
9675                    passed = evaluation.passed,
9676                    confidence = evaluation.confidence,
9677                    attempts = attempts + 1,
9678                    "Reflection evaluation complete"
9679                );
9680                let reflection_metadata = Some(
9681                    ReflectionMetadata::new(evaluation)
9682                        .with_attempts(attempts + 1)
9683                        .with_history(history),
9684                );
9685                return Ok((content, reflection_metadata));
9686            }
9687
9688            debug!(
9689                attempt = attempts + 1,
9690                failed_criteria = evaluation.failed_criteria().count(),
9691                "Response did not meet criteria, retrying"
9692            );
9693
9694            history.push(
9695                ReflectionAttempt::new(&content, evaluation.clone())
9696                    .with_feedback("Response did not meet quality criteria"),
9697            );
9698
9699            let feedback: Vec<String> = evaluation
9700                .failed_criteria()
9701                .map(|c| format!("- {}", c.criterion))
9702                .collect();
9703
9704            let retry_prompt = format!(
9705                "Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response.",
9706                feedback.join("\n")
9707            );
9708
9709            self.memory
9710                .add_message(ChatMessage::user(&retry_prompt))
9711                .await?;
9712
9713            let retry_messages = self.build_messages().await?;
9714            let retry_response = self
9715                .observe_purpose(
9716                    ObservationPurpose::ReflectionEvaluation,
9717                    llm.complete(&retry_messages, None),
9718                )
9719                .await
9720                .map_err(|e| AgentError::LLM(e.to_string()))?;
9721
9722            content = retry_response.content.trim().to_string();
9723            attempts += 1;
9724        }
9725    }
9726
9727    /// Record the assistant turn, evaluate transitions, and decide what to do next.
9728    /// Returns PostLoopResult so callers can apply_post_loop_result for re-dispatch.
9729    async fn post_loop_processing(
9730        &self,
9731        processed_input: &str,
9732        content: String,
9733    ) -> Result<PostLoopResult> {
9734        // Do NOT add the assistant message to memory yet.
9735        // evaluate_transitions receives content as a direct parameter, so the message does not need to be in memory for transitions to evaluate correctly.
9736        // For NeedsRedispatch we skip adding the stale old-state response entirely, keeping memory clean for the re-dispatched handler.
9737
9738        self.increment_turn();
9739
9740        // Run context extractors so guards can check freshly-extracted values.
9741        self.run_context_extractors(processed_input).await;
9742
9743        let transitioned = self.evaluate_transitions(processed_input, &content).await?;
9744
9745        if !transitioned {
9746            self.memory
9747                .add_message(ChatMessage::assistant(&content))
9748                .await?;
9749            self.check_memory_compression().await?;
9750            return Ok(PostLoopResult::NoTransition(content));
9751        }
9752
9753        // Check if we should skip re-generation after this transition.
9754        if !self.should_regenerate_after_transition() {
9755            self.memory
9756                .add_message(ChatMessage::assistant(&content))
9757                .await?;
9758            self.check_memory_compression().await?;
9759            return Ok(PostLoopResult::Transitioned(content));
9760        }
9761
9762        // Check if the new state needs full dispatch.
9763        // Orchestration states (concurrent, group_chat, pipeline, handoff, delegate) need their dedicated handlers.
9764        // 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.
9765        if self.needs_redispatch_for_new_state() {
9766            info!("Post-transition NeedsRedispatch: new state requires full dispatch");
9767            // Stale old-state content is NOT added to memory.
9768            // apply_post_loop_result will increment redispatch_depth and call run_loop_internal, which produces the correct response and adds it.
9769            return Ok(PostLoopResult::NeedsRedispatch);
9770        }
9771
9772        // Normal post-transition re-generation (plain LLM with optional tool calls).
9773        // Add the stale response to history so the new LLM sees the conversation.
9774        self.memory
9775            .add_message(ChatMessage::assistant(&content))
9776            .await?;
9777        self.check_memory_compression().await?;
9778
9779        // If a transition fired, on_enter actions already executed (e.g., HTTP calls).
9780        // The current content was generated in the OLD state context and is stale.
9781        // Re-generate in the new state context so the LLM can reference on_enter results.
9782        // If the LLM responds with a tool call, execute it in a mini-loop so the
9783        // result is not returned as raw JSON text.
9784        let new_llm = self.get_state_llm()?;
9785        let mut final_content;
9786
9787        for post_iter in 0..self.max_iterations {
9788            let protocol = self.main_tool_protocol(new_llm.as_ref(), false).await?;
9789            let new_messages = self
9790                .build_messages_internal(true, None, protocol.choice.is_none())
9791                .await?;
9792            if post_iter == 0
9793                && let Some(system_msg) = new_messages.first()
9794                && system_msg.role == ai_agents_core::Role::System
9795            {
9796                debug!(
9797                    prompt_preview =
9798                        &system_msg.content[system_msg.content.len().saturating_sub(200)..],
9799                    "Post-transition system prompt (last 200 chars)"
9800                );
9801            }
9802
9803            let new_response = self
9804                .complete_main_llm_with_recovery(Arc::clone(&new_llm), &new_messages, &protocol)
9805                .await?;
9806            final_content = new_response.content.trim().to_string();
9807
9808            // Check if the post-transition response contains tool calls.
9809            // If so, execute them and loop so the LLM can summarize the result.
9810            if let Some(tool_calls) = self.parse_main_tool_calls(&final_content, &protocol) {
9811                let native_tool_call = Self::is_native_tool_call_content(&final_content);
9812                debug!(
9813                    post_iter = post_iter,
9814                    tools = tool_calls.len(),
9815                    "Post-transition tool call detected, executing"
9816                );
9817
9818                self.memory
9819                    .add_message(ChatMessage::assistant(&final_content))
9820                    .await?;
9821
9822                let results = self.execute_tools_parallel(&tool_calls).await;
9823                for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
9824                    match result {
9825                        Ok(output) => {
9826                            self.memory
9827                                .add_message(Self::tool_result_message(
9828                                    tool_call,
9829                                    &output,
9830                                    native_tool_call,
9831                                ))
9832                                .await?;
9833                        }
9834                        Err(e) => {
9835                            self.memory
9836                                .add_message(Self::tool_result_message(
9837                                    tool_call,
9838                                    &format!("Error: {}", e),
9839                                    native_tool_call,
9840                                ))
9841                                .await?;
9842                        }
9843                    }
9844                }
9845                // Loop to let the LLM see the tool result and produce a text response.
9846                continue;
9847            }
9848
9849            // No tool call - this is the final text response.
9850            self.memory
9851                .add_message(ChatMessage::assistant(&final_content))
9852                .await?;
9853            return Ok(PostLoopResult::Transitioned(final_content));
9854        }
9855
9856        // Exhausted post-transition iterations (unlikely). Return last content.
9857        final_content = "Post-transition processing completed.".to_string();
9858        self.memory
9859            .add_message(ChatMessage::assistant(&final_content))
9860            .await?;
9861
9862        Ok(PostLoopResult::Transitioned(final_content))
9863    }
9864
9865    /// Build the final AgentResponse with all metadata.
9866    /// Check whether to re-generate a response after a state transition.
9867    fn should_regenerate_after_transition(&self) -> bool {
9868        if let Some(ref sm) = self.state_machine {
9869            // Global setting
9870            if !sm.config().regenerate_on_transition {
9871                return false;
9872            }
9873            // Per-state override on the new (current) state
9874            if let Some(def) = sm.current_definition()
9875                && let Some(regen) = def.regenerate_on_enter
9876            {
9877                return regen;
9878            }
9879        }
9880        true
9881    }
9882
9883    /// Return true when the new state requires full dispatch via run_loop_internal.
9884    /// Covers orchestration states and any non-None effective reasoning mode.
9885    fn needs_redispatch_for_new_state(&self) -> bool {
9886        if let Some(ref sm) = self.state_machine
9887            && let Some(def) = sm.current_definition()
9888        {
9889            if def.concurrent.is_some()
9890                || def.group_chat.is_some()
9891                || def.pipeline.is_some()
9892                || def.handoff.is_some()
9893                || def.delegate.is_some()
9894            {
9895                return true;
9896            }
9897            // Any non-None effective reasoning mode requires the main dispatch loop.
9898            let effective = self.get_effective_reasoning_config();
9899            if !matches!(effective.mode, ReasoningMode::None) {
9900                return true;
9901            }
9902        }
9903        false
9904    }
9905
9906    /// Consume a PostLoopResult. NeedsRedispatch re-enters run_loop_internal.
9907    /// The user message is already in memory - redispatch_depth suppresses re-adding it.
9908    async fn apply_post_loop_result(
9909        &self,
9910        processed_input: &str,
9911        result: PostLoopResult,
9912    ) -> Result<String> {
9913        match result {
9914            PostLoopResult::NoTransition(content) | PostLoopResult::Transitioned(content) => {
9915                Ok(content)
9916            }
9917            PostLoopResult::NeedsRedispatch => {
9918                const MAX_REDISPATCH_DEPTH: u32 = 3;
9919                let current_depth = *self.redispatch_depth.read();
9920                if current_depth >= MAX_REDISPATCH_DEPTH {
9921                    warn!(
9922                        depth = current_depth,
9923                        "Post-transition re-dispatch depth limit reached, returning empty response"
9924                    );
9925                    let content = String::new();
9926                    self.memory
9927                        .add_message(ChatMessage::assistant(&content))
9928                        .await?;
9929                    return Ok(content);
9930                }
9931                *self.redispatch_depth.write() += 1;
9932                if let Some(context) = self.active_turn_context.write().as_mut() {
9933                    context.enter_redispatch();
9934                }
9935                info!(
9936                    depth = current_depth + 1,
9937                    "Re-dispatching for new state after transition"
9938                );
9939                let resp = Box::pin(self.run_loop_internal(processed_input)).await;
9940                *self.redispatch_depth.write() -= 1;
9941                if let Some(context) = self.active_turn_context.write().as_mut() {
9942                    context.exit_redispatch();
9943                }
9944                resp.map(|r| r.content)
9945            }
9946        }
9947    }
9948
9949    /// Builds the final response metadata after committed output and transition handling complete.
9950    fn build_agent_response(&self, parts: AgentResponseParts) -> AgentResponse {
9951        let AgentResponseParts {
9952            content,
9953            all_tool_calls,
9954            reasoning_mode,
9955            auto_detected,
9956            iterations,
9957            thinking,
9958            reflection_metadata,
9959        } = parts;
9960        let reasoning_metadata = ReasoningMetadata::new(reasoning_mode.clone())
9961            .with_thinking(thinking.clone().unwrap_or_default())
9962            .with_iterations(iterations)
9963            .with_auto_detected(auto_detected);
9964
9965        let mut response = AgentResponse::new(&content);
9966        if !all_tool_calls.is_empty() {
9967            response = response.with_tool_calls(all_tool_calls);
9968        }
9969
9970        if let Some(state) = self.current_state() {
9971            response = response.with_metadata("current_state", serde_json::json!(state));
9972        }
9973
9974        response = response.with_metadata(
9975            "reasoning",
9976            serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
9977        );
9978
9979        if let Some(ref refl_meta) = reflection_metadata {
9980            response = response.with_metadata(
9981                "reflection",
9982                serde_json::to_value(refl_meta).unwrap_or_default(),
9983            );
9984        }
9985
9986        response
9987    }
9988
9989    // Handle delegation: forward user input to a registry agent.
9990    async fn handle_delegated_state(
9991        &self,
9992        input: &str,
9993        delegate_id: &str,
9994        state_def: &ai_agents_state::StateDefinition,
9995    ) -> Result<AgentResponse> {
9996        use std::time::Instant;
9997
9998        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
9999            AgentError::Config(format!(
10000                "State delegates to '{}' but no agent registry is configured. \
10001                 Add a spawner section with auto_spawn to your YAML.",
10002                delegate_id
10003            ))
10004        })?;
10005
10006        let state_name = self
10007            .state_machine
10008            .as_ref()
10009            .map(|sm| sm.current())
10010            .unwrap_or_else(|| "unknown".to_string());
10011
10012        self.hooks.on_delegate_start(delegate_id, &state_name).await;
10013        let start = Instant::now();
10014
10015        let delegate = registry.get(delegate_id).ok_or_else(|| {
10016            AgentError::Other(format!(
10017                "State '{}' delegates to '{}' but no agent with that ID exists in the registry.",
10018                state_name, delegate_id
10019            ))
10020        })?;
10021
10022        // Prepare input based on delegate_context mode.
10023        let context_mode = state_def.delegate_context.clone().unwrap_or_default();
10024        let effective_input = self
10025            .observe_purpose(
10026                ObservationPurpose::OrchestrationRouting,
10027                crate::orchestration::context::prepare_delegate_input(
10028                    input,
10029                    &context_mode,
10030                    &*self.memory,
10031                    self.llm_registry.get("router").ok().as_deref(),
10032                ),
10033            )
10034            .await?;
10035
10036        let response = delegate
10037            .chat_with_actor_context(&effective_input, self.outbound_actor_context())
10038            .await?;
10039
10040        let duration_ms = start.elapsed().as_millis() as u64;
10041        self.hooks
10042            .on_delegate_complete(delegate_id, &state_name, duration_ms)
10043            .await;
10044
10045        // Backward-compatible context key.
10046        let ctx_key = format!("delegation.{}.last_response", delegate_id);
10047        let _ = self.context_manager.set(
10048            &ctx_key,
10049            serde_json::Value::String(response.content.clone()),
10050        );
10051
10052        // Structured orchestration context.
10053        let _ = self.context_manager.set(
10054            "orchestration",
10055            serde_json::json!({
10056                "type": "delegate",
10057                "agent": delegate_id,
10058                "state": state_name,
10059                "response": response.content,
10060                "duration_ms": duration_ms,
10061            }),
10062        );
10063
10064        self.commit_root_user_message(input).await?;
10065
10066        // post_loop_processing records the assistant turn and evaluates transitions.
10067        // apply_post_loop_result handles NeedsRedispatch by re-entering run_loop_internal.
10068        let post_result = self
10069            .post_loop_processing(
10070                input,
10071                format!("[Delegated to {}]: {}", delegate_id, response.content),
10072            )
10073            .await?;
10074        let final_content = self.apply_post_loop_result(input, post_result).await?;
10075
10076        let mut result = AgentResponse::new(final_content);
10077
10078        let metadata = serde_json::json!({
10079            "orchestration": {
10080                "type": "delegate",
10081                "agent": delegate_id,
10082                "state": state_name,
10083                "response": response.content,
10084                "duration_ms": duration_ms,
10085            }
10086        });
10087        result.metadata = Some(
10088            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10089                metadata,
10090            )
10091            .unwrap_or_default(),
10092        );
10093
10094        self.finish_turn_if_root(&result).await?;
10095        Ok(result)
10096    }
10097
10098    // Handle concurrent execution: run multiple registry agents in parallel and aggregate.
10099    async fn handle_concurrent_state(
10100        &self,
10101        input: &str,
10102        config: &ai_agents_state::ConcurrentStateConfig,
10103    ) -> Result<AgentResponse> {
10104        use std::time::Instant;
10105
10106        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10107            AgentError::Config(
10108                "Concurrent state requires an agent registry. Add a spawner section.".into(),
10109            )
10110        })?;
10111
10112        // Render input template if provided, otherwise use the raw input.
10113        // Uses direct minijinja rendering (same approach as pipeline) so variables
10114        // are top-level: {{ user_input }}, not {{ context.user_input }}.
10115        // Enrich input with parent conversation history when context_mode is set.
10116        let context_mode = config.context_mode.clone().unwrap_or_default();
10117        let context_input = self
10118            .observe_purpose(
10119                ObservationPurpose::OrchestrationRouting,
10120                crate::orchestration::context::prepare_delegate_input(
10121                    input,
10122                    &context_mode,
10123                    &*self.memory,
10124                    self.llm_registry.get("router").ok().as_deref(),
10125                ),
10126            )
10127            .await?;
10128
10129        let effective_input = if let Some(ref tmpl) = config.input {
10130            render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10131                .unwrap_or_else(|_| context_input.clone())
10132        } else {
10133            context_input
10134        };
10135
10136        let start = Instant::now();
10137
10138        let llm_name = config
10139            .aggregation
10140            .synthesizer_llm
10141            .as_deref()
10142            .unwrap_or("router");
10143        let llm_provider = self.llm_registry.get(llm_name).ok();
10144
10145        let vote_parallelism = if self.runtime_config.optimization.enabled
10146            && self
10147                .runtime_config
10148                .optimization
10149                .parallel_orchestration_vote_extraction
10150        {
10151            Some(self.runtime_config.optimization.max_parallel_runtime_tasks)
10152        } else {
10153            None
10154        };
10155
10156        let result = self
10157            .observe_purpose(
10158                ObservationPurpose::OrchestrationAggregation,
10159                scope_actor_context(
10160                    self.outbound_actor_context(),
10161                    crate::orchestration::concurrent(
10162                        registry,
10163                        &effective_input,
10164                        &config.agents,
10165                        &config.aggregation,
10166                        llm_provider.as_deref(),
10167                        config.min_required,
10168                        config.timeout_ms,
10169                        config.on_partial_failure.clone(),
10170                        vote_parallelism,
10171                    ),
10172                ),
10173            )
10174            .await?;
10175
10176        let duration_ms = start.elapsed().as_millis() as u64;
10177        let agent_ids: Vec<String> = config.agents.iter().map(|a| a.id().to_string()).collect();
10178        let strategy = format!("{:?}", config.aggregation.strategy);
10179        self.hooks
10180            .on_concurrent_complete(&agent_ids, &strategy, duration_ms)
10181            .await;
10182
10183        // Backward-compatible context key.
10184        let _ = self.context_manager.set(
10185            "concurrent.result",
10186            serde_json::Value::String(result.response.content.clone()),
10187        );
10188
10189        // Build per-agent result data for context and metadata.
10190        let agents_json: Vec<serde_json::Value> = result
10191            .agent_results
10192            .iter()
10193            .map(|ar| {
10194                serde_json::json!({
10195                    "id": ar.agent_id,
10196                    "response": ar.response.as_ref().map(|r| r.content.as_str()),
10197                    "success": ar.success,
10198                    "error": ar.error,
10199                    "duration_ms": ar.duration_ms,
10200                })
10201            })
10202            .collect();
10203
10204        // Structured orchestration context with per-agent results.
10205        let _ = self.context_manager.set(
10206            "orchestration",
10207            serde_json::json!({
10208                "type": "concurrent",
10209                "result": result.response.content,
10210                "strategy": strategy,
10211                "agents": agents_json,
10212                "duration_ms": duration_ms,
10213            }),
10214        );
10215
10216        self.commit_root_user_message(input).await?;
10217
10218        let post_result = self
10219            .post_loop_processing(input, result.response.content.clone())
10220            .await?;
10221        let final_content = self.apply_post_loop_result(input, post_result).await?;
10222
10223        let mut response = AgentResponse::new(final_content);
10224        let metadata = serde_json::json!({
10225            "orchestration": {
10226                "type": "concurrent",
10227                "result": result.response.content,
10228                "strategy": strategy,
10229                "agents": agents_json,
10230                "duration_ms": duration_ms,
10231            }
10232        });
10233        response.metadata = Some(
10234            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10235                metadata,
10236            )
10237            .unwrap_or_default(),
10238        );
10239
10240        self.finish_turn_if_root(&response).await?;
10241        Ok(response)
10242    }
10243
10244    // Handle group chat: run a multi-turn multi-agent conversation.
10245    async fn handle_group_chat_state(
10246        &self,
10247        input: &str,
10248        config: &ai_agents_state::GroupChatStateConfig,
10249    ) -> Result<AgentResponse> {
10250        use std::time::Instant;
10251
10252        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10253            AgentError::Config(
10254                "Group chat state requires an agent registry. Add a spawner section.".into(),
10255            )
10256        })?;
10257
10258        let start = Instant::now();
10259
10260        let llm_provider = self.llm_registry.get("router").ok();
10261
10262        // Enrich input with parent conversation history when context_mode is set.
10263        let context_mode = config.context_mode.clone().unwrap_or_default();
10264        let context_input = self
10265            .observe_purpose(
10266                ObservationPurpose::OrchestrationRouting,
10267                crate::orchestration::context::prepare_delegate_input(
10268                    input,
10269                    &context_mode,
10270                    &*self.memory,
10271                    self.llm_registry.get("router").ok().as_deref(),
10272                ),
10273            )
10274            .await?;
10275
10276        // Render input template if provided, otherwise use the raw user message as topic.
10277        let effective_topic = if let Some(ref tmpl) = config.input {
10278            render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10279                .unwrap_or_else(|_| context_input.clone())
10280        } else {
10281            context_input
10282        };
10283
10284        let result = self
10285            .observe_purpose(
10286                ObservationPurpose::OrchestrationConversation,
10287                scope_actor_context(
10288                    self.outbound_actor_context(),
10289                    crate::orchestration::group_chat(
10290                        registry,
10291                        &effective_topic,
10292                        config,
10293                        llm_provider.as_deref(),
10294                        Some(&*self.hooks),
10295                    ),
10296                ),
10297            )
10298            .await?;
10299
10300        let duration_ms = start.elapsed().as_millis() as u64;
10301
10302        // Backward-compatible context key.
10303        let _ = self.context_manager.set(
10304            "group_chat.conclusion",
10305            serde_json::Value::String(result.response.content.clone()),
10306        );
10307
10308        // Build transcript data for context and metadata.
10309        let transcript_json: Vec<serde_json::Value> = result
10310            .transcript
10311            .iter()
10312            .map(|t| {
10313                serde_json::json!({
10314                    "speaker": t.speaker,
10315                    "round": t.round,
10316                    "content": t.content,
10317                })
10318            })
10319            .collect();
10320
10321        // Structured orchestration context with full transcript.
10322        let _ = self.context_manager.set(
10323            "orchestration",
10324            serde_json::json!({
10325                "type": "group_chat",
10326                "conclusion": result.response.content,
10327                "transcript": transcript_json,
10328                "rounds": result.rounds_completed,
10329                "termination": result.termination_reason,
10330                "duration_ms": duration_ms,
10331            }),
10332        );
10333
10334        self.commit_root_user_message(input).await?;
10335
10336        let post_result = self
10337            .post_loop_processing(input, result.response.content.clone())
10338            .await?;
10339        let final_content = self.apply_post_loop_result(input, post_result).await?;
10340
10341        let mut response = AgentResponse::new(final_content);
10342        let metadata = serde_json::json!({
10343            "orchestration": {
10344                "type": "group_chat",
10345                "conclusion": result.response.content,
10346                "transcript": transcript_json,
10347                "rounds": result.rounds_completed,
10348                "termination": result.termination_reason,
10349                "duration_ms": duration_ms,
10350            }
10351        });
10352        response.metadata = Some(
10353            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10354                metadata,
10355            )
10356            .unwrap_or_default(),
10357        );
10358
10359        self.finish_turn_if_root(&response).await?;
10360        Ok(response)
10361    }
10362
10363    // Handle pipeline: run agents sequentially with per-stage input templates.
10364    async fn handle_pipeline_state(
10365        &self,
10366        input: &str,
10367        config: &ai_agents_state::PipelineStateConfig,
10368    ) -> Result<AgentResponse> {
10369        use std::time::Instant;
10370
10371        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10372            AgentError::Config(
10373                "Pipeline state requires an agent registry. Add a spawner section.".into(),
10374            )
10375        })?;
10376
10377        let start = Instant::now();
10378
10379        let stages: Vec<crate::orchestration::PipelineStage> = config
10380            .stages
10381            .iter()
10382            .map(|entry| {
10383                let mut stage = crate::orchestration::PipelineStage::id(entry.id());
10384                if let Some(tmpl) = entry.input() {
10385                    stage = stage.with_input(tmpl);
10386                }
10387                stage
10388            })
10389            .collect();
10390
10391        // Enrich input with parent conversation history when context_mode is set.
10392        let context_mode = config.context_mode.clone().unwrap_or_default();
10393        let context_input = self
10394            .observe_purpose(
10395                ObservationPurpose::OrchestrationRouting,
10396                crate::orchestration::context::prepare_delegate_input(
10397                    input,
10398                    &context_mode,
10399                    &*self.memory,
10400                    self.llm_registry.get("router").ok().as_deref(),
10401                ),
10402            )
10403            .await?;
10404
10405        let context_values = self.build_context_with_overlays();
10406        let result = self
10407            .observe_purpose(
10408                ObservationPurpose::OrchestrationRouting,
10409                scope_actor_context(
10410                    self.outbound_actor_context(),
10411                    crate::orchestration::pipeline(
10412                        registry,
10413                        &context_input,
10414                        &stages,
10415                        config.timeout_ms,
10416                        Some(&*self.hooks),
10417                        Some(&context_values),
10418                    ),
10419                ),
10420            )
10421            .await?;
10422
10423        let duration_ms = start.elapsed().as_millis() as u64;
10424
10425        // Backward-compatible context key.
10426        let _ = self.context_manager.set(
10427            "pipeline.result",
10428            serde_json::Value::String(result.response.content.clone()),
10429        );
10430
10431        // Build per-stage data for context and metadata.
10432        let stages_json: Vec<serde_json::Value> = result
10433            .stage_outputs
10434            .iter()
10435            .map(|s| {
10436                serde_json::json!({
10437                    "agent_id": s.agent_id,
10438                    "output": s.output,
10439                    "duration_ms": s.duration_ms,
10440                    "skipped": s.skipped,
10441                })
10442            })
10443            .collect();
10444
10445        // Structured orchestration context.
10446        let _ = self.context_manager.set(
10447            "orchestration",
10448            serde_json::json!({
10449                "type": "pipeline",
10450                "result": result.response.content,
10451                "stages": stages_json,
10452                "duration_ms": duration_ms,
10453            }),
10454        );
10455
10456        self.commit_root_user_message(input).await?;
10457
10458        let post_result = self
10459            .post_loop_processing(input, result.response.content.clone())
10460            .await?;
10461        let final_content = self.apply_post_loop_result(input, post_result).await?;
10462
10463        let mut response = AgentResponse::new(final_content);
10464        let metadata = serde_json::json!({
10465            "orchestration": {
10466                "type": "pipeline",
10467                "result": result.response.content,
10468                "stages": stages_json,
10469                "duration_ms": duration_ms,
10470            }
10471        });
10472        response.metadata = Some(
10473            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10474                metadata,
10475            )
10476            .unwrap_or_default(),
10477        );
10478
10479        self.finish_turn_if_root(&response).await?;
10480        Ok(response)
10481    }
10482
10483    // Handle handoff: LLM-directed agent-to-agent control transfer.
10484    async fn handle_handoff_state(
10485        &self,
10486        input: &str,
10487        config: &ai_agents_state::HandoffStateConfig,
10488    ) -> Result<AgentResponse> {
10489        use std::time::Instant;
10490
10491        let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10492            AgentError::Config(
10493                "Handoff state requires an agent registry. Add a spawner section.".into(),
10494            )
10495        })?;
10496
10497        let llm = self
10498            .llm_registry
10499            .get("router")
10500            .map_err(|_| AgentError::Config("Handoff state requires a router LLM.".into()))?;
10501
10502        let start = Instant::now();
10503
10504        // Enrich input with parent conversation history when context_mode is set.
10505        let context_mode = config.context_mode.clone().unwrap_or_default();
10506        let context_input = self
10507            .observe_purpose(
10508                ObservationPurpose::OrchestrationRouting,
10509                crate::orchestration::context::prepare_delegate_input(
10510                    input,
10511                    &context_mode,
10512                    &*self.memory,
10513                    self.llm_registry.get("router").ok().as_deref(),
10514                ),
10515            )
10516            .await?;
10517
10518        // Render input template if provided, otherwise forward the raw user message.
10519        let effective_input = if let Some(ref tmpl) = config.input {
10520            render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10521                .unwrap_or_else(|_| context_input.clone())
10522        } else {
10523            context_input
10524        };
10525
10526        let result = self
10527            .observe_purpose(
10528                ObservationPurpose::OrchestrationRouting,
10529                scope_actor_context(
10530                    self.outbound_actor_context(),
10531                    crate::orchestration::handoff(
10532                        registry,
10533                        &effective_input,
10534                        &config.initial_agent,
10535                        &config.available_agents,
10536                        config.max_handoffs,
10537                        llm.as_ref(),
10538                        Some(&*self.hooks),
10539                    ),
10540                ),
10541            )
10542            .await?;
10543
10544        let duration_ms = start.elapsed().as_millis() as u64;
10545
10546        // Backward-compatible context key.
10547        let _ = self.context_manager.set(
10548            "handoff.result",
10549            serde_json::Value::String(result.response.content.clone()),
10550        );
10551
10552        // Build handoff chain data for context and metadata.
10553        let chain_json: Vec<serde_json::Value> = result
10554            .handoff_chain
10555            .iter()
10556            .map(|h| {
10557                serde_json::json!({
10558                    "from": h.from_agent,
10559                    "to": h.to_agent,
10560                    "reason": h.reason,
10561                })
10562            })
10563            .collect();
10564
10565        // Structured orchestration context.
10566        let _ = self.context_manager.set(
10567            "orchestration",
10568            serde_json::json!({
10569                "type": "handoff",
10570                "result": result.response.content,
10571                "final_agent": result.final_agent,
10572                "handoff_chain": chain_json,
10573                "duration_ms": duration_ms,
10574            }),
10575        );
10576
10577        self.commit_root_user_message(input).await?;
10578
10579        let post_result = self
10580            .post_loop_processing(input, result.response.content.clone())
10581            .await?;
10582        let final_content = self.apply_post_loop_result(input, post_result).await?;
10583
10584        let mut response = AgentResponse::new(final_content);
10585        let metadata = serde_json::json!({
10586            "orchestration": {
10587                "type": "handoff",
10588                "result": result.response.content,
10589                "final_agent": result.final_agent,
10590                "handoff_chain": chain_json,
10591                "duration_ms": duration_ms,
10592            }
10593        });
10594        response.metadata = Some(
10595            serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10596                metadata,
10597            )
10598            .unwrap_or_default(),
10599        );
10600
10601        self.finish_turn_if_root(&response).await?;
10602        Ok(response)
10603    }
10604
10605    // run_loop_internal: blocking (non-streaming) agent pipeline.
10606    async fn run_loop_internal(&self, input: &str) -> Result<AgentResponse> {
10607        self.begin_root_turn();
10608        // Resolve actor_id from context, reload facts if actor changed, bump counter.
10609        self.pre_turn_session_lifecycle().await;
10610
10611        let input_data = self.process_input(input).await?;
10612        self.update_active_turn_context(&input_data.content, input_data.context.clone());
10613
10614        // Inject process context (detect/extract results) into agent context
10615        // so system prompt templates can use {{ context.detected_language }} etc.
10616        for (key, value) in &input_data.context {
10617            let _ = self.context_manager.set(key, value.clone());
10618        }
10619
10620        if input_data.metadata.rejected {
10621            let reason = input_data
10622                .metadata
10623                .rejection_reason
10624                .unwrap_or_else(|| "Input rejected".to_string());
10625            warn!(reason = %reason, "Input rejected");
10626            let response = AgentResponse::new(reason);
10627            self.finish_turn_if_root(&response).await?;
10628            return Ok(response);
10629        }
10630
10631        let processed_input = &input_data.content;
10632
10633        if let Some(response) = self.try_pre_response_transition(processed_input).await? {
10634            return Ok(response);
10635        }
10636
10637        // Handle orchestration states (delegate, concurrent, group_chat, pipeline, handoff).
10638        if let Some(ref sm) = self.state_machine
10639            && let Some(def) = sm.current_definition()
10640        {
10641            if let Some(ref delegate_id) = def.delegate {
10642                return self
10643                    .handle_delegated_state(processed_input, delegate_id, &def)
10644                    .await;
10645            }
10646            if let Some(ref concurrent_config) = def.concurrent {
10647                return self
10648                    .handle_concurrent_state(processed_input, concurrent_config)
10649                    .await;
10650            }
10651            if let Some(ref group_chat_config) = def.group_chat {
10652                return self
10653                    .handle_group_chat_state(processed_input, group_chat_config)
10654                    .await;
10655            }
10656            if let Some(ref pipeline_config) = def.pipeline {
10657                return self
10658                    .handle_pipeline_state(processed_input, pipeline_config)
10659                    .await;
10660            }
10661            if let Some(ref handoff_config) = def.handoff {
10662                return self
10663                    .handle_handoff_state(processed_input, handoff_config)
10664                    .await;
10665            }
10666        }
10667
10668        //
10669        // The speculative future is boxed to keep the runtime future size manageable.
10670        // Removing the box can overflow small test stacks because this function is recursive through redispatch.
10671        //
10672        if let Some(response) =
10673            Box::pin(self.try_speculative_branches(processed_input, &input_data.context)).await?
10674        {
10675            return Ok(response);
10676        }
10677
10678        match self.try_skill_route(processed_input).await? {
10679            SkillRouteResult::Response { skill_id, content } => {
10680                self.commit_root_user_message(processed_input).await?;
10681                return self
10682                    .handle_skill_response(processed_input, &skill_id, content, &input_data.context)
10683                    .await;
10684            }
10685            SkillRouteResult::NeedsClarification {
10686                response,
10687                ownership,
10688            } => {
10689                let admission = self
10690                    .admit_optional_disambiguation_ownership(ownership)
10691                    .await?;
10692                self.commit_root_user_message(processed_input).await?;
10693                if let Some(q) = response
10694                    .metadata
10695                    .as_ref()
10696                    .and_then(|m| m.get("disambiguation"))
10697                    .and_then(|d| d.get("status"))
10698                    .and_then(|s| s.as_str())
10699                    && q == "awaiting_clarification"
10700                {
10701                    // Store the clarification question in memory so the next turn
10702                    // can be handled as a clarification response.
10703                    self.memory
10704                        .add_message(ChatMessage::assistant(&response.content))
10705                        .await?;
10706                }
10707                drop(admission);
10708                self.finish_turn_if_root(&response).await?;
10709                return Ok(response);
10710            }
10711            SkillRouteResult::NoMatch => {} // continue to normal LLM chat
10712        }
10713
10714        let effective_reasoning = self.get_effective_reasoning_config();
10715        let reasoning_mode = self.determine_reasoning_mode(processed_input).await?;
10716        let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
10717
10718        info!(
10719            reasoning_mode = ?reasoning_mode,
10720            auto_detected = auto_detected,
10721            reflection_enabled = ?self.reflection_config.enabled,
10722            "Reasoning mode determined"
10723        );
10724
10725        if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
10726            self.commit_root_user_message(processed_input).await?;
10727            return self
10728                .handle_plan_and_execute(processed_input, &input_data.context, auto_detected)
10729                .await;
10730        }
10731
10732        self.commit_root_user_message(processed_input).await?;
10733
10734        let mut iterations = 0u32;
10735        let mut all_tool_calls: Vec<ToolCall> = Vec::new();
10736        let mut thinking_content: Option<String> = None;
10737
10738        let llm = self.get_state_llm()?;
10739
10740        loop {
10741            // When reasoning is active, cap iterations at the reasoning-specific limit.
10742            let effective_max = if reasoning_mode != ReasoningMode::None {
10743                let rc = self.get_effective_reasoning_config();
10744                self.max_iterations.min(rc.max_iterations)
10745            } else {
10746                self.max_iterations
10747            };
10748
10749            if iterations >= effective_max {
10750                let err = AgentError::Other(format!("Max iterations ({}) exceeded", effective_max));
10751                self.hooks.on_error(&err).await;
10752                error!(iterations = iterations, "Max iterations exceeded");
10753                return Err(err);
10754            }
10755            iterations += 1;
10756            *self.iteration_count.write() = iterations;
10757
10758            debug!(iteration = iterations, max = effective_max, "LLM call");
10759
10760            let protocol = self.main_tool_protocol(llm.as_ref(), false).await?;
10761            let mut messages = self
10762                .build_messages_internal(true, None, protocol.choice.is_none())
10763                .await?;
10764            self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
10765
10766            self.hooks.on_llm_start(&messages).await;
10767            let llm_start = Instant::now();
10768            let response = self
10769                .complete_main_llm_with_recovery(Arc::clone(&llm), &messages, &protocol)
10770                .await?;
10771
10772            let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
10773            self.hooks.on_llm_complete(&response, llm_duration_ms).await;
10774
10775            let content = response.content.trim();
10776
10777            if let Some(tool_calls) = self.parse_main_tool_calls(content, &protocol) {
10778                match self
10779                    .handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
10780                    .await?
10781                {
10782                    ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
10783                    ToolCallOutcome::Rejected(resp) => {
10784                        self.finish_turn_if_root(&resp).await?;
10785                        return Ok(resp);
10786                    }
10787                }
10788            }
10789
10790            let (extracted_thinking, answer) = self.extract_thinking(content);
10791            if extracted_thinking.is_some() {
10792                thinking_content = extracted_thinking;
10793            }
10794
10795            let output_data = self.process_output(&answer, &input_data.context).await?;
10796
10797            let mut final_content = if output_data.metadata.rejected {
10798                output_data
10799                    .metadata
10800                    .rejection_reason
10801                    .unwrap_or_else(|| answer.to_string())
10802            } else {
10803                output_data.content
10804            };
10805
10806            // Run reflection (blocking LLM calls for retries)
10807            let reflection_metadata;
10808            (final_content, reflection_metadata) = self
10809                .run_reflection(&*llm, processed_input, final_content)
10810                .await?;
10811
10812            final_content =
10813                self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
10814
10815            // Post-loop: memory, transitions, post-transition re-generation.
10816            // apply_post_loop_result handles NeedsRedispatch by re-entering
10817            // run_loop_internal so the new state's full dispatch activates.
10818            let final_content = {
10819                let result = self
10820                    .post_loop_processing(processed_input, final_content)
10821                    .await?;
10822                self.apply_post_loop_result(processed_input, result).await?
10823            };
10824
10825            let reflected = reflection_metadata.is_some();
10826            let reasoning_mode_debug = format!("{:?}", reasoning_mode);
10827
10828            let response = self.build_agent_response(AgentResponseParts {
10829                content: final_content,
10830                all_tool_calls,
10831                reasoning_mode,
10832                auto_detected,
10833                iterations,
10834                thinking: thinking_content,
10835                reflection_metadata,
10836            });
10837
10838            self.finish_turn_if_root(&response).await?;
10839
10840            let tool_call_count = response.tool_calls.as_ref().map(|tc| tc.len()).unwrap_or(0);
10841            info!(
10842                tool_calls = tool_call_count,
10843                response_len = response.content.len(),
10844                reasoning_mode = %reasoning_mode_debug,
10845                reflected = reflected,
10846                "Chat completed"
10847            );
10848            return Ok(response);
10849        }
10850    }
10851
10852    async fn generate_buffered_streaming_draft(
10853        &self,
10854        processed_input: &str,
10855        routing_resolved: Arc<AtomicBool>,
10856    ) -> Result<StreamingDraftResult> {
10857        let llm = self.get_state_llm()?;
10858        if llm.configured_tool_choice().is_some() {
10859            let draft = self
10860                .generate_main_response_draft(processed_input, &ReasoningMode::None)
10861                .await?;
10862            return Ok(StreamingDraftResult::new(draft, Vec::new()));
10863        }
10864        let messages = self.build_messages_for_draft(processed_input).await?;
10865        let mut stream = self
10866            .observe_purpose(
10867                ObservationPurpose::MainResponse,
10868                llm.complete_stream(&messages, None),
10869            )
10870            .await
10871            .map_err(|e| AgentError::LLM(e.to_string()))?;
10872        let mut buffer = crate::optimization::StreamBranchBuffer::new(self.streaming.buffer_size)?;
10873        let mut chunks = Vec::new();
10874        let mut accumulated = String::new();
10875        while let Some(chunk_result) = stream.next().await {
10876            let chunk = chunk_result.map_err(|e| AgentError::LLM(e.to_string()))?;
10877            accumulated.push_str(&chunk.delta);
10878            let stream_chunk = StreamChunk::content(chunk.delta);
10879            if routing_resolved.load(Ordering::SeqCst) {
10880                chunks.push(stream_chunk);
10881            } else {
10882                buffer.push(stream_chunk)?;
10883            }
10884        }
10885        chunks.splice(0..0, buffer.drain());
10886        let content = accumulated.trim().to_string();
10887        let draft = if let Some(calls) = self.parse_tool_calls(&content) {
10888            MainResponseDraft::ToolCalls {
10889                raw_content: content,
10890                calls,
10891                thinking: None,
10892            }
10893        } else {
10894            MainResponseDraft::Text {
10895                raw_content: content,
10896                thinking: None,
10897            }
10898        };
10899        Ok(StreamingDraftResult::new(draft, chunks))
10900    }
10901
10902    async fn try_buffered_streaming_branches(
10903        &self,
10904        processed_input: &str,
10905        input_context: &HashMap<String, Value>,
10906    ) -> Result<Option<(AgentResponse, Vec<StreamChunk>)>> {
10907        let optimization = &self.runtime_config.optimization;
10908        if !optimization.enabled {
10909            return Ok(None);
10910        }
10911        let transition_enabled =
10912            optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
10913        if !transition_enabled {
10914            return Ok(None);
10915        }
10916        let mut branch_scheduler =
10917            TurnBranchScheduler::new(optimization.max_parallel_runtime_tasks)?;
10918        if !branch_scheduler.reserve_task() {
10919            return Ok(None);
10920        }
10921        if !self
10922            .reserve_active_speculative_llm_call(RuntimeOptimizationKind::BufferedStreamingRouting)
10923        {
10924            branch_scheduler.release_task();
10925            return Ok(None);
10926        }
10927        if !branch_scheduler.reserve_task() {
10928            branch_scheduler.release_task();
10929            return Ok(None);
10930        }
10931        let mut main_branch = RuntimeBranch::new(
10932            RuntimeTaskPurpose::MainResponse,
10933            RuntimeOptimizationKind::BufferedStreamingRouting,
10934            RuntimeTaskPriority::Normal,
10935            RuntimeCommitBehavior::FinalResponse,
10936        );
10937        let mut transition_branch = RuntimeBranch::new(
10938            RuntimeTaskPurpose::StateTransition,
10939            RuntimeOptimizationKind::ParallelStateTransition,
10940            RuntimeTaskPriority::Critical,
10941            RuntimeCommitBehavior::TransitionDecision,
10942        );
10943        let main_id = main_branch.branch_id();
10944        let transition_id = transition_branch.branch_id();
10945        let routing_resolved = Arc::new(AtomicBool::new(false));
10946        let mut main_future =
10947            Box::pin(crate::optimization::observability::with_branch_observation(
10948                &main_id,
10949                RuntimeOptimizationKind::BufferedStreamingRouting,
10950                RuntimeCommitBehavior::FinalResponse,
10951                self.generate_buffered_streaming_draft(
10952                    processed_input,
10953                    Arc::clone(&routing_resolved),
10954                ),
10955            ));
10956        let mut transition_future =
10957            Box::pin(crate::optimization::observability::with_branch_observation(
10958                &transition_id,
10959                RuntimeOptimizationKind::ParallelStateTransition,
10960                RuntimeCommitBehavior::TransitionDecision,
10961                self.select_parallel_transition_candidate(processed_input),
10962            ));
10963        let mut main_pending = true;
10964        let mut transition_pending = true;
10965        let mut main_result: Option<Result<StreamingDraftResult>> = None;
10966        let mut transition_finalized = false;
10967        let mut transition_candidate: Option<TransitionCandidate> = None;
10968        loop {
10969            if let Some(candidate) = transition_candidate.take() {
10970                if self
10971                    .approve_transition_target(&candidate.from_state, candidate.target())
10972                    .await?
10973                {
10974                    // Drop the stale stream future before transition side effects or redispatch reuse the provider.
10975                    drop(main_future);
10976                    drop(transition_future);
10977                    self.finalize_branch_loss(
10978                        &main_id,
10979                        RuntimeOptimizationKind::BufferedStreamingRouting,
10980                        RuntimeCommitBehavior::FinalResponse,
10981                        main_pending,
10982                        main_result.as_ref().map(|result| result.is_err()),
10983                    );
10984                    if !self
10985                        .apply_pre_response_transition_candidate(
10986                            &candidate,
10987                            &HashMap::new(),
10988                            processed_input,
10989                        )
10990                        .await?
10991                    {
10992                        self.finalize_optional_branch(
10993                            &transition_id,
10994                            RuntimeOptimizationKind::ParallelStateTransition,
10995                            RuntimeCommitBehavior::TransitionDecision,
10996                            "discarded",
10997                            false,
10998                        );
10999                        return Ok(None);
11000                    }
11001                    self.finalize_optional_branch(
11002                        &transition_id,
11003                        RuntimeOptimizationKind::ParallelStateTransition,
11004                        RuntimeCommitBehavior::TransitionDecision,
11005                        "committed",
11006                        true,
11007                    );
11008                    let response = self.redispatch_current_state(processed_input).await?;
11009                    return Ok(Some((
11010                        response.clone(),
11011                        vec![StreamChunk::content(response.content)],
11012                    )));
11013                }
11014                self.finalize_optional_branch(
11015                    &transition_id,
11016                    RuntimeOptimizationKind::ParallelStateTransition,
11017                    RuntimeCommitBehavior::TransitionDecision,
11018                    "discarded",
11019                    false,
11020                );
11021                routing_resolved.store(true, Ordering::SeqCst);
11022                transition_finalized = true;
11023            }
11024            if transition_finalized && let Some(result) = main_result.take() {
11025                let stream_draft = match result {
11026                    Ok(stream_draft) => stream_draft,
11027                    Err(error) => {
11028                        self.finalize_optional_branch(
11029                            &main_id,
11030                            RuntimeOptimizationKind::BufferedStreamingRouting,
11031                            RuntimeCommitBehavior::FinalResponse,
11032                            "failed",
11033                            false,
11034                        );
11035                        return Err(error);
11036                    }
11037                };
11038                let raw_draft_content = stream_draft.draft.raw_content().to_string();
11039                let buffered_chunks = stream_draft.chunks;
11040                self.finalize_optional_branch(
11041                    &main_id,
11042                    RuntimeOptimizationKind::BufferedStreamingRouting,
11043                    RuntimeCommitBehavior::FinalResponse,
11044                    "committed",
11045                    true,
11046                );
11047                let response = self
11048                    .commit_main_response_draft(
11049                        processed_input,
11050                        input_context,
11051                        stream_draft.draft,
11052                        ReasoningMode::None,
11053                        false,
11054                    )
11055                    .await?;
11056                let chunks = if response.content == raw_draft_content {
11057                    buffered_chunks
11058                } else {
11059                    vec![StreamChunk::content(response.content.clone())]
11060                };
11061                return Ok(Some((response, chunks)));
11062            }
11063            tokio::select! {
11064                result = &mut main_future, if main_pending => {
11065                    main_pending = false;
11066                    main_branch.transition_to(RuntimeBranchStatus::Completed)?;
11067                    main_result = Some(result);
11068                }
11069                result = &mut transition_future, if transition_pending => {
11070                    transition_pending = false;
11071                    transition_branch.transition_to(RuntimeBranchStatus::Completed)?;
11072                    match result {
11073                        Ok(ParallelTransitionSelection::Candidate(candidate)) => {
11074                            transition_candidate = Some(candidate)
11075                        }
11076                        Ok(ParallelTransitionSelection::NoMatch) => {
11077                            self.finalize_optional_branch(
11078                                &transition_id,
11079                                RuntimeOptimizationKind::ParallelStateTransition,
11080                                RuntimeCommitBehavior::TransitionDecision,
11081                                "discarded",
11082                                false,
11083                            );
11084                            routing_resolved.store(true, Ordering::SeqCst);
11085                            transition_finalized = true;
11086                        }
11087                        Ok(ParallelTransitionSelection::ReservationExhausted) => {
11088                            self.finalize_optional_branch(
11089                                &transition_id,
11090                                RuntimeOptimizationKind::ParallelStateTransition,
11091                                RuntimeCommitBehavior::TransitionDecision,
11092                                "cancelled",
11093                                false,
11094                            );
11095                            routing_resolved.store(true, Ordering::SeqCst);
11096                            self.finalize_branch_loss(
11097                                &main_id,
11098                                RuntimeOptimizationKind::BufferedStreamingRouting,
11099                                RuntimeCommitBehavior::FinalResponse,
11100                                main_pending,
11101                                main_result.as_ref().map(|result| result.is_err()),
11102                            );
11103                            return Ok(None);
11104                        }
11105                        Err(_) => {
11106                            self.finalize_optional_branch(
11107                                &transition_id,
11108                                RuntimeOptimizationKind::ParallelStateTransition,
11109                                RuntimeCommitBehavior::TransitionDecision,
11110                                "failed",
11111                                false,
11112                            );
11113                            routing_resolved.store(true, Ordering::SeqCst);
11114                            transition_finalized = true;
11115                        }
11116                    }
11117                }
11118            }
11119        }
11120    }
11121
11122    /// Streaming agent pipeline
11123    /// Uses all the same shared helpers as run_loop_internal.
11124    /// The ONLY difference: LLM calls use complete_stream() + yield deltas.
11125    fn run_loop_internal_stream<'a>(
11126        &'a self,
11127        input: &'a str,
11128        terminal: RuntimeStreamTerminalSlot,
11129    ) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
11130        let include_tool_events = self.streaming.include_tool_events;
11131        let include_state_events = self.streaming.include_state_events;
11132
11133        Box::pin(async_stream::stream! {
11134            self.begin_root_turn();
11135            // Parity with non-stream: resolve actor from context and load facts if changed.
11136            self.pre_turn_session_lifecycle().await;
11137
11138            let input_data = match self.process_input(input).await {
11139                Ok(data) => data,
11140                Err(e) => {
11141                    yield StreamChunk::error(e.to_string());
11142                    return;
11143                }
11144            };
11145            self.update_active_turn_context(&input_data.content, input_data.context.clone());
11146
11147            // Inject process context (detect/extract results) into agent context
11148            for (key, value) in &input_data.context {
11149                let _ = self.context_manager.set(key, value.clone());
11150            }
11151
11152            if input_data.metadata.rejected {
11153                let reason = input_data
11154                    .metadata
11155                    .rejection_reason
11156                    .unwrap_or_else(|| "Input rejected".to_string());
11157                warn!(reason = %reason, "Input rejected (stream)");
11158                yield StreamChunk::error(reason);
11159                return;
11160            }
11161
11162            let processed_input = &input_data.content;
11163
11164            if self.runtime_config.optimization.enabled
11165                && matches!(
11166                    self.runtime_config.optimization.streaming_policy,
11167                    crate::optimization::StreamingOptimizationPolicy::BufferUntilRoutingDone
11168                )
11169            {
11170                //
11171                // Buffered routing keeps stale stream output hidden until a branch winner is known.
11172                // The boxed future prevents this stream state machine from becoming too large.
11173                //
11174                match Box::pin(self.try_buffered_streaming_branches(processed_input, &input_data.context)).await {
11175                    Ok(Some((response, chunks))) => {
11176                        for chunk in chunks {
11177                            yield chunk;
11178                        }
11179                        record_runtime_stream_final(&terminal, response);
11180                        yield StreamChunk::Done {};
11181                        return;
11182                    }
11183                    Ok(None) => {}
11184                    Err(e) => {
11185                        yield StreamChunk::error(e.to_string());
11186                        return;
11187                    }
11188                }
11189            }
11190
11191            if self.runtime_config.optimization.enabled
11192                && matches!(
11193                    self.runtime_config.optimization.streaming_policy,
11194                    crate::optimization::StreamingOptimizationPolicy::PreflightOnly
11195                )
11196            {
11197                match self.try_pre_response_transition(processed_input).await {
11198                    Ok(Some(response)) => {
11199                        yield StreamChunk::content(&response.content);
11200                        record_runtime_stream_final(&terminal, response);
11201                        yield StreamChunk::Done {};
11202                        return;
11203                    }
11204                    Ok(None) => {}
11205                    Err(e) => {
11206                        yield StreamChunk::error(e.to_string());
11207                        return;
11208                    }
11209                }
11210            }
11211
11212            // Handle orchestration states in streaming mode.
11213            if let Some(ref sm) = self.state_machine
11214                && let Some(def) = sm.current_definition()
11215            {
11216                    let orchestration_result = if let Some(ref delegate_id) = def.delegate {
11217                        Some(self.handle_delegated_state(processed_input, delegate_id, &def).await)
11218                    } else if let Some(ref concurrent_config) = def.concurrent {
11219                        Some(self.handle_concurrent_state(processed_input, concurrent_config).await)
11220                    } else if let Some(ref group_chat_config) = def.group_chat {
11221                        Some(self.handle_group_chat_state(processed_input, group_chat_config).await)
11222                    } else if let Some(ref pipeline_config) = def.pipeline {
11223                        Some(self.handle_pipeline_state(processed_input, pipeline_config).await)
11224                    } else if let Some(ref handoff_config) = def.handoff {
11225                        Some(self.handle_handoff_state(processed_input, handoff_config).await)
11226                    } else {
11227                        None
11228                    };
11229
11230                    if let Some(result) = orchestration_result {
11231                        match result {
11232                            Ok(response) => {
11233                                yield StreamChunk::content(&response.content);
11234                                record_runtime_stream_final(&terminal, response);
11235                                yield StreamChunk::Done {};
11236                            }
11237                            Err(e) => {
11238                                yield StreamChunk::error(e.to_string());
11239                            }
11240                        }
11241                        return;
11242                    }
11243                }
11244
11245            // Skill routing
11246            match self.try_skill_route(processed_input).await {
11247                Ok(SkillRouteResult::Response { skill_id, content }) => {
11248                    if let Err(e) = self.commit_root_user_message(processed_input).await {
11249                        yield StreamChunk::error(e.to_string());
11250                        return;
11251                    }
11252                    match self.handle_skill_response(processed_input, &skill_id, content, &input_data.context).await {
11253                        Ok(resp) => {
11254                            yield StreamChunk::content(&resp.content);
11255                            record_runtime_stream_final(&terminal, resp);
11256                            yield StreamChunk::Done {};
11257                            return;
11258                        }
11259                        Err(e) => {
11260                            yield StreamChunk::error(e.to_string());
11261                            return;
11262                        }
11263                    }
11264                }
11265                Ok(SkillRouteResult::NeedsClarification {
11266                    response,
11267                    ownership,
11268                }) => {
11269                    let admission = match self
11270                        .admit_optional_disambiguation_ownership(ownership)
11271                        .await
11272                    {
11273                        Ok(admission) => admission,
11274                        Err(e) => {
11275                            yield StreamChunk::error(e.to_string());
11276                            return;
11277                        }
11278                    };
11279                    if let Err(e) = self.commit_root_user_message(processed_input).await {
11280                        yield StreamChunk::error(e.to_string());
11281                        return;
11282                    }
11283                    let _ = self.memory.add_message(ChatMessage::assistant(&response.content)).await;
11284                    drop(admission);
11285                    if let Err(e) = self.finish_turn_if_root(&response).await {
11286                        yield StreamChunk::error(e.to_string());
11287                        return;
11288                    }
11289                    yield StreamChunk::content(&response.content);
11290                    record_runtime_stream_final(&terminal, response);
11291                    yield StreamChunk::Done {};
11292                    return;
11293                }
11294                Ok(SkillRouteResult::NoMatch) => {} // no skill matched, continue
11295                Err(e) => {
11296                    yield StreamChunk::error(e.to_string());
11297                    return;
11298                }
11299            }
11300
11301            // Reasoning mode determination
11302            let effective_reasoning = self.get_effective_reasoning_config();
11303            let reasoning_mode = match self.determine_reasoning_mode(processed_input).await {
11304                Ok(mode) => mode,
11305                Err(e) => {
11306                    yield StreamChunk::error(e.to_string());
11307                    return;
11308                }
11309            };
11310            let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
11311
11312            info!(
11313                reasoning_mode = ?reasoning_mode,
11314                auto_detected = auto_detected,
11315                "Reasoning mode determined (stream)"
11316            );
11317
11318            // Plan-and-Execute: yield final result as single chunk (not token-by-token)
11319            if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
11320                if let Err(e) = self.commit_root_user_message(processed_input).await {
11321                    yield StreamChunk::error(e.to_string());
11322                    return;
11323                }
11324                match self.handle_plan_and_execute(processed_input, &input_data.context, auto_detected).await {
11325                    Ok(resp) => {
11326                        yield StreamChunk::content(&resp.content);
11327                        record_runtime_stream_final(&terminal, resp);
11328                        yield StreamChunk::Done {};
11329                        return;
11330                    }
11331                    Err(e) => {
11332                        yield StreamChunk::error(e.to_string());
11333                        return;
11334                    }
11335                }
11336            }
11337
11338            if let Err(e) = self.commit_root_user_message(processed_input).await {
11339                yield StreamChunk::error(e.to_string());
11340                return;
11341            }
11342
11343            let llm = match self.get_state_llm() {
11344                Ok(llm) => llm,
11345                Err(e) => {
11346                    yield StreamChunk::error(e.to_string());
11347                    return;
11348                }
11349            };
11350
11351            let mut iterations = 0u32;
11352            let mut all_tool_calls: Vec<ToolCall> = Vec::new();
11353            let mut thinking_content: Option<String> = None;
11354
11355            loop {
11356                // When reasoning is active, cap iterations at the reasoning-specific limit.
11357                let effective_max = if reasoning_mode != ReasoningMode::None {
11358                    let rc = self.get_effective_reasoning_config();
11359                    self.max_iterations.min(rc.max_iterations)
11360                } else {
11361                    self.max_iterations
11362                };
11363
11364                if iterations >= effective_max {
11365                    let err_msg = format!("Max iterations ({}) exceeded", effective_max);
11366                    let err = AgentError::Other(err_msg.clone());
11367                    self.hooks.on_error(&err).await;
11368                    error!(iterations = iterations, "Max iterations exceeded (stream)");
11369                    yield StreamChunk::error(err_msg);
11370                    return;
11371                }
11372                iterations += 1;
11373                *self.iteration_count.write() = iterations;
11374
11375                debug!(iteration = iterations, max = effective_max, "LLM call (stream)");
11376
11377                let protocol = match self.main_tool_protocol(llm.as_ref(), false).await {
11378                    Ok(protocol) => protocol,
11379                    Err(e) => {
11380                        yield StreamChunk::error(e.to_string());
11381                        return;
11382                    }
11383                };
11384                let mut messages = match self
11385                    .build_messages_internal(true, None, protocol.choice.is_none())
11386                    .await
11387                {
11388                    Ok(m) => m,
11389                    Err(e) => {
11390                        yield StreamChunk::error(e.to_string());
11391                        return;
11392                    }
11393                };
11394                self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
11395
11396                self.hooks.on_llm_start(&messages).await;
11397                let llm_start = Instant::now();
11398
11399                // Check if reflection is active — if so, suppress streaming for this LLM call
11400                // because we may need to retry and the user would see a stale first attempt.
11401                let reflection_active = self
11402                    .should_reflect(processed_input, "")
11403                    .await
11404                    .unwrap_or_default();
11405
11406                let buffered_decision = reflection_active || protocol.choice.is_some();
11407                let content = if buffered_decision {
11408                    //
11409                    // Explicit tool choice buffers the provider decision so no text or tool call is visible before the runtime validates and commits it.
11410                    //
11411                    let response = match self
11412                        .complete_main_llm_with_recovery(
11413                            Arc::clone(&llm),
11414                            &messages,
11415                            &protocol,
11416                        )
11417                        .await
11418                    {
11419                        Ok(r) => r,
11420                        Err(e) => {
11421                            yield StreamChunk::error(e.to_string());
11422                            return;
11423                        }
11424                    };
11425                    let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
11426                    self.hooks.on_llm_complete(&response, llm_duration_ms).await;
11427                    response.content.trim().to_string()
11428                } else {
11429                    // Streaming LLM call
11430                    let llm_stream = match self
11431                        .observe_purpose(
11432                            ObservationPurpose::MainResponse,
11433                            llm.complete_stream(&messages, None),
11434                        )
11435                        .await
11436                    {
11437                        Ok(s) => s,
11438                        Err(e) => {
11439                            yield StreamChunk::error(e.to_string());
11440                            return;
11441                        }
11442                    };
11443                    let mut accumulated = String::new();
11444                    let mut stream_inner = llm_stream;
11445                    while let Some(chunk_result) = stream_inner.next().await {
11446                        match chunk_result {
11447                            Ok(chunk) => {
11448                                accumulated.push_str(&chunk.delta);
11449                                yield StreamChunk::content(chunk.delta);
11450                            }
11451                            Err(e) => {
11452                                yield StreamChunk::error(e.to_string());
11453                                return;
11454                            }
11455                        }
11456                    }
11457                    let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
11458                    // Construct LLMResponse for hooks
11459                    let llm_response = ai_agents_core::LLMResponse::new(
11460                        accumulated.trim(),
11461                        ai_agents_core::FinishReason::Stop,
11462                    );
11463                    self.hooks.on_llm_complete(&llm_response, llm_duration_ms).await;
11464                    accumulated.trim().to_string()
11465                };
11466
11467                // Tool call handling
11468                if let Some(tool_calls) = self.parse_main_tool_calls(&content, &protocol) {
11469                    let native_tool_call = Self::is_native_tool_call_content(&content);
11470                    // Emit tool events for streaming
11471                    // First check transitions (same as blocking path)
11472                    let transition_fired = match self.evaluate_transitions(processed_input, &content).await {
11473                        Ok(v) => v,
11474                        Err(e) => {
11475                            yield StreamChunk::error(e.to_string());
11476                            return;
11477                        }
11478                    };
11479                    if transition_fired {
11480                        let _ = self.memory.add_message(ChatMessage::assistant(
11481                            "(Transitioned to new state — tool call handled by workflow)",
11482                        )).await;
11483
11484                        if include_state_events
11485                            && let Some(state) = self.current_state()
11486                        {
11487                                yield StreamChunk::state_transition(None, state);
11488                            }
11489                        continue;
11490                    }
11491
11492                    // Store the assistant's tool-call message (same as blocking path)
11493                    let _ = self.memory.add_message(ChatMessage::assistant(&content)).await;
11494
11495                    // Execute tools with streaming events
11496                    let results = self.execute_tools_parallel(&tool_calls).await;
11497
11498                    for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
11499                        if include_tool_events {
11500                            yield StreamChunk::tool_start(&tool_call.id, &tool_call.name);
11501                        }
11502
11503                        match result {
11504                            Ok(output) => {
11505                                if include_tool_events {
11506                                    yield StreamChunk::tool_result(
11507                                        &tool_call.id,
11508                                        &tool_call.name,
11509                                        &output,
11510                                        true,
11511                                    );
11512                                }
11513                                let _ = self.memory
11514                                    .add_message(Self::tool_result_message(
11515                                        tool_call,
11516                                        &output,
11517                                        native_tool_call,
11518                                    ))
11519                                    .await;
11520                            }
11521                            Err(e) => {
11522                                if matches!(e, AgentError::HITLRejected(_)) {
11523                                    let _ = self.memory.add_message(ChatMessage::assistant(
11524                                        format!("The operation was rejected by the approver: {}", e),
11525                                    )).await;
11526                                    let response = AgentResponse {
11527                                        content: format!("Operation cancelled: {}", e),
11528                                        metadata: None,
11529                                        tool_calls: Some(all_tool_calls.clone()),
11530                                    };
11531                                    if let Err(finalize_error) = self.finish_turn_if_root(&response).await {
11532                                        yield StreamChunk::error(finalize_error.to_string());
11533                                        return;
11534                                    }
11535                                    let legacy_error = response.content.clone();
11536                                    record_runtime_stream_final(&terminal, response);
11537                                    yield StreamChunk::error(legacy_error);
11538                                    yield StreamChunk::Done {};
11539                                    return;
11540                                }
11541                                if include_tool_events {
11542                                    yield StreamChunk::tool_result(
11543                                        &tool_call.id,
11544                                        &tool_call.name,
11545                                        e.to_string(),
11546                                        false,
11547                                    );
11548                                }
11549                                let _ = self.memory
11550                                    .add_message(Self::tool_result_message(
11551                                        tool_call,
11552                                        &format!("Error: {}", e),
11553                                        native_tool_call,
11554                                    ))
11555                                    .await;
11556                            }
11557                        }
11558                        all_tool_calls.push(tool_call.clone());
11559
11560                        if include_tool_events {
11561                            yield StreamChunk::tool_end(&tool_call.id);
11562                        }
11563                    }
11564                    continue;
11565                }
11566
11567                // Extract thinking, process output
11568                let (extracted_thinking, answer) = self.extract_thinking(&content);
11569                if extracted_thinking.is_some() {
11570                    thinking_content = extracted_thinking;
11571                }
11572
11573                let output_data = match self.process_output(&answer, &input_data.context).await {
11574                    Ok(d) => d,
11575                    Err(e) => {
11576                        yield StreamChunk::error(e.to_string());
11577                        return;
11578                    }
11579                };
11580
11581                let final_content = if output_data.metadata.rejected {
11582                    output_data
11583                        .metadata
11584                        .rejection_reason
11585                        .unwrap_or_else(|| answer.to_string())
11586                } else {
11587                    output_data.content
11588                };
11589
11590                // Reflection (uses blocking LLM calls for retries)
11591                let (final_content, reflection_metadata) = match self
11592                    .run_reflection(&*llm, processed_input, final_content)
11593                    .await
11594                {
11595                    Ok(r) => r,
11596                    Err(e) => {
11597                        yield StreamChunk::error(e.to_string());
11598                        return;
11599                    }
11600                };
11601
11602                let final_content = self.format_response_with_thinking(
11603                    thinking_content.as_deref(),
11604                    &final_content,
11605                );
11606
11607                // Buffered decisions emit only the accepted final text.
11608                if buffered_decision {
11609                    yield StreamChunk::content(&final_content);
11610                }
11611
11612                // Post-loop: memory, transitions, post-transition re-generation.
11613                // For NeedsRedispatch, run_loop_internal handles the new state's full
11614                // dispatch and its result is yielded as a single non-streamed chunk.
11615                let post_result = match self
11616                    .post_loop_processing(processed_input, final_content)
11617                    .await
11618                {
11619                    Ok(r) => r,
11620                    Err(e) => {
11621                        yield StreamChunk::error(e.to_string());
11622                        return;
11623                    }
11624                };
11625
11626                let (final_content, transitioned) = match post_result {
11627                    PostLoopResult::NoTransition(content) => (content, false),
11628                    PostLoopResult::Transitioned(content) => (content, true),
11629                    PostLoopResult::NeedsRedispatch => {
11630                        const MAX_REDISPATCH_DEPTH: u32 = 3;
11631                        let current_depth = *self.redispatch_depth.read();
11632                        let content = if current_depth >= MAX_REDISPATCH_DEPTH {
11633                            warn!(
11634                                depth = current_depth,
11635                                "Post-transition re-dispatch depth limit reached (stream)"
11636                            );
11637                            let c = String::new();
11638                            let _ = self.memory.add_message(ChatMessage::assistant(&c)).await;
11639                            c
11640                        } else {
11641                            *self.redispatch_depth.write() += 1;
11642                            if let Some(context) = self.active_turn_context.write().as_mut() {
11643                                context.enter_redispatch();
11644                            }
11645                            info!(
11646                                depth = current_depth + 1,
11647                                "Re-dispatching for new state after transition (stream)"
11648                            );
11649                            let result = self.run_loop_internal(processed_input).await;
11650                            *self.redispatch_depth.write() -= 1;
11651                            if let Some(context) = self.active_turn_context.write().as_mut() {
11652                                context.exit_redispatch();
11653                            }
11654                            match result {
11655                                Ok(resp) => resp.content,
11656                                Err(e) => {
11657                                    yield StreamChunk::error(e.to_string());
11658                                    return;
11659                                }
11660                            }
11661                        };
11662                        (content, true)
11663                    }
11664                };
11665
11666                if transitioned {
11667                    if include_state_events
11668                        && let Some(state) = self.current_state()
11669                    {
11670                            yield StreamChunk::state_transition(None, state);
11671                        }
11672                    // Yield the post-transition re-generated or re-dispatched content.
11673                    yield StreamChunk::content(&final_content);
11674                }
11675
11676                // Build and finalize the same authoritative response shape before exposing the terminal event.
11677                let final_response = self.build_agent_response(AgentResponseParts {
11678                    content: final_content,
11679                    all_tool_calls,
11680                    reasoning_mode,
11681                    auto_detected,
11682                    iterations,
11683                    thinking: thinking_content,
11684                    reflection_metadata,
11685                });
11686                if let Err(e) = self.finish_turn_if_root(&final_response).await {
11687                    yield StreamChunk::error(e.to_string());
11688                    return;
11689                }
11690
11691                record_runtime_stream_final(&terminal, final_response);
11692                yield StreamChunk::Done {};
11693                return;
11694            }
11695        })
11696    }
11697
11698    /// Streams the root turn while keeping clarification and confirmation questions as terminal responses for their turn.
11699    /// Pending manager and skill ownership must survive until explicit confirmation returns a resolved result for redispatch.
11700    fn run_loop_stream<'a>(
11701        &'a self,
11702        input: &'a str,
11703        terminal: RuntimeStreamTerminalSlot,
11704    ) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
11705        Box::pin(async_stream::stream! {
11706            self.begin_root_turn();
11707            let _root_cleanup = RootTurnCleanup::new(self);
11708            self.hooks.on_message_received(input).await;
11709
11710            // One-shot context initialization (mirrors run_loop)
11711            if !self.context_initialized.swap(true, Ordering::SeqCst) {
11712                if let Err(e) = self.context_manager.initialize().await {
11713                    yield StreamChunk::error(e.to_string());
11714                    return;
11715                }
11716                debug!("Context manager initialized (defaults, env, builtins)");
11717            }
11718
11719            if let Err(e) = self.check_turn_timeout().await {
11720                yield StreamChunk::error(e.to_string());
11721                return;
11722            }
11723            if let Err(e) = self.context_manager.refresh_per_turn().await {
11724                yield StreamChunk::error(e.to_string());
11725                return;
11726            }
11727
11728            // Clear stale disambiguation context from previous turns.
11729            self.clear_disambiguation_context();
11730
11731            // Disambiguation check (before input processing)
11732            if let Some(ref disambiguator) = self.disambiguation_manager {
11733                let disambiguation_context = match self.build_disambiguation_context().await {
11734                    Ok(ctx) => ctx,
11735                    Err(e) => {
11736                        yield StreamChunk::error(e.to_string());
11737                        return;
11738                    }
11739                };
11740
11741                let state_override = self
11742                    .state_machine
11743                    .as_ref()
11744                    .and_then(|sm| sm.current_definition())
11745                    .and_then(|def| def.disambiguation.clone());
11746
11747                let state_generation = self
11748                    .state_machine
11749                    .as_ref()
11750                    .map(|state_machine| state_machine.generation());
11751                let disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
11752                let mut result = match self
11753                    .observe_purpose(
11754                        ObservationPurpose::DisambiguationDetection,
11755                        disambiguator.process_input_with_override(
11756                            input,
11757                            &disambiguation_context,
11758                            state_override.as_ref(),
11759                            None,
11760                        ),
11761                    )
11762                    .await
11763                {
11764                    Ok(r) => r,
11765                    Err(e) => {
11766                        yield StreamChunk::error(e.to_string());
11767                        return;
11768                    }
11769                };
11770                let current_state_generation = self
11771                    .state_machine
11772                    .as_ref()
11773                    .map(|state_machine| state_machine.generation());
11774                if current_state_generation != state_generation
11775                    || self.disambiguation_epoch.load(Ordering::SeqCst) != disambiguation_epoch
11776                {
11777                    disambiguator.clear_pending().await;
11778                    *self.pending_skill_id.write() = None;
11779                    result = DisambiguationResult::Abandoned { new_input: None };
11780                    info!(
11781                        confirmation_event = "invalidated",
11782                        invalidation_reason = "state_generation_changed",
11783                        "Streaming disambiguation result invalidated before redispatch"
11784                    );
11785                }
11786                match result {
11787                    DisambiguationResult::Clear => {
11788                        debug!("Input is clear, proceeding normally (stream)");
11789                    }
11790                    DisambiguationResult::NeedsClarification {
11791                        question,
11792                        detection,
11793                    } => {
11794                        let admission = match self
11795                            .admit_disambiguation_redispatch(
11796                                disambiguation_epoch,
11797                                state_generation,
11798                            )
11799                            .await
11800                        {
11801                            Ok(admission) => admission,
11802                            Err(error) => {
11803                                *self.pending_skill_id.write() = None;
11804                                yield StreamChunk::error(error.to_string());
11805                                return;
11806                            }
11807                        };
11808                        let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
11809                        info!(
11810                            ambiguity_type = ?detection.ambiguity_type,
11811                            confidence = detection.confidence,
11812                            "Input requires clarification (stream)"
11813                        );
11814                        // Confirmation uses the same terminal branch so enriched input cannot stream before explicit agreement.
11815                        // Keep pending_skill_id intact for the later confirmed redispatch.
11816                        if let Err(e) = self.commit_root_user_message(input).await {
11817                            yield StreamChunk::error(e.to_string());
11818                            return;
11819                        }
11820                        let _ = self
11821                            .memory
11822                            .add_message(ChatMessage::assistant(&question.question))
11823                            .await;
11824                        let status = if awaiting_confirmation {
11825                            "awaiting_confirmation"
11826                        } else {
11827                            "awaiting_clarification"
11828                        };
11829                        let response = AgentResponse::new(&question.question).with_metadata(
11830                            "disambiguation",
11831                            serde_json::json!({ "status": status }),
11832                        );
11833                        drop(admission);
11834                        if let Err(e) = self.finish_turn_if_root(&response).await {
11835                            yield StreamChunk::error(e.to_string());
11836                            return;
11837                        }
11838                        yield StreamChunk::content(&question.question);
11839                        record_runtime_stream_final(&terminal, response);
11840                        yield StreamChunk::Done {};
11841                        return;
11842                    }
11843                    DisambiguationResult::Clarified {
11844                        enriched_input,
11845                        resolved,
11846                        ..
11847                    } => {
11848                        let admission = match self
11849                            .admit_disambiguation_redispatch(
11850                                disambiguation_epoch,
11851                                state_generation,
11852                            )
11853                            .await
11854                        {
11855                            Ok(admission) => admission,
11856                            Err(error) => {
11857                                *self.pending_skill_id.write() = None;
11858                                yield StreamChunk::error(error.to_string());
11859                                return;
11860                            }
11861                        };
11862                        info!(
11863                            resolved_count = resolved.len(),
11864                            enriched = %enriched_input,
11865                            "Input clarified (stream)"
11866                        );
11867                        for (key, value) in &resolved {
11868                            let context_key = format!("disambiguation.{}", key);
11869                            let _ = self.context_manager.set(&context_key, value.clone());
11870                        }
11871                        if let Some(intent) = resolved.get("intent") {
11872                            let _ = self.context_manager.set("resolved_intent", intent.clone());
11873                        }
11874                        let _ = self
11875                            .context_manager
11876                            .set("disambiguation.resolved", serde_json::Value::Bool(true));
11877
11878                        // Check if this clarification was triggered by a skill-level override.
11879                        // Re-run skill disambiguation to verify all required_clarity fields
11880                        // are present before executing.
11881                        let skill_id = self.pending_skill_id.read().clone();
11882                        if let Some(skill_id) = skill_id {
11883                            info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input (stream)");
11884                            drop(admission);
11885                            match self
11886                                .recheck_skill_disambiguation(
11887                                    &skill_id,
11888                                    &enriched_input,
11889                                    disambiguation_epoch,
11890                                    state_generation,
11891                                )
11892                                .await
11893                            {
11894                                Ok(resp) => {
11895                                    yield StreamChunk::content(&resp.content);
11896                                    record_runtime_stream_final(&terminal, resp);
11897                                    yield StreamChunk::Done {};
11898                                    return;
11899                                }
11900                                Err(e) => {
11901                                    yield StreamChunk::error(e.to_string());
11902                                    return;
11903                                }
11904                            }
11905                        }
11906
11907                        // Forward to internal stream with enriched input.
11908                        drop(admission);
11909                        let mut inner = self.run_loop_internal_stream(
11910                            &enriched_input,
11911                            Arc::clone(&terminal),
11912                        );
11913                        while let Some(chunk) = inner.next().await {
11914                            yield chunk;
11915                        }
11916                        return;
11917                    }
11918                    DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
11919                        info!("Proceeding with best guess (stream)");
11920
11921                        // Same skill-id re-check for best-guess path
11922                        let skill_id = self.pending_skill_id.read().clone();
11923                        if let Some(skill_id) = skill_id {
11924                            info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input (stream)");
11925                            match self
11926                                .recheck_skill_disambiguation(
11927                                    &skill_id,
11928                                    &enriched_input,
11929                                    disambiguation_epoch,
11930                                    state_generation,
11931                                )
11932                                .await
11933                            {
11934                                Ok(resp) => {
11935                                    yield StreamChunk::content(&resp.content);
11936                                    record_runtime_stream_final(&terminal, resp);
11937                                    yield StreamChunk::Done {};
11938                                    return;
11939                                }
11940                                Err(e) => {
11941                                    yield StreamChunk::error(e.to_string());
11942                                    return;
11943                                }
11944                            }
11945                        }
11946
11947                        let mut inner = self.run_loop_internal_stream(
11948                            &enriched_input,
11949                            Arc::clone(&terminal),
11950                        );
11951                        while let Some(chunk) = inner.next().await {
11952                            yield chunk;
11953                        }
11954                        return;
11955                    }
11956                    DisambiguationResult::GiveUp { reason } => {
11957                        *self.pending_skill_id.write() = None;
11958                        warn!(reason = %reason, "Disambiguation gave up (stream)");
11959                        let apology = self
11960                            .generate_localized_apology(
11961                                "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
11962                                &reason,
11963                            )
11964                            .await
11965                            .unwrap_or_else(|_| {
11966                                format!("I'm sorry, I couldn't understand your request: {}", reason)
11967                            });
11968                        let response = AgentResponse::new(&apology);
11969                        if let Err(e) = self.finish_turn_if_root(&response).await {
11970                            yield StreamChunk::error(e.to_string());
11971                            return;
11972                        }
11973                        yield StreamChunk::content(&apology);
11974                        record_runtime_stream_final(&terminal, response);
11975                        yield StreamChunk::Done {};
11976                        return;
11977                    }
11978                    DisambiguationResult::Escalate { reason } => {
11979                        *self.pending_skill_id.write() = None;
11980                        info!(reason = %reason, "Escalating to human (stream)");
11981                        if let Some(ref hitl) = self.hitl_engine {
11982                            let trigger =
11983                                ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
11984                            let mut context_map = HashMap::new();
11985                            context_map.insert("original_input".to_string(), serde_json::json!(input));
11986                            context_map.insert("reason".to_string(), serde_json::json!(&reason));
11987                            let check_result = HITLCheckResult::required(
11988                                trigger,
11989                                context_map,
11990                                format!("User request needs human assistance: {}", reason),
11991                                Some(hitl.config().default_timeout_seconds),
11992                            );
11993                            match self.request_hitl_approval(check_result).await {
11994                                Ok(ApprovalResult::Approved | ApprovalResult::Modified { .. }) => {
11995                                    let mut inner = self.run_loop_internal_stream(
11996                                        input,
11997                                        Arc::clone(&terminal),
11998                                    );
11999                                    while let Some(chunk) = inner.next().await {
12000                                        yield chunk;
12001                                    }
12002                                    return;
12003                                }
12004                                Ok(_) => {}
12005                                Err(e) => {
12006                                    yield StreamChunk::error(e.to_string());
12007                                    return;
12008                                }
12009                            }
12010                        }
12011                        let apology = self
12012                            .generate_localized_apology(
12013                                "Explain briefly that you're transferring the user to a human agent for help.",
12014                                &reason,
12015                            )
12016                            .await
12017                            .unwrap_or_else(|_| {
12018                                format!("I need human assistance to help with your request: {}", reason)
12019                            });
12020                        let response = AgentResponse::new(&apology);
12021                        if let Err(e) = self.finish_turn_if_root(&response).await {
12022                            yield StreamChunk::error(e.to_string());
12023                            return;
12024                        }
12025                        yield StreamChunk::content(&apology);
12026                        record_runtime_stream_final(&terminal, response);
12027                        yield StreamChunk::Done {};
12028                        return;
12029                    }
12030                    DisambiguationResult::Abandoned { new_input } => {
12031                        *self.pending_skill_id.write() = None;
12032
12033                        info!(
12034                            has_new_input = new_input.is_some(),
12035                            "Clarification abandoned by user (stream)"
12036                        );
12037
12038                        if let Err(e) = self.commit_root_user_message(input).await {
12039                            yield StreamChunk::error(e.to_string());
12040                            return;
12041                        }
12042
12043                        match new_input {
12044                            Some(fresh_input) => {
12045                                // Topic switch: forward to internal stream with fresh input.
12046                                let mut inner = self.run_loop_internal_stream(
12047                                    &fresh_input,
12048                                    Arc::clone(&terminal),
12049                                );
12050                                while let Some(chunk) = inner.next().await {
12051                                    yield chunk;
12052                                }
12053                                return;
12054                            }
12055                            None => {
12056                                // Pure abandonment: generate a brief acknowledgment.
12057                                let ack = self
12058                                    .generate_localized_apology(
12059                                        "The user changed their mind about their previous request. \
12060                                         Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
12061                                         Do NOT apologize excessively. Be concise.",
12062                                        "User abandoned clarification",
12063                                    )
12064                                    .await
12065                                    .unwrap_or_else(|_| {
12066                                        "OK, no problem. What else can I help with?".to_string()
12067                                    });
12068
12069                                let _ = self
12070                                    .memory
12071                                    .add_message(ChatMessage::assistant(&ack))
12072                                    .await;
12073
12074                                let response = AgentResponse::new(&ack);
12075                                if let Err(e) = self.finish_turn_if_root(&response).await {
12076                                    yield StreamChunk::error(e.to_string());
12077                                    return;
12078                                }
12079                                yield StreamChunk::content(&ack);
12080                                record_runtime_stream_final(&terminal, response);
12081                                yield StreamChunk::Done {};
12082                                return;
12083                            }
12084                        }
12085                    }
12086                }
12087            }
12088
12089            // No disambiguation or Clear result — proceed with internal stream
12090            let mut inner = self.run_loop_internal_stream(input, Arc::clone(&terminal));
12091            while let Some(chunk) = inner.next().await {
12092                yield chunk;
12093            }
12094        })
12095    }
12096
12097    pub fn info(&self) -> AgentInfo {
12098        self.info.clone()
12099    }
12100
12101    pub fn skills(&self) -> &[SkillDefinition] {
12102        &self.skills
12103    }
12104
12105    /// Clears conversation and pending runtime ownership through one reset contract.
12106    async fn reset_runtime_state(&self) -> Result<()> {
12107        let _admission = self.disambiguation_admission.write().await;
12108        if self.state_transition_reserved.load(Ordering::SeqCst) {
12109            return Err(AgentError::Other(
12110                "Cannot reset while a state transition is in progress".to_string(),
12111            ));
12112        }
12113        self.disambiguation_epoch.fetch_add(1, Ordering::SeqCst);
12114        *self.pending_skill_id.write() = None;
12115        if let Some(disambiguator) = self.disambiguation_manager.as_ref() {
12116            disambiguator.clear_pending().await;
12117        }
12118        self.memory.clear().await?;
12119        *self.iteration_count.write() = 0;
12120        self.tool_call_history.write().clear();
12121        if let Some(ref sm) = self.state_machine {
12122            sm.reset();
12123        }
12124        Ok(())
12125    }
12126
12127    /// Resets the runtime using the same cleanup path as the Agent trait.
12128    pub async fn reset(&self) -> Result<()> {
12129        self.reset_runtime_state().await
12130    }
12131
12132    pub fn max_context_tokens(&self) -> u32 {
12133        self.max_context_tokens
12134    }
12135
12136    pub fn llm_registry(&self) -> &Arc<LLMRegistry> {
12137        &self.llm_registry
12138    }
12139
12140    pub fn state_machine(&self) -> Option<&Arc<StateMachine>> {
12141        self.state_machine.as_ref()
12142    }
12143
12144    pub fn context_manager(&self) -> &Arc<ContextManager> {
12145        &self.context_manager
12146    }
12147
12148    pub fn tool_call_history(&self) -> Vec<ToolCallRecord> {
12149        self.tool_call_history.read().clone()
12150    }
12151
12152    pub fn memory_token_budget(&self) -> Option<&MemoryTokenBudget> {
12153        self.memory_token_budget.as_ref()
12154    }
12155
12156    pub fn parallel_tools_config(&self) -> &ParallelToolsConfig {
12157        &self.parallel_tools
12158    }
12159
12160    pub fn streaming_config(&self) -> &StreamingConfig {
12161        &self.streaming
12162    }
12163
12164    pub fn hooks(&self) -> &Arc<dyn AgentHooks> {
12165        &self.hooks
12166    }
12167
12168    pub fn hitl_engine(&self) -> Option<&HITLEngine> {
12169        self.hitl_engine.as_ref()
12170    }
12171
12172    pub fn approval_handler(&self) -> &Arc<dyn ApprovalHandler> {
12173        &self.approval_handler
12174    }
12175
12176    /// Build a context map with language hints from context_manager for HITL message localization.
12177    fn build_hitl_language_context(&self) -> HashMap<String, Value> {
12178        let mut ctx = HashMap::new();
12179        for key in &["user.language", "input.detected.language", "language"] {
12180            if let Some(val) = self.context_manager.get(key) {
12181                ctx.insert(key.to_string(), val);
12182            }
12183        }
12184        ctx
12185    }
12186
12187    /// Send a HITL check result through the approval flow and return the full ApprovalResult.
12188    async fn request_hitl_approval(&self, check_result: HITLCheckResult) -> Result<ApprovalResult> {
12189        let Some(request) = check_result.into_request() else {
12190            return Ok(ApprovalResult::Approved);
12191        };
12192
12193        self.hooks.on_approval_requested(&request).await;
12194
12195        let timeout = request.timeout;
12196
12197        let raw_result = if let Some(duration) = timeout {
12198            match tokio::time::timeout(
12199                duration,
12200                self.approval_handler.request_approval(request.clone()),
12201            )
12202            .await
12203            {
12204                Ok(result) => result,
12205                Err(_) => ApprovalResult::timeout(),
12206            }
12207        } else {
12208            self.approval_handler
12209                .request_approval(request.clone())
12210                .await
12211        };
12212
12213        self.hooks
12214            .on_approval_result(&request.id, &raw_result)
12215            .await;
12216
12217        let (outcome, effective_result): (ApprovalResolvedOutcome, Result<ApprovalResult>) =
12218            match &raw_result {
12219                ApprovalResult::Approved => (
12220                    ApprovalResolvedOutcome::Approved,
12221                    Ok(ApprovalResult::Approved),
12222                ),
12223                ApprovalResult::Rejected { reason } => (
12224                    ApprovalResolvedOutcome::Rejected {
12225                        reason: reason.clone(),
12226                    },
12227                    Ok(ApprovalResult::Rejected {
12228                        reason: reason.clone(),
12229                    }),
12230                ),
12231                ApprovalResult::Modified { changes } => (
12232                    ApprovalResolvedOutcome::Modified {
12233                        changes: changes.clone(),
12234                    },
12235                    Ok(ApprovalResult::Modified {
12236                        changes: changes.clone(),
12237                    }),
12238                ),
12239                ApprovalResult::Timeout => {
12240                    if let Some(ref engine) = self.hitl_engine {
12241                        match engine.config().on_timeout {
12242                            TimeoutAction::Approve => (
12243                                ApprovalResolvedOutcome::Approved,
12244                                Ok(ApprovalResult::Approved),
12245                            ),
12246                            TimeoutAction::Reject => {
12247                                let reason = Some("Timeout".to_string());
12248                                (
12249                                    ApprovalResolvedOutcome::Rejected {
12250                                        reason: reason.clone(),
12251                                    },
12252                                    Ok(ApprovalResult::Rejected { reason }),
12253                                )
12254                            }
12255                            TimeoutAction::Error => {
12256                                let message = "HITL approval timeout".to_string();
12257                                (
12258                                    ApprovalResolvedOutcome::Error {
12259                                        message: message.clone(),
12260                                    },
12261                                    Err(AgentError::Other(message)),
12262                                )
12263                            }
12264                        }
12265                    } else {
12266                        let reason = Some("Timeout (no engine)".to_string());
12267                        (
12268                            ApprovalResolvedOutcome::Rejected {
12269                                reason: reason.clone(),
12270                            },
12271                            Ok(ApprovalResult::Rejected { reason }),
12272                        )
12273                    }
12274                }
12275            };
12276
12277        self.hooks
12278            .on_approval_resolved(&request, &raw_result, &outcome)
12279            .await;
12280
12281        effective_result
12282    }
12283
12284    pub async fn check_state_hitl(&self, from: Option<&str>, to: &str) -> Result<bool> {
12285        if let Some(ref hitl_engine) = self.hitl_engine {
12286            let hitl_lang_ctx = self.build_hitl_language_context();
12287            let check_result = self
12288                .observe_purpose(
12289                    ObservationPurpose::HitlLocalization,
12290                    hitl_engine.check_state_transition_with_localization(
12291                        from,
12292                        to,
12293                        &hitl_lang_ctx,
12294                        self.approval_handler.as_ref(),
12295                        Some(&self.llm_registry),
12296                    ),
12297                )
12298                .await?;
12299            if check_result.is_required() {
12300                let result = self.request_hitl_approval(check_result).await?;
12301                return Ok(matches!(
12302                    result,
12303                    ApprovalResult::Approved | ApprovalResult::Modified { .. }
12304                ));
12305            }
12306        }
12307        Ok(true)
12308    }
12309
12310    /// Execute multiple tools in parallel
12311    async fn execute_tools_parallel(
12312        &self,
12313        tool_calls: &[ToolCall],
12314    ) -> Vec<(String, Result<String>)> {
12315        let can_run_parallel = tool_calls.iter().all(|tc| {
12316            self.tools
12317                .resolve(&tc.name)
12318                .map(|resolved| resolved.tool.classify_call(&tc.arguments).concurrency_safe)
12319                .unwrap_or(false)
12320        });
12321
12322        if !self.parallel_tools.enabled || tool_calls.len() <= 1 || !can_run_parallel {
12323            let mut results = Vec::new();
12324            for tc in tool_calls {
12325                let result = self
12326                    .observe_purpose(
12327                        current_observation_context()
12328                            .map(|context| context.purpose)
12329                            .unwrap_or_default(),
12330                        self.execute_tool_smart(tc),
12331                    )
12332                    .await;
12333                results.push((tc.id.clone(), result));
12334            }
12335            return results;
12336        }
12337
12338        let chunks: Vec<_> = tool_calls
12339            .chunks(self.parallel_tools.max_parallel)
12340            .collect();
12341
12342        let mut all_results = Vec::new();
12343
12344        for chunk in chunks {
12345            let futures: Vec<_> = chunk
12346                .iter()
12347                .map(|tc| {
12348                    let tc = tc.clone();
12349                    async move {
12350                        let result = self.execute_tool_smart(&tc).await;
12351                        (tc.id.clone(), result)
12352                    }
12353                })
12354                .collect();
12355
12356            let results = futures::future::join_all(futures).await;
12357            all_results.extend(results);
12358        }
12359
12360        all_results
12361    }
12362
12363    /// Streams one serialized root turn and releases its owned gate guard at Done or when the stream is dropped.
12364    ///
12365    /// 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.
12366    pub async fn chat_stream<'a>(
12367        &'a self,
12368        input: &'a str,
12369    ) -> Result<Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>>> {
12370        let RootTurnAdmission {
12371            guard: root_turn_guard,
12372            identity_stack,
12373        } = self.acquire_root_turn().await?;
12374        //
12375        // Streaming readiness runs inside the captured ownership stack while the gate is held so initialization cannot overlap or recursively enter another turn.
12376        //
12377        scope_runtime_gate_identity_stack(&identity_stack, self.init_storage()).await?;
12378        info!(input_len = input.len(), "Starting streaming chat");
12379        let terminal = new_runtime_stream_terminal_slot();
12380        let inner = self.run_loop_stream(input, terminal);
12381        let observation_context = self.build_observation_context(None);
12382        let stream: Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> =
12383            Box::pin(async_stream::stream! {
12384                let mut root_turn_guard = Some(root_turn_guard);
12385                let mut inner = inner;
12386                loop {
12387                    let next = scope_runtime_gate_identity_stack(&identity_stack, async {
12388                        if let Some(context) = observation_context.as_ref() {
12389                            with_observation_context(context.clone(), inner.next()).await
12390                        } else {
12391                            inner.next().await
12392                        }
12393                    })
12394                    .await;
12395                    match next {
12396                        Some(StreamChunk::Done {}) => {
12397                            while scope_runtime_gate_identity_stack(&identity_stack, async {
12398                                if let Some(context) = observation_context.as_ref() {
12399                                    with_observation_context(context.clone(), inner.next())
12400                                        .await
12401                                        .is_some()
12402                                } else {
12403                                    inner.next().await.is_some()
12404                                }
12405                            })
12406                            .await
12407                            {}
12408                            if observation_context.is_some() {
12409                                scope_runtime_gate_identity_stack(
12410                                    &identity_stack,
12411                                    self.export_observability_if_configured(),
12412                                )
12413                                .await;
12414                            }
12415                            drop(root_turn_guard.take());
12416                            yield StreamChunk::Done {};
12417                            return;
12418                        }
12419                        Some(chunk) => yield chunk,
12420                        None => {
12421                            if observation_context.is_some() {
12422                                scope_runtime_gate_identity_stack(
12423                                    &identity_stack,
12424                                    self.export_observability_if_configured(),
12425                                )
12426                                .await;
12427                            }
12428                            drop(root_turn_guard.take());
12429                            return;
12430                        }
12431                    }
12432                }
12433            });
12434        Ok(stream)
12435    }
12436
12437    /// Streams one serialized root turn and releases its owned gate guard at the authoritative terminal event or on drop.
12438    ///
12439    /// 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.
12440    pub async fn chat_stream_events<'a>(
12441        &'a self,
12442        input: &'a str,
12443    ) -> Result<Pin<Box<dyn Stream<Item = AgentStreamEvent> + Send + 'a>>> {
12444        let RootTurnAdmission {
12445            guard: root_turn_guard,
12446            identity_stack,
12447        } = self.acquire_root_turn().await?;
12448        //
12449        // Streaming readiness runs inside the captured ownership stack while the gate is held so initialization cannot overlap or recursively enter another turn.
12450        //
12451        scope_runtime_gate_identity_stack(&identity_stack, self.init_storage()).await?;
12452        info!(input_len = input.len(), "Starting streaming chat events");
12453        let terminal = new_runtime_stream_terminal_slot();
12454        let mut inner = self.run_loop_stream(input, Arc::clone(&terminal));
12455        let observation_context = self.build_observation_context(None);
12456        let stream: Pin<Box<dyn Stream<Item = AgentStreamEvent> + Send + 'a>> =
12457            Box::pin(async_stream::stream! {
12458                let mut root_turn_guard = Some(root_turn_guard);
12459                loop {
12460                    let next = scope_runtime_gate_identity_stack(&identity_stack, async {
12461                        if let Some(context) = observation_context.as_ref() {
12462                            with_observation_context(context.clone(), inner.next()).await
12463                        } else {
12464                            inner.next().await
12465                        }
12466                    })
12467                    .await;
12468                    match next {
12469                        Some(StreamChunk::Done {}) => {
12470                            let terminal_event = { terminal.write().take() };
12471                            if let Some(response) = terminal_event {
12472                                while scope_runtime_gate_identity_stack(&identity_stack, async {
12473                                    if let Some(context) = observation_context.as_ref() {
12474                                        with_observation_context(context.clone(), inner.next())
12475                                            .await
12476                                            .is_some()
12477                                    } else {
12478                                        inner.next().await.is_some()
12479                                    }
12480                                })
12481                                .await
12482                                {}
12483                                if observation_context.is_some() {
12484                                    scope_runtime_gate_identity_stack(
12485                                        &identity_stack,
12486                                        self.export_observability_if_configured(),
12487                                    )
12488                                    .await;
12489                                }
12490                                drop(root_turn_guard.take());
12491                                yield AgentStreamEvent::Final(response);
12492                                return;
12493                            }
12494                        }
12495                        Some(StreamChunk::Error { message }) => {
12496                            let finalized = { terminal.read().is_some() };
12497                            if finalized {
12498                                continue;
12499                            }
12500                            while scope_runtime_gate_identity_stack(&identity_stack, async {
12501                                if let Some(context) = observation_context.as_ref() {
12502                                    with_observation_context(context.clone(), inner.next())
12503                                        .await
12504                                        .is_some()
12505                                } else {
12506                                    inner.next().await.is_some()
12507                                }
12508                            })
12509                            .await
12510                            {}
12511                            if observation_context.is_some() {
12512                                scope_runtime_gate_identity_stack(
12513                                    &identity_stack,
12514                                    self.export_observability_if_configured(),
12515                                )
12516                                .await;
12517                            }
12518                            drop(root_turn_guard.take());
12519                            yield AgentStreamEvent::Chunk(StreamChunk::Error { message });
12520                            return;
12521                        }
12522                        Some(chunk) => yield AgentStreamEvent::Chunk(chunk),
12523                        None => {
12524                            if observation_context.is_some() {
12525                                scope_runtime_gate_identity_stack(
12526                                    &identity_stack,
12527                                    self.export_observability_if_configured(),
12528                                )
12529                                .await;
12530                            }
12531                            drop(root_turn_guard.take());
12532                            return;
12533                        }
12534                    }
12535                }
12536            });
12537        Ok(stream)
12538    }
12539}
12540
12541#[async_trait]
12542impl ToolInvoker for RuntimeAgent {
12543    async fn invoke_tool(&self, request: ToolExecutionRequest) -> Result<ToolExecutionRecord> {
12544        self.execute_tool_record(request).await
12545    }
12546}
12547
12548#[async_trait]
12549impl Agent for RuntimeAgent {
12550    /// Runs one blocking external root turn with task-local ownership visible through finalization, hooks, orchestration, and export.
12551    async fn chat(&self, input: &str) -> Result<AgentResponse> {
12552        let RootTurnAdmission {
12553            guard,
12554            identity_stack,
12555        } = self.acquire_root_turn().await?;
12556        let result = scope_runtime_gate_identity_stack(&identity_stack, async {
12557            let result = if let Some(context) = self.build_observation_context(None) {
12558                with_observation_context(context, self.run_loop(input)).await
12559            } else {
12560                self.run_loop(input).await
12561            };
12562            self.export_observability_if_configured().await;
12563            result
12564        })
12565        .await;
12566        drop(guard);
12567        result
12568    }
12569
12570    fn info(&self) -> AgentInfo {
12571        self.info.clone()
12572    }
12573
12574    /// Resets the runtime without leaving clarification or skill ownership behind.
12575    async fn reset(&self) -> Result<()> {
12576        self.reset_runtime_state().await
12577    }
12578}
12579
12580//
12581// Render a concurrent input template using direct minijinja.
12582// Same approach as pipeline's render_stage_template so variables are top-level.
12583//
12584// Available variables:
12585//   {{ user_input }}    - the user's actual message
12586//   {{ context.<key> }} - values from the context manager
12587//
12588/// Builds safe runtime tags for background maintenance lifecycle events.
12589fn background_maintenance_tags(
12590    label: &str,
12591    stage: &str,
12592    reason: Option<&str>,
12593    policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
12594) -> HashMap<String, String> {
12595    let mut tags = HashMap::new();
12596    tags.insert("runtime.background".to_string(), "true".to_string());
12597    tags.insert("runtime.maintenance".to_string(), label.to_string());
12598    tags.insert("runtime.maintenance_stage".to_string(), stage.to_string());
12599    if let Some(policy) = policy {
12600        tags.insert(
12601            "runtime.await_before_next_turn".to_string(),
12602            await_before_next_turn_label(policy.await_before_next_turn).to_string(),
12603        );
12604        tags.insert(
12605            "runtime.maintenance_mode".to_string(),
12606            maintenance_mode_label(policy.mode).to_string(),
12607        );
12608    }
12609    if let Some(reason) = reason {
12610        tags.insert("runtime.reason".to_string(), reason.to_string());
12611    }
12612    tags
12613}
12614
12615fn await_before_next_turn_label(policy: AwaitBeforeNextTurn) -> &'static str {
12616    match policy {
12617        AwaitBeforeNextTurn::Never => "never",
12618        AwaitBeforeNextTurn::SameActor => "same_actor",
12619        AwaitBeforeNextTurn::Always => "always",
12620    }
12621}
12622
12623fn maintenance_mode_label(mode: MaintenanceMode) -> &'static str {
12624    match mode {
12625        MaintenanceMode::InlineSerial => "inline_serial",
12626        MaintenanceMode::InlineParallel => "inline_parallel",
12627        MaintenanceMode::Background => "background",
12628    }
12629}
12630
12631/// Records a background maintenance lifecycle event when observability is enabled.
12632fn record_background_maintenance_event(
12633    manager: Option<&Arc<ObservabilityManager>>,
12634    label: &str,
12635    status: EventStatus,
12636    duration_ms: u64,
12637    stage: &str,
12638    reason: Option<String>,
12639    policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
12640) {
12641    if let Some(manager) = manager {
12642        manager.record_lifecycle_event(
12643            EventType::MemoryOperation {
12644                operation: format!("{}_background_{}", label, stage),
12645            },
12646            ObservationPurpose::Other(format!("{}_maintenance", label)),
12647            status,
12648            duration_ms,
12649            background_maintenance_tags(label, stage, reason.as_deref(), policy),
12650            None,
12651        );
12652    }
12653}
12654
12655fn effective_maintenance_mode(mode: MaintenanceMode, force_parallel: bool) -> MaintenanceMode {
12656    if force_parallel && matches!(mode, MaintenanceMode::InlineSerial) {
12657        MaintenanceMode::InlineParallel
12658    } else {
12659        mode
12660    }
12661}
12662
12663fn observation_purpose_for_process(hint: ProcessPurposeHint) -> ObservationPurpose {
12664    match hint {
12665        ProcessPurposeHint::Detect => ObservationPurpose::ProcessDetect,
12666        ProcessPurposeHint::Extract => ObservationPurpose::ProcessExtract,
12667        ProcessPurposeHint::Validate => ObservationPurpose::ProcessValidate,
12668        ProcessPurposeHint::Transform | ProcessPurposeHint::Other => {
12669            ObservationPurpose::ProcessTransform
12670        }
12671    }
12672}
12673
12674fn new_tool_resource_locks() -> ToolResourceLocks {
12675    Arc::new(RwLock::new(HashMap::new()))
12676}
12677
12678//
12679// Path-bound mutations share one conservative lock so aliases and parent-child paths cannot bypass serialization.
12680// Exact domain keys remain available for non-path side effects, and unbound effects share one fallback lock.
12681//
12682fn tool_resource_lock_keys(
12683    _canonical_id: &str,
12684    args: &Value,
12685    bindings: &ai_agents_core::ToolPolicyBindings,
12686    classification: &ai_agents_core::ToolCallClassification,
12687) -> Vec<String> {
12688    if classification.concurrency_safe {
12689        return Vec::new();
12690    }
12691
12692    let mut keys = Vec::new();
12693    let mut has_path_resource = false;
12694    for binding in &bindings.path_fields {
12695        let value = value_at_argument_path(args, &binding.field)
12696            .cloned()
12697            .or_else(|| {
12698                binding
12699                    .default_path
12700                    .as_ref()
12701                    .map(|path| Value::String(path.clone()))
12702            });
12703        if let Some(value) = value {
12704            collect_resource_strings(&value, |_| {
12705                has_path_resource = true;
12706            });
12707        }
12708    }
12709    for binding in &bindings.domain_fields {
12710        if let Some(value) = value_at_argument_path(args, &binding.field) {
12711            collect_resource_strings(value, |domain| {
12712                let normalized = if binding.is_url {
12713                    normalized_url_resource_key(domain)
12714                } else {
12715                    domain.trim().trim_end_matches('.').to_ascii_lowercase()
12716                };
12717                keys.push(format!("domain:{}", normalized));
12718            });
12719        }
12720    }
12721    for binding in &bindings.command_fields {
12722        if !matches!(binding.kind, ai_agents_core::CommandBindingKind::Cwd) {
12723            continue;
12724        }
12725        if let Some(value) = value_at_argument_path(args, &binding.field) {
12726            collect_resource_strings(value, |_| {
12727                has_path_resource = true;
12728            });
12729        }
12730    }
12731    if has_path_resource {
12732        keys.push("path-mutation:global".to_string());
12733    }
12734    if keys.is_empty() {
12735        keys.push("side-effect:unbound".to_string());
12736    }
12737    keys.sort();
12738    keys.dedup();
12739    keys
12740}
12741
12742fn value_at_argument_path<'a>(value: &'a Value, field: &str) -> Option<&'a Value> {
12743    let mut current = value;
12744    for segment in field.split('.') {
12745        if segment.is_empty() {
12746            return None;
12747        }
12748        current = current.get(segment)?;
12749    }
12750    Some(current)
12751}
12752
12753fn collect_resource_strings(value: &Value, mut collect: impl FnMut(&str)) {
12754    match value {
12755        Value::String(value) => collect(value),
12756        Value::Array(values) => {
12757            for value in values {
12758                if let Some(value) = value.as_str() {
12759                    collect(value);
12760                }
12761            }
12762        }
12763        _ => {}
12764    }
12765}
12766
12767fn normalized_url_resource_key(value: &str) -> String {
12768    let value = value.trim();
12769    let Some((scheme, remainder)) = value.split_once("://") else {
12770        return value.to_ascii_lowercase();
12771    };
12772    let authority_end = remainder.find(['/', '?', '#']).unwrap_or(remainder.len());
12773    let (authority, suffix) = remainder.split_at(authority_end);
12774    format!(
12775        "{}://{}{}",
12776        scheme.to_ascii_lowercase(),
12777        authority.to_ascii_lowercase(),
12778        suffix
12779    )
12780}
12781
12782fn render_concurrent_template(
12783    template: &str,
12784    user_input: &str,
12785    context_values: &std::collections::HashMap<String, serde_json::Value>,
12786) -> Result<String> {
12787    let mut env = minijinja::Environment::new();
12788    env.add_template("concurrent", template)
12789        .map_err(|e| AgentError::Other(format!("Concurrent template parse error: {}", e)))?;
12790
12791    let mut ctx = std::collections::BTreeMap::new();
12792    ctx.insert("user_input".to_string(), minijinja::Value::from(user_input));
12793
12794    // Expose context manager values under {{ context.<key> }}.
12795    let context_obj = minijinja::Value::from_serialize(context_values);
12796    ctx.insert("context".to_string(), context_obj);
12797
12798    let tmpl = env
12799        .get_template("concurrent")
12800        .map_err(|e| AgentError::Other(format!("Concurrent template error: {}", e)))?;
12801
12802    tmpl.render(minijinja::Value::from_serialize(&ctx))
12803        .map_err(|e| AgentError::Other(format!("Concurrent template render error: {}", e)))
12804}
12805
12806#[cfg(test)]
12807mod tests {
12808    use super::*;
12809    use crate::AgentBuilder;
12810    use ai_agents_core::{LLMChunk, LLMConfig, LLMError, LLMFeature, Tool};
12811    use ai_agents_llm::mock::MockLLMProvider;
12812    use ai_agents_skills::{SkillDefinition, SkillStep};
12813    use ai_agents_tools::{
12814        CalculatorTool, CopyPathTool, DeletePathTool, FileWriteTool, MovePathTool,
12815        WebFetchResolver, WebFetchTool, WebFetchTransport, WebFetchTransportRequest,
12816        WebFetchTransportResponse,
12817    };
12818
12819    fn mock_with_response(response: &str) -> MockLLMProvider {
12820        let mut mock = MockLLMProvider::new("test");
12821        mock.set_response(response);
12822        mock
12823    }
12824
12825    fn mock_with_responses(responses: Vec<&str>) -> MockLLMProvider {
12826        let mut mock = MockLLMProvider::new("test");
12827        mock.set_responses(responses.into_iter().map(String::from).collect(), true);
12828        mock
12829    }
12830
12831    /// Builds a two-state fixture so confirmation ownership can be invalidated by transition.
12832    fn disambiguation_state_machine(
12833        state_enabled: Option<bool>,
12834        require_confirmation: bool,
12835    ) -> Arc<StateMachine> {
12836        let definition = ai_agents_state::StateDefinition {
12837            prompt: Some("Handle the resolved request.".to_string()),
12838            disambiguation: Some(ai_agents_disambiguation::StateDisambiguationOverride {
12839                enabled: state_enabled,
12840                require_confirmation,
12841                ..Default::default()
12842            }),
12843            ..Default::default()
12844        };
12845        let review = ai_agents_state::StateDefinition {
12846            prompt: Some("Review a fresh request.".to_string()),
12847            ..Default::default()
12848        };
12849        Arc::new(
12850            StateMachine::new(ai_agents_state::StateConfig {
12851                initial: "active".to_string(),
12852                states: std::collections::HashMap::from([
12853                    ("active".to_string(), definition),
12854                    ("review".to_string(), review),
12855                ]),
12856                global_transitions: Vec::new(),
12857                fallback: None,
12858                max_no_transition: None,
12859                regenerate_on_transition: true,
12860            })
12861            .unwrap(),
12862        )
12863    }
12864
12865    /// Builds a state-aware disambiguation fixture without skills.
12866    fn state_disambiguation_agent(
12867        responses: Vec<&str>,
12868        manager_enabled: bool,
12869        state_enabled: Option<bool>,
12870        require_confirmation: bool,
12871    ) -> (RuntimeAgent, MockLLMProvider) {
12872        state_disambiguation_agent_with_skills(
12873            responses,
12874            manager_enabled,
12875            state_enabled,
12876            require_confirmation,
12877            Vec::new(),
12878        )
12879    }
12880
12881    /// Builds a state-aware disambiguation fixture with optional real skill routing.
12882    fn state_disambiguation_agent_with_skills(
12883        responses: Vec<&str>,
12884        manager_enabled: bool,
12885        state_enabled: Option<bool>,
12886        require_confirmation: bool,
12887        skills: Vec<SkillDefinition>,
12888    ) -> (RuntimeAgent, MockLLMProvider) {
12889        let mut mock = MockLLMProvider::new("state-confirmation");
12890        mock.set_responses(responses.into_iter().map(String::from).collect(), false);
12891        let observed = mock.clone();
12892        let agent = AgentBuilder::new()
12893            .system_prompt("Handle requests.")
12894            .llm(Arc::new(mock.clone()))
12895            .llm_alias("router", Arc::new(mock))
12896            .state_machine(disambiguation_state_machine(
12897                state_enabled,
12898                require_confirmation,
12899            ))
12900            .skills(skills)
12901            .build()
12902            .unwrap()
12903            .with_disambiguation(DisambiguationConfig {
12904                enabled: manager_enabled,
12905                ..Default::default()
12906            });
12907        (agent, observed)
12908    }
12909
12910    /// Defines a prompt skill whose provider call proves committed execution.
12911    fn confirmation_skill() -> SkillDefinition {
12912        SkillDefinition {
12913            id: "send_report".to_string(),
12914            description: "Send a report after clarification".to_string(),
12915            trigger: "When the user asks to send a report".to_string(),
12916            steps: vec![SkillStep::Prompt {
12917                prompt: "Execute confirmed report skill for: {{ input }}".to_string(),
12918                llm: None,
12919            }],
12920            reasoning: None,
12921            reflection: None,
12922            disambiguation: Some(ai_agents_disambiguation::SkillDisambiguationOverride {
12923                enabled: Some(true),
12924                ..Default::default()
12925            }),
12926        }
12927    }
12928
12929    /// Counts committed skill prompt calls without relying on final response wording.
12930    fn confirmation_skill_call_count(observed: &MockLLMProvider) -> usize {
12931        observed
12932            .call_history()
12933            .iter()
12934            .filter(|call| {
12935                call.messages
12936                    .iter()
12937                    .any(|message| message.content.contains("Execute confirmed report skill"))
12938            })
12939            .count()
12940    }
12941
12942    struct BlockingRuntimeConfirmationObserver {
12943        entered: tokio::sync::Barrier,
12944        release: tokio::sync::Notify,
12945    }
12946
12947    impl BlockingRuntimeConfirmationObserver {
12948        fn new() -> Self {
12949            Self {
12950                entered: tokio::sync::Barrier::new(2),
12951                release: tokio::sync::Notify::new(),
12952            }
12953        }
12954    }
12955
12956    struct ResetOnTransitionHooks {
12957        agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
12958        invoked: AtomicBool,
12959    }
12960
12961    #[async_trait]
12962    impl AgentHooks for ResetOnTransitionHooks {
12963        async fn on_state_transition(&self, _from: Option<&str>, _to: &str, _reason: &str) {
12964            if self.invoked.swap(true, Ordering::SeqCst) {
12965                return;
12966            }
12967            let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
12968            if let Some(agent) = agent {
12969                agent.reset().await.unwrap();
12970            }
12971        }
12972    }
12973
12974    impl ClarificationObserver for BlockingRuntimeConfirmationObserver {
12975        fn observe_question<'a>(
12976            &'a self,
12977            future: ClarificationQuestionFuture<'a>,
12978        ) -> ClarificationQuestionFuture<'a> {
12979            future
12980        }
12981
12982        fn observe_parse<'a>(
12983            &'a self,
12984            future: ClarificationParseFuture<'a>,
12985        ) -> ClarificationParseFuture<'a> {
12986            future
12987        }
12988
12989        fn observe_confirmation_parse<'a>(
12990            &'a self,
12991            future: ConfirmationParseFuture<'a>,
12992        ) -> ConfirmationParseFuture<'a> {
12993            Box::pin(async move {
12994                self.entered.wait().await;
12995                self.release.notified().await;
12996                future.await
12997            })
12998        }
12999    }
13000
13001    #[tokio::test]
13002    async fn state_confirmation_blocks_redispatch_until_explicit_agreement() {
13003        let (agent, observed) = state_disambiguation_agent(
13004            vec![
13005                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13006                r#"{"question":"What should I send?","options":null}"#,
13007                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13008                r#"{"question":"Should I send the report to Ada?"}"#,
13009                r#"{"status":"confirmed"}"#,
13010                "Request executed.",
13011            ],
13012            true,
13013            None,
13014            true,
13015        );
13016
13017        let clarification = agent.chat("Send it").await.unwrap();
13018        assert_eq!(clarification.content, "What should I send?");
13019        assert_eq!(observed.call_count(), 2);
13020
13021        let confirmation = agent.chat("The report to Ada").await.unwrap();
13022        assert_eq!(confirmation.content, "Should I send the report to Ada?");
13023        assert_eq!(
13024            confirmation
13025                .metadata
13026                .as_ref()
13027                .and_then(|metadata| metadata.get("disambiguation"))
13028                .and_then(|metadata| metadata.get("status"))
13029                .and_then(Value::as_str),
13030            Some("awaiting_confirmation")
13031        );
13032        assert_eq!(observed.call_count(), 4);
13033
13034        let completed = agent.chat("Yes").await.unwrap();
13035        assert_eq!(completed.content, "Request executed.");
13036        assert_eq!(observed.call_count(), 6);
13037    }
13038
13039    #[tokio::test]
13040    async fn streaming_state_confirmation_ends_the_turn_before_redispatch() {
13041        let (agent, observed) = state_disambiguation_agent(
13042            vec![
13043                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13044                r#"{"question":"What should I send?","options":null}"#,
13045                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13046                r#"{"question":"Should I send the report to Ada?"}"#,
13047                r#"{"status":"confirmed"}"#,
13048                "Request executed.",
13049            ],
13050            true,
13051            None,
13052            true,
13053        );
13054
13055        let mut clarification_stream = agent.chat_stream("Send it").await.unwrap();
13056        let mut clarification = String::new();
13057        while let Some(chunk) = clarification_stream.next().await {
13058            match chunk {
13059                StreamChunk::Content { text } => clarification.push_str(&text),
13060                StreamChunk::Done {} => break,
13061                StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
13062                _ => {}
13063            }
13064        }
13065        assert_eq!(clarification, "What should I send?");
13066        assert_eq!(observed.call_count(), 2);
13067
13068        let mut confirmation_stream = agent.chat_stream_events("The report to Ada").await.unwrap();
13069        let mut confirmation = None;
13070        while let Some(event) = confirmation_stream.next().await {
13071            match event {
13072                AgentStreamEvent::Final(response) => confirmation = Some(response),
13073                AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
13074                    panic!("unexpected stream error: {message}")
13075                }
13076                AgentStreamEvent::Chunk(_) => {}
13077            }
13078        }
13079        let confirmation = confirmation.expect("confirmation must finalize");
13080        assert_eq!(confirmation.content, "Should I send the report to Ada?");
13081        assert_eq!(
13082            confirmation
13083                .metadata
13084                .as_ref()
13085                .and_then(|metadata| metadata.get("disambiguation"))
13086                .and_then(|metadata| metadata.get("status"))
13087                .and_then(Value::as_str),
13088            Some("awaiting_confirmation")
13089        );
13090        assert_eq!(observed.call_count(), 4);
13091
13092        let mut completed_stream = agent.chat_stream("Yes").await.unwrap();
13093        let mut completed = String::new();
13094        while let Some(chunk) = completed_stream.next().await {
13095            match chunk {
13096                StreamChunk::Content { text } => completed.push_str(&text),
13097                StreamChunk::Done {} => break,
13098                StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
13099                _ => {}
13100            }
13101        }
13102        assert_eq!(completed, "Request executed.");
13103        assert_eq!(observed.call_count(), 6);
13104    }
13105
13106    /// Confirms a returned legacy stream blocks a blocking turn until drop and the event terminal releases the same gate.
13107    #[tokio::test]
13108    async fn root_turn_gate_serializes_blocking_and_streaming_entry_points() {
13109        let (complete_entered, mut complete_events) = tokio::sync::mpsc::unbounded_channel();
13110        let agent = Arc::new(
13111            AgentBuilder::new()
13112                .system_prompt("Serialize root turns.")
13113                .llm(Arc::new(RootTurnProbeProvider { complete_entered }))
13114                .build()
13115                .unwrap(),
13116        );
13117        let blocking_agent = Arc::clone(&agent);
13118
13119        let legacy_stream = agent.chat_stream("stream owner").await.unwrap();
13120        assert!(agent.root_turn_gate.try_lock().is_err());
13121        let blocking = tokio::spawn(async move { blocking_agent.chat("blocked").await.unwrap() });
13122        assert!(
13123            tokio::time::timeout(std::time::Duration::from_millis(50), complete_events.recv())
13124                .await
13125                .is_err(),
13126            "blocking turn reached the provider while the legacy stream owned the root gate"
13127        );
13128
13129        drop(legacy_stream);
13130        assert_eq!(
13131            tokio::time::timeout(std::time::Duration::from_secs(2), complete_events.recv())
13132                .await
13133                .expect("blocking turn did not enter after stream drop"),
13134            Some(())
13135        );
13136        let response = tokio::time::timeout(std::time::Duration::from_secs(2), blocking)
13137            .await
13138            .expect("blocking turn did not finish after stream drop")
13139            .unwrap();
13140        assert_eq!(response.content, "blocking complete");
13141
13142        let mut event_stream = agent.chat_stream_events("event terminal").await.unwrap();
13143        assert!(agent.root_turn_gate.try_lock().is_err());
13144        let mut saw_final = false;
13145        while let Some(event) = event_stream.next().await {
13146            if matches!(event, AgentStreamEvent::Final(_)) {
13147                saw_final = true;
13148                break;
13149            }
13150        }
13151        assert!(saw_final);
13152        assert!(
13153            agent.root_turn_gate.try_lock().is_ok(),
13154            "authoritative terminal event retained the root gate"
13155        );
13156    }
13157
13158    /// Confirms on_response receives a fail-fast error instead of deadlocking on the same runtime gate.
13159    #[tokio::test]
13160    async fn response_hook_rejects_same_runtime_chat_reentry() {
13161        let hooks = Arc::new(ResponseChatHooks {
13162            target: parking_lot::Mutex::new(None),
13163            invoked: AtomicBool::new(false),
13164            nested_result: parking_lot::Mutex::new(None),
13165        });
13166        let agent = Arc::new(
13167            AgentBuilder::new()
13168                .system_prompt("Reject response hook reentry.")
13169                .llm(Arc::new(mock_with_response("outer response")))
13170                .hooks(hooks.clone())
13171                .build()
13172                .unwrap(),
13173        );
13174        *hooks.target.lock() = Some(Arc::downgrade(&agent));
13175
13176        let response = tokio::time::timeout(
13177            std::time::Duration::from_secs(2),
13178            agent.chat("outer request"),
13179        )
13180        .await
13181        .expect("same-runtime response hook reentry must fail without deadlocking")
13182        .unwrap();
13183
13184        assert_eq!(response.content, "outer response");
13185        let nested_result = hooks
13186            .nested_result
13187            .lock()
13188            .clone()
13189            .expect("response hook must record its nested call");
13190        let error = nested_result.expect_err("same-runtime nested chat must be rejected");
13191        assert!(error.contains("reentrant root turn ownership"));
13192    }
13193
13194    /// Confirms a different runtime gate can nest while an A to B to A ownership cycle is rejected at A.
13195    #[tokio::test]
13196    async fn root_turn_gate_allows_nested_runtime_and_rejects_cycles() {
13197        let agent_a = AgentBuilder::new()
13198            .system_prompt("Runtime A.")
13199            .llm(Arc::new(mock_with_response("response A")))
13200            .build()
13201            .unwrap();
13202        let agent_b = AgentBuilder::new()
13203            .system_prompt("Runtime B.")
13204            .llm(Arc::new(mock_with_response("response B")))
13205            .build()
13206            .unwrap();
13207        let RootTurnAdmission {
13208            guard: guard_a,
13209            identity_stack: stack_a,
13210        } = agent_a.acquire_root_turn().await.unwrap();
13211
13212        let cycle_error = scope_runtime_gate_identity_stack(&stack_a, async {
13213            let RootTurnAdmission {
13214                guard: guard_b,
13215                identity_stack: stack_b,
13216            } = agent_b
13217                .acquire_root_turn()
13218                .await
13219                .expect("runtime B must acquire a different gate");
13220            let result =
13221                scope_runtime_gate_identity_stack(&stack_b, agent_a.acquire_root_turn()).await;
13222            drop(guard_b);
13223            match result {
13224                Err(error) => error,
13225                Ok(_) => panic!("runtime A accepted a repeated gate identity"),
13226            }
13227        })
13228        .await;
13229        drop(guard_a);
13230
13231        assert!(
13232            cycle_error
13233                .to_string()
13234                .contains("reentrant root turn ownership")
13235        );
13236    }
13237
13238    /// Confirms concurrent orchestration propagates root ancestry so an A to B to A cycle fails before waiting on A.
13239    #[tokio::test]
13240    async fn concurrent_orchestration_propagates_root_gate_ancestry() {
13241        let registry = Arc::new(crate::spawner::AgentRegistry::new());
13242        let hooks_a = Arc::new(ConcurrentResponseHooks {
13243            registry: Arc::downgrade(&registry),
13244            child_id: "runtime-b".to_string(),
13245            invoked: AtomicBool::new(false),
13246            nested_result: parking_lot::Mutex::new(None),
13247        });
13248        let hooks_b = Arc::new(ResponseChatHooks {
13249            target: parking_lot::Mutex::new(None),
13250            invoked: AtomicBool::new(false),
13251            nested_result: parking_lot::Mutex::new(None),
13252        });
13253        let agent_a = AgentBuilder::new()
13254            .system_prompt("Runtime A dispatches runtime B concurrently.")
13255            .llm(Arc::new(mock_with_response("response A")))
13256            .hooks(hooks_a.clone())
13257            .build()
13258            .unwrap();
13259        let agent_b = AgentBuilder::new()
13260            .system_prompt("Runtime B attempts to re-enter runtime A.")
13261            .llm(Arc::new(mock_with_response("response B")))
13262            .hooks(hooks_b.clone())
13263            .build()
13264            .unwrap();
13265        let spec_a = crate::spec::AgentSpec {
13266            name: "runtime-a".to_string(),
13267            system_prompt: "Runtime A dispatches runtime B concurrently.".to_string(),
13268            ..crate::spec::AgentSpec::default()
13269        };
13270        let spec_b = crate::spec::AgentSpec {
13271            name: "runtime-b".to_string(),
13272            system_prompt: "Runtime B attempts to re-enter runtime A.".to_string(),
13273            ..crate::spec::AgentSpec::default()
13274        };
13275        registry
13276            .register(crate::spawner::SpawnedAgent::from_runtime(
13277                "runtime-a".to_string(),
13278                agent_a,
13279                spec_a,
13280            ))
13281            .await
13282            .unwrap();
13283        registry
13284            .register(crate::spawner::SpawnedAgent::from_runtime(
13285                "runtime-b".to_string(),
13286                agent_b,
13287                spec_b,
13288            ))
13289            .await
13290            .unwrap();
13291        let runtime_a = registry.get("runtime-a").unwrap();
13292        *hooks_b.target.lock() = Some(Arc::downgrade(&runtime_a));
13293
13294        let response = tokio::time::timeout(
13295            std::time::Duration::from_secs(2),
13296            runtime_a.chat("outer concurrent request"),
13297        )
13298        .await
13299        .expect("concurrent orchestration cycle must fail without deadlocking")
13300        .unwrap();
13301
13302        assert_eq!(response.content, "response A");
13303        let child_result = hooks_a
13304            .nested_result
13305            .lock()
13306            .clone()
13307            .expect("runtime A hook must record runtime B completion");
13308        assert_eq!(child_result.unwrap(), "response B");
13309        let cycle_result = hooks_b
13310            .nested_result
13311            .lock()
13312            .clone()
13313            .expect("runtime B hook must record runtime A reentry");
13314        assert!(
13315            cycle_result
13316                .expect_err("runtime A accepted a repeated gate identity")
13317                .contains("reentrant root turn ownership")
13318        );
13319    }
13320
13321    /// Confirms that a pending skill remains inert until confirmation and then runs once.
13322    #[tokio::test]
13323    async fn confirmed_skill_route_executes_exactly_once() {
13324        let (agent, observed) = state_disambiguation_agent_with_skills(
13325            vec![
13326                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13327                "send_report",
13328                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13329                r#"{"question":"What should I send?","options":null}"#,
13330                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13331                r#"{"question":"Should I send the report to Ada?"}"#,
13332                r#"{"status":"confirmed"}"#,
13333                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"resolved","what_is_unclear":[],"detected_language":"en"}"#,
13334                "Report skill executed.",
13335            ],
13336            true,
13337            None,
13338            true,
13339            vec![confirmation_skill()],
13340        );
13341
13342        let clarification = agent.chat("Send it").await.unwrap();
13343        assert_eq!(clarification.content, "What should I send?");
13344        assert_eq!(confirmation_skill_call_count(&observed), 0);
13345
13346        let confirmation = agent.chat("The report to Ada").await.unwrap();
13347        assert_eq!(confirmation.content, "Should I send the report to Ada?");
13348        assert_eq!(
13349            confirmation
13350                .metadata
13351                .as_ref()
13352                .and_then(|metadata| metadata.get("disambiguation"))
13353                .and_then(|metadata| metadata.get("status"))
13354                .and_then(Value::as_str),
13355            Some("awaiting_confirmation")
13356        );
13357        assert_eq!(confirmation_skill_call_count(&observed), 0);
13358
13359        let completed = agent.chat("Yes").await.unwrap();
13360        assert_eq!(completed.content, "Report skill executed.");
13361        assert_eq!(confirmation_skill_call_count(&observed), 1);
13362        assert!(agent.pending_skill_id.read().is_none());
13363        let messages = agent.memory.get_messages(None).await.unwrap();
13364        assert!(!messages.iter().any(|message| message.content == "Yes"));
13365    }
13366
13367    /// Preserves a second clarification response after confirmation instead of overwriting its metadata.
13368    #[tokio::test]
13369    async fn confirmed_skill_recheck_preserves_new_clarification_metadata() {
13370        let (agent, observed) = state_disambiguation_agent_with_skills(
13371            vec![
13372                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13373                "send_report",
13374                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13375                r#"{"question":"What should I send?","options":null}"#,
13376                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13377                r#"{"question":"Should I send the report to Ada?"}"#,
13378                r#"{"status":"confirmed"}"#,
13379                r#"{"is_ambiguous":true,"confidence":0.3,"ambiguity_type":"missing_parameters","reasoning":"timing missing","what_is_unclear":["timing"],"detected_language":"en"}"#,
13380                r#"{"question":"When should I send it?","options":null}"#,
13381            ],
13382            true,
13383            None,
13384            true,
13385            vec![confirmation_skill()],
13386        );
13387
13388        agent.chat("Send it").await.unwrap();
13389        agent.chat("The report to Ada").await.unwrap();
13390        let follow_up = agent.chat("Yes").await.unwrap();
13391
13392        assert_eq!(follow_up.content, "When should I send it?");
13393        let metadata = follow_up
13394            .metadata
13395            .as_ref()
13396            .and_then(|metadata| metadata.get("disambiguation"))
13397            .unwrap();
13398        assert_eq!(
13399            metadata.get("status").and_then(Value::as_str),
13400            Some("awaiting_clarification")
13401        );
13402        assert_eq!(
13403            metadata.get("skill_id").and_then(Value::as_str),
13404            Some("send_report")
13405        );
13406        assert!(metadata.get("detection").is_some());
13407        assert_eq!(confirmation_skill_call_count(&observed), 0);
13408    }
13409
13410    /// Confirms rejection clears a pending skill without executing its prompt.
13411    #[tokio::test]
13412    async fn rejected_skill_confirmation_never_executes() {
13413        let (agent, observed) = state_disambiguation_agent_with_skills(
13414            vec![
13415                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13416                "send_report",
13417                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13418                r#"{"question":"What should I send?","options":null}"#,
13419                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13420                r#"{"question":"Should I send the report to Ada?"}"#,
13421                r#"{"status":"rejected"}"#,
13422                "Confirmation rejected.",
13423            ],
13424            true,
13425            None,
13426            true,
13427            vec![confirmation_skill()],
13428        );
13429
13430        agent.chat("Send it").await.unwrap();
13431        agent.chat("The report to Ada").await.unwrap();
13432        let rejected = agent.chat("No").await.unwrap();
13433
13434        assert_eq!(rejected.content, "Confirmation rejected.");
13435        assert_eq!(confirmation_skill_call_count(&observed), 0);
13436        assert!(agent.pending_skill_id.read().is_none());
13437    }
13438
13439    /// Confirms reset clears manager and skill ownership before a later streaming turn.
13440    #[tokio::test]
13441    async fn reset_invalidates_pending_skill_confirmation_before_streaming_input() {
13442        let (agent, observed) = state_disambiguation_agent_with_skills(
13443            vec![
13444                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13445                "send_report",
13446                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13447                r#"{"question":"What should I send?","options":null}"#,
13448                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13449                r#"{"question":"Should I send the report to Ada?"}"#,
13450                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"fresh input","what_is_unclear":[],"detected_language":"en"}"#,
13451                "none",
13452                "Fresh response.",
13453            ],
13454            true,
13455            None,
13456            true,
13457            vec![confirmation_skill()],
13458        );
13459
13460        agent.chat("Send it").await.unwrap();
13461        agent.chat("The report to Ada").await.unwrap();
13462        agent.reset().await.unwrap();
13463        assert!(agent.pending_skill_id.read().is_none());
13464        assert!(
13465            !agent
13466                .disambiguation_manager()
13467                .unwrap()
13468                .has_pending_clarification()
13469                .await
13470        );
13471
13472        let mut stream = agent.chat_stream("Yes").await.unwrap();
13473        let mut content = String::new();
13474        while let Some(chunk) = stream.next().await {
13475            match chunk {
13476                StreamChunk::Content { text } => content.push_str(&text),
13477                StreamChunk::Done {} => break,
13478                StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
13479                _ => {}
13480            }
13481        }
13482
13483        assert_eq!(content, "Fresh response.");
13484        assert_eq!(confirmation_skill_call_count(&observed), 0);
13485    }
13486
13487    /// Confirms the Agent trait reset uses the same pending ownership cleanup.
13488    #[tokio::test]
13489    async fn trait_reset_clears_pending_skill_confirmation() {
13490        let (agent, _) = state_disambiguation_agent_with_skills(
13491            vec![
13492                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13493                "send_report",
13494                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13495                r#"{"question":"What should I send?","options":null}"#,
13496                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13497                r#"{"question":"Should I send the report to Ada?"}"#,
13498            ],
13499            true,
13500            None,
13501            true,
13502            vec![confirmation_skill()],
13503        );
13504
13505        agent.chat("Send it").await.unwrap();
13506        agent.chat("The report to Ada").await.unwrap();
13507        <RuntimeAgent as Agent>::reset(&agent).await.unwrap();
13508
13509        assert!(agent.pending_skill_id.read().is_none());
13510        assert!(
13511            !agent
13512                .disambiguation_manager()
13513                .unwrap()
13514                .has_pending_clarification()
13515                .await
13516        );
13517    }
13518
13519    /// Confirms a state transition cancels stale skill execution before confirmation parsing.
13520    #[tokio::test]
13521    async fn state_change_invalidates_pending_skill_confirmation() {
13522        let (agent, observed) = state_disambiguation_agent_with_skills(
13523            vec![
13524                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13525                "send_report",
13526                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13527                r#"{"question":"What should I send?","options":null}"#,
13528                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13529                r#"{"question":"Should I send the report to Ada?"}"#,
13530                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"fresh input","what_is_unclear":[],"detected_language":"en"}"#,
13531                "none",
13532                "Fresh response.",
13533            ],
13534            true,
13535            None,
13536            true,
13537            vec![confirmation_skill()],
13538        );
13539
13540        agent.chat("Send it").await.unwrap();
13541        agent.chat("The report to Ada").await.unwrap();
13542        agent.transition_to("review").await.unwrap();
13543        let cancelled = agent.chat("Yes").await.unwrap();
13544
13545        assert_eq!(cancelled.content, "Fresh response.");
13546        assert_eq!(confirmation_skill_call_count(&observed), 0);
13547        assert!(agent.pending_skill_id.read().is_none());
13548    }
13549
13550    /// Confirms reset wins over a confirmation result that was already being parsed.
13551    #[tokio::test]
13552    async fn in_flight_confirmation_cannot_redispatch_after_reset() {
13553        let (mut agent, observed) = state_disambiguation_agent_with_skills(
13554            vec![
13555                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13556                "send_report",
13557                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13558                r#"{"question":"What should I send?","options":null}"#,
13559                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13560                r#"{"question":"Should I send the report to Ada?"}"#,
13561                r#"{"status":"confirmed"}"#,
13562                "Confirmation cancelled.",
13563            ],
13564            true,
13565            None,
13566            true,
13567            vec![confirmation_skill()],
13568        );
13569        let observer = Arc::new(BlockingRuntimeConfirmationObserver::new());
13570        let manager = agent
13571            .disambiguation_manager
13572            .take()
13573            .unwrap()
13574            .with_clarification_observer(observer.clone());
13575        agent.disambiguation_manager = Some(manager);
13576        let agent = Arc::new(agent);
13577
13578        agent.chat("Send it").await.unwrap();
13579        agent.chat("The report to Ada").await.unwrap();
13580
13581        let confirming_agent = Arc::clone(&agent);
13582        let confirmation = tokio::spawn(async move { confirming_agent.chat("Yes").await });
13583        observer.entered.wait().await;
13584        agent.reset().await.unwrap();
13585        observer.release.notify_one();
13586
13587        let response = confirmation.await.unwrap().unwrap();
13588        assert_eq!(response.content, "Confirmation cancelled.");
13589        assert_eq!(confirmation_skill_call_count(&observed), 0);
13590        assert!(agent.pending_skill_id.read().is_none());
13591    }
13592
13593    /// Confirms a queued reset prevents an older terminal question from being published afterward.
13594    #[tokio::test]
13595    async fn queued_reset_prevents_stale_confirmation_question_publication() {
13596        let (agent, observed) = state_disambiguation_agent(
13597            vec![
13598                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13599                r#"{"question":"What should I send?","options":null}"#,
13600                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13601                r#"{"question":"Should I send the report to Ada?"}"#,
13602            ],
13603            true,
13604            None,
13605            true,
13606        );
13607        let agent = Arc::new(agent);
13608        agent.chat("Send it").await.unwrap();
13609
13610        let admission = agent.disambiguation_admission.write().await;
13611        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
13612        let resetting_agent = Arc::clone(&agent);
13613        let reset = tokio::spawn(async move {
13614            let _ = started_tx.send(());
13615            resetting_agent.reset().await
13616        });
13617        started_rx.await.unwrap();
13618        tokio::task::yield_now().await;
13619
13620        let responding_agent = Arc::clone(&agent);
13621        let response =
13622            tokio::spawn(async move { responding_agent.chat("The report to Ada").await });
13623        tokio::time::timeout(std::time::Duration::from_secs(2), async {
13624            while observed.call_count() < 4 {
13625                tokio::task::yield_now().await;
13626            }
13627        })
13628        .await
13629        .expect("clarification processing must reach terminal publication");
13630        drop(admission);
13631
13632        reset.await.unwrap().unwrap();
13633        let error = response.await.unwrap().unwrap_err();
13634        assert!(error.to_string().contains("ownership changed"));
13635        assert!(
13636            !agent
13637                .disambiguation_manager()
13638                .unwrap()
13639                .has_pending_clarification()
13640                .await
13641        );
13642        assert!(agent.memory.get_messages(None).await.unwrap().is_empty());
13643    }
13644
13645    /// Confirms skill clarification consumers recheck ownership before publishing returned responses.
13646    #[tokio::test]
13647    async fn queued_reset_prevents_stale_skill_clarification_publication() {
13648        let (agent, observed) = state_disambiguation_agent_with_skills(
13649            vec![
13650                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13651                "send_report",
13652                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13653                r#"{"question":"What should I send?","options":null}"#,
13654            ],
13655            true,
13656            None,
13657            true,
13658            vec![confirmation_skill()],
13659        );
13660        let agent = Arc::new(agent);
13661        let admission = agent.disambiguation_admission.write().await;
13662        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
13663        let resetting_agent = Arc::clone(&agent);
13664        let reset = tokio::spawn(async move {
13665            let _ = started_tx.send(());
13666            resetting_agent.reset().await
13667        });
13668        started_rx.await.unwrap();
13669        tokio::task::yield_now().await;
13670
13671        let responding_agent = Arc::clone(&agent);
13672        let response = tokio::spawn(async move { responding_agent.chat("Send it").await });
13673        tokio::time::timeout(std::time::Duration::from_secs(2), async {
13674            while observed.call_count() < 4 {
13675                tokio::task::yield_now().await;
13676            }
13677        })
13678        .await
13679        .expect("skill clarification must reach terminal publication");
13680        drop(admission);
13681
13682        reset.await.unwrap().unwrap();
13683        let error = response.await.unwrap().unwrap_err();
13684        assert!(error.to_string().contains("ownership changed"));
13685        assert_eq!(confirmation_skill_call_count(&observed), 0);
13686        assert!(agent.pending_skill_id.read().is_none());
13687        assert!(agent.memory.get_messages(None).await.unwrap().is_empty());
13688    }
13689
13690    /// Confirms transition hooks can reenter reset after the state commit lock is released.
13691    #[tokio::test]
13692    async fn transition_hook_can_reset_without_admission_deadlock() {
13693        let hooks = Arc::new(ResetOnTransitionHooks {
13694            agent: parking_lot::Mutex::new(None),
13695            invoked: AtomicBool::new(false),
13696        });
13697        let agent = Arc::new(
13698            AgentBuilder::new()
13699                .system_prompt("Test transition hook reentrancy.")
13700                .llm(Arc::new(mock_with_response("done")))
13701                .state_machine(disambiguation_state_machine(None, false))
13702                .build()
13703                .unwrap()
13704                .with_hooks(hooks.clone()),
13705        );
13706        *hooks.agent.lock() = Some(Arc::downgrade(&agent));
13707
13708        let transitioned = tokio::time::timeout(
13709            std::time::Duration::from_secs(2),
13710            agent.apply_transition_target("active", "review", "test transition", None),
13711        )
13712        .await
13713        .expect("transition hook reset must not deadlock")
13714        .unwrap();
13715
13716        assert!(transitioned);
13717        assert!(hooks.invoked.load(Ordering::SeqCst));
13718        assert_eq!(agent.current_state().as_deref(), Some("active"));
13719    }
13720
13721    /// Confirms only the reserved transition can run source-state exit side effects.
13722    #[tokio::test]
13723    async fn concurrent_transition_cannot_duplicate_exit_actions() {
13724        let gate = PathMutationGate::new();
13725        let active = ai_agents_state::StateDefinition {
13726            on_exit: vec![StateAction::Tool {
13727                tool: "transition_exit".to_string(),
13728                args: Some(serde_json::json!({"path": "./transition-exit.txt"})),
13729            }],
13730            ..Default::default()
13731        };
13732        let state_machine = Arc::new(
13733            StateMachine::new(ai_agents_state::StateConfig {
13734                initial: "active".to_string(),
13735                states: HashMap::from([
13736                    ("active".to_string(), active),
13737                    (
13738                        "review".to_string(),
13739                        ai_agents_state::StateDefinition::default(),
13740                    ),
13741                ]),
13742                global_transitions: Vec::new(),
13743                fallback: None,
13744                max_no_transition: None,
13745                regenerate_on_transition: true,
13746            })
13747            .unwrap(),
13748        );
13749        let agent = Arc::new(
13750            AgentBuilder::new()
13751                .system_prompt("Test transition reservation.")
13752                .llm(Arc::new(mock_with_response("done")))
13753                .tool(Arc::new(BlockingPathMutationTool {
13754                    id: "transition_exit",
13755                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
13756                    gate: gate.clone(),
13757                }))
13758                .state_machine(state_machine)
13759                .build()
13760                .unwrap(),
13761        );
13762
13763        let first_agent = Arc::clone(&agent);
13764        let first = tokio::spawn(async move { first_agent.transition_to("review").await });
13765        tokio::time::timeout(std::time::Duration::from_secs(2), gate.wait_until_entered())
13766            .await
13767            .expect("reserved transition must enter its exit action");
13768
13769        let second = tokio::time::timeout(
13770            std::time::Duration::from_secs(2),
13771            agent.transition_to("review"),
13772        )
13773        .await
13774        .expect("competing transition must fail without waiting for the exit action")
13775        .unwrap_err();
13776        assert!(second.to_string().contains("already in progress"));
13777
13778        gate.release();
13779        first.await.unwrap().unwrap();
13780        assert_eq!(agent.current_state().as_deref(), Some("review"));
13781    }
13782
13783    /// Confirms a later transition cannot overtake the committed state's enter actions.
13784    #[tokio::test]
13785    async fn concurrent_transition_cannot_overtake_enter_actions() {
13786        let gate = PathMutationGate::new();
13787        let review = ai_agents_state::StateDefinition {
13788            on_enter: vec![StateAction::Tool {
13789                tool: "transition_enter".to_string(),
13790                args: Some(serde_json::json!({"path": "./transition-enter.txt"})),
13791            }],
13792            ..Default::default()
13793        };
13794        let state_machine = Arc::new(
13795            StateMachine::new(ai_agents_state::StateConfig {
13796                initial: "active".to_string(),
13797                states: HashMap::from([
13798                    (
13799                        "active".to_string(),
13800                        ai_agents_state::StateDefinition::default(),
13801                    ),
13802                    ("review".to_string(), review),
13803                ]),
13804                global_transitions: Vec::new(),
13805                fallback: None,
13806                max_no_transition: None,
13807                regenerate_on_transition: true,
13808            })
13809            .unwrap(),
13810        );
13811        let agent = Arc::new(
13812            AgentBuilder::new()
13813                .system_prompt("Test transition lifecycle reservation.")
13814                .llm(Arc::new(mock_with_response("done")))
13815                .tool(Arc::new(BlockingPathMutationTool {
13816                    id: "transition_enter",
13817                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
13818                    gate: gate.clone(),
13819                }))
13820                .state_machine(state_machine)
13821                .build()
13822                .unwrap(),
13823        );
13824
13825        let first_agent = Arc::clone(&agent);
13826        let first = tokio::spawn(async move { first_agent.transition_to("review").await });
13827        tokio::time::timeout(std::time::Duration::from_secs(2), gate.wait_until_entered())
13828            .await
13829            .expect("committed transition must enter its destination action");
13830
13831        let second = agent.transition_to("active").await.unwrap_err();
13832        assert!(second.to_string().contains("already in progress"));
13833        assert!(agent.reset().await.is_err());
13834
13835        gate.release();
13836        first.await.unwrap().unwrap();
13837        assert_eq!(agent.current_state().as_deref(), Some("review"));
13838    }
13839
13840    /// Confirms even a same-state restore invalidates manager and skill ownership.
13841    #[tokio::test]
13842    async fn same_state_restore_invalidates_pending_skill_confirmation() {
13843        let (agent, observed) = state_disambiguation_agent_with_skills(
13844            vec![
13845                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13846                "send_report",
13847                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13848                r#"{"question":"What should I send?","options":null}"#,
13849                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13850                r#"{"question":"Should I send the report to Ada?"}"#,
13851            ],
13852            true,
13853            None,
13854            true,
13855            vec![confirmation_skill()],
13856        );
13857
13858        agent.chat("Send it").await.unwrap();
13859        agent.chat("The report to Ada").await.unwrap();
13860        let snapshot = agent.save_state().await.unwrap();
13861        assert_eq!(agent.current_state().as_deref(), Some("active"));
13862
13863        agent.restore_state(snapshot).await.unwrap();
13864
13865        assert_eq!(agent.current_state().as_deref(), Some("active"));
13866        assert!(agent.pending_skill_id.read().is_none());
13867        assert!(
13868            !agent
13869                .disambiguation_manager()
13870                .unwrap()
13871                .has_pending_clarification()
13872                .await
13873        );
13874        assert_eq!(confirmation_skill_call_count(&observed), 0);
13875    }
13876
13877    /// Confirms direct state mutation cannot hide behind a return to the same state path.
13878    #[tokio::test]
13879    async fn direct_state_generation_change_invalidates_confirmation() {
13880        let (agent, observed) = state_disambiguation_agent_with_skills(
13881            vec![
13882                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13883                "send_report",
13884                r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13885                r#"{"question":"What should I send?","options":null}"#,
13886                r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13887                r#"{"question":"Should I send the report to Ada?"}"#,
13888                "Confirmation cancelled.",
13889            ],
13890            true,
13891            None,
13892            true,
13893            vec![confirmation_skill()],
13894        );
13895
13896        agent.chat("Send it").await.unwrap();
13897        agent.chat("The report to Ada").await.unwrap();
13898        let state_machine = agent.state_machine().unwrap();
13899        state_machine
13900            .transition_to("review", "external test")
13901            .unwrap();
13902        state_machine
13903            .transition_to("active", "external test")
13904            .unwrap();
13905
13906        let response = agent.chat("Yes").await.unwrap();
13907
13908        assert_eq!(response.content, "Confirmation cancelled.");
13909        assert_eq!(confirmation_skill_call_count(&observed), 0);
13910        assert!(agent.pending_skill_id.read().is_none());
13911    }
13912
13913    #[tokio::test]
13914    async fn state_confirmation_does_not_add_a_question_for_clear_input() {
13915        let (agent, observed) = state_disambiguation_agent(
13916            vec![
13917                r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"clear","what_is_unclear":[],"detected_language":"en"}"#,
13918                "Request executed.",
13919            ],
13920            true,
13921            None,
13922            true,
13923        );
13924
13925        let response = agent.chat("Send the report to Ada").await.unwrap();
13926
13927        assert_eq!(response.content, "Request executed.");
13928        assert_eq!(observed.call_count(), 2);
13929    }
13930
13931    #[tokio::test]
13932    async fn state_override_cannot_activate_a_disabled_top_level_manager() {
13933        let (agent, observed) =
13934            state_disambiguation_agent(vec!["Request executed."], false, Some(true), true);
13935
13936        assert!(!agent.has_disambiguation());
13937        let response = agent.chat("Send it").await.unwrap();
13938
13939        assert_eq!(response.content, "Request executed.");
13940        assert_eq!(observed.call_count(), 1);
13941    }
13942
13943    #[tokio::test]
13944    async fn native_required_choice_executes_through_the_shared_tool_path() {
13945        let mut mock = MockLLMProvider::new("native-required");
13946        mock.set_tool_choice(Some(ToolChoice::Required));
13947        mock.add_response(
13948            LLMResponse::new("", FinishReason::ToolCall)
13949                .with_tool_calls(vec![ToolCall {
13950                    id: "provider-call-1".to_string(),
13951                    name: "calculator".to_string(),
13952                    arguments: serde_json::json!({"expression": "2 + 2"}),
13953                }])
13954                .unwrap(),
13955        );
13956        mock.add_response(LLMResponse::new("The answer is 4.", FinishReason::Stop));
13957        let observed = mock.clone();
13958        let agent = AgentBuilder::new()
13959            .system_prompt("Use the calculator when needed.")
13960            .llm(Arc::new(mock))
13961            .tool(Arc::new(CalculatorTool::new()))
13962            .build()
13963            .unwrap();
13964
13965        let response = agent.chat("What is 2 + 2?").await.unwrap();
13966
13967        assert_eq!(response.content, "The answer is 4.");
13968        assert_eq!(
13969            response.tool_calls.as_ref().unwrap()[0].id,
13970            "provider-call-1"
13971        );
13972        let calls = observed.call_history();
13973        assert_eq!(calls.len(), 2);
13974        assert!(matches!(
13975            calls[0].request.as_ref().map(|request| &request.choice),
13976            Some(ToolChoice::Required)
13977        ));
13978        assert!(matches!(
13979            calls[1].request.as_ref().map(|request| &request.choice),
13980            Some(ToolChoice::Auto)
13981        ));
13982    }
13983
13984    #[tokio::test]
13985    async fn prompt_fallback_uses_one_corrective_retry() {
13986        let mut mock = MockLLMProvider::new("prompt-required");
13987        mock.set_tool_choice(Some(ToolChoice::Required));
13988        mock.set_native_tool_support(false);
13989        mock.set_responses(
13990            vec![
13991                "I can calculate that.".to_string(),
13992                r#"{"tool":"calculator","arguments":{"expression":"2 + 2"}}"#.to_string(),
13993                "The answer is 4.".to_string(),
13994            ],
13995            false,
13996        );
13997        let observed = mock.clone();
13998        let agent = AgentBuilder::new()
13999            .system_prompt("Use tools.")
14000            .llm(Arc::new(mock))
14001            .tool(Arc::new(CalculatorTool::new()))
14002            .build()
14003            .unwrap();
14004
14005        let response = agent.chat("What is 2 + 2?").await.unwrap();
14006
14007        assert_eq!(response.content, "The answer is 4.");
14008        assert_eq!(observed.call_count(), 3);
14009        let corrective = &observed.call_history()[1].messages;
14010        assert!(
14011            corrective
14012                .last()
14013                .unwrap()
14014                .content
14015                .contains("previous response")
14016        );
14017    }
14018
14019    #[tokio::test]
14020    async fn prompt_fallback_fails_after_one_noncompliant_retry() {
14021        let mut mock = MockLLMProvider::new("prompt-required-failure");
14022        mock.set_tool_choice(Some(ToolChoice::Required));
14023        mock.set_native_tool_support(false);
14024        mock.set_responses(
14025            vec!["No tool.".to_string(), "Still no tool.".to_string()],
14026            false,
14027        );
14028        let observed = mock.clone();
14029        let agent = AgentBuilder::new()
14030            .system_prompt("Use tools.")
14031            .llm(Arc::new(mock))
14032            .tool(Arc::new(CalculatorTool::new()))
14033            .build()
14034            .unwrap();
14035
14036        let error = agent.chat("What is 2 + 2?").await.unwrap_err();
14037
14038        assert!(error.to_string().contains("one corrective retry"));
14039        assert_eq!(observed.call_count(), 2);
14040    }
14041
14042    #[tokio::test]
14043    async fn specific_choice_cannot_widen_the_effective_grant() {
14044        let mut mock = MockLLMProvider::new("specific-outside-grant");
14045        mock.set_tool_choice(Some(ToolChoice::Specific("random".to_string())));
14046        let observed = mock.clone();
14047        let agent = AgentBuilder::new()
14048            .system_prompt("Use tools.")
14049            .llm(Arc::new(mock))
14050            .tool(Arc::new(CalculatorTool::new()))
14051            .build()
14052            .unwrap();
14053
14054        let error = agent.chat("Generate a value.").await.unwrap_err();
14055
14056        assert!(error.to_string().contains("is not registered"));
14057        assert_eq!(observed.call_count(), 0);
14058    }
14059
14060    #[tokio::test]
14061    async fn none_choice_exposes_no_tool_protocol() {
14062        let mut mock = MockLLMProvider::new("no-tools");
14063        mock.set_tool_choice(Some(ToolChoice::None));
14064        mock.set_response(r#"{"tool":"calculator","arguments":{"expression":"2 + 2"}}"#);
14065        let observed = mock.clone();
14066        let agent = AgentBuilder::new()
14067            .system_prompt("Answer directly.")
14068            .llm(Arc::new(mock))
14069            .tool(Arc::new(CalculatorTool::new()))
14070            .build()
14071            .unwrap();
14072
14073        let response = agent.chat("Hello").await.unwrap();
14074
14075        assert!(response.tool_calls.is_none());
14076        assert_eq!(observed.call_count(), 1);
14077        let call = observed.last_call().unwrap();
14078        assert!(call.request.is_none());
14079        assert!(
14080            call.messages
14081                .iter()
14082                .all(|message| !message.content.contains("Available tools:"))
14083        );
14084    }
14085
14086    struct RuntimeStorage {
14087        capabilities: Box<[StorageCapability]>,
14088        snapshots: RwLock<HashMap<String, AgentSnapshot>>,
14089        metadata: RwLock<HashMap<String, ai_agents_core::SessionMetadata>>,
14090        metadata_save_calls: AtomicU64,
14091        metadata_load_calls: AtomicU64,
14092        fail_metadata_save: AtomicBool,
14093        fail_metadata_load: AtomicBool,
14094    }
14095
14096    impl RuntimeStorage {
14097        fn new(capabilities: impl IntoIterator<Item = StorageCapability>) -> Self {
14098            Self {
14099                capabilities: capabilities.into_iter().collect(),
14100                snapshots: RwLock::new(HashMap::new()),
14101                metadata: RwLock::new(HashMap::new()),
14102                metadata_save_calls: AtomicU64::new(0),
14103                metadata_load_calls: AtomicU64::new(0),
14104                fail_metadata_save: AtomicBool::new(false),
14105                fail_metadata_load: AtomicBool::new(false),
14106            }
14107        }
14108    }
14109
14110    #[async_trait]
14111    impl AgentStorage for RuntimeStorage {
14112        fn supports(&self, capability: StorageCapability) -> bool {
14113            self.capabilities.contains(&capability)
14114        }
14115
14116        async fn save(&self, session_id: &str, snapshot: &AgentSnapshot) -> Result<()> {
14117            self.snapshots
14118                .write()
14119                .insert(session_id.to_string(), snapshot.clone());
14120            Ok(())
14121        }
14122
14123        async fn load(&self, session_id: &str) -> Result<Option<AgentSnapshot>> {
14124            Ok(self.snapshots.read().get(session_id).cloned())
14125        }
14126
14127        async fn delete(&self, session_id: &str) -> Result<()> {
14128            self.snapshots.write().remove(session_id);
14129            Ok(())
14130        }
14131
14132        async fn list_sessions(&self) -> Result<Vec<String>> {
14133            Ok(self.snapshots.read().keys().cloned().collect())
14134        }
14135
14136        async fn save_snapshot_with_metadata(
14137            &self,
14138            session_id: &str,
14139            snapshot: &AgentSnapshot,
14140            metadata: &ai_agents_core::SessionMetadata,
14141        ) -> Result<()> {
14142            self.metadata_save_calls.fetch_add(1, Ordering::SeqCst);
14143            if self.fail_metadata_save.load(Ordering::SeqCst) {
14144                return Err(AgentError::Persistence("metadata save failed".into()));
14145            }
14146            self.snapshots
14147                .write()
14148                .insert(session_id.to_string(), snapshot.clone());
14149            self.metadata
14150                .write()
14151                .insert(session_id.to_string(), metadata.clone());
14152            Ok(())
14153        }
14154
14155        async fn save_metadata(
14156            &self,
14157            session_id: &str,
14158            metadata: &ai_agents_core::SessionMetadata,
14159        ) -> Result<()> {
14160            self.metadata_save_calls.fetch_add(1, Ordering::SeqCst);
14161            if self.fail_metadata_save.load(Ordering::SeqCst) {
14162                return Err(AgentError::Persistence("metadata save failed".into()));
14163            }
14164            self.metadata
14165                .write()
14166                .insert(session_id.to_string(), metadata.clone());
14167            Ok(())
14168        }
14169
14170        async fn load_metadata(
14171            &self,
14172            session_id: &str,
14173        ) -> Result<Option<ai_agents_core::SessionMetadata>> {
14174            self.metadata_load_calls.fetch_add(1, Ordering::SeqCst);
14175            if self.fail_metadata_load.load(Ordering::SeqCst) {
14176                return Err(AgentError::Persistence("metadata load failed".into()));
14177            }
14178            Ok(self.metadata.read().get(session_id).cloned())
14179        }
14180    }
14181
14182    fn runtime_storage_agent() -> RuntimeAgent {
14183        AgentBuilder::new()
14184            .system_prompt("Test runtime storage integration.")
14185            .llm(Arc::new(mock_with_response("done")))
14186            .build()
14187            .unwrap()
14188    }
14189
14190    fn restore_spec(id: &str) -> crate::spec::AgentSpec {
14191        crate::spec::AgentSpec {
14192            name: id.to_string(),
14193            system_prompt: format!("Restore child {id}."),
14194            ..crate::spec::AgentSpec::default()
14195        }
14196    }
14197
14198    fn restore_entry(id: &str) -> ai_agents_core::SpawnedAgentEntry {
14199        ai_agents_core::SpawnedAgentEntry {
14200            id: id.to_string(),
14201            name: id.to_string(),
14202            spec_yaml: serde_yaml::to_string(&restore_spec(id)).unwrap(),
14203        }
14204    }
14205
14206    fn restore_spawner(
14207        storage: Arc<RuntimeStorage>,
14208        max_agents: usize,
14209    ) -> (
14210        Arc<crate::spawner::AgentSpawner>,
14211        Arc<crate::spawner::AgentRegistry>,
14212    ) {
14213        let mut llms = LLMRegistry::new();
14214        llms.register("default", Arc::new(mock_with_response("done")));
14215        (
14216            Arc::new(
14217                crate::spawner::AgentSpawner::new()
14218                    .with_shared_llms(llms)
14219                    .with_shared_storage(storage)
14220                    .with_max_agents(max_agents),
14221            ),
14222            Arc::new(crate::spawner::AgentRegistry::new()),
14223        )
14224    }
14225
14226    async fn save_restore_target(
14227        parent: &RuntimeAgent,
14228        storage: &RuntimeStorage,
14229        session_id: &str,
14230        entries: Vec<ai_agents_core::SpawnedAgentEntry>,
14231    ) {
14232        let mut snapshot = parent.save_state().await.unwrap();
14233        snapshot.spawned_agents = Some(entries);
14234        storage.save(session_id, &snapshot).await.unwrap();
14235        storage
14236            .save_metadata(session_id, &ai_agents_core::SessionMetadata::default())
14237            .await
14238            .unwrap();
14239    }
14240
14241    #[tokio::test]
14242    async fn storage_init_requires_storage_for_actor_facts() {
14243        let facts = ai_agents_facts::FactsConfig {
14244            enabled: true,
14245            ..Default::default()
14246        };
14247        let agent = runtime_storage_agent().with_facts_config(None, Some(facts));
14248
14249        let error = agent.init_storage().await.unwrap_err();
14250        assert!(matches!(
14251            error,
14252            AgentError::Config(message)
14253                if message.contains("actor facts or actor memory")
14254                    && message.contains("none is configured or injected")
14255        ));
14256    }
14257
14258    #[tokio::test]
14259    async fn storage_init_validates_actor_facts_capability() {
14260        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14261        let actor_memory = ai_agents_facts::ActorMemoryConfig {
14262            enabled: true,
14263            ..Default::default()
14264        };
14265        let agent = runtime_storage_agent()
14266            .with_storage(storage)
14267            .with_facts_config(Some(actor_memory), None);
14268
14269        assert!(matches!(
14270            agent.init_storage().await,
14271            Err(AgentError::UnsupportedStorageCapability(
14272                StorageCapability::ActorFacts
14273            ))
14274        ));
14275    }
14276
14277    #[tokio::test]
14278    async fn blocking_chat_rejects_unsupported_required_storage() {
14279        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14280        let facts = ai_agents_facts::FactsConfig {
14281            enabled: true,
14282            ..Default::default()
14283        };
14284        let agent = runtime_storage_agent()
14285            .with_storage(storage)
14286            .with_facts_config(None, Some(facts));
14287
14288        assert!(matches!(
14289            agent.chat("hello").await,
14290            Err(AgentError::UnsupportedStorageCapability(
14291                StorageCapability::ActorFacts
14292            ))
14293        ));
14294    }
14295
14296    #[tokio::test]
14297    async fn streaming_chat_rejects_unsupported_required_storage_before_stream_creation() {
14298        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14299        let config = ai_agents_relationships::RelationshipConfig {
14300            enabled: true,
14301            ..Default::default()
14302        };
14303        let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
14304        let agent = runtime_storage_agent()
14305            .with_storage(storage)
14306            .with_relationships(manager);
14307
14308        assert!(matches!(
14309            agent.chat_stream("hello").await,
14310            Err(AgentError::UnsupportedStorageCapability(
14311                StorageCapability::ActorRelationships
14312            ))
14313        ));
14314    }
14315
14316    #[tokio::test]
14317    async fn storage_init_completes_facts_for_injected_storage() {
14318        let storage = Arc::new(RuntimeStorage::new([
14319            StorageCapability::Snapshot,
14320            StorageCapability::ActorFacts,
14321        ]));
14322        let facts = ai_agents_facts::FactsConfig {
14323            enabled: true,
14324            ..Default::default()
14325        };
14326        let agent = runtime_storage_agent()
14327            .with_storage(storage)
14328            .with_facts_config(None, Some(facts));
14329
14330        agent.init_storage().await.unwrap();
14331        assert!(agent.fact_store().is_some());
14332    }
14333
14334    #[tokio::test]
14335    async fn storage_init_requires_storage_for_persistent_relationships() {
14336        let config = ai_agents_relationships::RelationshipConfig {
14337            enabled: true,
14338            ..Default::default()
14339        };
14340        let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
14341        let agent = runtime_storage_agent().with_relationships(manager);
14342
14343        let error = agent.init_storage().await.unwrap_err();
14344        assert!(matches!(
14345            error,
14346            AgentError::Config(message)
14347                if message.contains("persistent relationships")
14348                    && message.contains("none is configured or injected")
14349        ));
14350    }
14351
14352    #[tokio::test]
14353    async fn storage_init_validates_persistent_relationships_capability() {
14354        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14355        let config = ai_agents_relationships::RelationshipConfig {
14356            enabled: true,
14357            ..Default::default()
14358        };
14359        let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
14360        let agent = runtime_storage_agent()
14361            .with_storage(storage)
14362            .with_relationships(manager);
14363
14364        assert!(matches!(
14365            agent.init_storage().await,
14366            Err(AgentError::UnsupportedStorageCapability(
14367                StorageCapability::ActorRelationships
14368            ))
14369        ));
14370    }
14371
14372    #[tokio::test]
14373    async fn session_restore_updates_identity_and_clears_stale_actor_binding() {
14374        let storage = Arc::new(RuntimeStorage::new([
14375            StorageCapability::Snapshot,
14376            StorageCapability::SessionMetadata,
14377        ]));
14378        let agent = runtime_storage_agent().with_storage(storage.clone());
14379        agent.set_actor_id("old-actor").unwrap();
14380        agent.save_session("old").await.unwrap();
14381        storage
14382            .save("target", &agent.save_state().await.unwrap())
14383            .await
14384            .unwrap();
14385        storage
14386            .save_metadata("target", &ai_agents_core::SessionMetadata::default())
14387            .await
14388            .unwrap();
14389
14390        assert!(agent.load_session("target").await.unwrap());
14391
14392        assert_eq!(agent.current_session_id.read().as_deref(), Some("target"));
14393        assert_eq!(agent.actor_id(), None);
14394    }
14395
14396    #[tokio::test]
14397    async fn complete_restore_reconciles_growth_shrink_and_empty_topologies() {
14398        let storage = Arc::new(RuntimeStorage::new([
14399            StorageCapability::Snapshot,
14400            StorageCapability::SessionMetadata,
14401        ]));
14402        let (spawner, registry) = restore_spawner(storage.clone(), 3);
14403        let parent = runtime_storage_agent()
14404            .with_storage(storage.clone())
14405            .with_spawner_handles(Arc::clone(&spawner), Arc::clone(&registry));
14406
14407        for id in ["a", "b"] {
14408            let spawned = spawner
14409                .spawn_with_id(id.to_string(), restore_spec(id))
14410                .await
14411                .unwrap();
14412            spawned.agent.save_session("grow").await.unwrap();
14413            registry.register(spawned).await.unwrap();
14414        }
14415        let staged_c = crate::spawner::storage::NamespacedStorage::new(storage.clone(), "c");
14416        staged_c
14417            .save("grow", &AgentSnapshot::new("c".into()))
14418            .await
14419            .unwrap();
14420        staged_c
14421            .save_metadata("grow", &ai_agents_core::SessionMetadata::default())
14422            .await
14423            .unwrap();
14424        save_restore_target(
14425            &parent,
14426            storage.as_ref(),
14427            "grow",
14428            vec![restore_entry("a"), restore_entry("b"), restore_entry("c")],
14429        )
14430        .await;
14431
14432        assert_eq!(parent.restore_session_full("grow").await.unwrap(), 3);
14433        assert_eq!(registry.count(), 3);
14434        assert!(registry.contains("c"));
14435        assert_eq!(spawner.spawned_count(), 3);
14436
14437        for id in ["a", "b"] {
14438            registry
14439                .get(id)
14440                .unwrap()
14441                .save_session("shrink")
14442                .await
14443                .unwrap();
14444        }
14445        save_restore_target(
14446            &parent,
14447            storage.as_ref(),
14448            "shrink",
14449            vec![restore_entry("a"), restore_entry("b")],
14450        )
14451        .await;
14452
14453        assert_eq!(parent.restore_session_full("shrink").await.unwrap(), 2);
14454        assert_eq!(registry.count(), 2);
14455        assert!(!registry.contains("c"));
14456        assert_eq!(spawner.spawned_count(), 2);
14457
14458        save_restore_target(&parent, storage.as_ref(), "empty", Vec::new()).await;
14459
14460        assert_eq!(parent.restore_session_full("empty").await.unwrap(), 0);
14461        assert_eq!(registry.count(), 0);
14462        assert_eq!(spawner.spawned_count(), 0);
14463        assert_eq!(parent.current_session_id.read().as_deref(), Some("empty"));
14464    }
14465
14466    #[tokio::test]
14467    async fn storage_session_metadata_is_called_only_when_advertised() {
14468        let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14469        storage.fail_metadata_save.store(true, Ordering::SeqCst);
14470        storage.fail_metadata_load.store(true, Ordering::SeqCst);
14471        let agent = runtime_storage_agent().with_storage(storage.clone());
14472
14473        agent.save_session("session").await.unwrap();
14474        assert!(agent.load_session("session").await.unwrap());
14475        assert_eq!(storage.metadata_save_calls.load(Ordering::SeqCst), 0);
14476        assert_eq!(storage.metadata_load_calls.load(Ordering::SeqCst), 0);
14477    }
14478
14479    #[cfg(feature = "sqlite")]
14480    #[tokio::test]
14481    async fn sqlite_runtime_save_filter_reopen_and_reload_stay_consistent() {
14482        let directory =
14483            std::env::temp_dir().join(format!("ai-agents-runtime-sqlite-{}", uuid::Uuid::new_v4()));
14484        let path = directory.join("sessions.sqlite");
14485        let path_string = path.to_string_lossy().into_owned();
14486        let storage = Arc::new(
14487            ai_agents_storage::SqliteStorage::new(&path_string)
14488                .await
14489                .unwrap(),
14490        );
14491        let agent = runtime_storage_agent().with_storage(storage.clone());
14492        agent.set_session_metadata(ai_agents_core::SessionMetadata {
14493            tags: vec!["initial".into()],
14494            ..Default::default()
14495        });
14496        agent.chat("persist this turn").await.unwrap();
14497        agent.save_session("session").await.unwrap();
14498
14499        agent.set_session_metadata(ai_agents_core::SessionMetadata {
14500            tags: vec!["updated".into()],
14501            ..Default::default()
14502        });
14503        agent.save_session("session").await.unwrap();
14504        assert!(
14505            agent
14506                .list_sessions_filtered(&ai_agents_core::SessionFilter {
14507                    tags: Some(vec!["initial".into()]),
14508                    ..Default::default()
14509                })
14510                .await
14511                .unwrap()
14512                .is_empty()
14513        );
14514        assert_eq!(
14515            agent
14516                .list_sessions_filtered(&ai_agents_core::SessionFilter {
14517                    tags: Some(vec!["updated".into()]),
14518                    ..Default::default()
14519                })
14520                .await
14521                .unwrap()
14522                .len(),
14523            1
14524        );
14525        drop(agent);
14526        storage.close().await;
14527        drop(storage);
14528
14529        let reopened_storage = Arc::new(
14530            ai_agents_storage::SqliteStorage::new(&path_string)
14531                .await
14532                .unwrap(),
14533        );
14534        let restored = runtime_storage_agent().with_storage(reopened_storage.clone());
14535        assert!(restored.load_session("session").await.unwrap());
14536        assert_eq!(restored.session_metadata().tags, vec!["updated"]);
14537        assert_eq!(
14538            restored.current_session_id.read().as_deref(),
14539            Some("session")
14540        );
14541        assert!(restored.save_state().await.unwrap().memory.messages.len() >= 2);
14542        assert_eq!(
14543            restored
14544                .list_sessions_filtered(&ai_agents_core::SessionFilter {
14545                    tags: Some(vec!["updated".into()]),
14546                    ..Default::default()
14547                })
14548                .await
14549                .unwrap()
14550                .len(),
14551            1
14552        );
14553
14554        drop(restored);
14555        reopened_storage.close().await;
14556        drop(reopened_storage);
14557        crate::remove_sqlite_test_directory(&directory)
14558            .await
14559            .unwrap();
14560    }
14561
14562    #[tokio::test]
14563    async fn storage_session_metadata_backend_failures_propagate() {
14564        let storage = Arc::new(RuntimeStorage::new([
14565            StorageCapability::Snapshot,
14566            StorageCapability::SessionMetadata,
14567        ]));
14568        let agent = runtime_storage_agent().with_storage(storage.clone());
14569
14570        agent.save_session("session").await.unwrap();
14571        storage
14572            .save("target", &agent.save_state().await.unwrap())
14573            .await
14574            .unwrap();
14575        storage.fail_metadata_load.store(true, Ordering::SeqCst);
14576        assert!(matches!(
14577            agent.load_session("target").await,
14578            Err(AgentError::Persistence(message)) if message == "metadata load failed"
14579        ));
14580        assert_eq!(agent.current_session_id.read().as_deref(), Some("session"));
14581
14582        storage.fail_metadata_save.store(true, Ordering::SeqCst);
14583        assert!(matches!(
14584            agent.save_session("session").await,
14585            Err(AgentError::Persistence(message)) if message == "metadata save failed"
14586        ));
14587    }
14588
14589    struct ProviderFutureDropSignal {
14590        dropped: Arc<AtomicBool>,
14591    }
14592
14593    impl Drop for ProviderFutureDropSignal {
14594        fn drop(&mut self) {
14595            self.dropped.store(true, Ordering::SeqCst);
14596        }
14597    }
14598
14599    struct BufferedLockingProvider {
14600        lock: Arc<tokio::sync::Mutex<()>>,
14601        stream_started: Arc<tokio::sync::Notify>,
14602        stream_dropped: Arc<AtomicBool>,
14603        committed_after_drop: Arc<AtomicBool>,
14604    }
14605
14606    #[async_trait]
14607    impl LLMProvider for BufferedLockingProvider {
14608        async fn complete(
14609            &self,
14610            _messages: &[ChatMessage],
14611            _config: Option<&LLMConfig>,
14612        ) -> std::result::Result<LLMResponse, LLMError> {
14613            let _guard = self.lock.lock().await;
14614            self.committed_after_drop
14615                .store(self.stream_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14616            Ok(LLMResponse::new(
14617                "Committed technical response.",
14618                FinishReason::Stop,
14619            ))
14620        }
14621
14622        async fn complete_stream(
14623            &self,
14624            _messages: &[ChatMessage],
14625            _config: Option<&LLMConfig>,
14626        ) -> std::result::Result<
14627            Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14628            LLMError,
14629        > {
14630            let _guard = self.lock.lock().await;
14631            let _drop_signal = ProviderFutureDropSignal {
14632                dropped: Arc::clone(&self.stream_dropped),
14633            };
14634            self.stream_started.notify_one();
14635            std::future::pending().await
14636        }
14637
14638        fn provider_name(&self) -> &str {
14639            "buffered-locking"
14640        }
14641
14642        fn supports(&self, _feature: LLMFeature) -> bool {
14643            false
14644        }
14645    }
14646
14647    struct PendingDropStream {
14648        dropped: Arc<AtomicBool>,
14649        dropped_notify: Arc<tokio::sync::Notify>,
14650    }
14651
14652    impl Stream for PendingDropStream {
14653        type Item = std::result::Result<LLMChunk, LLMError>;
14654
14655        fn poll_next(
14656            self: Pin<&mut Self>,
14657            _cx: &mut std::task::Context<'_>,
14658        ) -> std::task::Poll<Option<Self::Item>> {
14659            std::task::Poll::Pending
14660        }
14661    }
14662
14663    impl Drop for PendingDropStream {
14664        fn drop(&mut self) {
14665            self.dropped.store(true, Ordering::SeqCst);
14666            self.dropped_notify.notify_one();
14667        }
14668    }
14669
14670    struct EstablishedStreamProvider {
14671        stream_started: Arc<tokio::sync::Notify>,
14672        stream_dropped: Arc<AtomicBool>,
14673        stream_dropped_notify: Arc<tokio::sync::Notify>,
14674        committed_after_drop: Arc<AtomicBool>,
14675    }
14676
14677    #[async_trait]
14678    impl LLMProvider for EstablishedStreamProvider {
14679        async fn complete(
14680            &self,
14681            _messages: &[ChatMessage],
14682            _config: Option<&LLMConfig>,
14683        ) -> std::result::Result<LLMResponse, LLMError> {
14684            if !self.stream_dropped.load(Ordering::SeqCst) {
14685                self.stream_dropped_notify.notified().await;
14686            }
14687            self.committed_after_drop
14688                .store(self.stream_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14689            Ok(LLMResponse::new(
14690                "Committed technical response.",
14691                FinishReason::Stop,
14692            ))
14693        }
14694
14695        async fn complete_stream(
14696            &self,
14697            _messages: &[ChatMessage],
14698            _config: Option<&LLMConfig>,
14699        ) -> std::result::Result<
14700            Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14701            LLMError,
14702        > {
14703            self.stream_started.notify_one();
14704            Ok(Box::new(PendingDropStream {
14705                dropped: Arc::clone(&self.stream_dropped),
14706                dropped_notify: Arc::clone(&self.stream_dropped_notify),
14707            }))
14708        }
14709
14710        fn provider_name(&self) -> &str {
14711            "established-stream"
14712        }
14713
14714        fn supports(&self, _feature: LLMFeature) -> bool {
14715            false
14716        }
14717    }
14718
14719    struct FirstCallLockingProvider {
14720        lock: Arc<tokio::sync::Mutex<()>>,
14721        first_started: Arc<tokio::sync::Notify>,
14722        first_dropped: Arc<AtomicBool>,
14723        committed_after_drop: Arc<AtomicBool>,
14724        calls: AtomicU64,
14725    }
14726
14727    #[async_trait]
14728    impl LLMProvider for FirstCallLockingProvider {
14729        async fn complete(
14730            &self,
14731            _messages: &[ChatMessage],
14732            _config: Option<&LLMConfig>,
14733        ) -> std::result::Result<LLMResponse, LLMError> {
14734            let _guard = self.lock.lock().await;
14735            let call = self.calls.fetch_add(1, Ordering::SeqCst);
14736            if call == 0 {
14737                let _drop_signal = ProviderFutureDropSignal {
14738                    dropped: Arc::clone(&self.first_dropped),
14739                };
14740                self.first_started.notify_one();
14741                return std::future::pending().await;
14742            }
14743            self.committed_after_drop
14744                .store(self.first_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14745            Ok(LLMResponse::new(
14746                "Committed technical response.",
14747                FinishReason::Stop,
14748            ))
14749        }
14750
14751        async fn complete_stream(
14752            &self,
14753            _messages: &[ChatMessage],
14754            _config: Option<&LLMConfig>,
14755        ) -> std::result::Result<
14756            Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14757            LLMError,
14758        > {
14759            Err(LLMError::Other(
14760                "streaming is not used in this test".to_string(),
14761            ))
14762        }
14763
14764        fn provider_name(&self) -> &str {
14765            "first-call-locking"
14766        }
14767
14768        fn supports(&self, _feature: LLMFeature) -> bool {
14769            false
14770        }
14771    }
14772
14773    struct RoutingAfterProviderStart {
14774        provider_started: Arc<tokio::sync::Notify>,
14775    }
14776
14777    #[async_trait]
14778    impl LLMProvider for RoutingAfterProviderStart {
14779        async fn complete(
14780            &self,
14781            _messages: &[ChatMessage],
14782            _config: Option<&LLMConfig>,
14783        ) -> std::result::Result<LLMResponse, LLMError> {
14784            self.provider_started.notified().await;
14785            Ok(LLMResponse::new("1", FinishReason::Stop))
14786        }
14787
14788        async fn complete_stream(
14789            &self,
14790            _messages: &[ChatMessage],
14791            _config: Option<&LLMConfig>,
14792        ) -> std::result::Result<
14793            Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14794            LLMError,
14795        > {
14796            Err(LLMError::Other(
14797                "streaming is not used in this test".to_string(),
14798            ))
14799        }
14800
14801        fn provider_name(&self) -> &str {
14802            "routing-after-start"
14803        }
14804
14805        fn supports(&self, _feature: LLMFeature) -> bool {
14806            false
14807        }
14808    }
14809
14810    /// Test hook that counts completed responses.
14811    struct ResponseCountingHooks {
14812        responses: Arc<std::sync::atomic::AtomicUsize>,
14813    }
14814
14815    /// Test provider that reports blocking entry and emits a deterministic event stream.
14816    struct RootTurnProbeProvider {
14817        complete_entered: tokio::sync::mpsc::UnboundedSender<()>,
14818    }
14819
14820    /// Calls a configured runtime from on_response and records whether nested root admission succeeded.
14821    struct ResponseChatHooks {
14822        target: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
14823        invoked: AtomicBool,
14824        nested_result: parking_lot::Mutex<Option<std::result::Result<String, String>>>,
14825    }
14826
14827    /// Dispatches one concurrent child from on_response and records the spawned orchestration result.
14828    struct ConcurrentResponseHooks {
14829        registry: Weak<crate::spawner::AgentRegistry>,
14830        child_id: String,
14831        invoked: AtomicBool,
14832        nested_result: parking_lot::Mutex<Option<std::result::Result<String, String>>>,
14833    }
14834
14835    /// Test tool that fails once and records the deadline observed by each invocation attempt.
14836    struct RetryDeadlineTool {
14837        calls: Arc<std::sync::atomic::AtomicUsize>,
14838        deadlines: Arc<parking_lot::Mutex<Vec<chrono::DateTime<chrono::Utc>>>>,
14839    }
14840
14841    /// Test hook that records shared-executor lifecycle order and authoritative records.
14842    struct ToolLifecycleRecordingHooks {
14843        events: parking_lot::Mutex<Vec<String>>,
14844        records: parking_lot::Mutex<Vec<ToolExecutionRecord>>,
14845    }
14846
14847    impl ToolLifecycleRecordingHooks {
14848        /// Creates an empty lifecycle recorder.
14849        fn new() -> Self {
14850            Self {
14851                events: parking_lot::Mutex::new(Vec::new()),
14852                records: parking_lot::Mutex::new(Vec::new()),
14853            }
14854        }
14855
14856        /// Returns a stable snapshot of recorded hook order.
14857        fn events(&self) -> Vec<String> {
14858            self.events.lock().clone()
14859        }
14860
14861        /// Returns a stable snapshot of authoritative execution records.
14862        fn records(&self) -> Vec<ToolExecutionRecord> {
14863            self.records.lock().clone()
14864        }
14865    }
14866
14867    /// Test tool that returns the execution context it received.
14868    struct ContextEchoTool;
14869
14870    #[async_trait]
14871    impl LLMProvider for RootTurnProbeProvider {
14872        async fn complete(
14873            &self,
14874            _messages: &[ChatMessage],
14875            _config: Option<&LLMConfig>,
14876        ) -> std::result::Result<LLMResponse, LLMError> {
14877            let _ = self.complete_entered.send(());
14878            Ok(LLMResponse::new("blocking complete", FinishReason::Stop))
14879        }
14880
14881        async fn complete_stream(
14882            &self,
14883            _messages: &[ChatMessage],
14884            _config: Option<&LLMConfig>,
14885        ) -> std::result::Result<
14886            Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14887            LLMError,
14888        > {
14889            Ok(Box::new(futures::stream::iter(vec![Ok(
14890                LLMChunk::final_chunk("stream complete", FinishReason::Stop, None),
14891            )])))
14892        }
14893
14894        fn provider_name(&self) -> &str {
14895            "root-turn-probe"
14896        }
14897
14898        fn supports(&self, feature: LLMFeature) -> bool {
14899            matches!(feature, LLMFeature::Streaming)
14900        }
14901    }
14902
14903    #[async_trait]
14904    impl ai_agents_core::Tool for ContextEchoTool {
14905        fn id(&self) -> &str {
14906            "context_echo"
14907        }
14908
14909        fn name(&self) -> &str {
14910            "Context Echo"
14911        }
14912
14913        fn description(&self) -> &str {
14914            "Returns selected execution context fields."
14915        }
14916
14917        fn input_schema(&self) -> Value {
14918            serde_json::json!({"type": "object"})
14919        }
14920
14921        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14922            ai_agents_core::ToolPolicyBindings {
14923                path_fields: vec![ai_agents_core::PathPolicyBinding::read("path")],
14924                result_limit_fields: vec![ai_agents_core::ResultLimitBinding::new(
14925                    "max_results",
14926                    ai_agents_core::ResultLimitKind::MaxResults,
14927                )],
14928                ..Default::default()
14929            }
14930        }
14931
14932        async fn execute(
14933            &self,
14934            _args: Value,
14935            ctx: ai_agents_core::ToolExecutionContext,
14936        ) -> ToolResult {
14937            ToolResult::ok(
14938                serde_json::json!({
14939                    "requested_name": ctx.requested_name,
14940                    "canonical_id": ctx.canonical_id,
14941                    "display_name": ctx.display_name,
14942                    "max_results": ctx.limits.max_results,
14943                    "custom_config": ctx.custom_config,
14944                })
14945                .to_string(),
14946            )
14947        }
14948    }
14949
14950    #[async_trait]
14951    impl ai_agents_core::Tool for RetryDeadlineTool {
14952        fn id(&self) -> &str {
14953            "retry_deadline"
14954        }
14955
14956        fn name(&self) -> &str {
14957            "Retry Deadline"
14958        }
14959
14960        fn description(&self) -> &str {
14961            "Records one deadline per retry invocation."
14962        }
14963
14964        fn input_schema(&self) -> Value {
14965            serde_json::json!({"type": "object"})
14966        }
14967
14968        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14969            ai_agents_core::ToolSafetyMetadata::compute()
14970        }
14971
14972        fn classify_call(&self, _args: &Value) -> ai_agents_core::ToolCallClassification {
14973            let mut classification =
14974                ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
14975            classification.timeout_ms = Some(1_000);
14976            classification.safely_retryable = true;
14977            classification
14978        }
14979
14980        async fn execute(
14981            &self,
14982            _args: Value,
14983            ctx: ai_agents_core::ToolExecutionContext,
14984        ) -> ToolResult {
14985            self.deadlines.lock().push(
14986                ctx.deadline
14987                    .expect("each invocation must receive a deadline"),
14988            );
14989            let call = self.calls.fetch_add(1, Ordering::SeqCst);
14990            if call == 0 {
14991                tokio::time::sleep(std::time::Duration::from_millis(20)).await;
14992                ToolResult::error("retry")
14993            } else {
14994                ToolResult::ok("done")
14995            }
14996        }
14997    }
14998
14999    /// Test tool that stays active long enough for runtime cancellation.
15000    struct SlowTool;
15001
15002    /// Test tool that fails once and must not be retried for writes.
15003    struct FlakyWriteTool {
15004        calls: Arc<std::sync::atomic::AtomicUsize>,
15005    }
15006
15007    /// Test tool that tracks concurrent execution on one path.
15008    struct LockedWriteTool {
15009        active: Arc<std::sync::atomic::AtomicUsize>,
15010        max_active: Arc<std::sync::atomic::AtomicUsize>,
15011    }
15012
15013    struct MultiResourceWriteTool {
15014        active: Arc<std::sync::atomic::AtomicUsize>,
15015        max_active: Arc<std::sync::atomic::AtomicUsize>,
15016    }
15017
15018    #[derive(Clone)]
15019    struct PathMutationGate {
15020        entered: Arc<AtomicBool>,
15021        entered_notify: Arc<tokio::sync::Notify>,
15022        release: Arc<tokio::sync::Notify>,
15023    }
15024
15025    impl PathMutationGate {
15026        fn new() -> Self {
15027            Self {
15028                entered: Arc::new(AtomicBool::new(false)),
15029                entered_notify: Arc::new(tokio::sync::Notify::new()),
15030                release: Arc::new(tokio::sync::Notify::new()),
15031            }
15032        }
15033
15034        async fn wait_until_entered(&self) {
15035            if !self.entered.load(Ordering::SeqCst) {
15036                self.entered_notify.notified().await;
15037            }
15038        }
15039
15040        fn release(&self) {
15041            self.release.notify_one();
15042        }
15043    }
15044
15045    struct BlockingPathMutationTool {
15046        id: &'static str,
15047        path_fields: Vec<ai_agents_core::PathPolicyBinding>,
15048        gate: PathMutationGate,
15049    }
15050
15051    struct NoBindingWriteTool {
15052        active: Arc<std::sync::atomic::AtomicUsize>,
15053        max_active: Arc<std::sync::atomic::AtomicUsize>,
15054    }
15055
15056    struct RecoveryTestTool {
15057        id: String,
15058        succeeds: bool,
15059        calls: Arc<std::sync::atomic::AtomicUsize>,
15060        max_output_chars: Option<usize>,
15061    }
15062
15063    struct BlockingApprovalHandler {
15064        entered: Arc<tokio::sync::Barrier>,
15065        release: Arc<tokio::sync::Notify>,
15066        result: ApprovalResult,
15067    }
15068
15069    struct CountingApprovalHandler {
15070        calls: Arc<std::sync::atomic::AtomicUsize>,
15071    }
15072
15073    struct RuntimeWebFetchTransport {
15074        calls: Arc<std::sync::atomic::AtomicUsize>,
15075    }
15076
15077    struct RuntimeWebFetchResolver;
15078
15079    struct ReentrantToolHooks {
15080        agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
15081        invoked: AtomicBool,
15082        nested_success: AtomicBool,
15083    }
15084
15085    #[async_trait]
15086    impl ai_agents_core::Tool for SlowTool {
15087        fn id(&self) -> &str {
15088            "slow"
15089        }
15090
15091        fn name(&self) -> &str {
15092            "Slow"
15093        }
15094
15095        fn description(&self) -> &str {
15096            "Waits until cancelled or timed out."
15097        }
15098
15099        fn input_schema(&self) -> Value {
15100            serde_json::json!({"type": "object"})
15101        }
15102
15103        async fn execute(
15104            &self,
15105            _args: Value,
15106            _ctx: ai_agents_core::ToolExecutionContext,
15107        ) -> ToolResult {
15108            tokio::time::sleep(std::time::Duration::from_secs(5)).await;
15109            ToolResult::ok("done")
15110        }
15111    }
15112
15113    #[async_trait]
15114    impl ai_agents_core::Tool for FlakyWriteTool {
15115        fn id(&self) -> &str {
15116            "flaky_write"
15117        }
15118
15119        fn name(&self) -> &str {
15120            "Flaky Write"
15121        }
15122
15123        fn description(&self) -> &str {
15124            "Fails on the first write attempt."
15125        }
15126
15127        fn input_schema(&self) -> Value {
15128            serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
15129        }
15130
15131        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15132            ai_agents_core::ToolPolicyBindings {
15133                path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15134                ..Default::default()
15135            }
15136        }
15137
15138        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15139            ai_agents_core::ToolSafetyMetadata {
15140                read_only: false,
15141                concurrency_safe: false,
15142                operation: ai_agents_core::ToolOperationKind::Write,
15143                side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
15144                requires_network: false,
15145                destructive: false,
15146                open_world: false,
15147                host_dependent: false,
15148                requires_user_interaction: false,
15149                supports_cancellation: true,
15150                default_requires_approval: false,
15151                should_defer_schema: false,
15152                max_output_chars: Some(1024),
15153                max_result_size_chars: Some(1024),
15154            }
15155        }
15156
15157        fn classify_call(&self, _args: &Value) -> ai_agents_core::ToolCallClassification {
15158            let mut classification =
15159                ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
15160            classification.safely_retryable = false;
15161            classification
15162        }
15163
15164        async fn execute(
15165            &self,
15166            _args: Value,
15167            _ctx: ai_agents_core::ToolExecutionContext,
15168        ) -> ToolResult {
15169            let call = self.calls.fetch_add(1, Ordering::SeqCst);
15170            if call == 0 {
15171                ToolResult::error("first failure")
15172            } else {
15173                ToolResult::ok("second success")
15174            }
15175        }
15176    }
15177
15178    #[async_trait]
15179    impl ai_agents_core::Tool for LockedWriteTool {
15180        fn id(&self) -> &str {
15181            "locked_write"
15182        }
15183
15184        fn name(&self) -> &str {
15185            "Locked Write"
15186        }
15187
15188        fn description(&self) -> &str {
15189            "Tracks concurrent execution on one resource."
15190        }
15191
15192        fn input_schema(&self) -> Value {
15193            serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
15194        }
15195
15196        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15197            ai_agents_core::ToolPolicyBindings {
15198                path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15199                ..Default::default()
15200            }
15201        }
15202
15203        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15204            ai_agents_core::ToolSafetyMetadata {
15205                read_only: false,
15206                concurrency_safe: false,
15207                operation: ai_agents_core::ToolOperationKind::Write,
15208                side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
15209                requires_network: false,
15210                destructive: false,
15211                open_world: false,
15212                host_dependent: false,
15213                requires_user_interaction: false,
15214                supports_cancellation: true,
15215                default_requires_approval: false,
15216                should_defer_schema: false,
15217                max_output_chars: Some(1024),
15218                max_result_size_chars: Some(1024),
15219            }
15220        }
15221
15222        async fn execute(
15223            &self,
15224            _args: Value,
15225            _ctx: ai_agents_core::ToolExecutionContext,
15226        ) -> ToolResult {
15227            let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
15228            loop {
15229                let current_max = self.max_active.load(Ordering::SeqCst);
15230                if active <= current_max {
15231                    break;
15232                }
15233                if self
15234                    .max_active
15235                    .compare_exchange(current_max, active, Ordering::SeqCst, Ordering::SeqCst)
15236                    .is_ok()
15237                {
15238                    break;
15239                }
15240            }
15241            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
15242            self.active.fetch_sub(1, Ordering::SeqCst);
15243            ToolResult::ok("done")
15244        }
15245    }
15246
15247    #[async_trait]
15248    impl ai_agents_core::Tool for MultiResourceWriteTool {
15249        fn id(&self) -> &str {
15250            "multi_resource_write"
15251        }
15252
15253        fn name(&self) -> &str {
15254            "Multi Resource Write"
15255        }
15256
15257        fn description(&self) -> &str {
15258            "Tracks concurrent execution across source and destination resources."
15259        }
15260
15261        fn input_schema(&self) -> Value {
15262            serde_json::json!({"type": "object"})
15263        }
15264
15265        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15266            ai_agents_core::ToolPolicyBindings {
15267                path_fields: vec![
15268                    ai_agents_core::PathPolicyBinding::read_write("source_path"),
15269                    ai_agents_core::PathPolicyBinding::write("destination_path"),
15270                ],
15271                ..Default::default()
15272            }
15273        }
15274
15275        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15276            LockedWriteTool {
15277                active: Arc::clone(&self.active),
15278                max_active: Arc::clone(&self.max_active),
15279            }
15280            .safety_metadata()
15281        }
15282
15283        async fn execute(
15284            &self,
15285            _args: Value,
15286            _ctx: ai_agents_core::ToolExecutionContext,
15287        ) -> ToolResult {
15288            let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
15289            self.max_active.fetch_max(active, Ordering::SeqCst);
15290            tokio::time::sleep(std::time::Duration::from_millis(75)).await;
15291            self.active.fetch_sub(1, Ordering::SeqCst);
15292            ToolResult::ok("done")
15293        }
15294    }
15295
15296    #[async_trait]
15297    impl ai_agents_core::Tool for BlockingPathMutationTool {
15298        fn id(&self) -> &str {
15299            self.id
15300        }
15301
15302        fn name(&self) -> &str {
15303            self.id
15304        }
15305
15306        fn description(&self) -> &str {
15307            "Blocks a path mutation until the test releases it."
15308        }
15309
15310        fn input_schema(&self) -> Value {
15311            serde_json::json!({"type": "object"})
15312        }
15313
15314        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15315            ai_agents_core::ToolPolicyBindings {
15316                path_fields: self.path_fields.clone(),
15317                ..Default::default()
15318            }
15319        }
15320
15321        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15322            ai_agents_core::ToolSafetyMetadata {
15323                read_only: false,
15324                concurrency_safe: false,
15325                operation: ai_agents_core::ToolOperationKind::Write,
15326                side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
15327                requires_network: false,
15328                destructive: false,
15329                open_world: false,
15330                host_dependent: false,
15331                requires_user_interaction: false,
15332                supports_cancellation: true,
15333                default_requires_approval: false,
15334                should_defer_schema: false,
15335                max_output_chars: Some(1024),
15336                max_result_size_chars: Some(1024),
15337            }
15338        }
15339
15340        async fn execute(
15341            &self,
15342            _args: Value,
15343            _ctx: ai_agents_core::ToolExecutionContext,
15344        ) -> ToolResult {
15345            self.gate.entered.store(true, Ordering::SeqCst);
15346            self.gate.entered_notify.notify_one();
15347            self.gate.release.notified().await;
15348            ToolResult::ok("done")
15349        }
15350    }
15351
15352    #[async_trait]
15353    impl ai_agents_core::Tool for NoBindingWriteTool {
15354        fn id(&self) -> &str {
15355            "no_binding_write"
15356        }
15357
15358        fn name(&self) -> &str {
15359            "No Binding Write"
15360        }
15361
15362        fn description(&self) -> &str {
15363            "Tracks concurrent execution without resource bindings."
15364        }
15365
15366        fn input_schema(&self) -> Value {
15367            serde_json::json!({"type": "object"})
15368        }
15369
15370        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15371            LockedWriteTool {
15372                active: Arc::clone(&self.active),
15373                max_active: Arc::clone(&self.max_active),
15374            }
15375            .safety_metadata()
15376        }
15377
15378        async fn execute(
15379            &self,
15380            _args: Value,
15381            _ctx: ai_agents_core::ToolExecutionContext,
15382        ) -> ToolResult {
15383            let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
15384            self.max_active.fetch_max(active, Ordering::SeqCst);
15385            tokio::time::sleep(std::time::Duration::from_millis(75)).await;
15386            self.active.fetch_sub(1, Ordering::SeqCst);
15387            ToolResult::ok("done")
15388        }
15389    }
15390
15391    #[async_trait]
15392    impl ai_agents_core::Tool for RecoveryTestTool {
15393        fn id(&self) -> &str {
15394            &self.id
15395        }
15396
15397        fn name(&self) -> &str {
15398            &self.id
15399        }
15400
15401        fn description(&self) -> &str {
15402            "Records recovery execution and returns a configured result."
15403        }
15404
15405        fn input_schema(&self) -> Value {
15406            serde_json::json!({"type": "object"})
15407        }
15408
15409        fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15410            ai_agents_core::ToolPolicyBindings {
15411                path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15412                ..Default::default()
15413            }
15414        }
15415
15416        /// Supplies a configurable output cap while preserving non-concurrent path mutation behavior.
15417        fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15418            ai_agents_core::ToolSafetyMetadata {
15419                read_only: false,
15420                concurrency_safe: false,
15421                operation: ai_agents_core::ToolOperationKind::Write,
15422                side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
15423                requires_network: false,
15424                destructive: false,
15425                open_world: false,
15426                host_dependent: false,
15427                requires_user_interaction: false,
15428                supports_cancellation: true,
15429                default_requires_approval: false,
15430                should_defer_schema: false,
15431                max_output_chars: Some(self.max_output_chars.unwrap_or(1024)),
15432                max_result_size_chars: Some(1024),
15433            }
15434        }
15435
15436        /// Returns deterministic output and metadata for recovery lifecycle assertions.
15437        async fn execute(
15438            &self,
15439            _args: Value,
15440            _ctx: ai_agents_core::ToolExecutionContext,
15441        ) -> ToolResult {
15442            self.calls.fetch_add(1, Ordering::SeqCst);
15443            let mut result = if self.succeeds {
15444                ToolResult::ok(format!("{} succeeded", self.id))
15445            } else {
15446                ToolResult::error(format!("{} failed", self.id))
15447            };
15448            result.metadata = Some(HashMap::from([(
15449                "recovery_test_tool".to_string(),
15450                Value::String(self.id.clone()),
15451            )]));
15452            result
15453        }
15454    }
15455
15456    #[async_trait]
15457    impl WebFetchTransport for RuntimeWebFetchTransport {
15458        /// Rejects unvalidated transport calls in the shared-executor fixture.
15459        async fn send(
15460            &self,
15461            _request: WebFetchTransportRequest,
15462        ) -> std::result::Result<WebFetchTransportResponse, String> {
15463            Err("validated addresses are required".to_string())
15464        }
15465
15466        /// Records transport only after runtime and tool policy validation complete.
15467        async fn send_validated(
15468            &self,
15469            _request: WebFetchTransportRequest,
15470            _addresses: &[std::net::SocketAddr],
15471        ) -> std::result::Result<WebFetchTransportResponse, String> {
15472            self.calls.fetch_add(1, Ordering::SeqCst);
15473            Ok(WebFetchTransportResponse {
15474                status: 200,
15475                content_type: Some("text/plain".to_string()),
15476                location: None,
15477                body: b"approved".to_vec(),
15478            })
15479        }
15480    }
15481
15482    #[async_trait]
15483    impl WebFetchResolver for RuntimeWebFetchResolver {
15484        /// Returns one public fixture address without external DNS.
15485        async fn resolve(
15486            &self,
15487            _host: &str,
15488            _port: u16,
15489        ) -> std::result::Result<Vec<std::net::IpAddr>, String> {
15490            Ok(vec![std::net::IpAddr::V4(std::net::Ipv4Addr::new(
15491                93, 184, 216, 34,
15492            ))])
15493        }
15494    }
15495
15496    #[async_trait]
15497    impl ApprovalHandler for BlockingApprovalHandler {
15498        async fn request_approval(
15499            &self,
15500            _request: ai_agents_hitl::ApprovalRequest,
15501        ) -> ApprovalResult {
15502            self.entered.wait().await;
15503            self.release.notified().await;
15504            self.result.clone()
15505        }
15506    }
15507
15508    #[async_trait]
15509    impl ApprovalHandler for CountingApprovalHandler {
15510        async fn request_approval(
15511            &self,
15512            _request: ai_agents_hitl::ApprovalRequest,
15513        ) -> ApprovalResult {
15514            self.calls.fetch_add(1, Ordering::SeqCst);
15515            ApprovalResult::Approved
15516        }
15517    }
15518
15519    #[async_trait]
15520    impl AgentHooks for ReentrantToolHooks {
15521        async fn on_tool_complete(&self, tool: &str, _result: &ToolResult, _duration_ms: u64) {
15522            if tool != "reentrant_write" || self.invoked.swap(true, Ordering::SeqCst) {
15523                return;
15524            }
15525            let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
15526            if let Some(agent) = agent {
15527                let result = agent
15528                    .invoke_tool(ToolExecutionRequest::new(
15529                        "nested-hook-call",
15530                        "reentrant_write",
15531                        serde_json::json!({"path": "./hook.txt"}),
15532                        ToolCallSource::Manual,
15533                    ))
15534                    .await;
15535                self.nested_success
15536                    .store(result.is_ok_and(|record| record.success), Ordering::SeqCst);
15537            }
15538        }
15539    }
15540
15541    #[async_trait]
15542    impl AgentHooks for ResponseCountingHooks {
15543        async fn on_response(&self, _response: &AgentResponse) {
15544            self.responses.fetch_add(1, Ordering::SeqCst);
15545        }
15546    }
15547
15548    #[async_trait]
15549    impl AgentHooks for ResponseChatHooks {
15550        /// Attempts one nested blocking turn without retaining the target mutex across the await.
15551        async fn on_response(&self, _response: &AgentResponse) {
15552            if self.invoked.swap(true, Ordering::SeqCst) {
15553                return;
15554            }
15555            let target = self.target.lock().as_ref().and_then(Weak::upgrade);
15556            let result = if let Some(target) = target {
15557                target
15558                    .chat("nested response hook call")
15559                    .await
15560                    .map(|response| response.content)
15561                    .map_err(|error| error.to_string())
15562            } else {
15563                Err("response hook target is unavailable".to_string())
15564            };
15565            *self.nested_result.lock() = Some(result);
15566        }
15567    }
15568
15569    #[async_trait]
15570    impl AgentHooks for ConcurrentResponseHooks {
15571        /// Runs one child through the real concurrent JoinSet boundary without retaining registry state across the await.
15572        async fn on_response(&self, _response: &AgentResponse) {
15573            if self.invoked.swap(true, Ordering::SeqCst) {
15574                return;
15575            }
15576            let Some(registry) = self.registry.upgrade() else {
15577                *self.nested_result.lock() =
15578                    Some(Err("concurrent registry is unavailable".to_string()));
15579                return;
15580            };
15581            let agents = [ai_agents_state::ConcurrentAgentRef::Id(
15582                self.child_id.clone(),
15583            )];
15584            let aggregation = ai_agents_state::AggregationConfig {
15585                strategy: ai_agents_state::AggregationStrategy::FirstWins,
15586                synthesizer_llm: None,
15587                synthesizer_prompt: None,
15588                vote: None,
15589            };
15590            let result = crate::orchestration::concurrent(
15591                &registry,
15592                "nested concurrent response hook call",
15593                &agents,
15594                &aggregation,
15595                None,
15596                Some(1),
15597                None,
15598                ai_agents_state::PartialFailureAction::Abort,
15599                None,
15600            )
15601            .await
15602            .map(|result| result.response.content)
15603            .map_err(|error| error.to_string());
15604            *self.nested_result.lock() = Some(result);
15605        }
15606    }
15607
15608    #[async_trait]
15609    impl AgentHooks for ToolLifecycleRecordingHooks {
15610        async fn on_tool_start(&self, tool: &str, _args: &Value) {
15611            self.events.lock().push(format!("start:{tool}"));
15612        }
15613
15614        async fn on_tool_complete(&self, tool: &str, result: &ToolResult, _duration_ms: u64) {
15615            self.events
15616                .lock()
15617                .push(format!("complete:{tool}:{}", result.success));
15618        }
15619
15620        async fn on_tool_execution_record(&self, record: &ToolExecutionRecord) {
15621            self.events.lock().push(format!(
15622                "record:{}:{}",
15623                record.canonical_id, record.executed
15624            ));
15625            self.records.lock().push(record.clone());
15626        }
15627
15628        /// Records error reporting so fallback cannot overtake the failed original lifecycle.
15629        async fn on_error(&self, _error: &AgentError) {
15630            self.events.lock().push("error".to_string());
15631        }
15632    }
15633
15634    struct ApprovalRecordingHooks {
15635        events: parking_lot::Mutex<Vec<String>>,
15636    }
15637
15638    impl ApprovalRecordingHooks {
15639        fn new() -> Self {
15640            Self {
15641                events: parking_lot::Mutex::new(Vec::new()),
15642            }
15643        }
15644
15645        fn events(&self) -> Vec<String> {
15646            self.events.lock().clone()
15647        }
15648    }
15649
15650    #[async_trait]
15651    impl AgentHooks for ApprovalRecordingHooks {
15652        async fn on_approval_result(&self, request_id: &str, result: &ApprovalResult) {
15653            self.events.lock().push(format!(
15654                "raw:{}:{}",
15655                request_id,
15656                approval_result_name(result)
15657            ));
15658        }
15659
15660        async fn on_approval_resolved(
15661            &self,
15662            request: &ai_agents_hitl::ApprovalRequest,
15663            raw_result: &ApprovalResult,
15664            outcome: &ApprovalResolvedOutcome,
15665        ) {
15666            self.events.lock().push(format!(
15667                "resolved:{}:{}:{}",
15668                request.id,
15669                approval_result_name(raw_result),
15670                approval_outcome_name(outcome)
15671            ));
15672        }
15673    }
15674
15675    fn approval_result_name(result: &ApprovalResult) -> &'static str {
15676        match result {
15677            ApprovalResult::Approved => "approved",
15678            ApprovalResult::Rejected { .. } => "rejected",
15679            ApprovalResult::Modified { .. } => "modified",
15680            ApprovalResult::Timeout => "timeout",
15681        }
15682    }
15683
15684    fn approval_outcome_name(outcome: &ApprovalResolvedOutcome) -> &'static str {
15685        match outcome {
15686            ApprovalResolvedOutcome::Approved => "approved",
15687            ApprovalResolvedOutcome::Rejected { .. } => "rejected",
15688            ApprovalResolvedOutcome::Modified { .. } => "modified",
15689            ApprovalResolvedOutcome::Error { .. } => "error",
15690        }
15691    }
15692
15693    fn assert_correlated_approval_events(
15694        events: &[String],
15695        raw_status: &str,
15696        outcome_status: &str,
15697    ) {
15698        assert_eq!(events.len(), 2);
15699        let raw: Vec<_> = events[0].split(':').collect();
15700        let resolved: Vec<_> = events[1].split(':').collect();
15701        assert_eq!(raw[0], "raw");
15702        assert_eq!(resolved[0], "resolved");
15703        assert_eq!(raw[1], resolved[1]);
15704        assert_eq!(raw[2], raw_status);
15705        assert_eq!(resolved[2], raw_status);
15706        assert_eq!(resolved[3], outcome_status);
15707    }
15708
15709    fn approval_security_config(policy_enabled: bool) -> ToolSecurityConfig {
15710        let mut security = ToolSecurityConfig {
15711            enabled: true,
15712            fail_closed: true,
15713            ..Default::default()
15714        };
15715        let policy = ai_agents_tools::ToolPolicyConfig {
15716            enabled: policy_enabled,
15717            write_paths: vec![".".to_string()],
15718            require_confirmation: true,
15719            ..Default::default()
15720        };
15721        security.tools.insert("locked_write".to_string(), policy);
15722        security
15723    }
15724
15725    struct MutationTestWorkspace {
15726        root: std::path::PathBuf,
15727    }
15728
15729    impl MutationTestWorkspace {
15730        fn new() -> Self {
15731            let root = std::env::temp_dir().join(format!(
15732                "ai-agents-runtime-mutation-{}",
15733                uuid::Uuid::new_v4()
15734            ));
15735            std::fs::create_dir_all(&root).unwrap();
15736            Self { root }
15737        }
15738    }
15739
15740    impl Drop for MutationTestWorkspace {
15741        fn drop(&mut self) {
15742            let _ = std::fs::remove_dir_all(&self.root);
15743        }
15744    }
15745
15746    async fn wait_for_resource_lock_strong_count(locks: &ToolResourceLocks, minimum: usize) {
15747        tokio::time::timeout(std::time::Duration::from_secs(2), async {
15748            loop {
15749                let strong_count = locks
15750                    .read()
15751                    .get("path-mutation:global")
15752                    .map_or(0, |lock| lock.strong_count());
15753                if strong_count >= minimum {
15754                    break;
15755                }
15756                tokio::task::yield_now().await;
15757            }
15758        })
15759        .await
15760        .expect("path mutation call did not reach the shared lock");
15761    }
15762
15763    async fn assert_path_mutation_pair_serialized(
15764        first_id: &'static str,
15765        first_fields: Vec<ai_agents_core::PathPolicyBinding>,
15766        first_args: Value,
15767        second_id: &'static str,
15768        second_fields: Vec<ai_agents_core::PathPolicyBinding>,
15769        second_args: Value,
15770    ) {
15771        let locks = new_tool_resource_locks();
15772        let first_gate = PathMutationGate::new();
15773        let second_gate = PathMutationGate::new();
15774        second_gate.release();
15775        let agent = Arc::new(
15776            AgentBuilder::new()
15777                .system_prompt("Test global path mutation locking.")
15778                .llm(Arc::new(mock_with_response("done")))
15779                .tool(Arc::new(BlockingPathMutationTool {
15780                    id: first_id,
15781                    path_fields: first_fields,
15782                    gate: first_gate.clone(),
15783                }))
15784                .tool(Arc::new(BlockingPathMutationTool {
15785                    id: second_id,
15786                    path_fields: second_fields,
15787                    gate: second_gate.clone(),
15788                }))
15789                .build()
15790                .unwrap()
15791                .with_shared_resource_locks(Arc::clone(&locks)),
15792        );
15793
15794        let first = {
15795            let agent = Arc::clone(&agent);
15796            tokio::spawn(async move {
15797                agent
15798                    .invoke_tool(ToolExecutionRequest::new(
15799                        format!("{}-first", first_id),
15800                        first_id,
15801                        first_args,
15802                        ToolCallSource::Manual,
15803                    ))
15804                    .await
15805                    .unwrap()
15806            })
15807        };
15808        first_gate.wait_until_entered().await;
15809
15810        let second = {
15811            let agent = Arc::clone(&agent);
15812            tokio::spawn(async move {
15813                agent
15814                    .invoke_tool(ToolExecutionRequest::new(
15815                        format!("{}-second", second_id),
15816                        second_id,
15817                        second_args,
15818                        ToolCallSource::Manual,
15819                    ))
15820                    .await
15821                    .unwrap()
15822            })
15823        };
15824        wait_for_resource_lock_strong_count(&locks, 2).await;
15825        assert!(!second_gate.entered.load(Ordering::SeqCst));
15826        assert!(!second.is_finished());
15827
15828        first_gate.release();
15829        let (first, second) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
15830            tokio::join!(first, second)
15831        })
15832        .await
15833        .expect("serialized path mutation calls did not finish");
15834        assert!(first.unwrap().success);
15835        assert!(second.unwrap().success);
15836        assert!(second_gate.entered.load(Ordering::SeqCst));
15837        assert!(locks.read().is_empty());
15838    }
15839
15840    #[derive(Clone, Copy)]
15841    enum MutationDenial {
15842        Policy,
15843        Approval,
15844    }
15845
15846    fn mutation_denial_security_config(
15847        tool_id: &str,
15848        workspace: &std::path::Path,
15849        denial: MutationDenial,
15850    ) -> ToolSecurityConfig {
15851        let workspace = workspace.to_string_lossy().into_owned();
15852        let mut policy = ai_agents_tools::ToolPolicyConfig {
15853            read_paths: vec![workspace.clone()],
15854            write_paths: vec![workspace.clone()],
15855            ..Default::default()
15856        };
15857        match denial {
15858            MutationDenial::Policy => policy.blocked_paths = vec![workspace],
15859            MutationDenial::Approval => policy.require_confirmation = true,
15860        }
15861
15862        let mut security = ToolSecurityConfig {
15863            enabled: true,
15864            fail_closed: true,
15865            ..Default::default()
15866        };
15867        security.tools.insert(tool_id.to_string(), policy);
15868        security
15869    }
15870
15871    async fn assert_path_mutation_denied(tool: Arc<dyn Tool>, denial: MutationDenial) {
15872        let workspace = MutationTestWorkspace::new();
15873        let tool_id = tool.id().to_string();
15874        let preserved = workspace.root.join(format!("{}-preserved.txt", tool_id));
15875        let destination = workspace.root.join(format!("{}-destination.txt", tool_id));
15876        std::fs::write(&preserved, "preserved").unwrap();
15877        let arguments = match tool_id.as_str() {
15878            "copy_path" | "move_path" => serde_json::json!({
15879                "source_path": preserved.to_string_lossy(),
15880                "destination_path": destination.to_string_lossy(),
15881                "dry_run": false
15882            }),
15883            "delete_path" => serde_json::json!({
15884                "path": preserved.to_string_lossy(),
15885                "recursive": false,
15886                "dry_run": false
15887            }),
15888            _ => panic!("unsupported mutation tool: {}", tool_id),
15889        };
15890        let security = mutation_denial_security_config(&tool_id, &workspace.root, denial);
15891        let builder = AgentBuilder::new()
15892            .system_prompt("Test mutation denial.")
15893            .llm(Arc::new(mock_with_response("done")))
15894            .tool(tool)
15895            .tool_security(ToolSecurityEngine::new(security));
15896        let builder = match denial {
15897            MutationDenial::Policy => builder,
15898            MutationDenial::Approval => builder
15899                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
15900                .approval_handler(Arc::new(RejectAllHandler::new())),
15901        };
15902        let agent = builder.build().unwrap();
15903
15904        let record = agent
15905            .invoke_tool(ToolExecutionRequest::new(
15906                format!("{}-denied", tool_id),
15907                tool_id.clone(),
15908                arguments,
15909                ToolCallSource::Manual,
15910            ))
15911            .await
15912            .unwrap();
15913
15914        assert!(!record.executed, "{} must not be invoked", tool_id);
15915        assert!(!record.success);
15916        match denial {
15917            MutationDenial::Policy => {
15918                assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
15919                assert!(record.approval.as_ref().is_some_and(|approval| matches!(
15920                    &approval.status,
15921                    ToolApprovalStatus::NotRequired
15922                )));
15923            }
15924            MutationDenial::Approval => {
15925                assert_eq!(record.policy.outcome, PermissionOutcome::RequiresApproval);
15926                assert!(record.approval.as_ref().is_some_and(|approval| matches!(
15927                    &approval.status,
15928                    ToolApprovalStatus::Rejected
15929                )));
15930            }
15931        }
15932        assert_eq!(std::fs::read_to_string(&preserved).unwrap(), "preserved");
15933        assert!(!destination.exists());
15934    }
15935
15936    fn recovery_manager_with_fallbacks(
15937        fallbacks: impl IntoIterator<Item = (String, String)>,
15938    ) -> RecoveryManager {
15939        use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
15940
15941        let per_tool = fallbacks
15942            .into_iter()
15943            .map(|(tool, fallback_tool)| {
15944                (
15945                    tool,
15946                    ToolRetryConfig {
15947                        max_retries: 0,
15948                        timeout_ms: Some(1_000),
15949                        on_failure: ToolFailureAction::Fallback { fallback_tool },
15950                    },
15951                )
15952            })
15953            .collect();
15954        RecoveryManager::new(ErrorRecoveryConfig {
15955            tools: ToolRecoveryConfig {
15956                per_tool,
15957                ..Default::default()
15958            },
15959            ..Default::default()
15960        })
15961    }
15962
15963    fn approval_check() -> HITLCheckResult {
15964        HITLCheckResult::required(
15965            ApprovalTrigger::tool("test", serde_json::json!({})),
15966            HashMap::new(),
15967            "Approve?",
15968            None,
15969        )
15970    }
15971
15972    fn agent_with_approval_result(
15973        raw_result: ApprovalResult,
15974        timeout_action: TimeoutAction,
15975        hooks: Arc<ApprovalRecordingHooks>,
15976    ) -> RuntimeAgent {
15977        use ai_agents_hitl::{CallbackHandler, HITLConfig};
15978
15979        let config = HITLConfig {
15980            on_timeout: timeout_action,
15981            ..Default::default()
15982        };
15983        let handler = CallbackHandler::new(move |_| raw_result.clone());
15984        AgentBuilder::new()
15985            .system_prompt("Test HITL hooks.")
15986            .llm(Arc::new(mock_with_response("done")))
15987            .build()
15988            .unwrap()
15989            .with_hooks(hooks)
15990            .with_hitl(HITLEngine::new(config), Arc::new(handler))
15991    }
15992
15993    #[tokio::test]
15994    async fn approval_hooks_expose_direct_effective_decisions_after_raw_results() {
15995        let cases = vec![
15996            (ApprovalResult::Approved, "approved"),
15997            (
15998                ApprovalResult::Rejected {
15999                    reason: Some("denied".to_string()),
16000                },
16001                "rejected",
16002            ),
16003            (
16004                ApprovalResult::Modified {
16005                    changes: HashMap::from([("value".to_string(), serde_json::json!(2))]),
16006                },
16007                "modified",
16008            ),
16009        ];
16010
16011        for (raw_result, expected) in cases {
16012            let hooks = Arc::new(ApprovalRecordingHooks::new());
16013            let agent =
16014                agent_with_approval_result(raw_result, TimeoutAction::Reject, hooks.clone());
16015
16016            let result = agent.request_hitl_approval(approval_check()).await.unwrap();
16017
16018            assert_eq!(approval_result_name(&result), expected);
16019            assert_correlated_approval_events(&hooks.events(), expected, expected);
16020        }
16021    }
16022
16023    #[tokio::test]
16024    async fn approval_hooks_expose_timeout_policy_decisions() {
16025        for (timeout_action, expected) in [
16026            (TimeoutAction::Approve, "approved"),
16027            (TimeoutAction::Reject, "rejected"),
16028        ] {
16029            let hooks = Arc::new(ApprovalRecordingHooks::new());
16030            let agent =
16031                agent_with_approval_result(ApprovalResult::Timeout, timeout_action, hooks.clone());
16032
16033            let result = agent.request_hitl_approval(approval_check()).await.unwrap();
16034
16035            assert_eq!(approval_result_name(&result), expected);
16036            assert_correlated_approval_events(&hooks.events(), "timeout", expected);
16037        }
16038    }
16039
16040    #[tokio::test]
16041    async fn timeout_error_fires_correlated_resolved_error_before_returning() {
16042        let hooks = Arc::new(ApprovalRecordingHooks::new());
16043        let agent = agent_with_approval_result(
16044            ApprovalResult::Timeout,
16045            TimeoutAction::Error,
16046            hooks.clone(),
16047        );
16048
16049        let error = agent
16050            .request_hitl_approval(approval_check())
16051            .await
16052            .unwrap_err();
16053
16054        assert!(error.to_string().contains("HITL approval timeout"));
16055        assert_correlated_approval_events(&hooks.events(), "timeout", "error");
16056    }
16057
16058    // Basic YAML → Build → Chat flow
16059    #[tokio::test]
16060    async fn test_integration_yaml_to_chat_basic() {
16061        let mock = mock_with_response("Hello! How can I help you?");
16062        let agent = AgentBuilder::new()
16063            .system_prompt("You are a test assistant.")
16064            .llm(Arc::new(mock))
16065            .build()
16066            .unwrap();
16067
16068        let response = agent.chat("Hi").await.unwrap();
16069        assert!(!response.content.is_empty());
16070        assert_eq!(response.content, "Hello! How can I help you?");
16071    }
16072
16073    #[tokio::test]
16074    async fn stream_events_emit_one_authoritative_final_without_legacy_done() {
16075        let agent = AgentBuilder::new()
16076            .system_prompt("You are a test assistant.")
16077            .llm(Arc::new(mock_with_response(
16078                "Hello from the final response.",
16079            )))
16080            .build()
16081            .unwrap();
16082
16083        let mut stream = agent.chat_stream_events("Hi").await.unwrap();
16084        let mut final_responses = Vec::new();
16085        let mut legacy_done = 0;
16086        while let Some(event) = stream.next().await {
16087            match event {
16088                AgentStreamEvent::Chunk(StreamChunk::Done {}) => legacy_done += 1,
16089                AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
16090                    panic!("unexpected stream error: {message}")
16091                }
16092                AgentStreamEvent::Final(response) => final_responses.push(response),
16093                AgentStreamEvent::Chunk(_) => {}
16094            }
16095        }
16096
16097        assert_eq!(legacy_done, 0);
16098        assert_eq!(final_responses.len(), 1);
16099        let response = final_responses.pop().unwrap();
16100        assert_eq!(response.content, "Hello from the final response.");
16101        assert!(
16102            response
16103                .metadata
16104                .as_ref()
16105                .is_some_and(|metadata| { metadata.contains_key("reasoning") })
16106        );
16107    }
16108
16109    #[tokio::test]
16110    async fn stream_final_content_includes_output_processing_after_provisional_chunks() {
16111        let yaml = r#"
16112name: ProcessedStreamAgent
16113system_prompt: "Answer directly."
16114process:
16115  output:
16116    - type: format
16117      config:
16118        template: "{{ response }} [finalized]"
16119streaming:
16120  enabled: true
16121"#;
16122        let agent = AgentBuilder::from_yaml(yaml)
16123            .unwrap()
16124            .llm(Arc::new(mock_with_response("provisional answer")))
16125            .auto_configure_features()
16126            .unwrap()
16127            .build()
16128            .unwrap();
16129
16130        let mut stream = agent.chat_stream_events("Hi").await.unwrap();
16131        let mut provisional = String::new();
16132        let mut final_content = None;
16133        while let Some(event) = stream.next().await {
16134            match event {
16135                AgentStreamEvent::Chunk(StreamChunk::Content { text }) => {
16136                    provisional.push_str(&text)
16137                }
16138                AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
16139                    panic!("unexpected stream error: {message}")
16140                }
16141                AgentStreamEvent::Final(response) => final_content = Some(response.content),
16142                AgentStreamEvent::Chunk(_) => {}
16143            }
16144        }
16145
16146        assert_eq!(provisional, "provisional answer");
16147        assert_eq!(
16148            final_content.as_deref(),
16149            Some("provisional answer [finalized]")
16150        );
16151    }
16152
16153    #[tokio::test]
16154    async fn stream_events_preserve_tool_progress_and_final_tool_calls() {
16155        let agent = AgentBuilder::new()
16156            .system_prompt("Use the echo tool once, then answer.")
16157            .llm(Arc::new(mock_with_responses(vec![
16158                r#"{"tool":"echo","arguments":{"message":"hello"}}"#,
16159                "Echo completed.",
16160            ])))
16161            .tool(Arc::new(ai_agents_tools::EchoTool::new()))
16162            .build()
16163            .unwrap();
16164
16165        let mut stream = agent.chat_stream_events("echo hello").await.unwrap();
16166        let mut starts = 0;
16167        let mut results = 0;
16168        let mut ends = 0;
16169        let mut final_response = None;
16170        while let Some(event) = stream.next().await {
16171            match event {
16172                AgentStreamEvent::Chunk(StreamChunk::ToolCallStart { name, .. }) => {
16173                    assert_eq!(name, "echo");
16174                    starts += 1;
16175                }
16176                AgentStreamEvent::Chunk(StreamChunk::ToolResult { name, success, .. }) => {
16177                    assert_eq!(name, "echo");
16178                    assert!(success);
16179                    results += 1;
16180                }
16181                AgentStreamEvent::Chunk(StreamChunk::ToolCallEnd { .. }) => ends += 1,
16182                AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
16183                    panic!("unexpected stream error: {message}")
16184                }
16185                AgentStreamEvent::Final(response) => final_response = Some(response),
16186                AgentStreamEvent::Chunk(_) => {}
16187            }
16188        }
16189
16190        assert_eq!((starts, results, ends), (1, 1, 1));
16191        let response = final_response.expect("tool stream must finalize");
16192        assert_eq!(response.content, "Echo completed.");
16193        assert_eq!(
16194            response.tool_calls.as_ref().map(|calls| calls
16195                .iter()
16196                .map(|call| call.name.as_str())
16197                .collect::<Vec<_>>()),
16198            Some(vec!["echo"])
16199        );
16200    }
16201
16202    #[tokio::test]
16203    async fn legacy_stream_still_emits_one_done_chunk() {
16204        let agent = AgentBuilder::new()
16205            .system_prompt("You are a test assistant.")
16206            .llm(Arc::new(mock_with_response(
16207                "Hello from the legacy stream.",
16208            )))
16209            .build()
16210            .unwrap();
16211
16212        let mut stream = agent.chat_stream("Hi").await.unwrap();
16213        let mut done = 0;
16214        while let Some(chunk) = stream.next().await {
16215            match chunk {
16216                StreamChunk::Done {} => done += 1,
16217                StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
16218                _ => {}
16219            }
16220        }
16221
16222        assert_eq!(done, 1);
16223    }
16224
16225    // Multi-turn conversation
16226    #[tokio::test]
16227    async fn test_integration_multi_turn_conversation() {
16228        let mock = mock_with_responses(vec![
16229            "Hello! I'm your assistant.",
16230            "The weather is sunny today.",
16231            "Goodbye!",
16232        ]);
16233        let agent = AgentBuilder::new()
16234            .system_prompt("You are helpful.")
16235            .llm(Arc::new(mock))
16236            .build()
16237            .unwrap();
16238
16239        let r1 = agent.chat("Hi").await.unwrap();
16240        assert_eq!(r1.content, "Hello! I'm your assistant.");
16241
16242        let r2 = agent.chat("What's the weather?").await.unwrap();
16243        assert_eq!(r2.content, "The weather is sunny today.");
16244
16245        let r3 = agent.chat("Bye").await.unwrap();
16246        assert_eq!(r3.content, "Goodbye!");
16247
16248        // Verify memory accumulated messages
16249        let messages = agent.memory.get_messages(None).await.unwrap();
16250        // 3 user + 3 assistant = 6 messages
16251        assert_eq!(messages.len(), 6);
16252    }
16253
16254    #[test]
16255    fn later_approval_preserves_modified_evidence() {
16256        let arguments = serde_json::json!({"dry_run": true});
16257        let mut record = Some(ToolApprovalRecord {
16258            status: ToolApprovalStatus::Modified,
16259            reason: None,
16260            modified_arguments: Some(arguments.clone()),
16261        });
16262
16263        merge_approved_record(&mut record);
16264
16265        let record = record.unwrap();
16266        assert!(matches!(record.status, ToolApprovalStatus::Modified));
16267        assert_eq!(record.modified_arguments, Some(arguments));
16268    }
16269
16270    #[test]
16271    fn approval_binding_rejects_replaced_tool_implementation() {
16272        let reviewed_tool: Arc<dyn ai_agents_core::Tool> = Arc::new(ContextEchoTool);
16273        let same_tool = Arc::clone(&reviewed_tool);
16274        let replacement_tool: Arc<dyn ai_agents_core::Tool> = Arc::new(ContextEchoTool);
16275        let arguments = serde_json::json!({"path": "."});
16276        let versions = ToolDecisionVersions {
16277            policy: 2,
16278            registry: 3,
16279            runtime_control: 4,
16280            state: Some(5),
16281        };
16282        let binding = ToolApprovalBinding {
16283            canonical_id: "context_echo".to_string(),
16284            arguments: arguments.clone(),
16285            confirmation_required: true,
16286            policy_version: versions.policy,
16287            runtime_control_version: versions.runtime_control,
16288            state_generation: versions.state,
16289            reviewed_tool,
16290        };
16291
16292        assert!(!binding.is_stale("context_echo", &arguments, true, versions, &same_tool,));
16293        assert!(binding.is_stale(
16294            "context_echo",
16295            &arguments,
16296            true,
16297            versions,
16298            &replacement_tool,
16299        ));
16300    }
16301
16302    #[tokio::test]
16303    async fn approved_mutation_to_dry_run_remains_executable() {
16304        use ai_agents_hitl::CallbackHandler;
16305
16306        let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
16307            changes: HashMap::from([("dry_run".to_string(), serde_json::json!(true))]),
16308        });
16309        let agent = AgentBuilder::new()
16310            .system_prompt("Test safer approval modifications.")
16311            .llm(Arc::new(mock_with_response("done")))
16312            .tool(Arc::new(ai_agents_tools::FileWriteTool::new()))
16313            .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16314            .approval_handler(Arc::new(handler))
16315            .build()
16316            .unwrap();
16317
16318        let record = agent
16319            .invoke_tool(ToolExecutionRequest::new(
16320                "approved-dry-run",
16321                "file_write",
16322                serde_json::json!({
16323                    "path": "./approval-dry-run.txt",
16324                    "content": "not written"
16325                }),
16326                ToolCallSource::Manual,
16327            ))
16328            .await
16329            .unwrap();
16330
16331        assert!(record.executed);
16332        assert!(record.success);
16333        assert_eq!(record.executed_arguments["dry_run"], true);
16334        assert!(matches!(
16335            record.approval.as_ref().map(|approval| &approval.status),
16336            Some(ToolApprovalStatus::Modified)
16337        ));
16338        let output: Value = serde_json::from_str(&record.output).unwrap();
16339        assert_eq!(output["mutation_performed"], false);
16340    }
16341
16342    /// Proves shared HITL approval evidence reaches the web fetch implementation unchanged.
16343    #[tokio::test]
16344    async fn shared_executor_approval_reaches_web_fetch_transport() {
16345        use ai_agents_hitl::{CallbackHandler, HITLConfig};
16346        use ai_agents_tools::{DomainPolicyConfig, ToolPolicyConfig};
16347
16348        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16349        let tool = WebFetchTool::with_transport_and_resolver(
16350            Arc::new(RuntimeWebFetchTransport {
16351                calls: Arc::clone(&calls),
16352            }),
16353            Arc::new(RuntimeWebFetchResolver),
16354        );
16355        let mut security = ToolSecurityConfig {
16356            enabled: true,
16357            fail_closed: true,
16358            ..Default::default()
16359        };
16360        security.tools.insert(
16361            "web_fetch".to_string(),
16362            ToolPolicyConfig {
16363                domains: DomainPolicyConfig {
16364                    requires_approval: vec!["approval.test".to_string()],
16365                    ..Default::default()
16366                },
16367                allowed_schemes: vec!["https".to_string()],
16368                allowed_ports: vec![443],
16369                ..Default::default()
16370            },
16371        );
16372        let handler = CallbackHandler::new(|_| ApprovalResult::Approved);
16373        let agent = AgentBuilder::new()
16374            .system_prompt("Test approved web fetch execution.")
16375            .llm(Arc::new(mock_with_response("done")))
16376            .tool(Arc::new(tool))
16377            .tool_security(ToolSecurityEngine::new(security))
16378            .build()
16379            .unwrap()
16380            .with_hitl(HITLEngine::new(HITLConfig::default()), Arc::new(handler));
16381
16382        let record = agent
16383            .invoke_tool(ToolExecutionRequest::new(
16384                "approved-web-fetch",
16385                "web_fetch",
16386                serde_json::json!({
16387                    "url": "https://approval.test/page",
16388                    "cache_ttl_seconds": 0
16389                }),
16390                ToolCallSource::Manual,
16391            ))
16392            .await
16393            .unwrap();
16394
16395        assert!(record.success);
16396        assert!(
16397            record
16398                .approval
16399                .as_ref()
16400                .is_some_and(|approval| matches!(approval.status, ToolApprovalStatus::Approved))
16401        );
16402        assert_eq!(calls.load(Ordering::SeqCst), 1);
16403    }
16404
16405    #[tokio::test]
16406    async fn context_preserves_requested_and_canonical_identity() {
16407        let mock = mock_with_response("hello");
16408        let mut tools = ai_agents_tools::ToolRegistry::new();
16409        tools.register(Arc::new(ContextEchoTool)).unwrap();
16410
16411        let mut security = ToolSecurityConfig {
16412            enabled: true,
16413            fail_closed: true,
16414            ..Default::default()
16415        };
16416        let mut policy = ai_agents_tools::ToolPolicyConfig {
16417            read_paths: vec![".".to_string()],
16418            max_results: Some(7),
16419            ..Default::default()
16420        };
16421        policy
16422            .config
16423            .insert("backend".to_string(), serde_json::json!("memory"));
16424        security.tools.insert("context_echo".to_string(), policy);
16425
16426        let agent = AgentBuilder::new()
16427            .system_prompt("You are helpful.")
16428            .llm(Arc::new(mock))
16429            .tools(tools)
16430            .tool_security(ToolSecurityEngine::new(security))
16431            .build()
16432            .unwrap();
16433
16434        let record = agent
16435            .invoke_tool(ToolExecutionRequest::new(
16436                "ctx-call",
16437                "Context Echo",
16438                serde_json::json!({"path": ".", "max_results": 99}),
16439                ToolCallSource::Manual,
16440            ))
16441            .await
16442            .unwrap();
16443
16444        assert!(record.success);
16445        assert!(matches!(&record.source, ToolCallSource::Manual));
16446        assert_eq!(record.requested_name, "Context Echo");
16447        assert_eq!(record.canonical_id, "context_echo");
16448        assert_eq!(record.policy.outcome, PermissionOutcome::Allow);
16449        assert_eq!(record.executed_arguments["max_results"], 7);
16450        let output: Value = serde_json::from_str(&record.output).unwrap();
16451        assert_eq!(output["requested_name"], "Context Echo");
16452        assert_eq!(output["canonical_id"], "context_echo");
16453        assert_eq!(output["max_results"], 7);
16454        assert_eq!(output["custom_config"]["backend"], "memory");
16455        assert!(record.metadata.contains_key("effective_limits"));
16456        assert!(record.metadata.contains_key("policy_snapshot"));
16457    }
16458
16459    #[tokio::test]
16460    async fn test_runtime_control_cancels_active_tool_call() {
16461        let mock = mock_with_response("hello");
16462        let agent = Arc::new(
16463            AgentBuilder::new()
16464                .system_prompt("You are helpful.")
16465                .llm(Arc::new(mock))
16466                .tool(Arc::new(SlowTool))
16467                .build()
16468                .unwrap(),
16469        );
16470        let control = agent.runtime_control();
16471        let running_agent = Arc::clone(&agent);
16472        let handle = tokio::spawn(async move {
16473            running_agent
16474                .invoke_tool(ToolExecutionRequest::new(
16475                    "slow-call",
16476                    "slow",
16477                    serde_json::json!({}),
16478                    ToolCallSource::Manual,
16479                ))
16480                .await
16481                .unwrap()
16482        });
16483
16484        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
16485        control.cancel_all();
16486        let record = handle.await.unwrap();
16487
16488        assert!(record.executed);
16489        assert!(record.cancelled);
16490        assert!(!record.success);
16491        assert!(record.cancellation_reason.is_some());
16492    }
16493
16494    #[tokio::test]
16495    async fn non_idempotent_tool_calls_are_not_retried() {
16496        use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
16497
16498        let mock = mock_with_response("hello");
16499        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16500        let agent = AgentBuilder::new()
16501            .system_prompt("You are helpful.")
16502            .llm(Arc::new(mock))
16503            .tool(Arc::new(FlakyWriteTool {
16504                calls: Arc::clone(&calls),
16505            }))
16506            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
16507                tools: ToolRecoveryConfig {
16508                    default: ToolRetryConfig {
16509                        max_retries: 2,
16510                        ..Default::default()
16511                    },
16512                    ..Default::default()
16513                },
16514                ..Default::default()
16515            }))
16516            .build()
16517            .unwrap();
16518
16519        let record = agent
16520            .invoke_tool(ToolExecutionRequest::new(
16521                "flaky-call",
16522                "flaky_write",
16523                serde_json::json!({"path": "./tmp.txt"}),
16524                ToolCallSource::Manual,
16525            ))
16526            .await
16527            .unwrap();
16528
16529        assert!(!record.success);
16530        assert_eq!(calls.load(Ordering::SeqCst), 1);
16531    }
16532
16533    #[tokio::test]
16534    async fn safely_retryable_tool_receives_a_fresh_deadline_per_attempt() {
16535        use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
16536
16537        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16538        let deadlines = Arc::new(parking_lot::Mutex::new(Vec::new()));
16539        let agent = AgentBuilder::new()
16540            .system_prompt("Test retry deadlines.")
16541            .llm(Arc::new(mock_with_response("done")))
16542            .tool(Arc::new(RetryDeadlineTool {
16543                calls: Arc::clone(&calls),
16544                deadlines: Arc::clone(&deadlines),
16545            }))
16546            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
16547                tools: ToolRecoveryConfig {
16548                    per_tool: HashMap::from([(
16549                        "retry_deadline".to_string(),
16550                        ToolRetryConfig {
16551                            max_retries: 1,
16552                            ..Default::default()
16553                        },
16554                    )]),
16555                    ..Default::default()
16556                },
16557                ..Default::default()
16558            }))
16559            .build()
16560            .unwrap();
16561
16562        let record = agent
16563            .invoke_tool(ToolExecutionRequest::new(
16564                "retry-deadline-call",
16565                "retry_deadline",
16566                serde_json::json!({}),
16567                ToolCallSource::Manual,
16568            ))
16569            .await
16570            .unwrap();
16571
16572        assert!(record.executed);
16573        assert!(record.success);
16574        assert_eq!(calls.load(Ordering::SeqCst), 2);
16575        let deadlines = deadlines.lock();
16576        assert_eq!(deadlines.len(), 2);
16577        assert!(
16578            deadlines[1] > deadlines[0],
16579            "retry inherited the first invocation deadline"
16580        );
16581    }
16582
16583    #[tokio::test]
16584    async fn side_effecting_tools_are_serialized_per_resource() {
16585        let mock = mock_with_response("hello");
16586        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16587        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16588        let agent = Arc::new(
16589            AgentBuilder::new()
16590                .system_prompt("You are helpful.")
16591                .llm(Arc::new(mock))
16592                .tool(Arc::new(LockedWriteTool {
16593                    active: Arc::clone(&active),
16594                    max_active: Arc::clone(&max_active),
16595                }))
16596                .build()
16597                .unwrap(),
16598        );
16599
16600        let left = {
16601            let agent = Arc::clone(&agent);
16602            tokio::spawn(async move {
16603                agent
16604                    .invoke_tool(ToolExecutionRequest::new(
16605                        "lock-1",
16606                        "locked_write",
16607                        serde_json::json!({"path": "./same.txt"}),
16608                        ToolCallSource::Manual,
16609                    ))
16610                    .await
16611                    .unwrap()
16612            })
16613        };
16614        let right = {
16615            let agent = Arc::clone(&agent);
16616            tokio::spawn(async move {
16617                agent
16618                    .invoke_tool(ToolExecutionRequest::new(
16619                        "lock-2",
16620                        "locked_write",
16621                        serde_json::json!({"path": "./same.txt"}),
16622                        ToolCallSource::Manual,
16623                    ))
16624                    .await
16625                    .unwrap()
16626            })
16627        };
16628
16629        let left = left.await.unwrap();
16630        let right = right.await.unwrap();
16631        assert!(left.success);
16632        assert!(right.success);
16633        assert_eq!(max_active.load(Ordering::SeqCst), 1);
16634    }
16635
16636    #[tokio::test]
16637    async fn path_resources_use_shared_global_lock_and_cleanup() {
16638        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16639        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16640        let bindings = ai_agents_core::ToolPolicyBindings {
16641            path_fields: vec![
16642                ai_agents_core::PathPolicyBinding::read_write("source_path"),
16643                ai_agents_core::PathPolicyBinding::write("destination_path"),
16644            ],
16645            ..Default::default()
16646        };
16647        let classification = ai_agents_core::ToolCallClassification::from_metadata(
16648            &MultiResourceWriteTool {
16649                active: Arc::clone(&active),
16650                max_active: Arc::clone(&max_active),
16651            }
16652            .safety_metadata(),
16653        );
16654        let left_args = serde_json::json!({
16655            "source_path": "./a/../first.txt",
16656            "destination_path": "./second.txt"
16657        });
16658        let right_args = serde_json::json!({
16659            "source_path": "./second.txt",
16660            "destination_path": "./first.txt"
16661        });
16662        let left_keys = tool_resource_lock_keys(
16663            "multi_resource_write",
16664            &left_args,
16665            &bindings,
16666            &classification,
16667        );
16668        let right_keys = tool_resource_lock_keys(
16669            "multi_resource_write",
16670            &right_args,
16671            &bindings,
16672            &classification,
16673        );
16674        assert_eq!(left_keys, right_keys);
16675        assert_eq!(left_keys, vec!["path-mutation:global".to_string()]);
16676
16677        let locks = new_tool_resource_locks();
16678        let build_agent = || {
16679            AgentBuilder::new()
16680                .system_prompt("Test shared resource locks.")
16681                .llm(Arc::new(mock_with_response("done")))
16682                .tool(Arc::new(MultiResourceWriteTool {
16683                    active: Arc::clone(&active),
16684                    max_active: Arc::clone(&max_active),
16685                }))
16686                .build()
16687                .unwrap()
16688                .with_shared_resource_locks(Arc::clone(&locks))
16689        };
16690        let left_agent = Arc::new(build_agent());
16691        let right_agent = Arc::new(build_agent());
16692        let left = tokio::spawn(async move {
16693            left_agent
16694                .invoke_tool(ToolExecutionRequest::new(
16695                    "multi-left",
16696                    "multi_resource_write",
16697                    left_args,
16698                    ToolCallSource::Manual,
16699                ))
16700                .await
16701                .unwrap()
16702        });
16703        let right = tokio::spawn(async move {
16704            right_agent
16705                .invoke_tool(ToolExecutionRequest::new(
16706                    "multi-right",
16707                    "multi_resource_write",
16708                    right_args,
16709                    ToolCallSource::Manual,
16710                ))
16711                .await
16712                .unwrap()
16713        });
16714        let (left, right) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
16715            tokio::join!(left, right)
16716        })
16717        .await
16718        .expect("reversed resource acquisition must not deadlock");
16719
16720        assert!(left.unwrap().success);
16721        assert!(right.unwrap().success);
16722        assert_eq!(max_active.load(Ordering::SeqCst), 1);
16723        assert!(locks.read().is_empty());
16724    }
16725
16726    #[tokio::test]
16727    async fn global_path_lock_serializes_copy_destination_with_file_write() {
16728        assert_path_mutation_pair_serialized(
16729            "copy_path",
16730            CopyPathTool::new().policy_bindings().path_fields,
16731            serde_json::json!({
16732                "source_path": "./source.txt",
16733                "destination_path": "./shared.txt"
16734            }),
16735            "file_write",
16736            FileWriteTool::new().policy_bindings().path_fields,
16737            serde_json::json!({"path": "./shared.txt"}),
16738        )
16739        .await;
16740    }
16741
16742    #[tokio::test]
16743    async fn parent_and_spawned_runtime_share_global_path_lock() {
16744        let workspace = MutationTestWorkspace::new();
16745        let destination = workspace.root.join("spawned.txt");
16746        let parent_gate = PathMutationGate::new();
16747        let parent = Arc::new(
16748            AgentBuilder::from_yaml(
16749                r#"
16750name: LockParent
16751system_prompt: parent
16752llm:
16753  default: default
16754tools:
16755  - parent_path_write
16756spawner:
16757  shared_llms: true
16758"#,
16759            )
16760            .unwrap()
16761            .llm(Arc::new(mock_with_response("done")))
16762            .auto_configure_spawner()
16763            .await
16764            .unwrap()
16765            .tool(Arc::new(BlockingPathMutationTool {
16766                id: "parent_path_write",
16767                path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
16768                gate: parent_gate.clone(),
16769            }))
16770            .build()
16771            .unwrap(),
16772        );
16773
16774        let mut child_spec = crate::spec::AgentSpec {
16775            name: "LockChild".to_string(),
16776            system_prompt: "child".to_string(),
16777            tools: Some(vec![crate::spec::ToolEntry::Simple(
16778                "file_write".to_string(),
16779            )]),
16780            ..Default::default()
16781        };
16782        child_spec.tool_security.enabled = true;
16783        child_spec.tool_security.fail_closed = true;
16784        let file_write_policy = ai_agents_tools::ToolPolicyConfig {
16785            write_paths: vec![workspace.root.to_string_lossy().into_owned()],
16786            allow_without_confirmation: true,
16787            ..Default::default()
16788        };
16789        child_spec
16790            .tool_security
16791            .tools
16792            .insert("file_write".to_string(), file_write_policy);
16793        let spawned = parent
16794            .spawner()
16795            .unwrap()
16796            .spawn_from_spec(child_spec)
16797            .await
16798            .unwrap();
16799        assert!(Arc::ptr_eq(
16800            &parent.resource_locks,
16801            &spawned.agent.resource_locks
16802        ));
16803        assert!(!Arc::ptr_eq(
16804            &parent.runtime_control,
16805            &spawned.agent.runtime_control
16806        ));
16807
16808        let parent_call = {
16809            let parent = Arc::clone(&parent);
16810            let destination = destination.clone();
16811            tokio::spawn(async move {
16812                parent
16813                    .invoke_tool(ToolExecutionRequest::new(
16814                        "parent-lock-holder",
16815                        "parent_path_write",
16816                        serde_json::json!({"path": destination}),
16817                        ToolCallSource::Manual,
16818                    ))
16819                    .await
16820                    .unwrap()
16821            })
16822        };
16823        parent_gate.wait_until_entered().await;
16824
16825        let child_call = {
16826            let child = Arc::clone(&spawned.agent);
16827            let destination = destination.clone();
16828            tokio::spawn(async move {
16829                child
16830                    .invoke_tool(ToolExecutionRequest::new(
16831                        "spawned-file-write",
16832                        "file_write",
16833                        serde_json::json!({
16834                            "path": destination,
16835                            "content": "spawned",
16836                            "dry_run": false
16837                        }),
16838                        ToolCallSource::Manual,
16839                    ))
16840                    .await
16841                    .unwrap()
16842            })
16843        };
16844        wait_for_resource_lock_strong_count(&parent.resource_locks, 2).await;
16845        assert!(!child_call.is_finished());
16846
16847        parent_gate.release();
16848        let (parent_record, child_record) =
16849            tokio::time::timeout(std::time::Duration::from_secs(2), async {
16850                tokio::join!(parent_call, child_call)
16851            })
16852            .await
16853            .expect("parent and spawned path mutations did not finish");
16854        assert!(parent_record.unwrap().success);
16855        assert!(child_record.unwrap().success);
16856        assert_eq!(std::fs::read_to_string(destination).unwrap(), "spawned");
16857        assert!(parent.resource_locks.read().is_empty());
16858    }
16859
16860    #[tokio::test]
16861    async fn cancelled_global_path_lock_waiter_does_not_retain_weak_entry() {
16862        let locks = new_tool_resource_locks();
16863        let holder_gate = PathMutationGate::new();
16864        let waiter_gate = PathMutationGate::new();
16865        waiter_gate.release();
16866        let holder = Arc::new(
16867            AgentBuilder::new()
16868                .system_prompt("Hold the global path lock.")
16869                .llm(Arc::new(mock_with_response("done")))
16870                .tool(Arc::new(BlockingPathMutationTool {
16871                    id: "holder_write",
16872                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
16873                    gate: holder_gate.clone(),
16874                }))
16875                .build()
16876                .unwrap()
16877                .with_shared_resource_locks(Arc::clone(&locks)),
16878        );
16879        let waiter = Arc::new(
16880            AgentBuilder::new()
16881                .system_prompt("Wait for the global path lock.")
16882                .llm(Arc::new(mock_with_response("done")))
16883                .tool(Arc::new(BlockingPathMutationTool {
16884                    id: "waiter_write",
16885                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
16886                    gate: waiter_gate.clone(),
16887                }))
16888                .build()
16889                .unwrap()
16890                .with_shared_resource_locks(Arc::clone(&locks)),
16891        );
16892
16893        let holder_call = {
16894            let holder = Arc::clone(&holder);
16895            tokio::spawn(async move {
16896                holder
16897                    .invoke_tool(ToolExecutionRequest::new(
16898                        "holder-call",
16899                        "holder_write",
16900                        serde_json::json!({"path": "./shared.txt"}),
16901                        ToolCallSource::Manual,
16902                    ))
16903                    .await
16904                    .unwrap()
16905            })
16906        };
16907        holder_gate.wait_until_entered().await;
16908
16909        let waiter_call = {
16910            let waiter = Arc::clone(&waiter);
16911            tokio::spawn(async move {
16912                waiter
16913                    .invoke_tool(ToolExecutionRequest::new(
16914                        "waiter-call",
16915                        "waiter_write",
16916                        serde_json::json!({"path": "./shared.txt"}),
16917                        ToolCallSource::Manual,
16918                    ))
16919                    .await
16920                    .unwrap()
16921            })
16922        };
16923        wait_for_resource_lock_strong_count(&locks, 2).await;
16924        waiter.runtime_control().cancel_all();
16925
16926        let waiter_record = tokio::time::timeout(std::time::Duration::from_secs(2), waiter_call)
16927            .await
16928            .expect("cancelled lock waiter did not finish")
16929            .unwrap();
16930        assert!(!waiter_record.executed);
16931        assert!(!waiter_gate.entered.load(Ordering::SeqCst));
16932        assert_eq!(
16933            locks
16934                .read()
16935                .get("path-mutation:global")
16936                .map_or(0, |lock| lock.strong_count()),
16937            1
16938        );
16939
16940        holder_gate.release();
16941        let holder_record = tokio::time::timeout(std::time::Duration::from_secs(2), holder_call)
16942            .await
16943            .expect("lock holder did not finish")
16944            .unwrap();
16945        assert!(holder_record.success);
16946        assert!(locks.read().is_empty());
16947    }
16948
16949    #[tokio::test]
16950    async fn path_mutation_policy_and_approval_denials_do_not_invoke_tools() {
16951        for denial in [MutationDenial::Policy, MutationDenial::Approval] {
16952            let tools: [Arc<dyn Tool>; 3] = [
16953                Arc::new(CopyPathTool::new()),
16954                Arc::new(MovePathTool::new()),
16955                Arc::new(DeletePathTool::new()),
16956            ];
16957            for tool in tools {
16958                assert_path_mutation_denied(tool, denial).await;
16959            }
16960        }
16961    }
16962
16963    #[tokio::test]
16964    async fn policy_denial_keeps_executor_hook_lifecycle_and_record_authority() {
16965        let workspace = MutationTestWorkspace::new();
16966        let target = workspace.root.join("denied.txt");
16967        let hooks = Arc::new(ToolLifecycleRecordingHooks::new());
16968        let agent = AgentBuilder::new()
16969            .system_prompt("Test denied tool hooks.")
16970            .llm(Arc::new(mock_with_response("done")))
16971            .tool(Arc::new(FileWriteTool::new()))
16972            .tool_security(ToolSecurityEngine::new(mutation_denial_security_config(
16973                "file_write",
16974                &workspace.root,
16975                MutationDenial::Policy,
16976            )))
16977            .hooks(hooks.clone())
16978            .build()
16979            .unwrap();
16980
16981        let record = agent
16982            .invoke_tool(ToolExecutionRequest::new(
16983                "denied-hook-call",
16984                "file_write",
16985                serde_json::json!({
16986                    "path": target.to_string_lossy(),
16987                    "content": "blocked"
16988                }),
16989                ToolCallSource::Manual,
16990            ))
16991            .await
16992            .unwrap();
16993
16994        assert!(!record.executed);
16995        assert!(!record.success);
16996        assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
16997        assert_eq!(
16998            hooks.events(),
16999            vec![
17000                "start:file_write",
17001                "complete:file_write:false",
17002                "record:file_write:false",
17003                "error"
17004            ]
17005        );
17006        assert!(!target.exists());
17007    }
17008
17009    #[tokio::test]
17010    async fn approval_argument_changes_are_rechecked_against_final_scope() {
17011        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17012        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17013        let entered = Arc::new(tokio::sync::Barrier::new(2));
17014        let release = Arc::new(tokio::sync::Notify::new());
17015        let handler = Arc::new(BlockingApprovalHandler {
17016            entered: Arc::clone(&entered),
17017            release: Arc::clone(&release),
17018            result: ApprovalResult::Modified {
17019                changes: HashMap::from([(
17020                    "path".to_string(),
17021                    Value::String("./after-approval.txt".to_string()),
17022                )]),
17023            },
17024        });
17025        let agent = Arc::new(
17026            AgentBuilder::new()
17027                .system_prompt("Test final scope validation.")
17028                .llm(Arc::new(mock_with_response("done")))
17029                .tool(Arc::new(LockedWriteTool {
17030                    active: Arc::clone(&active),
17031                    max_active: Arc::clone(&max_active),
17032                }))
17033                .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
17034                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17035                .approval_handler(handler)
17036                .build()
17037                .unwrap(),
17038        );
17039        let control = agent.runtime_control();
17040        let running = Arc::clone(&agent);
17041        let call = tokio::spawn(async move {
17042            running
17043                .invoke_tool(ToolExecutionRequest::new(
17044                    "approval-scope",
17045                    "locked_write",
17046                    serde_json::json!({"path": "./before-approval.txt"}),
17047                    ToolCallSource::Manual,
17048                ))
17049                .await
17050                .unwrap()
17051        });
17052        entered.wait().await;
17053        let expected_version = control.set_tool_scope(Vec::new());
17054        release.notify_one();
17055        let record = call.await.unwrap();
17056
17057        assert!(!record.executed);
17058        assert!(!record.success);
17059        assert_eq!(record.runtime_config_version, expected_version);
17060        assert_eq!(record.executed_arguments["path"], "./after-approval.txt");
17061        assert_eq!(max_active.load(Ordering::SeqCst), 0);
17062        assert_eq!(
17063            record.metadata["runtime_scope_snapshot"],
17064            serde_json::json!([])
17065        );
17066    }
17067
17068    #[tokio::test]
17069    async fn approval_is_rechecked_against_final_policy_snapshot() {
17070        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17071        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17072        let entered = Arc::new(tokio::sync::Barrier::new(2));
17073        let release = Arc::new(tokio::sync::Notify::new());
17074        let handler = Arc::new(BlockingApprovalHandler {
17075            entered: Arc::clone(&entered),
17076            release: Arc::clone(&release),
17077            result: ApprovalResult::Approved,
17078        });
17079        let agent = Arc::new(
17080            AgentBuilder::new()
17081                .system_prompt("Test final policy validation.")
17082                .llm(Arc::new(mock_with_response("done")))
17083                .tool(Arc::new(LockedWriteTool {
17084                    active: Arc::clone(&active),
17085                    max_active: Arc::clone(&max_active),
17086                }))
17087                .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
17088                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17089                .approval_handler(handler)
17090                .build()
17091                .unwrap(),
17092        );
17093        let control = agent.runtime_control();
17094        let running = Arc::clone(&agent);
17095        let call = tokio::spawn(async move {
17096            running
17097                .invoke_tool(ToolExecutionRequest::new(
17098                    "approval-policy",
17099                    "locked_write",
17100                    serde_json::json!({"path": "./policy.txt"}),
17101                    ToolCallSource::Manual,
17102                ))
17103                .await
17104                .unwrap()
17105        });
17106        entered.wait().await;
17107        let expected_version = control.set_tool_security(approval_security_config(false));
17108        release.notify_one();
17109        let record = call.await.unwrap();
17110
17111        assert!(!record.executed);
17112        assert!(!record.success);
17113        assert_eq!(record.runtime_config_version, expected_version);
17114        assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
17115        assert_eq!(max_active.load(Ordering::SeqCst), 0);
17116        assert!(record.metadata.contains_key("policy_snapshot"));
17117    }
17118
17119    #[test]
17120    fn invalid_live_policy_does_not_replace_snapshot_or_generation() {
17121        let agent = AgentBuilder::new()
17122            .system_prompt("Test runtime policy validation.")
17123            .llm(Arc::new(mock_with_response("done")))
17124            .build()
17125            .unwrap();
17126        let control = agent.runtime_control();
17127        let mut valid = ToolSecurityConfig::default();
17128        valid.tools.insert(
17129            "web_search".to_string(),
17130            ai_agents_tools::ToolPolicyConfig {
17131                max_results: Some(5),
17132                ..Default::default()
17133            },
17134        );
17135        let generation = control.try_set_tool_security(valid).unwrap();
17136
17137        let mut invalid = ToolSecurityConfig::default();
17138        invalid.tools.insert(
17139            "web_search".to_string(),
17140            ai_agents_tools::ToolPolicyConfig {
17141                max_results: Some(0),
17142                ..Default::default()
17143            },
17144        );
17145        let error = control.try_set_tool_security(invalid).unwrap_err();
17146
17147        assert!(
17148            error
17149                .to_string()
17150                .contains("max_results must be greater than 0")
17151        );
17152        assert_eq!(control.version(), generation);
17153        assert_eq!(
17154            control
17155                .state
17156                .tool_security_override
17157                .read()
17158                .as_ref()
17159                .unwrap()
17160                .config()
17161                .tools["web_search"]
17162                .max_results,
17163            Some(5)
17164        );
17165    }
17166
17167    #[tokio::test]
17168    async fn persistent_override_preserves_rate_history_within_generation() {
17169        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17170        let agent = AgentBuilder::new()
17171            .system_prompt("Test persistent policy overrides.")
17172            .llm(Arc::new(mock_with_response("done")))
17173            .tool(Arc::new(RecoveryTestTool {
17174                id: "limited_override".to_string(),
17175                succeeds: true,
17176                calls: Arc::clone(&calls),
17177                max_output_chars: None,
17178            }))
17179            .build()
17180            .unwrap();
17181        let mut security = ToolSecurityConfig {
17182            enabled: true,
17183            fail_closed: true,
17184            ..Default::default()
17185        };
17186        let policy = ai_agents_tools::ToolPolicyConfig {
17187            write_paths: vec![".".to_string()],
17188            rate_limit: Some(1),
17189            ..Default::default()
17190        };
17191        security
17192            .tools
17193            .insert("limited_override".to_string(), policy);
17194        let generation = agent.runtime_control().set_tool_security(security);
17195
17196        let first = agent
17197            .invoke_tool(ToolExecutionRequest::new(
17198                "limited-first",
17199                "limited_override",
17200                serde_json::json!({"path": "./limited.txt"}),
17201                ToolCallSource::Manual,
17202            ))
17203            .await
17204            .unwrap();
17205        let second = agent
17206            .invoke_tool(ToolExecutionRequest::new(
17207                "limited-second",
17208                "limited_override",
17209                serde_json::json!({"path": "./limited.txt"}),
17210                ToolCallSource::Manual,
17211            ))
17212            .await
17213            .unwrap();
17214
17215        assert!(first.success);
17216        assert_eq!(first.policy_version, generation);
17217        assert!(!second.executed);
17218        assert!(second.output.contains("Rate limit exceeded"));
17219        assert_eq!(second.policy_version, generation);
17220        assert_eq!(calls.load(Ordering::SeqCst), 1);
17221    }
17222
17223    #[tokio::test]
17224    async fn concurrent_rate_admission_consumes_capacity_atomically() {
17225        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17226        let tool = Arc::new(RecoveryTestTool {
17227            id: "atomic_rate".to_string(),
17228            succeeds: true,
17229            calls: Arc::clone(&calls),
17230            max_output_chars: None,
17231        });
17232        let arguments = serde_json::json!({"path": "./atomic-rate.txt"});
17233        let bindings = tool.policy_bindings();
17234        let classification = tool.classify_call(&arguments);
17235        let resource_keys =
17236            tool_resource_lock_keys(tool.id(), &arguments, &bindings, &classification);
17237        let mut security = ToolSecurityConfig {
17238            enabled: true,
17239            fail_closed: true,
17240            ..Default::default()
17241        };
17242        let policy = ai_agents_tools::ToolPolicyConfig {
17243            write_paths: vec![".".to_string()],
17244            rate_limit: Some(1),
17245            ..Default::default()
17246        };
17247        security.tools.insert(tool.id().to_string(), policy);
17248        let agent = Arc::new(
17249            AgentBuilder::new()
17250                .system_prompt("Test atomic rate admission.")
17251                .llm(Arc::new(mock_with_response("done")))
17252                .tool(tool)
17253                .tool_security(ToolSecurityEngine::new(security))
17254                .build()
17255                .unwrap(),
17256        );
17257        let held = agent
17258            .acquire_tool_resource_locks(&resource_keys)
17259            .await
17260            .unwrap();
17261        let left = {
17262            let agent = Arc::clone(&agent);
17263            let arguments = arguments.clone();
17264            tokio::spawn(async move {
17265                agent
17266                    .invoke_tool(ToolExecutionRequest::new(
17267                        "atomic-rate-left",
17268                        "atomic_rate",
17269                        arguments,
17270                        ToolCallSource::Manual,
17271                    ))
17272                    .await
17273                    .unwrap()
17274            })
17275        };
17276        let right = {
17277            let agent = Arc::clone(&agent);
17278            tokio::spawn(async move {
17279                agent
17280                    .invoke_tool(ToolExecutionRequest::new(
17281                        "atomic-rate-right",
17282                        "atomic_rate",
17283                        arguments,
17284                        ToolCallSource::Manual,
17285                    ))
17286                    .await
17287                    .unwrap()
17288            })
17289        };
17290        tokio::time::sleep(std::time::Duration::from_millis(25)).await;
17291        drop(held);
17292        let (left, right) = tokio::join!(left, right);
17293        let records = [left.unwrap(), right.unwrap()];
17294
17295        assert_eq!(records.iter().filter(|record| record.success).count(), 1);
17296        assert_eq!(records.iter().filter(|record| record.executed).count(), 1);
17297        assert!(
17298            records.iter().any(|record| {
17299                !record.executed && record.output.contains("Rate limit exceeded")
17300            })
17301        );
17302        assert_eq!(calls.load(Ordering::SeqCst), 1);
17303    }
17304
17305    #[tokio::test]
17306    async fn changed_policy_generation_invalidates_pending_approval() {
17307        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17308        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17309        let entered = Arc::new(tokio::sync::Barrier::new(2));
17310        let release = Arc::new(tokio::sync::Notify::new());
17311        let handler = Arc::new(BlockingApprovalHandler {
17312            entered: Arc::clone(&entered),
17313            release: Arc::clone(&release),
17314            result: ApprovalResult::Approved,
17315        });
17316        let agent = Arc::new(
17317            AgentBuilder::new()
17318                .system_prompt("Test stale approval denial.")
17319                .llm(Arc::new(mock_with_response("done")))
17320                .tool(Arc::new(LockedWriteTool {
17321                    active: Arc::clone(&active),
17322                    max_active: Arc::clone(&max_active),
17323                }))
17324                .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
17325                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17326                .approval_handler(handler)
17327                .build()
17328                .unwrap(),
17329        );
17330        let running = Arc::clone(&agent);
17331        let call = tokio::spawn(async move {
17332            running
17333                .invoke_tool(ToolExecutionRequest::new(
17334                    "stale-approval",
17335                    "locked_write",
17336                    serde_json::json!({"path": "./stale.txt"}),
17337                    ToolCallSource::Manual,
17338                ))
17339                .await
17340                .unwrap()
17341        });
17342        entered.wait().await;
17343        let generation = agent
17344            .runtime_control()
17345            .set_tool_security(approval_security_config(true));
17346        release.notify_one();
17347        let record = call.await.unwrap();
17348
17349        assert!(!record.executed);
17350        assert!(record.output.contains("Approval became stale"));
17351        assert_eq!(record.policy_version, generation);
17352        assert_eq!(max_active.load(Ordering::SeqCst), 0);
17353    }
17354
17355    #[tokio::test]
17356    async fn final_policy_reapplies_argument_caps_after_approval_changes() {
17357        use ai_agents_hitl::CallbackHandler;
17358
17359        let mut security = ToolSecurityConfig {
17360            enabled: true,
17361            fail_closed: true,
17362            ..Default::default()
17363        };
17364        let policy = ai_agents_tools::ToolPolicyConfig {
17365            read_paths: vec![".".to_string()],
17366            max_results: Some(5),
17367            require_confirmation: true,
17368            ..Default::default()
17369        };
17370        security.tools.insert("context_echo".to_string(), policy);
17371        let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
17372            changes: HashMap::from([("max_results".to_string(), serde_json::json!(99))]),
17373        });
17374        let agent = AgentBuilder::new()
17375            .system_prompt("Test final argument caps.")
17376            .llm(Arc::new(mock_with_response("done")))
17377            .tool(Arc::new(ContextEchoTool))
17378            .tool_security(ToolSecurityEngine::new(security))
17379            .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17380            .approval_handler(Arc::new(handler))
17381            .build()
17382            .unwrap();
17383
17384        let record = agent
17385            .invoke_tool(ToolExecutionRequest::new(
17386                "final-cap",
17387                "context_echo",
17388                serde_json::json!({"path": ".", "max_results": 1}),
17389                ToolCallSource::Manual,
17390            ))
17391            .await
17392            .unwrap();
17393
17394        assert!(record.success);
17395        assert_eq!(record.executed_arguments["max_results"], 5);
17396        assert_eq!(
17397            record.approval.unwrap().modified_arguments.unwrap()["max_results"],
17398            5
17399        );
17400    }
17401
17402    #[tokio::test]
17403    async fn no_binding_writes_use_canonical_fallback_lock() {
17404        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17405        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17406        let agent = Arc::new(
17407            AgentBuilder::new()
17408                .system_prompt("Test fallback resource locks.")
17409                .llm(Arc::new(mock_with_response("done")))
17410                .tool(Arc::new(NoBindingWriteTool {
17411                    active: Arc::clone(&active),
17412                    max_active: Arc::clone(&max_active),
17413                }))
17414                .build()
17415                .unwrap(),
17416        );
17417        let left = {
17418            let agent = Arc::clone(&agent);
17419            tokio::spawn(async move {
17420                agent
17421                    .invoke_tool(ToolExecutionRequest::new(
17422                        "no-binding-left",
17423                        "no_binding_write",
17424                        serde_json::json!({}),
17425                        ToolCallSource::Manual,
17426                    ))
17427                    .await
17428                    .unwrap()
17429            })
17430        };
17431        let right = {
17432            let agent = Arc::clone(&agent);
17433            tokio::spawn(async move {
17434                agent
17435                    .invoke_tool(ToolExecutionRequest::new(
17436                        "no-binding-right",
17437                        "no_binding_write",
17438                        serde_json::json!({}),
17439                        ToolCallSource::Manual,
17440                    ))
17441                    .await
17442                    .unwrap()
17443            })
17444        };
17445        let (left, right) = tokio::join!(left, right);
17446
17447        assert!(left.unwrap().success);
17448        assert!(right.unwrap().success);
17449        assert_eq!(max_active.load(Ordering::SeqCst), 1);
17450    }
17451
17452    #[tokio::test]
17453    async fn parent_and_child_paths_share_a_resource_lock() {
17454        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17455        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17456        let agent = Arc::new(
17457            AgentBuilder::new()
17458                .system_prompt("Test parent child resource locks.")
17459                .llm(Arc::new(mock_with_response("done")))
17460                .tool(Arc::new(LockedWriteTool {
17461                    active: Arc::clone(&active),
17462                    max_active: Arc::clone(&max_active),
17463                }))
17464                .build()
17465                .unwrap(),
17466        );
17467        let parent = format!("./lock-parent-{}", uuid::Uuid::new_v4());
17468        let child = format!("{}/child.txt", parent);
17469        let left = {
17470            let agent = Arc::clone(&agent);
17471            tokio::spawn(async move {
17472                agent
17473                    .invoke_tool(ToolExecutionRequest::new(
17474                        "parent-lock",
17475                        "locked_write",
17476                        serde_json::json!({"path": parent}),
17477                        ToolCallSource::Manual,
17478                    ))
17479                    .await
17480                    .unwrap()
17481            })
17482        };
17483        let right = {
17484            let agent = Arc::clone(&agent);
17485            tokio::spawn(async move {
17486                agent
17487                    .invoke_tool(ToolExecutionRequest::new(
17488                        "child-lock",
17489                        "locked_write",
17490                        serde_json::json!({"path": child}),
17491                        ToolCallSource::Manual,
17492                    ))
17493                    .await
17494                    .unwrap()
17495            })
17496        };
17497        let (left, right) = tokio::join!(left, right);
17498
17499        assert!(left.unwrap().success);
17500        assert!(right.unwrap().success);
17501        assert_eq!(max_active.load(Ordering::SeqCst), 1);
17502    }
17503
17504    #[tokio::test]
17505    async fn tool_hooks_can_reenter_after_resource_guards_are_dropped() {
17506        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17507        let hooks = Arc::new(ReentrantToolHooks {
17508            agent: parking_lot::Mutex::new(None),
17509            invoked: AtomicBool::new(false),
17510            nested_success: AtomicBool::new(false),
17511        });
17512        let agent = Arc::new(
17513            AgentBuilder::new()
17514                .system_prompt("Test hook reentrancy.")
17515                .llm(Arc::new(mock_with_response("done")))
17516                .tool(Arc::new(RecoveryTestTool {
17517                    id: "reentrant_write".to_string(),
17518                    succeeds: true,
17519                    calls: Arc::clone(&calls),
17520                    max_output_chars: None,
17521                }))
17522                .hooks(hooks.clone())
17523                .build()
17524                .unwrap(),
17525        );
17526        *hooks.agent.lock() = Some(Arc::downgrade(&agent));
17527        let record = tokio::time::timeout(
17528            std::time::Duration::from_secs(2),
17529            agent.invoke_tool(ToolExecutionRequest::new(
17530                "outer-hook-call",
17531                "reentrant_write",
17532                serde_json::json!({"path": "./hook.txt"}),
17533                ToolCallSource::Manual,
17534            )),
17535        )
17536        .await
17537        .expect("tool completion hook must not retain resource guards")
17538        .unwrap();
17539
17540        assert!(record.success);
17541        assert!(hooks.nested_success.load(Ordering::SeqCst));
17542        assert_eq!(calls.load(Ordering::SeqCst), 2);
17543    }
17544
17545    /// Confirms fallback starts only after the failed original lifecycle and resource ownership complete.
17546    #[tokio::test]
17547    async fn fallback_finalizes_original_record_before_shared_execution() {
17548        let primary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17549        let fallback_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17550        let hooks = Arc::new(ToolLifecycleRecordingHooks::new());
17551        let agent = AgentBuilder::new()
17552            .system_prompt("Test fallback execution.")
17553            .llm(Arc::new(mock_with_response("done")))
17554            .tool(Arc::new(RecoveryTestTool {
17555                id: "primary".to_string(),
17556                succeeds: false,
17557                calls: Arc::clone(&primary_calls),
17558                max_output_chars: None,
17559            }))
17560            .tool(Arc::new(RecoveryTestTool {
17561                id: "fallback".to_string(),
17562                succeeds: true,
17563                calls: Arc::clone(&fallback_calls),
17564                max_output_chars: None,
17565            }))
17566            .recovery_manager(recovery_manager_with_fallbacks([(
17567                "primary".to_string(),
17568                "fallback".to_string(),
17569            )]))
17570            .hooks(hooks.clone())
17571            .build()
17572            .unwrap();
17573        let record = tokio::time::timeout(
17574            std::time::Duration::from_secs(2),
17575            agent.invoke_tool(ToolExecutionRequest::new(
17576                "fallback-call",
17577                "primary",
17578                serde_json::json!({"path": "./shared.txt"}),
17579                ToolCallSource::Manual,
17580            )),
17581        )
17582        .await
17583        .expect("fallback must not retain the primary resource guard")
17584        .unwrap();
17585
17586        assert_eq!(
17587            hooks.events(),
17588            vec![
17589                "start:primary",
17590                "complete:primary:false",
17591                "record:primary:true",
17592                "error",
17593                "start:fallback",
17594                "complete:fallback:true",
17595                "record:fallback:true",
17596            ]
17597        );
17598        let records = hooks.records();
17599        assert_eq!(records.len(), 2);
17600        let original = &records[0];
17601        assert_eq!(original.canonical_id, "primary");
17602        assert!(matches!(original.source, ToolCallSource::Manual));
17603        assert!(original.executed);
17604        assert!(!original.success);
17605
17606        let fallback = &records[1];
17607        assert_eq!(fallback.canonical_id, "fallback");
17608        assert_eq!(fallback.call_id, "fallback-call");
17609        assert!(matches!(
17610            &fallback.source,
17611            ToolCallSource::Fallback { original_tool } if original_tool == "primary"
17612        ));
17613        assert!(fallback.executed);
17614        assert!(fallback.success);
17615        assert_eq!(record.canonical_id, fallback.canonical_id);
17616        assert_eq!(record.output, fallback.output);
17617
17618        let history = agent.tool_call_history();
17619        assert_eq!(
17620            history
17621                .iter()
17622                .map(|entry| entry.tool_id.as_str())
17623                .collect::<Vec<_>>(),
17624            vec!["primary", "fallback"]
17625        );
17626        assert_eq!(history[0].result.get("success"), Some(&Value::Bool(false)));
17627        assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
17628        assert_eq!(fallback_calls.load(Ordering::SeqCst), 1);
17629    }
17630
17631    #[tokio::test]
17632    async fn diagnostics_without_provider_records_unavailable_without_execution() {
17633        let mock = mock_with_response("hello");
17634        let yaml = r#"
17635name: DiagnosticsNoProviderAgent
17636system_prompt: "Review diagnostics."
17637tools: [diagnostics]
17638"#;
17639        let agent = AgentBuilder::from_yaml(yaml)
17640            .unwrap()
17641            .llm(Arc::new(mock))
17642            .auto_configure_features()
17643            .unwrap()
17644            .build()
17645            .unwrap();
17646
17647        let record = agent
17648            .invoke_tool(ToolExecutionRequest::new(
17649                "diagnostics-call",
17650                "diagnostics",
17651                serde_json::json!({}),
17652                ToolCallSource::Manual,
17653            ))
17654            .await
17655            .unwrap();
17656
17657        assert!(!record.executed);
17658        assert!(!record.success);
17659        assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
17660    }
17661
17662    #[tokio::test]
17663    async fn web_search_without_provider_records_unavailable_without_execution() {
17664        let mock = mock_with_response("hello");
17665        let yaml = r#"
17666name: WebSearchNoProviderAgent
17667system_prompt: "You search the web."
17668tools: [web_search]
17669"#;
17670        let agent = AgentBuilder::from_yaml(yaml)
17671            .unwrap()
17672            .llm(Arc::new(mock))
17673            .auto_configure_features()
17674            .unwrap()
17675            .build()
17676            .unwrap();
17677
17678        let record = agent
17679            .invoke_tool(ToolExecutionRequest::new(
17680                "web-search-call",
17681                "web_search",
17682                serde_json::json!({"query": "rust async"}),
17683                ToolCallSource::Manual,
17684            ))
17685            .await
17686            .unwrap();
17687
17688        assert!(!record.executed);
17689        assert!(!record.success);
17690        assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
17691    }
17692
17693    #[tokio::test]
17694    async fn unavailable_host_tool_does_not_request_approval() {
17695        let approvals = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17696        let handler = Arc::new(CountingApprovalHandler {
17697            calls: Arc::clone(&approvals),
17698        });
17699        let mut security = ToolSecurityConfig {
17700            enabled: true,
17701            fail_closed: true,
17702            ..Default::default()
17703        };
17704        security.tools.insert(
17705            "web_search".to_string(),
17706            ai_agents_tools::ToolPolicyConfig {
17707                enabled: true,
17708                require_confirmation: true,
17709                ..Default::default()
17710            },
17711        );
17712        let yaml = r#"
17713name: UnavailableApprovalAgent
17714system_prompt: "Search only with approval."
17715tools: [web_search]
17716"#;
17717        let agent = AgentBuilder::from_yaml(yaml)
17718            .unwrap()
17719            .llm(Arc::new(mock_with_response("done")))
17720            .auto_configure_features()
17721            .unwrap()
17722            .tool_security(ToolSecurityEngine::new(security))
17723            .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17724            .approval_handler(handler)
17725            .build()
17726            .unwrap();
17727
17728        let record = agent
17729            .invoke_tool(ToolExecutionRequest::new(
17730                "unavailable-before-approval",
17731                "web_search",
17732                serde_json::json!({"query": "rust async"}),
17733                ToolCallSource::Manual,
17734            ))
17735            .await
17736            .unwrap();
17737
17738        assert_eq!(approvals.load(Ordering::SeqCst), 0);
17739        assert!(!record.executed);
17740        assert!(!record.success);
17741        assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
17742        assert!(
17743            record
17744                .approval
17745                .as_ref()
17746                .is_some_and(|approval| matches!(approval.status, ToolApprovalStatus::Unavailable))
17747        );
17748    }
17749
17750    #[tokio::test]
17751    async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_omitted() {
17752        let mock = mock_with_response("hello");
17753        let yaml = r#"
17754name: SpawnerNoGrantAgent
17755system_prompt: "You manage agents."
17756spawner:
17757  max_agents: 2
17758"#;
17759        let agent = AgentBuilder::from_yaml(yaml)
17760            .unwrap()
17761            .llm(Arc::new(mock))
17762            .auto_configure_features()
17763            .unwrap()
17764            .auto_configure_spawner()
17765            .await
17766            .unwrap()
17767            .build()
17768            .unwrap();
17769
17770        let available = agent.get_available_tool_ids().await.unwrap();
17771        assert!(available.is_empty());
17772    }
17773
17774    #[tokio::test]
17775    async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_empty() {
17776        let mock = mock_with_response("hello");
17777        let yaml = r#"
17778name: EmptySpawnerNoGrantAgent
17779system_prompt: "You manage agents."
17780tools: []
17781spawner:
17782  max_agents: 2
17783"#;
17784        let agent = AgentBuilder::from_yaml(yaml)
17785            .unwrap()
17786            .llm(Arc::new(mock))
17787            .auto_configure_features()
17788            .unwrap()
17789            .auto_configure_spawner()
17790            .await
17791            .unwrap()
17792            .build()
17793            .unwrap();
17794
17795        let available = agent.get_available_tool_ids().await.unwrap();
17796        assert!(available.is_empty());
17797    }
17798
17799    #[tokio::test]
17800    async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_empty() {
17801        let mock = mock_with_response("hello");
17802        let yaml = r#"
17803name: ManagementGrantAgent
17804system_prompt: "You manage agents."
17805tools: []
17806spawner:
17807  management_tools: true
17808"#;
17809        let agent = AgentBuilder::from_yaml(yaml)
17810            .unwrap()
17811            .llm(Arc::new(mock))
17812            .auto_configure_features()
17813            .unwrap()
17814            .auto_configure_spawner()
17815            .await
17816            .unwrap()
17817            .build()
17818            .unwrap();
17819
17820        let available = agent.get_available_tool_ids().await.unwrap();
17821        assert_eq!(available.len(), 4);
17822        assert!(available.contains(&"spawn_agent".to_string()));
17823        assert!(available.contains(&"send_agent_message".to_string()));
17824        assert!(available.contains(&"list_agents".to_string()));
17825        assert!(available.contains(&"remove_agent".to_string()));
17826    }
17827
17828    #[tokio::test]
17829    async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_omitted() {
17830        let mock = mock_with_response("hello");
17831        let yaml = r#"
17832name: ManagementOmittedToolsGrantAgent
17833system_prompt: "You manage agents."
17834spawner:
17835  management_tools: true
17836"#;
17837        let agent = AgentBuilder::from_yaml(yaml)
17838            .unwrap()
17839            .llm(Arc::new(mock))
17840            .auto_configure_features()
17841            .unwrap()
17842            .auto_configure_spawner()
17843            .await
17844            .unwrap()
17845            .build()
17846            .unwrap();
17847
17848        let available = agent.get_available_tool_ids().await.unwrap();
17849        assert_eq!(available.len(), 4);
17850        assert!(available.contains(&"spawn_agent".to_string()));
17851        assert!(available.contains(&"send_agent_message".to_string()));
17852        assert!(available.contains(&"list_agents".to_string()));
17853        assert!(available.contains(&"remove_agent".to_string()));
17854    }
17855
17856    #[tokio::test]
17857    async fn test_management_tools_selected_grants_only_selected_tools() {
17858        let mock = mock_with_response("hello");
17859        let yaml = r#"
17860name: ManagementSelectedGrantAgent
17861system_prompt: "You manage agents."
17862tools: []
17863spawner:
17864  management_tools:
17865    - spawn_agent
17866    - send_agent_message
17867    - list_agents
17868"#;
17869        let agent = AgentBuilder::from_yaml(yaml)
17870            .unwrap()
17871            .llm(Arc::new(mock))
17872            .auto_configure_features()
17873            .unwrap()
17874            .auto_configure_spawner()
17875            .await
17876            .unwrap()
17877            .build()
17878            .unwrap();
17879
17880        let available = agent.get_available_tool_ids().await.unwrap();
17881        assert_eq!(available.len(), 3);
17882        assert!(available.contains(&"spawn_agent".to_string()));
17883        assert!(available.contains(&"send_agent_message".to_string()));
17884        assert!(available.contains(&"list_agents".to_string()));
17885        assert!(!available.contains(&"remove_agent".to_string()));
17886    }
17887
17888    #[tokio::test]
17889    async fn test_orchestration_tools_flag_grants_tools_when_top_level_tools_empty() {
17890        let mock = mock_with_response("hello");
17891        let yaml = r#"
17892name: OrchestrationGrantAgent
17893system_prompt: "You coordinate agents."
17894llms:
17895  default:
17896    provider: openai
17897    model: gpt-4
17898  router:
17899    provider: openai
17900    model: gpt-4
17901llm:
17902  default: default
17903  router: router
17904tools: []
17905spawner:
17906  orchestration_tools: true
17907"#;
17908        let agent = AgentBuilder::from_yaml(yaml)
17909            .unwrap()
17910            .llm(Arc::new(mock))
17911            .auto_configure_features()
17912            .unwrap()
17913            .auto_configure_spawner()
17914            .await
17915            .unwrap()
17916            .build()
17917            .unwrap();
17918
17919        let available = agent.get_available_tool_ids().await.unwrap();
17920        assert_eq!(available.len(), 5);
17921        assert!(available.contains(&"route_to_agent".to_string()));
17922        assert!(available.contains(&"pipeline_process".to_string()));
17923        assert!(available.contains(&"concurrent_ask".to_string()));
17924        assert!(available.contains(&"group_discussion".to_string()));
17925        assert!(available.contains(&"handoff_conversation".to_string()));
17926    }
17927
17928    #[tokio::test]
17929    async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_empty() {
17930        let mock = mock_with_response("hello");
17931        let yaml = r#"
17932name: PersonaGrantAgent
17933system_prompt: "You can evolve persona."
17934llm:
17935  provider: openai
17936  model: gpt-4
17937tools: []
17938persona:
17939  identity:
17940    name: "Guide"
17941    role: "Helper"
17942  evolution:
17943    enabled: true
17944    allow_llm_evolve: true
17945    mutable_fields:
17946      - traits.personality
17947"#;
17948        let agent = AgentBuilder::from_yaml(yaml)
17949            .unwrap()
17950            .llm(Arc::new(mock))
17951            .build()
17952            .unwrap();
17953
17954        let available = agent.get_available_tool_ids().await.unwrap();
17955        assert_eq!(available, vec!["persona_evolve".to_string()]);
17956    }
17957
17958    #[tokio::test]
17959    async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_omitted() {
17960        let mock = mock_with_response("hello");
17961        let yaml = r#"
17962name: PersonaOmittedToolsGrantAgent
17963system_prompt: "You can evolve persona."
17964llm:
17965  provider: openai
17966  model: gpt-4
17967persona:
17968  identity:
17969    name: "Guide"
17970    role: "Helper"
17971  evolution:
17972    enabled: true
17973    allow_llm_evolve: true
17974    mutable_fields:
17975      - traits.personality
17976"#;
17977        let agent = AgentBuilder::from_yaml(yaml)
17978            .unwrap()
17979            .llm(Arc::new(mock))
17980            .build()
17981            .unwrap();
17982
17983        let available = agent.get_available_tool_ids().await.unwrap();
17984        assert_eq!(available, vec!["persona_evolve".to_string()]);
17985    }
17986
17987    #[tokio::test]
17988    async fn test_omitted_yaml_tools_exposes_no_tools() {
17989        let mock = mock_with_response("hello");
17990        let yaml = r#"
17991name: NoToolsAgent
17992system_prompt: "You are helpful."
17993"#;
17994        let agent = AgentBuilder::from_yaml(yaml)
17995            .unwrap()
17996            .llm(Arc::new(mock))
17997            .auto_configure_features()
17998            .unwrap()
17999            .build()
18000            .unwrap();
18001
18002        let available = agent.get_available_tool_ids().await.unwrap();
18003        assert!(available.is_empty());
18004    }
18005
18006    #[tokio::test]
18007    async fn runtime_scope_cannot_widen_omitted_or_empty_yaml_grants() {
18008        for tools in ["", "tools: []"] {
18009            let yaml = format!(
18010                r#"
18011name: RuntimeScopeNoGrantAgent
18012system_prompt: "No ordinary tools are granted."
18013{tools}
18014"#
18015            );
18016            let agent = AgentBuilder::from_yaml(&yaml)
18017                .unwrap()
18018                .llm(Arc::new(mock_with_response("done")))
18019                .auto_configure_features()
18020                .unwrap()
18021                .build()
18022                .unwrap();
18023
18024            agent
18025                .runtime_control()
18026                .set_tool_scope(vec!["calculator".to_string()]);
18027
18028            assert!(agent.get_available_tool_ids().await.unwrap().is_empty());
18029        }
18030    }
18031
18032    #[tokio::test]
18033    async fn runtime_scope_widening_attempt_keeps_only_declared_tools() {
18034        let yaml = r#"
18035name: RuntimeScopeWideningAgent
18036system_prompt: "Runtime scope cannot add authority."
18037tools: [calculator]
18038"#;
18039        let agent = AgentBuilder::from_yaml(yaml)
18040            .unwrap()
18041            .llm(Arc::new(mock_with_response("done")))
18042            .auto_configure_features()
18043            .unwrap()
18044            .build()
18045            .unwrap();
18046
18047        agent
18048            .runtime_control()
18049            .set_tool_scope(vec!["calculator".to_string(), "datetime".to_string()]);
18050
18051        assert_eq!(
18052            agent.get_available_tool_ids().await.unwrap(),
18053            vec!["calculator".to_string()]
18054        );
18055    }
18056
18057    #[tokio::test]
18058    async fn runtime_scope_is_canonical_unique_ordered_and_clear_restores_declared_grant() {
18059        let yaml = r#"
18060name: RuntimeScopeIntersectionAgent
18061system_prompt: "Use only declared tools."
18062tools: [calculator, datetime]
18063"#;
18064        let agent = AgentBuilder::from_yaml(yaml)
18065            .unwrap()
18066            .llm(Arc::new(mock_with_response("done")))
18067            .auto_configure_features()
18068            .unwrap()
18069            .build()
18070            .unwrap();
18071        let mut aliases = ai_agents_tools::ToolAliases::default();
18072        aliases
18073            .names
18074            .insert("en".to_string(), "calculate_alias".to_string());
18075        agent.tools.set_tool_aliases("calculator", aliases);
18076        let control = agent.runtime_control();
18077
18078        control.set_tool_scope(vec![
18079            "datetime".to_string(),
18080            "calculate_alias".to_string(),
18081            "calculator".to_string(),
18082            "unknown".to_string(),
18083            "datetime".to_string(),
18084        ]);
18085        assert_eq!(
18086            agent.get_available_tool_ids().await.unwrap(),
18087            vec!["calculator".to_string(), "datetime".to_string()]
18088        );
18089
18090        control.set_tool_scope(vec!["datetime".to_string()]);
18091        assert_eq!(
18092            agent.get_available_tool_ids().await.unwrap(),
18093            vec!["datetime".to_string()]
18094        );
18095
18096        control.clear_tool_scope_override();
18097        assert_eq!(
18098            agent.get_available_tool_ids().await.unwrap(),
18099            vec!["calculator".to_string(), "datetime".to_string()]
18100        );
18101    }
18102
18103    #[tokio::test]
18104    async fn runtime_scope_preserves_programmatic_registration_as_declared_grant() {
18105        let agent = AgentBuilder::new()
18106            .system_prompt("Use registered tools.")
18107            .llm(Arc::new(mock_with_response("done")))
18108            .tool(Arc::new(ContextEchoTool))
18109            .tool(Arc::new(SlowTool))
18110            .build()
18111            .unwrap();
18112
18113        agent.runtime_control().set_tool_scope(vec![
18114            "Context Echo".to_string(),
18115            "context_echo".to_string(),
18116            "unknown".to_string(),
18117        ]);
18118
18119        assert_eq!(
18120            agent.get_available_tool_ids().await.unwrap(),
18121            vec!["context_echo".to_string()]
18122        );
18123    }
18124
18125    #[tokio::test]
18126    async fn nested_state_scopes_intersect_every_ancestor_with_aliases() {
18127        let yaml = r#"
18128name: NestedStateScopeAgent
18129system_prompt: "Honor every state scope."
18130tools: [calculator, datetime, echo]
18131states:
18132  initial: root
18133  states:
18134    root:
18135      tools: [calculate_alias, datetime]
18136      initial: middle
18137      states:
18138        middle:
18139          initial: leaf
18140          states:
18141            leaf:
18142              tools: [datetime_alias, echo]
18143"#;
18144        let agent = AgentBuilder::from_yaml(yaml)
18145            .unwrap()
18146            .llm(Arc::new(mock_with_response("done")))
18147            .auto_configure_features()
18148            .unwrap()
18149            .build()
18150            .unwrap();
18151        let mut calculator_aliases = ai_agents_tools::ToolAliases::default();
18152        calculator_aliases
18153            .names
18154            .insert("en".to_string(), "calculate_alias".to_string());
18155        agent
18156            .tools
18157            .set_tool_aliases("calculator", calculator_aliases);
18158        let mut datetime_aliases = ai_agents_tools::ToolAliases::default();
18159        datetime_aliases
18160            .names
18161            .insert("en".to_string(), "datetime_alias".to_string());
18162        agent.tools.set_tool_aliases("datetime", datetime_aliases);
18163        agent.runtime_control().set_tool_scope(vec![
18164            "unknown".to_string(),
18165            "datetime_alias".to_string(),
18166            "calculate_alias".to_string(),
18167            "datetime".to_string(),
18168        ]);
18169
18170        assert_eq!(agent.current_state().as_deref(), Some("root.middle.leaf"));
18171        assert_eq!(
18172            agent.get_available_tool_ids().await.unwrap(),
18173            vec!["datetime".to_string()]
18174        );
18175    }
18176
18177    #[tokio::test]
18178    async fn ancestor_empty_state_scope_denies_omitted_descendants() {
18179        let yaml = r#"
18180name: NestedEmptyStateScopeAgent
18181system_prompt: "An empty ancestor scope denies all tools."
18182tools: [calculator]
18183states:
18184  initial: root
18185  states:
18186    root:
18187      tools: []
18188      initial: middle
18189      states:
18190        middle:
18191          initial: leaf
18192          states:
18193            leaf: {}
18194"#;
18195        let agent = AgentBuilder::from_yaml(yaml)
18196            .unwrap()
18197            .llm(Arc::new(mock_with_response("done")))
18198            .auto_configure_features()
18199            .unwrap()
18200            .build()
18201            .unwrap();
18202
18203        assert!(agent.get_available_tool_ids().await.unwrap().is_empty());
18204    }
18205
18206    #[tokio::test]
18207    async fn state_change_during_approval_invalidates_the_reviewed_authority() {
18208        let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18209        let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18210        let entered = Arc::new(tokio::sync::Barrier::new(2));
18211        let release = Arc::new(tokio::sync::Notify::new());
18212        let handler = Arc::new(BlockingApprovalHandler {
18213            entered: Arc::clone(&entered),
18214            release: Arc::clone(&release),
18215            result: ApprovalResult::Approved,
18216        });
18217        let yaml = r#"
18218name: ApprovalStateGenerationAgent
18219system_prompt: "State authority may change during approval."
18220tools: [locked_write]
18221states:
18222  initial: first
18223  states:
18224    first:
18225      tools: [locked_write]
18226    second:
18227      tools: [locked_write]
18228"#;
18229        let agent = Arc::new(
18230            AgentBuilder::from_yaml(yaml)
18231                .unwrap()
18232                .llm(Arc::new(mock_with_response("done")))
18233                .tool(Arc::new(LockedWriteTool {
18234                    active: Arc::clone(&active),
18235                    max_active: Arc::clone(&max_active),
18236                }))
18237                .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
18238                .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
18239                .approval_handler(handler)
18240                .build()
18241                .unwrap(),
18242        );
18243        let running = Arc::clone(&agent);
18244        let call = tokio::spawn(async move {
18245            running
18246                .invoke_tool(ToolExecutionRequest::new(
18247                    "approval-state-generation",
18248                    "locked_write",
18249                    serde_json::json!({"path": "./state-generation.txt"}),
18250                    ToolCallSource::Manual,
18251                ))
18252                .await
18253                .unwrap()
18254        });
18255
18256        entered.wait().await;
18257        agent.transition_to("second").await.unwrap();
18258        release.notify_one();
18259        let record = call.await.unwrap();
18260
18261        assert!(!record.executed);
18262        assert!(record.output.contains("Approval became stale"));
18263        assert_eq!(max_active.load(Ordering::SeqCst), 0);
18264    }
18265
18266    #[tokio::test]
18267    async fn state_change_while_waiting_for_resource_lock_fails_final_admission() {
18268        let holder_gate = PathMutationGate::new();
18269        let waiter_gate = PathMutationGate::new();
18270        let yaml = r#"
18271name: LockedStateGenerationAgent
18272system_prompt: "State authority must remain stable through admission."
18273tools: [state_lock_holder, state_lock_waiter]
18274states:
18275  initial: first
18276  states:
18277    first:
18278      tools: [state_lock_holder, state_lock_waiter]
18279    second:
18280      tools: [state_lock_holder, state_lock_waiter]
18281"#;
18282        let agent = Arc::new(
18283            AgentBuilder::from_yaml(yaml)
18284                .unwrap()
18285                .llm(Arc::new(mock_with_response("done")))
18286                .tool(Arc::new(BlockingPathMutationTool {
18287                    id: "state_lock_holder",
18288                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
18289                    gate: holder_gate.clone(),
18290                }))
18291                .tool(Arc::new(BlockingPathMutationTool {
18292                    id: "state_lock_waiter",
18293                    path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
18294                    gate: waiter_gate.clone(),
18295                }))
18296                .build()
18297                .unwrap(),
18298        );
18299        let holder_call = {
18300            let agent = Arc::clone(&agent);
18301            tokio::spawn(async move {
18302                agent
18303                    .invoke_tool(ToolExecutionRequest::new(
18304                        "state-lock-holder",
18305                        "state_lock_holder",
18306                        serde_json::json!({"path": "./shared-state-path.txt"}),
18307                        ToolCallSource::Manual,
18308                    ))
18309                    .await
18310                    .unwrap()
18311            })
18312        };
18313        holder_gate.wait_until_entered().await;
18314        let waiter_call = {
18315            let agent = Arc::clone(&agent);
18316            tokio::spawn(async move {
18317                agent
18318                    .invoke_tool(ToolExecutionRequest::new(
18319                        "state-lock-waiter",
18320                        "state_lock_waiter",
18321                        serde_json::json!({"path": "./shared-state-path.txt"}),
18322                        ToolCallSource::Manual,
18323                    ))
18324                    .await
18325                    .unwrap()
18326            })
18327        };
18328
18329        wait_for_resource_lock_strong_count(&agent.resource_locks, 2).await;
18330        agent.transition_to("second").await.unwrap();
18331        holder_gate.release();
18332        let holder_record = holder_call.await.unwrap();
18333        let waiter_record = waiter_call.await.unwrap();
18334
18335        assert!(holder_record.success);
18336        assert!(!waiter_record.executed);
18337        assert!(
18338            waiter_record
18339                .output
18340                .contains("state scope changed before admission")
18341        );
18342        assert!(!waiter_gate.entered.load(Ordering::SeqCst));
18343    }
18344
18345    #[tokio::test]
18346    async fn test_state_tools_cannot_widen_top_level_grant() {
18347        let mock = mock_with_response("hello");
18348        let yaml = r#"
18349name: NarrowToolsAgent
18350system_prompt: "You are helpful."
18351tools:
18352  - calculator
18353states:
18354  initial: current
18355  states:
18356    current:
18357      tools: [datetime]
18358"#;
18359        let agent = AgentBuilder::from_yaml(yaml)
18360            .unwrap()
18361            .llm(Arc::new(mock))
18362            .auto_configure_features()
18363            .unwrap()
18364            .build()
18365            .unwrap();
18366
18367        let available = agent.get_available_tool_ids().await.unwrap();
18368        assert!(available.is_empty());
18369    }
18370
18371    // Tool execution in chat flow
18372    #[tokio::test]
18373    async fn test_integration_tool_execution() {
18374        // Mock LLM that returns a tool call then a final answer
18375        let mock = mock_with_responses(vec![
18376            // First response: tool call
18377            r#"I'll calculate that for you.
18378[TOOL_CALL: {"name": "calculator", "arguments": {"expression": "2+2"}}]"#,
18379            // After tool result: final answer
18380            "The answer is 4.",
18381        ]);
18382        let mut tools = ai_agents_tools::ToolRegistry::new();
18383        tools
18384            .register(Arc::new(ai_agents_tools::CalculatorTool))
18385            .unwrap();
18386
18387        let agent = AgentBuilder::new()
18388            .system_prompt("You are a calculator assistant.")
18389            .llm(Arc::new(mock))
18390            .tools(tools)
18391            .build()
18392            .unwrap();
18393
18394        let response = agent.chat("What is 2+2?").await.unwrap();
18395        // The agent should eventually produce a response
18396        assert!(!response.content.is_empty());
18397    }
18398
18399    #[tokio::test]
18400    async fn test_tool_hitl_rejection_finalizes_blocking_turn() {
18401        let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18402        let hooks = Arc::new(ResponseCountingHooks {
18403            responses: Arc::clone(&responses),
18404        });
18405        let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
18406        let yaml = r#"
18407name: ToolRejectAgent
18408system_prompt: "You use tools when requested."
18409tools:
18410  - echo
18411hitl:
18412  tools:
18413    echo:
18414      require_approval: true
18415      approval_message: "Approve echo?"
18416"#;
18417        let agent = AgentBuilder::from_yaml(yaml)
18418            .unwrap()
18419            .llm(Arc::new(mock))
18420            .auto_configure_features()
18421            .unwrap()
18422            .hooks(hooks)
18423            .build()
18424            .unwrap();
18425
18426        let response = agent.chat("echo hello").await.unwrap();
18427
18428        assert!(
18429            response.content.contains("Operation cancelled"),
18430            "unexpected response: {}",
18431            response.content
18432        );
18433        assert_eq!(responses.load(Ordering::SeqCst), 1);
18434        let messages = agent.memory.get_messages(None).await.unwrap();
18435        assert_eq!(messages.len(), 3);
18436        assert_eq!(messages[0].content, "echo hello");
18437        assert!(messages[1].content.contains("\"tool\":\"echo\""));
18438        assert!(messages[2].content.contains("rejected by the approver"));
18439    }
18440
18441    #[tokio::test]
18442    async fn test_tool_hitl_rejection_finalizes_streaming_turn() {
18443        use futures::StreamExt;
18444
18445        let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18446        let hooks = Arc::new(ResponseCountingHooks {
18447            responses: Arc::clone(&responses),
18448        });
18449        let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
18450        let yaml = r#"
18451name: ToolRejectStreamingAgent
18452system_prompt: "You use tools when requested."
18453tools:
18454  - echo
18455streaming:
18456  enabled: true
18457hitl:
18458  tools:
18459    echo:
18460      require_approval: true
18461      approval_message: "Approve echo?"
18462"#;
18463        let agent = AgentBuilder::from_yaml(yaml)
18464            .unwrap()
18465            .llm(Arc::new(mock))
18466            .auto_configure_features()
18467            .unwrap()
18468            .hooks(hooks)
18469            .build()
18470            .unwrap();
18471
18472        let mut stream = agent.chat_stream("echo hello").await.unwrap();
18473        let mut terminal_error = String::new();
18474        let mut done = false;
18475        while let Some(chunk) = stream.next().await {
18476            match chunk {
18477                StreamChunk::Error { message } => terminal_error = message,
18478                StreamChunk::Done {} => {
18479                    done = true;
18480                    break;
18481                }
18482                _ => {}
18483            }
18484        }
18485
18486        assert!(done);
18487        assert!(
18488            terminal_error.contains("Operation cancelled"),
18489            "unexpected terminal error: {}",
18490            terminal_error
18491        );
18492        assert_eq!(responses.load(Ordering::SeqCst), 1);
18493        let messages = agent.memory.get_messages(None).await.unwrap();
18494        assert_eq!(messages.len(), 3);
18495        assert_eq!(messages[0].content, "echo hello");
18496        assert!(messages[1].content.contains("\"tool\":\"echo\""));
18497        assert!(messages[2].content.contains("rejected by the approver"));
18498    }
18499
18500    #[tokio::test]
18501    async fn tool_hitl_rejection_preserves_legacy_error_but_finalizes_event_stream() {
18502        let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
18503        let yaml = r#"
18504name: ToolRejectEventAgent
18505system_prompt: "You use tools when requested."
18506tools:
18507  - echo
18508streaming:
18509  enabled: true
18510hitl:
18511  tools:
18512    echo:
18513      require_approval: true
18514      approval_message: "Approve echo?"
18515"#;
18516        let agent = AgentBuilder::from_yaml(yaml)
18517            .unwrap()
18518            .llm(Arc::new(mock))
18519            .auto_configure_features()
18520            .unwrap()
18521            .build()
18522            .unwrap();
18523
18524        let mut stream = agent.chat_stream_events("echo hello").await.unwrap();
18525        let mut error_seen = false;
18526        let mut final_response = None;
18527        while let Some(event) = stream.next().await {
18528            match event {
18529                AgentStreamEvent::Chunk(StreamChunk::Error { .. }) => error_seen = true,
18530                AgentStreamEvent::Final(response) => final_response = Some(response),
18531                AgentStreamEvent::Chunk(_) => {}
18532            }
18533        }
18534
18535        assert!(!error_seen);
18536        assert!(
18537            final_response
18538                .is_some_and(|response| { response.content.contains("Operation cancelled") })
18539        );
18540    }
18541
18542    #[tokio::test]
18543    async fn test_pre_response_guard_transition_skips_old_state_llm() {
18544        let mock = mock_with_response("Billing state response");
18545        let call_counter = mock.clone();
18546        let yaml = r#"
18547name: OptimizedStateAgent
18548system_prompt: "You route before answering."
18549runtime:
18550  optimization:
18551    enabled: true
18552    pre_response_deterministic_transitions: true
18553states:
18554  initial: greeting
18555  states:
18556    greeting:
18557      prompt: "Old state prompt that should be skipped."
18558      transitions:
18559        - to: billing
18560          guard:
18561            context:
18562              topic:
18563                eq: billing
18564          timing: pre_response
18565    billing:
18566      prompt: "Answer from the billing state."
18567"#;
18568        let agent = AgentBuilder::from_yaml(yaml)
18569            .unwrap()
18570            .llm(Arc::new(mock))
18571            .build()
18572            .unwrap();
18573        agent
18574            .set_context("topic", serde_json::json!("billing"))
18575            .unwrap();
18576
18577        let response = agent.chat("I need billing help").await.unwrap();
18578
18579        assert_eq!(agent.current_state().as_deref(), Some("billing"));
18580        assert_eq!(response.content, "Billing state response");
18581        assert_eq!(call_counter.call_count(), 1);
18582        assert_eq!(agent.actor_facts().len(), 0);
18583    }
18584
18585    #[tokio::test]
18586    async fn test_set_context_supports_dotted_paths_for_pre_response_guards() {
18587        let mock = mock_with_response("Billing state response");
18588        let call_counter = mock.clone();
18589        let yaml = r#"
18590name: OptimizedStateAgent
18591system_prompt: "You route before answering."
18592runtime:
18593  optimization:
18594    enabled: true
18595    pre_response_deterministic_transitions: true
18596context:
18597  request:
18598    type: runtime
18599    default:
18600      topic: general
18601states:
18602  initial: greeting
18603  states:
18604    greeting:
18605      prompt: "Old state prompt that should be skipped."
18606      transitions:
18607        - to: billing
18608          guard:
18609            context:
18610              request.topic:
18611                eq: billing
18612          timing: pre_response
18613    billing:
18614      prompt: "Answer from the billing state."
18615"#;
18616        let agent = AgentBuilder::from_yaml(yaml)
18617            .unwrap()
18618            .llm(Arc::new(mock))
18619            .build()
18620            .unwrap();
18621        agent
18622            .set_context("request.topic", serde_json::json!("billing"))
18623            .unwrap();
18624
18625        let response = agent.chat("I need billing help").await.unwrap();
18626
18627        assert_eq!(agent.current_state().as_deref(), Some("billing"));
18628        assert_eq!(response.content, "Billing state response");
18629        assert_eq!(call_counter.call_count(), 1);
18630        assert_eq!(
18631            agent.get_context().get("request"),
18632            Some(&serde_json::json!({"topic": "billing"}))
18633        );
18634    }
18635
18636    #[tokio::test]
18637    async fn test_pre_response_rejection_does_not_commit_staged_context_or_user() {
18638        let mock = mock_with_response("billing");
18639        let yaml = r#"
18640name: OptimizedStateAgent
18641system_prompt: "You route before answering."
18642runtime:
18643  optimization:
18644    enabled: true
18645    pre_response_deterministic_transitions: true
18646hitl:
18647  states:
18648    billing:
18649      on_enter: require_approval
18650      approval_message: "Approve billing route?"
18651states:
18652  initial: greeting
18653  states:
18654    greeting:
18655      prompt: "Old state prompt."
18656      extract:
18657        - key: topic
18658          description: "Support topic"
18659      transitions:
18660        - to: billing
18661          guard:
18662            context:
18663              topic:
18664                eq: billing
18665          timing: pre_response
18666          run_extractors: true
18667    billing:
18668      prompt: "Billing state."
18669"#;
18670        let agent = AgentBuilder::from_yaml(yaml)
18671            .unwrap()
18672            .llm(Arc::new(mock))
18673            .build()
18674            .unwrap();
18675
18676        let response = agent
18677            .try_pre_response_transition("billing please")
18678            .await
18679            .unwrap();
18680
18681        assert!(response.is_none());
18682        assert_eq!(agent.current_state().as_deref(), Some("greeting"));
18683        assert!(!agent.get_context().contains_key("topic"));
18684        assert_eq!(agent.memory.get_messages(None).await.unwrap().len(), 0);
18685    }
18686
18687    #[tokio::test]
18688    async fn test_pre_response_extractor_commits_context_on_winning_path() {
18689        let mock = mock_with_responses(vec!["billing", "Billing response"]);
18690        let yaml = r#"
18691name: OptimizedStateAgent
18692system_prompt: "You route before answering."
18693runtime:
18694  optimization:
18695    enabled: true
18696    pre_response_deterministic_transitions: true
18697states:
18698  initial: greeting
18699  states:
18700    greeting:
18701      prompt: "Old state prompt."
18702      extract:
18703        - key: topic
18704          description: "Support topic"
18705      transitions:
18706        - to: billing
18707          guard:
18708            context:
18709              topic:
18710                eq: billing
18711          timing: pre_response
18712          run_extractors: true
18713    billing:
18714      prompt: "Billing state."
18715"#;
18716        let agent = AgentBuilder::from_yaml(yaml)
18717            .unwrap()
18718            .llm(Arc::new(mock))
18719            .build()
18720            .unwrap();
18721
18722        let response = agent.chat("billing please").await.unwrap();
18723
18724        assert_eq!(agent.current_state().as_deref(), Some("billing"));
18725        assert_eq!(response.content, "Billing response");
18726        assert_eq!(
18727            agent.get_context().get("topic"),
18728            Some(&serde_json::json!("billing"))
18729        );
18730    }
18731
18732    #[tokio::test]
18733    async fn test_pre_response_extractor_miss_does_not_mutate_context() {
18734        let mock = mock_with_response("__NONE__");
18735        let yaml = r#"
18736name: OptimizedStateAgent
18737system_prompt: "You route before answering."
18738runtime:
18739  optimization:
18740    enabled: true
18741    pre_response_deterministic_transitions: true
18742states:
18743  initial: greeting
18744  states:
18745    greeting:
18746      prompt: "Old state prompt."
18747      extract:
18748        - key: topic
18749          description: "Support topic"
18750      transitions:
18751        - to: billing
18752          guard:
18753            context:
18754              topic:
18755                eq: billing
18756          timing: pre_response
18757          run_extractors: true
18758    billing:
18759      prompt: "Billing state."
18760"#;
18761        let agent = AgentBuilder::from_yaml(yaml)
18762            .unwrap()
18763            .llm(Arc::new(mock))
18764            .build()
18765            .unwrap();
18766
18767        let response = agent.try_pre_response_transition("hello").await.unwrap();
18768
18769        assert!(response.is_none());
18770        assert_eq!(agent.current_state().as_deref(), Some("greeting"));
18771        assert!(!agent.get_context().contains_key("topic"));
18772    }
18773
18774    #[tokio::test]
18775    async fn test_default_guard_transition_stays_post_response() {
18776        let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
18777        let call_counter = mock.clone();
18778        let yaml = r#"
18779name: TimingAgent
18780system_prompt: "You route carefully."
18781runtime:
18782  optimization:
18783    enabled: true
18784    pre_response_deterministic_transitions: true
18785states:
18786  initial: greeting
18787  states:
18788    greeting:
18789      prompt: "Old state prompt."
18790      transitions:
18791        - to: billing
18792          guard:
18793            context:
18794              topic:
18795                eq: billing
18796    billing:
18797      prompt: "Billing state."
18798"#;
18799        let agent = AgentBuilder::from_yaml(yaml)
18800            .unwrap()
18801            .llm(Arc::new(mock))
18802            .build()
18803            .unwrap();
18804        agent
18805            .set_context("topic", serde_json::json!("billing"))
18806            .unwrap();
18807
18808        let response = agent.chat("billing please").await.unwrap();
18809
18810        assert_eq!(agent.current_state().as_deref(), Some("billing"));
18811        assert_eq!(response.content, "Billing response");
18812        assert_eq!(call_counter.call_count(), 2);
18813    }
18814
18815    #[tokio::test]
18816    async fn test_explicit_post_response_guard_transition_stays_post_response() {
18817        let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
18818        let call_counter = mock.clone();
18819        let yaml = r#"
18820name: TimingAgent
18821system_prompt: "You route carefully."
18822runtime:
18823  optimization:
18824    enabled: true
18825    pre_response_deterministic_transitions: true
18826states:
18827  initial: greeting
18828  states:
18829    greeting:
18830      prompt: "Old state prompt."
18831      transitions:
18832        - to: billing
18833          guard:
18834            context:
18835              topic:
18836                eq: billing
18837          timing: post_response
18838    billing:
18839      prompt: "Billing state."
18840"#;
18841        let agent = AgentBuilder::from_yaml(yaml)
18842            .unwrap()
18843            .llm(Arc::new(mock))
18844            .build()
18845            .unwrap();
18846        agent
18847            .set_context("topic", serde_json::json!("billing"))
18848            .unwrap();
18849
18850        let response = agent.chat("billing please").await.unwrap();
18851
18852        assert_eq!(agent.current_state().as_deref(), Some("billing"));
18853        assert_eq!(response.content, "Billing response");
18854        assert_eq!(call_counter.call_count(), 2);
18855    }
18856
18857    #[tokio::test]
18858    async fn test_pre_response_extractors_are_transition_scoped() {
18859        let mock = mock_with_responses(vec!["billing", "Billing response"]);
18860        let yaml = r#"
18861name: ScopedExtractorAgent
18862system_prompt: "You route carefully."
18863runtime:
18864  optimization:
18865    enabled: true
18866    pre_response_deterministic_transitions: true
18867states:
18868  initial: greeting
18869  states:
18870    greeting:
18871      prompt: "Old state prompt."
18872      extract:
18873        - key: topic
18874          description: "Support topic"
18875      transitions:
18876        - to: wrong
18877          guard:
18878            context:
18879              topic:
18880                eq: billing
18881          timing: pre_response
18882        - to: billing
18883          guard:
18884            context:
18885              topic:
18886                eq: billing
18887          timing: pre_response
18888          run_extractors: true
18889    wrong:
18890      prompt: "Wrong state."
18891    billing:
18892      prompt: "Billing state."
18893"#;
18894        let agent = AgentBuilder::from_yaml(yaml)
18895            .unwrap()
18896            .llm(Arc::new(mock))
18897            .build()
18898            .unwrap();
18899
18900        let response = agent.chat("billing please").await.unwrap();
18901
18902        assert_eq!(agent.current_state().as_deref(), Some("billing"));
18903        assert_eq!(response.content, "Billing response");
18904    }
18905
18906    #[tokio::test]
18907    async fn test_pre_response_resolved_intent_routes_early() {
18908        let mock = mock_with_response("Billing response");
18909        let yaml = r#"
18910name: IntentAgent
18911system_prompt: "You route carefully."
18912runtime:
18913  optimization:
18914    enabled: true
18915    pre_response_deterministic_transitions: true
18916states:
18917  initial: greeting
18918  states:
18919    greeting:
18920      prompt: "Old state prompt."
18921      transitions:
18922        - to: billing
18923          intent: billing
18924          timing: pre_response
18925    billing:
18926      prompt: "Billing state."
18927"#;
18928        let agent = AgentBuilder::from_yaml(yaml)
18929            .unwrap()
18930            .llm(Arc::new(mock))
18931            .build()
18932            .unwrap();
18933        agent
18934            .set_context("resolved_intent", serde_json::json!("billing"))
18935            .unwrap();
18936
18937        let response = agent
18938            .try_pre_response_transition("I need billing help")
18939            .await
18940            .unwrap()
18941            .unwrap();
18942
18943        assert_eq!(agent.current_state().as_deref(), Some("billing"));
18944        assert_eq!(response.content, "Billing response");
18945    }
18946
18947    #[tokio::test]
18948    async fn test_background_overflow_error_surfaces() {
18949        let mut config = RuntimeConfig::default();
18950        config.optimization.enabled = true;
18951        config.optimization.post_turn.max_background_tasks = 1;
18952        config.optimization.post_turn.on_background_overflow = BackgroundOverflowPolicy::Error;
18953        let policy = crate::optimization::MaintenanceTaskPolicy {
18954            mode: MaintenanceMode::Background,
18955            await_before_next_turn: AwaitBeforeNextTurn::Always,
18956        };
18957        let agent = AgentBuilder::new()
18958            .system_prompt("You are helpful.")
18959            .llm(Arc::new(mock_with_response("ok")))
18960            .build()
18961            .unwrap()
18962            .with_runtime_config(config);
18963        agent
18964            .background_maintenance
18965            .spawn(None, async { std::future::pending::<Result<()>>().await })
18966            .unwrap();
18967
18968        let result = agent
18969            .spawn_or_handle_background(None, async { Ok(()) }, "facts", &policy)
18970            .await;
18971
18972        assert!(result.is_err());
18973    }
18974
18975    #[tokio::test]
18976    async fn test_speculative_reasoning_low_cap_uses_serial_reasoning() {
18977        let default_mock = mock_with_response("Plain draft response");
18978        let router_mock = mock_with_response("cot");
18979        let router_counter = router_mock.clone();
18980        let yaml = r#"
18981name: ReasoningReservationAgent
18982system_prompt: "You answer plainly unless reasoning wins."
18983llm:
18984  default: default
18985  router: router
18986observability:
18987  enabled: true
18988  export:
18989    write_raw_events: true
18990reasoning:
18991  mode: auto
18992  judge_llm: router
18993runtime:
18994  optimization:
18995    enabled: true
18996    max_speculative_llm_calls_per_turn: 1
18997    speculative_reasoning_auto: true
18998    max_parallel_runtime_tasks: 2
18999"#;
19000        let agent = AgentBuilder::from_yaml(yaml)
19001            .unwrap()
19002            .llm_alias("default", Arc::new(default_mock))
19003            .llm_alias("router", Arc::new(router_mock))
19004            .build()
19005            .unwrap();
19006
19007        let response = agent.chat("hello").await.unwrap();
19008
19009        assert_eq!(response.content, "Plain draft response");
19010        assert_eq!(router_counter.call_count(), 1);
19011        let events = agent.observability().unwrap().raw_events();
19012        assert!(!events.iter().any(|event| {
19013            event.dimensions.get("commit_behavior") == Some(&"reasoning_decision".to_string())
19014        }));
19015    }
19016
19017    #[tokio::test]
19018    async fn test_forced_reasoning_skips_plain_speculative_draft() {
19019        let mock = mock_with_response("Reasoned response");
19020        let yaml = r#"
19021name: ForcedReasoningAgent
19022system_prompt: "You reason before answering."
19023observability:
19024  enabled: true
19025  export:
19026    write_raw_events: true
19027reasoning:
19028  mode: cot
19029runtime:
19030  optimization:
19031    enabled: true
19032    max_speculative_llm_calls_per_turn: 2
19033    speculative_state_transitions: true
19034    max_parallel_runtime_tasks: 2
19035states:
19036  initial: triage
19037  states:
19038    triage:
19039      prompt: "Answer from triage."
19040      transitions:
19041        - to: billing
19042          guard:
19043            context:
19044              route:
19045                eq: billing
19046          timing: parallel
19047    billing:
19048      prompt: "Billing state."
19049"#;
19050        let agent = AgentBuilder::from_yaml(yaml)
19051            .unwrap()
19052            .llm(Arc::new(mock))
19053            .build()
19054            .unwrap();
19055
19056        let response = agent.chat("hello").await.unwrap();
19057
19058        assert_eq!(response.content, "Reasoned response");
19059        let events = agent.observability().unwrap().raw_events();
19060        assert!(
19061            !events
19062                .iter()
19063                .any(|event| event.dimensions.contains_key("branch_status"))
19064        );
19065    }
19066
19067    #[tokio::test]
19068    async fn test_speculative_skill_low_cap_uses_serial_skill_route() {
19069        let default_mock = mock_with_response("Skill committed response");
19070        let router_mock = mock_with_response("helper");
19071        let router_counter = router_mock.clone();
19072        let yaml = r#"
19073name: SkillReservationAgent
19074system_prompt: "Use skills when they match."
19075llm:
19076  default: default
19077  router: router
19078observability:
19079  enabled: true
19080  export:
19081    write_raw_events: true
19082runtime:
19083  optimization:
19084    enabled: true
19085    max_speculative_llm_calls_per_turn: 1
19086    speculative_skill_routing: true
19087    max_parallel_runtime_tasks: 2
19088skills:
19089  - id: helper
19090    description: "Answer helper requests"
19091    trigger: "User asks for helper"
19092    steps:
19093      - prompt: "Answer the helper request: {{ user_input }}"
19094"#;
19095        let agent = AgentBuilder::from_yaml(yaml)
19096            .unwrap()
19097            .llm_alias("default", Arc::new(default_mock))
19098            .llm_alias("router", Arc::new(router_mock))
19099            .build()
19100            .unwrap();
19101
19102        let response = agent.chat("please use helper").await.unwrap();
19103
19104        assert_eq!(response.content, "Skill committed response");
19105        assert_eq!(router_counter.call_count(), 1);
19106        let events = agent.observability().unwrap().raw_events();
19107        assert!(
19108            !events
19109                .iter()
19110                .any(|event| event.dimensions.contains_key("branch_status"))
19111        );
19112    }
19113
19114    #[tokio::test]
19115    async fn test_parallel_transition_low_cap_allows_deterministic_route() {
19116        let mock = mock_with_response("unused");
19117        let call_counter = mock.clone();
19118        let yaml = r#"
19119name: ParallelTransitionLowCapAgent
19120system_prompt: "Route before stale responses when safe."
19121runtime:
19122  optimization:
19123    enabled: true
19124    max_speculative_llm_calls_per_turn: 1
19125    speculative_state_transitions: true
19126    max_parallel_runtime_tasks: 2
19127states:
19128  initial: triage
19129  states:
19130    triage:
19131      prompt: "Triage state."
19132      transitions:
19133        - to: billing
19134          guard:
19135            context:
19136              route:
19137                eq: billing
19138          timing: parallel
19139    billing:
19140      prompt: "Billing state."
19141"#;
19142        let agent = AgentBuilder::from_yaml(yaml)
19143            .unwrap()
19144            .llm(Arc::new(mock))
19145            .build()
19146            .unwrap();
19147        agent
19148            .set_context("route", serde_json::json!("billing"))
19149            .unwrap();
19150        agent.update_active_turn_context("billing help", HashMap::new());
19151        assert!(
19152            agent.reserve_active_speculative_llm_call(
19153                RuntimeOptimizationKind::ParallelStateTransition
19154            )
19155        );
19156
19157        let selection = agent
19158            .select_parallel_transition_candidate("billing help")
19159            .await
19160            .unwrap();
19161        agent.end_root_turn();
19162
19163        match selection {
19164            ParallelTransitionSelection::Candidate(candidate) => {
19165                assert_eq!(candidate.target(), "billing");
19166            }
19167            ParallelTransitionSelection::NoMatch => panic!("deterministic route did not match"),
19168            ParallelTransitionSelection::ReservationExhausted => {
19169                panic!("deterministic route consumed LLM budget")
19170            }
19171        }
19172        assert_eq!(call_counter.call_count(), 0);
19173    }
19174
19175    #[tokio::test]
19176    async fn speculative_transition_drops_loser_before_state_actions() {
19177        let lock = Arc::new(tokio::sync::Mutex::new(()));
19178        let first_started = Arc::new(tokio::sync::Notify::new());
19179        let first_dropped = Arc::new(AtomicBool::new(false));
19180        let committed_after_drop = Arc::new(AtomicBool::new(false));
19181        let default = Arc::new(FirstCallLockingProvider {
19182            lock,
19183            first_started: Arc::clone(&first_started),
19184            first_dropped: Arc::clone(&first_dropped),
19185            committed_after_drop: Arc::clone(&committed_after_drop),
19186            calls: AtomicU64::new(0),
19187        });
19188        let router = Arc::new(RoutingAfterProviderStart {
19189            provider_started: first_started,
19190        });
19191        let yaml = r#"
19192name: SpeculativeCancellationAgent
19193system_prompt: "Route before committed work."
19194llm:
19195  default: default
19196  router: router
19197runtime:
19198  optimization:
19199    enabled: true
19200    max_speculative_llm_calls_per_turn: 2
19201    speculative_state_transitions: true
19202    max_parallel_runtime_tasks: 2
19203states:
19204  initial: triage
19205  states:
19206    triage:
19207      prompt: "Triage state."
19208      transitions:
19209        - to: technical
19210          when: "The request needs technical support"
19211          timing: parallel
19212    technical:
19213      prompt: "Technical state."
19214      on_enter:
19215        - prompt: "Prepare technical context."
19216          llm: default
19217          store_as: preparation
19218"#;
19219        let agent = AgentBuilder::from_yaml(yaml)
19220            .unwrap()
19221            .llm_alias("default", default)
19222            .llm_alias("router", router)
19223            .build()
19224            .unwrap();
19225
19226        let response = tokio::time::timeout(
19227            std::time::Duration::from_secs(2),
19228            agent.chat("I cannot log in because of AUTH-17."),
19229        )
19230        .await
19231        .expect("committed work must not wait on the losing provider future")
19232        .unwrap();
19233
19234        assert_eq!(response.content, "Committed technical response.");
19235        assert_eq!(agent.current_state().as_deref(), Some("technical"));
19236        assert!(first_dropped.load(Ordering::SeqCst));
19237        assert!(committed_after_drop.load(Ordering::SeqCst));
19238    }
19239
19240    #[tokio::test]
19241    async fn buffered_transition_drops_stale_stream_before_redispatch() {
19242        use futures::StreamExt;
19243
19244        let lock = Arc::new(tokio::sync::Mutex::new(()));
19245        let stream_started = Arc::new(tokio::sync::Notify::new());
19246        let stream_dropped = Arc::new(AtomicBool::new(false));
19247        let committed_after_drop = Arc::new(AtomicBool::new(false));
19248        let default = Arc::new(BufferedLockingProvider {
19249            lock,
19250            stream_started: Arc::clone(&stream_started),
19251            stream_dropped: Arc::clone(&stream_dropped),
19252            committed_after_drop: Arc::clone(&committed_after_drop),
19253        });
19254        let router = Arc::new(RoutingAfterProviderStart {
19255            provider_started: stream_started,
19256        });
19257        let yaml = r#"
19258name: BufferedCancellationAgent
19259system_prompt: "Hide stale streamed output."
19260llm:
19261  default: default
19262  router: router
19263streaming:
19264  enabled: true
19265  buffer_size: 8
19266runtime:
19267  optimization:
19268    enabled: true
19269    max_speculative_llm_calls_per_turn: 2
19270    speculative_state_transitions: true
19271    streaming_policy: buffer_until_routing_done
19272    max_parallel_runtime_tasks: 2
19273states:
19274  initial: triage
19275  states:
19276    triage:
19277      prompt: "Triage state."
19278      transitions:
19279        - to: technical
19280          when: "The request needs technical support"
19281          timing: parallel
19282    technical:
19283      prompt: "Technical state."
19284"#;
19285        let agent = AgentBuilder::from_yaml(yaml)
19286            .unwrap()
19287            .llm_alias("default", default)
19288            .llm_alias("router", router)
19289            .build()
19290            .unwrap();
19291
19292        let content = tokio::time::timeout(std::time::Duration::from_secs(2), async {
19293            let mut stream = agent
19294                .chat_stream("AUTH-17 needs technical help.")
19295                .await
19296                .unwrap();
19297            let mut content = String::new();
19298            while let Some(chunk) = stream.next().await {
19299                match chunk {
19300                    StreamChunk::Content { text } => content.push_str(&text),
19301                    StreamChunk::Done {} => break,
19302                    StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
19303                    _ => {}
19304                }
19305            }
19306            content
19307        })
19308        .await
19309        .expect("redispatch must not wait on the stale streaming future");
19310
19311        assert_eq!(content, "Committed technical response.");
19312        assert_eq!(agent.current_state().as_deref(), Some("technical"));
19313        assert!(stream_dropped.load(Ordering::SeqCst));
19314        assert!(committed_after_drop.load(Ordering::SeqCst));
19315    }
19316
19317    #[tokio::test]
19318    async fn buffered_transition_drops_established_stream_before_redispatch() {
19319        use futures::StreamExt;
19320
19321        let stream_started = Arc::new(tokio::sync::Notify::new());
19322        let stream_dropped = Arc::new(AtomicBool::new(false));
19323        let stream_dropped_notify = Arc::new(tokio::sync::Notify::new());
19324        let committed_after_drop = Arc::new(AtomicBool::new(false));
19325        let default = Arc::new(EstablishedStreamProvider {
19326            stream_started: Arc::clone(&stream_started),
19327            stream_dropped: Arc::clone(&stream_dropped),
19328            stream_dropped_notify,
19329            committed_after_drop: Arc::clone(&committed_after_drop),
19330        });
19331        let router = Arc::new(RoutingAfterProviderStart {
19332            provider_started: stream_started,
19333        });
19334        let yaml = r#"
19335name: EstablishedStreamCancellationAgent
19336system_prompt: "Hide stale streamed output."
19337llm:
19338  default: default
19339  router: router
19340streaming:
19341  enabled: true
19342  buffer_size: 8
19343runtime:
19344  optimization:
19345    enabled: true
19346    max_speculative_llm_calls_per_turn: 2
19347    speculative_state_transitions: true
19348    streaming_policy: buffer_until_routing_done
19349    max_parallel_runtime_tasks: 2
19350states:
19351  initial: triage
19352  states:
19353    triage:
19354      prompt: "Triage state."
19355      transitions:
19356        - to: technical
19357          when: "The request needs technical support"
19358          timing: parallel
19359    technical:
19360      prompt: "Technical state."
19361"#;
19362        let agent = AgentBuilder::from_yaml(yaml)
19363            .unwrap()
19364            .llm_alias("default", default)
19365            .llm_alias("router", router)
19366            .build()
19367            .unwrap();
19368
19369        let content = tokio::time::timeout(std::time::Duration::from_secs(2), async {
19370            let mut stream = agent
19371                .chat_stream("AUTH-17 needs technical help.")
19372                .await
19373                .unwrap();
19374            let mut content = String::new();
19375            while let Some(chunk) = stream.next().await {
19376                match chunk {
19377                    StreamChunk::Content { text } => content.push_str(&text),
19378                    StreamChunk::Done {} => break,
19379                    StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
19380                    _ => {}
19381                }
19382            }
19383            content
19384        })
19385        .await
19386        .expect("redispatch must wait for the established stale stream to be dropped");
19387
19388        assert_eq!(content, "Committed technical response.");
19389        assert_eq!(agent.current_state().as_deref(), Some("technical"));
19390        assert!(stream_dropped.load(Ordering::SeqCst));
19391        assert!(committed_after_drop.load(Ordering::SeqCst));
19392    }
19393
19394    #[tokio::test]
19395    async fn test_buffered_streaming_transition_reservation_falls_back() {
19396        use futures::StreamExt;
19397
19398        let mock = mock_with_responses(vec![
19399            "Serial streaming response",
19400            "Serial streaming response",
19401        ]);
19402        let router_mock = mock_with_response("1");
19403        let router_counter = router_mock.clone();
19404        let yaml = r#"
19405name: BufferedReservationFallbackAgent
19406system_prompt: "Stream normally if speculative routing cannot be evaluated."
19407llm:
19408  default: default
19409  router: router
19410observability:
19411  enabled: true
19412  export:
19413    write_raw_events: true
19414streaming:
19415  enabled: true
19416  buffer_size: 8
19417runtime:
19418  optimization:
19419    enabled: true
19420    max_speculative_llm_calls_per_turn: 1
19421    speculative_state_transitions: true
19422    streaming_policy: buffer_until_routing_done
19423    max_parallel_runtime_tasks: 2
19424states:
19425  initial: triage
19426  states:
19427    triage:
19428      prompt: "Triage state."
19429      transitions:
19430        - to: billing
19431          guard:
19432            context:
19433              route:
19434                eq: billing
19435          when: "User asks about billing"
19436          timing: parallel
19437    billing:
19438      prompt: "Billing state."
19439"#;
19440        let agent = AgentBuilder::from_yaml(yaml)
19441            .unwrap()
19442            .llm_alias("default", Arc::new(mock))
19443            .llm_alias("router", Arc::new(router_mock))
19444            .build()
19445            .unwrap();
19446
19447        let mut stream = agent.chat_stream("hello").await.unwrap();
19448        let mut content = String::new();
19449        let mut error = None;
19450        while let Some(chunk) = stream.next().await {
19451            match chunk {
19452                StreamChunk::Content { text } => content.push_str(&text),
19453                StreamChunk::Error { message } => error = Some(message),
19454                StreamChunk::Done {} => break,
19455                _ => {}
19456            }
19457        }
19458
19459        assert_eq!(error, None);
19460        assert_eq!(content, "Serial streaming response");
19461        assert_eq!(router_counter.call_count(), 0);
19462        let events = agent.observability().unwrap().raw_events();
19463        assert!(events.iter().any(|event| {
19464            event.dimensions.get("branch_status") == Some(&"cancelled".to_string())
19465                && event.dimensions.get("commit_behavior")
19466                    == Some(&"transition_decision".to_string())
19467        }));
19468    }
19469
19470    #[tokio::test]
19471    async fn test_blocking_error_cleanup_resets_root_turn_for_next_chat() {
19472        let mut mock = mock_with_response("Recovered response");
19473        mock.set_error("boom");
19474        let mut handle = mock.clone();
19475        let agent = AgentBuilder::new()
19476            .system_prompt("You are helpful.")
19477            .llm(Arc::new(mock))
19478            .build()
19479            .unwrap();
19480
19481        assert!(agent.chat("first").await.is_err());
19482        handle.clear_error();
19483        let response = agent.chat("second").await.unwrap();
19484
19485        assert_eq!(response.content, "Recovered response");
19486        let messages = agent.memory.get_messages(None).await.unwrap();
19487        let user_count = messages
19488            .iter()
19489            .filter(|message| message.role == ai_agents_core::Role::User)
19490            .count();
19491        assert_eq!(user_count, 2);
19492    }
19493
19494    #[tokio::test]
19495    async fn test_streaming_error_cleanup_resets_root_turn_for_next_chat() {
19496        use futures::StreamExt;
19497
19498        let mut mock = mock_with_response("Recovered response");
19499        mock.set_error("stream boom");
19500        let mut handle = mock.clone();
19501        let agent = AgentBuilder::new()
19502            .system_prompt("You are helpful.")
19503            .llm(Arc::new(mock))
19504            .build()
19505            .unwrap();
19506
19507        let mut stream = agent.chat_stream("first").await.unwrap();
19508        let mut saw_error = false;
19509        while let Some(chunk) = stream.next().await {
19510            if matches!(chunk, StreamChunk::Error { .. }) {
19511                saw_error = true;
19512            }
19513        }
19514        assert!(saw_error);
19515
19516        handle.clear_error();
19517        let response = agent.chat("second").await.unwrap();
19518
19519        assert_eq!(response.content, "Recovered response");
19520        let messages = agent.memory.get_messages(None).await.unwrap();
19521        let user_count = messages
19522            .iter()
19523            .filter(|message| message.role == ai_agents_core::Role::User)
19524            .count();
19525        assert_eq!(user_count, 2);
19526    }
19527
19528    #[tokio::test]
19529    async fn test_buffered_streaming_route_miss_releases_buffer_limit() {
19530        use futures::StreamExt;
19531
19532        let mut mock = mock_with_response("one two three");
19533        mock.set_latency(10);
19534        let yaml = r#"
19535name: BufferedMissAgent
19536system_prompt: "You stream safely."
19537llm:
19538  default: default
19539streaming:
19540  enabled: true
19541  buffer_size: 1
19542runtime:
19543  optimization:
19544    enabled: true
19545    max_speculative_llm_calls_per_turn: 2
19546    speculative_state_transitions: true
19547    streaming_policy: buffer_until_routing_done
19548    max_parallel_runtime_tasks: 2
19549states:
19550  initial: triage
19551  states:
19552    triage:
19553      prompt: "Answer from triage."
19554      transitions:
19555        - to: billing
19556          guard:
19557            context:
19558              route:
19559                eq: billing
19560          timing: parallel
19561    billing:
19562      prompt: "Billing state."
19563"#;
19564        let agent = AgentBuilder::from_yaml(yaml)
19565            .unwrap()
19566            .llm_alias("default", Arc::new(mock))
19567            .build()
19568            .unwrap();
19569
19570        let mut stream = agent.chat_stream("hello").await.unwrap();
19571        let mut content = String::new();
19572        let mut error = None;
19573        while let Some(chunk) = stream.next().await {
19574            match chunk {
19575                StreamChunk::Content { text } => content.push_str(&text),
19576                StreamChunk::Error { message } => error = Some(message),
19577                StreamChunk::Done {} => break,
19578                _ => {}
19579            }
19580        }
19581
19582        assert_eq!(error, None);
19583        assert_eq!(content, "one two three");
19584    }
19585
19586    #[tokio::test]
19587    async fn test_buffered_streaming_main_failure_finalizes_branch() {
19588        use futures::StreamExt;
19589
19590        let mock = mock_with_response("one two");
19591        let mut router_mock = mock_with_response("0");
19592        router_mock.set_latency(50);
19593        let yaml = r#"
19594name: BufferedFailureAgent
19595system_prompt: "You stream safely."
19596llm:
19597  default: default
19598  router: router
19599observability:
19600  enabled: true
19601  export:
19602    write_raw_events: true
19603streaming:
19604  enabled: true
19605  buffer_size: 1
19606runtime:
19607  optimization:
19608    enabled: true
19609    max_speculative_llm_calls_per_turn: 2
19610    speculative_state_transitions: true
19611    streaming_policy: buffer_until_routing_done
19612    max_parallel_runtime_tasks: 2
19613states:
19614  initial: triage
19615  states:
19616    triage:
19617      prompt: "Ask for the category."
19618      transitions:
19619        - to: billing
19620          when: "User asks about billing"
19621          timing: parallel
19622    billing:
19623      prompt: "Billing state."
19624"#;
19625        let agent = AgentBuilder::from_yaml(yaml)
19626            .unwrap()
19627            .llm_alias("default", Arc::new(mock))
19628            .llm_alias("router", Arc::new(router_mock))
19629            .build()
19630            .unwrap();
19631
19632        let mut stream = agent.chat_stream("hello").await.unwrap();
19633        let mut error = String::new();
19634        while let Some(chunk) = stream.next().await {
19635            if let StreamChunk::Error { message } = chunk {
19636                error = message;
19637            }
19638        }
19639
19640        assert!(
19641            error.contains("stream buffer filled"),
19642            "unexpected stream error: {}",
19643            error
19644        );
19645        let events = agent.observability().unwrap().raw_events();
19646        assert!(events.iter().any(|event| {
19647            event.dimensions.get("branch_status") == Some(&"failed".to_string())
19648                && event.dimensions.get("commit_behavior") == Some(&"final_response".to_string())
19649                && event.dimensions.get("optimization")
19650                    == Some(&"buffered_streaming_routing".to_string())
19651        }));
19652    }
19653
19654    #[tokio::test]
19655    async fn test_streaming_preflight_does_not_emit_old_state_content() {
19656        use futures::StreamExt;
19657
19658        let mock = mock_with_response("Billing streamed response");
19659        let yaml = r#"
19660name: StreamingOptimizedAgent
19661system_prompt: "You route before streaming."
19662runtime:
19663  optimization:
19664    enabled: true
19665    pre_response_deterministic_transitions: true
19666streaming:
19667  enabled: true
19668states:
19669  initial: greeting
19670  states:
19671    greeting:
19672      prompt: "OLD_STATE_SENTINEL"
19673      transitions:
19674        - to: billing
19675          guard:
19676            context:
19677              topic:
19678                eq: billing
19679          timing: pre_response
19680    billing:
19681      prompt: "Billing state."
19682"#;
19683        let agent = AgentBuilder::from_yaml(yaml)
19684            .unwrap()
19685            .llm(Arc::new(mock))
19686            .build()
19687            .unwrap();
19688        agent
19689            .set_context("topic", serde_json::json!("billing"))
19690            .unwrap();
19691
19692        let mut stream = agent.chat_stream("billing please").await.unwrap();
19693        let mut content = String::new();
19694        while let Some(chunk) = stream.next().await {
19695            match chunk {
19696                StreamChunk::Content { text } => content.push_str(&text),
19697                StreamChunk::Error { message } => panic!("stream error: {}", message),
19698                StreamChunk::Done {} => break,
19699                _ => {}
19700            }
19701        }
19702
19703        assert_eq!(agent.current_state().as_deref(), Some("billing"));
19704        assert!(content.contains("Billing streamed response"));
19705        assert!(!content.contains("OLD_STATE_SENTINEL"));
19706    }
19707
19708    // State machine transitions
19709    #[tokio::test]
19710    async fn test_integration_state_machine_basic() {
19711        let yaml = r#"
19712name: StateAgent
19713system_prompt: "You are a support agent."
19714states:
19715  initial: greeting
19716  states:
19717    greeting:
19718      prompt: "Welcome the user warmly."
19719      transitions:
19720        - to: support
19721          when: "User needs help"
19722          auto: true
19723    support:
19724      prompt: "Help solve the user's problem."
19725"#;
19726        let mock = mock_with_responses(vec![
19727            "Welcome! How can I help?", // greeting response
19728            "1",                        // transition evaluator picks first (index 0)
19729            "I'll help you with that.", // support response
19730        ]);
19731        let builder = AgentBuilder::from_yaml(yaml).unwrap();
19732        let agent = builder.llm(Arc::new(mock)).build().unwrap();
19733
19734        assert_eq!(agent.current_state(), Some("greeting".to_string()));
19735        let _ = agent.chat("I need help").await.unwrap();
19736        // After transition evaluation, state may or may not have changed
19737        // depending on mock evaluator response - the key is that it doesn't crash
19738    }
19739
19740    // State on_enter/on_exit actions
19741    #[tokio::test]
19742    async fn test_integration_state_on_enter_set_context() {
19743        let yaml = r#"
19744name: ActionAgent
19745system_prompt: "You are helpful."
19746states:
19747  initial: step1
19748  states:
19749    step1:
19750      prompt: "Step 1"
19751      on_exit:
19752        - set_context:
19753            step1_exited: true
19754      transitions:
19755        - to: step2
19756          when: "always"
19757          auto: true
19758    step2:
19759      prompt: "Step 2"
19760      on_enter:
19761        - set_context:
19762            step2_entered: true
19763"#;
19764        // The transition evaluator will pick the first transition (index 0)
19765        let mock = mock_with_responses(vec![
19766            "Processing step 1.",
19767            "0", // transition evaluator response: select first transition
19768        ]);
19769        let builder = AgentBuilder::from_yaml(yaml).unwrap();
19770        let agent = builder.llm(Arc::new(mock)).build().unwrap();
19771
19772        assert_eq!(agent.current_state(), Some("step1".to_string()));
19773
19774        // Manually transition to test on_enter/on_exit
19775        agent.transition_to("step2").await.unwrap();
19776
19777        assert_eq!(agent.current_state(), Some("step2".to_string()));
19778
19779        // Verify context was set by on_exit and on_enter actions
19780        let ctx = agent.get_context();
19781        assert_eq!(ctx.get("step1_exited"), Some(&serde_json::json!(true)));
19782        assert_eq!(ctx.get("step2_entered"), Some(&serde_json::json!(true)));
19783    }
19784
19785    #[tokio::test]
19786    async fn state_action_tool_preserves_source_in_stored_record() {
19787        let yaml = r#"
19788name: StateActionToolAgent
19789system_prompt: "You are helpful."
19790tools:
19791  - context_echo
19792states:
19793  initial: idle
19794  states:
19795    idle:
19796      prompt: "Idle"
19797    active:
19798      prompt: "Active"
19799      on_enter:
19800        - set_context:
19801            action_started: true
19802        - tool: context_echo
19803          args: {}
19804"#;
19805        let agent = AgentBuilder::from_yaml(yaml)
19806            .unwrap()
19807            .llm(Arc::new(mock_with_response("unused")))
19808            .tool(Arc::new(ContextEchoTool))
19809            .build()
19810            .unwrap();
19811
19812        agent.transition_to("active").await.unwrap();
19813
19814        let record: ToolExecutionRecord = serde_json::from_value(
19815            agent
19816                .get_context()
19817                .get("last_tool_record")
19818                .cloned()
19819                .expect("successful state action must store its execution record"),
19820        )
19821        .unwrap();
19822        assert!(record.executed);
19823        assert!(record.success);
19824        assert_eq!(record.canonical_id, "context_echo");
19825        assert!(matches!(
19826            &record.source,
19827            ToolCallSource::StateAction {
19828                state: Some(state),
19829                action_index: 1,
19830            } if state == "active"
19831        ));
19832    }
19833
19834    #[tokio::test]
19835    async fn test_ordinary_transition_uses_on_enter_then_on_reenter() {
19836        let yaml = r#"
19837name: OrdinaryLifecycleAgent
19838system_prompt: "You are helpful."
19839states:
19840  initial: intake
19841  regenerate_on_transition: false
19842  states:
19843    intake:
19844      prompt: "Intake"
19845      transitions:
19846        - to: drafting
19847          guard:
19848            context:
19849              route:
19850                eq: drafting
19851    drafting:
19852      prompt: "Drafting"
19853      on_enter:
19854        - set_context:
19855            draft_version: 1
19856      on_reenter:
19857        - set_context:
19858            draft_version: 2
19859      transitions:
19860        - to: review
19861          guard:
19862            context:
19863              route:
19864                eq: review
19865    review:
19866      prompt: "Review"
19867      on_enter:
19868        - set_context:
19869            review_entry: first
19870      transitions:
19871        - to: drafting
19872          guard:
19873            context:
19874              route:
19875                eq: drafting
19876"#;
19877        let agent = AgentBuilder::from_yaml(yaml)
19878            .unwrap()
19879            .llm(Arc::new(mock_with_responses(vec![
19880                "Intake response",
19881                "Draft response",
19882                "Review response",
19883            ])))
19884            .build()
19885            .unwrap();
19886
19887        agent
19888            .set_context("route", serde_json::json!("drafting"))
19889            .unwrap();
19890        agent.chat("Start a draft").await.unwrap();
19891        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19892        assert_eq!(
19893            agent.get_context().get("draft_version"),
19894            Some(&serde_json::json!(1))
19895        );
19896
19897        agent
19898            .set_context("route", serde_json::json!("review"))
19899            .unwrap();
19900        agent.chat("Review this").await.unwrap();
19901        assert_eq!(agent.current_state().as_deref(), Some("review"));
19902        assert_eq!(
19903            agent.get_context().get("review_entry"),
19904            Some(&serde_json::json!("first"))
19905        );
19906
19907        agent
19908            .set_context("route", serde_json::json!("drafting"))
19909            .unwrap();
19910        agent.chat("Revise this").await.unwrap();
19911        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19912        assert_eq!(
19913            agent.get_context().get("draft_version"),
19914            Some(&serde_json::json!(2))
19915        );
19916    }
19917
19918    #[tokio::test]
19919    async fn test_manual_transition_uses_on_enter_then_on_reenter() {
19920        let yaml = r#"
19921name: ManualLifecycleAgent
19922system_prompt: "You are helpful."
19923states:
19924  initial: intake
19925  states:
19926    intake:
19927      prompt: "Intake"
19928    drafting:
19929      prompt: "Drafting"
19930      on_enter:
19931        - set_context:
19932            draft_version: 1
19933      on_reenter:
19934        - set_context:
19935            draft_version: 2
19936    review:
19937      prompt: "Review"
19938"#;
19939        let agent = AgentBuilder::from_yaml(yaml)
19940            .unwrap()
19941            .llm(Arc::new(mock_with_response("unused")))
19942            .build()
19943            .unwrap();
19944
19945        assert!(!agent.get_context().contains_key("draft_version"));
19946        agent.transition_to("drafting").await.unwrap();
19947        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19948        assert_eq!(
19949            agent.get_context().get("draft_version"),
19950            Some(&serde_json::json!(1))
19951        );
19952
19953        agent.transition_to("review").await.unwrap();
19954        agent.transition_to("drafting").await.unwrap();
19955        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19956        assert_eq!(
19957            agent.get_context().get("draft_version"),
19958            Some(&serde_json::json!(2))
19959        );
19960    }
19961
19962    #[tokio::test]
19963    async fn test_timeout_transition_uses_on_enter_then_on_reenter() {
19964        let yaml = r#"
19965name: TimeoutLifecycleAgent
19966system_prompt: "You are helpful."
19967states:
19968  initial: intake
19969  regenerate_on_transition: false
19970  states:
19971    intake:
19972      prompt: "Intake"
19973      max_turns: 1
19974      timeout_to: drafting
19975    drafting:
19976      prompt: "Drafting"
19977      max_turns: 1
19978      timeout_to: review
19979      on_enter:
19980        - set_context:
19981            draft_version: 1
19982      on_reenter:
19983        - set_context:
19984            draft_version: 2
19985    review:
19986      prompt: "Review"
19987      max_turns: 1
19988      timeout_to: drafting
19989      on_enter:
19990        - set_context:
19991            review_entry: first
19992"#;
19993        let agent = AgentBuilder::from_yaml(yaml)
19994            .unwrap()
19995            .llm(Arc::new(mock_with_responses(vec![
19996                "Intake",
19997                "First draft",
19998                "Review",
19999                "Revised draft",
20000            ])))
20001            .build()
20002            .unwrap();
20003
20004        agent.chat("First turn").await.unwrap();
20005        assert_eq!(agent.current_state().as_deref(), Some("intake"));
20006        assert!(!agent.get_context().contains_key("draft_version"));
20007
20008        agent.chat("Second turn").await.unwrap();
20009        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
20010        assert_eq!(
20011            agent.get_context().get("draft_version"),
20012            Some(&serde_json::json!(1))
20013        );
20014
20015        agent.chat("Third turn").await.unwrap();
20016        assert_eq!(agent.current_state().as_deref(), Some("review"));
20017        assert_eq!(
20018            agent.get_context().get("review_entry"),
20019            Some(&serde_json::json!("first"))
20020        );
20021
20022        agent.chat("Fourth turn").await.unwrap();
20023        assert_eq!(agent.current_state().as_deref(), Some("drafting"));
20024        assert_eq!(
20025            agent.get_context().get("draft_version"),
20026            Some(&serde_json::json!(2))
20027        );
20028    }
20029
20030    // Process pipeline transforms input
20031    #[tokio::test]
20032    async fn test_integration_process_normalize() {
20033        let yaml = r#"
20034name: ProcessAgent
20035system_prompt: "You are helpful."
20036process:
20037  input:
20038    - type: normalize
20039      config:
20040        trim: true
20041        collapse_whitespace: true
20042"#;
20043        let mock = mock_with_response("Got your message.");
20044        let builder = AgentBuilder::from_yaml(yaml).unwrap();
20045        let agent = builder.llm(Arc::new(mock.clone())).build().unwrap();
20046
20047        let _ = agent.chat("  hello   world  ").await.unwrap();
20048
20049        // Verify the LLM received the normalized input (trimmed + collapsed whitespace)
20050        let history = mock.call_history();
20051        assert!(!history.is_empty());
20052        // The user message in LLM call should be normalized
20053        let last_call = history.last().unwrap();
20054        let user_msg = last_call
20055            .messages
20056            .iter()
20057            .find(|m| m.role == ai_agents_core::Role::User)
20058            .unwrap();
20059        assert_eq!(user_msg.content, "hello world");
20060    }
20061
20062    // ═══════════════════════════════════════════════════════════
20063    // Integration Test 2.1.7: Memory compression triggers
20064    // ═══════════════════════════════════════════════════════════
20065    #[tokio::test]
20066    async fn test_integration_memory_compression() {
20067        let yaml = r#"
20068name: MemoryAgent
20069system_prompt: "You are helpful."
20070memory:
20071  type: compacting
20072  max_messages: 100
20073  compress_threshold: 5
20074  max_recent_messages: 3
20075  summarize_batch_size: 2
20076"#;
20077        // Provide enough responses for compression to trigger
20078        let responses: Vec<&str> = (0..8).map(|_| "Response from assistant.").collect();
20079        let mock = mock_with_responses(responses);
20080        let builder = AgentBuilder::from_yaml(yaml).unwrap();
20081        let agent = builder.llm(Arc::new(mock)).build().unwrap();
20082
20083        // Send enough messages to trigger compression
20084        for i in 0..6 {
20085            let _ = agent.chat(&format!("Message {}", i)).await.unwrap();
20086        }
20087
20088        // Memory should have compressed - verify it didn't crash
20089        // and that messages are bounded
20090        let messages = agent.memory.get_messages(None).await.unwrap();
20091        // With compress_threshold=5 and max_recent_messages=3,
20092        // after 6 turns (12 messages), compression should have run
20093        assert!(messages.len() <= 12); // At most all messages if no compression, fewer if compressed
20094    }
20095
20096    // YAML with multiple LLMs
20097    #[tokio::test]
20098    async fn test_integration_multi_llm_registry() {
20099        let mut mock_default = MockLLMProvider::new("default");
20100        mock_default.set_response("Default LLM response.");
20101        let mut mock_router = MockLLMProvider::new("router");
20102        mock_router.set_response("Router response.");
20103
20104        let agent = AgentBuilder::new()
20105            .system_prompt("You are helpful.")
20106            .llm_alias("default", Arc::new(mock_default))
20107            .llm_alias("router", Arc::new(mock_router))
20108            .build()
20109            .unwrap();
20110
20111        let response = agent.chat("Hello").await.unwrap();
20112        assert_eq!(response.content, "Default LLM response.");
20113    }
20114
20115    // Agent reset clears state
20116    #[tokio::test]
20117    async fn test_integration_agent_reset() {
20118        let mock = mock_with_responses(vec!["Hello!", "Hello again!"]);
20119        let agent = AgentBuilder::new()
20120            .system_prompt("You are helpful.")
20121            .llm(Arc::new(mock))
20122            .build()
20123            .unwrap();
20124
20125        let _ = agent.chat("Hi").await.unwrap();
20126        let messages = agent.memory.get_messages(None).await.unwrap();
20127        assert_eq!(messages.len(), 2); // user + assistant
20128
20129        agent.reset().await.unwrap();
20130        let messages = agent.memory.get_messages(None).await.unwrap();
20131        assert_eq!(messages.len(), 0);
20132    }
20133
20134    // Process pipeline rejects input
20135    #[tokio::test]
20136    async fn test_integration_process_validate_reject() {
20137        use ai_agents_process::{ProcessConfig, ProcessProcessor};
20138
20139        let validate_config = ai_agents_process::ValidateStage {
20140            id: Some("length_check".to_string()),
20141            condition: None,
20142            config: ai_agents_process::ValidateConfig {
20143                rules: vec![ai_agents_process::ValidationRule::MinLength {
20144                    min_length: 10,
20145                    on_fail: ai_agents_process::ValidationAction {
20146                        action: ai_agents_process::ValidationActionType::Reject,
20147                        message: None,
20148                    },
20149                }],
20150                ..Default::default()
20151            },
20152        };
20153        let process_config = ProcessConfig {
20154            input: vec![ai_agents_process::ProcessStage::Validate(validate_config)],
20155            ..Default::default()
20156        };
20157        let processor = ProcessProcessor::new(process_config);
20158
20159        let mock = mock_with_response("Should not reach here.");
20160        let agent = AgentBuilder::new()
20161            .system_prompt("You are helpful.")
20162            .llm(Arc::new(mock))
20163            .process_processor(processor)
20164            .build()
20165            .unwrap();
20166
20167        let response = agent.chat("Hi").await.unwrap();
20168        // Rejected input should produce a rejection response, not call LLM
20169        assert!(
20170            response.content.contains("rejected")
20171                || response.content.contains("Input rejected")
20172                || response.content.contains("too short")
20173                || response.content.contains("Too short")
20174                || response.content.len() < 50, // rejection message is typically short
20175            "Expected rejection response, got: {}",
20176            response.content
20177        );
20178    }
20179
20180    // LLM fallback: primary fails, fallback LLM responds
20181    #[tokio::test]
20182    async fn test_llm_fallback_on_failure() {
20183        use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
20184
20185        let mut primary = MockLLMProvider::new("primary");
20186        primary.set_error("Primary LLM is unavailable");
20187
20188        let mut fallback = MockLLMProvider::new("fallback");
20189        fallback.set_response("Fallback response works!");
20190
20191        let agent = AgentBuilder::new()
20192            .system_prompt("You are helpful.")
20193            .llm_alias("default", Arc::new(primary))
20194            .llm_alias("backup", Arc::new(fallback))
20195            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
20196                llm: LLMRecoveryConfig {
20197                    on_failure: LLMFailureAction::FallbackLlm {
20198                        fallback_llm: "backup".to_string(),
20199                    },
20200                    ..Default::default()
20201                },
20202                ..Default::default()
20203            }))
20204            .build()
20205            .unwrap();
20206
20207        let response = agent.chat("Hello").await.unwrap();
20208        assert!(
20209            response.content.contains("Fallback response"),
20210            "Expected fallback response, got: {}",
20211            response.content
20212        );
20213    }
20214
20215    // LLM fallback: primary fails, static message returned
20216    #[tokio::test]
20217    async fn test_llm_fallback_response_static_message() {
20218        use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
20219
20220        let mut primary = MockLLMProvider::new("primary");
20221        primary.set_error("Primary LLM is unavailable");
20222
20223        let agent = AgentBuilder::new()
20224            .system_prompt("You are helpful.")
20225            .llm(Arc::new(primary))
20226            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
20227                llm: LLMRecoveryConfig {
20228                    on_failure: LLMFailureAction::FallbackResponse {
20229                        message: "I am temporarily unavailable. Please try again later."
20230                            .to_string(),
20231                    },
20232                    ..Default::default()
20233                },
20234                ..Default::default()
20235            }))
20236            .build()
20237            .unwrap();
20238
20239        let response = agent.chat("Hello").await.unwrap();
20240        assert!(
20241            response.content.contains("temporarily unavailable"),
20242            "Expected static fallback message, got: {}",
20243            response.content
20244        );
20245    }
20246
20247    // Tool skip: tool fails, on_failure: skip absorbs the error
20248    #[tokio::test]
20249    async fn test_tool_failure_skip() {
20250        use ai_agents_recovery::{
20251            ErrorRecoveryConfig, ToolFailureAction, ToolRecoveryConfig, ToolRetryConfig,
20252        };
20253
20254        // LLM requests a nonexistent tool, then responds after seeing the skip result
20255        let mock = mock_with_responses(vec![
20256            r#"I'll use the nonexistent tool.
20257[TOOL_CALL: {"name": "nonexistent_tool", "arguments": {}}]"#,
20258            "The tool was unavailable, but I can still help you.",
20259        ]);
20260
20261        let agent = AgentBuilder::new()
20262            .system_prompt("You are helpful.")
20263            .llm(Arc::new(mock))
20264            .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
20265                tools: ToolRecoveryConfig {
20266                    default: ToolRetryConfig {
20267                        max_retries: 0,
20268                        timeout_ms: None,
20269                        on_failure: ToolFailureAction::Skip,
20270                    },
20271                    ..Default::default()
20272                },
20273                ..Default::default()
20274            }))
20275            .build()
20276            .unwrap();
20277
20278        // The tool will fail (not found), but on_failure: skip absorbs the error
20279        let response = agent.chat("Use the nonexistent tool").await;
20280        assert!(
20281            response.is_ok(),
20282            "Expected Ok with skip policy, got: {:?}",
20283            response
20284        );
20285    }
20286}