agentlib_logger/

lib.rs

use agentlib_core::{Middleware, MiddlewareContext, MiddlewareScope, NextMiddleware, async_trait};
use anyhow::Result;
use chrono::Utc;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum LogLevel {
    Debug,
    Info,
    Warn,
    Error,
    Silent,
}

impl LogLevel {
    fn rank(&self) -> i32 {
        match self {
            LogLevel::Debug => 0,
            LogLevel::Info => 1,
            LogLevel::Warn => 2,
            LogLevel::Error => 3,
            LogLevel::Silent => 99,
        }
    }
}

#[derive(Debug, Clone, Serialize)]
pub struct LogEntry {
    pub level: LogLevel,
    pub scope: MiddlewareScope,
    pub agent_input: Option<String>,
    pub tool: Option<String>,
    pub timestamp: String,
    pub duration_ms: Option<u128>,
    pub meta: Option<serde_json::Value>,
}

pub type LogTransport = Box<dyn Fn(LogEntry) + Send + Sync>;

pub struct LoggerConfig {
    pub level: LogLevel,
    pub scopes: Option<Vec<MiddlewareScope>>,
    pub transport: Option<LogTransport>,
    pub timing: bool,
    pub prefix: String,
}

impl Default for LoggerConfig {
    fn default() -> Self {
        Self {
            level: LogLevel::Info,
            scopes: None,
            transport: None,
            timing: true,
            prefix: "[agentlib]".to_string(),
        }
    }
}

pub struct Logger {
    config: LoggerConfig,
    timers: std::sync::Arc<std::sync::Mutex<HashMap<String, std::time::Instant>>>,
}

impl Logger {
    pub fn new(config: LoggerConfig) -> Self {
        Self {
            config,
            timers: std::sync::Arc::new(std::sync::Mutex::new(HashMap::new())),
        }
    }

    fn should_log(&self, level: LogLevel) -> bool {
        level.rank() >= self.config.level.rank()
    }

    fn scope_to_level(scope: MiddlewareScope) -> LogLevel {
        match scope {
            MiddlewareScope::RunBefore | MiddlewareScope::RunAfter => LogLevel::Info,
            MiddlewareScope::StepBefore | MiddlewareScope::StepAfter => LogLevel::Debug,
            MiddlewareScope::ToolBefore | MiddlewareScope::ToolAfter => LogLevel::Debug,
        }
    }

    fn default_transport(prefix: &str, entry: LogEntry) {
        let mut parts = vec![
            prefix.to_string(),
            format!("[{}]", entry.timestamp),
            format!("[{:?}]", entry.level).to_uppercase(),
            format!("scope={:?}", entry.scope),
        ];

        if let Some(tool) = entry.tool {
            parts.push(format!("tool={}", tool));
        }

        if let Some(duration) = entry.duration_ms {
            parts.push(format!("duration={}ms", duration));
        }

        if let Some(meta) = entry.meta {
            parts.push(meta.to_string());
        }

        let line = parts.join(" ");

        match entry.level {
            LogLevel::Debug => println!("{}", line),
            LogLevel::Info => println!("{}", line),
            LogLevel::Warn => eprintln!("{}", line),
            LogLevel::Error => eprintln!("{}", line),
            LogLevel::Silent => {}
        }
    }
}

#[async_trait]
impl Middleware for Logger {
    fn name(&self) -> &str {
        "logger"
    }

    async fn run(
        &self,
        m_ctx: &mut MiddlewareContext<'_>,
        next: &mut dyn NextMiddleware,
    ) -> Result<()> {
        let scope = m_ctx.scope;

        // Scope filter
        if let Some(scopes) = &self.config.scopes {
            if !scopes.contains(&scope) {
                return next.run(m_ctx).await;
            }
        }

        let level = Self::scope_to_level(scope);
        if !self.should_log(level) {
            return next.run(m_ctx).await;
        }

        let timestamp = Utc::now().to_rfc3339();
        let tool_name = m_ctx.tool.as_ref().map(|t| t.name.clone());
        let timer_key = format!(
            "{:?}-{}-{}",
            scope,
            m_ctx.ctx.input.chars().take(20).collect::<String>(),
            tool_name.as_deref().unwrap_or("")
        );

        let is_before = match scope {
            MiddlewareScope::RunBefore
            | MiddlewareScope::StepBefore
            | MiddlewareScope::ToolBefore => true,
            _ => false,
        };

        if is_before {
            if self.config.timing {
                let mut timers = self.timers.lock().unwrap();
                timers.insert(timer_key.clone(), std::time::Instant::now());
            }

            let entry = LogEntry {
                level,
                scope,
                timestamp,
                agent_input: Some(m_ctx.ctx.input.clone()),
                tool: tool_name.clone(),
                duration_ms: None,
                meta: m_ctx
                    .tool
                    .as_ref()
                    .map(|t| serde_json::json!({ "args": t.args })),
            };

            if let Some(transport) = &self.config.transport {
                transport(entry);
            } else {
                Self::default_transport(&self.config.prefix, entry);
            }

            return next.run(m_ctx).await;
        }

        // AFTER scopes
        let start = if self.config.timing {
            let before_scope = match scope {
                MiddlewareScope::RunAfter => MiddlewareScope::RunBefore,
                MiddlewareScope::StepAfter => MiddlewareScope::StepBefore,
                MiddlewareScope::ToolAfter => MiddlewareScope::ToolBefore,
                _ => unreachable!(),
            };
            let before_key = format!(
                "{:?}-{}-{}",
                before_scope,
                m_ctx.ctx.input.chars().take(20).collect::<String>(),
                tool_name.as_deref().unwrap_or("")
            );
            let mut timers = self.timers.lock().unwrap();
            timers.remove(&before_key)
        } else {
            None
        };

        let res = next.run(m_ctx).await;

        let duration_ms = start.map(|s| s.elapsed().as_millis());
        let entry = LogEntry {
            level,
            scope,
            timestamp: Utc::now().to_rfc3339(),
            agent_input: Some(m_ctx.ctx.input.clone()),
            tool: tool_name,
            duration_ms,
            meta: m_ctx.tool.as_ref().and_then(|t| {
                t.result
                    .as_ref()
                    .map(|r| serde_json::json!({ "result": r }))
            }),
        };

        if let Some(transport) = &self.config.transport {
            transport(entry);
        } else {
            Self::default_transport(&self.config.prefix, entry);
        }

        res
    }
}

pub fn create_logger(config: LoggerConfig) -> Box<dyn Middleware> {
    Box::new(Logger::new(config))
}