Skip to main content

cuttlefish_host/
runner.rs

1//! The reactor loop: the host drives the guest, one command at a time.
2//!
3//! This module is the shape of the whole system. The guest never calls the host
4//! and waits — it returns a [`Command`], the host carries it out, and the host
5//! steps the guest again with the resulting [`Event`]. Two things follow from
6//! that inversion, and both are why it is worth the awkwardness:
7//!
8//! - **Cancellation needs no guest cooperation.** The host simply stops
9//!   stepping. A guest cannot ignore, delay, or trap its way out of it.
10//! - **Every iteration is observable.** Progress, token counts, and capability
11//!   decisions all pass through the host, even for a block whose internal loop
12//!   the DAG cannot see.
13//!
14//! The alternative — host functions the guest imports and blocks on — is not
15//! merely less tidy, it does not work: a single-threaded core-wasm guest offers
16//! no execution context for the host to call back into, and the wasmtime `Store`
17//! is `!Sync` while inference must run on a separate thread.
18
19use crate::caps::Capabilities;
20use crate::handles::Handles;
21use crate::infer::{InferBackend, InferRequest, InferResult};
22use cuttlefish_abi::{error_codes, Command, Envelope, Event, JobError, JobStatus, Usage};
23use std::sync::atomic::{AtomicBool, Ordering};
24use std::sync::Arc;
25use std::time::Instant;
26use tokio::sync::mpsc;
27use tokio_util::sync::CancellationToken;
28use wasmtime::{Engine, Instance, Linker, Memory, Module, Store, TypedFunc};
29
30/// Width, in pixels, that document pages render to.
31///
32/// Vision models work from a fixed-size input anyway, and a larger raster costs
33/// encode time and tokens without adding detail the model can use.
34const RENDER_WIDTH: u16 = 1024;
35
36/// Something worth telling a watcher about while a job runs.
37#[derive(Debug, Clone)]
38pub enum JobEvent {
39    /// One generated token.
40    Token(String),
41    /// Guest-supplied progress.
42    Progress(serde_json::Value),
43}
44
45/// Everything needed to run one job.
46pub struct JobSpec {
47    /// The checked graph, in topological order — safe to execute
48    /// front-to-back, threading `outputs` forward.
49    pub nodes: Vec<crate::dag::CheckedNode>,
50    /// Which nodes are exclusive to which branch decision+label — see
51    /// `crate::dag::BranchExclusivity`.
52    pub exclusive_to: std::collections::HashMap<String, crate::dag::BranchExclusivity>,
53    /// The job's input, handed to every entry node (a node with no `input`
54    /// expression — `node.input.is_none()`).
55    pub input: serde_json::Value,
56    /// What this job is permitted to reach.
57    pub caps: Capabilities,
58}
59
60/// Pointer width of a guest module, read from the module rather than assumed.
61///
62/// Only [`Abi::W32`] is supported today; 64-bit guests are rejected with a clear
63/// message. The enum exists anyway so that adding wasm64 later is a new arm plus
64/// a second set of [`TypedFunc`] signatures, rather than a hunt through this
65/// file for every place a pointer was assumed to be four bytes wide.
66#[derive(Debug, Clone, Copy, PartialEq, Eq)]
67pub enum Abi {
68    /// 32-bit linear memory.
69    W32,
70    /// 64-bit linear memory (memory64).
71    W64,
72}
73
74impl Abi {
75    /// Size of one pointer-sized field, and so half a descriptor.
76    fn ptr_size(self) -> usize {
77        match self {
78            Abi::W32 => 4,
79            Abi::W64 => 8,
80        }
81    }
82}
83
84struct Guest {
85    store: Store<()>,
86    memory: Memory,
87    abi: Abi,
88    alloc: TypedFunc<u32, u32>,
89    init: TypedFunc<(u32, u32), u32>,
90    step: TypedFunc<(u32, u32), u32>,
91    on_token: Option<TypedFunc<(u32, u32), i32>>,
92}
93
94impl Guest {
95    fn new(
96        engine: &Engine,
97        cache: &crate::module_cache::ModuleCache,
98        module_bytes: &[u8],
99    ) -> anyhow::Result<Self> {
100        let module = cache.compile(engine, module_bytes)?;
101
102        // An empty linker, deliberately. Guest blocks are built for
103        // `wasm32-unknown-unknown` and import nothing at all — a wasip1 guest
104        // would drag in `fd_write` and `proc_exit` through its panic path alone
105        // and fail to instantiate here.
106        let linker: Linker<()> = Linker::new(engine);
107        let mut store = Store::new(engine, ());
108        let instance: Instance = linker.instantiate(&mut store, &module)?;
109
110        let memory = instance
111            .get_memory(&mut store, "memory")
112            .ok_or_else(|| anyhow::anyhow!("guest exports no memory"))?;
113
114        // Width comes from the module itself. A 64-bit guest exports `cf_init`
115        // as `(i64, i64) -> i64`, so the typed lookups below would otherwise
116        // fail with a signature mismatch that says nothing about the real cause.
117        let abi = if memory.ty(&store).is_64() {
118            Abi::W64
119        } else {
120            Abi::W32
121        };
122        if abi == Abi::W64 {
123            anyhow::bail!("guest uses 64-bit memory; only 32-bit guests are supported");
124        }
125
126        Ok(Self {
127            alloc: instance.get_typed_func(&mut store, "cf_alloc")?,
128            init: instance.get_typed_func(&mut store, "cf_init")?,
129            step: instance.get_typed_func(&mut store, "cf_step")?,
130            // Optional: a block indifferent to streaming need not export it.
131            on_token: instance.get_typed_func(&mut store, "cf_on_token").ok(),
132            memory,
133            abi,
134            store,
135        })
136    }
137
138    fn write(&mut self, bytes: &[u8]) -> anyhow::Result<(u32, u32)> {
139        let len = bytes.len() as u32;
140        let ptr = self.alloc.call(&mut self.store, len)?;
141        self.memory.write(&mut self.store, ptr as usize, bytes)?;
142        Ok((ptr, len))
143    }
144
145    /// Read the descriptor the guest returned, then the payload it points at.
146    ///
147    /// Two reads rather than unpacking one integer — the cost of keeping these
148    /// signatures identical across pointer widths.
149    fn read_desc(&mut self, desc_ptr: u32) -> anyhow::Result<Vec<u8>> {
150        let w = self.abi.ptr_size();
151        let mut desc = vec![0u8; 2 * w];
152        self.memory
153            .read(&mut self.store, desc_ptr as usize, &mut desc)?;
154
155        let field = |bytes: &[u8]| -> u64 {
156            match w {
157                4 => u32::from_le_bytes(bytes.try_into().expect("4 bytes")) as u64,
158                _ => u64::from_le_bytes(bytes.try_into().expect("8 bytes")),
159            }
160        };
161        let ptr = field(&desc[..w]) as usize;
162        let len = field(&desc[w..]) as usize;
163
164        let mut buf = vec![0u8; len];
165        self.memory.read(&mut self.store, ptr, &mut buf)?;
166        Ok(buf)
167    }
168
169    fn call_init(&mut self, input: &serde_json::Value) -> anyhow::Result<Command> {
170        let bytes = serde_json::to_vec(input)?;
171        let (ptr, len) = self.write(&bytes)?;
172        let desc = self.init.call(&mut self.store, (ptr, len))?;
173        Ok(serde_json::from_slice(&self.read_desc(desc)?)?)
174    }
175
176    fn call_step(&mut self, event: &Event) -> anyhow::Result<Command> {
177        let bytes = serde_json::to_vec(event)?;
178        let (ptr, len) = self.write(&bytes)?;
179        let desc = self.step.call(&mut self.store, (ptr, len))?;
180        Ok(serde_json::from_slice(&self.read_desc(desc)?)?)
181    }
182
183    /// Ask the guest whether generation should continue.
184    fn call_on_token(&mut self, token: &str) -> anyhow::Result<bool> {
185        // Cloned rather than moved: wasmtime's TypedFunc is Clone but not Copy,
186        // and cloning also ends the borrow of `self` before `write` needs it
187        // mutably.
188        let Some(f) = self.on_token.clone() else {
189            return Ok(true);
190        };
191        let (ptr, len) = self.write(token.as_bytes())?;
192        Ok(f.call(&mut self.store, (ptr, len))? == 0)
193    }
194}
195
196/// Read a block's declared signature out of a compiled module.
197///
198/// Instantiates the module and calls its `cf_signature` export. That is heavier
199/// than parsing a sidecar file, and it is the point: the answer comes from the
200/// artifact that will actually run, so it cannot describe a different version of
201/// the block than the one being checked.
202///
203/// A block built before signatures existed has no such export. That is not an
204/// error — it reports the permissive default, so an older block still composes,
205/// just without the seam being checked.
206pub fn read_signature(
207    engine: &Engine,
208    module_bytes: &[u8],
209) -> anyhow::Result<cuttlefish_abi::Signature> {
210    let permissive = cuttlefish_abi::Signature {
211        input: cuttlefish_abi::Ty::Json,
212        output: cuttlefish_abi::Ty::Json,
213    };
214
215    // Deliberately does not go through `Guest`, which requires the whole reactor
216    // — alloc, init, step. Reading a declaration should not demand that a module
217    // be runnable: a block missing an export has a real problem, but it is one
218    // worth reporting when the job runs, with the job's error handling, rather
219    // than as a confusing failure during a typecheck.
220    let module = Module::new(engine, module_bytes)?;
221    let linker: Linker<()> = Linker::new(engine);
222    let mut store = Store::new(engine, ());
223    let instance = linker.instantiate(&mut store, &module)?;
224
225    let Ok(signature) = instance.get_typed_func::<(), u32>(&mut store, "cf_signature") else {
226        return Ok(permissive);
227    };
228    let Some(memory) = instance.get_memory(&mut store, "memory") else {
229        return Ok(permissive);
230    };
231
232    let desc_ptr = signature.call(&mut store, ())? as usize;
233    let mut desc = [0u8; 8];
234    memory.read(&mut store, desc_ptr, &mut desc)?;
235    let ptr = u32::from_le_bytes(desc[..4].try_into().expect("4 bytes")) as usize;
236    let len = u32::from_le_bytes(desc[4..].try_into().expect("4 bytes")) as usize;
237
238    let mut buf = vec![0u8; len];
239    memory.read(&mut store, ptr, &mut buf)?;
240    Ok(serde_json::from_slice(&buf)?)
241}
242
243fn fail(code: &str, message: impl Into<String>, usage: Usage) -> Envelope {
244    Envelope {
245        status: JobStatus::Failed,
246        result: None,
247        error: Some(JobError {
248            code: code.into(),
249            message: message.into(),
250        }),
251        usage,
252    }
253}
254
255/// Whether an `InputExpr` (transitively) references any node in `skipped`.
256fn references_any(
257    expr: &cuttlefish_core::graph::InputExpr,
258    skipped: &std::collections::HashSet<String>,
259) -> bool {
260    use cuttlefish_core::graph::InputExpr;
261    match expr {
262        InputExpr::FromNode(n) => skipped.contains(n),
263        InputExpr::Record(fields) => fields.values().any(|e| references_any(e, skipped)),
264        InputExpr::List(items) => items.iter().any(|e| references_any(e, skipped)),
265    }
266}
267
268/// Compose an `InputExpr` into an actual JSON value, by looking up each
269/// referenced node's already-produced output. Mirrors `dag::evaluate_expr_ty`
270/// (which does the same composition at the *type* level, at check time) —
271/// this is the runtime analogue.
272fn evaluate_input(
273    expr: &cuttlefish_core::graph::InputExpr,
274    outputs: &std::collections::HashMap<String, serde_json::Value>,
275) -> serde_json::Value {
276    use cuttlefish_core::graph::InputExpr;
277    match expr {
278        InputExpr::FromNode(n) => outputs.get(n).cloned().unwrap_or(serde_json::Value::Null),
279        InputExpr::Record(fields) => {
280            let mut map = serde_json::Map::new();
281            for (k, v) in fields {
282                map.insert(k.clone(), evaluate_input(v, outputs));
283            }
284            serde_json::Value::Object(map)
285        }
286        InputExpr::List(items) => {
287            serde_json::Value::Array(items.iter().map(|e| evaluate_input(e, outputs)).collect())
288        }
289    }
290}
291
292fn cancelled(usage: Usage, message: &str) -> Envelope {
293    Envelope {
294        status: JobStatus::Cancelled,
295        result: None,
296        error: Some(JobError {
297            code: error_codes::CANCELLED.into(),
298            message: message.into(),
299        }),
300        usage,
301    }
302}
303
304/// Drive one job to completion.
305///
306/// Always returns an [`Envelope`]; failures are values, not errors, because the
307/// caller has to report *something* to whoever submitted the job.
308///
309/// `ledger` is consulted before any resume-sensitive decision (branch-skip,
310/// transitive-skip, or actually running a node) and written to immediately
311/// after that decision is made, so that a process restart mid-job — a later
312/// task's concern, not this function's — can resume from exactly the state
313/// this function left behind. The whole body runs inside a single labeled
314/// block (`'run: { ... }`) so that every exit path, however it got there,
315/// still reaches the one `ledger.finish(...)` call at the end.
316pub async fn run_job(
317    engine: Arc<Engine>,
318    backend: Arc<dyn InferBackend>,
319    job: JobSpec,
320    events: mpsc::Sender<JobEvent>,
321    cancel: CancellationToken,
322    ledger: &crate::ledger::Ledger,
323    cache: &crate::module_cache::ModuleCache,
324) -> Envelope {
325    let started = Instant::now();
326    let mut usage = Usage {
327        model: backend.model_name(),
328        ..Usage::default()
329    };
330
331    // Dropped when this function returns, closing every file the job opened.
332    // That job-scoped lifetime is what makes handles unforgeable across jobs.
333    //
334    // Shared across stages on purpose: a handle produced by one block — a
335    // rendered page, say — stays usable by the next. Confining it to one stage
336    // would make a pipeline strictly weaker than a single block that did the
337    // same work, while adding nothing, since the job boundary is what the
338    // security property rests on.
339    let mut handles = Handles::default();
340
341    let envelope = 'run: {
342        if job.nodes.is_empty() {
343            usage.duration_ms = started.elapsed().as_millis() as u64;
344            break 'run fail(
345                error_codes::SCHEMA_VALIDATION_FAILED,
346                "this job has no nodes to run",
347                usage,
348            );
349        }
350
351        let total = job.nodes.len();
352        let mut outputs: std::collections::HashMap<String, serde_json::Value> =
353            std::collections::HashMap::new();
354        let mut skipped: std::collections::HashSet<String> = std::collections::HashSet::new();
355        let mut route_taken: std::collections::HashMap<String, String> =
356            std::collections::HashMap::new();
357
358        for (index, node) in job.nodes.iter().enumerate() {
359            // Resume: a node the ledger already marked skipped stays skipped
360            // — its branch decision is not re-evaluated.
361            match ledger.is_skipped(&node.name) {
362                Ok(true) => {
363                    skipped.insert(node.name.clone());
364                    continue;
365                }
366                Ok(false) => {}
367                Err(e) => {
368                    usage.duration_ms = started.elapsed().as_millis() as u64;
369                    break 'run fail(
370                        error_codes::SCHEMA_VALIDATION_FAILED,
371                        format!("reading ledger skip state for node `{}`: {e}", node.name),
372                        usage,
373                    );
374                }
375            }
376
377            // Resume: a completed checkpoint means reuse the cached output
378            // instead of re-running.
379            match ledger.get_completed(&node.name) {
380                Ok(Some(cached)) => {
381                    // If this was a branches decision node, its route must be
382                    // recorded in `route_taken` here too — not just after a
383                    // fresh `run_stage` call below — or a downstream,
384                    // not-yet-reached branch-exclusive node would see no
385                    // recorded decision at all on a resumed run and
386                    // (incorrectly) never be skipped. The checkpoint only
387                    // ever holds a value that already passed this same route
388                    // validation the first time it ran, so this is purely
389                    // re-deriving `route_taken`, never re-validating.
390                    if let Some(route) = cached.get("route").and_then(|v| v.as_str()) {
391                        route_taken.insert(node.name.clone(), route.to_string());
392                    }
393                    outputs.insert(node.name.clone(), cached);
394                    continue;
395                }
396                Ok(None) => {}
397                Err(e) => {
398                    usage.duration_ms = started.elapsed().as_millis() as u64;
399                    break 'run fail(
400                        error_codes::SCHEMA_VALIDATION_FAILED,
401                        format!("reading ledger checkpoint for node `{}`: {e}", node.name),
402                        usage,
403                    );
404                }
405            }
406
407            // Tell watchers which node is running. A pipeline that stalls is
408            // much easier to diagnose when the stream says where.
409            if total > 1 {
410                let _ = events
411                    .send(JobEvent::Progress(serde_json::json!({
412                        "stage": index + 1,
413                        "of": total,
414                        "node": node.name,
415                    })))
416                    .await;
417            }
418
419            // Fresh branch-skip decision (only reached if the ledger had no
420            // recorded state for this node — i.e. this is either a fresh
421            // run, or a resumed run that hadn't reached this node yet).
422            //
423            // Branch-skip: this node is exclusive to a decision+label, and
424            // that decision's chosen route (already recorded earlier in this
425            // same loop, since the branching node is always topologically
426            // before the nodes exclusive to its labels) doesn't match.
427            if let Some(ex) = job.exclusive_to.get(&node.name) {
428                if let Some(taken) = route_taken.get(&ex.decision) {
429                    if taken != &ex.label {
430                        skipped.insert(node.name.clone());
431                        if let Err(e) = ledger.write_skipped(&node.name) {
432                            usage.duration_ms = started.elapsed().as_millis() as u64;
433                            break 'run fail(
434                                error_codes::SCHEMA_VALIDATION_FAILED,
435                                format!("recording skip for node `{}` in ledger: {e}", node.name),
436                                usage,
437                            );
438                        }
439                        continue;
440                    }
441                }
442            }
443            // Transitive skip: this node's input needs a skipped node's
444            // output.
445            if let Some(expr) = &node.input {
446                if references_any(expr, &skipped) {
447                    skipped.insert(node.name.clone());
448                    if let Err(e) = ledger.write_skipped(&node.name) {
449                        usage.duration_ms = started.elapsed().as_millis() as u64;
450                        break 'run fail(
451                            error_codes::SCHEMA_VALIDATION_FAILED,
452                            format!("recording skip for node `{}` in ledger: {e}", node.name),
453                            usage,
454                        );
455                    }
456                    continue;
457                }
458            }
459
460            let node_input = match &node.input {
461                None => job.input.clone(),
462                Some(expr) => evaluate_input(expr, &outputs),
463            };
464
465            // Bounded loop: repeat_until re-invokes this node on its own
466            // prior output, up to max_iterations. A node without
467            // repeat_until runs exactly once (the loop's `None` arm breaks
468            // immediately).
469            let mut current_input = node_input;
470            let mut iterations: u32 = 0;
471            let result = loop {
472                let r = match run_stage(
473                    &engine,
474                    cache,
475                    &backend,
476                    &node.module_bytes,
477                    current_input.clone(),
478                    node.script.as_deref(),
479                    &job.caps,
480                    &mut handles,
481                    &events,
482                    &cancel,
483                    &mut usage,
484                    started,
485                    index,
486                )
487                .await
488                {
489                    Ok(v) => v,
490                    Err(envelope) => break 'run envelope,
491                };
492                match &node.repeat_until {
493                    None => break r,
494                    Some(field) => {
495                        iterations += 1;
496                        let done = r.get(field).and_then(|v| v.as_str()) == Some("done");
497                        if done {
498                            break r;
499                        }
500                        let max = node
501                            .max_iterations
502                            .expect("repeat_until requires max_iterations, enforced at parse time");
503                        if iterations >= max {
504                            usage.duration_ms = started.elapsed().as_millis() as u64;
505                            break 'run fail(
506                                error_codes::SCHEMA_VALIDATION_FAILED,
507                                format!(
508                                    "node `{}` did not reach repeat_until=\"done\" within max_iterations={max}",
509                                    node.name
510                                ),
511                                usage,
512                            );
513                        }
514                        current_input = r;
515                    }
516                }
517            };
518
519            // A branching node: read its `route` field and record the
520            // decision for later nodes' branch-skip check (top of this
521            // loop).
522            //
523            // Whether this node is actually a `branches` decision at all is
524            // determined from `job.exclusive_to` rather than by threading
525            // `Branches` itself into `JobSpec`: `dag::compute_branch_exclusivity`
526            // seeds an entry for every `(label, target)` pair of every declared
527            // decision, so the set of `label`s across all `exclusive_to` entries
528            // whose `decision` names this node IS exactly this decision's full,
529            // valid label set. An unmatched or missing `route` on a genuine
530            // decision node is a job failure — loud, not a silent no-op — per
531            // the design spec's "Conditional dispatch" section.
532            let valid_labels: Vec<&str> = job
533                .exclusive_to
534                .values()
535                .filter(|ex| ex.decision == node.name)
536                .map(|ex| ex.label.as_str())
537                .collect();
538            if !valid_labels.is_empty() {
539                match result.get("route").and_then(|v| v.as_str()) {
540                    Some(route) if valid_labels.contains(&route) => {
541                        route_taken.insert(node.name.clone(), route.to_string());
542                    }
543                    Some(route) => {
544                        usage.duration_ms = started.elapsed().as_millis() as u64;
545                        break 'run fail(
546                            error_codes::SCHEMA_VALIDATION_FAILED,
547                            format!(
548                                "node `{}` produced route \"{route}\", which doesn't match any \
549                                 declared label for this branches decision",
550                                node.name
551                            ),
552                            usage,
553                        );
554                    }
555                    None => {
556                        usage.duration_ms = started.elapsed().as_millis() as u64;
557                        break 'run fail(
558                            error_codes::SCHEMA_VALIDATION_FAILED,
559                            format!(
560                                "node `{}` is a branches decision but its output has no `route` field",
561                                node.name
562                            ),
563                            usage,
564                        );
565                    }
566                }
567            } else if let Some(route) = result.get("route").and_then(|v| v.as_str()) {
568                // Not a declared decision node, but its output happens to
569                // carry a route field anyway — harmless, record it
570                // defensively (no downstream node's exclusive_to can
571                // reference this decision name if it isn't a real decision,
572                // so this is inert either way, just avoids silently
573                // dropping information that happens to be present).
574                route_taken.insert(node.name.clone(), route.to_string());
575            }
576
577            // Checkpoint a genuinely-executed node's result before recording
578            // it in `outputs` — the ledger write must happen before the
579            // in-memory outputs map update so a crash between the two can't
580            // leave the ledger silently behind the in-memory state (the
581            // in-memory state doesn't survive a crash anyway, so
582            // ledger-first is the only order that matters for durability;
583            // the reverse order would just be a smaller window for the same
584            // underlying property, not correctness).
585            if let Err(e) = ledger.write_completed(&node.name, &result) {
586                usage.duration_ms = started.elapsed().as_millis() as u64;
587                break 'run fail(
588                    error_codes::SCHEMA_VALIDATION_FAILED,
589                    format!(
590                        "recording checkpoint for node `{}` in ledger: {e}",
591                        node.name
592                    ),
593                    usage,
594                );
595            }
596            outputs.insert(node.name.clone(), result);
597        }
598
599        usage.duration_ms = started.elapsed().as_millis() as u64;
600
601        // What is "the" job result for a graph? Convention: the output of the
602        // LAST node in topological order that wasn't skipped. This exactly
603        // matches today's linear-pipeline behavior in the degenerate case (a
604        // linear chain's last node in topo order is its sole sink), and
605        // generalizes sensibly to a branching graph (exactly one path executes,
606        // so there's still a well-defined "last node that actually ran") and to
607        // a fan-in graph without branches (the join/sink node is last in topo
608        // order, since nothing depends on it).
609        //
610        // Known limitation: for a graph with two or more independent sinks (no
611        // edge between them at all — `cuttlefishd`'s daemon path genuinely
612        // accepts such graphs, unlike `cuttlefish build`, which restricts itself
613        // to linear graphs), "last in topological order" is decided by the
614        // topological sort's tie-break (alphabetical node name among ready
615        // nodes), not by any deliberate semantic answer about which sink's
616        // output should represent the job. Don't mistake that tie-break for a
617        // considered multi-sink policy — it isn't one.
618        let result = job
619            .nodes
620            .iter()
621            .rev()
622            .find_map(|n| outputs.get(&n.name).cloned());
623
624        match result {
625            Some(value) => Envelope {
626                status: JobStatus::Completed,
627                result: Some(value),
628                error: None,
629                usage,
630            },
631            None => fail(
632                error_codes::SCHEMA_VALIDATION_FAILED,
633                "every node in this job was skipped; there is no result",
634                usage,
635            ),
636        }
637    };
638
639    let ledger_status = match envelope.status {
640        JobStatus::Completed => "completed",
641        JobStatus::Failed => "failed",
642        JobStatus::Cancelled => "cancelled",
643        _ => "running", // shouldn't occur — defensive only
644    };
645    if let Err(e) = ledger.finish(ledger_status) {
646        // A ledger write failure at this final step doesn't invalidate the
647        // job's already-computed, correct result — losing it here means a
648        // wrong Interrupted-detection on next restart, not a wrong result
649        // now. Loud enough to notice, not loud enough to discard real work
650        // that already succeeded or failed for its own, unrelated reason.
651        eprintln!("warning: failed to record job {ledger_status} status in ledger: {e}");
652    }
653    envelope
654}
655
656/// Run one block to completion, returning what it produced.
657///
658/// `Err` carries a finished [`Envelope`]: a stage that fails ends the whole job,
659/// because a later stage's input is the earlier one's output and there is
660/// nothing sensible to feed it.
661#[allow(clippy::too_many_arguments)]
662async fn run_stage(
663    engine: &Engine,
664    cache: &crate::module_cache::ModuleCache,
665    backend: &Arc<dyn InferBackend>,
666    module_bytes: &[u8],
667    input: serde_json::Value,
668    script: Option<&str>,
669    caps: &Capabilities,
670    handles: &mut Handles,
671    events: &mpsc::Sender<JobEvent>,
672    cancel: &CancellationToken,
673    usage: &mut Usage,
674    started: Instant,
675    stage_index: usize,
676) -> Result<serde_json::Value, Envelope> {
677    // Naming the stage turns "the job failed" into "the second block failed",
678    // which is the difference between a usable error and a hunt.
679    let blame = |message: String| -> String {
680        if stage_index == 0 {
681            message
682        } else {
683            format!("block {} of the pipeline: {message}", stage_index + 1)
684        }
685    };
686    // Documents are read from their path rather than their descriptor — both
687    // extraction and rendering want a file. Kept beside the handle table so the
688    // two are dropped together at the end of the job.
689    let mut doc_paths: std::collections::HashMap<u32, std::path::PathBuf> =
690        std::collections::HashMap::new();
691
692    let mut guest = match Guest::new(engine, cache, module_bytes) {
693        Ok(g) => g,
694        Err(e) => {
695            return Err(fail(
696                error_codes::WASM_TRAP,
697                blame(e.to_string()),
698                usage.clone(),
699            ))
700        }
701    };
702
703    // A Script-kind stage's `module_bytes` is always the shared interpreter
704    // (see `pipeline::resolve_and_load`), which expects its script text
705    // wrapped alongside the real job input — the interpreter itself never
706    // receives the raw input directly, and this is the one place per job
707    // where that wrapping actually happens, since the script is fixed at
708    // catalog time but the input is only known per job.
709    let input = match script {
710        Some(script) => serde_json::json!({
711            "__cuttlefish_script": script,
712            "input": input,
713        }),
714        None => input,
715    };
716
717    let mut command = match guest.call_init(&input) {
718        Ok(c) => c,
719        Err(e) => {
720            return Err(fail(
721                error_codes::WASM_TRAP,
722                blame(e.to_string()),
723                usage.clone(),
724            ))
725        }
726    };
727
728    loop {
729        if cancel.is_cancelled() {
730            usage.duration_ms = started.elapsed().as_millis() as u64;
731            return Err(cancelled(usage.clone(), "job cancelled"));
732        }
733
734        let event = match command {
735            // Ends this stage, not the job: the value becomes the next
736            // block's input, or the job's result if this was the last.
737            Command::Done { result } => return Ok(result),
738            Command::Fail { code, message } => {
739                usage.duration_ms = started.elapsed().as_millis() as u64;
740                return Err(fail(&code, blame(message), usage.clone()));
741            }
742            Command::Emit { progress } => {
743                let _ = events.send(JobEvent::Progress(progress)).await;
744                Event::Emitted
745            }
746
747            // The capability check lives here, at Open, and nowhere else. Slice
748            // takes a handle rather than a path, and handles are job-scoped, so
749            // there is no second place a path can enter the system.
750            Command::Open { path } => {
751                let p = std::path::PathBuf::from(&path);
752                if !caps.allows_read(&p) {
753                    usage.duration_ms = started.elapsed().as_millis() as u64;
754                    return Err(fail(
755                        error_codes::CAPABILITY_DENIED,
756                        format!("read not permitted: {path}"),
757                        usage.clone(),
758                    ));
759                }
760                match handles.open(&p) {
761                    Ok((handle, len, kind)) => {
762                        // A PDF's page count and text layer need the whole file,
763                        // which the handle layer deliberately does not read. Ask
764                        // the document layer, and fall back to the plain kind if
765                        // it cannot answer — a malformed PDF is still a file a
766                        // block may want to read bytes from.
767                        let kind = match kind {
768                            cuttlefish_abi::MediaKind::Document { .. } => {
769                                match crate::documents::inspect(&p) {
770                                    Ok(info) => cuttlefish_abi::MediaKind::Document {
771                                        pages: info.pages,
772                                        has_text_layer: info.has_text_layer,
773                                    },
774                                    Err(_) => cuttlefish_abi::MediaKind::Binary,
775                                }
776                            }
777                            other => other,
778                        };
779                        // Remember the path: rendering and text extraction work
780                        // from a file, not from the open descriptor.
781                        doc_paths.insert(handle, p.clone());
782                        Event::Opened { handle, len, kind }
783                    }
784                    Err(e) => {
785                        usage.duration_ms = started.elapsed().as_millis() as u64;
786                        return Err(fail(
787                            error_codes::CAPABILITY_DENIED,
788                            e.to_string(),
789                            usage.clone(),
790                        ));
791                    }
792                }
793            }
794
795            Command::Slice {
796                handle,
797                offset,
798                len,
799            } => match handles.slice(handle, offset, len) {
800                Ok(w) => Event::Sliced {
801                    text: w.text,
802                    next_offset: w.next_offset,
803                },
804                Err(e) => {
805                    usage.duration_ms = started.elapsed().as_millis() as u64;
806                    return Err(fail(
807                        error_codes::CAPABILITY_DENIED,
808                        e.to_string(),
809                        usage.clone(),
810                    ));
811                }
812            },
813
814            Command::SliceBytes {
815                handle,
816                offset,
817                len,
818            } => match handles.slice_bytes(handle, offset, len) {
819                Ok((bytes, next_offset)) => {
820                    use base64::Engine;
821                    Event::SlicedBytes {
822                        bytes_base64: base64::engine::general_purpose::STANDARD.encode(&bytes),
823                        next_offset,
824                    }
825                }
826                Err(e) => {
827                    usage.duration_ms = started.elapsed().as_millis() as u64;
828                    return Err(fail(
829                        error_codes::CAPABILITY_DENIED,
830                        e.to_string(),
831                        usage.clone(),
832                    ));
833                }
834            },
835
836            Command::PageText { handle, page } => {
837                let Some(path) = doc_paths.get(&handle).cloned() else {
838                    usage.duration_ms = started.elapsed().as_millis() as u64;
839                    return Err(fail(
840                        error_codes::CAPABILITY_DENIED,
841                        format!("no such handle: {handle}"),
842                        usage.clone(),
843                    ));
844                };
845                match crate::documents::page_text(&path, page) {
846                    Ok(text) => Event::PageTexted { text },
847                    Err(e) => {
848                        usage.duration_ms = started.elapsed().as_millis() as u64;
849                        return Err(fail(error_codes::UNSUPPORTED, e.to_string(), usage.clone()));
850                    }
851                }
852            }
853
854            Command::PageImage { handle, page } => {
855                let Some(path) = doc_paths.get(&handle).cloned() else {
856                    usage.duration_ms = started.elapsed().as_millis() as u64;
857                    return Err(fail(
858                        error_codes::CAPABILITY_DENIED,
859                        format!("no such handle: {handle}"),
860                        usage.clone(),
861                    ));
862                };
863                // A rendered page becomes a handle like any other, so it can be
864                // named in Infer exactly as a file-backed image would be.
865                match crate::documents::render_page(&path, page, RENDER_WIDTH) {
866                    Ok(png) => {
867                        let (handle, len) = handles.insert_bytes(
868                            png,
869                            cuttlefish_abi::MediaKind::Image {
870                                format: "png".into(),
871                            },
872                        );
873                        Event::PageImaged { handle, len }
874                    }
875                    Err(e) => {
876                        usage.duration_ms = started.elapsed().as_millis() as u64;
877                        return Err(fail(error_codes::UNSUPPORTED, e.to_string(), usage.clone()));
878                    }
879                }
880            }
881
882            Command::Infer {
883                prompt,
884                max_tokens,
885                images,
886            } => {
887                // Images are named by handle; the host loads the bytes, so they
888                // never pass through guest memory.
889                // Refuse rather than drop. Sending images to a backend that
890                // cannot use them produces a confident answer about nothing,
891                // which reads as a bad model rather than a misconfigured job —
892                // and the caller has no way to tell the difference.
893                if !images.is_empty() && !backend.supports_images() {
894                    usage.duration_ms = started.elapsed().as_millis() as u64;
895                    return Err(fail(
896                        error_codes::UNSUPPORTED,
897                        format!(
898                            "this job supplied {} image(s), but the backend serving `{}` cannot \
899                             accept them. Use a vision-capable model through the `ollama` \
900                             provider, or change the block to send text only.",
901                            images.len(),
902                            backend.model_name()
903                        ),
904                        usage.clone(),
905                    ));
906                }
907                let mut image_bytes = Vec::with_capacity(images.len());
908                for handle in &images {
909                    match handles.read_all(*handle) {
910                        Ok(bytes) => image_bytes.push(bytes),
911                        Err(e) => {
912                            usage.duration_ms = started.elapsed().as_millis() as u64;
913                            return Err(fail(
914                                error_codes::CAPABILITY_DENIED,
915                                e.to_string(),
916                                usage.clone(),
917                            ));
918                        }
919                    }
920                }
921
922                // Tokens must reach the guest *while* generation runs, because
923                // the guest's Stop verdict is what ends it early. The wasmtime
924                // Store is !Sync and cannot be touched from inside the backend's
925                // callback, so a channel carries tokens out and a shared flag
926                // carries the verdict back — without sharing the Store.
927                let (tx, mut rx) = mpsc::unbounded_channel::<String>();
928                let stop = Arc::new(AtomicBool::new(false));
929                let sink_stop = stop.clone();
930                let mut sink = move |t: &str| {
931                    tx.send(t.to_string()).is_ok() && !sink_stop.load(Ordering::Relaxed)
932                };
933
934                let mut trap: Option<String> = None;
935                let outcome: Option<anyhow::Result<InferResult>> = {
936                    let request = InferRequest {
937                        prompt: &prompt,
938                        max_tokens,
939                        images: &image_bytes,
940                    };
941                    let infer = backend.infer(request, &mut sink);
942                    tokio::pin!(infer);
943                    loop {
944                        tokio::select! {
945                            biased;
946                            _ = cancel.cancelled() => break None,
947                            Some(tok) = rx.recv() => {
948                                let _ = events.send(JobEvent::Token(tok.clone())).await;
949                                match guest.call_on_token(&tok) {
950                                    Ok(true) => {}
951                                    Ok(false) => stop.store(true, Ordering::Relaxed),
952                                    Err(e) => {
953                                        trap = Some(e.to_string());
954                                        break None;
955                                    }
956                                }
957                            }
958                            r = &mut infer => break Some(r),
959                        }
960                    }
961                };
962
963                if let Some(message) = trap {
964                    usage.duration_ms = started.elapsed().as_millis() as u64;
965                    return Err(fail(error_codes::WASM_TRAP, message, usage.clone()));
966                }
967
968                // Tokens generated in the same poll as the last one are still
969                // queued; forward them so the stream is complete.
970                while let Ok(tok) = rx.try_recv() {
971                    let _ = events.send(JobEvent::Token(tok)).await;
972                }
973
974                match outcome {
975                    None => {
976                        usage.duration_ms = started.elapsed().as_millis() as u64;
977                        return Err(cancelled(usage.clone(), "cancelled during inference"));
978                    }
979                    Some(Err(e)) => {
980                        usage.duration_ms = started.elapsed().as_millis() as u64;
981                        return Err(fail(
982                            error_codes::MODEL_LOAD_FAILED,
983                            e.to_string(),
984                            usage.clone(),
985                        ));
986                    }
987                    Some(Ok(r)) => {
988                        usage.tokens_in += r.tokens_in;
989                        usage.tokens_out += r.tokens_out;
990                        Event::InferDone {
991                            text: r.text,
992                            tokens_out: r.tokens_out,
993                        }
994                    }
995                }
996            }
997        };
998
999        command = match guest.call_step(&event) {
1000            Ok(c) => c,
1001            Err(e) => {
1002                usage.duration_ms = started.elapsed().as_millis() as u64;
1003                return Err(fail(
1004                    error_codes::WASM_TRAP,
1005                    blame(e.to_string()),
1006                    usage.clone(),
1007                ));
1008            }
1009        };
1010    }
1011}