Skip to main content

vtcode_core/tools/
apply_patch.rs

1//! Patch tool facade that exposes Codex-compatible patch parsing and application.
2//!
3//! Actual patch parsing logic lives in `tools::editing::patch` so future edit
4//! features can reuse the same primitives without depending on this facade.
5
6use anyhow::Context;
7use base64::Engine;
8use base64::engine::general_purpose::STANDARD as BASE64;
9use serde::{Deserialize, Serialize};
10use serde_json::Value;
11
12pub use crate::tools::editing::{Patch, PatchError, PatchHunk, PatchLine, PatchOperation};
13pub use vtcode_utility_tool_specs::{
14    APPLY_PATCH_ALIAS_DESCRIPTION, SEMANTIC_ANCHOR_GUIDANCE, with_semantic_anchor_guidance,
15};
16
17/// Input structure for the apply_patch tool
18#[derive(Debug, Deserialize, Serialize)]
19pub struct ApplyPatchInput {
20    pub input: String,
21}
22
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub struct DecodedApplyPatchInput {
25    pub text: String,
26    pub source_bytes: usize,
27    pub was_base64: bool,
28}
29
30pub fn patch_source_from_args(args: &Value) -> Option<&str> {
31    args.as_str()
32        .or_else(|| args.get("input").and_then(|value| value.as_str()))
33        .or_else(|| args.get("patch").and_then(|value| value.as_str()))
34}
35
36pub fn decode_apply_patch_input(args: &Value) -> anyhow::Result<Option<DecodedApplyPatchInput>> {
37    let Some(source) = patch_source_from_args(args) else {
38        return Ok(None);
39    };
40
41    let was_base64 = source.starts_with("base64:");
42    let cap = effective_max_payload_bytes();
43    let text = if was_base64 {
44        let decoded = BASE64
45            .decode(&source[7..])
46            .with_context(|| "Failed to decode base64 patch")?;
47        enforce_decoded_size_limit(decoded.len(), source.len(), was_base64, cap)?;
48        String::from_utf8(decoded).with_context(|| "Decoded patch is not valid UTF-8")?
49    } else {
50        enforce_decoded_size_limit(source.len(), source.len(), was_base64, cap)?;
51        source.to_string()
52    };
53
54    Ok(Some(DecodedApplyPatchInput {
55        text,
56        source_bytes: source.len(),
57        was_base64,
58    }))
59}
60
61fn enforce_decoded_size_limit(
62    decoded_bytes: usize,
63    source_bytes: usize,
64    was_base64: bool,
65    cap: usize,
66) -> anyhow::Result<()> {
67    // Cap applies to the *decoded* size to prevent base64 decompression-bomb-style attacks
68    // where a small source expands to a very large patch.
69    if decoded_bytes <= cap {
70        return Ok(());
71    }
72    anyhow::bail!(
73        "apply_patch payload too large after decoding: decoded={} bytes (source={} bytes, base64={}). \
74         The per-patch cap is {} bytes; split the change into smaller patches or raise {}.",
75        decoded_bytes,
76        source_bytes,
77        was_base64,
78        cap,
79        UNIFIED_FILE_MAX_PAYLOAD_BYTES_ENV
80    )
81}
82
83/// Hard upper bound for any single `apply_patch` payload (and `unified_file` `edit`/`patch`
84/// actions) after base64 decoding. The preflight cap mirrors this same value; both
85/// share the same env-var override.
86pub const UNIFIED_FILE_MAX_PAYLOAD_BYTES: usize = 1024 * 1024;
87
88/// Maximum allowed decoded patch size — same as `UNIFIED_FILE_MAX_PAYLOAD_BYTES` but
89/// exposed with a more specific name to clarify that it is enforced at decode time,
90/// not preflight time.
91pub const MAX_DECODED_PATCH_BYTES: usize = UNIFIED_FILE_MAX_PAYLOAD_BYTES;
92
93/// Env var name that overrides both the preflight cap and the post-decode cap.
94pub const UNIFIED_FILE_MAX_PAYLOAD_BYTES_ENV: &str = "VTCODE_UNIFIED_FILE_MAX_PAYLOAD_BYTES";
95
96/// Resolve the effective cap, honoring the env-var override. A 1 KiB safety
97/// floor is enforced so a sub-floor override can never silently disable the
98/// post-decode cap; values below the floor fall back to the default. The same
99/// floor is applied by the preflight resolver in `execution_kernel` so both
100/// stages agree on the effective cap.
101pub fn effective_max_payload_bytes() -> usize {
102    std::env::var(UNIFIED_FILE_MAX_PAYLOAD_BYTES_ENV)
103        .ok()
104        .and_then(|v| v.trim().parse::<usize>().ok())
105        .filter(|value| *value >= 1024)
106        .unwrap_or(UNIFIED_FILE_MAX_PAYLOAD_BYTES)
107}
108
109pub fn parameter_schema(input_description: &str) -> Value {
110    vtcode_utility_tool_specs::apply_patch_parameter_schema(input_description)
111}
112
113#[cfg(test)]
114mod tests {
115    use super::{
116        APPLY_PATCH_ALIAS_DESCRIPTION, SEMANTIC_ANCHOR_GUIDANCE, decode_apply_patch_input,
117        parameter_schema, patch_source_from_args, with_semantic_anchor_guidance,
118    };
119    use serde_json::json;
120
121    #[test]
122    fn patch_source_accepts_raw_string_and_object_fields() {
123        assert_eq!(
124            patch_source_from_args(&json!("*** Begin Patch\n*** End Patch\n")),
125            Some("*** Begin Patch\n*** End Patch\n")
126        );
127        assert_eq!(patch_source_from_args(&json!({"input": "x"})), Some("x"));
128        assert_eq!(patch_source_from_args(&json!({"patch": "y"})), Some("y"));
129    }
130
131    #[test]
132    fn decode_apply_patch_input_supports_base64_payloads() {
133        let payload = json!({
134            "patch": "base64:KioqIEJlZ2luIFBhdGNoCioqKiBFbmQgUGF0Y2gK"
135        });
136
137        let decoded = decode_apply_patch_input(&payload)
138            .expect("payload should decode")
139            .expect("payload should be present");
140
141        assert_eq!(decoded.text, "*** Begin Patch\n*** End Patch\n");
142        assert_eq!(decoded.source_bytes, 47);
143        assert!(decoded.was_base64);
144    }
145
146    #[test]
147    fn decode_apply_patch_input_rejects_invalid_base64() {
148        let error = decode_apply_patch_input(&json!({"input": "base64:not-valid"}))
149            .expect_err("invalid base64 should fail");
150
151        assert!(error.to_string().contains("Failed to decode base64 patch"));
152    }
153
154    #[test]
155    fn decode_apply_patch_input_caps_decoded_size() {
156        use base64::Engine;
157        // Build a 1.5 MiB decoded payload via base64. The default cap is 1 MiB so the
158        // decode must fail with the post-decode size error rather than producing the
159        // oversized text.
160        let large = "A".repeat(1_500_000);
161        let encoded = base64::engine::general_purpose::STANDARD.encode(large.as_bytes());
162        let payload = json!({ "input": format!("base64:{encoded}") });
163
164        let error = decode_apply_patch_input(&payload)
165            .expect_err("oversized decoded payload should be rejected");
166        let message = error.to_string();
167        assert!(
168            message.contains("apply_patch payload too large after decoding"),
169            "unexpected error message: {message}"
170        );
171        assert!(message.contains("base64=true"));
172    }
173
174    #[test]
175    fn decode_apply_patch_input_caps_raw_payload_size() {
176        // Even non-base64 inputs must respect the cap. Build a 1.5 MiB raw string
177        // and confirm it is rejected.
178        let large = "B".repeat(1_500_000);
179        let payload = json!({ "input": large });
180
181        let error = decode_apply_patch_input(&payload)
182            .expect_err("oversized raw payload should be rejected");
183        let message = error.to_string();
184        assert!(
185            message.contains("apply_patch payload too large after decoding"),
186            "unexpected error message: {message}"
187        );
188        assert!(message.contains("base64=false"));
189    }
190
191    #[test]
192    fn semantic_anchor_guidance_is_appended_once() {
193        let base = "Patch in VT Code format.";
194        let with_guidance = with_semantic_anchor_guidance(base);
195        assert!(with_guidance.contains(SEMANTIC_ANCHOR_GUIDANCE));
196        assert_eq!(
197            with_semantic_anchor_guidance(&with_guidance),
198            with_guidance,
199            "guidance should not be duplicated"
200        );
201    }
202
203    #[test]
204    fn parameter_schema_keeps_alias_and_guidance_consistent() {
205        let schema = parameter_schema("Patch in VT Code format");
206
207        assert_eq!(
208            schema["properties"]["patch"]["description"],
209            APPLY_PATCH_ALIAS_DESCRIPTION
210        );
211        let input_description = schema["properties"]["input"]["description"]
212            .as_str()
213            .expect("input description");
214        assert!(input_description.contains(SEMANTIC_ANCHOR_GUIDANCE));
215    }
216}