Skip to main content

flashmind_types/
session.rs

1//! Session identities and cumulative request usage.
2
3use std::fmt;
4use std::str::FromStr;
5
6use rust_decimal::Decimal;
7use serde::{Deserialize, Serialize};
8use thiserror::Error;
9use uuid::Uuid;
10
11use crate::{ModelPricing, TokenUsage};
12
13/// Reuse one ID for every request in a conversation, including after resume.
14/// Create a new ID for a new conversation or convert its existing UUID.
15#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
16#[serde(try_from = "String", into = "String")]
17pub struct SessionId(String);
18
19#[derive(Debug, Clone, PartialEq, Eq, Error)]
20pub enum SessionIdError {
21    #[error("session ID must not be empty")]
22    Empty,
23    #[error("session ID exceeds 256 characters")]
24    TooLong,
25}
26
27impl SessionId {
28    pub const MAX_LEN: usize = 256;
29
30    /// Generate an ID once when starting a conversation.
31    pub fn new() -> Self {
32        Uuid::new_v4().into()
33    }
34
35    pub fn as_str(&self) -> &str {
36        &self.0
37    }
38}
39
40impl Default for SessionId {
41    fn default() -> Self {
42        Self::new()
43    }
44}
45
46impl From<Uuid> for SessionId {
47    fn from(value: Uuid) -> Self {
48        Self(value.to_string())
49    }
50}
51
52impl TryFrom<String> for SessionId {
53    type Error = SessionIdError;
54
55    fn try_from(value: String) -> Result<Self, Self::Error> {
56        if value.is_empty() {
57            return Err(SessionIdError::Empty);
58        }
59        if value.chars().count() > Self::MAX_LEN {
60            return Err(SessionIdError::TooLong);
61        }
62        Ok(Self(value))
63    }
64}
65
66impl FromStr for SessionId {
67    type Err = SessionIdError;
68
69    fn from_str(value: &str) -> Result<Self, Self::Err> {
70        value.to_owned().try_into()
71    }
72}
73
74impl From<SessionId> for String {
75    fn from(value: SessionId) -> Self {
76        value.0
77    }
78}
79
80impl AsRef<str> for SessionId {
81    fn as_ref(&self) -> &str {
82        self.as_str()
83    }
84}
85
86impl fmt::Display for SessionId {
87    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
88        f.write_str(self.as_str())
89    }
90}
91
92/// Request totals for a session, independent of its current context window.
93#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
94pub struct SessionUsage {
95    /// Total input tokens, including cached input.
96    pub input_tokens: u64,
97    /// Cached input included in `input_tokens`. Older saves lack this breakdown.
98    #[serde(default)]
99    pub cached_input_tokens: Option<u64>,
100    /// Total generated tokens across all requests.
101    pub output_tokens: u64,
102    /// Known request costs in USD, serialized without rounding.
103    #[serde(with = "rust_decimal::serde::str")]
104    pub total_cost: Decimal,
105    /// Historical token totals were not recorded for this session.
106    pub tokens_unknown: bool,
107    /// At least one request did not have pricing available.
108    pub cost_unknown: bool,
109}
110
111impl Default for SessionUsage {
112    fn default() -> Self {
113        Self {
114            input_tokens: 0,
115            cached_input_tokens: Some(0),
116            output_tokens: 0,
117            total_cost: Decimal::ZERO,
118            tokens_unknown: false,
119            cost_unknown: false,
120        }
121    }
122}
123
124impl SessionUsage {
125    /// Add one request using its model pricing, including cache discounts.
126    pub fn add(&mut self, usage: &TokenUsage, pricing: &ModelPricing) {
127        self.input_tokens += u64::from(usage.prompt_tokens);
128        if let Some(cached) = &mut self.cached_input_tokens {
129            *cached += u64::from(usage.cache_read_tokens);
130        }
131        self.output_tokens += u64::from(usage.completion_tokens);
132        match usage.cost(pricing) {
133            Some(cost) => self.total_cost += cost,
134            None => self.cost_unknown = true,
135        }
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142
143    #[test]
144    fn accepts_existing_uuid_and_saved_id() {
145        let uuid = Uuid::new_v4();
146        let id = SessionId::from(uuid);
147        assert_eq!(id.as_str(), uuid.to_string());
148        let saved = serde_json::to_string(&id).unwrap();
149        assert_eq!(serde_json::from_str::<SessionId>(&saved).unwrap(), id);
150        assert_eq!(
151            "saved-session".parse::<SessionId>().unwrap().as_str(),
152            "saved-session"
153        );
154    }
155
156    #[test]
157    fn enforces_length_when_parsing_and_deserializing() {
158        assert_eq!("".parse::<SessionId>(), Err(SessionIdError::Empty));
159        assert!("x".repeat(256).parse::<SessionId>().is_ok());
160        assert!("é".repeat(256).parse::<SessionId>().is_ok());
161        let too_long = "x".repeat(257);
162        assert_eq!(too_long.parse::<SessionId>(), Err(SessionIdError::TooLong));
163        assert!(serde_json::from_value::<SessionId>(serde_json::json!(too_long)).is_err());
164    }
165}