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

llm/
context.rs
