Skip to main content

neuron_turn/
context.rs

1//! Context strategy for managing the conversation window.
2//!
3//! The [`ContextStrategy`] trait handles client-side context compaction.
4//! Provider-native truncation (e.g., OpenAI `truncation: auto`) is
5//! invisible to the strategy — handled by the Provider impl internally.
6
7use crate::types::ProviderMessage;
8
9/// Strategy for managing context window size.
10///
11/// Implementations: `NoCompaction` (passthrough), `SlidingWindow`
12/// (drop oldest messages), `Summarization` (future).
13pub trait ContextStrategy: Send + Sync {
14    /// Estimate token count for a message list.
15    fn token_estimate(&self, messages: &[ProviderMessage]) -> usize;
16
17    /// Whether compaction should run given the current messages and limit.
18    fn should_compact(&self, messages: &[ProviderMessage], limit: usize) -> bool;
19
20    /// Compact the message list. Returns a shorter list.
21    fn compact(&self, messages: Vec<ProviderMessage>) -> Vec<ProviderMessage>;
22}
23
24/// A no-op context strategy that never compacts.
25///
26/// Useful for short conversations or when the provider handles
27/// truncation natively.
28pub struct NoCompaction;
29
30impl ContextStrategy for NoCompaction {
31    fn token_estimate(&self, messages: &[ProviderMessage]) -> usize {
32        // Rough estimate: 4 chars per token
33        messages
34            .iter()
35            .flat_map(|m| &m.content)
36            .map(|part| {
37                use crate::types::ContentPart;
38                match part {
39                    ContentPart::Text { text } => text.len() / 4,
40                    ContentPart::ToolUse { input, .. } => input.to_string().len() / 4,
41                    ContentPart::ToolResult { content, .. } => content.len() / 4,
42                    ContentPart::Image { .. } => 1000, // rough image token estimate
43                }
44            })
45            .sum()
46    }
47
48    fn should_compact(&self, _messages: &[ProviderMessage], _limit: usize) -> bool {
49        false
50    }
51
52    fn compact(&self, messages: Vec<ProviderMessage>) -> Vec<ProviderMessage> {
53        messages
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60    use crate::types::{ContentPart, Role};
61
62    #[test]
63    fn no_compaction_never_compacts() {
64        let strategy = NoCompaction;
65        let messages = vec![ProviderMessage {
66            role: Role::User,
67            content: vec![ContentPart::Text {
68                text: "hello".into(),
69            }],
70        }];
71
72        assert!(!strategy.should_compact(&messages, 100));
73        let compacted = strategy.compact(messages.clone());
74        assert_eq!(compacted.len(), messages.len());
75    }
76
77    #[test]
78    fn no_compaction_estimates_tokens() {
79        let strategy = NoCompaction;
80        let messages = vec![ProviderMessage {
81            role: Role::User,
82            content: vec![ContentPart::Text {
83                text: "a".repeat(400),
84            }],
85        }];
86
87        let estimate = strategy.token_estimate(&messages);
88        assert_eq!(estimate, 100); // 400 chars / 4
89    }
90
91    #[test]
92    fn no_compaction_preserves_all_messages() {
93        let strategy = NoCompaction;
94        let messages = vec![
95            ProviderMessage {
96                role: Role::User,
97                content: vec![ContentPart::Text {
98                    text: "msg1".into(),
99                }],
100            },
101            ProviderMessage {
102                role: Role::Assistant,
103                content: vec![ContentPart::Text {
104                    text: "msg2".into(),
105                }],
106            },
107            ProviderMessage {
108                role: Role::User,
109                content: vec![ContentPart::Text {
110                    text: "msg3".into(),
111                }],
112            },
113        ];
114
115        let compacted = strategy.compact(messages.clone());
116        assert_eq!(compacted.len(), 3);
117        assert_eq!(compacted[0].content, messages[0].content);
118        assert_eq!(compacted[1].content, messages[1].content);
119        assert_eq!(compacted[2].content, messages[2].content);
120    }
121
122    #[test]
123    fn no_compaction_estimates_tool_use_tokens() {
124        let strategy = NoCompaction;
125        let messages = vec![ProviderMessage {
126            role: Role::Assistant,
127            content: vec![ContentPart::ToolUse {
128                id: "tu_1".into(),
129                name: "bash".into(),
130                input: serde_json::json!({"command": "ls"}),
131            }],
132        }];
133
134        let estimate = strategy.token_estimate(&messages);
135        // The JSON representation of the input will be tokenized
136        assert!(estimate > 0);
137    }
138
139    #[test]
140    fn no_compaction_estimates_tool_result_tokens() {
141        let strategy = NoCompaction;
142        let messages = vec![ProviderMessage {
143            role: Role::User,
144            content: vec![ContentPart::ToolResult {
145                tool_use_id: "tu_1".into(),
146                content: "a".repeat(200),
147                is_error: false,
148            }],
149        }];
150
151        let estimate = strategy.token_estimate(&messages);
152        assert_eq!(estimate, 50); // 200 chars / 4
153    }
154
155    #[test]
156    fn no_compaction_estimates_image_tokens() {
157        let strategy = NoCompaction;
158        let messages = vec![ProviderMessage {
159            role: Role::User,
160            content: vec![ContentPart::Image {
161                source: crate::types::ImageSource::Url {
162                    url: "https://example.com/img.png".into(),
163                },
164                media_type: "image/png".into(),
165            }],
166        }];
167
168        let estimate = strategy.token_estimate(&messages);
169        assert_eq!(estimate, 1000); // rough image estimate
170    }
171
172    #[test]
173    fn context_strategy_is_object_safe() {
174        fn _assert_object_safe(_: &dyn ContextStrategy) {}
175        let nc = NoCompaction;
176        _assert_object_safe(&nc);
177    }
178}