Skip to main content

codex_convert_proxy/proxy/
context_store.rs

1//! In-memory conversation context store for previous_response_id expansion.
2
3use std::collections::{HashMap, VecDeque};
4use std::sync::Mutex;
5
6use crate::types::chat_api::ChatMessage;
7
8/// Conversation snapshot used to reconstruct chat history.
9#[derive(Debug, Clone)]
10pub struct ConversationSnapshot {
11    pub instructions: Option<String>,
12    pub messages: Vec<ChatMessage>,
13}
14
15const MAX_CONVERSATION_ENTRIES: usize = 1024;
16
17#[derive(Debug, Default)]
18struct StoreInner {
19    map: HashMap<String, ConversationSnapshot>,
20    lru_order: VecDeque<String>,
21}
22
23/// Thread-safe in-memory store keyed by response id.
24///
25/// Uses `Mutex` rather than `RwLock`: every `get` mutates the LRU order, so
26/// there is no read-only path that could benefit from shared access.
27#[derive(Debug, Default)]
28pub struct ConversationStore {
29    inner: Mutex<StoreInner>,
30}
31
32impl ConversationStore {
33    pub fn new() -> Self {
34        Self::default()
35    }
36
37    pub fn get(&self, response_id: &str) -> Option<ConversationSnapshot> {
38        let mut guard = self.inner.lock().ok()?;
39        let snapshot = guard.map.get(response_id).cloned()?;
40        if let Some(pos) = guard.lru_order.iter().position(|k| k == response_id) {
41            guard.lru_order.remove(pos);
42        }
43        guard.lru_order.push_back(response_id.to_string());
44        Some(snapshot)
45    }
46
47    pub fn insert(&self, response_id: String, snapshot: ConversationSnapshot) {
48        if let Ok(mut guard) = self.inner.lock() {
49            if let Some(pos) = guard.lru_order.iter().position(|k| k == &response_id) {
50                guard.lru_order.remove(pos);
51            }
52            guard.lru_order.push_back(response_id.clone());
53            guard.map.insert(response_id, snapshot);
54
55            while guard.map.len() > MAX_CONVERSATION_ENTRIES {
56                if let Some(oldest_key) = guard.lru_order.pop_front() {
57                    guard.map.remove(&oldest_key);
58                } else {
59                    break;
60                }
61            }
62        }
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use super::*;
69    use crate::types::chat_api::{ChatMessage, Content, MessageRole};
70
71    fn snapshot(text: &str) -> ConversationSnapshot {
72        ConversationSnapshot {
73            instructions: Some("test".to_string()),
74            messages: vec![ChatMessage {
75                role: MessageRole::User,
76                content: Content::String(text.to_string()),
77                name: None,
78                annotations: None,
79                tool_calls: None,
80                function_call: None,
81                tool_call_id: None,
82                refusal: None,
83            }],
84        }
85    }
86
87    #[test]
88    fn test_lru_eviction_keeps_recent_entries() {
89        let store = ConversationStore::new();
90        for i in 0..=MAX_CONVERSATION_ENTRIES {
91            store.insert(format!("resp_{i}"), snapshot("x"));
92        }
93        assert!(store.get("resp_0").is_none());
94        assert!(store.get(&format!("resp_{}", MAX_CONVERSATION_ENTRIES)).is_some());
95    }
96
97    #[test]
98    fn test_get_refreshes_lru_order() {
99        let store = ConversationStore::new();
100        for i in 0..MAX_CONVERSATION_ENTRIES {
101            store.insert(format!("resp_{i}"), snapshot("x"));
102        }
103        // Touch oldest so it becomes recent.
104        assert!(store.get("resp_0").is_some());
105        // Insert one more, now resp_1 should be evicted first.
106        store.insert("resp_new".to_string(), snapshot("y"));
107        assert!(store.get("resp_0").is_some());
108        assert!(store.get("resp_1").is_none());
109    }
110}