1use std::collections::BTreeMap;
2use std::fmt::{Display, Formatter};
3use std::fs;
4use std::path::Path;
5
6use serde::{Deserialize, Serialize};
7
8use crate::json::{JsonError, JsonValue};
9use crate::usage::TokenUsage;
10
11#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
12#[serde(rename_all = "snake_case")]
13pub enum MessageRole {
14 System,
15 User,
16 Assistant,
17 Tool,
18}
19
20#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
21#[serde(tag = "type", rename_all = "snake_case")]
22pub enum ContentBlock {
23 Text {
24 text: String,
25 },
26 Image {
27 media_type: String,
28 data: String,
29 },
30 ToolUse {
31 id: String,
32 name: String,
33 input: String,
34 },
35 ToolResult {
36 tool_use_id: String,
37 tool_name: String,
38 output: String,
39 is_error: bool,
40 },
41}
42
43#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
44pub struct ConversationMessage {
45 pub role: MessageRole,
46 pub blocks: Vec<ContentBlock>,
47 pub usage: Option<TokenUsage>,
48}
49
50#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
51pub struct Session {
52 pub version: u32,
53 pub messages: Vec<ConversationMessage>,
54}
55
56#[derive(Debug)]
57pub enum SessionError {
58 Io(std::io::Error),
59 Json(JsonError),
60 Format(String),
61}
62
63impl Display for SessionError {
64 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
65 match self {
66 Self::Io(error) => write!(f, "{error}"),
67 Self::Json(error) => write!(f, "{error}"),
68 Self::Format(error) => write!(f, "{error}"),
69 }
70 }
71}
72
73impl std::error::Error for SessionError {}
74
75impl From<std::io::Error> for SessionError {
76 fn from(value: std::io::Error) -> Self {
77 Self::Io(value)
78 }
79}
80
81impl From<JsonError> for SessionError {
82 fn from(value: JsonError) -> Self {
83 Self::Json(value)
84 }
85}
86
87impl Session {
88 #[must_use]
89 pub fn new() -> Self {
90 Self {
91 version: 1,
92 messages: Vec::new(),
93 }
94 }
95
96 pub fn save_to_path(&self, path: impl AsRef<Path>) -> Result<(), SessionError> {
97 fs::write(path, self.to_json().render())?;
98 Ok(())
99 }
100
101 pub fn load_from_path(path: impl AsRef<Path>) -> Result<Self, SessionError> {
102 let contents = fs::read_to_string(path)?;
103 Self::from_json(&JsonValue::parse(&contents)?)
104 }
105
106 #[must_use]
107 pub fn to_json(&self) -> JsonValue {
108 let mut object = BTreeMap::new();
109 object.insert(
110 "version".to_string(),
111 JsonValue::Number(i64::from(self.version)),
112 );
113 object.insert(
114 "messages".to_string(),
115 JsonValue::Array(
116 self.messages
117 .iter()
118 .map(ConversationMessage::to_json)
119 .collect(),
120 ),
121 );
122 JsonValue::Object(object)
123 }
124
125 pub fn from_json(value: &JsonValue) -> Result<Self, SessionError> {
126 let object = value
127 .as_object()
128 .ok_or_else(|| SessionError::Format("session must be an object".to_string()))?;
129 let version = object
130 .get("version")
131 .and_then(JsonValue::as_i64)
132 .ok_or_else(|| SessionError::Format("missing version".to_string()))?;
133 let version = u32::try_from(version)
134 .map_err(|_| SessionError::Format("version out of range".to_string()))?;
135 let messages = object
136 .get("messages")
137 .and_then(JsonValue::as_array)
138 .ok_or_else(|| SessionError::Format("missing messages".to_string()))?
139 .iter()
140 .map(ConversationMessage::from_json)
141 .collect::<Result<Vec<_>, _>>()?;
142 Ok(Self { version, messages })
143 }
144}
145
146impl Default for Session {
147 fn default() -> Self {
148 Self::new()
149 }
150}
151
152impl ConversationMessage {
153 #[must_use]
154 pub fn user_text(text: impl Into<String>) -> Self {
155 Self {
156 role: MessageRole::User,
157 blocks: vec![ContentBlock::Text { text: text.into() }],
158 usage: None,
159 }
160 }
161
162 #[must_use]
163 pub fn user_blocks(blocks: Vec<ContentBlock>) -> Self {
164 Self {
165 role: MessageRole::User,
166 blocks,
167 usage: None,
168 }
169 }
170
171 #[must_use]
172 pub fn assistant(blocks: Vec<ContentBlock>) -> Self {
173 Self {
174 role: MessageRole::Assistant,
175 blocks,
176 usage: None,
177 }
178 }
179
180 #[must_use]
181 pub fn assistant_with_usage(blocks: Vec<ContentBlock>, usage: Option<TokenUsage>) -> Self {
182 Self {
183 role: MessageRole::Assistant,
184 blocks,
185 usage,
186 }
187 }
188
189 #[must_use]
190 pub fn tool_result(
191 tool_use_id: impl Into<String>,
192 tool_name: impl Into<String>,
193 output: impl Into<String>,
194 is_error: bool,
195 ) -> Self {
196 Self {
197 role: MessageRole::Tool,
198 blocks: vec![ContentBlock::ToolResult {
199 tool_use_id: tool_use_id.into(),
200 tool_name: tool_name.into(),
201 output: output.into(),
202 is_error,
203 }],
204 usage: None,
205 }
206 }
207
208 #[must_use]
209 pub fn to_json(&self) -> JsonValue {
210 let mut object = BTreeMap::new();
211 object.insert(
212 "role".to_string(),
213 JsonValue::String(
214 match self.role {
215 MessageRole::System => "system",
216 MessageRole::User => "user",
217 MessageRole::Assistant => "assistant",
218 MessageRole::Tool => "tool",
219 }
220 .to_string(),
221 ),
222 );
223 object.insert(
224 "blocks".to_string(),
225 JsonValue::Array(self.blocks.iter().map(ContentBlock::to_json).collect()),
226 );
227 if let Some(usage) = self.usage {
228 object.insert("usage".to_string(), usage_to_json(usage));
229 }
230 JsonValue::Object(object)
231 }
232
233 fn from_json(value: &JsonValue) -> Result<Self, SessionError> {
234 let object = value
235 .as_object()
236 .ok_or_else(|| SessionError::Format("message must be an object".to_string()))?;
237 let role = match object
238 .get("role")
239 .and_then(JsonValue::as_str)
240 .ok_or_else(|| SessionError::Format("missing role".to_string()))?
241 {
242 "system" => MessageRole::System,
243 "user" => MessageRole::User,
244 "assistant" => MessageRole::Assistant,
245 "tool" => MessageRole::Tool,
246 other => {
247 return Err(SessionError::Format(format!(
248 "unsupported message role: {other}"
249 )))
250 }
251 };
252 let blocks = object
253 .get("blocks")
254 .and_then(JsonValue::as_array)
255 .ok_or_else(|| SessionError::Format("missing blocks".to_string()))?
256 .iter()
257 .map(ContentBlock::from_json)
258 .collect::<Result<Vec<_>, _>>()?;
259 let usage = object.get("usage").map(usage_from_json).transpose()?;
260 Ok(Self {
261 role,
262 blocks,
263 usage,
264 })
265 }
266}
267
268impl ContentBlock {
269 #[must_use]
270 pub fn to_json(&self) -> JsonValue {
271 let mut object = BTreeMap::new();
272 match self {
273 Self::Text { text } => {
274 object.insert("type".to_string(), JsonValue::String("text".to_string()));
275 object.insert("text".to_string(), JsonValue::String(text.clone()));
276 }
277 Self::Image { media_type, data } => {
278 object.insert("type".to_string(), JsonValue::String("image".to_string()));
279 object.insert(
280 "media_type".to_string(),
281 JsonValue::String(media_type.clone()),
282 );
283 object.insert("data".to_string(), JsonValue::String(data.clone()));
284 }
285 Self::ToolUse { id, name, input } => {
286 object.insert(
287 "type".to_string(),
288 JsonValue::String("tool_use".to_string()),
289 );
290 object.insert("id".to_string(), JsonValue::String(id.clone()));
291 object.insert("name".to_string(), JsonValue::String(name.clone()));
292 object.insert("input".to_string(), JsonValue::String(input.clone()));
293 }
294 Self::ToolResult {
295 tool_use_id,
296 tool_name,
297 output,
298 is_error,
299 } => {
300 object.insert(
301 "type".to_string(),
302 JsonValue::String("tool_result".to_string()),
303 );
304 object.insert(
305 "tool_use_id".to_string(),
306 JsonValue::String(tool_use_id.clone()),
307 );
308 object.insert(
309 "tool_name".to_string(),
310 JsonValue::String(tool_name.clone()),
311 );
312 object.insert("output".to_string(), JsonValue::String(output.clone()));
313 object.insert("is_error".to_string(), JsonValue::Bool(*is_error));
314 }
315 }
316 JsonValue::Object(object)
317 }
318
319 fn from_json(value: &JsonValue) -> Result<Self, SessionError> {
320 let object = value
321 .as_object()
322 .ok_or_else(|| SessionError::Format("block must be an object".to_string()))?;
323 match object
324 .get("type")
325 .and_then(JsonValue::as_str)
326 .ok_or_else(|| SessionError::Format("missing block type".to_string()))?
327 {
328 "text" => Ok(Self::Text {
329 text: required_string(object, "text")?,
330 }),
331 "image" => Ok(Self::Image {
332 media_type: required_string(object, "media_type")?,
333 data: required_string(object, "data")?,
334 }),
335 "tool_use" => Ok(Self::ToolUse {
336 id: required_string(object, "id")?,
337 name: required_string(object, "name")?,
338 input: required_string(object, "input")?,
339 }),
340 "tool_result" => Ok(Self::ToolResult {
341 tool_use_id: required_string(object, "tool_use_id")?,
342 tool_name: required_string(object, "tool_name")?,
343 output: required_string(object, "output")?,
344 is_error: object
345 .get("is_error")
346 .and_then(JsonValue::as_bool)
347 .ok_or_else(|| SessionError::Format("missing is_error".to_string()))?,
348 }),
349 other => Err(SessionError::Format(format!(
350 "unsupported block type: {other}"
351 ))),
352 }
353 }
354}
355
356fn usage_to_json(usage: TokenUsage) -> JsonValue {
357 let mut object = BTreeMap::new();
358 object.insert(
359 "input_tokens".to_string(),
360 JsonValue::Number(i64::from(usage.input_tokens)),
361 );
362 object.insert(
363 "output_tokens".to_string(),
364 JsonValue::Number(i64::from(usage.output_tokens)),
365 );
366 object.insert(
367 "cache_creation_input_tokens".to_string(),
368 JsonValue::Number(i64::from(usage.cache_creation_input_tokens)),
369 );
370 object.insert(
371 "cache_read_input_tokens".to_string(),
372 JsonValue::Number(i64::from(usage.cache_read_input_tokens)),
373 );
374 JsonValue::Object(object)
375}
376
377fn usage_from_json(value: &JsonValue) -> Result<TokenUsage, SessionError> {
378 let object = value
379 .as_object()
380 .ok_or_else(|| SessionError::Format("usage must be an object".to_string()))?;
381 Ok(TokenUsage {
382 input_tokens: required_u32(object, "input_tokens")?,
383 output_tokens: required_u32(object, "output_tokens")?,
384 cache_creation_input_tokens: required_u32(object, "cache_creation_input_tokens")?,
385 cache_read_input_tokens: required_u32(object, "cache_read_input_tokens")?,
386 })
387}
388
389fn required_string(
390 object: &BTreeMap<String, JsonValue>,
391 key: &str,
392) -> Result<String, SessionError> {
393 object
394 .get(key)
395 .and_then(JsonValue::as_str)
396 .map(ToOwned::to_owned)
397 .ok_or_else(|| SessionError::Format(format!("missing {key}")))
398}
399
400fn required_u32(object: &BTreeMap<String, JsonValue>, key: &str) -> Result<u32, SessionError> {
401 let value = object
402 .get(key)
403 .and_then(JsonValue::as_i64)
404 .ok_or_else(|| SessionError::Format(format!("missing {key}")))?;
405 u32::try_from(value).map_err(|_| SessionError::Format(format!("{key} out of range")))
406}
407
408#[cfg(test)]
409mod tests {
410 use super::{ContentBlock, ConversationMessage, MessageRole, Session};
411 use crate::usage::TokenUsage;
412 use std::fs;
413 use std::path::Path;
414 use std::time::{SystemTime, UNIX_EPOCH};
415
416 #[test]
417 fn persists_and_restores_session_json() {
418 let mut session = Session::new();
419 session
420 .messages
421 .push(ConversationMessage::user_text("hello"));
422 session
423 .messages
424 .push(ConversationMessage::assistant_with_usage(
425 vec![
426 ContentBlock::Text {
427 text: "thinking".to_string(),
428 },
429 ContentBlock::ToolUse {
430 id: "tool-1".to_string(),
431 name: "bash".to_string(),
432 input: "echo hi".to_string(),
433 },
434 ],
435 Some(TokenUsage {
436 input_tokens: 10,
437 output_tokens: 4,
438 cache_creation_input_tokens: 1,
439 cache_read_input_tokens: 2,
440 }),
441 ));
442 session.messages.push(ConversationMessage::tool_result(
443 "tool-1", "bash", "hi", false,
444 ));
445
446 let nanos = SystemTime::now()
447 .duration_since(UNIX_EPOCH)
448 .expect("system time should be after epoch")
449 .as_nanos();
450 let path = std::env::temp_dir().join(format!("runtime-session-{nanos}.json"));
451 session.save_to_path(&path).expect("session should save");
452 let restored = Session::load_from_path(&path).expect("session should load");
453 fs::remove_file(&path).expect("temp file should be removable");
454
455 assert_eq!(restored, session);
456 assert_eq!(restored.messages[2].role, MessageRole::Tool);
457 assert_eq!(
458 restored.messages[1].usage.expect("usage").total_tokens(),
459 17
460 );
461 }
462
463 #[test]
464 fn round_trips_system_role_message() {
465 let json_str = r#"{"version":1,"messages":[{"role":"system","blocks":[{"type":"text","text":"sys prompt"}]}]}"#;
466 let parsed = crate::json::JsonValue::parse(json_str).unwrap();
467 let session = Session::from_json(&parsed).unwrap();
468 assert_eq!(session.messages[0].role, MessageRole::System);
469
470 let rendered = session.to_json();
471 let restored = Session::from_json(&rendered).unwrap();
472 assert_eq!(restored, session);
473 }
474
475 #[test]
476 fn rejects_unsupported_message_role() {
477 let json_str = r#"{"version":1,"messages":[{"role":"admin","blocks":[]}]}"#;
478 let parsed = crate::json::JsonValue::parse(json_str).unwrap();
479 let err = Session::from_json(&parsed).unwrap_err();
480 assert!(err.to_string().contains("unsupported message role"));
481 }
482
483 #[test]
484 fn rejects_unsupported_block_type() {
485 let json_str =
486 r#"{"version":1,"messages":[{"role":"user","blocks":[{"type":"video","url":"x"}]}]}"#;
487 let parsed = crate::json::JsonValue::parse(json_str).unwrap();
488 let err = Session::from_json(&parsed).unwrap_err();
489 assert!(err.to_string().contains("unsupported block type"));
490 }
491
492 #[test]
493 fn rejects_missing_version() {
494 let json_str = r#"{"messages":[]}"#;
495 let parsed = crate::json::JsonValue::parse(json_str).unwrap();
496 let err = Session::from_json(&parsed).unwrap_err();
497 assert!(err.to_string().contains("version"));
498 }
499
500 #[test]
501 fn rejects_non_object_root() {
502 let parsed = crate::json::JsonValue::parse("[1,2]").unwrap();
503 let err = Session::from_json(&parsed).unwrap_err();
504 assert!(err.to_string().contains("object"));
505 }
506
507 #[test]
508 fn load_from_nonexistent_path_returns_io_error() {
509 let err = Session::load_from_path(Path::new("/nonexistent/path/session.json")).unwrap_err();
510 assert!(matches!(err, super::SessionError::Io(_)));
511 }
512
513 #[test]
514 fn session_error_display_covers_all_variants() {
515 let io_err = super::SessionError::Io(std::io::Error::new(
516 std::io::ErrorKind::NotFound,
517 "not found",
518 ));
519 assert!(io_err.to_string().contains("not found"));
520
521 let json_err = super::SessionError::Json(crate::json::JsonError::new("bad"));
522 assert!(json_err.to_string().contains("bad"));
523
524 let fmt_err = super::SessionError::Format("bad format".into());
525 assert!(fmt_err.to_string().contains("bad format"));
526 }
527
528 #[test]
529 fn round_trips_image_content_block() {
530 let mut session = Session::new();
531 session.messages.push(ConversationMessage::user_blocks(vec![
532 ContentBlock::Text {
533 text: "describe this image".to_string(),
534 },
535 ContentBlock::Image {
536 media_type: "image/png".to_string(),
537 data: "iVBOR...".to_string(),
538 },
539 ]));
540
541 let nanos = SystemTime::now()
542 .duration_since(UNIX_EPOCH)
543 .expect("system time should be after epoch")
544 .as_nanos();
545 let path = std::env::temp_dir().join(format!("runtime-session-img-{nanos}.json"));
546 session.save_to_path(&path).expect("session should save");
547 let restored = Session::load_from_path(&path).expect("session should load");
548 fs::remove_file(&path).expect("temp file should be removable");
549
550 assert_eq!(restored, session);
551 assert_eq!(restored.messages[0].blocks.len(), 2);
552 match &restored.messages[0].blocks[1] {
553 ContentBlock::Image { media_type, data } => {
554 assert_eq!(media_type, "image/png");
555 assert_eq!(data, "iVBOR...");
556 }
557 other => panic!("expected Image block, got: {other:?}"),
558 }
559 }
560}