runifold_testkit/
fault.rs

1//! Deterministic fault injection shared by model, Tool, and recovery tests.
2
3use std::{
4    collections::{BTreeMap, VecDeque},
5    sync::{Arc, Mutex},
6};
7
8use futures_util::StreamExt;
9use runifold_model::{
10    ContentPart, Model, ModelCallContext, ModelCapabilities, ModelError, ModelEventStream,
11    ModelFuture, ModelRef, ModelRequest, ModelStreamEvent,
12};
13use runifold_tool::{Tool, ToolContext, ToolDescriptor, ToolError, ToolFuture, ToolOutput};
14use serde_json::Value;
15use thiserror::Error;
16
17/// Fault applied to one model invocation.
18#[derive(Clone, Debug)]
19#[non_exhaustive]
20pub enum ModelFault {
21    /// Executes the invocation without injection.
22    Pass,
23    /// Fails before a canonical stream is opened.
24    FailOpen(ModelError),
25    /// Ends the stream after the specified number of canonical events.
26    DisconnectAfterEvents(usize),
27    /// Ends the stream immediately after the first completed Tool call.
28    DisconnectAfterToolCall,
29}
30
31/// Fault applied to one named Tool invocation.
32#[derive(Clone, Debug)]
33pub struct ToolFault {
34    tool: String,
35    invocation: u64,
36    error: ToolError,
37}
38
39impl ToolFault {
40    /// Fails the selected one-based invocation of a named Tool.
41    pub fn on_invocation(tool: impl Into<String>, invocation: u64, error: ToolError) -> Self {
42        Self {
43            tool: tool.into(),
44            invocation,
45            error,
46        }
47    }
48}
49
50#[derive(Debug, Default)]
51struct FaultState {
52    model_faults: VecDeque<ModelFault>,
53    tool_faults: Vec<ToolFault>,
54    tool_invocations: BTreeMap<String, u64>,
55    tool_executions: BTreeMap<String, u64>,
56    runtime_restarts: u64,
57}
58
59/// Shared deterministic fault plan and execution counters.
60#[derive(Clone, Debug, Default)]
61pub struct FaultController {
62    state: Arc<Mutex<FaultState>>,
63}
64
65/// Fluent scenario name for [`FaultController`].
66pub type FaultScenario = FaultController;
67
68impl FaultController {
69    /// Creates an empty fault plan.
70    pub fn new() -> Self {
71        Self::default()
72    }
73
74    /// Queues a fault for the next model invocation.
75    #[must_use]
76    pub fn model_fault(self, fault: ModelFault) -> Self {
77        self.lock().model_faults.push_back(fault);
78        self
79    }
80
81    /// Disconnects the next model stream after a completed Tool call.
82    #[must_use]
83    pub fn disconnect_after_tool_call(self) -> Self {
84        self.model_fault(ModelFault::DisconnectAfterToolCall)
85    }
86
87    /// Disconnects the next model stream after `events` canonical events.
88    #[must_use]
89    pub fn disconnect_after_model_events(self, events: usize) -> Self {
90        self.model_fault(ModelFault::DisconnectAfterEvents(events))
91    }
92
93    /// Fails the next model invocation before opening its stream.
94    #[must_use]
95    pub fn fail_next_model(self, error: ModelError) -> Self {
96        self.model_fault(ModelFault::FailOpen(error))
97    }
98
99    /// Queues a named Tool fault.
100    #[must_use]
101    pub fn tool_fault(self, fault: ToolFault) -> Self {
102        self.lock().tool_faults.push(fault);
103        self
104    }
105
106    /// Fails the selected one-based invocation of a named Tool.
107    #[must_use]
108    pub fn fail_tool_on_invocation(
109        self,
110        tool: impl Into<String>,
111        invocation: u64,
112        error: ToolError,
113    ) -> Self {
114        self.tool_fault(ToolFault::on_invocation(tool, invocation, error))
115    }
116
117    /// Wraps a model with this plan.
118    pub fn model<M>(&self, model: M) -> FaultInjectingModel<M>
119    where
120        M: Model,
121    {
122        FaultInjectingModel {
123            inner: model,
124            controller: self.clone(),
125        }
126    }
127
128    /// Wraps a Tool with this plan.
129    pub fn tool(&self, tool: Arc<dyn Tool>) -> FaultInjectingTool {
130        FaultInjectingTool {
131            inner: tool,
132            controller: self.clone(),
133        }
134    }
135
136    /// Returns the number of observed invocations for one Tool.
137    pub fn tool_invocations(&self, tool: &str) -> u64 {
138        self.lock().tool_invocations.get(tool).copied().unwrap_or(0)
139    }
140
141    /// Returns the number of calls forwarded to the underlying Tool.
142    pub fn tool_executions(&self, tool: &str) -> u64 {
143        self.lock().tool_executions.get(tool).copied().unwrap_or(0)
144    }
145
146    /// Returns the number of explicit runtime reconstructions.
147    pub fn runtime_restarts(&self) -> u64 {
148        self.lock().runtime_restarts
149    }
150
151    /// Asserts exact Tool execution count with a typed failure.
152    ///
153    /// # Errors
154    ///
155    /// Returns [`ScenarioAssertionError`] when the observed count differs.
156    pub fn assert_tool_executed_exactly(
157        &self,
158        tool: &str,
159        expected: u64,
160    ) -> Result<(), ScenarioAssertionError> {
161        let actual = self.tool_executions(tool);
162        if actual == expected {
163            Ok(())
164        } else {
165            Err(ScenarioAssertionError::ToolInvocationCount {
166                tool: tool.into(),
167                expected,
168                actual,
169            })
170        }
171    }
172
173    fn lock(&self) -> std::sync::MutexGuard<'_, FaultState> {
174        self.state
175            .lock()
176            .unwrap_or_else(std::sync::PoisonError::into_inner)
177    }
178
179    fn next_model_fault(&self) -> ModelFault {
180        self.lock()
181            .model_faults
182            .pop_front()
183            .unwrap_or(ModelFault::Pass)
184    }
185
186    fn tool_result(&self, tool: &str) -> Option<ToolError> {
187        let mut state = self.lock();
188        let invocation = {
189            let count = state.tool_invocations.entry(tool.into()).or_default();
190            *count = count.saturating_add(1);
191            *count
192        };
193        let error = state
194            .tool_faults
195            .iter()
196            .find(|fault| fault.tool == tool && fault.invocation == invocation)
197            .map(|fault| fault.error.clone());
198        if error.is_none() {
199            let count = state.tool_executions.entry(tool.into()).or_default();
200            *count = count.saturating_add(1);
201        }
202        error
203    }
204
205    fn record_restart(&self) {
206        let mut state = self.lock();
207        state.runtime_restarts = state.runtime_restarts.saturating_add(1);
208    }
209}
210
211/// Model wrapper applying one queued fault per invocation.
212#[derive(Clone, Debug)]
213pub struct FaultInjectingModel<M> {
214    inner: M,
215    controller: FaultController,
216}
217
218impl<M> Model for FaultInjectingModel<M>
219where
220    M: Model,
221{
222    fn capabilities<'a>(
223        &'a self,
224        model: &'a ModelRef,
225    ) -> ModelFuture<'a, Result<ModelCapabilities, ModelError>> {
226        self.inner.capabilities(model)
227    }
228
229    fn stream(
230        &self,
231        request: ModelRequest,
232        context: ModelCallContext,
233    ) -> ModelFuture<'_, Result<ModelEventStream, ModelError>> {
234        let fault = self.controller.next_model_fault();
235        Box::pin(async move {
236            if let ModelFault::FailOpen(error) = fault {
237                return Err(error);
238            }
239            let stream = self.inner.stream(request, context).await?;
240            match fault {
241                ModelFault::Pass => Ok(stream),
242                ModelFault::DisconnectAfterEvents(limit) => Ok(Box::pin(stream.take(limit))),
243                ModelFault::DisconnectAfterToolCall => {
244                    let truncated = futures_util::stream::unfold(
245                        (stream, false),
246                        |(mut stream, stopped)| async move {
247                            if stopped {
248                                return None;
249                            }
250                            let item = stream.next().await?;
251                            let stop = matches!(
252                                &item,
253                                Ok(ModelStreamEvent::ContentPartCompleted {
254                                    part: ContentPart::ToolCall(_),
255                                    ..
256                                })
257                            );
258                            Some((item, (stream, stop)))
259                        },
260                    );
261                    Ok(Box::pin(truncated))
262                }
263                ModelFault::FailOpen(_) => unreachable!("handled before opening the stream"),
264            }
265        })
266    }
267}
268
269/// Tool wrapper applying named, one-based invocation faults.
270#[derive(Clone)]
271pub struct FaultInjectingTool {
272    inner: Arc<dyn Tool>,
273    controller: FaultController,
274}
275
276impl Tool for FaultInjectingTool {
277    fn descriptor(&self) -> &ToolDescriptor {
278        self.inner.descriptor()
279    }
280
281    fn invoke(
282        &self,
283        input: Value,
284        context: ToolContext,
285    ) -> ToolFuture<'_, Result<ToolOutput, ToolError>> {
286        let injected = self.controller.tool_result(&self.inner.descriptor().name);
287        if let Some(error) = injected {
288            return Box::pin(async move { Err(error) });
289        }
290        self.inner.invoke(input, context)
291    }
292}
293
294impl std::fmt::Debug for FaultInjectingTool {
295    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
296        formatter
297            .debug_struct("FaultInjectingTool")
298            .field("tool", &self.inner.descriptor().name)
299            .finish_non_exhaustive()
300    }
301}
302
303/// Reconstructs runtime state from one factory while preserving external
304/// fixtures such as stores, fault counters, and cassette servers.
305pub struct RecoveryHarness<T, F>
306where
307    F: FnMut() -> T,
308{
309    current: T,
310    factory: F,
311    controller: FaultController,
312}
313
314impl<T, F> RecoveryHarness<T, F>
315where
316    F: FnMut() -> T,
317{
318    /// Constructs the first runtime instance.
319    pub fn new(mut factory: F, controller: FaultController) -> Self {
320        let current = factory();
321        Self {
322            current,
323            factory,
324            controller,
325        }
326    }
327
328    /// Returns the active runtime fixture.
329    pub const fn current(&self) -> &T {
330        &self.current
331    }
332
333    /// Returns the active runtime fixture mutably.
334    pub fn current_mut(&mut self) -> &mut T {
335        &mut self.current
336    }
337
338    /// Drops and reconstructs runtime-owned state while retaining the shared
339    /// external fixture and fault controller.
340    pub fn restart(&mut self) -> &mut T {
341        self.current = (self.factory)();
342        self.controller.record_restart();
343        &mut self.current
344    }
345}
346
347/// Typed failure from a scenario assertion.
348#[derive(Clone, Debug, Error, Eq, PartialEq)]
349#[non_exhaustive]
350pub enum ScenarioAssertionError {
351    /// A Tool did not execute the expected number of times.
352    #[error("tool `{tool}` executed {actual} times; expected {expected}")]
353    ToolInvocationCount {
354        /// Tool name.
355        tool: String,
356        /// Expected count.
357        expected: u64,
358        /// Observed count.
359        actual: u64,
360    },
361}
362
363#[cfg(test)]
364mod tests {
365    use std::collections::BTreeMap;
366
367    use runifold_core::RetrySafety;
368    use runifold_model::{
369        ContentPart, FinishReason, Message, Model, ModelCallContext, ModelErrorKind, ModelRef,
370        ModelRequest, ModelStreamEvent, ToolCall,
371    };
372    use runifold_tool::{ToolError, ToolErrorKind};
373
374    use crate::ScriptedModel;
375
376    use super::{FaultController, ModelFault, RecoveryHarness, ToolFault};
377
378    #[test]
379    fn disconnect_after_tool_call_never_reaches_terminal_success() {
380        let model = ScriptedModel::new();
381        let model_ref = ModelRef::new("test", "model");
382        model.enqueue([
383            ModelStreamEvent::ResponseStarted {
384                id: Some("response".into()),
385                model: model_ref.clone(),
386            },
387            ModelStreamEvent::ContentPartCompleted {
388                index: 0,
389                part: ContentPart::ToolCall(ToolCall {
390                    id: "call".into(),
391                    name: "charge".into(),
392                    arguments: serde_json::json!({}),
393                    raw_arguments: Some("{}".into()),
394                    metadata: BTreeMap::new(),
395                }),
396            },
397            ModelStreamEvent::ResponseCompleted {
398                finish_reason: FinishReason::ToolCalls,
399                provider_metadata: BTreeMap::new(),
400            },
401        ]);
402        let controller = FaultController::new().model_fault(ModelFault::DisconnectAfterToolCall);
403        let model = controller.model(model);
404        let request = ModelRequest::new(model_ref, Message::user("run"));
405
406        let error =
407            futures_executor::block_on(model.invoke(request, ModelCallContext::new())).unwrap_err();
408
409        assert_eq!(error.kind, ModelErrorKind::Protocol);
410    }
411
412    #[test]
413    fn tool_faults_are_one_based_and_distinguish_attempts_from_execution() {
414        let mut error = ToolError::local(ToolErrorKind::Execution, "injected");
415        error.retry_safety = RetrySafety::Safe;
416        let controller =
417            FaultController::new().tool_fault(ToolFault::on_invocation("charge", 1, error.clone()));
418
419        assert_eq!(controller.tool_result("charge"), Some(error));
420        assert_eq!(controller.tool_result("charge"), None);
421        assert_eq!(controller.tool_invocations("charge"), 2);
422        assert_eq!(controller.tool_executions("charge"), 1);
423        controller
424            .assert_tool_executed_exactly("charge", 1)
425            .unwrap();
426    }
427
428    #[test]
429    fn recovery_harness_reconstructs_runtime_and_counts_restarts() {
430        let controller = FaultController::new();
431        let mut next = 0_u64;
432        let mut harness = RecoveryHarness::new(
433            || {
434                next += 1;
435                next
436            },
437            controller.clone(),
438        );
439
440        assert_eq!(*harness.current(), 1);
441        assert_eq!(*harness.restart(), 2);
442        assert_eq!(controller.runtime_restarts(), 1);
443    }
444}
runifold_testkit/fault.rs

runifold_testkit/
fault.rs
