1use std::fs::File;
2use std::io::{BufRead, BufReader, Seek, SeekFrom};
3
4use serde_json::Value;
5
6use crate::models;
7use crate::session::{ClaudeSession, SessionStatus, SubagentRollup, TelemetryStatus};
8use crate::transcript::{TranscriptBlock, TranscriptEvent, TranscriptRole, parse_line};
9
10#[derive(Default)]
11struct UsageRollup {
12 input_tokens: u64,
13 output_tokens: u64,
14 cache_read_tokens: u64,
15 cache_write_tokens: u64,
16 cache_write_1h_tokens: u64,
17 cost_usd: f64,
18 usage_metrics_available: bool,
19 cost_estimate_unverified: bool,
20}
21
22impl UsageRollup {
23 fn total_input_tokens(&self) -> u64 {
24 self.input_tokens + self.cache_read_tokens + self.cache_write_tokens
25 }
26}
27
28pub fn update_tokens(session: &mut ClaudeSession) {
30 let mut last_type = session.last_msg_type.clone();
32 let mut last_stop_reason = session.last_stop_reason.clone();
33 let mut is_waiting_for_task = session.is_waiting_for_task;
34 let mut saw_non_empty_line = false;
35 let mut recognized_events = 0usize;
36 let mut saw_parent_usage = false;
37 let jsonl_path = session.jsonl_path.clone();
38
39 match jsonl_path.as_ref() {
40 Some(path) => {
41 let mut file = match File::open(path) {
42 Ok(f) => f,
43 Err(_) => {
44 session.telemetry_status = TelemetryStatus::UnreadableTranscript;
45 finalize_usage(
46 session,
47 &last_type,
48 &last_stop_reason,
49 is_waiting_for_task,
50 false,
51 );
52 return;
53 }
54 };
55
56 let file_len = file.metadata().map(|m| m.len()).unwrap_or(0);
57
58 if file_len == 0 {
59 session.telemetry_status = TelemetryStatus::Pending;
60 } else {
61 if session.jsonl_offset > file_len {
62 session.jsonl_offset = 0;
63 session.own_input_tokens = 0;
64 session.own_output_tokens = 0;
65 session.own_cache_read_tokens = 0;
66 session.own_cache_write_tokens = 0;
67 session.own_cache_write_1h_tokens = 0;
68 session.last_usage_msg_id = None;
69 last_type.clear();
71 last_stop_reason.clear();
72 is_waiting_for_task = false;
73 }
74
75 if session.jsonl_offset < file_len {
76 if session.jsonl_offset > 0
77 && file.seek(SeekFrom::Start(session.jsonl_offset)).is_err()
78 {
79 finalize_usage(
80 session,
81 &last_type,
82 &last_stop_reason,
83 is_waiting_for_task,
84 false,
85 );
86 return;
87 }
88
89 let reader = BufReader::new(&file);
90
91 for line in reader.lines() {
92 let line = match line {
93 Ok(l) => l,
94 Err(_) => break,
95 };
96
97 if line.trim().is_empty() {
98 continue;
99 }
100 saw_non_empty_line = true;
101
102 let Some(event) = parse_line(&line) else {
103 continue;
104 };
105 recognized_events += 1;
106
107 match event {
108 TranscriptEvent::WaitingForTask => {
109 is_waiting_for_task = true;
110 }
111 TranscriptEvent::Message(message) => {
112 is_waiting_for_task = false;
113 last_type = match message.role {
114 TranscriptRole::Assistant => "assistant".to_string(),
115 TranscriptRole::User => "user".to_string(),
116 };
117
118 if let Some(reason) = message.stop_reason {
119 last_stop_reason = reason;
120 } else {
121 let has_tool_use = message
125 .content
126 .iter()
127 .any(|b| matches!(b, TranscriptBlock::ToolUse { .. }));
128 if has_tool_use {
129 last_stop_reason = "tool_use".to_string();
130 } else {
131 last_stop_reason.clear();
132 }
133 }
134
135 if let Some(usage) = message.usage {
136 let input = usage.input_tokens;
137 let cache_read = usage.cache_read_input_tokens;
138 let cache_create = usage.cache_creation_input_tokens;
139 let cache_create_1h = usage.cache_creation_1h_input_tokens;
140 let output = usage.output_tokens;
141
142 let repeat = match (&message.id, &session.last_usage_msg_id) {
156 (Some(id), Some(prev)) => id == prev,
157 _ => false,
158 };
159 if !repeat {
160 session.own_input_tokens +=
161 input + cache_read + cache_create;
162 session.own_output_tokens += output;
163 session.own_cache_read_tokens += cache_read;
164 session.own_cache_write_tokens += cache_create;
165 session.own_cache_write_1h_tokens += cache_create_1h;
166 if message.id.is_some() {
167 session.last_usage_msg_id = message.id.clone();
168 }
169 }
170 saw_parent_usage = true;
171
172 let context_size = input + cache_read + cache_create;
175 if context_size > 0 {
176 session.context_tokens = context_size;
177 }
178 }
179
180 if let Some(model) = message.model {
181 session.model = shorten_model(&model);
182 }
183
184 for block in message.content {
185 match &block {
186 TranscriptBlock::ToolUse { name, input } => {
187 record_tool_usage(name, input, session);
188 session.pending_tool_name = Some(name.clone());
190 session.pending_tool_input = input
191 .get("command")
192 .and_then(|v| v.as_str())
193 .map(|s| s.to_string());
194 session.pending_file_path = if matches!(
196 name.as_str(),
197 "Edit" | "Write" | "NotebookEdit"
198 ) {
199 input
200 .get("file_path")
201 .and_then(|v| v.as_str())
202 .map(|s| s.to_string())
203 } else {
204 None
205 };
206 }
207 TranscriptBlock::ToolResult {
208 is_error, content, ..
209 } => {
210 session.last_tool_error = *is_error;
211 if *is_error {
212 session.total_error_count += 1;
213 session.current_window_errors += 1;
214 let truncated = if content.len() > 256 {
215 format!(
216 "{}...",
217 crate::session::truncate_str(content, 256)
218 )
219 } else {
220 content.clone()
221 };
222 let tool_name = session
223 .pending_tool_name
224 .clone()
225 .unwrap_or_else(|| "?".into());
226 session.last_error_message =
227 Some(truncated.clone());
228 session.recent_errors.push(
229 crate::session::ErrorEntry {
230 tool_name,
231 message: truncated,
232 },
233 );
234 if session.recent_errors.len() > 5 {
235 session.recent_errors.remove(0);
236 }
237 } else {
238 session.last_error_message = None;
239 }
240 session.pending_tool_name = None;
242 session.pending_tool_input = None;
243 session.pending_file_path = None;
244 }
245 _ => {}
246 }
247 }
248 }
249 }
250 }
251 }
252
253 if recognized_events > 0 || session.telemetry_status.is_available() {
254 session.telemetry_status = TelemetryStatus::Available;
255 } else if saw_non_empty_line {
256 session.telemetry_status = TelemetryStatus::UnsupportedTranscript;
257 } else {
258 session.telemetry_status = TelemetryStatus::Pending;
259 }
260
261 session.jsonl_offset = file_len;
262 }
263
264 if let Ok(meta) = std::fs::metadata(path) {
265 if let Ok(modified) = meta.modified() {
266 let mtime_ms = modified
267 .duration_since(std::time::UNIX_EPOCH)
268 .unwrap_or_default()
269 .as_millis() as u64;
270 session.last_message_ts = mtime_ms;
271 }
272 }
273 }
274 None => {
275 session.telemetry_status = TelemetryStatus::MissingTranscript;
276 }
277 }
278
279 finalize_usage(
280 session,
281 &last_type,
282 &last_stop_reason,
283 is_waiting_for_task,
284 saw_parent_usage,
285 );
286}
287
288fn finalize_usage(
289 session: &mut ClaudeSession,
290 last_type: &str,
291 last_stop_reason: &str,
292 is_waiting_for_task: bool,
293 saw_parent_usage: bool,
294) {
295 let resolved_profile = models::resolve(&session.model);
296 session.context_max = resolved_profile.profile.context_max;
297 session.model_profile_source = resolved_profile.source.label().to_string();
298
299 let subagent_rollup = refresh_subagent_rollups(session);
300 session.subagent_input_tokens = subagent_rollup.total_input_tokens();
301 session.subagent_output_tokens = subagent_rollup.output_tokens;
302 session.subagent_cache_read_tokens = subagent_rollup.cache_read_tokens;
303 session.subagent_cache_write_tokens = subagent_rollup.cache_write_tokens;
304 session.subagent_cache_write_1h_tokens = subagent_rollup.cache_write_1h_tokens;
305 session.subagent_count = session.subagent_rollups.len();
306
307 session.total_input_tokens = session.own_input_tokens + session.subagent_input_tokens;
308 session.total_output_tokens = session.own_output_tokens + session.subagent_output_tokens;
309 session.cache_read_tokens = session.own_cache_read_tokens + session.subagent_cache_read_tokens;
310 session.cache_write_tokens =
311 session.own_cache_write_tokens + session.subagent_cache_write_tokens;
312
313 let own_usage_metrics_available = saw_parent_usage
314 || session.own_input_tokens > 0
315 || session.own_output_tokens > 0
316 || session.own_cache_read_tokens > 0
317 || session.own_cache_write_tokens > 0;
318 let (own_cost, own_cost_unverified) = estimate_cost_components(
319 &session.model,
320 session.own_input_tokens,
321 session.own_output_tokens,
322 session.own_cache_read_tokens,
323 session.own_cache_write_tokens,
324 session.own_cache_write_1h_tokens,
325 );
326 session.cost_usd = own_cost + subagent_rollup.cost_usd;
327 session.usage_metrics_available =
328 own_usage_metrics_available || subagent_rollup.usage_metrics_available;
329 session.cost_estimate_unverified = (own_usage_metrics_available && own_cost_unverified)
330 || subagent_rollup.cost_estimate_unverified;
331
332 session.last_msg_type = last_type.to_string();
334 session.last_stop_reason = last_stop_reason.to_string();
335 session.is_waiting_for_task = is_waiting_for_task;
336
337 infer_status(session, last_type, last_stop_reason, is_waiting_for_task);
338}
339
340pub fn infer_status(
341 session: &mut ClaudeSession,
342 last_msg_type: &str,
343 last_stop_reason: &str,
344 is_waiting_for_task: bool,
345) {
346 if session.status == SessionStatus::Finished {
357 return;
358 }
359
360 if session.cpu_percent > 5.0 {
363 session.status = SessionStatus::Processing;
364 return;
365 }
366
367 if is_waiting_for_task {
369 session.status = SessionStatus::NeedsInput;
370 return;
371 }
372
373 if !session.telemetry_status.is_available() && last_msg_type.is_empty() {
374 session.status = SessionStatus::Unknown;
375 return;
376 }
377
378 if last_msg_type == "assistant" && last_stop_reason == "end_turn" {
379 let now_ms = std::time::SystemTime::now()
382 .duration_since(std::time::UNIX_EPOCH)
383 .unwrap_or_default()
384 .as_millis() as u64;
385 let age_mins = (now_ms.saturating_sub(session.last_message_ts)) / 60_000;
386
387 if age_mins > 10 {
388 session.status = SessionStatus::Idle;
389 } else {
390 session.status = SessionStatus::WaitingInput;
391 }
392 return;
393 }
394
395 if last_msg_type == "assistant" && last_stop_reason == "tool_use" {
396 let has_pending_tool = session.pending_tool_name.is_some();
405
406 if session.cpu_percent < 2.0 && has_pending_tool {
407 session.status = SessionStatus::NeedsInput;
408 } else if session.cpu_percent < 2.0 {
409 let now_ms = std::time::SystemTime::now()
410 .duration_since(std::time::UNIX_EPOCH)
411 .unwrap_or_default()
412 .as_millis() as u64;
413 let age_secs = (now_ms.saturating_sub(session.last_message_ts)) / 1000;
414 if age_secs > 5 {
415 session.status = SessionStatus::NeedsInput;
416 } else {
417 session.status = SessionStatus::Processing;
418 }
419 } else {
420 session.status = SessionStatus::Processing;
421 }
422 return;
423 }
424
425 if last_msg_type == "user" {
426 if session.cpu_percent > 1.0 {
428 session.status = SessionStatus::Processing;
429 } else {
430 session.status = SessionStatus::Processing;
432 }
433 return;
434 }
435
436 session.status = SessionStatus::Idle;
437}
438
439#[allow(dead_code)]
441pub fn estimate_cost(session: &ClaudeSession) -> f64 {
442 estimate_cost_components(
443 &session.model,
444 session.total_input_tokens,
445 session.total_output_tokens,
446 session.cache_read_tokens,
447 session.cache_write_tokens,
448 session.own_cache_write_1h_tokens + session.subagent_cache_write_1h_tokens,
449 )
450 .0
451}
452
453pub fn model_context_max(model: &str) -> u64 {
455 models::resolve(model).profile.context_max
456}
457
458fn record_tool_usage(tool_name: &str, input: &Value, session: &mut ClaudeSession) {
460 if tool_name.is_empty() {
461 return;
462 }
463
464 session
465 .tool_usage
466 .entry(tool_name.to_string())
467 .or_default()
468 .calls += 1;
469
470 if crate::transcript::is_edit_tool(tool_name) {
471 if let Some(path) = crate::transcript::edited_path(input) {
472 *session.files_modified.entry(path.to_string()).or_insert(0) += 1;
473 session.file_reads_since_edit.remove(path);
475 }
476 let total_tokens = session.total_input_tokens + session.total_output_tokens;
478 session.total_tokens_at_edit_count += total_tokens;
479 session.edit_event_count += 1;
480 if session.baseline_tokens_per_edit.is_none() && session.edit_event_count >= 5 {
482 session.baseline_tokens_per_edit =
483 Some(session.total_tokens_at_edit_count as f64 / session.edit_event_count as f64);
484 }
485 }
486
487 if matches!(tool_name, "Read" | "Grep" | "Glob") {
489 if let Some(path) = input.get("file_path").and_then(|p| p.as_str()) {
490 *session
491 .file_reads_since_edit
492 .entry(path.to_string())
493 .or_insert(0) += 1;
494 }
495 }
496}
497
498pub fn shorten_model(model: &str) -> String {
499 models::shorten_model(model)
500}
501
502fn refresh_subagent_rollups(session: &mut ClaudeSession) -> UsageRollup {
503 for path in session.active_subagent_jsonl_paths.clone() {
504 let rollup = session.subagent_rollups.entry(path.clone()).or_default();
505 update_subagent_rollup(&path, rollup, &session.model);
506 }
507
508 let mut totals = UsageRollup::default();
509 for rollup in session.subagent_rollups.values() {
510 totals.input_tokens += rollup.input_tokens;
511 totals.output_tokens += rollup.output_tokens;
512 totals.cache_read_tokens += rollup.cache_read_tokens;
513 totals.cache_write_tokens += rollup.cache_write_tokens;
514 totals.cache_write_1h_tokens += rollup.cache_write_1h_tokens;
515 totals.cost_usd += rollup.cost_usd;
516 totals.usage_metrics_available |= rollup.usage_metrics_available;
517 totals.cost_estimate_unverified |= rollup.cost_estimate_unverified;
518 }
519 totals
520}
521
522fn update_subagent_rollup(
523 path: &std::path::Path,
524 rollup: &mut SubagentRollup,
525 default_model: &str,
526) {
527 let mut file = match File::open(path) {
528 Ok(file) => file,
529 Err(_) => return,
530 };
531
532 let file_len = file.metadata().map(|meta| meta.len()).unwrap_or(0);
533 if rollup.jsonl_offset > file_len {
534 *rollup = SubagentRollup::default();
535 }
536
537 if rollup.jsonl_offset >= file_len {
538 rollup.jsonl_offset = file_len;
539 return;
540 }
541
542 if rollup.jsonl_offset > 0 && file.seek(SeekFrom::Start(rollup.jsonl_offset)).is_err() {
543 return;
544 }
545
546 let mut current_model = if rollup.model.is_empty() {
547 default_model.to_string()
548 } else {
549 rollup.model.clone()
550 };
551
552 let reader = BufReader::new(&file);
553 for line in reader.lines() {
554 let Ok(line) = line else {
555 break;
556 };
557 let Some(TranscriptEvent::Message(message)) = parse_line(&line) else {
558 continue;
559 };
560
561 if let Some(model) = message.model {
562 current_model = shorten_model(&model);
563 rollup.model = current_model.clone();
564 }
565
566 let Some(usage) = message.usage else {
567 continue;
568 };
569
570 let repeat = match (&message.id, &rollup.last_usage_msg_id) {
575 (Some(id), Some(prev)) => id == prev,
576 _ => false,
577 };
578 if repeat {
579 continue;
580 }
581 if message.id.is_some() {
582 rollup.last_usage_msg_id = message.id.clone();
583 }
584
585 rollup.input_tokens += usage.input_tokens;
586 rollup.output_tokens += usage.output_tokens;
587 rollup.cache_read_tokens += usage.cache_read_input_tokens;
588 rollup.cache_write_tokens += usage.cache_creation_input_tokens;
589 rollup.cache_write_1h_tokens += usage.cache_creation_1h_input_tokens;
590 rollup.usage_metrics_available = true;
591
592 let input_with_cache =
593 usage.input_tokens + usage.cache_read_input_tokens + usage.cache_creation_input_tokens;
594 let model_for_cost = if current_model.is_empty() {
595 default_model
596 } else {
597 current_model.as_str()
598 };
599 let (delta_cost, unverified) = estimate_cost_components(
600 model_for_cost,
601 input_with_cache,
602 usage.output_tokens,
603 usage.cache_read_input_tokens,
604 usage.cache_creation_input_tokens,
605 usage.cache_creation_1h_input_tokens,
606 );
607 rollup.cost_usd += delta_cost;
608 rollup.cost_estimate_unverified |= unverified;
609 }
610
611 rollup.jsonl_offset = file_len;
612}
613
614fn estimate_cost_components(
615 model: &str,
616 total_input_tokens: u64,
617 total_output_tokens: u64,
618 cache_read_tokens: u64,
619 cache_write_tokens: u64,
620 cache_write_1h_tokens: u64,
621) -> (f64, bool) {
622 let plain_input = total_input_tokens
623 .saturating_sub(cache_read_tokens)
624 .saturating_sub(cache_write_tokens);
625 let resolved = models::resolve(model);
626
627 let write_1h = cache_write_1h_tokens.min(cache_write_tokens);
632 let write_5m = cache_write_tokens.saturating_sub(write_1h);
633
634 let cost = (plain_input as f64 / 1_000_000.0) * resolved.profile.input_per_m
635 + (total_output_tokens as f64 / 1_000_000.0) * resolved.profile.output_per_m
636 + (cache_read_tokens as f64 / 1_000_000.0) * resolved.profile.cache_read_per_m
637 + (write_5m as f64 / 1_000_000.0) * resolved.profile.cache_write_per_m
638 + (write_1h as f64 / 1_000_000.0) * resolved.profile.cache_write_1h_per_m();
639
640 (
641 cost,
642 resolved.source == models::ModelProfileSource::Fallback,
643 )
644}
645
646#[cfg(test)]
647mod edit_tool_tests {
648 use super::record_tool_usage;
649 use crate::session::{ClaudeSession, RawSession};
650 use serde_json::json;
651
652 fn session() -> ClaudeSession {
653 ClaudeSession::from_raw(RawSession {
654 pid: 100,
655 session_id: "test".into(),
656 cwd: "/repo".into(),
657 started_at: 0,
658 })
659 }
660
661 #[test]
664 fn every_edit_tool_counts_toward_files_modified() {
665 let mut s = session();
666 record_tool_usage("Edit", &json!({"file_path": "/repo/a.rs"}), &mut s);
667 record_tool_usage("Write", &json!({"file_path": "/repo/a.rs"}), &mut s);
668 record_tool_usage("MultiEdit", &json!({"file_path": "/repo/b.rs"}), &mut s);
669 record_tool_usage(
670 "NotebookEdit",
671 &json!({"notebook_path": "/repo/n.ipynb"}),
672 &mut s,
673 );
674 record_tool_usage("Read", &json!({"file_path": "/repo/c.rs"}), &mut s);
675 assert_eq!(s.files_modified.get("/repo/a.rs"), Some(&2));
676 assert_eq!(s.files_modified.get("/repo/b.rs"), Some(&1));
677 assert_eq!(s.files_modified.get("/repo/n.ipynb"), Some(&1));
678 assert_eq!(
679 s.files_modified.get("/repo/c.rs"),
680 None,
681 "reads are not edits"
682 );
683 assert_eq!(
684 s.tool_usage["Read"].calls, 1,
685 "but every tool is still counted"
686 );
687 }
688
689 #[test]
690 fn an_edit_resets_the_reads_since_edit_for_that_path() {
691 let mut s = session();
692 s.file_reads_since_edit.insert("/repo/b.rs".into(), 3);
693 record_tool_usage("MultiEdit", &json!({"file_path": "/repo/b.rs"}), &mut s);
694 assert_eq!(s.file_reads_since_edit.get("/repo/b.rs"), None);
695 }
696}