harn-vm 0.10.168

//! Ambient state inherited specifically by child-interpreter subtasks.

use std::sync::Arc;

use super::clone_via_swap;
use crate::call_budget::CallBudget;
use crate::event_log::AnyEventLog;
use crate::redact::NamedPattern;
use crate::security::SecurityPolicy;
use crate::vm::subtask::{swap_subtask_placement_context, SubtaskPlacement};

impl super::AmbientExecutionScope {
    pub(crate) fn with_prepared_identity(
        mut self,
        context: crate::prepared_run::PreparedIdentityContext,
    ) -> Self {
        self.prepared_identity = Some(Arc::new(context));
        self
    }
}

/// Mechanically checked projection of `thread_local_audit.toml`'s
/// `subtask_scope.captured` policy. Keep this beside the typed state that
/// implements the policy; the whole-crate audit test rejects drift either way.
#[cfg(test)]
pub(super) const CAPTURED_THREAD_LOCAL_PROJECTION: &[&str] = &[
    "EXECUTION_POLICY_STACK",
    "EXECUTION_APPROVAL_POLICY_STACK",
    "PREPARED_APPROVAL_POLICY",
    "PREPARED_IDENTITY_CONTEXT",
    "OPERATOR_APPROVAL_GRANT_STACK",
    "COMMAND_POLICY_STACK",
    "DYNAMIC_PERMISSION_STACK",
    "RUNTIME_CONTEXT_OVERLAY_STACK",
    "AUTONOMY_POLICY_STACK",
    "PERSONA_STACK",
    "STEP_STACK",
    "ACTIVE_CONTEXT_SUSPENSION_STACK",
    "LLM_RENDER_STACK",
    "TRANSCRIPT_DIR_STACK",
    "ACTIVE_HARN_CONNECTOR_CTX",
    "TRUSTED_BRIDGE_CALL_DEPTH",
    "COMMAND_POLICY_HOOK_DEPTH",
    "TOOL_PRECHECK_STACK",
    "TOOL_PRECHECK_DEPTH",
    "APPROVAL_REVIEWER_STACK",
    "APPROVAL_REVIEWER_DEPTH",
    "VM_EXECUTION_CONTEXT",
    "VM_SOURCE_DIR",
    "SCOPED_PERSISTENT_STATE_ROOT",
    "CURRENT_MUTATION_SESSION",
    "SESSION_ENVIRONMENT_CONTEXT",
    "PROCESS_ADMISSION_CONTEXT",
    "CURRENT_HOST_BRIDGE",
    "CURRENT_LOOP_SINKS",
    "CURRENT_AGENT_EVENT_TRANSPORT",
    "CURRENT_SESSION_STACK",
    "LLM_CONFIG_OVERRIDES_CONTEXT",
    "LLM_RUNTIME_PROVIDER_ENDPOINTS_CONTEXT",
    "LLM_CAPABILITY_OVERRIDES_CONTEXT",
    "LLM_MOCK_CONTEXT",
    "SCOPE",
    "EGRESS_POLICY_CONTEXT",
    "AMBIENT_BOUNDARY",
    "ACTIVE_EXECUTION_SCOPE_STACK",
    "RUN_EVENT_SINK_CONTEXT",
    "SUBTASK_PLACEMENT_CONTEXT",
    "SECURITY_POLICY_STACK",
    "REQUIRE_EXPLICIT_EGRESS_POLICY_DEPTH",
    "REQUIRE_SSRF_GUARD_DEPTH",
    "CUSTOM_PATTERNS",
    "ACTIVE_EVENT_LOG",
    "MCP_CALL_BUDGET",
    "PG_QUERY_BUDGET",
    "LLM_COST_BUDGET",
    "LLM_TOKEN_BUDGET",
    "ACTIVE_TOOL_CALL_CANCELLATION_REGISTRY",
    "ACTIVE_WORKER_REGISTRY",
    "ACTIVE_DAEMON_REGISTRY",
    "ACTIVE_TRIGGER_REGISTRY",
    "ACTIVE_SESSION_RUNTIME",
    "ACTIVE_TRACING_RUNTIME",
    "ACTIVE_AGENT_HOST_SESSION_RUNTIME",
];

