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_input: Option<serde_json::Value>,
25    pub tool_response: Option<serde_json::Value>,
26    pub event_index: Option<i32>,
27    pub transcript_lines: Option<Vec<serde_json::Value>>,
28    pub transcript_offset: Option<i64>,
29    pub model: Option<String>,
30    pub cwd: Option<String>,
31    pub final_stats: Option<SessionFinalStats>,
32}
33
34#[derive(Debug, Clone, Serialize, Deserialize)]
35pub struct SessionFinalStats {
36    pub duration_ms: Option<i64>,
37    pub total_tokens: Option<i64>,
38    pub input_tokens: Option<i64>,
39    pub output_tokens: Option<i64>,
40    pub cache_read_tokens: Option<i64>,
41    pub cache_write_tokens: Option<i64>,
42    pub user_messages: Option<i32>,
43    pub assistant_messages: Option<i32>,
44    pub total_tool_calls: Option<i32>,
45}
46
47#[derive(Debug, Clone, Serialize, Deserialize)]
48pub struct StreamEventResponse {
49    pub session_db_id: uuid::Uuid,
50    pub event_db_id: Option<uuid::Uuid>,
51    pub status: String,
52}
53
54#[derive(Debug, Clone, Serialize, Deserialize)]
55pub struct CommitPushRequest {
56    pub commit_sha: String,
57    pub branch: Option<String>,
58    pub author: String,
59    pub message: Option<String>,
60    pub diff_data: Option<serde_json::Value>,
61    pub committed_at: Option<DateTime<Utc>>,
62}
63
64#[derive(Debug, Clone, Serialize, Deserialize)]
65pub struct CommitPushResponse {
66    pub commit_db_id: uuid::Uuid,
67    pub attributions_count: i64,
68}
69
70#[derive(Debug, Clone)]
71pub struct ExtractedFileChange {
72    pub file_path: String,
73    pub change_type: String,
74    pub diff_text: Option<String>,
75    pub content_hash: Option<String>,
76}
77
78pub fn is_file_modifying_tool(tool_name: &str) -> bool {
79    matches!(tool_name, "Write" | "Edit" | "Bash")
80}
81
82impl StreamEventRequest {
83    /// Drop optional fields largest-first until the serialized payload is
84    /// under 512 KB. Prevents 413 errors on both real-time sends and flush.
85    pub fn truncate_large_fields(&mut self) {
86        const MAX_BYTES: usize = 512 * 1024;
87        if serde_json::to_string(self).map(|s| s.len()).unwrap_or(0) <= MAX_BYTES {
88            return;
89        }
90        self.transcript_lines = None;
91        if serde_json::to_string(self).map(|s| s.len()).unwrap_or(0) <= MAX_BYTES {
92            return;
93        }
94        self.tool_response = None;
95        if serde_json::to_string(self).map(|s| s.len()).unwrap_or(0) <= MAX_BYTES {
96            return;
97        }
98        self.tool_input = None;
99    }
100}
101
102pub fn extract_file_change(
103    tool_name: &str,
104    tool_input: &serde_json::Value,
105) -> Option<ExtractedFileChange> {
106    match tool_name {
107        "Write" => {
108            let file_path = tool_input.get("file_path")?.as_str()?.to_string();
109            let content = tool_input.get("content")?.as_str()?;
110            let mut hasher = Sha256::new();
111            hasher.update(content.as_bytes());
112            let hash = hex::encode(hasher.finalize());
113            let diff = content
114                .lines()
115                .map(|l| format!("+{l}"))
116                .collect::<Vec<_>>()
117                .join("\n");
118            Some(ExtractedFileChange {
119                file_path,
120                change_type: "create".to_string(),
121                diff_text: Some(diff),
122                content_hash: Some(hash),
123            })
124        }
125        "Edit" => {
126            let file_path = tool_input.get("file_path")?.as_str()?.to_string();
127            let old_string = tool_input.get("old_string")?.as_str()?;
128            let new_string = tool_input.get("new_string")?.as_str()?;
129            let diff = format!("--- {old_string}\n+++ {new_string}");
130            Some(ExtractedFileChange {
131                file_path,
132                change_type: "edit".to_string(),
133                diff_text: Some(diff),
134                content_hash: None,
135            })
136        }
137        _ => None,
138    }
139}