#![allow(
clippy::unwrap_used,
clippy::expect_used,
clippy::panic,
clippy::str_to_string
)]
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::time::sleep;
use async_trait::async_trait;
use serde_json::json;
use embacle::fallback::{is_empty_completion, FallbackProvider, QuotaCooldowns, ResponsePolicy};
use embacle::types::{
ChatMessage, ChatRequest, ChatResponse, ChatStream, ErrorKind, LlmCapabilities, LlmProvider,
RunnerError, ToolCallRequest,
};
enum Scripted {
Says(&'static str),
CallsATool,
Empty,
ProviderFault,
InvalidRequest,
QuotaRefusal,
}
struct Fake {
label: &'static str,
script: Scripted,
calls: Arc<AtomicUsize>,
models: Vec<String>,
}
impl Fake {
fn scripted(label: &'static str, script: Scripted) -> (Box<dyn LlmProvider>, Arc<AtomicUsize>) {
let calls = Arc::new(AtomicUsize::new(0));
let provider = Box::new(Self {
label,
script,
calls: Arc::clone(&calls),
models: vec!["test-model".to_owned()],
});
(provider, calls)
}
}
fn reply(content: &str, tool_calls: Option<Vec<ToolCallRequest>>) -> ChatResponse {
ChatResponse {
content: content.to_owned(),
model: "test-model".to_owned(),
usage: None,
finish_reason: Some("stop".to_owned()),
warnings: None,
tool_calls,
}
}
fn a_tool_call() -> ToolCallRequest {
ToolCallRequest {
id: "call_1".to_owned(),
function_name: "get_activities".to_owned(),
arguments: json!({"limit": 10}),
}
}
#[async_trait]
impl LlmProvider for Fake {
fn name(&self) -> &'static str {
self.label
}
fn display_name(&self) -> &str {
self.label
}
fn capabilities(&self) -> LlmCapabilities {
LlmCapabilities::FUNCTION_CALLING
}
fn default_model(&self) -> &str {
"test-model"
}
fn available_models(&self) -> &[String] {
&self.models
}
async fn complete(&self, _request: &ChatRequest) -> Result<ChatResponse, RunnerError> {
self.calls.fetch_add(1, Ordering::SeqCst);
match self.script {
Scripted::Says(text) => Ok(reply(text, None)),
Scripted::CallsATool => Ok(reply("", Some(vec![a_tool_call()]))),
Scripted::Empty => Ok(reply("", None)),
Scripted::ProviderFault => Err(RunnerError::external_service(
self.label,
"upstream returned 503",
)),
Scripted::InvalidRequest => Err(RunnerError::invalid_request(
"the request body was malformed",
)),
Scripted::QuotaRefusal => {
Err(RunnerError::rate_limit(self.label, "usage limit reached"))
}
}
}
async fn complete_stream(&self, _request: &ChatRequest) -> Result<ChatStream, RunnerError> {
Err(RunnerError::invalid_request(
"the scripted tier answers only complete()",
))
}
async fn health_check(&self) -> Result<bool, RunnerError> {
Ok(true)
}
}
fn strict_chain(
primary: Box<dyn LlmProvider>,
secondary: Box<dyn LlmProvider>,
) -> FallbackProvider {
FallbackProvider::new(vec![primary, secondary])
.expect("two tiers")
.with_fallthrough(ResponsePolicy::strict())
}
fn a_turn() -> ChatRequest {
ChatRequest::new(vec![
ChatMessage::system("coach persona"),
ChatMessage::user("how did my week go?"),
])
.with_model("primary-only-model")
.with_tools(vec![])
}
#[tokio::test]
async fn a_provider_fault_falls_through_to_the_secondary() {
let (primary, primary_calls) = Fake::scripted("primary", Scripted::ProviderFault);
let (secondary, secondary_calls) =
Fake::scripted("secondary", Scripted::Says("42 km this week."));
let response = strict_chain(primary, secondary)
.complete(&a_turn())
.await
.expect("a provider fault must be answered by the secondary");
assert_eq!(response.content, "42 km this week.");
assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
assert_eq!(
secondary_calls.load(Ordering::SeqCst),
1,
"the secondary is consulted exactly once"
);
}
#[tokio::test]
async fn an_empty_completion_falls_through_to_the_secondary() {
let (primary, primary_calls) = Fake::scripted("primary", Scripted::Empty);
let (secondary, secondary_calls) =
Fake::scripted("secondary", Scripted::Says("Tu as couru 42 km ce mois-ci."));
let response = strict_chain(primary, secondary)
.complete(&a_turn())
.await
.expect("an empty primary completion must be reissued against the secondary");
assert_eq!(response.content, "Tu as couru 42 km ce mois-ci.");
assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
assert_eq!(secondary_calls.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn an_empty_answer_carrying_a_tool_call_stays_with_the_primary() {
let (primary, primary_calls) = Fake::scripted("primary", Scripted::CallsATool);
let (secondary, secondary_calls) =
Fake::scripted("secondary", Scripted::Says("should not be used"));
let response = strict_chain(primary, secondary)
.complete(&a_turn())
.await
.expect("a tool-call turn is a success");
let calls = response
.tool_calls
.expect("the primary's tool calls must survive the chain");
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].function_name, "get_activities");
assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
assert_eq!(
secondary_calls.load(Ordering::SeqCst),
0,
"falling back on a tool-call turn would spend a completion on the common case"
);
}
#[tokio::test]
async fn an_invalid_request_is_not_rerouted() {
let (primary, primary_calls) = Fake::scripted("primary", Scripted::InvalidRequest);
let (secondary, secondary_calls) =
Fake::scripted("secondary", Scripted::Says("should not be used"));
let error = strict_chain(primary, secondary)
.complete(&a_turn())
.await
.expect_err("an invalid-request error does not fall through");
assert_eq!(error.kind, ErrorKind::InvalidRequest);
assert!(
error.message.contains("malformed"),
"the caller must see the primary's own diagnostic, got: {error}"
);
assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
assert_eq!(secondary_calls.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn a_primary_that_answers_is_the_answer() {
let (primary, primary_calls) = Fake::scripted("primary", Scripted::Says("Rest today."));
let (secondary, secondary_calls) =
Fake::scripted("secondary", Scripted::Says("should not be used"));
let response = strict_chain(primary, secondary)
.complete(&a_turn())
.await
.expect("a healthy primary answers");
assert_eq!(response.content, "Rest today.");
assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
assert_eq!(secondary_calls.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn strict_fall_through_follows_is_provider_fault_for_every_kind() {
let kinds = [
ErrorKind::Internal,
ErrorKind::ExternalService,
ErrorKind::Timeout,
ErrorKind::BinaryNotFound,
ErrorKind::AuthFailure,
ErrorKind::Config,
ErrorKind::Guardrail,
ErrorKind::ContextLength,
ErrorKind::ModelUnavailable,
ErrorKind::RateLimit,
ErrorKind::InvalidRequest,
];
for kind in kinds {
let primary = Box::new(Erring { kind });
let (secondary, secondary_calls) = Fake::scripted("secondary", Scripted::Says("answered"));
let outcome = strict_chain(primary, secondary).complete(&a_turn()).await;
if kind.is_provider_fault() {
assert_eq!(
outcome
.expect("a provider fault reaches the secondary")
.content,
"answered",
"{kind:?}"
);
assert_eq!(secondary_calls.load(Ordering::SeqCst), 1, "{kind:?}");
} else {
let err = outcome.expect_err("a request-side error propagates");
assert_eq!(err.kind, kind);
assert_eq!(
secondary_calls.load(Ordering::SeqCst),
0,
"{kind:?} must not spend a second tier"
);
}
}
}
struct Erring {
kind: ErrorKind,
}
#[async_trait]
impl LlmProvider for Erring {
fn name(&self) -> &'static str {
"erring"
}
fn display_name(&self) -> &str {
"erring"
}
fn capabilities(&self) -> LlmCapabilities {
LlmCapabilities::text_only()
}
fn default_model(&self) -> &str {
"test-model"
}
fn available_models(&self) -> &[String] {
&[]
}
async fn complete(&self, _request: &ChatRequest) -> Result<ChatResponse, RunnerError> {
Err(RunnerError {
kind: self.kind,
message: format!("{:?}", self.kind),
})
}
async fn complete_stream(&self, _request: &ChatRequest) -> Result<ChatStream, RunnerError> {
Err(RunnerError {
kind: self.kind,
message: format!("{:?}", self.kind),
})
}
async fn health_check(&self) -> Result<bool, RunnerError> {
Ok(true)
}
}
#[tokio::test]
async fn permissive_returns_an_empty_completion_as_success() {
let (primary, _) = Fake::scripted("primary", Scripted::Empty);
let (secondary, secondary_calls) = Fake::scripted("secondary", Scripted::Says("never asked"));
let chain = FallbackProvider::new(vec![primary, secondary])
.expect("two tiers")
.with_fallthrough(ResponsePolicy::permissive());
let response = chain.complete(&a_turn()).await.expect("Ok, however empty");
assert!(response.content.is_empty());
assert_eq!(secondary_calls.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn permissive_falls_through_on_every_error() {
let (primary, _) = Fake::scripted("primary", Scripted::InvalidRequest);
let (secondary, secondary_calls) = Fake::scripted("secondary", Scripted::Says("rerouted"));
let chain = FallbackProvider::new(vec![primary, secondary])
.expect("two tiers")
.with_fallthrough(ResponsePolicy::permissive());
let response = chain
.complete(&a_turn())
.await
.expect("the secondary answers");
assert_eq!(response.content, "rerouted");
assert_eq!(secondary_calls.load(Ordering::SeqCst), 1);
}
#[test]
fn the_default_policy_is_permissive() {
assert_eq!(ResponsePolicy::default(), ResponsePolicy::permissive());
assert_ne!(ResponsePolicy::default(), ResponsePolicy::strict());
}
#[tokio::test]
async fn an_empty_completion_from_the_last_tier_is_returned() {
let (primary, _) = Fake::scripted("primary", Scripted::ProviderFault);
let (secondary, secondary_calls) = Fake::scripted("secondary", Scripted::Empty);
let response = strict_chain(primary, secondary)
.complete(&a_turn())
.await
.expect("the last tier's empty answer is still an answer");
assert!(response.content.is_empty());
assert_eq!(secondary_calls.load(Ordering::SeqCst), 1);
}
#[test]
fn a_bare_empty_turn_is_empty() {
assert!(
is_empty_completion(&reply("", None)),
"an empty turn with no tool calls is the failure this exists to catch"
);
}
#[test]
fn whitespace_only_content_is_empty() {
assert!(is_empty_completion(&reply(" \n\t ", None)));
}
#[test]
fn an_empty_tool_call_vec_is_still_empty() {
assert!(is_empty_completion(&reply("", Some(vec![]))));
}
#[test]
fn empty_content_with_tool_calls_is_not_empty() {
assert!(
!is_empty_completion(&reply("", Some(vec![a_tool_call()]))),
"a tool-call turn is an ordinary mid-loop turn, not a lost one"
);
}
#[test]
fn whitespace_content_with_tool_calls_is_not_empty() {
assert!(!is_empty_completion(&reply(
" ",
Some(vec![a_tool_call()])
)));
}
#[test]
fn a_real_reply_is_not_empty() {
assert!(!is_empty_completion(&reply(
"Tu as couru 42 km ce mois-ci.",
None
)));
}
#[tokio::test]
async fn a_tier_that_refused_on_quota_is_passed_over_until_its_cooldown_ends() {
let (spent, spent_calls) = Fake::scripted("claude-code", Scripted::QuotaRefusal);
let (second, second_calls) = Fake::scripted("claude-code#2", Scripted::Says("pong"));
let chain = FallbackProvider::new(vec![spent, second])
.unwrap() .with_fallthrough(ResponsePolicy::strict())
.with_quota_cooldown(Duration::from_millis(200));
let request = ChatRequest::new(vec![ChatMessage::user("ping")]);
let first = chain.complete(&request).await.unwrap(); assert_eq!(first.content, "pong");
assert_eq!(spent_calls.load(Ordering::SeqCst), 1);
assert_eq!(second_calls.load(Ordering::SeqCst), 1);
for _ in 0..2 {
chain.complete(&request).await.unwrap(); }
assert_eq!(
spent_calls.load(Ordering::SeqCst),
1,
"a refused account is not asked again on every turn"
);
assert_eq!(second_calls.load(Ordering::SeqCst), 3);
sleep(Duration::from_millis(250)).await;
chain.complete(&request).await.unwrap(); assert_eq!(spent_calls.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn the_last_tier_is_asked_even_in_cooldown() {
let (only, calls) = Fake::scripted("claude-code", Scripted::QuotaRefusal);
let chain = FallbackProvider::new(vec![only])
.unwrap() .with_fallthrough(ResponsePolicy::strict())
.with_quota_cooldown(Duration::from_hours(1));
let request = ChatRequest::new(vec![ChatMessage::user("ping")]);
for _ in 0..2 {
let err = chain.complete(&request).await.unwrap_err();
assert_eq!(err.kind, ErrorKind::RateLimit);
}
assert_eq!(
calls.load(Ordering::SeqCst),
2,
"nowhere else to go: it is asked"
);
}
#[tokio::test]
async fn a_provider_fault_is_not_a_cooldown() {
let (flaky, flaky_calls) = Fake::scripted("claude-code", Scripted::ProviderFault);
let (second, _) = Fake::scripted("gemini", Scripted::Says("pong"));
let chain = FallbackProvider::new(vec![flaky, second])
.unwrap() .with_fallthrough(ResponsePolicy::strict())
.with_quota_cooldown(Duration::from_hours(1));
let request = ChatRequest::new(vec![ChatMessage::user("ping")]);
for _ in 0..3 {
chain.complete(&request).await.unwrap(); }
assert_eq!(
flaky_calls.load(Ordering::SeqCst),
3,
"a fault may clear on the next turn; only a quota refusal is held off"
);
}
#[tokio::test]
async fn chains_sharing_a_book_pass_over_a_tier_the_other_saw_refuse() {
let book = QuotaCooldowns::new(Duration::from_millis(200));
let (spent, spent_calls) = Fake::scripted("anthropic#2", Scripted::QuotaRefusal);
let (second, _) = Fake::scripted("claude-code", Scripted::Says("pong"));
let full = FallbackProvider::new(vec![spent, second])
.unwrap() .with_fallthrough(ResponsePolicy::strict())
.with_cooldowns(book.clone());
let (spent_again, spent_again_calls) = Fake::scripted("anthropic#2", Scripted::QuotaRefusal);
let (third, third_calls) = Fake::scripted("gemini", Scripted::Says("pong"));
let tail = FallbackProvider::new(vec![spent_again, third])
.unwrap() .with_fallthrough(ResponsePolicy::strict())
.with_cooldowns(book.clone());
let request = ChatRequest::new(vec![ChatMessage::user("ping")]);
full.complete(&request).await.unwrap(); assert_eq!(spent_calls.load(Ordering::SeqCst), 1);
assert!(book.is_cooling("anthropic#2"), "the refusal is in the book");
tail.complete(&request).await.unwrap(); assert_eq!(
spent_again_calls.load(Ordering::SeqCst),
0,
"the other chain does not ask an account the book says is cooling"
);
assert_eq!(third_calls.load(Ordering::SeqCst), 1);
sleep(Duration::from_millis(250)).await;
assert!(
!book.is_cooling("anthropic#2"),
"the window expires on its own"
);
tail.complete(&request).await.unwrap(); assert_eq!(
spent_again_calls.load(Ordering::SeqCst),
1,
"after the window the account is asked again"
);
}
#[tokio::test]
async fn chains_without_a_shared_book_keep_their_own_cooldowns() {
let (spent, spent_calls) = Fake::scripted("anthropic#2", Scripted::QuotaRefusal);
let (second, _) = Fake::scripted("claude-code", Scripted::Says("pong"));
let one = FallbackProvider::new(vec![spent, second])
.unwrap() .with_fallthrough(ResponsePolicy::strict());
let (spent_again, spent_again_calls) = Fake::scripted("anthropic#2", Scripted::QuotaRefusal);
let (third, _) = Fake::scripted("gemini", Scripted::Says("pong"));
let other = FallbackProvider::new(vec![spent_again, third])
.unwrap() .with_fallthrough(ResponsePolicy::strict());
let request = ChatRequest::new(vec![ChatMessage::user("ping")]);
one.complete(&request).await.unwrap(); assert!(one.cooldowns().is_cooling("anthropic#2"));
assert!(!other.cooldowns().is_cooling("anthropic#2"));
other.complete(&request).await.unwrap(); assert_eq!(spent_calls.load(Ordering::SeqCst), 1);
assert_eq!(
spent_again_calls.load(Ordering::SeqCst),
1,
"a private book does not hear about another chain's refusal"
);
}
#[tokio::test]
async fn a_refusal_recorded_outside_the_chain_is_honoured_by_it() {
let book = QuotaCooldowns::new(Duration::from_hours(1));
let (head, head_calls) = Fake::scripted("anthropic", Scripted::Says("head"));
let (second, _) = Fake::scripted("claude-code", Scripted::Says("second"));
let chain = FallbackProvider::new(vec![head, second])
.unwrap() .with_fallthrough(ResponsePolicy::strict())
.with_cooldowns(book.clone());
book.start("anthropic");
let reply = chain
.complete(&ChatRequest::new(vec![ChatMessage::user("ping")]))
.await
.unwrap(); assert_eq!(reply.content, "second");
assert_eq!(head_calls.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn a_chains_quota_cooldown_overrides_the_books_duration() {
let book = QuotaCooldowns::new(Duration::from_hours(1));
let (spent, spent_calls) = Fake::scripted("anthropic#2", Scripted::QuotaRefusal);
let (second, _) = Fake::scripted("claude-code", Scripted::Says("pong"));
let chain = FallbackProvider::new(vec![spent, second])
.unwrap() .with_fallthrough(ResponsePolicy::strict())
.with_cooldowns(book.clone())
.with_quota_cooldown(Duration::from_millis(100));
let request = ChatRequest::new(vec![ChatMessage::user("ping")]);
chain.complete(&request).await.unwrap(); let left = book.remaining("anthropic#2").unwrap(); assert!(left <= Duration::from_millis(100), "{left:?}");
sleep(Duration::from_millis(150)).await;
chain.complete(&request).await.unwrap(); assert_eq!(spent_calls.load(Ordering::SeqCst), 2);
}