aspect_std/

toolscope.rs

//! Tool-scope sandboxing aspect for agentic systems.
//!
//! Restricts an agent invocation to a declared tool allowlist. Wraps a
//! tool-dispatch join point and rejects calls whose tool name is not in
//! the per-session allowlist before the dispatch happens.
//!
//! The aspect is novel relative to classical AOP catalogues (logging,
//! caching, rate-limiting): it encodes a *capability boundary* that is
//! specific to agentic AI architectures, where an LLM may emit tool
//! calls that escape the intended scope of a session.

use aspect_core::{Aspect, AspectError, JoinPoint};
use parking_lot::RwLock;
use std::any::Any;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;

/// Identifier for a session whose tool calls are being scoped.
pub type SessionId = String;

/// Per-call context that the host system must populate before the
/// guarded function runs.
#[derive(Clone, Debug)]
pub struct ToolCall {
    /// The session whose allowlist applies.
    pub session: SessionId,
    /// Name of the tool being invoked.
    pub tool: String,
}

/// Sandbox aspect that enforces a per-session tool allowlist.
#[derive(Clone, Default)]
pub struct ToolScopeSandboxAspect {
    inner: Arc<Inner>,
}

#[derive(Default)]
struct Inner {
    allowlists: RwLock<HashMap<SessionId, HashSet<String>>>,
    pending: RwLock<HashMap<u64, ToolCall>>,
    /// Counter for assigning thread-local call keys; tests rely on it
    /// being monotonic but values are opaque to callers.
    next_key: RwLock<u64>,
}

impl ToolScopeSandboxAspect {
    /// Create an empty sandbox. Sessions must call `allow` before any
    /// guarded join points fire for them.
    pub fn new() -> Self {
        Self::default()
    }

    /// Declare the full allowlist for a session, replacing any prior set.
    pub fn allow<I, S>(&self, session: impl Into<SessionId>, tools: I)
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        let set: HashSet<String> = tools.into_iter().map(Into::into).collect();
        self.inner.allowlists.write().insert(session.into(), set);
    }

    /// Bind a tool call to the next guarded invocation. The host wraps
    /// each dispatch with `bind` ... `proceed` ... `release`. The
    /// returned key is opaque; pass it to `release`.
    pub fn bind(&self, call: ToolCall) -> u64 {
        let mut next = self.inner.next_key.write();
        *next += 1;
        let key = *next;
        self.inner.pending.write().insert(key, call);
        key
    }

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

    /// Inspect whether a `(session, tool)` pair would be admitted.
    pub fn is_allowed(&self, session: &str, tool: &str) -> bool {
        self.inner
            .allowlists
            .read()
            .get(session)
            .map(|s| s.contains(tool))
            .unwrap_or(false)
    }

    fn most_recent_call(&self) -> Option<ToolCall> {
        let pending = self.inner.pending.read();
        // Highest key wins; HashMap iteration order is non-deterministic.
        pending
            .iter()
            .max_by_key(|(k, _)| *k)
            .map(|(_, v)| v.clone())
    }
}

impl Aspect for ToolScopeSandboxAspect {
    fn before(&self, ctx: &JoinPoint) {
        if let Some(call) = self.most_recent_call() {
            if !self.is_allowed(&call.session, &call.tool) {
                // The around path is where we can actually deny; here we
                // only get a chance to log. Real enforcement lives in
                // `around` below.
                let _ = ctx;
            }
        }
    }

    fn around(
        &self,
        pjp: aspect_core::ProceedingJoinPoint,
    ) -> Result<Box<dyn Any>, AspectError> {
        if let Some(call) = self.most_recent_call() {
            if !self.is_allowed(&call.session, &call.tool) {
                return Err(AspectError::execution(format!(
                    "tool '{}' is out of scope for session '{}'",
                    call.tool, call.session
                )));
            }
        }
        pjp.proceed()
    }
}

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

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

    #[test]
    fn allowed_tool_proceeds() {
        let sandbox = ToolScopeSandboxAspect::new();
        sandbox.allow("s1", ["read_file", "list_dir"]);
        let key = sandbox.bind(ToolCall {
            session: "s1".into(),
            tool: "read_file".into(),
        });

        let pjp = ProceedingJoinPoint::new(|| Ok(Box::new(42_i32) as Box<dyn Any>), jp());
        let result = sandbox.around(pjp);
        sandbox.release(key);
        assert!(result.is_ok());
    }

    #[test]
    fn out_of_scope_tool_is_denied() {
        let sandbox = ToolScopeSandboxAspect::new();
        sandbox.allow("s1", ["read_file"]);
        let key = sandbox.bind(ToolCall {
            session: "s1".into(),
            tool: "delete_file".into(),
        });

        let pjp = ProceedingJoinPoint::new(|| Ok(Box::new(0_i32) as Box<dyn Any>), jp());
        let result = sandbox.around(pjp);
        sandbox.release(key);
        assert!(result.is_err());
    }

    #[test]
    fn unknown_session_is_denied() {
        let sandbox = ToolScopeSandboxAspect::new();
        let key = sandbox.bind(ToolCall {
            session: "anon".into(),
            tool: "read_file".into(),
        });

        let pjp = ProceedingJoinPoint::new(|| Ok(Box::new(0_i32) as Box<dyn Any>), jp());
        let result = sandbox.around(pjp);
        sandbox.release(key);
        assert!(result.is_err());
    }
}