1pub 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
66pub const CAPABILITY: &str = "embedding.embed";
68
69const 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
77const 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 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
182pub(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
231async 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 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 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
304struct 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 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 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
417fn 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
440async 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
454fn 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
504async 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 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
534async 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
551async 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
572fn 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
587fn 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
600fn 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
610fn 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
632fn 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), f64_field(), f64_field(), string_field(&intr), string_field(&intr), nullable_object_field(&intr), ],
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), nullable_object_field(&intr), f64_field(), nullable_object_field(&intr), f64_field(), string_field(&intr), ],
685 )
686}
687
688fn 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 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
776fn 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 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
831fn build_number_array(
833 caller: &mut wasmtime::Caller<'_, StoreData>,
834 row: &[u8],
835) -> wasmtime::Result<Val> {
836 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 #[derive(Clone, Copy)]
867 enum Outcome {
868 Reported,
870 TransportFailure,
872 Internal,
874 FailedAfter(u64),
876 CountMismatch,
878 }
879
880 struct MockProvider {
883 dimensions: usize,
884 max_input_bytes: u64,
885 outcome: Outcome,
886 embed_calls: AtomicUsize,
887 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 struct Run {
1035 result: wasmtime::Result<String>,
1036 host_attached_bytes: u64,
1037 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 #[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 #[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 #[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 #[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 #[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 #[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 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 #[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 #[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 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 #[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 #[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 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 assert_eq!(budget.used(), 128, "the reported usage is settled");
1481 assert_eq!(budget.held(), 0, "nothing stays held");
1482 }
1483
1484 #[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 #[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 #[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 #[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 #[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 #[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}