codex_convert_proxy/proxy/
context_store.rs1use std::collections::{HashMap, VecDeque};
4use std::sync::Mutex;
5
6use crate::types::chat_api::ChatMessage;
7
8#[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#[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 assert!(store.get("resp_0").is_some());
105 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}