netget/llm/
conversation_state.rs

1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3use std::collections::{HashSet, VecDeque};
4
5/// Manages conversation history with token-based size limits
6#[derive(Debug, Clone)]
7pub struct ConversationState {
8    /// Unique conversation ID
9    pub conversation_id: String,
10
11    /// Conversation messages (limited by token count)
12    pub messages: VecDeque<ConversationMessage>,
13
14    /// Maximum token/character size for history
15    pub max_token_size: usize,
16
17    /// Current total size in characters
18    pub current_size: usize,
19
20    /// Flag indicating if older messages were removed
21    pub truncated: bool,
22
23    /// Conversation metadata
24    pub started_at: DateTime<Utc>,
25    pub last_interaction: DateTime<Utc>,
26
27    /// Server protocols that have been documented in this conversation
28    /// (enables open_server for these protocols)
29    pub documented_server_protocols: HashSet<String>,
30
31    /// Client protocols that have been documented in this conversation
32    /// (enables open_client for these protocols)
33    pub documented_client_protocols: HashSet<String>,
34}
35
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct ConversationMessage {
38    pub timestamp: DateTime<Utc>,
39    pub role: MessageRole,
40    pub content: String,
41    pub message_type: MessageType,
42}
43
44#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
45pub enum MessageRole {
46    User,
47    Assistant,
48    System,
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize)]
52pub enum MessageType {
53    /// User input
54    UserInput(String),
55
56    /// LLM response (action JSON or raw if invalid)
57    LLMResponse {
58        action_json: Option<serde_json::Value>,
59        raw_output: String,
60    },
61
62    /// Retry instruction (e.g., "Invalid JSON, please retry")
63    RetryInstruction(String),
64
65    /// Tool call reference (without content)
66    ToolCall {
67        tool_name: String,
68        description: String, // Brief description, not full content
69    },
70}
71
72impl ConversationState {
73    /// Cap on a single stored message, as a fraction of `max_token_size`.
74    ///
75    /// `add_message` evicts from the front until the new message fits, so a single message
76    /// larger than the whole window emptied the history and was then stored in full anyway —
77    /// the window bounded the count of old messages but not the size of a new one. Capping
78    /// each message at half the window guarantees at least two messages of history survive
79    /// and that `current_size` can never exceed `max_token_size`.
80    const MESSAGE_FRACTION: usize = 2;
81    const MAX_MESSAGES: usize = 1024;
82    pub const MAX_DOCUMENTED_PROTOCOLS: usize = 1024;
83
84    /// Truncate a message payload to the per-message cap, char-safe, marking the cut so the
85    /// model is not shown a silently amputated value.
86    fn bound_message(&self, content: String) -> String {
87        let cap = (self.max_token_size / Self::MESSAGE_FRACTION)
88            .max(1)
89            .min(self.max_token_size);
90        if content.len() <= cap {
91            return content;
92        }
93        let marker = crate::utils::truncate::TRUNCATION_MARKER;
94        if cap < marker.len() {
95            return crate::utils::truncate_str(&content, cap).to_string();
96        }
97        crate::utils::truncate_with_suffix(&content, cap - marker.len(), marker)
98    }
99
100    /// Create a new conversation state with token size limit
101    pub fn new(max_token_size: usize) -> Self {
102        let now = Utc::now();
103        Self {
104            conversation_id: uuid::Uuid::new_v4().to_string(),
105            messages: VecDeque::new(),
106            max_token_size,
107            current_size: 0,
108            truncated: false,
109            started_at: now,
110            last_interaction: now,
111            documented_server_protocols: HashSet::new(),
112            documented_client_protocols: HashSet::new(),
113        }
114    }
115
116    /// Add a user input message
117    pub fn add_user_input(&mut self, input: String) {
118        let input = self.bound_message(input);
119        let message = ConversationMessage {
120            timestamp: Utc::now(),
121            role: MessageRole::User,
122            content: input.clone(),
123            message_type: MessageType::UserInput(input),
124        };
125        self.add_message(message);
126    }
127
128    /// Add an LLM response message
129    pub fn add_llm_response(&mut self, response: String, parsed_action: Option<serde_json::Value>) {
130        let response = self.bound_message(response);
131        // Parsed metadata is optional context, never a second unlimited history.
132        // Drop rejected structures iteratively so a programmatically built deeply
133        // nested Value cannot overflow the stack even during rejection.
134        let parsed_action = parsed_action.and_then(|value| {
135            if crate::utils::json_budget::within_budget(&value, self.max_token_size / 2, 4096, 64) {
136                Some(value)
137            } else {
138                crate::utils::json_budget::drop_iteratively(value);
139                None
140            }
141        });
142        let message = ConversationMessage {
143            timestamp: Utc::now(),
144            role: MessageRole::Assistant,
145            content: response.clone(),
146            message_type: MessageType::LLMResponse {
147                action_json: parsed_action,
148                raw_output: response,
149            },
150        };
151        self.add_message(message);
152    }
153
154    /// Add a retry instruction message
155    pub fn add_retry_instruction(&mut self, instruction: String) {
156        let instruction = self.bound_message(instruction);
157        let message = ConversationMessage {
158            timestamp: Utc::now(),
159            role: MessageRole::System,
160            content: instruction.clone(),
161            message_type: MessageType::RetryInstruction(instruction),
162        };
163        self.add_message(message);
164    }
165
166    /// Add a tool call reference
167    pub fn add_tool_call(&mut self, tool_name: String, brief_description: String) {
168        let tool_name =
169            crate::utils::truncate_str(&tool_name, 128.min(self.max_token_size)).to_owned();
170        let brief_description = self.bound_message(brief_description);
171        let content =
172            self.bound_message(format!("Tool Call - {} ({})", tool_name, brief_description));
173        let message = ConversationMessage {
174            timestamp: Utc::now(),
175            role: MessageRole::System,
176            content: content.clone(),
177            message_type: MessageType::ToolCall {
178                tool_name,
179                description: brief_description,
180            },
181        };
182        self.add_message(message);
183    }
184
185    /// Add a message and manage size limits
186    fn add_message(&mut self, message: ConversationMessage) {
187        let message_size = message.content.len();
188
189        // Empty entries consume no byte budget and would otherwise accumulate forever.
190        // A zero-sized history disables storage rather than retaining one oversized entry.
191        if message_size == 0 {
192            return;
193        }
194
195        // Remove oldest messages if needed to stay under token limit
196        while (self.current_size.saturating_add(message_size) > self.max_token_size
197            || self.messages.len() >= Self::MAX_MESSAGES)
198            && !self.messages.is_empty()
199        {
200            if let Some(removed) = self.messages.pop_front() {
201                self.current_size = self.current_size.saturating_sub(removed.content.len());
202                self.truncated = true;
203            }
204        }
205
206        // Add the new message
207        self.current_size += message_size;
208        self.messages.push_back(message);
209        self.last_interaction = Utc::now();
210    }
211
212    /// Get formatted history for inclusion in prompts
213    pub fn get_history_for_prompt(&self) -> String {
214        let mut history = String::new();
215
216        // Add truncation notice if needed
217        if self.truncated {
218            history.push_str("[Note: Earlier messages were removed due to size limits]\n");
219        }
220
221        // Include documented protocols at the start of history
222        // This ensures the LLM knows which protocols have been documented
223        if !self.documented_server_protocols.is_empty()
224            || !self.documented_client_protocols.is_empty()
225        {
226            history.push_str("<documented_protocols>\n");
227            if !self.documented_server_protocols.is_empty() {
228                let mut protocols: Vec<_> = self.documented_server_protocols.iter().collect();
229                protocols.sort();
230                history.push_str(&format!(
231                    "Server protocols documented: {}\n",
232                    protocols
233                        .into_iter()
234                        .cloned()
235                        .collect::<Vec<_>>()
236                        .join(", ")
237                ));
238            }
239            if !self.documented_client_protocols.is_empty() {
240                let mut protocols: Vec<_> = self.documented_client_protocols.iter().collect();
241                protocols.sort();
242                history.push_str(&format!(
243                    "Client protocols documented: {}\n",
244                    protocols
245                        .into_iter()
246                        .cloned()
247                        .collect::<Vec<_>>()
248                        .join(", ")
249                ));
250            }
251            history.push_str("</documented_protocols>\n");
252        }
253
254        // Format each message with appropriate tags
255        for message in &self.messages {
256            match &message.message_type {
257                MessageType::UserInput(input) => {
258                    history.push_str(&format!("<user>{}</user>\n", input));
259                }
260                MessageType::LLMResponse {
261                    action_json,
262                    raw_output,
263                } => {
264                    if action_json.is_some() {
265                        // Valid JSON response
266                        history.push_str(&format!("<assistant>{}</assistant>\n", raw_output));
267                    } else {
268                        // Invalid JSON, show raw output
269                        history.push_str(&format!("<assistant>{}</assistant>\n", raw_output));
270                    }
271                }
272                MessageType::RetryInstruction(instruction) => {
273                    history.push_str(&format!("<system>Retry - {}</system>\n", instruction));
274                }
275                MessageType::ToolCall {
276                    tool_name,
277                    description,
278                } => {
279                    history.push_str(&format!(
280                        "<system>Tool Call - {} ({})</system>\n",
281                        tool_name, description
282                    ));
283                }
284            }
285        }
286
287        history
288    }
289
290    /// Mark server protocols as documented in this conversation
291    pub fn mark_server_protocols_documented(&mut self, protocols: &[String]) {
292        for protocol in protocols {
293            if Self::valid_documented_protocol(protocol)
294                && self.documented_server_protocols.len() < Self::MAX_DOCUMENTED_PROTOCOLS
295            {
296                self.documented_server_protocols.insert(protocol.clone());
297            }
298        }
299    }
300
301    /// Mark client protocols as documented in this conversation
302    pub fn mark_client_protocols_documented(&mut self, protocols: &[String]) {
303        for protocol in protocols {
304            if Self::valid_documented_protocol(protocol)
305                && self.documented_client_protocols.len() < Self::MAX_DOCUMENTED_PROTOCOLS
306            {
307                self.documented_client_protocols.insert(protocol.clone());
308            }
309        }
310    }
311
312    /// Protocol identifiers carry neither markup nor arbitrary document content.
313    pub fn valid_documented_protocol(protocol: &str) -> bool {
314        !protocol.is_empty()
315            && protocol.len() <= 128
316            && protocol
317                .bytes()
318                .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-')
319    }
320
321    /// Check if any server protocols have been documented
322    pub fn has_server_docs(&self) -> bool {
323        !self.documented_server_protocols.is_empty()
324    }
325
326    /// Check if any client protocols have been documented
327    pub fn has_client_docs(&self) -> bool {
328        !self.documented_client_protocols.is_empty()
329    }
330
331    /// Get all documented server protocols
332    pub fn get_documented_server_protocols(&self) -> &HashSet<String> {
333        &self.documented_server_protocols
334    }
335
336    /// Get all documented client protocols
337    pub fn get_documented_client_protocols(&self) -> &HashSet<String> {
338        &self.documented_client_protocols
339    }
340
341    /// Clear all conversation history
342    pub fn clear_history(&mut self) {
343        self.messages.clear();
344        self.current_size = 0;
345        self.truncated = false;
346        self.last_interaction = Utc::now();
347    }
348
349    /// Get the current number of messages
350    pub fn message_count(&self) -> usize {
351        self.messages.len()
352    }
353
354    /// Check if history is empty
355    pub fn is_empty(&self) -> bool {
356        self.messages.is_empty()
357    }
358}
netget/llm/conversation_state.rs

netget/llm/
conversation_state.rs
