magi-code 0.80.2

//! Local prompt suggestions. All filesystem methods are worker-only.
use crate::persistence::{
    CrossProcessFileLock, atomic_write_with_permissions, open_regular_file,
    read_regular_file_bounded,
};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::{
    fs,
    path::{Path, PathBuf},
    time::{Duration, Instant},
};

const MAX_ENTRIES: usize = 200;
const MAX_PROMPT_BYTES: usize = 4096;
const MAX_STORE_BYTES: u64 = 4 * 1024 * 1024;

#[derive(Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) struct PromptHistory {
    entries: Vec<HistoryEntry>,
}

impl std::fmt::Debug for PromptHistory {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PromptHistory")
            .field("entries", &self.entries.len())
            .finish()
    }
}

#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
struct HistoryEntry {
    text: String,
    count: u64,
    last_used: i64,
}

pub(crate) struct PromptHistoryRecorder {
    path: PathBuf,
    initial_prompt: Option<String>,
}

impl PromptHistoryRecorder {
    pub(crate) fn new(path: PathBuf, prompt: String) -> Self {
        Self {
            path,
            initial_prompt: Some(prompt),
        }
    }

    // The initial input announcement follows required session persistence.
    pub(crate) fn record_event(
        &mut self,
        event: &crate::output::OutputEvent,
    ) -> anyhow::Result<Option<PromptHistory>> {
        use crate::output::OutputEvent;
        let prompt = match event {
            OutputEvent::UserPrompt { .. } => self.initial_prompt.take(),
            _ => None,
        };
        prompt
            .map(|prompt| PromptHistory::remember(&self.path, &prompt))
            .transpose()
    }

    pub(crate) fn record_steering_prompt(&self, prompt: &str) -> anyhow::Result<PromptHistory> {
        PromptHistory::remember(&self.path, prompt)
    }
}

pub(crate) fn enabled() -> bool {
    std::env::var_os("MC_PROMPT_HISTORY").is_none_or(|value| value != "0")
}

// Resolve Git identity without spawning Git or consulting worktree-specific absolute paths.
fn scope_identity(cwd: &Path) -> anyhow::Result<PathBuf> {
    let cwd = cwd.canonicalize()?;
    for parent in cwd.ancestors() {
        let marker = parent.join(".git");
        let git = if marker.is_dir() {
            marker
        } else if marker.is_file() {
            let text = read_text(&marker, 8192)?;
            let target = text
                .trim()
                .strip_prefix("gitdir:")
                .ok_or_else(|| anyhow::anyhow!("invalid Git directory marker"))?;
            parent.join(target.trim())
        } else {
            continue;
        };
        let common = git.join("commondir");
        let identity = if common.exists() {
            git.join(read_text(&common, 8192)?.trim())
        } else {
            git
        };
        return Ok(identity.canonicalize()?);
    }
    Ok(cwd)
}

fn read_text(path: &Path, limit: u64) -> anyhow::Result<String> {
    let file = open_regular_file(path, false)?;
    Ok(String::from_utf8(read_regular_file_bounded(&file, limit)?)?)
}

pub(crate) fn store_path(state_dir: &Path, cwd: &Path) -> anyhow::Result<PathBuf> {
    let identity = scope_identity(cwd)?;
    let digest = Sha256::digest(identity.as_os_str().as_encoded_bytes());
    let digest: String = digest.iter().map(|byte| format!("{byte:02x}")).collect();
    Ok(state_dir
        .join("prompt-history")
        .join(format!("{digest}.json")))
}

fn check_history_directory(path: &Path) -> anyhow::Result<()> {
    let parent = path
        .parent()
        .ok_or_else(|| anyhow::anyhow!("missing history directory"))?;
    match fs::symlink_metadata(parent) {
        Ok(metadata) if metadata.file_type().is_symlink() || !metadata.is_dir() => {
            anyhow::bail!("history directory must be a regular directory")
        }
        Err(error) if error.kind() != std::io::ErrorKind::NotFound => return Err(error.into()),
        _ => {}
    }
    Ok(())
}

impl PromptHistory {
    pub(crate) fn load(path: &Path) -> anyhow::Result<Self> {
        check_history_directory(path)?;
        let text = match read_text(path, MAX_STORE_BYTES) {
            Ok(text) => text,
            Err(error)
                if error
                    .downcast_ref::<std::io::Error>()
                    .is_some_and(|e| e.kind() == std::io::ErrorKind::NotFound) =>
            {
                return Ok(Self::default());
            }
            Err(error) => return Err(error),
        };
        let mut history: Self = serde_json::from_str(&text)?;
        history.entries.retain(|entry| eligible(&entry.text));
        history.bound();
        Ok(history)
    }

    pub(crate) fn remember(path: &Path, text: &str) -> anyhow::Result<Self> {
        check_history_directory(path)?;
        fs::create_dir_all(
            path.parent()
                .ok_or_else(|| anyhow::anyhow!("missing history directory"))?,
        )?;
        let _lock =
            CrossProcessFileLock::acquire_until(path, Instant::now() + Duration::from_millis(100))?;
        let mut history = Self::load(path)?;
        history.record(text, chrono::Utc::now().timestamp_millis());
        atomic_write_with_permissions(path, &serde_json::to_vec(&history)?, Some(0o600))?;
        Ok(history)
    }

    fn record(&mut self, text: &str, now: i64) {
        if !eligible(text) {
            return;
        }
        if let Some(entry) = self.entries.iter_mut().find(|entry| entry.text == text) {
            entry.count = entry.count.saturating_add(1);
            entry.last_used = now;
        } else {
            self.entries.push(HistoryEntry {
                text: text.to_owned(),
                count: 1,
                last_used: now,
            });
        }
        self.bound();
    }

