use crate::types::RunId;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use thiserror::Error;
pub mod sqlite;
pub use sqlite::SqliteRunTraceStore;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
pub struct TokenUsage {
pub input_tokens: u64,
pub output_tokens: u64,
pub total_tokens: u64,
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
pub struct WorkerStats {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub worker_kind: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub usage: Option<TokenUsage>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub num_turns: Option<u32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub adapter_data: Option<Value>,
}
impl WorkerStats {
pub fn is_empty(&self) -> bool {
self.worker_kind.is_none()
&& self.model.is_none()
&& self.usage.is_none()
&& self.num_turns.is_none()
&& self.adapter_data.is_none()
}
}
pub const TRACE_PAYLOAD_CAP_BYTES: usize = 8 * 1024;
pub const DEFAULT_TRACE_MAX_EVENTS_PER_RUN: usize = 10_000;
pub const DEFAULT_TRACE_LIST_LIMIT: usize = 1_000;
pub mod kind {
pub const RUN_STARTED: &str = "core.run_started";
pub const STEP_DISPATCHED: &str = "core.step_dispatched";
pub const STEP_COMPLETED: &str = "core.step_completed";
pub const RUN_FINISHED: &str = "core.run_finished";
pub const CANCEL_REQUESTED: &str = "core.cancel_requested";
pub const LONG_HOLD_WARN: &str = "mw.long_hold_warn";
pub const WORKER_DEGRADATION: &str = "worker.degradation";
}
#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
pub struct TraceEvent {
#[schemars(with = "String")]
pub run_id: RunId,
pub seq: u64,
pub ts_ms: i64,
pub kind: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub step_ref: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub attempt: Option<u32>,
pub payload: Value,
}
#[derive(Debug, Clone)]
pub struct TraceEventDraft {
pub kind: String,
pub step_ref: Option<String>,
pub attempt: Option<u32>,
pub payload: Value,
}
#[derive(Debug, Clone, Default)]
pub struct TraceQuery {
pub after: Option<u64>,
pub limit: Option<usize>,
pub latest: Option<usize>,
pub kinds: Vec<String>,
pub step_ref: Option<String>,
pub attempt: Option<u32>,
}
impl TraceQuery {
fn matches(&self, event: &TraceEvent) -> bool {
if !self.kinds.is_empty()
&& !self
.kinds
.iter()
.any(|k| event.kind == *k || event.kind.starts_with(k.as_str()))
{
return false;
}
if let Some(step_ref) = &self.step_ref {
if event.step_ref.as_deref() != Some(step_ref.as_str()) {
return false;
}
}
if let Some(attempt) = self.attempt {
if event.attempt != Some(attempt) {
return false;
}
}
true
}
fn page(&self, mut events: Vec<TraceEvent>) -> Vec<TraceEvent> {
if let Some(n) = self.latest {
let start = events.len().saturating_sub(n);
return events.split_off(start);
}
if let Some(after) = self.after {
events.retain(|e| e.seq > after);
}
let limit = self.limit.unwrap_or(DEFAULT_TRACE_LIST_LIMIT);
events.truncate(limit);
events
}
}
pub fn cap_payload(payload: Value) -> Value {
let serialized = payload.to_string();
if serialized.len() <= TRACE_PAYLOAD_CAP_BYTES {
return payload;
}
let head: String = serialized.chars().take(1024).collect();
serde_json::json!({
"truncated": true,
"size_bytes": serialized.len(),
"head": head,
})
}
#[derive(Debug, Error)]
pub enum TraceStoreError {
#[error("other: {0}")]
Other(String),
}
#[async_trait]
pub trait RunTraceStore: Send + Sync {
fn name(&self) -> &str;
async fn append(
&self,
run_id: &RunId,
draft: TraceEventDraft,
) -> Result<TraceEvent, TraceStoreError>;
async fn list(
&self,
run_id: &RunId,
query: &TraceQuery,
) -> Result<Vec<TraceEvent>, TraceStoreError>;
async fn delete_run(&self, run_id: &RunId) -> Result<u64, TraceStoreError>;
}
#[derive(Clone)]
pub struct TraceHandle {
run_id: RunId,
store: Arc<dyn RunTraceStore>,
}
impl TraceHandle {
pub fn new(run_id: RunId, store: Arc<dyn RunTraceStore>) -> Self {
Self { run_id, store }
}
pub fn run_id(&self) -> &RunId {
&self.run_id
}
pub async fn append(
&self,
kind: &str,
step_ref: Option<&str>,
attempt: Option<u32>,
payload: Value,
) {
let draft = TraceEventDraft {
kind: kind.to_string(),
step_ref: step_ref.map(str::to_string),
attempt,
payload,
};
if let Err(e) = self.store.append(&self.run_id, draft).await {
tracing::warn!(
run_id = %self.run_id,
kind = kind,
error = %e,
"TraceHandle::append failed (swallowed — trace is observational)"
);
}
}
}
impl std::fmt::Debug for TraceHandle {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TraceHandle")
.field("run_id", &self.run_id)
.field("store", &self.store.name())
.finish()
}
}
#[derive(Default)]
struct TraceInner {
events: HashMap<RunId, Vec<TraceEvent>>,
next_seq: HashMap<RunId, u64>,
}
pub struct InMemoryRunTraceStore {
inner: Mutex<TraceInner>,
max_events_per_run: usize,
}
impl InMemoryRunTraceStore {
pub fn new() -> Self {
Self {
inner: Mutex::new(TraceInner::default()),
max_events_per_run: DEFAULT_TRACE_MAX_EVENTS_PER_RUN,
}
}
pub fn with_max_events_per_run(max: usize) -> Self {
Self {
inner: Mutex::new(TraceInner::default()),
max_events_per_run: max,
}
}
}
impl Default for InMemoryRunTraceStore {
fn default() -> Self {
Self::new()
}
}
pub(crate) fn now_unix_ms() -> i64 {
use std::time::{SystemTime, UNIX_EPOCH};
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_millis() as i64)
.unwrap_or(0)
}
#[async_trait]
impl RunTraceStore for InMemoryRunTraceStore {
fn name(&self) -> &str {
"in-memory"
}
async fn append(
&self,
run_id: &RunId,
draft: TraceEventDraft,
) -> Result<TraceEvent, TraceStoreError> {
let mut inner = self.inner.lock().unwrap();
let seq_slot = inner.next_seq.entry(run_id.clone()).or_insert(0);
*seq_slot += 1;
let event = TraceEvent {
run_id: run_id.clone(),
seq: *seq_slot,
ts_ms: now_unix_ms(),
kind: draft.kind,
step_ref: draft.step_ref,
attempt: draft.attempt,
payload: cap_payload(draft.payload),
};
let list = inner.events.entry(run_id.clone()).or_default();
list.push(event.clone());
if list.len() > self.max_events_per_run {
let overflow = list.len() - self.max_events_per_run;
list.drain(..overflow);
}
Ok(event)
}
async fn list(
&self,
run_id: &RunId,
query: &TraceQuery,
) -> Result<Vec<TraceEvent>, TraceStoreError> {
let inner = self.inner.lock().unwrap();
let events: Vec<TraceEvent> = inner
.events
.get(run_id)
.map(|list| list.iter().filter(|e| query.matches(e)).cloned().collect())
.unwrap_or_default();
Ok(query.page(events))
}
async fn delete_run(&self, run_id: &RunId) -> Result<u64, TraceStoreError> {
let mut inner = self.inner.lock().unwrap();
inner.next_seq.remove(run_id);
Ok(inner
.events
.remove(run_id)
.map(|list| list.len() as u64)
.unwrap_or(0))
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn rid(s: &str) -> RunId {
RunId::parse(s).unwrap()
}
fn draft(kind: &str, step_ref: Option<&str>, attempt: Option<u32>) -> TraceEventDraft {
TraceEventDraft {
kind: kind.to_string(),
step_ref: step_ref.map(str::to_string),
attempt,
payload: json!({"k": kind}),
}
}
#[tokio::test]
async fn append_assigns_monotonic_seq_per_run() {
let s = InMemoryRunTraceStore::new();
let e1 = s
.append(&rid("R-1"), draft("core.run_started", None, None))
.await
.unwrap();
let e2 = s
.append(
&rid("R-1"),
draft("core.step_dispatched", Some("w"), Some(1)),
)
.await
.unwrap();
let other = s
.append(&rid("R-2"), draft("core.run_started", None, None))
.await
.unwrap();
assert_eq!(e1.seq, 1);
assert_eq!(e2.seq, 2);
assert_eq!(other.seq, 1, "seq is per-Run, not global");
assert!(e1.ts_ms > 0);
}
#[tokio::test]
async fn list_filters_by_kind_prefix_step_and_attempt() {
let s = InMemoryRunTraceStore::new();
let r = rid("R-1");
s.append(&r, draft("core.run_started", None, None))
.await
.unwrap();
s.append(&r, draft("core.step_dispatched", Some("a"), Some(1)))
.await
.unwrap();
s.append(&r, draft("mw.long_hold_warn", Some("a"), Some(1)))
.await
.unwrap();
s.append(&r, draft("core.step_completed", Some("b"), Some(2)))
.await
.unwrap();
let mw = s
.list(
&r,
&TraceQuery {
kinds: vec!["mw.".into()],
..Default::default()
},
)
.await
.unwrap();
assert_eq!(mw.len(), 1);
assert_eq!(mw[0].kind, "mw.long_hold_warn");
let step_a = s
.list(
&r,
&TraceQuery {
step_ref: Some("a".into()),
..Default::default()
},
)
.await
.unwrap();
assert_eq!(step_a.len(), 2);
let attempt2 = s
.list(
&r,
&TraceQuery {
attempt: Some(2),
..Default::default()
},
)
.await
.unwrap();
assert_eq!(attempt2.len(), 1);
assert_eq!(attempt2[0].step_ref.as_deref(), Some("b"));
}
#[tokio::test]
async fn list_paging_after_and_latest() {
let s = InMemoryRunTraceStore::new();
let r = rid("R-1");
for i in 0..5 {
s.append(&r, draft(&format!("core.e{i}"), None, None))
.await
.unwrap();
}
let after = s
.list(
&r,
&TraceQuery {
after: Some(3),
..Default::default()
},
)
.await
.unwrap();
assert_eq!(after.iter().map(|e| e.seq).collect::<Vec<_>>(), vec![4, 5]);
let latest = s
.list(
&r,
&TraceQuery {
latest: Some(2),
after: Some(1),
..Default::default()
},
)
.await
.unwrap();
assert_eq!(latest.iter().map(|e| e.seq).collect::<Vec<_>>(), vec![4, 5]);
let limited = s
.list(
&r,
&TraceQuery {
limit: Some(2),
..Default::default()
},
)
.await
.unwrap();
assert_eq!(
limited.iter().map(|e| e.seq).collect::<Vec<_>>(),
vec![1, 2]
);
}
#[tokio::test]
async fn retention_prunes_oldest_keeping_seq() {
let s = InMemoryRunTraceStore::with_max_events_per_run(3);
let r = rid("R-1");
for i in 0..5 {
s.append(&r, draft(&format!("core.e{i}"), None, None))
.await
.unwrap();
}
let all = s.list(&r, &TraceQuery::default()).await.unwrap();
assert_eq!(all.iter().map(|e| e.seq).collect::<Vec<_>>(), vec![3, 4, 5]);
let e6 = s.append(&r, draft("core.e5", None, None)).await.unwrap();
assert_eq!(e6.seq, 6);
}
#[tokio::test]
async fn delete_run_removes_stream() {
let s = InMemoryRunTraceStore::new();
let r = rid("R-1");
s.append(&r, draft("core.run_started", None, None))
.await
.unwrap();
s.append(&r, draft("core.run_finished", None, None))
.await
.unwrap();
assert_eq!(s.delete_run(&r).await.unwrap(), 2);
assert!(s.list(&r, &TraceQuery::default()).await.unwrap().is_empty());
assert_eq!(s.delete_run(&r).await.unwrap(), 0, "double delete is Ok(0)");
}
#[tokio::test]
async fn oversized_payload_is_truncated_with_marker() {
let s = InMemoryRunTraceStore::new();
let r = rid("R-1");
let big = "x".repeat(TRACE_PAYLOAD_CAP_BYTES + 100);
let e = s
.append(
&r,
TraceEventDraft {
kind: "worker.output".into(),
step_ref: None,
attempt: None,
payload: json!({"blob": big}),
},
)
.await
.unwrap();
assert_eq!(e.payload.get("truncated"), Some(&json!(true)));
assert!(e.payload.get("size_bytes").is_some());
}
#[tokio::test]
async fn trace_handle_appends_best_effort() {
let store: Arc<dyn RunTraceStore> = Arc::new(InMemoryRunTraceStore::new());
let handle = TraceHandle::new(rid("R-1"), store.clone());
handle
.append(kind::STEP_DISPATCHED, Some("w"), Some(1), json!({}))
.await;
let events = store
.list(&rid("R-1"), &TraceQuery::default())
.await
.unwrap();
assert_eq!(events.len(), 1);
assert_eq!(events[0].kind, kind::STEP_DISPATCHED);
}
#[test]
fn worker_stats_is_empty_reflects_fields() {
assert!(WorkerStats::default().is_empty());
let stats = WorkerStats {
usage: Some(TokenUsage {
input_tokens: 1,
output_tokens: 2,
total_tokens: 3,
}),
..Default::default()
};
assert!(!stats.is_empty());
}
}