Skip to main content

vtcode_core/tools/
continuation.rs

1use serde_json::{Value, json};
2
3pub const NEXT_CONTINUE_PROMPT: &str = "Still running — check again for more output.";
4pub const NEXT_READ_PROMPT: &str = "More output available — read the next chunk.";
5pub const DEFAULT_NEXT_READ_LIMIT: usize = 40;
6
7const SESSION_ID_KEY: &str = "session_id";
8const COMPACT_SESSION_ID_KEY: &str = "s";
9const PATH_KEY: &str = "path";
10const COMPACT_PATH_KEY: &str = "p";
11const OFFSET_KEY: &str = "offset";
12const COMPACT_OFFSET_KEY: &str = "o";
13const LIMIT_KEY: &str = "limit";
14const COMPACT_LIMIT_KEY: &str = "l";
15const OFFSET_BYTES_KEY: &str = "offset_bytes";
16const PAGE_SIZE_BYTES_KEY: &str = "page_size_bytes";
17
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub struct PtyContinuationArgs {
20    pub session_id: String,
21}
22
23impl PtyContinuationArgs {
24    pub fn new(session_id: impl Into<String>) -> Self {
25        Self { session_id: session_id.into() }
26    }
27
28    pub fn from_value(value: &Value) -> Option<Self> {
29        value
30            .get(SESSION_ID_KEY)
31            .or_else(|| value.get(COMPACT_SESSION_ID_KEY))
32            .and_then(Value::as_str)
33            .map(Self::new)
34    }
35
36    pub fn to_value(&self) -> Value {
37        json!({ SESSION_ID_KEY: self.session_id })
38    }
39
40    pub fn to_compact_value(&self) -> Value {
41        json!({ COMPACT_SESSION_ID_KEY: self.session_id })
42    }
43}
44
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub struct ReadChunkContinuationArgs {
47    pub path: String,
48    pub offset: usize,
49    pub limit: usize,
50    pub offset_bytes: Option<u64>,
51    pub page_size_bytes: Option<usize>,
52}
53
54impl ReadChunkContinuationArgs {
55    pub fn new(path: impl Into<String>, offset: usize, limit: usize) -> Self {
56        Self {
57            path: path.into(),
58            offset: offset.max(1),
59            limit: limit.max(1),
60            offset_bytes: None,
61            page_size_bytes: None,
62        }
63    }
64
65    /// Create byte-based continuation args.
66    pub fn new_byte_range(path: impl Into<String>, offset_bytes: u64, page_size_bytes: usize) -> Self {
67        Self {
68            path: path.into(),
69            offset: 1,
70            limit: DEFAULT_NEXT_READ_LIMIT,
71            offset_bytes: Some(offset_bytes),
72            page_size_bytes: Some(page_size_bytes),
73        }
74    }
75
76    pub fn from_value(value: &Value) -> Option<Self> {
77        let path = value
78            .get(PATH_KEY)
79            .or_else(|| value.get(COMPACT_PATH_KEY))
80            .and_then(Value::as_str)?
81            .to_string();
82
83        // Byte-based continuation
84        let offset_bytes = value.get(OFFSET_BYTES_KEY).and_then(Value::as_u64);
85        let page_size_bytes = value.get(PAGE_SIZE_BYTES_KEY).and_then(value_to_usize);
86
87        if offset_bytes.is_some() || page_size_bytes.is_some() {
88            return Some(Self {
89                path,
90                offset: 1,
91                limit: DEFAULT_NEXT_READ_LIMIT,
92                offset_bytes,
93                page_size_bytes,
94            });
95        }
96
97        // Line-based continuation
98        let offset = value
99            .get(OFFSET_KEY)
100            .or_else(|| value.get(COMPACT_OFFSET_KEY))
101            .and_then(value_to_usize)?
102            .max(1);
103        let limit = value
104            .get(LIMIT_KEY)
105            .or_else(|| value.get(COMPACT_LIMIT_KEY))
106            .and_then(value_to_usize)
107            .unwrap_or(DEFAULT_NEXT_READ_LIMIT)
108            .max(1);
109        Some(Self {
110            path,
111            offset,
112            limit,
113            offset_bytes: None,
114            page_size_bytes: None,
115        })
116    }
117
118    pub fn to_value(&self) -> Value {
119        self.serialize_inner(PATH_KEY, OFFSET_KEY, LIMIT_KEY)
120    }
121
122    pub fn to_compact_value(&self) -> Value {
123        self.serialize_inner(COMPACT_PATH_KEY, COMPACT_OFFSET_KEY, COMPACT_LIMIT_KEY)
124    }
125
126    fn serialize_inner(&self, path_key: &str, offset_key: &str, limit_key: &str) -> Value {
127        if self.offset_bytes.is_some() || self.page_size_bytes.is_some() {
128            let mut map = json!({ path_key: self.path });
129            if let Some(ob) = self.offset_bytes {
130                map[OFFSET_BYTES_KEY] = json!(ob);
131            }
132            if let Some(ps) = self.page_size_bytes {
133                map[PAGE_SIZE_BYTES_KEY] = json!(ps);
134            }
135            return map;
136        }
137        json!({
138            path_key: self.path,
139            offset_key: self.offset,
140            limit_key: self.limit
141        })
142    }
143}
144
145pub fn read_chunk_progress_from_result(result: &Value) -> Option<(usize, usize)> {
146    result
147        .get("next_read_args")
148        .and_then(ReadChunkContinuationArgs::from_value)
149        .map(|next_read_args| (next_read_args.offset, next_read_args.limit))
150}
151
152fn value_to_usize(value: &Value) -> Option<usize> {
153    value
154        .as_u64()
155        .and_then(|n| usize::try_from(n).ok())
156        .or_else(|| value.as_str().and_then(|s| s.parse::<usize>().ok()))
157}
158
159#[cfg(test)]
160mod tests {
161    use super::{PtyContinuationArgs, ReadChunkContinuationArgs, read_chunk_progress_from_result};
162    use serde_json::json;
163
164    #[test]
165    fn pty_continuation_round_trips() {
166        let args = PtyContinuationArgs::new("run-123");
167        let payload = args.to_value();
168        let parsed = PtyContinuationArgs::from_value(&payload).unwrap();
169
170        assert_eq!(parsed.session_id, "run-123");
171    }
172
173    #[test]
174    fn pty_continuation_accepts_compact_form() {
175        let parsed = PtyContinuationArgs::from_value(&json!({
176            "s": "run-123"
177        }))
178        .unwrap();
179
180        assert_eq!(parsed.session_id, "run-123");
181    }
182
183    #[test]
184    fn read_chunk_continuation_round_trips() {
185        let args = ReadChunkContinuationArgs::new("out.txt", 41, 40);
186        let payload = args.to_value();
187        let parsed = ReadChunkContinuationArgs::from_value(&payload).unwrap();
188
189        assert_eq!(parsed.path, "out.txt");
190        assert_eq!(parsed.offset, 41);
191        assert_eq!(parsed.limit, 40);
192    }
193
194    #[test]
195    fn read_chunk_continuation_accepts_string_numbers() {
196        let parsed = ReadChunkContinuationArgs::from_value(&json!({
197            "path": "out.txt",
198            "offset": "2",
199            "limit": "3"
200        }))
201        .unwrap();
202
203        assert_eq!(parsed.offset, 2);
204        assert_eq!(parsed.limit, 3);
205    }
206
207    #[test]
208    fn read_chunk_continuation_accepts_compact_form() {
209        let parsed = ReadChunkContinuationArgs::from_value(&json!({
210            "p": "out.txt",
211            "o": 2,
212            "l": 3
213        }))
214        .unwrap();
215
216        assert_eq!(parsed.path, "out.txt");
217        assert_eq!(parsed.offset, 2);
218        assert_eq!(parsed.limit, 3);
219    }
220
221    #[test]
222    fn read_chunk_progress_reads_canonical_args() {
223        let result = json!({
224            "next_read_args": {
225                "path": "out.txt",
226                "offset": 81,
227                "limit": 40
228            }
229        });
230        assert_eq!(read_chunk_progress_from_result(&result), Some((81, 40)));
231    }
232
233    #[test]
234    fn read_chunk_progress_reads_compact_args() {
235        let result = json!({
236            "next_read_args": {
237                "p": "out.txt",
238                "o": 81,
239                "l": 40
240            }
241        });
242        assert_eq!(read_chunk_progress_from_result(&result), Some((81, 40)));
243    }
244
245    #[test]
246    fn read_chunk_progress_requires_canonical_args() {
247        let result = json!({
248            "next_offset": "10",
249            "chunk_limit": "0"
250        });
251        assert_eq!(read_chunk_progress_from_result(&result), None);
252    }
253
254    #[test]
255    fn byte_range_continuation_round_trips() {
256        let args = ReadChunkContinuationArgs::new_byte_range("big.bin", 8192, 4096);
257        let payload = args.to_value();
258        let parsed = ReadChunkContinuationArgs::from_value(&payload).unwrap();
259
260        assert_eq!(parsed.path, "big.bin");
261        assert_eq!(parsed.offset_bytes, Some(8192));
262        assert_eq!(parsed.page_size_bytes, Some(4096));
263    }
264
265    #[test]
266    fn byte_range_continuation_from_value() {
267        let parsed = ReadChunkContinuationArgs::from_value(&json!({
268            "path": "data.log",
269            "offset_bytes": 1024,
270            "page_size_bytes": 2048
271        }))
272        .unwrap();
273
274        assert_eq!(parsed.path, "data.log");
275        assert_eq!(parsed.offset_bytes, Some(1024));
276        assert_eq!(parsed.page_size_bytes, Some(2048));
277    }
278
279    #[test]
280    fn byte_range_continuation_to_value_includes_byte_fields() {
281        let args = ReadChunkContinuationArgs::new_byte_range("out.bin", 0, 8192);
282        let value = args.to_value();
283
284        assert_eq!(value["path"], "out.bin");
285        assert_eq!(value["offset_bytes"], 0);
286        assert_eq!(value["page_size_bytes"], 8192);
287        // Line-based fields should not be present
288        assert!(value.get("offset").is_none());
289        assert!(value.get("limit").is_none());
290    }
291
292    #[test]
293    fn line_based_continuation_excludes_byte_fields() {
294        let args = ReadChunkContinuationArgs::new("out.txt", 41, 40);
295        let value = args.to_value();
296
297        assert_eq!(value["path"], "out.txt");
298        assert_eq!(value["offset"], 41);
299        assert_eq!(value["limit"], 40);
300        assert!(value.get("offset_bytes").is_none());
301        assert!(value.get("page_size_bytes").is_none());
302    }
303}