use std::{
collections::VecDeque,
sync::{
Arc, Mutex,
atomic::{AtomicU8, Ordering},
},
time::Duration,
};
use rust_decimal::Decimal;
use tokio::time::Instant;
use crate::{
ActionId, ChatError, Inference, LlmError, LlmFuture, LlmThread, SubmittedUpdate,
ToolOutput, UpdateSink, Updates, infer_with_retry,
};
type Trace = Arc<Mutex<Vec<(usize, String, u8)>>>;
struct RetryThread {
outcomes: VecDeque<Result<Inference, LlmError>>,
trace: Trace,
attempt: Arc<AtomicU8>,
}
impl LlmThread for RetryThread {
fn infer<'a>(&'a mut self, delta: &'a str) -> LlmFuture<'a> {
let identity = self as *const Self as usize;
self.trace.lock().unwrap().push((
identity,
delta.to_owned(),
self.attempt.load(Ordering::Relaxed),
));
let outcome = self.outcomes.pop_front().unwrap();
Box::pin(async move { outcome })
}
}
fn inference(text: &str) -> Inference {
Inference {
text: text.to_owned(),
calls: Vec::new(),
continue_inference: false,
}
}
#[tokio::test(start_paused = true)]
async fn transient_retries_keep_object_delta_timing_and_live_attempt() {
let attempt = Arc::new(AtomicU8::new(0));
let trace = Arc::new(Mutex::new(Vec::new()));
let outcomes = (1..=5)
.map(|number| Err(LlmError::Transient(format!("failure {number}"))))
.collect();
let start = Instant::now();
let task = tokio::spawn(infer_with_retry(
Box::new(RetryThread {
outcomes,
trace: trace.clone(),
attempt: attempt.clone(),
}),
"delta".to_owned(),
attempt.clone(),
));
tokio::task::yield_now().await;
assert_eq!(attempt.load(Ordering::Relaxed), 1);
let mut elapsed = 0;
for (next, wait) in [(2, 10), (3, 20), (4, 40), (5, 80)] {
tokio::time::advance(Duration::from_secs(wait - 1)).await;
tokio::task::yield_now().await;
assert_eq!(attempt.load(Ordering::Relaxed), next - 1);
assert_eq!(start.elapsed(), Duration::from_secs(elapsed + wait - 1));
tokio::time::advance(Duration::from_secs(1)).await;
tokio::task::yield_now().await;
elapsed += wait;
assert_eq!(attempt.load(Ordering::Relaxed), next);
}
let (_, result) = task.await.unwrap();
assert_eq!(result, Err("failure 5".to_owned()));
let trace = trace.lock().unwrap();
assert_eq!(trace.len(), 5);
let object = trace[0].0;
assert!(trace.iter().all(|(seen, delta, _)| {
*seen == object && delta == "delta"
}));
assert_eq!(
trace.iter().map(|(_, _, seen)| *seen).collect::<Vec<_>>(),
vec![1, 2, 3, 4, 5]
);
}
#[tokio::test(start_paused = true)]
async fn permanent_failure_has_no_wait_or_retry() {
let attempt = Arc::new(AtomicU8::new(0));
let trace = Arc::new(Mutex::new(Vec::new()));
let start = Instant::now();
let (_, result) = infer_with_retry(
Box::new(RetryThread {
outcomes: vec![Err(LlmError::Permanent("stop".to_owned()))].into(),
trace: trace.clone(),
attempt: attempt.clone(),
}),
"delta".to_owned(),
attempt.clone(),
)
.await;
assert_eq!(result, Err("stop".to_owned()));
assert_eq!(start.elapsed(), Duration::ZERO);
assert_eq!(attempt.load(Ordering::Relaxed), 1);
assert_eq!(trace.lock().unwrap().len(), 1);
}
#[tokio::test]
async fn success_returns_the_object_and_output() {
let attempt = Arc::new(AtomicU8::new(0));
let trace = Arc::new(Mutex::new(Vec::new()));
let (mut thread, result) = infer_with_retry(
Box::new(RetryThread {
outcomes: vec![Ok(inference("done")), Ok(inference("next"))].into(),
trace: trace.clone(),
attempt: attempt.clone(),
}),
"delta".to_owned(),
attempt,
)
.await;
assert_eq!(result, Ok(inference("done")));
assert_eq!(thread.infer("after").await, Ok(inference("next")));
let trace = trace.lock().unwrap();
assert_eq!(trace.len(), 2);
assert_eq!(trace[0].0, trace[1].0);
}
struct RecordingSink {
submissions: Mutex<Vec<(ActionId, u64, SubmittedUpdate)>>,
failure: Option<ChatError>,
}
impl UpdateSink for RecordingSink {
fn submit(
&self,
action: ActionId,
identity: u64,
update: SubmittedUpdate,
) -> Result<(), ChatError> {
if let Some(error) = &self.failure {
return Err(error.clone());
}
self.submissions
.lock()
.unwrap()
.push((action, identity, update));
Ok(())
}
}
#[test]
fn action_identity_output_and_update_delivery_are_exact() {
let first = ActionId::new([3; 12], 7);
let second = ActionId::new([3; 12], 8);
assert_eq!(first.session(), [3; 12]);
assert_eq!(first.sequence(), 7);
assert!(first < second);
assert_eq!(
ToolOutput {
text: "done".into(),
cost_cents: Decimal::new(31, 2),
},
ToolOutput {
text: "done".into(),
cost_cents: Decimal::new(31, 2),
}
);
let sink = Arc::new(RecordingSink {
submissions: Mutex::new(Vec::new()),
failure: None,
});
let updates = Updates::bind(first, sink.clone());
assert_eq!(updates.action_id(), first);
let activity = updates.activity("working".into());
let clone = activity.clone();
let append = updates.clone().append("answer".into());
activity.send().unwrap();
clone.send().unwrap();
append.send().unwrap();
assert_eq!(
*sink.submissions.lock().unwrap(),
vec![
(first, 1, SubmittedUpdate::Activity("working".into())),
(first, 1, SubmittedUpdate::Activity("working".into())),
(first, 2, SubmittedUpdate::Append("answer".into())),
]
);
}
#[test]
fn sink_failure_passes_through_unchanged() {
let sink = Arc::new(RecordingSink {
submissions: Mutex::new(Vec::new()),
failure: Some(ChatError::Closed),
});
let updates = Updates::bind(ActionId::new([1; 12], 2), sink);
assert_eq!(
updates.append("answer".into()).send(),
Err(ChatError::Closed)
);
}