use std::path::Path;
use lingshu_types::ToolError;
pub const DEFAULT_MAX_MUTATION_PAYLOAD_BYTES: usize = 32 * 1024;
pub const DEFAULT_MAX_MUTATION_PAYLOAD_KIB: usize = DEFAULT_MAX_MUTATION_PAYLOAD_BYTES / 1024;
pub const MIN_MUTATION_PAYLOAD_BYTES: usize = 8 * 1024;
pub const MAX_MUTATION_PAYLOAD_BYTES_CAP: usize = 256 * 1024;
pub const MAX_MUTATION_PAYLOAD_BYTES: usize = DEFAULT_MAX_MUTATION_PAYLOAD_BYTES;
pub const MAX_MUTATION_PAYLOAD_KIB: usize = DEFAULT_MAX_MUTATION_PAYLOAD_KIB;
pub fn clamp_write_limit_bytes(kib: usize) -> usize {
let bytes = kib.saturating_mul(1024);
bytes.clamp(MIN_MUTATION_PAYLOAD_BYTES, MAX_MUTATION_PAYLOAD_BYTES_CAP)
}
pub fn enforce_write_payload_limit(
tool_name: &str,
path: &str,
resolved: &Path,
content: &str,
) -> Result<(), ToolError> {
enforce_write_payload_limit_with_max(
tool_name,
path,
resolved,
content,
DEFAULT_MAX_MUTATION_PAYLOAD_BYTES,
)
}
pub fn enforce_write_payload_limit_with_max(
tool_name: &str,
path: &str,
resolved: &Path,
content: &str,
max_bytes: usize,
) -> Result<(), ToolError> {
let bytes = content.len();
if bytes <= max_bytes {
return Ok(());
}
let target_kind = if resolved.exists() {
"overwrite"
} else {
"creation"
};
let _ = target_kind;
Err(crate::recovery_catalog::mutation_payload_too_large(
tool_name,
path,
bytes,
max_bytes,
!resolved.exists(),
))
}
pub fn enforce_patch_payload_limit(
tool_name: &str,
path: &str,
payload_bytes: usize,
) -> Result<(), ToolError> {
enforce_patch_payload_limit_with_max(
tool_name,
path,
payload_bytes,
DEFAULT_MAX_MUTATION_PAYLOAD_BYTES,
)
}
pub fn enforce_patch_payload_limit_with_max(
tool_name: &str,
path: &str,
payload_bytes: usize,
max_bytes: usize,
) -> Result<(), ToolError> {
if payload_bytes <= max_bytes {
return Ok(());
}
Err(crate::recovery_catalog::mutation_payload_too_large(
tool_name,
path,
payload_bytes,
max_bytes,
false,
))
}
pub fn enforce_apply_patch_payload_limit(patch_text: &str) -> Result<(), ToolError> {
enforce_apply_patch_payload_limit_with_max(patch_text, DEFAULT_MAX_MUTATION_PAYLOAD_BYTES)
}
pub fn enforce_apply_patch_payload_limit_with_max(
patch_text: &str,
max_bytes: usize,
) -> Result<(), ToolError> {
let bytes = patch_text.len();
if bytes <= max_bytes {
return Ok(());
}
Err(crate::recovery_catalog::mutation_payload_too_large(
"apply_patch",
"(multi-file patch)",
bytes,
max_bytes,
false,
))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_clamp_write_limit_bytes() {
assert_eq!(clamp_write_limit_bytes(4), MIN_MUTATION_PAYLOAD_BYTES);
assert_eq!(clamp_write_limit_bytes(32), 32 * 1024);
assert_eq!(clamp_write_limit_bytes(64), 64 * 1024);
assert_eq!(clamp_write_limit_bytes(512), MAX_MUTATION_PAYLOAD_BYTES_CAP);
assert_eq!(clamp_write_limit_bytes(0), MIN_MUTATION_PAYLOAD_BYTES);
}
#[test]
fn test_enforce_write_limit_configurable() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test.txt");
std::fs::write(&path, "existing").unwrap();
let content = "x".repeat(10 * 1024);
assert!(
enforce_write_payload_limit_with_max(
"write_file",
"test.txt",
&path,
&content,
16 * 1024,
)
.is_ok()
);
let big = "x".repeat(20 * 1024);
let err =
enforce_write_payload_limit_with_max("write_file", "test.txt", &path, &big, 16 * 1024);
assert!(err.is_err());
let msg = format!("{}", err.unwrap_err());
assert!(msg.contains("16 KiB"));
}
#[test]
fn test_enforce_patch_limit_configurable() {
assert!(enforce_patch_payload_limit_with_max("patch", "f.rs", 10_000, 16 * 1024).is_ok());
assert!(enforce_patch_payload_limit_with_max("patch", "f.rs", 20_000, 16 * 1024).is_err());
}
#[test]
fn test_enforce_apply_patch_limit_configurable() {
let small = "x".repeat(10_000);
assert!(enforce_apply_patch_payload_limit_with_max(&small, 16 * 1024).is_ok());
let big = "x".repeat(20_000);
assert!(enforce_apply_patch_payload_limit_with_max(&big, 16 * 1024).is_err());
}
}