machi_protocol/
tokens.rs

1//! Rough token estimation for preflight overflow checks (W3.2).
2//!
3//! Maturity: **core**
4//!
5//! Heuristic: text bytes / 4; images fixed at [`IMAGE_TOKEN_COST`].
6
7/// Fixed token cost for one image part (provider-agnostic estimate).
8pub const IMAGE_TOKEN_COST: u32 = 765;
9
10/// Approximate tokens for a UTF-8 string: `ceil(bytes / 4)`.
11#[must_use]
12pub fn estimate_text_tokens(text: &str) -> u32 {
13    let bytes = u32::try_from(text.len()).unwrap_or(u32::MAX);
14    bytes.div_ceil(4)
15}
16
17/// Approximate tokens for one image attachment.
18#[must_use]
19pub const fn estimate_image_tokens() -> u32 {
20    IMAGE_TOKEN_COST
21}
22
23/// Framing overhead tokens per message (role + separators).
24pub const MESSAGE_FRAME_TOKENS: u32 = 4;
25
26/// Preflight decision before sampling.
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub enum PreflightOverflow {
29    /// Under threshold.
30    Ok {
31        /// Estimated tokens.
32        estimated: u32,
33    },
34    /// Over threshold; caller should compact or fail.
35    Overflow {
36        /// Estimated tokens.
37        estimated: u32,
38        /// Context window (tokens).
39        window: u32,
40        /// Soft threshold used (`window * threshold_ratio`).
41        limit: u32,
42    },
43}
44
45/// Check whether `estimated` exceeds `window * threshold_ratio` (clamped).
46///
47/// `threshold_ratio` is typically `0.85`–`0.95`.
48#[must_use]
49pub fn check_context_overflow(
50    estimated: u32,
51    window: u32,
52    threshold_ratio: f32,
53) -> PreflightOverflow {
54    if window == 0 {
55        return PreflightOverflow::Ok { estimated };
56    }
57    let ratio = threshold_ratio.clamp(0.1, 1.0);
58    // Integer math: ceil(window * ratio) = (window * numer + denom - 1) / denom
59    // with ratio ≈ numer/1000 for millis precision.
60    #[allow(
61        clippy::cast_possible_truncation,
62        clippy::cast_sign_loss,
63        clippy::cast_precision_loss,
64        reason = "token limits are u32; ratio is clamped to 0.1..=1.0"
65    )]
66    let numer = (ratio * 1000.0).round() as u64;
67    let numer = numer.clamp(100, 1000);
68    let limit = (u64::from(window).saturating_mul(numer).saturating_add(999) / 1000)
69        .min(u64::from(u32::MAX));
70    let limit = u32::try_from(limit).unwrap_or(u32::MAX);
71    if estimated > limit {
72        PreflightOverflow::Overflow {
73            estimated,
74            window,
75            limit,
76        }
77    } else {
78        PreflightOverflow::Ok { estimated }
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85
86    #[test]
87    fn text_estimate_bytes_over_four() {
88        assert_eq!(estimate_text_tokens("abcd"), 1);
89        assert_eq!(estimate_text_tokens("abcde"), 2);
90    }
91
92    #[test]
93    fn overflow_threshold() {
94        let est = 900;
95        match check_context_overflow(est, 1000, 0.85) {
96            PreflightOverflow::Overflow { limit, .. } => {
97                assert_eq!(limit, 850);
98            }
99            PreflightOverflow::Ok { .. } => unreachable!("expected overflow"),
100        }
101        assert!(matches!(
102            check_context_overflow(800, 1000, 0.85),
103            PreflightOverflow::Ok { .. }
104        ));
105    }
106}
machi_protocol/tokens.rs

machi_protocol/
tokens.rs
