pe_core/

internal_tool.rs

//! Internal tools -- cognitive-only tools for agent self-management.
//!
//! Internal tools operate inside the cognitive graph, reading from
//! [`CognitiveState`] to help the agent manage its working memory,
//! constraints, plans, and signals. They are NOT exposed to the outer
//! execution graph.
//!
//! Tools receive `&CognitiveState` (read-only) and return JSON describing
//! what should be applied. Actual state mutation happens via
//! [`CognitiveStateUpdate`](crate::CognitiveStateUpdate) at the node level,
//! following Pregel snapshot isolation.

use crate::cognitive::CognitiveState;
use crate::error::PeError;
use std::collections::HashMap;
use std::sync::Arc;

/// A tool available only inside the cognitive graph.
///
/// Internal tools operate on [`CognitiveState`] for self-management:
/// writing notes, recalling constraints, recording failures, etc.
/// They are NOT exposed to the outer execution graph.
///
/// # Example
///
/// ```
/// use pe_core::internal_tool::InternalTool;
/// use pe_core::cognitive::CognitiveState;
/// use pe_core::error::PeError;
///
/// struct MyTool;
///
/// #[async_trait::async_trait]
/// impl InternalTool for MyTool {
///     fn name(&self) -> &str { "my_tool" }
///     fn description(&self) -> &str { "A custom internal tool" }
///     async fn execute(
///         &self,
///         _state: &CognitiveState,
///         _args: serde_json::Value,
///     ) -> Result<serde_json::Value, PeError> {
///         Ok(serde_json::json!({"status": "ok"}))
///     }
/// }
/// ```
#[async_trait::async_trait]
pub trait InternalTool: Send + Sync {
    /// Unique tool name.
    fn name(&self) -> &str;

    /// Human-readable description for LLM tool-use prompts.
    fn description(&self) -> &str;

    /// Execute the tool with the current cognitive state and arguments.
    async fn execute(
        &self,
        state: &CognitiveState,
        args: serde_json::Value,
    ) -> Result<serde_json::Value, PeError>;
}

/// Registry of internal tools, keyed by name.
///
/// Stores tools as `Arc<dyn InternalTool>` for cheap cloning and
/// shared ownership across cognitive graph nodes.
///
/// # Example
///
/// ```
/// use pe_core::internal_tool::InternalToolRegistry;
///
/// let registry = InternalToolRegistry::new();
/// assert!(registry.is_empty());
/// assert_eq!(registry.len(), 0);
/// ```
pub struct InternalToolRegistry {
    tools: HashMap<String, Arc<dyn InternalTool>>,
}

impl InternalToolRegistry {
    /// Create an empty registry.
    pub fn new() -> Self {
        Self {
            tools: HashMap::new(),
        }
    }

    /// Register an internal tool. Overwrites any existing tool with the same name.
    pub fn register(&mut self, tool: Arc<dyn InternalTool>) -> &mut Self {
        self.tools.insert(tool.name().to_string(), tool);
        self
    }

    /// Remove a tool by name. Returns the removed tool if it existed.
    pub fn unregister(&mut self, name: &str) -> Option<Arc<dyn InternalTool>> {
        self.tools.remove(name)
    }

    /// Look up a tool by name.
    pub fn get(&self, name: &str) -> Option<&Arc<dyn InternalTool>> {
        self.tools.get(name)
    }

    /// List all registered tool names.
    pub fn list(&self) -> Vec<&str> {
        self.tools.keys().map(|k| k.as_str()).collect()
    }

    /// Number of registered tools.
    pub fn len(&self) -> usize {
        self.tools.len()
    }

    /// Whether the registry is empty.
    pub fn is_empty(&self) -> bool {
        self.tools.is_empty()
    }
}

impl Default for InternalToolRegistry {
    fn default() -> Self {
        Self::new()
    }
}

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

    struct DummyTool {
        tool_name: &'static str,
    }

    #[async_trait::async_trait]
    impl InternalTool for DummyTool {
        fn name(&self) -> &str {
            self.tool_name
        }
        fn description(&self) -> &str {
            "A dummy tool for testing"
        }
        async fn execute(
            &self,
            _state: &CognitiveState,
            _args: serde_json::Value,
        ) -> Result<serde_json::Value, PeError> {
            Ok(serde_json::json!({"status": "ok"}))
        }
    }

    #[test]
    fn test_registry_new_is_empty() {
        let registry = InternalToolRegistry::new();
        assert!(registry.is_empty());
        assert_eq!(registry.len(), 0);
        assert!(registry.list().is_empty());
    }

    #[test]
    fn test_registry_register_and_get() {
        let mut registry = InternalToolRegistry::new();
        registry.register(Arc::new(DummyTool { tool_name: "note" }));
        assert_eq!(registry.len(), 1);
        assert!(!registry.is_empty());

        let tool = registry.get("note").expect("tool should exist");
        assert_eq!(tool.name(), "note");
        assert_eq!(tool.description(), "A dummy tool for testing");
    }

    #[test]
    fn test_registry_get_missing_returns_none() {
        let registry = InternalToolRegistry::new();
        assert!(registry.get("nonexistent").is_none());
    }

    #[test]
    fn test_registry_overwrite_same_name() {
        let mut registry = InternalToolRegistry::new();
        registry.register(Arc::new(DummyTool { tool_name: "note" }));
        registry.register(Arc::new(DummyTool { tool_name: "note" }));
        assert_eq!(registry.len(), 1);
    }

    #[test]
    fn test_registry_list_returns_all_names() {
        let mut registry = InternalToolRegistry::new();
        registry.register(Arc::new(DummyTool { tool_name: "note" }));
        registry.register(Arc::new(DummyTool {
            tool_name: "read_notes",
        }));
        registry.register(Arc::new(DummyTool {
            tool_name: "signal",
        }));
        let mut names = registry.list();
        names.sort();
        assert_eq!(names, vec!["note", "read_notes", "signal"]);
    }

    #[test]
    fn test_registry_chained_register() {
        let mut registry = InternalToolRegistry::new();
        registry
            .register(Arc::new(DummyTool { tool_name: "a" }))
            .register(Arc::new(DummyTool { tool_name: "b" }));
        assert_eq!(registry.len(), 2);
    }

    #[tokio::test]
    async fn test_tool_execute() {
        let tool = DummyTool { tool_name: "test" };
        let state = CognitiveState::default();
        let result = tool
            .execute(&state, serde_json::json!({}))
            .await
            .expect("execute should succeed");
        assert_eq!(result, serde_json::json!({"status": "ok"}));
    }
}