synaptic_core/
context_budget.rs

1use crate::token_counter::TokenCounter;
2use crate::Message;
3use std::sync::Arc;
4
5/// Priority level for context slots. Lower values = higher priority.
6#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
7pub struct Priority(pub u8);
8
9impl Priority {
10    pub const CRITICAL: Priority = Priority(0);
11    pub const HIGH: Priority = Priority(64);
12    pub const NORMAL: Priority = Priority(128);
13    pub const LOW: Priority = Priority(192);
14}
15
16/// A slot of context to include in the budget.
17pub struct ContextSlot {
18    pub name: String,
19    pub priority: Priority,
20    pub messages: Vec<Message>,
21    /// Minimum reserved tokens for this slot (guaranteed if budget allows).
22    pub reserved_tokens: usize,
23}
24
25/// Assembles messages from multiple context slots within a token budget.
26///
27/// Slots are sorted by priority (lowest value = highest priority).
28/// Higher-priority slots are included first. Lower-priority slots are
29/// dropped if the budget is exceeded.
30pub struct ContextBudget {
31    max_tokens: usize,
32    counter: Arc<dyn TokenCounter>,
33}
34
35impl ContextBudget {
36    pub fn new(max_tokens: usize, counter: Arc<dyn TokenCounter>) -> Self {
37        Self {
38            max_tokens,
39            counter,
40        }
41    }
42
43    /// Assemble messages from slots that fit within the token budget.
44    ///
45    /// Slots are processed in priority order (CRITICAL first, LOW last).
46    /// Each slot's messages are included if they fit. Slots with
47    /// `reserved_tokens > 0` are guaranteed inclusion (if total reserved
48    /// fits within budget).
49    pub fn assemble(&self, mut slots: Vec<ContextSlot>) -> Vec<Message> {
50        // Sort by priority (lower value = higher priority)
51        slots.sort_by_key(|s| s.priority);
52
53        let mut result = Vec::new();
54        let mut used_tokens = 0;
55
56        for slot in slots {
57            let slot_tokens = self.counter.count_messages(&slot.messages);
58
59            if slot.reserved_tokens > 0 {
60                // Reserved slots are always included (up to budget)
61                if used_tokens + slot_tokens <= self.max_tokens {
62                    used_tokens += slot_tokens;
63                    result.extend(slot.messages);
64                }
65            } else if used_tokens + slot_tokens <= self.max_tokens {
66                used_tokens += slot_tokens;
67                result.extend(slot.messages);
68            }
69            // If doesn't fit and not reserved, skip
70        }
71
72        result
73    }
74}
synaptic_core/context_budget.rs

synaptic_core/
context_budget.rs
