magi-code 0.97.5

//! Bounded, display-only applied previews. Never part of tool result/provider replay.
use super::{SessionEvent, SessionEventKind};
use crate::output::sanitize_display_text;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::{HashMap, HashSet};

const PREVIEW_BYTES: usize = 8192;
const PREVIEW_LINES: usize = 128;

#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct AppliedPreview {
    schema_version: u8,
    diff: String,
    truncated: bool,
}
impl AppliedPreview {
    pub(crate) fn select(diff: &str, already_truncated: bool) -> Self {
        // Redact multiline secrets before line/byte selection can cut their delimiters.
        let safe_diff = sanitize_display_text(diff);
        let mut selected = Self {
            schema_version: 1,
            diff: String::new(),
            truncated: already_truncated,
        };
        for (index, safe) in safe_diff.lines().enumerate() {
            if index == PREVIEW_LINES {
                selected.truncated = true;
                break;
            }
            let mut end = safe.len().min(1024);
            while !safe.is_char_boundary(end) {
                end -= 1;
            }
            selected.truncated |= end != safe.len();
            let previous = selected.diff.len();
            selected.diff.push_str(&safe[..end]);
            selected.diff.push('\n');
            if serde_json::to_vec(&selected)
                .expect("selected preview")
                .len()
                > PREVIEW_BYTES - 64
            {
                selected.diff.truncate(previous);
                selected.truncated = true;
                break;
            }
        }
        selected
    }
    pub(crate) fn apply(&self, result: &mut crate::tools::ToolResult) {
        result.display.edit_diff = Some(self.diff.clone());
        result.metadata["diff_truncated"] =
            Value::Bool(self.truncated || result.metadata["diff_truncated"] == true);
    }
    pub(crate) fn into_value(self) -> Value {
        serde_json::to_value(self).expect("selected preview")
    }
    fn parse(value: &Value) -> Option<Self> {
        if value.as_object()?.len() != 3
            || value["schema_version"].as_u64()? != 1
            || value["diff"].as_str()?.len() > PREVIEW_BYTES
            || value["truncated"].as_bool().is_none()
        {
            return None;
        }
        if serde_json::to_vec(value).ok()?.len() > PREVIEW_BYTES {
            return None;
        }
        let preview: Self = serde_json::from_value(value.clone()).ok()?;
        (preview.schema_version == 1
            && preview.diff.lines().count() <= PREVIEW_LINES
            && preview.diff.lines().all(|line| line.len() <= 1024)
            && sanitize_display_text(&preview.diff) == preview.diff)
            .then_some(preview)
    }
}

/// Unique result IDs, not adjacency, pair display records despite background writes.
/// Ambiguous/duplicate previews are unavailable, never last-writer-wins.
pub(crate) struct RecordedPreviews<'a> {
    records: HashMap<(&'a str, &'a str), Option<(usize, &'a SessionEvent)>>,
    result_ids: HashSet<(&'a str, &'a str)>,
    duplicates: HashSet<(&'a str, &'a str)>,
}
impl<'a> RecordedPreviews<'a> {
    pub(crate) fn new(events: &'a [SessionEvent]) -> Self {
        let mut selected = Self {
            records: HashMap::new(),
            result_ids: HashSet::new(),
            duplicates: HashSet::new(),
        };
        for (index, event) in events.iter().enumerate() {
            let Some(id) = event.payload["display_id"]
                .as_str()
                .filter(|id| !id.is_empty() && id.len() <= 64)
            else {
                continue;
            };
            let key = (event.session_id.as_str(), id);
            match event.kind() {
                Some(
                    SessionEventKind::ToolResult
                    | SessionEventKind::ToolDisplayResult
                    | SessionEventKind::CodeModeToolResult,
                ) => {
                    if !selected.result_ids.insert(key) {
                        selected.duplicates.insert(key);
                    }
                }
                Some(SessionEventKind::ToolAppliedPreview) => {
                    selected
                        .records
                        .entry(key)
                        .and_modify(|record| *record = None)
                        .or_insert(Some((index, event)));
                }
                _ => {}
            }
        }
        selected
    }
    pub(crate) fn apply(
        &self,
        index: usize,
        event: &SessionEvent,
        result: &mut crate::tools::ToolResult,
    ) {
        let Some(id) = event.payload["display_id"].as_str() else {
            return;
        };
        let key = (event.session_id.as_str(), id);
        if self.duplicates.contains(&key) {
            return;
        }
        let Some(Some((record_index, record))) = self.records.get(&key) else {
            return;
        };
        if *record_index <= index
            || record.payload["parent_call_id"] != event.payload["parent_call_id"]
            || record.payload["sequence"] != event.payload["sequence"]
        {
            return;
        }
        if !matches!(result.tool_name.as_str(), "hash_edit" | "write") {
            return;
        }
        if let Some(preview) = AppliedPreview::parse(&record.payload["preview"]) {
            preview.apply(result);
        }
    }
}