llm/
context.rs

1use serde::{Deserialize, Serialize};
2
3use crate::catalog::LlmModel;
4use crate::model_settings::ModelSettings;
5use crate::reasoning::ReasoningEffort;
6use crate::types::IsoString;
7
8use super::{AssistantReply, ChatMessage, MessageId, ToolCallError, ToolCallResult, ToolDefinition};
9
10#[doc = include_str!("docs/context.md")]
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct Context {
13    messages: Vec<ChatMessage>,
14    tools: Vec<ToolDefinition>,
15    #[serde(skip)]
16    reasoning_effort: ReasoningEffort,
17    #[serde(skip)]
18    model_settings: ModelSettings,
19    #[serde(skip)]
20    prompt_cache_key: Option<String>,
21    #[serde(skip)]
22    session_affinity_key: Option<String>,
23}
24
25impl Context {
26    pub fn new(messages: Vec<ChatMessage>, tools: Vec<ToolDefinition>) -> Self {
27        Self {
28            messages,
29            tools,
30            reasoning_effort: ReasoningEffort::Default,
31            model_settings: ModelSettings::default(),
32            prompt_cache_key: None,
33            session_affinity_key: None,
34        }
35    }
36
37    pub fn prompt_cache_key(&self) -> Option<&str> {
38        self.prompt_cache_key.as_deref()
39    }
40
41    pub fn set_prompt_cache_key(&mut self, key: Option<String>) {
42        self.prompt_cache_key = key;
43    }
44
45    pub fn session_affinity_key(&self) -> Option<&str> {
46        self.session_affinity_key.as_deref()
47    }
48
49    pub fn set_session_affinity_key(&mut self, key: Option<String>) {
50        self.session_affinity_key = key;
51    }
52
53    pub fn reasoning_effort(&self) -> ReasoningEffort {
54        self.reasoning_effort
55    }
56
57    pub fn set_reasoning_effort(&mut self, effort: ReasoningEffort) {
58        self.reasoning_effort = effort;
59    }
60
61    pub fn model_settings(&self) -> &ModelSettings {
62        &self.model_settings
63    }
64
65    pub fn set_model_settings(&mut self, settings: ModelSettings) {
66        self.model_settings = settings;
67    }
68
69    pub fn add_message(&mut self, message: ChatMessage) {
70        self.messages.push(message);
71    }
72
73    pub fn set_tools(&mut self, tools: Vec<ToolDefinition>) {
74        self.tools = tools;
75    }
76
77    pub fn set_system_content(&mut self, content: String) {
78        if let Some(ChatMessage::System { content: existing, .. }) = self.messages.first_mut() {
79            *existing = content;
80        } else {
81            self.messages.insert(0, ChatMessage::system(content));
82        }
83    }
84
85    pub fn system_content(&self) -> Option<&str> {
86        match self.messages.first() {
87            Some(ChatMessage::System { content, .. }) if !content.is_empty() => Some(content),
88            _ => None,
89        }
90    }
91
92    pub fn messages(&self) -> &Vec<ChatMessage> {
93        &self.messages
94    }
95
96    pub fn tools(&self) -> &Vec<ToolDefinition> {
97        &self.tools
98    }
99
100    /// Returns the number of messages in the context
101    pub fn message_count(&self) -> usize {
102        self.messages.len()
103    }
104
105    /// Estimate total token count using the ~4 bytes/token heuristic.
106    /// Includes messages and tool definitions. Used for pre-flight overflow detection.
107    pub fn estimated_token_count(&self) -> u32 {
108        let message_bytes: usize = self.messages.iter().map(ChatMessage::estimated_bytes).sum();
109        let tool_bytes: usize = self
110            .tools
111            .iter()
112            .map(|tool| tool.name.len() + tool.description.len() + tool.parameters.to_string().len())
113            .sum();
114        let total_bytes = message_bytes + tool_bytes;
115        u32::try_from(total_bytes / 4).unwrap_or(u32::MAX)
116    }
117
118    pub fn push_assistant_turn(
119        &mut self,
120        message_id: MessageId,
121        model: Option<LlmModel>,
122        mut content: AssistantReply,
123        results: Vec<Result<ToolCallResult, ToolCallError>>,
124    ) -> &ChatMessage {
125        content.pair_tool_calls(&results);
126        let reply_index = self.messages.len();
127        self.messages.push(ChatMessage::Assistant { message_id, model, content, timestamp: IsoString::now() });
128        self.messages.extend(results.into_iter().map(ChatMessage::ToolCallResult));
129        &self.messages[reply_index]
130    }
131
132    /// Clear all non-system messages, retaining only system prompts.
133    pub fn clear_conversation(&mut self) {
134        self.messages.retain(super::chat_message::ChatMessage::is_system);
135    }
136
137    /// Replace all non-system messages while preserving the system prompt and runtime state.
138    pub fn replace_conversation(&mut self, messages: Vec<ChatMessage>) {
139        self.messages = self
140            .messages
141            .drain(..)
142            .filter(ChatMessage::is_system)
143            .chain(messages.into_iter().filter(|m| !m.is_system()))
144            .collect();
145    }
146
147    /// Get all non-system messages for summarization
148    pub fn messages_for_summary(&self) -> Vec<&ChatMessage> {
149        self.messages.iter().filter(|msg| !msg.is_system()).collect()
150    }
151
152    /// Create a new context with all messages replaced by a summary.
153    /// Preserves the system prompt and tools.
154    pub fn with_compacted_summary(&self, message_id: MessageId, summary: &str) -> Context {
155        let system_messages: Vec<_> = self.messages.iter().filter(|msg| msg.is_system()).cloned().collect();
156
157        let non_system_count = self.messages.len() - system_messages.len();
158
159        let mut messages = system_messages;
160        if non_system_count > 0 {
161            messages.push(ChatMessage::Summary {
162                message_id,
163                content: summary.to_string(),
164                timestamp: IsoString::now(),
165                messages_compacted: non_system_count,
166            });
167        }
168
169        Context { messages, ..self.clone() }
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176
177    use crate::AssistantContent;
178    use crate::testing::assistant_message;
179
180    fn create_test_context() -> Context {
181        let messages = vec![
182            ChatMessage::system("You are a helpful assistant."),
183            ChatMessage::user("Hello"),
184            assistant_message().text("Hi there!").build(),
185            ChatMessage::ToolCallResult(Ok(ToolCallResult {
186                id: "1".to_string(),
187                name: "tool1".to_string(),
188                arguments: "{}".to_string(),
189                result: "Result 1".to_string(),
190            })),
191            ChatMessage::ToolCallResult(Ok(ToolCallResult {
192                id: "2".to_string(),
193                name: "tool2".to_string(),
194                arguments: "{}".to_string(),
195                result: "Result 2".to_string(),
196            })),
197            ChatMessage::ToolCallResult(Ok(ToolCallResult {
198                id: "3".to_string(),
199                name: "tool3".to_string(),
200                arguments: "{}".to_string(),
201                result: "Result 3".to_string(),
202            })),
203        ];
204        Context::new(messages, vec![])
205    }
206
207    #[test]
208    fn replace_conversation_preserves_system_message() {
209        let mut ctx = create_test_context();
210        ctx.replace_conversation(vec![ChatMessage::user("new")]);
211
212        assert_eq!(ctx.message_count(), 2);
213        assert!(ctx.messages()[0].is_system());
214        assert!(matches!(ctx.messages()[1], ChatMessage::User { .. }));
215    }
216
217    #[test]
218    fn replace_conversation_replaces_old_non_system_messages() {
219        let mut ctx = create_test_context();
220        ctx.replace_conversation(vec![assistant_message().text("replacement").build()]);
221
222        assert_eq!(ctx.message_count(), 2);
223        assert!(
224            ctx.messages()
225                .iter()
226                .all(|message| { !matches!(message, ChatMessage::User { .. } | ChatMessage::ToolCallResult(_)) })
227        );
228        assert!(
229            matches!(&ctx.messages()[1], ChatMessage::Assistant { content, .. } if content.text() == "replacement")
230        );
231    }
232
233    #[test]
234    fn replace_conversation_filters_incoming_system_messages() {
235        let mut ctx = create_test_context();
236        ctx.replace_conversation(vec![ChatMessage::system("wrong system"), ChatMessage::user("kept")]);
237
238        assert_eq!(ctx.message_count(), 2);
239        assert!(
240            matches!(ctx.messages()[0], ChatMessage::System { ref content, .. } if content == "You are a helpful assistant.")
241        );
242        assert!(matches!(ctx.messages()[1], ChatMessage::User { .. }));
243    }
244
245    #[test]
246    fn replace_conversation_does_not_change_tools() {
247        let tool = ToolDefinition::new("read_file", "Reads a file", serde_json::json!({}));
248        let mut ctx = Context::new(vec![ChatMessage::system("system")], vec![tool.clone()]);
249        ctx.replace_conversation(vec![ChatMessage::user("new")]);
250
251        assert_eq!(ctx.tools(), &vec![tool]);
252    }
253
254    #[test]
255    fn test_message_count() {
256        let ctx = create_test_context();
257        assert_eq!(ctx.message_count(), 6);
258    }
259
260    #[test]
261    fn test_with_compacted_summary_preserves_system_prompt() {
262        let ctx = create_test_context();
263        let compacted = ctx.with_compacted_summary(MessageId::new(), "This is a summary of previous conversation.");
264
265        assert_eq!(compacted.message_count(), 2);
266        assert!(compacted.messages()[0].is_system());
267        assert!(compacted.messages()[1].is_summary());
268    }
269
270    #[test]
271    fn test_with_compacted_summary_empty_context() {
272        let ctx = Context::new(vec![ChatMessage::system("System")], vec![]);
273        let compacted = ctx.with_compacted_summary(MessageId::new(), "Summary");
274
275        assert_eq!(compacted.message_count(), 1);
276    }
277
278    #[test]
279    fn test_messages_for_summary() {
280        let ctx = create_test_context();
281        let msgs = ctx.messages_for_summary();
282
283        assert_eq!(msgs.len(), 5);
284        assert!(msgs.iter().all(|m| !m.is_system()));
285    }
286
287    #[test]
288    fn test_prompt_cache_key_default_is_none() {
289        let ctx = create_test_context();
290        assert_eq!(ctx.prompt_cache_key(), None);
291    }
292
293    #[test]
294    fn test_prompt_cache_key_set_and_get() {
295        let mut ctx = create_test_context();
296        ctx.set_prompt_cache_key(Some("session-123".to_string()));
297        assert_eq!(ctx.prompt_cache_key(), Some("session-123"));
298
299        ctx.set_prompt_cache_key(None);
300        assert_eq!(ctx.prompt_cache_key(), None);
301    }
302
303    #[test]
304    fn test_prompt_cache_key_preserved_through_compaction() {
305        let mut ctx = create_test_context();
306        ctx.set_prompt_cache_key(Some("session-abc".to_string()));
307        let compacted = ctx.with_compacted_summary(MessageId::new(), "Summary");
308        assert_eq!(compacted.prompt_cache_key(), Some("session-abc"));
309    }
310
311    #[test]
312    fn session_affinity_key_is_runtime_metadata_preserved_by_compaction() {
313        let mut context = create_test_context();
314        assert_eq!(context.session_affinity_key(), None);
315
316        context.set_session_affinity_key(Some("conversation-123".to_string()));
317
318        assert_eq!(context.session_affinity_key(), Some("conversation-123"));
319        assert_eq!(
320            context.with_compacted_summary(MessageId::new(), "Summary").session_affinity_key(),
321            Some("conversation-123")
322        );
323    }
324
325    #[test]
326    fn test_reasoning_effort_initializes_to_default() {
327        let ctx = create_test_context();
328        assert_eq!(ctx.reasoning_effort(), ReasoningEffort::Default);
329    }
330
331    #[test]
332    fn test_reasoning_effort_set_and_get() {
333        let mut ctx = create_test_context();
334        ctx.set_reasoning_effort(crate::ReasoningEffort::High);
335        assert_eq!(ctx.reasoning_effort(), crate::ReasoningEffort::High);
336
337        ctx.set_reasoning_effort(ReasoningEffort::Default);
338        assert_eq!(ctx.reasoning_effort(), ReasoningEffort::Default);
339    }
340
341    #[test]
342    fn test_reasoning_effort_preserved_through_compaction() {
343        let mut ctx = create_test_context();
344        ctx.set_reasoning_effort(crate::ReasoningEffort::Disabled);
345        let compacted = ctx.with_compacted_summary(MessageId::new(), "Summary");
346        assert_eq!(compacted.reasoning_effort(), crate::ReasoningEffort::Disabled);
347        assert_eq!(ctx.clone().reasoning_effort(), ReasoningEffort::Disabled);
348        ctx.replace_conversation(vec![]);
349        assert_eq!(ctx.reasoning_effort(), ReasoningEffort::Disabled);
350    }
351
352    #[test]
353    fn test_estimated_token_count() {
354        use crate::ToolDefinition;
355
356        // "You are a helpful assistant." = 28 bytes
357        // "Hello" = 5 bytes
358        // "Hi there!" = 9 bytes (assistant, no reasoning, no tool calls)
359        // 3 tool results: "Result 1" (8) + "tool1" (5) + "{}" (2) = 15 each = 45 total
360        // Total message bytes = 28 + 5 + 9 + 45 = 87
361        let ctx = create_test_context();
362        let base_estimate = ctx.estimated_token_count();
363
364        // With no tools, estimate = message_bytes / 4
365        assert_eq!(base_estimate, 87 / 4);
366
367        let tool = ToolDefinition::new("read_file", "Reads a file", serde_json::json!({}));
368        let ctx_with_tools = Context::new(ctx.messages().clone(), vec![tool]);
369        let with_tools_estimate = ctx_with_tools.estimated_token_count();
370        assert_eq!(with_tools_estimate, (87 + 9 + 12 + 2) / 4);
371        assert!(with_tools_estimate > base_estimate);
372    }
373
374    #[test]
375    fn compaction_drops_encrypted_reasoning() {
376        let ctx = Context::new(
377            vec![
378                ChatMessage::user("Hello"),
379                assistant_message()
380                    .model("anthropic:claude-opus-4-6")
381                    .reasoning("thinking")
382                    .encrypted_reasoning("r_test", "blob")
383                    .text("I see.")
384                    .build(),
385            ],
386            vec![],
387        );
388        let compacted = ctx.with_compacted_summary(MessageId::new(), "Summary of conversation");
389
390        assert!(compacted.messages().iter().all(|message| encrypted_reasoning(message).is_empty()));
391    }
392
393    #[test]
394    fn assistant_turn_pairs_every_tool_call_with_a_result() {
395        let mut ctx = Context::new(vec![], vec![]);
396        let model: LlmModel = "anthropic:claude-opus-4-6".parse().unwrap();
397        let content = assistant_message()
398            .tool_call("call_1", "read", "{}")
399            .tool_call("call_cancelled", "read", "{}")
400            .build_reply();
401
402        let result = ToolCallResult {
403            id: "call_1".to_string(),
404            name: "read".to_string(),
405            arguments: "{}".to_string(),
406            result: "ok".to_string(),
407        };
408        let error = ToolCallError {
409            id: "call_retried".to_string(),
410            name: "bash".to_string(),
411            arguments: None,
412            error: "failed".to_string(),
413        };
414
415        let recorded = ctx
416            .push_assistant_turn(MessageId::new(), Some(model.clone()), content, vec![Ok(result), Err(error)])
417            .clone();
418
419        assert_eq!(ctx.messages()[0], recorded);
420        assert_eq!(ctx.message_count(), 3);
421        let ChatMessage::Assistant { model: author, content, .. } = recorded else {
422            panic!("expected assistant message")
423        };
424        let ids: Vec<_> = content.tool_calls().map(|call| call.id.as_str()).collect();
425        assert_eq!(ids, ["call_1", "call_retried"]);
426        assert_eq!(author, Some(model));
427    }
428
429    fn encrypted_reasoning(message: &ChatMessage) -> Vec<&str> {
430        let ChatMessage::Assistant { content, .. } = message else { return vec![] };
431        content
432            .blocks()
433            .iter()
434            .filter_map(|block| match block {
435                AssistantContent::Reasoning { encrypted: Some(encrypted), .. } => Some(encrypted.content.as_str()),
436                _ => None,
437            })
438            .collect()
439    }
440}
llm/context.rs

llm/
context.rs