/// One structural contract for state that must cross a subtask thread boundary.
///
/// Absence is part of each slot's explicit contract: for example, `None`
/// means no dispatch call budget or event log is installed. Keeping the slots
/// together prevents a new child-interpreter path from capturing only a subset
/// of the security, audit, placement, and shared-budget state.
#[derive(Clone, Default)]
pub(super) struct SubtaskAmbientState {
    placement: Option<SubtaskPlacement>,
    explicit_egress_policy_depth: usize,
    ssrf_guard_depth: usize,
    security_policy: Vec<SecurityPolicy>,
    redaction_patterns: Vec<NamedPattern>,
    event_log: Option<Arc<AnyEventLog>>,
    mcp_call_budget: Option<CallBudget>,
    pg_query_budget: Option<CallBudget>,
    llm_cost_budget: Option<crate::llm::cost::budget::LlmCostBudget>,
    llm_token_budget: Option<crate::llm::cost::budget::LlmTokenBudget>,
    tool_call_cancellations: Arc<crate::tool_call_cancellations::CancellationRegistry>,
    worker_registry: Arc<crate::stdlib::agents::agents_workers::WorkerRegistry>,
    daemon_registry: Arc<crate::stdlib::agents_daemon::DaemonRegistry>,
    trigger_registry: Arc<crate::triggers::registry::TriggerRegistryRuntime>,
    session_runtime: Arc<crate::agent_sessions::AgentSessionRuntime>,
    tracing_runtime: Arc<crate::tracing::TracingRuntime>,
    agent_host_session_runtime: Arc<crate::llm::agent_session_host::AgentHostSessionRuntime>,
}

impl SubtaskAmbientState {
    pub(super) fn capture() -> Self {
        Self {
            placement: clone_via_swap(swap_subtask_placement_context),
            explicit_egress_policy_depth: clone_via_swap(
                crate::egress::swap_require_explicit_egress_policy_depth,
            ),
            ssrf_guard_depth: clone_via_swap(crate::egress::swap_require_ssrf_guard_depth),
            security_policy: clone_via_swap(crate::security::swap_security_policy_stack),
            redaction_patterns: clone_via_swap(crate::redact::swap_custom_patterns),
            event_log: clone_via_swap(crate::event_log::swap_active_event_log),
            mcp_call_budget: clone_via_swap(crate::call_budget::swap_mcp_call_budget),
            pg_query_budget: clone_via_swap(crate::call_budget::swap_pg_query_budget),
            llm_cost_budget: crate::llm::cost::budget::capture_llm_cost_budget(),
            llm_token_budget: crate::llm::cost::budget::capture_llm_token_budget(),
            tool_call_cancellations: clone_via_swap(
                crate::tool_call_cancellations::swap_active_registry,
            ),
            worker_registry: clone_via_swap(
                crate::stdlib::agents::agents_workers::swap_active_worker_registry,
            ),
            daemon_registry: clone_via_swap(
                crate::stdlib::agents_daemon::swap_active_daemon_registry,
            ),
            trigger_registry: clone_via_swap(
                crate::triggers::registry::swap_active_trigger_registry,
            ),
            session_runtime: clone_via_swap(crate::agent_sessions::swap_active_session_runtime),
            tracing_runtime: clone_via_swap(crate::tracing::swap_active_tracing_runtime),
            agent_host_session_runtime: clone_via_swap(
                crate::llm::agent_session_host::swap_active_agent_host_session_runtime,
            ),
        }
    }

    pub(super) fn set_placement(&mut self, placement: Option<SubtaskPlacement>) {
        self.placement = placement;
    }

    pub(super) fn set_tool_call_cancellations(
        &mut self,
        registry: Arc<crate::tool_call_cancellations::CancellationRegistry>,
    ) {
        self.tool_call_cancellations = registry;
    }

    pub(super) fn set_worker_registry(
        &mut self,
        registry: Arc<crate::stdlib::agents::agents_workers::WorkerRegistry>,
    ) {
        self.worker_registry = registry;
    }

    pub(super) fn set_daemon_registry(
        &mut self,
        registry: Arc<crate::stdlib::agents_daemon::DaemonRegistry>,
    ) {
        self.daemon_registry = registry;
    }

    pub(super) fn set_trigger_registry(
        &mut self,
        registry: Arc<crate::triggers::registry::TriggerRegistryRuntime>,
    ) {
        self.trigger_registry = registry;
    }

    pub(super) fn set_session_runtime(
        &mut self,
        runtime: Arc<crate::agent_sessions::AgentSessionRuntime>,
    ) {
        self.session_runtime = runtime;
    }

    pub(super) fn set_tracing_runtime(&mut self, runtime: Arc<crate::tracing::TracingRuntime>) {
        self.tracing_runtime = runtime;
    }

    pub(super) fn set_agent_host_session_runtime(
        &mut self,
        runtime: Arc<crate::llm::agent_session_host::AgentHostSessionRuntime>,
    ) {
        self.agent_host_session_runtime = runtime;
    }

