use pe_core::state::State;
use std::marker::PhantomData;
use std::sync::Arc;
pub struct InjectedState<S: State> {
state: S,
_phantom: PhantomData<S>,
}
impl<S: State> InjectedState<S> {
pub fn new(state: S) -> Self {
Self {
state,
_phantom: PhantomData,
}
}
pub fn get(&self) -> &S {
&self.state
}
pub fn into_inner(self) -> S {
self.state
}
}
pub struct InjectedStore {
store: Arc<dyn std::any::Any + Send + Sync>,
}
impl InjectedStore {
pub fn new(store: Arc<dyn std::any::Any + Send + Sync>) -> Self {
Self { store }
}
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);
}
}