use chrono::{DateTime, Utc};
use klieo_core::ids::RunId;
use serde::{Deserialize, Serialize};
use std::time::Duration;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum RunStatus {
Running,
Completed,
Failed,
Cancelled,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum StepKind {
LlmCall,
ToolCall,
SummaryCheckpoint,
OpsEvent,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct Step {
pub idx: u32,
pub kind: StepKind,
pub name: Option<String>,
pub input: serde_json::Value,
pub output: serde_json::Value,
pub error: Option<String>,
#[serde(with = "duration_millis")]
#[cfg_attr(feature = "schemars", schemars(with = "u64"))]
pub latency: Duration,
pub span_id: Option<String>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct LlmIo {
pub prompt: String,
pub completion: String,
pub model: Option<String>,
#[serde(default)]
pub provider: Option<String>,
#[serde(default)]
pub prompt_tokens: Option<u32>,
#[serde(default)]
pub completion_tokens: Option<u32>,
}
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct Usage {
pub prompt_tokens: u32,
pub completion_tokens: u32,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct Cost {
pub prompt_usd: f64,
pub completion_usd: f64,
pub total_usd: f64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct RunLog {
#[cfg_attr(feature = "schemars", schemars(with = "String"))]
pub run_id: RunId,
pub agent: String,
pub started_at: DateTime<Utc>,
pub finished_at: Option<DateTime<Utc>>,
pub status: RunStatus,
pub steps: Vec<Step>,
pub tokens: Usage,
pub cost_estimate: Option<Cost>,
}
mod duration_millis {
use serde::{Deserialize, Deserializer, Serializer};
use std::time::Duration;
pub fn serialize<S: Serializer>(d: &Duration, s: S) -> Result<S::Ok, S::Error> {
s.serialize_u64(d.as_millis() as u64)
}
pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Duration, D::Error> {
let ms = u64::deserialize(d)?;
Ok(Duration::from_millis(ms))
}
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::Utc;
use klieo_core::ids::RunId;
use std::time::Duration;
#[test]
fn run_status_serialises_snake_case() {
let json = serde_json::to_string(&RunStatus::Running).unwrap();
assert_eq!(json, "\"running\"");
}
#[test]
fn step_kind_round_trip() {
for k in [StepKind::LlmCall, StepKind::ToolCall] {
let json = serde_json::to_string(&k).unwrap();
let back: StepKind = serde_json::from_str(&json).unwrap();
assert_eq!(k, back);
}
}
#[test]
fn step_round_trip() {
let step = Step {
idx: 0,
kind: StepKind::ToolCall,
name: Some("calc".into()),
input: serde_json::json!({"a": 1}),
output: serde_json::json!({"b": 2}),
error: None,
latency: Duration::from_millis(42),
span_id: Some("span-xyz".into()),
};
let json = serde_json::to_string(&step).unwrap();
let back: Step = serde_json::from_str(&json).unwrap();
assert_eq!(step.idx, back.idx);
assert_eq!(step.latency, back.latency);
assert_eq!(step.name, back.name);
}
#[test]
fn run_log_round_trip_empty_steps() {
let log = RunLog {
run_id: RunId::new(),
agent: "writer".into(),
started_at: Utc::now(),
finished_at: None,
status: RunStatus::Running,
steps: vec![],
tokens: Usage::default(),
cost_estimate: None,
};
let json = serde_json::to_string(&log).unwrap();
assert!(json.contains("\"steps\":[]"));
let _back: RunLog = serde_json::from_str(&json).unwrap();
}
#[test]
fn usage_default_is_zero() {
let u = Usage::default();
assert_eq!(u.prompt_tokens, 0);
assert_eq!(u.completion_tokens, 0);
}
#[test]
fn cost_round_trip() {
let c = Cost {
prompt_usd: 0.001,
completion_usd: 0.002,
total_usd: 0.003,
};
let json = serde_json::to_string(&c).unwrap();
let back: Cost = serde_json::from_str(&json).unwrap();
assert!((back.total_usd - 0.003).abs() < 1e-9);
}
}