capo_agent/

user_message.rs

#![cfg_attr(test, allow(clippy::expect_used, clippy::unwrap_used))]

//! Public `UserMessage` and `Attachment` types plus the internal
//! `prepare_user_message` pipeline that turns them into a
//! `motosan_agent_loop::Message` for provider dispatch.
//!
//! See spec §2 / §3 in
//! `docs/superpowers/specs/2026-05-20-capo-v0.6-design.md`.

use std::path::PathBuf;

use base64::{engine::general_purpose::STANDARD as B64, Engine as _};
use serde::{Deserialize, Serialize};

/// A user message comprising a text body and zero-or-more attachments.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct UserMessage {
    pub text: String,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub attachments: Vec<Attachment>,
}

impl UserMessage {
    /// Convenience for the common text-only case.
    pub fn text(text: impl Into<String>) -> Self {
        Self {
            text: text.into(),
            attachments: Vec::new(),
        }
    }
}

/// A user-message attachment. `non_exhaustive` because future variants
/// (inline-bytes, skill, mention) must not break downstream pattern-match.
#[non_exhaustive]
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum Attachment {
    /// Local image file. Read, sniffed, base64-encoded by `prepare_user_message`
    /// before reaching the provider.
    Image { path: PathBuf },
}

/// Discriminator for `AttachmentError`. Carried over the wire on
/// `UiEvent::AttachmentError` so RPC clients can branch programmatically.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AttachmentErrorKind {
    NotFound,
    UnsupportedExtension,
    TooLarge,
    UnreadableImage,
}

/// Why `prepare_user_message` rejected an attachment.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AttachmentError {
    NotFound { path: PathBuf },
    UnsupportedExtension { path: PathBuf, ext: String },
    TooLarge { path: PathBuf, size: u64 },
    UnreadableImage { path: PathBuf },
}

impl AttachmentError {
    pub fn kind(&self) -> AttachmentErrorKind {
        match self {
            Self::NotFound { .. } => AttachmentErrorKind::NotFound,
            Self::UnsupportedExtension { .. } => AttachmentErrorKind::UnsupportedExtension,
            Self::TooLarge { .. } => AttachmentErrorKind::TooLarge,
            Self::UnreadableImage { .. } => AttachmentErrorKind::UnreadableImage,
        }
    }
}

impl std::fmt::Display for AttachmentError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::NotFound { path } => {
                write!(f, "image not found: {}", path.display())
            }
            Self::UnsupportedExtension { path, ext } => write!(
                f,
                "image has unsupported extension '.{}': {} (supported: png, jpg, jpeg, gif, webp)",
                ext,
                path.display(),
            ),
            Self::TooLarge { path, size } => write!(
                f,
                "image is {} bytes; capo caps images at 5 MiB (5242880 bytes): {}",
                size,
                path.display(),
            ),
            Self::UnreadableImage { path } => {
                write!(
                    f,
                    "image could not be read or recognised: {}",
                    path.display()
                )
            }
        }
    }
}

impl std::error::Error for AttachmentError {}

/// Maximum image size (raw file bytes, pre-base64). 5 MiB.
pub(crate) const MAX_IMAGE_BYTES: u64 = 5 * 1024 * 1024;

const SUPPORTED_EXTENSIONS: &[&str] = &["png", "jpg", "jpeg", "gif", "webp"];

/// Return the MIME type for a candidate image file, or `None` if neither
/// magic-byte sniff nor extension lookup recognises it.
///
/// Magic-byte sniff is attempted only when `bytes.len() >= 12`. Extension
/// match is case-insensitive.
pub(crate) fn sniff_mime(bytes: &[u8], extension: &str) -> Option<&'static str> {
    if bytes.len() >= 12 {
        // PNG: 89 50 4E 47 0D 0A 1A 0A
        if bytes.starts_with(&[0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]) {
            return Some("image/png");
        }
        // JPEG: FF D8 FF
        if bytes.starts_with(&[0xFF, 0xD8, 0xFF]) {
            return Some("image/jpeg");
        }
        // GIF: "GIF87a" or "GIF89a"
        if bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a") {
            return Some("image/gif");
        }
        // WEBP: "RIFF????WEBP" — 4 bytes RIFF, 4 byte size, then "WEBP"
        if bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP" {
            return Some("image/webp");
        }
    }
    // Extension fallback.
    match extension.to_ascii_lowercase().as_str() {
        "png" => Some("image/png"),
        "jpg" | "jpeg" => Some("image/jpeg"),
        "gif" => Some("image/gif"),
        "webp" => Some("image/webp"),
        _ => None,
    }
}

