Skip to main content

lex_runtime/handler/
mod.rs

1//! Native effect handlers, dispatched at runtime through the VM's
2//! `EffectHandler` trait. The handler also re-checks the runtime policy
3//! per spec §7.4 (the static check is necessary but not sufficient: a fn
4//! declared `[fs_read("/data")]` that's allowed at startup still has to
5//! pass the path check at the point of dispatch).
6
7use lex_bytecode::vm::{EffectHandler, Vm};
8use lex_bytecode::{Program, Value};
9use std::path::PathBuf;
10use std::sync::atomic::{AtomicI64, AtomicU64, Ordering};
11use std::sync::{Mutex, OnceLock};
12use std::sync::Arc;
13use std::time::{SystemTime, UNIX_EPOCH};
14
15use crate::builtins::{call_pure_builtin, is_pure_call};
16use crate::policy::Policy;
17
18mod approval;
19mod dispatch;
20mod fs;
21mod http_client;
22mod http_serve;
23mod kv;
24mod llm;
25mod logging;
26mod proc;
27mod redis_store;
28mod sql;
29mod udp;
30
31pub use approval::{ApprovalSink, NullApprovalSink, StdinApprovalSink};
32pub use http_serve::TlsConfig;
33#[cfg(feature = "quic")]
34pub(crate) use http_serve::{
35    build_request_value_parts, dispatch_route, stamp_path_params, unpack_response, ResponseBodyOut,
36    RouteSeg, ServeOpts, UnpackedResponse,
37};
38use http_client::*;
39use http_serve::*;
40use kv::*;
41use llm::*;
42use redis_store::*;
43use sql::*;
44use udp::*;
45
46/// Output sink used by `io.print`. Tests inject a buffer; production prints
47/// to stdout.
48pub trait IoSink: Send {
49    fn print_line(&mut self, s: &str);
50}
51
52pub struct StdoutSink;
53impl IoSink for StdoutSink {
54    fn print_line(&mut self, s: &str) {
55        use std::io::Write;
56        println!("{s}");
57        let _ = std::io::stdout().flush();
58    }
59}
60
61#[derive(Default)]
62pub struct CapturedSink { pub lines: Vec<String> }
63impl IoSink for CapturedSink {
64    fn print_line(&mut self, s: &str) { self.lines.push(s.to_string()); }
65}
66
67/// `agent.cloud_stream` registry: per-handle producer iterators
68/// keyed by opaque handle id (#305 slice 3).
69pub type StreamRegistry =
70    std::collections::HashMap<String, Box<dyn Iterator<Item = String> + Send>>;
71
72/// `requested_exit`'s "nothing requested" value. Outside the `i32`
73/// range on purpose, so no real exit status can collide with it.
74pub const NO_EXIT: i64 = i64::MIN;
75
76pub struct DefaultHandler {
77    policy: Policy,
78    pub sink: Box<dyn IoSink>,
79    /// Optional read root for `io.read` — when set, `io.read("p")` resolves
80    /// to `read_root.join(p)`. Lets tests run without touching the real fs.
81    pub read_root: Option<PathBuf>,
82    /// Per-run budget pool (#225). `Arc<AtomicU64>` so parallel
83    /// branches share one counter without locking. Initialized to
84    /// the policy ceiling at handler construction; each call to a
85    /// function with declared `[budget(N)]` deducts N atomically
86    /// via `note_call_budget`. Cloning the handler is intentional
87    /// for net.serve / chat handlers — they share the same pool.
88    pub budget_remaining: Arc<AtomicU64>,
89    /// The original ceiling that `budget_remaining` started at, kept
90    /// for diagnostics so a `BudgetExceeded` error can report
91    /// `(used, ceiling)` rather than just "exceeded by N".
92    pub budget_ceiling: Option<u64>,
93    /// Shared reference to the program, needed by `net.serve` so the
94    /// handler can spin up fresh VMs to dispatch incoming requests.
95    /// `None` if the handler was constructed without a program.
96    pub program: Option<Arc<Program>>,
97    /// Chat registry; populated by `net.serve_ws`'s per-message
98    /// dispatch so `chat.broadcast` / `chat.send` work from inside
99    /// a handler invocation.
100    pub chat_registry: Option<Arc<crate::ws::ChatRegistry>>,
101    /// LRU cache of `agent.call_mcp` clients keyed by the
102    /// command-line string (#197). Avoids spawn-per-call cost
103    /// when an agent invokes the same MCP server in tight loops.
104    /// Capped — when the cache is full, the least-recently-used
105    /// entry is dropped (its subprocess is reaped on Drop).
106    pub mcp_clients: crate::mcp_client::McpClientCache,
107    /// Stream registry for `agent.cloud_stream` / `stream.next` /
108    /// `stream.collect` (#305 slice 3). Keyed by an opaque handle
109    /// id; values are the producer iterators. Wrapped in
110    /// `Arc<Mutex<…>>` so par_map workers can share the same
111    /// stream pool (when slice-2's per-worker handler split chains
112    /// the registry through).
113    pub streams: Arc<std::sync::Mutex<StreamRegistry>>,
114    /// Monotonic counter for handing out fresh stream handle ids.
115    pub next_stream_id: Arc<std::sync::atomic::AtomicU64>,
116    /// Stack of per-request arenas (#463 scaffolding). One entry
117    /// per active request scope; `net.serve_fn`'s request loop
118    /// pushes on entry, pops on exit. Today nothing reads from the
119    /// arenas — they're scaffolding for the Value-rep follow-on
120    /// that routes `MakeRecord` / `MakeList` allocations into the
121    /// active arena. See `crates/lex-runtime/src/arena.rs`.
122    ///
123    /// Held by value (not Arc) so worker-clone handlers
124    /// (`spawn_for_worker`) get a fresh empty stack rather than
125    /// sharing the parent's arenas — worker-thread allocations
126    /// have a different lifetime than the request that spawned
127    /// them.
128    arena_stack: Vec<(u64, crate::arena::Arena)>,
129    /// Monotonic counter for the scope ids handed out by
130    /// `enter_request_scope`. `enter` returns a fresh id; `exit`
131    /// finds and removes the matching entry. Plain `u64`, not
132    /// shared — each handler instance has its own counter.
133    next_scope_id: u64,
134    /// The status `std.process.exit` asked to terminate with (#754).
135    ///
136    /// `Arc`-shared for the same reason `budget_remaining` is: a
137    /// `par_map` worker gets a cloned handler, and an exit called from
138    /// a worker has to reach the VM that will actually return. Held
139    /// per-handler instead, an exit inside parallel work would be
140    /// dropped on the floor — the failure mode being that a program
141    /// signalling failure silently exits 0.
142    ///
143    /// [`NO_EXIT`] means "not requested". The first writer wins, so a
144    /// racing second exit cannot overwrite the status the program
145    /// already stopped with.
146    pub requested_exit: Arc<AtomicI64>,
147    /// Arguments passed after `--` in `lex run <file> -- [args...]`.
148    /// Returned by `io.argv()` so Lex `main` functions can read CLI flags.
149    pub program_args: Vec<String>,
150    /// Host boundary for `approval.request`. Defaults to
151    /// `NullApprovalSink` (always refuses) so a handler must opt in
152    /// via `with_approval_sink` before `[approval]` calls can succeed.
153    pub approval_sink: Box<dyn ApprovalSink>,
154}
155
156impl DefaultHandler {
157    pub fn new(policy: Policy) -> Self {
158        // If the caller supplied a ceiling, the pool starts at that
159        // ceiling and counts down. No ceiling = `u64::MAX` so calls
160        // never refuse on budget grounds (existing behavior).
161        let ceiling = policy.budget;
162        let initial = ceiling.unwrap_or(u64::MAX);
163        Self {
164            policy,
165            sink: Box::new(StdoutSink),
166            read_root: None,
167            budget_remaining: Arc::new(AtomicU64::new(initial)),
168            budget_ceiling: ceiling,
169            program: None,
170            chat_registry: None,
171            mcp_clients: crate::mcp_client::McpClientCache::with_capacity(16),
172            streams: Arc::new(std::sync::Mutex::new(std::collections::HashMap::new())),
173            next_stream_id: Arc::new(std::sync::atomic::AtomicU64::new(0)),
174            arena_stack: Vec::new(),
175            next_scope_id: 1,
176            requested_exit: Arc::new(AtomicI64::new(NO_EXIT)),
177            program_args: Vec::new(),
178            approval_sink: Box::new(NullApprovalSink),
179        }
180    }
181
182    pub fn with_approval_sink(mut self, sink: Box<dyn ApprovalSink>) -> Self {
183        self.approval_sink = sink; self
184    }
185
186    /// Read-only access to the currently-active request arena, if
187    /// any. `None` outside a request scope. The follow-on slice
188    /// that routes `Value` allocations consults this from the VM
189    /// path; today it has no callers in tree but is exercised in
190    /// tests.
191    pub fn active_arena(&self) -> Option<&crate::arena::Arena> {
192        self.arena_stack.last().map(|(_, a)| a)
193    }
194
195    /// Test-only: depth of the arena stack. Lets tests confirm the
196    /// `net.serve_fn` request loop pushes/pops symmetrically.
197    pub fn arena_stack_depth(&self) -> usize {
198        self.arena_stack.len()
199    }
200
201    pub fn with_program(mut self, program: Arc<Program>) -> Self {
202        self.program = Some(program); self
203    }
204
205    pub fn with_chat_registry(mut self, registry: Arc<crate::ws::ChatRegistry>) -> Self {
206        self.chat_registry = Some(registry); self
207    }
208
209    pub fn with_sink(mut self, sink: Box<dyn IoSink>) -> Self {
210        self.sink = sink; self
211    }
212
213    pub fn with_read_root(mut self, root: PathBuf) -> Self {
214        self.read_root = Some(root); self
215    }
216
217    pub fn with_program_args(mut self, args: Vec<String>) -> Self {
218        self.program_args = args; self
219    }
220
221    fn ensure_kind_allowed(&self, kind: &str) -> Result<(), String> {
222        if self.policy.allow_effects.contains(kind) {
223            Ok(())
224        } else {
225            Err(format!("effect `{kind}` not in --allow-effects"))
226        }
227    }
228
229    fn resolve_read_path(&self, p: &str) -> PathBuf {
230        match &self.read_root {
231            Some(root) => root.join(p.trim_start_matches('/')),
232            None => PathBuf::from(p),
233        }
234    }
235
236    /// Enforce `--allow-net-host` against an outgoing URL. Empty
237    /// allowlist = any host. Non-empty = the URL's host must match
238    /// (substring; port-agnostic) at least one entry.
239    fn ensure_host_allowed(&self, url: &str) -> Result<(), String> {
240        if self.policy.allow_net_host.is_empty() { return Ok(()); }
241        let host = extract_host(url).unwrap_or("");
242        if self.policy.allow_net_host.iter().any(|h| host == h) {
243            Ok(())
244        } else {
245            Err(format!(
246                "net call to host `{host}` not in --allow-net-host {:?}",
247                self.policy.allow_net_host,
248            ))
249        }
250    }
251}
252
253fn extract_host(url: &str) -> Option<&str> {
254    let rest = url
255        .strip_prefix("http://")
256        .or_else(|| url.strip_prefix("https://"))
257        .or_else(|| url.strip_prefix("redis://"))
258        .or_else(|| url.strip_prefix("rediss://"))
259        // `@user:pass@host:port` — strip auth prefix if present
260        .map(|r| r.split_once('@').map(|(_, after)| after).unwrap_or(r))?;
261    let host_port = match rest.find('/') {
262        Some(i) => &rest[..i],
263        None => rest,
264    };
265    Some(match host_port.rsplit_once(':') {
266        Some((h, _)) => h,
267        None => host_port,
268    })
269}
270
271fn expect_record(v: Option<&Value>) -> Result<&indexmap::IndexMap<smol_str::SmolStr, Value>, String> {
272    match v {
273        Some(Value::Record { fields: r, .. }) => Ok(r),
274        Some(other) => Err(format!("expected Record, got {other:?}")),
275        None => Err("missing Record argument".into()),
276    }
277}
278
279fn err_value(msg: String) -> Value {
280    Value::Variant { name: "Err".into(), args: vec![Value::Str(msg.into())] }
281}
282
283fn expect_str(v: Option<&Value>) -> Result<&str, String> {
284    match v {
285        Some(Value::Str(s)) => Ok(s),
286        Some(other) => Err(format!("expected Str arg, got {other:?}")),
287        None => Err("missing argument".into()),
288    }
289}
290
291fn expect_int(v: Option<&Value>) -> Result<i64, String> {
292    match v {
293        Some(Value::Int(n)) => Ok(*n),
294        Some(other) => Err(format!("expected Int arg, got {other:?}")),
295        None => Err("missing argument".into()),
296    }
297}
298
299fn ok(v: Value) -> Value {
300    Value::Variant { name: "Ok".into(), args: vec![v] }
301}
302fn err(v: Value) -> Value {
303    Value::Variant { name: "Err".into(), args: vec![v] }
304}
305
306// Root of the process content store for std.vcs (#5). Matches the store-using
307// CLI commands (branch/op): $LEX_STORE_ROOT override, else ~/.lex/store.
308fn vcs_store_root() -> std::path::PathBuf {
309    if let Ok(p) = std::env::var("LEX_STORE_ROOT") {
310        return std::path::PathBuf::from(p);
311    }
312    let home = std::env::var("HOME")
313        .map(std::path::PathBuf::from)
314        .unwrap_or_else(|_| std::path::PathBuf::from("."));
315    home.join(".lex/store")
316}
317
318fn decode_unicode_escapes(s: &str) -> String {
319    let mut result = String::with_capacity(s.len());
320    let mut chars = s.chars().peekable();
321    while let Some(c) = chars.next() {
322        if c != '\\' {
323            result.push(c);
324            continue;
325        }
326        match chars.peek() {
327            Some('u') => {
328                chars.next();
329                let hex: String = (0..4).filter_map(|_| chars.next()).collect();
330                if hex.len() == 4 {
331                    if let Ok(n) = u32::from_str_radix(&hex, 16) {
332                        if let Some(ch) = char::from_u32(n) {
333                            result.push(ch);
334                            continue;
335                        }
336                    }
337                }
338                result.push('\\');
339                result.push('u');
340                result.push_str(&hex);
341            }
342            _ => result.push(c),
343        }
344    }
345    result
346}
347
348fn some(v: Value) -> Value {
349    Value::Variant { name: "Some".into(), args: vec![v] }
350}
351fn none() -> Value {
352    Value::Variant { name: "None".into(), args: vec![] }
353}
354
355fn expect_bytes(v: Option<&Value>) -> Result<&Vec<u8>, String> {
356    match v {
357        Some(Value::Bytes(b)) => Ok(b),
358        Some(other) => Err(format!("expected Bytes arg, got {other:?}")),
359        None => Err("missing argument".into()),
360    }
361}
362
363#[allow(dead_code)]
364fn expect_str_list(v: Option<&Value>) -> Result<Vec<String>, String> {
365    match v {
366        Some(Value::List(items)) => items.iter().map(|x| match x {
367            Value::Str(s) => Ok(s.to_string()),
368            other => Err(format!("expected List[Str] element, got {other:?}")),
369        }).collect(),
370        Some(other) => Err(format!("expected List[Str], got {other:?}")),
371        None => Err("missing List[Str] argument".into()),
372    }
373}
374