use serde::{Deserialize, Serialize};
use std::path::PathBuf;
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
pub struct Session {
pub path: PathBuf,
pub format: String,
pub metadata: SessionMetadata,
pub turns: Vec<Turn>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub parent_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub agent_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub subagent_type: Option<String>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
pub struct SessionMetadata {
pub session_id: Option<String>,
pub timestamp: Option<String>,
pub provider: Option<String>,
pub model: Option<String>,
pub project: Option<String>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
pub struct Turn {
pub messages: Vec<Message>,
pub token_usage: Option<TokenUsage>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
pub struct Message {
pub role: Role,
pub content: Vec<ContentBlock>,
pub timestamp: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "lowercase")]
pub enum Role {
User,
Assistant,
System,
Tool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ContentBlock {
Text { text: String },
ToolUse {
id: String,
name: String,
input: serde_json::Value,
},
ToolResult {
tool_use_id: String,
content: String,
is_error: bool,
},
Thinking { text: String },
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
pub struct TokenUsage {
pub input: u64,
pub output: u64,
pub cache_read: Option<u64>,
pub cache_create: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
}
impl Session {
pub fn new(path: PathBuf, format: impl Into<String>) -> Self {
Self {
path,
format: format.into(),
metadata: SessionMetadata::default(),
turns: Vec::new(),
parent_id: None,
agent_id: None,
subagent_type: None,
}
}
pub fn is_subagent(&self) -> bool {
self.parent_id.is_some()
}
pub fn message_count(&self) -> usize {
self.turns.iter().map(|t| t.messages.len()).sum()
}
pub fn messages_by_role(&self, role: Role) -> usize {
self.turns
.iter()
.flat_map(|t| &t.messages)
.filter(|m| m.role == role)
.count()
}
pub fn tool_uses(&self) -> impl Iterator<Item = (&str, &serde_json::Value)> {
self.turns.iter().flat_map(|t| &t.messages).flat_map(|m| {
m.content.iter().filter_map(|block| match block {
ContentBlock::ToolUse { name, input, .. } => Some((name.as_str(), input)),
_ => None,
})
})
}
pub fn tool_results(&self) -> impl Iterator<Item = (&str, bool)> {
self.turns.iter().flat_map(|t| &t.messages).flat_map(|m| {
m.content.iter().filter_map(|block| match block {
ContentBlock::ToolResult {
content, is_error, ..
} => Some((content.as_str(), *is_error)),
_ => None,
})
})
}
pub fn total_tokens(&self) -> TokenUsage {
let mut total = TokenUsage::default();
for turn in &self.turns {
if let Some(usage) = &turn.token_usage {
total.input += usage.input;
total.output += usage.output;
if let Some(cache_read) = usage.cache_read {
*total.cache_read.get_or_insert(0) += cache_read;
}
if let Some(cache_create) = usage.cache_create {
*total.cache_create.get_or_insert(0) += cache_create;
}
}
}
total
}
}
impl std::fmt::Display for Role {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Role::User => write!(f, "user"),
Role::Assistant => write!(f, "assistant"),
Role::System => write!(f, "system"),
Role::Tool => write!(f, "tool"),
}
}
}