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/// An original operation used in reservation settlement, never a corrected price.
24#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
25pub struct RunOperationUsage {
26    /// Stable operation identity within the Run, suitable for host deduplication.
27    pub operation_id: String,
28    /// Original immutable measurement, including provider/model or tool attribution.
29    pub usage: OperationUsage,
30    /// Preparation has no durable result; the usage is a conservative estimate.
31    pub pending: bool,
32}
33
34/// Privileged report that replaces the derived measurement, never its original event.
35#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
36#[serde(deny_unknown_fields)]
37pub struct MeteringCorrection {
38    /// Stable idempotency identity, unique within the Session.
39    pub id: MeteringCorrectionId,
40    /// Terminal Run whose measurement is corrected.
41    pub run_id: RunId,
42    /// Original stable operation identity.
43    pub operation: String,
44    /// Exact measurement revision the reporter observed.
45    pub expected_revision: u64,
46    /// Corrected prompt tokens.
47    pub prompt_tokens: u64,
48    /// Corrected completion tokens.
49    pub completion_tokens: u64,
50    /// Corrected resource units; no monetary meaning in Factory.
51    pub cost_units: u64,
52    /// Reported attribution, preserving the original operation identity.
53    pub metering: MeteringDetails,
54    /// Short audit reason, without credentials or provider response bodies.
55    pub reason: String,
56}
57
58/// Accepted correction with server-derived actor and immutable event position.
59#[derive(Debug, Clone, PartialEq, Eq)]
60pub struct AppliedMeteringCorrection {
61    /// Caller-supplied report.
62    pub correction: MeteringCorrection,
63    /// Authenticated reporting subject.
64    pub actor_id: SubjectId,
65    /// Session event sequence of first acceptance.
66    pub seq: u64,
67}
68
69impl MeteringCorrection {
70    /// Validate bounded, non-overflowing reported measurements before any write.
71    pub fn validate(&self) -> Result<(), EventError> {
72        self.metering.validate()?;
73        let valid_text = |value: &str| {
74            !value.trim().is_empty() && value.len() <= 512 && !value.chars().any(char::is_control)
75        };
76        if !valid_text(self.id.as_str())
77            || !valid_text(&self.operation)
78            || !valid_text(&self.reason)
79            || self.expected_revision == u64::MAX
80            || self
81                .prompt_tokens
82                .checked_add(self.completion_tokens)
83                .and_then(|v| v.checked_add(self.cost_units))
84                .is_none_or(|v| v > i64::MAX as u64)
85            || !matches!(
86                self.metering,
87                MeteringDetails::Model {
88                    source: MeteringSource::Reported,
89                    ..
90                } | MeteringDetails::Tool {
91                    source: MeteringSource::Reported,
92                    ..
93                }
94            )
95        {
96            return Err(EventError::InvalidMetering);
97        }
98        Ok(())
99    }
100}
101
102impl SessionProjection {
103    /// Original settlement evidence in stable operation order, including pending estimates.
104    /// Corrections are deliberately separate (`operation_usage`); they must never
105    /// silently reprice an already settled reservation. Legacy facts lack attribution.
106    /// Read through an authorized AgentStore projection; this value grants no access.
107    pub fn run_usage_operations(&self, run_id: &RunId) -> Vec<RunOperationUsage> {
108        let mut result: std::collections::BTreeMap<String, RunOperationUsage> = self
109            .pending_usage_operations
110            .iter()
111            .filter(|((id, _), _)| id == run_id)
112            .map(|((_, operation), (prompt, completion, cost, metering))| {
113                (
114                    operation.clone(),
115                    RunOperationUsage {
116                        operation_id: operation.clone(),
117                        usage: OperationUsage {
118                            prompt_tokens: *prompt,
119                            completion_tokens: *completion,
120                            cost_units: *cost,
121                            metering: metering.clone(),
122                            revision: 0,
123                        },
124                        pending: true,
125                    },
126                )
127            })
128            .collect();
129        for ((id, operation), (prompt, completion, cost, metering)) in &self.usage_operations {
130            if id == run_id {
131                result.insert(
132                    operation.clone(),
133                    RunOperationUsage {
134                        operation_id: operation.clone(),
135                        usage: OperationUsage {
136                            prompt_tokens: *prompt,
137                            completion_tokens: *completion,
138                            cost_units: *cost,
139                            metering: metering.clone(),
140                            revision: 0,
141                        },
142                        pending: false,
143                    },
144                );
145            }
146        }
147        result.into_values().collect()
148    }
149
150    /// Latest authoritative measurement. Original quota settlement is queried separately.
151    pub fn operation_usage(&self, run_id: &RunId, operation: &str) -> Option<OperationUsage> {
152        let key = (run_id.clone(), operation.to_owned());
153        if let Some(value) = self.corrected_usage.get(&key) {
154            return Some(value.clone());
155        }
156        if let Some((prompt_tokens, completion_tokens, cost_units, metering)) =
157            self.usage_operations.get(&key)
158        {
159            return Some(OperationUsage {
160                prompt_tokens: *prompt_tokens,
161                completion_tokens: *completion_tokens,
162                cost_units: *cost_units,
163                metering: metering.clone(),
164                revision: 0,
165            });
166        }
167        self.pending_usage_operations.get(&key).map(
168            |(prompt_tokens, completion_tokens, cost_units, metering)| OperationUsage {
169                prompt_tokens: *prompt_tokens,
170                completion_tokens: *completion_tokens,
171                cost_units: *cost_units,
172                metering: metering.clone(),
173                revision: 0,
174            },
175        )
176    }
177
178    /// Original receipt for an idempotently accepted correction.
179    pub fn metering_correction(
180        &self,
181        id: &MeteringCorrectionId,
182    ) -> Option<&AppliedMeteringCorrection> {
183        self.metering_corrections.get(id)
184    }
185
186    pub(super) fn apply_metering_correction(
187        &mut self,
188        correction: &MeteringCorrection,
189        actor_id: &SubjectId,
190        seq: u64,
191    ) -> Result<(), EventError> {
192        correction.validate()?;
193        if let Some(applied) = self.metering_corrections.get(&correction.id) {
194            return if applied.correction == *correction && applied.actor_id == *actor_id {
195                Ok(())
196            } else {
197                Err(EventError::UsageConflict(correction.operation.clone()))
198            };
199        }
200        let previous = self
201            .operation_usage(&correction.run_id, &correction.operation)
202            .ok_or_else(|| EventError::UsageConflict(correction.operation.clone()))?;
203        if !self
204            .run_status
205            .get(&correction.run_id)
206            .is_some_and(|status| status.is_terminal())
207            || previous.revision != correction.expected_revision
208            || !previous
209                .metering
210                .as_ref()
211                .is_some_and(|details| correction.metering.completes(details))
212        {
213            return Err(EventError::UsageConflict(correction.operation.clone()));
214        }
215        self.corrected_usage.insert(
216            (correction.run_id.clone(), correction.operation.clone()),
217            OperationUsage {
218                prompt_tokens: correction.prompt_tokens,
219                completion_tokens: correction.completion_tokens,
220                cost_units: correction.cost_units,
221                metering: Some(correction.metering.clone()),
222                revision: previous.revision + 1,
223            },
224        );
225        self.metering_corrections.insert(
226            correction.id.clone(),
227            AppliedMeteringCorrection {
228                correction: correction.clone(),
229                actor_id: actor_id.clone(),
230                seq,
231            },
232        );
233        Ok(())
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240    use crate::{MeteringOutcome, RunState, RunStatus};
241
242    #[test]
243    fn original_evidence_keeps_mixed_models_pending_tools_and_excludes_corrections() {
244        let mut projection = SessionProjection::default();
245        let run: RunId = "run".parse().unwrap();
246        for (operation, model) in [("model-a", "model-one"), ("model-b", "model-two")] {
247            projection.usage_operations.insert(
248                (run.clone(), operation.into()),
249                (
250                    3,
251                    2,
252                    0,
253                    Some(MeteringDetails::Model {
254                        model: model.into(),
255                        provider: Some("provider".into()),
256                        provider_attempt_id: operation.parse().unwrap(),
257                        source: MeteringSource::Reported,
258                        outcome: MeteringOutcome::Completed,
259                    }),
260                ),
261            );
262        }
263        projection.pending_usage_operations.insert(
264            (run.clone(), "tool".into()),
265            (
266                0,
267                0,
268                7,
269                Some(MeteringDetails::Tool {
270                    call_id: "call".parse().unwrap(),
271                    name: "lookup".into(),
272                    source: MeteringSource::Estimated,
273                    outcome: MeteringOutcome::Unknown,
274                }),
275            ),
276        );
277        projection.usage_operations.insert(
278            ("another-run".parse().unwrap(), "model-a".into()),
279            (99, 0, 0, None),
280        );
281        let original = projection.run_usage_operations(&run);
282        assert_eq!(original.len(), 3);
283        assert!(!original[0].pending);
284        assert!(original[2].pending);
285        assert_ne!(original[0].usage.metering, original[1].usage.metering);
286        assert_eq!(
287            original
288                .iter()
289                .map(|fact| fact.usage.prompt_tokens
290                    + fact.usage.completion_tokens
291                    + fact.usage.cost_units)
292                .sum::<u64>(),
293            projection.billable_units_for("run")
294        );
295        projection.corrected_usage.insert(
296            (run.clone(), "model-a".into()),
297            OperationUsage {
298                prompt_tokens: 1,
299                completion_tokens: 1,
300                cost_units: 0,
301                metering: None,
302                revision: 1,
303            },
304        );
305        assert_eq!(projection.run_usage_operations(&run), original);
306        assert_eq!(projection.clone().run_usage_operations(&run), original);
307        assert!(projection
308            .run_usage_operations(&"missing".parse().unwrap())
309            .is_empty());
310    }
311
312    fn report() -> MeteringCorrection {
313        MeteringCorrection {
314            id: "correction".parse().unwrap(),
315            run_id: "run".parse().unwrap(),
316            operation: "operation".into(),
317            expected_revision: 0,
318            prompt_tokens: 1,
319            completion_tokens: 2,
320            cost_units: 0,
321            metering: MeteringDetails::Model {
322                model: "m".into(),
323                provider: Some("p".into()),
324                provider_attempt_id: "attempt".parse().unwrap(),
325                source: MeteringSource::Reported,
326                outcome: MeteringOutcome::Completed,
327            },
328            reason: "provider report".into(),
329        }
330    }
331
332    #[test]
333    fn invalid_reports_and_mismatched_identities_fail_closed() {
334        let good = report();
335        good.validate().unwrap();
336        for mutation in 0..7 {
337            let mut invalid = good.clone();
338            match mutation {
339                0 => invalid.reason.clear(),
340                1 => invalid.operation = "a".repeat(513),
341                2 => invalid.prompt_tokens = u64::MAX,
342                3 => invalid.expected_revision = u64::MAX,
343                4 => invalid.id = "bad\nid".parse().unwrap(),
344                5 => {
345                    if let MeteringDetails::Model { source, .. } = &mut invalid.metering {
346                        *source = MeteringSource::Estimated;
347                    }
348                }
349                _ => {
350                    if let MeteringDetails::Model { model, .. } = &mut invalid.metering {
351                        model.clear();
352                    }
353                }
354            }
355            assert!(invalid.validate().is_err());
356        }
357        let mut projection = SessionProjection::default();
358        let actor = "admin".parse().unwrap();
359        assert!(projection
360            .apply_metering_correction(&good, &actor, 10)
361            .is_err());
362        projection
363            .run_status
364            .insert(good.run_id.clone(), RunState::Terminal(RunStatus::Failed));
365        let mut pending = good.metering.clone();
366        if let MeteringDetails::Model {
367            source, outcome, ..
368        } = &mut pending
369        {
370            *source = MeteringSource::Estimated;
371            *outcome = MeteringOutcome::Unknown;
372        }
373        projection.pending_usage_operations.insert(
374            (good.run_id.clone(), good.operation.clone()),
375            (100, 20, 0, Some(pending)),
376        );
377        assert_eq!(
378            projection
379                .operation_usage(&good.run_id, &good.operation)
380                .unwrap()
381                .prompt_tokens,
382            100
383        );
384        let mut wrong = good.clone();
385        if let MeteringDetails::Model { provider, .. } = &mut wrong.metering {
386            *provider = Some("wrong".into());
387        }
388        assert!(projection
389            .apply_metering_correction(&wrong, &actor, 10)
390            .is_err());
391        projection
392            .apply_metering_correction(&good, &actor, 10)
393            .unwrap();
394        projection
395            .apply_metering_correction(&good, &actor, 11)
396            .unwrap();
397        assert_eq!(projection.metering_correction(&good.id).unwrap().seq, 10);
398        assert!(projection
399            .apply_metering_correction(&good, &"other".parse().unwrap(), 12)
400            .is_err());
401        assert_eq!(
402            projection
403                .operation_usage(&good.run_id, &good.operation)
404                .unwrap()
405                .prompt_tokens,
406            1
407        );
408        assert_eq!(projection.billable_units_for("run"), 120);
409        assert!(projection
410            .operation_usage(&good.run_id, "missing")
411            .is_none());
412    }
413}
af_agent_session/corrections.rs

af_agent_session/
corrections.rs
