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 cost_usd: f64,
17 usage_metrics_available: bool,
18 cost_estimate_unverified: bool,
19}
20
21impl UsageRollup {
22 fn total_input_tokens(&self) -> u64 {
23 self.input_tokens + self.cache_read_tokens + self.cache_write_tokens
24 }
25}
26
27pub fn update_tokens(session: &mut ClaudeSession) {
29 let mut last_type = session.last_msg_type.clone();
31 let mut last_stop_reason = session.last_stop_reason.clone();
32 let mut is_waiting_for_task = session.is_waiting_for_task;
33 let mut saw_non_empty_line = false;
34 let mut recognized_events = 0usize;
35 let mut saw_parent_usage = false;
36 let jsonl_path = session.jsonl_path.clone();
37
38 match jsonl_path.as_ref() {
39 Some(path) => {
40 let mut file = match File::open(path) {
41 Ok(f) => f,
42 Err(_) => {
43 session.telemetry_status = TelemetryStatus::UnreadableTranscript;
44 finalize_usage(
45 session,
46 &last_type,
47 &last_stop_reason,
48 is_waiting_for_task,
49 false,
50 );
51 return;
52 }
53 };
54
55 let file_len = file.metadata().map(|m| m.len()).unwrap_or(0);
56
57 if file_len == 0 {
58 session.telemetry_status = TelemetryStatus::Pending;
59 } else {
60 if session.jsonl_offset > file_len {
61 session.jsonl_offset = 0;
62 session.own_input_tokens = 0;
63 session.own_output_tokens = 0;
64 session.own_cache_read_tokens = 0;
65 session.own_cache_write_tokens = 0;
66 last_type.clear();
68 last_stop_reason.clear();
69 is_waiting_for_task = false;
70 }
71
72 if session.jsonl_offset < file_len {
73 if session.jsonl_offset > 0
74 && file.seek(SeekFrom::Start(session.jsonl_offset)).is_err()
75 {
76 finalize_usage(
77 session,
78 &last_type,
79 &last_stop_reason,
80 is_waiting_for_task,
81 false,
82 );
83 return;
84 }
85
86 let reader = BufReader::new(&file);
87
88 for line in reader.lines() {
89 let line = match line {
90 Ok(l) => l,
91 Err(_) => break,
92 };
93
94 if line.trim().is_empty() {
95 continue;
96 }
97 saw_non_empty_line = true;
98
99 let Some(event) = parse_line(&line) else {
100 continue;
101 };
102 recognized_events += 1;
103
104 match event {
105 TranscriptEvent::WaitingForTask => {
106 is_waiting_for_task = true;
107 }
108 TranscriptEvent::Message(message) => {
109 is_waiting_for_task = false;
110 last_type = match message.role {
111 TranscriptRole::Assistant => "assistant".to_string(),
112 TranscriptRole::User => "user".to_string(),
113 };
114
115 if let Some(reason) = message.stop_reason {
116 last_stop_reason = reason;
117 } else {
118 let has_tool_use = message
122 .content
123 .iter()
124 .any(|b| matches!(b, TranscriptBlock::ToolUse { .. }));
125 if has_tool_use {
126 last_stop_reason = "tool_use".to_string();
127 } else {
128 last_stop_reason.clear();
129 }
130 }
131
132 if let Some(usage) = message.usage {
133 let input = usage.input_tokens;
134 let cache_read = usage.cache_read_input_tokens;
135 let cache_create = usage.cache_creation_input_tokens;
136 let output = usage.output_tokens;
137
138 session.own_input_tokens += input + cache_read + cache_create;
139 session.own_output_tokens += output;
140 session.own_cache_read_tokens += cache_read;
141 session.own_cache_write_tokens += cache_create;
142 saw_parent_usage = true;
143
144 let context_size = input + cache_read + cache_create;
147 if context_size > 0 {
148 session.context_tokens = context_size;
149 }
150 }
151
152 if let Some(model) = message.model {
153 session.model = shorten_model(&model);
154 }
155
156 for block in message.content {
157 match &block {
158 TranscriptBlock::ToolUse { name, input } => {
159 record_tool_usage(name, input, session);
160 session.pending_tool_name = Some(name.clone());
162 session.pending_tool_input = input
163 .get("command")
164 .and_then(|v| v.as_str())
165 .map(|s| s.to_string());
166 session.pending_file_path = if matches!(
168 name.as_str(),
169 "Edit" | "Write" | "NotebookEdit"
170 ) {
171 input
172 .get("file_path")
173 .and_then(|v| v.as_str())
174 .map(|s| s.to_string())
175 } else {
176 None
177 };
178 }
179 TranscriptBlock::ToolResult {
180 is_error, content, ..
181 } => {
182 session.last_tool_error = *is_error;
183 if *is_error {
184 session.total_error_count += 1;
185 session.current_window_errors += 1;
186 let truncated = if content.len() > 256 {
187 format!(
188 "{}...",
189 crate::session::truncate_str(content, 256)
190 )
191 } else {
192 content.clone()
193 };
194 let tool_name = session
195 .pending_tool_name
196 .clone()
197 .unwrap_or_else(|| "?".into());
198 session.last_error_message =
199 Some(truncated.clone());
200 session.recent_errors.push(
201 crate::session::ErrorEntry {
202 tool_name,
203 message: truncated,
204 },
205 );
206 if session.recent_errors.len() > 5 {
207 session.recent_errors.remove(0);
208 }
209 } else {
210 session.last_error_message = None;
211 }
212 session.pending_tool_name = None;
214 session.pending_tool_input = None;
215 session.pending_file_path = None;
216 }
217 _ => {}
218 }
219 }
220 }
221 }
222 }
223 }
224
225 if recognized_events > 0 || session.telemetry_status.is_available() {
226 session.telemetry_status = TelemetryStatus::Available;
227 } else if saw_non_empty_line {
228 session.telemetry_status = TelemetryStatus::UnsupportedTranscript;
229 } else {
230 session.telemetry_status = TelemetryStatus::Pending;
231 }
232
233 session.jsonl_offset = file_len;
234 }
235
236 if let Ok(meta) = std::fs::metadata(path) {
237 if let Ok(modified) = meta.modified() {
238 let mtime_ms = modified
239 .duration_since(std::time::UNIX_EPOCH)
240 .unwrap_or_default()
241 .as_millis() as u64;
242 session.last_message_ts = mtime_ms;
243 }
244 }
245 }
246 None => {
247 session.telemetry_status = TelemetryStatus::MissingTranscript;
248 }
249 }
250
251 finalize_usage(
252 session,
253 &last_type,
254 &last_stop_reason,
255 is_waiting_for_task,
256 saw_parent_usage,
257 );
258}
259
260fn finalize_usage(
261 session: &mut ClaudeSession,
262 last_type: &str,
263 last_stop_reason: &str,
264 is_waiting_for_task: bool,
265 saw_parent_usage: bool,
266) {
267 let resolved_profile = models::resolve(&session.model);
268 session.context_max = resolved_profile.profile.context_max;
269 session.model_profile_source = resolved_profile.source.label().to_string();
270
271 let subagent_rollup = refresh_subagent_rollups(session);
272 session.subagent_input_tokens = subagent_rollup.total_input_tokens();
273 session.subagent_output_tokens = subagent_rollup.output_tokens;
274 session.subagent_cache_read_tokens = subagent_rollup.cache_read_tokens;
275 session.subagent_cache_write_tokens = subagent_rollup.cache_write_tokens;
276 session.subagent_count = session.subagent_rollups.len();
277
278 session.total_input_tokens = session.own_input_tokens + session.subagent_input_tokens;
279 session.total_output_tokens = session.own_output_tokens + session.subagent_output_tokens;
280 session.cache_read_tokens = session.own_cache_read_tokens + session.subagent_cache_read_tokens;
281 session.cache_write_tokens =
282 session.own_cache_write_tokens + session.subagent_cache_write_tokens;
283
284 let own_usage_metrics_available = saw_parent_usage
285 || session.own_input_tokens > 0
286 || session.own_output_tokens > 0
287 || session.own_cache_read_tokens > 0
288 || session.own_cache_write_tokens > 0;
289 let (own_cost, own_cost_unverified) = estimate_cost_components(
290 &session.model,
291 session.own_input_tokens,
292 session.own_output_tokens,
293 session.own_cache_read_tokens,
294 session.own_cache_write_tokens,
295 );
296 session.cost_usd = own_cost + subagent_rollup.cost_usd;
297 session.usage_metrics_available =
298 own_usage_metrics_available || subagent_rollup.usage_metrics_available;
299 session.cost_estimate_unverified = (own_usage_metrics_available && own_cost_unverified)
300 || subagent_rollup.cost_estimate_unverified;
301
302 session.last_msg_type = last_type.to_string();
304 session.last_stop_reason = last_stop_reason.to_string();
305 session.is_waiting_for_task = is_waiting_for_task;
306
307 infer_status(session, last_type, last_stop_reason, is_waiting_for_task);
308}
309
310pub fn infer_status(
311 session: &mut ClaudeSession,
312 last_msg_type: &str,
313 last_stop_reason: &str,
314 is_waiting_for_task: bool,
315) {
316 if session.cpu_percent > 5.0 {
319 session.status = SessionStatus::Processing;
320 return;
321 }
322
323 if is_waiting_for_task {
325 session.status = SessionStatus::NeedsInput;
326 return;
327 }
328
329 if !session.telemetry_status.is_available() && last_msg_type.is_empty() {
330 session.status = SessionStatus::Unknown;
331 return;
332 }
333
334 if last_msg_type == "assistant" && last_stop_reason == "end_turn" {
335 let now_ms = std::time::SystemTime::now()
338 .duration_since(std::time::UNIX_EPOCH)
339 .unwrap_or_default()
340 .as_millis() as u64;
341 let age_mins = (now_ms.saturating_sub(session.last_message_ts)) / 60_000;
342
343 if age_mins > 10 {
344 session.status = SessionStatus::Idle;
345 } else {
346 session.status = SessionStatus::WaitingInput;
347 }
348 return;
349 }
350
351 if last_msg_type == "assistant" && last_stop_reason == "tool_use" {
352 let has_pending_tool = session.pending_tool_name.is_some();
361
362 if session.cpu_percent < 2.0 && has_pending_tool {
363 session.status = SessionStatus::NeedsInput;
364 } else if session.cpu_percent < 2.0 {
365 let now_ms = std::time::SystemTime::now()
366 .duration_since(std::time::UNIX_EPOCH)
367 .unwrap_or_default()
368 .as_millis() as u64;
369 let age_secs = (now_ms.saturating_sub(session.last_message_ts)) / 1000;
370 if age_secs > 5 {
371 session.status = SessionStatus::NeedsInput;
372 } else {
373 session.status = SessionStatus::Processing;
374 }
375 } else {
376 session.status = SessionStatus::Processing;
377 }
378 return;
379 }
380
381 if last_msg_type == "user" {
382 if session.cpu_percent > 1.0 {
384 session.status = SessionStatus::Processing;
385 } else {
386 session.status = SessionStatus::Processing;
388 }
389 return;
390 }
391
392 session.status = SessionStatus::Idle;
393}
394
395#[allow(dead_code)]
397pub fn estimate_cost(session: &ClaudeSession) -> f64 {
398 estimate_cost_components(
399 &session.model,
400 session.total_input_tokens,
401 session.total_output_tokens,
402 session.cache_read_tokens,
403 session.cache_write_tokens,
404 )
405 .0
406}
407
408pub fn model_context_max(model: &str) -> u64 {
410 models::resolve(model).profile.context_max
411}
412
413fn record_tool_usage(tool_name: &str, input: &Value, session: &mut ClaudeSession) {
415 if tool_name.is_empty() {
416 return;
417 }
418
419 session
420 .tool_usage
421 .entry(tool_name.to_string())
422 .or_default()
423 .calls += 1;
424
425 if matches!(tool_name, "Edit" | "Write" | "NotebookEdit") {
426 if let Some(path) = input.get("file_path").and_then(|p| p.as_str()) {
427 *session.files_modified.entry(path.to_string()).or_insert(0) += 1;
428 session.file_reads_since_edit.remove(path);
430 }
431 let total_tokens = session.total_input_tokens + session.total_output_tokens;
433 session.total_tokens_at_edit_count += total_tokens;
434 session.edit_event_count += 1;
435 if session.baseline_tokens_per_edit.is_none() && session.edit_event_count >= 5 {
437 session.baseline_tokens_per_edit =
438 Some(session.total_tokens_at_edit_count as f64 / session.edit_event_count as f64);
439 }
440 }
441
442 if matches!(tool_name, "Read" | "Grep" | "Glob") {
444 if let Some(path) = input.get("file_path").and_then(|p| p.as_str()) {
445 *session
446 .file_reads_since_edit
447 .entry(path.to_string())
448 .or_insert(0) += 1;
449 }
450 }
451}
452
453pub fn shorten_model(model: &str) -> String {
454 models::shorten_model(model)
455}
456
457fn refresh_subagent_rollups(session: &mut ClaudeSession) -> UsageRollup {
458 for path in session.active_subagent_jsonl_paths.clone() {
459 let rollup = session.subagent_rollups.entry(path.clone()).or_default();
460 update_subagent_rollup(&path, rollup, &session.model);
461 }
462
463 let mut totals = UsageRollup::default();
464 for rollup in session.subagent_rollups.values() {
465 totals.input_tokens += rollup.input_tokens;
466 totals.output_tokens += rollup.output_tokens;
467 totals.cache_read_tokens += rollup.cache_read_tokens;
468 totals.cache_write_tokens += rollup.cache_write_tokens;
469 totals.cost_usd += rollup.cost_usd;
470 totals.usage_metrics_available |= rollup.usage_metrics_available;
471 totals.cost_estimate_unverified |= rollup.cost_estimate_unverified;
472 }
473 totals
474}
475
476fn update_subagent_rollup(
477 path: &std::path::Path,
478 rollup: &mut SubagentRollup,
479 default_model: &str,
480) {
481 let mut file = match File::open(path) {
482 Ok(file) => file,
483 Err(_) => return,
484 };
485
486 let file_len = file.metadata().map(|meta| meta.len()).unwrap_or(0);
487 if rollup.jsonl_offset > file_len {
488 *rollup = SubagentRollup::default();
489 }
490
491 if rollup.jsonl_offset >= file_len {
492 rollup.jsonl_offset = file_len;
493 return;
494 }
495
496 if rollup.jsonl_offset > 0 && file.seek(SeekFrom::Start(rollup.jsonl_offset)).is_err() {
497 return;
498 }
499
500 let mut current_model = if rollup.model.is_empty() {
501 default_model.to_string()
502 } else {
503 rollup.model.clone()
504 };
505
506 let reader = BufReader::new(&file);
507 for line in reader.lines() {
508 let Ok(line) = line else {
509 break;
510 };
511 let Some(TranscriptEvent::Message(message)) = parse_line(&line) else {
512 continue;
513 };
514
515 if let Some(model) = message.model {
516 current_model = shorten_model(&model);
517 rollup.model = current_model.clone();
518 }
519
520 let Some(usage) = message.usage else {
521 continue;
522 };
523
524 rollup.input_tokens += usage.input_tokens;
525 rollup.output_tokens += usage.output_tokens;
526 rollup.cache_read_tokens += usage.cache_read_input_tokens;
527 rollup.cache_write_tokens += usage.cache_creation_input_tokens;
528 rollup.usage_metrics_available = true;
529
530 let input_with_cache =
531 usage.input_tokens + usage.cache_read_input_tokens + usage.cache_creation_input_tokens;
532 let model_for_cost = if current_model.is_empty() {
533 default_model
534 } else {
535 current_model.as_str()
536 };
537 let (delta_cost, unverified) = estimate_cost_components(
538 model_for_cost,
539 input_with_cache,
540 usage.output_tokens,
541 usage.cache_read_input_tokens,
542 usage.cache_creation_input_tokens,
543 );
544 rollup.cost_usd += delta_cost;
545 rollup.cost_estimate_unverified |= unverified;
546 }
547
548 rollup.jsonl_offset = file_len;
549}
550
551fn estimate_cost_components(
552 model: &str,
553 total_input_tokens: u64,
554 total_output_tokens: u64,
555 cache_read_tokens: u64,
556 cache_write_tokens: u64,
557) -> (f64, bool) {
558 let plain_input = total_input_tokens
559 .saturating_sub(cache_read_tokens)
560 .saturating_sub(cache_write_tokens);
561 let resolved = models::resolve(model);
562
563 let cost = (plain_input as f64 / 1_000_000.0) * resolved.profile.input_per_m
564 + (total_output_tokens as f64 / 1_000_000.0) * resolved.profile.output_per_m
565 + (cache_read_tokens as f64 / 1_000_000.0) * resolved.profile.cache_read_per_m
566 + (cache_write_tokens as f64 / 1_000_000.0) * resolved.profile.cache_write_per_m;
567
568 (
569 cost,
570 resolved.source == models::ModelProfileSource::Fallback,
571 )
572}