deepstrike_core/context/
measurement.rs1use serde::{Deserialize, Serialize};
9
10#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
14#[serde(tag = "kind", rename_all = "snake_case")]
15pub enum MeasurementSource {
16 Native { provider: String },
19 LocalExact { tokenizer: String },
23 Postflight,
26 Heuristic,
28 HostProvided,
30}
31
32#[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 use sha2::{Digest as _, Sha256};
46 let digest = Sha256::digest(serde_json::to_vec(message).expect("message is serializable"));
47 let hex = digest
48 .iter()
49 .map(|b| format!("{b:02x}"))
50 .collect::<String>();
51 Self {
52 fingerprint: format!("sha256:{hex}"),
53 tokens,
54 source: MeasurementSource::HostProvided,
55 confidence: MeasurementConfidence::HighConfidence,
56 }
57 }
58}
59
60#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
63#[serde(deny_unknown_fields)]
64pub struct ToolMeasurement {
65 pub call_id: String,
66 pub tokens: u32,
67}
68
69impl ToolMeasurement {
70 pub fn new(call_id: impl Into<String>, tokens: u32) -> Self {
71 Self {
72 call_id: call_id.into(),
73 tokens,
74 }
75 }
76}
77
78#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
83#[serde(rename_all = "snake_case")]
84pub enum MeasurementConfidence {
85 Exact,
86 HighConfidence,
87 LowConfidence,
88}
89
90#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
92#[serde(deny_unknown_fields)]
93pub struct PromptMeasurement {
94 pub input_tokens: u32,
95 pub source: MeasurementSource,
96 pub confidence: MeasurementConfidence,
97}
98
99#[cfg(test)]
100mod tests {
101 use super::*;
102
103 #[test]
104 fn constructs_all_three_measurement_source_variants() {
105 let native = MeasurementSource::Native {
106 provider: "anthropic".to_string(),
107 };
108 let local_exact = MeasurementSource::LocalExact {
109 tokenizer: "cl100k_base".to_string(),
110 };
111 let heuristic = MeasurementSource::Heuristic;
112
113 assert_ne!(native, local_exact);
114 assert_ne!(local_exact, heuristic);
115 }
116
117 #[test]
118 fn prompt_measurement_round_trips_through_json() {
119 let m = PromptMeasurement {
120 input_tokens: 1234,
121 source: MeasurementSource::Native {
122 provider: "openai".to_string(),
123 },
124 confidence: MeasurementConfidence::Exact,
125 };
126 let json = serde_json::to_string(&m).unwrap();
127 let back: PromptMeasurement = serde_json::from_str(&json).unwrap();
128 assert_eq!(m, back);
129 }
130
131 #[test]
132 fn postflight_source_round_trips_and_differs_from_preflight_kinds() {
133 let m = PromptMeasurement {
135 input_tokens: 1000,
136 source: MeasurementSource::Postflight,
137 confidence: MeasurementConfidence::Exact,
138 };
139 let json = serde_json::to_string(&m).unwrap();
140 let back: PromptMeasurement = serde_json::from_str(&json).unwrap();
141 assert_eq!(m, back);
142 assert_eq!(
143 serde_json::to_string(&m.source).unwrap(),
144 r#"{"kind":"postflight"}"#
145 );
146 }
147
148 #[test]
149 fn unknown_field_is_rejected() {
150 let raw = r#"{"input_tokens": 10, "source": {"kind": "heuristic"}, "confidence": "exact", "extra": true}"#;
151 let result: Result<PromptMeasurement, _> = serde_json::from_str(raw);
152 assert!(
153 result.is_err(),
154 "deny_unknown_fields must reject stray keys"
155 );
156 }
157
158 #[test]
159 fn tool_measurement_is_independent_from_tool_result_state() {
160 let measurement = ToolMeasurement::new("call-1", 42);
161 let json = serde_json::to_string(&measurement).unwrap();
162 let decoded: ToolMeasurement = serde_json::from_str(&json).unwrap();
163 assert_eq!(decoded, measurement);
164 }
165}