use atm_core::{SessionDomain, StatusLineData};
use serde::Deserialize;
#[derive(Debug, Clone, Deserialize)]
pub struct RawStatusLine {
pub session_id: String,
#[serde(default)]
pub transcript_path: Option<String>,
#[serde(default)]
pub cwd: Option<String>,
#[serde(default)]
pub model: Option<RawModel>,
#[serde(default)]
pub workspace: Option<RawWorkspace>,
#[serde(default)]
pub version: Option<String>,
#[serde(default)]
pub cost: Option<RawCost>,
#[serde(default)]
pub context_window: Option<RawContextWindow>,
#[serde(default)]
pub exceeds_200k_tokens: Option<bool>,
#[serde(default)]
pub pid: Option<u32>,
#[serde(default)]
pub tmux_pane: Option<String>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct RawModel {
pub id: String,
#[serde(default)]
pub display_name: Option<String>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct RawWorkspace {
#[serde(default)]
pub current_dir: Option<String>,
#[serde(default)]
pub project_dir: Option<String>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct RawCost {
pub total_cost_usd: f64,
pub total_duration_ms: u64,
#[serde(default)]
pub total_api_duration_ms: u64,
#[serde(default)]
pub total_lines_added: u64,
#[serde(default)]
pub total_lines_removed: u64,
}
#[derive(Debug, Clone, Deserialize)]
pub struct RawContextWindow {
#[serde(default)]
pub total_input_tokens: u64,
#[serde(default)]
pub total_output_tokens: u64,
#[serde(default = "default_context_window_size")]
pub context_window_size: u32,
#[serde(default)]
pub used_percentage: Option<f64>,
#[serde(default)]
pub remaining_percentage: Option<f64>,
#[serde(default)]
pub current_usage: Option<RawCurrentUsage>,
}
fn default_context_window_size() -> u32 {
200_000
}
#[derive(Debug, Clone, Deserialize)]
pub struct RawCurrentUsage {
#[serde(default)]
pub input_tokens: u64,
#[serde(default)]
pub output_tokens: u64,
#[serde(default)]
pub cache_creation_input_tokens: u64,
#[serde(default)]
pub cache_read_input_tokens: u64,
}
impl RawStatusLine {
pub fn to_status_line_data(&self) -> Option<StatusLineData> {
let model = self.model.as_ref()?;
let cost = self.cost.as_ref();
let context = self.context_window.as_ref();
let current = context.and_then(|c| c.current_usage.as_ref());
Some(StatusLineData {
session_id: self.session_id.clone(),
model_id: model.id.clone(),
model_display_name: model.display_name.clone(),
cost_usd: cost.map(|c| c.total_cost_usd).unwrap_or(0.0),
total_duration_ms: cost.map(|c| c.total_duration_ms).unwrap_or(0),
api_duration_ms: cost.map(|c| c.total_api_duration_ms).unwrap_or(0),
lines_added: cost.map(|c| c.total_lines_added).unwrap_or(0),
lines_removed: cost.map(|c| c.total_lines_removed).unwrap_or(0),
total_input_tokens: context.map(|c| c.total_input_tokens).unwrap_or(0),
total_output_tokens: context.map(|c| c.total_output_tokens).unwrap_or(0),
context_window_size: context.map(|c| c.context_window_size).unwrap_or(200_000),
current_input_tokens: current.map(|c| c.input_tokens).unwrap_or(0),
current_output_tokens: current.map(|c| c.output_tokens).unwrap_or(0),
cache_creation_tokens: current.map(|c| c.cache_creation_input_tokens).unwrap_or(0),
cache_read_tokens: current.map(|c| c.cache_read_input_tokens).unwrap_or(0),
cwd: self.cwd.clone(),
version: self.version.clone(),
})
}
pub fn to_session_domain(&self) -> Option<SessionDomain> {
let data = self.to_status_line_data()?;
Some(SessionDomain::from_status_line(&data))
}
pub fn update_session(&self, session: &mut SessionDomain) -> bool {
use atm_core::Model;
if let Some(model) = &self.model {
let parsed = Model::from_id(&model.id);
session.model = parsed;
if parsed.is_unknown() && !model.id.is_empty() {
session.model_display_override = Some(
model
.display_name
.clone()
.unwrap_or_else(|| atm_core::derive_display_name(&model.id)),
);
} else {
session.model_display_override = None;
}
}
let cost = self.cost.as_ref();
let context = self.context_window.as_ref();
let current = context.and_then(|c| c.current_usage.as_ref());
let data = StatusLineData {
session_id: self.session_id.clone(),
model_id: String::new(), model_display_name: None, cost_usd: cost.map(|c| c.total_cost_usd).unwrap_or(0.0),
total_duration_ms: cost.map(|c| c.total_duration_ms).unwrap_or(0),
api_duration_ms: cost.map(|c| c.total_api_duration_ms).unwrap_or(0),
lines_added: cost.map(|c| c.total_lines_added).unwrap_or(0),
lines_removed: cost.map(|c| c.total_lines_removed).unwrap_or(0),
total_input_tokens: context.map(|c| c.total_input_tokens).unwrap_or(0),
total_output_tokens: context.map(|c| c.total_output_tokens).unwrap_or(0),
context_window_size: context.map(|c| c.context_window_size).unwrap_or(200_000),
current_input_tokens: current.map(|c| c.input_tokens).unwrap_or(0),
current_output_tokens: current.map(|c| c.output_tokens).unwrap_or(0),
cache_creation_tokens: current.map(|c| c.cache_creation_input_tokens).unwrap_or(0),
cache_read_tokens: current.map(|c| c.cache_read_input_tokens).unwrap_or(0),
cwd: self.cwd.clone(),
version: self.version.clone(),
};
session.update_from_status_line(&data)
}
}
#[cfg(test)]
mod tests {
use super::*;
use atm_core::Model;
#[test]
fn test_raw_status_line_parsing() {
let json = r#"{
"session_id": "test-123",
"model": {"id": "claude-opus-4-5-20251101", "display_name": "Opus 4.5"},
"cost": {"total_cost_usd": 0.35, "total_duration_ms": 35000},
"context_window": {"total_input_tokens": 5000, "context_window_size": 200000}
}"#;
let raw: RawStatusLine = serde_json::from_str(json).unwrap();
let session = raw.to_session_domain().expect("should create session");
assert_eq!(session.id.as_str(), "test-123");
assert_eq!(session.model, Model::Opus45);
assert!((session.cost.as_usd() - 0.35).abs() < 0.001);
assert_eq!(session.context.total_input_tokens.as_u64(), 5000);
}
#[test]
fn test_raw_status_line_with_current_usage() {
let json = r#"{
"session_id": "test-456",
"model": {"id": "claude-sonnet-4-20250514"},
"cost": {"total_cost_usd": 0.10, "total_duration_ms": 10000},
"context_window": {
"total_input_tokens": 1000,
"total_output_tokens": 500,
"context_window_size": 200000,
"current_usage": {
"input_tokens": 200,
"output_tokens": 100,
"cache_creation_input_tokens": 50,
"cache_read_input_tokens": 25
}
}
}"#;
let raw: RawStatusLine = serde_json::from_str(json).unwrap();
let session = raw.to_session_domain().expect("should create session");
assert_eq!(session.context.current_input_tokens.as_u64(), 200);
assert_eq!(session.context.cache_creation_tokens.as_u64(), 50);
}
#[test]
fn test_raw_status_line_context_from_current_usage() {
let json = r#"{
"session_id": "test-pct",
"model": {"id": "claude-sonnet-4-20250514"},
"context_window": {
"total_input_tokens": 50000,
"total_output_tokens": 10000,
"context_window_size": 200000,
"current_usage": {
"input_tokens": 1000,
"output_tokens": 500,
"cache_creation_input_tokens": 2000,
"cache_read_input_tokens": 40000
}
}
}"#;
let raw: RawStatusLine = serde_json::from_str(json).unwrap();
let session = raw.to_session_domain().expect("should create session");
assert_eq!(session.context.context_tokens().as_u64(), 43_000);
assert!((session.context.usage_percentage() - 21.5).abs() < 0.01);
}
#[test]
fn test_raw_status_line_zero_without_current_usage() {
let json = r#"{
"session_id": "test-fallback",
"model": {"id": "claude-sonnet-4-20250514"},
"context_window": {
"total_input_tokens": 50000,
"total_output_tokens": 10000,
"context_window_size": 200000
}
}"#;
let raw: RawStatusLine = serde_json::from_str(json).unwrap();
let session = raw.to_session_domain().expect("should create session");
assert_eq!(session.context.context_tokens().as_u64(), 0);
assert!((session.context.usage_percentage() - 0.0).abs() < 0.01);
}
#[test]
fn test_raw_status_line_missing_model_returns_none() {
let json = r#"{"session_id": "test-789"}"#;
let raw: RawStatusLine = serde_json::from_str(json).unwrap();
assert!(raw.to_session_domain().is_none());
}
#[test]
fn test_update_session_fills_in_model() {
use atm_core::{AgentType, SessionDomain};
let mut session = SessionDomain::new(
atm_core::SessionId::new("test-discovered"),
AgentType::GeneralPurpose,
Model::Unknown,
);
assert_eq!(session.model, Model::Unknown);
let json = r#"{
"session_id": "test-discovered",
"model": {"id": "claude-opus-4-5-20251101"},
"cost": {"total_cost_usd": 0.50, "total_duration_ms": 10000}
}"#;
let raw: RawStatusLine = serde_json::from_str(json).unwrap();
raw.update_session(&mut session);
assert_eq!(session.model, Model::Opus45);
assert!(session.model_display_override.is_none());
}
#[test]
fn test_update_session_unknown_model_with_display_name() {
use atm_core::{AgentType, SessionDomain};
let mut session = SessionDomain::new(
atm_core::SessionId::new("test-non-anthropic"),
AgentType::GeneralPurpose,
Model::Unknown,
);
let json = r#"{
"session_id": "test-non-anthropic",
"model": {"id": "gpt-4o", "display_name": "GPT-4o"}
}"#;
let raw: RawStatusLine = serde_json::from_str(json).unwrap();
raw.update_session(&mut session);
assert_eq!(session.model, Model::Unknown);
assert_eq!(session.model_display_override.as_deref(), Some("GPT-4o"));
}
#[test]
fn test_update_session_unknown_model_without_display_name() {
use atm_core::{AgentType, SessionDomain};
let mut session = SessionDomain::new(
atm_core::SessionId::new("test-unknown"),
AgentType::GeneralPurpose,
Model::Unknown,
);
let json = r#"{
"session_id": "test-unknown",
"model": {"id": "gemini-1.5-pro"}
}"#;
let raw: RawStatusLine = serde_json::from_str(json).unwrap();
raw.update_session(&mut session);
assert_eq!(session.model, Model::Unknown);
assert_eq!(
session.model_display_override.as_deref(),
Some("gemini-1.5-pro")
);
}
#[test]
fn test_new_session_opus46() {
let json = r#"{
"session_id": "test-opus46",
"model": {"id": "claude-opus-4-6"}
}"#;
let raw: RawStatusLine = serde_json::from_str(json).unwrap();
let session = raw.to_session_domain().expect("should create session");
assert_eq!(session.model, Model::Opus46);
assert!(session.model_display_override.is_none());
}
#[test]
fn test_new_session_non_anthropic_model() {
let json = r#"{
"session_id": "test-gpt",
"model": {"id": "gpt-4o", "display_name": "GPT-4o"}
}"#;
let raw: RawStatusLine = serde_json::from_str(json).unwrap();
let session = raw.to_session_domain().expect("should create session");
assert_eq!(session.model, Model::Unknown);
assert_eq!(session.model_display_override.as_deref(), Some("GPT-4o"));
}
#[test]
fn test_raw_status_line_partial_data() {
let json = r#"{
"session_id": "test-partial",
"model": {"id": "claude-sonnet-4-20250514"}
}"#;
let raw: RawStatusLine = serde_json::from_str(json).unwrap();
let session = raw
.to_session_domain()
.expect("should create with defaults");
assert_eq!(session.id.as_str(), "test-partial");
assert!((session.cost.as_usd() - 0.0).abs() < 0.001);
assert_eq!(session.context.total_input_tokens.as_u64(), 0);
}
}