Skip to main content

deepstrike_core/context/
partitions.rs

1use super::config::ContextConfig;
2use super::task_state::TaskState;
3use super::token_engine::ContextTokenEngine;
4use crate::types::message::Message;
5
6/// A single context partition — a named bucket of messages with a token counter.
7#[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
37/// Three-partition context model aligned with LLM API slots:
38///
39///   Slot 1 — Identity  (system):    who the agent is; role, rules, constraints.
40///                                    Maps to: Anthropic system[0] cache_control, OpenAI system role.
41///                                    Never changes within a run.
42///
43///   Slot 2 — Knowledge (knowledge): what the agent knows; memory retrievals, skill
44///                                    definitions, artifacts. Low-frequency changes.
45///                                    Maps to: Anthropic system[1] cache_control.
46///
47///   Slot 3 — State     (task_state + signals): what the agent is doing right now.
48///                                    task_state = goal/plan/progress (structured).
49///                                    signals = runtime events (rollback notes, interrupts).
50///                                    Maps to: messages[0] user turn, rebuilt every call.
51///
52///   Slot 4 — History   (history):   what the agent has done; conversation turns,
53///                                    tool calls and results. Compression pipeline target.
54///                                    Maps to: messages[1..N].
55pub struct ContextPartitions {
56    pub system: Partition,
57    pub knowledge: Partition,
58    pub task_state: TaskState,
59    /// Runtime signals injected into the current turn (rollback notes, interrupts).
60    /// Cleared after each render — signals are ephemeral per-turn events.
61    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    /// Total token count across all slots.
77    /// task_state tokens are measured from its rendered compact form.
78    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}