Skip to main content

harn_vm/stdlib/template/
llm_context.rs

1//! Ambient render-time LLM context exposed to `.harn.prompt` templates.
2//!
3//! When `render()` / `render_prompt()` / `render_string()` is invoked from
4//! within an LLM-aware frame (`llm_call`, the default handler stack, or
5//! `agent_loop`), the active provider/model/family/capabilities are
6//! published as the reserved `llm` scope key so authors can write
7//! capability-aware partials without manual plumbing:
8//!
9//! ```text
10//! {{ if llm.capabilities.native_tools }}
11//!   Call `finish_task` when done.
12//! {{ else }}
13//!   When done, output: `<<DONE>>`
14//! {{ end }}
15//! ```
16//!
17//! Bare `render()` calls outside any LLM frame leave `llm = nil`, so
18//! templates branch on `{{ if llm }}` for the doc-gen / CI paths.
19//!
20//! The context lives in a thread-local stack so concurrent agent_loop
21//! iterations on different threads stay isolated; nested
22//! push/pop pairs (e.g. an inner `llm_call` from a middleware handler)
23//! shadow the outer frame for the duration of the inner render.
24
25use crate::value::VmDictExt;
26use std::cell::RefCell;
27use std::collections::BTreeMap;
28use std::sync::atomic::{AtomicU64, Ordering};
29
30use crate::value::VmValue;
31
32/// Resolved provider/model identity plus the corresponding capability
33/// snapshot, materialized at LLM-frame entry and injected as the `llm`
34/// binding during any `render()` call inside that frame.
35#[derive(Debug, Clone)]
36pub struct LlmRenderContext {
37    pub provider: String,
38    pub model: String,
39    pub family: String,
40    /// Snapshot of `harness.llm.provider_capabilities(provider, model)` — a
41    /// `VmValue::Dict` shaped exactly like the builtin's return value.
42    pub capabilities: VmValue,
43}
44
45impl LlmRenderContext {
46    /// Build a context from resolved provider/model strings, looking up
47    /// the capability snapshot and deriving the canonical model family.
48    pub fn resolve(provider: &str, model: &str) -> Self {
49        let (provider, model) = crate::llm::managed_supply::logical_route(provider, model)
50            .unwrap_or_else(|_| (provider.to_string(), model.to_string()));
51        let caps = crate::llm::capabilities::lookup(&provider, &model);
52        let capabilities =
53            crate::llm::config_builtins::capabilities_to_vm_value(&provider, &model, &caps);
54        Self {
55            family: crate::llm_config::model_family(&provider, &model),
56            provider,
57            model,
58            capabilities,
59        }
60    }
61
62    /// Materialize the context as the `llm` scope value:
63    /// `{provider, model, family, capabilities: <provider_capabilities dict>}`.
64    pub fn to_vm_value(&self) -> VmValue {
65        let mut dict = BTreeMap::new();
66        dict.put_str("provider", self.provider.as_str());
67        dict.put_str("model", self.model.as_str());
68        dict.put_str("family", self.family.as_str());
69        dict.insert("capabilities".to_string(), self.capabilities.clone());
70        VmValue::dict(dict)
71    }
72}
73
74thread_local! {
75    static LLM_RENDER_STACK: RefCell<Vec<LlmRenderContextFrame>> = const { RefCell::new(Vec::new()) };
76}
77
78static NEXT_LLM_RENDER_FRAME_ID: AtomicU64 = AtomicU64::new(1);
79
80/// A stack entry tagged so an RAII guard can clean up the frame it owns without
81/// accidentally popping an unrelated ambient scope when an async task is
82/// cancelled while its task-local stack is swapped out.
83#[derive(Debug, Clone)]
84pub(crate) struct LlmRenderContextFrame {
85    id: u64,
86    context: LlmRenderContext,
87}
88
89fn next_frame_id() -> u64 {
90    NEXT_LLM_RENDER_FRAME_ID.fetch_add(1, Ordering::Relaxed)
91}
92
93fn push_llm_render_context_frame(ctx: LlmRenderContext) -> u64 {
94    let id = next_frame_id();
95    LLM_RENDER_STACK.with(|stack| {
96        stack
97            .borrow_mut()
98            .push(LlmRenderContextFrame { id, context: ctx });
99    });
100    id
101}
102
103/// Push a frame onto the ambient render-context stack. Pair with
104/// `pop_llm_render_context` (or use `LlmRenderContextGuard`) so the
105/// stack stays balanced even on the unwind path.
106pub fn push_llm_render_context(ctx: LlmRenderContext) {
107    push_llm_render_context_frame(ctx);
108}
109
110/// Pop the most recently pushed frame. Returns `None` (rather than
111/// panicking) if the stack was empty, since the host may legitimately
112/// unwind through a balanced push/pop sequence.
113pub fn pop_llm_render_context() -> Option<LlmRenderContext> {
114    LLM_RENDER_STACK.with(|stack| stack.borrow_mut().pop().map(|frame| frame.context))
115}
116
117/// Return a clone of the active frame, or `None` if no LLM context is
118/// in scope. Render entry-points use this to decide whether to inject
119/// the `llm` binding.
120pub fn current_llm_render_context() -> Option<LlmRenderContext> {
121    LLM_RENDER_STACK.with(|stack| stack.borrow().last().map(|frame| frame.context.clone()))
122}
123
124/// Per-task ambient-scope swap of the LLM render-context stack. See
125/// `orchestration::ambient_scope`: this frame is pushed for the duration of an
126/// `llm_call` (held across the model HTTP `.await`), so concurrent fan-out
127/// children each running their own `llm_call` would otherwise render templates
128/// with a sibling's provider/model/capabilities.
129pub(crate) fn swap_llm_render_stack(
130    next: Vec<LlmRenderContextFrame>,
131) -> Vec<LlmRenderContextFrame> {
132    LLM_RENDER_STACK.with(|stack| std::mem::replace(&mut *stack.borrow_mut(), next))
133}
134
135/// Reset the stack — wired into `reset_thread_local_state` so tests
136/// and serialized adapter sessions start clean.
137pub(crate) fn reset_llm_render_stack() {
138    LLM_RENDER_STACK.with(|stack| stack.borrow_mut().clear());
139}
140
141/// RAII guard that pushes a context on construction and pops on drop.
142/// Use this in Rust hosts (e.g. `llm_call_impl`) so the stack stays
143/// balanced across `?`-shortcircuits and panics.
144pub struct LlmRenderContextGuard {
145    frame_id: u64,
146    /// Tagged so a misuse (drop-order inversion across nested guards)
147    /// surfaces as a `debug_assert` instead of silently popping the
148    /// wrong frame. Carries no runtime cost in release builds.
149    expected_depth: usize,
150}
151
152impl LlmRenderContextGuard {
153    pub fn enter(ctx: LlmRenderContext) -> Self {
154        let frame_id = push_llm_render_context_frame(ctx);
155        let depth = LLM_RENDER_STACK.with(|stack| stack.borrow().len());
156        Self {
157            frame_id,
158            expected_depth: depth,
159        }
160    }
161}
162
163impl Drop for LlmRenderContextGuard {
164    fn drop(&mut self) {
165        LLM_RENDER_STACK.with(|stack| {
166            let mut stack = stack.borrow_mut();
167            if stack.last().is_some_and(|frame| frame.id == self.frame_id) {
168                stack.pop();
169                return;
170            }
171
172            if let Some(index) = stack.iter().position(|frame| frame.id == self.frame_id) {
173                debug_assert_eq!(
174                    stack.len(),
175                    self.expected_depth,
176                    "LlmRenderContextGuard nested-drop order violated",
177                );
178                debug_assert_eq!(index + 1, stack.len(), "nested-drop order violated");
179                stack.remove(index);
180            }
181        });
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188
189    fn derive_family(provider: &str, model: &str) -> String {
190        crate::llm_config::model_family(provider, model)
191    }
192
193    #[test]
194    fn family_from_model_id_takes_precedence() {
195        assert_eq!(
196            derive_family("openrouter", "anthropic/claude-3-5-sonnet"),
197            "anthropic-claude"
198        );
199        assert_eq!(derive_family("openrouter", "openai/gpt-4o"), "openai-gpt");
200        assert_eq!(
201            derive_family("openrouter", "google/gemini-1.5-pro"),
202            "google-gemini"
203        );
204        assert_eq!(derive_family("llamacpp", "qwen3.6-35b-a3b"), "qwen");
205    }
206
207    #[test]
208    fn family_falls_back_to_provider_alias() {
209        assert_eq!(
210            derive_family("anthropic", "unknown-future-model"),
211            "anthropic-claude"
212        );
213        assert_eq!(derive_family("azure", "deployment-xyz"), "openai-gpt");
214        assert_eq!(derive_family("vertex", "model-xyz"), "google-gemini");
215        assert_eq!(derive_family("local", "anonymous-snapshot"), "local");
216        assert_eq!(derive_family("", ""), "unknown");
217    }
218
219    #[test]
220    fn push_pop_stack_round_trip() {
221        reset_llm_render_stack();
222        assert!(current_llm_render_context().is_none());
223        push_llm_render_context(LlmRenderContext::resolve("anthropic", "claude-3-5-sonnet"));
224        assert_eq!(
225            current_llm_render_context().map(|c| c.family),
226            Some("anthropic-claude".to_string()),
227        );
228        push_llm_render_context(LlmRenderContext::resolve("openai", "gpt-4o"));
229        assert_eq!(
230            current_llm_render_context().map(|c| c.family),
231            Some("openai-gpt".to_string()),
232        );
233        pop_llm_render_context();
234        assert_eq!(
235            current_llm_render_context().map(|c| c.family),
236            Some("anthropic-claude".to_string()),
237        );
238        pop_llm_render_context();
239        assert!(current_llm_render_context().is_none());
240    }
241
242    #[test]
243    fn guard_pops_on_drop() {
244        reset_llm_render_stack();
245        {
246            let _guard = LlmRenderContextGuard::enter(LlmRenderContext::resolve(
247                "anthropic",
248                "claude-3-5-sonnet",
249            ));
250            assert!(current_llm_render_context().is_some());
251        }
252        assert!(current_llm_render_context().is_none());
253    }
254}