magi-code 0.96.1

use crate::cancellation::AgentCancellation;
use serde_json::Value;
use std::io::{self, Write};

/// Count before allocating; serialization of discovery/results cannot grow an
/// unbounded temporary string outside the JavaScript allocator.
pub(crate) fn serialized_size(
    value: &Value,
    limit: usize,
    cancellation: &AgentCancellation,
) -> anyhow::Result<usize> {
    struct Counter<'a> {
        bytes: usize,
        limit: usize,
        cancellation: &'a AgentCancellation,
    }
    impl Write for Counter<'_> {
        fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
            if self.cancellation.is_canceled() {
                return Err(io::Error::other("Code Mode canceled during serialization"));
            }
            self.bytes = self.bytes.saturating_add(bytes.len());
            if self.bytes > self.limit {
                return Err(io::Error::other("Code Mode result limit exceeded"));
            }
            Ok(bytes.len())
        }
        fn flush(&mut self) -> io::Result<()> {
            Ok(())
        }
    }
    let mut counter = Counter {
        bytes: 0,
        limit,
        cancellation,
    };
    let serialized = serde_json::to_writer(&mut counter, value);
    cancellation.check()?;
    if counter.bytes > limit {
        return Err(super::ResourceLimitExceeded("result").into());
    }
    serialized?;
    Ok(counter.bytes)
}

/// Conservative storage reservation for serde_json strings, vectors, map nodes,
/// and parser scratch space. Punctuation inside strings is not a container node.
pub(super) fn parse_reservation(
    text: &str,
    cancellation: &AgentCancellation,
) -> anyhow::Result<usize> {
    let mut bytes = text.len().saturating_mul(2).saturating_add(4096);
    let mut quoted = false;
    let mut escaped = false;
    for (index, byte) in text.bytes().enumerate() {
        if index % 4096 == 0 {
            cancellation.check()?;
        }
        if quoted {
            if escaped {
                escaped = false;
            } else if byte == b'\\' {
                escaped = true;
            } else if byte == b'"' {
                quoted = false;
            }
        } else if byte == b'"' {
            quoted = true;
        } else if matches!(byte, b'[' | b'{' | b',' | b':') {
            bytes = bytes.saturating_add(256);
        }
    }
    Ok(bytes)
}

/// Conservative charge for native JSON trees before cloning SDK projections.
pub(crate) fn value_reservation(
    value: &Value,
    cancellation: &AgentCancellation,
) -> anyhow::Result<usize> {
    cancellation.check()?;
    let mut bytes = 256usize;
    match value {
        Value::String(text) => bytes = bytes.saturating_add(text.len().saturating_mul(2)),
        Value::Array(values) => {
            for value in values {
                bytes = bytes.saturating_add(value_reservation(value, cancellation)?);
            }
        }
        Value::Object(values) => {
            for (key, value) in values {
                bytes = bytes
                    .saturating_add(key.len().saturating_mul(2))
                    .saturating_add(value_reservation(value, cancellation)?);
            }
        }
        _ => {}
    }
    Ok(bytes)
}