Skip to main content

superfluid_daemon/
files.rs

1//! File store for the OpenAI `/v1/files` API (the substrate for `/v1/batches`).
2
3use std::path::{Path, PathBuf};
4use std::sync::atomic::{AtomicU64, Ordering};
5use std::sync::Mutex;
6
7use serde::{Deserialize, Serialize};
8
9use crate::DaemonError;
10
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct FileMeta {
13    pub id: String,
14    #[serde(default = "file_object")]
15    pub object: String,
16    pub bytes: u64,
17    pub created_at: u64,
18    pub filename: String,
19    pub purpose: String,
20    #[serde(default, skip_serializing_if = "Option::is_none")]
21    pub owner: Option<String>,
22}
23
24fn file_object() -> String {
25    "file".to_string()
26}
27
28pub struct FileStore {
29    dir: PathBuf,
30    max_bytes: Option<u64>,
31    expiry_secs: Option<u64>,
32    admit: Mutex<()>,
33}
34
35fn now_secs() -> u64 {
36    std::time::SystemTime::now()
37        .duration_since(std::time::UNIX_EPOCH)
38        .map(|d| d.as_secs())
39        .unwrap_or(0)
40}
41
42impl FileStore {
43    pub fn new(dir: PathBuf) -> FileStore {
44        Self::with_limits(dir, None, None)
45    }
46
47    pub fn with_limits(
48        dir: PathBuf,
49        max_bytes: Option<u64>,
50        expiry_secs: Option<u64>,
51    ) -> FileStore {
52        let _ = std::fs::create_dir_all(&dir);
53        FileStore { dir, max_bytes, expiry_secs, admit: Mutex::new(()) }
54    }
55
56    pub fn stored_bytes(&self) -> u64 {
57        self.list().iter().map(|m| m.bytes).sum()
58    }
59
60    pub fn sweep(&self) -> usize {
61        let Some(max_age) = self.expiry_secs else {
62            return 0;
63        };
64        let now = now_secs();
65        let mut removed = 0;
66        for m in self.list() {
67            if now.saturating_sub(m.created_at) > max_age && self.delete(&m.id) {
68                removed += 1;
69            }
70        }
71        removed
72    }
73
74    fn new_id() -> String {
75        static N: AtomicU64 = AtomicU64::new(0);
76        let nanos = std::time::SystemTime::now()
77            .duration_since(std::time::UNIX_EPOCH)
78            .map(|d| d.as_nanos() as u64)
79            .unwrap_or(0);
80        format!("file-{:x}{:x}", nanos, N.fetch_add(1, Ordering::Relaxed))
81    }
82
83    fn meta_path(&self, id: &str) -> PathBuf {
84        self.dir.join(format!("{id}.json"))
85    }
86    fn blob_path(&self, id: &str) -> PathBuf {
87        self.dir.join(id)
88    }
89
90    fn valid_id(id: &str) -> bool {
91        id.starts_with("file-") && id.len() > 5 && id[5..].chars().all(|c| c.is_ascii_hexdigit())
92    }
93
94    pub fn put(
95        &self,
96        filename: &str,
97        purpose: &str,
98        bytes: &[u8],
99        owner: Option<&str>,
100    ) -> Result<FileMeta, DaemonError> {
101        let _admit = self.max_bytes.map(|_| self.admit.lock().expect("file store"));
102        if let Some(max) = self.max_bytes {
103            self.sweep();
104            let stored = self.stored_bytes();
105            let incoming = bytes.len() as u64;
106            if stored.saturating_add(incoming) > max {
107                return Err(DaemonError::FilesQuota { stored, incoming, max });
108            }
109        }
110        let id = Self::new_id();
111        std::fs::write(self.blob_path(&id), bytes).map_err(DaemonError::Io)?;
112        let meta = FileMeta {
113            id: id.clone(),
114            object: "file".to_string(),
115            bytes: bytes.len() as u64,
116            created_at: now_secs(),
117            filename: filename.to_string(),
118            purpose: purpose.to_string(),
119            owner: owner.map(str::to_string),
120        };
121        let json = serde_json::to_vec(&meta).map_err(|_| DaemonError::Protocol("file meta encode"))?;
122        std::fs::write(self.meta_path(&id), json).map_err(DaemonError::Io)?;
123        Ok(meta)
124    }
125
126    pub fn get(&self, id: &str) -> Option<FileMeta> {
127        if !Self::valid_id(id) {
128            return None;
129        }
130        let bytes = std::fs::read(self.meta_path(id)).ok()?;
131        serde_json::from_slice(&bytes).ok()
132    }
133
134    pub fn content(&self, id: &str) -> Option<Vec<u8>> {
135        if !Self::valid_id(id) {
136            return None;
137        }
138        std::fs::read(self.blob_path(id)).ok()
139    }
140
141    pub fn list(&self) -> Vec<FileMeta> {
142        let mut out = Vec::new();
143        if let Ok(rd) = std::fs::read_dir(&self.dir) {
144            for e in rd.flatten() {
145                let p = e.path();
146                if p.extension().and_then(|x| x.to_str()) == Some("json") {
147                    if let Ok(b) = std::fs::read(&p) {
148                        if let Ok(m) = serde_json::from_slice::<FileMeta>(&b) {
149                            out.push(m);
150                        }
151                    }
152                }
153            }
154        }
155        out.sort_by_key(|m| std::cmp::Reverse(m.created_at));
156        out
157    }
158
159    pub fn delete(&self, id: &str) -> bool {
160        if !Self::valid_id(id) || self.get(id).is_none() {
161            return false;
162        }
163        let _ = std::fs::remove_file(self.blob_path(id));
164        let _ = std::fs::remove_file(self.meta_path(id));
165        true
166    }
167}
168
169pub fn files_dir_for(sessions_dir: &Path) -> PathBuf {
170    sessions_dir.join("files")
171}