Skip to main content

kaizen/visualization/
build.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2
3use super::activity::activity;
4use super::types::*;
5use crate::core::event::SessionRecord;
6use crate::store::{SessionSearchQuery, Store};
7use anyhow::{Result, ensure};
8use serde::{Deserialize, Serialize};
9
10const ACTIVE_TTL_MS: u64 = 5 * 60_000;
11
12#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
13pub struct VisualizationLimits {
14    /// Maximum latest sessions materialized for summaries.
15    pub sessions: usize,
16    /// Maximum latest events materialized for selected detail.
17    pub selected_events: usize,
18    /// Maximum spans materialized for selected detail.
19    pub selected_spans: usize,
20    /// Maximum files materialized for selected detail.
21    pub selected_files: usize,
22}
23
24#[derive(Clone, Debug, Serialize, Deserialize)]
25pub struct VisualizationQuery {
26    pub workspace: String,
27    pub selected_session_id: Option<String>,
28    pub now_ms: u64,
29    /// Compute heatmap bins only for surfaces that render them.
30    pub include_activity: bool,
31    /// Fall back to latest session when requested selection is absent or invalid.
32    pub select_latest: bool,
33    #[serde(default)]
34    pub session_search: SessionSearchInput,
35    pub limits: VisualizationLimits,
36}
37
38#[derive(Clone, Copy, Debug, Eq, PartialEq)]
39pub(crate) struct MaterializedRows {
40    pub(crate) sessions: usize,
41    pub(crate) selected_events: usize,
42    pub(crate) selected_spans: usize,
43    pub(crate) selected_files: usize,
44}
45
46pub(crate) struct BuiltReport {
47    pub(crate) report: VisualizationReport,
48    pub(crate) materialized: MaterializedRows,
49}
50
51pub fn build_report(store: &Store, query: VisualizationQuery) -> Result<VisualizationReport> {
52    Ok(build_report_observed(store, query)?.report)
53}
54
55pub(crate) fn build_report_observed(
56    store: &Store,
57    query: VisualizationQuery,
58) -> Result<BuiltReport> {
59    validate(&query.limits)?;
60    let active_since_ms = query.now_ms.saturating_sub(ACTIVE_TTL_MS);
61    let (totals, quality) = store.visualization_totals(&query.workspace, active_since_ms)?;
62    let search = SessionSearchQuery::new(
63        &query.workspace,
64        &query.session_search.q,
65        query.session_search.offset,
66        query.limits.sessions,
67        query.now_ms,
68    )?;
69    let (sessions, session_page) = store.visualization_sessions(&search, query.now_ms)?;
70    let selected = selected_detail(store, &query, sessions.first())?;
71    let activity = activity_report(store, &query)?;
72    let materialized = counts(&sessions, &selected);
73    Ok(BuiltReport {
74        report: report(
75            query,
76            totals,
77            quality,
78            sessions,
79            session_page,
80            selected,
81            activity,
82        ),
83        materialized,
84    })
85}
86
87fn validate(limits: &VisualizationLimits) -> Result<()> {
88    ensure_positive(limits.sessions, "session")?;
89    ensure_positive(limits.selected_events, "event")?;
90    ensure_positive(limits.selected_spans, "span")?;
91    ensure_positive(limits.selected_files, "file")?;
92    Ok(())
93}
94
95fn ensure_positive(limit: usize, kind: &str) -> Result<()> {
96    ensure!(limit > 0, "visualization {kind} limit must be positive");
97    Ok(())
98}
99
100fn activity_report(store: &Store, query: &VisualizationQuery) -> Result<ActivityReport> {
101    if query.include_activity {
102        activity(store, &query.workspace, query.now_ms)
103    } else {
104        Ok(Default::default())
105    }
106}
107
108fn selected_detail(
109    store: &Store,
110    query: &VisualizationQuery,
111    latest: Option<&TraceSummary>,
112) -> Result<Option<TraceDetail>> {
113    let Some(session) = selected_session(store, query, latest)? else {
114        return Ok(None);
115    };
116    let id = session.id.clone();
117    Ok(Some(TraceDetail {
118        session,
119        prompt: None,
120        events: store.list_latest_events_for_session(&id, query.limits.selected_events)?,
121        spans: store.limited_session_span_tree(&id, query.limits.selected_spans)?,
122        files: store.limited_files_for_session(&id, query.limits.selected_files)?,
123    }))
124}
125
126fn selected_session(
127    store: &Store,
128    query: &VisualizationQuery,
129    latest: Option<&TraceSummary>,
130) -> Result<Option<SessionRecord>> {
131    if let Some(session) = requested_session(store, query)? {
132        return Ok(Some(session));
133    }
134    if !query.select_latest {
135        return Ok(None);
136    }
137    latest
138        .map(|summary| store.get_session(&summary.id))
139        .transpose()
140        .map(Option::flatten)
141}
142
143fn requested_session(store: &Store, query: &VisualizationQuery) -> Result<Option<SessionRecord>> {
144    let Some(id) = query.selected_session_id.as_deref() else {
145        return Ok(None);
146    };
147    Ok(store
148        .get_session(id)?
149        .filter(|session| session.workspace == query.workspace))
150}
151
152fn counts(sessions: &[TraceSummary], selected: &Option<TraceDetail>) -> MaterializedRows {
153    MaterializedRows {
154        sessions: sessions.len(),
155        selected_events: selected.as_ref().map_or(0, |detail| detail.events.len()),
156        selected_spans: selected
157            .as_ref()
158            .map_or(0, |detail| span_count(&detail.spans)),
159        selected_files: selected.as_ref().map_or(0, |detail| detail.files.len()),
160    }
161}
162
163fn span_count(spans: &[crate::store::SpanNode]) -> usize {
164    spans
165        .iter()
166        .map(|node| 1 + span_count(&node.children))
167        .sum()
168}
169
170fn report(
171    query: VisualizationQuery,
172    totals: VisualizationTotals,
173    quality: DataQuality,
174    sessions: Vec<TraceSummary>,
175    session_page: SessionPageMeta,
176    selected: Option<TraceDetail>,
177    activity: ActivityReport,
178) -> VisualizationReport {
179    VisualizationReport {
180        generated_at_ms: query.now_ms,
181        workspace: query.workspace,
182        totals,
183        activity,
184        sessions,
185        session_page,
186        selected,
187        quality,
188    }
189}