pe_tools/
inject.rs

1//! Dependency injection markers — types that inject runtime context into tools.
2//!
3//! These are zero-cost marker types. When the `#[tool]` macro (Plan 010)
4//! processes a tool function, it recognizes these types and excludes them
5//! from the JSON schema sent to the LLM. Instead, they are injected by
6//! the runtime before tool execution.
7//!
8//! In Plan 006, we define the types. The macro integration comes in Plan 010.
9
10use pe_core::state::State;
11use std::marker::PhantomData;
12use std::sync::Arc;
13
14/// Marker: this tool parameter receives the current graph state.
15///
16/// The LLM never sees this parameter — it's injected by the runtime.
17/// The value is a clone of the state at the time `ToolNode::call()` executes.
18///
19/// # Future usage (with `#[tool]` macro, Plan 010)
20///
21/// ```ignore
22/// #[tool]
23/// async fn my_tool(
24///     query: String,
25///     state: InjectedState<AgentState>,
26/// ) -> Result<Value, PeError> {
27///     let messages = state.get().messages();
28///     // ... use state for context ...
29/// }
30/// ```
31pub struct InjectedState<S: State> {
32    state: S,
33    _phantom: PhantomData<S>,
34}
35
36impl<S: State> InjectedState<S> {
37    /// Create a new injected state wrapper.
38    pub fn new(state: S) -> Self {
39        Self {
40            state,
41            _phantom: PhantomData,
42        }
43    }
44
45    /// Access the wrapped state.
46    pub fn get(&self) -> &S {
47        &self.state
48    }
49
50    /// Consume the wrapper and return the inner state.
51    pub fn into_inner(self) -> S {
52        self.state
53    }
54}
55
56/// Marker: this tool parameter receives the long-term memory store.
57///
58/// The LLM never sees this parameter — it's injected by the runtime.
59/// The value is an `Arc<dyn Store>` from the pe-memory crate.
60///
61/// # Future usage (with `#[tool]` macro, Plan 010)
62///
63/// ```ignore
64/// #[tool]
65/// async fn memory_search(
66///     query: String,
67///     store: InjectedStore,
68/// ) -> Result<Value, PeError> {
69///     let results = store.get().search(&query).await?;
70///     // ...
71/// }
72/// ```
73///
74/// Note: The `Store` trait is defined in pe-memory. We store `Arc<dyn Any>`
75/// here to avoid a circular dependency. The runtime layer casts it back.
76pub struct InjectedStore {
77    store: Arc<dyn std::any::Any + Send + Sync>,
78}
79
80impl InjectedStore {
81    /// Create a new injected store wrapper.
82    pub fn new(store: Arc<dyn std::any::Any + Send + Sync>) -> Self {
83        Self { store }
84    }
85
86    /// Access the wrapped store as `Arc<dyn Any>`.
87    /// The runtime layer downcasts this to the concrete `Store` impl.
88    pub fn get(&self) -> &Arc<dyn std::any::Any + Send + Sync> {
89        &self.store
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    use pe_core::message::Message;
97    use pe_core::state::{CoreState, ExecutionContext, StateUpdate};
98    use serde::{Deserialize, Serialize};
99
100    #[derive(Debug, Clone, Serialize, Deserialize)]
101    struct TestState {
102        messages: Vec<Message>,
103        iterations: u32,
104        thread_id: String,
105        context: ExecutionContext,
106    }
107
108    #[derive(Debug, Clone, Default, Serialize, Deserialize)]
109    struct TestUpdate;
110    impl StateUpdate for TestUpdate {}
111
112    impl pe_core::state::State for TestState {
113        type Update = TestUpdate;
114        fn apply(&mut self, _update: Self::Update) {}
115    }
116
117    impl CoreState for TestState {
118        fn messages(&self) -> &[Message] {
119            &self.messages
120        }
121        fn messages_mut(&mut self) -> &mut Vec<Message> {
122            &mut self.messages
123        }
124        fn iterations(&self) -> u32 {
125            self.iterations
126        }
127        fn set_iterations(&mut self, n: u32) {
128            self.iterations = n;
129        }
130        fn thread_id(&self) -> &str {
131            &self.thread_id
132        }
133        fn context(&self) -> &ExecutionContext {
134            &self.context
135        }
136        fn context_mut(&mut self) -> &mut ExecutionContext {
137            &mut self.context
138        }
139    }
140
141    #[test]
142    fn injected_state_wraps_and_unwraps() {
143        let state = TestState {
144            messages: vec![Message::human("hello")],
145            iterations: 0,
146            thread_id: "t1".into(),
147            context: ExecutionContext::new("agent-1"),
148        };
149
150        let injected = InjectedState::new(state.clone());
151        assert_eq!(injected.get().messages().len(), 1);
152        assert_eq!(injected.get().thread_id(), "t1");
153
154        let inner = injected.into_inner();
155        assert_eq!(inner.iterations(), 0);
156    }
157
158    #[test]
159    fn injected_store_wraps_any() {
160        let data: Arc<dyn std::any::Any + Send + Sync> = Arc::new(42_u32);
161        let injected = InjectedStore::new(data);
162        let store = injected.get();
163        let value = store.downcast_ref::<u32>().unwrap();
164        assert_eq!(*value, 42);
165    }
166}
pe_tools/inject.rs

pe_tools/
inject.rs
