mcp_usage_core/
price.rs

1//! The price book: how many billable units a delivered call is worth.
2//!
3//! Units are integers, not currency. Converting units to money is the billing
4//! provider's job (meter events, an invoice line, or another billing API), and
5//! keeping money out of this crate keeps rounding, currency, and tax out of the
6//! attribution engine where they would only cause trouble.
7//!
8//! Pricing is keyed on the value of `Mcp-Name`, which covers a tool name, a
9//! prompt name, or a resource URI depending on the method. One map therefore
10//! prices all three, which is what operators actually want: an expensive tool and
11//! an expensive resource are the same kind of problem.
12
13use serde::{Deserialize, Serialize};
14use std::collections::BTreeMap;
15
16use crate::Method;
17
18const fn default_units() -> u64 {
19    1
20}
21
22/// Per-name unit pricing with a fallback.
23///
24/// `BTreeMap` rather than `HashMap` so a serialized price book round trips in a
25/// stable order, which makes config diffs readable and tests deterministic.
26#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
27pub struct PriceBook {
28    /// Charged for a delivered call with no entry in [`Self::names`].
29    #[serde(default = "default_units")]
30    pub default_units: u64,
31
32    /// Per-name overrides, keyed on the decoded `Mcp-Name` value.
33    ///
34    /// A value of `0` makes that name free, which is the supported way to expose
35    /// a loss-leader tool without carving it out of the billing path.
36    #[serde(default)]
37    pub names: BTreeMap<String, u64>,
38
39    /// Per-method overrides, applied when no name matched.
40    ///
41    /// Mainly useful for pricing `resources/read` differently from `tools/call`
42    /// across the board.
43    #[serde(default)]
44    pub methods: BTreeMap<String, u64>,
45}
46
47impl Default for PriceBook {
48    fn default() -> Self {
49        Self {
50            default_units: default_units(),
51            names: BTreeMap::new(),
52            methods: BTreeMap::new(),
53        }
54    }
55}
56
57impl PriceBook {
58    /// A price book charging `units` for everything.
59    #[must_use]
60    pub fn flat(units: u64) -> Self {
61        Self {
62            default_units: units,
63            ..Self::default()
64        }
65    }
66
67    /// Set the price of one name. Chainable, for building books in tests and config.
68    #[must_use]
69    pub fn with_name(mut self, name: impl Into<String>, units: u64) -> Self {
70        self.names.insert(name.into(), units);
71        self
72    }
73
74    /// Set the price of one method.
75    #[must_use]
76    pub fn with_method(mut self, method: &Method, units: u64) -> Self {
77        self.methods.insert(method.as_str().to_owned(), units);
78        self
79    }
80
81    /// Units owed for a delivered call.
82    ///
83    /// Resolution order is most specific first: the name, then the method, then
84    /// the default. `name` must already be decoded from the `Mcp-Name` sentinel
85    /// form (see [`crate::name::decode`]); passing the raw header value here
86    /// silently prices every non-ASCII-named tool at the default.
87    #[must_use]
88    pub fn units_for(&self, method: &Method, name: Option<&str>) -> u64 {
89        if let Some(units) = name.and_then(|n| self.names.get(n)) {
90            return *units;
91        }
92        if let Some(units) = self.methods.get(method.as_str()) {
93            return *units;
94        }
95        self.default_units
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102
103    #[test]
104    fn defaults_to_one_unit_per_delivered_call() {
105        let book = PriceBook::default();
106        assert_eq!(book.units_for(&Method::ToolsCall, Some("anything")), 1);
107        assert_eq!(book.units_for(&Method::ToolsCall, None), 1);
108    }
109
110    #[test]
111    fn name_beats_method_beats_default() {
112        let book = PriceBook::flat(1)
113            .with_method(&Method::ToolsCall, 5)
114            .with_name("expensive_tool", 100);
115
116        assert_eq!(
117            book.units_for(&Method::ToolsCall, Some("expensive_tool")),
118            100
119        );
120        assert_eq!(book.units_for(&Method::ToolsCall, Some("other_tool")), 5);
121        assert_eq!(book.units_for(&Method::ToolsCall, None), 5);
122        assert_eq!(book.units_for(&Method::ResourcesRead, Some("other")), 1);
123    }
124
125    #[test]
126    fn zero_units_makes_a_name_free() {
127        let book = PriceBook::flat(10).with_name("free_probe", 0);
128        assert_eq!(book.units_for(&Method::ToolsCall, Some("free_probe")), 0);
129        assert_eq!(book.units_for(&Method::ToolsCall, Some("paid")), 10);
130    }
131
132    #[test]
133    fn one_book_prices_tools_prompts_and_resources_alike() {
134        // Mcp-Name carries a tool name, a prompt name, or a resource URI, so a
135        // single name map covers all three.
136        let book = PriceBook::flat(1)
137            .with_name("get_weather", 3)
138            .with_name("file:///big/dataset.csv", 50);
139
140        assert_eq!(book.units_for(&Method::ToolsCall, Some("get_weather")), 3);
141        assert_eq!(
142            book.units_for(&Method::ResourcesRead, Some("file:///big/dataset.csv")),
143            50
144        );
145    }
146
147    #[test]
148    fn round_trips_through_json_in_stable_order() {
149        let book = PriceBook::flat(2)
150            .with_name("zebra", 9)
151            .with_name("alpha", 1)
152            .with_method(&Method::ResourcesRead, 4);
153
154        let encoded = serde_json::to_string(&book).expect("serializes");
155        let decoded: PriceBook = serde_json::from_str(&encoded).expect("deserializes");
156        assert_eq!(book, decoded);
157        // BTreeMap ordering keeps config diffs readable.
158        assert!(
159            encoded.find("alpha").unwrap() < encoded.find("zebra").unwrap(),
160            "names should serialize in sorted order: {encoded}"
161        );
162    }
163
164    #[test]
165    fn a_minimal_config_deserializes() {
166        let book: PriceBook = serde_json::from_str("{}").expect("empty config is valid");
167        assert_eq!(book, PriceBook::default());
168
169        let book: PriceBook =
170            serde_json::from_str(r#"{"default_units":7}"#).expect("partial config is valid");
171        assert_eq!(book.default_units, 7);
172        assert!(book.names.is_empty());
173    }
174}
mcp_usage_core/price.rs

mcp_usage_core/
price.rs
