1use std::sync::Arc;
4
5use tokio::sync::Mutex;
6use tracing::instrument;
7
8use crate::errors::{Result, WinxError};
9use crate::state::bash_state::BashState;
10use crate::tools::file_write_or_edit::{commit_edit, plan_revision_edit};
11use crate::types::{normalize_thread_id, ApplyPatch, LinePatch};
12
13const MAX_PATCHES: usize = 256;
14
15#[derive(Debug)]
16pub struct ApplyPatchOutcome {
17 pub text: String,
18 pub revision: String,
19}
20
21#[instrument(level = "info", skip(bash_state_arc, request))]
22pub async fn handle_tool_call(
23 bash_state_arc: &Arc<Mutex<Option<BashState>>>,
24 request: ApplyPatch,
25) -> Result<ApplyPatchOutcome> {
26 validate_revision(&request.expected_revision)?;
27 validate_patch_count(&request.patches)?;
28
29 let mut bash_state_guard = bash_state_arc.lock().await;
30 {
31 let state = bash_state_guard.as_ref().ok_or(WinxError::BashStateNotInitialized)?;
32 let thread_id = normalize_thread_id(&request.thread_id);
33 if thread_id != state.current_thread_id {
34 return Err(WinxError::ThreadIdMismatch(thread_id));
35 }
36 }
37
38 let mut state = bash_state_guard.take().ok_or(WinxError::BashStateNotInitialized)?;
39 let recovery_state = state.clone();
40 let joined = tokio::task::spawn_blocking(move || {
41 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
42 let required_ranges = required_read_ranges(&request.patches);
43 let planned = plan_revision_edit(
44 &state,
45 &request.file_path,
46 &request.expected_revision,
47 &required_ranges,
48 |content| apply_line_patches(content, &request.patches),
49 )?;
50 let revision = planned.new_revision();
51 let text = commit_edit(&mut state, planned)?;
52 Ok(ApplyPatchOutcome { text, revision })
53 }))
54 .unwrap_or_else(|_| {
55 Err(WinxError::CommandExecutionError(
56 "ApplyPatch panicked on the blocking worker".to_string(),
57 ))
58 });
59 (state, result)
60 })
61 .await;
62
63 match joined {
64 Ok((state, result)) => {
65 *bash_state_guard = Some(state);
66 result
67 }
68 Err(error) => {
69 *bash_state_guard = Some(recovery_state);
70 Err(WinxError::CommandExecutionError(format!(
71 "ApplyPatch blocking task failed: {error}"
72 )))
73 }
74 }
75}
76
77fn validate_revision(revision: &str) -> Result<()> {
78 let hash = revision.strip_prefix("sha256:").unwrap_or_default();
79 if hash.len() != 64 || !hash.bytes().all(|byte| byte.is_ascii_hexdigit()) {
80 return Err(WinxError::ParameterValidationError {
81 field: "expected_revision".to_string(),
82 message: "must be the exact sha256:<64 lowercase hex> token returned by ReadFiles"
83 .to_string(),
84 });
85 }
86 if hash.bytes().any(|byte| byte.is_ascii_uppercase()) {
87 return Err(WinxError::ParameterValidationError {
88 field: "expected_revision".to_string(),
89 message: "revision tokens are case-sensitive; copy the ReadFiles value exactly"
90 .to_string(),
91 });
92 }
93 Ok(())
94}
95
96fn validate_patch_count(patches: &[LinePatch]) -> Result<()> {
97 if patches.is_empty() || patches.len() > MAX_PATCHES {
98 return Err(WinxError::ParameterValidationError {
99 field: "patches".to_string(),
100 message: format!("must contain between 1 and {MAX_PATCHES} patches"),
101 });
102 }
103 Ok(())
104}
105
106fn required_read_ranges(patches: &[LinePatch]) -> Vec<(usize, usize)> {
107 patches
108 .iter()
109 .map(|patch| {
110 if patch.delete_lines == 0 {
111 (patch.start_line, patch.start_line)
112 } else {
113 (
114 patch.start_line,
115 patch.start_line.saturating_add(patch.delete_lines).saturating_sub(1),
116 )
117 }
118 })
119 .collect()
120}
121
122fn apply_line_patches(content: &str, patches: &[LinePatch]) -> Result<String> {
123 let spans = line_spans(content);
124 validate_patch_coordinates(patches, spans.len())?;
125 let mut output = content.to_string();
126 for patch in patches.iter().rev() {
127 let start_index = patch.start_line - 1;
128 let byte_start = spans.get(start_index).map_or(content.len(), |(start, _)| *start);
129 let byte_end = if patch.delete_lines == 0 {
130 byte_start
131 } else {
132 spans[start_index + patch.delete_lines - 1].1
133 };
134 output.replace_range(byte_start..byte_end, &patch.replacement);
135 }
136 Ok(output)
137}
138
139fn validate_patch_coordinates(patches: &[LinePatch], total_lines: usize) -> Result<()> {
140 let mut previous_start = 0;
141 let mut previous_end = 0;
142 for (index, patch) in patches.iter().enumerate() {
143 if patch.start_line == 0 || patch.start_line > total_lines.saturating_add(1) {
144 return invalid_patch(index, format!("start_line must be in 1..={}", total_lines + 1));
145 }
146 let start = patch.start_line - 1;
147 let end = start.checked_add(patch.delete_lines).ok_or_else(|| {
148 WinxError::ParameterValidationError {
149 field: format!("patches[{index}].delete_lines"),
150 message: "line range overflow".to_string(),
151 }
152 })?;
153 if end > total_lines {
154 return invalid_patch(index, "delete_lines extends past the original file".to_string());
155 }
156 if index > 0 && (start < previous_end || start == previous_start) {
157 return invalid_patch(
158 index,
159 "patches must be strictly ordered and non-overlapping in the original revision"
160 .to_string(),
161 );
162 }
163 previous_start = start;
164 previous_end = end;
165 }
166 Ok(())
167}
168
169fn invalid_patch<T>(index: usize, message: String) -> Result<T> {
170 Err(WinxError::ParameterValidationError { field: format!("patches[{index}]"), message })
171}
172
173fn line_spans(content: &str) -> Vec<(usize, usize)> {
174 let mut spans = Vec::new();
175 let mut start = 0;
176 for (index, byte) in content.bytes().enumerate() {
177 if byte == b'\n' {
178 spans.push((start, index + 1));
179 start = index + 1;
180 }
181 }
182 if start < content.len() {
183 spans.push((start, content.len()));
184 }
185 spans
186}
187
188#[cfg(test)]
189mod tests {
190 use std::sync::Arc;
191
192 use super::*;
193 use tokio::sync::Mutex;
194
195 use crate::errors::ReadRequirement;
196 use crate::types::{Modes, ReadFiles};
197
198 fn patch(start_line: usize, delete_lines: usize, replacement: &str) -> LinePatch {
199 LinePatch { start_line, delete_lines, replacement: replacement.to_string() }
200 }
201
202 #[test]
203 fn applies_multiple_original_revision_patches_from_the_end() -> Result<()> {
204 let output = apply_line_patches(
205 "one\ntwo\nthree\nfour\n",
206 &[patch(2, 1, "TWO\n"), patch(4, 0, "before four\n")],
207 )?;
208 assert_eq!(output, "one\nTWO\nthree\nbefore four\nfour\n");
209 Ok(())
210 }
211
212 #[test]
213 fn supports_append_and_empty_file_insert() -> Result<()> {
214 assert_eq!(apply_line_patches("one\n", &[patch(2, 0, "two\n")])?, "one\ntwo\n");
215 assert_eq!(apply_line_patches("", &[patch(1, 0, "first\n")])?, "first\n");
216 Ok(())
217 }
218
219 #[test]
220 fn rejects_overlapping_or_out_of_order_coordinates() {
221 assert!(apply_line_patches("one\ntwo\nthree\n", &[patch(2, 2, "x\n"), patch(3, 1, "y\n")])
222 .is_err());
223 assert!(apply_line_patches("one\n", &[patch(3, 0, "x")]).is_err());
224 }
225
226 fn state_for(root: &std::path::Path, thread_id: &str) -> Arc<Mutex<Option<BashState>>> {
227 let mut state = BashState::new();
228 state.initialized = true;
229 state.cwd = root.to_path_buf();
230 state.workspace_root = root.to_path_buf();
231 state.current_thread_id = normalize_thread_id(thread_id);
232 state.mode = Modes::Wcgw;
233 Arc::new(Mutex::new(Some(state)))
234 }
235
236 #[tokio::test]
237 async fn revision_patch_uses_only_visible_lines_and_stale_retry_is_safe() -> Result<()> {
238 let temp = tempfile::tempdir()?;
239 let path = temp.path().join("sample.txt");
240 std::fs::write(&path, "one\ntwo\nthree\n")?;
241 let state = state_for(temp.path(), "patch-test");
242 let read = crate::tools::read_files::handle_tool_call_detailed(
243 &state,
244 ReadFiles {
245 file_paths: vec![path.to_string_lossy().into_owned()],
246 thread_id: "patch-test".to_string(),
247 start_line_nums: vec![Some(2)],
248 end_line_nums: vec![Some(2)],
249 },
250 )
251 .await?;
252 let revision = read.files[0].revision.clone();
253
254 let outcome = handle_tool_call(
255 &state,
256 ApplyPatch {
257 file_path: path.to_string_lossy().into_owned(),
258 expected_revision: revision.clone(),
259 patches: vec![patch(2, 1, "TWO\n")],
260 thread_id: "patch-test".to_string(),
261 },
262 )
263 .await?;
264 assert_ne!(outcome.revision, revision);
265 assert_eq!(std::fs::read_to_string(&path)?, "one\nTWO\nthree\n");
266 assert!(state
267 .lock()
268 .await
269 .as_ref()
270 .is_some_and(|state| state.whitelist_for_overwrite.is_empty()));
271
272 let stale_result = handle_tool_call(
273 &state,
274 ApplyPatch {
275 file_path: path.to_string_lossy().into_owned(),
276 expected_revision: revision,
277 patches: vec![patch(2, 1, "WRONG\n")],
278 thread_id: "patch-test".to_string(),
279 },
280 )
281 .await;
282 assert!(matches!(stale_result, Err(WinxError::FileRevisionMismatch { .. })));
283 assert_eq!(std::fs::read_to_string(&path)?, "one\nTWO\nthree\n");
284 Ok(())
285 }
286
287 #[tokio::test]
288 async fn patch_cannot_touch_an_unread_line() -> Result<()> {
289 let temp = tempfile::tempdir()?;
290 let path = temp.path().join("sample.txt");
291 std::fs::write(&path, "one\ntwo\nthree\n")?;
292 let state = state_for(temp.path(), "coverage-test");
293 let read = crate::tools::read_files::handle_tool_call_detailed(
294 &state,
295 ReadFiles {
296 file_paths: vec![path.to_string_lossy().into_owned()],
297 thread_id: "coverage-test".to_string(),
298 start_line_nums: vec![Some(2)],
299 end_line_nums: vec![Some(2)],
300 },
301 )
302 .await?;
303 let patch_result = handle_tool_call(
304 &state,
305 ApplyPatch {
306 file_path: path.to_string_lossy().into_owned(),
307 expected_revision: read.files[0].revision.clone(),
308 patches: vec![patch(3, 1, "THREE\n")],
309 thread_id: "coverage-test".to_string(),
310 },
311 )
312 .await;
313 assert!(matches!(
314 patch_result,
315 Err(WinxError::FileReadRequired { reason: ReadRequirement::InsufficientCoverage, .. })
316 ));
317 assert_eq!(std::fs::read_to_string(path)?, "one\ntwo\nthree\n");
318 Ok(())
319 }
320}