Skip to main content

agtrace_sdk/query/
session.rs

1//! Session query types for search, list, and get operations.
2
3use chrono::{DateTime, Utc};
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7
8use super::Provider;
9use super::filters::EventType;
10use crate::types::{AgentEvent, AgentSession, EventPayload, truncate};
11
12// ============================================================================
13// Cursor Pagination
14// ============================================================================
15
16#[derive(Debug, Serialize, Deserialize)]
17pub struct Cursor {
18    pub offset: usize,
19}
20
21impl Cursor {
22    pub fn encode(&self) -> String {
23        let json = serde_json::to_string(self).unwrap_or_default();
24        base64::Engine::encode(&base64::engine::general_purpose::STANDARD, json.as_bytes())
25    }
26
27    pub fn decode(cursor: &str) -> Option<Self> {
28        let bytes =
29            base64::Engine::decode(&base64::engine::general_purpose::STANDARD, cursor).ok()?;
30        let json = String::from_utf8(bytes).ok()?;
31        serde_json::from_str(&json).ok()
32    }
33}
34
35// ============================================================================
36// Search Events API
37// ============================================================================
38
39#[derive(Debug, Serialize, Deserialize, JsonSchema)]
40pub struct SearchEventsArgs {
41    pub query: String,
42    pub session_id: Option<String>,
43    #[serde(default)]
44    pub limit: Option<usize>,
45    #[serde(default)]
46    pub cursor: Option<String>,
47    pub provider: Option<Provider>,
48    pub event_type: Option<EventType>,
49    pub project_root: Option<String>,
50    pub project_hash: Option<String>,
51}
52
53impl SearchEventsArgs {
54    pub fn limit(&self) -> usize {
55        self.limit.unwrap_or(20).min(50)
56    }
57}
58
59#[derive(Debug, Serialize)]
60pub struct SearchEventsResponse {
61    pub matches: Vec<EventMatch>,
62    #[serde(skip_serializing_if = "Option::is_none")]
63    pub next_cursor: Option<String>,
64}
65
66#[derive(Debug, Serialize)]
67pub struct EventMatch {
68    pub session_id: String,
69    pub event_index: usize,
70    pub turn_index: usize,
71    pub step_index: usize,
72    pub event_type: EventType,
73    pub preview: String,
74    pub timestamp: DateTime<Utc>,
75}
76
77impl EventMatch {
78    pub fn new(
79        session_id: String,
80        event_index: usize,
81        turn_index: usize,
82        step_index: usize,
83        event: &AgentEvent,
84    ) -> Self {
85        let event_type = EventType::from_payload(&event.payload);
86        let preview = Self::extract_preview(&event.payload);
87
88        Self {
89            session_id,
90            event_index,
91            turn_index,
92            step_index,
93            event_type,
94            preview,
95            timestamp: event.timestamp,
96        }
97    }
98
99    fn extract_preview(payload: &EventPayload) -> String {
100        let text = match payload {
101            EventPayload::ToolCall(tc) => {
102                serde_json::to_string(tc).unwrap_or_else(|_| String::new())
103            }
104            EventPayload::ToolResult(tr) => {
105                serde_json::to_string(tr).unwrap_or_else(|_| String::new())
106            }
107            EventPayload::User(u) => u.text.clone(),
108            EventPayload::Message(m) => m.text.clone(),
109            EventPayload::Reasoning(r) => r.text.clone(),
110            EventPayload::TokenUsage(tu) => {
111                format!("tokens: in={} out={}", tu.input.total(), tu.output.total())
112            }
113            EventPayload::Notification(n) => {
114                serde_json::to_string(n).unwrap_or_else(|_| String::new())
115            }
116            EventPayload::SlashCommand(sc) => {
117                if let Some(args) = &sc.args {
118                    format!("{} {}", sc.name, args)
119                } else {
120                    sc.name.clone()
121                }
122            }
123            EventPayload::QueueOperation(qo) => {
124                format!(
125                    "{}{}",
126                    qo.operation,
127                    qo.content
128                        .as_ref()
129                        .map(|c| format!(": {}", c))
130                        .unwrap_or_default()
131                )
132            }
133            EventPayload::Summary(s) => s.summary.clone(),
134        };
135
136        if text.len() > 200 {
137            let truncated: String = text.chars().take(200).collect();
138            format!("{}...", truncated)
139        } else {
140            text
141        }
142    }
143}
144
145// ============================================================================
146// List Turns API
147// ============================================================================
148
149#[derive(Debug, Serialize, Deserialize, JsonSchema)]
150pub struct ListTurnsArgs {
151    pub session_id: String,
152    #[serde(default)]
153    pub limit: Option<usize>,
154    #[serde(default)]
155    pub cursor: Option<String>,
156}
157
158impl ListTurnsArgs {
159    pub fn limit(&self) -> usize {
160        self.limit.unwrap_or(50).min(100)
161    }
162}
163
164#[derive(Debug, Serialize)]
165pub struct ListTurnsResponse {
166    pub session_id: String,
167    pub total_turns: usize,
168    pub turns: Vec<TurnMetadata>,
169    #[serde(skip_serializing_if = "Option::is_none")]
170    pub next_cursor: Option<String>,
171}
172
173#[derive(Debug, Serialize)]
174pub struct TurnMetadata {
175    pub turn_index: usize,
176    pub user_content: String,
177    #[serde(skip_serializing_if = "Option::is_none")]
178    pub slash_command: Option<SlashCommandDetail>,
179    pub status: TurnStatus,
180    pub step_count: usize,
181    pub duration_ms: u64,
182    pub total_tokens: u64,
183    pub tools_used: HashMap<String, usize>,
184}
185
186#[derive(Debug, Serialize, PartialEq)]
187#[serde(rename_all = "lowercase")]
188pub enum TurnStatus {
189    Completed,
190    Failed,
191}
192
193impl ListTurnsResponse {
194    pub fn new(
195        session: AgentSession,
196        offset: usize,
197        limit: usize,
198        next_cursor: Option<String>,
199    ) -> Self {
200        let total_turns = session.turns.len();
201        let turns: Vec<_> = session
202            .turns
203            .into_iter()
204            .enumerate()
205            .skip(offset)
206            .take(limit)
207            .map(|(idx, turn)| {
208                let step_count = turn.steps.len();
209                let duration_ms = turn.stats.duration_ms as u64;
210                let total_tokens = turn.stats.total_tokens as u64;
211
212                let mut tools_used: HashMap<String, usize> = HashMap::new();
213                for step in &turn.steps {
214                    for tool in &step.tools {
215                        *tools_used
216                            .entry(tool.call.content.name().to_string())
217                            .or_insert(0) += 1;
218                    }
219                }
220
221                let status = if turn
222                    .steps
223                    .iter()
224                    .any(|s| s.tools.iter().any(|t| t.is_error))
225                {
226                    TurnStatus::Failed
227                } else {
228                    TurnStatus::Completed
229                };
230
231                let user_content = truncate(&turn.user.content.text, 100);
232
233                let slash_command =
234                    turn.user
235                        .slash_command
236                        .as_ref()
237                        .map(|cmd| SlashCommandDetail {
238                            name: cmd.name.clone(),
239                            args: cmd.args.clone(),
240                        });
241
242                TurnMetadata {
243                    turn_index: idx,
244                    user_content,
245                    slash_command,
246                    status,
247                    step_count,
248                    duration_ms,
249                    total_tokens,
250                    tools_used,
251                }
252            })
253            .collect();
254
255        Self {
256            session_id: session.session_id.to_string(),
257            total_turns,
258            turns,
259            next_cursor,
260        }
261    }
262}
263
264// ============================================================================
265// Get Turns API
266// ============================================================================
267
268// Data-driven defaults based on actual distribution analysis:
269// - 3,000 chars covers P90 (2,552 bytes) without truncation for most events
270// - 30 steps covers P50 (15) comfortably and approaches P75 (38), allowing ~70% of turns to be read fully
271// Worst case: 3,000 × 30 = 90k chars ≈ 22k tokens (still safe for 1-2 turns per request)
272// Average case: 1,176 × 30 ≈ 35k chars ≈ 8k tokens (very safe)
273const DEFAULT_MAX_CHARS_PER_FIELD: usize = 3_000;
274const DEFAULT_MAX_STEPS_LIMIT: usize = 30;
275
276#[derive(Debug, Serialize, Deserialize, JsonSchema)]
277pub struct GetTurnsArgs {
278    pub session_id: String,
279    pub turn_indices: Vec<usize>,
280    #[serde(default = "default_true")]
281    pub truncate: Option<bool>,
282    #[serde(default)]
283    pub max_chars_per_field: Option<usize>,
284    #[serde(default)]
285    pub max_steps_limit: Option<usize>,
286}
287
288fn default_true() -> Option<bool> {
289    Some(true)
290}
291
292impl GetTurnsArgs {
293    pub fn should_truncate(&self) -> bool {
294        self.truncate.unwrap_or(true)
295    }
296
297    pub fn max_chars(&self) -> usize {
298        self.max_chars_per_field
299            .unwrap_or(DEFAULT_MAX_CHARS_PER_FIELD)
300    }
301
302    pub fn max_steps(&self) -> usize {
303        self.max_steps_limit.unwrap_or(DEFAULT_MAX_STEPS_LIMIT)
304    }
305}
306
307#[derive(Debug, Serialize)]
308pub struct GetTurnsResponse {
309    pub turns: Vec<TurnDetail>,
310}
311
312#[derive(Debug, Serialize)]
313pub struct TurnDetail {
314    pub turn_index: usize,
315    pub user_content: String,
316    #[serde(skip_serializing_if = "Option::is_none")]
317    pub slash_command: Option<SlashCommandDetail>,
318    pub steps: Vec<StepDetail>,
319    #[serde(skip_serializing_if = "Option::is_none")]
320    pub steps_truncated: Option<bool>,
321}
322
323#[derive(Debug, Serialize)]
324pub struct SlashCommandDetail {
325    pub name: String,
326    #[serde(skip_serializing_if = "Option::is_none")]
327    pub args: Option<String>,
328}
329
330#[derive(Debug, Serialize)]
331pub struct StepDetail {
332    #[serde(skip_serializing_if = "Option::is_none")]
333    pub reasoning: Option<String>,
334    #[serde(skip_serializing_if = "Option::is_none")]
335    pub message: Option<String>,
336    pub tools: Vec<ToolDetail>,
337    pub is_failed: bool,
338}
339
340#[derive(Debug, Serialize)]
341pub struct ToolDetail {
342    pub name: String,
343    pub args: String,
344    pub result: Option<String>,
345    pub is_error: bool,
346}
347
348impl GetTurnsResponse {
349    pub fn new(session: AgentSession, args: &GetTurnsArgs) -> Result<Self, String> {
350        let should_truncate = args.should_truncate();
351        let max_chars = args.max_chars();
352        let max_steps = args.max_steps();
353
354        let mut turns = Vec::new();
355
356        for &turn_index in &args.turn_indices {
357            if turn_index >= session.turns.len() {
358                return Err(format!(
359                    "Turn index {} out of range (session has {} turns)",
360                    turn_index,
361                    session.turns.len()
362                ));
363            }
364
365            let turn = &session.turns[turn_index];
366            let user_content = if should_truncate {
367                truncate(&turn.user.content.text, max_chars)
368            } else {
369                turn.user.content.text.clone()
370            };
371
372            let total_steps = turn.steps.len();
373            let steps_to_process = if should_truncate && total_steps > max_steps {
374                &turn.steps[..max_steps]
375            } else {
376                &turn.steps[..]
377            };
378
379            let steps: Vec<StepDetail> = steps_to_process
380                .iter()
381                .map(|step| {
382                    let reasoning = step.reasoning.as_ref().map(|r| {
383                        if should_truncate {
384                            truncate(&r.content.text, max_chars)
385                        } else {
386                            r.content.text.clone()
387                        }
388                    });
389
390                    let message = step.message.as_ref().map(|m| {
391                        if should_truncate {
392                            truncate(&m.content.text, max_chars)
393                        } else {
394                            m.content.text.clone()
395                        }
396                    });
397
398                    let tools: Vec<ToolDetail> = step
399                        .tools
400                        .iter()
401                        .map(|tool| {
402                            let args_json = serde_json::to_string(&tool.call.content)
403                                .unwrap_or_else(|_| String::from("{}"));
404                            let result_json = tool.result.as_ref().map(|r| {
405                                serde_json::to_string(r).unwrap_or_else(|_| String::from("{}"))
406                            });
407
408                            ToolDetail {
409                                name: tool.call.content.name().to_string(),
410                                args: if should_truncate {
411                                    truncate(&args_json, max_chars)
412                                } else {
413                                    args_json
414                                },
415                                result: result_json.map(|r| {
416                                    if should_truncate {
417                                        truncate(&r, max_chars)
418                                    } else {
419                                        r
420                                    }
421                                }),
422                                is_error: tool.is_error,
423                            }
424                        })
425                        .collect();
426
427                    StepDetail {
428                        reasoning,
429                        message,
430                        tools,
431                        is_failed: step.is_failed,
432                    }
433                })
434                .collect();
435
436            let steps_truncated = if should_truncate && total_steps > max_steps {
437                Some(true)
438            } else {
439                None
440            };
441
442            let slash_command = turn
443                .user
444                .slash_command
445                .as_ref()
446                .map(|cmd| SlashCommandDetail {
447                    name: cmd.name.clone(),
448                    args: cmd.args.clone(),
449                });
450
451            turns.push(TurnDetail {
452                turn_index,
453                user_content,
454                slash_command,
455                steps,
456                steps_truncated,
457            });
458        }
459
460        Ok(Self { turns })
461    }
462}