Skip to main content

vtcode_commons/
message_metadata.rs

1//! Per-message metadata for conversation history.
2//!
3//! Each message in the conversation carries metadata about its origin,
4//! importance, compression state, and resource usage. This enables smart
5//! context pruning (drop low-importance messages first), compression
6//! tracking, and latency analysis.
7//!
8//! Following the "state as a first-class citizen" principle (Hitchhiker's
9//! Guide to Agentic AI, Section 18.6.1), metadata is the foundation for
10//! conversation state quality-of-service decisions.
11
12use serde::{Deserialize, Serialize};
13
14/// Metadata attached to every message in the conversation history.
15///
16/// Skipped during serialization when `None` to preserve backward compatibility
17/// with all existing persistence formats (session archives, snapshots, etc.).
18#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
19pub struct MessageMetadata {
20    /// Unix millisecond timestamp when the message was created.
21    pub timestamp: u64,
22
23    /// Importance score in [0.0, 1.0]: 0.0 = low (safe to drop first),
24    /// 1.0 = high (preserve as long as possible).
25    ///
26    /// Initialised to 0.5 (neutral) and adjusted by the compression/pruning
27    /// system or by explicit agent reflection.
28    pub importance_score: f64,
29
30    /// Current compression status of this message.
31    pub compression_status: CompressionStatus,
32
33    /// Cached token estimate for this message. Populated on creation and
34    /// updated after compression.
35    pub estimated_tokens: usize,
36
37    /// Origin of this message: "user_input", "llm_response", "tool_result",
38    /// "system", or "synthetic".
39    pub source: Option<String>,
40}
41
42impl MessageMetadata {
43    /// Create metadata for a message originating from a user.
44    pub fn user_input(timestamp: u64, estimated_tokens: usize) -> Self {
45        Self {
46            timestamp,
47            importance_score: 0.5,
48            compression_status: CompressionStatus::Uncompressed,
49            estimated_tokens,
50            source: Some("user_input".into()),
51        }
52    }
53
54    /// Create metadata for a message originating from an LLM response.
55    pub fn llm_response(timestamp: u64, estimated_tokens: usize) -> Self {
56        Self {
57            timestamp,
58            importance_score: 0.6,
59            compression_status: CompressionStatus::Uncompressed,
60            estimated_tokens,
61            source: Some("llm_response".into()),
62        }
63    }
64
65    /// Create metadata for a tool result message.
66    pub fn tool_result(timestamp: u64, estimated_tokens: usize) -> Self {
67        Self {
68            timestamp,
69            importance_score: 0.4,
70            compression_status: CompressionStatus::Uncompressed,
71            estimated_tokens,
72            source: Some("tool_result".into()),
73        }
74    }
75
76    /// Create metadata for a system message.
77    pub fn system(timestamp: u64, estimated_tokens: usize) -> Self {
78        Self {
79            timestamp,
80            importance_score: 1.0,
81            compression_status: CompressionStatus::Uncompressed,
82            estimated_tokens,
83            source: Some("system".into()),
84        }
85    }
86
87    /// Create metadata for a synthetic (e.g., recovery/injected) message.
88    pub fn synthetic(timestamp: u64, estimated_tokens: usize) -> Self {
89        Self {
90            timestamp,
91            importance_score: 0.3,
92            compression_status: CompressionStatus::Uncompressed,
93            estimated_tokens,
94            source: Some("synthetic".into()),
95        }
96    }
97
98    /// Mark this message as compressed, recording the original and new token counts.
99    pub fn mark_compressed(&mut self, original_tokens: usize, compressed_tokens: usize) {
100        self.compression_status = CompressionStatus::Compressed {
101            original_token_count: original_tokens,
102            summary_token_count: compressed_tokens,
103        };
104        self.estimated_tokens = compressed_tokens;
105    }
106
107    /// Mark this message as summarized.
108    pub fn mark_summarized(&mut self, original_tokens: usize, summary_tokens: usize) {
109        self.compression_status = CompressionStatus::Summarized {
110            original_token_count: original_tokens,
111            summary_token_count: summary_tokens,
112        };
113        self.estimated_tokens = summary_tokens;
114    }
115
116    /// Set the importance score (clamped to [0.0, 1.0]).
117    pub fn set_importance(&mut self, score: f64) {
118        self.importance_score = score.clamp(0.0, 1.0);
119    }
120
121    /// Returns the original (pre-compression) token count, or the current count
122    /// if the message was never compressed.
123    pub fn original_token_count(&self) -> usize {
124        match self.compression_status {
125            CompressionStatus::Uncompressed => self.estimated_tokens,
126            CompressionStatus::Compressed { original_token_count, .. }
127            | CompressionStatus::Summarized { original_token_count, .. } => original_token_count,
128            CompressionStatus::Dropped => 0,
129        }
130    }
131
132    /// Returns the effective (post-compression) token count.
133    pub fn effective_token_count(&self) -> usize {
134        match self.compression_status {
135            CompressionStatus::Uncompressed => self.estimated_tokens,
136            CompressionStatus::Compressed { summary_token_count, .. }
137            | CompressionStatus::Summarized { summary_token_count, .. } => summary_token_count,
138            CompressionStatus::Dropped => 0,
139        }
140    }
141}
142
143/// Tracks the compression state of a single message in conversation history.
144#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
145#[serde(rename_all = "snake_case")]
146pub enum CompressionStatus {
147    /// Message is in its original uncompressed form.
148    Uncompressed,
149    /// Message has been compressed with token-level preservation of information.
150    Compressed {
151        original_token_count: usize,
152        summary_token_count: usize,
153    },
154    /// Message has been semantically summarized (lossy compression).
155    Summarized {
156        original_token_count: usize,
157        summary_token_count: usize,
158    },
159    /// Message has been dropped from the active context but may be in long-term
160    /// memory.
161    Dropped,
162}
163
164#[allow(clippy::derivable_impls)]
165impl Default for CompressionStatus {
166    fn default() -> Self {
167        Self::Uncompressed
168    }
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174
175    #[test]
176    fn test_create_user_metadata() {
177        let meta = MessageMetadata::user_input(1000, 50);
178        assert_eq!(meta.timestamp, 1000);
179        assert!((meta.importance_score - 0.5).abs() < f64::EPSILON);
180        assert_eq!(meta.compression_status, CompressionStatus::Uncompressed);
181        assert_eq!(meta.estimated_tokens, 50);
182        assert_eq!(meta.source.as_deref(), Some("user_input"));
183    }
184
185    #[test]
186    fn test_create_llm_response_metadata() {
187        let meta = MessageMetadata::llm_response(2000, 150);
188        assert!((meta.importance_score - 0.6).abs() < f64::EPSILON);
189    }
190
191    #[test]
192    fn test_mark_compressed() {
193        let mut meta = MessageMetadata::user_input(1000, 200);
194        meta.mark_compressed(200, 50);
195        assert_eq!(meta.estimated_tokens, 50);
196        assert_eq!(meta.effective_token_count(), 50);
197        assert_eq!(meta.original_token_count(), 200);
198    }
199
200    #[test]
201    fn test_mark_summarized() {
202        let mut meta = MessageMetadata::user_input(1000, 300);
203        meta.mark_summarized(300, 30);
204        assert_eq!(meta.effective_token_count(), 30);
205        assert_eq!(meta.original_token_count(), 300);
206    }
207
208    #[test]
209    fn test_set_importance_clamps() {
210        let mut meta = MessageMetadata::user_input(1000, 50);
211        meta.set_importance(1.5);
212        assert!((meta.importance_score - 1.0).abs() < f64::EPSILON);
213        meta.set_importance(-0.5);
214        assert!((meta.importance_score - 0.0).abs() < f64::EPSILON);
215    }
216
217    #[test]
218    fn test_compression_status_serde_roundtrip() {
219        let status =
220            CompressionStatus::Compressed { original_token_count: 200, summary_token_count: 50 };
221        let json = serde_json::to_string(&status).unwrap();
222        let deserialized: CompressionStatus = serde_json::from_str(&json).unwrap();
223        assert_eq!(status, deserialized);
224    }
225}