Skip to main content

runifold_model/
artifact.rs

1use std::{
2    collections::BTreeMap,
3    future::Future,
4    pin::Pin,
5    sync::{Arc, Mutex},
6    time::{SystemTime, UNIX_EPOCH},
7};
8
9use base64::{Engine as _, engine::general_purpose::STANDARD};
10use serde::{Deserialize, Deserializer, Serialize};
11use sha2::{Digest, Sha256};
12use thiserror::Error;
13
14use crate::{
15    ContentPart, MediaSource, Model, ModelCallContext, ModelCapabilities, ModelError,
16    ModelErrorKind, ModelEventStream, ModelFuture, ModelRef, ModelRequest, ProviderModel,
17};
18
19/// Maximum artifact size accepted by the reference store and resolver.
20pub const DEFAULT_MAX_ARTIFACT_BYTES: usize = 16 * 1024 * 1024;
21/// Maximum number of artifact references returned or deleted per operation.
22pub const MAX_ARTIFACT_PAGE_SIZE: u32 = 1_000;
23/// Maximum UTF-8 byte length of an artifact idempotency key.
24pub const MAX_ARTIFACT_IDEMPOTENCY_KEY_BYTES: usize = 512;
25/// Maximum UTF-8 byte length of an artifact display name.
26pub const MAX_ARTIFACT_NAME_BYTES: usize = 256;
27
28/// Validated tenant or application isolation boundary for artifacts.
29#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd, Serialize)]
30#[serde(transparent)]
31pub struct ArtifactScope(String);
32
33impl ArtifactScope {
34    /// Parses a bounded, storage-safe scope.
35    ///
36    /// # Errors
37    ///
38    /// Returns [`ArtifactError::InvalidInput`] when the value is empty, too
39    /// long, or contains characters that are unsafe for storage boundaries.
40    pub fn parse(value: impl Into<String>) -> Result<Self, ArtifactError> {
41        let value = value.into();
42        if value.is_empty()
43            || value.len() > 128
44            || !value.bytes().all(|byte| {
45                byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'-' | b':')
46            })
47        {
48            return Err(ArtifactError::InvalidInput(
49                "artifact scope must contain 1..=128 safe ASCII characters".into(),
50            ));
51        }
52        Ok(Self(value))
53    }
54
55    /// Returns the validated scope text.
56    pub fn as_str(&self) -> &str {
57        &self.0
58    }
59}
60
61impl<'de> Deserialize<'de> for ArtifactScope {
62    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
63    where
64        D: Deserializer<'de>,
65    {
66        let value = String::deserialize(deserializer)?;
67        Self::parse(value).map_err(serde::de::Error::custom)
68    }
69}
70
71/// A stable, integrity-bound reference to binary artifact content.
72#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
73pub struct ArtifactRef {
74    /// Tenant or application isolation boundary.
75    pub scope: ArtifactScope,
76    /// Store-owned stable identity.
77    pub artifact_id: String,
78    /// Trusted MIME type associated with the stored bytes.
79    pub media_type: String,
80    /// Raw content length.
81    pub size_bytes: u64,
82    /// Lowercase SHA-256 digest.
83    pub sha256: String,
84    /// Optional neutral display name.
85    pub name: Option<String>,
86    /// Store-assigned creation time in Unix milliseconds.
87    pub created_at_unix_ms: u64,
88    /// Optional expiration time in Unix milliseconds.
89    pub expires_at_unix_ms: Option<u64>,
90}
91
92impl ArtifactRef {
93    /// Converts this reference into canonical lazy media.
94    pub fn media_source(&self) -> MediaSource {
95        MediaSource::Artifact {
96            reference: self.clone(),
97        }
98    }
99}
100
101/// One bounded page of artifact references in stable identity order.
102#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
103pub struct ArtifactPage {
104    /// References visible in the requested scope.
105    pub items: Vec<ArtifactRef>,
106    /// Exclusive cursor for the next page.
107    pub next_cursor: Option<String>,
108}
109
110/// Bytes loaded from an artifact store together with verified metadata.
111#[derive(Clone, Debug, Eq, PartialEq)]
112pub struct Artifact {
113    /// Integrity-bound reference.
114    pub reference: ArtifactRef,
115    /// Raw bytes.
116    pub bytes: Vec<u8>,
117}
118
119/// Idempotent artifact creation request.
120#[derive(Clone, Debug)]
121pub struct ArtifactWrite {
122    /// Tenant or application isolation boundary.
123    scope: ArtifactScope,
124    /// Stable key reused across retries of one logical write.
125    idempotency_key: String,
126    /// Trusted MIME type.
127    media_type: String,
128    /// Optional neutral display name.
129    name: Option<String>,
130    /// Raw bytes.
131    bytes: Vec<u8>,
132    /// Optional retention deadline.
133    expires_at_unix_ms: Option<u64>,
134}
135
136impl ArtifactWrite {
137    /// Creates a validated write request.
138    ///
139    /// # Errors
140    ///
141    /// Returns [`ArtifactError`] for blank identities, invalid MIME values, or
142    /// empty content.
143    pub fn new(
144        scope: ArtifactScope,
145        idempotency_key: impl Into<String>,
146        media_type: impl Into<String>,
147        bytes: Vec<u8>,
148    ) -> Result<Self, ArtifactError> {
149        let idempotency_key = idempotency_key.into();
150        let media_type = media_type.into();
151        if idempotency_key.trim().is_empty()
152            || idempotency_key.len() > MAX_ARTIFACT_IDEMPOTENCY_KEY_BYTES
153            || idempotency_key.chars().any(char::is_control)
154        {
155            return Err(ArtifactError::InvalidInput(format!(
156                "artifact idempotency key must contain 1..={MAX_ARTIFACT_IDEMPOTENCY_KEY_BYTES} non-control UTF-8 bytes"
157            )));
158        }
159        if !valid_media_type(&media_type) {
160            return Err(ArtifactError::InvalidInput(
161                "artifact MIME type is invalid".into(),
162            ));
163        }
164        if bytes.is_empty() {
165            return Err(ArtifactError::InvalidInput(
166                "artifact content cannot be empty".into(),
167            ));
168        }
169        if !content_matches_media_type(&media_type, &bytes) {
170            return Err(ArtifactError::InvalidInput(
171                "artifact bytes do not match the declared MIME type".into(),
172            ));
173        }
174        Ok(Self {
175            scope,
176            idempotency_key,
177            media_type,
178            name: None,
179            bytes,
180            expires_at_unix_ms: None,
181        })
182    }
183
184    /// Sets an absolute retention deadline portable across durable stores.
185    ///
186    /// # Errors
187    ///
188    /// Returns [`ArtifactError::InvalidInput`] when the timestamp exceeds the
189    /// signed 64-bit millisecond representation shared by `SQLite` and
190    /// `PostgreSQL`.
191    pub fn with_expires_at_unix_ms(
192        mut self,
193        expires_at_unix_ms: u64,
194    ) -> Result<Self, ArtifactError> {
195        if expires_at_unix_ms > i64::MAX as u64 {
196            return Err(ArtifactError::InvalidInput(
197                "artifact expiration exceeds the portable i64 millisecond range".into(),
198            ));
199        }
200        self.expires_at_unix_ms = Some(expires_at_unix_ms);
201        Ok(self)
202    }
203
204    /// Returns the isolation scope.
205    pub const fn scope(&self) -> &ArtifactScope {
206        &self.scope
207    }
208
209    /// Sets a bounded neutral display name.
210    ///
211    /// # Errors
212    ///
213    /// Returns [`ArtifactError::InvalidInput`] for a blank, oversized, or
214    /// control-character-containing name.
215    pub fn with_name(mut self, name: impl Into<String>) -> Result<Self, ArtifactError> {
216        let name = name.into();
217        if name.trim().is_empty()
218            || name.len() > MAX_ARTIFACT_NAME_BYTES
219            || name.chars().any(char::is_control)
220        {
221            return Err(ArtifactError::InvalidInput(format!(
222                "artifact name must contain 1..={MAX_ARTIFACT_NAME_BYTES} non-control UTF-8 bytes"
223            )));
224        }
225        self.name = Some(name);
226        Ok(self)
227    }
228
229    /// Returns the stable idempotency key.
230    pub fn idempotency_key(&self) -> &str {
231        &self.idempotency_key
232    }
233
234    /// Returns the trusted MIME type.
235    pub fn media_type(&self) -> &str {
236        &self.media_type
237    }
238
239    /// Returns the optional display name.
240    pub fn name(&self) -> Option<&str> {
241        self.name.as_deref()
242    }
243
244    /// Returns the raw artifact bytes.
245    pub fn bytes(&self) -> &[u8] {
246        &self.bytes
247    }
248
249    /// Returns the optional retention deadline.
250    pub const fn expires_at_unix_ms(&self) -> Option<u64> {
251        self.expires_at_unix_ms
252    }
253
254    /// Returns whether a stored reference represents this exact immutable
255    /// write, excluding only the store-assigned creation time.
256    pub fn matches_immutable_reference(&self, reference: &ArtifactRef) -> bool {
257        reference.scope == self.scope
258            && reference.artifact_id == artifact_identity(&self.media_type, &self.bytes)
259            && reference.media_type == self.media_type
260            && reference.size_bytes == self.bytes.len() as u64
261            && reference.sha256 == sha256(&self.bytes)
262            && reference.name == self.name
263            && reference.expires_at_unix_ms == self.expires_at_unix_ms
264    }
265}
266
267/// Typed artifact persistence or integrity failure.
268#[derive(Clone, Debug, Error, Eq, PartialEq)]
269#[non_exhaustive]
270pub enum ArtifactError {
271    /// Caller supplied an invalid request.
272    #[error("invalid artifact input: {0}")]
273    InvalidInput(String),
274    /// Artifact identity was not found.
275    #[error("artifact `{0}` was not found")]
276    NotFound(String),
277    /// One idempotency key was reused for different content.
278    #[error("artifact idempotency conflict for `{0}`")]
279    IdempotencyConflict(String),
280    /// The same content address already exists with different immutable metadata.
281    #[error("artifact metadata conflict for `{0}`")]
282    MetadataConflict(String),
283    /// Stored content no longer matches its integrity metadata.
284    #[error("artifact integrity check failed for `{0}`")]
285    Integrity(String),
286    /// Artifact exists but is past its retention deadline.
287    #[error("artifact `{0}` has expired")]
288    Expired(String),
289    /// Backend storage failed.
290    #[error("artifact storage failed: {0}")]
291    Storage(String),
292}
293
294/// A boxed artifact-store future.
295#[cfg(not(target_arch = "wasm32"))]
296pub type ArtifactFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
297
298/// A boxed artifact-store future for single-threaded WASM.
299#[cfg(target_arch = "wasm32")]
300pub type ArtifactFuture<'a, T> = Pin<Box<dyn Future<Output = T> + 'a>>;
301
302/// Binary artifact persistence boundary.
303pub trait ArtifactStore: Send + Sync {
304    /// Idempotently writes content and returns its integrity-bound reference.
305    fn put(&self, write: ArtifactWrite) -> ArtifactFuture<'_, Result<ArtifactRef, ArtifactError>>;
306
307    /// Loads and verifies one artifact.
308    fn get(&self, reference: &ArtifactRef) -> ArtifactFuture<'_, Result<Artifact, ArtifactError>>;
309
310    /// Lists one stable, bounded page inside one scope.
311    fn list(
312        &self,
313        scope: &ArtifactScope,
314        after: Option<&str>,
315        limit: u32,
316    ) -> ArtifactFuture<'_, Result<ArtifactPage, ArtifactError>>;
317
318    /// Idempotently deletes one artifact and its idempotency records.
319    fn delete(
320        &self,
321        scope: &ArtifactScope,
322        artifact_id: &str,
323    ) -> ArtifactFuture<'_, Result<bool, ArtifactError>>;
324
325    /// Deletes at most `limit` expired artifacts in one scope.
326    fn purge_expired(
327        &self,
328        scope: &ArtifactScope,
329        now_unix_ms: u64,
330        limit: u32,
331    ) -> ArtifactFuture<'_, Result<u32, ArtifactError>>;
332}
333
334impl<T> ArtifactStore for Arc<T>
335where
336    T: ArtifactStore + ?Sized,
337{
338    fn put(&self, write: ArtifactWrite) -> ArtifactFuture<'_, Result<ArtifactRef, ArtifactError>> {
339        (**self).put(write)
340    }
341
342    fn get(&self, reference: &ArtifactRef) -> ArtifactFuture<'_, Result<Artifact, ArtifactError>> {
343        (**self).get(reference)
344    }
345
346    fn list(
347        &self,
348        scope: &ArtifactScope,
349        after: Option<&str>,
350        limit: u32,
351    ) -> ArtifactFuture<'_, Result<ArtifactPage, ArtifactError>> {
352        (**self).list(scope, after, limit)
353    }
354
355    fn delete(
356        &self,
357        scope: &ArtifactScope,
358        artifact_id: &str,
359    ) -> ArtifactFuture<'_, Result<bool, ArtifactError>> {
360        (**self).delete(scope, artifact_id)
361    }
362
363    fn purge_expired(
364        &self,
365        scope: &ArtifactScope,
366        now_unix_ms: u64,
367        limit: u32,
368    ) -> ArtifactFuture<'_, Result<u32, ArtifactError>> {
369        (**self).purge_expired(scope, now_unix_ms, limit)
370    }
371}
372
373/// Cloneable bounded in-memory artifact store for tests and local processes.
374#[derive(Clone, Debug)]
375pub struct InMemoryArtifactStore {
376    state: Arc<Mutex<MemoryState>>,
377    max_artifact_bytes: usize,
378}
379
380#[derive(Debug, Default)]
381struct MemoryState {
382    artifacts: BTreeMap<(ArtifactScope, String), Artifact>,
383    idempotency: BTreeMap<(ArtifactScope, String), String>,
384}
385
386impl InMemoryArtifactStore {
387    /// Creates a store with the default 16 MiB per-artifact limit.
388    pub fn new() -> Self {
389        Self {
390            state: Arc::new(Mutex::new(MemoryState::default())),
391            max_artifact_bytes: DEFAULT_MAX_ARTIFACT_BYTES,
392        }
393    }
394
395    /// Replaces the per-artifact byte limit.
396    #[must_use]
397    pub const fn with_max_artifact_bytes(mut self, limit: usize) -> Self {
398        self.max_artifact_bytes = limit;
399        self
400    }
401}
402
403impl Default for InMemoryArtifactStore {
404    fn default() -> Self {
405        Self::new()
406    }
407}
408
409impl ArtifactStore for InMemoryArtifactStore {
410    fn put(&self, write: ArtifactWrite) -> ArtifactFuture<'_, Result<ArtifactRef, ArtifactError>> {
411        Box::pin(async move {
412            if write.bytes.len() > self.max_artifact_bytes {
413                return Err(ArtifactError::InvalidInput(format!(
414                    "artifact is {} bytes and exceeds the {}-byte limit",
415                    write.bytes.len(),
416                    self.max_artifact_bytes
417                )));
418            }
419            let digest = sha256(&write.bytes);
420            let artifact_id = artifact_identity(&write.media_type, &write.bytes);
421            let scope = write.scope.clone();
422            let artifact_key = (scope.clone(), artifact_id.clone());
423            let idempotency_key = (scope.clone(), write.idempotency_key.clone());
424            let mut state = self
425                .state
426                .lock()
427                .map_err(|_| ArtifactError::Storage("artifact store lock is poisoned".into()))?;
428            if let Some(existing_id) = state.idempotency.get(&idempotency_key) {
429                if existing_id != &artifact_id {
430                    return Err(ArtifactError::IdempotencyConflict(write.idempotency_key));
431                }
432                let reference = state
433                    .artifacts
434                    .get(&(scope, existing_id.clone()))
435                    .map(|artifact| artifact.reference.clone())
436                    .ok_or_else(|| ArtifactError::Integrity(existing_id.clone()))?;
437                if !write.matches_immutable_reference(&reference) {
438                    return Err(ArtifactError::IdempotencyConflict(write.idempotency_key));
439                }
440                return Ok(reference);
441            }
442            if let Some(existing) = state.artifacts.get(&artifact_key) {
443                if !write.matches_immutable_reference(&existing.reference) {
444                    return Err(ArtifactError::MetadataConflict(artifact_id));
445                }
446                let reference = existing.reference.clone();
447                state.idempotency.insert(idempotency_key, artifact_id);
448                return Ok(reference);
449            }
450            let reference = ArtifactRef {
451                scope,
452                artifact_id: artifact_id.clone(),
453                media_type: write.media_type,
454                size_bytes: write.bytes.len() as u64,
455                sha256: digest,
456                name: write.name,
457                created_at_unix_ms: unix_time_ms()?,
458                expires_at_unix_ms: write.expires_at_unix_ms,
459            };
460            state.artifacts.insert(
461                artifact_key,
462                Artifact {
463                    reference: reference.clone(),
464                    bytes: write.bytes,
465                },
466            );
467            state.idempotency.insert(idempotency_key, artifact_id);
468            Ok(reference)
469        })
470    }
471
472    fn get(&self, reference: &ArtifactRef) -> ArtifactFuture<'_, Result<Artifact, ArtifactError>> {
473        let reference = reference.clone();
474        Box::pin(async move {
475            let artifact = self
476                .state
477                .lock()
478                .map_err(|_| ArtifactError::Storage("artifact store lock is poisoned".into()))?
479                .artifacts
480                .get(&(reference.scope.clone(), reference.artifact_id.clone()))
481                .cloned()
482                .ok_or_else(|| ArtifactError::NotFound(reference.artifact_id.clone()))?;
483            verify_artifact(&artifact)?;
484            verify_reference(&reference, &artifact.reference)?;
485            ensure_not_expired(&artifact.reference, unix_time_ms()?)?;
486            Ok(artifact)
487        })
488    }
489
490    fn list(
491        &self,
492        scope: &ArtifactScope,
493        after: Option<&str>,
494        limit: u32,
495    ) -> ArtifactFuture<'_, Result<ArtifactPage, ArtifactError>> {
496        let scope = scope.clone();
497        let after = after.map(str::to_owned);
498        Box::pin(async move {
499            validate_page_limit(limit)?;
500            let state = self
501                .state
502                .lock()
503                .map_err(|_| ArtifactError::Storage("artifact store lock is poisoned".into()))?;
504            let mut items = state
505                .artifacts
506                .iter()
507                .filter(|((item_scope, id), _)| {
508                    item_scope == &scope && after.as_ref().is_none_or(|cursor| id > cursor)
509                })
510                .map(|(_, artifact)| artifact.reference.clone())
511                .take(limit as usize + 1)
512                .collect::<Vec<_>>();
513            let next_cursor = if items.len() > limit as usize {
514                items.pop();
515                items.last().map(|item| item.artifact_id.clone())
516            } else {
517                None
518            };
519            Ok(ArtifactPage { items, next_cursor })
520        })
521    }
522
523    fn delete(
524        &self,
525        scope: &ArtifactScope,
526        artifact_id: &str,
527    ) -> ArtifactFuture<'_, Result<bool, ArtifactError>> {
528        let scope = scope.clone();
529        let artifact_id = artifact_id.to_owned();
530        Box::pin(async move {
531            let mut state = self
532                .state
533                .lock()
534                .map_err(|_| ArtifactError::Storage("artifact store lock is poisoned".into()))?;
535            let removed = state
536                .artifacts
537                .remove(&(scope.clone(), artifact_id.clone()))
538                .is_some();
539            if removed {
540                state
541                    .idempotency
542                    .retain(|(item_scope, _), id| item_scope != &scope || id != &artifact_id);
543            }
544            Ok(removed)
545        })
546    }
547
548    fn purge_expired(
549        &self,
550        scope: &ArtifactScope,
551        now_unix_ms: u64,
552        limit: u32,
553    ) -> ArtifactFuture<'_, Result<u32, ArtifactError>> {
554        let scope = scope.clone();
555        Box::pin(async move {
556            validate_page_limit(limit)?;
557            let mut state = self
558                .state
559                .lock()
560                .map_err(|_| ArtifactError::Storage("artifact store lock is poisoned".into()))?;
561            let ids = state
562                .artifacts
563                .iter()
564                .filter(|((item_scope, _), artifact)| {
565                    item_scope == &scope
566                        && artifact
567                            .reference
568                            .expires_at_unix_ms
569                            .is_some_and(|expires| expires <= now_unix_ms)
570                })
571                .map(|((_, id), _)| id.clone())
572                .take(limit as usize)
573                .collect::<Vec<_>>();
574            for id in &ids {
575                state.artifacts.remove(&(scope.clone(), id.clone()));
576                state.idempotency.retain(|(item_scope, _), artifact_id| {
577                    item_scope != &scope || artifact_id != id
578                });
579            }
580            u32::try_from(ids.len())
581                .map_err(|_| ArtifactError::Storage("purge count exceeds u32".into()))
582        })
583    }
584}
585
586/// Model decorator that resolves lazy artifact references immediately before
587/// provider transport while keeping transcripts and checkpoints reference-only.
588#[derive(Clone, Debug)]
589pub struct ArtifactResolvingModel<M, S> {
590    inner: M,
591    scope: ArtifactScope,
592    store: S,
593    max_artifact_bytes: usize,
594}
595
596impl<M, S> ArtifactResolvingModel<M, S> {
597    /// Wraps a model with an artifact store.
598    pub fn new(inner: M, scope: ArtifactScope, store: S) -> Self {
599        Self {
600            inner,
601            scope,
602            store,
603            max_artifact_bytes: DEFAULT_MAX_ARTIFACT_BYTES,
604        }
605    }
606
607    /// Replaces the maximum resolved artifact size.
608    #[must_use]
609    pub const fn with_max_artifact_bytes(mut self, limit: usize) -> Self {
610        self.max_artifact_bytes = limit;
611        self
612    }
613}
614
615impl<M, S> Model for ArtifactResolvingModel<M, S>
616where
617    M: Model,
618    S: ArtifactStore,
619{
620    fn capabilities<'a>(
621        &'a self,
622        model: &'a ModelRef,
623    ) -> ModelFuture<'a, Result<ModelCapabilities, ModelError>> {
624        self.inner.capabilities(model)
625    }
626
627    fn stream(
628        &self,
629        mut request: ModelRequest,
630        context: ModelCallContext,
631    ) -> ModelFuture<'_, Result<ModelEventStream, ModelError>> {
632        Box::pin(async move {
633            resolve_request_artifacts(
634                &mut request,
635                &self.scope,
636                &self.store,
637                self.max_artifact_bytes,
638            )
639            .await?;
640            self.inner.stream(request, context).await
641        })
642    }
643}
644
645impl<M, S> ProviderModel for ArtifactResolvingModel<M, S>
646where
647    M: ProviderModel,
648    S: ArtifactStore,
649{
650    fn provider(&self) -> &str {
651        self.inner.provider()
652    }
653}
654
655async fn resolve_request_artifacts<S: ArtifactStore>(
656    request: &mut ModelRequest,
657    scope: &ArtifactScope,
658    store: &S,
659    max_artifact_bytes: usize,
660) -> Result<(), ModelError> {
661    for message in &mut request.messages {
662        for part in &mut message.content {
663            resolve_content_artifacts(part, scope, store, max_artifact_bytes).await?;
664        }
665    }
666    Ok(())
667}
668
669async fn resolve_content_artifacts<S: ArtifactStore>(
670    part: &mut ContentPart,
671    scope: &ArtifactScope,
672    store: &S,
673    max_artifact_bytes: usize,
674) -> Result<(), ModelError> {
675    match part {
676        ContentPart::Image { source }
677        | ContentPart::Audio { source }
678        | ContentPart::Document { source, .. } => {
679            resolve_source(source, scope, store, max_artifact_bytes).await
680        }
681        ContentPart::ToolResult(result) => {
682            for content in &mut result.content {
683                Box::pin(resolve_content_artifacts(
684                    content,
685                    scope,
686                    store,
687                    max_artifact_bytes,
688                ))
689                .await?;
690            }
691            Ok(())
692        }
693        _ => Ok(()),
694    }
695}
696
697async fn resolve_source<S: ArtifactStore>(
698    source: &mut MediaSource,
699    scope: &ArtifactScope,
700    store: &S,
701    max_artifact_bytes: usize,
702) -> Result<(), ModelError> {
703    let MediaSource::Artifact { reference } = source else {
704        return Ok(());
705    };
706    if &reference.scope != scope {
707        return Err(ModelError::local(
708            ModelErrorKind::InvalidRequest,
709            "artifact reference belongs to a different scope",
710        ));
711    }
712    let artifact = store
713        .get(reference)
714        .await
715        .map_err(|error| artifact_model_error(&error))?;
716    if artifact.bytes.len() > max_artifact_bytes {
717        return Err(ModelError::local(
718            ModelErrorKind::InvalidRequest,
719            format!(
720                "artifact `{}` is {} bytes and exceeds the {max_artifact_bytes}-byte resolution limit",
721                reference.artifact_id,
722                artifact.bytes.len()
723            ),
724        ));
725    }
726    *source = MediaSource::Base64 {
727        media_type: artifact.reference.media_type,
728        data: STANDARD.encode(artifact.bytes),
729    };
730    Ok(())
731}
732
733fn verify_artifact(artifact: &Artifact) -> Result<(), ArtifactError> {
734    if artifact.reference.size_bytes != artifact.bytes.len() as u64
735        || artifact.reference.sha256 != sha256(&artifact.bytes)
736        || artifact.reference.artifact_id
737            != artifact_identity(&artifact.reference.media_type, &artifact.bytes)
738    {
739        return Err(ArtifactError::Integrity(
740            artifact.reference.artifact_id.clone(),
741        ));
742    }
743    Ok(())
744}
745
746fn verify_reference(expected: &ArtifactRef, actual: &ArtifactRef) -> Result<(), ArtifactError> {
747    if expected != actual {
748        return Err(ArtifactError::Integrity(expected.artifact_id.clone()));
749    }
750    Ok(())
751}
752
753fn ensure_not_expired(reference: &ArtifactRef, now_unix_ms: u64) -> Result<(), ArtifactError> {
754    if reference
755        .expires_at_unix_ms
756        .is_some_and(|expires| expires <= now_unix_ms)
757    {
758        return Err(ArtifactError::Expired(reference.artifact_id.clone()));
759    }
760    Ok(())
761}
762
763fn validate_page_limit(limit: u32) -> Result<(), ArtifactError> {
764    if limit == 0 || limit > MAX_ARTIFACT_PAGE_SIZE {
765        return Err(ArtifactError::InvalidInput(format!(
766            "artifact page limit must be between 1 and {MAX_ARTIFACT_PAGE_SIZE}"
767        )));
768    }
769    Ok(())
770}
771
772fn unix_time_ms() -> Result<u64, ArtifactError> {
773    let elapsed = SystemTime::now()
774        .duration_since(UNIX_EPOCH)
775        .map_err(|error| ArtifactError::Storage(error.to_string()))?;
776    u64::try_from(elapsed.as_millis())
777        .map_err(|_| ArtifactError::Storage("system time exceeds u64 milliseconds".into()))
778}
779
780fn sha256(bytes: &[u8]) -> String {
781    hex_digest(Sha256::digest(bytes))
782}
783
784fn artifact_identity(media_type: &str, bytes: &[u8]) -> String {
785    let mut digest = Sha256::new();
786    digest.update(media_type.as_bytes());
787    digest.update([0]);
788    digest.update(bytes);
789    format!("sha256:{}", hex_digest(digest.finalize()))
790}
791
792fn hex_digest(digest: impl AsRef<[u8]>) -> String {
793    const HEX: &[u8; 16] = b"0123456789abcdef";
794    let bytes = digest.as_ref();
795    let mut output = String::with_capacity(bytes.len().saturating_mul(2));
796    for byte in bytes {
797        output.push(char::from(HEX[usize::from(byte >> 4)]));
798        output.push(char::from(HEX[usize::from(byte & 0x0f)]));
799    }
800    output
801}
802
803fn valid_media_type(value: &str) -> bool {
804    let Some((kind, subtype)) = value.split_once('/') else {
805        return false;
806    };
807    valid_media_token(kind) && valid_media_token(subtype)
808}
809
810fn valid_media_token(value: &str) -> bool {
811    !value.is_empty()
812        && value.bytes().all(|byte| {
813            byte.is_ascii_alphanumeric()
814                || matches!(
815                    byte,
816                    b'!' | b'#'
817                        | b'$'
818                        | b'%'
819                        | b'&'
820                        | b'\''
821                        | b'*'
822                        | b'+'
823                        | b'-'
824                        | b'.'
825                        | b'^'
826                        | b'_'
827                        | b'`'
828                        | b'|'
829                        | b'~'
830                )
831        })
832}
833
834fn content_matches_media_type(media_type: &str, bytes: &[u8]) -> bool {
835    match media_type {
836        "image/png" => bytes.starts_with(b"\x89PNG\r\n\x1a\n"),
837        "image/jpeg" => bytes.starts_with(&[0xff, 0xd8, 0xff]),
838        "image/gif" => bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a"),
839        "image/webp" => bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP",
840        "audio/wav" => bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WAVE",
841        "audio/mpeg" => {
842            bytes.starts_with(b"ID3")
843                || bytes
844                    .get(..2)
845                    .is_some_and(|prefix| prefix[0] == 0xff && prefix[1] & 0xe0 == 0xe0)
846        }
847        "audio/ogg" => bytes.starts_with(b"OggS"),
848        "application/pdf" => bytes.starts_with(b"%PDF-"),
849        value if value.starts_with("text/") => std::str::from_utf8(bytes).is_ok(),
850        _ => true,
851    }
852}
853
854fn artifact_model_error(error: &ArtifactError) -> ModelError {
855    let kind = match error {
856        ArtifactError::NotFound(_)
857        | ArtifactError::InvalidInput(_)
858        | ArtifactError::Expired(_)
859        | ArtifactError::IdempotencyConflict(_)
860        | ArtifactError::MetadataConflict(_) => ModelErrorKind::InvalidRequest,
861        ArtifactError::Integrity(_) | ArtifactError::Storage(_) => ModelErrorKind::Protocol,
862    };
863    ModelError::local(kind, error.to_string())
864}
865
866#[cfg(test)]
867mod tests {
868    use std::collections::BTreeMap;
869
870    use futures_executor::block_on;
871
872    use super::*;
873
874    const PNG: &[u8] = b"\x89PNG\r\n\x1a\npng";
875
876    fn scope() -> ArtifactScope {
877        ArtifactScope::parse("tenant.test").unwrap()
878    }
879
880    #[test]
881    fn in_memory_store_is_content_addressed_idempotent_and_integrity_bound() {
882        let store = InMemoryArtifactStore::new();
883        let first =
884            block_on(store.put(
885                ArtifactWrite::new(scope(), "turn-1:image", "image/png", PNG.to_vec()).unwrap(),
886            ))
887            .unwrap();
888        let replay =
889            block_on(store.put(
890                ArtifactWrite::new(scope(), "turn-1:image", "image/png", PNG.to_vec()).unwrap(),
891            ))
892            .unwrap();
893        assert_eq!(first, replay);
894        assert_eq!(block_on(store.get(&first)).unwrap().bytes, PNG);
895
896        let conflict = block_on(
897            store.put(
898                ArtifactWrite::new(
899                    scope(),
900                    "turn-1:image",
901                    "image/png",
902                    b"\x89PNG\r\n\x1a\nother".to_vec(),
903                )
904                .unwrap(),
905            ),
906        )
907        .unwrap_err();
908        assert!(matches!(conflict, ArtifactError::IdempotencyConflict(_)));
909    }
910
911    #[test]
912    fn request_resolution_keeps_storage_references_outside_provider_transport() {
913        let store = InMemoryArtifactStore::new();
914        let reference =
915            block_on(store.put(
916                ArtifactWrite::new(scope(), "turn-2:image", "image/png", PNG.to_vec()).unwrap(),
917            ))
918            .unwrap();
919        let result = crate::ToolResult {
920            call_id: "call-1".into(),
921            name: Some("render".into()),
922            content: vec![ContentPart::Image {
923                source: reference.media_source(),
924            }],
925            structured_content: None,
926            is_error: false,
927            metadata: BTreeMap::new(),
928        };
929        let message =
930            crate::Message::new(crate::Role::Tool, vec![ContentPart::ToolResult(result)]).unwrap();
931        let mut request = ModelRequest::new(ModelRef::new("test", "vision"), message);
932
933        block_on(resolve_request_artifacts(
934            &mut request,
935            &scope(),
936            &store,
937            DEFAULT_MAX_ARTIFACT_BYTES,
938        ))
939        .unwrap();
940
941        let ContentPart::ToolResult(result) = &request.messages[0].content[0] else {
942            panic!("tool result must remain canonical");
943        };
944        assert!(matches!(
945            &result.content[0],
946            ContentPart::Image {
947                source: MediaSource::Base64 { media_type, data }
948            } if media_type == "image/png" && data == &STANDARD.encode(PNG)
949        ));
950    }
951
952    #[test]
953    fn rejects_known_media_with_mismatched_magic_bytes() {
954        let error = ArtifactWrite::new(scope(), "turn-3:image", "image/png", b"not-png".to_vec())
955            .unwrap_err();
956
957        assert!(matches!(error, ArtifactError::InvalidInput(_)));
958    }
959
960    #[test]
961    fn deserialization_and_write_metadata_cannot_bypass_validation() {
962        assert!(serde_json::from_str::<ArtifactScope>("\"../tenant\"").is_err());
963        assert!(
964            ArtifactWrite::new(
965                scope(),
966                "x".repeat(MAX_ARTIFACT_IDEMPOTENCY_KEY_BYTES + 1),
967                "text/plain",
968                b"text".to_vec(),
969            )
970            .is_err()
971        );
972        assert!(
973            ArtifactWrite::new(scope(), "key", "text/plain", b"text".to_vec())
974                .unwrap()
975                .with_name("bad\nname")
976                .is_err()
977        );
978        assert!(ArtifactWrite::new(scope(), "key", "image//png", b"text".to_vec()).is_err());
979        assert!(
980            ArtifactWrite::new(scope(), "key", "text/plain", b"text".to_vec())
981                .unwrap()
982                .with_expires_at_unix_ms(u64::MAX)
983                .is_err()
984        );
985    }
986
987    #[test]
988    fn immutable_metadata_is_bound_to_content_and_idempotency() {
989        let store = InMemoryArtifactStore::new();
990        let original = ArtifactWrite::new(scope(), "first", "text/plain", b"same".to_vec())
991            .unwrap()
992            .with_name("original")
993            .unwrap();
994        let reference = block_on(store.put(original)).unwrap();
995
996        let replay_with_changed_expiry =
997            ArtifactWrite::new(scope(), "first", "text/plain", b"same".to_vec())
998                .unwrap()
999                .with_name("original")
1000                .unwrap()
1001                .with_expires_at_unix_ms(i64::MAX as u64)
1002                .unwrap();
1003        assert!(matches!(
1004            block_on(store.put(replay_with_changed_expiry)),
1005            Err(ArtifactError::IdempotencyConflict(_))
1006        ));
1007
1008        let alias_with_changed_name =
1009            ArtifactWrite::new(scope(), "second", "text/plain", b"same".to_vec())
1010                .unwrap()
1011                .with_name("changed")
1012                .unwrap();
1013        assert!(matches!(
1014            block_on(store.put(alias_with_changed_name)),
1015            Err(ArtifactError::MetadataConflict(_))
1016        ));
1017        assert_eq!(
1018            block_on(store.get(&reference))
1019                .unwrap()
1020                .reference
1021                .name
1022                .as_deref(),
1023            Some("original")
1024        );
1025    }
1026
1027    #[test]
1028    fn scopes_pagination_expiration_and_deletion_are_enforced() {
1029        let store = InMemoryArtifactStore::new();
1030        let left = scope();
1031        let right = ArtifactScope::parse("tenant.other").unwrap();
1032        let expired = block_on(
1033            store.put(
1034                ArtifactWrite::new(left.clone(), "expired", "image/png", PNG.to_vec())
1035                    .unwrap()
1036                    .with_expires_at_unix_ms(1)
1037                    .unwrap(),
1038            ),
1039        )
1040        .unwrap();
1041        let active = block_on(store.put(
1042            ArtifactWrite::new(left.clone(), "active", "text/plain", b"active".to_vec()).unwrap(),
1043        ))
1044        .unwrap();
1045        let isolated = block_on(store.put(
1046            ArtifactWrite::new(right.clone(), "active", "text/plain", b"other".to_vec()).unwrap(),
1047        ))
1048        .unwrap();
1049
1050        assert!(matches!(
1051            block_on(store.get(&expired)),
1052            Err(ArtifactError::Expired(_))
1053        ));
1054        assert_eq!(block_on(store.list(&left, None, 1)).unwrap().items.len(), 1);
1055        assert_eq!(
1056            block_on(store.list(&right, None, 10)).unwrap().items,
1057            [isolated]
1058        );
1059        assert_eq!(
1060            block_on(store.purge_expired(&left, u64::MAX, 10)).unwrap(),
1061            1
1062        );
1063        assert!(block_on(store.delete(&left, &active.artifact_id)).unwrap());
1064        assert!(!block_on(store.delete(&left, &active.artifact_id)).unwrap());
1065    }
1066}