vtcode_core/exec/
sdk_ipc.rs1use crate::utils::file_utils::{parse_json_with_context, read_file_with_context, write_file_with_context};
48use anyhow::{Context, Result};
49use std::path::PathBuf;
50use std::sync::Arc;
51use std::time::Duration;
52use tokio::fs;
53use tokio::time::sleep;
54use uuid::Uuid;
55
56use crate::tools::request_response::{ToolCallRequest, ToolCallResponse};
57
58const REQUEST_POLL_INTERVAL: Duration = Duration::from_millis(50);
59
60pub type ToolRequest = ToolCallRequest;
62
63pub type ToolResponse = ToolCallResponse;
65
66pub struct ToolIpcHandler {
68 ipc_dir: PathBuf,
69 pii_tokenizer: Option<Arc<crate::exec::PiiTokenizer>>,
70}
71
72impl ToolIpcHandler {
73 pub fn new(ipc_dir: PathBuf) -> Self {
75 Self { ipc_dir, pii_tokenizer: None }
76 }
77
78 pub fn with_pii_protection(ipc_dir: PathBuf) -> Result<Self> {
80 Ok(Self {
81 ipc_dir,
82 pii_tokenizer: Some(Arc::new(crate::exec::PiiTokenizer::new()?)),
83 })
84 }
85
86 pub fn enable_pii_protection(&mut self) -> Result<()> {
88 self.pii_tokenizer = Some(Arc::new(crate::exec::PiiTokenizer::new()?));
89 Ok(())
90 }
91
92 pub async fn read_request(&self) -> Result<Option<ToolRequest>> {
94 let request_file = self.ipc_dir.join("request.json");
95
96 if !fs::try_exists(&request_file).await.unwrap_or(false) {
97 return Ok(None);
98 }
99
100 let content = read_file_with_context(&request_file, "request file").await?;
101 let request: ToolRequest = parse_json_with_context(&content, "request JSON")?;
102
103 let _ = fs::remove_file(&request_file).await;
105
106 Ok(Some(request))
107 }
108
109 pub fn process_request_for_pii(&self, request: &mut ToolRequest) -> Result<()> {
111 if let Some(tokenizer) = &self.pii_tokenizer {
112 let args_str = serde_json::to_string(&request.args).context("failed to serialize request args")?;
113 let (tokenized, _) = tokenizer.tokenize_string(&args_str).context("PII tokenization failed")?;
114 request.args = parse_json_with_context(&tokenized, "tokenized args")?;
115 }
116 Ok(())
117 }
118
119 pub fn process_response_for_pii(&self, response: &mut ToolResponse) -> Result<()> {
121 if let Some(tokenizer) = &self.pii_tokenizer
122 && let Some(result) = &response.result
123 {
124 let result_str = serde_json::to_string(result).context("failed to serialize response result")?;
125 let detokenized = tokenizer.detokenize_string(&result_str).context("PII de-tokenization failed")?;
126 response.result = Some(parse_json_with_context(&detokenized, "de-tokenized result")?);
127 }
128 Ok(())
129 }
130
131 pub async fn write_response(&self, mut response: ToolResponse) -> Result<()> {
133 self.process_response_for_pii(&mut response)?;
135
136 let response_file = self.ipc_dir.join("response.json");
137
138 let json = serde_json::to_string(&response).context("failed to serialize response")?;
139
140 write_file_with_context(&response_file, &json, "response file").await?;
141
142 Ok(())
143 }
144
145 pub async fn wait_for_request(&self, timeout: Duration) -> Result<Option<ToolRequest>> {
147 let start = std::time::Instant::now();
148
149 loop {
150 if let Some(request) = self.read_request().await? {
151 return Ok(Some(request));
152 }
153
154 let Some(remaining_timeout) = timeout.checked_sub(start.elapsed()) else {
155 return Ok(None);
156 };
157
158 sleep(remaining_timeout.min(REQUEST_POLL_INTERVAL)).await;
159 }
160 }
161
162 pub fn new_request_id() -> String {
164 Uuid::new_v4().to_string()
165 }
166}
167
168#[cfg(test)]
169mod tests {
170 use super::*;
171 use serde_json::json;
172 use tempfile::tempdir;
173 use tokio::time::Instant;
174
175 #[test]
176 fn serialize_tool_request() {
177 let request = ToolRequest {
178 id: "test-id".into(),
179 tool_name: "read_file".into(),
180 args: json!({"path": "/test"}),
181 metadata: None,
182 };
183
184 let json = serde_json::to_string(&request).expect("ToolRequest should serialize");
185 assert!(json.contains("test-id"));
186 assert!(json.contains("read_file"));
187 }
188
189 #[test]
190 fn serialize_success_response() {
191 let response = ToolResponse {
192 id: "test-id".into(),
193 success: true,
194 result: Some(json!({"data": "test"})),
195 error: None,
196 duration_ms: None,
197 cache_hit: None,
198 };
199
200 let json = serde_json::to_string(&response).expect("ToolResponse should serialize");
201 assert!(json.contains("test-id"));
202 assert!(json.contains("true"));
203 assert!(!json.contains("error"));
204 }
205
206 #[test]
207 fn serialize_error_response() {
208 let response = ToolResponse {
209 id: "test-id".into(),
210 success: false,
211 result: None,
212 error: Some("File not found".into()),
213 duration_ms: None,
214 cache_hit: None,
215 };
216
217 let json = serde_json::to_string(&response).expect("ToolResponse should serialize");
218 assert!(json.contains("test-id"));
219 assert!(json.contains("false"));
220 assert!(json.contains("File not found"));
221 }
222
223 #[tokio::test]
224 async fn wait_for_request_reads_delayed_request() {
225 let temp_dir = tempdir().expect("temp dir should create");
226 let handler = ToolIpcHandler::new(temp_dir.path().to_path_buf());
227 let request = ToolRequest {
228 id: "test-id".into(),
229 tool_name: "read_file".into(),
230 args: json!({"path": "/tmp/test"}),
231 metadata: None,
232 };
233 let request_json = serde_json::to_string(&request).expect("request should serialize to JSON");
234 let request_path = temp_dir.path().join("request.json");
235
236 tokio::spawn(async move {
237 sleep(Duration::from_millis(10)).await;
238 fs::write(request_path, request_json).await.expect("request file should write");
239 });
240
241 let received = handler
242 .wait_for_request(Duration::from_millis(200))
243 .await
244 .expect("request wait should succeed");
245
246 assert_eq!(received.expect("request should arrive").id, "test-id");
247 }
248
249 #[tokio::test]
250 async fn wait_for_request_respects_short_timeout() {
251 let temp_dir = tempdir().expect("temp dir should create");
252 let handler = ToolIpcHandler::new(temp_dir.path().to_path_buf());
253 let start = Instant::now();
254
255 let received = handler
256 .wait_for_request(Duration::from_millis(5))
257 .await
258 .expect("request wait should succeed");
259
260 assert!(received.is_none());
261 assert!(start.elapsed() < Duration::from_millis(40));
262 }
263}