magi-code 0.82.1

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

use std::collections::BTreeMap;

use serde_json::Value;

use super::{ProviderConversationItem, ProviderToolResult};
use crate::tools::skill_provenance::SkillReadProvenance;

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>,
}

struct Witness<'a> {
    provenance: &'a SkillReadProvenance,
    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_some_and(|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.skill_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> {
    let mut descriptors = result.skill_reads.iter().collect::<Vec<_>>();
    if descriptors
        .iter()
        .any(|read| !read.validates(&result.output))
    {
        return None;
    }
    descriptors.sort_by_key(|read| read.start);
    // Reject overlapping metadata as a unit; otherwise an apparent witness could
    // itself be removed by another descriptor in the same multi-read output.
    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.provenance.identity == read.identity
                && witness.provenance.sha256 == read.sha256
                // Do not rely on a digest alone to justify removing bytes.
                && witness.body == body
        });
        if let Some(witness) = witness {
            let marker = format!(
                "[Unchanged skill body: see {:?}, tool call {:?}, skill section {}.]",
                witness.provenance.source, 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 {
                provenance: 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>();
    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(),
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::providers::{
        ProviderRequest, anthropic::anthropic_messages_body, codex_responses_body,
        openai_compatible_chat_completions_body,
    };
    use serde_json::json;
    use sha2::{Digest, Sha256};

    fn read_call(id: &str) -> ProviderConversationItem {
        ProviderConversationItem::ResponseItem(json!({
            "type": "function_call", "call_id": id, "name": "read",
            "arguments": "{\"paths\":[\"skill://example\"]}"
        }))
    }

    fn skill_result(id: &str, body: &str) -> ProviderConversationItem {
        let header = "--- FILE: skill://example ---\n";
        ProviderConversationItem::ToolResult(ProviderToolResult {
            call_id: id.into(),
            tool_name: "read".into(),
            success: true,
            output: format!("{header}{body}"),
            skill_reads: vec![SkillReadProvenance {
                version: 1,
                identity: crate::hex::lower_hex(Sha256::digest(b"resolved-example-skill")),
                source: "skill://example".into(),
                start: header.len(),
                end: header.len() + body.len(),
                sha256: crate::hex::lower_hex(Sha256::digest(body.as_bytes())),
            }],
        })
    }

    fn serialized_outputs(request: &ProviderRequest, call_id: &str) -> [String; 3] {
        let codex = codex_responses_body("model", request);
        let chat = openai_compatible_chat_completions_body("model", request);
        let anthropic = anthropic_messages_body("model", request, None);
        [
            codex["input"]
                .as_array()
                .unwrap()
                .iter()
                .find(|item| item["type"] == "function_call_output" && item["call_id"] == call_id)
                .unwrap()["output"]
                .as_str()
                .unwrap()
                .to_owned(),
            chat["messages"]
                .as_array()
                .unwrap()
                .iter()
                .find(|item| item["role"] == "tool" && item["tool_call_id"] == call_id)
                .unwrap()["content"]
                .as_str()
                .unwrap()
                .to_owned(),
            anthropic["messages"]
                .as_array()
                .unwrap()
                .iter()
                .flat_map(|message| message["content"].as_array().unwrap())
                .find(|block| block["type"] == "tool_result" && block["tool_use_id"] == call_id)
                .unwrap()["content"]
                .as_str()
                .unwrap()
                .to_owned(),
        ]
    }

    #[test]
    fn serializer_duplicate_call_ids_cannot_remove_skill_body_witnesses() {
        let body = "Use scoped reads and preserve the complete user request.\n".repeat(30);
        let raw = vec![
            read_call("first"),
            skill_result("first", &body),
            read_call("reread"),
            skill_result("reread", &body),
        ];
        let request = ProviderRequest::from_conversation("model", raw.clone());
        for output in serialized_outputs(&request, "reread") {
            assert!(output.contains("[Unchanged skill body:"));
            assert!(output.contains("tool call \"first\""));
            assert!(!output.contains(&body));
        }

        let collisions = [
            ProviderConversationItem::ToolResult(ProviderToolResult {
                call_id: "first".into(),
                tool_name: "read".into(),
                success: true,
                output: "Earlier output without the skill body".into(),
                skill_reads: Vec::new(),
            }),
            ProviderConversationItem::ResponseItem(json!({
                "type": "function_call_output", "call_id": "first",
                "output": "Earlier output without the skill body"
            })),
            ProviderConversationItem::ResponseItem(json!({
                "type": "function_call", "call_id": "first", "name": "bash", "arguments": "{}"
            })),
            ProviderConversationItem::ResponseItem(json!({
                "role": "tool", "tool_call_id": "first", "content": "Legacy output"
            })),
            ProviderConversationItem::ResponseItem(json!({
                "role": "assistant", "content": " ", "tool_calls": [{
                    "id": "first", "type": "function", "function": {"name": "read", "arguments": "{}"}
                }]
            })),
            ProviderConversationItem::ResponseItem(json!({
                "type": "message", "role": "assistant", "content": " ", "tool_calls": [{
                    "id": "first", "type": "function", "function": {"name": "read", "arguments": "{}"}
                }]
            })),
        ];
        for collision in collisions {
            for insertion_index in [0, 2] {
                let mut items = raw.clone();
                items.insert(insertion_index, collision.clone());
                let request = ProviderRequest::from_conversation("model", items);
                for output in serialized_outputs(&request, "reread") {
                    assert_eq!(output, format!("--- FILE: skill://example ---\n{body}"));
                }
            }
        }
    }

    #[test]
    fn serializers_retain_bodies_with_missing_malformed_or_truncated_provenance() {
        let body = "é Preserve complete skill instructions and exact file bytes.\n".repeat(30);
        for malformed_index in [1, 3] {
            for case in [
                "missing",
                "version",
                "identity",
                "digest",
                "source",
                "reversed",
                "past_end",
                "utf8",
                "overlap",
                "truncated",
                "failed",
            ] {
                let mut items = vec![
                    read_call("first"),
                    skill_result("first", &body),
                    read_call("reread"),
                    skill_result("reread", &body),
                ];
                let ProviderConversationItem::ToolResult(result) = &mut items[malformed_index]
                else {
                    unreachable!();
                };
                match case {
                    "missing" => result.skill_reads.clear(),
                    "version" => result.skill_reads[0].version += 1,
                    "identity" => result.skill_reads[0].identity = "untrusted identity".into(),
                    "digest" => result.skill_reads[0].sha256 = "0".repeat(64),
                    "source" => result.skill_reads[0].source.push('\n'),
                    "reversed" => result.skill_reads[0].start = result.skill_reads[0].end + 1,
                    "past_end" => result.skill_reads[0].end = result.output.len() + 1,
                    "utf8" => result.skill_reads[0].start += 1,
                    "overlap" => result.skill_reads.push(result.skill_reads[0].clone()),
                    "truncated" => result.output.truncate(result.output.len() - 1),
                    "failed" => result.success = false,
                    _ => unreachable!(),
                }
                let ProviderConversationItem::ToolResult(reread) = &items[3] else {
                    unreachable!();
                };
                let expected = reread.output.clone();
                let request = ProviderRequest::from_conversation("model", items);
                for output in serialized_outputs(&request, "reread") {
                    assert_eq!(output, expected, "{case}, item {malformed_index}");
                }
            }
        }
    }

    #[test]
    fn snapshots_and_reasoning_resets_restore_raw_bodies_when_read_context_is_removed() {
        use crate::thinking::ThinkingLevel;
        use std::sync::Arc;

        let selection = |effort| ProviderConversationItem::ReasoningSelection {
            provider: "openai-codex".into(),
            model: "model".into(),
            effort,
        };
        let body = "Keep full skill instructions available after rebuilding context.\n".repeat(30);
        let raw = [
            selection(ThinkingLevel::High),
            read_call("first"),
            skill_result("first", &body),
            selection(ThinkingLevel::Default),
            selection(ThinkingLevel::Low),
            read_call("reread"),
            skill_result("reread", &body),
        ];
        let shared = ProviderRequest::from_shared_conversation(
            "model",
            Arc::from(raw[..4].to_vec()),
            &raw[4..],
        );
        for request in [
            shared.clone(),
            shared.to_owned_request(),
            shared.clone().with_reasoning_updates("openai-codex"),
            shared
                .with_reasoning_updates("openai-codex")
                .to_owned_request(),
        ] {
            let expected = format!("--- FILE: skill://example ---\n{body}");
            for output in serialized_outputs(&request, "first") {
                assert_eq!(output, expected);
            }
            for output in serialized_outputs(&request, "reread") {
                assert!(output.contains("tool call \"first\""));
                assert!(!output.contains(&body));
            }

            let without_calls = request.clone().with_response_items(Vec::new());
            let without_witness_call =
                without_calls
                    .clone()
                    .with_response_items(vec![match read_call("reread") {
                        ProviderConversationItem::ResponseItem(item) => item,
                        _ => unreachable!(),
                    }]);
            for removed in [without_calls, without_witness_call] {
                for output in serialized_outputs(&removed, "reread") {
                    assert_eq!(output, expected);
                }
                let restored = removed.with_response_items(
                    [read_call("first"), read_call("reread")]
                        .into_iter()
                        .map(|item| match item {
                            ProviderConversationItem::ResponseItem(item) => item,
                            _ => unreachable!(),
                        })
                        .collect(),
                );
                for output in serialized_outputs(&restored, "reread") {
                    assert!(output.contains("tool call \"first\""));
                    assert!(!output.contains(&body));
                }
            }
            let ProviderConversationItem::ToolResult(reread) = raw.last().unwrap() else {
                unreachable!();
            };
            let without_witness_result = request.with_tool_results(vec![reread.clone()]);
            for output in serialized_outputs(&without_witness_result, "reread") {
                assert_eq!(output, expected);
            }
        }
    }

    #[test]
    fn reassembled_projected_results_cannot_witness_another_read() {
        let body = "A previously emitted marker is not a retained skill body.\n".repeat(30);
        let request = ProviderRequest::from_conversation(
            "model",
            vec![
                read_call("first"),
                skill_result("first", &body),
                read_call("second"),
                skill_result("second", &body),
            ],
        );
        let mut items = request
            .conversation_items_iter()
            .skip(2)
            .cloned()
            .collect::<Vec<_>>();
        items.extend([read_call("reread"), skill_result("reread", &body)]);
        let rebuilt = ProviderRequest::from_conversation("model", items);
        for output in serialized_outputs(&rebuilt, "reread") {
            assert_eq!(output, format!("--- FILE: skill://example ---\n{body}"));
        }
    }

    #[test]
    fn serializers_keep_large_mixed_reads_raw_when_overlay_budget_is_exhausted() {
        let body = "Preserve the unrelated file contents in a mixed skill read.\n".repeat(30);
        let suffix = format!("\n--- FILE: large.txt ---\n{}", "x".repeat(600 * 1024));
        let mixed = |id| {
            let ProviderConversationItem::ToolResult(mut result) = skill_result(id, &body) else {
                unreachable!();
            };
            result.output.push_str(&suffix);
            ProviderConversationItem::ToolResult(result)
        };
        let request = ProviderRequest::from_conversation(
            "model",
            vec![
                read_call("first"),
                skill_result("first", &body),
                read_call("mixed"),
                mixed("mixed"),
                read_call("reread"),
                mixed("reread"),
                read_call("small"),
                skill_result("small", &body),
            ],
        );
        for output in serialized_outputs(&request, "mixed") {
            assert!(output.contains("tool call \"first\""));
            assert!(output.ends_with(&suffix));
            assert!(!output.contains(&body));
        }
        for output in serialized_outputs(&request, "reread") {
            assert_eq!(
                output,
                format!("--- FILE: skill://example ---\n{body}{suffix}")
            );
        }
        for output in serialized_outputs(&request, "small") {
            assert!(output.contains("tool call \"first\""));
            assert!(!output.contains(&body));
        }
    }
}