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