Skip to main content

tracevault_core/
streaming.rs

1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3use sha2::{Digest, Sha256};
4
5#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
6#[serde(rename_all = "snake_case")]
7pub enum StreamEventType {
8    ToolUse,
9    Transcript,
10    SessionStart,
11    SessionEnd,
12    /// Client declares the start of a validation window. Only the most recent
13    /// event of this type per session is meaningful — re-opening is idempotent.
14    ValidationWindowStart,
15}
16
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct StreamEventRequest {
19    pub protocol_version: u32,
20    #[serde(default)]
21    pub tool: Option<String>,
22    pub event_type: StreamEventType,
23    pub session_id: String,
24    pub timestamp: DateTime<Utc>,
25    pub hook_event_name: Option<String>,
26    pub tool_name: Option<String>,
27    pub tool_use_id: Option<String>,
28    pub tool_input: Option<serde_json::Value>,
29    pub tool_response: Option<serde_json::Value>,
30    pub tool_is_error: Option<bool>,
31    pub event_index: Option<i32>,
32    pub transcript_lines: Option<Vec<serde_json::Value>>,
33    pub transcript_offset: Option<i64>,
34    pub model: Option<String>,
35    pub cwd: Option<String>,
36    pub final_stats: Option<SessionFinalStats>,
37}
38
39#[derive(Debug, Clone, Serialize, Deserialize)]
40pub struct SessionFinalStats {
41    pub duration_ms: Option<i64>,
42    pub total_tokens: Option<i64>,
43    pub input_tokens: Option<i64>,
44    pub output_tokens: Option<i64>,
45    pub cache_read_tokens: Option<i64>,
46    pub cache_write_tokens: Option<i64>,
47    pub user_messages: Option<i32>,
48    pub assistant_messages: Option<i32>,
49    pub total_tool_calls: Option<i32>,
50}
51
52#[derive(Debug, Clone, Serialize, Deserialize)]
53pub struct StreamEventResponse {
54    pub session_db_id: uuid::Uuid,
55    pub event_db_id: Option<uuid::Uuid>,
56    pub status: String,
57}
58
59#[derive(Debug, Clone, Serialize, Deserialize)]
60pub struct CommitPushRequest {
61    pub commit_sha: String,
62    pub branch: Option<String>,
63    pub author: String,
64    pub message: Option<String>,
65    pub diff_data: Option<serde_json::Value>,
66    pub committed_at: Option<DateTime<Utc>>,
67}
68
69#[derive(Debug, Clone, Serialize, Deserialize)]
70pub struct CommitPushResponse {
71    pub commit_db_id: uuid::Uuid,
72    pub attributions_count: i64,
73}
74
75#[derive(Debug, Clone)]
76pub struct ExtractedFileChange {
77    pub file_path: String,
78    pub change_type: String,
79    pub diff_text: Option<String>,
80    pub content_hash: Option<String>,
81}
82
83pub fn is_file_modifying_tool(tool_name: &str) -> bool {
84    matches!(tool_name, "Write" | "Edit" | "Bash")
85}
86
87impl StreamEventRequest {
88    /// Drop optional fields largest-first until the serialized payload is
89    /// under 512 KB. Prevents 413 errors on both real-time sends and flush.
90    pub fn truncate_large_fields(&mut self) {
91        const MAX_BYTES: usize = 512 * 1024;
92        if serde_json::to_string(self).map(|s| s.len()).unwrap_or(0) <= MAX_BYTES {
93            return;
94        }
95        self.transcript_lines = None;
96        if serde_json::to_string(self).map(|s| s.len()).unwrap_or(0) <= MAX_BYTES {
97            return;
98        }
99        self.tool_response = None;
100        if serde_json::to_string(self).map(|s| s.len()).unwrap_or(0) <= MAX_BYTES {
101            return;
102        }
103        self.tool_input = None;
104    }
105}
106
107/// Scan transcript lines for a tool_result whose tool_use_id matches the given id
108/// and return its is_error flag. Returns None if not found or no transcript available.
109pub fn extract_is_error_from_transcript(
110    tool_use_id: &str,
111    transcript_lines: &[serde_json::Value],
112) -> Option<bool> {
113    for line in transcript_lines {
114        // Transcript lines are full message objects; content is an array of blocks.
115        // Use explicit matching so a malformed line is skipped rather than aborting the scan.
116        let content = match line
117            .get("message")
118            .and_then(|m| m.get("content"))
119            .and_then(|c| c.as_array())
120        {
121            Some(c) => c,
122            None => continue,
123        };
124        for block in content {
125            let block_type = block.get("type").and_then(|v| v.as_str());
126            let block_uid = block.get("tool_use_id").and_then(|v| v.as_str());
127            if block_type == Some("tool_result") && block_uid == Some(tool_use_id) {
128                return block.get("is_error").and_then(|v| v.as_bool());
129            }
130        }
131    }
132    None
133}
134
135pub fn extract_file_change(
136    tool_name: &str,
137    tool_input: &serde_json::Value,
138) -> Option<ExtractedFileChange> {
139    match tool_name {
140        "Write" => {
141            let file_path = tool_input.get("file_path")?.as_str()?.to_string();
142            let content = tool_input.get("content")?.as_str()?;
143            let mut hasher = Sha256::new();
144            hasher.update(content.as_bytes());
145            let hash = hex::encode(hasher.finalize());
146            let diff = content
147                .lines()
148                .map(|l| format!("+{l}"))
149                .collect::<Vec<_>>()
150                .join("\n");
151            Some(ExtractedFileChange {
152                file_path,
153                change_type: "create".to_string(),
154                diff_text: Some(diff),
155                content_hash: Some(hash),
156            })
157        }
158        "Edit" => {
159            let file_path = tool_input.get("file_path")?.as_str()?.to_string();
160            let old_string = tool_input.get("old_string")?.as_str()?;
161            let new_string = tool_input.get("new_string")?.as_str()?;
162            let diff = format!("--- {old_string}\n+++ {new_string}");
163            Some(ExtractedFileChange {
164                file_path,
165                change_type: "edit".to_string(),
166                diff_text: Some(diff),
167                content_hash: None,
168            })
169        }
170        _ => None,
171    }
172}