tensor_wasm_wasi_gpu/streaming.rs
1// SPDX-License-Identifier: Apache-2.0
2// Copyright 2026 Craton Software Company
3
4//! `wasi:tensor/host` streaming surface (roadmap feature #2).
5//!
6//! Lets guests emit chunks of output that the host gateway flushes to the
7//! HTTP client as SSE or chunked-transfer bytes. v0.3.7 lands the WIT
8//! contract + a buffered in-memory channel; v0.4 lands actual streaming
9//! through the axum response path.
10//!
11//! ## Surface
12//!
13//! [`StreamingContext`] owns an optional Tokio `mpsc::Sender<Vec<u8>>`
14//! channel. The host functions wrapped by [`add_streaming_to_linker`]
15//! call into [`StreamingContext::emit_chunk`] and
16//! [`StreamingContext::flush`]; the API gateway holds the matching
17//! `mpsc::Receiver` and forwards the chunks into the axum SSE / chunked
18//! response body.
19//!
20//! ## Caps
21//!
22//! Two hard caps are enforced before any chunk is forwarded:
23//!
24//! * [`MAX_CHUNK_BYTES`] (64 KiB) — single-chunk size cap. Currently the
25//! parser-side check is left for the v0.4 host-fn implementation (the
26//! wasi-tensor WIT signature does not yet carry a per-call cap distinct
27//! from the total cap); the constant is exported so the host-fn wrapper
28//! and tests share a single source of truth.
29//! * [`MAX_TOTAL_STREAM_BYTES`] (64 MiB) — total bytes a single
30//! invocation may emit before the host returns the documented
31//! `-2 = cap exceeded` code from `emit-chunk`.
32//!
33//! Caps are intentionally conservative for the scaffold so a runaway
34//! guest cannot exhaust gateway memory while v0.4 is in flight.
35//!
36//! ## Error codes
37//!
38//! Mirrors the WIT contract in `wit/wasi-tensor.wit`. Acceptance is
39//! all-or-nothing — the host forwards the whole chunk or none of it, so
40//! success is a flat `0` rather than a byte count:
41//! * `0` — chunk fully accepted (the entire payload was forwarded).
42//! * `-1` — streaming not enabled for this invocation.
43//! * `-2` — guest tried to emit past the documented size cap.
44//! * `-3` — downstream client disconnected (receiver dropped).
45//!
46//! Tests in `tests/streaming_scaffold.rs` exercise every branch.
47
48use std::sync::atomic::{AtomicU64, Ordering};
49use std::sync::Arc;
50use tokio::sync::mpsc;
51use wasmtime::{Caller, Linker};
52
53use crate::abi::AbiError;
54
55/// Maximum size, in bytes, of a single `emit-chunk` call. 64 KiB matches
56/// typical HTTP chunk-encoder buffer sizes; guests producing larger
57/// payloads should call `emit-chunk` repeatedly. Enforced by the v0.4
58/// host-fn wrapper; exported here so wrapper and tests share a constant.
59pub const MAX_CHUNK_BYTES: usize = 64 * 1024;
60
61/// Maximum total bytes a single invocation may emit across all
62/// `emit-chunk` calls before the host starts returning the documented
63/// `-2 = cap exceeded` code. 64 MiB is generous for LLM-style token
64/// streams while keeping the worst-case per-invocation memory footprint
65/// bounded.
66pub const MAX_TOTAL_STREAM_BYTES: u64 = 64 * 1024 * 1024;
67
68/// Per-invocation streaming context. Owns the producer end of the
69/// `mpsc::Sender<Vec<u8>>` channel into which guest-emitted chunks are
70/// pushed; the API gateway holds the matching receiver and forwards the
71/// bytes onto the wire.
72///
73/// Cloning is cheap (a refcount bump on each of the two `Arc` fields).
74/// The atomic byte-counter is shared across clones so the
75/// per-invocation cap remains a single source of truth even if the
76/// context is duplicated across host-fn closures.
77#[derive(Debug, Clone)]
78pub struct StreamingContext {
79 /// `None` means streaming is disabled for this invocation — every
80 /// `emit_chunk` returns the `-1` error code immediately. `Some(_)`
81 /// means a downstream receiver is attached.
82 sender: Option<Arc<mpsc::Sender<Vec<u8>>>>,
83 /// Running total of bytes successfully accepted (i.e. forwarded
84 /// onto the channel). Compared against `max_total` on every emit.
85 bytes_emitted: Arc<AtomicU64>,
86 /// Cap on `bytes_emitted`. Zero when streaming is disabled — emit
87 /// returns `-1` before this is consulted, but a zero cap is the
88 /// safe scaffold default for the disabled path.
89 max_total: u64,
90}
91
92impl StreamingContext {
93 /// Construct a context with streaming disabled. Every `emit_chunk`
94 /// returns `-1`; useful as the default for invocations the gateway
95 /// did not opt into the streaming response path.
96 pub fn disabled() -> Self {
97 Self {
98 sender: None,
99 bytes_emitted: Arc::new(AtomicU64::new(0)),
100 max_total: 0,
101 }
102 }
103
104 /// Construct a context wired to `sender`. Bytes pushed via
105 /// `emit_chunk` are forwarded onto the channel; the gateway-side
106 /// receiver drives the SSE / chunked response body.
107 ///
108 /// The total-bytes cap defaults to [`MAX_TOTAL_STREAM_BYTES`] —
109 /// embedders that need a different cap should construct the context
110 /// directly via [`Self::with_channel_and_cap`].
111 pub fn with_channel(sender: mpsc::Sender<Vec<u8>>) -> Self {
112 Self::with_channel_and_cap(sender, MAX_TOTAL_STREAM_BYTES)
113 }
114
115 /// Construct a context wired to `sender` with an explicit total-bytes
116 /// cap. Intended primarily for tests that exercise the cap branch
117 /// without having to emit 64 MiB of data.
118 pub fn with_channel_and_cap(sender: mpsc::Sender<Vec<u8>>, max_total: u64) -> Self {
119 Self {
120 sender: Some(Arc::new(sender)),
121 bytes_emitted: Arc::new(AtomicU64::new(0)),
122 max_total,
123 }
124 }
125
126 /// Forward `bytes` to the downstream receiver if streaming is
127 /// enabled.
128 ///
129 /// Returns:
130 /// * `0` on success (bytes accepted). The chunk's contribution is
131 /// added to [`Self::bytes_emitted`].
132 /// * `-1` if streaming is disabled (no channel attached).
133 /// * `-2` if accepting this chunk would push the total past
134 /// `max_total`. The counter is rolled back afterward.
135 /// * `-3` if the receiver has been dropped — the downstream HTTP
136 /// client disconnected. The counter is rolled back symmetrically.
137 ///
138 /// Concurrency note: the running total is maintained with an
139 /// optimistic `fetch_add` followed by a rollback on the failure
140 /// branches, so it reflects only successfully forwarded bytes *once
141 /// quiescent*. While concurrent emits are in flight, a reader of
142 /// [`Self::bytes_emitted`] may briefly observe a value above the bytes
143 /// actually forwarded — each in-flight emit can overshoot by its own
144 /// chunk size before its rollback lands, so with the per-call cap the
145 /// transient overshoot is bounded by roughly `2 * MAX_CHUNK_BYTES` per
146 /// racing emit. The overshoot is reconciled (rolled back) before each
147 /// failing call returns; it never causes the cap to under-count
148 /// accepted bytes.
149 pub async fn emit_chunk(&self, bytes: Vec<u8>) -> i32 {
150 let Some(s) = &self.sender else {
151 return -1;
152 };
153 let added = bytes.len() as u64;
154 // Optimistic-add + rollback-on-failure keeps the success path a
155 // single atomic op while keeping the running total accurate once
156 // quiescent. Under concurrency the total can transiently overshoot
157 // (each racing emit adds its chunk size before any rollback lands),
158 // bounded by ~2*MAX_CHUNK_BYTES per racing emit — see the
159 // `emit_chunk` doc comment.
160 let new_total = self.bytes_emitted.fetch_add(added, Ordering::SeqCst) + added;
161 if new_total > self.max_total {
162 self.bytes_emitted.fetch_sub(added, Ordering::SeqCst);
163 return -2;
164 }
165 match s.send(bytes).await {
166 Ok(_) => 0,
167 Err(_) => {
168 // Receiver dropped — roll back the accounting so a
169 // subsequent attempt (if the gateway swaps in a new
170 // receiver) sees an accurate total.
171 self.bytes_emitted.fetch_sub(added, Ordering::SeqCst);
172 -3
173 }
174 }
175 }
176
177 /// Flush any buffered chunks. The scaffold uses an unbuffered
178 /// `mpsc::Sender` whose `send` already delivers per-call, so this
179 /// is a no-op. v0.4 may introduce a coalescing buffer in front of
180 /// the channel; the flush hook is reserved for that.
181 ///
182 /// Returns `0` on success, `-1` when streaming is disabled.
183 pub fn flush(&self) -> i32 {
184 if self.sender.is_none() {
185 return -1;
186 }
187 0
188 }
189
190 /// Snapshot of bytes successfully emitted so far. Reads through
191 /// [`Ordering::SeqCst`] for symmetry with [`Self::emit_chunk`].
192 pub fn bytes_emitted(&self) -> u64 {
193 self.bytes_emitted.load(Ordering::SeqCst)
194 }
195
196 /// `true` if a receiver channel is attached (i.e. the invocation
197 /// was dispatched through the streaming response path).
198 pub fn is_enabled(&self) -> bool {
199 self.sender.is_some()
200 }
201}
202
203/// Trait implemented by store data types that can hand out a
204/// [`StreamingContext`]. Parallels [`crate::host::HasWasiCuda`]; the
205/// executor's `InstanceState` will implement this once the v0.4 wiring
206/// lands.
207pub trait HasStreaming {
208 /// Borrow the streaming context.
209 fn streaming(&self) -> &StreamingContext;
210}
211
212/// Module name used to register the wasi-tensor host functions:
213/// `wasi:tensor/host@0.1.0`.
214///
215/// The on-the-wire string carries the `@x.y.z` version segment so it
216/// matches the `package wasi:tensor@0.1.0;` declaration in
217/// `wit/wasi-tensor.wit`. wit-bindgen-generated guests for
218/// `wasi:tensor@0.1.0` emit imports named `wasi:tensor/host@0.1.0`, so the
219/// version segment is required for those bindings to resolve against this
220/// host module. The version is kept in lockstep with the WIT package
221/// declaration — bumping one without the other will cause guests generated
222/// from the WIT to fail to link, mirroring the `wasi:cuda/host@0.2.0`
223/// ([`crate::abi::MODULE`]) and `wasi:scheduler/host@0.1.0`
224/// ([`crate::scheduler::SCHEDULER_MODULE`]) conventions.
225pub const STREAMING_MODULE: &str = "wasi:tensor/host@0.1.0";
226
227/// Host-function name for `emit-chunk`.
228pub const FN_EMIT_CHUNK: &str = "emit-chunk";
229
230/// Host-function name for `flush`.
231pub const FN_FLUSH: &str = "flush";
232
233/// Host-function name for `input-len`.
234pub const FN_INPUT_LEN: &str = "input-len";
235
236/// Host-function name for `read-input`.
237pub const FN_READ_INPUT: &str = "read-input";
238
239/// Register the wasi-tensor host functions on a wasmtime `Linker`.
240///
241/// `T` is the store data type and must implement [`HasStreaming`].
242///
243/// The host-fn wrappers registered here read the `(buf_ptr, buf_len)`
244/// argument pair from the guest, bounds-check the region against the
245/// guest's linear memory exported as `"memory"`, then forward the bytes
246/// into [`StreamingContext::emit_chunk`].
247///
248/// Single-chunk size cap: emits whose `buf_len` exceeds
249/// [`MAX_CHUNK_BYTES`] return `-2` (cap exceeded) without touching the
250/// channel — chunks above the cap would interleave badly with the
251/// downstream SSE / chunked-transfer framing, so we refuse them at the
252/// boundary. The total-bytes cap is still enforced inside
253/// [`StreamingContext::emit_chunk`] on top of this per-call check.
254///
255/// Returns the same `wasmtime::Result` shape as [`crate::host::add_to_linker`]
256/// for parity. Idempotency: registering the same `(module, name)`
257/// twice on the same `Linker` is an error — callers should construct a
258/// fresh `Linker` per build, just as the existing wasi-cuda registration
259/// expects.
260pub fn add_streaming_to_linker<T: HasStreaming + Send + 'static>(
261 linker: &mut Linker<T>,
262) -> wasmtime::Result<()> {
263 // `func_wrap_async` mirrors the wasi-cuda launch path: emit_chunk
264 // performs an `.await` on the `mpsc::Sender`, so it must run on a
265 // wasmtime async fiber. Even a synchronous wrapper would have to
266 // poll the future to completion; using `func_wrap_async` keeps the
267 // backpressure semantics honest.
268 linker.func_wrap_async(
269 STREAMING_MODULE,
270 FN_EMIT_CHUNK,
271 |mut caller: Caller<'_, T>,
272 (buf_ptr, buf_len): (i32, i32)|
273 -> Box<dyn std::future::Future<Output = i32> + Send + '_> {
274 // Synchronously: validate the (ptr, len) pair, copy the
275 // bytes out of guest linear memory, and clone the streaming
276 // context out of store data. The `Caller`'s memory borrow
277 // cannot survive an `.await` (wasmtime's `Memory::data`
278 // borrow may be invalidated by any await on a different
279 // host fn or by guest re-entry), so we materialise the
280 // `Vec<u8>` up front and emit on the cloned context.
281 //
282 // Cloning `StreamingContext` is two refcount bumps — see
283 // the type's `Clone` impl — so the cost is negligible
284 // compared to the byte copy.
285 let prep = prepare_emit_chunk(&mut caller, buf_ptr, buf_len);
286 Box::new(async move {
287 match prep {
288 Ok((bytes, ctx)) => ctx.emit_chunk(bytes).await,
289 Err(code) => code,
290 }
291 })
292 },
293 )?;
294
295 linker.func_wrap(STREAMING_MODULE, FN_FLUSH, |caller: Caller<'_, T>| -> i32 {
296 caller.data().streaming().flush()
297 })?;
298
299 Ok(())
300}
301
302/// Synchronous preamble for [`add_streaming_to_linker`]'s `emit-chunk`
303/// host function. Validates the `(buf_ptr, buf_len)` pair against the
304/// guest's exported `"memory"`, copies the bytes out of linear memory,
305/// and clones the [`StreamingContext`] out of store data.
306///
307/// Returns `Err(-2)` (cap-exceeded / invalid pointer) on any failure
308/// path — both `MAX_CHUNK_BYTES` overflow and out-of-bounds region
309/// surface the same documented `-2` code so the guest cannot
310/// fingerprint the host's memory layout from the return value. The
311/// disabled / receiver-dropped branches (`-1` / `-3`) are reached only
312/// once the bytes have been forwarded onto the channel.
313fn prepare_emit_chunk<T: HasStreaming>(
314 caller: &mut Caller<'_, T>,
315 buf_ptr: i32,
316 buf_len: i32,
317) -> Result<(Vec<u8>, StreamingContext), i32> {
318 if buf_len < 0 || buf_ptr < 0 {
319 return Err(-2);
320 }
321 if (buf_len as usize) > MAX_CHUNK_BYTES {
322 return Err(-2);
323 }
324 let memory = caller
325 .get_export("memory")
326 .and_then(|e| e.into_memory())
327 .ok_or(-2_i32)?;
328 let start = buf_ptr as usize;
329 let end = start.checked_add(buf_len as usize).ok_or(-2_i32)?;
330 // `Memory::data` returns a slice tied to the borrow of the
331 // store, so we scope the borrow tightly and copy out before the
332 // subsequent `caller.data()` call re-borrows for the
333 // streaming-context clone. Mirrors the pattern used by
334 // `wasi-cuda`'s `read_bytes` helper in `src/host.rs`.
335 let bytes: Vec<u8> = {
336 let data = memory.data(&caller);
337 if end > data.len() {
338 return Err(-2);
339 }
340 data[start..end].to_vec()
341 };
342 let ctx = caller.data().streaming().clone();
343 Ok((bytes, ctx))
344}
345
346// ---------------------------------------------------------------------------
347// Guest input channel (pull model): `input-len` / `read-input`
348// ---------------------------------------------------------------------------
349
350/// Per-invocation input context. Owns the bytes the host staged for the
351/// guest to pull via the `wasi:tensor/host` `input-len` / `read-input`
352/// host functions.
353///
354/// This is the input counterpart to [`StreamingContext`] (which is
355/// output-only, guest → host). The gateway stages the request payload —
356/// e.g. the OpenAI completions shim stages the assembled prompt bytes —
357/// before the guest runs, and the guest copies them into its own linear
358/// memory.
359///
360/// Cloning is cheap: the bytes live behind an `Arc<[u8]>`, so a clone is
361/// a single refcount bump. An empty context (`InputContext::empty`) is
362/// the default for invocations the gateway did not stage input for —
363/// `len()` is `0` and `read-input` copies nothing.
364#[derive(Debug, Clone, Default)]
365pub struct InputContext {
366 /// Bytes staged for the guest. Empty by default.
367 bytes: Arc<[u8]>,
368}
369
370impl InputContext {
371 /// Construct an empty context. `len()` returns `0`; `read-input`
372 /// copies nothing. This is the default for invocations the gateway
373 /// did not stage any input for.
374 pub fn empty() -> Self {
375 Self::default()
376 }
377
378 /// Construct a context wrapping the given staged input bytes.
379 pub fn new(bytes: impl Into<Arc<[u8]>>) -> Self {
380 Self {
381 bytes: bytes.into(),
382 }
383 }
384
385 /// Borrow the staged input bytes.
386 pub fn bytes(&self) -> &[u8] {
387 &self.bytes
388 }
389
390 /// Number of staged input bytes. Clamped to `u32::MAX` so the value
391 /// fits the WIT `input-len() -> u32` return type — a staged buffer
392 /// above 4 GiB is implausible (the API body-size cap is far below
393 /// that) but we saturate rather than wrap on the cast.
394 pub fn len_u32(&self) -> u32 {
395 u32::try_from(self.bytes.len()).unwrap_or(u32::MAX)
396 }
397
398 /// `true` when no input bytes are staged.
399 pub fn is_empty(&self) -> bool {
400 self.bytes.is_empty()
401 }
402}
403
404/// Trait implemented by store data types that can hand out an
405/// [`InputContext`]. Parallels [`HasStreaming`]; the executor's
406/// `InstanceState` implements it so the `input-len` / `read-input` host
407/// functions reach the per-instance staged buffer.
408pub trait HasInput {
409 /// Borrow the input context.
410 fn input(&self) -> &InputContext;
411}
412
413/// Register the `wasi:tensor/host` guest-input host functions
414/// (`input-len`, `read-input`) on a wasmtime `Linker`.
415///
416/// `T` is the store data type and must implement [`HasInput`]. Mirrors
417/// [`add_streaming_to_linker`]: both register host functions on the same
418/// [`STREAMING_MODULE`] (`wasi:tensor/host@0.1.0`) so a guest importing
419/// either surface links against one host module. They are split into two
420/// registration functions because the input surface needs only
421/// [`HasInput`] while the streaming surface needs [`HasStreaming`]; the
422/// executor's `InstanceState` implements both and calls both.
423///
424/// `read-input` bounds-checks the `(ptr, len)` destination region against
425/// the guest's exported `"memory"` with the same `checked_add` pattern as
426/// `prepare_emit_chunk` and the wasi-cuda `read_bytes` path, returning
427/// [`AbiError::InvalidPointer`] (`-2`) on any out-of-bounds / overflow /
428/// missing-memory failure. A successful copy returns the number of bytes
429/// written (`min(len, input-len())`); `0` means "nothing staged" or
430/// `len == 0`, never an error.
431pub fn add_input_to_linker<T: HasInput + Send + 'static>(
432 linker: &mut Linker<T>,
433) -> wasmtime::Result<()> {
434 linker.func_wrap(
435 STREAMING_MODULE,
436 FN_INPUT_LEN,
437 |caller: Caller<'_, T>| -> u32 { caller.data().input().len_u32() },
438 )?;
439
440 linker.func_wrap(
441 STREAMING_MODULE,
442 FN_READ_INPUT,
443 |mut caller: Caller<'_, T>, ptr: u32, len: u32| -> i32 {
444 // Clone the staged bytes out of store data before borrowing
445 // memory mutably for the write. The `Arc<[u8]>` clone is a
446 // refcount bump.
447 let input = caller.data().input().clone();
448 let staged = input.bytes();
449 if staged.is_empty() || len == 0 {
450 return 0;
451 }
452 let to_copy = std::cmp::min(staged.len(), len as usize);
453 let memory = match caller.get_export("memory").and_then(|e| e.into_memory()) {
454 Some(m) => m,
455 None => return AbiError::InvalidPointer.code(),
456 };
457 // Bounds-check the destination region against the current
458 // linear-memory size with the same `checked_add` guard the
459 // emit-chunk / read_bytes paths use: without it a guest could
460 // pass `(ptr = u32::MAX - 1, len = 4)` and wrap to a small
461 // `end` that looks in-bounds.
462 let start = ptr as usize;
463 let end = match start.checked_add(to_copy) {
464 Some(e) => e,
465 None => return AbiError::InvalidPointer.code(),
466 };
467 let mem_len = memory.data(&caller).len();
468 if end > mem_len {
469 return AbiError::InvalidPointer.code();
470 }
471 // Copy out of the Arc-backed buffer into an owned Vec so the
472 // subsequent mutable `memory.write` borrow does not alias the
473 // `input` borrow (mirrors `last_error_copy` in `host.rs`).
474 let buf = staged[..to_copy].to_vec();
475 if memory.write(&mut caller, start, &buf).is_err() {
476 return AbiError::InvalidPointer.code();
477 }
478 // `to_copy <= len <= u32::MAX` and `to_copy <= staged.len()`,
479 // and the WIT contract caps a single read at `input-len()`
480 // which is itself clamped to `u32::MAX`, so the `as i32` cast
481 // is non-negative for any realistic buffer. Guard the
482 // (implausible) >2 GiB case so we never hand back a negative
483 // count that a guest would read as an error.
484 i32::try_from(to_copy).unwrap_or(i32::MAX)
485 },
486 )?;
487
488 Ok(())
489}
490
491#[cfg(test)]
492mod tests {
493 use super::*;
494
495 #[tokio::test]
496 async fn disabled_context_returns_minus_one() {
497 let ctx = StreamingContext::disabled();
498 assert!(!ctx.is_enabled());
499 assert_eq!(ctx.emit_chunk(vec![1, 2, 3]).await, -1);
500 assert_eq!(ctx.flush(), -1);
501 assert_eq!(ctx.bytes_emitted(), 0);
502 }
503
504 #[tokio::test]
505 async fn enabled_context_forwards_bytes() {
506 let (tx, mut rx) = mpsc::channel::<Vec<u8>>(4);
507 let ctx = StreamingContext::with_channel(tx);
508 assert!(ctx.is_enabled());
509 assert_eq!(ctx.emit_chunk(vec![0xAA, 0xBB, 0xCC]).await, 0);
510 assert_eq!(ctx.bytes_emitted(), 3);
511 let chunk = rx.recv().await.expect("chunk delivered");
512 assert_eq!(chunk, vec![0xAA, 0xBB, 0xCC]);
513 assert_eq!(ctx.flush(), 0);
514 }
515
516 #[test]
517 fn input_context_empty_by_default() {
518 let ctx = InputContext::empty();
519 assert!(ctx.is_empty());
520 assert_eq!(ctx.len_u32(), 0);
521 assert_eq!(ctx.bytes(), b"");
522 // Default impl agrees with `empty()`.
523 assert!(InputContext::default().is_empty());
524 }
525
526 #[test]
527 fn input_context_carries_staged_bytes() {
528 let ctx = InputContext::new(b"hello prompt".to_vec());
529 assert!(!ctx.is_empty());
530 assert_eq!(ctx.len_u32(), 12);
531 assert_eq!(ctx.bytes(), b"hello prompt");
532 // Cheap clone shares the same backing buffer.
533 let cloned = ctx.clone();
534 assert_eq!(cloned.bytes(), ctx.bytes());
535 }
536
537 #[test]
538 fn streaming_module_is_versioned() {
539 assert!(STREAMING_MODULE.contains("wasi:tensor"));
540 assert!(STREAMING_MODULE.ends_with("@0.1.0"));
541 }
542
543 /// Pin the host [`STREAMING_MODULE`] string against drift from
544 /// `wit/wasi-tensor.wit`.
545 ///
546 /// Mirrors `abi.rs::module_version_matches_wit_package_decl`. The WIT
547 /// file is the authoritative spec; the host's import-module name has to
548 /// carry the same `@x.y.z` segment so wit-bindgen-generated guests for
549 /// `wasi:tensor@x.y.z` (which emit imports named
550 /// `wasi:tensor/host@x.y.z`) can link against this host. Parse the
551 /// version out of the WIT `package wasi:tensor@x.y.z;` line and compare
552 /// it to the suffix of [`STREAMING_MODULE`]. If somebody bumps one
553 /// without the other, this trips before the linker error reaches
554 /// downstream users.
555 #[test]
556 fn streaming_module_version_matches_wit_package_decl() {
557 // Path is relative to this source file
558 // (`crates/tensor-wasm-wasi-gpu/src/streaming.rs`):
559 // ../ -> crates/tensor-wasm-wasi-gpu/, where the in-crate
560 // `wit/wasi-tensor.wit` copy lives. Kept in-crate so the
561 // published tarball is self-contained (`cargo publish`
562 // rejects paths that escape the crate root).
563 const WIT: &str = include_str!("../wit/wasi-tensor.wit");
564 let pkg_line = WIT
565 .lines()
566 .find(|l| l.trim_start().starts_with("package wasi:tensor@"))
567 .expect("wit/wasi-tensor.wit must declare `package wasi:tensor@x.y.z;`");
568 let version = pkg_line
569 .trim()
570 .trim_start_matches("package wasi:tensor@")
571 .trim_end_matches(';')
572 .trim();
573 assert!(
574 !version.is_empty(),
575 "could not parse a version out of: {pkg_line:?}"
576 );
577 let expected_suffix = format!("@{version}");
578 assert!(
579 STREAMING_MODULE.ends_with(&expected_suffix),
580 "STREAMING_MODULE ({STREAMING_MODULE:?}) drifted from \
581 wit/wasi-tensor.wit package version ({version:?}); update \
582 src/streaming.rs::STREAMING_MODULE or the WIT file so they agree."
583 );
584 }
585}