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    /// `ttl` must outlive the longest task the upstream server can run, not
33    /// the client's polling interval. A task that outlives it bills zero. See
34    /// the [crate-level notes](crate#choosing-a-ttl).
35    ///
36    /// # Errors
37    ///
38    /// Returns a sanitized error for invalid configuration or connection
39    /// failure. Connection URLs are never retained for debug output.
40    pub async fn connect(
41        connection_url: &str,
42        key_prefix: impl Into<String>,
43        ttl: Duration,
44    ) -> Result<Self, StoreConfigError> {
45        Self::connect_with_timeout(connection_url, key_prefix, ttl, DEFAULT_OPERATION_TIMEOUT).await
46    }
47
48    /// Connect with an explicit bound for connection and command operations.
49    ///
50    /// # Errors
51    ///
52    /// Returns a sanitized error for invalid configuration or connection
53    /// failure. Connection URLs are never retained for debug output.
54    pub async fn connect_with_timeout(
55        connection_url: &str,
56        key_prefix: impl Into<String>,
57        ttl: Duration,
58        operation_timeout: Duration,
59    ) -> Result<Self, StoreConfigError> {
60        let key_prefix = key_prefix.into();
61        if !valid_prefix(&key_prefix) {
62            return Err(StoreConfigError::InvalidPrefix);
63        }
64        let ttl_seconds = i64::try_from(ttl.as_secs()).map_err(|_| StoreConfigError::InvalidTtl)?;
65        if ttl_seconds == 0 {
66            return Err(StoreConfigError::InvalidTtl);
67        }
68        if operation_timeout.is_zero() {
69            return Err(StoreConfigError::InvalidTimeout);
70        }
71        let client =
72            redis::Client::open(connection_url).map_err(|_| StoreConfigError::Connection)?;
73        let config = redis::aio::ConnectionManagerConfig::new()
74            .set_connection_timeout(Some(operation_timeout))
75            .set_response_timeout(Some(operation_timeout));
76        let connection = client
77            .get_connection_manager_with_config(config)
78            .await
79            .map_err(|_| StoreConfigError::Connection)?;
80        Ok(Self {
81            connection,
82            key_prefix,
83            ttl_seconds,
84        })
85    }
86
87    fn key(&self, tenant_id: &str, task_id: &str) -> String {
88        format!(
89            "{}:{}:{}",
90            self.key_prefix,
91            encode_hash(tenant_id),
92            encode_hash(task_id)
93        )
94    }
95}
96
97impl TaskAttributionStore for RedisTaskStore {
98    fn insert<'a>(
99        &'a self,
100        tenant_id: &'a str,
101        task_id: &'a str,
102        attribution: TaskAttribution,
103    ) -> TaskStoreFuture<'a, ()> {
104        Box::pin(async move {
105            let mut connection = self.connection.clone();
106            let _: Option<String> = redis::cmd("SET")
107                .arg(self.key(tenant_id, task_id))
108                .arg(&encode_attribution(attribution))
109                .arg("NX")
110                .arg("EX")
111                .arg(self.ttl_seconds)
112                .query_async(&mut connection)
113                .await
114                .map_err(|_| TaskStoreError::BackendUnavailable)?;
115            Ok(())
116        })
117    }
118
119    fn get<'a>(
120        &'a self,
121        tenant_id: &'a str,
122        task_id: &'a str,
123    ) -> TaskStoreFuture<'a, Option<TaskAttribution>> {
124        Box::pin(async move {
125            let mut connection = self.connection.clone();
126            let payload: Option<Vec<u8>> = redis::cmd("GET")
127                .arg(self.key(tenant_id, task_id))
128                .query_async(&mut connection)
129                .await
130                .map_err(|_| TaskStoreError::BackendUnavailable)?;
131            payload.as_deref().map(decode_attribution).transpose()
132        })
133    }
134
135    fn claim<'a>(
136        &'a self,
137        tenant_id: &'a str,
138        task_id: &'a str,
139    ) -> TaskStoreFuture<'a, Option<TaskAttribution>> {
140        Box::pin(async move {
141            let mut connection = self.connection.clone();
142            let payload: Option<Vec<u8>> = redis::cmd("EVAL")
143                .arg(CLAIM_SCRIPT)
144                .arg(1)
145                .arg(self.key(tenant_id, task_id))
146                .query_async(&mut connection)
147                .await
148                .map_err(|_| TaskStoreError::BackendUnavailable)?;
149            payload.as_deref().map(decode_attribution).transpose()
150        })
151    }
152
153    fn remove<'a>(&'a self, tenant_id: &'a str, task_id: &'a str) -> TaskStoreFuture<'a, ()> {
154        Box::pin(async move {
155            let mut connection = self.connection.clone();
156            let _: u64 = redis::cmd("DEL")
157                .arg(self.key(tenant_id, task_id))
158                .query_async(&mut connection)
159                .await
160                .map_err(|_| TaskStoreError::BackendUnavailable)?;
161            Ok(())
162        })
163    }
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169    use mcp_usage_core::TaskOriginKind;
170
171    #[test]
172    fn payload_contains_only_fixed_category_and_resolved_units() {
173        let attribution = TaskAttribution::new(TaskOriginKind::Other, 42);
174        let payload = encode_attribution(attribution);
175        assert_eq!(payload.len(), 10);
176        assert_eq!(decode_attribution(&payload).unwrap(), attribution);
177    }
178
179    #[tokio::test]
180    async fn rejects_zero_operation_timeout_before_connecting() {
181        let error = RedisTaskStore::connect_with_timeout(
182            "redis://127.0.0.1/",
183            "usagekit:tasks",
184            Duration::from_secs(60),
185            Duration::ZERO,
186        )
187        .await
188        .unwrap_err();
189
190        assert_eq!(error, StoreConfigError::InvalidTimeout);
191    }
192
193    #[tokio::test]
194    async fn rejects_ttl_larger_than_redis_signed_integer_before_connecting() {
195        let oversized = u64::try_from(i64::MAX).unwrap() + 1;
196        let error = RedisTaskStore::connect(
197            "redis://127.0.0.1:1/",
198            "usagekit:tasks",
199            Duration::from_secs(oversized),
200        )
201        .await
202        .unwrap_err();
203
204        assert_eq!(error, StoreConfigError::InvalidTtl);
205    }
206}
mcp_usage_store/redis.rs

mcp_usage_store/
redis.rs
