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

llm/
context.rs
