huskies: merge 1017
This commit is contained in:
dave
@@ -0,0 +1,333 @@
|
||||
//! TransitionFired subscriber that snapshots token cost into the CostRollup register
|
||||
//! on terminal pipeline stage transitions.
|
||||
//!
|
||||
//! When a story reaches Done, Archived, Abandoned, Superseded, or Rejected,
|
||||
//! this subscriber reads the accumulated token records from disk, aggregates
|
||||
//! them by agent, and writes the result into
|
||||
//! [`crate::service::agents::cost_rollup`] so that status renderers and the
|
||||
//! cost command can read from the register instead of re-walking the JSONL file.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use crate::pipeline_state::Stage;
|
||||
use crate::slog;
|
||||
use crate::slog_warn;
|
||||
|
||||
/// Spawn a background task that maintains the CostRollup register.
|
||||
///
|
||||
/// On every terminal stage transition (Done, Archived, Abandoned, Superseded,
|
||||
/// Rejected), reads the token records for that story and writes an aggregated
|
||||
/// [`crate::service::agents::cost_rollup::CostRollup`] to the global register.
|
||||
pub(crate) fn spawn_cost_rollup_subscriber(project_root: PathBuf) {
|
||||
let mut rx = crate::pipeline_state::subscribe_transitions();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
match rx.recv().await {
|
||||
Ok(fired) => {
|
||||
on_terminal_transition(&project_root, &fired).await;
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
|
||||
slog_warn!(
|
||||
"[cost-rollup-sub] Subscriber lagged, skipped {n} event(s). \
|
||||
Some story cost rollups may be stale until next terminal transition."
|
||||
);
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Returns `true` if `stage` is a terminal pipeline stage.
|
||||
///
|
||||
/// Terminal stages are those from which no further work is expected:
|
||||
/// Done, Archived, Abandoned, Superseded, Rejected.
|
||||
/// MergeFailure variants are NOT terminal — stories can recover from them.
|
||||
fn is_terminal(stage: &Stage) -> bool {
|
||||
matches!(
|
||||
stage,
|
||||
Stage::Done { .. }
|
||||
| Stage::Archived { .. }
|
||||
| Stage::Abandoned { .. }
|
||||
| Stage::Superseded { .. }
|
||||
| Stage::Rejected { .. }
|
||||
)
|
||||
}
|
||||
|
||||
/// Snapshot the cost data for `fired.story_id` into the register when
|
||||
/// `fired.after` is a terminal stage.
|
||||
pub(crate) async fn on_terminal_transition(
|
||||
project_root: &Path,
|
||||
fired: &crate::pipeline_state::TransitionFired,
|
||||
) {
|
||||
if !is_terminal(&fired.after) {
|
||||
return;
|
||||
}
|
||||
|
||||
let story_id = &fired.story_id.0;
|
||||
let records = match crate::agents::token_usage::read_all(project_root) {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
slog_warn!("[cost-rollup-sub] Failed to read token records for '{story_id}': {e}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let summary = crate::service::agents::token::aggregate_for_story(&records, story_id);
|
||||
|
||||
crate::service::agents::cost_rollup::set_rollup(
|
||||
project_root,
|
||||
story_id,
|
||||
crate::service::agents::cost_rollup::CostRollup {
|
||||
story_id: story_id.clone(),
|
||||
total_cost_usd: summary.total_cost_usd,
|
||||
agents: summary.agents,
|
||||
recorded_at: fired.at,
|
||||
},
|
||||
);
|
||||
|
||||
slog!(
|
||||
"[cost-rollup-sub] Rolled up cost for '{}': ${:.4}",
|
||||
story_id,
|
||||
summary.total_cost_usd
|
||||
);
|
||||
}
|
||||
|
||||
// ── Tests ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::agents::TokenUsage;
|
||||
use crate::agents::token_usage::TokenUsageRecord;
|
||||
use crate::pipeline_state::{BranchName, PipelineEvent, Stage, StoryId, TransitionFired};
|
||||
use chrono::Utc;
|
||||
use std::num::NonZeroU32;
|
||||
|
||||
fn make_record(story_id: &str, agent: &str, cost: f64) -> TokenUsageRecord {
|
||||
TokenUsageRecord {
|
||||
story_id: story_id.to_string(),
|
||||
agent_name: agent.to_string(),
|
||||
timestamp: Utc::now().to_rfc3339(),
|
||||
model: None,
|
||||
usage: TokenUsage {
|
||||
input_tokens: 100,
|
||||
output_tokens: 50,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
total_cost_usd: cost,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn write_records(root: &std::path::Path, records: &[TokenUsageRecord]) {
|
||||
for r in records {
|
||||
crate::agents::token_usage::append_record(root, r).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
fn fired_done(story_id: &str) -> TransitionFired {
|
||||
TransitionFired {
|
||||
story_id: StoryId(story_id.to_string()),
|
||||
before: Stage::Merge {
|
||||
feature_branch: BranchName("feature/test".to_string()),
|
||||
commits_ahead: NonZeroU32::new(1).unwrap(),
|
||||
claim: None,
|
||||
},
|
||||
after: Stage::Done {
|
||||
merged_at: Utc::now(),
|
||||
merge_commit: crate::pipeline_state::GitSha("abc123".to_string()),
|
||||
},
|
||||
event: PipelineEvent::MergeSucceeded {
|
||||
merge_commit: crate::pipeline_state::GitSha("abc123".to_string()),
|
||||
},
|
||||
at: Utc::now(),
|
||||
}
|
||||
}
|
||||
|
||||
fn fired_abandoned(story_id: &str) -> TransitionFired {
|
||||
TransitionFired {
|
||||
story_id: StoryId(story_id.to_string()),
|
||||
before: Stage::Coding { claim: None },
|
||||
after: Stage::Abandoned { ts: Utc::now() },
|
||||
event: PipelineEvent::Abandon,
|
||||
at: Utc::now(),
|
||||
}
|
||||
}
|
||||
|
||||
fn fired_coding(story_id: &str) -> TransitionFired {
|
||||
TransitionFired {
|
||||
story_id: StoryId(story_id.to_string()),
|
||||
before: Stage::Backlog,
|
||||
after: Stage::Coding { claim: None },
|
||||
event: PipelineEvent::DepsMet,
|
||||
at: Utc::now(),
|
||||
}
|
||||
}
|
||||
|
||||
// ── AC1: subscriber writes the register on terminal transitions ──────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn done_transition_writes_rollup() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let story_id = "1017_sub_done";
|
||||
|
||||
write_records(
|
||||
tmp.path(),
|
||||
&[
|
||||
make_record(story_id, "coder-1", 1.50),
|
||||
make_record(story_id, "qa-1", 0.50),
|
||||
],
|
||||
);
|
||||
|
||||
let fired = fired_done(story_id);
|
||||
on_terminal_transition(tmp.path(), &fired).await;
|
||||
|
||||
let rollup = crate::service::agents::cost_rollup::get_rollup(tmp.path(), story_id)
|
||||
.expect("rollup must exist after Done transition");
|
||||
assert!(
|
||||
(rollup.total_cost_usd - 2.0).abs() < f64::EPSILON,
|
||||
"total must be sum of all agents: {}",
|
||||
rollup.total_cost_usd
|
||||
);
|
||||
assert_eq!(rollup.agents.len(), 2, "must have 2 agent entries");
|
||||
}
|
||||
|
||||
// ── AC4: multiple agents produce correct breakdown ────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn multiple_agents_correct_breakdown() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let story_id = "1017_sub_multi";
|
||||
|
||||
write_records(
|
||||
tmp.path(),
|
||||
&[
|
||||
make_record(story_id, "coder-1", 1.00),
|
||||
make_record(story_id, "coder-1", 0.50), // second run
|
||||
make_record(story_id, "qa-1", 0.75),
|
||||
make_record(story_id, "mergemaster", 0.25),
|
||||
// Different story — must not leak into rollup
|
||||
make_record("9999_other_story", "coder-1", 99.0),
|
||||
],
|
||||
);
|
||||
|
||||
on_terminal_transition(tmp.path(), &fired_done(story_id)).await;
|
||||
|
||||
let rollup = crate::service::agents::cost_rollup::get_rollup(tmp.path(), story_id)
|
||||
.expect("rollup must exist");
|
||||
|
||||
// Total: 1.00 + 0.50 + 0.75 + 0.25 = 2.50
|
||||
assert!(
|
||||
(rollup.total_cost_usd - 2.50).abs() < f64::EPSILON,
|
||||
"total must be 2.50, got {}",
|
||||
rollup.total_cost_usd
|
||||
);
|
||||
|
||||
// coder-1 should be aggregated: 1.00 + 0.50 = 1.50
|
||||
let coder = rollup
|
||||
.agents
|
||||
.iter()
|
||||
.find(|a| a.agent_name == "coder-1")
|
||||
.expect("coder-1 must be in agents");
|
||||
assert!(
|
||||
(coder.total_cost_usd - 1.50).abs() < f64::EPSILON,
|
||||
"coder-1 cost must be 1.50, got {}",
|
||||
coder.total_cost_usd
|
||||
);
|
||||
|
||||
// Other story must NOT appear in this rollup
|
||||
assert!(
|
||||
rollup.agents.iter().all(|a| a.agent_name != "9999_other"),
|
||||
"other story's agent must not leak into rollup"
|
||||
);
|
||||
}
|
||||
|
||||
// ── Non-terminal transitions are ignored ──────────────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn non_terminal_transition_does_not_write_rollup() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let story_id = "1017_sub_nonterminal";
|
||||
|
||||
write_records(tmp.path(), &[make_record(story_id, "coder-1", 5.0)]);
|
||||
|
||||
// Coding transition — not terminal
|
||||
on_terminal_transition(tmp.path(), &fired_coding(story_id)).await;
|
||||
|
||||
assert!(
|
||||
crate::service::agents::cost_rollup::get_rollup(tmp.path(), story_id).is_none(),
|
||||
"non-terminal transition must not write rollup"
|
||||
);
|
||||
}
|
||||
|
||||
// ── Abandoned transition is also terminal ─────────────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn abandoned_transition_writes_rollup() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let story_id = "1017_sub_abandoned";
|
||||
|
||||
write_records(tmp.path(), &[make_record(story_id, "coder-1", 0.30)]);
|
||||
on_terminal_transition(tmp.path(), &fired_abandoned(story_id)).await;
|
||||
|
||||
let rollup = crate::service::agents::cost_rollup::get_rollup(tmp.path(), story_id)
|
||||
.expect("rollup must exist after Abandoned transition");
|
||||
assert!(
|
||||
(rollup.total_cost_usd - 0.30).abs() < f64::EPSILON,
|
||||
"abandoned story cost must be 0.30, got {}",
|
||||
rollup.total_cost_usd
|
||||
);
|
||||
}
|
||||
|
||||
// ── Story with no records gets a zero rollup ──────────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn zero_cost_story_gets_empty_rollup() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let story_id = "1017_sub_zero";
|
||||
|
||||
// No records written for this story
|
||||
on_terminal_transition(tmp.path(), &fired_done(story_id)).await;
|
||||
|
||||
let rollup = crate::service::agents::cost_rollup::get_rollup(tmp.path(), story_id)
|
||||
.expect("rollup must be written even for zero-cost story");
|
||||
assert_eq!(rollup.total_cost_usd, 0.0);
|
||||
assert!(rollup.agents.is_empty());
|
||||
}
|
||||
|
||||
// ── Via broadcast subscriber ──────────────────────────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn broadcast_subscriber_reacts_to_terminal_transition() {
|
||||
crate::crdt_state::init_for_test();
|
||||
crate::db::ensure_content_store();
|
||||
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let story_id = "1017_sub_broadcast";
|
||||
|
||||
crate::db::write_item_with_content(
|
||||
story_id,
|
||||
"2_current",
|
||||
"---\nname: Test\n---\n",
|
||||
crate::db::ItemMeta::named("Test"),
|
||||
);
|
||||
|
||||
write_records(tmp.path(), &[make_record(story_id, "coder-1", 1.23)]);
|
||||
|
||||
spawn_cost_rollup_subscriber(tmp.path().to_path_buf());
|
||||
|
||||
crate::agents::lifecycle::abandon_story(story_id).expect("abandon must succeed");
|
||||
|
||||
// Give the subscriber task time to process.
|
||||
tokio::time::sleep(std::time::Duration::from_millis(200)).await;
|
||||
|
||||
let rollup = crate::service::agents::cost_rollup::get_rollup(tmp.path(), story_id)
|
||||
.expect("rollup must be written via broadcast subscriber");
|
||||
assert!(
|
||||
(rollup.total_cost_usd - 1.23).abs() < f64::EPSILON,
|
||||
"rollup cost must match records: {}",
|
||||
rollup.total_cost_usd
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,7 @@
|
||||
//! Agent pool — manages the set of active agents across all pipeline stages.
|
||||
pub(crate) mod auto_assign;
|
||||
/// TransitionFired subscriber that snapshots token cost on terminal stage transitions.
|
||||
pub(crate) mod cost_rollup_subscriber;
|
||||
mod pipeline;
|
||||
mod process;
|
||||
mod query;
|
||||
|
||||
Reference in New Issue
Block a user