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#![forbid(unsafe_code)]
10#![warn(missing_docs, clippy::pedantic)]
11#![allow(clippy::module_name_repetitions)]
12
13#[cfg(any(feature = "postgres", feature = "redis", test))]
14use mcp_usage_core::{TaskAttribution, TaskOriginKind};
15#[cfg(any(feature = "postgres", feature = "redis", test))]
16use mcp_usage_tower::TaskStoreError;
17#[cfg(any(feature = "postgres", feature = "redis", test))]
18use sha2::{Digest, Sha256};
19use thiserror::Error;
20
21#[cfg(feature = "postgres")]
22mod postgres;
23#[cfg(feature = "redis")]
24mod redis;
25
26#[cfg(feature = "postgres")]
27pub use postgres::PostgresTaskStore;
28#[cfg(feature = "redis")]
29pub use redis::RedisTaskStore;
30
31/// Configuration or connection failure while constructing a store.
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
33pub enum StoreConfigError {
34    /// The time-to-live must be greater than zero and fit the backend format.
35    #[error("task-attribution TTL is outside the supported range")]
36    InvalidTtl,
37    /// Backend operations require a nonzero timeout.
38    #[error("task-attribution operation timeout must be greater than zero")]
39    InvalidTimeout,
40    /// A key prefix was empty, too long, or contained unsupported characters.
41    #[error("task-attribution key prefix is invalid")]
42    InvalidPrefix,
43    /// The backend connection could not be established.
44    #[error("task-attribution backend connection failed")]
45    Connection,
46    /// The backend schema could not be installed or verified.
47    #[error("task-attribution backend schema initialization failed")]
48    Schema,
49}
50
51#[cfg(any(feature = "postgres", feature = "redis", test))]
52fn identifier_hash(value: &str) -> [u8; 32] {
53    Sha256::digest(value.as_bytes()).into()
54}
55
56#[cfg(any(feature = "postgres", feature = "redis", test))]
57fn encode_attribution(attribution: TaskAttribution) -> [u8; 10] {
58    let mut encoded = [0_u8; 10];
59    encoded[0] = 1;
60    encoded[1] = match attribution.origin_kind() {
61        TaskOriginKind::ToolsCall => 1,
62        TaskOriginKind::ResourcesRead => 2,
63        TaskOriginKind::PromptsGet => 3,
64        TaskOriginKind::Other => 255,
65    };
66    encoded[2..].copy_from_slice(&attribution.units().to_be_bytes());
67    encoded
68}
69
70#[cfg(any(feature = "postgres", feature = "redis", test))]
71fn decode_attribution(payload: &[u8]) -> Result<TaskAttribution, TaskStoreError> {
72    let [1, kind, units @ ..] = payload else {
73        return Err(TaskStoreError::InvalidRecord);
74    };
75    let origin_kind = match kind {
76        1 => TaskOriginKind::ToolsCall,
77        2 => TaskOriginKind::ResourcesRead,
78        3 => TaskOriginKind::PromptsGet,
79        255 => TaskOriginKind::Other,
80        _ => return Err(TaskStoreError::InvalidRecord),
81    };
82    let units = u64::from_be_bytes(
83        units
84            .try_into()
85            .map_err(|_| TaskStoreError::InvalidRecord)?,
86    );
87    Ok(TaskAttribution::new(origin_kind, units))
88}
89
90#[cfg(any(feature = "redis", test))]
91fn encode_hash(value: &str) -> String {
92    const HEX: &[u8; 16] = b"0123456789abcdef";
93    let digest = identifier_hash(value);
94    let mut encoded = String::with_capacity(digest.len() * 2);
95    for byte in digest {
96        encoded.push(char::from(HEX[usize::from(byte >> 4)]));
97        encoded.push(char::from(HEX[usize::from(byte & 0x0f)]));
98    }
99    encoded
100}
101
102#[cfg(any(feature = "redis", test))]
103fn valid_prefix(prefix: &str) -> bool {
104    !prefix.is_empty()
105        && prefix.len() <= 64
106        && prefix
107            .bytes()
108            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b':' | b'_' | b'-'))
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114
115    #[test]
116    fn identifiers_are_fixed_length_and_absent_from_keys() {
117        let encoded = encode_hash("tenant:private@example.test");
118        assert_eq!(encoded.len(), 64);
119        assert!(!encoded.contains("private"));
120    }
121
122    #[test]
123    fn prefixes_are_bounded_and_ascii_only() {
124        assert!(valid_prefix("mcp-usage:tasks"));
125        assert!(!valid_prefix(""));
126        assert!(!valid_prefix("spaces are rejected"));
127        assert!(!valid_prefix(&"a".repeat(65)));
128    }
129
130    #[test]
131    fn attribution_encoding_is_fixed_length_and_validated() {
132        for kind in [
133            TaskOriginKind::ToolsCall,
134            TaskOriginKind::ResourcesRead,
135            TaskOriginKind::PromptsGet,
136            TaskOriginKind::Other,
137        ] {
138            let attribution = TaskAttribution::new(kind, u64::MAX);
139            let encoded = encode_attribution(attribution);
140            assert_eq!(encoded.len(), 10);
141            assert_eq!(decode_attribution(&encoded).unwrap(), attribution);
142        }
143        assert_eq!(
144            decode_attribution(&[1, 2]),
145            Err(TaskStoreError::InvalidRecord)
146        );
147        assert_eq!(
148            decode_attribution(&[2, 1, 0, 0, 0, 0, 0, 0, 0, 1]),
149            Err(TaskStoreError::InvalidRecord)
150        );
151    }
152}
mcp_usage_store/lib.rs

mcp_usage_store/
lib.rs
