aspect-std 0.1.2

//! Action-audit-trail aspect for agentic systems.
//!
//! Records an append-only trail of tool invocations: timestamp, session,
//! tool name, outcome (Ok/Err). The trail is the artifact regulatory
//! frameworks (EU AI Act, OWASP LLM-Top-10) call for; modularizing it
//! as an aspect rather than inlining it at every call site is the
//! contribution this work measures.

use aspect_core::{Aspect, AspectError, JoinPoint, ProceedingJoinPoint};
use parking_lot::{Mutex, RwLock};
use std::any::Any;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};

/// Identifier for a session whose actions are being audited.
pub type SessionId = String;

/// Per-call context the host populates before the audited join point.
#[derive(Clone, Debug)]
pub struct AuditContext {
    /// Session whose action is being recorded.
    pub session: SessionId,
    /// Name of the tool being invoked.
    pub tool: String,
}

/// One audit-trail entry.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AuditRecord {
    /// Unix-millisecond timestamp of the call.
    pub timestamp_ms: u128,
    /// Session id.
    pub session: SessionId,
    /// Tool name.
    pub tool: String,
    /// Function being guarded.
    pub function: String,
    /// Whether the call returned `Ok`.
    pub ok: bool,
    /// Error message, if the call failed.
    pub error: Option<String>,
}

/// Audit-trail aspect.
///
/// The host calls `bind(ctx)` before each guarded invocation. The
/// aspect appends one record per invocation, in arrival order.
#[derive(Clone, Default)]
pub struct ActionAuditTrailAspect {
    inner: Arc<Inner>,
}

#[derive(Default)]
struct Inner {
    trail: Mutex<Vec<AuditRecord>>,
    pending: RwLock<HashMap<u64, AuditContext>>,
    next_key: RwLock<u64>,
}

impl ActionAuditTrailAspect {
    /// Create an empty audit-trail aspect.
    pub fn new() -> Self {
        Self::default()
    }

    /// Bind an audit context to the next guarded invocation.
    pub fn bind(&self, ctx: AuditContext) -> u64 {
        let mut next = self.inner.next_key.write();
        *next += 1;
        let key = *next;
        self.inner.pending.write().insert(key, ctx);
        key
    }

    /// Clear a previously bound context. Idempotent.
    pub fn release(&self, key: u64) {
        self.inner.pending.write().remove(&key);
    }

    /// Snapshot the full audit trail.
    pub fn snapshot(&self) -> Vec<AuditRecord> {
        self.inner.trail.lock().clone()
    }

    /// Number of recorded entries.
    pub fn len(&self) -> usize {
        self.inner.trail.lock().len()
    }

    /// Whether the trail is empty.
    pub fn is_empty(&self) -> bool {
        self.inner.trail.lock().is_empty()
    }

    fn most_recent_context(&self) -> Option<AuditContext> {
        self.inner
            .pending
            .read()
            .iter()
            .max_by_key(|(k, _)| *k)
            .map(|(_, v)| v.clone())
    }

    fn now_ms() -> u128 {
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_millis())
            .unwrap_or(0)
    }

    fn append(&self, record: AuditRecord) {
        self.inner.trail.lock().push(record);
    }
}

impl Aspect for ActionAuditTrailAspect {
    fn around(&self, pjp: ProceedingJoinPoint) -> Result<Box<dyn Any>, AspectError> {
        let ctx = self.most_recent_context();
        let function = pjp.context().function_name.to_string();
        let ts = Self::now_ms();
        let result = pjp.proceed();
        let (ok, err) = match &result {
            Ok(_) => (true, None),
            Err(e) => (false, Some(format!("{:?}", e))),
        };
        let (session, tool) = ctx
            .map(|c| (c.session, c.tool))
            .unwrap_or_else(|| ("unknown".into(), "unknown".into()));
        self.append(AuditRecord {
            timestamp_ms: ts,
            session,
            tool,
            function,
            ok,
            error: err,
        });
        result
    }

    fn after(&self, _ctx: &JoinPoint, _result: &dyn Any) {
        // around-based path is authoritative; this hook is a no-op so that
        // composing with aspects that drive `before`/`after` does not
        // produce duplicate audit records.
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use aspect_core::Location;

    fn jp() -> JoinPoint {
        JoinPoint {
            function_name: "execute_tool_call",
            module_path: "agent::dispatch",
            location: Location {
                file: "agent.rs",
                line: 1,
            },
        }
    }

    #[test]
    fn records_a_successful_call() {
        let audit = ActionAuditTrailAspect::new();
        let key = audit.bind(AuditContext {
            session: "s1".into(),
            tool: "read_file".into(),
        });
        let pjp = ProceedingJoinPoint::new(|| Ok(Box::new(7_i32) as Box<dyn Any>), jp());
        let _ = audit.around(pjp);
        audit.release(key);

        let trail = audit.snapshot();
        assert_eq!(trail.len(), 1);
        assert_eq!(trail[0].session, "s1");
        assert_eq!(trail[0].tool, "read_file");
        assert!(trail[0].ok);
        assert_eq!(trail[0].error, None);
    }

    #[test]
    fn records_an_error_call() {
        let audit = ActionAuditTrailAspect::new();
        let key = audit.bind(AuditContext {
            session: "s1".into(),
            tool: "delete_file".into(),
        });
        let pjp = ProceedingJoinPoint::new(
            || Err::<Box<dyn Any>, AspectError>(AspectError::execution("denied")),
            jp(),
        );
        let _ = audit.around(pjp);
        audit.release(key);

        let trail = audit.snapshot();
        assert_eq!(trail.len(), 1);
        assert!(!trail[0].ok);
        assert!(trail[0].error.as_deref().unwrap().contains("denied"));
    }

    #[test]
    fn trail_is_append_only_across_calls() {
        let audit = ActionAuditTrailAspect::new();
        for i in 0..3 {
            let key = audit.bind(AuditContext {
                session: format!("s{}", i),
                tool: "t".into(),
            });
            let pjp = ProceedingJoinPoint::new(|| Ok(Box::new(0_i32) as Box<dyn Any>), jp());
            let _ = audit.around(pjp);
            audit.release(key);
        }
        let trail = audit.snapshot();
        assert_eq!(trail.len(), 3);
        assert_eq!(trail[0].session, "s0");
        assert_eq!(trail[2].session, "s2");
    }
}