flashmind_types/
session.rs1use 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#[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 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#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
94pub struct SessionUsage {
95 pub input_tokens: u64,
97 #[serde(default)]
99 pub cached_input_tokens: Option<u64>,
100 pub output_tokens: u64,
102 #[serde(with = "rust_decimal::serde::str")]
104 pub total_cost: Decimal,
105 pub tokens_unknown: bool,
107 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 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}