1use std::fs::{File, OpenOptions};
26use std::io::{BufWriter, Write};
27use std::path::{Path, PathBuf};
28use std::sync::{Arc, Mutex};
29
30use serde::{Deserialize, Serialize};
31use serde_json::Value;
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
35#[serde(rename_all = "snake_case")]
36pub enum ToolAuditStatus {
37 Success,
39 Failure,
41 Timeout,
43 Cancelled,
45 Blocked,
47}
48
49impl ToolAuditStatus {
50 #[must_use]
52 pub fn as_str(self) -> &'static str {
53 match self {
54 Self::Success => "success",
55 Self::Failure => "failure",
56 Self::Timeout => "timeout",
57 Self::Cancelled => "cancelled",
58 Self::Blocked => "blocked",
59 }
60 }
61}
62
63#[derive(Debug, Clone, Serialize, Deserialize)]
69pub struct ToolAuditEntry {
70 timestamp_unix_ms: u64,
72 session_id: String,
74 turn_id: String,
76 tool_call_id: String,
78 tool_name: String,
80 arguments_hash: String,
82 #[serde(skip_serializing_if = "Option::is_none")]
84 arguments_redacted: Option<Value>,
85 result_hash: String,
87 #[serde(skip_serializing_if = "Option::is_none")]
89 result_summary: Option<String>,
90 duration_ms: u64,
92 status: ToolAuditStatus,
94 #[serde(skip_serializing_if = "Option::is_none")]
97 sandbox_policy: Option<String>,
98 #[serde(skip_serializing_if = "Option::is_none")]
100 transport: Option<String>,
101 #[serde(skip_serializing_if = "Option::is_none")]
103 server_address: Option<String>,
104 #[serde(skip_serializing_if = "Option::is_none")]
106 server_port: Option<u16>,
107 #[serde(skip_serializing_if = "Option::is_none")]
109 model_id: Option<String>,
110 prompt_injection_flagged: bool,
113 #[serde(skip_serializing_if = "Option::is_none")]
115 reason: Option<String>,
116}
117
118pub trait ToolAuditSink: Send + Sync {
120 fn write(&self, entry: &ToolAuditEntry);
123
124 fn flush(&self) {}
126}
127
128#[derive(Debug, Default, Clone, Copy)]
130pub struct NullSink;
131
132impl ToolAuditSink for NullSink {
133 fn write(&self, _entry: &ToolAuditEntry) {}
134}
135
136#[derive(Debug, Default, Clone)]
138pub struct InMemorySink {
139 entries: Arc<Mutex<Vec<ToolAuditEntry>>>,
140}
141
142impl InMemorySink {
143 #[must_use]
145 fn new() -> Self {
146 Self::default()
147 }
148
149 fn entries(&self) -> Vec<ToolAuditEntry> {
151 self.entries.lock().expect("in-memory sink poisoned").clone()
152 }
153
154 fn len(&self) -> usize {
156 self.entries.lock().expect("in-memory sink poisoned").len()
157 }
158
159 fn is_empty(&self) -> bool {
161 self.len() == 0
162 }
163
164 fn clear(&self) {
166 self.entries.lock().expect("in-memory sink poisoned").clear();
167 }
168}
169
170impl ToolAuditSink for InMemorySink {
171 fn write(&self, entry: &ToolAuditEntry) {
172 self.entries.lock().expect("in-memory sink poisoned").push(entry.clone());
173 }
174}
175
176pub struct JsonlFileSink {
183 path: PathBuf,
184 max_size_bytes: u64,
185 max_files: usize,
186 state: Mutex<JsonlFileState>,
187}
188
189struct JsonlFileState {
190 writer: Option<BufWriter<File>>,
191 bytes_written: u64,
192}
193
194impl std::fmt::Debug for JsonlFileSink {
195 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
196 f.debug_struct("JsonlFileSink")
197 .field("path", &self.path)
198 .field("max_size_bytes", &self.max_size_bytes)
199 .field("max_files", &self.max_files)
200 .finish_non_exhaustive()
201 }
202}
203
204impl JsonlFileSink {
205 fn open(path: impl Into<PathBuf>, max_size_bytes: u64, max_files: usize) -> std::io::Result<Self> {
208 let path = path.into();
209 if let Some(parent) = path.parent() {
210 std::fs::create_dir_all(parent)?;
211 }
212 let file = OpenOptions::new().create(true).append(true).open(&path)?;
213 let bytes_written = file.metadata().map(|m| m.len()).unwrap_or(0);
214 Ok(Self {
215 path,
216 max_size_bytes,
217 max_files: max_files.max(1),
218 state: Mutex::new(JsonlFileState { writer: Some(BufWriter::new(file)), bytes_written }),
219 })
220 }
221
222 #[must_use]
224 pub fn path(&self) -> &Path {
225 &self.path
226 }
227
228 fn rotate_if_needed(
229 state: &mut JsonlFileState,
230 path: &Path,
231 max_size_bytes: u64,
232 max_files: usize,
233 ) -> std::io::Result<()> {
234 if state.bytes_written < max_size_bytes {
235 return Ok(());
236 }
237 if let Some(mut writer) = state.writer.take() {
239 let _ = writer.flush();
240 }
241 for index in (1..max_files).rev() {
243 let from = rotated_path(path, index);
244 let to = rotated_path(path, index + 1);
245 if from.exists() {
246 let _ = std::fs::rename(&from, &to);
247 }
248 }
249 if path.exists() {
250 std::fs::rename(path, rotated_path(path, 1))?;
251 }
252 let file = OpenOptions::new().create(true).append(true).open(path)?;
253 state.writer = Some(BufWriter::new(file));
254 state.bytes_written = 0;
255 Ok(())
256 }
257}
258
259fn rotated_path(path: &Path, index: usize) -> PathBuf {
260 let mut s = path.as_os_str().to_owned();
261 s.push(format!(".{index}"));
262 PathBuf::from(s)
263}
264
265impl ToolAuditSink for JsonlFileSink {
266 fn write(&self, entry: &ToolAuditEntry) {
267 let mut state = match self.state.lock() {
268 Ok(state) => state,
269 Err(poisoned) => poisoned.into_inner(),
270 };
271
272 let serialized = match serde_json::to_string(entry) {
273 Ok(serialized) => serialized,
274 Err(err) => {
275 tracing::warn!(error = %err, "JsonlFileSink: failed to serialize audit entry");
276 return;
277 }
278 };
279 let line_length = serialized.len() as u64 + 1; if let Err(err) = Self::rotate_if_needed(&mut state, &self.path, self.max_size_bytes, self.max_files) {
282 tracing::warn!(error = %err, path = %self.path.display(), "JsonlFileSink: rotation failed");
283 }
284 if let Some(writer) = state.writer.as_mut() {
285 if let Err(err) = writeln!(writer, "{serialized}") {
286 tracing::warn!(error = %err, path = %self.path.display(), "JsonlFileSink: write failed");
287 return;
288 }
289 state.bytes_written = state.bytes_written.saturating_add(line_length);
290 }
291 }
292
293 fn flush(&self) {
294 let mut state = match self.state.lock() {
295 Ok(state) => state,
296 Err(poisoned) => poisoned.into_inner(),
297 };
298 if let Some(writer) = state.writer.as_mut() {
299 let _ = writer.flush();
300 }
301 }
302}
303
304impl Drop for JsonlFileSink {
305 fn drop(&mut self) {
306 self.flush();
307 }
308}
309
310pub struct MultiSink {
315 inner: Vec<Arc<dyn ToolAuditSink>>,
316}
317
318impl std::fmt::Debug for MultiSink {
319 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
320 f.debug_struct("MultiSink").field("sink_count", &self.inner.len()).finish()
321 }
322}
323
324impl MultiSink {
325 #[must_use]
329 fn new(sinks: Vec<Arc<dyn ToolAuditSink>>) -> Self {
330 Self { inner: sinks }
331 }
332
333 #[must_use]
335 pub fn len(&self) -> usize {
336 self.inner.len()
337 }
338
339 pub fn is_empty(&self) -> bool {
341 self.inner.is_empty()
342 }
343}
344
345impl ToolAuditSink for MultiSink {
346 fn write(&self, entry: &ToolAuditEntry) {
347 for sink in &self.inner {
348 sink.write(entry);
349 }
350 }
351
352 fn flush(&self) {
353 for sink in &self.inner {
354 sink.flush();
355 }
356 }
357}
358
359#[derive(Clone)]
364pub struct ToolAuditLogger {
365 sink: Arc<dyn ToolAuditSink>,
366}
367
368impl std::fmt::Debug for ToolAuditLogger {
369 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
370 f.debug_struct("ToolAuditLogger").finish_non_exhaustive()
371 }
372}
373
374impl ToolAuditLogger {
375 #[must_use]
377 fn new(sink: Arc<dyn ToolAuditSink>) -> Self {
378 Self { sink }
379 }
380
381 #[must_use]
383 fn disabled() -> Self {
384 Self::new(Arc::new(NullSink))
385 }
386
387 fn record(&self, entry: ToolAuditEntry) {
389 self.sink.write(&entry);
390 }
391
392 fn flush(&self) {
394 self.sink.flush();
395 }
396
397 #[must_use]
399 pub fn sink(&self) -> &Arc<dyn ToolAuditSink> {
400 &self.sink
401 }
402}
403
404impl Default for ToolAuditLogger {
405 fn default() -> Self {
406 Self::disabled()
407 }
408}
409
410#[must_use]
415fn sha256_hex(bytes: &[u8]) -> String {
416 use sha2::{Digest, Sha256};
417 let mut hasher = Sha256::new();
418 hasher.update(bytes);
419 let digest = hasher.finalize();
420 let mut out = String::with_capacity(digest.len() * 2);
421 for byte in digest {
422 use std::fmt::Write;
423 let _ = write!(&mut out, "{byte:02x}");
424 }
425 out
426}
427
428#[cfg(test)]
429mod tests {
430 use super::*;
431 use tempfile::TempDir;
432
433 fn sample_entry(suffix: &str) -> ToolAuditEntry {
434 ToolAuditEntry {
435 timestamp_unix_ms: 1_700_000_000_000 + u64::from(suffix.bytes().next().unwrap_or(b'a')),
436 session_id: format!("session-{suffix}"),
437 turn_id: format!("turn-{suffix}"),
438 tool_call_id: format!("call-{suffix}"),
439 tool_name: "mcp::fetch::fetch".to_owned(),
440 arguments_hash: sha256_hex(suffix.as_bytes()),
441 arguments_redacted: None,
442 result_hash: sha256_hex(format!("result-{suffix}").as_bytes()),
443 result_summary: Some(format!("first line of result {suffix}")),
444 duration_ms: 42,
445 status: ToolAuditStatus::Success,
446 sandbox_policy: None,
447 transport: Some("stdio".to_owned()),
448 server_address: None,
449 server_port: None,
450 model_id: Some("test-model".to_owned()),
451 prompt_injection_flagged: false,
452 reason: None,
453 }
454 }
455
456 #[test]
457 fn in_memory_sink_records_entries() {
458 let sink = InMemorySink::new();
459 sink.write(&sample_entry("a"));
460 sink.write(&sample_entry("b"));
461 assert_eq!(sink.len(), 2);
462 assert_eq!(sink.entries()[0].tool_name, "mcp::fetch::fetch");
463 sink.clear();
464 assert!(sink.is_empty());
465 }
466
467 #[test]
468 fn null_sink_accepts_without_recording() {
469 let sink = NullSink;
470 sink.write(&sample_entry("x"));
471 }
473
474 #[test]
475 fn jsonl_file_sink_appends_and_flushes_on_drop() {
476 let dir = TempDir::new().expect("tempdir");
477 let path = dir.path().join("audit.jsonl");
478 let sink = JsonlFileSink::open(&path, 1024 * 1024, 4).expect("open sink");
479
480 sink.write(&sample_entry("a"));
481 sink.write(&sample_entry("b"));
482 sink.flush();
483
484 let body = std::fs::read_to_string(&path).expect("read back");
485 let lines: Vec<&str> = body.lines().collect();
486 assert_eq!(lines.len(), 2);
487 for line in lines {
488 let value: Value = serde_json::from_str(line).expect("line is valid JSON");
489 assert_eq!(value["tool_name"], "mcp::fetch::fetch");
490 }
491 }
492
493 #[test]
494 fn jsonl_file_sink_rotates_when_threshold_exceeded() {
495 let dir = TempDir::new().expect("tempdir");
496 let path = dir.path().join("audit.jsonl");
497 let sink = JsonlFileSink::open(&path, 60, 3).expect("open sink");
499
500 sink.write(&sample_entry("a"));
501 sink.flush();
502 sink.write(&sample_entry("b"));
503 sink.flush();
504
505 let active = std::fs::read_to_string(&path).expect("active");
508 assert!(active.contains("\"call-b\""), "expected rotated active file to contain call-b, got: {active}");
509 let rotated = std::fs::read_to_string(dir.path().join("audit.jsonl.1")).expect("rotated");
510 assert!(rotated.contains("\"call-a\""), "expected rotated file to contain call-a, got: {rotated}");
511 }
512
513 #[test]
514 fn multi_sink_forwards_to_every_inner_sink() {
515 let a = Arc::new(InMemorySink::new());
516 let b = Arc::new(InMemorySink::new());
517 let multi = MultiSink::new(vec![a.clone(), b.clone()]);
518 multi.write(&sample_entry("z"));
519 assert_eq!(a.len(), 1);
520 assert_eq!(b.len(), 1);
521 }
522
523 #[test]
524 fn tool_audit_logger_record_routes_through_sink() {
525 let sink = Arc::new(InMemorySink::new());
526 let logger = ToolAuditLogger::new(sink.clone());
527 logger.record(sample_entry("k"));
528 logger.flush();
529 assert_eq!(sink.len(), 1);
530 }
531
532 #[test]
533 fn sha256_hex_is_stable() {
534 assert_eq!(sha256_hex(b"hello"), "2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824");
535 assert_eq!(sha256_hex(b"hello"), sha256_hex(b"hello"));
536 }
537}