/// Validate every attachment, read+sniff+base64-encode each image, and
/// assemble a `motosan_agent_loop::Message` (text-then-images, in declaration
/// order). If `msg.text` is empty, no text part is emitted.
///
/// Errors short-circuit on the first bad attachment; later attachments are
/// not inspected. Callers must surface the error before starting a turn.
pub(crate) fn prepare_user_message(
    msg: &UserMessage,
) -> Result<motosan_agent_loop::Message, AttachmentError> {
    use motosan_agent_loop::ContentPart;

    let mut parts: Vec<ContentPart> = Vec::with_capacity(msg.attachments.len() + 1);
    if !msg.text.is_empty() {
        parts.push(ContentPart::text(&msg.text));
    }

    for att in &msg.attachments {
        match att {
            Attachment::Image { path } => {
                // 1. Existence.
                let metadata = match std::fs::metadata(path) {
                    Ok(m) if m.is_file() => m,
                    Ok(_) | Err(_) => {
                        return Err(AttachmentError::NotFound { path: path.clone() });
                    }
                };

                // 2. Extension.
                let ext = path
                    .extension()
                    .and_then(|s| s.to_str())
                    .map(|s| s.to_ascii_lowercase())
                    .unwrap_or_default();
                if !SUPPORTED_EXTENSIONS.contains(&ext.as_str()) {
                    return Err(AttachmentError::UnsupportedExtension {
                        path: path.clone(),
                        ext,
                    });
                }

                // 3. Size.
                let size = metadata.len();
                if size > MAX_IMAGE_BYTES {
                    return Err(AttachmentError::TooLarge {
                        path: path.clone(),
                        size,
                    });
                }

                // 4. Read + sniff + encode.
                let bytes = std::fs::read(path)
                    .map_err(|_| AttachmentError::UnreadableImage { path: path.clone() })?;
                let mime = sniff_mime(&bytes, &ext)
                    .ok_or_else(|| AttachmentError::UnreadableImage { path: path.clone() })?;
                let data = B64.encode(&bytes);

                parts.push(ContentPart::image_base64(mime.to_string(), data));
            }
        }
    }

    Ok(motosan_agent_loop::Message::user_with_parts(parts))
}

#[cfg(test)]
mod tests {
    use super::*;

    use motosan_agent_loop::{ContentPart, Message, Role};
    use std::io::Write;

    /// Helper: write `bytes` to a temp file with `extension`, return its path.
    /// The file is left in the OS tempdir; tests do not need cleanup.
    fn tempfile_with(extension: &str, bytes: &[u8]) -> PathBuf {
        let mut path = std::env::temp_dir();
        let name = format!(
            "capo-v06-test-{}-{}.{}",
            std::process::id(),
            uuid_like_suffix(),
            extension,
        );
        path.push(name);
        let mut f = std::fs::File::create(&path).expect("create tempfile");
        f.write_all(bytes).expect("write tempfile");
        path
    }

    fn uuid_like_suffix() -> String {
        use std::sync::atomic::{AtomicU64, Ordering};
        static COUNTER: AtomicU64 = AtomicU64::new(0);
        let n = COUNTER.fetch_add(1, Ordering::Relaxed);
        format!("{n:016x}")
    }

