use std::fs;
use std::path::PathBuf;
use anyhow::{Context, Result};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use crate::config::data_dir;
const ETU_WEIGHT_INPUT: f64 = 1.0;
const ETU_WEIGHT_CACHE_READ: f64 = 0.1;
const ETU_WEIGHT_CACHE_WRITE: f64 = 1.25;
const ETU_WEIGHT_OUTPUT: f64 = 5.0;
#[allow(clippy::cast_precision_loss)]
pub(crate) fn compute_etu(input: u64, cache_read: u64, cache_write: u64, output: u64) -> f64 {
ETU_WEIGHT_INPUT * input as f64
+ ETU_WEIGHT_CACHE_READ * cache_read as f64
+ ETU_WEIGHT_CACHE_WRITE * cache_write as f64
+ ETU_WEIGHT_OUTPUT * output as f64
}
#[derive(Debug, Clone, Default, Serialize, PartialEq)]
pub struct AiStats {
pub provider: String,
pub model: String,
pub input_tokens: u64,
pub output_tokens: u64,
pub duration_ms: u64,
#[serde(default)]
pub cost_usd: Option<f64>,
#[serde(default)]
pub fallback_provider: Option<String>,
#[serde(default)]
pub prompt_chars: usize,
#[serde(default)]
pub cache_read_tokens: u64,
#[serde(default)]
pub cache_write_tokens: u64,
#[serde(default)]
pub effective_token_units: f64,
#[serde(default)]
#[serde(skip_serializing_if = "Option::is_none")]
pub trace_id: Option<String>,
}
impl AiStats {
#[must_use]
pub fn with_computed_etu(mut self) -> Self {
self.effective_token_units = compute_etu(
self.input_tokens,
self.cache_read_tokens,
self.cache_write_tokens,
self.output_tokens,
);
self
}
}
impl<'de> Deserialize<'de> for AiStats {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
#[derive(Deserialize)]
struct Helper {
#[serde(default)]
provider: String,
#[serde(default)]
model: String,
#[serde(default)]
input_tokens: u64,
#[serde(default)]
output_tokens: u64,
#[serde(default)]
duration_ms: u64,
#[serde(default)]
cost_usd: Option<f64>,
#[serde(default)]
fallback_provider: Option<String>,
#[serde(default)]
prompt_chars: usize,
#[serde(default)]
cache_read_tokens: u64,
#[serde(default)]
cache_write_tokens: u64,
#[serde(default)]
#[allow(dead_code)]
effective_token_units: f64,
#[serde(default)]
trace_id: Option<String>,
}
let h = Helper::deserialize(deserializer)?;
Ok(AiStats {
provider: h.provider,
model: h.model,
input_tokens: h.input_tokens,
output_tokens: h.output_tokens,
duration_ms: h.duration_ms,
cost_usd: h.cost_usd,
fallback_provider: h.fallback_provider,
prompt_chars: h.prompt_chars,
cache_read_tokens: h.cache_read_tokens,
cache_write_tokens: h.cache_write_tokens,
effective_token_units: 0.0,
trace_id: h.trace_id,
}
.with_computed_etu())
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum ContributionStatus {
#[default]
Pending,
Accepted,
Rejected,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Contribution {
pub id: Uuid,
pub repo: String,
pub issue: u64,
pub action: String,
pub timestamp: DateTime<Utc>,
pub comment_url: String,
#[serde(default)]
pub status: ContributionStatus,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub ai_stats: Option<AiStats>,
}
#[derive(Debug, Default, Serialize, Deserialize)]
pub struct HistoryData {
pub contributions: Vec<Contribution>,
}
impl HistoryData {
#[must_use]
pub fn total_tokens(&self) -> u64 {
self.contributions
.iter()
.filter_map(|c| c.ai_stats.as_ref())
.map(|stats| stats.input_tokens + stats.output_tokens)
.sum()
}
#[must_use]
pub fn total_cost(&self) -> f64 {
self.contributions
.iter()
.filter_map(|c| c.ai_stats.as_ref())
.filter_map(|stats| stats.cost_usd)
.sum()
}
#[must_use]
#[allow(clippy::cast_precision_loss)]
pub fn avg_tokens_per_triage(&self) -> f64 {
let contributions_with_stats: Vec<_> = self
.contributions
.iter()
.filter_map(|c| c.ai_stats.as_ref())
.collect();
if contributions_with_stats.is_empty() {
return 0.0;
}
let total: u64 = contributions_with_stats
.iter()
.map(|stats| stats.input_tokens + stats.output_tokens)
.sum();
total as f64 / contributions_with_stats.len() as f64
}
#[must_use]
pub fn cost_by_model(&self) -> std::collections::HashMap<String, f64> {
let mut costs = std::collections::HashMap::new();
for contribution in &self.contributions {
if let Some(stats) = &contribution.ai_stats
&& let Some(cost) = stats.cost_usd
{
*costs.entry(stats.model.clone()).or_insert(0.0) += cost;
}
}
costs
}
}
#[must_use]
pub fn history_file_path() -> PathBuf {
data_dir().join("history.json")
}
pub fn load() -> Result<HistoryData> {
let path = history_file_path();
if !path.exists() {
return Ok(HistoryData::default());
}
let contents = fs::read_to_string(&path)
.with_context(|| format!("Failed to read history file: {}", path.display()))?;
let data: HistoryData = serde_json::from_str(&contents)
.with_context(|| format!("Failed to parse history file: {}", path.display()))?;
Ok(data)
}
pub fn save(data: &HistoryData) -> Result<()> {
let path = history_file_path();
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)
.with_context(|| format!("Failed to create directory: {}", parent.display()))?;
}
let contents =
serde_json::to_string_pretty(data).context("Failed to serialize history data")?;
fs::write(&path, contents)
.with_context(|| format!("Failed to write history file: {}", path.display()))?;
Ok(())
}
pub fn add_contribution(contribution: Contribution) -> Result<()> {
let mut data = load()?;
data.contributions.push(contribution);
save(&data)?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
fn test_contribution() -> Contribution {
Contribution {
id: Uuid::new_v4(),
repo: "owner/repo".to_string(),
issue: 123,
action: "triage".to_string(),
timestamp: Utc::now(),
comment_url: "https://github.com/owner/repo/issues/123#issuecomment-1".to_string(),
status: ContributionStatus::Pending,
ai_stats: None,
}
}
#[test]
fn test_contribution_serialization_roundtrip() {
let contribution = test_contribution();
let json = serde_json::to_string(&contribution).expect("serialize");
let parsed: Contribution = serde_json::from_str(&json).expect("deserialize");
assert_eq!(contribution.id, parsed.id);
assert_eq!(contribution.repo, parsed.repo);
assert_eq!(contribution.issue, parsed.issue);
assert_eq!(contribution.action, parsed.action);
assert_eq!(contribution.comment_url, parsed.comment_url);
assert_eq!(contribution.status, parsed.status);
}
#[test]
fn test_history_data_serialization_roundtrip() {
let data = HistoryData {
contributions: vec![test_contribution(), test_contribution()],
};
let json = serde_json::to_string_pretty(&data).expect("serialize");
let parsed: HistoryData = serde_json::from_str(&json).expect("deserialize");
assert_eq!(parsed.contributions.len(), 2);
}
#[test]
fn test_contribution_status_default() {
let status = ContributionStatus::default();
assert_eq!(status, ContributionStatus::Pending);
}
#[test]
fn test_contribution_status_serialization() {
assert_eq!(
serde_json::to_string(&ContributionStatus::Pending).unwrap(),
"\"pending\""
);
assert_eq!(
serde_json::to_string(&ContributionStatus::Accepted).unwrap(),
"\"accepted\""
);
assert_eq!(
serde_json::to_string(&ContributionStatus::Rejected).unwrap(),
"\"rejected\""
);
}
#[test]
fn test_empty_history_default() {
let data = HistoryData::default();
assert!(data.contributions.is_empty());
}
#[test]
fn test_ai_stats_serialization_roundtrip() {
let stats = AiStats {
provider: "openrouter".to_string(),
model: "mistralai/mistral-small-2603".to_string(),
input_tokens: 1000,
output_tokens: 500,
duration_ms: 1500,
cost_usd: Some(0.0),
fallback_provider: None,
prompt_chars: 0,
cache_read_tokens: 0,
cache_write_tokens: 0,
effective_token_units: 0.0,
trace_id: None,
};
let json = serde_json::to_string(&stats).expect("serialize");
let parsed: AiStats = serde_json::from_str(&json).expect("deserialize");
assert_eq!(stats.provider, parsed.provider);
assert_eq!(stats.model, parsed.model);
assert_eq!(stats.input_tokens, parsed.input_tokens);
assert_eq!(stats.output_tokens, parsed.output_tokens);
assert_eq!(stats.duration_ms, parsed.duration_ms);
assert_eq!(stats.cost_usd, parsed.cost_usd);
assert_eq!(stats.fallback_provider, parsed.fallback_provider);
assert_eq!(stats.prompt_chars, parsed.prompt_chars);
assert_eq!(stats.cache_read_tokens, parsed.cache_read_tokens);
assert_eq!(stats.cache_write_tokens, parsed.cache_write_tokens);
assert_eq!(stats.trace_id, parsed.trace_id);
assert!((parsed.effective_token_units - 3500.0).abs() < f64::EPSILON);
}
#[test]
fn test_contribution_with_ai_stats() {
let mut contribution = test_contribution();
contribution.ai_stats = Some(AiStats {
provider: "openrouter".to_string(),
model: "mistralai/mistral-small-2603".to_string(),
input_tokens: 1000,
output_tokens: 500,
duration_ms: 1500,
cost_usd: Some(0.0),
fallback_provider: None,
prompt_chars: 0,
cache_read_tokens: 0,
cache_write_tokens: 0,
effective_token_units: 0.0,
trace_id: None,
});
let json = serde_json::to_string(&contribution).expect("serialize");
let parsed: Contribution = serde_json::from_str(&json).expect("deserialize");
assert!(parsed.ai_stats.is_some());
assert_eq!(
parsed.ai_stats.unwrap().model,
"mistralai/mistral-small-2603"
);
}
#[test]
fn test_contribution_without_ai_stats_backward_compat() {
let json = r#"{
"id": "550e8400-e29b-41d4-a716-446655440000",
"repo": "owner/repo",
"issue": 123,
"action": "triage",
"timestamp": "2024-01-01T00:00:00Z",
"comment_url": "https://github.com/owner/repo/issues/123#issuecomment-1",
"status": "pending"
}"#;
let parsed: Contribution = serde_json::from_str(json).expect("deserialize");
assert!(parsed.ai_stats.is_none());
}
#[test]
fn test_total_tokens() {
let mut data = HistoryData::default();
let mut c1 = test_contribution();
c1.ai_stats = Some(AiStats {
provider: "openrouter".to_string(),
model: "model1".to_string(),
input_tokens: 100,
output_tokens: 50,
duration_ms: 1000,
cost_usd: Some(0.01),
fallback_provider: None,
prompt_chars: 0,
cache_read_tokens: 0,
cache_write_tokens: 0,
effective_token_units: 0.0,
trace_id: None,
});
let mut c2 = test_contribution();
c2.ai_stats = Some(AiStats {
provider: "openrouter".to_string(),
model: "model2".to_string(),
input_tokens: 200,
output_tokens: 100,
duration_ms: 2000,
cost_usd: Some(0.02),
fallback_provider: None,
prompt_chars: 0,
cache_read_tokens: 0,
cache_write_tokens: 0,
effective_token_units: 0.0,
trace_id: None,
});
data.contributions.push(c1);
data.contributions.push(c2);
data.contributions.push(test_contribution());
assert_eq!(data.total_tokens(), 450);
}
#[test]
fn test_total_cost() {
let mut data = HistoryData::default();
let mut c1 = test_contribution();
c1.ai_stats = Some(AiStats {
provider: "openrouter".to_string(),
model: "model1".to_string(),
input_tokens: 100,
output_tokens: 50,
duration_ms: 1000,
cost_usd: Some(0.01),
fallback_provider: None,
prompt_chars: 0,
cache_read_tokens: 0,
cache_write_tokens: 0,
effective_token_units: 0.0,
trace_id: None,
});
let mut c2 = test_contribution();
c2.ai_stats = Some(AiStats {
provider: "openrouter".to_string(),
model: "model2".to_string(),
input_tokens: 200,
output_tokens: 100,
duration_ms: 2000,
cost_usd: Some(0.02),
fallback_provider: None,
prompt_chars: 0,
cache_read_tokens: 0,
cache_write_tokens: 0,
effective_token_units: 0.0,
trace_id: None,
});
data.contributions.push(c1);
data.contributions.push(c2);
assert!((data.total_cost() - 0.03).abs() < f64::EPSILON);
}
#[test]
fn test_avg_tokens_per_triage() {
let mut data = HistoryData::default();
let mut c1 = test_contribution();
c1.ai_stats = Some(AiStats {
provider: "openrouter".to_string(),
model: "model1".to_string(),
input_tokens: 100,
output_tokens: 50,
duration_ms: 1000,
cost_usd: Some(0.01),
fallback_provider: None,
prompt_chars: 0,
cache_read_tokens: 0,
cache_write_tokens: 0,
effective_token_units: 0.0,
trace_id: None,
});
let mut c2 = test_contribution();
c2.ai_stats = Some(AiStats {
provider: "openrouter".to_string(),
model: "model2".to_string(),
input_tokens: 200,
output_tokens: 100,
duration_ms: 2000,
cost_usd: Some(0.02),
fallback_provider: None,
prompt_chars: 0,
cache_read_tokens: 0,
cache_write_tokens: 0,
effective_token_units: 0.0,
trace_id: None,
});
data.contributions.push(c1);
data.contributions.push(c2);
assert!((data.avg_tokens_per_triage() - 225.0).abs() < f64::EPSILON);
}
#[test]
fn test_avg_tokens_per_triage_empty() {
let data = HistoryData::default();
assert!((data.avg_tokens_per_triage() - 0.0).abs() < f64::EPSILON);
}
#[test]
fn test_cost_by_model() {
let mut data = HistoryData::default();
let mut c1 = test_contribution();
c1.ai_stats = Some(AiStats {
provider: "openrouter".to_string(),
model: "model1".to_string(),
input_tokens: 100,
output_tokens: 50,
duration_ms: 1000,
cost_usd: Some(0.01),
fallback_provider: None,
prompt_chars: 0,
cache_read_tokens: 0,
cache_write_tokens: 0,
effective_token_units: 0.0,
trace_id: None,
});
let mut c2 = test_contribution();
c2.ai_stats = Some(AiStats {
provider: "openrouter".to_string(),
model: "model1".to_string(),
input_tokens: 200,
output_tokens: 100,
duration_ms: 2000,
cost_usd: Some(0.02),
fallback_provider: None,
prompt_chars: 0,
cache_read_tokens: 0,
cache_write_tokens: 0,
effective_token_units: 0.0,
trace_id: None,
});
let mut c3 = test_contribution();
c3.ai_stats = Some(AiStats {
provider: "openrouter".to_string(),
model: "model2".to_string(),
input_tokens: 150,
output_tokens: 75,
duration_ms: 1500,
cost_usd: Some(0.015),
fallback_provider: None,
prompt_chars: 0,
cache_read_tokens: 0,
cache_write_tokens: 0,
effective_token_units: 0.0,
trace_id: None,
});
data.contributions.push(c1);
data.contributions.push(c2);
data.contributions.push(c3);
let costs = data.cost_by_model();
assert_eq!(costs.len(), 2);
assert!((costs.get("model1").unwrap() - 0.03).abs() < f64::EPSILON);
assert!((costs.get("model2").unwrap() - 0.015).abs() < f64::EPSILON);
}
#[test]
fn test_ai_stats_cache_tokens_roundtrip() {
let stats = AiStats {
provider: "anthropic".to_string(),
model: "claude-sonnet-4-6".to_string(),
input_tokens: 1000,
output_tokens: 500,
duration_ms: 1500,
cost_usd: Some(0.05),
fallback_provider: None,
prompt_chars: 5000,
cache_read_tokens: 100,
cache_write_tokens: 50,
effective_token_units: 0.0,
trace_id: None,
};
let json = serde_json::to_string(&stats).expect("serialize");
let parsed: AiStats = serde_json::from_str(&json).expect("deserialize");
assert_eq!(stats.provider, parsed.provider);
assert_eq!(stats.model, parsed.model);
assert_eq!(stats.input_tokens, parsed.input_tokens);
assert_eq!(stats.output_tokens, parsed.output_tokens);
assert_eq!(stats.duration_ms, parsed.duration_ms);
assert_eq!(stats.cost_usd, parsed.cost_usd);
assert_eq!(stats.fallback_provider, parsed.fallback_provider);
assert_eq!(stats.prompt_chars, parsed.prompt_chars);
assert_eq!(stats.cache_read_tokens, 100);
assert_eq!(stats.cache_write_tokens, 50);
assert_eq!(parsed.cache_read_tokens, 100);
assert_eq!(parsed.cache_write_tokens, 50);
assert!((parsed.effective_token_units - 3572.5).abs() < f64::EPSILON);
}
#[test]
fn test_ai_stats_cache_tokens_default() {
let json = r#"{
"provider": "openrouter",
"model": "mistralai/mistral-small-2603",
"input_tokens": 1000,
"output_tokens": 500,
"duration_ms": 1500,
"cost_usd": 0.0,
"fallback_provider": null,
"prompt_chars": 0
}"#;
let parsed: AiStats = serde_json::from_str(json).expect("deserialize");
assert_eq!(parsed.cache_read_tokens, 0);
assert_eq!(parsed.cache_write_tokens, 0);
}
#[test]
fn test_etu_formula() {
let stats = AiStats {
input_tokens: 1000,
output_tokens: 200,
cache_read_tokens: 500,
cache_write_tokens: 100,
..AiStats::default()
}
.with_computed_etu();
assert!((stats.effective_token_units - 2175.0).abs() < f64::EPSILON);
}
#[test]
fn test_etu_zero_on_default() {
let stats = AiStats::default().with_computed_etu();
assert_eq!(stats.effective_token_units, 0.0);
}
#[test]
fn test_etu_recomputed_on_deserialize() {
let json = r#"{
"provider": "anthropic",
"model": "claude-sonnet-4-6",
"input_tokens": 1000,
"output_tokens": 200,
"cache_read_tokens": 500,
"cache_write_tokens": 100,
"effective_token_units": 99999.0
}"#;
let stats: AiStats = serde_json::from_str(json).unwrap();
assert!((stats.effective_token_units - 2175.0).abs() < f64::EPSILON);
}
}