use std::collections::BTreeMap;
use jiff::civil::Date;
use serde::Serialize;
use crate::analysis::aggregate::Filter;
use crate::analysis::dedup::dedup_session;
use crate::model::{Session, Usage};
use crate::pricing::{self, PricingTable};
const UNKNOWN_PROJECT: &str = "(unknown project)";
const UNKNOWN_MODEL: &str = "(unknown model)";
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize)]
pub enum FlameMetric {
#[default]
Tokens,
Cost,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
pub enum Dim {
Project,
Session,
Model,
Type,
}
impl Dim {
pub const DEFAULT: [Dim; 4] = [Dim::Project, Dim::Session, Dim::Model, Dim::Type];
}
fn short_session(id: &str) -> String {
id.chars().take(8).collect()
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct FoldedStack {
pub frames: Vec<String>,
pub value: u64,
}
impl FoldedStack {
pub fn to_folded_line(&self) -> String {
let mut line = String::new();
for (i, frame) in self.frames.iter().enumerate() {
if i > 0 {
line.push(';');
}
line.push_str(&sanitize_frame(frame));
}
line.push(' ');
line.push_str(&self.value.to_string());
line
}
}
#[derive(Debug, Clone, Serialize)]
pub struct FlameData {
pub agent: String,
pub since: Option<Date>,
pub metric: FlameMetric,
pub unit: &'static str,
pub stacks: Vec<FoldedStack>,
pub total_value: u64,
pub unpriced_requests: u64,
}
struct Leaf {
label: &'static str,
tokens: u64,
micro_usd: f64,
}
fn sanitize_frame(frame: &str) -> String {
frame
.chars()
.map(|c| match c {
';' => ':',
c if c.is_ascii_whitespace() => '_',
c => c,
})
.collect()
}
fn cache_write_micro_usd(usage: &Usage, p: &pricing::ModelPricing) -> f64 {
match (usage.cache_creation_5m, usage.cache_creation_1h) {
(None, None) => usage.cache_creation.unwrap_or(0) as f64 * p.cache_write_5m,
(m5, h1) => {
m5.unwrap_or(0) as f64 * p.cache_write_5m + h1.unwrap_or(0) as f64 * p.cache_write_1h
}
}
}
fn leaves(usage: &Usage, pricing: Option<&pricing::ModelPricing>) -> [Leaf; 5] {
let input = usage.input.unwrap_or(0);
let output = usage.output.unwrap_or(0);
let cache_read = usage.cache_read.unwrap_or(0);
let cache_write = usage.cache_creation.unwrap_or(0);
let thinking = usage.thinking.unwrap_or(0);
let (c_in, c_out, c_read, c_write, c_think) = match pricing {
Some(p) => (
input as f64 * p.input,
output as f64 * p.output,
cache_read as f64 * p.cache_read,
cache_write_micro_usd(usage, p),
thinking as f64 * p.output,
),
None => (0.0, 0.0, 0.0, 0.0, 0.0),
};
[
Leaf {
label: "input",
tokens: input,
micro_usd: c_in,
},
Leaf {
label: "output",
tokens: output,
micro_usd: c_out,
},
Leaf {
label: "cache-read",
tokens: cache_read,
micro_usd: c_read,
},
Leaf {
label: "cache-write",
tokens: cache_write,
micro_usd: c_write,
},
Leaf {
label: "thinking",
tokens: thinking,
micro_usd: c_think,
},
]
}
pub fn fold(
sessions: &[Session],
filter: &Filter,
agent: &str,
metric: FlameMetric,
dims: &[Dim],
pricing_table: &PricingTable,
) -> FlameData {
let mut paths: BTreeMap<Vec<String>, f64> = BTreeMap::new();
let mut unpriced_requests = 0u64;
for session in sessions {
let (records, _stats) = dedup_session(session, filter.since);
for rec in records {
let pricing = match metric {
FlameMetric::Tokens => None,
FlameMetric::Cost => {
match rec.model.as_deref().and_then(|m| pricing_table.lookup(m)) {
Some(p) => Some(p),
None => {
unpriced_requests += 1;
continue;
}
}
}
};
let project = rec
.project
.clone()
.unwrap_or_else(|| UNKNOWN_PROJECT.into());
let session_label =
short_session(rec.parent_session.as_deref().unwrap_or(&rec.session_id));
let model = rec.model.clone().unwrap_or_else(|| UNKNOWN_MODEL.into());
for leaf in leaves(&rec.usage, pricing) {
let weight = match metric {
FlameMetric::Tokens => leaf.tokens as f64,
FlameMetric::Cost => leaf.micro_usd,
};
if weight <= 0.0 {
continue;
}
let frames: Vec<String> = dims
.iter()
.map(|dim| match dim {
Dim::Project => project.clone(),
Dim::Session => session_label.clone(),
Dim::Model => model.clone(),
Dim::Type => leaf.label.to_string(),
})
.collect();
*paths.entry(frames).or_insert(0.0) += weight;
}
}
}
let mut total_value = 0u64;
let mut stacks = Vec::with_capacity(paths.len());
for (frames, weight) in paths {
let value = weight.round() as u64;
if value == 0 {
continue;
}
total_value += value;
stacks.push(FoldedStack { frames, value });
}
FlameData {
agent: agent.to_string(),
since: filter.since,
metric,
unit: match metric {
FlameMetric::Tokens => "tokens",
FlameMetric::Cost => "µ$",
},
stacks,
total_value,
unpriced_requests,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::model::{Event, EventKind, Usage};
const DIMS: &[Dim] = &Dim::DEFAULT;
fn turn(rid: &str, usage: Usage) -> Event {
Event {
kind: EventKind::Assistant,
ts: Some("2026-06-02T10:00:00Z".parse().unwrap()),
request_id: Some(rid.into()),
model: Some("claude-sonnet-4-5".into()),
usage: Some(usage),
tool_calls: Vec::new(),
sidechain: false,
content_summary: None,
content_chars: 0,
thinking_chars: 0,
has_thinking: false,
tool_use_id: None,
attachment_kind: None,
item_count: 0,
}
}
fn session(id: &str, parent: Option<&str>, events: Vec<Event>) -> Session {
Session {
id: id.into(),
agent: "claude-code".into(),
project: Some("proj".into()),
model: Some("claude-sonnet-4-5".into()),
parent_session: parent.map(str::to_string),
started_at: None,
ended_at: None,
events,
sub_agents: Vec::new(),
skipped_lines: 0,
}
}
fn usage(input: u64, output: u64, cache_read: u64, cache_write: u64) -> Usage {
Usage {
input: Some(input),
output: Some(output),
cache_read: Some(cache_read),
cache_creation: Some(cache_write),
..Usage::default()
}
}
fn leaf_value(
d: &FlameData,
project: &str,
session: &str,
model: &str,
ty: &str,
) -> Option<u64> {
let want = [project, session, model, ty];
d.stacks.iter().find(|s| s.frames == want).map(|s| s.value)
}
#[test]
fn basic_tokens_fold_produces_per_type_leaves() {
let s = session("s", None, vec![turn("r", usage(1_000, 200, 5_000, 300))]);
let d = fold(
&[s],
&Filter::default(),
"claude-code",
FlameMetric::Tokens,
DIMS,
&PricingTable::embedded(),
);
assert_eq!(d.unit, "tokens");
assert_eq!(d.unpriced_requests, 0);
let m = "claude-sonnet-4-5";
assert_eq!(leaf_value(&d, "proj", "s", m, "input"), Some(1_000));
assert_eq!(leaf_value(&d, "proj", "s", m, "output"), Some(200));
assert_eq!(leaf_value(&d, "proj", "s", m, "cache-read"), Some(5_000));
assert_eq!(leaf_value(&d, "proj", "s", m, "cache-write"), Some(300));
assert_eq!(d.stacks.len(), 4);
assert!(leaf_value(&d, "proj", "s", m, "thinking").is_none());
assert_eq!(d.total_value, 1_000 + 200 + 5_000 + 300);
let stack_sum: u64 = d.stacks.iter().map(|s| s.value).sum();
assert_eq!(d.total_value, stack_sum);
}
#[test]
fn dedup_collapses_repeated_lines() {
let u = usage(1_000, 200, 0, 0);
let s = session(
"s",
None,
vec![turn("req", u), turn("req", u), turn("req", u)],
);
let d = fold(
&[s],
&Filter::default(),
"claude-code",
FlameMetric::Tokens,
DIMS,
&PricingTable::embedded(),
);
let m = "claude-sonnet-4-5";
assert_eq!(leaf_value(&d, "proj", "s", m, "input"), Some(1_000));
assert_eq!(leaf_value(&d, "proj", "s", m, "output"), Some(200));
assert_eq!(d.total_value, 1_200, "3 lines -> 1 request, not 3x");
}
#[test]
fn subagent_folds_into_parent_session_frame() {
let parent = session("parent", None, vec![turn("p", usage(1_000, 0, 0, 0))]);
let child = session(
"agent-x",
Some("parent"),
vec![turn("c", usage(500, 0, 0, 0))],
);
let d = fold(
&[parent, child],
&Filter::default(),
"claude-code",
FlameMetric::Tokens,
DIMS,
&PricingTable::embedded(),
);
let m = "claude-sonnet-4-5";
assert_eq!(
leaf_value(&d, "proj", "parent", m, "input"),
Some(1_500),
"parent + child fold into one leaf"
);
assert!(leaf_value(&d, "proj", "agent-x", m, "input").is_none());
assert_eq!(d.total_value, 1_500);
}
#[test]
fn cost_metric_prices_each_leaf() {
let u = usage(1_000, 100, 2_000, 400);
let s = session("s", None, vec![turn("r", u)]);
let d = fold(
&[s],
&Filter::default(),
"claude-code",
FlameMetric::Cost,
DIMS,
&PricingTable::embedded(),
);
assert_eq!(d.unit, "µ$");
assert_eq!(d.unpriced_requests, 0);
let m = "claude-sonnet-4-5";
assert_eq!(leaf_value(&d, "proj", "s", m, "input"), Some(3_000)); assert_eq!(leaf_value(&d, "proj", "s", m, "output"), Some(1_500)); assert_eq!(leaf_value(&d, "proj", "s", m, "cache-read"), Some(600)); assert_eq!(leaf_value(&d, "proj", "s", m, "cache-write"), Some(1_500));
let expected = pricing::cost_usd(Some("claude-sonnet-4-5"), &u).unwrap() * 1e6;
let diff = (d.total_value as f64 - expected).abs();
assert!(
diff <= 3.0,
"flame cost {} vs pricer {expected}",
d.total_value
);
}
#[test]
fn cost_metric_thinking_bills_at_output_rate() {
let u = Usage {
thinking: Some(1_000),
..Usage::default()
};
let s = session("s", None, vec![turn("r", u)]);
let d = fold(
&[s],
&Filter::default(),
"claude-code",
FlameMetric::Cost,
DIMS,
&PricingTable::embedded(),
);
assert_eq!(
leaf_value(&d, "proj", "s", "claude-sonnet-4-5", "thinking"),
Some(15_000)
);
}
#[test]
fn cost_metric_skips_and_counts_unpriced() {
let mut unp = turn("u", usage(5_000, 100, 0, 0));
unp.model = Some("claude-opus-5-0".into()); let only_unpriced = session("u", None, vec![unp.clone()]);
let d = fold(
&[only_unpriced],
&Filter::default(),
"claude-code",
FlameMetric::Cost,
DIMS,
&PricingTable::embedded(),
);
assert!(d.stacks.is_empty());
assert_eq!(d.total_value, 0);
assert_eq!(d.unpriced_requests, 1);
let priced = turn("p", usage(1_000, 0, 0, 0)); let mixed = session("mix", None, vec![unp, priced]);
let d = fold(
&[mixed],
&Filter::default(),
"claude-code",
FlameMetric::Cost,
DIMS,
&PricingTable::embedded(),
);
assert_eq!(d.unpriced_requests, 1);
assert_eq!(d.stacks.len(), 1, "only the priced request contributes");
assert_eq!(
leaf_value(&d, "proj", "mix", "claude-sonnet-4-5", "input"),
Some(3_000)
);
}
#[test]
fn tokens_metric_ignores_pricing() {
let mut unp = turn("u", usage(5_000, 100, 0, 0));
unp.model = Some("claude-opus-5-0".into());
let s = session("s", None, vec![unp]);
let d = fold(
&[s],
&Filter::default(),
"claude-code",
FlameMetric::Tokens,
DIMS,
&PricingTable::embedded(),
);
assert_eq!(d.unpriced_requests, 0);
assert_eq!(
leaf_value(&d, "proj", "s", "claude-opus-5-0", "input"),
Some(5_000)
);
}
#[test]
fn since_filter_excludes_out_of_window_requests() {
let mut old = turn("old", usage(9_999, 0, 0, 0));
old.ts = Some("2026-06-01T10:00:00Z".parse().unwrap());
let mut new = turn("new", usage(100, 0, 0, 0));
new.ts = Some("2026-06-02T10:00:00Z".parse().unwrap());
let s = session("s", None, vec![old, new]);
let filter = Filter {
since: Some("2026-06-02".parse().unwrap()),
};
let d = fold(
&[s],
&filter,
"claude-code",
FlameMetric::Tokens,
DIMS,
&PricingTable::embedded(),
);
assert_eq!(d.since, filter.since);
assert_eq!(d.total_value, 100, "only the in-window request remains");
assert_eq!(
leaf_value(&d, "proj", "s", "claude-sonnet-4-5", "input"),
Some(100)
);
}
#[test]
fn empty_input_is_an_empty_result() {
let d = fold(
&[],
&Filter::default(),
"claude-code",
FlameMetric::Tokens,
DIMS,
&PricingTable::embedded(),
);
assert!(d.stacks.is_empty());
assert_eq!(d.total_value, 0);
assert_eq!(d.unpriced_requests, 0);
}
#[test]
fn to_folded_line_formats_and_sanitizes() {
let clean = FoldedStack {
frames: vec!["a".into(), "b".into(), "c".into(), "input".into()],
value: 1_234,
};
assert_eq!(clean.to_folded_line(), "a;b;c;input 1234");
let dirty = FoldedStack {
frames: vec!["a;b".into(), "c d".into()],
value: 7,
};
assert_eq!(dirty.to_folded_line(), "a:b;c_d 7");
}
#[test]
fn stacks_are_emitted_in_sorted_order() {
let mut a = session("sa", None, vec![turn("ra", usage(10, 20, 30, 40))]);
a.project = Some("zeta".into());
let mut b = session("sb", None, vec![turn("rb", usage(1, 2, 3, 4))]);
b.project = Some("alpha".into());
let d = fold(
&[a, b],
&Filter::default(),
"claude-code",
FlameMetric::Tokens,
DIMS,
&PricingTable::embedded(),
);
assert!(d.stacks.len() >= 2);
for w in d.stacks.windows(2) {
assert!(w[0].frames <= w[1].frames, "stacks must be sorted");
}
assert_eq!(d.stacks[0].frames[0], "alpha");
}
#[test]
fn session_frame_is_shortened_to_prefix() {
let id = "3e9d2c41-7b5a-4f2e-9c1d-2f6b8a1c0e11";
let s = session(id, None, vec![turn("r", usage(1_000, 0, 0, 0))]);
let d = fold(
&[s],
&Filter::default(),
"claude-code",
FlameMetric::Tokens,
DIMS,
&PricingTable::embedded(),
);
let m = "claude-sonnet-4-5";
assert_eq!(leaf_value(&d, "proj", "3e9d2c41", m, "input"), Some(1_000));
assert!(
!d.stacks.iter().any(|s| s.frames.iter().any(|f| f == id)),
"the full UUID must never appear as a frame"
);
assert_eq!(d.total_value, 1_000);
}
#[test]
fn dims_drop_session_level() {
let id = "3e9d2c41-7b5a-4f2e-9c1d-2f6b8a1c0e11";
let s = session(id, None, vec![turn("r", usage(1_000, 200, 0, 0))]);
let dims = &[Dim::Project, Dim::Model, Dim::Type];
let d = fold(
&[s],
&Filter::default(),
"claude-code",
FlameMetric::Tokens,
dims,
&PricingTable::embedded(),
);
assert!(d.stacks.iter().all(|s| s.frames.len() == 3));
assert!(
!d.stacks
.iter()
.any(|s| s.frames.iter().any(|f| f == "3e9d2c41")),
"session level dropped"
);
let m = "claude-sonnet-4-5";
let input = d.stacks.iter().find(|s| s.frames == ["proj", m, "input"]);
assert_eq!(input.map(|s| s.value), Some(1_000));
assert_eq!(d.total_value, 1_200);
}
#[test]
fn omitting_type_sums_token_types() {
let s = session("s", None, vec![turn("r", usage(1_000, 200, 50, 0))]);
let dims = &[Dim::Project, Dim::Model];
let d = fold(
&[s],
&Filter::default(),
"claude-code",
FlameMetric::Tokens,
dims,
&PricingTable::embedded(),
);
let m = "claude-sonnet-4-5";
assert_eq!(d.stacks.len(), 1);
assert_eq!(d.stacks[0].frames, ["proj", m]);
assert_eq!(d.stacks[0].value, 1_250, "1000 + 200 + 50 summed");
assert_eq!(d.total_value, 1_250);
}
}