magi-code 0.104.0

use super::{MissionControlApp, completion_worker::CompletionWorker};
use crate::tui::{
    input::PromptViewport, prompt_editor::PromptEditResult, state::MissionControlState,
};
use anyhow::Context;
use cap_fs_ext::{DirExt, FollowSymlinks, OpenOptionsFollowExt};
#[cfg(unix)]
use cap_std::fs::{OpenOptionsExt, PermissionsExt};
use cap_std::{
    ambient_authority,
    fs::{Dir, OpenOptions},
};
use image::ImageEncoder;
use std::{
    io::{self, Write},
    path::{Path, PathBuf},
};

const MAX_CLIPBOARD_IMAGE_PIXELS: usize = 16 * 1024 * 1024;
const CLIPBOARD_PASTE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);

// One request per conversation; retain completion even when its wake is lost.
#[derive(Default)]
pub(super) struct ClipboardPasteRuntime {
    worker: Option<CompletionWorker<ClipboardContent>>,
    insertion: Option<PendingInsertion>,
}

struct PendingInsertion {
    session_id: Option<String>,
    generation: u64,
    editor: (u64, usize, Option<(usize, usize)>),
    viewport: PromptViewport,
    started: std::time::Instant,
}

enum ClipboardContent {
    Image { path: String, id: u64 },
    Text(String),
}

impl ClipboardPasteRuntime {
    pub(super) fn is_pending(&self) -> bool {
        self.insertion.is_some()
    }

    pub(super) fn is_running(&self) -> bool {
        self.worker.is_some()
    }

    pub(super) fn cleanup(&mut self) -> Vec<String> {
        self.insertion = None;
        let Some(worker) = self.worker.take() else {
            return Vec::new();
        };
        match worker.join_with_timeout(super::super::WORKER_EXIT_JOIN_TIMEOUT) {
            Some(Ok(_)) => Vec::new(),
            Some(Err(error)) => vec![format!("clipboard paste: {error}")],
            None => vec!["clipboard paste did not finish before shutdown".into()],
        }
    }
}

impl MissionControlApp {
    pub(super) fn start_clipboard_paste(
        &mut self,
        ui: &mut MissionControlState,
        viewport: PromptViewport,
    ) {
        if self.clipboard_paste.is_running() {
            ui.status = "clipboard worker is still busy; please wait".into();
            return;
        }
        let root = self.config.paths.root.clone();
        let id = ui.prompt_editor.next_paste_id();
        let max_bytes = self.settings.tools.view_image.max_image_bytes;
        let insertion = PendingInsertion {
            session_id: self.state.active_session_id().map(str::to_owned),
            generation: self.session_generation,
            editor: ui.prompt_editor.paste_insertion_state(),
            viewport,
            started: std::time::Instant::now(),
        };
        match CompletionWorker::spawn("magi-clipboard-paste", self.events.clone(), move || {
            read_clipboard(&root, id, max_bytes).map_err(|error| error.to_string())
        }) {
            Ok(worker) => {
                self.clipboard_paste.worker = Some(worker);
                self.clipboard_paste.insertion = Some(insertion);
                ui.status = "Reading clipboard…".into();
                self.mark_initial_prompt_input(true, None, ui);
            }
            Err(error) => ui.status = format!("clipboard paste: {error}"),
        }
    }

    pub(super) fn poll_clipboard_paste(&mut self, ui: &mut MissionControlState) -> bool {
        let Some(worker) = self.clipboard_paste.worker.as_mut() else {
            return false;
        };
        let result = worker.take_result();
        let expired = result.is_none()
            && self
                .clipboard_paste
                .insertion
                .as_ref()
                .is_some_and(|pending| pending.started.elapsed() >= CLIPBOARD_PASTE_TIMEOUT);
        let changed = result.is_some() || expired;
        if expired {
            self.clipboard_paste.insertion = None;
            ui.status = "Clipboard paste timed out; draft unchanged".into();
        }
        if worker.ready_to_reap() {
            let _ = self
                .clipboard_paste
                .worker
                .take()
                .map(CompletionWorker::join);
        }
        if let Some(result) = result
            && let Some(insertion) = self.clipboard_paste.insertion.take()
        {
            if insertion.session_id.as_deref() != self.state.active_session_id()
                || insertion.generation != self.session_generation
            {
                return changed;
            }
            // Late native clipboard work remains owned; its result never changes a newer draft.
            if insertion.editor != ui.prompt_editor.paste_insertion_state() {
                ui.status = "Clipboard paste discarded: prompt changed; press Alt-V again".into();
                return changed;
            }
            match result {
                Ok(ClipboardContent::Image { path, id }) => {
                    let result = ui.insert_prompt_image(&path, id);
                    if result == PromptEditResult::Changed {
                        ui.status = format!("Pasted image #{id}");
                    }
                }
                Ok(ClipboardContent::Text(text)) => {
                    let result = ui.insert_prompt_str(
                        &text,
                        insertion.viewport.visible_rows,
                        insertion.viewport.wrap_width,
                    );
                    if matches!(
                        result,
                        PromptEditResult::Changed | PromptEditResult::Unchanged
                    ) {
                        ui.status = if result == PromptEditResult::Changed {
                            "Pasted clipboard text"
                        } else {
                            "Clipboard text is empty"
                        }
                        .into();
                    }
                }
                Err(error) => ui.status = format!("clipboard paste: {error}"),
            }
        }
        changed
    }
}

