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