codecraft/server/src/media.rs

//! Media overview of a chat: every image, video and audio file that was shown in it –
//! generated images, images in tool results (screenshots), media files the agent wrote,
//! media referenced in answers, and the user's own attachments. Includes subagents.

use crate::app::App;
use crate::types::{Message, Part, Role};
use serde::Serialize;
use serde_json::Value;
use std::collections::HashSet;
use std::path::{Path, PathBuf};

#[derive(Debug, Serialize, PartialEq)]
pub struct MediaItem {
    /// "image", "video", "audio" or "model" (3D).
    pub kind: &'static str,
    pub url: String,
    pub name: String,
    /// "generated", "tool", "written", "answer" or "upload".
    pub source: &'static str,
    pub created_at: i64,
    /// Message to jump to (only for messages of the chat itself, not of subagents).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub message_id: Option<String>,
    /// 3D models: absolute path (opened in the 3D view) and a preview image.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub path: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub preview: Option<String>,
}

const IMAGE: &[&str] = &["png", "jpg", "jpeg", "gif", "webp", "svg", "bmp", "avif"];
const VIDEO: &[&str] = &["mp4", "webm", "mov", "m4v", "ogv"];
const AUDIO: &[&str] = &["mp3", "wav", "ogg", "oga", "m4a", "flac", "opus", "aac"];
const MODEL: &[&str] = &["glb", "gltf", "stl", "3mf", "obj", "step", "stp"];

/// Kind of a media file by its extension.
pub fn kind_of(path: &str) -> Option<&'static str> {
    let ext = Path::new(path.split(['?', '#']).next().unwrap_or(path)).extension()?.to_str()?.to_ascii_lowercase();
    if IMAGE.contains(&ext.as_str()) {
        Some("image")
    } else if VIDEO.contains(&ext.as_str()) {
        Some("video")
    } else if AUDIO.contains(&ext.as_str()) {
        Some("audio")
    } else if MODEL.contains(&ext.as_str()) {
        Some("model")
    } else {
        None
    }
}

fn kind_of_mime(mime: &str) -> &'static str {
    if mime.starts_with("video/") {
        "video"
    } else if mime.starts_with("audio/") {
        "audio"
    } else {
        "image"
    }
}

/// Like JavaScript's encodeURIComponent.
fn encode(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    for b in s.bytes() {
        if b.is_ascii_alphanumeric() || b"-_.!~*'()".contains(&b) {
            out.push(b as char);
        } else {
            out.push_str(&format!("%{b:02X}"));
        }
    }
    out
}

/// Absolute, normalised path of `p` (relative to `cwd`).
fn resolve(p: &str, cwd: &Path) -> PathBuf {
    let joined = if p.starts_with('/') { PathBuf::from(p) } else { cwd.join(p) };
    crate::tools::normalize(&joined)
}

/// Media file paths in a text: Markdown images/links and plain paths with a media extension.
fn paths_in(text: &str) -> Vec<String> {
    let mut out = Vec::new();
    for token in text.split(|c: char| c.is_whitespace() || "()[]<>\"'`,;".contains(c)) {
        let t = token.trim_end_matches(['.', ':', '!', '?']);
        if t.is_empty() || t.contains("://") || t.starts_with("data:") {
            continue;
        }
        if kind_of(t).is_some() && t.chars().any(|c| c == '/' || c == '.') {
            out.push(t.to_string());
        }
    }
    out
}

/// All media of a chat and its subagents, newest first.
pub fn collect(app: &App, session_id: &str) -> anyhow::Result<Vec<MediaItem>> {
    let mut sessions = vec![session_id.to_string()];
    // Subagents (and theirs).
    let mut i = 0;
    while i < sessions.len() && sessions.len() < 500 {
        let children = app.store.child_sessions(&sessions[i])?;
        sessions.extend(children);
        i += 1;
    }
    let mut out: Vec<MediaItem> = Vec::new();
    let mut seen: HashSet<String> = HashSet::new();
    let mut push = |item: MediaItem, out: &mut Vec<MediaItem>| {
        if seen.insert(item.url.clone()) {
            out.push(item);
        }
    };
    for sid in &sessions {
        let Some(session) = app.store.get_session(sid)? else { continue };
        let cwd = PathBuf::from(&session.cwd);
        let own = sid == session_id;
        // Tool results are shown in the card of the call: jump to that message.
        let mut call_msg: Option<String> = None;
        for m in app.store.messages(sid)? {
            collect_message(&m, &cwd, own, call_msg.as_deref(), &mut |item| push(item, &mut out));
            if m.role == Role::Assistant {
                call_msg = Some(m.id.clone());
            }
        }
    }
    // Generated images whose tool result was not found (e.g. removed by a rewind).
    for img in app.store.images_of_sessions(&sessions)? {
        push(
            MediaItem { kind: "image", url: format!("/api/images/{}/file", img.id), name: file_name(&img.path), source: "generated", created_at: img.created_at, message_id: None, path: None, preview: None },
            &mut out,
        );
    }
    // Names of generated images (their files).
    let names: std::collections::HashMap<String, String> = app.store.images_of_sessions(&sessions)?.into_iter().map(|i| (format!("/api/images/{}/file", i.id), file_name(&i.path))).collect();
    for item in &mut out {
        if let Some(n) = names.get(&item.url) {
            item.name = n.clone();
        }
    }
    out.sort_by(|a, b| b.created_at.cmp(&a.created_at));
    Ok(out)
}

