orchestral_runtime/
in_memory_blob.rs1use 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}