use crate::types::MCPTool;
use chrono::{DateTime, Utc};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::debug;
#[derive(Debug, Clone)]
#[allow(dead_code)] pub struct CacheEntry {
pub tools: Vec<MCPTool>,
pub cached_at: DateTime<Utc>,
pub ttl_seconds: u64,
}
#[allow(dead_code)] impl CacheEntry {
pub fn new(tools: Vec<MCPTool>, ttl_seconds: u64) -> Self {
Self {
tools,
cached_at: Utc::now(),
ttl_seconds,
}
}
pub fn is_expired(&self) -> bool {
let now = Utc::now();
let expiry_time = self.cached_at + chrono::Duration::seconds(self.ttl_seconds as i64);
now > expiry_time
}
pub fn remaining_ttl(&self) -> i64 {
let now = Utc::now();
let expiry_time = self.cached_at + chrono::Duration::seconds(self.ttl_seconds as i64);
(expiry_time - now).num_seconds()
}
}
#[derive(Debug, Clone)]
#[allow(dead_code)] pub struct ToolCache {
cache: Arc<RwLock<HashMap<String, CacheEntry>>>,
default_ttl: u64,
}
#[allow(dead_code)] impl ToolCache {
pub fn new(default_ttl: u64) -> Self {
Self {
cache: Arc::new(RwLock::new(HashMap::new())),
default_ttl,
}
}
pub async fn get(&self, url: &str) -> Option<Vec<MCPTool>> {
let cache = self.cache.read().await;
if let Some(entry) = cache.get(url) {
if !entry.is_expired() {
debug!(
"Cache hit for {}: {} tools, TTL remaining: {}s",
url,
entry.tools.len(),
entry.remaining_ttl()
);
return Some(entry.tools.clone());
} else {
debug!("Cache entry for {} has expired", url);
}
}
None
}
pub async fn put(&self, url: String, tools: Vec<MCPTool>) {
self.put_with_ttl(url, tools, self.default_ttl).await;
}
pub async fn put_with_ttl(&self, url: String, tools: Vec<MCPTool>, ttl_seconds: u64) {
let entry = CacheEntry::new(tools.clone(), ttl_seconds);
let mut cache = self.cache.write().await;
cache.insert(url.clone(), entry);
debug!(
"Cached {} tools for {} with TTL {}s",
tools.len(),
url,
ttl_seconds
);
}
pub async fn remove(&self, url: &str) -> bool {
let mut cache = self.cache.write().await;
let removed = cache.remove(url).is_some();
if removed {
debug!("Removed cache entry for {}", url);
}
removed
}
pub async fn clear(&self) {
let mut cache = self.cache.write().await;
let count = cache.len();
cache.clear();
debug!("Cleared {} cache entries", count);
}
pub async fn cleanup_expired(&self) -> usize {
let mut cache = self.cache.write().await;
let initial_count = cache.len();
cache.retain(|url, entry| {
if entry.is_expired() {
debug!("Removing expired cache entry for {}", url);
false
} else {
true
}
});
let removed_count = initial_count - cache.len();
if removed_count > 0 {
debug!("Cleaned up {} expired cache entries", removed_count);
}
removed_count
}
pub async fn stats(&self) -> CacheStats {
let cache = self.cache.read().await;
let total_entries = cache.len();
let mut expired_entries = 0;
let mut total_tools = 0;
for entry in cache.values() {
if entry.is_expired() {
expired_entries += 1;
}
total_tools += entry.tools.len();
}
CacheStats {
total_entries,
expired_entries,
active_entries: total_entries - expired_entries,
total_tools,
default_ttl: self.default_ttl,
}
}
pub async fn contains(&self, url: &str) -> bool {
let cache = self.cache.read().await;
cache.contains_key(url)
}
pub async fn get_cached_urls(&self) -> Vec<String> {
let cache = self.cache.read().await;
cache.keys().cloned().collect()
}
}
#[derive(Debug, Clone)]
#[allow(dead_code)] pub struct CacheStats {
pub total_entries: usize,
pub expired_entries: usize,
pub active_entries: usize,
pub total_tools: usize,
pub default_ttl: u64,
}
impl Default for ToolCache {
fn default() -> Self {
Self::new(3600)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::MCPTool;
use std::collections::HashMap;
use tokio::time::{sleep, Duration};
fn create_test_tool(name: &str) -> MCPTool {
MCPTool {
name: name.to_string(),
description: Some(format!("Test tool {name}")),
input_schema: None,
output_schema: None,
parameters: HashMap::new(),
category: None,
tags: vec![],
deprecated: false,
raw_json: None,
}
}
#[tokio::test]
async fn test_cache_basic_operations() {
let cache = ToolCache::new(60); let url = "http://example.com";
let tools = vec![create_test_tool("test1"), create_test_tool("test2")];
assert!(cache.get(url).await.is_none());
cache.put(url.to_string(), tools.clone()).await;
let cached_tools = cache.get(url).await.unwrap();
assert_eq!(cached_tools.len(), 2);
assert_eq!(cached_tools[0].name, "test1");
assert!(cache.remove(url).await);
assert!(cache.get(url).await.is_none());
}
#[tokio::test]
async fn test_cache_expiry() {
let cache = ToolCache::new(1); let url = "http://example.com";
let tools = vec![create_test_tool("test")];
cache.put(url.to_string(), tools).await;
assert!(cache.get(url).await.is_some());
sleep(Duration::from_secs(2)).await;
assert!(cache.get(url).await.is_none());
}
#[tokio::test]
async fn test_cache_stats() {
let cache = ToolCache::new(60);
let tools = vec![create_test_tool("test")];
cache.put("url1".to_string(), tools.clone()).await;
cache.put("url2".to_string(), tools).await;
let stats = cache.stats().await;
assert_eq!(stats.total_entries, 2);
assert_eq!(stats.active_entries, 2);
assert_eq!(stats.expired_entries, 0);
assert_eq!(stats.total_tools, 2);
}
}