fn read_clipboard(root: &Path, id: u64, max_bytes: u64) -> anyhow::Result<ClipboardContent> {
    let mut clipboard = arboard::Clipboard::new()
        .context("native clipboard unavailable (requires a local desktop clipboard)")?;
    match clipboard.get_image() {
        Ok(image) => {
            let path = save_clipboard_image(root, id, image, max_bytes)?;
            Ok(ClipboardContent::Image { path, id })
        }
        Err(arboard::Error::ContentNotAvailable) => {
            let text = clipboard
                .get_text()
                .context("clipboard contains no supported image or text")?;
            let text = crate::tui::prompt_editor::normalize_prompt_text_bounded(&text)
                .map_err(|_| anyhow::anyhow!(crate::tui::prompt_editor::PROMPT_TOO_LARGE_STATUS))?;
            Ok(ClipboardContent::Text(text))
        }
        Err(error) => Err(error).context("could not read clipboard image"),
    }
}

fn save_clipboard_image(
    root: &Path,
    id: u64,
    image: arboard::ImageData<'_>,
    max_bytes: u64,
) -> anyhow::Result<String> {
    let pixels = image
        .width
        .checked_mul(image.height)
        .context("clipboard image dimensions overflow")?;
    anyhow::ensure!(
        pixels > 0 && pixels <= MAX_CLIPBOARD_IMAGE_PIXELS,
        "clipboard image exceeds 16 megapixels or is empty"
    );
    anyhow::ensure!(
        pixels.checked_mul(4) == Some(image.bytes.len()),
        "clipboard image has invalid RGBA data"
    );
    let width = u32::try_from(image.width)?;
    let height = u32::try_from(image.height)?;
    let root = root.canonicalize()?;
    let home = Dir::open_ambient_dir(&root, ambient_authority())?;
    match home.create_dir("tmp") {
        Ok(()) => {}
        Err(error) if error.kind() == io::ErrorKind::AlreadyExists => {}
        Err(error) => return Err(error.into()),
    }
    // Held directory handle prevents tmp symlink replacement from redirecting writes.
    let directory = home
        .open_dir_nofollow("tmp")
        .context("clipboard tmp must be a real directory, not a symlink")?;
    #[cfg(unix)]
    directory.set_permissions(".", cap_std::fs::Permissions::from_mode(0o700))?;
    let name = crate::tui::pasted_image::image_filename(id);
    let path: PathBuf = root.join("tmp").join(&name);
    let path = path
        .to_str()
        .context("clipboard image path must be UTF-8")?;
    anyhow::ensure!(
        !path.contains(']') && !path.chars().any(char::is_control),
        "clipboard image path cannot contain closing brackets or control characters"
    );
    let mut options = OpenOptions::new();
    options
        .write(true)
        .create_new(true)
        .follow(FollowSymlinks::No);
    #[cfg(unix)]
    options.mode(0o600);
    let mut file = directory.open_with(&name, &options)?;
    let result = (|| -> anyhow::Result<()> {
        let writer = LimitedImageWriter {
            writer: &mut file,
            remaining: max_bytes,
        };
        image::codecs::png::PngEncoder::new(writer).write_image(
            &image.bytes,
            width,
            height,
            image::ExtendedColorType::Rgba8,
        )?;
        file.sync_all()?;
        Ok(())
    })();
    drop(file);
    if let Err(error) = result {
        let _ = directory.remove_file(&name);
        return Err(error).context(format!(
            "could not save clipboard PNG (maximum {max_bytes} bytes)"
        ));
    }
    // Retain saved images across process exit so session references remain usable.
    Ok(path.to_owned())
}

struct LimitedImageWriter<W> {
    writer: W,
    remaining: u64,
}

impl<W: Write> Write for LimitedImageWriter<W> {
    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
        if bytes.len() as u64 > self.remaining {
            return Err(io::Error::other(
                "clipboard PNG exceeds view_image.max_image_bytes",
            ));
        }
        let written = self.writer.write(bytes)?;
        self.remaining -= written as u64;
        Ok(written)
    }
    fn flush(&mut self) -> io::Result<()> {
        self.writer.flush()
    }
}