Skip to main content

vtcode_core/exec/
sdk_ipc.rs

1//! Inter-process communication for calling MCP tools from executing code.
2//!
3//! This module provides a file-based IPC mechanism that allows code running in
4//! a separate process to call MCP tools. The code writes tool requests to a file, and
5//! the executor reads and processes them, writing back results.
6//!
7//! Optionally supports PII (Personally Identifiable Information) protection by
8//! tokenizing sensitive data in requests before tool execution and de-tokenizing
9//! responses before returning to the code.
10//!
11//! # Protocol
12//!
13//! Requests (code → executor):
14//! ```json
15//! {
16//!   "id": "uuid",
17//!   "tool_name": "search_tools",
18//!   "args": {"keyword": "file"}
19//! }
20//! ```
21//!
22//! Responses (executor → code):
23//! ```json
24//! {
25//!   "id": "uuid",
26//!   "success": true,
27//!   "result": {...}
28//! }
29//! ```
30//! or
31//! ```json
32//! {
33//!   "id": "uuid",
34//!   "success": false,
35//!   "error": "Tool not found"
36//! }
37//! ```
38//!
39//! # PII Protection
40//!
41//! When enabled, the handler automatically:
42//! 1. Detects PII patterns in request arguments
43//! 2. Tokenizes sensitive data before tool execution
44//! 3. De-tokenizes responses before returning to code
45//! 4. Maintains token mapping for the session
46
47use 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
60/// IPC request from executing code to executor.
61pub type ToolRequest = ToolCallRequest;
62
63/// IPC response from executor to executing code.
64pub type ToolResponse = ToolCallResponse;
65
66/// IPC handler for tool invocation between code and executor.
67pub struct ToolIpcHandler {
68    ipc_dir: PathBuf,
69    pii_tokenizer: Option<Arc<crate::exec::PiiTokenizer>>,
70}
71
72impl ToolIpcHandler {
73    /// Create a new IPC handler with the given directory.
74    pub fn new(ipc_dir: PathBuf) -> Self {
75        Self { ipc_dir, pii_tokenizer: None }
76    }
77
78    /// Create a new IPC handler with PII protection enabled.
79    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    /// Enable PII protection on existing handler.
87    pub fn enable_pii_protection(&mut self) -> Result<()> {
88        self.pii_tokenizer = Some(Arc::new(crate::exec::PiiTokenizer::new()?));
89        Ok(())
90    }
91
92    /// Read a tool request from the code.
93    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        // Clean up request file
104        let _ = fs::remove_file(&request_file).await;
105
106        Ok(Some(request))
107    }
108
109    /// Process request for PII (tokenize if enabled).
110    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    /// Process response for PII (de-tokenize if enabled).
120    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    /// Write a tool response back to the code.
132    pub async fn write_response(&self, mut response: ToolResponse) -> Result<()> {
133        // De-tokenize response before writing back to code
134        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    /// Wait for a request with timeout.
146    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    /// Create a request ID.
163    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}