vv_agent/budget/
wire.rs

1use std::collections::BTreeMap;
2
3use serde::{Deserialize, Deserializer};
4
5use super::{
6    dimension_precedence, BudgetDimension, BudgetEnforcementBoundary, BudgetExhaustion,
7    BudgetExhaustionReason, BudgetUnavailableDimension, BudgetUnavailableReason,
8    BudgetUsageSnapshot, HostCost, RunBudgetLimits, UnavailableMetricPolicy,
9};
10
11#[derive(Deserialize)]
12struct HostCostWire {
13    unit: String,
14    #[serde(default)]
15    currency: Option<String>,
16    amount_microunits: u64,
17}
18
19impl<'de> Deserialize<'de> for HostCost {
20    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
21    where
22        D: Deserializer<'de>,
23    {
24        let wire = HostCostWire::deserialize(deserializer)?;
25        let cost = Self {
26            unit: wire.unit,
27            currency: wire.currency,
28            amount_microunits: wire.amount_microunits,
29        };
30        cost.validate().map_err(serde::de::Error::custom)?;
31        Ok(cost)
32    }
33}
34
35#[derive(Deserialize)]
36struct BudgetUnavailableDimensionWire {
37    dimension: BudgetDimension,
38    reason: BudgetUnavailableReason,
39    #[serde(default)]
40    expected_unit: Option<String>,
41    #[serde(default)]
42    observed_unit: Option<String>,
43    #[serde(default)]
44    expected_currency: Option<String>,
45    #[serde(default)]
46    observed_currency: Option<String>,
47}
48
49impl<'de> Deserialize<'de> for BudgetUnavailableDimension {
50    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
51    where
52        D: Deserializer<'de>,
53    {
54        let wire = BudgetUnavailableDimensionWire::deserialize(deserializer)?;
55        let observation = Self {
56            dimension: wire.dimension,
57            reason: wire.reason,
58            expected_unit: wire.expected_unit,
59            observed_unit: wire.observed_unit,
60            expected_currency: wire.expected_currency,
61            observed_currency: wire.observed_currency,
62        };
63        observation.validate().map_err(serde::de::Error::custom)?;
64        Ok(observation)
65    }
66}
67
68#[derive(Default, Deserialize)]
69struct RunBudgetLimitsWire {
70    #[serde(default)]
71    max_total_tokens: Option<u64>,
72    #[serde(default)]
73    max_uncached_input_tokens: Option<u64>,
74    #[serde(default)]
75    max_tool_calls: Option<u64>,
76    #[serde(default)]
77    max_tool_calls_by_name: BTreeMap<String, u64>,
78    #[serde(default)]
79    max_wall_time_ms: Option<u64>,
80    #[serde(default)]
81    max_host_cost: Option<HostCost>,
82    #[serde(default)]
83    unavailable_metric_policy: UnavailableMetricPolicy,
84}
85
86impl<'de> Deserialize<'de> for RunBudgetLimits {
87    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
88    where
89        D: Deserializer<'de>,
90    {
91        let wire = RunBudgetLimitsWire::deserialize(deserializer)?;
92        let limits = Self {
93            max_total_tokens: wire.max_total_tokens,
94            max_uncached_input_tokens: wire.max_uncached_input_tokens,
95            max_tool_calls: wire.max_tool_calls,
96            max_tool_calls_by_name: wire.max_tool_calls_by_name,
97            max_wall_time_ms: wire.max_wall_time_ms,
98            max_host_cost: wire.max_host_cost,
99            unavailable_metric_policy: wire.unavailable_metric_policy,
100        };
101        limits.validate().map_err(serde::de::Error::custom)?;
102        Ok(limits)
103    }
104}
105
106#[derive(Deserialize)]
107struct BudgetUsageSnapshotWire {
108    cycles: u64,
109    total_tokens: Option<u64>,
110    uncached_input_tokens: Option<u64>,
111    tool_calls: u64,
112    tool_calls_by_name: BTreeMap<String, u64>,
113    elapsed_ms: u64,
114    host_cost: Option<HostCost>,
115    unavailable_dimensions: Vec<BudgetUnavailableDimension>,
116}
117
118impl<'de> Deserialize<'de> for BudgetUsageSnapshot {
119    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
120    where
121        D: Deserializer<'de>,
122    {
123        let wire = BudgetUsageSnapshotWire::deserialize(deserializer)?;
124        let mut snapshot = Self {
125            cycles: wire.cycles,
126            total_tokens: wire.total_tokens,
127            uncached_input_tokens: wire.uncached_input_tokens,
128            tool_calls: wire.tool_calls,
129            tool_calls_by_name: wire.tool_calls_by_name,
130            elapsed_ms: wire.elapsed_ms,
131            host_cost: wire.host_cost,
132            unavailable_dimensions: wire.unavailable_dimensions,
133        };
134        snapshot
135            .unavailable_dimensions
136            .sort_by_key(|item| dimension_precedence(item.dimension));
137        snapshot.validate().map_err(serde::de::Error::custom)?;
138        Ok(snapshot)
139    }
140}
141
142#[derive(Deserialize)]
143struct BudgetExhaustionWire {
144    dimension: BudgetDimension,
145    #[serde(default)]
146    tool_name: Option<String>,
147    reason: BudgetExhaustionReason,
148    limit: u64,
149    observed: Option<u64>,
150    attempted_increment: Option<u64>,
151    overshoot: Option<u64>,
152    unit: String,
153    #[serde(default)]
154    currency: Option<String>,
155    enforcement_boundary: BudgetEnforcementBoundary,
156    #[serde(default)]
157    unavailable_reason: Option<BudgetUnavailableReason>,
158}
159
160impl<'de> Deserialize<'de> for BudgetExhaustion {
161    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
162    where
163        D: Deserializer<'de>,
164    {
165        let wire = BudgetExhaustionWire::deserialize(deserializer)?;
166        let exhaustion = Self {
167            dimension: wire.dimension,
168            tool_name: wire.tool_name,
169            reason: wire.reason,
170            limit: wire.limit,
171            observed: wire.observed,
172            attempted_increment: wire.attempted_increment,
173            overshoot: wire.overshoot,
174            unit: wire.unit,
175            currency: wire.currency,
176            enforcement_boundary: wire.enforcement_boundary,
177            unavailable_reason: wire.unavailable_reason,
178        };
179        exhaustion.validate().map_err(serde::de::Error::custom)?;
180        Ok(exhaustion)
181    }
182}
vv_agent/budget/wire.rs

vv_agent/budget/
wire.rs
