Skip to main content

submilli_engine/stdlib/embedding/
mod.rs

1//! `submilli:embedding` — remote text embeddings from inside a Submilli program.
2//!
3//! Rust host functions registered directly under the package name. Each op runs
4//! the security gate before anything leaves the process; the one gated
5//! capability is `embedding.embed`, cataloged in [`crate::stdlib::capabilities`]
6//! — keep it in sync when adding or removing a gate (see AGENTS.md).
7//!
8//! Dispatch is the embedder's: [`StoreData::embedding_provider`] holds the
9//! provider, and a runtime with none refuses every op rather than inventing a
10//! vector — the rule `submilli:llm` follows.
11//!
12//! **One capability, per-alias filtering.** `embed` and `models()` share the
13//! grant, with `input_count` in the filter context. `models()` filters each
14//! candidate through the same `model` filter that gates `embed`, so a listing
15//! never offers an alias the caller would be denied at call time.
16//!
17//! **No input text or vector crosses this boundary in metadata.** Not in the
18//! filter context, not in an error, not in a log. The context carries `model`
19//! and `input_count`.
20//!
21//! **Vectors stay compact, in the GC heap.** A result is a backing struct whose
22//! hidden field holds the vectors as a packed `i8` array: little-endian `f32`
23//! bytes, row-major, `count × dimensions × 4` bytes — the storage a `Uint8Array`
24//! uses. The array is allocated through the GC limiter, which collects inside
25//! host calls, so a discarded result is reclaimed before the heap grows: the
26//! vectors live in the GC heap, never in host memory the collector cannot see.
27//! The field has no getter and its type is unreachable from the guest, so only
28//! `vector(i)` and `bytes(i)` read it, copying one row out on demand; a
29//! plain `number[][]` would cost about 60 bytes per number.
30
31pub mod declaration;
32
33use crate::runtime::effects::{self, NativeError};
34use std::sync::Arc;
35
36use wasmtime::{FuncType, HeapType, Linker, RefType, StructType, Val, ValType};
37
38use crate::runtime::ModelCallUsage;
39use crate::runtime::StoreData;
40use crate::runtime::call_log::{ModelUsage, Payload, Side, record_payload, record_usage};
41use crate::runtime::decision::CallTicket;
42use crate::runtime::embedding::{
43    EmbeddingBatch, EmbeddingBoundKind, EmbeddingError, EmbeddingLimits, EmbeddingMalformedReason,
44    EmbeddingModel, EmbeddingProvider, EmbeddingTokenBudget, Purpose, SubBatchSettlement,
45};
46use crate::runtime::fuel;
47use crate::runtime::host::{
48    abi_arg, abi_result, fatal_host_error, quota_exceeded_error, range_error, read_string_arg,
49    register_host_fn, register_host_value_async, type_error, write_boxed_number_struct,
50    write_submilli_array_struct_precharged, write_submilli_string_struct,
51    write_submilli_uint8array_struct, write_uint8_array_precharged,
52};
53use crate::runtime::intrinsic_types::build_intrinsic_types;
54use crate::stdlib::abi::{
55    self, HostView, backing_struct, f64_field, install_field_getters, nullable_object_field,
56    raw_bytes_field, string_field,
57};
58use crate::stdlib::shared::{
59    audit_quota_denial, check_and_approve_call, check_candidate, filters_candidate, mark_filtered,
60    optional_number, preflight_models, sanitize_description,
61};
62
63pub use crate::runtime::EMBEDDING_MODULE_NAME as MODULE_NAME;
64pub use declaration::package_declaration;
65
66/// The single capability gating every op in this package.
67pub const CAPABILITY: &str = "embedding.embed";
68
69// `$EmbeddingsBacking` field indices (0 is the vtable).
70const E_VECTORS: usize = 1;
71const E_COUNT: usize = 2;
72const E_DIMENSIONS: usize = 3;
73const E_IDENTITY: usize = 4;
74const E_MODEL: usize = 5;
75const E_INPUT_TOKENS: usize = 6;
76
77// `$EmbeddingModelBacking` field indices (0 is the vtable).
78const M_NAME: usize = 1;
79const M_DESCRIPTION: usize = 2;
80const M_DIMENSIONS: usize = 3;
81const M_MAX_INPUT_TOKENS: usize = 4;
82const M_MAX_INPUT_BYTES: usize = 5;
83const M_IDENTITY: usize = 6;
84
85#[derive(serde::Serialize)]
86struct EmbedRequest {
87    model: String,
88    texts: Vec<String>,
89    purpose: Purpose,
90}
91
92pub fn install(linker: &mut Linker<StoreData>) -> wasmtime::Result<()> {
93    let engine = linker.engine().clone();
94    let intr = build_intrinsic_types(&engine)?;
95    let string = ValType::Ref(RefType::new(
96        false,
97        HeapType::ConcreteStruct(intr.string.clone()),
98    ));
99    let array = ValType::Ref(RefType::new(
100        false,
101        HeapType::ConcreteStruct(intr.array.clone()),
102    ));
103    let uint8 = ValType::Ref(RefType::new(
104        false,
105        HeapType::ConcreteStruct(intr.uint8_array.clone()),
106    ));
107    // `Embeddings` crosses the boundary as the universal `(ref null $Object)`
108    // lowering; a `Direct` receiver is the non-null `(ref $Object)`.
109    let nullable_object = ValType::Ref(RefType::new(
110        true,
111        HeapType::ConcreteStruct(intr.object.clone()),
112    ));
113    let receiver = ValType::Ref(RefType::new(
114        false,
115        HeapType::ConcreteStruct(intr.object.clone()),
116    ));
117
118    register_host_value_async(
119        linker,
120        MODULE_NAME,
121        crate::mangle::package_symbol(MODULE_NAME, "embed"),
122        FuncType::new(
123            &engine,
124            [string.clone(), array.clone(), string.clone()],
125            [nullable_object.clone()],
126        ),
127        |caller, params| {
128            Box::pin(async move {
129                let model =
130                    read_string_arg(&mut *caller, abi_arg(params, 0)?, "embedding.embed (model)")?;
131                let purpose = read_string_arg(
132                    &mut *caller,
133                    abi_arg(params, 2)?,
134                    "embedding.embed (purpose)",
135                )?;
136                embed(caller, &model, abi_arg(params, 1)?, &purpose).await
137            })
138        },
139    )?;
140
141    register_host_value_async(
142        linker,
143        MODULE_NAME,
144        crate::mangle::package_symbol(MODULE_NAME, "models"),
145        FuncType::new(&engine, [], [array.clone()]),
146        |caller, _params| Box::pin(async move { models(caller).await }),
147    )?;
148
149    for view in HOST_VIEWS {
150        install_field_getters(linker, &engine, &receiver, view)?;
151    }
152    install_methods(linker, &engine, receiver, array, uint8)
153}
154
155const EMBEDDINGS_VIEW: HostView = HostView {
156    module: MODULE_NAME,
157    interface: "Embeddings",
158    backing: embeddings_backing_struct,
159    properties: &[
160        ("count", E_COUNT),
161        ("dimensions", E_DIMENSIONS),
162        ("identity", E_IDENTITY),
163        ("model", E_MODEL),
164        ("inputTokens", E_INPUT_TOKENS),
165    ],
166};
167
168const EMBEDDING_MODEL_VIEW: HostView = HostView {
169    module: MODULE_NAME,
170    interface: "EmbeddingModel",
171    backing: model_backing_struct,
172    properties: &[
173        ("name", M_NAME),
174        ("description", M_DESCRIPTION),
175        ("dimensions", M_DIMENSIONS),
176        ("maxInputTokens", M_MAX_INPUT_TOKENS),
177        ("maxInputBytes", M_MAX_INPUT_BYTES),
178        ("identity", M_IDENTITY),
179    ],
180};
181
182/// The host-backed interfaces this package's values are.
183pub(crate) const HOST_VIEWS: &[&HostView] = &[&EMBEDDINGS_VIEW, &EMBEDDING_MODEL_VIEW];
184
185fn install_methods(
186    linker: &mut Linker<StoreData>,
187    engine: &wasmtime::Engine,
188    receiver: ValType,
189    array: ValType,
190    uint8: ValType,
191) -> wasmtime::Result<()> {
192    let embeddings_key = crate::mangle::package_symbol(MODULE_NAME, "Embeddings");
193
194    register_host_fn(
195        linker,
196        MODULE_NAME,
197        crate::mangle::extend(&embeddings_key, "vector"),
198        FuncType::new(engine, [receiver.clone(), ValType::F64], [array]),
199        |caller, params, results| {
200            let row = read_row(
201                caller,
202                abi_arg(params, 0)?,
203                abi_arg(params, 1)?,
204                "embedding.vector",
205            )?;
206            *abi_result(results, 0)? = build_number_array(caller, &row)?;
207            Ok(())
208        },
209    )?;
210
211    register_host_fn(
212        linker,
213        MODULE_NAME,
214        crate::mangle::extend(&embeddings_key, "bytes"),
215        FuncType::new(engine, [receiver, ValType::F64], [uint8]),
216        |caller, params, results| {
217            let row = read_row(
218                caller,
219                abi_arg(params, 0)?,
220                abi_arg(params, 1)?,
221                "embedding.bytes",
222            )?;
223            *abi_result(results, 0)? = Val::AnyRef(Some(
224                write_submilli_uint8array_struct(caller, &row)?.to_anyref(),
225            ));
226            Ok(())
227        },
228    )
229}
230
231/// One `embed`: gate, bound, resolve, reserve, dispatch, settle.
232///
233/// The ordering is the contract and is load-bearing at every step.
234///
235/// 1. **Gate before anything leaves.** A denial must cost nothing and reveal
236///    nothing, so the capability check runs ahead of the provider — which is
237///    what knows whether an alias exists.
238/// 2. **Per-call caps, then alias, then per-input limit.** Each refusal happens
239///    before the budget is touched, so a refused call charges nothing and the
240///    provider is never called.
241/// 3. **Reserve before dispatching.** The whole estimate is reserved up front.
242/// 4. **Settle exactly once** on success and on failure, from the settlements the
243///    provider recorded; a call that never sent releases its whole reservation.
244async fn embed(
245    caller: &mut wasmtime::Caller<'_, StoreData>,
246    model: &str,
247    texts: &Val,
248    purpose: &str,
249) -> wasmtime::Result<Val> {
250    let elements = crate::runtime::prelude::collection::read_array_vals(caller, texts)?;
251    let ticket = gate(caller, model, elements.len()).await?;
252
253    let budget = execution_budget(caller);
254    let limits = budget.limits();
255    let texts = read_texts(caller, model, &elements, &limits)?;
256    // The texts are first in hand once the per-call bounds pass, so the request is
257    // recorded here: a call refused for size has no copy to keep.
258    record_payload(&*caller, ticket, Side::Request, || {
259        let body = serde_json::to_vec(&texts).unwrap_or_default();
260        let bytes: usize = texts.iter().map(String::len).sum();
261        Payload::meta(serde_json::json!({
262            "op": "embed",
263            "model": model,
264            "purpose": purpose,
265            "count": texts.len(),
266        }))
267        .with_owned_body(body)
268        .with_size(bytes as u64)
269    });
270    let purpose = parse_purpose(model, purpose)?;
271
272    let request = EmbedRequest {
273        model: model.to_owned(),
274        texts,
275        purpose,
276    };
277    let mut response = effects::execute(
278        caller,
279        "embedding.embed",
280        request,
281        move |caller, _, request| Box::pin(call_provider(caller, ticket, request)),
282    )
283    .await?;
284    // Charge even a partially failed batch before exposing its result or error.
285    if let Some(settlement) = &mut response.settlement {
286        settlement.charge();
287    } else {
288        budget.restore_spend(response.used, response.held, response.requests);
289    }
290    let batch = response.result.map_err(NativeError::into_error)?;
291    build_embeddings(caller, batch)
292}
293
294#[derive(serde::Serialize, serde::Deserialize)]
295struct EmbeddingResponse {
296    result: Result<EmbeddingBatch, NativeError>,
297    used: u64,
298    held: u64,
299    requests: u64,
300    #[serde(skip)]
301    settlement: Option<EmbeddingSettlement>,
302}
303
304/// Retain live reservation cleanup if storing the outcome fails. Replay has no
305/// live reservation and restores only the recorded per-run counters.
306struct EmbeddingSettlement {
307    budget: Option<Arc<EmbeddingTokenBudget>>,
308    reserved: u64,
309    settlements: Vec<SubBatchSettlement>,
310}
311
312impl EmbeddingSettlement {
313    fn charge(&mut self) {
314        if let Some(budget) = self.budget.take() {
315            budget.settle(self.reserved, &self.settlements);
316        }
317    }
318}
319
320impl Drop for EmbeddingSettlement {
321    fn drop(&mut self) {
322        self.charge();
323    }
324}
325
326async fn call_provider(
327    caller: &mut wasmtime::Caller<'_, StoreData>,
328    ticket: Option<CallTicket>,
329    request: EmbedRequest,
330) -> wasmtime::Result<EmbeddingResponse> {
331    let model = request.model.as_str();
332    let texts = request.texts;
333    let purpose = request.purpose;
334    let budget = execution_budget(caller);
335    let provider = provider(caller, model)?;
336    check_inputs(caller, &*provider, model, &texts).await?;
337
338    let sent: usize = texts.iter().map(String::len).sum();
339    fuel::charge(&mut *caller, fuel::IO, sent as u64)?;
340
341    let estimate = provider.estimate_tokens(model, &texts);
342    budget
343        .reserve(model, estimate)
344        .map_err(|error| quota_throw(caller, ticket, model, error))?;
345    let requests_before = budget.requests();
346    let dispatched = provider.embed(model, &texts, purpose, &budget).await;
347    let requests = budget.requests().saturating_sub(requests_before);
348    let settlements = match &dispatched {
349        Ok(batch) => batch.settlements(),
350        Err(error) => error.settlements(),
351    };
352    let (used, held) = settlements.iter().fold((0u64, 0u64), |(used, held), part| {
353        (
354            used.saturating_add(part.reported)
355                .saturating_add(part.indeterminate),
356            held.saturating_add(part.indeterminate),
357        )
358    });
359    let settlement = EmbeddingSettlement {
360        budget: Some(budget),
361        reserved: estimate,
362        settlements: settlements.to_vec(),
363    };
364    let result = match dispatched {
365        // The provider returned a batch, so the dispatch ran. A short or long
366        // row count is a guest error, not an absent call: the budget still
367        // settles the batch's tokens.
368        Ok(batch) if batch.count() != texts.len() => {
369            note_embedding_usage(caller, ticket, model, batch.input_tokens());
370            Err(throw(EmbeddingError::Malformed {
371                alias: model.to_string(),
372                reason: EmbeddingMalformedReason::CountMismatch,
373                settlements: Vec::new(),
374            }))
375        }
376        Ok(batch) => {
377            let received = (batch.values().len() as u64).saturating_mul(4);
378            record_payload(&*caller, ticket, Side::Response, || {
379                Payload::meta(serde_json::json!({
380                    "count": batch.count(),
381                    "dimensions": batch.dimensions(),
382                    "identity": batch.identity(),
383                    "model": model,
384                    "inputTokens": batch.input_tokens(),
385                }))
386                .with_size(received)
387            });
388            note_embedding_usage(caller, ticket, model, batch.input_tokens());
389            fuel::settle(&mut *caller, fuel::IO, received)?;
390            Ok(batch)
391        }
392        Err(error) => {
393            // Some sub-batches may have completed before the provider failed.
394            let reported = error
395                .settlements()
396                .iter()
397                .fold(0u64, |sum, part| sum.saturating_add(part.reported));
398            if reported > 0 {
399                note_embedding_usage(caller, ticket, model, Some(reported));
400            }
401            Err(quota_throw(caller, ticket, model, error))
402        }
403    };
404    let result = match result {
405        Err(error) if crate::runtime::host::ends_the_run(&error) => return Err(error),
406        result => result.map_err(|error| NativeError::capture(&error)),
407    };
408    Ok(EmbeddingResponse {
409        result,
410        used,
411        held,
412        requests,
413        settlement: Some(settlement),
414    })
415}
416
417/// Record the usage the call log keeps, and the same counts for the embedder.
418///
419/// Only `call_provider` may call this. Replay restores the recorded outcome
420/// without entering that action.
421fn note_embedding_usage(
422    caller: &mut wasmtime::Caller<'_, StoreData>,
423    ticket: Option<CallTicket>,
424    model: &str,
425    input_tokens: Option<u64>,
426) {
427    record_usage(
428        &*caller,
429        ticket,
430        ModelUsage {
431            input_tokens,
432            output_tokens: None,
433        },
434    );
435    caller
436        .data_mut()
437        .note_model_call(ModelCallUsage::embedding(model, input_tokens));
438}
439
440/// The capability check, run before any bytes leave the process. The context is
441/// exactly `model` and `input_count`; input text is never in it.
442async fn gate(
443    caller: &mut wasmtime::Caller<'_, StoreData>,
444    model: &str,
445    input_count: usize,
446) -> wasmtime::Result<Option<CallTicket>> {
447    check_and_approve_call(caller, CAPABILITY, gate_context(model, input_count)).await
448}
449
450fn gate_context(model: &str, input_count: usize) -> serde_json::Value {
451    serde_json::json!({ "model": model, "input_count": input_count })
452}
453
454/// Bound the call by count and total bytes, then read the texts. The count is
455/// checked on the element list, before any string is copied out of the guest.
456fn read_texts(
457    caller: &mut wasmtime::Caller<'_, StoreData>,
458    model: &str,
459    elements: &[Val],
460    limits: &EmbeddingLimits,
461) -> wasmtime::Result<Vec<String>> {
462    let count = elements.len() as u64;
463    if count > limits.max_texts_per_call {
464        return Err(throw(EmbeddingError::BoundsExceeded {
465            alias: model.to_string(),
466            kind: EmbeddingBoundKind::TextCount,
467            actual: count,
468            limit: limits.max_texts_per_call,
469        }));
470    }
471    if elements.is_empty() {
472        return Err(range_error(format!(
473            "embedding.embed(\"{model}\"): at least one text is required — pass a non-empty array"
474        )));
475    }
476    let mut texts = Vec::with_capacity(elements.len());
477    let mut total = 0u64;
478    for element in elements {
479        let text = read_string_arg(caller, element, "embedding.embed (texts)")?;
480        total = total.saturating_add(text.len() as u64);
481        if total > limits.max_bytes_per_call {
482            return Err(throw(EmbeddingError::BoundsExceeded {
483                alias: model.to_string(),
484                kind: EmbeddingBoundKind::TotalBytes,
485                actual: total,
486                limit: limits.max_bytes_per_call,
487            }));
488        }
489        texts.push(text);
490    }
491    Ok(texts)
492}
493
494fn parse_purpose(model: &str, purpose: &str) -> wasmtime::Result<Purpose> {
495    match purpose {
496        "query" => Ok(Purpose::Query),
497        "document" => Ok(Purpose::Document),
498        _ => Err(range_error(format!(
499            "embedding.embed(\"{model}\"): purpose must be \"query\" or \"document\""
500        ))),
501    }
502}
503
504/// Resolve the alias and refuse any input over its byte limit, before the budget
505/// is touched. Numbering is from 0, like the vectors.
506async fn check_inputs(
507    caller: &mut wasmtime::Caller<'_, StoreData>,
508    provider: &dyn EmbeddingProvider,
509    model: &str,
510    texts: &[String],
511) -> wasmtime::Result<()> {
512    let Some(limit) = provider.max_input_bytes(model) else {
513        let available = available_aliases(caller, provider).await?;
514        return Err(throw(EmbeddingError::UnknownModel {
515            alias: model.to_string(),
516            available,
517        }));
518    };
519    // The provider repeats this check. Doing it here too keeps the refusal ahead
520    // of the budget reservation (see the ordering in [`embed`]) and does not
521    // depend on every provider implementing it.
522    let too_long = texts.iter().position(|text| text.len() as u64 > limit);
523    if let Some(index) = too_long {
524        return Err(throw(EmbeddingError::InputTooLong {
525            alias: model.to_string(),
526            index: Some(index),
527            limit: Some(limit),
528            settlements: Vec::new(),
529        }));
530    }
531    Ok(())
532}
533
534/// The aliases this caller may use, for an unknown-alias error. Filtered through
535/// the capability like `models()`, so the error cannot enumerate aliases the
536/// policy hides.
537async fn available_aliases(
538    caller: &mut wasmtime::Caller<'_, StoreData>,
539    provider: &dyn EmbeddingProvider,
540) -> wasmtime::Result<Vec<String>> {
541    let candidates = provider.models().await.map_err(throw)?;
542    let mut names = Vec::new();
543    for candidate in candidates {
544        if may_embed(caller, &candidate.name)? {
545            names.push(candidate.name);
546        }
547    }
548    Ok(names)
549}
550
551/// `models()`: check runtime invariants, then gate each candidate with the same
552/// `model` filter that gates `embed`.
553///
554/// Filtering acts **only** on a policy denial. An invariant denial means the
555/// check itself could not be made, and swallowing it would turn a runtime
556/// refusal into a silently short listing.
557async fn models(caller: &mut wasmtime::Caller<'_, StoreData>) -> wasmtime::Result<Val> {
558    preflight_models(caller, CAPABILITY, "input_count")?;
559    let candidates = effects::execute(caller, "embedding.models", (), |caller, _, ()| {
560        Box::pin(async move { provider(caller, "")?.models().await.map_err(throw) })
561    })
562    .await?;
563    let mut built = Vec::with_capacity(candidates.len());
564    for candidate in candidates {
565        if may_embed(caller, &candidate.name)? {
566            built.push(build_model(caller, candidate)?);
567        }
568    }
569    abi::new_array(caller, &built)
570}
571
572/// The per-candidate gate. A denial omits the alias rather than failing the
573/// call: a listing that threw on the first forbidden alias would itself disclose
574/// that the operator configured it.
575fn may_embed(caller: &mut wasmtime::Caller<'_, StoreData>, model: &str) -> wasmtime::Result<bool> {
576    let keeps = filters_candidate(check_candidate(
577        &mut *caller,
578        CAPABILITY,
579        gate_context(model, 0),
580    ))?;
581    if !keeps {
582        mark_filtered(&*caller);
583    }
584    Ok(keeps)
585}
586
587/// Clone the provider out of the store before any `await`. Runs *after* the
588/// capability check, deliberately: see [`embed`].
589fn provider(
590    caller: &wasmtime::Caller<'_, StoreData>,
591    model: &str,
592) -> wasmtime::Result<Arc<dyn EmbeddingProvider>> {
593    caller.data().embedding_provider.clone().ok_or_else(|| {
594        throw(EmbeddingError::NotConfigured {
595            alias: model.to_string(),
596        })
597    })
598}
599
600/// The execution's budget, or an unmetered one in the pure-interpreter path. The
601/// per-call caps in its limits still apply there.
602fn execution_budget(caller: &mut wasmtime::Caller<'_, StoreData>) -> Arc<EmbeddingTokenBudget> {
603    caller
604        .data_mut()
605        .embedding_budget
606        .get_or_insert_with(|| Arc::new(EmbeddingTokenBudget::unmetered()))
607        .clone()
608}
609
610/// [`throw`] plus an audit record when the failure is a budget refusal.
611fn quota_throw(
612    caller: &wasmtime::Caller<'_, StoreData>,
613    ticket: Option<CallTicket>,
614    model: &str,
615    error: EmbeddingError,
616) -> wasmtime::Error {
617    if !error.is_budget_exceeded() {
618        return throw(error);
619    }
620    if let Err(denial) = audit_quota_denial(
621        caller,
622        ticket,
623        CAPABILITY,
624        model,
625        "embedding-token budget exceeded",
626    ) {
627        return denial;
628    }
629    throw(error)
630}
631
632/// Every failure but an internal one reaches the guest as a catchable error.
633/// The `Display` impls exclude input text and vectors, so the message passes
634/// through whole. An [`EmbeddingError::Internal`] is a host invariant failure
635/// and ends the run, like every other internal host failure.
636///
637/// Budget refusals are quota errors, count/byte/length bounds are argument range
638/// errors, a malformed provider response is a type error, and the rest keep the
639/// base error type.
640fn throw(error: EmbeddingError) -> wasmtime::Error {
641    let message = error.to_string();
642    match error {
643        EmbeddingError::BudgetExceeded { .. } => quota_exceeded_error(message),
644        EmbeddingError::BoundsExceeded { .. } | EmbeddingError::InputTooLong { .. } => {
645            range_error(message)
646        }
647        EmbeddingError::Malformed { .. } => type_error(message),
648        EmbeddingError::Internal { .. } => fatal_host_error(message),
649        EmbeddingError::NotConfigured { .. }
650        | EmbeddingError::UnknownModel { .. }
651        | EmbeddingError::Unauthorized { .. }
652        | EmbeddingError::Provider { .. } => wasmtime::Error::msg(message),
653    }
654}
655
656fn embeddings_backing_struct(engine: &wasmtime::Engine) -> wasmtime::Result<StructType> {
657    let intr = build_intrinsic_types(engine)?;
658    backing_struct(
659        engine,
660        &intr,
661        vec![
662            raw_bytes_field(&intr),       // vectors (hidden: no getter)
663            f64_field(),                  // count
664            f64_field(),                  // dimensions
665            string_field(&intr),          // identity
666            string_field(&intr),          // model
667            nullable_object_field(&intr), // inputTokens
668        ],
669    )
670}
671
672fn model_backing_struct(engine: &wasmtime::Engine) -> wasmtime::Result<StructType> {
673    let intr = build_intrinsic_types(engine)?;
674    backing_struct(
675        engine,
676        &intr,
677        vec![
678            string_field(&intr),          // name
679            nullable_object_field(&intr), // description
680            f64_field(),                  // dimensions
681            nullable_object_field(&intr), // maxInputTokens
682            f64_field(),                  // maxInputBytes
683            string_field(&intr),          // identity
684        ],
685    )
686}
687
688/// Seal `batch`: copy its vectors into a packed byte array in the GC heap and
689/// hand the guest a backing struct that holds it in a hidden field.
690///
691/// Charge `count × dimensions × 4` bytes once, before encoding. A recorded
692/// native batch can be replayed if this conversion runs out of fuel.
693/// At most two host copies of the vectors are live at once (the encoded bytes
694/// and the engine's internal copy); the batch's own floats are gone before the
695/// array is allocated.
696fn build_embeddings(
697    caller: &mut wasmtime::Caller<'_, StoreData>,
698    batch: EmbeddingBatch,
699) -> wasmtime::Result<Val> {
700    let byte_count = batch
701        .values()
702        .len()
703        .checked_mul(4)
704        .ok_or_else(|| fatal_host_error("embedding.embed: result size overflows"))?;
705    fuel::charge(&mut *caller, fuel::COPY, byte_count as u64)?;
706    let mut bytes = Vec::new();
707    bytes
708        .try_reserve_exact(byte_count)
709        .map_err(|error| fatal_host_error(format!("embedding.embed: {error}")))?;
710    for value in batch.values() {
711        bytes.extend_from_slice(&value.to_le_bytes());
712    }
713    let count = batch.count() as f64;
714    let dimensions = batch.dimensions() as f64;
715    let identity = batch.identity().to_string();
716    let model = batch.model().to_string();
717    let input_tokens = batch.input_tokens().map(|n| n as f64);
718    drop(batch);
719
720    // Allocated through the GC limiter, which collects before the heap grows.
721    let array_ty = build_intrinsic_types(caller.engine())?.raw_uint8_array;
722    let vectors = write_uint8_array_precharged(&mut *caller, array_ty, &bytes)?;
723    drop(bytes);
724
725    let identity = string_val(caller, &identity)?;
726    let model = string_val(caller, &model)?;
727    let input_tokens = optional_number(caller, input_tokens)?;
728
729    let ty = embeddings_backing_struct(caller.engine())?;
730    abi::new_backing(
731        caller,
732        ty,
733        &[
734            Val::AnyRef(Some(vectors.to_anyref())),
735            Val::F64(count.to_bits()),
736            Val::F64(dimensions.to_bits()),
737            identity,
738            model,
739            input_tokens,
740        ],
741    )
742}
743
744fn build_model(
745    caller: &mut wasmtime::Caller<'_, StoreData>,
746    model: EmbeddingModel,
747) -> wasmtime::Result<Val> {
748    let name = string_val(caller, &model.name)?;
749    let description = match model.description.as_deref().and_then(sanitize_description) {
750        Some(text) => string_val(caller, &text)?,
751        None => crate::runtime::prelude::undefined::value(caller)?,
752    };
753    let max_input_tokens = optional_number(caller, model.max_input_tokens.map(|n| n as f64))?;
754    let identity = string_val(caller, &model.identity)?;
755    let ty = model_backing_struct(caller.engine())?;
756    abi::new_backing(
757        caller,
758        ty,
759        &[
760            name,
761            description,
762            Val::F64((model.dimensions as f64).to_bits()),
763            max_input_tokens,
764            Val::F64((model.max_input_bytes as f64).to_bits()),
765            identity,
766        ],
767    )
768}
769
770fn string_val(caller: &mut wasmtime::Caller<'_, StoreData>, text: &str) -> wasmtime::Result<Val> {
771    Ok(Val::AnyRef(Some(
772        write_submilli_string_struct(caller, text)?.to_anyref(),
773    )))
774}
775
776/// Copy row `index` out of the sealed result. The index must be an integer below
777/// `count`; anything else is the program's mistake, a `RangeError`.
778fn read_row(
779    caller: &mut wasmtime::Caller<'_, StoreData>,
780    receiver: &Val,
781    index: &Val,
782    ctx: &str,
783) -> wasmtime::Result<Vec<u8>> {
784    let Val::F64(bits) = index else {
785        return Err(fatal_host_error(format!("{ctx}: index is not a number")));
786    };
787    let index = f64::from_bits(*bits);
788
789    let st = abi::backing_receiver(caller, receiver)?;
790    let Val::AnyRef(Some(vectors)) = st.field(&mut *caller, E_VECTORS)? else {
791        return Err(fatal_host_error(format!(
792            "{ctx}: result vectors are missing"
793        )));
794    };
795    let vectors = vectors
796        .as_array(&mut *caller)?
797        .ok_or_else(|| fatal_host_error(format!("{ctx}: result vectors are not an array")))?;
798    let (Val::F64(count), Val::F64(dimensions)) = (
799        st.field(&mut *caller, E_COUNT)?,
800        st.field(&mut *caller, E_DIMENSIONS)?,
801    ) else {
802        return Err(fatal_host_error(format!("{ctx}: result shape is missing")));
803    };
804    let (count, dimensions) = (f64::from_bits(count), f64::from_bits(dimensions));
805    let in_range = index.is_finite() && index.fract() == 0.0 && index >= 0.0 && index < count;
806    if !in_range {
807        return Err(range_error(format!(
808            "{ctx}: index {index} is out of range for {count} vectors — use an integer from 0 to \
809             count - 1"
810        )));
811    }
812    let row_bytes = (dimensions as usize)
813        .checked_mul(4)
814        .ok_or_else(|| fatal_host_error(format!("{ctx}: row size overflows")))?;
815    let offset = (index as usize)
816        .checked_mul(row_bytes)
817        .and_then(|offset| u32::try_from(offset).ok())
818        .ok_or_else(|| fatal_host_error(format!("{ctx}: row offset overflows")))?;
819    // The row copy-out is charged by size, before it is made.
820    fuel::charge(&mut *caller, fuel::COPY, row_bytes as u64)?;
821    let mut row = Vec::new();
822    row.try_reserve_exact(row_bytes)
823        .map_err(|error| fatal_host_error(format!("{ctx}: {error}")))?;
824    row.resize(row_bytes, 0);
825    vectors
826        .read_i8(&mut *caller, offset, &mut row)
827        .map_err(|error| fatal_host_error(format!("{ctx}: {error}")))?;
828    Ok(row)
829}
830
831/// A `number[]` of the row's components, widened from `f32`.
832fn build_number_array(
833    caller: &mut wasmtime::Caller<'_, StoreData>,
834    row: &[u8],
835) -> wasmtime::Result<Val> {
836    // Admit the whole array before boxing any element.
837    fuel::charge(&mut *caller, fuel::ELEM, (row.len() / 4) as u64)?;
838    let mut boxed = Vec::with_capacity(row.len() / 4);
839    for chunk in row.as_chunks::<4>().0 {
840        let value = f32::from_le_bytes(*chunk);
841        boxed.push(Val::AnyRef(Some(
842            write_boxed_number_struct(caller, f64::from(value))?.to_anyref(),
843        )));
844    }
845    Ok(Val::AnyRef(Some(
846        write_submilli_array_struct_precharged(caller, &boxed)?.to_anyref(),
847    )))
848}
849
850#[cfg(test)]
851mod tests {
852    use std::pin::Pin;
853    use std::sync::atomic::{AtomicUsize, Ordering};
854    use std::sync::{Arc, Mutex};
855
856    use crate::runtime::{
857        CheckOutcome, EmbeddingBatch, EmbeddingError, EmbeddingFailureReason, EmbeddingLimits,
858        EmbeddingModel, EmbeddingProvider, EmbeddingTokenBudget, Purpose, RuntimeConfig,
859        SecurityCheck, SharedTokenBudget, StoreData, SubBatchSettlement, Vfs, dispatch_main_async,
860        install_runtime_async,
861    };
862
863    type BoxFuture<'a, T> = Pin<Box<dyn std::future::Future<Output = T> + Send + 'a>>;
864
865    /// What the mock answers a send with.
866    #[derive(Clone, Copy)]
867    enum Outcome {
868        /// Succeeds and reports usage equal to the estimate.
869        Reported,
870        /// Fails after sending; the estimate stays held.
871        TransportFailure,
872        /// Sends, then reports an internal bookkeeping failure.
873        Internal,
874        /// Sent, then failed, having reported this many tokens.
875        FailedAfter(u64),
876        /// Returns one row for however many texts were sent, with 50 input tokens.
877        CountMismatch,
878    }
879
880    /// Counts `embed` calls and answers deterministically. The `embed` count is
881    /// what proves "the provider was not called".
882    struct MockProvider {
883        dimensions: usize,
884        max_input_bytes: u64,
885        outcome: Outcome,
886        embed_calls: AtomicUsize,
887        /// Provider-reported input tokens on a successful batch. `None` leaves
888        /// the batch's count unreported, which is the older mock's behavior.
889        input_tokens: Option<u64>,
890    }
891
892    impl MockProvider {
893        fn new(dimensions: usize, max_input_bytes: u64, outcome: Outcome) -> Arc<Self> {
894            Arc::new(Self {
895                dimensions,
896                max_input_bytes,
897                outcome,
898                embed_calls: AtomicUsize::new(0),
899                input_tokens: None,
900            })
901        }
902
903        fn reporting(tokens: u64) -> Arc<Self> {
904            Arc::new(Self {
905                dimensions: 2,
906                max_input_bytes: 1_000,
907                outcome: Outcome::Reported,
908                embed_calls: AtomicUsize::new(0),
909                input_tokens: Some(tokens),
910            })
911        }
912
913        fn calls(&self) -> usize {
914            self.embed_calls.load(Ordering::Relaxed)
915        }
916    }
917
918    impl EmbeddingProvider for MockProvider {
919        fn embed<'a>(
920            &'a self,
921            alias: &'a str,
922            texts: &'a [String],
923            _purpose: Purpose,
924            budget: &'a EmbeddingTokenBudget,
925        ) -> BoxFuture<'a, Result<EmbeddingBatch, EmbeddingError>> {
926            self.embed_calls.fetch_add(1, Ordering::Relaxed);
927            Box::pin(async move {
928                let estimate = self.estimate_tokens(alias, texts);
929                budget.mark_sent(alias, estimate)?;
930                if matches!(self.outcome, Outcome::Internal) {
931                    return Err(EmbeddingError::Internal {
932                        alias: alias.to_string(),
933                        settlements: vec![SubBatchSettlement {
934                            estimate,
935                            reported: 0,
936                            indeterminate: estimate,
937                        }],
938                    });
939                }
940                if matches!(self.outcome, Outcome::TransportFailure) {
941                    return Err(EmbeddingError::Provider {
942                        alias: alias.to_string(),
943                        reason: EmbeddingFailureReason::Transport,
944                        settlements: vec![SubBatchSettlement {
945                            estimate,
946                            reported: 0,
947                            indeterminate: estimate,
948                        }],
949                    });
950                }
951                if matches!(self.outcome, Outcome::CountMismatch) {
952                    let values = vec![0.5f32; self.dimensions];
953                    return EmbeddingBatch::new(values, 1, self.dimensions, "emb1:mock", alias)
954                        .map(|batch| {
955                            batch.with_input_tokens(Some(50)).with_settlements(vec![
956                                SubBatchSettlement {
957                                    estimate,
958                                    reported: 50,
959                                    indeterminate: 0,
960                                },
961                            ])
962                        })
963                        .map_err(|_| EmbeddingError::Malformed {
964                            alias: alias.to_string(),
965                            reason: crate::runtime::EmbeddingMalformedReason::CountMismatch,
966                            settlements: Vec::new(),
967                        });
968                }
969                if let Outcome::FailedAfter(reported) = self.outcome {
970                    return Err(EmbeddingError::Provider {
971                        alias: alias.to_string(),
972                        reason: EmbeddingFailureReason::Transport,
973                        settlements: vec![SubBatchSettlement {
974                            estimate,
975                            reported,
976                            indeterminate: estimate.saturating_sub(reported),
977                        }],
978                    });
979                }
980                let values = vec![0.5f32; texts.len() * self.dimensions];
981                EmbeddingBatch::new(values, texts.len(), self.dimensions, "emb1:mock", alias)
982                    .map(|batch| {
983                        batch
984                            .with_input_tokens(self.input_tokens)
985                            .with_settlements(vec![SubBatchSettlement {
986                                estimate,
987                                reported: self.input_tokens.unwrap_or(estimate),
988                                indeterminate: 0,
989                            }])
990                    })
991                    .map_err(|_| EmbeddingError::Malformed {
992                        alias: alias.to_string(),
993                        reason: crate::runtime::EmbeddingMalformedReason::InvalidBody,
994                        settlements: Vec::new(),
995                    })
996            })
997        }
998
999        fn models<'a>(&'a self) -> BoxFuture<'a, Result<Vec<EmbeddingModel>, EmbeddingError>> {
1000            let model = EmbeddingModel {
1001                name: "mock".to_string(),
1002                description: None,
1003                dimensions: self.dimensions as u64,
1004                max_input_tokens: None,
1005                max_input_bytes: self.max_input_bytes,
1006                identity: "emb1:mock".to_string(),
1007            };
1008            Box::pin(async move { Ok(vec![model]) })
1009        }
1010
1011        fn max_input_bytes(&self, alias: &str) -> Option<u64> {
1012            (alias == "mock").then_some(self.max_input_bytes)
1013        }
1014    }
1015
1016    struct RecordingPolicy(Arc<Mutex<Vec<serde_json::Value>>>);
1017
1018    impl SecurityCheck for RecordingPolicy {
1019        fn check(
1020            &self,
1021            _caller: &str,
1022            _capability: &str,
1023            context: &serde_json::Value,
1024        ) -> CheckOutcome {
1025            if let Ok(mut contexts) = self.0.lock() {
1026                contexts.push(context.clone());
1027            }
1028            CheckOutcome::Allow { rule: None }
1029        }
1030    }
1031
1032    /// Runs `source` against the provider and returns what the host observed:
1033    /// the `main` result and the host-attached bytes still held at the end.
1034    struct Run {
1035        result: wasmtime::Result<String>,
1036        host_attached_bytes: u64,
1037        /// GC heap growth the limiter has observed for the whole run.
1038        observed_bytes: u64,
1039        model_calls: Vec<crate::runtime::ModelCallUsage>,
1040    }
1041
1042    async fn run(
1043        source: &str,
1044        provider: Option<Arc<MockProvider>>,
1045        budget: Option<Arc<EmbeddingTokenBudget>>,
1046        policy: Option<Arc<dyn SecurityCheck>>,
1047    ) -> Run {
1048        let compiled = crate::compile_script(source, "test.ts", crate::FileId(0), &[], &[])
1049            .expect("compile clean");
1050        let cfg = RuntimeConfig::default();
1051        let engine = cfg.engine().expect("engine");
1052        let mut data = StoreData::with_vfs("test-run".into(), Vfs::tempdir().expect("tempdir"));
1053        data.install_type_info(compiled.type_info.clone());
1054        data.embedding_provider = provider.map(|p| p as Arc<dyn EmbeddingProvider>);
1055        data.embedding_budget = budget;
1056        if let Some(policy) = policy {
1057            data.security_check = policy;
1058        }
1059        let mut store = cfg.store_async(&engine, data).expect("store");
1060        crate::runtime::install_tenant_limits(&mut store);
1061        let module = wasmtime::Module::new(&engine, &compiled.wasm).expect("module");
1062        let mut linker = wasmtime::Linker::<StoreData>::new(&engine);
1063        install_runtime_async(&mut linker, &mut store)
1064            .await
1065            .expect("install");
1066        let inst = linker
1067            .instantiate_async(&mut store, &module)
1068            .await
1069            .expect("instantiate");
1070        let result = dispatch_main_async(&mut store, &inst)
1071            .await
1072            .map(Option::unwrap_or_default);
1073        let host_attached_bytes = store.data().tenant_limits.host_attached_bytes();
1074        let observed_bytes = store.data().tenant_limits.observed_bytes();
1075        let model_calls = store.data().model_calls().to_vec();
1076        Run {
1077            result,
1078            host_attached_bytes,
1079            observed_bytes,
1080            model_calls,
1081        }
1082    }
1083
1084    fn budget(limits: EmbeddingLimits) -> Arc<EmbeddingTokenBudget> {
1085        Arc::new(EmbeddingTokenBudget::new(
1086            limits,
1087            SharedTokenBudget::new(u64::MAX),
1088        ))
1089    }
1090
1091    const ROUND_TRIP: &str = r#"import embedding from "submilli:embedding";
1092        function main(): void {
1093          const r = embedding.embed("mock", ["a", "b", "c"], "document");
1094          assert(r.count === 3, "three vectors");
1095        }"#;
1096
1097    /// `Embeddings` and `EmbeddingModel` stringify through their properties;
1098    /// the vectors, read through methods, are not among them.
1099    #[tokio::test]
1100    async fn embeddings_and_models_stringify_their_properties() {
1101        let provider = MockProvider::new(4, 10, Outcome::Reported);
1102        let run = run(
1103            r#"import embedding from "submilli:embedding";
1104               function main(): string {
1105                 const r = embedding.embed("mock", ["a"], "document");
1106                 return JSON.stringify(r)! + "|" + JSON.stringify(embedding.models())!;
1107               }"#,
1108            Some(provider),
1109            None,
1110            None,
1111        )
1112        .await;
1113        let out = run.result.expect("program completes");
1114        let (embeddings, models) = out.split_once('|').expect("two values");
1115        let embeddings: serde_json::Value = serde_json::from_str(embeddings).expect("JSON");
1116        assert_eq!(embeddings["count"], 1, "{embeddings}");
1117        assert_eq!(embeddings["dimensions"], 4, "{embeddings}");
1118        assert_eq!(embeddings["model"], "mock", "{embeddings}");
1119        assert!(embeddings.get("vectors").is_none(), "{embeddings}");
1120        assert_eq!(
1121            models,
1122            r#"[{"name":"mock","dimensions":4,"maxInputBytes":10,"identity":"emb1:mock"}]"#
1123        );
1124    }
1125
1126    /// AE1: an over-length input is refused before the provider is called, and
1127    /// the budget is unchanged.
1128    #[tokio::test]
1129    async fn ae1_over_length_input_is_refused_before_the_provider_and_the_budget() {
1130        let provider = MockProvider::new(4, 10, Outcome::Reported);
1131        let budget = budget(EmbeddingLimits::default());
1132        let run = run(
1133            r#"import embedding from "submilli:embedding";
1134               function main(): void {
1135                 let message = "";
1136                 try {
1137                   embedding.embed("mock", ["ok", "ok", "this one is too long"], "document");
1138                 } catch (e: RangeError) {
1139                   message = e.message;
1140                 }
1141                 assert(message.indexOf("text 2") >= 0, "names index 2: " + message);
1142               }"#,
1143            Some(Arc::clone(&provider)),
1144            Some(Arc::clone(&budget)),
1145            None,
1146        )
1147        .await;
1148        run.result.expect("program completes");
1149        assert_eq!(provider.calls(), 0, "the provider is never called");
1150        assert_eq!(budget.used(), 0, "nothing is charged");
1151        assert_eq!(budget.held(), 0, "nothing is held");
1152        assert_eq!(budget.requests(), 0, "nothing counts as sent");
1153    }
1154
1155    /// AE6: an exhausted per-run budget is a `QuotaExceededError` and nothing is
1156    /// sent.
1157    #[tokio::test]
1158    async fn ae6_exhausted_budget_refuses_without_calling_the_provider() {
1159        let provider = MockProvider::new(4, 1_000, Outcome::Reported);
1160        let budget = budget(EmbeddingLimits {
1161            per_execution_tokens: 1,
1162            ..EmbeddingLimits::default()
1163        });
1164        let run = run(
1165            r#"import embedding from "submilli:embedding";
1166               function main(): void {
1167                 let quota = false;
1168                 try {
1169                   embedding.embed("mock", ["more than one token of text"], "document");
1170                 } catch (e: QuotaExceededError) {
1171                   quota = true;
1172                 }
1173                 assert(quota, "the call is a QuotaExceededError");
1174               }"#,
1175            Some(Arc::clone(&provider)),
1176            Some(Arc::clone(&budget)),
1177            None,
1178        )
1179        .await;
1180        run.result.expect("program completes");
1181        assert_eq!(provider.calls(), 0, "the provider is never called");
1182        assert_eq!(budget.used(), 0, "the refusal charged nothing");
1183        assert_eq!(budget.requests(), 0);
1184    }
1185
1186    /// Settlement is applied once on success: the reported usage is committed and
1187    /// nothing stays held.
1188    #[tokio::test]
1189    async fn success_settles_reported_usage() {
1190        let provider = MockProvider::new(4, 1_000, Outcome::Reported);
1191        let budget = budget(EmbeddingLimits::default());
1192        let run = run(ROUND_TRIP, Some(provider), Some(Arc::clone(&budget)), None).await;
1193        run.result.expect("program completes");
1194        assert_eq!(
1195            budget.used(),
1196            3,
1197            "three one-byte texts estimate one token each"
1198        );
1199        assert_eq!(budget.held(), 0, "reported usage leaves nothing held");
1200        assert_eq!(budget.requests(), 1);
1201    }
1202
1203    /// A send that fails after leaving keeps its estimate held (unknown spend is
1204    /// indeterminate, not free), and the program sees a catchable error.
1205    #[tokio::test]
1206    async fn failure_after_send_settles_the_estimate_as_held() {
1207        let provider = MockProvider::new(4, 1_000, Outcome::TransportFailure);
1208        let budget = budget(EmbeddingLimits::default());
1209        let run = run(
1210            r#"import embedding from "submilli:embedding";
1211               function main(): void {
1212                 let message = "";
1213                 try {
1214                   embedding.embed("mock", ["secret text"], "document");
1215                 } catch (e: Error) {
1216                   message = e.message;
1217                 }
1218                 assert(message.indexOf("transport") >= 0, message);
1219                 assert(message.indexOf("secret text") < 0, "no input text in the error");
1220               }"#,
1221            Some(provider),
1222            Some(Arc::clone(&budget)),
1223            None,
1224        )
1225        .await;
1226        run.result.expect("program completes");
1227        assert_eq!(budget.held(), 4, "ceil(11 / 3) tokens stay held");
1228        assert_eq!(budget.used(), 4, "and count against the run");
1229    }
1230
1231    /// AE8: a 128 x 3,072 result is held at 4 bytes per number in the GC heap —
1232    /// not the ~60 bytes a plain array costs — and reading one vector or its
1233    /// bytes works.
1234    #[tokio::test]
1235    async fn ae8_results_are_held_at_four_bytes_per_number() {
1236        const COUNT: u64 = 128;
1237        const DIMENSIONS: u64 = 3072;
1238        let source = |texts: u32| {
1239            format!(
1240                r#"import embedding from "submilli:embedding";
1241               function main(): void {{
1242                 const texts: string[] = [];
1243                 for (let i = 0; i < {texts}; i = i + 1) {{
1244                   texts.push("t");
1245                 }}
1246                 const r = embedding.embed("mock", texts, "document");
1247                 assert(r.count === {texts} && r.dimensions === 3072, "shape");
1248                 assert(r.vector(0).length === 3072, "reading vector 0 returns 3,072 numbers");
1249                 assert(r.bytes(0).length === 12288, "exporting it returns 12,288 bytes");
1250                 if ({texts} > 5) {{
1251                   assert(r.vector(5).length === 3072, "reading vector 5 returns 3,072 numbers");
1252                   assert(r.bytes(5).length === 12288, "exporting it returns 12,288 bytes");
1253                 }}
1254                 let range = false;
1255                 try {{
1256                   r.vector({texts});
1257                 }} catch (e: RangeError) {{
1258                   range = true;
1259                 }}
1260                 assert(range, "reading vector {texts} is a RangeError");
1261               }}"#
1262            )
1263        };
1264        let provider = |_| MockProvider::new(DIMENSIONS as usize, 1_000, Outcome::Reported);
1265        let one = run(&source(1), Some(provider(())), None, None).await;
1266        one.result.expect("one-vector program completes");
1267        let many = run(&source(COUNT as u32), Some(provider(())), None, None).await;
1268        many.result.expect("program completes");
1269
1270        // The two runs differ only in the result: 127 more vectors. The GC heap
1271        // growth the limiter observed grows by that many vectors at 4 bytes per
1272        // number (within 10%: the heap grows in steps, and 127 one-character texts ride along).
1273        let expected = (COUNT - 1) * DIMENSIONS * 4;
1274        let grown = many.observed_bytes.saturating_sub(one.observed_bytes);
1275        assert!(
1276            grown >= expected - expected / 10 && grown <= expected + expected / 10,
1277            "the result costs about count x dimensions x 4 bytes of GC heap: grew {grown}, \
1278             expected about {expected}"
1279        );
1280        assert_eq!(
1281            many.host_attached_bytes, 0,
1282            "the vectors are not host-attached bytes"
1283        );
1284    }
1285
1286    /// Discarded results are reclaimed by collection: 100 results of 128 x 3,072
1287    /// numbers (1.5 MiB each, 150 MiB in all) fit the default 50 MB store when
1288    /// each loop pass drops the last one. The texts are built once so the loop
1289    /// allocates almost nothing else and only the results press on the heap.
1290    #[tokio::test]
1291    async fn discarded_results_are_collected_under_the_default_cap() {
1292        let provider = MockProvider::new(3072, 1_000, Outcome::Reported);
1293        let run = run(
1294            r#"import embedding from "submilli:embedding";
1295               function main(): void {
1296                 const texts: string[] = [];
1297                 for (let i = 0; i < 128; i = i + 1) {
1298                   texts.push("t");
1299                 }
1300                 let total = 0;
1301                 for (let round = 0; round < 100; round = round + 1) {
1302                   const r = embedding.embed("mock", texts, "document");
1303                   total = total + r.count;
1304                 }
1305                 assert(total === 12800, "every round returned its vectors");
1306               }"#,
1307            Some(provider),
1308            Some(budget(EmbeddingLimits::default())),
1309            None,
1310        )
1311        .await;
1312        run.result
1313            .expect("program completes without memory exhaustion");
1314    }
1315
1316    /// The hidden vectors array is sealed: no declared member reaches it, and a
1317    /// cast to `Uint8Array` is refused at run time rather than aliasing it.
1318    #[tokio::test]
1319    async fn the_hidden_vectors_are_not_reachable_from_the_guest() {
1320        for member in ["vectors", "handle", "data", "buffer"] {
1321            let source = format!(
1322                r#"import embedding from "submilli:embedding";
1323                   function main(): void {{
1324                     const r = embedding.embed("mock", ["a"], "document");
1325                     const hidden = r.{member};
1326                   }}"#
1327            );
1328            assert!(
1329                crate::compile_script(&source, "test.ts", crate::FileId(0), &[], &[]).is_err(),
1330                "`Embeddings.{member}` must not type-check"
1331            );
1332        }
1333        let provider = MockProvider::new(4, 1_000, Outcome::Reported);
1334        let cast = run(
1335            r#"import embedding from "submilli:embedding";
1336               function main(): void {
1337                 const r = embedding.embed("mock", ["a"], "document");
1338                 const u = r as unknown as Uint8Array;
1339                 u[0] = 255;
1340               }"#,
1341            Some(provider),
1342            None,
1343            None,
1344        )
1345        .await;
1346        assert!(
1347            cast.result.is_err(),
1348            "casting the result to a Uint8Array must fail, not expose the vectors"
1349        );
1350
1351        // Hashing the result as bytes must never read the vectors. The cast to
1352        // `string | Uint8Array` refuses at run time (the value is neither); the
1353        // host-level refusal behind it is pinned by
1354        // `a_host_read_of_the_result_as_a_uint8array_is_refused`.
1355        for call in [
1356            "crypto.sha256(x)",
1357            "crypto.hmacSha256(crypto.randomBytes(4), x)",
1358        ] {
1359            let provider = MockProvider::new(4, 1_000, Outcome::Reported);
1360            let source = format!(
1361                r#"import embedding from "submilli:embedding";
1362                   import crypto from "submilli:crypto";
1363                   function main(): void {{
1364                     const r = embedding.embed("mock", ["a"], "document");
1365                     const x = r as unknown as (string | Uint8Array);
1366                     crypto.sha256(crypto.randomBytes(1));
1367                     {call};
1368                   }}"#
1369            );
1370            let outcome = run(&source, Some(provider), None, None).await;
1371            assert!(
1372                outcome.result.is_err(),
1373                "`{call}` on the result must fail, not hash the vectors"
1374            );
1375        }
1376    }
1377
1378    /// The host-side seal behind every `Uint8Array` argument: a backing struct
1379    /// with a hidden byte field is not a `$Uint8Array`, so reading it as one is
1380    /// refused even if a cast were bypassed.
1381    #[tokio::test]
1382    async fn a_host_read_of_the_result_as_a_uint8array_is_refused() {
1383        let compiled = crate::compile_script(
1384            "function main(): void {}",
1385            "test.ts",
1386            crate::FileId(0),
1387            &[],
1388            &[],
1389        )
1390        .expect("compile clean");
1391        let cfg = RuntimeConfig::default();
1392        let engine = cfg.engine().expect("engine");
1393        let mut data = StoreData::with_vfs("test-run".into(), Vfs::tempdir().expect("tempdir"));
1394        data.install_type_info(compiled.type_info.clone());
1395        let mut store = cfg.store_async(&engine, data).expect("store");
1396        crate::runtime::install_tenant_limits(&mut store);
1397        let module = wasmtime::Module::new(&engine, &compiled.wasm).expect("module");
1398        let mut linker = wasmtime::Linker::<StoreData>::new(&engine);
1399        install_runtime_async(&mut linker, &mut store)
1400            .await
1401            .expect("install");
1402        linker
1403            .instantiate_async(&mut store, &module)
1404            .await
1405            .expect("instantiate");
1406
1407        let probe = wasmtime::Func::new(
1408            &mut store,
1409            wasmtime::FuncType::new(&engine, [], []),
1410            |mut caller, _, _| {
1411                let batch = EmbeddingBatch::new(vec![1.0, 2.0], 1, 2, "id", "mock")
1412                    .map_err(wasmtime::Error::new)?;
1413                let sealed = super::build_embeddings(&mut caller, batch)?;
1414                let refused =
1415                    crate::runtime::host::read_uint8_array_arg(&mut caller, &sealed, "probe");
1416                match refused {
1417                    Err(error) if error.to_string().contains("expects a Uint8Array") => Ok(()),
1418                    other => Err(wasmtime::Error::msg(format!(
1419                        "the sealed result was not refused: {other:?}"
1420                    ))),
1421                }
1422            },
1423        );
1424        probe
1425            .call_async(&mut store, &[], &mut [])
1426            .await
1427            .expect("the host refuses to read the sealed result as bytes");
1428    }
1429
1430    /// The vectors count against the run's memory cap like any GC object: a store
1431    /// too small for the result ends the run.
1432    #[tokio::test]
1433    async fn the_result_charge_counts_against_the_run_memory_cap() {
1434        let provider = MockProvider::new(3072, 1_000, Outcome::Reported);
1435        let compiled = crate::compile_script(
1436            r#"import embedding from "submilli:embedding";
1437               function main(): void {
1438                 const texts: string[] = [];
1439                 for (let i = 0; i < 128; i = i + 1) {
1440                   texts.push("t");
1441                 }
1442                 embedding.embed("mock", texts, "document");
1443               }"#,
1444            "test.ts",
1445            crate::FileId(0),
1446            &[],
1447            &[],
1448        )
1449        .expect("compile clean");
1450        let cfg = RuntimeConfig::default();
1451        let engine = cfg.engine().expect("engine");
1452        let mut data = StoreData::with_vfs("test-run".into(), Vfs::tempdir().expect("tempdir"));
1453        data.install_type_info(compiled.type_info.clone());
1454        data.embedding_provider = Some(provider);
1455        let budget = budget(EmbeddingLimits::default());
1456        data.embedding_budget = Some(Arc::clone(&budget));
1457        let mut store = cfg.store_async(&engine, data).expect("store");
1458        crate::runtime::install_tenant_limits(&mut store);
1459        let module = wasmtime::Module::new(&engine, &compiled.wasm).expect("module");
1460        let mut linker = wasmtime::Linker::<StoreData>::new(&engine);
1461        install_runtime_async(&mut linker, &mut store)
1462            .await
1463            .expect("install");
1464        let inst = linker
1465            .instantiate_async(&mut store, &module)
1466            .await
1467            .expect("instantiate");
1468        // Leave the program less than the 128 x 3,072 x 4 = 1.5 MiB of
1469        // vectors it is about to ask for.
1470        let cap = store.data().tenant_limits.observed_bytes() + 100_000;
1471        store.data_mut().tenant_limits.max_total_bytes = cap;
1472        let result = dispatch_main_async(&mut store, &inst).await;
1473        let error = result.expect_err("the cap refuses the result");
1474        assert!(
1475            crate::runtime::is_memory_exhausted(&error),
1476            "a memory refusal ends the run: {error:?}"
1477        );
1478        // The result is built after settlement: the provider's spend is
1479        // committed even though the guest never sees it.
1480        assert_eq!(budget.used(), 128, "the reported usage is settled");
1481        assert_eq!(budget.held(), 0, "nothing stays held");
1482    }
1483
1484    /// An internal provider failure ends the run rather than becoming a catchable
1485    /// error, and the sent sub-batch stays held.
1486    #[tokio::test]
1487    async fn an_internal_provider_failure_ends_the_run_and_settles_the_budget() {
1488        let provider = MockProvider::new(4, 1_000, Outcome::Internal);
1489        let budget = budget(EmbeddingLimits::default());
1490        let run = run(
1491            r#"import embedding from "submilli:embedding";
1492               function main(): void {
1493                 try {
1494                   embedding.embed("mock", ["secret text"], "document");
1495                 } catch (e: Error) {
1496                   assert(false, "an internal failure must not be catchable");
1497                 }
1498               }"#,
1499            Some(provider),
1500            Some(Arc::clone(&budget)),
1501            None,
1502        )
1503        .await;
1504        let error = run.result.expect_err("the run ends");
1505        assert!(
1506            crate::runtime::host::ends_the_run(&error),
1507            "an internal failure is fatal: {error:?}"
1508        );
1509        assert_eq!(budget.held(), 4, "ceil(11 / 3) tokens stay held");
1510        assert_eq!(budget.used(), 4);
1511    }
1512
1513    /// The filter context is exactly `model` and `input_count`; input text never
1514    /// reaches policy, and `models()` asks with `input_count` 0.
1515    #[tokio::test]
1516    async fn the_filter_context_carries_the_numbers_and_never_the_text() {
1517        const SECRET: &str = "the patient's diagnosis is confidential";
1518        let provider = MockProvider::new(4, 1_000, Outcome::Reported);
1519        let contexts = Arc::new(Mutex::new(Vec::new()));
1520        let policy: Arc<dyn SecurityCheck> = Arc::new(RecordingPolicy(Arc::clone(&contexts)));
1521        let run = run(
1522            &format!(
1523                r#"import embedding from "submilli:embedding";
1524                   function main(): void {{
1525                     embedding.embed("mock", ["{SECRET}", "b"], "query");
1526                     embedding.models();
1527                   }}"#
1528            ),
1529            Some(provider),
1530            None,
1531            Some(policy),
1532        )
1533        .await;
1534        run.result.expect("program completes");
1535
1536        let contexts = contexts.lock().expect("contexts");
1537        assert_eq!(contexts.len(), 2, "one check per embed, one per candidate");
1538        for context in contexts.iter() {
1539            let mut keys: Vec<&str> = context
1540                .as_object()
1541                .expect("an object context")
1542                .keys()
1543                .map(String::as_str)
1544                .collect();
1545            keys.sort_unstable();
1546            assert_eq!(keys, ["input_count", "model"]);
1547            assert!(!context.to_string().contains("patient"), "{context}");
1548        }
1549        assert_eq!(contexts[0]["input_count"], 2, "embed: the text count");
1550        assert_eq!(contexts[1]["input_count"], 0, "models: zero");
1551    }
1552
1553    /// An unknown alias is refused after the gate and before the budget, naming
1554    /// the aliases that exist.
1555    #[tokio::test]
1556    async fn an_unknown_alias_names_the_available_ones_and_charges_nothing() {
1557        let provider = MockProvider::new(4, 1_000, Outcome::Reported);
1558        let budget = budget(EmbeddingLimits::default());
1559        let run = run(
1560            r#"import embedding from "submilli:embedding";
1561               function main(): void {
1562                 let message = "";
1563                 try {
1564                   embedding.embed("ghost", ["a"], "document");
1565                 } catch (e: Error) {
1566                   message = e.message;
1567                 }
1568                 assert(message.indexOf("mock") >= 0, message);
1569               }"#,
1570            Some(Arc::clone(&provider)),
1571            Some(Arc::clone(&budget)),
1572            None,
1573        )
1574        .await;
1575        run.result.expect("program completes");
1576        assert_eq!(provider.calls(), 0);
1577        assert_eq!(budget.used(), 0);
1578    }
1579
1580    /// A successful embed records the alias the program passed and the tokens
1581    /// the provider reported. That record is not an llm call of the same name.
1582    #[tokio::test]
1583    async fn an_embed_records_input_tokens_for_the_alias() {
1584        let provider = MockProvider::reporting(42);
1585        let run = run(
1586            r#"import embedding from "submilli:embedding";
1587               function main(): void {
1588                 embedding.embed("mock", ["a", "b"], "document");
1589               }"#,
1590            Some(provider),
1591            None,
1592            None,
1593        )
1594        .await;
1595        run.result.expect("program completes");
1596        let calls = run.model_calls;
1597        assert_eq!(
1598            calls,
1599            vec![crate::runtime::ModelCallUsage::embedding("mock", Some(42))]
1600        );
1601        assert_eq!(
1602            crate::runtime::ModelCallUsage::summed_by_model(&calls),
1603            vec![crate::runtime::ModelCallUsage::embedding("mock", Some(42))]
1604        );
1605    }
1606
1607    /// A failed embed that already reported tokens still records them. The call
1608    /// log does, and the embedder's bill must not disagree.
1609    #[tokio::test]
1610    async fn a_failed_embed_that_reported_tokens_keeps_them() {
1611        let provider = MockProvider::new(2, 1_000, Outcome::FailedAfter(4));
1612        let run = run(
1613            r#"import embedding from "submilli:embedding";
1614               function main(): void {
1615                 embedding.embed("mock", ["a"], "document");
1616               }"#,
1617            Some(provider),
1618            None,
1619            None,
1620        )
1621        .await;
1622        assert!(run.result.is_err(), "the failed embed reaches the guest");
1623        assert_eq!(
1624            run.model_calls,
1625            vec![crate::runtime::ModelCallUsage::embedding("mock", Some(4))]
1626        );
1627    }
1628
1629    /// A provider `Ok` whose row count does not match still ran. The guest sees
1630    /// an error; the embedder still sees the tokens the batch reported.
1631    #[tokio::test]
1632    async fn a_count_mismatch_still_records_reported_tokens() {
1633        let provider = MockProvider::new(2, 1_000, Outcome::CountMismatch);
1634        let run = run(
1635            r#"import embedding from "submilli:embedding";
1636               function main(): void {
1637                 embedding.embed("mock", ["a", "b"], "document");
1638               }"#,
1639            Some(provider),
1640            None,
1641            None,
1642        )
1643        .await;
1644        assert!(run.result.is_err(), "a short batch is a guest error");
1645        assert_eq!(
1646            run.model_calls,
1647            vec![crate::runtime::ModelCallUsage::embedding("mock", Some(50))]
1648        );
1649    }
1650}