mcp_usage_store/
lib.rs

1//! Distributed durable-task attribution for `UsageKit` for MCP.
2//!
3//! Backends store SHA-256 digests of tenant and task identifiers, never their
4//! plaintext values. Values contain only a fixed method category and the price
5//! resolved when the task was created, never a name, URI, or extension method
6//! string. Inserts are first-writer-wins and claims are atomic, so a completed
7//! task can be accounted for by at most one application instance.
8//!
9//! # Choosing a TTL
10//!
11//! **The TTL must outlive the longest task the upstream server can run.** It is
12//! the deadline on the whole task, not on the polling interval: the attribution
13//! is written when the task is created and read when it completes, so a task
14//! that takes longer than the TTL loses the price resolved at creation.
15//!
16//! Nothing fails loudly when that happens. The completing poll finds no
17//! attribution, classifies as [`FreeReason::MissingTaskAttribution`], and the
18//! call bills zero. From the customer's side the work was delivered and never
19//! charged, so the symptom is silent under-billing that scales with how long
20//! tasks run. Watch `mcp_usage_free_deliveries_by_reason_total{reason=
21//! "missing_task_attribution"}`: a nonzero rate on a server that issues durable
22//! tasks is the signal that this TTL is too short.
23//!
24//! Set it from the upstream's own task timeout with headroom, not from a guess
25//! about client behaviour. Over-provisioning is cheap - rows are 10 bytes of
26//! value behind two 32-byte digests, and expiry reclaims them - while
27//! under-provisioning loses revenue.
28//!
29//! A per-task TTL would remove the guess, since the creating call knows what it
30//! commissioned. It is not available yet: [`TaskAttribution`] carries only the
31//! origin category and the resolved units, and
32//! [`TaskAttributionStore::insert`] takes no deadline, so expressing one is a
33//! breaking change to that trait rather than a change inside these backends.
34//!
35//! [`FreeReason::MissingTaskAttribution`]: mcp_usage_core::FreeReason::MissingTaskAttribution
36//! [`TaskAttributionStore::insert`]: mcp_usage_tower::TaskAttributionStore::insert
37
38#![forbid(unsafe_code)]
39#![warn(missing_docs, clippy::pedantic)]
40#![allow(clippy::module_name_repetitions)]
41
42#[cfg(any(feature = "postgres", feature = "redis", test))]
43use mcp_usage_core::{TaskAttribution, TaskOriginKind};
44#[cfg(any(feature = "postgres", feature = "redis", test))]
45use mcp_usage_tower::TaskStoreError;
46#[cfg(any(feature = "postgres", feature = "redis", test))]
47use sha2::{Digest, Sha256};
48use thiserror::Error;
49
50#[cfg(feature = "postgres")]
51mod postgres;
52#[cfg(feature = "redis")]
53mod redis;
54
55#[cfg(feature = "postgres")]
56pub use postgres::PostgresTaskStore;
57#[cfg(feature = "redis")]
58pub use redis::RedisTaskStore;
59
60/// Configuration or connection failure while constructing a store.
61#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
62pub enum StoreConfigError {
63    /// The time-to-live must be greater than zero and fit the backend format.
64    #[error("task-attribution TTL is outside the supported range")]
65    InvalidTtl,
66    /// Backend operations require a nonzero timeout.
67    #[error("task-attribution operation timeout must be greater than zero")]
68    InvalidTimeout,
69    /// A key prefix was empty, too long, or contained unsupported characters.
70    #[error("task-attribution key prefix is invalid")]
71    InvalidPrefix,
72    /// The backend connection could not be established.
73    #[error("task-attribution backend connection failed")]
74    Connection,
75    /// The backend schema could not be installed or verified.
76    #[error("task-attribution backend schema initialization failed")]
77    Schema,
78}
79
80#[cfg(any(feature = "postgres", feature = "redis", test))]
81fn identifier_hash(value: &str) -> [u8; 32] {
82    Sha256::digest(value.as_bytes()).into()
83}
84
85#[cfg(any(feature = "postgres", feature = "redis", test))]
86fn encode_attribution(attribution: TaskAttribution) -> [u8; 10] {
87    let mut encoded = [0_u8; 10];
88    encoded[0] = 1;
89    encoded[1] = match attribution.origin_kind() {
90        TaskOriginKind::ToolsCall => 1,
91        TaskOriginKind::ResourcesRead => 2,
92        TaskOriginKind::PromptsGet => 3,
93        TaskOriginKind::Other => 255,
94    };
95    encoded[2..].copy_from_slice(&attribution.units().to_be_bytes());
96    encoded
97}
98
99#[cfg(any(feature = "postgres", feature = "redis", test))]
100fn decode_attribution(payload: &[u8]) -> Result<TaskAttribution, TaskStoreError> {
101    let [1, kind, units @ ..] = payload else {
102        return Err(TaskStoreError::InvalidRecord);
103    };
104    let origin_kind = match kind {
105        1 => TaskOriginKind::ToolsCall,
106        2 => TaskOriginKind::ResourcesRead,
107        3 => TaskOriginKind::PromptsGet,
108        255 => TaskOriginKind::Other,
109        _ => return Err(TaskStoreError::InvalidRecord),
110    };
111    let units = u64::from_be_bytes(
112        units
113            .try_into()
114            .map_err(|_| TaskStoreError::InvalidRecord)?,
115    );
116    Ok(TaskAttribution::new(origin_kind, units))
117}
118
119#[cfg(any(feature = "redis", test))]
120fn encode_hash(value: &str) -> String {
121    const HEX: &[u8; 16] = b"0123456789abcdef";
122    let digest = identifier_hash(value);
123    let mut encoded = String::with_capacity(digest.len() * 2);
124    for byte in digest {
125        encoded.push(char::from(HEX[usize::from(byte >> 4)]));
126        encoded.push(char::from(HEX[usize::from(byte & 0x0f)]));
127    }
128    encoded
129}
130
131#[cfg(any(feature = "redis", test))]
132fn valid_prefix(prefix: &str) -> bool {
133    !prefix.is_empty()
134        && prefix.len() <= 64
135        && prefix
136            .bytes()
137            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b':' | b'_' | b'-'))
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143
144    #[test]
145    fn identifiers_are_fixed_length_and_absent_from_keys() {
146        let encoded = encode_hash("tenant:private@example.test");
147        assert_eq!(encoded.len(), 64);
148        assert!(!encoded.contains("private"));
149    }
150
151    #[test]
152    fn prefixes_are_bounded_and_ascii_only() {
153        assert!(valid_prefix("mcp-usage:tasks"));
154        assert!(!valid_prefix(""));
155        assert!(!valid_prefix("spaces are rejected"));
156        assert!(!valid_prefix(&"a".repeat(65)));
157    }
158
159    #[test]
160    fn attribution_encoding_is_fixed_length_and_validated() {
161        for kind in [
162            TaskOriginKind::ToolsCall,
163            TaskOriginKind::ResourcesRead,
164            TaskOriginKind::PromptsGet,
165            TaskOriginKind::Other,
166        ] {
167            let attribution = TaskAttribution::new(kind, u64::MAX);
168            let encoded = encode_attribution(attribution);
169            assert_eq!(encoded.len(), 10);
170            assert_eq!(decode_attribution(&encoded).unwrap(), attribution);
171        }
172        assert_eq!(
173            decode_attribution(&[1, 2]),
174            Err(TaskStoreError::InvalidRecord)
175        );
176        assert_eq!(
177            decode_attribution(&[2, 1, 0, 0, 0, 0, 0, 0, 0, 1]),
178            Err(TaskStoreError::InvalidRecord)
179        );
180    }
181}
mcp_usage_store/lib.rs

mcp_usage_store/
lib.rs
