Skip to main content

vtcode_core/prompts/
system_prompt_cache.rs

1use lru::LruCache;
2use std::num::NonZeroUsize;
3use std::sync::{LazyLock, Mutex};
4
5/// Maximum cache size (increased from 32 to 128 for multi-project workflows)
6const MAX_CACHE_SIZE: usize = 128;
7
8/// Task categories for prompt generation
9#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
10pub enum TaskType {
11    System,
12    Lightweight,
13    Specialized,
14}
15
16/// Providers can expose a cache key describing the prompt variant.
17pub trait PromptProvider {
18    fn cache_key(&self) -> String;
19    fn task_type(&self) -> TaskType;
20}
21
22/// Simple in-memory prompt cache keyed by provider + task type.
23/// Uses LRU eviction to manage memory in multi-project workflows.
24pub struct SystemPromptCache {
25    entries: Mutex<LruCache<String, String>>,
26}
27
28impl Default for SystemPromptCache {
29    fn default() -> Self {
30        Self::new()
31    }
32}
33
34impl SystemPromptCache {
35    pub fn new() -> Self {
36        let cache_size = NonZeroUsize::new(MAX_CACHE_SIZE).unwrap_or(NonZeroUsize::MIN);
37        Self {
38            entries: Mutex::new(LruCache::new(cache_size)),
39        }
40    }
41
42    /// Get cached prompt or build it if missing.
43    /// Uses LRU eviction when cache is full (no more clear-on-overflow).
44    pub fn get_or_insert_with<F>(&self, key: &str, builder: F) -> String
45    where
46        F: FnOnce() -> String,
47    {
48        // Check if already cached (LruCache.get() requires mutable access for LRU update)
49        {
50            let mut store = self.entries.lock().unwrap_or_else(|p| p.into_inner());
51            if let Some(value) = store.get(key) {
52                return value.clone();
53            }
54        }
55
56        // Not cached - build value
57        let value = builder();
58
59        // Insert into cache
60        {
61            let mut store = self.entries.lock().unwrap_or_else(|p| p.into_inner());
62            store.put(key.to_string(), value.clone());
63            // LRU cache automatically evicts least recently used entry when full
64        }
65
66        value
67    }
68
69    /// Get cached value, returning None on miss (for async callers that want to build outside the lock)
70    pub fn get(&self, key: &str) -> Option<String> {
71        let mut store = self.entries.lock().unwrap_or_else(|p| p.into_inner());
72        store.get(key).cloned()
73    }
74
75    /// Insert a value into the cache
76    pub fn insert(&self, key: String, value: String) {
77        let mut store = self.entries.lock().unwrap_or_else(|p| p.into_inner());
78        store.put(key, value);
79        // LRU cache automatically evicts least recently used entry when full
80    }
81
82    pub fn clear(&self) {
83        if let Ok(mut store) = self.entries.lock() {
84            store.clear();
85        }
86    }
87}
88
89/// Global prompt cache shared across runs.
90pub static PROMPT_CACHE: LazyLock<SystemPromptCache> = LazyLock::new(SystemPromptCache::new);