pe-tools 0.1.0

Tool registry and MCP adapter for Potential Expectations — schema-driven tool nodes and protocol bridge
Documentation
//! Dependency injection markers — types that inject runtime context into tools.
//!
//! These are zero-cost marker types. When the `#[tool]` macro (Plan 010)
//! processes a tool function, it recognizes these types and excludes them
//! from the JSON schema sent to the LLM. Instead, they are injected by
//! the runtime before tool execution.
//!
//! In Plan 006, we define the types. The macro integration comes in Plan 010.

use pe_core::state::State;
use std::marker::PhantomData;
use std::sync::Arc;

/// Marker: this tool parameter receives the current graph state.
///
/// The LLM never sees this parameter — it's injected by the runtime.
/// The value is a clone of the state at the time `ToolNode::call()` executes.
///
/// # Future usage (with `#[tool]` macro, Plan 010)
///
/// ```ignore
/// #[tool]
/// async fn my_tool(
///     query: String,
///     state: InjectedState<AgentState>,
/// ) -> Result<Value, PeError> {
///     let messages = state.get().messages();
///     // ... use state for context ...
/// }
/// ```
pub struct InjectedState<S: State> {
    state: S,
    _phantom: PhantomData<S>,
}

impl<S: State> InjectedState<S> {
    /// Create a new injected state wrapper.
    pub fn new(state: S) -> Self {
        Self {
            state,
            _phantom: PhantomData,
        }
    }

    /// Access the wrapped state.
    pub fn get(&self) -> &S {
        &self.state
    }

    /// Consume the wrapper and return the inner state.
    pub fn into_inner(self) -> S {
        self.state
    }
}

/// Marker: this tool parameter receives the long-term memory store.
///
/// The LLM never sees this parameter — it's injected by the runtime.
/// The value is an `Arc<dyn Store>` from the pe-memory crate.
///
/// # Future usage (with `#[tool]` macro, Plan 010)
///
/// ```ignore
/// #[tool]
/// async fn memory_search(
///     query: String,
///     store: InjectedStore,
/// ) -> Result<Value, PeError> {
///     let results = store.get().search(&query).await?;
///     // ...
/// }
/// ```
///
/// Note: The `Store` trait is defined in pe-memory. We store `Arc<dyn Any>`
/// here to avoid a circular dependency. The runtime layer casts it back.
pub struct InjectedStore {
    store: Arc<dyn std::any::Any + Send + Sync>,
}

impl InjectedStore {
    /// Create a new injected store wrapper.
    pub fn new(store: Arc<dyn std::any::Any + Send + Sync>) -> Self {
        Self { store }
    }

    /// Access the wrapped store as `Arc<dyn Any>`.
    /// The runtime layer downcasts this to the concrete `Store` impl.
    pub fn get(&self) -> &Arc<dyn std::any::Any + Send + Sync> {
        &self.store
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use pe_core::message::Message;
    use pe_core::state::{CoreState, ExecutionContext, StateUpdate};
    use serde::{Deserialize, Serialize};

    #[derive(Debug, Clone, Serialize, Deserialize)]
    struct TestState {
        messages: Vec<Message>,
        iterations: u32,
        thread_id: String,
        context: ExecutionContext,
    }

    #[derive(Debug, Clone, Default, Serialize, Deserialize)]
    struct TestUpdate;
    impl StateUpdate for TestUpdate {}

    impl pe_core::state::State for TestState {
        type Update = TestUpdate;
        fn apply(&mut self, _update: Self::Update) {}
    }

    impl CoreState for TestState {
        fn messages(&self) -> &[Message] {
            &self.messages
        }
        fn messages_mut(&mut self) -> &mut Vec<Message> {
            &mut self.messages
        }
        fn iterations(&self) -> u32 {
            self.iterations
        }
        fn set_iterations(&mut self, n: u32) {
            self.iterations = n;
        }
        fn thread_id(&self) -> &str {
            &self.thread_id
        }
        fn context(&self) -> &ExecutionContext {
            &self.context
        }
        fn context_mut(&mut self) -> &mut ExecutionContext {
            &mut self.context
        }
    }

    #[test]
    fn injected_state_wraps_and_unwraps() {
        let state = TestState {
            messages: vec![Message::human("hello")],
            iterations: 0,
            thread_id: "t1".into(),
            context: ExecutionContext::new("agent-1"),
        };

        let injected = InjectedState::new(state.clone());
        assert_eq!(injected.get().messages().len(), 1);
        assert_eq!(injected.get().thread_id(), "t1");

        let inner = injected.into_inner();
        assert_eq!(inner.iterations(), 0);
    }

    #[test]
    fn injected_store_wraps_any() {
        let data: Arc<dyn std::any::Any + Send + Sync> = Arc::new(42_u32);
        let injected = InjectedStore::new(data);
        let store = injected.get();
        let value = store.downcast_ref::<u32>().unwrap();
        assert_eq!(*value, 42);
    }
}