1use 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#[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 pub error_boundary: Option<String>,
57}
58
59#[derive(Debug, Default, Clone)]
60pub struct PersonaDefinition {
61 pub name: String,
62 pub stages: Vec<StageDecl>,
67 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#[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 pub span_id: u64,
108 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#[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
180pub 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
210pub 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
269fn 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
294fn 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
317pub 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
412pub 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
561pub(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
619fn 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 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
725pub 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 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
785pub 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
833pub 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
906pub fn with_active_step<R>(f: impl FnOnce(&ActiveStep) -> R) -> Option<R> {
909 STEP_STACK.with(|stack| stack.borrow().last().map(f))
910}
911
912pub 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
917pub fn active_step_frame_depth() -> Option<usize> {
921 STEP_STACK.with(|stack| stack.borrow().last().map(|s| s.frame_depth))
922}
923
924pub 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
935pub 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
1035pub 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
1055pub 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
1079pub fn drain_completed_steps() -> Vec<CompletedStep> {
1082 COMPLETED_STEPS.with(|completed| std::mem::take(&mut *completed.borrow_mut()))
1083}
1084
1085pub fn peek_completed_steps() -> Vec<CompletedStep> {
1088 COMPLETED_STEPS.with(|completed| completed.borrow().clone())
1089}
1090
1091pub fn completed_step_to_json(step: &CompletedStep) -> JsonValue {
1094 serde_json::to_value(step).unwrap_or(JsonValue::Null)
1095}
1096
1097pub 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 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 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 assert!(matches!(persona_output_style(vec![]), VmValue::Nil));
1184 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_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 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 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}