Skip to main content

batty_cli/team/
cost.rs

1use std::collections::{BTreeSet, HashMap};
2use std::fs::{self, File};
3use std::io::{BufRead, BufReader};
4use std::path::{Path, PathBuf};
5
6use anyhow::{Context, Result, bail};
7use serde::Deserialize;
8use serde_json::Value;
9use tracing::warn;
10
11use super::config::{ModelPricing, TeamConfig};
12use super::hierarchy::{self, MemberInstance};
13use super::{daemon_state_path, status, team_config_path};
14
15#[derive(Debug, Clone, Default, PartialEq, Eq)]
16struct TokenUsage {
17    input_tokens: u64,
18    cached_input_tokens: u64,
19    cache_creation_input_tokens: u64,
20    cache_creation_5m_input_tokens: u64,
21    cache_creation_1h_input_tokens: u64,
22    cache_read_input_tokens: u64,
23    output_tokens: u64,
24    reasoning_output_tokens: u64,
25}
26
27impl TokenUsage {
28    fn total_tokens(&self) -> u64 {
29        self.input_tokens
30            + self.cached_input_tokens
31            + self.cache_creation_input_tokens
32            + self.cache_read_input_tokens
33            + self.output_tokens
34            + self.reasoning_output_tokens
35    }
36
37    fn display_cache_tokens(&self) -> u64 {
38        self.cached_input_tokens + self.cache_creation_input_tokens + self.cache_read_input_tokens
39    }
40}
41
42#[derive(Debug, Clone, PartialEq, Eq)]
43struct SessionUsage {
44    model: Option<String>,
45    usage: TokenUsage,
46}
47
48#[derive(Debug, Clone, PartialEq)]
49struct CostEntry {
50    member_name: String,
51    agent: String,
52    model: String,
53    task: String,
54    session_file: PathBuf,
55    usage: TokenUsage,
56    estimated_cost_usd: Option<f64>,
57}
58
59#[derive(Debug, Clone, PartialEq)]
60struct CostReport {
61    team_name: String,
62    entries: Vec<CostEntry>,
63    total_estimated_cost_usd: f64,
64    unpriced_models: BTreeSet<String>,
65}
66
67#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
68enum SessionAgent {
69    Codex,
70    Claude,
71}
72
73#[derive(Debug, Clone)]
74struct SessionRoots {
75    codex_sessions_root: PathBuf,
76    claude_projects_root: PathBuf,
77}
78
79#[derive(Debug, Deserialize, Default)]
80struct LaunchIdentityRecord {
81    #[serde(default)]
82    session_id: Option<String>,
83}
84
85#[derive(Debug, Deserialize, Default)]
86struct PersistedDaemonStateCostView {
87    #[serde(default)]
88    active_tasks: HashMap<String, u32>,
89}
90
91pub fn show_cost(project_root: &Path) -> Result<()> {
92    let config_path = team_config_path(project_root);
93    if !config_path.exists() {
94        bail!("no team config found at {}", config_path.display());
95    }
96
97    let team_config = TeamConfig::load(&config_path)?;
98    let report = collect_cost_report(project_root, &team_config, &SessionRoots::default())?;
99
100    println!("Run cost estimate for team {}", report.team_name);
101    if report.entries.is_empty() {
102        println!("No agent session files with token usage were found.");
103        return Ok(());
104    }
105
106    println!();
107    println!(
108        "{:<20} {:<12} {:<20} {:<10} {:>10} {:>10} {:>10} {:>12}",
109        "MEMBER", "AGENT", "MODEL", "TASK", "INPUT", "CACHE", "OUTPUT", "COST"
110    );
111    println!("{}", "-".repeat(112));
112    for entry in &report.entries {
113        println!(
114            "{:<20} {:<12} {:<20} {:<10} {:>10} {:>10} {:>10} {:>12}",
115            entry.member_name,
116            entry.agent,
117            truncate_model(&entry.model),
118            entry.task,
119            entry.usage.input_tokens,
120            entry.usage.display_cache_tokens(),
121            entry.usage.output_tokens + entry.usage.reasoning_output_tokens,
122            entry
123                .estimated_cost_usd
124                .map(|cost| format!("${cost:.4}"))
125                .unwrap_or_else(|| "n/a".to_string()),
126        );
127    }
128
129    println!();
130    println!("Estimated total: ${:.4}", report.total_estimated_cost_usd);
131    if !report.unpriced_models.is_empty() {
132        println!(
133            "Unpriced models: {}",
134            report
135                .unpriced_models
136                .iter()
137                .cloned()
138                .collect::<Vec<_>>()
139                .join(", ")
140        );
141    }
142
143    Ok(())
144}
145
146impl Default for SessionRoots {
147    fn default() -> Self {
148        let home = std::env::var_os("HOME")
149            .map(PathBuf::from)
150            .unwrap_or_else(|| PathBuf::from("/"));
151        Self {
152            codex_sessions_root: home.join(".codex").join("sessions"),
153            claude_projects_root: home.join(".claude").join("projects"),
154        }
155    }
156}
157
158fn collect_cost_report(
159    project_root: &Path,
160    team_config: &TeamConfig,
161    session_roots: &SessionRoots,
162) -> Result<CostReport> {
163    let members = hierarchy::resolve_hierarchy(team_config)?;
164    let launch_state = load_launch_state(project_root);
165    let active_tasks = load_active_tasks(project_root);
166    let owned_task_buckets = status::owned_task_buckets(project_root, &members);
167    let mut session_target_counts = HashMap::<(SessionAgent, PathBuf), usize>::new();
168    for member in &members {
169        let Some((agent_kind, session_cwd, _)) = member_session_target(project_root, member) else {
170            continue;
171        };
172        *session_target_counts
173            .entry((agent_kind, session_cwd))
174            .or_insert(0usize) += 1;
175    }
176    let mut entries = Vec::new();
177    let mut total_estimated_cost_usd = 0.0;
178    let mut unpriced_models = BTreeSet::new();
179
180    for member in &members {
181        let Some((agent_kind, session_cwd, agent_label)) =
182            member_session_target(project_root, member)
183        else {
184            continue;
185        };
186        let session_id = launch_state
187            .get(&member.name)
188            .and_then(|identity| identity.session_id.as_deref());
189        let allow_cwd_fallback = session_id.is_some()
190            || session_target_counts
191                .get(&(agent_kind, session_cwd.clone()))
192                .copied()
193                .unwrap_or(0)
194                <= 1;
195        let session_file = match agent_kind {
196            SessionAgent::Codex => discover_codex_session_file(
197                &session_roots.codex_sessions_root,
198                &session_cwd,
199                session_id,
200                allow_cwd_fallback,
201            )?,
202            SessionAgent::Claude => discover_claude_session_file(
203                &session_roots.claude_projects_root,
204                &session_cwd,
205                session_id,
206                allow_cwd_fallback,
207            )?,
208        };
209        let Some(session_file) = session_file else {
210            continue;
211        };
212
213        let session_usage = match agent_kind {
214            SessionAgent::Codex => parse_codex_session_usage(&session_file)?,
215            SessionAgent::Claude => parse_claude_session_usage(&session_file)?,
216        };
217        let Some(session_usage) = session_usage else {
218            continue;
219        };
220        if session_usage.usage.total_tokens() == 0 {
221            continue;
222        }
223
224        let model = session_usage.model.unwrap_or_else(|| "unknown".to_string());
225        let estimated_cost_usd = pricing_for_model(&team_config.cost.models, &model)
226            .map(|pricing| estimate_cost_usd(&session_usage.usage, &pricing));
227        if let Some(cost) = estimated_cost_usd {
228            total_estimated_cost_usd += cost;
229        } else {
230            unpriced_models.insert(model.clone());
231        }
232
233        let task = active_tasks
234            .get(&member.name)
235            .map(|task_id| format!("#{task_id}"))
236            .or_else(|| {
237                owned_task_buckets
238                    .get(&member.name)
239                    .map(|buckets| status::format_owned_tasks_summary(&buckets.active))
240            })
241            .unwrap_or_else(|| "-".to_string());
242
243        entries.push(CostEntry {
244            member_name: member.name.clone(),
245            agent: agent_label.to_string(),
246            model,
247            task,
248            session_file,
249            usage: session_usage.usage,
250            estimated_cost_usd,
251        });
252    }
253
254    entries.sort_by(|left, right| left.member_name.cmp(&right.member_name));
255
256    Ok(CostReport {
257        team_name: team_config.name.clone(),
258        entries,
259        total_estimated_cost_usd,
260        unpriced_models,
261    })
262}
263
264fn truncate_model(model: &str) -> String {
265    const MAX_LEN: usize = 20;
266    if model.chars().count() <= MAX_LEN {
267        model.to_string()
268    } else {
269        let short = model.chars().take(MAX_LEN - 3).collect::<String>();
270        format!("{short}...")
271    }
272}
273
274fn member_session_target(
275    project_root: &Path,
276    member: &MemberInstance,
277) -> Option<(SessionAgent, PathBuf, &'static str)> {
278    if member.role_type == super::config::RoleType::User {
279        return None;
280    }
281
282    let work_dir = if member.use_worktrees {
283        project_root
284            .join(".batty")
285            .join("worktrees")
286            .join(&member.name)
287    } else {
288        project_root.to_path_buf()
289    };
290
291    match member.agent.as_deref() {
292        Some("codex") | Some("codex-cli") => Some((
293            SessionAgent::Codex,
294            work_dir
295                .join(".batty")
296                .join("codex-context")
297                .join(&member.name),
298            "codex",
299        )),
300        Some("claude") | Some("claude-code") | None => {
301            Some((SessionAgent::Claude, work_dir, "claude"))
302        }
303        _ => None,
304    }
305}
306
307fn load_launch_state(project_root: &Path) -> HashMap<String, LaunchIdentityRecord> {
308    let path = project_root.join(".batty").join("launch-state.json");
309    let Ok(content) = fs::read_to_string(&path) else {
310        return HashMap::new();
311    };
312    match serde_json::from_str::<HashMap<String, LaunchIdentityRecord>>(&content) {
313        Ok(state) => state,
314        Err(error) => {
315            warn!(path = %path.display(), error = %error, "failed to parse launch state for cost reporting");
316            HashMap::new()
317        }
318    }
319}
320
321fn load_active_tasks(project_root: &Path) -> HashMap<String, u32> {
322    let path = daemon_state_path(project_root);
323    let Ok(content) = fs::read_to_string(&path) else {
324        return HashMap::new();
325    };
326    match serde_json::from_str::<PersistedDaemonStateCostView>(&content) {
327        Ok(state) => state.active_tasks,
328        Err(error) => {
329            warn!(path = %path.display(), error = %error, "failed to parse daemon state for cost reporting");
330            HashMap::new()
331        }
332    }
333}
334
335fn pricing_for_model(
336    overrides: &HashMap<String, ModelPricing>,
337    model: &str,
338) -> Option<ModelPricing> {
339    let normalized = model.to_ascii_lowercase();
340    if let Some(pricing) = overrides.get(&normalized) {
341        return Some(pricing.clone());
342    }
343    if let Some(pricing) = overrides.get(model) {
344        return Some(pricing.clone());
345    }
346    built_in_model_pricing(&normalized)
347}
348
349fn built_in_model_pricing(model: &str) -> Option<ModelPricing> {
350    // Defaults keep the command useful out of the box. team.yaml can override any model entry.
351    if model.starts_with("gpt-5.4") {
352        return Some(ModelPricing {
353            input_usd_per_mtok: 2.5,
354            cached_input_usd_per_mtok: 0.25,
355            cache_creation_input_usd_per_mtok: None,
356            cache_creation_5m_input_usd_per_mtok: None,
357            cache_creation_1h_input_usd_per_mtok: None,
358            cache_read_input_usd_per_mtok: 0.0,
359            output_usd_per_mtok: 15.0,
360            reasoning_output_usd_per_mtok: Some(15.0),
361        });
362    }
363    if model.starts_with("claude-opus-4") {
364        return Some(ModelPricing {
365            input_usd_per_mtok: 15.0,
366            cached_input_usd_per_mtok: 0.0,
367            cache_creation_input_usd_per_mtok: None,
368            cache_creation_5m_input_usd_per_mtok: Some(18.75),
369            cache_creation_1h_input_usd_per_mtok: Some(30.0),
370            cache_read_input_usd_per_mtok: 1.5,
371            output_usd_per_mtok: 75.0,
372            reasoning_output_usd_per_mtok: None,
373        });
374    }
375    if model.starts_with("claude-sonnet-4") {
376        return Some(ModelPricing {
377            input_usd_per_mtok: 3.0,
378            cached_input_usd_per_mtok: 0.0,
379            cache_creation_input_usd_per_mtok: None,
380            cache_creation_5m_input_usd_per_mtok: Some(3.75),
381            cache_creation_1h_input_usd_per_mtok: Some(6.0),
382            cache_read_input_usd_per_mtok: 0.3,
383            output_usd_per_mtok: 15.0,
384            reasoning_output_usd_per_mtok: None,
385        });
386    }
387    None
388}
389
390fn estimate_cost_usd(usage: &TokenUsage, pricing: &ModelPricing) -> f64 {
391    let classified_cache_creation =
392        usage.cache_creation_5m_input_tokens + usage.cache_creation_1h_input_tokens;
393    let unclassified_cache_creation = usage
394        .cache_creation_input_tokens
395        .saturating_sub(classified_cache_creation);
396    let reasoning_rate = pricing
397        .reasoning_output_usd_per_mtok
398        .unwrap_or(pricing.output_usd_per_mtok);
399    let cache_creation_generic_rate = pricing
400        .cache_creation_input_usd_per_mtok
401        .or(pricing.cache_creation_5m_input_usd_per_mtok)
402        .unwrap_or(pricing.input_usd_per_mtok);
403
404    let total_usd = (usage.input_tokens as f64 * pricing.input_usd_per_mtok)
405        + (usage.cached_input_tokens as f64 * pricing.cached_input_usd_per_mtok)
406        + (usage.cache_creation_5m_input_tokens as f64
407            * pricing
408                .cache_creation_5m_input_usd_per_mtok
409                .unwrap_or(cache_creation_generic_rate))
410        + (usage.cache_creation_1h_input_tokens as f64
411            * pricing
412                .cache_creation_1h_input_usd_per_mtok
413                .unwrap_or(cache_creation_generic_rate))
414        + (unclassified_cache_creation as f64 * cache_creation_generic_rate)
415        + (usage.cache_read_input_tokens as f64 * pricing.cache_read_input_usd_per_mtok)
416        + (usage.output_tokens as f64 * pricing.output_usd_per_mtok)
417        + (usage.reasoning_output_tokens as f64 * reasoning_rate);
418
419    total_usd / 1_000_000.0
420}
421
422fn discover_codex_session_file(
423    sessions_root: &Path,
424    cwd: &Path,
425    session_id: Option<&str>,
426    allow_cwd_fallback: bool,
427) -> Result<Option<PathBuf>> {
428    if !sessions_root.exists() {
429        return Ok(None);
430    }
431
432    let mut newest: Option<(std::time::SystemTime, PathBuf)> = None;
433    for year in read_dir_paths(sessions_root)? {
434        for month in read_dir_paths(&year)? {
435            for day in read_dir_paths(&month)? {
436                for entry in read_dir_paths(&day)? {
437                    if entry.extension().and_then(|ext| ext.to_str()) != Some("jsonl") {
438                        continue;
439                    }
440                    let Some(meta) = read_codex_session_meta(&entry)? else {
441                        continue;
442                    };
443                    if meta.cwd.as_deref() != Some(cwd.as_os_str()) {
444                        continue;
445                    }
446                    if let Some(wanted) = session_id
447                        && meta.id.as_deref() == Some(wanted)
448                    {
449                        return Ok(Some(entry));
450                    }
451                    let modified = fs::metadata(&entry)
452                        .and_then(|metadata| metadata.modified())
453                        .unwrap_or(std::time::SystemTime::UNIX_EPOCH);
454                    match &newest {
455                        Some((current, _)) if modified <= *current => {}
456                        _ => newest = Some((modified, entry)),
457                    }
458                }
459            }
460        }
461    }
462
463    Ok(allow_cwd_fallback
464        .then_some(newest)
465        .flatten()
466        .map(|(_, path)| path))
467}
468
469fn discover_claude_session_file(
470    projects_root: &Path,
471    cwd: &Path,
472    session_id: Option<&str>,
473    allow_cwd_fallback: bool,
474) -> Result<Option<PathBuf>> {
475    if !projects_root.exists() {
476        return Ok(None);
477    }
478
479    let preferred_dir = projects_root.join(cwd.to_string_lossy().replace('/', "-"));
480    if let Some(session_id) = session_id {
481        let exact = preferred_dir.join(format!("{session_id}.jsonl"));
482        if exact.is_file() {
483            return Ok(Some(exact));
484        }
485    }
486
487    let mut newest: Option<(std::time::SystemTime, PathBuf)> = None;
488    if preferred_dir.is_dir() {
489        for entry in read_dir_paths(&preferred_dir)? {
490            if entry.extension().and_then(|ext| ext.to_str()) != Some("jsonl") {
491                continue;
492            }
493            let modified = fs::metadata(&entry)
494                .and_then(|metadata| metadata.modified())
495                .unwrap_or(std::time::SystemTime::UNIX_EPOCH);
496            match &newest {
497                Some((current, _)) if modified <= *current => {}
498                _ => newest = Some((modified, entry)),
499            }
500        }
501    }
502
503    Ok(allow_cwd_fallback
504        .then_some(newest)
505        .flatten()
506        .map(|(_, path)| path))
507}
508
509#[derive(Debug)]
510struct CodexSessionMeta {
511    id: Option<String>,
512    cwd: Option<std::ffi::OsString>,
513}
514
515fn read_codex_session_meta(path: &Path) -> Result<Option<CodexSessionMeta>> {
516    let file = File::open(path)?;
517    let reader = BufReader::new(file);
518    for line in reader.lines() {
519        let line = line?;
520        if line.trim().is_empty() {
521            continue;
522        }
523        let Ok(entry) = serde_json::from_str::<Value>(&line) else {
524            continue;
525        };
526        if entry.get("type").and_then(Value::as_str) != Some("session_meta") {
527            continue;
528        }
529        let payload = entry.get("payload");
530        return Ok(Some(CodexSessionMeta {
531            id: payload
532                .and_then(|payload| payload.get("id"))
533                .and_then(Value::as_str)
534                .map(str::to_string),
535            cwd: payload
536                .and_then(|payload| payload.get("cwd"))
537                .and_then(Value::as_str)
538                .map(std::ffi::OsString::from),
539        }));
540    }
541    Ok(None)
542}
543
544fn parse_codex_session_usage(path: &Path) -> Result<Option<SessionUsage>> {
545    let file = File::open(path)
546        .with_context(|| format!("failed to open codex session {}", path.display()))?;
547    let reader = BufReader::new(file);
548    let mut usage = TokenUsage::default();
549    let mut model = None;
550
551    for line in reader.lines() {
552        let line = line?;
553        if line.trim().is_empty() {
554            continue;
555        }
556        let Ok(entry) = serde_json::from_str::<Value>(&line) else {
557            continue;
558        };
559
560        if entry.get("type").and_then(Value::as_str) == Some("turn_context")
561            && let Some(value) = entry
562                .get("payload")
563                .and_then(|payload| payload.get("model"))
564                .and_then(Value::as_str)
565        {
566            model = Some(value.to_string());
567        }
568
569        if entry.get("type").and_then(Value::as_str) != Some("event_msg")
570            || entry
571                .get("payload")
572                .and_then(|payload| payload.get("type"))
573                .and_then(Value::as_str)
574                != Some("token_count")
575        {
576            continue;
577        }
578
579        let Some(last_usage) = entry
580            .get("payload")
581            .and_then(|payload| payload.get("info"))
582            .and_then(|info| info.get("last_token_usage"))
583        else {
584            continue;
585        };
586
587        usage.input_tokens += json_u64(last_usage.get("input_tokens"));
588        usage.cached_input_tokens += json_u64(last_usage.get("cached_input_tokens"));
589        usage.output_tokens += json_u64(last_usage.get("output_tokens"));
590        usage.reasoning_output_tokens += json_u64(last_usage.get("reasoning_output_tokens"));
591    }
592
593    if model.is_none() && usage.total_tokens() == 0 {
594        return Ok(None);
595    }
596
597    Ok(Some(SessionUsage { model, usage }))
598}
599
600fn parse_claude_session_usage(path: &Path) -> Result<Option<SessionUsage>> {
601    let file = File::open(path)
602        .with_context(|| format!("failed to open claude session {}", path.display()))?;
603    let reader = BufReader::new(file);
604    let mut usage = TokenUsage::default();
605    let mut model = None;
606
607    for line in reader.lines() {
608        let line = line?;
609        if line.trim().is_empty() {
610            continue;
611        }
612        let Ok(entry) = serde_json::from_str::<Value>(&line) else {
613            continue;
614        };
615
616        let Some(message) = entry.get("message") else {
617            continue;
618        };
619        if let Some(value) = message.get("model").and_then(Value::as_str) {
620            model = Some(value.to_string());
621        }
622        let Some(usage_value) = message.get("usage") else {
623            continue;
624        };
625
626        usage.input_tokens += json_u64(usage_value.get("input_tokens"));
627        usage.output_tokens += json_u64(usage_value.get("output_tokens"));
628        usage.cache_creation_input_tokens +=
629            json_u64(usage_value.get("cache_creation_input_tokens"));
630        usage.cache_read_input_tokens += json_u64(usage_value.get("cache_read_input_tokens"));
631
632        let cache_creation = usage_value.get("cache_creation");
633        usage.cache_creation_5m_input_tokens +=
634            json_u64(cache_creation.and_then(|value| value.get("ephemeral_5m_input_tokens")));
635        usage.cache_creation_1h_input_tokens +=
636            json_u64(cache_creation.and_then(|value| value.get("ephemeral_1h_input_tokens")));
637    }
638
639    if model.is_none() && usage.total_tokens() == 0 {
640        return Ok(None);
641    }
642
643    Ok(Some(SessionUsage { model, usage }))
644}
645
646fn json_u64(value: Option<&Value>) -> u64 {
647    value.and_then(Value::as_u64).unwrap_or(0)
648}
649
650fn read_dir_paths(dir: &Path) -> Result<Vec<PathBuf>> {
651    let mut paths = Vec::new();
652    for entry in fs::read_dir(dir)? {
653        let entry = entry?;
654        paths.push(entry.path());
655    }
656    Ok(paths)
657}
658
659#[cfg(test)]
660mod tests {
661    use super::*;
662    use crate::team::config::{CostConfig, RoleType, TeamConfig, WorkflowMode};
663
664    fn test_team_config(models: HashMap<String, ModelPricing>) -> TeamConfig {
665        TeamConfig {
666            name: "batty".to_string(),
667            agent: None,
668            workflow_mode: WorkflowMode::Legacy,
669            board: Default::default(),
670            standup: Default::default(),
671            automation: Default::default(),
672            automation_sender: None,
673            external_senders: Vec::new(),
674            orchestrator_pane: true,
675            orchestrator_position: Default::default(),
676            layout: None,
677            workflow_policy: Default::default(),
678            cost: CostConfig { models },
679            event_log_max_bytes: crate::team::DEFAULT_EVENT_LOG_MAX_BYTES,
680            retro_min_duration_secs: 60,
681            roles: vec![
682                crate::team::config::RoleDef {
683                    name: "architect".to_string(),
684                    role_type: RoleType::Architect,
685                    agent: Some("claude".to_string()),
686                    instances: 1,
687                    prompt: None,
688                    talks_to: Vec::new(),
689                    channel: None,
690                    channel_config: None,
691                    nudge_interval_secs: None,
692                    receives_standup: None,
693                    standup_interval_secs: None,
694                    owns: Vec::new(),
695                    use_worktrees: false,
696                },
697                crate::team::config::RoleDef {
698                    name: "engineer".to_string(),
699                    role_type: RoleType::Engineer,
700                    agent: Some("codex".to_string()),
701                    instances: 1,
702                    prompt: None,
703                    talks_to: Vec::new(),
704                    channel: None,
705                    channel_config: None,
706                    nudge_interval_secs: None,
707                    receives_standup: None,
708                    standup_interval_secs: None,
709                    owns: Vec::new(),
710                    use_worktrees: true,
711                },
712            ],
713        }
714    }
715
716    #[test]
717    fn parse_codex_session_usage_sums_last_token_usage() {
718        let tmp = tempfile::tempdir().unwrap();
719        let path = tmp.path().join("codex.jsonl");
720        fs::write(
721            &path,
722            concat!(
723                "{\"type\":\"session_meta\",\"payload\":{\"id\":\"abc\",\"cwd\":\"/tmp/repo\"}}\n",
724                "{\"type\":\"turn_context\",\"payload\":{\"model\":\"gpt-5.4\"}}\n",
725                "{\"type\":\"event_msg\",\"payload\":{\"type\":\"token_count\",\"info\":{\"last_token_usage\":{\"input_tokens\":100,\"cached_input_tokens\":25,\"output_tokens\":10,\"reasoning_output_tokens\":5}}}}\n",
726                "{\"type\":\"event_msg\",\"payload\":{\"type\":\"token_count\",\"info\":{\"last_token_usage\":{\"input_tokens\":50,\"cached_input_tokens\":5,\"output_tokens\":4,\"reasoning_output_tokens\":1}}}}\n"
727            ),
728        )
729        .unwrap();
730
731        let usage = parse_codex_session_usage(&path).unwrap().unwrap();
732        assert_eq!(usage.model.as_deref(), Some("gpt-5.4"));
733        assert_eq!(usage.usage.input_tokens, 150);
734        assert_eq!(usage.usage.cached_input_tokens, 30);
735        assert_eq!(usage.usage.output_tokens, 14);
736        assert_eq!(usage.usage.reasoning_output_tokens, 6);
737    }
738
739    #[test]
740    fn parse_claude_session_usage_sums_message_usage() {
741        let tmp = tempfile::tempdir().unwrap();
742        let path = tmp.path().join("claude.jsonl");
743        fs::write(
744            &path,
745            concat!(
746                "{\"message\":{\"model\":\"claude-opus-4-6\",\"usage\":{\"input_tokens\":10,\"output_tokens\":2,\"cache_creation_input_tokens\":20,\"cache_read_input_tokens\":3,\"cache_creation\":{\"ephemeral_5m_input_tokens\":5,\"ephemeral_1h_input_tokens\":15}}}}\n",
747                "{\"message\":{\"model\":\"claude-opus-4-6\",\"usage\":{\"input_tokens\":7,\"output_tokens\":4,\"cache_creation_input_tokens\":8,\"cache_read_input_tokens\":1}}}\n"
748            ),
749        )
750        .unwrap();
751
752        let usage = parse_claude_session_usage(&path).unwrap().unwrap();
753        assert_eq!(usage.model.as_deref(), Some("claude-opus-4-6"));
754        assert_eq!(usage.usage.input_tokens, 17);
755        assert_eq!(usage.usage.output_tokens, 6);
756        assert_eq!(usage.usage.cache_creation_input_tokens, 28);
757        assert_eq!(usage.usage.cache_read_input_tokens, 4);
758        assert_eq!(usage.usage.cache_creation_5m_input_tokens, 5);
759        assert_eq!(usage.usage.cache_creation_1h_input_tokens, 15);
760    }
761
762    #[test]
763    fn estimate_cost_uses_pricing_breakdown() {
764        let usage = TokenUsage {
765            input_tokens: 1_000_000,
766            cached_input_tokens: 2_000_000,
767            cache_creation_input_tokens: 1_000_000,
768            cache_creation_5m_input_tokens: 400_000,
769            cache_creation_1h_input_tokens: 500_000,
770            cache_read_input_tokens: 300_000,
771            output_tokens: 100_000,
772            reasoning_output_tokens: 50_000,
773        };
774        let pricing = ModelPricing {
775            input_usd_per_mtok: 2.0,
776            cached_input_usd_per_mtok: 0.5,
777            cache_creation_input_usd_per_mtok: Some(3.0),
778            cache_creation_5m_input_usd_per_mtok: Some(4.0),
779            cache_creation_1h_input_usd_per_mtok: Some(5.0),
780            cache_read_input_usd_per_mtok: 0.25,
781            output_usd_per_mtok: 8.0,
782            reasoning_output_usd_per_mtok: Some(10.0),
783        };
784
785        let estimated = estimate_cost_usd(&usage, &pricing);
786        let expected = 2.0 + 1.0 + 1.6 + 2.5 + 0.3 + 0.075 + 0.8 + 0.5;
787        assert!((estimated - expected).abs() < 1e-9);
788    }
789
790    #[test]
791    fn built_in_pricing_supports_common_models() {
792        assert!(pricing_for_model(&HashMap::new(), "gpt-5.4").is_some());
793        assert!(pricing_for_model(&HashMap::new(), "claude-opus-4-6").is_some());
794        assert!(pricing_for_model(&HashMap::new(), "claude-sonnet-4").is_some());
795    }
796
797    #[test]
798    fn collect_cost_report_maps_members_to_current_tasks() {
799        let tmp = tempfile::tempdir().unwrap();
800        let project_root = tmp.path();
801        fs::create_dir_all(
802            project_root
803                .join(".batty")
804                .join("worktrees")
805                .join("engineer"),
806        )
807        .unwrap();
808        fs::create_dir_all(
809            project_root
810                .join(".batty")
811                .join("team_config")
812                .join("board")
813                .join("tasks"),
814        )
815        .unwrap();
816
817        fs::write(
818            project_root.join(".batty").join("launch-state.json"),
819            r#"{
820  "architect": {"session_id": "claude-session"},
821  "engineer": {"session_id": "codex-session"}
822}"#,
823        )
824        .unwrap();
825        fs::write(
826            daemon_state_path(project_root),
827            r#"{"active_tasks":{"engineer":100}}"#,
828        )
829        .unwrap();
830        fs::write(
831            project_root
832                .join(".batty")
833                .join("team_config")
834                .join("board")
835                .join("tasks")
836                .join("100-task.md"),
837            concat!(
838                "---\n",
839                "id: 100\n",
840                "title: Cost task\n",
841                "status: in-progress\n",
842                "claimed_by: engineer\n",
843                "---\n"
844            ),
845        )
846        .unwrap();
847
848        let codex_root = project_root.join("codex-sessions");
849        let codex_day = codex_root.join("2026").join("03").join("21");
850        fs::create_dir_all(&codex_day).unwrap();
851        fs::write(
852            codex_day.join("rollout.jsonl"),
853            format!(
854                "{{\"type\":\"session_meta\",\"payload\":{{\"id\":\"codex-session\",\"cwd\":\"{}\"}}}}\n{}\n{}\n",
855                project_root
856                    .join(".batty")
857                    .join("worktrees")
858                    .join("engineer")
859                    .join(".batty")
860                    .join("codex-context")
861                    .join("engineer")
862                    .display(),
863                r#"{"type":"turn_context","payload":{"model":"gpt-5.4"}}"#,
864                r#"{"type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":1000,"cached_input_tokens":250,"output_tokens":100,"reasoning_output_tokens":10}}}}"#,
865            ),
866        )
867        .unwrap();
868
869        let claude_root = project_root.join("claude-projects");
870        let claude_dir = claude_root.join(project_root.to_string_lossy().replace('/', "-"));
871        fs::create_dir_all(&claude_dir).unwrap();
872        fs::write(
873            claude_dir.join("claude-session.jsonl"),
874            "{\"message\":{\"model\":\"claude-opus-4-6\",\"usage\":{\"input_tokens\":10,\"output_tokens\":2,\"cache_creation_input_tokens\":20,\"cache_read_input_tokens\":3,\"cache_creation\":{\"ephemeral_5m_input_tokens\":5,\"ephemeral_1h_input_tokens\":15}}}}\n",
875        )
876        .unwrap();
877
878        let report = collect_cost_report(
879            project_root,
880            &test_team_config(HashMap::new()),
881            &SessionRoots {
882                codex_sessions_root: codex_root,
883                claude_projects_root: claude_root,
884            },
885        )
886        .unwrap();
887
888        assert_eq!(report.entries.len(), 2);
889        let engineer = report
890            .entries
891            .iter()
892            .find(|entry| entry.member_name == "engineer")
893            .unwrap();
894        assert_eq!(engineer.task, "#100");
895        assert_eq!(engineer.model, "gpt-5.4");
896        assert!(engineer.estimated_cost_usd.unwrap() > 0.0);
897    }
898
899    #[test]
900    fn collect_cost_report_skips_user_roles_and_shared_cwd_fallbacks() {
901        let tmp = tempfile::tempdir().unwrap();
902        let project_root = tmp.path();
903        let mut config = test_team_config(HashMap::new());
904        config.roles.insert(
905            0,
906            crate::team::config::RoleDef {
907                name: "human".to_string(),
908                role_type: RoleType::User,
909                agent: None,
910                instances: 1,
911                prompt: None,
912                talks_to: Vec::new(),
913                channel: None,
914                channel_config: None,
915                nudge_interval_secs: None,
916                receives_standup: None,
917                standup_interval_secs: None,
918                owns: Vec::new(),
919                use_worktrees: false,
920            },
921        );
922        config.roles.push(crate::team::config::RoleDef {
923            name: "manager".to_string(),
924            role_type: RoleType::Manager,
925            agent: Some("claude".to_string()),
926            instances: 1,
927            prompt: None,
928            talks_to: Vec::new(),
929            channel: None,
930            channel_config: None,
931            nudge_interval_secs: None,
932            receives_standup: None,
933            standup_interval_secs: None,
934            owns: Vec::new(),
935            use_worktrees: false,
936        });
937
938        let claude_root = project_root.join("claude-projects");
939        let claude_dir = claude_root.join(project_root.to_string_lossy().replace('/', "-"));
940        fs::create_dir_all(&claude_dir).unwrap();
941        fs::write(
942            claude_dir.join("shared.jsonl"),
943            "{\"message\":{\"model\":\"claude-opus-4-6\",\"usage\":{\"input_tokens\":10,\"output_tokens\":2}}}\n",
944        )
945        .unwrap();
946
947        let report = collect_cost_report(
948            project_root,
949            &config,
950            &SessionRoots {
951                codex_sessions_root: project_root.join("codex-sessions"),
952                claude_projects_root: claude_root,
953            },
954        )
955        .unwrap();
956
957        assert!(report.entries.is_empty());
958    }
959
960    // ── TokenUsage ───────────────────────────────────────────────
961
962    #[test]
963    fn token_usage_total_tokens_sums_all_fields() {
964        let usage = TokenUsage {
965            input_tokens: 100,
966            cached_input_tokens: 50,
967            cache_creation_input_tokens: 30,
968            cache_creation_5m_input_tokens: 0,
969            cache_creation_1h_input_tokens: 0,
970            cache_read_input_tokens: 20,
971            output_tokens: 10,
972            reasoning_output_tokens: 5,
973        };
974        assert_eq!(usage.total_tokens(), 100 + 50 + 30 + 20 + 10 + 5);
975    }
976
977    #[test]
978    fn token_usage_total_tokens_zero_when_default() {
979        let usage = TokenUsage::default();
980        assert_eq!(usage.total_tokens(), 0);
981    }
982
983    #[test]
984    fn token_usage_display_cache_tokens_sums_cache_fields() {
985        let usage = TokenUsage {
986            cached_input_tokens: 100,
987            cache_creation_input_tokens: 200,
988            cache_read_input_tokens: 50,
989            ..TokenUsage::default()
990        };
991        assert_eq!(usage.display_cache_tokens(), 350);
992    }
993
994    // ── truncate_model ───────────────────────────────────────────
995
996    #[test]
997    fn truncate_model_short_name_unchanged() {
998        assert_eq!(truncate_model("gpt-5.4"), "gpt-5.4");
999    }
1000
1001    #[test]
1002    fn truncate_model_exact_limit_unchanged() {
1003        let model = "a".repeat(20);
1004        assert_eq!(truncate_model(&model), model);
1005    }
1006
1007    #[test]
1008    fn truncate_model_long_name_truncated() {
1009        let model = "a".repeat(25);
1010        let result = truncate_model(&model);
1011        assert_eq!(result.len(), 20);
1012        assert!(result.ends_with("..."));
1013    }
1014
1015    // ── member_session_target ────────────────────────────────────
1016
1017    #[test]
1018    fn member_session_target_returns_none_for_user_role() {
1019        let member = MemberInstance {
1020            name: "human".to_string(),
1021            role_name: "human".to_string(),
1022            role_type: super::super::config::RoleType::User,
1023            agent: None,
1024            prompt: None,
1025            reports_to: None,
1026            use_worktrees: false,
1027        };
1028        assert!(member_session_target(Path::new("/tmp"), &member).is_none());
1029    }
1030
1031    #[test]
1032    fn member_session_target_codex_agent() {
1033        let member = MemberInstance {
1034            name: "eng-1".to_string(),
1035            role_name: "eng".to_string(),
1036            role_type: super::super::config::RoleType::Engineer,
1037            agent: Some("codex".to_string()),
1038            prompt: None,
1039            reports_to: None,
1040            use_worktrees: false,
1041        };
1042        let (agent, cwd, label) = member_session_target(Path::new("/tmp/repo"), &member).unwrap();
1043        assert!(matches!(agent, SessionAgent::Codex));
1044        assert_eq!(label, "codex");
1045        assert!(cwd.to_string_lossy().contains("codex-context"));
1046    }
1047
1048    #[test]
1049    fn member_session_target_claude_agent() {
1050        let member = MemberInstance {
1051            name: "architect".to_string(),
1052            role_name: "architect".to_string(),
1053            role_type: super::super::config::RoleType::Architect,
1054            agent: Some("claude".to_string()),
1055            prompt: None,
1056            reports_to: None,
1057            use_worktrees: false,
1058        };
1059        let (agent, cwd, label) = member_session_target(Path::new("/tmp/repo"), &member).unwrap();
1060        assert!(matches!(agent, SessionAgent::Claude));
1061        assert_eq!(label, "claude");
1062        assert_eq!(cwd, Path::new("/tmp/repo"));
1063    }
1064
1065    #[test]
1066    fn member_session_target_claude_code_agent() {
1067        let member = MemberInstance {
1068            name: "eng-2".to_string(),
1069            role_name: "eng".to_string(),
1070            role_type: super::super::config::RoleType::Engineer,
1071            agent: Some("claude-code".to_string()),
1072            prompt: None,
1073            reports_to: None,
1074            use_worktrees: false,
1075        };
1076        let (agent, _, label) = member_session_target(Path::new("/tmp/repo"), &member).unwrap();
1077        assert!(matches!(agent, SessionAgent::Claude));
1078        assert_eq!(label, "claude");
1079    }
1080
1081    #[test]
1082    fn member_session_target_none_agent_defaults_to_claude() {
1083        let member = MemberInstance {
1084            name: "eng-3".to_string(),
1085            role_name: "eng".to_string(),
1086            role_type: super::super::config::RoleType::Engineer,
1087            agent: None,
1088            prompt: None,
1089            reports_to: None,
1090            use_worktrees: false,
1091        };
1092        let (agent, _, label) = member_session_target(Path::new("/tmp/repo"), &member).unwrap();
1093        assert!(matches!(agent, SessionAgent::Claude));
1094        assert_eq!(label, "claude");
1095    }
1096
1097    #[test]
1098    fn member_session_target_unknown_agent_returns_none() {
1099        let member = MemberInstance {
1100            name: "eng-4".to_string(),
1101            role_name: "eng".to_string(),
1102            role_type: super::super::config::RoleType::Engineer,
1103            agent: Some("gemini".to_string()),
1104            prompt: None,
1105            reports_to: None,
1106            use_worktrees: false,
1107        };
1108        assert!(member_session_target(Path::new("/tmp/repo"), &member).is_none());
1109    }
1110
1111    #[test]
1112    fn member_session_target_worktree_path_for_codex() {
1113        let member = MemberInstance {
1114            name: "eng-1".to_string(),
1115            role_name: "eng".to_string(),
1116            role_type: super::super::config::RoleType::Engineer,
1117            agent: Some("codex-cli".to_string()),
1118            prompt: None,
1119            reports_to: None,
1120            use_worktrees: true,
1121        };
1122        let (_, cwd, _) = member_session_target(Path::new("/tmp/repo"), &member).unwrap();
1123        assert!(cwd.starts_with("/tmp/repo/.batty/worktrees/eng-1"));
1124    }
1125
1126    #[test]
1127    fn member_session_target_worktree_path_for_claude() {
1128        let member = MemberInstance {
1129            name: "eng-1".to_string(),
1130            role_name: "eng".to_string(),
1131            role_type: super::super::config::RoleType::Engineer,
1132            agent: Some("claude".to_string()),
1133            prompt: None,
1134            reports_to: None,
1135            use_worktrees: true,
1136        };
1137        let (_, cwd, _) = member_session_target(Path::new("/tmp/repo"), &member).unwrap();
1138        assert_eq!(cwd, Path::new("/tmp/repo/.batty/worktrees/eng-1"));
1139    }
1140
1141    // ── load_launch_state ────────────────────────────────────────
1142
1143    #[test]
1144    fn load_launch_state_returns_empty_when_file_missing() {
1145        let tmp = tempfile::tempdir().unwrap();
1146        let state = load_launch_state(tmp.path());
1147        assert!(state.is_empty());
1148    }
1149
1150    #[test]
1151    fn load_launch_state_parses_valid_json() {
1152        let tmp = tempfile::tempdir().unwrap();
1153        let batty_dir = tmp.path().join(".batty");
1154        fs::create_dir_all(&batty_dir).unwrap();
1155        fs::write(
1156            batty_dir.join("launch-state.json"),
1157            r#"{"eng-1": {"session_id": "abc123"}, "eng-2": {}}"#,
1158        )
1159        .unwrap();
1160
1161        let state = load_launch_state(tmp.path());
1162        assert_eq!(state.len(), 2);
1163        assert_eq!(
1164            state.get("eng-1").unwrap().session_id.as_deref(),
1165            Some("abc123")
1166        );
1167        assert!(state.get("eng-2").unwrap().session_id.is_none());
1168    }
1169
1170    #[test]
1171    fn load_launch_state_returns_empty_on_invalid_json() {
1172        let tmp = tempfile::tempdir().unwrap();
1173        let batty_dir = tmp.path().join(".batty");
1174        fs::create_dir_all(&batty_dir).unwrap();
1175        fs::write(batty_dir.join("launch-state.json"), "not json").unwrap();
1176
1177        let state = load_launch_state(tmp.path());
1178        assert!(state.is_empty());
1179    }
1180
1181    // ── load_active_tasks ────────────────────────────────────────
1182
1183    #[test]
1184    fn load_active_tasks_returns_empty_when_file_missing() {
1185        let tmp = tempfile::tempdir().unwrap();
1186        let tasks = load_active_tasks(tmp.path());
1187        assert!(tasks.is_empty());
1188    }
1189
1190    #[test]
1191    fn load_active_tasks_parses_valid_state() {
1192        let tmp = tempfile::tempdir().unwrap();
1193        let state_path = daemon_state_path(tmp.path());
1194        fs::create_dir_all(state_path.parent().unwrap()).unwrap();
1195        fs::write(&state_path, r#"{"active_tasks":{"eng-1":42,"eng-2":99}}"#).unwrap();
1196
1197        let tasks = load_active_tasks(tmp.path());
1198        assert_eq!(tasks.len(), 2);
1199        assert_eq!(tasks.get("eng-1"), Some(&42));
1200        assert_eq!(tasks.get("eng-2"), Some(&99));
1201    }
1202
1203    #[test]
1204    fn load_active_tasks_returns_empty_on_invalid_json() {
1205        let tmp = tempfile::tempdir().unwrap();
1206        let state_path = daemon_state_path(tmp.path());
1207        fs::create_dir_all(state_path.parent().unwrap()).unwrap();
1208        fs::write(&state_path, "garbage").unwrap();
1209
1210        let tasks = load_active_tasks(tmp.path());
1211        assert!(tasks.is_empty());
1212    }
1213
1214    // ── pricing_for_model ────────────────────────────────────────
1215
1216    #[test]
1217    fn pricing_for_model_uses_override_exact_match() {
1218        let mut overrides = HashMap::new();
1219        overrides.insert(
1220            "custom-model".to_string(),
1221            ModelPricing {
1222                input_usd_per_mtok: 99.0,
1223                cached_input_usd_per_mtok: 0.0,
1224                cache_creation_input_usd_per_mtok: None,
1225                cache_creation_5m_input_usd_per_mtok: None,
1226                cache_creation_1h_input_usd_per_mtok: None,
1227                cache_read_input_usd_per_mtok: 0.0,
1228                output_usd_per_mtok: 99.0,
1229                reasoning_output_usd_per_mtok: None,
1230            },
1231        );
1232        let pricing = pricing_for_model(&overrides, "custom-model").unwrap();
1233        assert!((pricing.input_usd_per_mtok - 99.0).abs() < f64::EPSILON);
1234    }
1235
1236    #[test]
1237    fn pricing_for_model_normalized_case_match() {
1238        let mut overrides = HashMap::new();
1239        overrides.insert(
1240            "my-model".to_string(),
1241            ModelPricing {
1242                input_usd_per_mtok: 5.0,
1243                cached_input_usd_per_mtok: 0.0,
1244                cache_creation_input_usd_per_mtok: None,
1245                cache_creation_5m_input_usd_per_mtok: None,
1246                cache_creation_1h_input_usd_per_mtok: None,
1247                cache_read_input_usd_per_mtok: 0.0,
1248                output_usd_per_mtok: 5.0,
1249                reasoning_output_usd_per_mtok: None,
1250            },
1251        );
1252        let pricing = pricing_for_model(&overrides, "My-Model").unwrap();
1253        assert!((pricing.input_usd_per_mtok - 5.0).abs() < f64::EPSILON);
1254    }
1255
1256    #[test]
1257    fn pricing_for_model_returns_none_for_unknown() {
1258        assert!(pricing_for_model(&HashMap::new(), "totally-unknown-model").is_none());
1259    }
1260
1261    // ── built_in_model_pricing ───────────────────────────────────
1262
1263    #[test]
1264    fn built_in_pricing_gpt54_has_reasoning_rate() {
1265        let pricing = built_in_model_pricing("gpt-5.4").unwrap();
1266        assert!(pricing.reasoning_output_usd_per_mtok.is_some());
1267    }
1268
1269    #[test]
1270    fn built_in_pricing_opus_has_cache_tiers() {
1271        let pricing = built_in_model_pricing("claude-opus-4-6").unwrap();
1272        assert!(pricing.cache_creation_5m_input_usd_per_mtok.is_some());
1273        assert!(pricing.cache_creation_1h_input_usd_per_mtok.is_some());
1274    }
1275
1276    #[test]
1277    fn built_in_pricing_sonnet_has_cache_tiers() {
1278        let pricing = built_in_model_pricing("claude-sonnet-4-6").unwrap();
1279        assert!(pricing.cache_creation_5m_input_usd_per_mtok.is_some());
1280        assert!(pricing.cache_creation_1h_input_usd_per_mtok.is_some());
1281    }
1282
1283    #[test]
1284    fn built_in_pricing_unknown_returns_none() {
1285        assert!(built_in_model_pricing("llama-3").is_none());
1286    }
1287
1288    // ── estimate_cost_usd edge cases ─────────────────────────────
1289
1290    #[test]
1291    fn estimate_cost_zero_usage_returns_zero() {
1292        let usage = TokenUsage::default();
1293        let pricing = ModelPricing {
1294            input_usd_per_mtok: 10.0,
1295            cached_input_usd_per_mtok: 5.0,
1296            cache_creation_input_usd_per_mtok: None,
1297            cache_creation_5m_input_usd_per_mtok: None,
1298            cache_creation_1h_input_usd_per_mtok: None,
1299            cache_read_input_usd_per_mtok: 1.0,
1300            output_usd_per_mtok: 20.0,
1301            reasoning_output_usd_per_mtok: None,
1302        };
1303        assert!((estimate_cost_usd(&usage, &pricing)).abs() < f64::EPSILON);
1304    }
1305
1306    #[test]
1307    fn estimate_cost_uses_output_rate_when_no_reasoning_rate() {
1308        let usage = TokenUsage {
1309            reasoning_output_tokens: 1_000_000,
1310            ..TokenUsage::default()
1311        };
1312        let pricing = ModelPricing {
1313            input_usd_per_mtok: 0.0,
1314            cached_input_usd_per_mtok: 0.0,
1315            cache_creation_input_usd_per_mtok: None,
1316            cache_creation_5m_input_usd_per_mtok: None,
1317            cache_creation_1h_input_usd_per_mtok: None,
1318            cache_read_input_usd_per_mtok: 0.0,
1319            output_usd_per_mtok: 10.0,
1320            reasoning_output_usd_per_mtok: None,
1321        };
1322        // reasoning falls back to output rate: 1M tokens * 10/M = $10
1323        assert!((estimate_cost_usd(&usage, &pricing) - 10.0).abs() < f64::EPSILON);
1324    }
1325
1326    #[test]
1327    fn estimate_cost_unclassified_cache_creation_uses_generic_rate() {
1328        let usage = TokenUsage {
1329            cache_creation_input_tokens: 1_000_000,
1330            // No 5m or 1h breakdown → all goes to unclassified
1331            ..TokenUsage::default()
1332        };
1333        let pricing = ModelPricing {
1334            input_usd_per_mtok: 0.0,
1335            cached_input_usd_per_mtok: 0.0,
1336            cache_creation_input_usd_per_mtok: Some(5.0),
1337            cache_creation_5m_input_usd_per_mtok: None,
1338            cache_creation_1h_input_usd_per_mtok: None,
1339            cache_read_input_usd_per_mtok: 0.0,
1340            output_usd_per_mtok: 0.0,
1341            reasoning_output_usd_per_mtok: None,
1342        };
1343        assert!((estimate_cost_usd(&usage, &pricing) - 5.0).abs() < f64::EPSILON);
1344    }
1345
1346    #[test]
1347    fn estimate_cost_generic_rate_falls_back_to_5m_then_input() {
1348        let usage = TokenUsage {
1349            cache_creation_input_tokens: 1_000_000,
1350            ..TokenUsage::default()
1351        };
1352        // No generic rate → falls back to 5m rate
1353        let pricing = ModelPricing {
1354            input_usd_per_mtok: 2.0,
1355            cached_input_usd_per_mtok: 0.0,
1356            cache_creation_input_usd_per_mtok: None,
1357            cache_creation_5m_input_usd_per_mtok: Some(4.0),
1358            cache_creation_1h_input_usd_per_mtok: None,
1359            cache_read_input_usd_per_mtok: 0.0,
1360            output_usd_per_mtok: 0.0,
1361            reasoning_output_usd_per_mtok: None,
1362        };
1363        // unclassified = 1M, generic rate = 5m rate = 4.0 → $4
1364        assert!((estimate_cost_usd(&usage, &pricing) - 4.0).abs() < f64::EPSILON);
1365
1366        // No generic rate AND no 5m rate → falls back to input rate
1367        let pricing2 = ModelPricing {
1368            input_usd_per_mtok: 2.0,
1369            cached_input_usd_per_mtok: 0.0,
1370            cache_creation_input_usd_per_mtok: None,
1371            cache_creation_5m_input_usd_per_mtok: None,
1372            cache_creation_1h_input_usd_per_mtok: None,
1373            cache_read_input_usd_per_mtok: 0.0,
1374            output_usd_per_mtok: 0.0,
1375            reasoning_output_usd_per_mtok: None,
1376        };
1377        assert!((estimate_cost_usd(&usage, &pricing2) - 2.0).abs() < f64::EPSILON);
1378    }
1379
1380    // ── parse session files edge cases ───────────────────────────
1381
1382    #[test]
1383    fn parse_codex_session_returns_none_for_empty_file() {
1384        let tmp = tempfile::tempdir().unwrap();
1385        let path = tmp.path().join("empty.jsonl");
1386        fs::write(&path, "").unwrap();
1387
1388        assert!(parse_codex_session_usage(&path).unwrap().is_none());
1389    }
1390
1391    #[test]
1392    fn parse_codex_session_skips_malformed_lines() {
1393        let tmp = tempfile::tempdir().unwrap();
1394        let path = tmp.path().join("mixed.jsonl");
1395        fs::write(
1396            &path,
1397            concat!(
1398                "not valid json\n",
1399                "{\"type\":\"turn_context\",\"payload\":{\"model\":\"gpt-5.4\"}}\n",
1400                "{\"type\":\"event_msg\",\"payload\":{\"type\":\"token_count\",\"info\":{\"last_token_usage\":{\"input_tokens\":100}}}}\n",
1401            ),
1402        )
1403        .unwrap();
1404
1405        let usage = parse_codex_session_usage(&path).unwrap().unwrap();
1406        assert_eq!(usage.model.as_deref(), Some("gpt-5.4"));
1407        assert_eq!(usage.usage.input_tokens, 100);
1408    }
1409
1410    #[test]
1411    fn parse_claude_session_returns_none_for_empty_file() {
1412        let tmp = tempfile::tempdir().unwrap();
1413        let path = tmp.path().join("empty.jsonl");
1414        fs::write(&path, "").unwrap();
1415
1416        assert!(parse_claude_session_usage(&path).unwrap().is_none());
1417    }
1418
1419    #[test]
1420    fn parse_claude_session_skips_entries_without_message() {
1421        let tmp = tempfile::tempdir().unwrap();
1422        let path = tmp.path().join("mixed.jsonl");
1423        fs::write(
1424            &path,
1425            concat!(
1426                "{\"not_message\":true}\n",
1427                "{\"message\":{\"model\":\"claude-sonnet-4\",\"usage\":{\"input_tokens\":50,\"output_tokens\":10}}}\n",
1428            ),
1429        )
1430        .unwrap();
1431
1432        let usage = parse_claude_session_usage(&path).unwrap().unwrap();
1433        assert_eq!(usage.model.as_deref(), Some("claude-sonnet-4"));
1434        assert_eq!(usage.usage.input_tokens, 50);
1435    }
1436
1437    // ── json_u64 ─────────────────────────────────────────────────
1438
1439    #[test]
1440    fn json_u64_returns_value_for_number() {
1441        let v: Value = serde_json::json!(42);
1442        assert_eq!(json_u64(Some(&v)), 42);
1443    }
1444
1445    #[test]
1446    fn json_u64_returns_zero_for_none() {
1447        assert_eq!(json_u64(None), 0);
1448    }
1449
1450    #[test]
1451    fn json_u64_returns_zero_for_non_number() {
1452        let v: Value = serde_json::json!("not a number");
1453        assert_eq!(json_u64(Some(&v)), 0);
1454    }
1455
1456    // ── read_codex_session_meta ──────────────────────────────────
1457
1458    #[test]
1459    fn read_codex_session_meta_parses_session_meta_line() {
1460        let tmp = tempfile::tempdir().unwrap();
1461        let path = tmp.path().join("session.jsonl");
1462        fs::write(
1463            &path,
1464            "{\"type\":\"session_meta\",\"payload\":{\"id\":\"sess-123\",\"cwd\":\"/tmp/work\"}}\n",
1465        )
1466        .unwrap();
1467
1468        let meta = read_codex_session_meta(&path).unwrap().unwrap();
1469        assert_eq!(meta.id.as_deref(), Some("sess-123"));
1470        assert_eq!(meta.cwd.as_deref(), Some(std::ffi::OsStr::new("/tmp/work")));
1471    }
1472
1473    #[test]
1474    fn read_codex_session_meta_returns_none_for_no_meta() {
1475        let tmp = tempfile::tempdir().unwrap();
1476        let path = tmp.path().join("no-meta.jsonl");
1477        fs::write(
1478            &path,
1479            "{\"type\":\"event_msg\",\"payload\":{\"type\":\"token_count\"}}\n",
1480        )
1481        .unwrap();
1482
1483        assert!(read_codex_session_meta(&path).unwrap().is_none());
1484    }
1485
1486    #[test]
1487    fn read_codex_session_meta_skips_blank_and_malformed_lines() {
1488        let tmp = tempfile::tempdir().unwrap();
1489        let path = tmp.path().join("messy.jsonl");
1490        fs::write(
1491            &path,
1492            concat!(
1493                "\n",
1494                "not json\n",
1495                "{\"type\":\"session_meta\",\"payload\":{\"id\":\"found\"}}\n",
1496            ),
1497        )
1498        .unwrap();
1499
1500        let meta = read_codex_session_meta(&path).unwrap().unwrap();
1501        assert_eq!(meta.id.as_deref(), Some("found"));
1502    }
1503
1504    // ── discover session files ───────────────────────────────────
1505
1506    #[test]
1507    fn discover_codex_session_returns_none_when_root_missing() {
1508        let result = discover_codex_session_file(
1509            Path::new("/nonexistent/path"),
1510            Path::new("/tmp/cwd"),
1511            None,
1512            true,
1513        )
1514        .unwrap();
1515        assert!(result.is_none());
1516    }
1517
1518    #[test]
1519    fn discover_claude_session_returns_none_when_root_missing() {
1520        let result = discover_claude_session_file(
1521            Path::new("/nonexistent/path"),
1522            Path::new("/tmp/cwd"),
1523            None,
1524            true,
1525        )
1526        .unwrap();
1527        assert!(result.is_none());
1528    }
1529
1530    #[test]
1531    fn discover_claude_session_finds_exact_session_id() {
1532        let tmp = tempfile::tempdir().unwrap();
1533        let projects_root = tmp.path().join("projects");
1534        let project_dir = projects_root.join("-tmp-myrepo");
1535        fs::create_dir_all(&project_dir).unwrap();
1536        fs::write(
1537            project_dir.join("session-abc.jsonl"),
1538            "{\"message\":{\"model\":\"claude\"}}\n",
1539        )
1540        .unwrap();
1541
1542        let result = discover_claude_session_file(
1543            &projects_root,
1544            Path::new("/tmp/myrepo"),
1545            Some("session-abc"),
1546            true,
1547        )
1548        .unwrap();
1549        assert!(result.is_some());
1550        assert!(result.unwrap().ends_with("session-abc.jsonl"));
1551    }
1552
1553    #[test]
1554    fn discover_claude_session_cwd_fallback_finds_newest() {
1555        let tmp = tempfile::tempdir().unwrap();
1556        let projects_root = tmp.path().join("projects");
1557        let project_dir = projects_root.join("-tmp-myrepo");
1558        fs::create_dir_all(&project_dir).unwrap();
1559
1560        fs::write(project_dir.join("old.jsonl"), "old\n").unwrap();
1561        // Give a small delay or touch to make newer file distinguishable
1562        fs::write(project_dir.join("new.jsonl"), "new\n").unwrap();
1563
1564        let result =
1565            discover_claude_session_file(&projects_root, Path::new("/tmp/myrepo"), None, true)
1566                .unwrap();
1567        assert!(result.is_some());
1568    }
1569
1570    #[test]
1571    fn discover_claude_session_no_cwd_fallback_when_disabled() {
1572        let tmp = tempfile::tempdir().unwrap();
1573        let projects_root = tmp.path().join("projects");
1574        let project_dir = projects_root.join("-tmp-myrepo");
1575        fs::create_dir_all(&project_dir).unwrap();
1576        fs::write(project_dir.join("session.jsonl"), "data\n").unwrap();
1577
1578        let result = discover_claude_session_file(
1579            &projects_root,
1580            Path::new("/tmp/myrepo"),
1581            None,
1582            false, // cwd fallback disabled
1583        )
1584        .unwrap();
1585        assert!(result.is_none());
1586    }
1587
1588    // ── read_dir_paths ───────────────────────────────────────────
1589
1590    #[test]
1591    fn read_dir_paths_returns_entries() {
1592        let tmp = tempfile::tempdir().unwrap();
1593        fs::write(tmp.path().join("a.txt"), "a").unwrap();
1594        fs::write(tmp.path().join("b.txt"), "b").unwrap();
1595        let paths = read_dir_paths(tmp.path()).unwrap();
1596        assert_eq!(paths.len(), 2);
1597    }
1598
1599    #[test]
1600    fn read_dir_paths_returns_empty_for_empty_dir() {
1601        let tmp = tempfile::tempdir().unwrap();
1602        let paths = read_dir_paths(tmp.path()).unwrap();
1603        assert!(paths.is_empty());
1604    }
1605
1606    // ── SessionRoots default ─────────────────────────────────────
1607
1608    #[test]
1609    fn session_roots_default_uses_home() {
1610        let roots = SessionRoots::default();
1611        let home = std::env::var("HOME").unwrap_or_default();
1612        assert!(roots.codex_sessions_root.starts_with(&home));
1613        assert!(roots.claude_projects_root.starts_with(&home));
1614    }
1615
1616    // ── collect_cost_report empty board ──────────────────────────
1617
1618    #[test]
1619    fn collect_cost_report_empty_when_no_sessions() {
1620        let tmp = tempfile::tempdir().unwrap();
1621        let project_root = tmp.path();
1622        fs::create_dir_all(
1623            project_root
1624                .join(".batty")
1625                .join("team_config")
1626                .join("board")
1627                .join("tasks"),
1628        )
1629        .unwrap();
1630
1631        let report = collect_cost_report(
1632            project_root,
1633            &test_team_config(HashMap::new()),
1634            &SessionRoots {
1635                codex_sessions_root: project_root.join("no-codex"),
1636                claude_projects_root: project_root.join("no-claude"),
1637            },
1638        )
1639        .unwrap();
1640
1641        assert!(report.entries.is_empty());
1642        assert!((report.total_estimated_cost_usd).abs() < f64::EPSILON);
1643        assert!(report.unpriced_models.is_empty());
1644    }
1645
1646    #[test]
1647    fn collect_cost_report_tracks_unpriced_models() {
1648        let tmp = tempfile::tempdir().unwrap();
1649        let project_root = tmp.path();
1650        fs::create_dir_all(
1651            project_root
1652                .join(".batty")
1653                .join("team_config")
1654                .join("board")
1655                .join("tasks"),
1656        )
1657        .unwrap();
1658
1659        // Create a claude session with an unknown model
1660        let claude_root = project_root.join("claude-projects");
1661        let claude_dir = claude_root.join(project_root.to_string_lossy().replace('/', "-"));
1662        fs::create_dir_all(&claude_dir).unwrap();
1663        fs::write(
1664            claude_dir.join("session.jsonl"),
1665            "{\"message\":{\"model\":\"totally-unknown-model\",\"usage\":{\"input_tokens\":100,\"output_tokens\":10}}}\n",
1666        )
1667        .unwrap();
1668
1669        let report = collect_cost_report(
1670            project_root,
1671            &test_team_config(HashMap::new()),
1672            &SessionRoots {
1673                codex_sessions_root: project_root.join("no-codex"),
1674                claude_projects_root: claude_root,
1675            },
1676        )
1677        .unwrap();
1678
1679        assert!(report.unpriced_models.contains("totally-unknown-model"));
1680    }
1681}