Skip to main content

skiagram_core/analysis/
context.rs

1//! Context-bloat attribution — the v0.2 headline feature (CLAUDE.md §2.2, §11).
2//!
3//! Answers "*why is my context window full, and what's filling it?*" by breaking
4//! a session's window down by source. Two kinds of number, kept strictly apart so
5//! the user is never misled (CLAUDE.md §8):
6//!
7//! 1. **MEASURED** (real, agent-reported tokens, from `Usage`):
8//!    - **startup overhead** — on a *cold-start* first request (`cache_read == 0`),
9//!      `input + cache_creation` is the fixed prefix cached before any real work:
10//!      system prompt + tool definitions + memory/CLAUDE.md + the first user turn.
11//!      VERIFIED on real files: a fresh session's first request shows e.g.
12//!      `input=3, cache_creation=18104, cache_read=0` → ~18 k tokens sitting in the
13//!      window before you type. A warm *resume* (`cache_read > 0` on the first
14//!      request) can't isolate this floor, so it's reported as unknown, not zero.
15//!    - **peak / final input context** — `max`/last of
16//!      `input + cache_read + cache_creation` across requests = how full the window
17//!      got. The agent does not log per-tool definition sizes, so this measured
18//!      bundle cannot be split further; we report the *inventory* instead (below).
19//!
20//! 2. **ESTIMATED** (from on-disk content sizes ÷ [`EST_CHARS_PER_TOKEN`], NEVER
21//!    billed): the *relative* composition of transcript content by
22//!    [`ContextSource`], per-MCP-server tool footprint, and the heaviest single
23//!    contributors (the context "fat tail" — one giant `Read`/MCP result can
24//!    dominate the window).
25//!
26//! Plus an **exact inventory** from `attachment` lines and tool calls: how many
27//! MCP servers are in play, how many tools were *deferred* (available but NOT
28//! loaded into the window — so they do *not* bloat it, a common misconception),
29//! how many skills were listed, and how heavy the MCP-instruction blocks are.
30//!
31//! Inputs already captured by the model: [`Usage`] on assistant events,
32//! [`Event::content_chars`]/[`Event::thinking_chars`], [`ToolCall::input_bytes`] +
33//! [`ToolCall::server`], [`Event::tool_use_id`] (links a `ToolResult` back to the
34//! tool that produced it), and [`Event::attachment_kind`]/[`Event::item_count`].
35//!
36//! [`Usage`]: crate::model::Usage
37//! [`Event::content_chars`]: crate::model::Event::content_chars
38//! [`Event::thinking_chars`]: crate::model::Event::thinking_chars
39//! [`ToolCall::input_bytes`]: crate::model::ToolCall::input_bytes
40//! [`ToolCall::server`]: crate::model::ToolCall::server
41//! [`Event::tool_use_id`]: crate::model::Event::tool_use_id
42//! [`EST_CHARS_PER_TOKEN`]: super::EST_CHARS_PER_TOKEN
43
44use std::collections::{BTreeMap, BTreeSet, HashMap};
45
46use jiff::civil::Date;
47use serde::Serialize;
48
49use crate::model::{EventKind, Session, Usage};
50
51/// Bucket a piece of transcript content is attributed to. The split between
52/// `AssistantText` and `Thinking` uses [`crate::model::Event::thinking_chars`]
53/// (a subset of `content_chars`); `ToolCalls` is the tool-use input JSON.
54#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
55pub enum ContextSource {
56    /// User-typed prompts.
57    UserPrompts,
58    /// Assistant visible text (excludes thinking and tool-call JSON).
59    AssistantText,
60    /// Extended-thinking text (measurable only; encrypted thinking is invisible).
61    Thinking,
62    /// Tool-use input JSON the assistant emitted.
63    ToolCalls,
64    /// Tool/command results returned into the window (usually the #1 filler).
65    ToolResults,
66    /// Injected non-message context: deferred-tool/skill listings, MCP
67    /// instruction blocks, IDE/file context, reminders (`attachment` lines).
68    Attachments,
69}
70
71/// Estimated share of transcript content from one [`ContextSource`].
72#[derive(Debug, Clone, Serialize)]
73pub struct SourceBreakdown {
74    pub source: ContextSource,
75    /// Raw content chars attributed to this source (measured from the transcript).
76    pub chars: u64,
77    /// Estimated tokens (`chars / EST_CHARS_PER_TOKEN`) — never a billed figure.
78    pub est_tokens: u64,
79    /// Fraction of total transcript content chars, `0.0..=1.0`.
80    pub share: f64,
81}
82
83/// Per-MCP-server context footprint: the tool-use JSON the assistant emitted to
84/// call the server's tools, plus the result payloads that came back (attributed
85/// via [`crate::model::Event::tool_use_id`]). Plain built-in tools (Read/Edit/…)
86/// are grouped under the synthetic server name [`ServerBreakdown::BUILTIN`].
87#[derive(Debug, Clone, Serialize)]
88pub struct ServerBreakdown {
89    /// MCP server name (`mcp__<server>__…`), or [`ServerBreakdown::BUILTIN`].
90    pub server: String,
91    pub calls: u64,
92    /// Chars of tool-use input JSON sent to this server's tools.
93    pub call_chars: u64,
94    /// Chars of tool-result payloads returned by this server's tools.
95    pub result_chars: u64,
96    /// Estimated tokens for `call_chars + result_chars`.
97    pub est_tokens: u64,
98}
99
100impl ServerBreakdown {
101    /// Synthetic server name for plain (non-MCP) tools and for results whose
102    /// originating tool call could not be located.
103    pub const BUILTIN: &'static str = "(built-in)";
104}
105
106/// One unusually large individual context contributor (the context "fat tail").
107#[derive(Debug, Clone, Serialize)]
108pub struct HeavyItem {
109    pub source: ContextSource,
110    /// Tool name / MCP server / short content summary, when known.
111    pub label: Option<String>,
112    /// Session the item belongs to.
113    pub session_id: String,
114    pub chars: u64,
115    pub est_tokens: u64,
116}
117
118/// Context-bloat profile for one session (each session — parent or sub-agent —
119/// is its own window, so each gets its own row; nothing is folded, nothing is
120/// double-counted).
121#[derive(Debug, Clone, Serialize)]
122pub struct SessionContext {
123    pub id: String,
124    pub project: Option<String>,
125    pub model: Option<String>,
126    pub sidechain: bool,
127    /// MEASURED real tokens cached on a cold-start first request (system prompt +
128    /// tool defs + memory + first user turn). `None` for a warm resume.
129    pub startup_overhead_tokens: Option<u64>,
130    /// True when a cold-start request (`cache_read == 0`) anchored the overhead.
131    pub cold_start: bool,
132    /// MEASURED: peak `input + cache_read + cache_creation` over requests.
133    pub peak_input_tokens: Option<u64>,
134    /// MEASURED: the last request's `input + cache_read + cache_creation`.
135    pub final_input_tokens: Option<u64>,
136    /// ESTIMATED transcript-content breakdown, sorted by `chars` desc.
137    pub sources: Vec<SourceBreakdown>,
138    pub total_content_chars: u64,
139    pub compactions: u64,
140}
141
142/// The whole context report: per-session profiles plus cross-session rollups.
143#[derive(Debug, Clone, Serialize)]
144pub struct ContextReport {
145    pub agent: String,
146    pub since: Option<Date>,
147    pub sessions_profiled: u64,
148    /// Per-session profiles, sorted by `peak_input_tokens` desc (then id).
149    pub sessions: Vec<SessionContext>,
150    /// ESTIMATED transcript-content breakdown across all sessions, sorted desc.
151    pub by_source: Vec<SourceBreakdown>,
152    /// MCP/tool footprint across all sessions, sorted by `est_tokens` desc.
153    pub by_server: Vec<ServerBreakdown>,
154    /// Largest individual contributors across all sessions, sorted desc.
155    pub heaviest: Vec<HeavyItem>,
156    /// Distinct MCP servers actually used (from `mcp__<server>__…` tool calls).
157    pub mcp_servers: Vec<String>,
158    /// Tools available but DEFERRED — not loaded into the window, so not bloat
159    /// (the max set size seen across `deferred_tools_delta` attachments).
160    pub deferred_tools: u64,
161    /// Skills listed in the window (max across `skill_listing` attachments).
162    pub skills_listed: u64,
163    /// Chars of MCP-server instruction blocks injected into the window.
164    pub mcp_instruction_chars: u64,
165    /// Total chars across all `attachment` lines.
166    pub attachment_chars: u64,
167    pub compactions: u64,
168    /// Largest cold-start startup overhead seen (measured), with its session id.
169    pub max_startup_overhead: Option<(String, u64)>,
170}
171
172/// Cap on how many [`HeavyItem`] entries the "fat tail" view keeps.
173const HEAVIEST_N: usize = 12;
174
175/// Window predicate shared with [`super::dedup`] / [`super::aggregate`]:
176/// when `since` is set, an event is in range only if it has a timestamp whose
177/// UTC date is on/after `since`. Undated events are excluded once a filter is on.
178fn in_window(since: Option<Date>, event: &crate::model::Event) -> bool {
179    match (since, event.ts) {
180        (None, _) => true,
181        (Some(since), Some(ts)) => super::utc_date(ts) >= since,
182        (Some(_), None) => false,
183    }
184}
185
186/// The MEASURED context size of one assistant request: everything that was in
187/// the window for that turn = fresh input + cache read + cache write.
188fn ctx_tokens(u: &Usage) -> u64 {
189    u.input.unwrap_or(0) + u.cache_read.unwrap_or(0) + u.cache_creation.unwrap_or(0)
190}
191
192/// Per-source char accumulators for one scope (a session, or the whole report).
193#[derive(Default)]
194struct SourceTally {
195    user_prompts: u64,
196    assistant_text: u64,
197    thinking: u64,
198    tool_calls: u64,
199    tool_results: u64,
200    attachments: u64,
201}
202
203impl SourceTally {
204    fn total(&self) -> u64 {
205        self.user_prompts
206            + self.assistant_text
207            + self.thinking
208            + self.tool_calls
209            + self.tool_results
210            + self.attachments
211    }
212
213    /// Emit `SourceBreakdown` rows for the non-empty sources, shares taken over
214    /// `total` (0.0 when nothing accumulated), sorted by chars desc then source.
215    fn breakdown(&self) -> Vec<SourceBreakdown> {
216        let total = self.total();
217        let mut rows: Vec<SourceBreakdown> = [
218            (ContextSource::UserPrompts, self.user_prompts),
219            (ContextSource::AssistantText, self.assistant_text),
220            (ContextSource::Thinking, self.thinking),
221            (ContextSource::ToolCalls, self.tool_calls),
222            (ContextSource::ToolResults, self.tool_results),
223            (ContextSource::Attachments, self.attachments),
224        ]
225        .into_iter()
226        .filter(|(_, chars)| *chars > 0)
227        .map(|(source, chars)| SourceBreakdown {
228            source,
229            chars,
230            est_tokens: super::est_tokens(chars),
231            share: if total == 0 {
232                0.0
233            } else {
234                chars as f64 / total as f64
235            },
236        })
237        .collect();
238        rows.sort_by(|a, b| {
239            b.chars
240                .cmp(&a.chars)
241                .then_with(|| source_rank(a.source).cmp(&source_rank(b.source)))
242        });
243        rows
244    }
245}
246
247/// Stable tiebreak ordering for sources that have equal chars.
248fn source_rank(s: ContextSource) -> u8 {
249    match s {
250        ContextSource::UserPrompts => 0,
251        ContextSource::AssistantText => 1,
252        ContextSource::Thinking => 2,
253        ContextSource::ToolCalls => 3,
254        ContextSource::ToolResults => 4,
255        ContextSource::Attachments => 5,
256    }
257}
258
259/// Σ of `input_bytes` over an event's tool calls.
260fn tool_call_bytes(event: &crate::model::Event) -> u64 {
261    event.tool_calls.iter().map(|c| c.input_bytes).sum()
262}
263
264/// Build a [`ContextReport`] from parsed sessions.
265///
266/// `since` filters events by UTC date (inclusive), consistent with
267/// [`super::dedup`] / [`super::aggregate`]: when set, events without a timestamp
268/// are excluded (they cannot be proven in range).
269///
270/// Computation (see module docs for the why):
271/// - **startup_overhead** / **cold_start**: from the first in-window assistant
272///   event that carries `Usage`; cold when its `cache_read` is `Some(0)`, in
273///   which case overhead = `input + cache_creation` (unknown fields count as 0).
274/// - **peak/final input**: over assistant events with usage,
275///   `input + cache_read + cache_creation`.
276/// - **by source** (estimated): `UserPrompts` = `User` event chars;
277///   `ToolResults` = `ToolResult` event chars; `Attachments` = `Attachment` event
278///   chars; `ToolCalls` = Σ `ToolCall::input_bytes`; `Thinking` = Σ
279///   `thinking_chars`; `AssistantText` = `content_chars − thinking_chars −
280///   tool-call bytes` (saturating). `share` is over the session's total.
281/// - **by server**: first map every `ToolCall::id → server`; attribute each
282///   tool-use's `input_bytes` to its server and each `ToolResult`'s chars to the
283///   server of its `tool_use_id` (unmatched → [`ServerBreakdown::BUILTIN`]).
284/// - **heaviest**: every tool result, assistant text chunk, thinking block, and
285///   tool call as a candidate; keep the top N by chars.
286/// - **inventory**: `deferred_tools`/`skills_listed` = max `item_count` over the
287///   respective attachment kinds; `mcp_instruction_chars` = Σ chars of
288///   `mcp_instructions_delta` attachments; `mcp_servers` = distinct tool-call
289///   servers; `compactions` = `Compaction` events.
290pub fn profile(sessions: &[Session], agent: &str, since: Option<Date>) -> ContextReport {
291    let refs: Vec<&Session> = sessions.iter().collect();
292    profile_refs(&refs, agent, since)
293}
294
295/// Same as [`profile`] but over borrowed sessions, so a caller that already holds
296/// `&Session` references (e.g. the drill-down, which profiles one session group at
297/// a time) need not clone whole transcripts. [`profile`] is the thin owned-slice
298/// wrapper.
299pub fn profile_refs(sessions: &[&Session], agent: &str, since: Option<Date>) -> ContextReport {
300    let mut session_rows: Vec<SessionContext> = Vec::new();
301
302    // Cross-session accumulators.
303    let mut global_sources = SourceTally::default();
304    let mut heaviest: Vec<HeavyItem> = Vec::new();
305
306    // Per-server footprint, two passes over in-window events.
307    // Pass 1 builds id -> server / id -> tool-name from the tool calls; pass 2
308    // attributes ToolResult bytes via tool_use_id.
309    struct ServerAcc {
310        calls: u64,
311        call_chars: u64,
312        result_chars: u64,
313    }
314    let mut servers: BTreeMap<String, ServerAcc> = BTreeMap::new();
315    let mut id_to_server: HashMap<String, String> = HashMap::new();
316    let mut id_to_tool: HashMap<String, String> = HashMap::new();
317
318    // Inventory accumulators.
319    let mut mcp_servers: BTreeSet<String> = BTreeSet::new();
320    let mut deferred_tools: u64 = 0;
321    let mut skills_listed: u64 = 0;
322    let mut mcp_instruction_chars: u64 = 0;
323    let mut attachment_chars: u64 = 0;
324    let mut total_compactions: u64 = 0;
325    let mut max_startup_overhead: Option<(String, u64)> = None;
326
327    // Pass 1: per-session profile + global source tally + the id->server map +
328    // inventory. (Tool-call servers/maps are global because results in any
329    // session reference ids; in practice ids are session-local but the map is
330    // keyed by id so this is safe.)
331    for session in sessions {
332        let mut tally = SourceTally::default();
333        let mut compactions: u64 = 0;
334
335        // MEASURED state, threaded over assistant-with-usage events in order.
336        let mut first_usage: Option<Usage> = None;
337        let mut peak: Option<u64> = None;
338        let mut final_ctx: Option<u64> = None;
339
340        for event in &session.events {
341            if !in_window(since, event) {
342                continue;
343            }
344
345            // --- inventory + per-server pass 1 (over all in-window events) ---
346            for call in &event.tool_calls {
347                if let Some(server) = &call.server {
348                    mcp_servers.insert(server.clone());
349                }
350                let key = call
351                    .server
352                    .clone()
353                    .unwrap_or_else(|| ServerBreakdown::BUILTIN.to_string());
354                let acc = servers.entry(key.clone()).or_insert(ServerAcc {
355                    calls: 0,
356                    call_chars: 0,
357                    result_chars: 0,
358                });
359                acc.calls += 1;
360                acc.call_chars += call.input_bytes;
361                id_to_server.insert(call.id.clone(), key);
362                id_to_tool.insert(call.id.clone(), call.name.clone());
363            }
364
365            // --- MEASURED (assistant events with usage) ---
366            if event.kind == EventKind::Assistant {
367                if let Some(u) = event.usage {
368                    if first_usage.is_none() {
369                        first_usage = Some(u);
370                    }
371                    let c = ctx_tokens(&u);
372                    peak = Some(peak.map_or(c, |p| p.max(c)));
373                    final_ctx = Some(c);
374                }
375            }
376
377            // --- ESTIMATED source attribution + heaviest candidates ---
378            match event.kind {
379                EventKind::User => {
380                    tally.user_prompts += event.content_chars;
381                    global_sources.user_prompts += event.content_chars;
382                    if event.content_chars > 0 {
383                        heaviest.push(HeavyItem {
384                            source: ContextSource::UserPrompts,
385                            label: event.content_summary.clone(),
386                            session_id: session.id.clone(),
387                            chars: event.content_chars,
388                            est_tokens: 0,
389                        });
390                    }
391                }
392                EventKind::ToolResult => {
393                    tally.tool_results += event.content_chars;
394                    global_sources.tool_results += event.content_chars;
395                    if event.content_chars > 0 {
396                        let label = event
397                            .tool_use_id
398                            .as_ref()
399                            .and_then(|id| id_to_tool.get(id).cloned());
400                        heaviest.push(HeavyItem {
401                            source: ContextSource::ToolResults,
402                            label,
403                            session_id: session.id.clone(),
404                            chars: event.content_chars,
405                            est_tokens: 0,
406                        });
407                    }
408                }
409                EventKind::Attachment => {
410                    tally.attachments += event.content_chars;
411                    global_sources.attachments += event.content_chars;
412                    attachment_chars += event.content_chars;
413                    match event.attachment_kind.as_deref() {
414                        Some("deferred_tools_delta") => {
415                            deferred_tools = deferred_tools.max(event.item_count);
416                        }
417                        Some("skill_listing") => {
418                            skills_listed = skills_listed.max(event.item_count);
419                        }
420                        Some("mcp_instructions_delta") => {
421                            mcp_instruction_chars += event.content_chars;
422                        }
423                        _ => {}
424                    }
425                }
426                EventKind::Assistant | EventKind::ToolCall => {
427                    let tc = tool_call_bytes(event);
428                    let text = event
429                        .content_chars
430                        .saturating_sub(event.thinking_chars + tc);
431                    tally.tool_calls += tc;
432                    tally.thinking += event.thinking_chars;
433                    tally.assistant_text += text;
434                    global_sources.tool_calls += tc;
435                    global_sources.thinking += event.thinking_chars;
436                    global_sources.assistant_text += text;
437
438                    if text > 0 {
439                        heaviest.push(HeavyItem {
440                            source: ContextSource::AssistantText,
441                            label: event.content_summary.clone(),
442                            session_id: session.id.clone(),
443                            chars: text,
444                            est_tokens: 0,
445                        });
446                    }
447                    if event.thinking_chars > 0 {
448                        heaviest.push(HeavyItem {
449                            source: ContextSource::Thinking,
450                            label: None,
451                            session_id: session.id.clone(),
452                            chars: event.thinking_chars,
453                            est_tokens: 0,
454                        });
455                    }
456                    for call in &event.tool_calls {
457                        if call.input_bytes > 0 {
458                            heaviest.push(HeavyItem {
459                                source: ContextSource::ToolCalls,
460                                label: Some(call.name.clone()),
461                                session_id: session.id.clone(),
462                                chars: call.input_bytes,
463                                est_tokens: 0,
464                            });
465                        }
466                    }
467                }
468                EventKind::Compaction => {
469                    compactions += 1;
470                }
471                EventKind::System | EventKind::SubAgentSpawn => {}
472            }
473        }
474
475        total_compactions += compactions;
476
477        let cold_start = first_usage.is_some_and(|u| u.cache_read == Some(0));
478        let startup_overhead_tokens = if cold_start {
479            first_usage.map(|u| u.input.unwrap_or(0) + u.cache_creation.unwrap_or(0))
480        } else {
481            None
482        };
483        if let Some(v) = startup_overhead_tokens {
484            match &max_startup_overhead {
485                Some((_, best)) if *best >= v => {}
486                _ => max_startup_overhead = Some((session.id.clone(), v)),
487            }
488        }
489
490        let total_content_chars = tally.total();
491        // No usage and no content => nothing to profile; skip the empty row.
492        if first_usage.is_none() && total_content_chars == 0 {
493            continue;
494        }
495
496        session_rows.push(SessionContext {
497            id: session.id.clone(),
498            project: session.project.clone(),
499            model: session.model.clone(),
500            sidechain: session.parent_session.is_some(),
501            startup_overhead_tokens,
502            cold_start,
503            peak_input_tokens: peak,
504            final_input_tokens: final_ctx,
505            sources: tally.breakdown(),
506            total_content_chars,
507            compactions,
508        });
509    }
510
511    // --- per-server pass 2: attribute ToolResult chars by tool_use_id ---
512    for session in sessions {
513        for event in &session.events {
514            if !in_window(since, event) {
515                continue;
516            }
517            if event.kind != EventKind::ToolResult {
518                continue;
519            }
520            let key = event
521                .tool_use_id
522                .as_ref()
523                .and_then(|id| id_to_server.get(id).cloned())
524                .unwrap_or_else(|| ServerBreakdown::BUILTIN.to_string());
525            servers
526                .entry(key)
527                .or_insert(ServerAcc {
528                    calls: 0,
529                    call_chars: 0,
530                    result_chars: 0,
531                })
532                .result_chars += event.content_chars;
533        }
534    }
535
536    let mut by_server: Vec<ServerBreakdown> = servers
537        .into_iter()
538        .map(|(server, acc)| ServerBreakdown {
539            server,
540            calls: acc.calls,
541            call_chars: acc.call_chars,
542            result_chars: acc.result_chars,
543            est_tokens: super::est_tokens(acc.call_chars + acc.result_chars),
544        })
545        .collect();
546    by_server.sort_by(|a, b| {
547        b.est_tokens
548            .cmp(&a.est_tokens)
549            .then_with(|| a.server.cmp(&b.server))
550    });
551
552    // Heaviest: rank by chars desc, keep top N, then fill est_tokens.
553    heaviest.sort_by(|a, b| {
554        b.chars
555            .cmp(&a.chars)
556            .then_with(|| a.session_id.cmp(&b.session_id))
557    });
558    heaviest.truncate(HEAVIEST_N);
559    for item in &mut heaviest {
560        item.est_tokens = super::est_tokens(item.chars);
561    }
562
563    // Per-session rows sorted by peak desc (None last) then id.
564    session_rows.sort_by(|a, b| match (b.peak_input_tokens, a.peak_input_tokens) {
565        (Some(x), Some(y)) => x.cmp(&y).then_with(|| a.id.cmp(&b.id)),
566        (Some(_), None) => std::cmp::Ordering::Greater,
567        (None, Some(_)) => std::cmp::Ordering::Less,
568        (None, None) => a.id.cmp(&b.id),
569    });
570
571    ContextReport {
572        agent: agent.to_string(),
573        since,
574        sessions_profiled: session_rows.len() as u64,
575        sessions: session_rows,
576        by_source: global_sources.breakdown(),
577        by_server,
578        heaviest,
579        mcp_servers: mcp_servers.into_iter().collect(),
580        deferred_tools,
581        skills_listed,
582        mcp_instruction_chars,
583        attachment_chars,
584        compactions: total_compactions,
585        max_startup_overhead,
586    }
587}
588
589#[cfg(test)]
590mod tests {
591    use super::*;
592    use crate::model::{Event, ToolCall};
593    use jiff::Timestamp;
594
595    const TS: &str = "2026-06-02T10:00:00Z";
596
597    /// A bare event of a given kind, timestamped, everything else empty.
598    fn ev(kind: EventKind) -> Event {
599        Event {
600            kind,
601            ts: Some(TS.parse().unwrap()),
602            request_id: None,
603            model: None,
604            usage: None,
605            tool_calls: Vec::new(),
606            sidechain: false,
607            content_summary: None,
608            content_chars: 0,
609            thinking_chars: 0,
610            has_thinking: false,
611            tool_use_id: None,
612            attachment_kind: None,
613            item_count: 0,
614        }
615    }
616
617    fn usage(input: u64, cc: u64, cr: Option<u64>) -> Usage {
618        Usage {
619            input: Some(input),
620            output: Some(0),
621            cache_creation: Some(cc),
622            cache_read: cr,
623            ..Usage::default()
624        }
625    }
626
627    /// Assistant event carrying only usage (no content).
628    fn asst_usage(u: Usage) -> Event {
629        let mut e = ev(EventKind::Assistant);
630        e.usage = Some(u);
631        e
632    }
633
634    fn call(id: &str, name: &str, bytes: u64) -> ToolCall {
635        ToolCall {
636            id: id.into(),
637            name: name.into(),
638            input_bytes: bytes,
639            server: ToolCall::server_from_name(name),
640        }
641    }
642
643    fn session(id: &str, parent: Option<&str>, events: Vec<Event>) -> Session {
644        Session {
645            id: id.into(),
646            agent: "claude-code".into(),
647            project: Some("proj".into()),
648            model: Some("claude-sonnet-4-5".into()),
649            parent_session: parent.map(str::to_string),
650            started_at: None,
651            ended_at: None,
652            events,
653            sub_agents: Vec::new(),
654            skipped_lines: 0,
655        }
656    }
657
658    fn src(b: &[SourceBreakdown], s: ContextSource) -> Option<&SourceBreakdown> {
659        b.iter().find(|x| x.source == s)
660    }
661
662    fn server<'a>(b: &'a [ServerBreakdown], name: &str) -> Option<&'a ServerBreakdown> {
663        b.iter().find(|x| x.server == name)
664    }
665
666    // --- MEASURED: cold start vs warm resume vs unknown ---
667
668    #[test]
669    fn cold_start_isolates_startup_overhead() {
670        // First request: cache_read == Some(0) => cold. overhead = 3 + 18104.
671        let s = session("s1", None, vec![asst_usage(usage(3, 18104, Some(0)))]);
672        let r = profile(&[s], "claude-code", None);
673        let sc = &r.sessions[0];
674        assert!(sc.cold_start);
675        assert_eq!(sc.startup_overhead_tokens, Some(18107));
676        // ctx = 3 + 0 (cr) + 18104 = 18107 for peak & final.
677        assert_eq!(sc.peak_input_tokens, Some(18107));
678        assert_eq!(sc.final_input_tokens, Some(18107));
679        assert_eq!(r.max_startup_overhead, Some(("s1".to_string(), 18107)));
680    }
681
682    #[test]
683    fn warm_resume_reports_overhead_as_unknown() {
684        // First request already reads cache (Some(>0)) => not cold, overhead None.
685        let s = session("s1", None, vec![asst_usage(usage(10, 0, Some(500)))]);
686        let r = profile(&[s], "claude-code", None);
687        let sc = &r.sessions[0];
688        assert!(!sc.cold_start);
689        assert_eq!(sc.startup_overhead_tokens, None);
690        // ctx = 10 + 500 + 0 = 510.
691        assert_eq!(sc.peak_input_tokens, Some(510));
692        assert_eq!(r.max_startup_overhead, None);
693    }
694
695    #[test]
696    fn unknown_cache_read_is_not_cold_start() {
697        // cache_read None => unknown, must NOT be treated as cold (Some(0)).
698        let s = session("s1", None, vec![asst_usage(usage(10, 0, None))]);
699        let r = profile(&[s], "claude-code", None);
700        let sc = &r.sessions[0];
701        assert!(!sc.cold_start);
702        assert_eq!(sc.startup_overhead_tokens, None);
703        // ctx = 10 + 0 (cr None -> 0) + 0 = 10.
704        assert_eq!(sc.final_input_tokens, Some(10));
705    }
706
707    #[test]
708    fn peak_is_max_and_final_is_last_request() {
709        // ctx values: 100, 5000, 800 -> peak 5000, final 800.
710        let s = session(
711            "s1",
712            None,
713            vec![
714                asst_usage(usage(100, 0, Some(0))),
715                asst_usage(usage(0, 0, Some(5000))),
716                asst_usage(usage(300, 0, Some(500))),
717            ],
718        );
719        let r = profile(&[s], "claude-code", None);
720        let sc = &r.sessions[0];
721        assert_eq!(sc.peak_input_tokens, Some(5000));
722        assert_eq!(sc.final_input_tokens, Some(800));
723        // First request cache_read Some(0) -> cold.
724        assert!(sc.cold_start);
725        assert_eq!(sc.startup_overhead_tokens, Some(100));
726    }
727
728    // --- ESTIMATED: source split within one assistant event ---
729
730    #[test]
731    fn assistant_event_splits_text_thinking_and_tool_calls() {
732        // content_chars=100, thinking=30, one tool call=20 bytes.
733        // AssistantText = 100 - (30 + 20) = 50; Thinking=30; ToolCalls=20.
734        let mut a = ev(EventKind::Assistant);
735        a.content_chars = 100;
736        a.thinking_chars = 30;
737        a.tool_calls = vec![call("t1", "Read", 20)];
738        let s = session("s1", None, vec![a]);
739        let r = profile(&[s], "claude-code", None);
740        let sc = &r.sessions[0];
741
742        assert_eq!(sc.total_content_chars, 100, "50 text + 30 think + 20 call");
743        assert_eq!(
744            src(&sc.sources, ContextSource::AssistantText)
745                .unwrap()
746                .chars,
747            50
748        );
749        assert_eq!(src(&sc.sources, ContextSource::Thinking).unwrap().chars, 30);
750        assert_eq!(
751            src(&sc.sources, ContextSource::ToolCalls).unwrap().chars,
752            20
753        );
754        // shares: 50/100, 30/100, 20/100.
755        let at = src(&sc.sources, ContextSource::AssistantText).unwrap();
756        assert!((at.share - 0.5).abs() < 1e-9);
757        assert_eq!(at.est_tokens, 12, "50 / 4 = 12");
758        // sorted by chars desc: AssistantText(50), Thinking(30), ToolCalls(20).
759        assert_eq!(sc.sources[0].source, ContextSource::AssistantText);
760        assert_eq!(sc.sources[1].source, ContextSource::Thinking);
761        assert_eq!(sc.sources[2].source, ContextSource::ToolCalls);
762    }
763
764    #[test]
765    fn source_split_across_event_kinds() {
766        let mut user = ev(EventKind::User);
767        user.content_chars = 40;
768        let mut result = ev(EventKind::ToolResult);
769        result.content_chars = 1000;
770        result.tool_use_id = Some("t1".into());
771        let mut attach = ev(EventKind::Attachment);
772        attach.content_chars = 8;
773        let mut a = ev(EventKind::Assistant);
774        a.content_chars = 12; // pure text, no thinking, no calls
775        let s = session("s1", None, vec![user, a, result, attach]);
776        let r = profile(&[s], "claude-code", None);
777        let sc = &r.sessions[0];
778
779        assert_eq!(sc.total_content_chars, 40 + 12 + 1000 + 8);
780        assert_eq!(
781            src(&sc.sources, ContextSource::UserPrompts).unwrap().chars,
782            40
783        );
784        assert_eq!(
785            src(&sc.sources, ContextSource::ToolResults).unwrap().chars,
786            1000
787        );
788        assert_eq!(
789            src(&sc.sources, ContextSource::Attachments).unwrap().chars,
790            8
791        );
792        assert_eq!(
793            src(&sc.sources, ContextSource::AssistantText)
794                .unwrap()
795                .chars,
796            12
797        );
798        // Largest is the tool result.
799        assert_eq!(sc.sources[0].source, ContextSource::ToolResults);
800        // Cross-session by_source mirrors single session here.
801        assert_eq!(
802            src(&r.by_source, ContextSource::ToolResults).unwrap().chars,
803            1000
804        );
805    }
806
807    // --- by server, including unmatched result -> BUILTIN ---
808
809    #[test]
810    fn server_attribution_groups_calls_and_results() {
811        // One MCP call (github) with a result, one builtin Read call with a
812        // result, and one result whose tool_use_id matches no call -> BUILTIN.
813        let mut a = ev(EventKind::Assistant);
814        a.tool_calls = vec![
815            call("g1", "mcp__github__search", 100),
816            call("r1", "Read", 30),
817        ];
818        let mut res_github = ev(EventKind::ToolResult);
819        res_github.tool_use_id = Some("g1".into());
820        res_github.content_chars = 4000;
821        let mut res_read = ev(EventKind::ToolResult);
822        res_read.tool_use_id = Some("r1".into());
823        res_read.content_chars = 200;
824        let mut res_orphan = ev(EventKind::ToolResult);
825        res_orphan.tool_use_id = Some("missing".into());
826        res_orphan.content_chars = 7;
827        let s = session("s1", None, vec![a, res_github, res_read, res_orphan]);
828        let r = profile(&[s], "claude-code", None);
829
830        let gh = server(&r.by_server, "github").unwrap();
831        assert_eq!(gh.calls, 1);
832        assert_eq!(gh.call_chars, 100);
833        assert_eq!(gh.result_chars, 4000);
834        assert_eq!(gh.est_tokens, (100 + 4000) / 4);
835
836        let bi = server(&r.by_server, ServerBreakdown::BUILTIN).unwrap();
837        assert_eq!(bi.calls, 1, "the Read call");
838        assert_eq!(bi.call_chars, 30);
839        // 200 (matched Read result) + 7 (orphan result) both land on BUILTIN.
840        assert_eq!(bi.result_chars, 207);
841
842        // Sorted by est_tokens desc -> github (1025) before builtin (59).
843        assert_eq!(r.by_server[0].server, "github");
844        // Inventory: github is the only mcp server.
845        assert_eq!(r.mcp_servers, vec!["github".to_string()]);
846    }
847
848    // --- heaviest ordering ---
849
850    #[test]
851    fn heaviest_ranks_biggest_contributor_first() {
852        let mut a = ev(EventKind::Assistant);
853        a.content_chars = 60; // pure text 60
854        a.content_summary = Some("hello".into());
855        let mut big = ev(EventKind::ToolResult);
856        big.tool_use_id = Some("t1".into());
857        big.content_chars = 9000;
858        let mut small = ev(EventKind::ToolResult);
859        small.content_chars = 100;
860        let mut a2 = ev(EventKind::Assistant);
861        a2.tool_calls = vec![call("t1", "Grep", 25)];
862        let s = session("s1", None, vec![a, a2, big, small]);
863        let r = profile(&[s], "claude-code", None);
864
865        assert_eq!(r.heaviest[0].chars, 9000);
866        assert_eq!(r.heaviest[0].source, ContextSource::ToolResults);
867        // Label resolved via tool_use_id -> the Grep call name.
868        assert_eq!(r.heaviest[0].label.as_deref(), Some("Grep"));
869        assert_eq!(r.heaviest[0].est_tokens, 2250);
870        // The set: 9000 (result), 100 (result), 60 (asst text), 25 (tool call).
871        let chars: Vec<u64> = r.heaviest.iter().map(|h| h.chars).collect();
872        assert_eq!(chars, vec![9000, 100, 60, 25]);
873    }
874
875    // --- inventory counts ---
876
877    #[test]
878    fn inventory_counts_deferred_skills_servers_and_compactions() {
879        let mut deferred1 = ev(EventKind::Attachment);
880        deferred1.attachment_kind = Some("deferred_tools_delta".into());
881        deferred1.item_count = 30;
882        deferred1.content_chars = 5;
883        let mut deferred2 = ev(EventKind::Attachment);
884        deferred2.attachment_kind = Some("deferred_tools_delta".into());
885        deferred2.item_count = 42; // max wins
886        deferred2.content_chars = 5;
887        let mut skills = ev(EventKind::Attachment);
888        skills.attachment_kind = Some("skill_listing".into());
889        skills.item_count = 9;
890        skills.content_chars = 5;
891        let mut mcp_instr = ev(EventKind::Attachment);
892        mcp_instr.attachment_kind = Some("mcp_instructions_delta".into());
893        mcp_instr.content_chars = 800;
894        let mut a = ev(EventKind::Assistant);
895        a.tool_calls = vec![
896            call("g1", "mcp__github__x", 1),
897            call("c1", "mcp__canva__y", 1),
898        ];
899        let comp = ev(EventKind::Compaction);
900        let s = session(
901            "s1",
902            None,
903            vec![deferred1, deferred2, skills, mcp_instr, a, comp],
904        );
905        let r = profile(&[s], "claude-code", None);
906
907        assert_eq!(r.deferred_tools, 42, "max item_count across deltas");
908        assert_eq!(r.skills_listed, 9);
909        assert_eq!(r.mcp_instruction_chars, 800);
910        // attachment_chars = 5 + 5 + 5 + 800.
911        assert_eq!(r.attachment_chars, 815);
912        assert_eq!(r.compactions, 1);
913        assert_eq!(r.sessions[0].compactions, 1);
914        assert_eq!(
915            r.mcp_servers,
916            vec!["canva".to_string(), "github".to_string()]
917        );
918    }
919
920    // --- sidechain flag + empty-row skipping ---
921
922    #[test]
923    fn sidechain_flag_and_empty_sessions_skipped() {
924        let child = session(
925            "agent-x",
926            Some("parent"),
927            vec![asst_usage(usage(5, 0, Some(0)))],
928        );
929        // A session with neither usage nor content => no row.
930        let empty = session("empty", None, vec![ev(EventKind::System)]);
931        let r = profile(&[child, empty], "claude-code", None);
932        assert_eq!(r.sessions_profiled, 1, "empty session produced no row");
933        assert_eq!(r.sessions[0].id, "agent-x");
934        assert!(r.sessions[0].sidechain);
935    }
936
937    // --- since filter ---
938
939    #[test]
940    fn since_filter_excludes_older_and_undated_events() {
941        let mut old = asst_usage(usage(1000, 0, Some(0)));
942        old.ts = Some("2026-06-01T10:00:00Z".parse::<Timestamp>().unwrap());
943        let mut new = asst_usage(usage(50, 0, Some(200)));
944        new.ts = Some("2026-06-02T09:00:00Z".parse::<Timestamp>().unwrap());
945        // Undated event must be excluded when a filter is set.
946        let mut undated = ev(EventKind::ToolResult);
947        undated.ts = None;
948        undated.content_chars = 99999;
949
950        let since: Date = "2026-06-02".parse().unwrap();
951        let s = session("s1", None, vec![old, new, undated]);
952        let r = profile(&[s], "claude-code", Some(since));
953        let sc = &r.sessions[0];
954
955        // Only the 2026-06-02 assistant event survives.
956        assert_eq!(
957            sc.peak_input_tokens,
958            Some(250),
959            "50 + 200, old 1000 excluded"
960        );
961        assert_eq!(sc.final_input_tokens, Some(250));
962        // Its cache_read is Some(200) -> warm, not cold.
963        assert!(!sc.cold_start);
964        // Undated ToolResult's 99999 chars excluded.
965        assert_eq!(sc.total_content_chars, 0);
966        assert!(r.heaviest.is_empty(), "undated heavy result filtered out");
967        assert_eq!(r.since, Some(since));
968    }
969
970    #[test]
971    fn sessions_sorted_by_peak_desc_then_id() {
972        let a = session("aaa", None, vec![asst_usage(usage(100, 0, Some(0)))]);
973        let b = session("bbb", None, vec![asst_usage(usage(0, 0, Some(9000)))]);
974        // No-usage-but-has-content session: peak None -> sorts last.
975        let mut txt = ev(EventKind::User);
976        txt.content_chars = 10;
977        let c = session("ccc", None, vec![txt]);
978        let r = profile(&[a, b, c], "claude-code", None);
979        let ids: Vec<&str> = r.sessions.iter().map(|s| s.id.as_str()).collect();
980        assert_eq!(ids, vec!["bbb", "aaa", "ccc"], "peak 9000, 100, then None");
981        assert_eq!(r.sessions[2].peak_input_tokens, None);
982    }
983}