magi-code 0.102.2

//! Provider-only read-body deduplication. Raw conversation storage is never rewritten.

use std::collections::BTreeMap;

use serde_json::Value;

use super::{ProviderConversationItem, ProviderToolResult};

const MAX_WITNESSES: usize = 256;
const MAX_DESCRIPTORS: usize = 8;
const MAX_CALL_ID_BYTES: usize = 128;
// Bound ambiguity bookkeeping too. Larger requests retain every raw body.
const MAX_CANDIDATE_CALLS: usize = 4096;
// Bound owned overlay payloads, including mixed outputs and copied provenance. Map
// nodes and descriptor structs are separately bounded by MAX_CANDIDATE_CALLS.
const MAX_PROJECTED_BYTES: usize = 1024 * 1024;

#[derive(Default)]
struct CallEvidence {
    calls: u8,
    results: u8,
    read_call_index: Option<usize>,
}

#[derive(Clone, Copy)]
struct ReadBody<'a> {
    identity: &'a str,
    start: usize,
    end: usize,
    skill_source: Option<&'a str>,
    tag: Option<&'a str>,
}

struct Witness<'a> {
    read: ReadBody<'a>,
    body: &'a str,
    call_id: &'a str,
    section: usize,
}

pub(super) fn project<'a>(
    items: impl Iterator<Item = &'a ProviderConversationItem> + Clone,
) -> BTreeMap<usize, ProviderConversationItem> {
    let mut calls = BTreeMap::<&str, CallEvidence>::new();
    for item in items.clone() {
        if let ProviderConversationItem::ToolResult(result) = item
            && eligible_result(result)
            && !calls.contains_key(result.call_id.as_str())
        {
            if calls.len() == MAX_CANDIDATE_CALLS {
                return BTreeMap::new();
            }
            calls.insert(&result.call_id, CallEvidence::default());
        }
    }
    if calls.is_empty() {
        return BTreeMap::new();
    }

    // Count the whole assembled request, including later duplicates: serializers keep
    // the first output for a call ID, which might not contain a usable full body.
    for (index, item) in items.clone().enumerate() {
        match item {
            ProviderConversationItem::ToolResult(result) => {
                record_result(&mut calls, &result.call_id);
            }
            ProviderConversationItem::ResponseItem(item) => {
                match item.get("type").and_then(Value::as_str) {
                    Some("function_call") => {
                        if let Some(call) = item
                            .get("call_id")
                            .and_then(Value::as_str)
                            .and_then(|id| calls.get_mut(id))
                        {
                            call.calls = call.calls.saturating_add(1);
                            if item.get("name").and_then(Value::as_str) == Some("read") {
                                call.read_call_index = Some(index);
                            }
                        }
                    }
                    Some("function_call_output") => {
                        if let Some(id) = item.get("call_id").and_then(Value::as_str) {
                            record_result(&mut calls, id);
                        }
                    }
                    _ => {
                        // Legacy Chat messages can also consume IDs before canonical
                        // function calls/results. They are not portable witnesses.
                        if item.get("role").and_then(Value::as_str) == Some("tool")
                            && let Some(id) = item.get("tool_call_id").and_then(Value::as_str)
                        {
                            record_result(&mut calls, id);
                        }
                        if item.get("role").and_then(Value::as_str) == Some("assistant") {
                            for call in item
                                .get("tool_calls")
                                .and_then(Value::as_array)
                                .into_iter()
                                .flatten()
                            {
                                if let Some(evidence) = call
                                    .get("id")
                                    .and_then(Value::as_str)
                                    .and_then(|id| calls.get_mut(id))
                                {
                                    evidence.calls = evidence.calls.saturating_add(1);
                                }
                            }
                        }
                    }
                }
            }
            _ => {}
        }
    }

    let mut witnesses = Vec::new();
    let mut replacements = BTreeMap::new();
    let mut remaining_bytes = MAX_PROJECTED_BYTES;
    for (index, item) in items.enumerate() {
        let ProviderConversationItem::ToolResult(result) = item else {
            continue;
        };
        if !eligible_result(result) {
            continue;
        }
        let Some(evidence) = calls.get(result.call_id.as_str()) else {
            continue;
        };
        if evidence.calls != 1
            || evidence.results != 1
            || evidence
                .read_call_index
                .is_none_or(|call_index| call_index >= index)
        {
            continue;
        }
        if let Some(projected) = project_result(result, &mut witnesses, &mut remaining_bytes) {
            replacements.insert(index, ProviderConversationItem::ToolResult(projected));
        }
    }
    replacements
}

