use tracing::{error, info, warn};
pub trait OutputHandler: Send + Sync {
fn on_stdout(&self, line: &str);
fn on_stderr(&self, line: &str);
fn on_agent_stderr(&self, line: &str);
fn on_info(&self, message: &str);
fn on_warn(&self, message: &str);
fn on_error(&self, message: &str);
fn on_success(&self, message: &str);
}
#[derive(Debug, Clone, Default)]
pub struct LogOutputHandler;
impl LogOutputHandler {
pub fn new() -> Self {
Self
}
}
impl OutputHandler for LogOutputHandler {
fn on_stdout(&self, line: &str) {
info!(target: "orchestrator::output", "{}", line);
}
fn on_stderr(&self, line: &str) {
warn!(target: "orchestrator::output", "{}", line);
}
fn on_agent_stderr(&self, line: &str) {
info!(target: "orchestrator::output", "{}", line);
}
fn on_info(&self, message: &str) {
info!("{}", message);
}
fn on_warn(&self, message: &str) {
warn!("{}", message);
}
fn on_error(&self, message: &str) {
error!("{}", message);
}
fn on_success(&self, message: &str) {
info!("{}", message);
}
}
#[derive(Debug, Clone, Default)]
#[allow(dead_code)] pub struct NullOutputHandler;
impl NullOutputHandler {
#[allow(dead_code)] pub fn new() -> Self {
Self
}
}
impl OutputHandler for NullOutputHandler {
fn on_stdout(&self, _line: &str) {}
fn on_stderr(&self, _line: &str) {}
fn on_agent_stderr(&self, _line: &str) {}
fn on_info(&self, _message: &str) {}
fn on_warn(&self, _message: &str) {}
fn on_error(&self, _message: &str) {}
fn on_success(&self, _message: &str) {}
}
#[derive(Clone)]
pub struct ChannelOutputHandler<F>
where
F: Fn(OutputMessage) + Send + Sync,
{
callback: F,
}
#[derive(Debug, Clone)]
pub enum OutputMessage {
Stdout(String),
Stderr(String),
AgentStderr(String),
Info(String),
Warn(String),
Error(String),
Success(String),
}
impl<F> ChannelOutputHandler<F>
where
F: Fn(OutputMessage) + Send + Sync,
{
pub fn new(callback: F) -> Self {
Self { callback }
}
}
impl<F> OutputHandler for ChannelOutputHandler<F>
where
F: Fn(OutputMessage) + Send + Sync,
{
fn on_stdout(&self, line: &str) {
(self.callback)(OutputMessage::Stdout(line.to_string()));
}
fn on_stderr(&self, line: &str) {
(self.callback)(OutputMessage::Stderr(line.to_string()));
}
fn on_agent_stderr(&self, line: &str) {
(self.callback)(OutputMessage::AgentStderr(line.to_string()));
}
fn on_info(&self, message: &str) {
(self.callback)(OutputMessage::Info(message.to_string()));
}
fn on_warn(&self, message: &str) {
(self.callback)(OutputMessage::Warn(message.to_string()));
}
fn on_error(&self, message: &str) {
(self.callback)(OutputMessage::Error(message.to_string()));
}
fn on_success(&self, message: &str) {
(self.callback)(OutputMessage::Success(message.to_string()));
}
}
pub struct ContextualOutputHandler<H: OutputHandler> {
inner: H,
operation: std::sync::Arc<std::sync::RwLock<String>>,
}
impl<H: OutputHandler> ContextualOutputHandler<H> {
pub fn new(inner: H, operation: std::sync::Arc<std::sync::RwLock<String>>) -> Self {
Self { inner, operation }
}
#[allow(dead_code)]
pub fn set_operation(&self, operation: impl Into<String>) {
*self.operation.write().unwrap() = operation.into();
}
#[allow(dead_code)]
pub fn operation(&self) -> String {
self.operation.read().unwrap().clone()
}
}
impl<H: OutputHandler> OutputHandler for ContextualOutputHandler<H> {
fn on_stdout(&self, line: &str) {
self.inner.on_stdout(line);
}
fn on_stderr(&self, line: &str) {
self.inner.on_stderr(line);
}
fn on_agent_stderr(&self, line: &str) {
self.inner.on_agent_stderr(line);
}
fn on_info(&self, message: &str) {
self.inner.on_info(message);
}
fn on_warn(&self, message: &str) {
self.inner.on_warn(message);
}
fn on_error(&self, message: &str) {
self.inner.on_error(message);
}
fn on_success(&self, message: &str) {
self.inner.on_success(message);
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::{Arc, Mutex};
#[derive(Default)]
struct TestOutputHandler {
messages: Arc<Mutex<Vec<(String, String)>>>,
}
impl TestOutputHandler {
fn new() -> Self {
Self {
messages: Arc::new(Mutex::new(Vec::new())),
}
}
fn get_messages(&self) -> Vec<(String, String)> {
self.messages.lock().unwrap().clone()
}
}
impl OutputHandler for TestOutputHandler {
fn on_stdout(&self, line: &str) {
self.messages
.lock()
.unwrap()
.push(("stdout".to_string(), line.to_string()));
}
fn on_stderr(&self, line: &str) {
self.messages
.lock()
.unwrap()
.push(("stderr".to_string(), line.to_string()));
}
fn on_agent_stderr(&self, line: &str) {
self.messages
.lock()
.unwrap()
.push(("agent_stderr".to_string(), line.to_string()));
}
fn on_info(&self, message: &str) {
self.messages
.lock()
.unwrap()
.push(("info".to_string(), message.to_string()));
}
fn on_warn(&self, message: &str) {
self.messages
.lock()
.unwrap()
.push(("warn".to_string(), message.to_string()));
}
fn on_error(&self, message: &str) {
self.messages
.lock()
.unwrap()
.push(("error".to_string(), message.to_string()));
}
fn on_success(&self, message: &str) {
self.messages
.lock()
.unwrap()
.push(("success".to_string(), message.to_string()));
}
}
#[test]
fn test_test_output_handler() {
let handler = TestOutputHandler::new();
handler.on_stdout("stdout line");
handler.on_stderr("stderr line");
handler.on_agent_stderr("agent stderr line");
handler.on_info("info message");
handler.on_warn("warn message");
handler.on_error("error message");
handler.on_success("success message");
let messages = handler.get_messages();
assert_eq!(messages.len(), 7);
assert_eq!(
messages[0],
("stdout".to_string(), "stdout line".to_string())
);
assert_eq!(
messages[1],
("stderr".to_string(), "stderr line".to_string())
);
assert_eq!(
messages[2],
("agent_stderr".to_string(), "agent stderr line".to_string())
);
assert_eq!(
messages[3],
("info".to_string(), "info message".to_string())
);
assert_eq!(
messages[4],
("warn".to_string(), "warn message".to_string())
);
assert_eq!(
messages[5],
("error".to_string(), "error message".to_string())
);
assert_eq!(
messages[6],
("success".to_string(), "success message".to_string())
);
}
#[test]
fn test_null_output_handler() {
let handler = NullOutputHandler::new();
handler.on_stdout("stdout");
handler.on_stderr("stderr");
handler.on_agent_stderr("agent stderr");
handler.on_info("info");
handler.on_warn("warn");
handler.on_error("error");
handler.on_success("success");
}
#[test]
fn test_log_output_handler() {
let handler = LogOutputHandler::new();
handler.on_stdout("stdout");
handler.on_stderr("stderr");
handler.on_agent_stderr("agent stderr");
handler.on_info("info");
handler.on_warn("warn");
handler.on_error("error");
handler.on_success("success");
}
#[test]
fn test_agent_stderr_is_distinct_from_stderr() {
let handler = TestOutputHandler::new();
handler.on_stderr("internal warning");
handler.on_agent_stderr("agent progress output");
let messages = handler.get_messages();
assert_eq!(messages.len(), 2);
assert_eq!(
messages[0],
("stderr".to_string(), "internal warning".to_string())
);
assert_eq!(
messages[1],
(
"agent_stderr".to_string(),
"agent progress output".to_string()
)
);
}
#[test]
fn test_channel_output_handler_agent_stderr() {
use std::sync::{Arc, Mutex};
let received = Arc::new(Mutex::new(Vec::new()));
let received_clone = received.clone();
let handler = ChannelOutputHandler::new(move |msg| {
received_clone.lock().unwrap().push(msg);
});
handler.on_agent_stderr("agent output");
let messages = received.lock().unwrap();
assert_eq!(messages.len(), 1);
assert!(matches!(&messages[0], OutputMessage::AgentStderr(s) if s == "agent output"));
}
#[test]
fn test_contextual_output_handler_delegates_agent_stderr() {
let inner = TestOutputHandler::new();
let operation = std::sync::Arc::new(std::sync::RwLock::new("apply".to_string()));
let handler = ContextualOutputHandler::new(inner, operation);
handler.on_agent_stderr("agent line");
}
}