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}