Skip to main content

warden/reports/
query.rs

1//! `warden query` — a general rollup over named dimensions.
2//!
3//! This is the escape valve from the fixed report set: `--group-by
4//! project,model` answers a question no named report does. It is deliberately
5//! narrow — a closed list of dimensions, no filters beyond the global ones —
6//! because a real query language would freeze the record shape — warden offers a
7//! fixed rollup, not arbitrary queries.
8
9use std::collections::BTreeMap;
10
11use crate::output::{Cell, Report, Table};
12use crate::store::Event;
13use crate::store::Scanner;
14
15use super::{by_weight_desc, count, day_of, scan, ReportCtx, ReportError, Totals};
16
17/// The dimensions `--group-by` accepts, in the order the error lists them.
18pub const DIMENSIONS: [&str; 7] = [
19    "project", "model", "agent", "provider", "day", "session", "role",
20];
21
22/// Used when `--group-by` is omitted: the rollup people mean by default.
23pub const DEFAULT_GROUP_BY: &str = "project";
24
25/// Shown in a table cell when an event does not carry this dimension. The JSON
26/// keeps `null`, so a consumer never has to parse this string.
27const ABSENT: &str = "(none)";
28
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub struct Dimension(&'static str);
31
32impl Dimension {
33    pub fn name(self) -> &'static str {
34        self.0
35    }
36
37    fn value(self, event: &Event) -> Option<String> {
38        match self.0 {
39            "project" => event.project.clone(),
40            "model" => event.model.clone(),
41            "agent" => Some(event.agent.clone()),
42            "provider" => Some(event.provider.clone()),
43            "day" => day_of(event.ts),
44            "session" => event.session_id.clone(),
45            "role" => Some(event.role.clone()),
46            _ => None,
47        }
48    }
49}
50
51/// Parse a comma-separated `--group-by`, rejecting anything not in
52/// [`DIMENSIONS`] and de-duplicating repeats.
53pub fn parse_dimensions(spec: &str) -> Result<Vec<Dimension>, ReportError> {
54    let mut dims: Vec<Dimension> = Vec::new();
55    for raw in spec.split(',') {
56        let name = raw.trim();
57        if name.is_empty() {
58            continue;
59        }
60        let known = DIMENSIONS
61            .iter()
62            .find(|dim| **dim == name)
63            .ok_or_else(|| ReportError::UnknownDimension(name.to_string()))?;
64        let dim = Dimension(known);
65        if !dims.contains(&dim) {
66            dims.push(dim);
67        }
68    }
69    if dims.is_empty() {
70        return Err(ReportError::UnknownDimension(spec.trim().to_string()));
71    }
72    Ok(dims)
73}
74
75pub fn build(
76    scanner: &Scanner,
77    ctx: &ReportCtx,
78    dims: &[Dimension],
79) -> Result<Report, ReportError> {
80    let scanned = scan(scanner, ctx)?;
81
82    let mut buckets: BTreeMap<Vec<Option<String>>, Totals> = BTreeMap::new();
83    for event in &scanned.events {
84        let key: Vec<Option<String>> = dims.iter().map(|dim| dim.value(event)).collect();
85        buckets.entry(key).or_default().add(event, &ctx.pricing);
86    }
87
88    let mut headers: Vec<String> = dims.iter().map(|dim| dim.name().to_string()).collect();
89    headers.extend(
90        ["requests", "sessions", "in", "out", "cache r", "est. cost"]
91            .into_iter()
92            .map(String::from),
93    );
94    let mut table = Table::new(headers);
95    let mut rows = Vec::new();
96
97    for (key, totals) in by_weight_desc(buckets) {
98        let mut row: Vec<Cell> = key
99            .iter()
100            .map(|value| Cell::text(value.clone().unwrap_or_else(|| ABSENT.to_string())))
101            .collect();
102        row.push(count(totals.requests));
103        row.push(count(totals.sessions.len() as u64));
104        row.extend(totals.tail_cells());
105        table.push(row);
106
107        let mut json = serde_json::Map::new();
108        for (dim, value) in dims.iter().zip(&key) {
109            json.insert(dim.name().to_string(), serde_json::json!(value));
110        }
111        json.insert("sessions".into(), serde_json::json!(totals.sessions.len()));
112        totals.write_json(&mut json);
113        rows.push(serde_json::Value::Object(json));
114    }
115
116    let mut notes = scanned.notes;
117    notes.push(format!(
118        "grouped by {}; a dimension an event does not carry is {ABSENT} in the table and null in \
119         --json",
120        dims.iter()
121            .map(|dim| dim.name())
122            .collect::<Vec<_>>()
123            .join(", ")
124    ));
125
126    Ok(Report::new("query", ctx.window, table)
127        .with_json_rows(rows)
128        .with_notes(notes.finish()))
129}
130
131#[cfg(test)]
132mod tests {
133    use super::super::testkit::*;
134    use super::*;
135    use crate::cli::TimeWindow;
136    use crate::output::Style;
137
138    fn report(spec: &str) -> Report {
139        let mut sonnet = used("c", ms(2026, 8, 4, 9), "acme", "sonnet", 10, 1);
140        sonnet.session_id = Some("s2".into());
141        let (_dir, paths) = store(&[
142            priced(
143                used("a", ms(2026, 8, 4, 8), "acme", "opus", 1_000, 100),
144                5.0,
145            ),
146            priced(sonnet, 0.1),
147            priced(used("d", ms(2026, 8, 5, 9), "dotfiles", "opus", 20, 2), 0.2),
148        ]);
149        build(
150            &Scanner::new(paths),
151            &ReportCtx::new(TimeWindow::all(), None, true),
152            &parse_dimensions(spec).unwrap(),
153        )
154        .unwrap()
155    }
156
157    #[test]
158    fn groups_by_several_dimensions_at_once() {
159        let report = report("project,model");
160        assert_eq!(report.json_rows.len(), 3);
161        let first = &report.json_rows[0];
162        assert_eq!(first["project"], "acme");
163        assert_eq!(first["model"], "opus");
164        assert_eq!(first["input_tok"], 1_000);
165        assert_eq!(first["cost_est"], 5.0);
166
167        let rendered = report.table.render(Style::plain());
168        assert!(rendered.starts_with("PROJECT   MODEL"), "{rendered}");
169    }
170
171    #[test]
172    fn supports_every_documented_dimension() {
173        for dim in DIMENSIONS {
174            let report = report(dim);
175            assert!(!report.json_rows.is_empty(), "{dim}");
176            assert!(report.json_rows[0].get(dim).is_some(), "{dim}");
177        }
178    }
179
180    #[test]
181    fn rejects_an_unknown_dimension_and_lists_the_real_ones() {
182        let err = parse_dimensions("project,colour").unwrap_err();
183        let msg = err.to_string();
184        assert!(
185            msg.contains("unknown --group-by dimension \"colour\""),
186            "{msg}"
187        );
188        for dim in DIMENSIONS {
189            assert!(msg.contains(dim), "{msg} is missing {dim}");
190        }
191        assert!(parse_dimensions("").is_err());
192        assert!(parse_dimensions(" , ").is_err());
193    }
194
195    #[test]
196    fn repeated_dimensions_collapse() {
197        let dims = parse_dimensions("project, model ,project").unwrap();
198        assert_eq!(
199            dims.iter().map(|d| d.name()).collect::<Vec<_>>(),
200            ["project", "model"]
201        );
202    }
203
204    #[test]
205    fn a_dimension_an_event_lacks_is_null_in_json() {
206        let mut orphan = used("x", ms(2026, 8, 4, 8), "acme", "m", 5, 1);
207        orphan.project = None;
208        let (_dir, paths) = store(&[orphan]);
209        let report = build(
210            &Scanner::new(paths),
211            &ReportCtx::new(TimeWindow::all(), None, true),
212            &parse_dimensions("project").unwrap(),
213        )
214        .unwrap();
215        assert!(report.json_rows[0]["project"].is_null());
216        assert!(report.table.render(Style::plain()).contains(ABSENT));
217    }
218}