magi-code 0.104.0

//! Append-only passive summaries; never used for replay or compaction.
mod status;
use super::{Session, SessionEvent};
use crate::{
    config::{EffectiveConfig, Settings},
    persistence::{CrossProcessFileLock, read_regular_file_bounded},
    progressive_session_maintenance::{Trigger, is_tracked},
};
use serde::{Deserialize, Serialize};
pub(crate) use status::{Stage, Status};
use std::{
    collections::HashSet,
    fs::File,
    io::{BufRead, BufReader, Read},
    path::Path,
};

const MAX_INPUT_BYTES: usize = 8 * 1024 * 1024;
const MAX_SIDECAR_BYTES: u64 = 64 * 1024 * 1024;
const MAX_SCAN_BYTES: u64 = 512 * 1024 * 1024;
const MAX_SCAN_RECORDS: usize = 1_000_000;
const MAX_BATCH_RECORDS: usize = 1000;

#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct Entry {
    event_id: String,
    #[serde(rename = "type")]
    event_type: String,
    summary: String,
    response_id: Option<String>,
}

#[derive(Serialize, Deserialize)]
struct Batch {
    schema_version: u8,
    source_event_ids: Vec<String>,
    entries: Vec<Entry>,
}

struct Record {
    event: SessionEvent,
    raw: String,
    turn: usize,
    // A terminal status allows a failed/cancelled call to have no response.
    turn_ended: bool,
}

/// Returns whether a scheduled batch is waiting for an in-flight tool response.
pub(crate) fn maintain(
    session: &Session,
    trigger: &Trigger,
    config: &EffectiveConfig,
    settings: &Settings,
    cwd: &Path,
    status: &mut Status,
) -> anyhow::Result<bool> {
    let root = session
        .path()
        .parent()
        .ok_or_else(|| anyhow::anyhow!("session root missing"))?;
    super::store::prepare_session_root(root)?;
    let name = format!("{}.psm.ndjson", session.id());
    let path = root.join(&name);
    status.stage(Stage::WaitingForLock);
    let _guard = CrossProcessFileLock::acquire(&path)?;
    status.stage(Stage::ReadingSummaries);
    let acknowledged = read_acknowledged(root, &name)?;
    status.stage(Stage::ReadingHistory);
    let records = collect_records(session, trigger, &acknowledged)?;
    status.pending(records.len());
    let mut start = 0;
    while start < records.len() {
        let mut end = start;
        let mut bytes = 2;
        while end < records.len() && end - start < MAX_BATCH_RECORDS {
            let next = records[end].raw.len() + 1;
            if bytes + next > MAX_INPUT_BYTES {
                break;
            }
            bytes += next;
            end += 1;
        }
        anyhow::ensure!(end > start, "PSM record exceeds full-record request limit");
        // Never commit a call without a response that may still arrive. Keep the
        // entire suffix for the next trigger, even across a size/count boundary.
        while end > start {
            let incomplete = records[start..end].iter().position(|record| {
                record.event.event_type == "tool_call"
                    && response_id(record, &records[start..end]).is_none()
                    && (!record.turn_ended || response_id(record, &records).is_some())
            });
            match incomplete {
                Some(index) => end = start + index,
                None => break,
            }
        }
        if end == start {
            anyhow::ensure!(
                response_id(&records[start], &records).is_none(),
                "PSM call/response span exceeds batch limits"
            );
            return Ok(true);
        }
        let pending = &records[start..end];
        let entries = generate_entries(config, settings, cwd, pending, status)?;
        let source_event_ids = pending
            .iter()
            .filter_map(|record| record.event.event_id.clone())
            .collect();
        let batch = Batch {
            schema_version: 1,
            source_event_ids,
            entries,
        };
        let mut line = serde_json::to_vec(&batch)?;
        status.stage(Stage::WritingSummaries);
        line.push(b'\n');
        let mut file = super::store::open_append_named(root, &name)?;
        let previous_len = file.metadata()?.len();
        anyhow::ensure!(
            previous_len + line.len() as u64 <= MAX_SIDECAR_BYTES,
            "PSM sidecar exceeds size limit"
        );
        super::write::append_event_line_durably(&mut file, &line, previous_len)?;
        crate::persistence::sync_parent_dir(root)?;
        status.saved(pending.len());
        start = end;
    }
    Ok(false)
}

