context_creator/mcp_server/
cache.rs

1//! Cache implementation for MCP server
2
3use moka::future::Cache;
4use serde::{Deserialize, Serialize};
5use std::hash::Hash;
6use std::path::PathBuf;
7use std::time::Duration;
8
9/// Cache key for process_local_codebase requests
10#[derive(Debug, Clone, Hash, PartialEq, Eq)]
11pub struct ProcessLocalCacheKey {
12    pub prompt: String,
13    pub path: PathBuf,
14    pub include_patterns: Vec<String>,
15    pub ignore_patterns: Vec<String>,
16    pub include_imports: bool,
17    pub max_tokens: Option<u32>,
18    pub llm_tool: Option<String>,
19}
20
21impl ProcessLocalCacheKey {
22    pub fn from_request(request: &super::ProcessLocalRequest) -> Self {
23        Self {
24            prompt: request.prompt.clone(),
25            path: request.path.clone(),
26            include_patterns: request.include_patterns.clone(),
27            ignore_patterns: request.ignore_patterns.clone(),
28            include_imports: request.include_imports,
29            max_tokens: request.max_tokens,
30            llm_tool: request.llm_tool.clone(),
31        }
32    }
33}
34
35/// Cached response for process_local_codebase
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct ProcessLocalCacheValue {
38    pub answer: String,
39    pub markdown: String,
40    pub file_count: usize,
41    pub token_count: usize,
42    pub llm_tool: String,
43}
44
45/// Cache key for process_remote_repo requests
46#[derive(Debug, Clone, Hash, PartialEq, Eq)]
47pub struct ProcessRemoteCacheKey {
48    pub prompt: String,
49    pub repo_url: String,
50    pub include_patterns: Vec<String>,
51    pub ignore_patterns: Vec<String>,
52    pub include_imports: bool,
53    pub max_tokens: Option<u32>,
54    pub llm_tool: Option<String>,
55}
56
57impl ProcessRemoteCacheKey {
58    pub fn from_request(request: &super::ProcessRemoteRequest) -> Self {
59        Self {
60            prompt: request.prompt.clone(),
61            repo_url: request.repo_url.clone(),
62            include_patterns: request.include_patterns.clone(),
63            ignore_patterns: request.ignore_patterns.clone(),
64            include_imports: request.include_imports,
65            max_tokens: request.max_tokens,
66            llm_tool: request.llm_tool.clone(),
67        }
68    }
69}
70
71/// Cached response for process_remote_repo
72#[derive(Debug, Clone, Serialize, Deserialize)]
73pub struct ProcessRemoteCacheValue {
74    pub answer: String,
75    pub markdown: String,
76    pub file_count: usize,
77    pub token_count: usize,
78    pub repo_name: String,
79    pub llm_tool: String,
80}
81
82/// MCP server cache
83#[derive(Debug)]
84pub struct McpCache {
85    process_local_cache: Cache<ProcessLocalCacheKey, ProcessLocalCacheValue>,
86    process_remote_cache: Cache<ProcessRemoteCacheKey, ProcessRemoteCacheValue>,
87}
88
89impl McpCache {
90    /// Create a new cache with default configuration
91    pub fn new() -> Self {
92        let process_local_cache = Cache::builder()
93            .max_capacity(100) // Cache up to 100 requests
94            .time_to_live(Duration::from_secs(300)) // 5 minutes TTL
95            .build();
96
97        let process_remote_cache = Cache::builder()
98            .max_capacity(50) // Cache up to 50 remote requests
99            .time_to_live(Duration::from_secs(600)) // 10 minutes TTL
100            .build();
101
102        Self {
103            process_local_cache,
104            process_remote_cache,
105        }
106    }
107
108    /// Get a cached process_local_codebase response
109    pub async fn get_process_local(
110        &self,
111        key: &ProcessLocalCacheKey,
112    ) -> Option<ProcessLocalCacheValue> {
113        self.process_local_cache.get(key).await
114    }
115
116    /// Cache a process_local_codebase response
117    pub async fn set_process_local(
118        &self,
119        key: ProcessLocalCacheKey,
120        value: ProcessLocalCacheValue,
121    ) {
122        self.process_local_cache.insert(key, value).await;
123    }
124
125    /// Get a cached process_remote_repo response
126    pub async fn get_process_remote(
127        &self,
128        key: &ProcessRemoteCacheKey,
129    ) -> Option<ProcessRemoteCacheValue> {
130        self.process_remote_cache.get(key).await
131    }
132
133    /// Cache a process_remote_repo response
134    pub async fn set_process_remote(
135        &self,
136        key: ProcessRemoteCacheKey,
137        value: ProcessRemoteCacheValue,
138    ) {
139        self.process_remote_cache.insert(key, value).await;
140    }
141
142    /// Clear all caches
143    pub fn clear_all(&self) {
144        self.process_local_cache.invalidate_all();
145        self.process_remote_cache.invalidate_all();
146    }
147}
148
149impl Default for McpCache {
150    fn default() -> Self {
151        Self::new()
152    }
153}
context_creator/mcp_server/cache.rs

context_creator/mcp_server/
cache.rs
