a3s_code_core/
trace.rs

1//! Runtime trace primitives.
2//!
3//! Trace events are compact execution facts emitted by the harness. They are
4//! separate from model-visible tool output and from large artifacts.
5
6use serde::{Deserialize, Serialize};
7use std::sync::{Arc, RwLock};
8use std::time::Duration;
9
10pub const TRACE_EVENT_SCHEMA: &str = "a3s.trace_event.v1";
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
13#[serde(rename_all = "snake_case")]
14pub enum TraceEventKind {
15    ToolExecution,
16    ProgramExecution,
17}
18
19#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
20pub struct TraceEvent {
21    pub schema: String,
22    pub kind: TraceEventKind,
23    pub name: String,
24    pub success: bool,
25    pub exit_code: i32,
26    pub duration_ms: u64,
27    pub output_bytes: usize,
28    #[serde(default, skip_serializing_if = "Vec::is_empty")]
29    pub metadata_keys: Vec<String>,
30    #[serde(default, skip_serializing_if = "Vec::is_empty")]
31    pub artifact_uris: Vec<String>,
32    #[serde(default, skip_serializing_if = "Option::is_none")]
33    pub details: Option<serde_json::Value>,
34}
35
36impl TraceEvent {
37    pub fn tool_execution(
38        name: impl Into<String>,
39        success: bool,
40        exit_code: i32,
41        duration: Duration,
42        output_bytes: usize,
43        metadata: Option<&serde_json::Value>,
44    ) -> Self {
45        Self {
46            schema: TRACE_EVENT_SCHEMA.to_string(),
47            kind: TraceEventKind::ToolExecution,
48            name: name.into(),
49            success,
50            exit_code,
51            duration_ms: duration.as_millis().min(u128::from(u64::MAX)) as u64,
52            output_bytes,
53            metadata_keys: metadata_keys(metadata),
54            artifact_uris: artifact_uris(metadata),
55            details: None,
56        }
57    }
58
59    pub fn program_execution(
60        name: impl Into<String>,
61        success: bool,
62        exit_code: i32,
63        duration: Duration,
64        output_bytes: usize,
65        metadata: Option<&serde_json::Value>,
66    ) -> Self {
67        let details = metadata
68            .and_then(|metadata| metadata.get("trace"))
69            .map(program_trace_summary);
70
71        Self {
72            schema: TRACE_EVENT_SCHEMA.to_string(),
73            kind: TraceEventKind::ProgramExecution,
74            name: name.into(),
75            success,
76            exit_code,
77            duration_ms: duration.as_millis().min(u128::from(u64::MAX)) as u64,
78            output_bytes,
79            metadata_keys: metadata_keys(metadata),
80            artifact_uris: artifact_uris(metadata),
81            details,
82        }
83    }
84}
85
86pub trait TraceSink: Send + Sync {
87    fn record(&self, event: TraceEvent);
88}
89
90#[derive(Debug, Clone, Default)]
91pub struct InMemoryTraceSink {
92    events: Arc<RwLock<Vec<TraceEvent>>>,
93}
94
95impl InMemoryTraceSink {
96    pub fn events(&self) -> Vec<TraceEvent> {
97        self.events.read().unwrap().clone()
98    }
99
100    pub fn replace_events(&self, events: Vec<TraceEvent>) {
101        *self.events.write().unwrap() = events;
102    }
103
104    pub fn clear(&self) {
105        self.events.write().unwrap().clear();
106    }
107}
108
109impl TraceSink for InMemoryTraceSink {
110    fn record(&self, event: TraceEvent) {
111        self.events.write().unwrap().push(event);
112    }
113}
114
115#[derive(Debug, Clone, Copy, Default)]
116pub struct NoopTraceSink;
117
118impl TraceSink for NoopTraceSink {
119    fn record(&self, _event: TraceEvent) {}
120}
121
122fn metadata_keys(metadata: Option<&serde_json::Value>) -> Vec<String> {
123    let Some(serde_json::Value::Object(object)) = metadata else {
124        return Vec::new();
125    };
126
127    let mut keys = object.keys().cloned().collect::<Vec<_>>();
128    keys.sort();
129    keys
130}
131
132fn artifact_uris(metadata: Option<&serde_json::Value>) -> Vec<String> {
133    let mut uris = Vec::new();
134    if let Some(metadata) = metadata {
135        collect_artifact_uris(metadata, &mut uris);
136    }
137    uris.sort();
138    uris.dedup();
139    uris
140}
141
142fn collect_artifact_uris(value: &serde_json::Value, uris: &mut Vec<String>) {
143    match value {
144        serde_json::Value::Object(object) => {
145            if let Some(uri) = object.get("artifact_uri").and_then(|value| value.as_str()) {
146                uris.push(uri.to_string());
147            }
148            for value in object.values() {
149                collect_artifact_uris(value, uris);
150            }
151        }
152        serde_json::Value::Array(items) => {
153            for value in items {
154                collect_artifact_uris(value, uris);
155            }
156        }
157        _ => {}
158    }
159}
160
161fn program_trace_summary(trace: &serde_json::Value) -> serde_json::Value {
162    serde_json::json!({
163        "program_name": trace.get("program_name").cloned().unwrap_or_default(),
164        "success": trace.get("success").cloned().unwrap_or_default(),
165        "step_count": trace.get("step_count").cloned().unwrap_or_default(),
166        "failed_steps": trace.get("failed_steps").cloned().unwrap_or_default(),
167    })
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    #[test]
175    fn in_memory_trace_sink_records_events() {
176        let sink = InMemoryTraceSink::default();
177        sink.record(TraceEvent::tool_execution(
178            "read",
179            true,
180            0,
181            Duration::from_millis(7),
182            12,
183            Some(&serde_json::json!({
184                "artifact": {
185                    "artifact_uri": "a3s://tool-output/read/abc"
186                },
187                "file_path": "src/lib.rs"
188            })),
189        ));
190
191        let events = sink.events();
192
193        assert_eq!(events.len(), 1);
194        assert_eq!(events[0].schema, TRACE_EVENT_SCHEMA);
195        assert_eq!(events[0].kind, TraceEventKind::ToolExecution);
196        assert_eq!(events[0].metadata_keys, vec!["artifact", "file_path"]);
197        assert_eq!(events[0].artifact_uris, vec!["a3s://tool-output/read/abc"]);
198    }
199
200    #[test]
201    fn program_trace_event_stores_compact_summary() {
202        let event = TraceEvent::program_execution(
203            "program",
204            true,
205            0,
206            Duration::from_millis(3),
207            42,
208            Some(&serde_json::json!({
209                "trace": {
210                    "program_name": "program_repo_map",
211                    "success": true,
212                    "step_count": 7,
213                    "failed_steps": 0,
214                    "steps": [{"output": "not copied into event"}]
215                }
216            })),
217        );
218
219        assert_eq!(event.kind, TraceEventKind::ProgramExecution);
220        assert_eq!(
221            event.details.as_ref().unwrap()["program_name"],
222            "program_repo_map"
223        );
224        assert!(event.details.as_ref().unwrap().get("steps").is_none());
225    }
226}
a3s_code_core/trace.rs

a3s_code_core/
trace.rs
