huskies: merge 1017

2026-05-13 23:51:12 +00:00
parent 29e800da21
commit 52180bc402
8 changed files with 651 additions and 35 deletions
@@ -0,0 +1,266 @@
+//! In-memory CostRollup register, keyed by project root.
+//!
+//! Keying by `project_root` provides test isolation: each test's `TempDir`
+//! maps to its own slice of the store, so parallel tests cannot pollute
+//! each other's data.
+//!
+//! Populated on server startup from existing JSONL records
+//! ([`init_from_disk`]) and kept current by the cost-rollup subscriber
+//! (`agents::pool::cost_rollup_subscriber`) which fires on every terminal
+//! pipeline stage transition.
+//!
+//! All readers call [`get_rollup`] or [`all_rollups`] instead of
+//! re-walking `token_usage.jsonl`.
+
+use chrono::{DateTime, Utc};
+use std::collections::HashMap;
+use std::path::{Path, PathBuf};
+use std::sync::{OnceLock, RwLock};
+
+use super::token::{AgentTokenCost, TokenCostSummary};
+
+// ── Types ─────────────────────────────────────────────────────────────────────
+
+/// Snapshotted cost totals for a story at the time it reached a terminal stage.
+#[derive(Debug, Clone)]
+pub struct CostRollup {
+    /// The story this rollup belongs to.
+    pub story_id: String,
+    /// Total USD spend across all agents.
+    pub total_cost_usd: f64,
+    /// Per-agent breakdown, sorted alphabetically by agent name.
+    pub agents: Vec<AgentTokenCost>,
+    /// When the rollup was written (terminal transition timestamp or, for
+    /// startup-bootstrapped entries, the latest token record timestamp).
+    pub recorded_at: DateTime<Utc>,
+}
+
+impl CostRollup {
+    /// Convert to a [`TokenCostSummary`] for callers that expect the legacy type.
+    pub fn as_summary(&self) -> TokenCostSummary {
+        TokenCostSummary {
+            total_cost_usd: self.total_cost_usd,
+            agents: self.agents.clone(),
+        }
+    }
+}
+
+// ── Global store ──────────────────────────────────────────────────────────────
+
+/// Composite key: (canonical project root, story_id).
+type StoreKey = (PathBuf, String);
+
+static STORE: OnceLock<RwLock<HashMap<StoreKey, CostRollup>>> = OnceLock::new();
+
+fn store() -> &'static RwLock<HashMap<StoreKey, CostRollup>> {
+    STORE.get_or_init(|| RwLock::new(HashMap::new()))
+}
+
+fn canonical(p: &Path) -> PathBuf {
+    p.canonicalize().unwrap_or_else(|_| p.to_path_buf())
+}
+
+/// Look up the cost rollup for `story_id` within `project_root`.
+///
+/// Returns `None` for stories that have not yet reached a terminal stage
+/// (or have never had any token records).
+pub fn get_rollup(project_root: &Path, story_id: &str) -> Option<CostRollup> {
+    let key = (canonical(project_root), story_id.to_string());
+    store().read().ok()?.get(&key).cloned()
+}
+
+/// Write or overwrite the cost rollup for `story_id` within `project_root`.
+pub fn set_rollup(project_root: &Path, story_id: impl Into<String>, rollup: CostRollup) {
+    let key = (canonical(project_root), story_id.into());
+    if let Ok(mut guard) = store().write() {
+        guard.insert(key, rollup);
+    }
+}
+
+/// Return a snapshot of all rollups for `project_root`, in arbitrary order.
+pub fn all_rollups(project_root: &Path) -> Vec<CostRollup> {
+    let root = canonical(project_root);
+    store()
+        .read()
+        .map(|g| {
+            g.iter()
+                .filter(|((r, _), _)| r == &root)
+                .map(|(_, v)| v.clone())
+                .collect()
+        })
+        .unwrap_or_default()
+}
+
+// ── Startup bootstrap ─────────────────────────────────────────────────────────
+
+/// Pre-populate the register from `token_usage.jsonl`.
+///
+/// Called once at server startup before the live subscriber starts listening.
+/// Stories with zero cost are skipped.  `recorded_at` is set to the latest
+/// record timestamp for each story so the 24 h window in the cost command
+/// reflects when work was actually done rather than when the server restarted.
+pub fn init_from_disk(project_root: &Path) {
+    let records = match crate::agents::token_usage::read_all(project_root) {
+        Ok(r) => r,
+        Err(e) => {
+            crate::slog_warn!("[cost-rollup] init_from_disk failed to read token records: {e}");
+            return;
+        }
+    };
+
+    if records.is_empty() {
+        return;
+    }
+
+    // Collect unique story IDs.
+    let mut story_ids: Vec<&str> = records.iter().map(|r| r.story_id.as_str()).collect();
+    story_ids.sort_unstable();
+    story_ids.dedup();
+
+    for sid in story_ids {
+        let summary = super::token::aggregate_for_story(&records, sid);
+        if summary.total_cost_usd == 0.0 && summary.agents.is_empty() {
+            continue;
+        }
+
+        // Use the latest record timestamp as recorded_at.
+        let recorded_at = records
+            .iter()
+            .filter(|r| r.story_id == sid)
+            .filter_map(|r| chrono::DateTime::parse_from_rfc3339(&r.timestamp).ok())
+            .map(|ts| ts.with_timezone(&Utc))
+            .max()
+            .unwrap_or_else(Utc::now);
+
+        set_rollup(
+            project_root,
+            sid,
+            CostRollup {
+                story_id: sid.to_string(),
+                total_cost_usd: summary.total_cost_usd,
+                agents: summary.agents,
+                recorded_at,
+            },
+        );
+    }
+}
+
+// ── Tests ─────────────────────────────────────────────────────────────────────
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::agents::TokenUsage;
+    use crate::agents::token_usage::TokenUsageRecord;
+
+    fn dummy_root() -> PathBuf {
+        std::env::temp_dir().join("cost_rollup_dummy_root")
+    }
+
+    fn make_record(story_id: &str, agent: &str, cost: f64, ts: &str) -> TokenUsageRecord {
+        TokenUsageRecord {
+            story_id: story_id.to_string(),
+            agent_name: agent.to_string(),
+            timestamp: ts.to_string(),
+            model: None,
+            usage: TokenUsage {
+                input_tokens: 100,
+                output_tokens: 50,
+                cache_creation_input_tokens: 0,
+                cache_read_input_tokens: 0,
+                total_cost_usd: cost,
+            },
+        }
+    }
+
+    #[test]
+    fn get_rollup_returns_none_for_unknown_story() {
+        let root = dummy_root();
+        assert!(get_rollup(&root, "999_no_such_story_xyzzy").is_none());
+    }
+
+    #[test]
+    fn set_and_get_rollup_roundtrip() {
+        let root = dummy_root();
+        let rollup = CostRollup {
+            story_id: "1017_test_roundtrip".to_string(),
+            total_cost_usd: 3.50,
+            agents: vec![],
+            recorded_at: Utc::now(),
+        };
+        set_rollup(&root, "1017_test_roundtrip", rollup);
+        let got = get_rollup(&root, "1017_test_roundtrip").expect("rollup must exist after set");
+        assert!((got.total_cost_usd - 3.50).abs() < f64::EPSILON);
+    }
+
+    #[test]
+    fn init_from_disk_populates_store() {
+        let tmp = tempfile::tempdir().unwrap();
+        let records = vec![
+            make_record("1017_init_a", "coder-1", 1.0, "2026-01-01T00:00:00Z"),
+            make_record("1017_init_a", "qa-1", 0.5, "2026-01-02T00:00:00Z"),
+            make_record("1017_init_b", "coder-1", 2.0, "2026-01-01T12:00:00Z"),
+        ];
+        for r in &records {
+            crate::agents::token_usage::append_record(tmp.path(), r).unwrap();
+        }
+        init_from_disk(tmp.path());
+
+        let a = get_rollup(tmp.path(), "1017_init_a").expect("story a must be in store");
+        assert!((a.total_cost_usd - 1.5).abs() < f64::EPSILON);
+        assert_eq!(a.agents.len(), 2);
+
+        let b = get_rollup(tmp.path(), "1017_init_b").expect("story b must be in store");
+        assert!((b.total_cost_usd - 2.0).abs() < f64::EPSILON);
+    }
+
+    #[test]
+    fn init_from_disk_uses_latest_timestamp_as_recorded_at() {
+        let tmp = tempfile::tempdir().unwrap();
+        let records = vec![
+            make_record("1017_ts_test", "coder-1", 1.0, "2026-01-01T00:00:00Z"),
+            make_record("1017_ts_test", "qa-1", 0.5, "2026-01-03T00:00:00Z"),
+        ];
+        for r in &records {
+            crate::agents::token_usage::append_record(tmp.path(), r).unwrap();
+        }
+        init_from_disk(tmp.path());
+
+        let rollup = get_rollup(tmp.path(), "1017_ts_test").expect("rollup must exist");
+        // recorded_at should be 2026-01-03 (the later timestamp)
+        assert_eq!(rollup.recorded_at.date_naive().to_string(), "2026-01-03");
+    }
+
+    #[test]
+    fn all_rollups_isolated_by_project_root() {
+        let tmp_a = tempfile::tempdir().unwrap();
+        let tmp_b = tempfile::tempdir().unwrap();
+        set_rollup(
+            tmp_a.path(),
+            "1017_all_a",
+            CostRollup {
+                story_id: "1017_all_a".to_string(),
+                total_cost_usd: 1.0,
+                agents: vec![],
+                recorded_at: Utc::now(),
+            },
+        );
+        set_rollup(
+            tmp_b.path(),
+            "1017_all_b",
+            CostRollup {
+                story_id: "1017_all_b".to_string(),
+                total_cost_usd: 2.0,
+                agents: vec![],
+                recorded_at: Utc::now(),
+            },
+        );
+        let all_a = all_rollups(tmp_a.path());
+        assert!(all_a.iter().any(|r| r.story_id == "1017_all_a"));
+        assert!(!all_a.iter().any(|r| r.story_id == "1017_all_b"));
+
+        let all_b = all_rollups(tmp_b.path());
+        assert!(all_b.iter().any(|r| r.story_id == "1017_all_b"));
+        assert!(!all_b.iter().any(|r| r.story_id == "1017_all_a"));
+    }
+}
@@ -6,6 +6,8 @@
 //! into `AgentPool` or the filesystem.
 //!
 //! Conventions: `docs/architecture/service-modules.md`
+/// In-memory cost rollup register — written by the cost-rollup subscriber.
+pub mod cost_rollup;
 mod io;
 /// Agent selection heuristics — pick the best agent for a story.
 pub mod selection;
@@ -215,13 +217,20 @@ pub fn get_test_results(
    io::read_test_results_from_file(project_root, story_id)
 }

-/// Get the aggregated token cost for a specific story.
+/// Get the aggregated token cost for a specific story from the rollup register.
+///
+/// Returns a zero-cost summary when no rollup has been recorded yet (story
+/// still in progress or no token records exist).
 pub fn get_work_item_token_cost(
    project_root: &Path,
    story_id: &str,
 ) -> Result<TokenCostSummary, Error> {
-    let records = io::read_token_records(project_root)?;
-    Ok(token::aggregate_for_story(&records, story_id))
+    Ok(cost_rollup::get_rollup(project_root, story_id)
+        .map(|r| r.as_summary())
+        .unwrap_or(TokenCostSummary {
+            total_cost_usd: 0.0,
+            agents: vec![],
+        }))
 }

 /// Get all token usage records across all stories.