harn_vm/

step_runtime.rs

//! Per-step runtime state for `@step`-annotated persona functions.
//!
//! The compiler emits a call to the `__register_step` builtin after each
//! `@step` declaration so the runtime can dispatch on the step's metadata
//! when its function is invoked. While a step's frame is on the call
//! stack, an [`ActiveStep`] entry tracks per-step LLM usage, defaults
//! `llm_call`'s model when the call site doesn't override it, and bounds
//! cumulative token and cost spend against the step's budget.
//!
//! This module owns three thread-locals (a per-program registry, a stack
//! of currently-active steps, and a log of completed step summaries) but
//! exposes only narrow helpers — `current_active_step_*` /
//! `record_step_llm_usage` / etc. — so the call sites in
//! `crates/harn-vm/src/llm/`, `crates/harn-vm/src/vm/`, and the compiler
//! stay focused.

use std::cell::{Cell, RefCell};
use std::collections::BTreeMap;
use std::rc::Rc;

use serde::Serialize;
use serde_json::Value as JsonValue;

use crate::orchestration::{
    current_execution_policy, pop_execution_policy, push_execution_policy, CapabilityPolicy,
    HookEvent,
};
use crate::personas::StageDecl;
use crate::value::{VmClosure, VmError, VmValue};

fn vm_str(value: &VmValue) -> Option<&str> {
    match value {
        VmValue::String(s) => Some(s.as_ref()),
        _ => None,
    }
}

/// Static metadata captured from a `@step(...)` attribute.
///
/// Populated by the `__register_step` builtin (see [`register_step_from_dict`])
/// when the program first runs, then consulted by `llm_call` and the
/// frame-pop hooks while the step is active.
#[derive(Debug, Default, Clone)]
pub struct StepDefinition {
    pub name: String,
    pub function: String,
    pub model: Option<String>,
    pub max_tokens: Option<u64>,
    pub max_usd: Option<f64>,
    /// One of "fail" (default), "continue", "escalate". Drives how a
    /// `budget_exceeded` error propagating out of the step is handled —
    /// see `crates/harn-vm/src/vm/execution.rs`.
    pub error_boundary: Option<String>,
}

#[derive(Debug, Default, Clone)]
pub struct PersonaDefinition {
    pub name: String,
    /// Per-stage tool/side-effect scoping. Keyed lookups by stage name happen
    /// every step entry; the list is small (a handful of stages per persona)
    /// so a `Vec` keeps insertion order and matches the manifest's authored
    /// ordering.
    pub stages: Vec<StageDecl>,
}