fn read_acknowledged(root: &Path, name: &str) -> anyhow::Result<HashSet<String>> {
    let Some(file) = super::store::open_existing_named(root, name)? else {
        return Ok(HashSet::new());
    };
    let bytes = read_regular_file_bounded(&file, MAX_SIDECAR_BYTES)?;
    let complete = bytes
        .iter()
        .rposition(|byte| *byte == b'\n')
        .map_or(0, |index| index + 1);
    if complete != bytes.len() {
        // A process killed during an append leaves an uncommitted tail, never a
        // successful cursor. Repair only that tail under the sidecar lock.
        let file = super::store::open_append_named(root, name)?;
        file.set_len(complete as u64)?;
        file.sync_all()?;
    }
    let mut ids = HashSet::new();
    for line in bytes[..complete]
        .split(|byte| *byte == b'\n')
        .filter(|line| !line.is_empty())
    {
        let batch: Batch = serde_json::from_slice(line)?;
        anyhow::ensure!(batch.schema_version == 1, "unsupported PSM sidecar version");
        ids.extend(batch.source_event_ids);
    }
    Ok(ids)
}

fn source_files(session: &Session) -> anyhow::Result<Vec<File>> {
    // Hold append/rotation lock only while opening immutable archives and the
    // current inode. No model request or history scan holds a session write lock.
    let append_lock = super::write::session_append_lock(session.path())?;
    let _guard = append_lock
        .lock()
        .map_err(|_| anyhow::anyhow!("session append lock poisoned"))?;
    let _file_guard = CrossProcessFileLock::acquire(session.path())?;
    let root = session
        .path()
        .parent()
        .ok_or_else(|| anyhow::anyhow!("session root missing"))?;
    let history_parent = root.join(".history");
    let mut files = Vec::new();
    if history_parent.try_exists()? {
        super::store::prepare_session_root(&history_parent)?;
        let history = history_parent.join(session.id());
        if history.try_exists()? {
            super::store::prepare_session_root(&history)?;
            let mut generations = Vec::new();
            for (index, entry) in std::fs::read_dir(&history)?.take(129).enumerate() {
                anyhow::ensure!(index < 128, "too many PSM history directory entries");
                let entry = entry?;
                if let Some(generation) = entry
                    .file_name()
                    .to_str()
                    .and_then(|name| name.strip_suffix(".jsonl"))
                    .and_then(|name| name.parse::<u64>().ok())
                {
                    generations.push(generation);
                }
            }
            generations.sort_unstable();
            for generation in generations {
                if let Some(file) =
                    super::store::open_existing_named(&history, &format!("{generation}.jsonl"))?
                {
                    files.push(file);
                }
            }
        }
    }
    if let Some(file) = super::store::open_existing_primary(root, session.id())? {
        files.push(file);
    }
    Ok(files)
}

