use std::fs;
use std::path::Path;
use super::{
BudgetSpend, RESPONSE_FILE, STEP_SEQ_WIDTH, STEPS_DIR, context_window, last_usage,
spend_from_bytes,
};
const REQUEST_FILE: &str = "request.json";
const META_FILE: &str = "meta.json";
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Scope {
Agent(String),
Tree(String),
Workspace,
}
impl Scope {
pub fn wants(&self, conv_id: &str) -> bool {
match self {
Self::Workspace => true,
Self::Agent(id) => conv_id == id,
Self::Tree(root) => {
conv_id == root
|| (conv_id.starts_with(root.as_str())
&& conv_id.as_bytes().get(root.len()) == Some(&b'-'))
}
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StepBill {
pub conv: String,
pub seq: String,
pub model: Option<String>,
pub spend: BudgetSpend,
pub last_usage: BudgetSpend,
pub window: Option<u64>,
pub wall_secs: u64,
}
pub fn bills(workspace: &Path, scope: &Scope) -> Vec<StepBill> {
let Ok(entries) = fs::read_dir(workspace.join(STEPS_DIR)) else {
return Vec::new();
};
let mut out = Vec::new();
for entry in entries.flatten() {
let raw = entry.file_name();
let conv = raw.to_string_lossy();
if scope.wants(&conv) {
conv_bills(&entry.path(), &conv, &mut out);
}
}
out
}
pub fn total(bills: &[StepBill]) -> BudgetSpend {
let mut total = BudgetSpend::default();
for bill in bills {
total.add(bill.spend);
}
total
}
pub fn wall(bills: &[StepBill]) -> u64 {
bills.iter().map(|b| b.wall_secs).sum()
}
fn conv_bills(conv_dir: &Path, conv: &str, out: &mut Vec<StepBill>) {
let Ok(entries) = fs::read_dir(conv_dir) else {
return;
};
for entry in entries.flatten() {
let raw = entry.file_name();
let name = raw.to_string_lossy();
if name.len() == STEP_SEQ_WIDTH && name.bytes().all(|b| b.is_ascii_digit()) {
out.push(step_bill(&entry.path(), conv, &name));
}
}
}
fn step_bill(step_dir: &Path, conv: &str, seq: &str) -> StepBill {
let bytes = fs::read(step_dir.join(RESPONSE_FILE)).unwrap_or_default();
StepBill {
conv: conv.to_owned(),
seq: seq.to_owned(),
model: step_model(step_dir),
spend: spend_from_bytes(&bytes),
last_usage: last_usage(&bytes),
window: context_window(&bytes),
wall_secs: step_wall(step_dir),
}
}
fn step_wall(step_dir: &Path) -> u64 {
let bytes = fs::read(step_dir.join(META_FILE)).unwrap_or_default();
let Ok(meta) = serde_json::from_slice::<serde_json::Value>(&bytes) else {
return 0;
};
let at = |key: &str| {
crate::ui_state::epoch_from_iso8601(meta.get(key).and_then(serde_json::Value::as_str)?)
};
let (Some(start), Some(end)) = (at("started_at"), at("ended_at")) else {
return 0;
};
u64::try_from(end - start).unwrap_or(0)
}
fn step_model(step_dir: &Path) -> Option<String> {
let bytes = fs::read(step_dir.join(REQUEST_FILE)).ok()?;
let value: serde_json::Value = serde_json::from_slice(&bytes).ok()?;
Some(value.get("model")?.as_str()?.to_owned())
}
#[cfg(test)]
mod tests;