Skip to main content

af_agent_session/
metering.rs

1//! Attribution attached to immutable operation usage. Monetary policy stays in the host.
2use af_context::{ProviderAttemptId, ToolCallId};
3use serde::{Deserialize, Serialize};
4
5/// Origin of token or resource measurements.
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
7#[serde(rename_all = "snake_case")]
8pub enum MeteringSource {
9    /// The executing provider supplied the measurement.
10    Reported,
11    /// Factory estimated usage or applied a declared resource weight.
12    Estimated,
13}
14
15/// Result known at the time this immutable measurement was recorded.
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
17#[serde(rename_all = "snake_case")]
18pub enum MeteringOutcome {
19    /// The operation returned successfully.
20    Completed,
21    /// The operation returned an error.
22    Failed,
23    /// The caller cancelled the operation.
24    Cancelled,
25    /// New input interrupted the operation.
26    Steered,
27    /// Deadline elapsed before the result was known.
28    TimedOut,
29    /// Admission failed before provider dispatch.
30    NotDispatched,
31    /// A durable preparation survived but the result did not.
32    Unknown,
33}
34
35/// Stable attribution for one operation; contains no request, secret or price.
36#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
37#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
38pub enum MeteringDetails {
39    /// One model invocation (including summarization).
40    Model {
41        /// Exact requested model.
42        model: String,
43        /// Registered provider; absent for an unlabelled custom adapter.
44        provider: Option<String>,
45        /// Stable identity recorded before dispatch.
46        provider_attempt_id: ProviderAttemptId,
47        /// Whether usage came from the provider.
48        source: MeteringSource,
49        /// Result known when recording usage.
50        outcome: MeteringOutcome,
51    },
52    /// One tool invocation.
53    Tool {
54        /// Stable tool call identity.
55        call_id: ToolCallId,
56        /// Registered tool name.
57        name: String,
58        /// Declared resource weights are estimates, not prices.
59        source: MeteringSource,
60        /// Result known when recording usage.
61        outcome: MeteringOutcome,
62    },
63}
64
65impl MeteringDetails {
66    /// A result must preserve the preparation's identity; unknown provider may become known.
67    pub fn completes(&self, prepared: &Self) -> bool {
68        match (self, prepared) {
69            (
70                Self::Model {
71                    model,
72                    provider,
73                    provider_attempt_id,
74                    ..
75                },
76                Self::Model {
77                    model: expected_model,
78                    provider: expected_provider,
79                    provider_attempt_id: expected_attempt,
80                    ..
81                },
82            ) => {
83                model == expected_model
84                    && provider_attempt_id == expected_attempt
85                    && expected_provider
86                        .as_ref()
87                        .is_none_or(|expected| provider.as_ref() == Some(expected))
88            }
89            (
90                Self::Tool { call_id, name, .. },
91                Self::Tool {
92                    call_id: expected_call,
93                    name: expected_name,
94                    ..
95                },
96            ) => call_id == expected_call && name == expected_name,
97            _ => false,
98        }
99    }
100
101    /// Reject malformed identifiers and unbounded attribution at the write boundary.
102    pub fn validate(&self) -> Result<(), crate::EventError> {
103        let valid =
104            |s: &str| !s.trim().is_empty() && s.len() <= 512 && !s.chars().any(char::is_control);
105        let ok = match self {
106            Self::Model {
107                model,
108                provider,
109                provider_attempt_id,
110                ..
111            } => {
112                valid(model)
113                    && provider.as_deref().is_none_or(valid)
114                    && valid(provider_attempt_id.as_str())
115            }
116            Self::Tool { call_id, name, .. } => valid(call_id.as_str()) && valid(name),
117        };
118        if ok {
119            Ok(())
120        } else {
121            Err(crate::EventError::InvalidMetering)
122        }
123    }
124}
125
126impl MeteringSource {
127    /// Stable protocol spelling.
128    pub fn as_str(self) -> &'static str {
129        match self {
130            Self::Reported => "reported",
131            Self::Estimated => "estimated",
132        }
133    }
134}
135impl MeteringOutcome {
136    /// Stable protocol spelling.
137    pub fn as_str(self) -> &'static str {
138        match self {
139            Self::Completed => "completed",
140            Self::Failed => "failed",
141            Self::Cancelled => "cancelled",
142            Self::Steered => "steered",
143            Self::TimedOut => "timed_out",
144            Self::NotDispatched => "not_dispatched",
145            Self::Unknown => "unknown",
146        }
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153    #[test]
154    fn attribution_is_bounded_and_preserves_unknown_provider_and_outcomes() {
155        for outcome in [
156            MeteringOutcome::Completed,
157            MeteringOutcome::Failed,
158            MeteringOutcome::Cancelled,
159            MeteringOutcome::Steered,
160            MeteringOutcome::TimedOut,
161            MeteringOutcome::NotDispatched,
162            MeteringOutcome::Unknown,
163        ] {
164            let mut details = MeteringDetails::Model {
165                model: "m".into(),
166                provider: None,
167                provider_attempt_id: "attempt".parse().unwrap(),
168                source: MeteringSource::Estimated,
169                outcome,
170            };
171            details.validate().unwrap();
172            assert_eq!(serde_json::to_value(outcome).unwrap(), outcome.as_str());
173            assert_eq!(
174                serde_json::from_value::<MeteringDetails>(serde_json::to_value(&details).unwrap())
175                    .unwrap(),
176                details
177            );
178            if let MeteringDetails::Model { provider, .. } = &mut details {
179                *provider = Some("\n".into());
180            }
181            assert!(details.validate().is_err());
182        }
183        for name in ["".to_owned(), "a".repeat(513), "a\nb".into()] {
184            assert!(MeteringDetails::Tool {
185                name,
186                call_id: "call".parse().unwrap(),
187                source: MeteringSource::Estimated,
188                outcome: MeteringOutcome::Unknown
189            }
190            .validate()
191            .is_err());
192        }
193    }
194}