Skip to main content

tracevault_core/
streaming.rs

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