use crate::output::NormalizedUsageSnapshot;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
pub(crate) const USAGE_ACTIVITY_PREFIX: &str = "session-usage/";
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct SessionUsageRecord {
pub(crate) id: String,
pub(crate) usage: NormalizedUsageSnapshot,
pub(crate) request_sequence: u64,
pub(crate) final_usage: bool,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub(crate) struct SessionUsageLedger {
records: HashMap<String, SessionUsageRecord>,
#[serde(with = "wide_counter")]
input: u128,
#[serde(with = "wide_counter")]
output: u128,
#[serde(with = "wide_counter")]
cache_read: u128,
unknown_cache: usize,
recorded: bool,
pub(crate) incomplete: bool,
}
impl SessionUsageLedger {
pub(crate) fn from_checkpoint(value: serde_json::Value) -> anyhow::Result<Self> {
let ledger: Self = serde_json::from_value(value)?;
let mut active = Self::default();
for (id, record) in &ledger.records {
anyhow::ensure!(
id == &record.id,
"usage checkpoint record identity mismatch"
);
active.observe(record.clone());
}
anyhow::ensure!(
ledger.input >= active.input
&& ledger.output >= active.output
&& ledger.cache_read >= active.cache_read
&& ledger.unknown_cache >= active.unknown_cache,
"usage checkpoint totals are smaller than active records"
);
anyhow::ensure!(
ledger.recorded
|| (ledger.input == 0
&& ledger.output == 0
&& ledger.cache_read == 0
&& ledger.records.is_empty()),
"usage checkpoint activity marker is missing"
);
Ok(ledger)
}
pub(crate) fn observe(&mut self, record: SessionUsageRecord) {
if let Some(previous) = self.records.get(&record.id) {
if record.request_sequence < previous.request_sequence
|| (record.request_sequence == previous.request_sequence && previous.final_usage)
{
return;
}
if record.request_sequence > previous.request_sequence && !previous.final_usage {
self.incomplete = true;
}
self.input -= u128::from(previous.usage.effective_input);
self.output -= u128::from(previous.usage.output);
self.cache_read -= u128::from(previous.usage.cache_read);
self.unknown_cache -= usize::from(!previous.usage.cache_known);
}
self.input += u128::from(record.usage.effective_input);
self.output += u128::from(record.usage.output);
self.cache_read += u128::from(record.usage.cache_read);
self.unknown_cache += usize::from(!record.usage.cache_known);
self.recorded = true;
self.records.insert(record.id.clone(), record);
}
pub(crate) fn is_complete(&self) -> bool {
!self.incomplete && self.records.values().all(|record| record.final_usage)
}
pub(crate) fn fold_completed_run(&mut self) {
self.incomplete |= !self.is_complete();
self.records.clear();
}
pub(crate) fn merge(&mut self, other: &Self) {
self.incomplete |= other.incomplete;
for record in other.records.values() {
self.observe(record.clone());
}
}
pub(crate) fn pending_request(
&self,
id: String,
request_sequence: u64,
) -> Option<SessionUsageRecord> {
let previous = self.records.get(&id);
if previous.is_some_and(|record| record.request_sequence >= request_sequence) {
return None;
}
Some(SessionUsageRecord {
id,
usage: previous.map(|record| record.usage).unwrap_or_default(),
request_sequence,
final_usage: false,
})
}
pub(crate) fn has_recorded_activity(&self, activity_id: &str) -> bool {
let suffix = format!("/activity/{activity_id}");
self.records
.keys()
.any(|id| id.starts_with(USAGE_ACTIVITY_PREFIX) && id.ends_with(&suffix))
}
pub(crate) fn totals(&self) -> NormalizedUsageSnapshot {
NormalizedUsageSnapshot {
effective_input: self.input.min(u128::from(u64::MAX)) as u64,
output: self.output.min(u128::from(u64::MAX)) as u64,
cache_read: self.cache_read.min(u128::from(u64::MAX)) as u64,
cache_known: self.recorded && self.unknown_cache == 0,
}
}
}
pub(super) fn checkpoint_usage(
reader: impl std::io::Read,
prefix_bytes: u64,
) -> anyhow::Result<serde_json::Value> {
use std::io::BufRead;
anyhow::ensure!(
prefix_bytes <= super::read::MAX_METADATA_VISIT_BYTES as u64,
"session usage checkpoint exceeds byte limit"
);
let mut reader = std::io::BufReader::new(reader.take(prefix_bytes));
let mut ledger = SessionUsageLedger::default();
let mut primary_usage_since_output = false;
let mut active_run = None;
let mut line = Vec::new();
for index in 0..=super::read::MAX_METADATA_VISIT_LINES {
line.clear();
if reader.read_until(b'\n', &mut line)? == 0 {
break;
}
anyhow::ensure!(
index < super::read::MAX_METADATA_VISIT_LINES,
"session usage checkpoint exceeds line limit"
);
let Ok(event) = serde_json::from_slice::<super::SessionEvent>(&line) else {
ledger.incomplete = true;
continue;
};
match event.kind() {
Some(super::SessionEventKind::Compaction) => {
if let Some(value) = event.payload.get("session_usage") {
ledger = SessionUsageLedger::from_checkpoint(value.clone())?;
active_run = ledger.records.keys().find_map(|id| {
id.strip_prefix(USAGE_ACTIVITY_PREFIX)
.and_then(|id| id.split('/').next())
.map(str::to_owned)
});
primary_usage_since_output = event
.payload
.get("primary_usage_since_output")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false);
} else {
ledger.incomplete = true;
}
}
Some(super::SessionEventKind::SessionUsage) => {
match serde_json::from_value::<SessionUsageRecord>(event.payload) {
Ok(record) => {
let run = record
.id
.strip_prefix(USAGE_ACTIVITY_PREFIX)
.and_then(|id| id.split('/').next())
.map(str::to_owned);
if active_run.is_some() && active_run != run {
ledger.fold_completed_run();
}
active_run = run;
primary_usage_since_output |= record.id.contains("/primary/");
ledger.observe(record);
}
Err(_) => ledger.incomplete = true,
}
}
Some(super::SessionEventKind::AssistantOutput) => {
if !primary_usage_since_output {
let tokens = |key| {
event
.payload
.get("usage")
.and_then(|usage| usage.get(key))
.and_then(serde_json::Value::as_u64)
.unwrap_or_default()
};
ledger.incomplete = true;
ledger.observe(SessionUsageRecord {
id: format!("legacy/primary/{index}"),
usage: NormalizedUsageSnapshot {
effective_input: tokens("input"),
output: tokens("output"),
cache_read: tokens("cache_read"),
cache_known: false,
},
request_sequence: 0,
final_usage: true,
});
ledger.records.remove(&format!("legacy/primary/{index}"));
}
primary_usage_since_output = false;
}
Some(super::SessionEventKind::TurnStatus)
if event.turn_status_payload().is_some_and(|status| {
matches!(
status.status,
super::TurnStatus::Cancelled | super::TurnStatus::Failed
)
}) =>
{
ledger.incomplete |= !primary_usage_since_output;
primary_usage_since_output = false;
}
_ => {}
}
}
Ok(
serde_json::json!({"session_usage": ledger, "primary_usage_since_output": primary_usage_since_output}),
)
}
mod wide_counter {
use serde::{Deserialize, Deserializer, Serializer};
pub(super) fn serialize<S: Serializer>(value: &u128, serializer: S) -> Result<S::Ok, S::Error> {
serializer.serialize_str(&value.to_string())
}
pub(super) fn deserialize<'de, D: Deserializer<'de>>(
deserializer: D,
) -> Result<u128, D::Error> {
String::deserialize(deserializer)?
.parse()
.map_err(serde::de::Error::custom)
}
}