Skip to main content

vtcode_core/tools/
continuation.rs

1use serde_json::{Value, json};
2
3pub const NEXT_CONTINUE_PROMPT: &str = "Reuse `next_continue_args`.";
4pub const NEXT_READ_PROMPT: &str = "Reuse `next_read_args`.";
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";
15
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub struct PtyContinuationArgs {
18    pub session_id: String,
19}
20
21impl PtyContinuationArgs {
22    pub fn new(session_id: impl Into<String>) -> Self {
23        Self {
24            session_id: session_id.into(),
25        }
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}
51
52impl ReadChunkContinuationArgs {
53    pub fn new(path: impl Into<String>, offset: usize, limit: usize) -> Self {
54        Self {
55            path: path.into(),
56            offset: offset.max(1),
57            limit: limit.max(1),
58        }
59    }
60
61    pub fn from_value(value: &Value) -> Option<Self> {
62        let path = value
63            .get(PATH_KEY)
64            .or_else(|| value.get(COMPACT_PATH_KEY))
65            .and_then(Value::as_str)?
66            .to_string();
67        let offset = value
68            .get(OFFSET_KEY)
69            .or_else(|| value.get(COMPACT_OFFSET_KEY))
70            .and_then(value_to_usize)?
71            .max(1);
72        let limit = value
73            .get(LIMIT_KEY)
74            .or_else(|| value.get(COMPACT_LIMIT_KEY))
75            .and_then(value_to_usize)
76            .unwrap_or(DEFAULT_NEXT_READ_LIMIT)
77            .max(1);
78        Some(Self {
79            path,
80            offset,
81            limit,
82        })
83    }
84
85    pub fn to_value(&self) -> Value {
86        json!({
87            PATH_KEY: self.path,
88            OFFSET_KEY: self.offset,
89            LIMIT_KEY: self.limit
90        })
91    }
92
93    pub fn to_compact_value(&self) -> Value {
94        json!({
95            COMPACT_PATH_KEY: self.path,
96            COMPACT_OFFSET_KEY: self.offset,
97            COMPACT_LIMIT_KEY: self.limit
98        })
99    }
100}
101
102pub fn read_chunk_progress_from_result(result: &Value) -> Option<(usize, usize)> {
103    result
104        .get("next_read_args")
105        .and_then(ReadChunkContinuationArgs::from_value)
106        .map(|next_read_args| (next_read_args.offset, next_read_args.limit))
107}
108
109fn value_to_usize(value: &Value) -> Option<usize> {
110    value
111        .as_u64()
112        .and_then(|n| usize::try_from(n).ok())
113        .or_else(|| value.as_str().and_then(|s| s.parse::<usize>().ok()))
114}
115
116#[cfg(test)]
117mod tests {
118    use super::{PtyContinuationArgs, ReadChunkContinuationArgs, read_chunk_progress_from_result};
119    use serde_json::json;
120
121    #[test]
122    fn pty_continuation_round_trips() {
123        let args = PtyContinuationArgs::new("run-123");
124        let payload = args.to_value();
125        let parsed = PtyContinuationArgs::from_value(&payload).unwrap();
126
127        assert_eq!(parsed.session_id, "run-123");
128    }
129
130    #[test]
131    fn pty_continuation_accepts_compact_form() {
132        let parsed = PtyContinuationArgs::from_value(&json!({
133            "s": "run-123"
134        }))
135        .unwrap();
136
137        assert_eq!(parsed.session_id, "run-123");
138    }
139
140    #[test]
141    fn read_chunk_continuation_round_trips() {
142        let args = ReadChunkContinuationArgs::new("out.txt", 41, 40);
143        let payload = args.to_value();
144        let parsed = ReadChunkContinuationArgs::from_value(&payload).unwrap();
145
146        assert_eq!(parsed.path, "out.txt");
147        assert_eq!(parsed.offset, 41);
148        assert_eq!(parsed.limit, 40);
149    }
150
151    #[test]
152    fn read_chunk_continuation_accepts_string_numbers() {
153        let parsed = ReadChunkContinuationArgs::from_value(&json!({
154            "path": "out.txt",
155            "offset": "2",
156            "limit": "3"
157        }))
158        .unwrap();
159
160        assert_eq!(parsed.offset, 2);
161        assert_eq!(parsed.limit, 3);
162    }
163
164    #[test]
165    fn read_chunk_continuation_accepts_compact_form() {
166        let parsed = ReadChunkContinuationArgs::from_value(&json!({
167            "p": "out.txt",
168            "o": 2,
169            "l": 3
170        }))
171        .unwrap();
172
173        assert_eq!(parsed.path, "out.txt");
174        assert_eq!(parsed.offset, 2);
175        assert_eq!(parsed.limit, 3);
176    }
177
178    #[test]
179    fn read_chunk_progress_reads_canonical_args() {
180        let result = json!({
181            "next_read_args": {
182                "path": "out.txt",
183                "offset": 81,
184                "limit": 40
185            }
186        });
187        assert_eq!(read_chunk_progress_from_result(&result), Some((81, 40)));
188    }
189
190    #[test]
191    fn read_chunk_progress_reads_compact_args() {
192        let result = json!({
193            "next_read_args": {
194                "p": "out.txt",
195                "o": 81,
196                "l": 40
197            }
198        });
199        assert_eq!(read_chunk_progress_from_result(&result), Some((81, 40)));
200    }
201
202    #[test]
203    fn read_chunk_progress_requires_canonical_args() {
204        let result = json!({
205            "next_offset": "10",
206            "chunk_limit": "0"
207        });
208        assert_eq!(read_chunk_progress_from_result(&result), None);
209    }
210}