use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use tokio_util::sync::CancellationToken;
use crate::error::EngineError;
use crate::traits::context::ContextBuilder;
use crate::traits::processor::OutputProcessor;
use crate::traits::thinker::Thinker;
use crate::types::signal::Signal;
trait ErasedThinker<Ctx, Out>: Send + Sync {
fn think_erased<'a>(
&'a self,
ctx: &'a Ctx,
) -> Pin<Box<dyn Future<Output = Result<Out, EngineError>> + Send + 'a>>;
}
trait ErasedContextBuilder<Ctx>: Send + Sync {
fn build_erased<'a>(
&'a self,
ctx: Ctx,
) -> Pin<Box<dyn Future<Output = Result<Ctx, EngineError>> + Send + 'a>>;
}
trait ErasedOutputProcessor<Ctx, Out>: Send + Sync {
fn process_erased<'a>(
&'a self,
output: &'a Out,
ctx: &'a mut Ctx,
) -> Pin<Box<dyn Future<Output = Result<Signal, EngineError>> + Send + 'a>>;
}
impl<T, Ctx, Out> ErasedThinker<Ctx, Out> for T
where
T: Thinker<Context = Ctx, Output = Out> + 'static,
Ctx: Send + Sync,
Out: Send,
{
fn think_erased<'a>(
&'a self,
ctx: &'a Ctx,
) -> Pin<Box<dyn Future<Output = Result<Out, EngineError>> + Send + 'a>> {
Box::pin(T::think(self, ctx))
}
}
impl<T, Ctx> ErasedContextBuilder<Ctx> for T
where
T: ContextBuilder<Context = Ctx> + 'static,
Ctx: Send,
{
fn build_erased<'a>(
&'a self,
ctx: Ctx,
) -> Pin<Box<dyn Future<Output = Result<Ctx, EngineError>> + Send + 'a>> {
Box::pin(T::build(self, ctx))
}
}
impl<T, Ctx, Out> ErasedOutputProcessor<Ctx, Out> for T
where
T: OutputProcessor<Context = Ctx, Output = Out> + 'static,
Ctx: Send,
Out: Sync,
{
fn process_erased<'a>(
&'a self,
output: &'a Out,
ctx: &'a mut Ctx,
) -> Pin<Box<dyn Future<Output = Result<Signal, EngineError>> + Send + 'a>> {
Box::pin(T::process(self, output, ctx))
}
}
pub struct DynThinker<Ctx, Out>(Arc<dyn ErasedThinker<Ctx, Out>>);
pub struct DynContextBuilder<Ctx>(Arc<dyn ErasedContextBuilder<Ctx>>);
pub struct DynOutputProcessor<Ctx, Out>(Arc<dyn ErasedOutputProcessor<Ctx, Out>>);
impl<Ctx, Out> DynThinker<Ctx, Out> {
pub fn new<T>(thinker: T) -> Self
where
T: Thinker<Context = Ctx, Output = Out> + 'static,
Ctx: Send + Sync,
Out: Send,
{
DynThinker(Arc::new(thinker))
}
}
impl<Ctx> DynContextBuilder<Ctx> {
pub fn new<T>(builder: T) -> Self
where
T: ContextBuilder<Context = Ctx> + 'static,
Ctx: Send,
{
DynContextBuilder(Arc::new(builder))
}
}
impl<Ctx, Out> DynOutputProcessor<Ctx, Out> {
pub fn new<T>(processor: T) -> Self
where
T: OutputProcessor<Context = Ctx, Output = Out> + 'static,
Ctx: Send,
Out: Sync,
{
DynOutputProcessor(Arc::new(processor))
}
}
impl<Ctx, Out> Thinker for DynThinker<Ctx, Out>
where
Ctx: Sync,
{
type Context = Ctx;
type Output = Out;
async fn think(&self, ctx: &Self::Context) -> Result<Self::Output, EngineError> {
self.0.think_erased(ctx).await
}
}
impl<Ctx> ContextBuilder for DynContextBuilder<Ctx>
where
Ctx: Send,
{
type Context = Ctx;
async fn build(&self, ctx: Self::Context) -> Result<Self::Context, EngineError> {
self.0.build_erased(ctx).await
}
}
impl<Ctx, Out> OutputProcessor for DynOutputProcessor<Ctx, Out>
where
Ctx: Send,
Out: Sync,
{
type Context = Ctx;
type Output = Out;
async fn process(
&self,
output: &Self::Output,
ctx: &mut Self::Context,
) -> Result<Signal, EngineError> {
self.0.process_erased(output, ctx).await
}
}
#[derive(Debug, Clone)]
pub struct TurnResult<Out> {
pub output: Out,
pub stopped: bool,
}
pub struct CoreEngine<Ctx, Out> {
context_builders: Vec<DynContextBuilder<Ctx>>,
thinker: DynThinker<Ctx, Out>,
output_processors: Vec<DynOutputProcessor<Ctx, Out>>,
cancel: CancellationToken,
}
impl<Ctx, Out> CoreEngine<Ctx, Out>
where
Ctx: Clone + Send + Sync,
Out: Send + Sync,
{
pub async fn run(&self, initial: Ctx) -> Result<Ctx, EngineError> {
let (ctx, _last) = self.run_with_output(initial).await?;
Ok(ctx)
}
pub async fn run_with_output(&self, initial: Ctx) -> Result<(Ctx, Option<Out>), EngineError> {
let mut ctx = initial;
let mut last_output: Option<Out> = None;
loop {
if self.cancel.is_cancelled() {
return Err(EngineError::Cancelled);
}
for builder in &self.context_builders {
ctx = builder.build(ctx).await?;
}
let output = self.thinker.think(&ctx).await?;
if self.output_processors.is_empty() {
last_output = Some(output);
break;
}
let mut stop = false;
for processor in &self.output_processors {
match processor.process(&output, &mut ctx).await? {
Signal::Stop => {
stop = true;
break;
}
Signal::Continue => {}
}
}
last_output = Some(output);
if stop {
break;
}
}
Ok((ctx, last_output))
}
pub async fn run_once(&self, ctx: &mut Ctx) -> Result<TurnResult<Out>, EngineError> {
if self.cancel.is_cancelled() {
return Err(EngineError::Cancelled);
}
for builder in &self.context_builders {
let next = builder.build(ctx.clone()).await?;
*ctx = next;
}
let output = self.thinker.think(ctx).await?;
let mut stopped = self.output_processors.is_empty();
for processor in &self.output_processors {
match processor.process(&output, ctx).await? {
Signal::Stop => {
stopped = true;
break;
}
Signal::Continue => {}
}
}
Ok(TurnResult { output, stopped })
}
pub fn cancel_handle(&self) -> CancellationToken {
self.cancel.clone()
}
}
pub struct EngineBuilder<Ctx, Out> {
thinker: Option<DynThinker<Ctx, Out>>,
context_builders: Vec<DynContextBuilder<Ctx>>,
output_processors: Vec<DynOutputProcessor<Ctx, Out>>,
cancel: CancellationToken,
}
impl<Ctx, Out> Default for EngineBuilder<Ctx, Out> {
fn default() -> Self {
Self::new()
}
}
impl<Ctx, Out> EngineBuilder<Ctx, Out> {
pub fn new() -> Self {
EngineBuilder {
thinker: None,
context_builders: Vec::new(),
output_processors: Vec::new(),
cancel: CancellationToken::new(),
}
}
}
impl<Ctx, Out> EngineBuilder<Ctx, Out>
where
Ctx: Send + Sync,
Out: Send + Sync,
{
pub fn thinker(mut self, thinker: impl Thinker<Context = Ctx, Output = Out> + 'static) -> Self {
self.thinker = Some(DynThinker::new(thinker));
self
}
pub fn context(mut self, builder: impl ContextBuilder<Context = Ctx> + 'static) -> Self {
self.context_builders.push(DynContextBuilder::new(builder));
self
}
pub fn processor(
mut self,
processor: impl OutputProcessor<Context = Ctx, Output = Out> + 'static,
) -> Self {
self.output_processors.push(DynOutputProcessor::new(processor));
self
}
pub fn cancel(mut self, token: CancellationToken) -> Self {
self.cancel = token;
self
}
pub fn build(self) -> Result<CoreEngine<Ctx, Out>, EngineError> {
let thinker = self
.thinker
.ok_or_else(|| EngineError::Config("thinker is required".into()))?;
Ok(CoreEngine {
context_builders: self.context_builders,
thinker,
output_processors: self.output_processors,
cancel: self.cancel,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc as StdArc;
use std::sync::atomic::{AtomicUsize, Ordering};
struct EchoThinker;
impl Thinker for EchoThinker {
type Context = String;
type Output = String;
async fn think(&self, ctx: &Self::Context) -> Result<Self::Output, EngineError> {
Ok(ctx.clone())
}
}
struct FixedThinker {
output: String,
}
impl Thinker for FixedThinker {
type Context = String;
type Output = String;
async fn think(&self, _ctx: &Self::Context) -> Result<Self::Output, EngineError> {
Ok(self.output.clone())
}
}
struct AppendSuffix {
suffix: String,
}
impl ContextBuilder for AppendSuffix {
type Context = String;
async fn build(&self, mut ctx: Self::Context) -> Result<Self::Context, EngineError> {
ctx.push_str(&self.suffix);
Ok(ctx)
}
}
struct CountingBuilder {
count: StdArc<AtomicUsize>,
}
impl CountingBuilder {
fn new(count: StdArc<AtomicUsize>) -> Self {
CountingBuilder { count }
}
}
impl ContextBuilder for CountingBuilder {
type Context = String;
async fn build(&self, ctx: Self::Context) -> Result<Self::Context, EngineError> {
self.count.fetch_add(1, Ordering::SeqCst);
Ok(ctx)
}
}
struct ContinueProcessor;
impl OutputProcessor for ContinueProcessor {
type Context = String;
type Output = String;
async fn process(
&self,
_output: &Self::Output,
_ctx: &mut Self::Context,
) -> Result<Signal, EngineError> {
Ok(Signal::Continue)
}
}
struct StopProcessor;
impl OutputProcessor for StopProcessor {
type Context = String;
type Output = String;
async fn process(
&self,
_output: &Self::Output,
_ctx: &mut Self::Context,
) -> Result<Signal, EngineError> {
Ok(Signal::Stop)
}
}
struct AppendOutput;
impl OutputProcessor for AppendOutput {
type Context = String;
type Output = String;
async fn process(
&self,
output: &Self::Output,
ctx: &mut Self::Context,
) -> Result<Signal, EngineError> {
ctx.push_str(output);
Ok(Signal::Continue)
}
}
struct StopOnMatch {
keyword: &'static str,
}
impl OutputProcessor for StopOnMatch {
type Context = String;
type Output = String;
async fn process(
&self,
output: &Self::Output,
_ctx: &mut Self::Context,
) -> Result<Signal, EngineError> {
if output.contains(self.keyword) { Ok(Signal::Stop) } else { Ok(Signal::Continue) }
}
}
struct CountingEchoThinker {
count: StdArc<AtomicUsize>,
}
impl Thinker for CountingEchoThinker {
type Context = String;
type Output = String;
async fn think(&self, ctx: &Self::Context) -> Result<Self::Output, EngineError> {
self.count.fetch_add(1, Ordering::SeqCst);
Ok(ctx.clone())
}
}
struct RecordingBuilder {
id: &'static str,
count: StdArc<AtomicUsize>,
}
impl ContextBuilder for RecordingBuilder {
type Context = String;
async fn build(&self, mut ctx: Self::Context) -> Result<Self::Context, EngineError> {
self.count.fetch_add(1, Ordering::SeqCst);
ctx.push_str(&format!("|{}", self.id));
Ok(ctx)
}
}
struct CountingSignalProcessor {
count: StdArc<AtomicUsize>,
signal: Signal,
}
impl OutputProcessor for CountingSignalProcessor {
type Context = String;
type Output = String;
async fn process(
&self,
_output: &Self::Output,
_ctx: &mut Self::Context,
) -> Result<Signal, EngineError> {
self.count.fetch_add(1, Ordering::SeqCst);
Ok(self.signal)
}
}
#[test]
fn test_builder_missing_thinker_returns_err() {
let result: Result<CoreEngine<String, String>, _> = EngineBuilder::new().build();
assert!(result.is_err(), "build without thinker should fail");
}
#[test]
fn test_builder_with_thinker_returns_ok() {
let result = EngineBuilder::new().thinker(EchoThinker).build();
assert!(result.is_ok());
}
#[tokio::test]
async fn test_minimal_engine_single_iteration() {
let engine: CoreEngine<String, String> =
EngineBuilder::new().thinker(EchoThinker).processor(StopProcessor).build().unwrap();
let result = engine.run("hello".into()).await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_context_builder_chain() {
let count = StdArc::new(AtomicUsize::new(0));
let engine: CoreEngine<String, String> = EngineBuilder::new()
.context(AppendSuffix { suffix: " world".into() })
.context(CountingBuilder::new(StdArc::clone(&count)))
.context(AppendSuffix { suffix: "!".into() })
.thinker(EchoThinker)
.processor(StopProcessor)
.build()
.unwrap();
let result = engine.run("hello".into()).await;
assert!(result.is_ok());
assert_eq!(count.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_context_builder_modifies_context() {
let engine: CoreEngine<String, String> = EngineBuilder::new()
.context(AppendSuffix { suffix: " world".into() })
.thinker(EchoThinker)
.processor(StopProcessor)
.build()
.unwrap();
let result = engine.run("hello".into()).await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_processor_continues_loop() {
let iteration_count = StdArc::new(AtomicUsize::new(0));
let count_clone = StdArc::clone(&iteration_count);
struct StopAfterN {
count: StdArc<AtomicUsize>,
limit: usize,
}
impl OutputProcessor for StopAfterN {
type Context = String;
type Output = String;
async fn process(
&self,
_output: &Self::Output,
_ctx: &mut Self::Context,
) -> Result<Signal, EngineError> {
let current = self.count.fetch_add(1, Ordering::SeqCst) + 1;
if current >= self.limit { Ok(Signal::Stop) } else { Ok(Signal::Continue) }
}
}
let engine: CoreEngine<String, String> = EngineBuilder::new()
.thinker(EchoThinker)
.processor(StopAfterN { count: StdArc::clone(&count_clone), limit: 3 })
.build()
.unwrap();
let result = engine.run("hello".into()).await;
assert!(result.is_ok());
assert_eq!(count_clone.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn test_stop_signal_breaks_loop() {
let engine: CoreEngine<String, String> =
EngineBuilder::new().thinker(EchoThinker).processor(StopProcessor).build().unwrap();
let result = engine.run("test".into()).await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_cancellation_returns_cancelled_error() {
let token = CancellationToken::new();
token.cancel();
let engine: CoreEngine<String, String> =
EngineBuilder::new().thinker(EchoThinker).cancel(token).build().unwrap();
let result = engine.run("test".into()).await;
assert!(result.is_err());
match result.unwrap_err() {
EngineError::Cancelled => {} other => panic!("expected Cancelled, got {other:?}"),
}
}
#[tokio::test]
async fn test_cancel_handle() {
let engine: CoreEngine<String, String> =
EngineBuilder::new().thinker(EchoThinker).processor(ContinueProcessor).build().unwrap();
let handle = engine.cancel_handle();
assert!(!handle.is_cancelled());
handle.cancel();
assert!(handle.is_cancelled());
let result = engine.run("test".into()).await;
assert!(matches!(result.unwrap_err(), EngineError::Cancelled));
}
#[tokio::test]
async fn test_processor_chain_order() {
let engine: CoreEngine<String, String> = EngineBuilder::new()
.thinker(EchoThinker)
.processor(ContinueProcessor)
.processor(StopProcessor)
.build()
.unwrap();
let result = engine.run("test".into()).await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_processor_modifies_context() {
let count = StdArc::new(AtomicUsize::new(0));
let count2 = StdArc::clone(&count);
struct StopAfterOne {
count: StdArc<AtomicUsize>,
}
impl OutputProcessor for StopAfterOne {
type Context = String;
type Output = String;
async fn process(
&self,
_output: &Self::Output,
_ctx: &mut Self::Context,
) -> Result<Signal, EngineError> {
let n = self.count.fetch_add(1, Ordering::SeqCst) + 1;
if n >= 2 { Ok(Signal::Stop) } else { Ok(Signal::Continue) }
}
}
let engine: CoreEngine<String, String> = EngineBuilder::new()
.thinker(EchoThinker)
.processor(AppendOutput)
.processor(StopAfterOne { count: StdArc::clone(&count2) })
.build()
.unwrap();
let result = engine.run("hello".into()).await;
assert!(result.is_ok());
assert_eq!(count2.load(Ordering::SeqCst), 2);
assert_eq!(result.unwrap(), "hellohellohellohello");
}
#[tokio::test]
async fn test_multiple_context_builders() {
let engine: CoreEngine<String, String> = EngineBuilder::new()
.context(AppendSuffix { suffix: " world".into() })
.context(AppendSuffix { suffix: "!".into() })
.thinker(FixedThinker { output: "done".into() })
.processor(StopOnMatch { keyword: "done" })
.build()
.unwrap();
let result = engine.run("hello".into()).await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_dyn_context_builder_clone() {
let builder = DynContextBuilder::<String>::new(AppendSuffix { suffix: " test".into() });
let ctx: String = builder.build("hello".into()).await.unwrap();
assert_eq!(ctx, "hello test");
}
#[tokio::test]
async fn test_dyn_thinker_clone() {
let thinker = DynThinker::<String, String>::new(EchoThinker);
let result = thinker.think(&"input".to_string()).await.unwrap();
assert_eq!(result, "input");
}
#[tokio::test]
async fn test_dyn_output_processor_clone() {
let processor = DynOutputProcessor::<String, String>::new(StopProcessor);
let mut ctx = String::from("test");
let result = processor.process(&"output".into(), &mut ctx).await.unwrap();
assert_eq!(result, Signal::Stop);
}
#[tokio::test]
async fn test_complete_loop_with_mock() {
let cb_count = StdArc::new(AtomicUsize::new(0));
let t_count = StdArc::new(AtomicUsize::new(0));
let op_count = StdArc::new(AtomicUsize::new(0));
let engine: CoreEngine<String, String> = EngineBuilder::new()
.context(CountingBuilder::new(StdArc::clone(&cb_count)))
.thinker(CountingEchoThinker { count: StdArc::clone(&t_count) })
.processor(CountingSignalProcessor {
count: StdArc::clone(&op_count),
signal: Signal::Stop,
})
.build()
.unwrap();
let result = engine.run("hello".into()).await;
assert!(result.is_ok());
assert_eq!(cb_count.load(Ordering::SeqCst), 1);
assert_eq!(t_count.load(Ordering::SeqCst), 1);
assert_eq!(op_count.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_three_context_builders_executed() {
let count_a = StdArc::new(AtomicUsize::new(0));
let count_b = StdArc::new(AtomicUsize::new(0));
let count_c = StdArc::new(AtomicUsize::new(0));
let engine: CoreEngine<String, String> = EngineBuilder::new()
.context(RecordingBuilder { id: "A", count: StdArc::clone(&count_a) })
.context(RecordingBuilder { id: "B", count: StdArc::clone(&count_b) })
.context(RecordingBuilder { id: "C", count: StdArc::clone(&count_c) })
.thinker(FixedThinker { output: "done".into() })
.processor(StopOnMatch { keyword: "done" })
.build()
.unwrap();
let result = engine.run("init".into()).await;
assert!(result.is_ok());
assert_eq!(count_a.load(Ordering::SeqCst), 1);
assert_eq!(count_b.load(Ordering::SeqCst), 1);
assert_eq!(count_c.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_output_processor_stop_signal() {
let t_count = StdArc::new(AtomicUsize::new(0));
let op_count = StdArc::new(AtomicUsize::new(0));
let engine: CoreEngine<String, String> = EngineBuilder::new()
.thinker(CountingEchoThinker { count: StdArc::clone(&t_count) })
.processor(CountingSignalProcessor {
count: StdArc::clone(&op_count),
signal: Signal::Stop,
})
.build()
.unwrap();
let result = engine.run("test".into()).await;
assert!(result.is_ok());
assert_eq!(t_count.load(Ordering::SeqCst), 1);
assert_eq!(op_count.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_continue_signal_continues() {
let continue_count = StdArc::new(AtomicUsize::new(0));
let stop_count = StdArc::new(AtomicUsize::new(0));
let engine: CoreEngine<String, String> = EngineBuilder::new()
.thinker(EchoThinker)
.processor(CountingSignalProcessor {
count: StdArc::clone(&continue_count),
signal: Signal::Continue,
})
.processor(CountingSignalProcessor {
count: StdArc::clone(&stop_count),
signal: Signal::Stop,
})
.build()
.unwrap();
let result = engine.run("test".into()).await;
assert!(result.is_ok());
assert_eq!(continue_count.load(Ordering::SeqCst), 1);
assert_eq!(stop_count.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_empty_context_and_no_processors() {
let token = CancellationToken::new();
token.cancel();
let engine: CoreEngine<String, String> =
EngineBuilder::new().thinker(EchoThinker).cancel(token).build().unwrap();
let result = engine.run("data".into()).await;
assert!(
matches!(result, Err(EngineError::Cancelled)),
"expected Cancelled, got {result:?}"
);
}
#[tokio::test]
async fn test_run_returns_final_context() {
let engine: CoreEngine<String, String> = EngineBuilder::new()
.context(AppendSuffix { suffix: " world".into() })
.thinker(EchoThinker)
.processor(AppendOutput)
.processor(StopProcessor)
.build()
.unwrap();
let final_ctx = engine.run("hi".into()).await.unwrap();
assert_eq!(final_ctx, "hi worldhi world");
}
#[tokio::test]
async fn test_no_processors_stops_after_one_turn() {
let t_count = StdArc::new(AtomicUsize::new(0));
let engine: CoreEngine<String, String> = EngineBuilder::new()
.thinker(CountingEchoThinker { count: StdArc::clone(&t_count) })
.build()
.unwrap();
let final_ctx = engine.run("once".into()).await.unwrap();
assert_eq!(final_ctx, "once");
assert_eq!(t_count.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_iterator_counting() {
let iter_count = StdArc::new(AtomicUsize::new(0));
let count_clone = StdArc::clone(&iter_count);
struct StopAfterN {
count: StdArc<AtomicUsize>,
limit: usize,
}
impl OutputProcessor for StopAfterN {
type Context = String;
type Output = String;
async fn process(
&self,
_output: &Self::Output,
_ctx: &mut Self::Context,
) -> Result<Signal, EngineError> {
let current = self.count.fetch_add(1, Ordering::SeqCst) + 1;
if current >= self.limit { Ok(Signal::Stop) } else { Ok(Signal::Continue) }
}
}
let engine: CoreEngine<String, String> = EngineBuilder::new()
.thinker(EchoThinker)
.processor(StopAfterN { count: StdArc::clone(&count_clone), limit: 5 })
.build()
.unwrap();
let result = engine.run("start".into()).await;
assert!(result.is_ok());
assert_eq!(count_clone.load(Ordering::SeqCst), 5);
}
}