    fn bound(&mut self) {
        self.entries.sort_by(|a, b| {
            b.count
                .cmp(&a.count)
                .then_with(|| b.last_used.cmp(&a.last_used))
                .then_with(|| a.text.cmp(&b.text))
        });
        self.entries.truncate(MAX_ENTRIES);
    }

    pub(crate) fn suffix(&self, prefix: &str) -> Option<&str> {
        if prefix.is_empty() || prefix.contains(['\n', '\r']) {
            return None;
        }
        self.entries.iter().find_map(|entry| {
            entry
                .text
                .strip_prefix(prefix)
                .filter(|suffix| !suffix.trim().is_empty())
        })
    }
}

fn eligible(text: &str) -> bool {
    !text.trim().is_empty() && text.len() <= MAX_PROMPT_BYTES
        && !text.chars().any(char::is_control)
        && !text.starts_with(['/', '!'])
        && crate::output::redact_sensitive_text(text) == text
        // Conservative: never replay local absolute paths into another worktree.
        && !text.split(|ch: char| ch.is_whitespace() || "=\"'([{@`".contains(ch)).any(|word| {
            word.starts_with('/') || word.starts_with("~/") || word.contains(":\\")
                || (word.as_bytes().get(1..3) == Some(b":/"))
        })
}

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

    #[test]
    fn exact_prefix_ranks_frequency_then_recency_and_preserves_unicode() {
        let mut history = PromptHistory::default();
        history.record("fix café tests", 1);
        history.record("fix café tests", 2);
        history.record("fix café docs", 3);
        assert_eq!(history.suffix("fix café "), Some("tests"));
        history.record("fix café docs", 4);
        assert_eq!(history.suffix("fix café "), Some("docs"));
        assert_eq!(history.suffix("Fix"), None);
        assert_eq!(history.suffix(""), None);
        assert_eq!(history.suffix("fix café docs"), None);
    }

    #[test]
    fn unsafe_and_multiline_prompts_are_not_suggested() {
        let mut history = PromptHistory::default();
        for text in [
            "read /Users/me/repo/file",
            "read @/tmp/file",
            "two\nlines",
            "!rm file",
            "/settings",
            " ",
        ] {
            history.record(text, 1);
        }
        assert!(history.entries.is_empty());
    }

    #[test]
    fn worktrees_share_identity_and_non_git_directories_do_not() {
        let temp = tempfile::tempdir().unwrap();
        let main = temp.path().join("main");
        let linked = temp.path().join("linked");
        fs::create_dir_all(main.join(".git/worktrees/linked")).unwrap();
        fs::create_dir_all(&linked).unwrap();
        fs::write(
            linked.join(".git"),
            format!("gitdir: {}", main.join(".git/worktrees/linked").display()),
        )
        .unwrap();
        fs::write(main.join(".git/worktrees/linked/commondir"), "../..\n").unwrap();
        assert_eq!(
            scope_identity(&main).unwrap(),
            scope_identity(&linked).unwrap()
        );
        assert_ne!(
            scope_identity(temp.path()).unwrap(),
            scope_identity(&main).unwrap()
        );
    }

    #[test]
    fn store_survives_reload_and_is_bounded() {
        let temp = tempfile::tempdir().unwrap();
        let path = temp.path().join("history.json");
        PromptHistory::remember(&path, "check tests").unwrap();
        assert_eq!(
            PromptHistory::load(&path).unwrap().suffix("check "),
            Some("tests")
        );
        let mut history = PromptHistory::default();
        for n in 0..250 {
            history.record(&format!("prompt {n}"), n);
        }
        assert_eq!(history.entries.len(), MAX_ENTRIES);
        assert_eq!(history.suffix("prompt "), Some("249"));
    }
    #[test]
    fn independent_writers_merge_and_private_files_reject_symlinks() {
        let temp = tempfile::tempdir().unwrap();
        let path = temp.path().join("history.json");
        let other = path.clone();
        let worker = std::thread::spawn(move || PromptHistory::remember(&other, "check tests"));
        PromptHistory::remember(&path, "check docs").unwrap();
        worker.join().unwrap().unwrap();
        let history = PromptHistory::load(&path).unwrap();
        assert_eq!(history.suffix("check t"), Some("ests"));
        assert_eq!(history.suffix("check d"), Some("ocs"));
        #[cfg(unix)]
        {
            use std::os::unix::fs::{PermissionsExt, symlink};
            assert_eq!(
                fs::metadata(&path).unwrap().permissions().mode() & 0o777,
                0o600
            );
            let link = temp.path().join("link.json");
            symlink(&path, &link).unwrap();
            assert!(PromptHistory::remember(&link, "unwanted").is_err());
            assert!(PromptHistory::load(&link).is_err());
            let directory_link = temp.path().join("linked-directory");
            symlink(temp.path(), &directory_link).unwrap();
            assert!(PromptHistory::load(&directory_link.join("history.json")).is_err());
            assert!(
                PromptHistory::remember(&directory_link.join("history.json"), "unwanted").is_err()
            );
        }
    }

    #[test]
    fn oversized_history_is_rejected_without_overwriting_it() {
        let temp = tempfile::tempdir().unwrap();
        let path = temp.path().join("history.json");
        let file = fs::File::create(&path).unwrap();
        file.set_len(MAX_STORE_BYTES + 1).unwrap();
        assert!(PromptHistory::remember(&path, "check tests").is_err());
        assert_eq!(fs::metadata(path).unwrap().len(), MAX_STORE_BYTES + 1);
    }
}