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