    fn png_header_bytes() -> Vec<u8> {
        // Minimal valid PNG: 8-byte magic + IHDR (13 bytes) + IEND (12 bytes).
        // We don't need a parseable PNG, only correct magic for sniff + non-zero size.
        let mut v = vec![0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
        v.extend_from_slice(&[0u8; 64]); // padding for total len > 12 bytes
        v
    }

    /// Helper: extract content parts from a `Message` known to be a `User`.
    fn parts(msg: &Message) -> &[ContentPart] {
        match msg {
            Message::User { content, .. } => content.as_slice(),
            other => panic!("expected User message, got {other:?}"),
        }
    }

    #[test]
    fn user_message_text_only_serializes_without_attachments_key() {
        let msg = UserMessage::text("hi");
        let json = serde_json::to_string(&msg).expect("serialize");
        assert_eq!(json, r#"{"text":"hi"}"#);
    }

    #[test]
    fn user_message_text_only_deserializes_when_attachments_absent() {
        let msg: UserMessage = serde_json::from_str(r#"{"text":"hi"}"#).expect("deserialize");
        assert_eq!(msg, UserMessage::text("hi"));
    }

    #[test]
    fn user_message_with_image_attachment_round_trips() {
        let msg = UserMessage {
            text: "look".into(),
            attachments: vec![Attachment::Image {
                path: PathBuf::from("/tmp/foo.png"),
            }],
        };
        let json = serde_json::to_string(&msg).expect("serialize");
        assert_eq!(
            json,
            r#"{"text":"look","attachments":[{"type":"image","path":"/tmp/foo.png"}]}"#
        );
        let back: UserMessage = serde_json::from_str(&json).expect("deserialize");
        assert_eq!(back, msg);
    }

    #[test]
    fn attachment_error_kind_serializes_to_stable_wire_strings() {
        let cases = [
            (AttachmentErrorKind::NotFound, "\"not_found\""),
            (
                AttachmentErrorKind::UnsupportedExtension,
                "\"unsupported_extension\"",
            ),
            (AttachmentErrorKind::TooLarge, "\"too_large\""),
            (AttachmentErrorKind::UnreadableImage, "\"unreadable_image\""),
        ];
        for (kind, expected) in cases {
            let got = serde_json::to_string(&kind).expect("serialize");
            assert_eq!(got, expected, "kind {kind:?} wire form");
        }
    }

    #[test]
    fn attachment_unknown_type_is_rejected() {
        // Locks the discriminator: future variants must use a NEW type tag,
        // not silently parse a misspelled "image" as something else.
        let err = serde_json::from_str::<Attachment>(r#"{"type":"img","path":"/tmp/x.png"}"#);
        assert!(err.is_err(), "unknown type tag must fail to deserialize");
    }

    #[test]
    fn prepare_text_only_produces_single_text_part() {
        let msg = UserMessage::text("hello");
        let out = prepare_user_message(&msg).expect("ok");
        assert_eq!(out.role(), Role::User);
        let p = parts(&out);
        assert_eq!(p.len(), 1);
        assert!(matches!(&p[0], ContentPart::Text { text, .. } if text == "hello"));
    }

    #[test]
    fn prepare_text_plus_one_image_produces_text_then_image() {
        let path = tempfile_with("png", &png_header_bytes());
        let msg = UserMessage {
            text: "look".into(),
            attachments: vec![Attachment::Image { path }],
        };
        let out = prepare_user_message(&msg).expect("ok");
        let p = parts(&out);
        assert_eq!(p.len(), 2);
        assert!(matches!(&p[0], ContentPart::Text { text, .. } if text == "look"));
        assert!(matches!(&p[1], ContentPart::Image { .. }));
    }

    #[test]
    fn prepare_text_plus_two_images_preserves_declared_order() {
        let p1 = tempfile_with("png", &png_header_bytes());
        let p2 = tempfile_with("png", &png_header_bytes());
        let msg = UserMessage {
            text: "compare".into(),
            attachments: vec![
                Attachment::Image { path: p1 },
                Attachment::Image { path: p2 },
            ],
        };
        let out = prepare_user_message(&msg).expect("ok");
        let p = parts(&out);
        assert_eq!(p.len(), 3);
        assert!(matches!(&p[0], ContentPart::Text { .. }));
        assert!(matches!(&p[1], ContentPart::Image { .. }));
        assert!(matches!(&p[2], ContentPart::Image { .. }));
    }

    #[test]
    fn prepare_empty_text_plus_one_image_omits_text_part() {
        let path = tempfile_with("png", &png_header_bytes());
        let msg = UserMessage {
            text: "".into(),
            attachments: vec![Attachment::Image { path }],
        };
        let out = prepare_user_message(&msg).expect("ok");
        let p = parts(&out);
        assert_eq!(p.len(), 1, "no leading empty text part");
        assert!(matches!(&p[0], ContentPart::Image { .. }));
    }

    #[test]
    fn prepare_user_message_rejects_missing_path() {
        let msg = UserMessage {
            text: "look".into(),
            attachments: vec![Attachment::Image {
                path: PathBuf::from("/tmp/definitely-does-not-exist-capo-v06.png"),
            }],
        };
        let err = prepare_user_message(&msg).expect_err("should fail");
        assert!(matches!(err, AttachmentError::NotFound { .. }));
        assert_eq!(err.kind(), AttachmentErrorKind::NotFound);
    }

    #[test]
    fn prepare_user_message_rejects_unsupported_extension() {
        let path = tempfile_with("txt", b"hello world");
        let msg = UserMessage {
            text: "".into(),
            attachments: vec![Attachment::Image { path }],
        };
        let err = prepare_user_message(&msg).expect_err("should fail");
        assert!(
            matches!(err, AttachmentError::UnsupportedExtension { ref ext, .. } if ext == "txt")
        );
    }

    #[test]
    fn prepare_user_message_rejects_oversize_file() {
        let big = vec![0u8; 5 * 1024 * 1024 + 1]; // 5 MiB + 1 byte
        let path = tempfile_with("png", &big);
        let msg = UserMessage {
            text: "".into(),
            attachments: vec![Attachment::Image { path }],
        };
        let err = prepare_user_message(&msg).expect_err("should fail");
        assert!(
            matches!(err, AttachmentError::TooLarge { size, .. } if size == 5 * 1024 * 1024 + 1)
        );
    }

    #[test]
    fn sniff_mime_recognises_png_magic_bytes() {
        let png_header: [u8; 12] = [
            0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x00, 0x00, 0x00,
        ];
        let got = sniff_mime(&png_header, "png");
        assert_eq!(got, Some("image/png"));
    }

    #[test]
    fn sniff_mime_recognises_jpeg_magic_bytes() {
        let mut bytes = [0u8; 12];
        bytes[..3].copy_from_slice(&[0xFF, 0xD8, 0xFF]);
        let got = sniff_mime(&bytes, "jpg");
        assert_eq!(got, Some("image/jpeg"));
    }

    #[test]
    fn sniff_mime_recognises_gif_magic_bytes() {
        // "GIF87a..."
        let bytes = b"GIF87a\0\0\0\0\0\0";
        let got = sniff_mime(bytes, "gif");
        assert_eq!(got, Some("image/gif"));
    }

    #[test]
    fn sniff_mime_recognises_webp_magic_bytes() {
        // "RIFF????WEBP"
        let bytes = b"RIFF\0\0\0\0WEBP";
        let got = sniff_mime(bytes, "webp");
        assert_eq!(got, Some("image/webp"));
    }

    #[test]
    fn sniff_mime_falls_back_to_extension_when_magic_inconclusive() {
        // All zero bytes — no magic matches.
        let bytes = [0u8; 12];
        assert_eq!(sniff_mime(&bytes, "png"), Some("image/png"));
        assert_eq!(sniff_mime(&bytes, "JPG"), Some("image/jpeg"));
        assert_eq!(sniff_mime(&bytes, "Jpeg"), Some("image/jpeg"));
        assert_eq!(sniff_mime(&bytes, "gif"), Some("image/gif"));
        assert_eq!(sniff_mime(&bytes, "webp"), Some("image/webp"));
    }

    #[test]
    fn sniff_mime_handles_files_shorter_than_12_bytes_via_extension_only() {
        let bytes = [0u8; 4];
        // Magic sniff skipped because slice < 12 bytes; extension lookup wins.
        assert_eq!(sniff_mime(&bytes, "png"), Some("image/png"));
    }

    #[test]
    fn sniff_mime_returns_none_when_both_magic_and_extension_unknown() {
        let bytes = [0u8; 12];
        assert_eq!(sniff_mime(&bytes, "txt"), None);
        assert_eq!(sniff_mime(&[], "bmp"), None);
    }
}