use chrono::Utc;
use serde::{Deserialize, Serialize};
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::Path;
pub struct AgentOutputLogger {
file: BufWriter<File>,
iteration: u32,
hat: String,
}
#[derive(Debug, Serialize, Deserialize, PartialEq)]
pub struct AgentOutputEntry {
pub ts: String,
pub iteration: u32,
pub hat: String,
#[serde(flatten)]
pub content: AgentOutputContent,
}
#[derive(Debug, Serialize, Deserialize, PartialEq)]
#[serde(tag = "type")]
pub enum AgentOutputContent {
#[serde(rename = "text")]
Text { text: String },
#[serde(rename = "tool_call")]
ToolCall {
name: String,
id: String,
input: serde_json::Value,
},
#[serde(rename = "tool_result")]
ToolResult { id: String, output: String },
#[serde(rename = "error")]
Error { message: String },
#[serde(rename = "complete")]
Complete {
input_tokens: Option<u64>,
output_tokens: Option<u64>,
},
}
impl AgentOutputLogger {
pub fn new(session_dir: &Path) -> std::io::Result<Self> {
let file_path = session_dir.join("agent-output.jsonl");
let file = File::create(file_path)?;
Ok(Self {
file: BufWriter::new(file),
iteration: 0,
hat: String::new(),
})
}
pub fn set_context(&mut self, iteration: u32, hat: &str) {
self.iteration = iteration;
self.hat = hat.to_string();
}
pub fn log(&mut self, content: AgentOutputContent) -> std::io::Result<()> {
let entry = AgentOutputEntry {
ts: Utc::now().to_rfc3339(),
iteration: self.iteration,
hat: self.hat.clone(),
content,
};
let json = serde_json::to_string(&entry)?;
writeln!(self.file, "{}", json)?;
self.file.flush()?;
Ok(())
}
pub fn flush(&mut self) -> std::io::Result<()> {
self.file.flush()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::{BufRead, BufReader};
use tempfile::TempDir;
#[test]
fn test_logger_creates_file() {
let temp = TempDir::new().unwrap();
let _logger = AgentOutputLogger::new(temp.path()).unwrap();
let file_path = temp.path().join("agent-output.jsonl");
assert!(file_path.exists());
}
#[test]
fn test_log_writes_valid_jsonl() {
let temp = TempDir::new().unwrap();
let mut logger = AgentOutputLogger::new(temp.path()).unwrap();
logger.set_context(1, "ralph");
logger
.log(AgentOutputContent::Text {
text: "Hello".to_string(),
})
.unwrap();
logger
.log(AgentOutputContent::ToolCall {
name: "Read".to_string(),
id: "tool_1".to_string(),
input: serde_json::json!({"file": "test.rs"}),
})
.unwrap();
drop(logger);
let file = File::open(temp.path().join("agent-output.jsonl")).unwrap();
let reader = BufReader::new(file);
let lines: Vec<String> = reader.lines().map(|l| l.unwrap()).collect();
assert_eq!(lines.len(), 2);
let entry1: AgentOutputEntry = serde_json::from_str(&lines[0]).unwrap();
assert_eq!(entry1.iteration, 1);
assert_eq!(entry1.hat, "ralph");
assert!(matches!(entry1.content, AgentOutputContent::Text { .. }));
let entry2: AgentOutputEntry = serde_json::from_str(&lines[1]).unwrap();
assert_eq!(entry2.iteration, 1);
assert_eq!(entry2.hat, "ralph");
assert!(matches!(
entry2.content,
AgentOutputContent::ToolCall { .. }
));
}
#[test]
fn test_immediate_flush() {
let temp = TempDir::new().unwrap();
let mut logger = AgentOutputLogger::new(temp.path()).unwrap();
logger.set_context(1, "ralph");
logger
.log(AgentOutputContent::Text {
text: "Test".to_string(),
})
.unwrap();
let file = File::open(temp.path().join("agent-output.jsonl")).unwrap();
let reader = BufReader::new(file);
let lines: Vec<String> = reader.lines().map(|l| l.unwrap()).collect();
assert_eq!(lines.len(), 1);
}
#[test]
fn test_all_content_types_serialize() {
let temp = TempDir::new().unwrap();
let mut logger = AgentOutputLogger::new(temp.path()).unwrap();
logger.set_context(2, "builder");
logger
.log(AgentOutputContent::Text {
text: "Building...".to_string(),
})
.unwrap();
logger
.log(AgentOutputContent::ToolCall {
name: "Execute".to_string(),
id: "t1".to_string(),
input: serde_json::json!({"cmd": "cargo test"}),
})
.unwrap();
logger
.log(AgentOutputContent::ToolResult {
id: "t1".to_string(),
output: "Tests passed".to_string(),
})
.unwrap();
logger
.log(AgentOutputContent::Error {
message: "Parse failed".to_string(),
})
.unwrap();
logger
.log(AgentOutputContent::Complete {
input_tokens: Some(1500),
output_tokens: Some(800),
})
.unwrap();
drop(logger);
let file = File::open(temp.path().join("agent-output.jsonl")).unwrap();
let reader = BufReader::new(file);
let lines: Vec<String> = reader.lines().map(|l| l.unwrap()).collect();
assert_eq!(lines.len(), 5);
for line in lines {
let entry: AgentOutputEntry = serde_json::from_str(&line).unwrap();
assert_eq!(entry.iteration, 2);
assert_eq!(entry.hat, "builder");
}
}
}