bcode_session 0.0.1-alpha.0

Session lifecycle and attachment management for bcode
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137

//! Incremental accounting over canonical request observations.

use bcode_session_models::{SessionCostEstimate, SessionTokenUsage, SessionUsageSummary};

/// Decide whether a request observation replaces its current contribution.
/// Identical and stale deliveries are harmless; conflicting duplicates are not.
pub fn accepts(
    current: Option<&SessionTokenUsage>,
    next: &SessionTokenUsage,
) -> Result<bool, String> {
    let Some(current) = current else {
        return Ok(true);
    };
    let mut current_facts = current.clone();
    let mut next_facts = next.clone();
    current_facts.cost = None;
    next_facts.cost = None;
    if next_facts == current_facts || next.observation_ordinal < current.observation_ordinal {
        return Ok(false);
    }
    if let (Some(current_request), Some(next_request)) = (&current.request, &next.request)
        && current_request != next_request
    {
        return Err("request attribution changed within one billing attempt".to_owned());
    }
    if current.terminal || next.observation_ordinal == current.observation_ordinal {
        return Err("conflicting usage observation for an existing request".to_owned());
    }
    Ok(true)
}

/// Replace exactly one accepted request contribution, without replaying other requests.
pub fn replace(
    summary: &mut SessionUsageSummary,
    previous: Option<&SessionTokenUsage>,
    next: &SessionTokenUsage,
) -> Result<(), String> {
    let mut updated = summary.clone();
    if let Some(previous) = previous {
        contribution(&mut updated, previous, false)?;
    }
    contribution(&mut updated, next, true)?;
    let mut next = next.clone();
    // The latest-usage snapshot carries facts, not a second cost projection.
    next.cost = None;
    updated.latest_usage = Some(next);
    *summary = updated;
    Ok(())
}

fn adjust(value: &mut u64, amount: u64, add: bool) -> Result<(), String> {
    *value = if add {
        value.checked_add(amount)
    } else {
        value.checked_sub(amount)
    }
    .ok_or_else(|| "session usage arithmetic is inconsistent or exceeds capacity".to_owned())?;
    Ok(())
}

fn contribution(
    summary: &mut SessionUsageSummary,
    usage: &SessionTokenUsage,
    add: bool,
) -> Result<(), String> {
    adjust(&mut summary.observed_usage_count, 1, add)?;
    adjust(
        &mut summary.cumulative_metered_tokens,
        u64::from(usage.metered_total_tokens().unwrap_or_default()),
        add,
    )?;
    match &usage.cost {
        Some(SessionCostEstimate::Estimated {
            currency,
            total_micros,
            ..
        }) => {
            adjust(
                summary.totals_micros.entry(currency.clone()).or_default(),
                *total_micros,
                add,
            )?;
            adjust(&mut summary.estimated_usage_count, 1, add)?;
        }
        Some(SessionCostEstimate::Unavailable { .. }) => {
            adjust(&mut summary.unavailable_usage_count, 1, add)?;
        }
        None => {}
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn costs_do_not_constrain_usage_fact_lifecycle() {
        let mut current = SessionTokenUsage {
            request_id: Some("request".into()),
            observation_id: Some("request:0".into()),
            cost: Some(SessionCostEstimate::Estimated {
                currency: "USD".into(),
                total_micros: 10,
                components: Vec::new(),
                source: "fixture".into(),
                revision: None,
            }),
            ..Default::default()
        };
        let mut next = current.clone();
        next.observation_ordinal = 1;
        next.cost = None;
        assert!(accepts(Some(&current), &next).is_ok());
        current.terminal = true;
        next.cost = current.cost.clone();
        assert!(accepts(Some(&current), &next).is_err());
    }

    #[test]
    fn observations_are_idempotent_and_conflicts_fail_closed() {
        let initial = SessionTokenUsage {
            request_id: Some("request".into()),
            observation_id: Some("request:0".into()),
            ..SessionTokenUsage::default()
        };
        let mut final_usage = initial.clone();
        final_usage.observation_ordinal = 1;
        final_usage.terminal = true;
        assert_eq!(accepts(Some(&initial), &initial), Ok(false));
        assert_eq!(accepts(Some(&initial), &final_usage), Ok(true));
        assert_eq!(accepts(Some(&final_usage), &initial), Ok(false));
        let mut conflict = final_usage.clone();
        conflict.output_tokens = Some(1);
        assert!(accepts(Some(&final_usage), &conflict).is_err());
    }
}