trybench_sdk/
simulation.rs

1use crate::{Application, Bench, EvaluationContext, EvaluationOptions, EvaluationReport};
2use serde_json::Value;
3use std::{
4    collections::HashMap,
5    future::Future,
6    pin::Pin,
7    sync::{Arc, Mutex},
8    time::Duration,
9};
10
11type ValueFuture = Pin<Box<dyn Future<Output = Result<Value, String>> + Send>>;
12type CloseFuture = Pin<Box<dyn Future<Output = Result<(), String>> + Send>>;
13type Turn = Arc<dyn Fn(Value, EvaluationContext) -> ValueFuture + Send + Sync>;
14type Observe = Arc<dyn Fn() -> ValueFuture + Send + Sync>;
15type Close = Arc<dyn Fn() -> CloseFuture + Send + Sync>;
16
17pub struct SimulationSession {
18    turn: Turn,
19    observe: Observe,
20    close: Option<Close>,
21}
22impl SimulationSession {
23    pub fn new<T, TF, O, OF, C, CF>(turn: T, observe: O, close: C) -> Self
24    where
25        T: Fn(Value, EvaluationContext) -> TF + Send + Sync + 'static,
26        TF: Future<Output = Result<Value, String>> + Send + 'static,
27        O: Fn() -> OF + Send + Sync + 'static,
28        OF: Future<Output = Result<Value, String>> + Send + 'static,
29        C: Fn() -> CF + Send + Sync + 'static,
30        CF: Future<Output = Result<(), String>> + Send + 'static,
31    {
32        Self {
33            turn: Arc::new(move |value, ctx| Box::pin(turn(value, ctx))),
34            observe: Arc::new(move || Box::pin(observe())),
35            close: Some(Arc::new(move || Box::pin(close()))),
36        }
37    }
38    async fn close(&mut self) -> Result<(), String> {
39        if let Some(close) = self.close.take() {
40            tokio::time::timeout(Duration::from_secs(5), close())
41                .await
42                .map_err(|_| "Session cleanup timed out.")??;
43        }
44        Ok(())
45    }
46}
47impl Drop for SimulationSession {
48    fn drop(&mut self) {
49        if let Some(close) = self.close.take() {
50            if let Ok(handle) = tokio::runtime::Handle::try_current() {
51                handle.spawn(async move {
52                    let _ = tokio::time::timeout(Duration::from_secs(5), close()).await;
53                });
54            }
55        }
56    }
57}
58
59impl Bench {
60    /// Scripted turns keep the real app and inspect fixture state before cleanup.
61    pub async fn simulate_system<F, Fut>(
62        &self,
63        options: EvaluationOptions,
64        create_session: F,
65    ) -> Result<EvaluationReport, String>
66    where
67        F: Fn(Value, EvaluationContext) -> Fut + Send + Sync + 'static,
68        Fut: Future<Output = Result<SimulationSession, String>> + Send + 'static,
69    {
70        for case in &options.cases {
71            let turns = case
72                .input
73                .get("turns")
74                .and_then(Value::as_array)
75                .ok_or("Provide scripted user turns.")?;
76            if turns.is_empty()
77                || turns.len() > 20
78                || case.input.get("initialState").is_none()
79                || case.expected_state.is_none()
80            {
81                return Err(
82                    "Provide initialState, 1 to 20 turns and an expected business state.".into(),
83                );
84            }
85        }
86        let factory = Arc::new(create_session);
87        let observations = Arc::new(Mutex::new(HashMap::new()));
88        let observed = observations.clone();
89        let app = Application::new(move |input, context| {
90            let factory = factory.clone();
91            let observations = observations.clone();
92            async move {
93                let mut session = factory(input["initialState"].clone(), context.clone()).await?;
94                let outcome = async {
95                    let mut reply = Value::Null;
96                    for message in input["turns"].as_array().ok_or("Invalid turns.")? {
97                        if context.is_cancelled() {
98                            return Err("Simulation stopped.".into());
99                        }
100                        reply = (session.turn)(message.clone(), context.clone()).await?;
101                    }
102                    if context.is_cancelled() {
103                        return Err("Simulation stopped.".into());
104                    }
105                    let state = (session.observe)().await?;
106                    observations
107                        .lock()
108                        .unwrap_or_else(|e| e.into_inner())
109                        .insert(context.case_id.clone(), state);
110                    Ok::<_, String>(reply)
111                }
112                .await;
113                let cleanup = session.close().await;
114                let reply = outcome?;
115                cleanup?;
116                Ok(reply)
117            }
118        })
119        .with_observer(move |context| {
120            let observed = observed.clone();
121            async move {
122                observed
123                    .lock()
124                    .unwrap_or_else(|e| e.into_inner())
125                    .remove(&context.case_id)
126                    .ok_or_else(|| "Missing observed state.".into())
127            }
128        });
129        self.evaluate_system(options, app).await
130    }
131}
trybench_sdk/simulation.rs

trybench_sdk/
simulation.rs