fn collect_records(
    session: &Session,
    trigger: &Trigger,
    acknowledged: &HashSet<String>,
) -> anyhow::Result<Vec<Record>> {
    let mut seen = HashSet::new();
    let mut records: Vec<Record> = Vec::new();
    let mut scanned_bytes = 0;
    let mut retained_bytes = 0;
    let mut scanned_records = 0;
    let mut turn_start = 0;
    let mut turn = 0;
    for file in source_files(session)? {
        // Read only the captured length; concurrent appends belong to a later trigger.
        let len = file.metadata()?.len();
        let mut reader = BufReader::new(file.take(len));
        loop {
            let mut line = Vec::new();
            let count = reader
                .by_ref()
                .take(MAX_INPUT_BYTES as u64 + 1)
                .read_until(b'\n', &mut line)?;
            if count == 0 {
                break;
            }
            scanned_bytes += count as u64;
            scanned_records += 1;
            anyhow::ensure!(
                count <= MAX_INPUT_BYTES
                    && scanned_bytes <= MAX_SCAN_BYTES
                    && scanned_records <= MAX_SCAN_RECORDS,
                "PSM history scan exceeds limits"
            );
            anyhow::ensure!(
                line.last() == Some(&b'\n'),
                "incomplete session record during PSM scan"
            );
            let event: SessionEvent = serde_json::from_slice(&line)?;
            anyhow::ensure!(
                event.session_id == session.id(),
                "PSM source session mismatch"
            );
            let Some(id) = event.event_id.as_ref() else {
                continue;
            };
            let through = id == &trigger.through_event_id;
            if seen.insert(id.clone()) {
                if (event.event_type == "turn_status"
                    && event.payload["status"]
                        .as_str()
                        .is_some_and(|status| status != "incomplete"))
                    || (event.event_type == "user_input"
                        && event.payload["origin"]
                            .as_str()
                            .is_none_or(|origin| origin == "user"))
                {
                    for record in &mut records[turn_start..] {
                        record.turn_ended = true;
                    }
                    turn_start = records.len();
                    turn += 1;
                }
                if is_tracked(&event.event_type) && !acknowledged.contains(id) {
                    retained_bytes += line.len();
                    anyhow::ensure!(
                        retained_bytes <= 64 * 1024 * 1024,
                        "PSM pending records exceed memory limit"
                    );
                    records.push(Record {
                        event,
                        raw: String::from_utf8(line)?,
                        turn,
                        turn_ended: false,
                    });
                }
            }
            if through {
                if trigger.turn_ended {
                    for record in &mut records[turn_start..] {
                        record.turn_ended = true;
                    }
                }
                return Ok(records);
            }
        }
    }
    anyhow::bail!("PSM trigger record not found in session history")
}

fn response_id<'a>(call: &Record, records: &'a [Record]) -> Option<&'a str> {
    let call_id = call.event.payload["id"].as_str()?;
    let call_index = records
        .iter()
        .position(|record| record.event.event_id == call.event.event_id)?;
    records[call_index + 1..]
        .iter()
        .take_while(|record| record.turn == call.turn)
        .find(|record| {
            record.event.event_type == "tool_result"
                && record.event.payload["call_id"].as_str() == Some(call_id)
        })
        .and_then(|record| record.event.event_id.as_deref())
}

fn generate_entries(
    config: &EffectiveConfig,
    settings: &Settings,
    cwd: &Path,
    records: &[Record],
    status: &mut Status,
) -> anyhow::Result<Vec<Entry>> {
    let raw_records: Vec<_> = records.iter().map(|record| record.raw.trim_end()).collect();
    let raw = crate::progressive_session_maintenance::summarize_records(
        config,
        settings,
        cwd,
        &raw_records,
        status,
    )?;
    status.stage(Stage::ValidatingSummary);
    let mut entries: Vec<Entry> = serde_json::from_str(raw.trim())?;
    let sources: Vec<_> = records
        .iter()
        .filter(|record| record.event.event_type != "tool_result")
        .collect();
    anyhow::ensure!(
        entries.len() == sources.len(),
        "PSM response must cover every summary event exactly once"
    );
    for (entry, source) in entries.iter_mut().zip(sources) {
        anyhow::ensure!(
            Some(entry.event_id.as_str()) == source.event.event_id.as_deref()
                && entry.event_type == source.event.event_type,
            "PSM event identity or order mismatch"
        );
        let expected_response = if entry.event_type == "tool_call" {
            response_id(source, records)
        } else {
            None
        };
        anyhow::ensure!(
            entry.response_id.as_deref() == expected_response,
            "PSM response link mismatch"
        );
        let mut preserved_text = false;
        if matches!(entry.event_type.as_str(), "user_input" | "assistant_output")
            && let Some(text) = source.event.payload["text"]
                .as_str()
                .filter(|text| text.chars().count() <= 240)
        {
            entry.summary = text.to_string();
            preserved_text = true;
        }
        anyhow::ensure!(
            (preserved_text || !entry.summary.trim().is_empty())
                && entry.summary.chars().count() <= 2000,
            "PSM summary is empty or too long"
        );
        entry.summary = crate::output::redact_sensitive_text(&entry.summary);
    }
    Ok(entries)
}