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