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

llm/
context.rs