    pub(super) fn swap_in_place(&mut self) {
        fn swap_slot<T: Default>(slot: &mut T, swap: impl FnOnce(T) -> T) {
            *slot = swap(std::mem::take(slot));
        }

        swap_slot(&mut self.placement, swap_subtask_placement_context);
        swap_slot(
            &mut self.explicit_egress_policy_depth,
            crate::egress::swap_require_explicit_egress_policy_depth,
        );
        swap_slot(
            &mut self.ssrf_guard_depth,
            crate::egress::swap_require_ssrf_guard_depth,
        );
        swap_slot(
            &mut self.security_policy,
            crate::security::swap_security_policy_stack,
        );
        swap_slot(
            &mut self.redaction_patterns,
            crate::redact::swap_custom_patterns,
        );
        swap_slot(&mut self.event_log, crate::event_log::swap_active_event_log);
        swap_slot(
            &mut self.mcp_call_budget,
            crate::call_budget::swap_mcp_call_budget,
        );
        swap_slot(
            &mut self.pg_query_budget,
            crate::call_budget::swap_pg_query_budget,
        );
        swap_slot(
            &mut self.tool_call_cancellations,
            crate::tool_call_cancellations::swap_active_registry,
        );
        swap_slot(
            &mut self.llm_cost_budget,
            crate::llm::cost::budget::swap_llm_cost_budget,
        );
        swap_slot(
            &mut self.llm_token_budget,
            crate::llm::cost::budget::swap_llm_token_budget,
        );
        swap_slot(
            &mut self.worker_registry,
            crate::stdlib::agents::agents_workers::swap_active_worker_registry,
        );
        swap_slot(
            &mut self.daemon_registry,
            crate::stdlib::agents_daemon::swap_active_daemon_registry,
        );
        swap_slot(
            &mut self.trigger_registry,
            crate::triggers::registry::swap_active_trigger_registry,
        );
        swap_slot(
            &mut self.session_runtime,
            crate::agent_sessions::swap_active_session_runtime,
        );
        swap_slot(
            &mut self.tracing_runtime,
            crate::tracing::swap_active_tracing_runtime,
        );
        swap_slot(
            &mut self.agent_host_session_runtime,
            crate::llm::agent_session_host::swap_active_agent_host_session_runtime,
        );
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeSet;

    use serde::Deserialize;

    use super::CAPTURED_THREAD_LOCAL_PROJECTION;

    #[derive(Deserialize)]
    struct Audit {
        subtask_scope: SubtaskScope,
    }

    #[derive(Deserialize)]
    struct SubtaskScope {
        captured: Vec<String>,
    }

    #[test]
    fn capture_projection_matches_manifest_policy() {
        let audit: Audit = toml::from_str(include_str!("../../../thread_local_audit.toml"))
            .expect("parse thread-local audit");
        let policy: BTreeSet<_> = audit.subtask_scope.captured.into_iter().collect();
        let projection: BTreeSet<_> = CAPTURED_THREAD_LOCAL_PROJECTION
            .iter()
            .map(|name| (*name).to_string())
            .collect();
        assert_eq!(projection, policy);
    }

    #[tokio::test(flavor = "current_thread")]
    async fn captured_child_scope_retains_transport_after_same_session_rebind() {
        use crate::agent_events::{
            emit_event, register_sink, AgentEvent, AgentEventSink, AgentEventTransport,
        };
        use std::sync::atomic::{AtomicUsize, Ordering};
        use std::sync::Arc;

        struct Count(Arc<AtomicUsize>);
        impl AgentEventSink for Count {
            fn handle_event(&self, _: &AgentEvent) {
                self.0.fetch_add(1, Ordering::SeqCst);
            }
        }
        let transport = || {
            let count = Arc::new(AtomicUsize::new(0));
            (
                AgentEventTransport::new(Arc::new(Count(count.clone()))),
                count,
            )
        };
        let session = format!("transport-child-{}", uuid::Uuid::now_v7());
        let event = AgentEvent::IterationStart {
            session_id: session.clone(),
            iteration: 1,
            provider: String::new(),
            model: String::new(),
        };
        let (old, old_count) = transport();
        let (new, new_count) = transport();
        let durable = Arc::new(AtomicUsize::new(0));
        register_sink(session.clone(), Arc::new(Count(durable.clone())));
        let old_scope = old
            .scope(async { crate::orchestration::AmbientExecutionScope::capture_inherited() })
            .await;
        new.scope(async {
            emit_event(&event);
            crate::orchestration::scope_ambient(old_scope, async {
                emit_event(&event);
                tokio::task::yield_now().await;
                emit_event(&event);
            })
            .await;
            emit_event(&event);
        })
        .await;
        crate::agent_events::clear_session_sinks(&session);
        assert_eq!(
            (
                old_count.load(Ordering::SeqCst),
                new_count.load(Ordering::SeqCst)
            ),
            (2, 2)
        );
        assert_eq!(durable.load(Ordering::SeqCst), 4);
    }
}