fn eligible_result(result: &ProviderToolResult) -> bool {
    result.success
        && result.tool_name == "read"
        && !(result.skill_reads.is_empty() && result.file_reads.is_empty())
        && result.skill_reads.len() + result.file_reads.len() <= MAX_DESCRIPTORS
        && !result.call_id.is_empty()
        && result.call_id.len() <= MAX_CALL_ID_BYTES
        && !result
            .call_id
            .chars()
            .any(|c| c.is_whitespace() || c.is_control())
}

fn record_result(calls: &mut BTreeMap<&str, CallEvidence>, id: &str) {
    if let Some(call) = calls.get_mut(id) {
        call.results = call.results.saturating_add(1);
    }
}

fn project_result<'a>(
    result: &'a ProviderToolResult,
    witnesses: &mut Vec<Witness<'a>>,
    remaining_bytes: &mut usize,
) -> Option<ProviderToolResult> {
    if result
        .skill_reads
        .iter()
        .any(|read| !read.validates(&result.output))
        || result
            .file_reads
            .iter()
            .any(|read| !read.validates(&result.output))
    {
        return None;
    }
    let mut descriptors = result
        .skill_reads
        .iter()
        .map(|read| ReadBody {
            identity: &read.identity,
            start: read.start,
            end: read.end,
            skill_source: Some(&read.source),
            tag: None,
        })
        .chain(result.file_reads.iter().map(|read| ReadBody {
            identity: &read.identity,
            start: read.start,
            end: read.end,
            skill_source: None,
            tag: Some(&read.tag),
        }))
        .collect::<Vec<_>>();
    descriptors.sort_by_key(|read| read.start);
    // Overlaps could remove a section that was registered as a retained witness.
    if descriptors
        .windows(2)
        .any(|pair| pair[0].end > pair[1].start)
    {
        return None;
    }

    let mut edits = Vec::new();
    let mut output_len = result.output.len();
    for (section, read) in descriptors.into_iter().enumerate() {
        let body = result.output.get(read.start..read.end)?;
        let witness = witnesses.iter().find(|witness| {
            witness.read.identity == read.identity
                && witness.read.tag == read.tag
                // Do not rely on a digest alone to justify removing bytes.
                && witness.body == body
        });
        if let Some(witness) = witness {
            let marker = if let Some(tag) = witness.read.tag {
                format!(
                    "[unchanged: identical to read result {} above, tag #{tag}; anchors from that result remain valid]",
                    witness.call_id
                )
            } else {
                format!(
                    "[Unchanged skill body: see {:?}, tool call {:?}, skill section {}.]",
                    witness.read.skill_source.unwrap_or_default(),
                    witness.call_id,
                    witness.section
                )
            };
            if marker.len() < body.len() {
                output_len -= body.len() - marker.len();
                edits.push((read.start, read.end, marker));
            }
        } else if witnesses.len() < MAX_WITNESSES {
            // Only full, retained bodies can witness later sections or results.
            witnesses.push(Witness {
                read,
                body,
                call_id: &result.call_id,
                section: section + 1,
            });
        }
    }
    if edits.is_empty() {
        return None;
    }
    let metadata_bytes = result.call_id.len()
        + result.tool_name.len()
        + result
            .skill_reads
            .iter()
            .map(|read| read.identity.len() + read.source.len() + read.sha256.len())
            .sum::<usize>()
        + result
            .file_reads
            .iter()
            .map(|read| read.identity.len() + read.tag.len() + read.sha256.len())
            .sum::<usize>();
    let projected_bytes = output_len.checked_add(metadata_bytes)?;
    if projected_bytes > *remaining_bytes {
        // Nothing was removed: all bodies registered above still remain visible.
        return None;
    }
    let mut output = String::with_capacity(output_len);
    let mut cursor = 0;
    for (start, end, marker) in edits {
        output.push_str(&result.output[cursor..start]);
        output.push_str(&marker);
        cursor = end;
    }
    output.push_str(&result.output[cursor..]);
    *remaining_bytes -= projected_bytes;
    Some(ProviderToolResult {
        call_id: result.call_id.clone(),
        tool_name: result.tool_name.clone(),
        success: result.success,
        output,
        // Keep raw offsets: a reassembled view must still pass full span/digest
        // validation before it can witness another read.
        skill_reads: result.skill_reads.clone(),
        file_reads: result.file_reads.clone(),
    })
}