Skip to main content

vtcode_core/prompts/
system_prompt_cache.rs

1use lru::LruCache;
2use parking_lot::RwLock;
3use std::collections::hash_map::DefaultHasher;
4use std::hash::{Hash, Hasher};
5use std::num::NonZeroUsize;
6use std::sync::LazyLock;
7
8use crate::prompts::system::SystemPromptReport;
9
10/// Maximum cache size per shard. With N shards the total capacity is
11/// N * MAX_SHARD_SIZE, but entries are distributed by key so the effective
12/// capacity is approximately MAX_SHARD_SIZE per project.
13const MAX_SHARD_SIZE: usize = 32;
14const NUM_SHARDS: usize = 16;
15const SHARD_MASK: usize = NUM_SHARDS - 1;
16
17/// Sharded in-memory prompt cache.
18///
19/// Splits the cache across N `RwLock<LruCache>` shards so that reads on
20/// different shards proceed concurrently without contending on a single
21/// global mutex. Shard selection is by Fx hash of the key.
22///
23/// The fast path uses `LruCache::peek()` under a read-lock — no LRU
24/// promotion on reads, but fully concurrent. The slow path (insert) uses
25/// a write-lock on the affected shard only.
26pub struct SystemPromptCache<V: Clone> {
27    shards: [RwLock<LruCache<String, V>>; NUM_SHARDS],
28}
29
30impl<V: Clone> Default for SystemPromptCache<V> {
31    fn default() -> Self {
32        Self::new()
33    }
34}
35
36impl<V: Clone> SystemPromptCache<V> {
37    pub fn new() -> Self {
38        let shard_size = NonZeroUsize::new(MAX_SHARD_SIZE).unwrap_or(NonZeroUsize::MIN);
39        let shard = || RwLock::new(LruCache::new(shard_size));
40        Self { shards: [(); NUM_SHARDS].map(|_| shard()) }
41    }
42
43    #[inline]
44    fn shard_index(key: &str) -> usize {
45        let mut hasher = DefaultHasher::new();
46        key.hash(&mut hasher);
47        (hasher.finish() as usize) & SHARD_MASK
48    }
49
50    /// Get cached value, returning None on miss.
51    pub fn get(&self, key: &str) -> Option<V> {
52        let shard = self.shards[Self::shard_index(key)].read();
53        shard.peek(key).cloned()
54    }
55
56    /// Insert a value into the cache.
57    pub fn insert(&self, key: String, value: V) {
58        let idx = Self::shard_index(&key);
59        let mut shard = self.shards[idx].write();
60        shard.put(key, value);
61    }
62}
63
64/// Global prompt cache shared across runs. Caches the composed prompt string
65/// together with its [`SystemPromptReport`] so cache hits surface the same
66/// token-budget report a cache miss would have computed.
67///
68/// Sharded into 16 `RwLock<LruCache>` shards to eliminate mutex contention
69/// in multi-project workflows.
70pub static PROMPT_CACHE: LazyLock<SystemPromptCache<(String, SystemPromptReport)>> =
71    LazyLock::new(SystemPromptCache::new);