deepstrike_core/context/
partitions.rs1use super::config::ContextConfig;
2use super::task_state::TaskState;
3use super::token_engine::ContextTokenEngine;
4use crate::types::message::Message;
5
6#[derive(Debug, Clone)]
8pub struct Partition {
9 pub messages: Vec<Message>,
10 pub token_count: u32,
11}
12
13impl Partition {
14 pub fn new() -> Self {
15 Self { messages: Vec::new(), token_count: 0 }
16 }
17
18 pub fn push(&mut self, mut msg: Message, token_count: u32) {
19 msg.token_count = Some(token_count);
20 self.token_count += token_count;
21 self.messages.push(msg);
22 }
23
24 pub fn clear(&mut self) {
25 self.messages.clear();
26 self.token_count = 0;
27 }
28
29 pub fn len(&self) -> usize { self.messages.len() }
30 pub fn is_empty(&self) -> bool { self.messages.is_empty() }
31}
32
33impl Default for Partition {
34 fn default() -> Self { Self::new() }
35}
36
37pub struct ContextPartitions {
56 pub system: Partition,
57 pub knowledge: Partition,
58 pub task_state: TaskState,
59 pub signals: Vec<String>,
62 pub history: Partition,
63}
64
65impl ContextPartitions {
66 pub fn new(_config: &ContextConfig) -> Self {
67 Self {
68 system: Partition::new(),
69 knowledge: Partition::new(),
70 task_state: TaskState::default(),
71 signals: Vec::new(),
72 history: Partition::new(),
73 }
74 }
75
76 pub fn total_tokens(&self, engine: &ContextTokenEngine) -> u32 {
79 self.system.token_count
80 + self.knowledge.token_count
81 + engine.count(&self.task_state.format_compact())
82 + self.history.token_count
83 }
84}
85
86impl Default for ContextPartitions {
87 fn default() -> Self {
88 Self::new(&ContextConfig::default())
89 }
90}
91
92#[cfg(test)]
93mod tests {
94 use super::*;
95 use crate::context::config::ContextConfig;
96 use crate::context::token_engine::ContextTokenEngine;
97 use crate::types::message::Message;
98
99 fn engine() -> ContextTokenEngine { ContextTokenEngine::char_approx() }
100
101 #[test]
102 fn push_updates_token_count() {
103 let mut ctx = ContextPartitions::new(&ContextConfig::default());
104 let base = ctx.total_tokens(&engine());
105 ctx.system.push(Message::system("rules"), 10);
106 ctx.history.push(Message::user("hello"), 5);
107 assert_eq!(ctx.total_tokens(&engine()), base + 15);
108 }
109
110 #[test]
111 fn task_state_tokens_included_in_total() {
112 use crate::context::task_state::TaskState;
113 let mut ctx = ContextPartitions::new(&ContextConfig::default());
114 let before = ctx.total_tokens(&engine());
115 ctx.task_state = TaskState { goal: "do something important".to_string(), ..Default::default() };
116 let after = ctx.total_tokens(&engine());
117 assert!(after > before, "task_state should contribute to total_tokens");
118 }
119
120 #[test]
121 fn knowledge_tokens_included_in_total() {
122 let mut ctx = ContextPartitions::new(&ContextConfig::default());
123 let before = ctx.total_tokens(&engine());
124 ctx.knowledge.push(Message::system("skill: debug"), 20);
125 assert_eq!(ctx.total_tokens(&engine()), before + 20);
126 }
127}