kaizen/core/
trace_span.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2//! Additive trace-span model for Datadog/OTLP-style session timelines.
3
4use serde::{Deserialize, Serialize};
5use serde_json::{Value, json};
6
7#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
8pub enum TraceSpanKind {
9    Session,
10    Agent,
11    Step,
12    Llm,
13    Tool,
14    Permission,
15}
16
17impl TraceSpanKind {
18    pub fn as_str(&self) -> &'static str {
19        match self {
20            Self::Session => "session",
21            Self::Agent => "agent",
22            Self::Step => "step",
23            Self::Llm => "llm",
24            Self::Tool => "tool",
25            Self::Permission => "permission",
26        }
27    }
28
29    pub fn parse(s: &str) -> Self {
30        match s {
31            "session" => Self::Session,
32            "agent" => Self::Agent,
33            "step" => Self::Step,
34            "tool" => Self::Tool,
35            "permission" => Self::Permission,
36            _ => Self::Llm,
37        }
38    }
39}
40
41#[derive(Debug, Clone, Serialize, Deserialize)]
42pub struct TraceSpanRecord {
43    pub span_id: String,
44    pub trace_id: String,
45    pub parent_span_id: Option<String>,
46    pub session_id: String,
47    pub kind: TraceSpanKind,
48    pub name: String,
49    pub status: String,
50    pub started_at_ms: Option<u64>,
51    pub ended_at_ms: Option<u64>,
52    pub duration_ms: Option<u32>,
53    pub model: Option<String>,
54    pub tool: Option<String>,
55    pub tokens_in: Option<u32>,
56    pub tokens_out: Option<u32>,
57    pub reasoning_tokens: Option<u32>,
58    pub cost_usd_e6: Option<i64>,
59    pub context_used_tokens: Option<u32>,
60    pub context_max_tokens: Option<u32>,
61    pub payload: Value,
62}
63
64impl TraceSpanRecord {
65    pub fn llm_proxy(
66        session_id: &str,
67        seq: u64,
68        started_at_ms: u64,
69        ended_at_ms: u64,
70        model: Option<String>,
71        payload: Value,
72    ) -> Self {
73        let duration = ended_at_ms.saturating_sub(started_at_ms);
74        Self {
75            span_id: format!("llm-{session_id}-{seq}"),
76            trace_id: trace_id_for_session(session_id),
77            parent_span_id: Some(format!("step-{session_id}-{seq}")),
78            session_id: session_id.to_string(),
79            kind: Self::llm_kind(),
80            name: "llm.proxy".into(),
81            status: "ok".into(),
82            started_at_ms: Some(started_at_ms),
83            ended_at_ms: Some(ended_at_ms),
84            duration_ms: u32::try_from(duration).ok(),
85            model,
86            tool: None,
87            tokens_in: None,
88            tokens_out: None,
89            reasoning_tokens: None,
90            cost_usd_e6: None,
91            context_used_tokens: None,
92            context_max_tokens: None,
93            payload,
94        }
95    }
96
97    fn llm_kind() -> TraceSpanKind {
98        TraceSpanKind::Llm
99    }
100}
101
102pub fn trace_id_for_session(session_id: &str) -> String {
103    let hash = blake3::hash(session_id.as_bytes());
104    hex::encode(&hash.as_bytes()[..16])
105}
106
107pub fn span_payload(provider: &str, stream: bool, request_id: Option<&str>) -> Value {
108    json!({"provider": provider, "stream": stream, "request_id": request_id})
109}
kaizen/core/trace_span.rs

kaizen/core/
trace_span.rs
