moly_kit/protocol/
attachment.rs

1//! This module contains the `Attachment` abstraction to exchange files back and forth
2//! between users and AI.
3//!
4//! See [`Attachment`] for more details.
5
6use crate::utils::asynchronous::BoxPlatformSendFuture;
7#[cfg(target_arch = "wasm32")]
8use crate::utils::asynchronous::ThreadToken;
9#[cfg(target_arch = "wasm32")]
10use std::sync::atomic::{AtomicU64, Ordering};
11
12use std::sync::Arc;
13
14#[cfg(feature = "json")]
15use serde::{Deserialize, Serialize};
16
17/// Private `rfd::FileHandle` wrapper with a runtime generated ID for partial equality.
18#[cfg(target_arch = "wasm32")]
19#[derive(Clone)]
20struct WebFileHandle {
21    id: u64,
22    rfd_handle: rfd::FileHandle,
23}
24
25#[cfg(target_arch = "wasm32")]
26impl From<rfd::FileHandle> for WebFileHandle {
27    fn from(handle: rfd::FileHandle) -> Self {
28        static NEXT_ID: AtomicU64 = AtomicU64::new(0);
29        let id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
30        WebFileHandle {
31            id,
32            rfd_handle: handle,
33        }
34    }
35}
36
37#[cfg(target_arch = "wasm32")]
38impl PartialEq for WebFileHandle {
39    fn eq(&self, other: &Self) -> bool {
40        self.id == other.id
41    }
42}
43
44#[derive(Clone)]
45struct PersistedAttachmentHandle {
46    reader: Arc<
47        dyn Fn(&str) -> BoxPlatformSendFuture<'static, std::io::Result<Arc<[u8]>>> + Send + Sync,
48    >,
49    key: String,
50}
51
52impl std::fmt::Debug for PersistedAttachmentHandle {
53    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54        f.debug_struct("PersistedAttachmentHandle")
55            .field("key", &self.key)
56            .field("reader", &format_args!("{:p}", Arc::as_ptr(&self.reader)))
57            .finish()
58    }
59}
60
61/// Private type that points to wherever the attachment content is stored.
62///
63/// Comparision is done by pointer, file path, file handle, etc. Not by content.
64#[derive(Debug, Clone)]
65enum AttachmentContentHandle {
66    InMemory(Arc<[u8]>),
67    #[cfg(not(target_arch = "wasm32"))]
68    FilePick(std::path::PathBuf),
69    #[cfg(target_arch = "wasm32")]
70    FilePick(ThreadToken<WebFileHandle>),
71    ErasedPersisted(String),
72    Persisted(PersistedAttachmentHandle),
73}
74
75#[cfg(feature = "json")]
76impl serde::Serialize for AttachmentContentHandle {
77    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
78    where
79        S: serde::Serializer,
80    {
81        match self {
82            AttachmentContentHandle::ErasedPersisted(key) => serializer.serialize_str(key),
83            AttachmentContentHandle::Persisted(persisted) => {
84                serializer.serialize_str(&persisted.key)
85            }
86            _ => serializer.serialize_none(),
87        }
88    }
89}
90
91#[cfg(feature = "json")]
92impl<'de> serde::Deserialize<'de> for AttachmentContentHandle {
93    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
94    where
95        D: serde::Deserializer<'de>,
96    {
97        let opt_key: Option<String> = Option::deserialize(deserializer)?;
98        match opt_key {
99            Some(key) => Ok(AttachmentContentHandle::ErasedPersisted(key)),
100            None => Err(serde::de::Error::custom(
101                "AttachmentContentHandle cannot be deserialized from null",
102            )),
103        }
104    }
105}
106
107impl PartialEq for AttachmentContentHandle {
108    fn eq(&self, other: &Self) -> bool {
109        match (self, other) {
110            (AttachmentContentHandle::InMemory(a), AttachmentContentHandle::InMemory(b)) => {
111                Arc::ptr_eq(a, b)
112            }
113            #[cfg(not(target_arch = "wasm32"))]
114            (AttachmentContentHandle::FilePick(a), AttachmentContentHandle::FilePick(b)) => a == b,
115            #[cfg(target_arch = "wasm32")]
116            (AttachmentContentHandle::FilePick(a), AttachmentContentHandle::FilePick(b)) => {
117                let a_id = a.peek(|handle| handle.id);
118                let b_id = b.peek(|handle| handle.id);
119                a_id == b_id
120            }
121            (
122                AttachmentContentHandle::ErasedPersisted(a),
123                AttachmentContentHandle::ErasedPersisted(b),
124            ) => a == b,
125            (AttachmentContentHandle::Persisted(a), AttachmentContentHandle::Persisted(b)) => {
126                a.key == b.key
127            }
128
129            _ => false,
130        }
131    }
132}
133
134impl Eq for AttachmentContentHandle {}
135
136impl std::hash::Hash for AttachmentContentHandle {
137    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
138        match self {
139            AttachmentContentHandle::InMemory(content) => {
140                Arc::as_ptr(content).hash(state);
141            }
142            #[cfg(not(target_arch = "wasm32"))]
143            AttachmentContentHandle::FilePick(path) => path.hash(state),
144            #[cfg(target_arch = "wasm32")]
145            AttachmentContentHandle::FilePick(handle) => handle.peek(|h| h.id).hash(state),
146            AttachmentContentHandle::ErasedPersisted(key) => key.hash(state),
147            AttachmentContentHandle::Persisted(persisted) => persisted.key.hash(state),
148        }
149    }
150}
151
152impl From<&[u8]> for AttachmentContentHandle {
153    fn from(bytes: &[u8]) -> Self {
154        AttachmentContentHandle::InMemory(Arc::from(bytes))
155    }
156}
157
158#[cfg(not(target_arch = "wasm32"))]
159impl From<std::path::PathBuf> for AttachmentContentHandle {
160    fn from(path: std::path::PathBuf) -> Self {
161        AttachmentContentHandle::FilePick(path)
162    }
163}
164
165#[cfg(target_arch = "wasm32")]
166impl From<rfd::FileHandle> for AttachmentContentHandle {
167    fn from(handle: rfd::FileHandle) -> Self {
168        AttachmentContentHandle::FilePick(ThreadToken::new(WebFileHandle::from(handle)))
169    }
170}
171
172impl AttachmentContentHandle {
173    async fn read(&self) -> std::io::Result<Arc<[u8]>> {
174        match self {
175            AttachmentContentHandle::InMemory(content) => Ok(content.clone()),
176            #[cfg(not(target_arch = "wasm32"))]
177            AttachmentContentHandle::FilePick(path) => {
178                let content = tokio::fs::read(path).await?;
179                Ok(Arc::from(content))
180            }
181            #[cfg(target_arch = "wasm32")]
182            AttachmentContentHandle::FilePick(handle) => {
183                let handle = handle.clone_inner();
184                let content = handle.rfd_handle.read().await;
185                Ok(Arc::from(content))
186            }
187            AttachmentContentHandle::ErasedPersisted(_) => Err(std::io::Error::new(
188                std::io::ErrorKind::Other,
189                "Cannot read erased persisted attachment. Please restore the reader with `set_persistence_reader` first.",
190            )),
191            AttachmentContentHandle::Persisted(persisted) => {
192                (persisted.reader)(persisted.key.as_str()).await
193            }
194        }
195    }
196}
197
198/// Represents a file/image/document sent or received as part of a message.
199///
200/// ## Examples
201///
202/// - An attachment sent by the user.
203/// - An image generated by the AI.
204///
205/// ## Equality
206///
207/// When comparing, two [`Attachment`]s are considered equal if they have the same
208/// metadata (name, content type, etc), and they **point** to the same data.
209///
210/// This means:
211/// - For in-memory attachments, the content is compared by reference (pointer equality).
212/// - For attachments picked from a native file system, the path is compared.
213/// - For attachments picked on the web, the (wrapped) file handle must be the same.
214/// - For persisted attachments, the storage key is compared (independent of the reader).
215///
216/// The content itself is never compared, because not all attachments can be read
217/// synchronously, and it would be expensive to do so.
218///
219/// ## Serialization
220///
221/// Unless a persistence key is configured, when serializing this type, the "pointer" to
222/// data is skipped and the attachment will become "unavailable" when deserialized back.
223///
224/// Two unavailable attachments are considered equal if they have the same metadata.
225///
226/// If a persistence key is configured, the attachment will be serialized with the key.
227/// That key will be restored when deserializing, however, you will need to manually
228/// restore its reader implementation.
229///
230/// ## Abstraction details
231///
232/// Different than other abstraction in [`crate::protocol`], this one not only
233/// acts as "data model", but also as system and I/O interface. This coupling was
234/// originally intended to give "pragmatic" access to methods like `read()` and `pick_multiple()`,
235/// but as everything mixing concerns, this now causes some issues like making the persistence
236/// feature uglier to integrate. So this abstraction is likely to change in the future.
237#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
238#[cfg_attr(feature = "json", derive(Serialize, Deserialize))]
239pub struct Attachment {
240    /// Normally the original filename.
241    pub name: String,
242    /// Mime type of the content, if known.
243    pub content_type: Option<String>,
244    content: Option<AttachmentContentHandle>,
245}
246
247// File type filters for the file picker dialog
248const SUPPORTED_IMAGE_EXTENSIONS: &[&str] = &["png", "jpg", "jpeg", "gif", "webp", "bmp", "svg"];
249const SUPPORTED_DOCUMENT_EXTENSIONS: &[&str] = &["pdf", "txt", "md", "html", "htm"];
250const SUPPORTED_ALL_EXTENSIONS: &[&str] = &[
251    "png", "jpg", "jpeg", "gif", "webp", "bmp", "svg", // Images
252    "pdf", // Documents
253    "txt", "md", "html", "htm", "xml", "json", "yaml", "yml", "csv", "log", "ini", "cfg",
254    "conf", // Text files
255    "js", "ts", "tsx", "jsx", "py", "rs", "go", "java", "c", "cpp", "h", "hpp", "cs", "rb",
256    "php", // Code files
257    "swift", "kt", "scala", "r", "m", "sql", "sh", "bash", "zsh", "fish", "ps1", "bat", "cmd",
258    "css", "scss", "sass", "less", "toml", "env", // Other text formats
259];
260
261impl Attachment {
262    /// Crate private utility to pick files from the file system.
263    ///
264    /// - On web, async API is required to pick files.
265    /// - On macos, sync API is required and must be called from the main UI thread.
266    ///   - This is the reason why it takes a closure instead of returning a Future.
267    ///     Because on native `spawn` may run in a separate thread. So we can't generalize.
268    /// - We follow macos requirements on all native platforms just in case.
269    pub(crate) fn pick_multiple(cb: impl FnOnce(Result<Vec<Attachment>, ()>) + 'static) {
270        cfg_if::cfg_if! {
271            if #[cfg(target_arch = "wasm32")] {
272                crate::utils::asynchronous::spawn(async move {
273                    let Some(handles) = rfd::AsyncFileDialog::new()
274                        .add_filter("Supported Files", SUPPORTED_ALL_EXTENSIONS)
275                        .add_filter("Images", SUPPORTED_IMAGE_EXTENSIONS)
276                        .add_filter("Documents", SUPPORTED_DOCUMENT_EXTENSIONS)
277                        .add_filter("All Files", &["*"])
278                        .pick_files()
279                        .await
280                    else {
281                        cb(Err(()));
282                        return;
283                    };
284
285                    let mut attachments = Vec::with_capacity(handles.len());
286                    for handle in handles {
287                        let name = handle.file_name();
288                        let content_type = mime_guess::from_path(&name)
289                            .first()
290                            .map(|m| m.to_string());
291                        attachments.push(Attachment {
292                            name,
293                            content_type,
294                            content: Some(handle.into()),
295                        });
296                    }
297
298                    cb(Ok(attachments));
299                });
300            } else if #[cfg(any(target_os = "windows", target_os = "macos", target_os = "linux"))] {
301                let Some(paths) = rfd::FileDialog::new()
302                    .add_filter("Supported Files", SUPPORTED_ALL_EXTENSIONS)
303                    .add_filter("Images", SUPPORTED_IMAGE_EXTENSIONS)
304                    .add_filter("Documents", SUPPORTED_DOCUMENT_EXTENSIONS)
305                    .add_filter("All Files", &["*"])
306                    .pick_files()
307                else {
308                    cb(Err(()));
309                    return;
310                };
311
312                let mut attachments = Vec::with_capacity(paths.len());
313                for path in paths {
314                    let name = path.file_name().unwrap_or_default().to_string_lossy().to_string();
315                    let content_type = mime_guess::from_path(&name)
316                        .first()
317                        .map(|m| m.to_string());
318
319                    attachments.push(Attachment {
320                        name,
321                        content_type,
322                        content: Some(path.into()),
323                    });
324                }
325                cb(Ok(attachments));
326            } else {
327                ::log::warn!("Attachment picking is not supported on this platform");
328                cb(Err(()));
329            }
330        }
331    }
332
333    /// Creates a new in-memory attachment from the given bytes.
334    pub fn from_bytes(name: String, content_type: Option<String>, content: &[u8]) -> Self {
335        Attachment {
336            name,
337            content_type,
338            content: Some(content.into()),
339        }
340    }
341
342    /// Creates a new in-memory attachment from a base64 encoded string.
343    pub fn from_base64(
344        name: String,
345        content_type: Option<String>,
346        base64_content: &str,
347    ) -> std::io::Result<Self> {
348        use base64::Engine;
349        let content = base64::engine::general_purpose::STANDARD
350            .decode(base64_content)
351            .map_err(|_| {
352                std::io::Error::new(std::io::ErrorKind::InvalidData, "Invalid base64 content")
353            })?;
354
355        Ok(Attachment::from_bytes(name, content_type, &content))
356    }
357
358    pub fn is_available(&self) -> bool {
359        self.content.is_some()
360    }
361
362    pub fn is_image(&self) -> bool {
363        if let Some(content_type) = &self.content_type {
364            content_type.starts_with("image/")
365        } else {
366            false
367        }
368    }
369
370    pub fn is_pdf(&self) -> bool {
371        self.content_type.as_deref() == Some("application/pdf")
372    }
373
374    pub async fn read(&self) -> std::io::Result<Arc<[u8]>> {
375        if let Some(content) = &self.content {
376            content.read().await
377        } else {
378            Err(std::io::Error::new(
379                std::io::ErrorKind::NotFound,
380                "Attachment content not available",
381            ))
382        }
383    }
384
385    pub async fn read_base64(&self) -> std::io::Result<String> {
386        use base64::Engine;
387        let content = self.read().await?;
388        Ok(base64::engine::general_purpose::STANDARD.encode(content))
389    }
390
391    /// Crate private utility to save/download the attachment to the file system.
392    pub(crate) fn save(&self) {
393        ::log::info!("Downloading attachment: {}", self.name);
394
395        if self.content.is_none() {
396            ::log::warn!("Attachment content not available for saving: {}", self.name);
397            return;
398        }
399
400        self.save_impl();
401    }
402
403    #[cfg(target_arch = "wasm32")]
404    fn save_impl(&self) {
405        let self_clone = self.clone();
406        crate::utils::asynchronous::spawn(async move {
407            let Ok(content) = self_clone.content.as_ref().unwrap().read().await else {
408                ::log::warn!(
409                    "Failed to read attachment content for saving: {}",
410                    self_clone.name
411                );
412                return;
413            };
414
415            use crate::utils::platform::{create_scoped_blob_url, trigger_download};
416            create_scoped_blob_url(&content, self_clone.content_type.as_deref(), |url| {
417                trigger_download(url, &self_clone.name);
418            });
419        });
420    }
421
422    #[cfg(any(target_os = "windows", target_os = "macos", target_os = "linux"))]
423    fn save_impl(&self) {
424        let content_handle = self.content.as_ref().unwrap();
425
426        // Although we could read this content asynchronously, we would still need
427        // to open the save dialog synchronously from the main thread, which would
428        // complicate things.
429        let content = match futures::executor::block_on(content_handle.read()) {
430            Ok(content) => content,
431            Err(err) => {
432                ::log::warn!(
433                    "Failed to read attachment content for saving {}: {}",
434                    self.name,
435                    err
436                );
437                return;
438            }
439        };
440
441        crate::utils::platform::trigger_save_as(&content, Some(self.name.as_str()));
442    }
443
444    #[cfg(not(any(
445        target_arch = "wasm32",
446        target_os = "windows",
447        target_os = "macos",
448        target_os = "linux"
449    )))]
450    fn save_impl(&self) {
451        ::log::warn!("Attachment saving is not supported on this platform");
452    }
453
454    /// Get the content type or "application/octet-stream" if not set.
455    pub fn content_type_or_octet_stream(&self) -> &str {
456        self.content_type
457            .as_deref()
458            .unwrap_or("application/octet-stream")
459    }
460
461    /// Get the persistence key if set.
462    pub fn get_persistence_key(&self) -> Option<&str> {
463        match &self.content {
464            Some(AttachmentContentHandle::Persisted(persisted)) => Some(&persisted.key),
465            Some(AttachmentContentHandle::ErasedPersisted(key)) => Some(key),
466            _ => None,
467        }
468    }
469
470    /// Check if the attachment has a persistence key set, indicating it was persisted.
471    pub fn has_persistence_key(&self) -> bool {
472        self.get_persistence_key().is_some()
473    }
474
475    /// Check if this attachment has a reader implementation set.
476    pub fn has_persistence_reader(&self) -> bool {
477        matches!(&self.content, Some(AttachmentContentHandle::Persisted(_)))
478    }
479
480    /// Give this attachment a custom persistence key.
481    ///
482    /// This means you persisted this attachment somewhere and you will take care
483    /// of how it's read.
484    ///
485    /// You should set this key only after you really persisted (wrote) the attachment.
486    ///
487    /// If you call this, you should also call [`Self::set_persistence_reader`]
488    /// to configure how this attachment will be read using this key.
489    pub fn set_persistence_key(&mut self, key: String) {
490        match &self.content {
491            Some(AttachmentContentHandle::Persisted(persisted)) => {
492                self.content = Some(AttachmentContentHandle::Persisted(
493                    PersistedAttachmentHandle {
494                        reader: persisted.reader.clone(),
495                        key,
496                    },
497                ));
498            }
499            Some(AttachmentContentHandle::ErasedPersisted(_)) => {
500                self.content = Some(AttachmentContentHandle::ErasedPersisted(key));
501            }
502            _ => {
503                self.content = Some(AttachmentContentHandle::ErasedPersisted(key));
504            }
505        }
506    }
507
508    /// Gives this attachment a custom implementation to read the persisted content.
509    ///
510    /// Can only be used after setting a persistence key with [`Self::set_persistence_key`].
511    pub fn set_persistence_reader(
512        &mut self,
513        reader: impl Fn(&str) -> BoxPlatformSendFuture<'static, std::io::Result<Arc<[u8]>>>
514        + Send
515        + Sync
516        + 'static,
517    ) {
518        match &self.content {
519            Some(AttachmentContentHandle::Persisted(persisted)) => {
520                self.content = Some(AttachmentContentHandle::Persisted(
521                    PersistedAttachmentHandle {
522                        reader: Arc::new(reader),
523                        key: persisted.key.clone(),
524                    },
525                ));
526            }
527            Some(AttachmentContentHandle::ErasedPersisted(key)) => {
528                self.content = Some(AttachmentContentHandle::Persisted(
529                    PersistedAttachmentHandle {
530                        reader: Arc::new(reader),
531                        key: key.clone(),
532                    },
533                ));
534            }
535            _ => {
536                ::log::warn!("Cannot set persistence reader on a non-persisted attachment");
537            }
538        }
539    }
540}
moly_kit/protocol/attachment.rs

moly_kit/protocol/
attachment.rs
