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
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 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
160pub(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
209async 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 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 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
282struct 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 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 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
395fn 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
418async 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
432fn 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
482async 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 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
512async 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
529async 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
550fn 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
565fn 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
578fn 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
588fn 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
610fn 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), f64_field(), f64_field(), string_field(&intr), string_field(&intr), nullable_object_field(&intr), ],
647 )
648}
649
650fn 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 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
744fn 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 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
799fn build_number_array(
801 caller: &mut wasmtime::Caller<'_, StoreData>,
802 row: &[u8],
803) -> wasmtime::Result<Val> {
804 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 #[derive(Clone, Copy)]
835 enum Outcome {
836 Reported,
838 TransportFailure,
840 Internal,
842 FailedAfter(u64),
844 CountMismatch,
846 }
847
848 struct MockProvider {
851 dimensions: usize,
852 max_input_bytes: u64,
853 outcome: Outcome,
854 embed_calls: AtomicUsize,
855 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 struct Run {
1003 result: wasmtime::Result<String>,
1004 host_attached_bytes: u64,
1005 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 #[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 #[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 #[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 #[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 #[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 #[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 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 #[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 #[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 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 #[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 #[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 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 assert_eq!(budget.used(), 128, "the reported usage is settled");
1449 assert_eq!(budget.held(), 0, "nothing stays held");
1450 }
1451
1452 #[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 #[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 #[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 #[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 #[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 #[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}