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            grafana: Default::default(),
680            use_shim: false,
681            use_sdk_mode: false,
682            auto_respawn_on_crash: false,
683            shim_health_check_interval_secs: 60,
684            shim_health_timeout_secs: 120,
685            shim_shutdown_timeout_secs: 30,
686            shim_working_state_timeout_secs: 1800,
687            pending_queue_max_age_secs: 600,
688            event_log_max_bytes: crate::team::DEFAULT_EVENT_LOG_MAX_BYTES,
689            retro_min_duration_secs: 60,
690            roles: vec![
691                crate::team::config::RoleDef {
692                    name: "architect".to_string(),
693                    role_type: RoleType::Architect,
694                    agent: Some("claude".to_string()),
695                    auth_mode: None,
696                    auth_env: vec![],
697                    instances: 1,
698                    prompt: None,
699                    talks_to: Vec::new(),
700                    channel: None,
701                    channel_config: None,
702                    nudge_interval_secs: None,
703                    receives_standup: None,
704                    standup_interval_secs: None,
705                    owns: Vec::new(),
706                    barrier_group: None,
707                    use_worktrees: false,
708                    ..Default::default()
709                },
710                crate::team::config::RoleDef {
711                    name: "engineer".to_string(),
712                    role_type: RoleType::Engineer,
713                    agent: Some("codex".to_string()),
714                    auth_mode: None,
715                    auth_env: vec![],
716                    instances: 1,
717                    prompt: None,
718                    talks_to: Vec::new(),
719                    channel: None,
720                    channel_config: None,
721                    nudge_interval_secs: None,
722                    receives_standup: None,
723                    standup_interval_secs: None,
724                    owns: Vec::new(),
725                    barrier_group: None,
726                    use_worktrees: true,
727                    ..Default::default()
728                },
729            ],
730        }
731    }
732
733    #[test]
734    fn parse_codex_session_usage_sums_last_token_usage() {
735        let tmp = tempfile::tempdir().unwrap();
736        let path = tmp.path().join("codex.jsonl");
737        fs::write(
738            &path,
739            concat!(
740                "{\"type\":\"session_meta\",\"payload\":{\"id\":\"abc\",\"cwd\":\"/tmp/repo\"}}\n",
741                "{\"type\":\"turn_context\",\"payload\":{\"model\":\"gpt-5.4\"}}\n",
742                "{\"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",
743                "{\"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"
744            ),
745        )
746        .unwrap();
747
748        let usage = parse_codex_session_usage(&path).unwrap().unwrap();
749        assert_eq!(usage.model.as_deref(), Some("gpt-5.4"));
750        assert_eq!(usage.usage.input_tokens, 150);
751        assert_eq!(usage.usage.cached_input_tokens, 30);
752        assert_eq!(usage.usage.output_tokens, 14);
753        assert_eq!(usage.usage.reasoning_output_tokens, 6);
754    }
755
756    #[test]
757    fn parse_claude_session_usage_sums_message_usage() {
758        let tmp = tempfile::tempdir().unwrap();
759        let path = tmp.path().join("claude.jsonl");
760        fs::write(
761            &path,
762            concat!(
763                "{\"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",
764                "{\"message\":{\"model\":\"claude-opus-4-6\",\"usage\":{\"input_tokens\":7,\"output_tokens\":4,\"cache_creation_input_tokens\":8,\"cache_read_input_tokens\":1}}}\n"
765            ),
766        )
767        .unwrap();
768
769        let usage = parse_claude_session_usage(&path).unwrap().unwrap();
770        assert_eq!(usage.model.as_deref(), Some("claude-opus-4-6"));
771        assert_eq!(usage.usage.input_tokens, 17);
772        assert_eq!(usage.usage.output_tokens, 6);
773        assert_eq!(usage.usage.cache_creation_input_tokens, 28);
774        assert_eq!(usage.usage.cache_read_input_tokens, 4);
775        assert_eq!(usage.usage.cache_creation_5m_input_tokens, 5);
776        assert_eq!(usage.usage.cache_creation_1h_input_tokens, 15);
777    }
778
779    #[test]
780    fn estimate_cost_uses_pricing_breakdown() {
781        let usage = TokenUsage {
782            input_tokens: 1_000_000,
783            cached_input_tokens: 2_000_000,
784            cache_creation_input_tokens: 1_000_000,
785            cache_creation_5m_input_tokens: 400_000,
786            cache_creation_1h_input_tokens: 500_000,
787            cache_read_input_tokens: 300_000,
788            output_tokens: 100_000,
789            reasoning_output_tokens: 50_000,
790        };
791        let pricing = ModelPricing {
792            input_usd_per_mtok: 2.0,
793            cached_input_usd_per_mtok: 0.5,
794            cache_creation_input_usd_per_mtok: Some(3.0),
795            cache_creation_5m_input_usd_per_mtok: Some(4.0),
796            cache_creation_1h_input_usd_per_mtok: Some(5.0),
797            cache_read_input_usd_per_mtok: 0.25,
798            output_usd_per_mtok: 8.0,
799            reasoning_output_usd_per_mtok: Some(10.0),
800        };
801
802        let estimated = estimate_cost_usd(&usage, &pricing);
803        let expected = 2.0 + 1.0 + 1.6 + 2.5 + 0.3 + 0.075 + 0.8 + 0.5;
804        assert!((estimated - expected).abs() < 1e-9);
805    }
806
807    #[test]
808    fn built_in_pricing_supports_common_models() {
809        assert!(pricing_for_model(&HashMap::new(), "gpt-5.4").is_some());
810        assert!(pricing_for_model(&HashMap::new(), "claude-opus-4-6").is_some());
811        assert!(pricing_for_model(&HashMap::new(), "claude-sonnet-4").is_some());
812    }
813
814    #[test]
815    fn collect_cost_report_maps_members_to_current_tasks() {
816        let tmp = tempfile::tempdir().unwrap();
817        let project_root = tmp.path();
818        fs::create_dir_all(
819            project_root
820                .join(".batty")
821                .join("worktrees")
822                .join("engineer"),
823        )
824        .unwrap();
825        fs::create_dir_all(
826            project_root
827                .join(".batty")
828                .join("team_config")
829                .join("board")
830                .join("tasks"),
831        )
832        .unwrap();
833
834        fs::write(
835            project_root.join(".batty").join("launch-state.json"),
836            r#"{
837  "architect": {"session_id": "claude-session"},
838  "engineer": {"session_id": "codex-session"}
839}"#,
840        )
841        .unwrap();
842        fs::write(
843            daemon_state_path(project_root),
844            r#"{"active_tasks":{"engineer":100}}"#,
845        )
846        .unwrap();
847        fs::write(
848            project_root
849                .join(".batty")
850                .join("team_config")
851                .join("board")
852                .join("tasks")
853                .join("100-task.md"),
854            concat!(
855                "---\n",
856                "id: 100\n",
857                "title: Cost task\n",
858                "status: in-progress\n",
859                "claimed_by: engineer\n",
860                "---\n"
861            ),
862        )
863        .unwrap();
864
865        let codex_root = project_root.join("codex-sessions");
866        let codex_day = codex_root.join("2026").join("03").join("21");
867        fs::create_dir_all(&codex_day).unwrap();
868        fs::write(
869            codex_day.join("rollout.jsonl"),
870            format!(
871                "{{\"type\":\"session_meta\",\"payload\":{{\"id\":\"codex-session\",\"cwd\":\"{}\"}}}}\n{}\n{}\n",
872                project_root
873                    .join(".batty")
874                    .join("worktrees")
875                    .join("engineer")
876                    .join(".batty")
877                    .join("codex-context")
878                    .join("engineer")
879                    .display(),
880                r#"{"type":"turn_context","payload":{"model":"gpt-5.4"}}"#,
881                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}}}}"#,
882            ),
883        )
884        .unwrap();
885
886        let claude_root = project_root.join("claude-projects");
887        let claude_dir = claude_root.join(project_root.to_string_lossy().replace('/', "-"));
888        fs::create_dir_all(&claude_dir).unwrap();
889        fs::write(
890            claude_dir.join("claude-session.jsonl"),
891            "{\"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",
892        )
893        .unwrap();
894
895        let report = collect_cost_report(
896            project_root,
897            &test_team_config(HashMap::new()),
898            &SessionRoots {
899                codex_sessions_root: codex_root,
900                claude_projects_root: claude_root,
901            },
902        )
903        .unwrap();
904
905        assert_eq!(report.entries.len(), 2);
906        let engineer = report
907            .entries
908            .iter()
909            .find(|entry| entry.member_name == "engineer")
910            .unwrap();
911        assert_eq!(engineer.task, "#100");
912        assert_eq!(engineer.model, "gpt-5.4");
913        assert!(engineer.estimated_cost_usd.unwrap() > 0.0);
914    }
915
916    #[test]
917    fn collect_cost_report_skips_user_roles_and_shared_cwd_fallbacks() {
918        let tmp = tempfile::tempdir().unwrap();
919        let project_root = tmp.path();
920        let mut config = test_team_config(HashMap::new());
921        config.roles.insert(
922            0,
923            crate::team::config::RoleDef {
924                name: "human".to_string(),
925                role_type: RoleType::User,
926                agent: None,
927                auth_mode: None,
928                auth_env: vec![],
929                instances: 1,
930                prompt: None,
931                talks_to: Vec::new(),
932                channel: None,
933                channel_config: None,
934                nudge_interval_secs: None,
935                receives_standup: None,
936                standup_interval_secs: None,
937                owns: Vec::new(),
938                barrier_group: None,
939                use_worktrees: false,
940                ..Default::default()
941            },
942        );
943        config.roles.push(crate::team::config::RoleDef {
944            name: "manager".to_string(),
945            role_type: RoleType::Manager,
946            agent: Some("claude".to_string()),
947            auth_mode: None,
948            auth_env: vec![],
949            instances: 1,
950            prompt: None,
951            talks_to: Vec::new(),
952            channel: None,
953            channel_config: None,
954            nudge_interval_secs: None,
955            receives_standup: None,
956            standup_interval_secs: None,
957            owns: Vec::new(),
958            barrier_group: None,
959            use_worktrees: false,
960            ..Default::default()
961        });
962
963        let claude_root = project_root.join("claude-projects");
964        let claude_dir = claude_root.join(project_root.to_string_lossy().replace('/', "-"));
965        fs::create_dir_all(&claude_dir).unwrap();
966        fs::write(
967            claude_dir.join("shared.jsonl"),
968            "{\"message\":{\"model\":\"claude-opus-4-6\",\"usage\":{\"input_tokens\":10,\"output_tokens\":2}}}\n",
969        )
970        .unwrap();
971
972        let report = collect_cost_report(
973            project_root,
974            &config,
975            &SessionRoots {
976                codex_sessions_root: project_root.join("codex-sessions"),
977                claude_projects_root: claude_root,
978            },
979        )
980        .unwrap();
981
982        assert!(report.entries.is_empty());
983    }
984
985    // ── TokenUsage ───────────────────────────────────────────────
986
987    #[test]
988    fn token_usage_total_tokens_sums_all_fields() {
989        let usage = TokenUsage {
990            input_tokens: 100,
991            cached_input_tokens: 50,
992            cache_creation_input_tokens: 30,
993            cache_creation_5m_input_tokens: 0,
994            cache_creation_1h_input_tokens: 0,
995            cache_read_input_tokens: 20,
996            output_tokens: 10,
997            reasoning_output_tokens: 5,
998        };
999        assert_eq!(usage.total_tokens(), 100 + 50 + 30 + 20 + 10 + 5);
1000    }
1001
1002    #[test]
1003    fn token_usage_total_tokens_zero_when_default() {
1004        let usage = TokenUsage::default();
1005        assert_eq!(usage.total_tokens(), 0);
1006    }
1007
1008    #[test]
1009    fn token_usage_display_cache_tokens_sums_cache_fields() {
1010        let usage = TokenUsage {
1011            cached_input_tokens: 100,
1012            cache_creation_input_tokens: 200,
1013            cache_read_input_tokens: 50,
1014            ..TokenUsage::default()
1015        };
1016        assert_eq!(usage.display_cache_tokens(), 350);
1017    }
1018
1019    // ── truncate_model ───────────────────────────────────────────
1020
1021    #[test]
1022    fn truncate_model_short_name_unchanged() {
1023        assert_eq!(truncate_model("gpt-5.4"), "gpt-5.4");
1024    }
1025
1026    #[test]
1027    fn truncate_model_exact_limit_unchanged() {
1028        let model = "a".repeat(20);
1029        assert_eq!(truncate_model(&model), model);
1030    }
1031
1032    #[test]
1033    fn truncate_model_long_name_truncated() {
1034        let model = "a".repeat(25);
1035        let result = truncate_model(&model);
1036        assert_eq!(result.len(), 20);
1037        assert!(result.ends_with("..."));
1038    }
1039
1040    // ── member_session_target ────────────────────────────────────
1041
1042    #[test]
1043    fn member_session_target_returns_none_for_user_role() {
1044        let member = MemberInstance {
1045            name: "human".to_string(),
1046            role_name: "human".to_string(),
1047            role_type: super::super::config::RoleType::User,
1048            agent: None,
1049            prompt: None,
1050            reports_to: None,
1051            use_worktrees: false,
1052            ..Default::default()
1053        };
1054        assert!(member_session_target(Path::new("/tmp"), &member).is_none());
1055    }
1056
1057    #[test]
1058    fn member_session_target_codex_agent() {
1059        let member = MemberInstance {
1060            name: "eng-1".to_string(),
1061            role_name: "eng".to_string(),
1062            role_type: super::super::config::RoleType::Engineer,
1063            agent: Some("codex".to_string()),
1064            prompt: None,
1065            reports_to: None,
1066            use_worktrees: false,
1067            ..Default::default()
1068        };
1069        let (agent, cwd, label) = member_session_target(Path::new("/tmp/repo"), &member).unwrap();
1070        assert!(matches!(agent, SessionAgent::Codex));
1071        assert_eq!(label, "codex");
1072        assert!(cwd.to_string_lossy().contains("codex-context"));
1073    }
1074
1075    #[test]
1076    fn member_session_target_claude_agent() {
1077        let member = MemberInstance {
1078            name: "architect".to_string(),
1079            role_name: "architect".to_string(),
1080            role_type: super::super::config::RoleType::Architect,
1081            agent: Some("claude".to_string()),
1082            prompt: None,
1083            reports_to: None,
1084            use_worktrees: false,
1085            ..Default::default()
1086        };
1087        let (agent, cwd, label) = member_session_target(Path::new("/tmp/repo"), &member).unwrap();
1088        assert!(matches!(agent, SessionAgent::Claude));
1089        assert_eq!(label, "claude");
1090        assert_eq!(cwd, Path::new("/tmp/repo"));
1091    }
1092
1093    #[test]
1094    fn member_session_target_claude_code_agent() {
1095        let member = MemberInstance {
1096            name: "eng-2".to_string(),
1097            role_name: "eng".to_string(),
1098            role_type: super::super::config::RoleType::Engineer,
1099            agent: Some("claude-code".to_string()),
1100            prompt: None,
1101            reports_to: None,
1102            use_worktrees: false,
1103            ..Default::default()
1104        };
1105        let (agent, _, label) = member_session_target(Path::new("/tmp/repo"), &member).unwrap();
1106        assert!(matches!(agent, SessionAgent::Claude));
1107        assert_eq!(label, "claude");
1108    }
1109
1110    #[test]
1111    fn member_session_target_none_agent_defaults_to_claude() {
1112        let member = MemberInstance {
1113            name: "eng-3".to_string(),
1114            role_name: "eng".to_string(),
1115            role_type: super::super::config::RoleType::Engineer,
1116            agent: None,
1117            prompt: None,
1118            reports_to: None,
1119            use_worktrees: false,
1120            ..Default::default()
1121        };
1122        let (agent, _, label) = member_session_target(Path::new("/tmp/repo"), &member).unwrap();
1123        assert!(matches!(agent, SessionAgent::Claude));
1124        assert_eq!(label, "claude");
1125    }
1126
1127    #[test]
1128    fn member_session_target_unknown_agent_returns_none() {
1129        let member = MemberInstance {
1130            name: "eng-4".to_string(),
1131            role_name: "eng".to_string(),
1132            role_type: super::super::config::RoleType::Engineer,
1133            agent: Some("gemini".to_string()),
1134            prompt: None,
1135            reports_to: None,
1136            use_worktrees: false,
1137            ..Default::default()
1138        };
1139        assert!(member_session_target(Path::new("/tmp/repo"), &member).is_none());
1140    }
1141
1142    #[test]
1143    fn member_session_target_worktree_path_for_codex() {
1144        let member = MemberInstance {
1145            name: "eng-1".to_string(),
1146            role_name: "eng".to_string(),
1147            role_type: super::super::config::RoleType::Engineer,
1148            agent: Some("codex-cli".to_string()),
1149            prompt: None,
1150            reports_to: None,
1151            use_worktrees: true,
1152            ..Default::default()
1153        };
1154        let (_, cwd, _) = member_session_target(Path::new("/tmp/repo"), &member).unwrap();
1155        assert!(cwd.starts_with("/tmp/repo/.batty/worktrees/eng-1"));
1156    }
1157
1158    #[test]
1159    fn member_session_target_worktree_path_for_claude() {
1160        let member = MemberInstance {
1161            name: "eng-1".to_string(),
1162            role_name: "eng".to_string(),
1163            role_type: super::super::config::RoleType::Engineer,
1164            agent: Some("claude".to_string()),
1165            prompt: None,
1166            reports_to: None,
1167            use_worktrees: true,
1168            ..Default::default()
1169        };
1170        let (_, cwd, _) = member_session_target(Path::new("/tmp/repo"), &member).unwrap();
1171        assert_eq!(cwd, Path::new("/tmp/repo/.batty/worktrees/eng-1"));
1172    }
1173
1174    // ── load_launch_state ────────────────────────────────────────
1175
1176    #[test]
1177    fn load_launch_state_returns_empty_when_file_missing() {
1178        let tmp = tempfile::tempdir().unwrap();
1179        let state = load_launch_state(tmp.path());
1180        assert!(state.is_empty());
1181    }
1182
1183    #[test]
1184    fn load_launch_state_parses_valid_json() {
1185        let tmp = tempfile::tempdir().unwrap();
1186        let batty_dir = tmp.path().join(".batty");
1187        fs::create_dir_all(&batty_dir).unwrap();
1188        fs::write(
1189            batty_dir.join("launch-state.json"),
1190            r#"{"eng-1": {"session_id": "abc123"}, "eng-2": {}}"#,
1191        )
1192        .unwrap();
1193
1194        let state = load_launch_state(tmp.path());
1195        assert_eq!(state.len(), 2);
1196        assert_eq!(
1197            state.get("eng-1").unwrap().session_id.as_deref(),
1198            Some("abc123")
1199        );
1200        assert!(state.get("eng-2").unwrap().session_id.is_none());
1201    }
1202
1203    #[test]
1204    fn load_launch_state_returns_empty_on_invalid_json() {
1205        let tmp = tempfile::tempdir().unwrap();
1206        let batty_dir = tmp.path().join(".batty");
1207        fs::create_dir_all(&batty_dir).unwrap();
1208        fs::write(batty_dir.join("launch-state.json"), "not json").unwrap();
1209
1210        let state = load_launch_state(tmp.path());
1211        assert!(state.is_empty());
1212    }
1213
1214    // ── load_active_tasks ────────────────────────────────────────
1215
1216    #[test]
1217    fn load_active_tasks_returns_empty_when_file_missing() {
1218        let tmp = tempfile::tempdir().unwrap();
1219        let tasks = load_active_tasks(tmp.path());
1220        assert!(tasks.is_empty());
1221    }
1222
1223    #[test]
1224    fn load_active_tasks_parses_valid_state() {
1225        let tmp = tempfile::tempdir().unwrap();
1226        let state_path = daemon_state_path(tmp.path());
1227        fs::create_dir_all(state_path.parent().unwrap()).unwrap();
1228        fs::write(&state_path, r#"{"active_tasks":{"eng-1":42,"eng-2":99}}"#).unwrap();
1229
1230        let tasks = load_active_tasks(tmp.path());
1231        assert_eq!(tasks.len(), 2);
1232        assert_eq!(tasks.get("eng-1"), Some(&42));
1233        assert_eq!(tasks.get("eng-2"), Some(&99));
1234    }
1235
1236    #[test]
1237    fn load_active_tasks_returns_empty_on_invalid_json() {
1238        let tmp = tempfile::tempdir().unwrap();
1239        let state_path = daemon_state_path(tmp.path());
1240        fs::create_dir_all(state_path.parent().unwrap()).unwrap();
1241        fs::write(&state_path, "garbage").unwrap();
1242
1243        let tasks = load_active_tasks(tmp.path());
1244        assert!(tasks.is_empty());
1245    }
1246
1247    // ── pricing_for_model ────────────────────────────────────────
1248
1249    #[test]
1250    fn pricing_for_model_uses_override_exact_match() {
1251        let mut overrides = HashMap::new();
1252        overrides.insert(
1253            "custom-model".to_string(),
1254            ModelPricing {
1255                input_usd_per_mtok: 99.0,
1256                cached_input_usd_per_mtok: 0.0,
1257                cache_creation_input_usd_per_mtok: None,
1258                cache_creation_5m_input_usd_per_mtok: None,
1259                cache_creation_1h_input_usd_per_mtok: None,
1260                cache_read_input_usd_per_mtok: 0.0,
1261                output_usd_per_mtok: 99.0,
1262                reasoning_output_usd_per_mtok: None,
1263            },
1264        );
1265        let pricing = pricing_for_model(&overrides, "custom-model").unwrap();
1266        assert!((pricing.input_usd_per_mtok - 99.0).abs() < f64::EPSILON);
1267    }
1268
1269    #[test]
1270    fn pricing_for_model_normalized_case_match() {
1271        let mut overrides = HashMap::new();
1272        overrides.insert(
1273            "my-model".to_string(),
1274            ModelPricing {
1275                input_usd_per_mtok: 5.0,
1276                cached_input_usd_per_mtok: 0.0,
1277                cache_creation_input_usd_per_mtok: None,
1278                cache_creation_5m_input_usd_per_mtok: None,
1279                cache_creation_1h_input_usd_per_mtok: None,
1280                cache_read_input_usd_per_mtok: 0.0,
1281                output_usd_per_mtok: 5.0,
1282                reasoning_output_usd_per_mtok: None,
1283            },
1284        );
1285        let pricing = pricing_for_model(&overrides, "My-Model").unwrap();
1286        assert!((pricing.input_usd_per_mtok - 5.0).abs() < f64::EPSILON);
1287    }
1288
1289    #[test]
1290    fn pricing_for_model_returns_none_for_unknown() {
1291        assert!(pricing_for_model(&HashMap::new(), "totally-unknown-model").is_none());
1292    }
1293
1294    // ── built_in_model_pricing ───────────────────────────────────
1295
1296    #[test]
1297    fn built_in_pricing_gpt54_has_reasoning_rate() {
1298        let pricing = built_in_model_pricing("gpt-5.4").unwrap();
1299        assert!(pricing.reasoning_output_usd_per_mtok.is_some());
1300    }
1301
1302    #[test]
1303    fn built_in_pricing_opus_has_cache_tiers() {
1304        let pricing = built_in_model_pricing("claude-opus-4-6").unwrap();
1305        assert!(pricing.cache_creation_5m_input_usd_per_mtok.is_some());
1306        assert!(pricing.cache_creation_1h_input_usd_per_mtok.is_some());
1307    }
1308
1309    #[test]
1310    fn built_in_pricing_sonnet_has_cache_tiers() {
1311        let pricing = built_in_model_pricing("claude-sonnet-4-6").unwrap();
1312        assert!(pricing.cache_creation_5m_input_usd_per_mtok.is_some());
1313        assert!(pricing.cache_creation_1h_input_usd_per_mtok.is_some());
1314    }
1315
1316    #[test]
1317    fn built_in_pricing_unknown_returns_none() {
1318        assert!(built_in_model_pricing("llama-3").is_none());
1319    }
1320
1321    // ── estimate_cost_usd edge cases ─────────────────────────────
1322
1323    #[test]
1324    fn estimate_cost_zero_usage_returns_zero() {
1325        let usage = TokenUsage::default();
1326        let pricing = ModelPricing {
1327            input_usd_per_mtok: 10.0,
1328            cached_input_usd_per_mtok: 5.0,
1329            cache_creation_input_usd_per_mtok: None,
1330            cache_creation_5m_input_usd_per_mtok: None,
1331            cache_creation_1h_input_usd_per_mtok: None,
1332            cache_read_input_usd_per_mtok: 1.0,
1333            output_usd_per_mtok: 20.0,
1334            reasoning_output_usd_per_mtok: None,
1335        };
1336        assert!((estimate_cost_usd(&usage, &pricing)).abs() < f64::EPSILON);
1337    }
1338
1339    #[test]
1340    fn estimate_cost_uses_output_rate_when_no_reasoning_rate() {
1341        let usage = TokenUsage {
1342            reasoning_output_tokens: 1_000_000,
1343            ..TokenUsage::default()
1344        };
1345        let pricing = ModelPricing {
1346            input_usd_per_mtok: 0.0,
1347            cached_input_usd_per_mtok: 0.0,
1348            cache_creation_input_usd_per_mtok: None,
1349            cache_creation_5m_input_usd_per_mtok: None,
1350            cache_creation_1h_input_usd_per_mtok: None,
1351            cache_read_input_usd_per_mtok: 0.0,
1352            output_usd_per_mtok: 10.0,
1353            reasoning_output_usd_per_mtok: None,
1354        };
1355        // reasoning falls back to output rate: 1M tokens * 10/M = $10
1356        assert!((estimate_cost_usd(&usage, &pricing) - 10.0).abs() < f64::EPSILON);
1357    }
1358
1359    #[test]
1360    fn estimate_cost_unclassified_cache_creation_uses_generic_rate() {
1361        let usage = TokenUsage {
1362            cache_creation_input_tokens: 1_000_000,
1363            // No 5m or 1h breakdown → all goes to unclassified
1364            ..TokenUsage::default()
1365        };
1366        let pricing = ModelPricing {
1367            input_usd_per_mtok: 0.0,
1368            cached_input_usd_per_mtok: 0.0,
1369            cache_creation_input_usd_per_mtok: Some(5.0),
1370            cache_creation_5m_input_usd_per_mtok: None,
1371            cache_creation_1h_input_usd_per_mtok: None,
1372            cache_read_input_usd_per_mtok: 0.0,
1373            output_usd_per_mtok: 0.0,
1374            reasoning_output_usd_per_mtok: None,
1375        };
1376        assert!((estimate_cost_usd(&usage, &pricing) - 5.0).abs() < f64::EPSILON);
1377    }
1378
1379    #[test]
1380    fn estimate_cost_generic_rate_falls_back_to_5m_then_input() {
1381        let usage = TokenUsage {
1382            cache_creation_input_tokens: 1_000_000,
1383            ..TokenUsage::default()
1384        };
1385        // No generic rate → falls back to 5m rate
1386        let pricing = ModelPricing {
1387            input_usd_per_mtok: 2.0,
1388            cached_input_usd_per_mtok: 0.0,
1389            cache_creation_input_usd_per_mtok: None,
1390            cache_creation_5m_input_usd_per_mtok: Some(4.0),
1391            cache_creation_1h_input_usd_per_mtok: None,
1392            cache_read_input_usd_per_mtok: 0.0,
1393            output_usd_per_mtok: 0.0,
1394            reasoning_output_usd_per_mtok: None,
1395        };
1396        // unclassified = 1M, generic rate = 5m rate = 4.0 → $4
1397        assert!((estimate_cost_usd(&usage, &pricing) - 4.0).abs() < f64::EPSILON);
1398
1399        // No generic rate AND no 5m rate → falls back to input rate
1400        let pricing2 = ModelPricing {
1401            input_usd_per_mtok: 2.0,
1402            cached_input_usd_per_mtok: 0.0,
1403            cache_creation_input_usd_per_mtok: None,
1404            cache_creation_5m_input_usd_per_mtok: None,
1405            cache_creation_1h_input_usd_per_mtok: None,
1406            cache_read_input_usd_per_mtok: 0.0,
1407            output_usd_per_mtok: 0.0,
1408            reasoning_output_usd_per_mtok: None,
1409        };
1410        assert!((estimate_cost_usd(&usage, &pricing2) - 2.0).abs() < f64::EPSILON);
1411    }
1412
1413    // ── parse session files edge cases ───────────────────────────
1414
1415    #[test]
1416    fn parse_codex_session_returns_none_for_empty_file() {
1417        let tmp = tempfile::tempdir().unwrap();
1418        let path = tmp.path().join("empty.jsonl");
1419        fs::write(&path, "").unwrap();
1420
1421        assert!(parse_codex_session_usage(&path).unwrap().is_none());
1422    }
1423
1424    #[test]
1425    fn parse_codex_session_skips_malformed_lines() {
1426        let tmp = tempfile::tempdir().unwrap();
1427        let path = tmp.path().join("mixed.jsonl");
1428        fs::write(
1429            &path,
1430            concat!(
1431                "not valid json\n",
1432                "{\"type\":\"turn_context\",\"payload\":{\"model\":\"gpt-5.4\"}}\n",
1433                "{\"type\":\"event_msg\",\"payload\":{\"type\":\"token_count\",\"info\":{\"last_token_usage\":{\"input_tokens\":100}}}}\n",
1434            ),
1435        )
1436        .unwrap();
1437
1438        let usage = parse_codex_session_usage(&path).unwrap().unwrap();
1439        assert_eq!(usage.model.as_deref(), Some("gpt-5.4"));
1440        assert_eq!(usage.usage.input_tokens, 100);
1441    }
1442
1443    #[test]
1444    fn parse_claude_session_returns_none_for_empty_file() {
1445        let tmp = tempfile::tempdir().unwrap();
1446        let path = tmp.path().join("empty.jsonl");
1447        fs::write(&path, "").unwrap();
1448
1449        assert!(parse_claude_session_usage(&path).unwrap().is_none());
1450    }
1451
1452    #[test]
1453    fn parse_claude_session_skips_entries_without_message() {
1454        let tmp = tempfile::tempdir().unwrap();
1455        let path = tmp.path().join("mixed.jsonl");
1456        fs::write(
1457            &path,
1458            concat!(
1459                "{\"not_message\":true}\n",
1460                "{\"message\":{\"model\":\"claude-sonnet-4\",\"usage\":{\"input_tokens\":50,\"output_tokens\":10}}}\n",
1461            ),
1462        )
1463        .unwrap();
1464
1465        let usage = parse_claude_session_usage(&path).unwrap().unwrap();
1466        assert_eq!(usage.model.as_deref(), Some("claude-sonnet-4"));
1467        assert_eq!(usage.usage.input_tokens, 50);
1468    }
1469
1470    // ── json_u64 ─────────────────────────────────────────────────
1471
1472    #[test]
1473    fn json_u64_returns_value_for_number() {
1474        let v: Value = serde_json::json!(42);
1475        assert_eq!(json_u64(Some(&v)), 42);
1476    }
1477
1478    #[test]
1479    fn json_u64_returns_zero_for_none() {
1480        assert_eq!(json_u64(None), 0);
1481    }
1482
1483    #[test]
1484    fn json_u64_returns_zero_for_non_number() {
1485        let v: Value = serde_json::json!("not a number");
1486        assert_eq!(json_u64(Some(&v)), 0);
1487    }
1488
1489    // ── read_codex_session_meta ──────────────────────────────────
1490
1491    #[test]
1492    fn read_codex_session_meta_parses_session_meta_line() {
1493        let tmp = tempfile::tempdir().unwrap();
1494        let path = tmp.path().join("session.jsonl");
1495        fs::write(
1496            &path,
1497            "{\"type\":\"session_meta\",\"payload\":{\"id\":\"sess-123\",\"cwd\":\"/tmp/work\"}}\n",
1498        )
1499        .unwrap();
1500
1501        let meta = read_codex_session_meta(&path).unwrap().unwrap();
1502        assert_eq!(meta.id.as_deref(), Some("sess-123"));
1503        assert_eq!(meta.cwd.as_deref(), Some(std::ffi::OsStr::new("/tmp/work")));
1504    }
1505
1506    #[test]
1507    fn read_codex_session_meta_returns_none_for_no_meta() {
1508        let tmp = tempfile::tempdir().unwrap();
1509        let path = tmp.path().join("no-meta.jsonl");
1510        fs::write(
1511            &path,
1512            "{\"type\":\"event_msg\",\"payload\":{\"type\":\"token_count\"}}\n",
1513        )
1514        .unwrap();
1515
1516        assert!(read_codex_session_meta(&path).unwrap().is_none());
1517    }
1518
1519    #[test]
1520    fn read_codex_session_meta_skips_blank_and_malformed_lines() {
1521        let tmp = tempfile::tempdir().unwrap();
1522        let path = tmp.path().join("messy.jsonl");
1523        fs::write(
1524            &path,
1525            concat!(
1526                "\n",
1527                "not json\n",
1528                "{\"type\":\"session_meta\",\"payload\":{\"id\":\"found\"}}\n",
1529            ),
1530        )
1531        .unwrap();
1532
1533        let meta = read_codex_session_meta(&path).unwrap().unwrap();
1534        assert_eq!(meta.id.as_deref(), Some("found"));
1535    }
1536
1537    // ── discover session files ───────────────────────────────────
1538
1539    #[test]
1540    fn discover_codex_session_returns_none_when_root_missing() {
1541        let result = discover_codex_session_file(
1542            Path::new("/nonexistent/path"),
1543            Path::new("/tmp/cwd"),
1544            None,
1545            true,
1546        )
1547        .unwrap();
1548        assert!(result.is_none());
1549    }
1550
1551    #[test]
1552    fn discover_claude_session_returns_none_when_root_missing() {
1553        let result = discover_claude_session_file(
1554            Path::new("/nonexistent/path"),
1555            Path::new("/tmp/cwd"),
1556            None,
1557            true,
1558        )
1559        .unwrap();
1560        assert!(result.is_none());
1561    }
1562
1563    #[test]
1564    fn discover_claude_session_finds_exact_session_id() {
1565        let tmp = tempfile::tempdir().unwrap();
1566        let projects_root = tmp.path().join("projects");
1567        let project_dir = projects_root.join("-tmp-myrepo");
1568        fs::create_dir_all(&project_dir).unwrap();
1569        fs::write(
1570            project_dir.join("session-abc.jsonl"),
1571            "{\"message\":{\"model\":\"claude\"}}\n",
1572        )
1573        .unwrap();
1574
1575        let result = discover_claude_session_file(
1576            &projects_root,
1577            Path::new("/tmp/myrepo"),
1578            Some("session-abc"),
1579            true,
1580        )
1581        .unwrap();
1582        assert!(result.is_some());
1583        assert!(result.unwrap().ends_with("session-abc.jsonl"));
1584    }
1585
1586    #[test]
1587    fn discover_claude_session_cwd_fallback_finds_newest() {
1588        let tmp = tempfile::tempdir().unwrap();
1589        let projects_root = tmp.path().join("projects");
1590        let project_dir = projects_root.join("-tmp-myrepo");
1591        fs::create_dir_all(&project_dir).unwrap();
1592
1593        fs::write(project_dir.join("old.jsonl"), "old\n").unwrap();
1594        // Give a small delay or touch to make newer file distinguishable
1595        fs::write(project_dir.join("new.jsonl"), "new\n").unwrap();
1596
1597        let result =
1598            discover_claude_session_file(&projects_root, Path::new("/tmp/myrepo"), None, true)
1599                .unwrap();
1600        assert!(result.is_some());
1601    }
1602
1603    #[test]
1604    fn discover_claude_session_no_cwd_fallback_when_disabled() {
1605        let tmp = tempfile::tempdir().unwrap();
1606        let projects_root = tmp.path().join("projects");
1607        let project_dir = projects_root.join("-tmp-myrepo");
1608        fs::create_dir_all(&project_dir).unwrap();
1609        fs::write(project_dir.join("session.jsonl"), "data\n").unwrap();
1610
1611        let result = discover_claude_session_file(
1612            &projects_root,
1613            Path::new("/tmp/myrepo"),
1614            None,
1615            false, // cwd fallback disabled
1616        )
1617        .unwrap();
1618        assert!(result.is_none());
1619    }
1620
1621    // ── read_dir_paths ───────────────────────────────────────────
1622
1623    #[test]
1624    fn read_dir_paths_returns_entries() {
1625        let tmp = tempfile::tempdir().unwrap();
1626        fs::write(tmp.path().join("a.txt"), "a").unwrap();
1627        fs::write(tmp.path().join("b.txt"), "b").unwrap();
1628        let paths = read_dir_paths(tmp.path()).unwrap();
1629        assert_eq!(paths.len(), 2);
1630    }
1631
1632    #[test]
1633    fn read_dir_paths_returns_empty_for_empty_dir() {
1634        let tmp = tempfile::tempdir().unwrap();
1635        let paths = read_dir_paths(tmp.path()).unwrap();
1636        assert!(paths.is_empty());
1637    }
1638
1639    // ── SessionRoots default ─────────────────────────────────────
1640
1641    #[test]
1642    fn session_roots_default_uses_home() {
1643        let roots = SessionRoots::default();
1644        let home = std::env::var("HOME").unwrap_or_default();
1645        assert!(roots.codex_sessions_root.starts_with(&home));
1646        assert!(roots.claude_projects_root.starts_with(&home));
1647    }
1648
1649    // ── collect_cost_report empty board ──────────────────────────
1650
1651    #[test]
1652    fn collect_cost_report_empty_when_no_sessions() {
1653        let tmp = tempfile::tempdir().unwrap();
1654        let project_root = tmp.path();
1655        fs::create_dir_all(
1656            project_root
1657                .join(".batty")
1658                .join("team_config")
1659                .join("board")
1660                .join("tasks"),
1661        )
1662        .unwrap();
1663
1664        let report = collect_cost_report(
1665            project_root,
1666            &test_team_config(HashMap::new()),
1667            &SessionRoots {
1668                codex_sessions_root: project_root.join("no-codex"),
1669                claude_projects_root: project_root.join("no-claude"),
1670            },
1671        )
1672        .unwrap();
1673
1674        assert!(report.entries.is_empty());
1675        assert!((report.total_estimated_cost_usd).abs() < f64::EPSILON);
1676        assert!(report.unpriced_models.is_empty());
1677    }
1678
1679    #[test]
1680    fn collect_cost_report_tracks_unpriced_models() {
1681        let tmp = tempfile::tempdir().unwrap();
1682        let project_root = tmp.path();
1683        fs::create_dir_all(
1684            project_root
1685                .join(".batty")
1686                .join("team_config")
1687                .join("board")
1688                .join("tasks"),
1689        )
1690        .unwrap();
1691
1692        // Create a claude session with an unknown model
1693        let claude_root = project_root.join("claude-projects");
1694        let claude_dir = claude_root.join(project_root.to_string_lossy().replace('/', "-"));
1695        fs::create_dir_all(&claude_dir).unwrap();
1696        fs::write(
1697            claude_dir.join("session.jsonl"),
1698            "{\"message\":{\"model\":\"totally-unknown-model\",\"usage\":{\"input_tokens\":100,\"output_tokens\":10}}}\n",
1699        )
1700        .unwrap();
1701
1702        let report = collect_cost_report(
1703            project_root,
1704            &test_team_config(HashMap::new()),
1705            &SessionRoots {
1706                codex_sessions_root: project_root.join("no-codex"),
1707                claude_projects_root: claude_root,
1708            },
1709        )
1710        .unwrap();
1711
1712        assert!(report.unpriced_models.contains("totally-unknown-model"));
1713    }
1714}