Skip to main content

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}