Skip to main content

orchestral_runtime/
in_memory_blob.rs

1//! Minimal in-memory BlobStore for embedded composition and deterministic tests.
2
3use std::collections::HashMap;
4
5use async_trait::async_trait;
6use bytes::Bytes;
7use futures_util::StreamExt;
8use tokio::sync::RwLock;
9
10use orchestral_core::agent_protocol::wire::Digest;
11use orchestral_core::io::{
12    BlobHead, BlobId, BlobIoError, BlobMeta, BlobRead, BlobStore, BlobWriteRequest,
13};
14
15#[derive(Clone)]
16struct InMemoryBlobObject {
17    meta: BlobMeta,
18    bytes: Vec<u8>,
19}
20
21#[derive(Default)]
22pub struct InMemoryBlobStore {
23    objects: RwLock<HashMap<String, InMemoryBlobObject>>,
24}
25
26#[async_trait]
27impl BlobStore for InMemoryBlobStore {
28    async fn write(&self, mut request: BlobWriteRequest) -> Result<BlobMeta, BlobIoError> {
29        let blob_id = BlobId::from(uuid::Uuid::new_v4().to_string());
30        let mut data = Vec::new();
31        while let Some(chunk) = request.body.next().await {
32            data.extend_from_slice(&chunk?);
33        }
34        if data.is_empty() {
35            return Err(BlobIoError::Invalid("empty blob payload".to_owned()));
36        }
37        let now = chrono::Utc::now();
38        let meta = BlobMeta {
39            id: blob_id,
40            file_name: request.file_name.take(),
41            mime_type: request.mime_type.take(),
42            byte_size: data.len() as u64,
43            checksum_sha256: Some(Digest::sha256(&data).to_string()),
44            metadata: if request.metadata.is_null() {
45                serde_json::json!({})
46            } else {
47                request.metadata
48            },
49            created_at: now,
50            updated_at: now,
51        };
52        self.objects.write().await.insert(
53            meta.id.to_string(),
54            InMemoryBlobObject {
55                meta: meta.clone(),
56                bytes: data,
57            },
58        );
59        Ok(meta)
60    }
61
62    async fn read(&self, blob_id: &BlobId) -> Result<BlobRead, BlobIoError> {
63        let object = self
64            .objects
65            .read()
66            .await
67            .get(blob_id.as_str())
68            .cloned()
69            .ok_or_else(|| BlobIoError::NotFound(blob_id.to_string()))?;
70        let body = Box::pin(futures_util::stream::once(async move {
71            Ok(Bytes::from(object.bytes))
72        }));
73        Ok(BlobRead {
74            meta: object.meta,
75            body,
76        })
77    }
78
79    async fn head(&self, blob_id: &BlobId) -> Result<BlobHead, BlobIoError> {
80        let object = self
81            .objects
82            .read()
83            .await
84            .get(blob_id.as_str())
85            .cloned()
86            .ok_or_else(|| BlobIoError::NotFound(blob_id.to_string()))?;
87        Ok(BlobHead {
88            byte_size: object.meta.byte_size,
89            etag: None,
90            last_modified: Some(object.meta.updated_at),
91        })
92    }
93
94    async fn delete(&self, blob_id: &BlobId) -> Result<bool, BlobIoError> {
95        Ok(self
96            .objects
97            .write()
98            .await
99            .remove(blob_id.as_str())
100            .is_some())
101    }
102}