agtrace_providers/gemini/
parser.rs

1use agtrace_types::*;
2use anyhow::Result;
3use chrono::DateTime;
4use std::path::Path;
5use uuid::Uuid;
6
7use crate::builder::{EventBuilder, SemanticSuffix};
8use crate::gemini::schema::{GeminiMessage, GeminiSession};
9
10/// Normalize Gemini session to events
11/// Unfolds nested structure (thoughts[], toolCalls[]) into event stream
12pub(crate) fn normalize_gemini_session(
13    session: &GeminiSession,
14    raw_messages: Vec<serde_json::Value>,
15) -> Vec<AgentEvent> {
16    // Create session_id UUID from session_id string (deterministic)
17    let session_id_uuid = Uuid::new_v5(&Uuid::NAMESPACE_OID, session.session_id.as_bytes());
18    let mut builder = EventBuilder::new(session_id_uuid);
19    let mut events = Vec::new();
20
21    for (idx, msg) in session.messages.iter().enumerate() {
22        let raw_value = raw_messages
23            .get(idx)
24            .cloned()
25            .unwrap_or(serde_json::Value::Null);
26
27        match msg {
28            GeminiMessage::User(user_msg) => {
29                // Skip numeric IDs (legacy CLI events)
30                if user_msg.id.parse::<u32>().is_ok() {
31                    continue;
32                }
33
34                let timestamp = parse_timestamp(&user_msg.timestamp);
35                builder.build_and_push(
36                    &mut events,
37                    &user_msg.id,
38                    SemanticSuffix::User,
39                    timestamp,
40                    EventPayload::User(UserPayload {
41                        text: user_msg.content.clone(),
42                    }),
43                    Some(raw_value),
44                    StreamId::Main,
45                );
46            }
47
48            GeminiMessage::Gemini(gemini_msg) => {
49                let timestamp = parse_timestamp(&gemini_msg.timestamp);
50                let base_id = &gemini_msg.id;
51
52                // 1. Reasoning events (thoughts)
53                for (idx, thought) in gemini_msg.thoughts.iter().enumerate() {
54                    let indexed_base_id = format!("{}-thought-{}", base_id, idx);
55                    builder.build_and_push(
56                        &mut events,
57                        &indexed_base_id,
58                        SemanticSuffix::Reasoning,
59                        timestamp,
60                        EventPayload::Reasoning(ReasoningPayload {
61                            text: format!("{}: {}", thought.subject, thought.description),
62                        }),
63                        Some(raw_value.clone()),
64                        StreamId::Main,
65                    );
66                }
67
68                // 2. Tool calls and results
69                for (idx, tool_call) in gemini_msg.tool_calls.iter().enumerate() {
70                    let indexed_base_id = format!("{}-tool-{}", base_id, idx);
71
72                    // ToolCall event
73                    let tool_call_uuid = builder.build_and_push(
74                        &mut events,
75                        &indexed_base_id,
76                        SemanticSuffix::ToolCall,
77                        timestamp,
78                        EventPayload::ToolCall(super::mapper::normalize_gemini_tool_call(
79                            tool_call.name.clone(),
80                            tool_call.args.clone(),
81                            Some(tool_call.id.clone()),
82                        )),
83                        Some(raw_value.clone()),
84                        StreamId::Main,
85                    );
86
87                    // Register tool call ID mapping (provider ID -> UUID)
88                    builder.register_tool_call(tool_call.id.clone(), tool_call_uuid);
89
90                    // ToolResult event (if result exists)
91                    if !tool_call.result.is_empty() {
92                        let output = tool_call
93                            .result_display
94                            .clone()
95                            .unwrap_or_else(|| format!("{:?}", tool_call.result));
96
97                        let is_error = tool_call
98                            .status
99                            .as_ref()
100                            .map(|s| s == "error")
101                            .unwrap_or(false);
102
103                        builder.build_and_push(
104                            &mut events,
105                            &indexed_base_id,
106                            SemanticSuffix::ToolResult,
107                            timestamp,
108                            EventPayload::ToolResult(ToolResultPayload {
109                                output,
110                                tool_call_id: tool_call_uuid, // Reference to ToolCall UUID
111                                is_error,
112                            }),
113                            Some(raw_value.clone()),
114                            StreamId::Main,
115                        );
116                    }
117                }
118
119                // 3. Message event (assistant response)
120                builder.build_and_push(
121                    &mut events,
122                    base_id,
123                    SemanticSuffix::Message,
124                    timestamp,
125                    EventPayload::Message(MessagePayload {
126                        text: gemini_msg.content.clone(),
127                    }),
128                    Some(raw_value.clone()),
129                    StreamId::Main,
130                );
131
132                // 4. TokenUsage event (sidecar attached to message)
133                // Gemini returns turn-level totals, so we attach to the last generation event
134                builder.build_and_push(
135                    &mut events,
136                    base_id,
137                    SemanticSuffix::TokenUsage,
138                    timestamp,
139                    EventPayload::TokenUsage(TokenUsagePayload {
140                        input_tokens: gemini_msg.tokens.input as i32,
141                        output_tokens: gemini_msg.tokens.output as i32,
142                        total_tokens: gemini_msg.tokens.total as i32,
143                        details: Some(TokenUsageDetails {
144                            cache_creation_input_tokens: None, // Gemini doesn't track cache creation separately
145                            cache_read_input_tokens: Some(gemini_msg.tokens.cached as i32),
146                            reasoning_output_tokens: Some(gemini_msg.tokens.thoughts as i32),
147                        }),
148                    }),
149                    Some(raw_value),
150                    StreamId::Main,
151                );
152            }
153
154            GeminiMessage::Info(info_msg) => {
155                let timestamp = parse_timestamp(&info_msg.timestamp);
156                builder.build_and_push(
157                    &mut events,
158                    &info_msg.id,
159                    SemanticSuffix::Notification,
160                    timestamp,
161                    EventPayload::Notification(NotificationPayload {
162                        text: info_msg.content.clone(),
163                        level: Some("info".to_string()),
164                    }),
165                    Some(raw_value),
166                    StreamId::Main,
167                );
168            }
169        }
170    }
171
172    events
173}
174
175/// Parse Gemini timestamp to DateTime<Utc>
176fn parse_timestamp(ts: &str) -> DateTime<chrono::Utc> {
177    DateTime::parse_from_rfc3339(ts)
178        .map(|dt| dt.with_timezone(&chrono::Utc))
179        .unwrap_or_else(|_| chrono::Utc::now())
180}
181
182/// Gemini session parser implementation
183pub struct GeminiParser;
184
185impl crate::traits::SessionParser for GeminiParser {
186    fn parse_file(&self, path: &Path) -> Result<Vec<AgentEvent>> {
187        super::io::normalize_gemini_file(path)
188    }
189
190    fn parse_record(&self, content: &str) -> Result<Option<AgentEvent>> {
191        // Gemini uses JSON format (not JSONL), parse as AgentEvent
192        match serde_json::from_str::<AgentEvent>(content) {
193            Ok(event) => Ok(Some(event)),
194            Err(_) => Ok(None), // Skip malformed lines
195        }
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202    use crate::gemini::schema::{GeminiAssistantMessage, TokenUsage, UserMessage};
203
204    #[test]
205    fn test_normalize_user_message() {
206        let session = GeminiSession {
207            session_id: "test-session".to_string(),
208            project_hash: agtrace_types::ProjectHash::from("test-hash"),
209            start_time: "2024-01-01T00:00:00Z".to_string(),
210            last_updated: "2024-01-01T00:00:00Z".to_string(),
211            messages: vec![GeminiMessage::User(UserMessage {
212                id: "uuid-123".to_string(),
213                timestamp: "2024-01-01T00:00:00Z".to_string(),
214                content: "Hello".to_string(),
215            })],
216        };
217
218        let events = normalize_gemini_session(&session, vec![]);
219        assert_eq!(events.len(), 1);
220
221        match &events[0].payload {
222            EventPayload::User(payload) => assert_eq!(payload.text, "Hello"),
223            _ => panic!("Expected User payload"),
224        }
225        assert_eq!(events[0].parent_id, None);
226    }
227
228    #[test]
229    fn test_normalize_assistant_with_tokens() {
230        let session = GeminiSession {
231            session_id: "test-session".to_string(),
232            project_hash: agtrace_types::ProjectHash::from("test-hash"),
233            start_time: "2024-01-01T00:00:00Z".to_string(),
234            last_updated: "2024-01-01T00:00:00Z".to_string(),
235            messages: vec![GeminiMessage::Gemini(GeminiAssistantMessage {
236                id: "uuid-456".to_string(),
237                timestamp: "2024-01-01T00:00:01Z".to_string(),
238                content: "Hello back!".to_string(),
239                model: "gemini-2.0-flash".to_string(),
240                thoughts: vec![],
241                tool_calls: vec![],
242                tokens: TokenUsage {
243                    input: 10,
244                    output: 5,
245                    total: 15,
246                    cached: 2,
247                    thoughts: 1,
248                    tool: 0,
249                },
250            })],
251        };
252
253        let events = normalize_gemini_session(&session, vec![]);
254        // Should have: Message + TokenUsage (2 events)
255        assert_eq!(events.len(), 2);
256
257        match &events[0].payload {
258            EventPayload::Message(payload) => assert_eq!(payload.text, "Hello back!"),
259            _ => panic!("Expected Message payload"),
260        }
261
262        match &events[1].payload {
263            EventPayload::TokenUsage(payload) => {
264                assert_eq!(payload.input_tokens, 10);
265                assert_eq!(payload.output_tokens, 5);
266                assert_eq!(payload.total_tokens, 15);
267                assert_eq!(
268                    payload.details.as_ref().unwrap().cache_read_input_tokens,
269                    Some(2)
270                );
271                assert_eq!(
272                    payload.details.as_ref().unwrap().reasoning_output_tokens,
273                    Some(1)
274                );
275            }
276            _ => panic!("Expected TokenUsage payload"),
277        }
278    }
279}
agtrace_providers/gemini/parser.rs

agtrace_providers/gemini/
parser.rs
