zeroclaw_runtime/rpc/
attachments.rs

1//! File attachment processing for the RPC transport.
2//!
3//! Handles base64-encoded uploads and local-path reads, SHA-256
4//! content-addressed deduplication, workspace storage, and marker
5//! generation. Used by both `file/attach` and `session/prompt`
6//! (inline attachments).
7
8use super::session::SessionStore;
9// FileSource is only referenced from the `#[cfg(test)] mod tests` below,
10// which re-imports via `use super::*;`. Quiet the non-test "unused" warning
11// without splitting the import into two cfg-gated lines.
12#[cfg_attr(not(test), allow(unused_imports))]
13use super::types::{FileEntry, FileEntryResult, FileSource};
14use zeroclaw_api::jsonrpc::JsonRpcError;
15use zeroclaw_api::jsonrpc::error_codes::*;
16
17/// Per-file size limit (decoded bytes).
18pub const MAX_FILE_BYTES: u64 = 10 * 1024 * 1024;
19
20/// Per-request total size limit (decoded bytes).
21pub const MAX_REQUEST_BYTES: u64 = 20 * 1024 * 1024;
22
23fn rpc_err(code: i32, msg: impl Into<String>) -> JsonRpcError {
24    JsonRpcError {
25        code,
26        message: msg.into(),
27        data: None,
28    }
29}
30
31/// Process a single [`FileEntry`] — resolve bytes, dedup, write to the
32/// upload root, and return a [`FileEntryResult`].
33///
34/// `upload_root` is the directory under which a `uploads/` subdir is
35/// created and bytes are written. Callers should pass the per-agent
36/// workspace dir, NOT the session cwd — uploads belong to the agent,
37/// not to whatever directory the user happened to launch the TUI from.
38pub async fn process_file_entry(
39    entry: &FileEntry,
40    session_id: &str,
41    upload_root: &str,
42    is_wss: bool,
43    sessions: &SessionStore,
44) -> Result<FileEntryResult, JsonRpcError> {
45    use base64::{Engine, engine::general_purpose::STANDARD};
46    use sha2::{Digest, Sha256};
47
48    // 1. Resolve bytes + filename + mime_type.
49    let (bytes, filename, mime_type, original_path) = if let Some(ref b64) = entry.data_b64 {
50        let decoded = STANDARD
51            .decode(b64)
52            .map_err(|e| rpc_err(INVALID_PARAMS, format!("Invalid base64: {e}")))?;
53        if decoded.len() as u64 > MAX_FILE_BYTES {
54            return Err(rpc_err(
55                INVALID_PARAMS,
56                format!(
57                    "File exceeds {} MB limit ({} bytes)",
58                    MAX_FILE_BYTES / (1024 * 1024),
59                    decoded.len()
60                ),
61            ));
62        }
63        let fname = entry.filename.as_deref().unwrap_or("upload").to_string();
64        let mime = entry
65            .mime_type
66            .clone()
67            .unwrap_or_else(|| mime_from_filename(&fname));
68        (decoded, fname, mime, None)
69    } else if let Some(ref path) = entry.path {
70        if is_wss {
71            return Err(rpc_err(
72                INVALID_PARAMS,
73                "Path mode is not available over WSS; send data_b64 instead",
74            ));
75        }
76        let p = std::path::Path::new(path);
77        if !p.is_absolute() {
78            return Err(rpc_err(INVALID_PARAMS, "Path must be absolute"));
79        }
80        let bytes = tokio::fs::read(p)
81            .await
82            .map_err(|e| rpc_err(INVALID_PARAMS, format!("Cannot read file: {e}")))?;
83        if bytes.len() as u64 > MAX_FILE_BYTES {
84            return Err(rpc_err(
85                INVALID_PARAMS,
86                format!(
87                    "File exceeds {} MB limit ({} bytes)",
88                    MAX_FILE_BYTES / (1024 * 1024),
89                    bytes.len()
90                ),
91            ));
92        }
93        let fname = p
94            .file_name()
95            .map(|n| n.to_string_lossy().to_string())
96            .unwrap_or_else(|| "upload".to_string());
97        let mime = entry
98            .mime_type
99            .clone()
100            .unwrap_or_else(|| mime_from_filename(&fname));
101        (bytes, fname, mime, Some(path.clone()))
102    } else {
103        return Err(rpc_err(
104            INVALID_PARAMS,
105            "Each file entry must have either `data_b64` or `path`",
106        ));
107    };
108
109    // 2. SHA-256 → ref_id.
110    let hash = Sha256::digest(&bytes);
111    let hex = format!("{hash:x}");
112    let ref_id = format!("sha256:{hex}");
113
114    // 3. Dedup check.
115    if let Some(existing) = sessions.get_upload(session_id, &ref_id).await {
116        return Ok(FileEntryResult {
117            ref_id: existing.ref_id,
118            marker: existing.marker,
119            workspace_path: existing.workspace_path,
120            size_bytes: existing.size_bytes,
121            deduplicated: true,
122        });
123    }
124
125    // 4. Sanitize filename.
126    let sanitized = sanitize_filename(&filename);
127
128    // 5. Determine extension + write to workspace.
129    let ext = std::path::Path::new(&sanitized)
130        .extension()
131        .map(|e| e.to_string_lossy().to_string())
132        .unwrap_or_default();
133    let storage_name = if ext.is_empty() {
134        hex[..16].to_string()
135    } else {
136        format!("{}.{ext}", &hex[..16])
137    };
138    let upload_dir = std::path::Path::new(upload_root).join("uploads");
139    tokio::fs::create_dir_all(&upload_dir)
140        .await
141        .map_err(|e| rpc_err(INTERNAL_ERROR, format!("Cannot create upload dir: {e}")))?;
142    let dest = upload_dir.join(&storage_name);
143    tokio::fs::write(&dest, &bytes)
144        .await
145        .map_err(|e| rpc_err(INTERNAL_ERROR, format!("Cannot write upload: {e}")))?;
146
147    // Canonicalize so the marker always contains an absolute path —
148    // upload_root may be relative (e.g. ".") when no path was provided.
149    let canonical = tokio::fs::canonicalize(&dest)
150        .await
151        .unwrap_or_else(|_| dest.clone());
152    let workspace_path = canonical.to_string_lossy().to_string();
153
154    // 6. Build marker.
155    //
156    // Images use `[IMAGE:path]` so the multimodal processor can inline them
157    // as data URIs for vision models. Non-image files use a prose format
158    // matching the channel attachment style (`[Document: name] path`) so the
159    // LLM sees a readable path it can access with file-reading tools.
160    //
161    // Regardless of source (file pick vs clipboard paste) and regardless of
162    // transport (Unix path vs WSS base64), the canonical workspace path is
163    // ALWAYS a valid local file that the multimodal pipeline can load — the
164    // bytes were just written above. Emitting `[IMAGE:<workspace_path>]` for
165    // every source ensures vision models receive the actual image data.
166    //
167    // (A previous implementation emitted `[IMAGE from clipboard]` for the
168    // Clipboard source. That marker had no path, so the multimodal loader
169    // silently produced no inline image part and the model received text
170    // only — observed as the agent hallucinating about prior screenshots.)
171    //
172    // The `display_path` preference is the user's original path only for
173    // stable file picks (Unix transport, non-clipboard). Clipboard pastes
174    // use a /tmp path that the TUI deletes after the turn completes, so
175    // on the next turn the multimodal pipeline would find the file gone
176    // and emit a WARN. Always use the workspace /uploads/ copy for clipboard.
177    let kind = attachment_kind(&mime_type);
178    let is_clipboard = matches!(entry.source, FileSource::Clipboard);
179    let display_path = if is_clipboard {
180        &workspace_path
181    } else {
182        original_path.as_deref().unwrap_or(&workspace_path)
183    };
184    let marker = if kind == "IMAGE" {
185        format!("[IMAGE:{display_path}]")
186    } else {
187        // Non-image: prose format with workspace path so the agent can
188        // read the file with its tools regardless of transport.
189        format!("[Document: {filename}] {workspace_path}")
190    };
191
192    let size_bytes = bytes.len() as u64;
193
194    // 7. Index in session upload map.
195    sessions
196        .insert_upload(
197            session_id,
198            super::session::UploadEntry {
199                ref_id: ref_id.clone(),
200                marker: marker.clone(),
201                workspace_path: workspace_path.clone(),
202                size_bytes,
203            },
204        )
205        .await;
206
207    Ok(FileEntryResult {
208        ref_id,
209        marker,
210        workspace_path,
211        size_bytes,
212        deduplicated: false,
213    })
214}
215
216/// Sanitize a filename: strip path separators and null bytes.
217fn sanitize_filename(name: &str) -> String {
218    name.replace(['/', '\\', '\0'], "_")
219}
220
221/// Derive MIME type from filename extension via `mime_guess`.
222/// Falls back to `application/octet-stream` for unknown extensions.
223fn mime_from_filename(name: &str) -> String {
224    mime_guess::from_path(name)
225        .first_or_octet_stream()
226        .to_string()
227}
228
229/// Map MIME type to attachment kind for markers.
230fn attachment_kind(mime: &str) -> &'static str {
231    if mime.starts_with("image/") {
232        "IMAGE"
233    } else if mime == "application/pdf" {
234        "DOCUMENT"
235    } else {
236        "FILE"
237    }
238}
239
240#[cfg(test)]
241mod tests {
242    use super::*;
243    use serde_json::json;
244
245    #[test]
246    fn mime_from_filename_common_types() {
247        assert_eq!(mime_from_filename("photo.png"), "image/png");
248        assert_eq!(mime_from_filename("photo.jpg"), "image/jpeg");
249        assert_eq!(mime_from_filename("doc.pdf"), "application/pdf");
250        assert_eq!(mime_from_filename("data.csv"), "text/csv");
251        assert_eq!(
252            mime_from_filename("unknown.zzzzz"),
253            "application/octet-stream"
254        );
255        assert_eq!(mime_from_filename("noext"), "application/octet-stream");
256    }
257
258    #[test]
259    fn attachment_kind_maps_correctly() {
260        assert_eq!(attachment_kind("image/png"), "IMAGE");
261        assert_eq!(attachment_kind("image/jpeg"), "IMAGE");
262        assert_eq!(attachment_kind("image/svg+xml"), "IMAGE");
263        assert_eq!(attachment_kind("application/pdf"), "DOCUMENT");
264        assert_eq!(attachment_kind("application/zip"), "FILE");
265        assert_eq!(attachment_kind("text/plain"), "FILE");
266    }
267
268    #[test]
269    fn sanitize_filename_strips_separators() {
270        assert_eq!(sanitize_filename("normal.txt"), "normal.txt");
271        assert_eq!(sanitize_filename("path/to/file.txt"), "path_to_file.txt");
272        assert_eq!(sanitize_filename("back\\slash.txt"), "back_slash.txt");
273        assert_eq!(sanitize_filename("null\0byte.txt"), "null_byte.txt");
274    }
275
276    #[test]
277    fn file_source_default_is_file() {
278        let source: FileSource = Default::default();
279        assert!(matches!(source, FileSource::File));
280    }
281
282    #[test]
283    fn file_entry_deserialize_data_mode() {
284        let v = json!({
285            "filename": "screenshot.png",
286            "mime_type": "image/png",
287            "data_b64": "aGVsbG8="
288        });
289        let entry: FileEntry = serde_json::from_value(v).unwrap();
290        assert_eq!(entry.filename.as_deref(), Some("screenshot.png"));
291        assert_eq!(entry.data_b64.as_deref(), Some("aGVsbG8="));
292        assert!(entry.path.is_none());
293        assert!(matches!(entry.source, FileSource::File));
294    }
295
296    #[test]
297    fn file_entry_deserialize_path_mode() {
298        let v = json!({
299            "path": "/home/user/doc.pdf",
300            "source": "file"
301        });
302        let entry: FileEntry = serde_json::from_value(v).unwrap();
303        assert_eq!(entry.path.as_deref(), Some("/home/user/doc.pdf"));
304        assert!(entry.data_b64.is_none());
305    }
306
307    #[test]
308    fn file_entry_deserialize_clipboard_source() {
309        let v = json!({
310            "filename": "paste.png",
311            "mime_type": "image/png",
312            "data_b64": "aGVsbG8=",
313            "source": "clipboard"
314        });
315        let entry: FileEntry = serde_json::from_value(v).unwrap();
316        assert!(matches!(entry.source, FileSource::Clipboard));
317    }
318
319    // ── Integration tests against process_file_entry ─────────────
320
321    fn make_session_store(max: usize) -> SessionStore {
322        SessionStore::new(
323            max,
324            std::sync::Arc::new(zeroclaw_infra::session_queue::SessionActorQueue::new(
325                4, 10, 60,
326            )),
327        )
328    }
329
330    fn make_test_agent() -> crate::agent::agent::Agent {
331        use crate::agent::dispatcher::NativeToolDispatcher;
332
333        let mem_cfg = zeroclaw_config::schema::MemoryConfig {
334            backend: "none".into(),
335            ..zeroclaw_config::schema::MemoryConfig::default()
336        };
337        let mem = std::sync::Arc::from(
338            zeroclaw_memory::create_memory(&mem_cfg, &std::env::temp_dir(), None).unwrap(),
339        );
340
341        crate::agent::agent::Agent::builder()
342            .model_provider(Box::new(StubProvider))
343            .tools(vec![])
344            .memory(mem)
345            .observer(std::sync::Arc::new(crate::observability::NoopObserver {})
346                as std::sync::Arc<dyn crate::observability::Observer>)
347            .tool_dispatcher(Box::new(NativeToolDispatcher))
348            .workspace_dir(std::env::temp_dir())
349            .build()
350            .unwrap()
351    }
352
353    struct StubProvider;
354
355    #[async_trait::async_trait]
356    impl zeroclaw_providers::ModelProvider for StubProvider {
357        async fn chat_with_system(
358            &self,
359            _: Option<&str>,
360            _: &str,
361            _: &str,
362            _: Option<f64>,
363        ) -> anyhow::Result<String> {
364            Ok(String::new())
365        }
366        async fn chat(
367            &self,
368            _: zeroclaw_providers::ChatRequest<'_>,
369            _: &str,
370            _: Option<f64>,
371        ) -> anyhow::Result<zeroclaw_providers::ChatResponse> {
372            Ok(zeroclaw_providers::ChatResponse {
373                text: Some("stub".into()),
374                tool_calls: vec![],
375                usage: None,
376                reasoning_content: None,
377            })
378        }
379    }
380    impl zeroclaw_api::attribution::Attributable for StubProvider {
381        fn role(&self) -> zeroclaw_api::attribution::Role {
382            zeroclaw_api::attribution::Role::Provider(
383                zeroclaw_api::attribution::ProviderKind::Model(
384                    zeroclaw_api::attribution::ModelProviderKind::Custom,
385                ),
386            )
387        }
388        fn alias(&self) -> &str {
389            "stub"
390        }
391    }
392
393    async fn setup_store(workspace: &str) -> SessionStore {
394        let store = make_session_store(4);
395        store
396            .insert(
397                "s1".into(),
398                super::super::session::RpcSession::new(
399                    make_test_agent(),
400                    "a",
401                    workspace,
402                    crate::rpc::types::ChatMode::Chat,
403                ),
404            )
405            .await
406            .unwrap();
407        store
408    }
409
410    #[tokio::test]
411    async fn clipboard_image() {
412        use base64::{Engine, engine::general_purpose::STANDARD};
413
414        let tmp = tempfile::tempdir().unwrap();
415        let ws = tmp.path().to_string_lossy().to_string();
416        let store = setup_store(&ws).await;
417
418        let png_bytes = b"fake-png-data";
419        let entry = FileEntry {
420            path: None,
421            data_b64: Some(STANDARD.encode(png_bytes)),
422            filename: Some("screenshot.png".into()),
423            mime_type: Some("image/png".into()),
424            source: FileSource::Clipboard,
425        };
426
427        let r = process_file_entry(&entry, "s1", &ws, false, &store)
428            .await
429            .unwrap();
430
431        assert!(r.ref_id.starts_with("sha256:"));
432        // Clipboard images: marker must contain the workspace path so the
433        // multimodal pipeline can load and inline the image bytes. The
434        // previous `[IMAGE from clipboard]` marker had no path and silently
435        // produced text-only requests (model never saw the image).
436        assert!(
437            r.marker.starts_with("[IMAGE:") && r.marker.ends_with(']'),
438            "marker = {}",
439            r.marker
440        );
441        assert!(
442            r.marker.contains("/uploads/"),
443            "clipboard image marker should reference workspace uploads path: {}",
444            r.marker
445        );
446        assert!(!r.deduplicated);
447        assert_eq!(r.size_bytes, png_bytes.len() as u64);
448        assert!(std::path::Path::new(&r.workspace_path).exists());
449    }
450
451    #[tokio::test]
452    async fn file_pdf() {
453        use base64::{Engine, engine::general_purpose::STANDARD};
454
455        let tmp = tempfile::tempdir().unwrap();
456        let ws = tmp.path().to_string_lossy().to_string();
457        let store = setup_store(&ws).await;
458
459        let entry = FileEntry {
460            path: None,
461            data_b64: Some(STANDARD.encode(b"%PDF-1.4 fake")),
462            filename: Some("report.pdf".into()),
463            mime_type: Some("application/pdf".into()),
464            source: FileSource::File,
465        };
466
467        let r = process_file_entry(&entry, "s1", &ws, false, &store)
468            .await
469            .unwrap();
470
471        // data_b64 mode: non-image uses prose format with workspace path.
472        assert!(
473            r.marker.starts_with("[Document: report.pdf]"),
474            "marker = {}",
475            r.marker
476        );
477        assert!(
478            r.marker.contains("/uploads/"),
479            "marker should include workspace uploads path: {}",
480            r.marker
481        );
482        assert!(!r.deduplicated);
483    }
484
485    #[tokio::test]
486    async fn deduplication() {
487        use base64::{Engine, engine::general_purpose::STANDARD};
488
489        let tmp = tempfile::tempdir().unwrap();
490        let ws = tmp.path().to_string_lossy().to_string();
491        let store = setup_store(&ws).await;
492
493        let b64 = STANDARD.encode(b"identical-bytes");
494
495        let entry = FileEntry {
496            path: None,
497            data_b64: Some(b64.clone()),
498            filename: Some("img.png".into()),
499            mime_type: Some("image/png".into()),
500            source: FileSource::Clipboard,
501        };
502
503        let r1 = process_file_entry(&entry, "s1", &ws, false, &store)
504            .await
505            .unwrap();
506        assert!(!r1.deduplicated);
507
508        let entry2 = FileEntry {
509            path: None,
510            data_b64: Some(b64),
511            filename: Some("img2.png".into()),
512            mime_type: Some("image/png".into()),
513            source: FileSource::Clipboard,
514        };
515
516        let r2 = process_file_entry(&entry2, "s1", &ws, false, &store)
517            .await
518            .unwrap();
519        assert!(r2.deduplicated);
520        assert_eq!(r1.ref_id, r2.ref_id);
521    }
522
523    #[tokio::test]
524    async fn malformed_base64() {
525        let tmp = tempfile::tempdir().unwrap();
526        let ws = tmp.path().to_string_lossy().to_string();
527        let store = setup_store(&ws).await;
528
529        let entry = FileEntry {
530            path: None,
531            data_b64: Some("not-valid-base64!!!".into()),
532            filename: Some("bad.png".into()),
533            mime_type: Some("image/png".into()),
534            source: FileSource::File,
535        };
536
537        let err = process_file_entry(&entry, "s1", &ws, false, &store)
538            .await
539            .unwrap_err();
540        assert_eq!(err.code, INVALID_PARAMS);
541        assert!(err.message.contains("base64"));
542    }
543
544    #[tokio::test]
545    async fn rejects_path_over_wss() {
546        let tmp = tempfile::tempdir().unwrap();
547        let ws = tmp.path().to_string_lossy().to_string();
548        let store = setup_store(&ws).await;
549
550        let entry = FileEntry {
551            path: Some("/home/user/file.txt".into()),
552            data_b64: None,
553            filename: None,
554            mime_type: None,
555            source: FileSource::File,
556        };
557
558        let err = process_file_entry(&entry, "s1", &ws, true, &store)
559            .await
560            .unwrap_err();
561        assert_eq!(err.code, INVALID_PARAMS);
562        assert!(err.message.contains("WSS"));
563    }
564
565    #[tokio::test]
566    async fn rejects_no_data_and_no_path() {
567        let tmp = tempfile::tempdir().unwrap();
568        let ws = tmp.path().to_string_lossy().to_string();
569        let store = setup_store(&ws).await;
570
571        let entry = FileEntry {
572            path: None,
573            data_b64: None,
574            filename: Some("orphan.txt".into()),
575            mime_type: None,
576            source: FileSource::File,
577        };
578
579        let err = process_file_entry(&entry, "s1", &ws, false, &store)
580            .await
581            .unwrap_err();
582        assert_eq!(err.code, INVALID_PARAMS);
583        assert!(err.message.contains("data_b64"));
584    }
585
586    #[tokio::test]
587    async fn path_mode() {
588        let tmp = tempfile::tempdir().unwrap();
589        let ws = tmp.path().to_string_lossy().to_string();
590        let store = setup_store(&ws).await;
591
592        let file_path = tmp.path().join("testfile.pdf");
593        std::fs::write(&file_path, b"%PDF-1.4 test content").unwrap();
594
595        let entry = FileEntry {
596            path: Some(file_path.to_string_lossy().to_string()),
597            data_b64: None,
598            filename: None,
599            mime_type: None,
600            source: FileSource::File,
601        };
602
603        let r = process_file_entry(&entry, "s1", &ws, false, &store)
604            .await
605            .unwrap();
606
607        assert!(r.ref_id.starts_with("sha256:"));
608        // Non-image path mode: prose format with original filename and workspace path.
609        assert!(
610            r.marker.starts_with("[Document: testfile.pdf]"),
611            "marker = {}",
612            r.marker
613        );
614        assert!(
615            r.marker.contains("/uploads/"),
616            "marker should include workspace path: {}",
617            r.marker
618        );
619        assert!(!r.deduplicated);
620        assert!(std::path::Path::new(&r.workspace_path).exists());
621    }
622}
zeroclaw_runtime/rpc/attachments.rs

zeroclaw_runtime/rpc/
attachments.rs
