Skip to main content

deepstrike_core/context/
measurement.rs

1//! spc_011-C-02: prompt token measurement — the three-way split the spec's §6.1 invariant
2//! ("measure provider-visible input, never underestimate") requires between what a request is
3//! *estimated* to cost before it is sent (`PromptMeasurement`, this module), what the provider
4//! actually reports back (`ProviderUsage`, Node/Python Host layer — not this crate). This module
5//! only defines a host-side preflight fact for provider-visible request preparation.
6
7use serde::{Deserialize, Serialize};
8
9/// Where a token count came from — never just a bare `u32`, so a caller can tell "the provider's
10/// own count API said so" apart from "we guessed." The provenance remains part of the reserved
11/// contract, but no runtime currently persists or consumes this fact.
12#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
13#[serde(tag = "kind", rename_all = "snake_case")]
14pub enum MeasurementSource {
15    /// The provider's own token-counting endpoint answered (e.g. Anthropic
16    /// `messages/count_tokens`, OpenAI Responses token counting, Gemini `countTokens`).
17    Native { provider: String },
18    /// A real BPE tokenizer ran locally, but against a vocabulary that may not be the target
19    /// provider's own (e.g. cl100k standing in for a non-OpenAI vendor) — see
20    /// `FallbackEstimator` (spc_011-C-01) for the concrete counter this describes.
21    LocalExact { tokenizer: String },
22    /// The provider's own postflight usage for this exact request, fed back after execution
23    /// (spc_024-06). The authority a replay reuses — never a preflight guess.
24    Postflight,
25    /// No tokenizer ran at all; this is a coarse guess with a generous safety margin.
26    Heuristic,
27    /// A host supplied measurement tied to a canonical object fingerprint.
28    HostProvided,
29}
30
31/// Host-side token evidence for one canonical runtime object. This is deliberately separate from
32/// `CoreMessage`: a measurement may be refreshed or invalidated without changing semantics.
33#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
34#[serde(deny_unknown_fields)]
35pub struct TokenMeasurement {
36    pub fingerprint: String,
37    pub tokens: u32,
38    pub source: MeasurementSource,
39    pub confidence: MeasurementConfidence,
40}
41
42impl TokenMeasurement {
43    pub fn for_message(message: &crate::types::message::CoreMessage, tokens: u32) -> Self {
44        Self {
45            fingerprint: Self::message_fingerprint(message),
46            tokens,
47            source: MeasurementSource::HostProvided,
48            confidence: MeasurementConfidence::HighConfidence,
49        }
50    }
51
52    pub fn matches_message(&self, message: &crate::types::message::CoreMessage) -> bool {
53        self.fingerprint == Self::message_fingerprint(message)
54    }
55
56    fn message_fingerprint(message: &crate::types::message::CoreMessage) -> String {
57        use sha2::{Digest as _, Sha256};
58        let material =
59            super::execution::message_material(message, &crate::mm::handle::HandleTable::new());
60        let digest =
61            Sha256::digest(serde_json::to_vec(&material).expect("message is serializable"));
62        let hex = digest
63            .iter()
64            .map(|b| format!("{b:02x}"))
65            .collect::<String>();
66        format!("sha256:{hex}")
67    }
68}
69
70/// Host-owned accounting fact for one completed tool call. It is an input to settlement and
71/// context accounting, never a field on the runtime `ToolResult` message.
72#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
73#[serde(deny_unknown_fields)]
74pub struct ToolMeasurement {
75    pub call_id: String,
76    pub tokens: u32,
77}
78
79impl ToolMeasurement {
80    pub fn new(call_id: impl Into<String>, tokens: u32) -> Self {
81        Self {
82            call_id: call_id.into(),
83            tokens,
84        }
85    }
86}
87
88/// How much to trust `PromptMeasurement.input_tokens` when deciding whether to compress. Kept
89/// separate from `MeasurementSource` — the *reason* a number exists and how much a caller should
90/// *lean on it* are different questions (a `LocalExact` cl100k count for an Anthropic request is
91/// still nontrivially uncertain, but it is not a bare guess either).
92#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
93#[serde(rename_all = "snake_case")]
94pub enum MeasurementConfidence {
95    Exact,
96    HighConfidence,
97    LowConfidence,
98}
99
100/// A single preflight token-count fact about a candidate render, for a specific provider/model.
101#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
102#[serde(deny_unknown_fields)]
103pub struct PromptMeasurement {
104    pub input_tokens: u32,
105    pub source: MeasurementSource,
106    pub confidence: MeasurementConfidence,
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn constructs_all_three_measurement_source_variants() {
115        let native = MeasurementSource::Native {
116            provider: "anthropic".to_string(),
117        };
118        let local_exact = MeasurementSource::LocalExact {
119            tokenizer: "cl100k_base".to_string(),
120        };
121        let heuristic = MeasurementSource::Heuristic;
122
123        assert_ne!(native, local_exact);
124        assert_ne!(local_exact, heuristic);
125    }
126
127    #[test]
128    fn prompt_measurement_round_trips_through_json() {
129        let m = PromptMeasurement {
130            input_tokens: 1234,
131            source: MeasurementSource::Native {
132                provider: "openai".to_string(),
133            },
134            confidence: MeasurementConfidence::Exact,
135        };
136        let json = serde_json::to_string(&m).unwrap();
137        let back: PromptMeasurement = serde_json::from_str(&json).unwrap();
138        assert_eq!(m, back);
139    }
140
141    #[test]
142    fn postflight_source_round_trips_and_differs_from_preflight_kinds() {
143        // spc_024-06: observed usage fed back after execution — the authority a replay reuses.
144        let m = PromptMeasurement {
145            input_tokens: 1000,
146            source: MeasurementSource::Postflight,
147            confidence: MeasurementConfidence::Exact,
148        };
149        let json = serde_json::to_string(&m).unwrap();
150        let back: PromptMeasurement = serde_json::from_str(&json).unwrap();
151        assert_eq!(m, back);
152        assert_eq!(
153            serde_json::to_string(&m.source).unwrap(),
154            r#"{"kind":"postflight"}"#
155        );
156    }
157
158    #[test]
159    fn unknown_field_is_rejected() {
160        let raw = r#"{"input_tokens": 10, "source": {"kind": "heuristic"}, "confidence": "exact", "extra": true}"#;
161        let result: Result<PromptMeasurement, _> = serde_json::from_str(raw);
162        assert!(
163            result.is_err(),
164            "deny_unknown_fields must reject stray keys"
165        );
166    }
167
168    #[test]
169    fn tool_measurement_is_independent_from_tool_result_state() {
170        let measurement = ToolMeasurement::new("call-1", 42);
171        let json = serde_json::to_string(&measurement).unwrap();
172        let decoded: ToolMeasurement = serde_json::from_str(&json).unwrap();
173        assert_eq!(decoded, measurement);
174    }
175}