fn raw_url(p: &Path) -> String {
    format!("/api/files/raw?path={}", encode(&p.display().to_string()))
}

fn file_name(p: &str) -> String {
    Path::new(p).file_name().and_then(|n| n.to_str()).unwrap_or(p).to_string()
}

fn collect_message(m: &Message, cwd: &Path, own: bool, call_msg: Option<&str>, add: &mut dyn FnMut(MediaItem)) {
    let mid = own.then(|| m.id.clone());
    let result_mid = own.then(|| call_msg.unwrap_or(&m.id).to_string());
    let at = m.created_at;
    let file = |p: &str, source: &'static str| -> Option<MediaItem> {
        let kind = kind_of(p)?;
        let abs = resolve(p, cwd);
        abs.is_file().then(|| MediaItem { kind, url: raw_url(&abs), name: file_name(p), source, created_at: at, message_id: mid.clone(), path: Some(abs.display().to_string()), preview: None })
    };
    for p in &m.parts {
        match p {
            Part::Image { image } if m.role == Role::User => {
                if let Some(h) = &image.blob {
                    add(MediaItem { kind: kind_of_mime(&image.media_type), url: format!("/api/blobs/{h}"), name: "image".into(), source: "upload", created_at: at, message_id: mid.clone(), path: None, preview: None });
                }
            }
            // generate_3d: the model with its preview (not the preview as a separate image).
            Part::ToolResult { meta: Some(meta), .. } if meta["model_file"].is_string() => {
                let file = PathBuf::from(meta["model_file"].as_str().unwrap_or_default());
                if file.is_file() {
                    let preview = meta["preview"].as_str().map(PathBuf::from).filter(|p| p.is_file()).map(|p| raw_url(&p));
                    add(MediaItem { kind: "model", url: raw_url(&file), name: file_name(&file.display().to_string()), source: "generated", created_at: at, message_id: result_mid.clone(), path: Some(file.display().to_string()), preview });
                }
            }
            Part::ToolResult { images, meta, .. } => {
                let ids: Vec<&str> = meta.as_ref().and_then(|v| v["image_ids"].as_array()).map(|a| a.iter().filter_map(Value::as_str).collect()).unwrap_or_default();
                if !ids.is_empty() {
                    for id in ids {
                        add(MediaItem { kind: "image", url: format!("/api/images/{id}/file"), name: "image".into(), source: "generated", created_at: at, message_id: result_mid.clone(), path: None, preview: None });
                    }
                } else {
                    for img in images {
                        if let Some(h) = &img.blob {
                            add(MediaItem { kind: kind_of_mime(&img.media_type), url: format!("/api/blobs/{h}"), name: "screenshot".into(), source: "tool", created_at: at, message_id: result_mid.clone(), path: None, preview: None });
                        }
                    }
                }
            }
            Part::ToolCall { name, input, .. } if matches!(name.as_str(), "write" | "edit") => {
                if let Some(item) = input["path"].as_str().and_then(|p| file(p, "written")) {
                    add(item);
                }
            }
            Part::Text { text } if m.role == Role::Assistant => {
                for p in paths_in(text) {
                    if let Some(item) = file(&p, "answer") {
                        add(item);
                    }
                }
            }
            _ => {}
        }
    }
}

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

    #[test]
    fn kinds_paths_and_encoding() {
        assert_eq!(kind_of("a/b.PNG"), Some("image"));
        assert_eq!(kind_of("clip.mp4?v=1"), Some("video"));
        assert_eq!(kind_of("song.flac"), Some("audio"));
        assert_eq!(kind_of("main.rs"), None);
        assert_eq!(paths_in("Saved to `art/final/hero.png` and ![x](shots/a.webp). See https://x.org/y.png"), vec!["art/final/hero.png", "shots/a.webp"]);
        assert_eq!(encode("/w/ä b.png"), "%2Fw%2F%C3%A4%20b.png");
    }
}