llm/
context.rs

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

llm/
context.rs
