af_agent_session/
corrections.rs

1//! Append-only corrections to measurement facts. Reservation settlement remains immutable.
2use af_context::{MeteringCorrectionId, RunId, SubjectId};
3use serde::{Deserialize, Serialize};
4
5use crate::{EventError, MeteringDetails, MeteringSource, SessionProjection};
6
7/// Latest measurement for one stable Run operation; revision zero is its original fact.
8#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
9#[serde(deny_unknown_fields)]
10pub struct OperationUsage {
11    /// Prompt tokens measured.
12    pub prompt_tokens: u64,
13    /// Completion tokens measured.
14    pub completion_tokens: u64,
15    /// Provider-neutral resource weight, never a price.
16    pub cost_units: u64,
17    /// Attribution; absent only for legacy data.
18    pub metering: Option<MeteringDetails>,
19    /// Number of accepted corrections to this operation.
20    pub revision: u64,
21}
22
23/// Privileged report that replaces the derived measurement, never its original event.
24#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
25#[serde(deny_unknown_fields)]
26pub struct MeteringCorrection {
27    /// Stable idempotency identity, unique within the Session.
28    pub id: MeteringCorrectionId,
29    /// Terminal Run whose measurement is corrected.
30    pub run_id: RunId,
31    /// Original stable operation identity.
32    pub operation: String,
33    /// Exact measurement revision the reporter observed.
34    pub expected_revision: u64,
35    /// Corrected prompt tokens.
36    pub prompt_tokens: u64,
37    /// Corrected completion tokens.
38    pub completion_tokens: u64,
39    /// Corrected resource units; no monetary meaning in Factory.
40    pub cost_units: u64,
41    /// Reported attribution, preserving the original operation identity.
42    pub metering: MeteringDetails,
43    /// Short audit reason, without credentials or provider response bodies.
44    pub reason: String,
45}
46
47/// Accepted correction with server-derived actor and immutable event position.
48#[derive(Debug, Clone, PartialEq, Eq)]
49pub struct AppliedMeteringCorrection {
50    /// Caller-supplied report.
51    pub correction: MeteringCorrection,
52    /// Authenticated reporting subject.
53    pub actor_id: SubjectId,
54    /// Session event sequence of first acceptance.
55    pub seq: u64,
56}
57
58impl MeteringCorrection {
59    /// Validate bounded, non-overflowing reported measurements before any write.
60    pub fn validate(&self) -> Result<(), EventError> {
61        self.metering.validate()?;
62        let valid_text = |value: &str| {
63            !value.trim().is_empty() && value.len() <= 512 && !value.chars().any(char::is_control)
64        };
65        if !valid_text(self.id.as_str())
66            || !valid_text(&self.operation)
67            || !valid_text(&self.reason)
68            || self.expected_revision == u64::MAX
69            || self
70                .prompt_tokens
71                .checked_add(self.completion_tokens)
72                .and_then(|v| v.checked_add(self.cost_units))
73                .is_none_or(|v| v > i64::MAX as u64)
74            || !matches!(
75                self.metering,
76                MeteringDetails::Model {
77                    source: MeteringSource::Reported,
78                    ..
79                } | MeteringDetails::Tool {
80                    source: MeteringSource::Reported,
81                    ..
82                }
83            )
84        {
85            return Err(EventError::InvalidMetering);
86        }
87        Ok(())
88    }
89}
90
91impl SessionProjection {
92    /// Latest authoritative measurement. Original quota settlement is queried separately.
93    pub fn operation_usage(&self, run_id: &RunId, operation: &str) -> Option<OperationUsage> {
94        let key = (run_id.clone(), operation.to_owned());
95        if let Some(value) = self.corrected_usage.get(&key) {
96            return Some(value.clone());
97        }
98        if let Some((prompt_tokens, completion_tokens, cost_units, metering)) =
99            self.usage_operations.get(&key)
100        {
101            return Some(OperationUsage {
102                prompt_tokens: *prompt_tokens,
103                completion_tokens: *completion_tokens,
104                cost_units: *cost_units,
105                metering: metering.clone(),
106                revision: 0,
107            });
108        }
109        self.pending_usage_operations.get(&key).map(
110            |(prompt_tokens, completion_tokens, cost_units, metering)| OperationUsage {
111                prompt_tokens: *prompt_tokens,
112                completion_tokens: *completion_tokens,
113                cost_units: *cost_units,
114                metering: metering.clone(),
115                revision: 0,
116            },
117        )
118    }
119
120    /// Original receipt for an idempotently accepted correction.
121    pub fn metering_correction(
122        &self,
123        id: &MeteringCorrectionId,
124    ) -> Option<&AppliedMeteringCorrection> {
125        self.metering_corrections.get(id)
126    }
127
128    pub(super) fn apply_metering_correction(
129        &mut self,
130        correction: &MeteringCorrection,
131        actor_id: &SubjectId,
132        seq: u64,
133    ) -> Result<(), EventError> {
134        correction.validate()?;
135        if let Some(applied) = self.metering_corrections.get(&correction.id) {
136            return if applied.correction == *correction && applied.actor_id == *actor_id {
137                Ok(())
138            } else {
139                Err(EventError::UsageConflict(correction.operation.clone()))
140            };
141        }
142        let previous = self
143            .operation_usage(&correction.run_id, &correction.operation)
144            .ok_or_else(|| EventError::UsageConflict(correction.operation.clone()))?;
145        if !self
146            .run_status
147            .get(&correction.run_id)
148            .is_some_and(|status| status.is_terminal())
149            || previous.revision != correction.expected_revision
150            || !previous
151                .metering
152                .as_ref()
153                .is_some_and(|details| correction.metering.completes(details))
154        {
155            return Err(EventError::UsageConflict(correction.operation.clone()));
156        }
157        self.corrected_usage.insert(
158            (correction.run_id.clone(), correction.operation.clone()),
159            OperationUsage {
160                prompt_tokens: correction.prompt_tokens,
161                completion_tokens: correction.completion_tokens,
162                cost_units: correction.cost_units,
163                metering: Some(correction.metering.clone()),
164                revision: previous.revision + 1,
165            },
166        );
167        self.metering_corrections.insert(
168            correction.id.clone(),
169            AppliedMeteringCorrection {
170                correction: correction.clone(),
171                actor_id: actor_id.clone(),
172                seq,
173            },
174        );
175        Ok(())
176    }
177}
178
179#[cfg(test)]
180mod tests {
181    use super::*;
182    use crate::{MeteringOutcome, RunState, RunStatus};
183
184    fn report() -> MeteringCorrection {
185        MeteringCorrection {
186            id: "correction".parse().unwrap(),
187            run_id: "run".parse().unwrap(),
188            operation: "operation".into(),
189            expected_revision: 0,
190            prompt_tokens: 1,
191            completion_tokens: 2,
192            cost_units: 0,
193            metering: MeteringDetails::Model {
194                model: "m".into(),
195                provider: Some("p".into()),
196                provider_attempt_id: "attempt".parse().unwrap(),
197                source: MeteringSource::Reported,
198                outcome: MeteringOutcome::Completed,
199            },
200            reason: "provider report".into(),
201        }
202    }
203
204    #[test]
205    fn invalid_reports_and_mismatched_identities_fail_closed() {
206        let good = report();
207        good.validate().unwrap();
208        for mutation in 0..7 {
209            let mut invalid = good.clone();
210            match mutation {
211                0 => invalid.reason.clear(),
212                1 => invalid.operation = "a".repeat(513),
213                2 => invalid.prompt_tokens = u64::MAX,
214                3 => invalid.expected_revision = u64::MAX,
215                4 => invalid.id = "bad\nid".parse().unwrap(),
216                5 => {
217                    if let MeteringDetails::Model { source, .. } = &mut invalid.metering {
218                        *source = MeteringSource::Estimated;
219                    }
220                }
221                _ => {
222                    if let MeteringDetails::Model { model, .. } = &mut invalid.metering {
223                        model.clear();
224                    }
225                }
226            }
227            assert!(invalid.validate().is_err());
228        }
229        let mut projection = SessionProjection::default();
230        let actor = "admin".parse().unwrap();
231        assert!(projection
232            .apply_metering_correction(&good, &actor, 10)
233            .is_err());
234        projection
235            .run_status
236            .insert(good.run_id.clone(), RunState::Terminal(RunStatus::Failed));
237        let mut pending = good.metering.clone();
238        if let MeteringDetails::Model {
239            source, outcome, ..
240        } = &mut pending
241        {
242            *source = MeteringSource::Estimated;
243            *outcome = MeteringOutcome::Unknown;
244        }
245        projection.pending_usage_operations.insert(
246            (good.run_id.clone(), good.operation.clone()),
247            (100, 20, 0, Some(pending)),
248        );
249        assert_eq!(
250            projection
251                .operation_usage(&good.run_id, &good.operation)
252                .unwrap()
253                .prompt_tokens,
254            100
255        );
256        let mut wrong = good.clone();
257        if let MeteringDetails::Model { provider, .. } = &mut wrong.metering {
258            *provider = Some("wrong".into());
259        }
260        assert!(projection
261            .apply_metering_correction(&wrong, &actor, 10)
262            .is_err());
263        projection
264            .apply_metering_correction(&good, &actor, 10)
265            .unwrap();
266        projection
267            .apply_metering_correction(&good, &actor, 11)
268            .unwrap();
269        assert_eq!(projection.metering_correction(&good.id).unwrap().seq, 10);
270        assert!(projection
271            .apply_metering_correction(&good, &"other".parse().unwrap(), 12)
272            .is_err());
273        assert_eq!(
274            projection
275                .operation_usage(&good.run_id, &good.operation)
276                .unwrap()
277                .prompt_tokens,
278            1
279        );
280        assert_eq!(projection.billable_units_for("run"), 120);
281        assert!(projection
282            .operation_usage(&good.run_id, "missing")
283            .is_none());
284    }
285}
af_agent_session/corrections.rs

af_agent_session/
corrections.rs
