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::RwLock;
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#[derive(Debug, Default)]
25pub struct ConversationStore {
26    inner: RwLock<StoreInner>,
27}
28
29impl ConversationStore {
30    pub fn new() -> Self {
31        Self::default()
32    }
33
34    pub fn get(&self, response_id: &str) -> Option<ConversationSnapshot> {
35        let mut guard = self.inner.write().ok()?;
36        let snapshot = guard.map.get(response_id).cloned()?;
37        if let Some(pos) = guard.lru_order.iter().position(|k| k == response_id) {
38            guard.lru_order.remove(pos);
39        }
40        guard.lru_order.push_back(response_id.to_string());
41        Some(snapshot)
42    }
43
44    pub fn insert(&self, response_id: String, snapshot: ConversationSnapshot) {
45        if let Ok(mut guard) = self.inner.write() {
46            if let Some(pos) = guard.lru_order.iter().position(|k| k == &response_id) {
47                guard.lru_order.remove(pos);
48            }
49            guard.lru_order.push_back(response_id.clone());
50            guard.map.insert(response_id, snapshot);
51
52            while guard.map.len() > MAX_CONVERSATION_ENTRIES {
53                if let Some(oldest_key) = guard.lru_order.pop_front() {
54                    guard.map.remove(&oldest_key);
55                } else {
56                    break;
57                }
58            }
59        }
60    }
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66    use crate::types::chat_api::{ChatMessage, Content, MessageRole};
67
68    fn snapshot(text: &str) -> ConversationSnapshot {
69        ConversationSnapshot {
70            instructions: Some("test".to_string()),
71            messages: vec![ChatMessage {
72                role: MessageRole::User,
73                content: Content::String(text.to_string()),
74                name: None,
75                annotations: None,
76                tool_calls: None,
77                function_call: None,
78                tool_call_id: None,
79                refusal: None,
80            }],
81        }
82    }
83
84    #[test]
85    fn test_lru_eviction_keeps_recent_entries() {
86        let store = ConversationStore::new();
87        for i in 0..=MAX_CONVERSATION_ENTRIES {
88            store.insert(format!("resp_{i}"), snapshot("x"));
89        }
90        assert!(store.get("resp_0").is_none());
91        assert!(store.get(&format!("resp_{}", MAX_CONVERSATION_ENTRIES)).is_some());
92    }
93
94    #[test]
95    fn test_get_refreshes_lru_order() {
96        let store = ConversationStore::new();
97        for i in 0..MAX_CONVERSATION_ENTRIES {
98            store.insert(format!("resp_{i}"), snapshot("x"));
99        }
100        // Touch oldest so it becomes recent.
101        assert!(store.get("resp_0").is_some());
102        // Insert one more, now resp_1 should be evicted first.
103        store.insert("resp_new".to_string(), snapshot("y"));
104        assert!(store.get("resp_0").is_some());
105        assert!(store.get("resp_1").is_none());
106    }
107}
codex_convert_proxy/proxy/context_store.rs

codex_convert_proxy/proxy/
context_store.rs
