mcp_usage_store/
redis.rs

1use std::time::Duration;
2
3use mcp_usage_core::TaskAttribution;
4use mcp_usage_tower::{TaskAttributionStore, TaskStoreError, TaskStoreFuture};
5
6use crate::{StoreConfigError, decode_attribution, encode_attribution, encode_hash, valid_prefix};
7
8const CLAIM_SCRIPT: &str = "local value=redis.call('GET',KEYS[1]);if value then redis.call('DEL',KEYS[1]);end;return value";
9const DEFAULT_OPERATION_TIMEOUT: Duration = Duration::from_secs(2);
10
11/// Redis or Valkey durable-task attribution with atomic, once-only claims.
12#[derive(Clone)]
13pub struct RedisTaskStore {
14    connection: redis::aio::ConnectionManager,
15    key_prefix: String,
16    ttl_seconds: i64,
17}
18
19impl std::fmt::Debug for RedisTaskStore {
20    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
21        formatter
22            .debug_struct("RedisTaskStore")
23            .field("key_prefix", &self.key_prefix)
24            .field("ttl_seconds", &self.ttl_seconds)
25            .finish_non_exhaustive()
26    }
27}
28
29impl RedisTaskStore {
30    /// Connect with a reconnecting, multiplexed async connection manager.
31    ///
32    /// # Errors
33    ///
34    /// Returns a sanitized error for invalid configuration or connection
35    /// failure. Connection URLs are never retained for debug output.
36    pub async fn connect(
37        connection_url: &str,
38        key_prefix: impl Into<String>,
39        ttl: Duration,
40    ) -> Result<Self, StoreConfigError> {
41        Self::connect_with_timeout(connection_url, key_prefix, ttl, DEFAULT_OPERATION_TIMEOUT).await
42    }
43
44    /// Connect with an explicit bound for connection and command operations.
45    ///
46    /// # Errors
47    ///
48    /// Returns a sanitized error for invalid configuration or connection
49    /// failure. Connection URLs are never retained for debug output.
50    pub async fn connect_with_timeout(
51        connection_url: &str,
52        key_prefix: impl Into<String>,
53        ttl: Duration,
54        operation_timeout: Duration,
55    ) -> Result<Self, StoreConfigError> {
56        let key_prefix = key_prefix.into();
57        if !valid_prefix(&key_prefix) {
58            return Err(StoreConfigError::InvalidPrefix);
59        }
60        let ttl_seconds = i64::try_from(ttl.as_secs()).map_err(|_| StoreConfigError::InvalidTtl)?;
61        if ttl_seconds == 0 {
62            return Err(StoreConfigError::InvalidTtl);
63        }
64        if operation_timeout.is_zero() {
65            return Err(StoreConfigError::InvalidTimeout);
66        }
67        let client =
68            redis::Client::open(connection_url).map_err(|_| StoreConfigError::Connection)?;
69        let config = redis::aio::ConnectionManagerConfig::new()
70            .set_connection_timeout(Some(operation_timeout))
71            .set_response_timeout(Some(operation_timeout));
72        let connection = client
73            .get_connection_manager_with_config(config)
74            .await
75            .map_err(|_| StoreConfigError::Connection)?;
76        Ok(Self {
77            connection,
78            key_prefix,
79            ttl_seconds,
80        })
81    }
82
83    fn key(&self, tenant_id: &str, task_id: &str) -> String {
84        format!(
85            "{}:{}:{}",
86            self.key_prefix,
87            encode_hash(tenant_id),
88            encode_hash(task_id)
89        )
90    }
91}
92
93impl TaskAttributionStore for RedisTaskStore {
94    fn insert<'a>(
95        &'a self,
96        tenant_id: &'a str,
97        task_id: &'a str,
98        attribution: TaskAttribution,
99    ) -> TaskStoreFuture<'a, ()> {
100        Box::pin(async move {
101            let mut connection = self.connection.clone();
102            let _: Option<String> = redis::cmd("SET")
103                .arg(self.key(tenant_id, task_id))
104                .arg(&encode_attribution(attribution))
105                .arg("NX")
106                .arg("EX")
107                .arg(self.ttl_seconds)
108                .query_async(&mut connection)
109                .await
110                .map_err(|_| TaskStoreError::BackendUnavailable)?;
111            Ok(())
112        })
113    }
114
115    fn get<'a>(
116        &'a self,
117        tenant_id: &'a str,
118        task_id: &'a str,
119    ) -> TaskStoreFuture<'a, Option<TaskAttribution>> {
120        Box::pin(async move {
121            let mut connection = self.connection.clone();
122            let payload: Option<Vec<u8>> = redis::cmd("GET")
123                .arg(self.key(tenant_id, task_id))
124                .query_async(&mut connection)
125                .await
126                .map_err(|_| TaskStoreError::BackendUnavailable)?;
127            payload.as_deref().map(decode_attribution).transpose()
128        })
129    }
130
131    fn claim<'a>(
132        &'a self,
133        tenant_id: &'a str,
134        task_id: &'a str,
135    ) -> TaskStoreFuture<'a, Option<TaskAttribution>> {
136        Box::pin(async move {
137            let mut connection = self.connection.clone();
138            let payload: Option<Vec<u8>> = redis::cmd("EVAL")
139                .arg(CLAIM_SCRIPT)
140                .arg(1)
141                .arg(self.key(tenant_id, task_id))
142                .query_async(&mut connection)
143                .await
144                .map_err(|_| TaskStoreError::BackendUnavailable)?;
145            payload.as_deref().map(decode_attribution).transpose()
146        })
147    }
148
149    fn remove<'a>(&'a self, tenant_id: &'a str, task_id: &'a str) -> TaskStoreFuture<'a, ()> {
150        Box::pin(async move {
151            let mut connection = self.connection.clone();
152            let _: u64 = redis::cmd("DEL")
153                .arg(self.key(tenant_id, task_id))
154                .query_async(&mut connection)
155                .await
156                .map_err(|_| TaskStoreError::BackendUnavailable)?;
157            Ok(())
158        })
159    }
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165    use mcp_usage_core::TaskOriginKind;
166
167    #[test]
168    fn payload_contains_only_fixed_category_and_resolved_units() {
169        let attribution = TaskAttribution::new(TaskOriginKind::Other, 42);
170        let payload = encode_attribution(attribution);
171        assert_eq!(payload.len(), 10);
172        assert_eq!(decode_attribution(&payload).unwrap(), attribution);
173    }
174
175    #[tokio::test]
176    async fn rejects_zero_operation_timeout_before_connecting() {
177        let error = RedisTaskStore::connect_with_timeout(
178            "redis://127.0.0.1/",
179            "usagekit:tasks",
180            Duration::from_secs(60),
181            Duration::ZERO,
182        )
183        .await
184        .unwrap_err();
185
186        assert_eq!(error, StoreConfigError::InvalidTimeout);
187    }
188
189    #[tokio::test]
190    async fn rejects_ttl_larger_than_redis_signed_integer_before_connecting() {
191        let oversized = u64::try_from(i64::MAX).unwrap() + 1;
192        let error = RedisTaskStore::connect(
193            "redis://127.0.0.1:1/",
194            "usagekit:tasks",
195            Duration::from_secs(oversized),
196        )
197        .await
198        .unwrap_err();
199
200        assert_eq!(error, StoreConfigError::InvalidTtl);
201    }
202}
mcp_usage_store/redis.rs

mcp_usage_store/
redis.rs