impl StepDefinition {
    pub fn boundary(&self) -> StepErrorBoundary {
        match self.error_boundary.as_deref() {
            Some("continue") => StepErrorBoundary::Continue,
            Some("escalate") => StepErrorBoundary::Escalate,
            _ => StepErrorBoundary::Fail,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StepErrorBoundary {
    Fail,
    Continue,
    Escalate,
}

/// Tracks one in-flight step. The `frame_depth` is `Vm::frames.len()`
/// captured immediately after `push_closure_frame` returns, so an
/// `ActiveStep` is "alive" while `Vm::frames.len() >= frame_depth`.
#[derive(Debug, Clone)]
pub struct ActiveStep {
    pub frame_depth: usize,
    pub definition: Rc<StepDefinition>,
    pub persona: Option<String>,
    pub args: Vec<VmValue>,
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub cost_usd: f64,
    pub llm_calls: u32,
    pub last_model: Option<String>,
    /// Tracing span id opened when the step's frame was pushed; ended on
    /// completion. 0 when tracing was disabled at push time, in which
    /// case `span_end` is a no-op anyway.
    pub span_id: u64,
    /// True when this step pushed a per-stage `CapabilityPolicy` onto the
    /// execution policy stack. The runtime pops it when the step's frame
    /// unwinds, mirroring the RAII guard pattern in
    /// `crates/harn-serve/src/adapters/acp/modes.rs`.
    pub stage_policy_pushed: bool,
}

impl ActiveStep {
    fn new(
        frame_depth: usize,
        definition: Rc<StepDefinition>,
        persona: Option<String>,
        args: Vec<VmValue>,
        span_id: u64,
        stage_policy_pushed: bool,
    ) -> Self {
        Self {
            frame_depth,
            definition,
            persona,
            args,
            input_tokens: 0,
            output_tokens: 0,
            cost_usd: 0.0,
            llm_calls: 0,
            last_model: None,
            span_id,
            stage_policy_pushed,
        }
    }

    fn total_tokens(&self) -> u64 {
        self.input_tokens.saturating_add(self.output_tokens)
    }
}

#[derive(Debug, Clone)]
pub struct ActivePersona {
    pub frame_depth: usize,
    pub definition: Rc<PersonaDefinition>,
}

/// Snapshot persisted into [`COMPLETED_STEPS`] when the step's frame
/// unwinds. Receipts and `harn persona inspect`-style downstream consumers
/// read it back via [`drain_completed_steps`].
#[derive(Debug, Clone, Serialize)]
pub struct CompletedStep {
    pub name: String,
    pub function: String,
    pub model: Option<String>,
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub cost_usd: f64,
    pub llm_calls: u32,
    pub status: String,
    pub error: Option<String>,
}

thread_local! {
    static STEP_REGISTRY: RefCell<BTreeMap<String, Rc<StepDefinition>>> =
        const { RefCell::new(BTreeMap::new()) };
    static PERSONA_REGISTRY: RefCell<BTreeMap<String, Rc<PersonaDefinition>>> =
        const { RefCell::new(BTreeMap::new()) };
    static STEP_REGISTRY_LEN: Cell<usize> = const { Cell::new(0) };
    static PERSONA_REGISTRY_LEN: Cell<usize> = const { Cell::new(0) };
    static PERSONA_STACK: RefCell<Vec<ActivePersona>> = const { RefCell::new(Vec::new()) };
    static STEP_STACK: RefCell<Vec<ActiveStep>> = const { RefCell::new(Vec::new()) };
    static COMPLETED_STEPS: RefCell<Vec<CompletedStep>> = const { RefCell::new(Vec::new()) };
    static PERSONA_HOOKS: RefCell<Vec<PersonaHookRegistration>> = const { RefCell::new(Vec::new()) };
}

/// Reset every thread-local owned by this module. Called between test
/// runs and at the start of each top-level program execution so leftover
/// registrations don't leak across runs.
pub fn reset_thread_local_state() {
    STEP_REGISTRY.with(|r| r.borrow_mut().clear());
    PERSONA_REGISTRY.with(|r| r.borrow_mut().clear());
    STEP_REGISTRY_LEN.with(|len| len.set(0));
    PERSONA_REGISTRY_LEN.with(|len| len.set(0));
    PERSONA_STACK.with(|s| s.borrow_mut().clear());
    STEP_STACK.with(|s| s.borrow_mut().clear());
    COMPLETED_STEPS.with(|c| c.borrow_mut().clear());
    PERSONA_HOOKS.with(|h| h.borrow_mut().clear());
}

#[inline]
fn step_registry_empty() -> bool {
    STEP_REGISTRY_LEN.with(|len| len.get() == 0)
}

#[inline]
fn persona_registry_empty() -> bool {
    PERSONA_REGISTRY_LEN.with(|len| len.get() == 0)
}

#[inline]
fn tracked_registries_empty() -> bool {
    step_registry_empty() && persona_registry_empty()
}

/// Bind a `@step` function name to its declared metadata. Idempotent: a
/// second call replaces the prior definition (matches re-evaluation
/// semantics of `harn run` and the conformance harness).
pub fn register_step(function: &str, definition: StepDefinition) {
    let inserted = STEP_REGISTRY.with(|registry| {
        registry
            .borrow_mut()
            .insert(function.to_string(), Rc::new(definition))
            .is_none()
    });
    if inserted {
        STEP_REGISTRY_LEN.with(|len| len.set(len.get() + 1));
    }
}

pub fn register_persona(function: &str, definition: PersonaDefinition) {
    let inserted = PERSONA_REGISTRY.with(|registry| {
        registry
            .borrow_mut()
            .insert(function.to_string(), Rc::new(definition))
            .is_none()
    });
    if inserted {
        PERSONA_REGISTRY_LEN.with(|len| len.set(len.get() + 1));
    }
}

pub fn register_persona_from_dict(args: Vec<VmValue>) -> Result<VmValue, VmError> {
    let function = args
        .first()
        .and_then(vm_str)
        .map(|s| s.to_string())
        .ok_or_else(|| {
            VmError::Thrown(VmValue::String(Rc::from(
                "__register_persona: expected (function_name, metadata_dict)",
            )))
        })?;
    let meta = args
        .get(1)
        .and_then(VmValue::as_dict)
        .cloned()
        .ok_or_else(|| {
            VmError::Thrown(VmValue::String(Rc::from(
                "__register_persona: metadata argument must be a dict",
            )))
        })?;
    let definition = PersonaDefinition {
        name: meta
            .get("name")
            .and_then(vm_str)
            .map(str::to_string)
            .unwrap_or_else(|| function.clone()),
        stages: parse_stage_decls(meta.get("stages"))?,
    };
    register_persona(&function, definition);
    Ok(VmValue::Nil)
}

fn parse_stage_decls(value: Option<&VmValue>) -> Result<Vec<StageDecl>, VmError> {
    let Some(value) = value else {
        return Ok(Vec::new());
    };
    let entries = match value {
        VmValue::Nil => return Ok(Vec::new()),
        VmValue::List(list) => list.as_ref(),
        _ => {
            return Err(VmError::Thrown(VmValue::String(Rc::from(
                "__register_persona: stages argument must be a list of dicts",
            ))));
        }
    };
    let mut out = Vec::with_capacity(entries.len());
    for entry in entries {
        let dict = entry.as_dict().ok_or_else(|| {
            VmError::Thrown(VmValue::String(Rc::from(
                "__register_persona: each stage entry must be a dict",
            )))
        })?;
        let Some(name) = dict.get("name").and_then(vm_str) else {
            return Err(VmError::Thrown(VmValue::String(Rc::from(
                "__register_persona: stage dict missing required 'name'",
            ))));
        };
        let allowed_tools = match dict.get("allowed_tools") {
            None | Some(VmValue::Nil) => None,
            Some(VmValue::List(items)) => Some(
                items
                    .iter()
                    .map(|item| {
                        vm_str(item).map(str::to_string).ok_or_else(|| {
                            VmError::Thrown(VmValue::String(Rc::from(
                                "__register_persona: stage allowed_tools entries must be strings",
                            )))
                        })
                    })
                    .collect::<Result<Vec<_>, _>>()?,
            ),
            _ => {
                return Err(VmError::Thrown(VmValue::String(Rc::from(
                    "__register_persona: stage allowed_tools must be a list of strings",
                ))));
            }
        };
        let side_effect_level = dict
            .get("side_effect_level")
            .and_then(vm_str)
            .map(str::to_string)
            .filter(|s| !s.is_empty());
        let max_iterations = match dict.get("max_iterations") {
            Some(VmValue::Int(n)) if *n >= 0 => Some(*n as u32),
            Some(VmValue::Float(f)) if f.is_finite() && *f >= 0.0 => Some(*f as u32),
            _ => None,
        };
        out.push(StageDecl {
            name: name.to_string(),
            allowed_tools,
            side_effect_level,
            max_iterations,
            on_exit: None,
        });
    }
    Ok(out)
}

/// Builtin entry point invoked by compiler-emitted bytecode after every
/// `@step` function declaration. Accepts a dict mirroring
/// `harn_modules::PersonaStepMetadata`.
pub fn register_step_from_dict(args: Vec<VmValue>) -> Result<VmValue, VmError> {
    let function = args
        .first()
        .and_then(vm_str)
        .map(|s| s.to_string())
        .ok_or_else(|| {
            VmError::Thrown(VmValue::String(Rc::from(
                "__register_step: expected (function_name, metadata_dict)",
            )))
        })?;
    let meta = args
        .get(1)
        .and_then(VmValue::as_dict)
        .cloned()
        .ok_or_else(|| {
            VmError::Thrown(VmValue::String(Rc::from(
                "__register_step: metadata argument must be a dict",
            )))
        })?;

    let mut definition = StepDefinition {
        function: function.clone(),
        ..StepDefinition::default()
    };
    definition.name = meta
        .get("name")
        .and_then(vm_str)
        .map(|s| s.to_string())
        .unwrap_or_else(|| function.clone());
    definition.model = meta
        .get("model")
        .and_then(vm_str)
        .map(|s| s.to_string())
        .filter(|s| !s.is_empty());
    definition.error_boundary = meta
        .get("error_boundary")
        .and_then(vm_str)
        .map(|s| s.to_string());

    if let Some(VmValue::Dict(budget)) = meta.get("budget") {
        if let Some(value) = budget.get("max_tokens") {
            definition.max_tokens = match value {
                VmValue::Int(n) if *n > 0 => Some(*n as u64),
                VmValue::Float(f) if f.is_finite() && *f > 0.0 => Some(*f as u64),
                _ => None,
            };
        }
        if let Some(value) = budget.get("max_usd") {
            definition.max_usd = match value {
                VmValue::Float(f) if f.is_finite() && *f >= 0.0 => Some(*f),
                VmValue::Int(n) if *n >= 0 => Some(*n as f64),
                _ => None,
            };
        }
    }

    register_step(&function, definition);
    Ok(VmValue::Nil)
}

#[derive(Clone)]
pub struct PersonaHookRegistration {
    pub persona_pattern: String,
    pub step_name: Option<String>,
    pub event: HookEvent,
    pub threshold_pct: Option<f64>,
    pub handler: Rc<VmClosure>,
}

impl std::fmt::Debug for PersonaHookRegistration {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PersonaHookRegistration")
            .field("persona_pattern", &self.persona_pattern)
            .field("step_name", &self.step_name)
            .field("event", &self.event)
            .field("threshold_pct", &self.threshold_pct)
            .field("handler", &"..")
            .finish()
    }
}

#[derive(Debug, Clone)]
pub struct PersonaHookInvocation {
    pub handler: Rc<VmClosure>,
    pub event: HookEvent,
}

pub fn register_persona_hook(
    persona_pattern: impl Into<String>,
    event: HookEvent,
    threshold_pct: Option<f64>,
    handler: Rc<VmClosure>,
) {
    PERSONA_HOOKS.with(|hooks| {
        hooks.borrow_mut().push(PersonaHookRegistration {
            persona_pattern: persona_pattern.into(),
            step_name: None,
            event,
            threshold_pct,
            handler,
        });
    });
}

pub fn register_step_hook(
    persona_pattern: impl Into<String>,
    step_name: impl Into<String>,
    event: HookEvent,
    threshold_pct: Option<f64>,
    handler: Rc<VmClosure>,
) {
    PERSONA_HOOKS.with(|hooks| {
        hooks.borrow_mut().push(PersonaHookRegistration {
            persona_pattern: persona_pattern.into(),
            step_name: Some(step_name.into()),
            event,
            threshold_pct,
            handler,
        });
    });
}

pub fn clear_persona_hooks() {
    PERSONA_HOOKS.with(|hooks| hooks.borrow_mut().clear());
}

pub struct ActiveContextSnapshot {
    steps: Vec<ActiveStep>,
    personas: Vec<ActivePersona>,
}

pub fn take_active_context() -> ActiveContextSnapshot {
    ActiveContextSnapshot {
        steps: STEP_STACK.with(|stack| std::mem::take(&mut *stack.borrow_mut())),
        personas: PERSONA_STACK.with(|stack| std::mem::take(&mut *stack.borrow_mut())),
    }
}

pub fn restore_active_context(snapshot: ActiveContextSnapshot) {
    STEP_STACK.with(|stack| *stack.borrow_mut() = snapshot.steps);
    PERSONA_STACK.with(|stack| *stack.borrow_mut() = snapshot.personas);
}

pub fn is_tracked_function(function_name: &str) -> bool {
    if tracked_registries_empty() {
        return false;
    }
    (!step_registry_empty()
        && STEP_REGISTRY.with(|registry| registry.borrow().contains_key(function_name)))
        || (!persona_registry_empty()
            && PERSONA_REGISTRY.with(|registry| registry.borrow().contains_key(function_name)))
}

pub fn step_definition_for_function(function_name: &str) -> Option<Rc<StepDefinition>> {
    if step_registry_empty() {
        return None;
    }
    STEP_REGISTRY.with(|registry| registry.borrow().get(function_name).cloned())
}

pub fn current_persona_name() -> Option<String> {
    PERSONA_STACK.with(|stack| stack.borrow().last().map(|p| p.definition.name.clone()))
}

/// Resolve the per-stage policy for `step_name` against the currently
/// active persona's stage declarations. Returns `None` when no persona is
/// active or no stage matches the step name. Caller pushes the result onto
/// `EXECUTION_POLICY_STACK`.
///
/// When an ambient policy is already active, the stage policy is
/// intersected with it so a stage can only ever tighten the tool surface
/// and side-effect ceiling — never widen them.
fn stage_policy_for_active_step(step_name: &str) -> Option<CapabilityPolicy> {
    let stage_policy = PERSONA_STACK.with(|stack| {
        let stack = stack.borrow();
        let persona = stack.last()?;
        let stage = persona
            .definition
            .stages
            .iter()
            .find(|stage| stage.name == step_name)?;
        Some(stage_decl_to_policy(stage))
    })?;
    let Some(parent) = current_execution_policy() else {
        return Some(stage_policy);
    };
    // `intersect` is conservative: failure means the stage referenced a tool
    // the ambient policy already denied. Fall back to a narrowed copy that
    // drops those entries so the stage can never widen the ceiling.
    Some(parent.intersect(&stage_policy).unwrap_or_else(|_| {
        let intersected_tools: Vec<String> = stage_policy
            .tools
            .iter()
            .filter(|tool| parent.tools.is_empty() || parent.tools.contains(*tool))
            .cloned()
            .collect();
        CapabilityPolicy {
            tools: intersected_tools,
            ..stage_policy
        }
    }))
}

fn stage_decl_to_policy(stage: &StageDecl) -> CapabilityPolicy {
    CapabilityPolicy {
        tools: stage.allowed_tools.clone().unwrap_or_default(),
        side_effect_level: stage.side_effect_level.clone(),
        ..CapabilityPolicy::default()
    }
}

fn persona_matches(pattern: &str, persona: &str) -> bool {
    crate::orchestration::glob_match(pattern, persona)
}

pub fn matching_hooks(
    event: HookEvent,
    persona: Option<&str>,
    step_name: Option<&str>,
    budget_pct: Option<f64>,
) -> Vec<PersonaHookInvocation> {
    let persona = persona.unwrap_or("");
    PERSONA_HOOKS.with(|hooks| {
        hooks
            .borrow()
            .iter()
            .filter(|hook| hook.event == event)
            .filter(|hook| persona_matches(&hook.persona_pattern, persona))
            .filter(|hook| match (&hook.step_name, step_name) {
                (Some(expected), Some(actual)) => expected == actual,
                (Some(_), None) => false,
                (None, _) => true,
            })
            .filter(|hook| match (hook.threshold_pct, budget_pct) {
                (Some(threshold), Some(pct)) => pct >= threshold,
                (Some(_), None) => false,
                (None, _) => true,
            })
            .map(|hook| PersonaHookInvocation {
                handler: hook.handler.clone(),
                event: hook.event,
            })
            .collect()
    })
}

pub fn maybe_push_active_persona(function_name: &str, frame_depth: usize) -> bool {
    if persona_registry_empty() {
        return false;
    }
    let definition =
        PERSONA_REGISTRY.with(|registry| registry.borrow().get(function_name).cloned());
    let Some(definition) = definition else {
        return false;
    };
    PERSONA_STACK.with(|stack| {
        stack.borrow_mut().push(ActivePersona {
            frame_depth,
            definition,
        });
    });
    true
}

/// Push an active step onto the stack iff `function_name` has metadata
/// registered. Returns `true` when a frame was pushed so the call site
/// can record that fact. Called from `Vm::push_closure_frame` after the
/// new frame has been added.
pub fn maybe_push_active_step(function_name: &str, frame_depth: usize, args: &[VmValue]) -> bool {
    if step_registry_empty() {
        return false;
    }
    let definition = STEP_REGISTRY.with(|registry| registry.borrow().get(function_name).cloned());
    let Some(definition) = definition else {
        return false;
    };
    let persona = current_persona_name();
    let span_id =
        crate::tracing::span_start(crate::tracing::SpanKind::Step, definition.name.clone());
    if let Some(persona_name) = persona.as_deref() {
        crate::tracing::span_set_metadata(
            span_id,
            "persona",
            serde_json::Value::String(persona_name.to_string()),
        );
    }
    if let Some(model) = definition.model.as_deref() {
        crate::tracing::span_set_metadata(
            span_id,
            "model",
            serde_json::Value::String(model.to_string()),
        );
    }
    let step_name = definition.name.clone();
    STEP_STACK.with(|stack| {
        stack.borrow_mut().push(ActiveStep::new(
            frame_depth,
            definition,
            persona,
            args.to_vec(),
            span_id,
            false,
        ));
    });
    if let Some(policy) = stage_policy_for_active_step(&step_name) {
        push_execution_policy(policy);
        STEP_STACK.with(|stack| {
            if let Some(top) = stack.borrow_mut().last_mut() {
                top.stage_policy_pushed = true;
            }
        });
    }
    true
}

/// Drop any step entries whose owning frame has already been unwound,
/// recording a `CompletedStep` summary for each. The `current_frame_depth`
/// is `Vm::frames.len()` at the call site — entries with
/// `frame_depth > current_frame_depth` are stale.
pub fn prune_below_frame(current_frame_depth: usize) {
    let mut popped: Vec<ActiveStep> = Vec::new();
    STEP_STACK.with(|stack| {
        let mut stack = stack.borrow_mut();
        while let Some(top) = stack.last() {
            if top.frame_depth > current_frame_depth {
                popped.push(stack.pop().unwrap());
            } else {
                break;
            }
        }
    });
    for step in popped {
        finish_step(step, "completed", None);
    }
    PERSONA_STACK.with(|stack| {
        let mut stack = stack.borrow_mut();
        while stack
            .last()
            .is_some_and(|persona| persona.frame_depth > current_frame_depth)
        {
            stack.pop();
        }
    });
}

pub fn take_active_step(current_frame_depth: usize) -> Option<ActiveStep> {
    STEP_STACK.with(|stack| {
        let mut stack = stack.borrow_mut();
        if stack
            .last()
            .is_some_and(|step| step.frame_depth == current_frame_depth)
        {
            stack.pop()
        } else {
            None
        }
    })
}

pub fn finish_active_step(step: ActiveStep, status: &str, error: Option<String>) {
    finish_step(step, status, error);
}

/// Pop the topmost active step (if its frame is the current one) and
/// record an explicit completion status. Used when an error boundary
/// rewrites or absorbs an in-flight error so the receipt log reflects the
/// outcome the persona actually saw.
pub fn pop_and_record(current_frame_depth: usize, status: &str, error: Option<String>) -> bool {
    let popped = STEP_STACK.with(|stack| {
        let mut stack = stack.borrow_mut();
        if stack
            .last()
            .map(|step| step.frame_depth == current_frame_depth)
            .unwrap_or(false)
        {
            stack.pop()
        } else {
            None
        }
    });
    let Some(step) = popped else {
        return false;
    };
    finish_step(step, status, error);
    true
}

fn finish_step(step: ActiveStep, status: &str, error: Option<String>) {
    if step.stage_policy_pushed {
        pop_execution_policy();
    }
    crate::tracing::span_set_metadata(
        step.span_id,
        "status",
        serde_json::Value::String(status.to_string()),
    );
    crate::tracing::span_set_metadata(
        step.span_id,
        "llm_calls",
        serde_json::Value::Number(step.llm_calls.into()),
    );
    crate::tracing::span_set_metadata(
        step.span_id,
        "input_tokens",
        serde_json::Value::Number(step.input_tokens.into()),
    );
    crate::tracing::span_set_metadata(
        step.span_id,
        "output_tokens",
        serde_json::Value::Number(step.output_tokens.into()),
    );
    if let Some(cost_n) = serde_json::Number::from_f64(step.cost_usd) {
        crate::tracing::span_set_metadata(
            step.span_id,
            "cost_usd",
            serde_json::Value::Number(cost_n),
        );
    }
    crate::tracing::span_end(step.span_id);
    let summary = CompletedStep {
        name: step.definition.name.clone(),
        function: step.definition.function.clone(),
        model: step
            .last_model
            .clone()
            .or_else(|| step.definition.model.clone()),
        input_tokens: step.input_tokens,
        output_tokens: step.output_tokens,
        cost_usd: step.cost_usd,
        llm_calls: step.llm_calls,
        status: status.to_string(),
        error,
    };
    COMPLETED_STEPS.with(|completed| completed.borrow_mut().push(summary));
}

/// Get a snapshot of the topmost active step, if any. Used by the
/// llm_call path to fill in defaults — never for mutation.
pub fn with_active_step<R>(f: impl FnOnce(&ActiveStep) -> R) -> Option<R> {
    STEP_STACK.with(|stack| stack.borrow().last().map(f))
}

/// Mutate the topmost active step (typically to attribute LLM usage).
pub fn with_active_step_mut<R>(f: impl FnOnce(&mut ActiveStep) -> R) -> Option<R> {
    STEP_STACK.with(|stack| stack.borrow_mut().last_mut().map(f))
}

/// Frame depth of the topmost active step, or `None` when no step is
/// active. Used by `handle_error` to detect "this throw is exiting a
/// step's frame".
pub fn active_step_frame_depth() -> Option<usize> {
    STEP_STACK.with(|stack| stack.borrow().last().map(|s| s.frame_depth))
}

/// Default model the topmost active step should impose on `llm_call`
/// invocations whose options dict didn't pin a model.
pub fn active_step_model_default() -> Option<String> {
    STEP_STACK.with(|stack| {
        stack
            .borrow()
            .last()
            .and_then(|step| step.definition.model.clone())
    })
}

/// Record that `llm_call` consumed `input_tokens` / `output_tokens` for
/// `cost_usd`. Updates the active step's running totals and returns a
/// budget-exhaustion error if the step's ceiling is now breached.
///
/// The check is performed AFTER the call so the test fixture's first
/// call (which fits under budget) succeeds and subsequent calls trip the
/// limit. This matches the existing `accumulate_cost_for_provider`
/// pattern where global budget is also checked post-hoc.
pub fn record_step_llm_usage(
    model: &str,
    input_tokens: i64,
    output_tokens: i64,
    cost_usd: f64,
) -> Result<(), VmError> {
    let exhausted = STEP_STACK.with(|stack| -> Option<VmError> {
        let mut stack = stack.borrow_mut();
        let step = stack.last_mut()?;
        step.input_tokens = step.input_tokens.saturating_add(input_tokens.max(0) as u64);
        step.output_tokens = step
            .output_tokens
            .saturating_add(output_tokens.max(0) as u64);
        step.cost_usd += cost_usd;
        step.llm_calls = step.llm_calls.saturating_add(1);
        if !model.is_empty() {
            step.last_model = Some(model.to_string());
        }

        if let Some(max_tokens) = step.definition.max_tokens {
            if step.total_tokens() > max_tokens {
                return Some(budget_exhausted_error(
                    &step.definition,
                    "max_tokens",
                    max_tokens as f64,
                    step.total_tokens() as f64,
                    step.cost_usd,
                ));
            }
        }
        if let Some(max_usd) = step.definition.max_usd {
            if step.cost_usd > max_usd {
                return Some(budget_exhausted_error(
                    &step.definition,
                    "max_usd",
                    max_usd,
                    step.total_tokens() as f64,
                    step.cost_usd,
                ));
            }
        }
        None
    });
    if let Some(err) = exhausted {
        return Err(err);
    }
    Ok(())
}

fn budget_exhausted_error(
    definition: &StepDefinition,
    limit: &str,
    limit_value: f64,
    consumed_tokens: f64,
    consumed_cost_usd: f64,
) -> VmError {
    let mut dict: BTreeMap<String, VmValue> = BTreeMap::new();
    dict.insert(
        "category".to_string(),
        VmValue::String(Rc::from("budget_exceeded")),
    );
    dict.insert(
        "kind".to_string(),
        VmValue::String(Rc::from("budget_exhausted")),
    );
    dict.insert(
        "reason".to_string(),
        VmValue::String(Rc::from("step_budget_exhausted")),
    );
    dict.insert(
        "step".to_string(),
        VmValue::String(Rc::from(definition.name.clone())),
    );
    dict.insert(
        "function".to_string(),
        VmValue::String(Rc::from(definition.function.clone())),
    );
    dict.insert(
        "limit".to_string(),
        VmValue::String(Rc::from(limit.to_string())),
    );
    dict.insert("limit_value".to_string(), VmValue::Float(limit_value));
    dict.insert(
        "consumed_tokens".to_string(),
        VmValue::Float(consumed_tokens),
    );
    dict.insert(
        "consumed_cost_usd".to_string(),
        VmValue::Float(consumed_cost_usd),
    );
    dict.insert(
        "error_boundary".to_string(),
        VmValue::String(Rc::from(
            definition
                .error_boundary
                .clone()
                .unwrap_or_else(|| "fail".to_string()),
        )),
    );
    dict.insert(
        "message".to_string(),
        VmValue::String(Rc::from(format!(
            "step `{}` exceeded {} budget ({} > {})",
            definition.name, limit, consumed_tokens as i64, limit_value as i64
        ))),
    );
    VmError::Thrown(VmValue::Dict(Rc::new(dict)))
}

/// Returns true if the thrown value looks like a budget-exhausted
/// error — either our typed step-budget dict or the existing
/// `crates/harn-vm/src/llm/cost.rs::budget_exceeded_error` shape.
/// Either form is treated identically by `error_boundary` because the
/// per-step budget machinery layers onto the existing envelope; a step
/// whose budget the preflight projection rejects is still a budget
/// exhaustion the step authored.
pub fn is_step_budget_exhausted(err: &VmError) -> bool {
    let VmError::Thrown(VmValue::Dict(dict)) = err else {
        return false;
    };
    let category = dict.get("category").and_then(vm_str);
    let kind = dict.get("kind").and_then(vm_str);
    let reason = dict.get("reason").and_then(vm_str);
    if matches!(kind, Some("budget_exhausted")) && matches!(reason, Some("step_budget_exhausted")) {
        return true;
    }
    matches!(category, Some("budget_exceeded"))
}

/// Annotate an existing budget-exhausted error with `escalated: true`
/// and the step's identity so the persona body / handoff receiver can
/// route on it. Returns the original error if it isn't a thrown dict.
/// Ensures `step` and `function` keys reflect the just-finished step
/// even when the underlying error was raised by the preflight budget
/// machinery (which doesn't know which step it's running under).
pub fn mark_escalated(err: VmError, step_name: Option<&str>, function: Option<&str>) -> VmError {
    let VmError::Thrown(VmValue::Dict(dict)) = err else {
        return err;
    };
    let mut next = (*dict).clone();
    next.insert("escalated".to_string(), VmValue::Bool(true));
    next.insert(
        "category".to_string(),
        VmValue::String(Rc::from("handoff_escalation")),
    );
    if let Some(step) = step_name {
        next.entry("step".to_string())
            .or_insert_with(|| VmValue::String(Rc::from(step.to_string())));
    }
    if let Some(function) = function {
        next.entry("function".to_string())
            .or_insert_with(|| VmValue::String(Rc::from(function.to_string())));
    }
    VmError::Thrown(VmValue::Dict(Rc::new(next)))
}

/// Drain the completed-step log. Used by receipt builders that want a
/// per-step model + token + cost breakdown for the just-finished run.
pub fn drain_completed_steps() -> Vec<CompletedStep> {
    COMPLETED_STEPS.with(|completed| std::mem::take(&mut *completed.borrow_mut()))
}

/// Read the completed-step log without clearing it. Use when callers
/// want a peek without disturbing the global record stream.
pub fn peek_completed_steps() -> Vec<CompletedStep> {
    COMPLETED_STEPS.with(|completed| completed.borrow().clone())
}

/// Lower a [`CompletedStep`] into JSON for embedding in receipts /
/// inspect output.
pub fn completed_step_to_json(step: &CompletedStep) -> JsonValue {
    serde_json::to_value(step).unwrap_or(JsonValue::Null)
}

/// Register the `__register_step` host builtin. Compiler-emitted
/// bytecode after every `@step` declaration calls it with
/// `(function_name, metadata_dict)` so the runtime can later dispatch on
/// the step's metadata when its function is invoked.
pub fn register_step_builtins(vm: &mut crate::vm::Vm) {
    vm.register_builtin("__register_step", |args, _out| {
        register_step_from_dict(args.to_vec())
    });
    vm.register_builtin("__register_persona", |args, _out| {
        register_persona_from_dict(args.to_vec())
    });
}

#[cfg(test)]
mod tests {
    use super::*;

    fn fresh_state() {
        reset_thread_local_state();
    }

    #[test]
    fn registers_and_pops_step_from_dict() {
        fresh_state();
        let mut budget: BTreeMap<String, VmValue> = BTreeMap::new();
        budget.insert("max_tokens".to_string(), VmValue::Int(100));
        budget.insert("max_usd".to_string(), VmValue::Float(0.05));
        let mut meta: BTreeMap<String, VmValue> = BTreeMap::new();
        meta.insert("name".to_string(), VmValue::String(Rc::from("plan")));
        meta.insert(
            "model".to_string(),
            VmValue::String(Rc::from("claude-haiku-4-5")),
        );
        meta.insert(
            "error_boundary".to_string(),
            VmValue::String(Rc::from("continue")),
        );
        meta.insert("budget".to_string(), VmValue::Dict(Rc::new(budget)));

        register_step_from_dict(vec![
            VmValue::String(Rc::from("plan_step")),
            VmValue::Dict(Rc::new(meta)),
        ])
        .expect("registration succeeds");

        assert!(maybe_push_active_step("plan_step", 3, &[]));
        assert_eq!(active_step_frame_depth(), Some(3));
        assert_eq!(
            active_step_model_default().as_deref(),
            Some("claude-haiku-4-5")
        );

        record_step_llm_usage("claude-haiku-4-5", 10, 20, 0.001).expect("under budget");
        with_active_step(|step| {
            assert_eq!(step.input_tokens, 10);
            assert_eq!(step.output_tokens, 20);
            assert!((step.cost_usd - 0.001).abs() < 1e-9);
        });

        let err =
            record_step_llm_usage("claude-haiku-4-5", 50, 50, 0.0).expect_err("should exhaust");
        assert!(is_step_budget_exhausted(&err));

        prune_below_frame(2);
        let completed = drain_completed_steps();
        assert_eq!(completed.len(), 1);
        assert_eq!(completed[0].llm_calls, 2);
    }

    #[test]
    fn unregistered_function_does_not_push() {
        fresh_state();
        assert!(!maybe_push_active_step("not_a_step", 1, &[]));
        assert!(active_step_frame_depth().is_none());
    }

    #[test]
    fn tracked_registry_empty_fast_path_tracks_registrations_and_reset() {
        fresh_state();
        assert!(tracked_registries_empty());
        assert!(!is_tracked_function("plan_step"));

        register_step(
            "plan_step",
            StepDefinition {
                name: "plan".to_string(),
                function: "plan_step".to_string(),
                ..StepDefinition::default()
            },
        );
        assert!(!tracked_registries_empty());
        assert!(is_tracked_function("plan_step"));
        assert!(step_definition_for_function("plan_step").is_some());

        register_step(
            "plan_step",
            StepDefinition {
                name: "plan_v2".to_string(),
                function: "plan_step".to_string(),
                ..StepDefinition::default()
            },
        );
        assert!(is_tracked_function("plan_step"));

        fresh_state();
        assert!(tracked_registries_empty());
        assert!(!is_tracked_function("plan_step"));
    }

    #[test]
    fn stage_policy_narrows_but_does_not_widen_parent_policy() {
        fresh_state();
        let mut meta: BTreeMap<String, VmValue> = BTreeMap::new();
        meta.insert("name".to_string(), VmValue::String(Rc::from("research")));
        register_step_from_dict(vec![
            VmValue::String(Rc::from("research_step")),
            VmValue::Dict(Rc::new(meta)),
        ])
        .expect("step registration");

        let mut stage_dict: BTreeMap<String, VmValue> = BTreeMap::new();
        stage_dict.insert("name".to_string(), VmValue::String(Rc::from("research")));
        // Stage tries to add `edit` on top of a parent that only allowed `read`.
        stage_dict.insert(
            "allowed_tools".to_string(),
            VmValue::List(Rc::new(vec![
                VmValue::String(Rc::from("read")),
                VmValue::String(Rc::from("edit")),
            ])),
        );
        let mut persona_meta: BTreeMap<String, VmValue> = BTreeMap::new();
        persona_meta.insert("name".to_string(), VmValue::String(Rc::from("scoped")));
        persona_meta.insert(
            "stages".to_string(),
            VmValue::List(Rc::new(vec![VmValue::Dict(Rc::new(stage_dict))])),
        );
        register_persona_from_dict(vec![
            VmValue::String(Rc::from("scoped_persona")),
            VmValue::Dict(Rc::new(persona_meta)),
        ])
        .expect("persona registration");

        push_execution_policy(CapabilityPolicy {
            tools: vec!["read".to_string()],
            ..CapabilityPolicy::default()
        });
        assert!(maybe_push_active_persona("scoped_persona", 1));
        assert!(maybe_push_active_step("research_step", 2, &[]));
        let policy = current_execution_policy().expect("stage policy active");
        // `edit` is filtered out because the parent already denied it.
        assert_eq!(policy.tools, vec!["read".to_string()]);

        prune_below_frame(0);
        pop_execution_policy();
        assert!(current_execution_policy().is_none());
    }

    #[test]
    fn stage_policy_is_pushed_and_popped_around_step() {
        fresh_state();
        let mut meta: BTreeMap<String, VmValue> = BTreeMap::new();
        meta.insert("name".to_string(), VmValue::String(Rc::from("research")));
        register_step_from_dict(vec![
            VmValue::String(Rc::from("research_step")),
            VmValue::Dict(Rc::new(meta)),
        ])
        .expect("step registration succeeds");

        let mut stage_dict: BTreeMap<String, VmValue> = BTreeMap::new();
        stage_dict.insert("name".to_string(), VmValue::String(Rc::from("research")));
        stage_dict.insert(
            "allowed_tools".to_string(),
            VmValue::List(Rc::new(vec![VmValue::String(Rc::from("read"))])),
        );
        let mut persona_meta: BTreeMap<String, VmValue> = BTreeMap::new();
        persona_meta.insert("name".to_string(), VmValue::String(Rc::from("scoped")));
        persona_meta.insert(
            "stages".to_string(),
            VmValue::List(Rc::new(vec![VmValue::Dict(Rc::new(stage_dict))])),
        );
        register_persona_from_dict(vec![
            VmValue::String(Rc::from("scoped_persona")),
            VmValue::Dict(Rc::new(persona_meta)),
        ])
        .expect("persona registration succeeds");

        assert!(maybe_push_active_persona("scoped_persona", 1));
        assert!(crate::orchestration::current_execution_policy().is_none());
        assert!(maybe_push_active_step("research_step", 2, &[]));
        let policy = crate::orchestration::current_execution_policy()
            .expect("stage policy is active inside step");
        assert_eq!(policy.tools, vec!["read".to_string()]);

        prune_below_frame(0);
        assert!(crate::orchestration::current_execution_policy().is_none());
    }
}