use serde::{Deserialize, Serialize};
use std::collections::VecDeque;
#[derive(Debug)]
pub struct ExecutionTracer {
#[allow(dead_code)]
agent_id: [u8; 16],
events: VecDeque<TraceEvent>,
max_events: usize,
filter: Option<TraceFilter>,
enabled: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum TraceEvent {
FunctionCall {
timestamp: u64,
function_name: String,
arguments: Vec<String>,
},
FunctionReturn {
timestamp: u64,
function_name: String,
return_value: Option<String>,
},
MemoryAllocation {
timestamp: u64,
address: usize,
size: usize,
},
MemoryDeallocation { timestamp: u64, address: usize },
StateChange {
timestamp: u64,
variable_name: String,
old_value: String,
new_value: String,
},
MessageSent {
timestamp: u64,
destination: [u8; 16],
size: usize,
},
MessageReceived {
timestamp: u64,
source: [u8; 16],
size: usize,
},
Custom {
timestamp: u64,
event_type: String,
data: String,
},
}
impl TraceEvent {
pub fn timestamp(&self) -> u64 {
match self {
Self::FunctionCall { timestamp, .. }
| Self::FunctionReturn { timestamp, .. }
| Self::MemoryAllocation { timestamp, .. }
| Self::MemoryDeallocation { timestamp, .. }
| Self::StateChange { timestamp, .. }
| Self::MessageSent { timestamp, .. }
| Self::MessageReceived { timestamp, .. }
| Self::Custom { timestamp, .. } => *timestamp,
}
}
pub fn event_type(&self) -> &str {
match self {
Self::FunctionCall { .. } => "FunctionCall",
Self::FunctionReturn { .. } => "FunctionReturn",
Self::MemoryAllocation { .. } => "MemoryAllocation",
Self::MemoryDeallocation { .. } => "MemoryDeallocation",
Self::StateChange { .. } => "StateChange",
Self::MessageSent { .. } => "MessageSent",
Self::MessageReceived { .. } => "MessageReceived",
Self::Custom { .. } => "Custom",
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TraceFilter {
pub include_function_calls: bool,
pub include_memory_events: bool,
pub include_state_changes: bool,
pub include_messages: bool,
pub include_custom_events: bool,
pub function_name_pattern: Option<String>,
}
impl Default for TraceFilter {
fn default() -> Self {
Self {
include_function_calls: true,
include_memory_events: true,
include_state_changes: true,
include_messages: true,
include_custom_events: true,
function_name_pattern: None,
}
}
}
impl TraceFilter {
pub fn matches(&self, event: &TraceEvent) -> bool {
match event {
TraceEvent::FunctionCall { .. } | TraceEvent::FunctionReturn { .. } => {
self.include_function_calls
}
TraceEvent::MemoryAllocation { .. } | TraceEvent::MemoryDeallocation { .. } => {
self.include_memory_events
}
TraceEvent::StateChange { .. } => self.include_state_changes,
TraceEvent::MessageSent { .. } | TraceEvent::MessageReceived { .. } => {
self.include_messages
}
TraceEvent::Custom { .. } => self.include_custom_events,
}
}
}
impl ExecutionTracer {
pub fn new(agent_id: [u8; 16], max_events: usize) -> Self {
Self {
agent_id,
events: VecDeque::with_capacity(max_events.min(10000)),
max_events,
filter: None,
enabled: false,
}
}
pub fn enable(&mut self) {
self.enabled = true;
}
pub fn disable(&mut self) {
self.enabled = false;
}
pub fn set_filter(&mut self, filter: TraceFilter) {
self.filter = Some(filter);
}
pub fn clear_filter(&mut self) {
self.filter = None;
}
pub fn record(&mut self, event: TraceEvent) {
if !self.enabled {
return;
}
if let Some(ref filter) = self.filter {
if !filter.matches(&event) {
return;
}
}
self.events.push_back(event);
if self.events.len() > self.max_events {
self.events.pop_front();
}
}
pub fn events(&self) -> Vec<&TraceEvent> {
self.events.iter().collect()
}
pub fn events_in_range(&self, start: u64, end: u64) -> Vec<&TraceEvent> {
self.events
.iter()
.filter(|e| {
let ts = e.timestamp();
ts >= start && ts <= end
})
.collect()
}
pub fn clear(&mut self) {
self.events.clear();
}
pub fn event_count(&self) -> usize {
self.events.len()
}
}
#[derive(Debug)]
pub struct TraceRecorder {
tracers: std::collections::HashMap<[u8; 16], ExecutionTracer>,
}
impl TraceRecorder {
pub fn new() -> Self {
Self {
tracers: std::collections::HashMap::new(),
}
}
pub fn add_tracer(&mut self, agent_id: [u8; 16], max_events: usize) {
self.tracers
.insert(agent_id, ExecutionTracer::new(agent_id, max_events));
}
pub fn get_tracer(&self, agent_id: &[u8; 16]) -> Option<&ExecutionTracer> {
self.tracers.get(agent_id)
}
pub fn get_tracer_mut(&mut self, agent_id: &[u8; 16]) -> Option<&mut ExecutionTracer> {
self.tracers.get_mut(agent_id)
}
pub fn remove_tracer(&mut self, agent_id: &[u8; 16]) -> Option<ExecutionTracer> {
self.tracers.remove(agent_id)
}
pub fn traced_agents(&self) -> Vec<[u8; 16]> {
self.tracers.keys().copied().collect()
}
pub fn total_events(&self) -> usize {
self.tracers.values().map(|t| t.event_count()).sum()
}
}
impl Default for TraceRecorder {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn current_timestamp() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
#[test]
fn test_execution_tracer_creation() {
let agent_id = [1u8; 16];
let tracer = ExecutionTracer::new(agent_id, 1000);
assert_eq!(tracer.agent_id, agent_id);
assert!(!tracer.enabled);
assert_eq!(tracer.event_count(), 0);
}
#[test]
fn test_execution_tracer_record() {
let agent_id = [1u8; 16];
let mut tracer = ExecutionTracer::new(agent_id, 1000);
tracer.enable();
let event = TraceEvent::FunctionCall {
timestamp: current_timestamp(),
function_name: "test_fn".to_string(),
arguments: vec!["arg1".to_string()],
};
tracer.record(event);
assert_eq!(tracer.event_count(), 1);
}
#[test]
fn test_execution_tracer_filter() {
let agent_id = [1u8; 16];
let mut tracer = ExecutionTracer::new(agent_id, 1000);
tracer.enable();
let filter = TraceFilter {
include_function_calls: false,
..Default::default()
};
tracer.set_filter(filter);
tracer.record(TraceEvent::FunctionCall {
timestamp: current_timestamp(),
function_name: "test".to_string(),
arguments: vec![],
});
assert_eq!(tracer.event_count(), 0);
tracer.record(TraceEvent::StateChange {
timestamp: current_timestamp(),
variable_name: "x".to_string(),
old_value: "1".to_string(),
new_value: "2".to_string(),
});
assert_eq!(tracer.event_count(), 1);
}
#[test]
fn test_trace_recorder() {
let mut recorder = TraceRecorder::new();
let agent_id = [1u8; 16];
recorder.add_tracer(agent_id, 1000);
assert!(recorder.get_tracer(&agent_id).is_some());
let removed = recorder.remove_tracer(&agent_id);
assert!(removed.is_some());
}
#[test]
fn test_trace_event_type() {
let event = TraceEvent::FunctionCall {
timestamp: current_timestamp(),
function_name: "test".to_string(),
arguments: vec![],
};
assert_eq!(event.event_type(), "FunctionCall");
}
}