use std::collections::BTreeMap;
use std::fmt;
use serde::{Deserialize, Serialize};
use crate::error::RStructorError;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct TokenUsage {
pub model: String,
pub input_tokens: u64,
pub cached_input_tokens: u64,
pub cache_write_input_tokens: u64,
pub output_tokens: u64,
}
impl TokenUsage {
pub fn new(model: impl Into<String>, input_tokens: u64, output_tokens: u64) -> Self {
Self {
model: model.into(),
input_tokens,
cached_input_tokens: 0,
cache_write_input_tokens: 0,
output_tokens,
}
}
#[must_use]
pub fn with_cache_tokens(
mut self,
cached_input_tokens: u64,
cache_write_input_tokens: u64,
) -> Self {
self.cached_input_tokens = cached_input_tokens;
self.cache_write_input_tokens = cache_write_input_tokens;
self
}
pub fn total_tokens(&self) -> u64 {
self.input_tokens + self.output_tokens
}
}
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
#[non_exhaustive]
pub struct RunUsage {
pub reported_attempts: usize,
pub input_tokens: u64,
pub cached_input_tokens: u64,
pub cache_write_input_tokens: u64,
pub output_tokens: u64,
pub by_model: BTreeMap<String, TokenUsage>,
pub overflowed: bool,
}
impl RunUsage {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn from_response(usage: TokenUsage) -> Self {
let mut total = Self::new();
total.record(usage);
total
}
pub fn record(&mut self, usage: TokenUsage) {
self.reported_attempts = match self.reported_attempts.checked_add(1) {
Some(attempts) => attempts,
None => {
self.overflowed = true;
usize::MAX
}
};
self.input_tokens =
saturating_add(&mut self.overflowed, self.input_tokens, usage.input_tokens);
self.cached_input_tokens = saturating_add(
&mut self.overflowed,
self.cached_input_tokens,
usage.cached_input_tokens,
);
self.cache_write_input_tokens = saturating_add(
&mut self.overflowed,
self.cache_write_input_tokens,
usage.cache_write_input_tokens,
);
self.output_tokens = saturating_add(
&mut self.overflowed,
self.output_tokens,
usage.output_tokens,
);
let model_usage = self
.by_model
.entry(usage.model.clone())
.or_insert_with(|| TokenUsage::new(usage.model, 0, 0));
model_usage.input_tokens = saturating_add(
&mut self.overflowed,
model_usage.input_tokens,
usage.input_tokens,
);
model_usage.cached_input_tokens = saturating_add(
&mut self.overflowed,
model_usage.cached_input_tokens,
usage.cached_input_tokens,
);
model_usage.cache_write_input_tokens = saturating_add(
&mut self.overflowed,
model_usage.cache_write_input_tokens,
usage.cache_write_input_tokens,
);
model_usage.output_tokens = saturating_add(
&mut self.overflowed,
model_usage.output_tokens,
usage.output_tokens,
);
if self.input_tokens.checked_add(self.output_tokens).is_none()
|| model_usage
.input_tokens
.checked_add(model_usage.output_tokens)
.is_none()
{
self.overflowed = true;
}
}
#[must_use]
pub fn total_tokens(&self) -> u64 {
self.input_tokens.saturating_add(self.output_tokens)
}
}
fn saturating_add(overflowed: &mut bool, left: u64, right: u64) -> u64 {
match left.checked_add(right) {
Some(total) => total,
None => {
*overflowed = true;
u64::MAX
}
}
}
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum AttemptKind {
Semantic,
Transport,
}
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum RetryDisposition {
Retried,
BudgetExhausted,
NonRetryable,
}
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum AttemptOutcome {
Succeeded,
Failed {
message: String,
disposition: RetryDisposition,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct AttemptRecord {
pub number: usize,
pub kind: AttemptKind,
pub outcome: AttemptOutcome,
pub usage: Option<TokenUsage>,
pub response: Option<crate::ResponseMetadata>,
}
impl AttemptRecord {
#[cfg(any(test, feature = "mock"))]
pub(crate) fn succeeded(number: usize, usage: Option<TokenUsage>) -> Self {
Self::succeeded_with_response(number, usage, None)
}
#[cfg(any(test, feature = "_client", feature = "mock"))]
pub(crate) fn succeeded_with_response(
number: usize,
usage: Option<TokenUsage>,
response: Option<crate::ResponseMetadata>,
) -> Self {
Self {
number,
kind: AttemptKind::Semantic,
outcome: AttemptOutcome::Succeeded,
usage,
response,
}
}
#[cfg(any(test, feature = "mock"))]
pub(crate) fn failed(
number: usize,
kind: AttemptKind,
error: &RStructorError,
disposition: RetryDisposition,
usage: Option<TokenUsage>,
) -> Self {
Self::failed_with_response(number, kind, error, disposition, usage, None)
}
#[cfg(any(test, feature = "_client", feature = "mock"))]
pub(crate) fn failed_with_response(
number: usize,
kind: AttemptKind,
error: &RStructorError,
disposition: RetryDisposition,
usage: Option<TokenUsage>,
response: Option<crate::ResponseMetadata>,
) -> Self {
Self {
number,
kind,
outcome: AttemptOutcome::Failed {
message: error.to_string(),
disposition,
},
usage,
response,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct ExtractionReport {
pub final_usage: Option<TokenUsage>,
pub cumulative_usage: Option<RunUsage>,
pub attempts: Vec<AttemptRecord>,
pub attempts_complete: bool,
}
impl ExtractionReport {
fn from_success<T>(report: MaterializeReport<T>) -> (T, Self) {
let MaterializeReport {
data,
final_usage,
cumulative_usage,
attempts,
attempts_complete,
} = report;
(
data,
Self {
final_usage,
cumulative_usage,
attempts,
attempts_complete,
},
)
}
fn from_failure(failure: MaterializeFailure) -> (Box<RStructorError>, Self) {
let MaterializeFailure {
error,
cumulative_usage,
attempts,
attempts_complete,
} = failure;
let final_usage = attempts.last().and_then(|attempt| attempt.usage.clone());
(
error,
Self {
final_usage,
cumulative_usage,
attempts,
attempts_complete,
},
)
}
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct Extraction<T> {
pub data: T,
pub report: ExtractionReport,
}
impl<T> Extraction<T> {
pub fn map<U, F: FnOnce(T) -> U>(self, f: F) -> Extraction<U> {
Extraction {
data: f(self.data),
report: self.report,
}
}
#[must_use]
pub fn into_data(self) -> T {
self.data
}
}
impl<T> From<MaterializeReport<T>> for Extraction<T> {
fn from(report: MaterializeReport<T>) -> Self {
let (data, report) = ExtractionReport::from_success(report);
Self { data, report }
}
}
#[derive(Debug)]
pub struct ExtractionError {
error: Box<RStructorError>,
pub report: ExtractionReport,
}
impl ExtractionError {
#[must_use]
pub fn error(&self) -> &RStructorError {
&self.error
}
#[must_use]
pub fn into_error(self) -> RStructorError {
*self.error
}
}
impl From<MaterializeFailure> for ExtractionError {
fn from(failure: MaterializeFailure) -> Self {
let (error, report) = ExtractionReport::from_failure(failure);
Self { error, report }
}
}
impl fmt::Display for ExtractionError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
self.error().fmt(formatter)
}
}
impl std::error::Error for ExtractionError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
Some(self.error())
}
}
pub type ExtractionResult<T> = std::result::Result<Extraction<T>, ExtractionError>;
#[non_exhaustive]
#[derive(Debug, Clone)]
pub struct MaterializeReport<T> {
pub data: T,
pub final_usage: Option<TokenUsage>,
pub cumulative_usage: Option<RunUsage>,
pub attempts: Vec<AttemptRecord>,
pub attempts_complete: bool,
}
impl<T> MaterializeReport<T> {
#[cfg(any(test, feature = "_client", feature = "mock"))]
pub(crate) fn new(
data: T,
final_usage: Option<TokenUsage>,
cumulative_usage: Option<RunUsage>,
attempts: Vec<AttemptRecord>,
) -> Self {
Self {
data,
final_usage,
cumulative_usage,
attempts,
attempts_complete: true,
}
}
#[cfg(feature = "mock")]
pub(crate) fn from_fixture_parts(
data: T,
final_usage: Option<TokenUsage>,
cumulative_usage: Option<RunUsage>,
attempts: Vec<AttemptRecord>,
attempts_complete: bool,
) -> Self {
Self {
data,
final_usage,
cumulative_usage,
attempts,
attempts_complete,
}
}
#[must_use]
pub fn from_result(result: MaterializeResult<T>) -> Self {
Self {
data: result.data,
final_usage: result.usage,
cumulative_usage: None,
attempts: Vec::new(),
attempts_complete: false,
}
}
pub fn map<U, F: FnOnce(T) -> U>(self, f: F) -> MaterializeReport<U> {
MaterializeReport {
data: f(self.data),
final_usage: self.final_usage,
cumulative_usage: self.cumulative_usage,
attempts: self.attempts,
attempts_complete: self.attempts_complete,
}
}
#[must_use]
pub fn into_result(self) -> MaterializeResult<T> {
MaterializeResult::new(self.data, self.final_usage)
}
}
#[non_exhaustive]
#[derive(Debug)]
pub struct MaterializeFailure {
error: Box<RStructorError>,
pub cumulative_usage: Option<RunUsage>,
pub attempts: Vec<AttemptRecord>,
pub attempts_complete: bool,
}
impl MaterializeFailure {
#[cfg(any(test, feature = "_client", feature = "mock"))]
pub(crate) fn new(
error: RStructorError,
cumulative_usage: Option<RunUsage>,
attempts: Vec<AttemptRecord>,
) -> Self {
Self {
error: Box::new(error),
cumulative_usage,
attempts,
attempts_complete: true,
}
}
#[cfg(feature = "mock")]
pub(crate) fn from_fixture_parts(
error: RStructorError,
cumulative_usage: Option<RunUsage>,
attempts: Vec<AttemptRecord>,
attempts_complete: bool,
) -> Self {
Self {
error: Box::new(error),
cumulative_usage,
attempts,
attempts_complete,
}
}
#[must_use]
pub fn from_error(error: RStructorError) -> Self {
Self {
error: Box::new(error),
cumulative_usage: None,
attempts: Vec::new(),
attempts_complete: false,
}
}
#[must_use]
pub fn error(&self) -> &RStructorError {
&self.error
}
#[must_use]
pub fn into_error(self) -> RStructorError {
*self.error
}
}
impl fmt::Display for MaterializeFailure {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
self.error().fmt(formatter)
}
}
impl std::error::Error for MaterializeFailure {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
Some(self.error())
}
}
#[derive(Debug, Clone)]
pub struct MaterializeResult<T> {
pub data: T,
pub usage: Option<TokenUsage>,
}
impl<T> MaterializeResult<T> {
pub fn new(data: T, usage: Option<TokenUsage>) -> Self {
Self { data, usage }
}
pub fn from_data(data: T) -> Self {
Self { data, usage: None }
}
pub fn map<U, F: FnOnce(T) -> U>(self, f: F) -> MaterializeResult<U> {
MaterializeResult {
data: f(self.data),
usage: self.usage,
}
}
}
#[derive(Debug, Clone)]
pub struct GenerateResult {
pub text: String,
pub usage: Option<TokenUsage>,
}
impl GenerateResult {
pub fn new(text: String, usage: Option<TokenUsage>) -> Self {
Self { text, usage }
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn run_usage_groups_exact_provider_model_versions() {
let mut usage = RunUsage::new();
usage.record(TokenUsage::new("gpt-5.6-2026-07-01", 120, 30).with_cache_tokens(80, 0));
usage.record(TokenUsage::new("gpt-5.6-2026-07-01", 180, 45).with_cache_tokens(100, 20));
usage.record(TokenUsage::new("gpt-5.6-2026-07-15", 200, 50).with_cache_tokens(0, 150));
assert_eq!(usage.reported_attempts, 3);
assert_eq!(usage.input_tokens, 500);
assert_eq!(usage.cached_input_tokens, 180);
assert_eq!(usage.cache_write_input_tokens, 170);
assert_eq!(usage.output_tokens, 125);
assert_eq!(usage.total_tokens(), 625);
assert!(!usage.overflowed);
assert_eq!(
usage.by_model["gpt-5.6-2026-07-01"],
TokenUsage::new("gpt-5.6-2026-07-01", 300, 75).with_cache_tokens(180, 20)
);
assert_eq!(
usage.by_model["gpt-5.6-2026-07-15"],
TokenUsage::new("gpt-5.6-2026-07-15", 200, 50).with_cache_tokens(0, 150)
);
}
#[test]
fn run_usage_saturates_and_flags_untrusted_counter_overflow() {
let mut usage = RunUsage::new();
usage.record(
TokenUsage::new("hostile-compatible-endpoint", u64::MAX, 1)
.with_cache_tokens(u64::MAX, u64::MAX),
);
usage.record(
TokenUsage::new("hostile-compatible-endpoint", 1, u64::MAX).with_cache_tokens(1, 1),
);
assert!(usage.overflowed);
assert_eq!(usage.reported_attempts, 2);
assert_eq!(usage.input_tokens, u64::MAX);
assert_eq!(usage.cached_input_tokens, u64::MAX);
assert_eq!(usage.cache_write_input_tokens, u64::MAX);
assert_eq!(usage.output_tokens, u64::MAX);
assert_eq!(usage.total_tokens(), u64::MAX);
assert_eq!(
usage.by_model["hostile-compatible-endpoint"].input_tokens,
u64::MAX
);
assert_eq!(
usage.by_model["hostile-compatible-endpoint"].output_tokens,
u64::MAX
);
assert_eq!(
usage.by_model["hostile-compatible-endpoint"].cached_input_tokens,
u64::MAX
);
assert_eq!(
usage.by_model["hostile-compatible-endpoint"].cache_write_input_tokens,
u64::MAX
);
}
#[test]
fn custom_client_result_preserves_final_usage_without_inventing_attempts() {
let final_usage = TokenUsage::new("mock-risk-model", 42, 11);
let report =
MaterializeReport::from_result(MaterializeResult::new("portfolio", Some(final_usage)));
assert_eq!(report.data, "portfolio");
assert_eq!(report.final_usage.as_ref().unwrap().total_tokens(), 53);
assert!(report.cumulative_usage.is_none());
assert!(report.attempts.is_empty());
assert!(!report.attempts_complete);
}
#[test]
fn unknown_custom_client_failure_does_not_invent_a_provider_attempt() {
let failure =
MaterializeFailure::from_error(RStructorError::SchemaError("bad schema".into()));
assert!(failure.attempts.is_empty());
assert!(failure.cumulative_usage.is_none());
assert!(!failure.attempts_complete);
assert!(matches!(failure.error(), RStructorError::SchemaError(_)));
}
#[test]
fn extraction_success_uses_the_shared_report_and_maps_data() {
let final_usage = TokenUsage::new("risk-model-v2", 120, 18);
let cumulative_usage = RunUsage::from_response(final_usage.clone());
let legacy = MaterializeReport::new(
"HF-ALPHA-001",
Some(final_usage.clone()),
Some(cumulative_usage.clone()),
vec![AttemptRecord::succeeded(1, Some(final_usage.clone()))],
);
let extraction = Extraction::from(legacy).map(str::len);
assert_eq!(extraction.data, 12);
assert_eq!(extraction.report.final_usage, Some(final_usage));
assert_eq!(extraction.report.cumulative_usage, Some(cumulative_usage));
assert_eq!(extraction.report.attempts.len(), 1);
assert!(extraction.report.attempts_complete);
}
#[test]
fn extraction_failure_uses_the_same_report_shape_and_final_attempt_usage() {
let first_usage = TokenUsage::new("risk-model-v1", 90, 15);
let final_usage = TokenUsage::new("risk-model-v2", 100, 12);
let mut cumulative_usage = RunUsage::new();
cumulative_usage.record(first_usage.clone());
cumulative_usage.record(final_usage.clone());
let final_error = RStructorError::OutputDecodeError {
path: "$.positions[1].quantity".into(),
message: "invalid type: string, expected i64".into(),
};
let failure = MaterializeFailure::new(
final_error,
Some(cumulative_usage.clone()),
vec![
AttemptRecord::failed(
1,
AttemptKind::Semantic,
&RStructorError::ValidationError("invalid quantity".into()),
RetryDisposition::Retried,
Some(first_usage),
),
AttemptRecord::failed(
2,
AttemptKind::Semantic,
&RStructorError::ValidationError("invalid quantity".into()),
RetryDisposition::BudgetExhausted,
Some(final_usage.clone()),
),
],
);
let error = ExtractionError::from(failure);
assert!(matches!(
error.error(),
RStructorError::OutputDecodeError { path, .. }
if path == "$.positions[1].quantity"
));
assert_eq!(error.report.final_usage, Some(final_usage));
assert_eq!(error.report.cumulative_usage, Some(cumulative_usage));
assert_eq!(error.report.attempts.len(), 2);
assert!(error.report.attempts_complete);
}
}