Skip to main content

harn_vm/
step_runtime.rs

1//! Per-step runtime state for `@step`-annotated persona functions.
2//!
3//! The compiler emits a call to the `__register_step` builtin after each
4//! `@step` declaration so the runtime can dispatch on the step's metadata
5//! when its function is invoked. While a step's frame is on the call
6//! stack, an [`ActiveStep`] entry tracks per-step LLM usage, defaults
7//! `llm_call`'s model when the call site doesn't override it, and bounds
8//! cumulative token and cost spend against the step's budget.
9//!
10//! This module owns three thread-locals (a per-program registry, a stack
11//! of currently-active steps, and a log of completed step summaries) but
12//! exposes only narrow helpers — `current_active_step_*` /
13//! `record_step_llm_usage` / etc. — so the call sites in
14//! `crates/harn-vm/src/llm/`, `crates/harn-vm/src/vm/`, and the compiler
15//! stay focused.
16
17use crate::value::VmDictExt;
18use std::cell::{Cell, RefCell};
19use std::collections::BTreeMap;
20use std::sync::atomic::{AtomicU64, Ordering};
21use std::sync::Arc;
22
23use serde::Serialize;
24use serde_json::Value as JsonValue;
25
26use crate::orchestration::{
27    current_execution_policy, pop_execution_policy, push_execution_policy, CapabilityPolicy,
28    HookEvent,
29};
30use crate::personas::StageDecl;
31use crate::stdlib::macros::{harn_builtin, VmBuiltinDef};
32use crate::value::{VmClosure, VmError, VmValue};
33
34fn vm_str(value: &VmValue) -> Option<&str> {
35    match value {
36        VmValue::String(s) => Some(s.as_ref()),
37        _ => None,
38    }
39}
40
41/// Static metadata captured from a `@step(...)` attribute.
42///
43/// Populated by the `__register_step` builtin (see [`register_step_from_dict`])
44/// when the program first runs, then consulted by `llm_call` and the
45/// frame-pop hooks while the step is active.
46#[derive(Debug, Default, Clone)]
47pub struct StepDefinition {
48    pub name: String,
49    pub function: String,
50    pub model: Option<String>,
51    pub max_tokens: Option<u64>,
52    pub max_usd: Option<f64>,
53    /// One of "fail" (default), "continue", "escalate". Drives how a
54    /// `budget_exceeded` error propagating out of the step is handled —
55    /// see `crates/harn-vm/src/vm/execution.rs`.
56    pub error_boundary: Option<String>,
57}
58
59#[derive(Debug, Default, Clone)]
60pub struct PersonaDefinition {
61    pub name: String,
62    /// Per-stage tool/side-effect scoping. Keyed lookups by stage name happen
63    /// every step entry; the list is small (a handful of stages per persona)
64    /// so a `Vec` keeps insertion order and matches the manifest's authored
65    /// ordering.
66    pub stages: Vec<StageDecl>,
67    /// The persona's declared output style (how it shapes its prose), surfaced
68    /// to Harn via `persona_output_style()`. `None` when the persona declares
69    /// no style.
70    pub output_style: Option<harn_modules::personas::PersonaOutputStyle>,
71}
72
73impl StepDefinition {
74    pub fn boundary(&self) -> StepErrorBoundary {
75        match self.error_boundary.as_deref() {
76            Some("continue") => StepErrorBoundary::Continue,
77            Some("escalate") => StepErrorBoundary::Escalate,
78            _ => StepErrorBoundary::Fail,
79        }
80    }
81}
82
83#[derive(Debug, Clone, Copy, PartialEq, Eq)]
84pub enum StepErrorBoundary {
85    Fail,
86    Continue,
87    Escalate,
88}
89
90/// Tracks one in-flight step. The `frame_depth` is `Vm::frames.len()`
91/// captured immediately after `push_closure_frame` returns, so an
92/// `ActiveStep` is "alive" while `Vm::frames.len() >= frame_depth`.
93#[derive(Debug, Clone)]
94pub struct ActiveStep {
95    pub frame_depth: usize,
96    pub definition: Arc<StepDefinition>,
97    pub persona: Option<String>,
98    pub args: Vec<VmValue>,
99    pub input_tokens: u64,
100    pub output_tokens: u64,
101    pub cost_usd: f64,
102    pub llm_calls: u32,
103    pub last_model: Option<String>,
104    /// Tracing span id opened when the step's frame was pushed; ended on
105    /// completion. 0 when tracing was disabled at push time, in which
106    /// case `span_end` is a no-op anyway.
107    pub span_id: u64,
108    /// True when this step pushed a per-stage `CapabilityPolicy` onto the
109    /// execution policy stack. The runtime pops it when the step's frame
110    /// unwinds, mirroring the RAII guard pattern in
111    /// `crates/harn-serve/src/adapters/acp/modes.rs`.
112    pub stage_policy_pushed: bool,
113}
114
115impl ActiveStep {
116    fn new(
117        frame_depth: usize,
118        definition: Arc<StepDefinition>,
119        persona: Option<String>,
120        args: Vec<VmValue>,
121        span_id: u64,
122        stage_policy_pushed: bool,
123    ) -> Self {
124        Self {
125            frame_depth,
126            definition,
127            persona,
128            args,
129            input_tokens: 0,
130            output_tokens: 0,
131            cost_usd: 0.0,
132            llm_calls: 0,
133            last_model: None,
134            span_id,
135            stage_policy_pushed,
136        }
137    }
138
139    fn total_tokens(&self) -> u64 {
140        self.input_tokens.saturating_add(self.output_tokens)
141    }
142}
143
144#[derive(Debug, Clone)]
145pub struct ActivePersona {
146    pub frame_depth: usize,
147    pub definition: Arc<PersonaDefinition>,
148}
149
150/// Snapshot persisted into [`COMPLETED_STEPS`] when the step's frame
151/// unwinds. Receipts and `harn persona inspect`-style downstream consumers
152/// read it back via [`drain_completed_steps`].
153#[derive(Debug, Clone, Serialize)]
154pub struct CompletedStep {
155    pub name: String,
156    pub function: String,
157    pub model: Option<String>,
158    pub input_tokens: u64,
159    pub output_tokens: u64,
160    pub cost_usd: f64,
161    pub llm_calls: u32,
162    pub status: String,
163    pub error: Option<String>,
164}
165
166thread_local! {
167    static STEP_REGISTRY: RefCell<BTreeMap<String, Arc<StepDefinition>>> =
168        const { RefCell::new(std::collections::BTreeMap::new()) };
169    static PERSONA_REGISTRY: RefCell<BTreeMap<String, Arc<PersonaDefinition>>> =
170        const { RefCell::new(std::collections::BTreeMap::new()) };
171    static STEP_REGISTRY_LEN: Cell<usize> = const { Cell::new(0) };
172    static PERSONA_REGISTRY_LEN: Cell<usize> = const { Cell::new(0) };
173    static PERSONA_STACK: RefCell<Vec<ActivePersona>> = const { RefCell::new(Vec::new()) };
174    static STEP_STACK: RefCell<Vec<ActiveStep>> = const { RefCell::new(Vec::new()) };
175    static ACTIVE_CONTEXT_SUSPENSION_STACK: RefCell<Vec<u64>> = const { RefCell::new(Vec::new()) };
176    static COMPLETED_STEPS: RefCell<Vec<CompletedStep>> = const { RefCell::new(Vec::new()) };
177    static PERSONA_HOOKS: RefCell<Vec<PersonaHookRegistration>> = const { RefCell::new(Vec::new()) };
178}
179
180/// Reset every thread-local owned by this module. Called between test
181/// runs and at the start of each top-level program execution so leftover
182/// registrations don't leak across runs.
183pub fn reset_thread_local_state() {
184    STEP_REGISTRY.with(|r| r.borrow_mut().clear());
185    PERSONA_REGISTRY.with(|r| r.borrow_mut().clear());
186    STEP_REGISTRY_LEN.with(|len| len.set(0));
187    PERSONA_REGISTRY_LEN.with(|len| len.set(0));
188    PERSONA_STACK.with(|s| s.borrow_mut().clear());
189    STEP_STACK.with(|s| s.borrow_mut().clear());
190    ACTIVE_CONTEXT_SUSPENSION_STACK.with(|s| s.borrow_mut().clear());
191    COMPLETED_STEPS.with(|c| c.borrow_mut().clear());
192    PERSONA_HOOKS.with(|h| h.borrow_mut().clear());
193}
194
195#[inline]
196fn step_registry_empty() -> bool {
197    STEP_REGISTRY_LEN.with(|len| len.get() == 0)
198}
199
200#[inline]
201fn persona_registry_empty() -> bool {
202    PERSONA_REGISTRY_LEN.with(|len| len.get() == 0)
203}
204
205#[inline]
206fn tracked_registries_empty() -> bool {
207    step_registry_empty() && persona_registry_empty()
208}
209
210/// Bind a `@step` function name to its declared metadata. Idempotent: a
211/// second call replaces the prior definition (matches re-evaluation
212/// semantics of `harn run` and the conformance harness).
213pub fn register_step(function: &str, definition: StepDefinition) {
214    let inserted = STEP_REGISTRY.with(|registry| {
215        registry
216            .borrow_mut()
217            .insert(function.to_string(), Arc::new(definition))
218            .is_none()
219    });
220    if inserted {
221        STEP_REGISTRY_LEN.with(|len| len.set(len.get() + 1));
222    }
223}
224
225pub fn register_persona(function: &str, definition: PersonaDefinition) {
226    let inserted = PERSONA_REGISTRY.with(|registry| {
227        registry
228            .borrow_mut()
229            .insert(function.to_string(), Arc::new(definition))
230            .is_none()
231    });
232    if inserted {
233        PERSONA_REGISTRY_LEN.with(|len| len.set(len.get() + 1));
234    }
235}
236
237pub fn register_persona_from_dict(args: Vec<VmValue>) -> Result<VmValue, VmError> {
238    let function = args
239        .first()
240        .and_then(vm_str)
241        .map(|s| s.to_string())
242        .ok_or_else(|| {
243            VmError::Thrown(VmValue::String(arcstr::ArcStr::from(
244                "__register_persona: expected (function_name, metadata_dict)",
245            )))
246        })?;
247    let meta = args
248        .get(1)
249        .and_then(VmValue::as_dict)
250        .cloned()
251        .ok_or_else(|| {
252            VmError::Thrown(VmValue::String(arcstr::ArcStr::from(
253                "__register_persona: metadata argument must be a dict",
254            )))
255        })?;
256    let definition = PersonaDefinition {
257        name: meta
258            .get("name")
259            .and_then(vm_str)
260            .map(str::to_string)
261            .unwrap_or_else(|| function.clone()),
262        stages: parse_stage_decls(meta.get("stages"))?,
263        output_style: parse_output_style(meta.get("output_style")),
264    };
265    register_persona(&function, definition);
266    Ok(VmValue::Nil)
267}
268
269/// Parse an `output_style` metadata value into a [`PersonaOutputStyle`].
270/// Accepts a bare string (a named style) or a dict with `name`/`instructions`.
271/// Returns `None` for nil or an empty style.
272fn parse_output_style(
273    value: Option<&VmValue>,
274) -> Option<harn_modules::personas::PersonaOutputStyle> {
275    use harn_modules::personas::PersonaOutputStyle;
276    let style = match value? {
277        VmValue::Nil => return None,
278        VmValue::String(name) => PersonaOutputStyle::from_name(name.to_string()),
279        VmValue::Dict(_) => {
280            let dict = value?.as_dict()?;
281            PersonaOutputStyle {
282                name: dict.get("name").and_then(vm_str).map(str::to_string),
283                instructions: dict
284                    .get("instructions")
285                    .and_then(vm_str)
286                    .map(str::to_string),
287            }
288        }
289        _ => return None,
290    };
291    (!style.is_empty()).then_some(style)
292}
293
294/// Build the Harn dict shape for a persona output style.
295fn output_style_to_vm(style: &harn_modules::personas::PersonaOutputStyle) -> VmValue {
296    use crate::value::{intern_key, DictMap};
297    let mut map = DictMap::new();
298    map.insert(
299        intern_key("name"),
300        style
301            .name
302            .as_deref()
303            .map(|name| VmValue::String(arcstr::ArcStr::from(name)))
304            .unwrap_or(VmValue::Nil),
305    );
306    map.insert(
307        intern_key("instructions"),
308        style
309            .instructions
310            .as_deref()
311            .map(|text| VmValue::String(arcstr::ArcStr::from(text)))
312            .unwrap_or(VmValue::Nil),
313    );
314    VmValue::dict(map)
315}
316
317/// Look up a persona's declared output style. With no argument (or nil), reads
318/// the currently-active persona (top of the persona stack); with a persona
319/// function name, reads that persona from the registry. Returns
320/// `{name, instructions}` or nil.
321pub fn persona_output_style(args: Vec<VmValue>) -> VmValue {
322    if let Some(function) = args.first().and_then(vm_str) {
323        return PERSONA_REGISTRY.with(|registry| {
324            registry
325                .borrow()
326                .get(function)
327                .and_then(|definition| definition.output_style.as_ref().map(output_style_to_vm))
328                .unwrap_or(VmValue::Nil)
329        });
330    }
331    PERSONA_STACK.with(|stack| {
332        stack
333            .borrow()
334            .last()
335            .and_then(|active| {
336                active
337                    .definition
338                    .output_style
339                    .as_ref()
340                    .map(output_style_to_vm)
341            })
342            .unwrap_or(VmValue::Nil)
343    })
344}
345
346fn parse_stage_decls(value: Option<&VmValue>) -> Result<Vec<StageDecl>, VmError> {
347    let Some(value) = value else {
348        return Ok(Vec::new());
349    };
350    let entries = match value {
351        VmValue::Nil => return Ok(Vec::new()),
352        VmValue::List(list) => list.as_ref(),
353        _ => {
354            return Err(VmError::Thrown(VmValue::String(arcstr::ArcStr::from(
355                "__register_persona: stages argument must be a list of dicts",
356            ))));
357        }
358    };
359    let mut out = Vec::with_capacity(entries.len());
360    for entry in entries {
361        let dict = entry.as_dict().ok_or_else(|| {
362            VmError::Thrown(VmValue::String(arcstr::ArcStr::from(
363                "__register_persona: each stage entry must be a dict",
364            )))
365        })?;
366        let Some(name) = dict.get("name").and_then(vm_str) else {
367            return Err(VmError::Thrown(VmValue::String(arcstr::ArcStr::from(
368                "__register_persona: stage dict missing required 'name'",
369            ))));
370        };
371        let allowed_tools = match dict.get("allowed_tools") {
372            None | Some(VmValue::Nil) => None,
373            Some(VmValue::List(items)) => Some(
374                items
375                    .iter()
376                    .map(|item| {
377                        vm_str(item).map(str::to_string).ok_or_else(|| {
378                            VmError::Thrown(VmValue::String(arcstr::ArcStr::from(
379                                "__register_persona: stage allowed_tools entries must be strings",
380                            )))
381                        })
382                    })
383                    .collect::<Result<Vec<_>, _>>()?,
384            ),
385            _ => {
386                return Err(VmError::Thrown(VmValue::String(arcstr::ArcStr::from(
387                    "__register_persona: stage allowed_tools must be a list of strings",
388                ))));
389            }
390        };
391        let side_effect_level = dict
392            .get("side_effect_level")
393            .and_then(vm_str)
394            .map(str::to_string)
395            .filter(|s| !s.is_empty());
396        let max_iterations = match dict.get("max_iterations") {
397            Some(VmValue::Int(n)) if *n >= 0 => Some(*n as u32),
398            Some(VmValue::Float(f)) if f.is_finite() && *f >= 0.0 => Some(*f as u32),
399            _ => None,
400        };
401        out.push(StageDecl {
402            name: name.to_string(),
403            allowed_tools,
404            side_effect_level,
405            max_iterations,
406            on_exit: None,
407        });
408    }
409    Ok(out)
410}
411
412/// Builtin entry point invoked by compiler-emitted bytecode after every
413/// `@step` function declaration. Accepts a dict mirroring
414/// `harn_modules::PersonaStepMetadata`.
415pub fn register_step_from_dict(args: Vec<VmValue>) -> Result<VmValue, VmError> {
416    let function = args
417        .first()
418        .and_then(vm_str)
419        .map(|s| s.to_string())
420        .ok_or_else(|| {
421            VmError::Thrown(VmValue::String(arcstr::ArcStr::from(
422                "__register_step: expected (function_name, metadata_dict)",
423            )))
424        })?;
425    let meta = args
426        .get(1)
427        .and_then(VmValue::as_dict)
428        .cloned()
429        .ok_or_else(|| {
430            VmError::Thrown(VmValue::String(arcstr::ArcStr::from(
431                "__register_step: metadata argument must be a dict",
432            )))
433        })?;
434
435    let mut definition = StepDefinition {
436        function: function.clone(),
437        ..StepDefinition::default()
438    };
439    definition.name = meta
440        .get("name")
441        .and_then(vm_str)
442        .map(|s| s.to_string())
443        .unwrap_or_else(|| function.clone());
444    definition.model = meta
445        .get("model")
446        .and_then(vm_str)
447        .map(|s| s.to_string())
448        .filter(|s| !s.is_empty());
449    definition.error_boundary = meta
450        .get("error_boundary")
451        .and_then(vm_str)
452        .map(|s| s.to_string());
453
454    if let Some(VmValue::Dict(budget)) = meta.get("budget") {
455        if let Some(value) = budget.get("max_tokens") {
456            definition.max_tokens = match value {
457                VmValue::Int(n) if *n > 0 => Some(*n as u64),
458                VmValue::Float(f) if f.is_finite() && *f > 0.0 => Some(*f as u64),
459                _ => None,
460            };
461        }
462        if let Some(value) = budget.get("max_usd") {
463            definition.max_usd = match value {
464                VmValue::Float(f) if f.is_finite() && *f >= 0.0 => Some(*f),
465                VmValue::Int(n) if *n >= 0 => Some(*n as f64),
466                _ => None,
467            };
468        }
469    }
470
471    register_step(&function, definition);
472    Ok(VmValue::Nil)
473}
474
475#[derive(Clone)]
476pub struct PersonaHookRegistration {
477    pub persona_pattern: String,
478    pub step_name: Option<String>,
479    pub event: HookEvent,
480    pub threshold_pct: Option<f64>,
481    pub handler: Arc<VmClosure>,
482}
483
484impl std::fmt::Debug for PersonaHookRegistration {
485    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
486        f.debug_struct("PersonaHookRegistration")
487            .field("persona_pattern", &self.persona_pattern)
488            .field("step_name", &self.step_name)
489            .field("event", &self.event)
490            .field("threshold_pct", &self.threshold_pct)
491            .field("handler", &"..")
492            .finish()
493    }
494}
495
496#[derive(Debug, Clone)]
497pub struct PersonaHookInvocation {
498    pub handler: Arc<VmClosure>,
499    pub event: HookEvent,
500}
501
502pub fn register_persona_hook(
503    persona_pattern: impl Into<String>,
504    event: HookEvent,
505    threshold_pct: Option<f64>,
506    handler: Arc<VmClosure>,
507) {
508    PERSONA_HOOKS.with(|hooks| {
509        hooks.borrow_mut().push(PersonaHookRegistration {
510            persona_pattern: persona_pattern.into(),
511            step_name: None,
512            event,
513            threshold_pct,
514            handler,
515        });
516    });
517}
518
519pub fn register_step_hook(
520    persona_pattern: impl Into<String>,
521    step_name: impl Into<String>,
522    event: HookEvent,
523    threshold_pct: Option<f64>,
524    handler: Arc<VmClosure>,
525) {
526    PERSONA_HOOKS.with(|hooks| {
527        hooks.borrow_mut().push(PersonaHookRegistration {
528            persona_pattern: persona_pattern.into(),
529            step_name: Some(step_name.into()),
530            event,
531            threshold_pct,
532            handler,
533        });
534    });
535}
536
537pub fn clear_persona_hooks() {
538    PERSONA_HOOKS.with(|hooks| hooks.borrow_mut().clear());
539}
540
541#[derive(Clone, Default)]
542pub(crate) struct ActiveContextSnapshot {
543    steps: Vec<ActiveStep>,
544    personas: Vec<ActivePersona>,
545}
546
547pub(crate) fn swap_active_context(snapshot: ActiveContextSnapshot) -> ActiveContextSnapshot {
548    ActiveContextSnapshot {
549        steps: STEP_STACK.with(|stack| std::mem::replace(&mut *stack.borrow_mut(), snapshot.steps)),
550        personas: PERSONA_STACK
551            .with(|stack| std::mem::replace(&mut *stack.borrow_mut(), snapshot.personas)),
552    }
553}
554
555pub(crate) fn swap_active_context_suspension_stack(next: Vec<u64>) -> Vec<u64> {
556    ACTIVE_CONTEXT_SUSPENSION_STACK.with(|stack| std::mem::replace(&mut *stack.borrow_mut(), next))
557}
558
559static NEXT_ACTIVE_CONTEXT_SUSPENSION_ID: AtomicU64 = AtomicU64::new(1);
560
561/// Temporarily clear the active step/persona context and restore it on drop.
562///
563/// The guard is deliberately held across callback/module futures. If one of
564/// those futures is cancelled, its drop path restores the caller's context
565/// instead of stranding the empty or nested context in the thread-local slots.
566pub(crate) fn suspend_active_context() -> ActiveContextGuard {
567    let id = NEXT_ACTIVE_CONTEXT_SUSPENSION_ID.fetch_add(1, Ordering::Relaxed);
568    ACTIVE_CONTEXT_SUSPENSION_STACK.with(|stack| stack.borrow_mut().push(id));
569    ActiveContextGuard {
570        id,
571        outer: Some(swap_active_context(ActiveContextSnapshot::default())),
572    }
573}
574
575pub(crate) struct ActiveContextGuard {
576    id: u64,
577    outer: Option<ActiveContextSnapshot>,
578}
579
580impl Drop for ActiveContextGuard {
581    fn drop(&mut self) {
582        let owns_current_scope = ACTIVE_CONTEXT_SUSPENSION_STACK.with(|stack| {
583            let mut stack = stack.borrow_mut();
584            if stack.last() == Some(&self.id) {
585                stack.pop();
586                true
587            } else {
588                false
589            }
590        });
591        if owns_current_scope {
592            let outer = self.outer.take().expect("active context snapshot");
593            let _ = swap_active_context(outer);
594        }
595    }
596}
597
598pub fn is_tracked_function(function_name: &str) -> bool {
599    if tracked_registries_empty() {
600        return false;
601    }
602    (!step_registry_empty()
603        && STEP_REGISTRY.with(|registry| registry.borrow().contains_key(function_name)))
604        || (!persona_registry_empty()
605            && PERSONA_REGISTRY.with(|registry| registry.borrow().contains_key(function_name)))
606}
607
608pub fn step_definition_for_function(function_name: &str) -> Option<Arc<StepDefinition>> {
609    if step_registry_empty() {
610        return None;
611    }
612    STEP_REGISTRY.with(|registry| registry.borrow().get(function_name).cloned())
613}
614
615pub fn current_persona_name() -> Option<String> {
616    PERSONA_STACK.with(|stack| stack.borrow().last().map(|p| p.definition.name.clone()))
617}
618
619/// Resolve the per-stage policy for `step_name` against the currently
620/// active persona's stage declarations. Returns `None` when no persona is
621/// active or no stage matches the step name. Caller pushes the result onto
622/// `EXECUTION_POLICY_STACK`.
623///
624/// When an ambient policy is already active, the stage policy is
625/// intersected with it so a stage can only ever tighten the tool surface
626/// and side-effect ceiling — never widen them.
627fn stage_policy_for_active_step(step_name: &str) -> Option<CapabilityPolicy> {
628    let stage_policy = PERSONA_STACK.with(|stack| {
629        let stack = stack.borrow();
630        let persona = stack.last()?;
631        let stage = persona
632            .definition
633            .stages
634            .iter()
635            .find(|stage| stage.name == step_name)?;
636        Some(stage_decl_to_policy(stage))
637    })?;
638    let Some(parent) = current_execution_policy() else {
639        return Some(stage_policy);
640    };
641    let mut stage_policy = stage_policy;
642    if stage_policy.tools_are_restricted() {
643        let tools = stage_policy
644            .tools
645            .iter()
646            .filter(|tool| parent.tool_pattern_allows(tool))
647            .cloned()
648            .collect();
649        stage_policy.restrict_tools(tools);
650    }
651    Some(
652        parent
653            .intersect(&stage_policy)
654            .expect("pre-narrowed stage policy must fit the parent ceiling"),
655    )
656}
657
658fn stage_decl_to_policy(stage: &StageDecl) -> CapabilityPolicy {
659    // A stage declares tools and a side-effect ceiling; it says nothing about
660    // filesystem confinement. Overlay, not a complete policy — see
661    // `CapabilityPolicy::neutral`.
662    let mut policy = CapabilityPolicy {
663        side_effect_level: stage.side_effect_level.clone(),
664        ..CapabilityPolicy::neutral()
665    };
666    if let Some(tools) = &stage.allowed_tools {
667        policy.restrict_tools(tools.clone());
668    }
669    policy
670}
671
672fn persona_matches(pattern: &str, persona: &str) -> bool {
673    crate::orchestration::glob_match(pattern, persona)
674}
675
676pub fn matching_hooks(
677    event: HookEvent,
678    persona: Option<&str>,
679    step_name: Option<&str>,
680    budget_pct: Option<f64>,
681) -> Vec<PersonaHookInvocation> {
682    let persona = persona.unwrap_or("");
683    PERSONA_HOOKS.with(|hooks| {
684        hooks
685            .borrow()
686            .iter()
687            .filter(|hook| hook.event == event)
688            .filter(|hook| persona_matches(&hook.persona_pattern, persona))
689            .filter(|hook| match (&hook.step_name, step_name) {
690                (Some(expected), Some(actual)) => expected == actual,
691                (Some(_), None) => false,
692                (None, _) => true,
693            })
694            .filter(|hook| match (hook.threshold_pct, budget_pct) {
695                (Some(threshold), Some(pct)) => pct >= threshold,
696                (Some(_), None) => false,
697                (None, _) => true,
698            })
699            .map(|hook| PersonaHookInvocation {
700                handler: hook.handler.clone(),
701                event: hook.event,
702            })
703            .collect()
704    })
705}
706
707pub fn maybe_push_active_persona(function_name: &str, frame_depth: usize) -> bool {
708    if persona_registry_empty() {
709        return false;
710    }
711    let definition =
712        PERSONA_REGISTRY.with(|registry| registry.borrow().get(function_name).cloned());
713    let Some(definition) = definition else {
714        return false;
715    };
716    PERSONA_STACK.with(|stack| {
717        stack.borrow_mut().push(ActivePersona {
718            frame_depth,
719            definition,
720        });
721    });
722    true
723}
724
725/// Push an active step onto the stack iff `function_name` has metadata
726/// registered. Returns `true` when a frame was pushed so the call site
727/// can record that fact. Called from `Vm::push_closure_frame` after the
728/// new frame has been added.
729pub fn maybe_push_active_step(function_name: &str, frame_depth: usize, args: &[VmValue]) -> bool {
730    if step_registry_empty() {
731        return false;
732    }
733    let definition = STEP_REGISTRY.with(|registry| registry.borrow().get(function_name).cloned());
734    let Some(definition) = definition else {
735        return false;
736    };
737    let persona = current_persona_name();
738    let span_id =
739        crate::tracing::span_start(crate::tracing::SpanKind::Step, definition.name.clone());
740    if let Some(persona_name) = persona.as_deref() {
741        crate::tracing::span_set_metadata(
742            span_id,
743            "persona",
744            serde_json::Value::String(persona_name.to_string()),
745        );
746    }
747    if let Some(model) = definition.model.as_deref() {
748        crate::tracing::span_set_metadata(
749            span_id,
750            "model",
751            serde_json::Value::String(model.to_string()),
752        );
753    }
754    let step_name = definition.name.clone();
755    // The root Harness is invocation authority, not domain input. Keep it out
756    // of step transcripts so explicit authority does not pollute replay,
757    // assertions, or durable step identities with a non-serializable handle.
758    let recorded_args = match args.first() {
759        Some(VmValue::Harness(handle)) if handle.kind() == crate::harness::HarnessKind::Root => {
760            &args[1..]
761        }
762        _ => args,
763    };
764    STEP_STACK.with(|stack| {
765        stack.borrow_mut().push(ActiveStep::new(
766            frame_depth,
767            definition,
768            persona,
769            recorded_args.to_vec(),
770            span_id,
771            false,
772        ));
773    });
774    if let Some(policy) = stage_policy_for_active_step(&step_name) {
775        push_execution_policy(policy);
776        STEP_STACK.with(|stack| {
777            if let Some(top) = stack.borrow_mut().last_mut() {
778                top.stage_policy_pushed = true;
779            }
780        });
781    }
782    true
783}
784
785/// Drop any step entries whose owning frame has already been unwound,
786/// recording a `CompletedStep` summary for each. The `current_frame_depth`
787/// is `Vm::frames.len()` at the call site — entries with
788/// `frame_depth > current_frame_depth` are stale.
789pub fn prune_below_frame(current_frame_depth: usize) {
790    let mut popped: Vec<ActiveStep> = Vec::new();
791    STEP_STACK.with(|stack| {
792        let mut stack = stack.borrow_mut();
793        while let Some(top) = stack.last() {
794            if top.frame_depth > current_frame_depth {
795                popped.push(stack.pop().unwrap());
796            } else {
797                break;
798            }
799        }
800    });
801    for step in popped {
802        finish_step(step, "completed", None);
803    }
804    PERSONA_STACK.with(|stack| {
805        let mut stack = stack.borrow_mut();
806        while stack
807            .last()
808            .is_some_and(|persona| persona.frame_depth > current_frame_depth)
809        {
810            stack.pop();
811        }
812    });
813}
814
815pub fn take_active_step(current_frame_depth: usize) -> Option<ActiveStep> {
816    STEP_STACK.with(|stack| {
817        let mut stack = stack.borrow_mut();
818        if stack
819            .last()
820            .is_some_and(|step| step.frame_depth == current_frame_depth)
821        {
822            stack.pop()
823        } else {
824            None
825        }
826    })
827}
828
829pub fn finish_active_step(step: ActiveStep, status: &str, error: Option<String>) {
830    finish_step(step, status, error);
831}
832
833/// Pop the topmost active step (if its frame is the current one) and
834/// record an explicit completion status. Used when an error boundary
835/// rewrites or absorbs an in-flight error so the receipt log reflects the
836/// outcome the persona actually saw.
837pub fn pop_and_record(current_frame_depth: usize, status: &str, error: Option<String>) -> bool {
838    let popped = STEP_STACK.with(|stack| {
839        let mut stack = stack.borrow_mut();
840        if stack
841            .last()
842            .map(|step| step.frame_depth == current_frame_depth)
843            .unwrap_or(false)
844        {
845            stack.pop()
846        } else {
847            None
848        }
849    });
850    let Some(step) = popped else {
851        return false;
852    };
853    finish_step(step, status, error);
854    true
855}
856
857fn finish_step(step: ActiveStep, status: &str, error: Option<String>) {
858    if step.stage_policy_pushed {
859        pop_execution_policy();
860    }
861    crate::tracing::span_set_metadata(
862        step.span_id,
863        "status",
864        serde_json::Value::String(status.to_string()),
865    );
866    crate::tracing::span_set_metadata(
867        step.span_id,
868        "llm_calls",
869        serde_json::Value::Number(step.llm_calls.into()),
870    );
871    crate::tracing::span_set_metadata(
872        step.span_id,
873        "input_tokens",
874        serde_json::Value::Number(step.input_tokens.into()),
875    );
876    crate::tracing::span_set_metadata(
877        step.span_id,
878        "output_tokens",
879        serde_json::Value::Number(step.output_tokens.into()),
880    );
881    if let Some(cost_n) = serde_json::Number::from_f64(step.cost_usd) {
882        crate::tracing::span_set_metadata(
883            step.span_id,
884            "cost_usd",
885            serde_json::Value::Number(cost_n),
886        );
887    }
888    crate::tracing::span_end(step.span_id);
889    let summary = CompletedStep {
890        name: step.definition.name.clone(),
891        function: step.definition.function.clone(),
892        model: step
893            .last_model
894            .clone()
895            .or_else(|| step.definition.model.clone()),
896        input_tokens: step.input_tokens,
897        output_tokens: step.output_tokens,
898        cost_usd: step.cost_usd,
899        llm_calls: step.llm_calls,
900        status: status.to_string(),
901        error,
902    };
903    COMPLETED_STEPS.with(|completed| completed.borrow_mut().push(summary));
904}
905
906/// Get a snapshot of the topmost active step, if any. Used by the
907/// llm_call path to fill in defaults — never for mutation.
908pub fn with_active_step<R>(f: impl FnOnce(&ActiveStep) -> R) -> Option<R> {
909    STEP_STACK.with(|stack| stack.borrow().last().map(f))
910}
911
912/// Mutate the topmost active step (typically to attribute LLM usage).
913pub fn with_active_step_mut<R>(f: impl FnOnce(&mut ActiveStep) -> R) -> Option<R> {
914    STEP_STACK.with(|stack| stack.borrow_mut().last_mut().map(f))
915}
916
917/// Frame depth of the topmost active step, or `None` when no step is
918/// active. Used by `handle_error` to detect "this throw is exiting a
919/// step's frame".
920pub fn active_step_frame_depth() -> Option<usize> {
921    STEP_STACK.with(|stack| stack.borrow().last().map(|s| s.frame_depth))
922}
923
924/// Default model the topmost active step should impose on `llm_call`
925/// invocations whose options dict didn't pin a model.
926pub fn active_step_model_default() -> Option<String> {
927    STEP_STACK.with(|stack| {
928        stack
929            .borrow()
930            .last()
931            .and_then(|step| step.definition.model.clone())
932    })
933}
934
935/// Record that `llm_call` consumed `input_tokens` / `output_tokens` for
936/// `cost_usd`. Updates the active step's running totals and returns a
937/// budget-exhaustion error if the step's ceiling is now breached.
938///
939/// The check is performed AFTER the call so the test fixture's first
940/// call (which fits under budget) succeeds and subsequent calls trip the
941/// limit. This matches the existing `accumulate_cost_for_provider`
942/// pattern where global budget is also checked post-hoc.
943pub fn record_step_llm_usage(
944    model: &str,
945    input_tokens: i64,
946    output_tokens: i64,
947    cost_usd: f64,
948) -> Result<(), VmError> {
949    let exhausted = STEP_STACK.with(|stack| -> Option<VmError> {
950        let mut stack = stack.borrow_mut();
951        let step = stack.last_mut()?;
952        step.input_tokens = step.input_tokens.saturating_add(input_tokens.max(0) as u64);
953        step.output_tokens = step
954            .output_tokens
955            .saturating_add(output_tokens.max(0) as u64);
956        step.cost_usd += cost_usd;
957        step.llm_calls = step.llm_calls.saturating_add(1);
958        if !model.is_empty() {
959            step.last_model = Some(model.to_string());
960        }
961
962        if let Some(max_tokens) = step.definition.max_tokens {
963            if step.total_tokens() > max_tokens {
964                return Some(budget_exhausted_error(
965                    &step.definition,
966                    "max_tokens",
967                    max_tokens as f64,
968                    step.total_tokens() as f64,
969                    step.cost_usd,
970                ));
971            }
972        }
973        if let Some(max_usd) = step.definition.max_usd {
974            if step.cost_usd > max_usd {
975                return Some(budget_exhausted_error(
976                    &step.definition,
977                    "max_usd",
978                    max_usd,
979                    step.total_tokens() as f64,
980                    step.cost_usd,
981                ));
982            }
983        }
984        None
985    });
986    if let Some(err) = exhausted {
987        return Err(err);
988    }
989    Ok(())
990}
991
992fn budget_exhausted_error(
993    definition: &StepDefinition,
994    limit: &str,
995    limit_value: f64,
996    consumed_tokens: f64,
997    consumed_cost_usd: f64,
998) -> VmError {
999    let mut dict: crate::value::DictMap = crate::value::DictMap::new();
1000    dict.put_str("category", "budget_exceeded");
1001    dict.put_str("kind", "budget_exhausted");
1002    dict.put_str("reason", "step_budget_exhausted");
1003    dict.put_str("step", definition.name.clone());
1004    dict.put_str("function", definition.function.clone());
1005    dict.put_str("limit", limit);
1006    dict.insert(
1007        crate::value::intern_key("limit_value"),
1008        VmValue::Float(limit_value),
1009    );
1010    dict.insert(
1011        crate::value::intern_key("consumed_tokens"),
1012        VmValue::Float(consumed_tokens),
1013    );
1014    dict.insert(
1015        crate::value::intern_key("consumed_cost_usd"),
1016        VmValue::Float(consumed_cost_usd),
1017    );
1018    dict.put_str(
1019        "error_boundary",
1020        definition
1021            .error_boundary
1022            .clone()
1023            .unwrap_or_else(|| "fail".to_string()),
1024    );
1025    dict.put_str(
1026        "message",
1027        format!(
1028            "step `{}` exceeded {} budget ({} > {})",
1029            definition.name, limit, consumed_tokens as i64, limit_value as i64
1030        ),
1031    );
1032    VmError::Thrown(VmValue::dict(dict))
1033}
1034
1035/// Returns true if the thrown value looks like a budget-exhausted
1036/// error — either our typed step-budget dict or the existing
1037/// `crates/harn-vm/src/llm/cost.rs::budget_exceeded_error` shape.
1038/// Either form is treated identically by `error_boundary` because the
1039/// per-step budget machinery layers onto the existing envelope; a step
1040/// whose budget the preflight projection rejects is still a budget
1041/// exhaustion the step authored.
1042pub fn is_step_budget_exhausted(err: &VmError) -> bool {
1043    let VmError::Thrown(VmValue::Dict(dict)) = err else {
1044        return false;
1045    };
1046    let category = dict.get("category").and_then(vm_str);
1047    let kind = dict.get("kind").and_then(vm_str);
1048    let reason = dict.get("reason").and_then(vm_str);
1049    if matches!(kind, Some("budget_exhausted")) && matches!(reason, Some("step_budget_exhausted")) {
1050        return true;
1051    }
1052    matches!(category, Some("budget_exceeded"))
1053}
1054
1055/// Annotate an existing budget-exhausted error with `escalated: true`
1056/// and the step's identity so the persona body / handoff receiver can
1057/// route on it. Returns the original error if it isn't a thrown dict.
1058/// Ensures `step` and `function` keys reflect the just-finished step
1059/// even when the underlying error was raised by the preflight budget
1060/// machinery (which doesn't know which step it's running under).
1061pub fn mark_escalated(err: VmError, step_name: Option<&str>, function: Option<&str>) -> VmError {
1062    let VmError::Thrown(VmValue::Dict(dict)) = err else {
1063        return err;
1064    };
1065    let mut next = (*dict).clone();
1066    next.insert(crate::value::intern_key("escalated"), VmValue::Bool(true));
1067    next.put_str("category", "handoff_escalation");
1068    if let Some(step) = step_name {
1069        next.entry(crate::value::intern_key("step"))
1070            .or_insert_with(|| VmValue::String(arcstr::ArcStr::from(step.to_string())));
1071    }
1072    if let Some(function) = function {
1073        next.entry(crate::value::intern_key("function"))
1074            .or_insert_with(|| VmValue::String(arcstr::ArcStr::from(function.to_string())));
1075    }
1076    VmError::Thrown(VmValue::dict(next))
1077}
1078
1079/// Drain the completed-step log. Used by receipt builders that want a
1080/// per-step model + token + cost breakdown for the just-finished run.
1081pub fn drain_completed_steps() -> Vec<CompletedStep> {
1082    COMPLETED_STEPS.with(|completed| std::mem::take(&mut *completed.borrow_mut()))
1083}
1084
1085/// Read the completed-step log without clearing it. Use when callers
1086/// want a peek without disturbing the global record stream.
1087pub fn peek_completed_steps() -> Vec<CompletedStep> {
1088    COMPLETED_STEPS.with(|completed| completed.borrow().clone())
1089}
1090
1091/// Lower a [`CompletedStep`] into JSON for embedding in receipts /
1092/// inspect output.
1093pub fn completed_step_to_json(step: &CompletedStep) -> JsonValue {
1094    serde_json::to_value(step).unwrap_or(JsonValue::Null)
1095}
1096
1097/// Register the `__register_step` and `__register_persona` host builtins.
1098/// Compiler-emitted bytecode after every `@step` / persona declaration
1099/// calls these with `(function_name, metadata_dict)` so the runtime can
1100/// later dispatch on the step's metadata when its function is invoked.
1101pub fn register_step_builtins(vm: &mut crate::vm::Vm) {
1102    for def in MODULE_BUILTINS {
1103        vm.register_builtin_def(def);
1104    }
1105}
1106
1107pub(crate) const MODULE_BUILTINS: &[&VmBuiltinDef] = &[
1108    &__REGISTER_STEP_DEF,
1109    &__REGISTER_PERSONA_DEF,
1110    &__PERSONA_OUTPUT_STYLE_DEF,
1111];
1112
1113#[harn_builtin(
1114    exposure = "runtime_internal",
1115    effects = [],
1116    category = "step_runtime", runtime_only = true
1117)]
1118fn __register_step(args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
1119    register_step_from_dict(args.to_vec())
1120}
1121
1122#[harn_builtin(
1123    exposure = "runtime_internal",
1124    effects = [],
1125    category = "step_runtime", runtime_only = true
1126)]
1127fn __register_persona(args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
1128    register_persona_from_dict(args.to_vec())
1129}
1130
1131#[harn_builtin(
1132    exposure = "runtime_internal",
1133    effects = [],
1134    sig = "__persona_output_style(function?: string) -> dict",
1135    category = "step_runtime",
1136    runtime_only = true
1137)]
1138fn __persona_output_style(args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
1139    Ok(persona_output_style(args.to_vec()))
1140}
1141
1142#[cfg(test)]
1143mod tests {
1144    use super::*;
1145    use crate::value::VmDictExt;
1146
1147    fn fresh_state() {
1148        reset_thread_local_state();
1149    }
1150
1151    #[test]
1152    fn persona_output_style_reads_registry_and_active_stack() {
1153        use crate::value::{intern_key, DictMap};
1154        fresh_state();
1155
1156        // Register a persona whose metadata carries a table-form output style.
1157        let mut style = DictMap::new();
1158        style.put_str("name", "concise");
1159        style.put_str("instructions", "Be terse.");
1160        let mut meta = DictMap::new();
1161        meta.put_str("name", "Reviewer");
1162        meta.insert(intern_key("output_style"), VmValue::dict(style));
1163        register_persona_from_dict(vec![
1164            VmValue::String(arcstr::ArcStr::from("reviewer_fn")),
1165            VmValue::dict(meta),
1166        ])
1167        .expect("persona registers");
1168
1169        // Lookup by function name returns the declared style.
1170        let by_name =
1171            persona_output_style(vec![VmValue::String(arcstr::ArcStr::from("reviewer_fn"))]);
1172        let dict = by_name.as_dict().expect("dict");
1173        assert_eq!(
1174            dict.get("name").map(VmValue::display).as_deref(),
1175            Some("concise")
1176        );
1177        assert_eq!(
1178            dict.get("instructions").map(VmValue::display).as_deref(),
1179            Some("Be terse.")
1180        );
1181
1182        // No active persona on the stack → nil.
1183        assert!(matches!(persona_output_style(vec![]), VmValue::Nil));
1184        // Unknown persona → nil.
1185        assert!(matches!(
1186            persona_output_style(vec![VmValue::String(arcstr::ArcStr::from("nope"))]),
1187            VmValue::Nil
1188        ));
1189    }
1190
1191    #[test]
1192    fn registers_and_pops_step_from_dict() {
1193        fresh_state();
1194        let mut budget: crate::value::DictMap = crate::value::DictMap::new();
1195        budget.insert(crate::value::intern_key("max_tokens"), VmValue::Int(100));
1196        budget.insert(crate::value::intern_key("max_usd"), VmValue::Float(0.05));
1197        let mut meta: crate::value::DictMap = crate::value::DictMap::new();
1198        meta.put_str("name", "plan");
1199        meta.put_str("model", "claude-haiku-4-5");
1200        meta.put_str("error_boundary", "continue");
1201        meta.insert(crate::value::intern_key("budget"), VmValue::dict(budget));
1202
1203        register_step_from_dict(vec![
1204            VmValue::String(arcstr::ArcStr::from("plan_step")),
1205            VmValue::dict(meta),
1206        ])
1207        .expect("registration succeeds");
1208
1209        assert!(maybe_push_active_step("plan_step", 3, &[]));
1210        assert_eq!(active_step_frame_depth(), Some(3));
1211        assert_eq!(
1212            active_step_model_default().as_deref(),
1213            Some("claude-haiku-4-5")
1214        );
1215
1216        record_step_llm_usage("claude-haiku-4-5", 10, 20, 0.001).expect("under budget");
1217        with_active_step(|step| {
1218            assert_eq!(step.input_tokens, 10);
1219            assert_eq!(step.output_tokens, 20);
1220            assert!((step.cost_usd - 0.001).abs() < 1e-9);
1221        });
1222
1223        let err =
1224            record_step_llm_usage("claude-haiku-4-5", 50, 50, 0.0).expect_err("should exhaust");
1225        assert!(is_step_budget_exhausted(&err));
1226
1227        prune_below_frame(2);
1228        let completed = drain_completed_steps();
1229        assert_eq!(completed.len(), 1);
1230        assert_eq!(completed[0].llm_calls, 2);
1231    }
1232
1233    #[test]
1234    fn unregistered_function_does_not_push() {
1235        fresh_state();
1236        assert!(!maybe_push_active_step("not_a_step", 1, &[]));
1237        assert!(active_step_frame_depth().is_none());
1238    }
1239
1240    #[test]
1241    fn tracked_registry_empty_fast_path_tracks_registrations_and_reset() {
1242        fresh_state();
1243        assert!(tracked_registries_empty());
1244        assert!(!is_tracked_function("plan_step"));
1245
1246        register_step(
1247            "plan_step",
1248            StepDefinition {
1249                name: "plan".to_string(),
1250                function: "plan_step".to_string(),
1251                ..StepDefinition::default()
1252            },
1253        );
1254        assert!(!tracked_registries_empty());
1255        assert!(is_tracked_function("plan_step"));
1256        assert!(step_definition_for_function("plan_step").is_some());
1257
1258        register_step(
1259            "plan_step",
1260            StepDefinition {
1261                name: "plan_v2".to_string(),
1262                function: "plan_step".to_string(),
1263                ..StepDefinition::default()
1264            },
1265        );
1266        assert!(is_tracked_function("plan_step"));
1267
1268        fresh_state();
1269        assert!(tracked_registries_empty());
1270        assert!(!is_tracked_function("plan_step"));
1271    }
1272
1273    #[test]
1274    fn stage_policy_narrows_but_does_not_widen_parent_policy() {
1275        fresh_state();
1276        let mut meta: crate::value::DictMap = crate::value::DictMap::new();
1277        meta.put_str("name", "research");
1278        register_step_from_dict(vec![
1279            VmValue::String(arcstr::ArcStr::from("research_step")),
1280            VmValue::dict(meta),
1281        ])
1282        .expect("step registration");
1283
1284        let mut stage_dict: crate::value::DictMap = crate::value::DictMap::new();
1285        stage_dict.put_str("name", "research");
1286        // Stage tries to add `edit` on top of a parent that only allowed `read`.
1287        stage_dict.insert(
1288            crate::value::intern_key("allowed_tools"),
1289            VmValue::List(std::sync::Arc::new(vec![
1290                VmValue::String(arcstr::ArcStr::from("read")),
1291                VmValue::String(arcstr::ArcStr::from("edit")),
1292            ])),
1293        );
1294        let mut persona_meta: crate::value::DictMap = crate::value::DictMap::new();
1295        persona_meta.put_str("name", "scoped");
1296        persona_meta.insert(
1297            crate::value::intern_key("stages"),
1298            VmValue::List(std::sync::Arc::new(vec![VmValue::Dict(
1299                std::sync::Arc::new(stage_dict),
1300            )])),
1301        );
1302        register_persona_from_dict(vec![
1303            VmValue::String(arcstr::ArcStr::from("scoped_persona")),
1304            VmValue::dict(persona_meta),
1305        ])
1306        .expect("persona registration");
1307
1308        push_execution_policy(CapabilityPolicy {
1309            tools: vec!["read".to_string()],
1310            capabilities: std::collections::BTreeMap::from([(
1311                "workspace".to_string(),
1312                vec!["read_text".to_string()],
1313            )]),
1314            workspace_roots: vec!["/workspace".to_string()],
1315            side_effect_level: Some("workspace_read".to_string()),
1316            ..CapabilityPolicy::default()
1317        });
1318        assert!(maybe_push_active_persona("scoped_persona", 1));
1319        assert!(maybe_push_active_step("research_step", 2, &[]));
1320        let policy = current_execution_policy().expect("stage policy active");
1321        // `edit` is filtered out because the parent already denied it.
1322        assert_eq!(policy.tools, vec!["read".to_string()]);
1323        assert_eq!(
1324            policy.capabilities,
1325            std::collections::BTreeMap::from([(
1326                "workspace".to_string(),
1327                vec!["read_text".to_string()],
1328            )])
1329        );
1330        assert_eq!(policy.workspace_roots, vec!["/workspace".to_string()]);
1331        assert_eq!(policy.side_effect_level.as_deref(), Some("workspace_read"));
1332
1333        prune_below_frame(0);
1334        pop_execution_policy();
1335        assert!(current_execution_policy().is_none());
1336    }
1337
1338    #[test]
1339    fn stage_policy_does_not_confine_a_run_that_chose_no_confinement() {
1340        fresh_state();
1341        let mut meta: crate::value::DictMap = crate::value::DictMap::new();
1342        meta.put_str("name", "survey");
1343        register_step_from_dict(vec![
1344            VmValue::String(arcstr::ArcStr::from("survey_step")),
1345            VmValue::dict(meta),
1346        ])
1347        .expect("step registration");
1348
1349        let mut stage_dict: crate::value::DictMap = crate::value::DictMap::new();
1350        stage_dict.put_str("name", "survey");
1351        stage_dict.insert(
1352            crate::value::intern_key("allowed_tools"),
1353            VmValue::List(std::sync::Arc::new(vec![VmValue::String(
1354                arcstr::ArcStr::from("read"),
1355            )])),
1356        );
1357        let mut persona_meta: crate::value::DictMap = crate::value::DictMap::new();
1358        persona_meta.put_str("name", "ambient");
1359        persona_meta.insert(
1360            crate::value::intern_key("stages"),
1361            VmValue::List(std::sync::Arc::new(vec![VmValue::Dict(
1362                std::sync::Arc::new(stage_dict),
1363            )])),
1364        );
1365        register_persona_from_dict(vec![
1366            VmValue::String(arcstr::ArcStr::from("ambient_persona")),
1367            VmValue::dict(persona_meta),
1368        ])
1369        .expect("persona registration");
1370
1371        // No parent policy: the run is deliberately unsandboxed. A stage
1372        // declares tools, not filesystem scope, so entering it must not
1373        // conjure workspace-root enforcement out of the stage declaration.
1374        assert!(current_execution_policy().is_none());
1375        assert!(maybe_push_active_persona("ambient_persona", 1));
1376        assert!(maybe_push_active_step("survey_step", 2, &[]));
1377        let policy = current_execution_policy().expect("stage policy active");
1378        assert_eq!(policy.tools, vec!["read".to_string()]);
1379        assert_eq!(
1380            policy.sandbox_profile,
1381            crate::orchestration::SandboxProfile::Unrestricted
1382        );
1383        assert!(policy.workspace_roots.is_empty());
1384
1385        prune_below_frame(0);
1386        pop_execution_policy();
1387        assert!(current_execution_policy().is_none());
1388    }
1389
1390    #[test]
1391    fn explicit_empty_stage_tool_list_denies_every_tool() {
1392        let policy = stage_decl_to_policy(&StageDecl {
1393            name: "observe".to_string(),
1394            allowed_tools: Some(Vec::new()),
1395            ..StageDecl::default()
1396        });
1397
1398        assert!(policy.tools_are_restricted());
1399        assert!(policy.tools_deny_all());
1400    }
1401
1402    #[test]
1403    fn stage_policy_is_pushed_and_popped_around_step() {
1404        fresh_state();
1405        let mut meta: crate::value::DictMap = crate::value::DictMap::new();
1406        meta.put_str("name", "research");
1407        register_step_from_dict(vec![
1408            VmValue::String(arcstr::ArcStr::from("research_step")),
1409            VmValue::dict(meta),
1410        ])
1411        .expect("step registration succeeds");
1412
1413        let mut stage_dict: crate::value::DictMap = crate::value::DictMap::new();
1414        stage_dict.put_str("name", "research");
1415        stage_dict.insert(
1416            crate::value::intern_key("allowed_tools"),
1417            VmValue::List(std::sync::Arc::new(vec![VmValue::String(
1418                arcstr::ArcStr::from("read"),
1419            )])),
1420        );
1421        let mut persona_meta: crate::value::DictMap = crate::value::DictMap::new();
1422        persona_meta.put_str("name", "scoped");
1423        persona_meta.insert(
1424            crate::value::intern_key("stages"),
1425            VmValue::List(std::sync::Arc::new(vec![VmValue::Dict(
1426                std::sync::Arc::new(stage_dict),
1427            )])),
1428        );
1429        register_persona_from_dict(vec![
1430            VmValue::String(arcstr::ArcStr::from("scoped_persona")),
1431            VmValue::dict(persona_meta),
1432        ])
1433        .expect("persona registration succeeds");
1434
1435        assert!(maybe_push_active_persona("scoped_persona", 1));
1436        assert!(crate::orchestration::current_execution_policy().is_none());
1437        assert!(maybe_push_active_step("research_step", 2, &[]));
1438        let policy = crate::orchestration::current_execution_policy()
1439            .expect("stage policy is active inside step");
1440        assert_eq!(policy.tools, vec!["read".to_string()]);
1441
1442        prune_below_frame(0);
1443        assert!(crate::orchestration::current_execution_policy().is_none());
1444    }
1445}