1use roder_api::media::{
2 MediaArtifact, MediaArtifactId, MediaDimensions, MediaGenerationMetadata, MediaKind,
3 MediaPreview, MediaPreviewStrategy,
4};
5use std::collections::hash_map::DefaultHasher;
6use std::hash::{Hash, Hasher};
7use std::path::{Path, PathBuf};
8use time::OffsetDateTime;
9
10#[derive(Debug, Clone)]
12pub struct GeneratedMediaSpec<'a> {
13 pub prompt: &'a str,
14 pub kind: MediaKind,
15 pub mime_type: &'a str,
16 pub provider: &'a str,
17 pub bytes: &'a [u8],
18 pub dimensions: Option<MediaDimensions>,
19 pub duration_millis: Option<u64>,
20 pub generation: Option<MediaGenerationMetadata>,
21}
22
23#[derive(Debug, Clone)]
24pub struct MediaArtifactStore {
25 root: PathBuf,
26 max_read_bytes: u64,
27}
28
29impl MediaArtifactStore {
30 pub fn new(root: impl Into<PathBuf>) -> Self {
31 Self {
32 root: root.into(),
33 max_read_bytes: 10 * 1024 * 1024,
34 }
35 }
36
37 pub fn with_max_read_bytes(mut self, max_read_bytes: u64) -> Self {
38 self.max_read_bytes = max_read_bytes;
39 self
40 }
41
42 pub fn root(&self) -> &Path {
43 &self.root
44 }
45
46 pub fn write_generated(
47 &self,
48 spec: &GeneratedMediaSpec<'_>,
49 ) -> anyhow::Result<(MediaArtifact, MediaPreview)> {
50 std::fs::create_dir_all(&self.root)?;
51 let ext = extension_for_mime(spec.mime_type);
52 let id = format!("media-{}", uuid::Uuid::new_v4());
53 let store_path = self.root.join(format!("{id}.{ext}"));
54 std::fs::write(&store_path, spec.bytes)?;
55 let artifact = MediaArtifact {
56 id: id.clone(),
57 kind: spec.kind.clone(),
58 mime_type: spec.mime_type.to_string(),
59 dimensions: spec.dimensions.clone(),
60 duration_millis: spec.duration_millis,
61 byte_size: spec.bytes.len() as u64,
62 provider: spec.provider.to_string(),
63 prompt_hash: prompt_hash(spec.prompt),
64 store_path: store_path.display().to_string(),
65 thumbnail_path: Some(store_path.display().to_string()),
66 generation: spec.generation.clone(),
67 created_at: OffsetDateTime::now_utc(),
68 roder_owned: true,
69 };
70 self.write_metadata(&artifact)?;
71 let preview = preview_for(&artifact);
72 Ok((artifact, preview))
73 }
74
75 pub fn list(&self) -> anyhow::Result<Vec<MediaArtifact>> {
76 let mut artifacts = Vec::new();
77 if !self.root.exists() {
78 return Ok(artifacts);
79 }
80 for entry in std::fs::read_dir(&self.root)? {
81 let path = entry?.path();
82 if path.extension().and_then(|ext| ext.to_str()) != Some("json") {
83 continue;
84 }
85 let text = std::fs::read_to_string(path)?;
86 artifacts.push(serde_json::from_str(&text)?);
87 }
88 artifacts.sort_by(|left: &MediaArtifact, right| left.id.cmp(&right.id));
89 Ok(artifacts)
90 }
91
92 pub fn read(
93 &self,
94 artifact_id: &MediaArtifactId,
95 max_bytes: Option<u64>,
96 ) -> anyhow::Result<(MediaArtifact, Vec<u8>)> {
97 let artifact = self.get(artifact_id)?;
98 let limit = max_bytes.unwrap_or(self.max_read_bytes);
99 if artifact.byte_size > limit {
100 anyhow::bail!(
101 "media artifact {} is {} bytes, over read limit {}",
102 artifact.id,
103 artifact.byte_size,
104 limit
105 );
106 }
107 let bytes = std::fs::read(&artifact.store_path)?;
108 Ok((artifact, bytes))
109 }
110
111 pub fn get(&self, artifact_id: &MediaArtifactId) -> anyhow::Result<MediaArtifact> {
112 let path = self.metadata_path(artifact_id);
113 let text = std::fs::read_to_string(&path)?;
114 Ok(serde_json::from_str(&text)?)
115 }
116
117 pub fn delete(&self, artifact_id: &MediaArtifactId) -> anyhow::Result<bool> {
118 let artifact = self.get(artifact_id)?;
119 if !artifact.roder_owned {
120 anyhow::bail!(
121 "refusing to delete non-Roder-owned artifact {}",
122 artifact.id
123 );
124 }
125 let mut deleted = false;
126 let data_path = PathBuf::from(&artifact.store_path);
127 if data_path.starts_with(&self.root) && data_path.exists() {
128 std::fs::remove_file(data_path)?;
129 deleted = true;
130 }
131 let metadata = self.metadata_path(artifact_id);
132 if metadata.exists() {
133 std::fs::remove_file(metadata)?;
134 deleted = true;
135 }
136 Ok(deleted)
137 }
138
139 pub fn preview(&self, artifact_id: &MediaArtifactId) -> anyhow::Result<MediaPreview> {
140 Ok(preview_for(&self.get(artifact_id)?))
141 }
142
143 fn write_metadata(&self, artifact: &MediaArtifact) -> anyhow::Result<()> {
144 let path = self.metadata_path(&artifact.id);
145 std::fs::write(path, serde_json::to_string_pretty(artifact)?)?;
146 Ok(())
147 }
148
149 fn metadata_path(&self, artifact_id: &MediaArtifactId) -> PathBuf {
150 self.root.join(format!("{artifact_id}.json"))
151 }
152}
153
154pub fn default_media_artifact_dir() -> anyhow::Result<PathBuf> {
155 let data_dir = std::env::var_os("RODER_DATA_DIR")
156 .or_else(|| std::env::var_os("RODER_CONFIG_DIR"))
157 .map(PathBuf::from)
158 .or_else(|| std::env::var_os("HOME").map(|home| PathBuf::from(home).join(".roder")))
159 .ok_or_else(|| {
160 anyhow::anyhow!("could not resolve Roder data directory for media artifacts")
161 })?;
162 Ok(data_dir.join("artifacts"))
163}
164
165fn preview_for(artifact: &MediaArtifact) -> MediaPreview {
166 let strategy = match artifact.kind {
167 MediaKind::Image => MediaPreviewStrategy::Thumbnail,
168 _ => MediaPreviewStrategy::MetadataOnly,
169 };
170 MediaPreview {
171 artifact_id: artifact.id.clone(),
172 strategy,
173 thumbnail_path: artifact.thumbnail_path.clone(),
174 fallback_label: format!(
175 "{} {} ({} bytes)",
176 artifact.provider, artifact.mime_type, artifact.byte_size
177 ),
178 warning: None,
179 }
180}
181
182fn prompt_hash(prompt: &str) -> String {
183 let mut hasher = DefaultHasher::new();
184 prompt.hash(&mut hasher);
185 format!("{:016x}", hasher.finish())
186}
187
188fn extension_for_mime(mime_type: &str) -> &'static str {
189 match mime_type {
190 "image/png" => "png",
191 "image/jpeg" => "jpg",
192 "image/gif" => "gif",
193 "video/mp4" => "mp4",
194 _ => "bin",
195 }
196}
197
198#[cfg(test)]
199mod tests {
200 use super::*;
201
202 #[test]
203 fn media_artifact_store_writes_reads_previews_and_deletes_owned_artifacts() {
204 let root =
205 std::env::temp_dir().join(format!("roder-media-artifacts-{}", uuid::Uuid::new_v4()));
206 let store = MediaArtifactStore::new(&root).with_max_read_bytes(1024);
207
208 let (artifact, preview) = store
209 .write_generated(&GeneratedMediaSpec {
210 prompt: "tiny image",
211 kind: MediaKind::Image,
212 mime_type: "image/png",
213 provider: "fake",
214 bytes: b"abc",
215 dimensions: Some(MediaDimensions {
216 width: 1,
217 height: 1,
218 }),
219 duration_millis: None,
220 generation: None,
221 })
222 .unwrap();
223
224 assert!(
225 artifact
226 .store_path
227 .starts_with(root.to_string_lossy().as_ref())
228 );
229 assert_eq!(preview.artifact_id, artifact.id);
230 assert_eq!(store.list().unwrap().len(), 1);
231 assert_eq!(store.read(&artifact.id, None).unwrap().1, b"abc");
232 assert!(store.delete(&artifact.id).unwrap());
233 assert!(store.list().unwrap().is_empty());
234 }
235
236 #[test]
237 fn media_artifact_store_persists_image_generation_metadata() {
238 let root = std::env::temp_dir().join(format!(
239 "roder-media-artifacts-meta-{}",
240 uuid::Uuid::new_v4()
241 ));
242 let store = MediaArtifactStore::new(&root);
243
244 let (artifact, _) = store
245 .write_generated(&GeneratedMediaSpec {
246 prompt: "watermarked",
247 kind: MediaKind::Image,
248 mime_type: "image/png",
249 provider: "google",
250 bytes: b"abc",
251 dimensions: None,
252 duration_millis: None,
253 generation: Some(MediaGenerationMetadata {
254 provider: "google".to_string(),
255 model: Some("gemini-3.1-flash-image".to_string()),
256 revised_prompt: None,
257 watermark: Some("synthid".to_string()),
258 safety: None,
259 provider_response_id: None,
260 }),
261 })
262 .unwrap();
263
264 let reloaded = store.get(&artifact.id).unwrap();
265 let generation = reloaded.generation.expect("generation metadata persists");
266 assert_eq!(generation.model.as_deref(), Some("gemini-3.1-flash-image"));
267 assert_eq!(generation.watermark.as_deref(), Some("synthid"));
268
269 store.delete(&artifact.id).unwrap();
270 }
271}