use futures::StreamExt;
use llmshim::{
client::ShimClient,
policy::{
AttemptAccounting, AttemptEvent, AttemptIdentity, AttemptKind, AttemptOutcome,
AttemptPolicy, AttemptPolicyError, AttemptPolicyErrorKind, AttemptPolicyFuture,
AttemptPolicyRefusal, AttemptPolicyRefusalKind, AttemptUsageObservation,
DispatchPolicyContext, PreparedAttempt,
},
provider::{Provider, ProviderRequest},
providers::{
anthropic::Anthropic,
chatgpt::{ChatGpt, ChatGptAuth},
gemini::Gemini,
openai::OpenAi,
openai_compat::OpenAiCompatible,
openrouter::OpenRouter,
xai::Xai,
},
streaming::StreamRetentionLimits,
};
use serde_json::{json, Value};
use std::sync::{
atomic::{AtomicBool, Ordering},
Arc, Mutex,
};
struct InvalidHeaderProvider {
endpoint: String,
}
struct CancellationPolicy {
acquired: AtomicBool,
usage_started: tokio::sync::Notify,
abandoned: AtomicBool,
}
impl CancellationPolicy {
fn new() -> Self {
Self {
acquired: AtomicBool::new(false),
usage_started: tokio::sync::Notify::new(),
abandoned: AtomicBool::new(false),
}
}
}
impl AttemptPolicy for CancellationPolicy {
fn acquire<'a>(
&'a self,
_attempt: &'a PreparedAttempt<'a>,
) -> AttemptPolicyFuture<'a, Result<(), AttemptPolicyRefusal>> {
Box::pin(async move {
self.acquired.store(true, Ordering::SeqCst);
Ok(())
})
}
fn observe<'a>(
&'a self,
_attempt: &'a AttemptIdentity,
event: AttemptEvent<'a>,
) -> AttemptPolicyFuture<'a, Result<(), AttemptPolicyError>> {
Box::pin(async move {
if matches!(event, AttemptEvent::Usage { .. }) {
self.usage_started.notify_one();
futures::future::pending::<()>().await;
}
Ok(())
})
}
fn observe_abandoned(
&self,
_attempt: &AttemptIdentity,
_outcome: AttemptOutcome,
) -> Result<(), AttemptPolicyError> {
self.abandoned.store(true, Ordering::SeqCst);
Ok(())
}
}
impl Provider for InvalidHeaderProvider {
fn name(&self) -> &str {
"invalid-header"
}
fn transform_request(
&self,
model: &str,
_request: &Value,
) -> llmshim::error::Result<ProviderRequest> {
Ok(ProviderRequest {
url: self.endpoint.clone(),
headers: vec![("x-invalid".into(), "line one\nline two".into())],
body: json!({"model": model}),
})
}
fn transform_response(&self, _model: &str, response: Value) -> llmshim::error::Result<Value> {
Ok(response)
}
fn transform_stream_chunk(
&self,
_model: &str,
_chunk: &str,
) -> llmshim::error::Result<Option<String>> {
Ok(None)
}
}
#[derive(Clone, Debug)]
enum RecordedEvent {
Acquired {
id: uuid::Uuid,
kind: AttemptKind,
provider_name: String,
resolved_model: String,
native_model: String,
wire: llmshim::reasoning::WireFormat,
account_fingerprint: Option<String>,
endpoint: String,
method: String,
native_body: Value,
},
ResponseHeaders {
id: uuid::Uuid,
status: u16,
},
Usage {
id: uuid::Uuid,
usage: Value,
terminal: bool,
counters_complete: bool,
explicit_zero: bool,
},
Finished {
id: uuid::Uuid,
outcome: AttemptOutcome,
},
Abandoned {
id: uuid::Uuid,
outcome: AttemptOutcome,
},
}
impl RecordedEvent {
fn id(&self) -> uuid::Uuid {
match self {
Self::Acquired { id, .. }
| Self::ResponseHeaders { id, .. }
| Self::Usage { id, .. }
| Self::Finished { id, .. }
| Self::Abandoned { id, .. } => *id,
}
}
}
struct RecordingPolicy {
events: Mutex<Vec<RecordedEvent>>,
refuse_acquire_number: Option<usize>,
refusal_kind: AttemptPolicyRefusalKind,
failing_event: Option<FailingEvent>,
fail_abandonment: bool,
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum FailingEvent {
Headers,
Usage,
Finished,
}
impl Default for RecordingPolicy {
fn default() -> Self {
Self {
events: Mutex::new(Vec::new()),
refuse_acquire_number: None,
refusal_kind: AttemptPolicyRefusalKind::Budget,
failing_event: None,
fail_abandonment: false,
}
}
}
impl RecordingPolicy {
fn refusing(acquire_number: usize) -> Self {
Self {
events: Mutex::new(Vec::new()),
refuse_acquire_number: Some(acquire_number),
refusal_kind: AttemptPolicyRefusalKind::Budget,
failing_event: None,
fail_abandonment: false,
}
}
fn refusing_with_kind(acquire_number: usize, refusal_kind: AttemptPolicyRefusalKind) -> Self {
Self {
events: Mutex::new(Vec::new()),
refuse_acquire_number: Some(acquire_number),
refusal_kind,
failing_event: None,
fail_abandonment: false,
}
}
fn failing(failing_event: FailingEvent) -> Self {
Self {
events: Mutex::new(Vec::new()),
refuse_acquire_number: None,
refusal_kind: AttemptPolicyRefusalKind::Budget,
failing_event: Some(failing_event),
fail_abandonment: false,
}
}
fn failing_abandonment() -> Self {
Self {
events: Mutex::new(Vec::new()),
refuse_acquire_number: None,
refusal_kind: AttemptPolicyRefusalKind::Budget,
failing_event: None,
fail_abandonment: true,
}
}
fn events(&self) -> Vec<RecordedEvent> {
self.events.lock().unwrap().clone()
}
}
impl AttemptPolicy for RecordingPolicy {
fn acquire<'a>(
&'a self,
attempt: &'a PreparedAttempt<'a>,
) -> AttemptPolicyFuture<'a, Result<(), AttemptPolicyRefusal>> {
Box::pin(async move {
let identity = attempt.identity();
let mut events = self.events.lock().unwrap();
let acquire_number = events
.iter()
.filter(|event| matches!(event, RecordedEvent::Acquired { .. }))
.count()
+ 1;
events.push(RecordedEvent::Acquired {
id: identity.id(),
kind: identity.kind(),
provider_name: identity.provider_name().to_owned(),
resolved_model: identity.resolved_model().to_owned(),
native_model: identity.native_model().to_owned(),
wire: identity.wire(),
account_fingerprint: identity.account_fingerprint().map(str::to_owned),
endpoint: identity.endpoint().to_owned(),
method: attempt.method().to_owned(),
native_body: attempt.native_body().clone(),
});
if self.refuse_acquire_number == Some(acquire_number) {
return Err(AttemptPolicyRefusal::new(self.refusal_kind, None));
}
Ok(())
})
}
fn observe<'a>(
&'a self,
attempt: &'a AttemptIdentity,
event: AttemptEvent<'a>,
) -> AttemptPolicyFuture<'a, Result<(), AttemptPolicyError>> {
Box::pin(async move {
let (recorded, event_kind) = match event {
AttemptEvent::ResponseHeaders { status } => (
RecordedEvent::ResponseHeaders {
id: attempt.id(),
status,
},
FailingEvent::Headers,
),
AttemptEvent::Usage { usage } => (
RecordedEvent::Usage {
id: attempt.id(),
usage: usage.clone(),
terminal: false,
counters_complete: false,
explicit_zero: false,
},
FailingEvent::Usage,
),
AttemptEvent::Finished(outcome) => (
RecordedEvent::Finished {
id: attempt.id(),
outcome,
},
FailingEvent::Finished,
),
};
self.events.lock().unwrap().push(recorded);
if self.failing_event == Some(event_kind) {
return Err(AttemptPolicyError::new(
AttemptPolicyErrorKind::CoordinatorUnavailable,
));
}
Ok(())
})
}
fn observe_usage<'a>(
&'a self,
attempt: &'a AttemptIdentity,
observation: AttemptUsageObservation<'a>,
) -> AttemptPolicyFuture<'a, Result<(), AttemptPolicyError>> {
Box::pin(async move {
self.events.lock().unwrap().push(RecordedEvent::Usage {
id: attempt.id(),
usage: observation.usage().clone(),
terminal: observation.terminal(),
counters_complete: observation.counters_complete(),
explicit_zero: observation.explicit_zero(),
});
if self.failing_event == Some(FailingEvent::Usage) {
return Err(AttemptPolicyError::new(
AttemptPolicyErrorKind::CoordinatorUnavailable,
));
}
Ok(())
})
}
fn observe_abandoned(
&self,
attempt: &AttemptIdentity,
outcome: AttemptOutcome,
) -> Result<(), AttemptPolicyError> {
self.events.lock().unwrap().push(RecordedEvent::Abandoned {
id: attempt.id(),
outcome,
});
if self.fail_abandonment {
Err(AttemptPolicyError::new(
AttemptPolicyErrorKind::CoordinatorUnavailable,
))
} else {
Ok(())
}
}
}
fn context(policy: Arc<RecordingPolicy>) -> DispatchPolicyContext {
DispatchPolicyContext::new(policy)
}
fn request(model: &str) -> Value {
json!({
"model": model,
"messages": [{"role": "user", "content": "answer"}],
})
}
fn structured_request(model: &str) -> Value {
json!({
"model": model,
"messages": [{"role": "user", "content": "answer"}],
"response_format": {
"type": "json_schema",
"json_schema": {
"name": "answer",
"schema": {"type": "integer", "minimum": 3}
}
},
"x-shim": {"structured_output": "prompt"}
})
}
fn response(content: Value) -> String {
json!({
"id": "response",
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": content},
"finish_reason": "stop"
}],
"usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}
})
.to_string()
}
fn successful_stream(content: &str) -> String {
format!(
"data: {}\n\ndata: {}\n\ndata: {}\n\ndata: [DONE]\n\n",
json!({"id":"response","choices":[{"index":0,"delta":{"content":content},"finish_reason":null}]}),
json!({"id":"response","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}),
json!({"choices":[],"usage":{"prompt_tokens":7,"completion_tokens":3,"total_tokens":10}}),
)
}
async fn record_chat_stream(events: Vec<Value>) -> Vec<RecordedEvent> {
let mut server = mockito::Server::new_async().await;
let mut body = events
.into_iter()
.map(|event| format!("data: {event}\n\n"))
.collect::<String>();
body.push_str("data: [DONE]\n\n");
let upstream = server
.mock("POST", "/chat/completions")
.with_header("content-type", "text/event-stream")
.with_body(body)
.expect(1)
.create_async()
.await;
let provider = OpenRouter::new("test-key".into()).with_base_url(server.url());
let policy = Arc::new(RecordingPolicy::default());
let chunks: Vec<_> = ShimClient::new()
.stream_with_policy(
&provider,
"x-ai/grok-4.7",
&request("openrouter/x-ai/grok-4.7"),
&context(policy.clone()),
)
.await
.unwrap()
.collect()
.await;
assert!(chunks.iter().all(Result::is_ok), "{chunks:?}");
upstream.assert_async().await;
policy.events()
}
async fn record_gemini_stream(events: Vec<Value>) -> Vec<RecordedEvent> {
let mut server = mockito::Server::new_async().await;
let body = events
.into_iter()
.map(|event| format!("data: {event}\n\n"))
.collect::<String>();
let upstream = server
.mock("POST", mockito::Matcher::Regex("/models/.*".into()))
.with_header("content-type", "text/event-stream")
.with_body(body)
.expect(1)
.create_async()
.await;
let provider = Gemini::new("test-key".into()).with_base_url(server.url());
let policy = Arc::new(RecordingPolicy::default());
let chunks: Vec<_> = ShimClient::new()
.stream_with_policy(
&provider,
"gemini-3.8-flash",
&request("gemini/gemini-3.8-flash"),
&context(policy.clone()),
)
.await
.unwrap()
.collect()
.await;
assert!(chunks.iter().all(Result::is_ok), "{chunks:?}");
upstream.assert_async().await;
policy.events()
}
#[tokio::test]
async fn retained_state_failure_observes_current_usage_before_stream_failure() {
let mut server = mockito::Server::new_async().await;
let body = format!(
"data: {}\n\ndata: [DONE]\n\n",
json!({
"choices":[{"index":0,"delta":{},"finish_reason":"stop"}],
"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2,"cost":0.25}
})
);
let upstream = server
.mock("POST", "/chat/completions")
.with_header("content-type", "text/event-stream")
.with_body(body)
.expect(1)
.create_async()
.await;
let provider = OpenRouter::new("test-key".into()).with_base_url(server.url());
let policy = Arc::new(RecordingPolicy::default());
let limits =
llmshim::streaming::StreamRetentionLimits::new(16 * 1024, 64, 16 * 1024, 1).unwrap();
let chunks: Vec<_> = ShimClient::new()
.with_stream_retention_limits(limits)
.unwrap()
.stream_with_policy(
&provider,
"x-ai/grok-4.7",
&request("openrouter/x-ai/grok-4.7"),
&context(policy.clone()),
)
.await
.unwrap()
.collect()
.await;
upstream.assert_async().await;
assert_eq!(
chunks.last().unwrap().as_ref().unwrap_err().to_string(),
"stream error: upstream stream retained state exceeds limit"
);
let events = policy.events();
let usage_position = events
.iter()
.position(
|event| matches!(event, RecordedEvent::Usage { usage, .. } if usage["cost"] == 0.25),
)
.expect("provider usage must be observed");
let failure_position = events
.iter()
.position(|event| {
matches!(
event,
RecordedEvent::Finished {
outcome: AttemptOutcome::StreamFailure { .. },
..
}
)
})
.expect("retention failure must finish the attempt as a stream failure");
assert!(usage_position < failure_position);
assert!(!events
.iter()
.any(|event| { matches!(event, RecordedEvent::Usage { terminal: true, .. }) }));
}
#[tokio::test]
async fn anthropic_retention_failure_observes_merged_current_provider_bill() {
let mut server = mockito::Server::new_async().await;
let body = format!(
"data: {}\n\ndata: {}\n\ndata: [DONE]\n\n",
json!({"type":"message_start","message":{"usage":{"input_tokens":1,"cost":0.25}}}),
json!({"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":2,"cost":0.75}})
);
let upstream = server
.mock("POST", "/messages")
.with_header("content-type", "text/event-stream")
.with_body(body)
.expect(1)
.create_async()
.await;
let provider = Anthropic::new("test-key".into()).with_base_url(server.url());
let policy = Arc::new(RecordingPolicy::default());
let limits =
llmshim::streaming::StreamRetentionLimits::new(16 * 1024, 64, 16 * 1024, 4).unwrap();
let chunks: Vec<_> = ShimClient::new()
.with_stream_retention_limits(limits)
.unwrap()
.stream_with_policy(
&provider,
"claude-sonnet-5",
&request("anthropic/claude-sonnet-5"),
&context(policy.clone()),
)
.await
.unwrap()
.collect()
.await;
upstream.assert_async().await;
assert_eq!(
chunks.last().unwrap().as_ref().unwrap_err().to_string(),
"stream error: upstream stream retained state exceeds limit"
);
let events = policy.events();
let usage: Vec<_> = events
.iter()
.filter_map(|event| match event {
RecordedEvent::Usage {
usage, terminal, ..
} => Some((usage, terminal)),
_ => None,
})
.collect();
assert_eq!(usage.len(), 2);
assert_eq!(usage[0].0["cost"], 0.25);
assert_eq!(usage[1].0["cost"], 0.75);
assert_eq!(usage[1].0["cost_usd"], 0.75);
assert_eq!(usage[1].0["cost_source"], "provider");
assert_eq!(usage[1].0["prompt_tokens"], 1);
assert_eq!(usage[1].0["completion_tokens"], 2);
assert!(!*usage[1].1);
let failure_position = events
.iter()
.position(|event| {
matches!(
event,
RecordedEvent::Finished {
outcome: AttemptOutcome::StreamFailure { .. },
..
}
)
})
.unwrap();
let second_usage_position = events
.iter()
.rposition(|event| matches!(event, RecordedEvent::Usage { .. }))
.unwrap();
assert!(second_usage_position < failure_position);
}
fn terminal_usage_count(events: &[RecordedEvent]) -> usize {
events
.iter()
.filter(|event| matches!(event, RecordedEvent::Usage { terminal: true, .. }))
.count()
}
#[tokio::test]
async fn a_refused_retry_never_reaches_the_network() {
let mut server = mockito::Server::new_async().await;
let upstream = server
.mock("POST", "/chat/completions")
.with_status(500)
.with_body("retryable")
.expect(1)
.create_async()
.await;
let provider = OpenAiCompatible::new("local", server.url(), None);
let policy = Arc::new(RecordingPolicy::refusing(2));
let error = ShimClient::new()
.completion_with_policy(
&provider,
"test",
&request("local/test"),
&context(policy.clone()),
)
.await
.unwrap_err();
assert!(error.to_string().contains("attempt budget exhausted"));
upstream.assert_async().await;
let events = policy.events();
assert_eq!(
events
.iter()
.filter(|event| matches!(event, RecordedEvent::Acquired { .. }))
.count(),
2
);
assert!(events.iter().all(|event| event.id() != uuid::Uuid::nil()));
assert!(events
.iter()
.any(|event| matches!(event, RecordedEvent::ResponseHeaders { status: 500, .. })));
assert!(events.iter().any(|event| matches!(
event,
RecordedEvent::Finished {
outcome: AttemptOutcome::HttpFailure {
status: 500,
accounting: AttemptAccounting::Unknown
},
..
}
)));
}
#[tokio::test]
async fn request_build_failure_happens_before_attempt_acquisition() {
let mut server = mockito::Server::new_async().await;
let upstream = server.mock("POST", "/never").expect(0).create_async().await;
let provider = InvalidHeaderProvider {
endpoint: format!("{}/never", server.url()),
};
let policy = Arc::new(RecordingPolicy::default());
let result = ShimClient::new()
.completion_with_policy(
&provider,
"test",
&request("invalid-header/test"),
&context(policy.clone()),
)
.await;
assert!(result.is_err());
assert!(policy.events().is_empty());
upstream.assert_async().await;
}
#[tokio::test]
async fn structured_refusals_keep_typed_provenance_on_unary_and_borrowed_streams() {
let mut server = mockito::Server::new_async().await;
let upstream = server
.mock("POST", "/chat/completions")
.expect(0)
.create_async()
.await;
let provider = OpenAiCompatible::new("local", server.url(), None);
for stream in [false, true] {
let policy = Arc::new(RecordingPolicy::refusing(1));
let client = ShimClient::new();
let result = if stream {
client
.stream_with_policy(
&provider,
"test",
&structured_request("local/test"),
&context(policy),
)
.await
.map(|_| Value::Null)
} else {
client
.completion_with_policy(
&provider,
"test",
&structured_request("local/test"),
&context(policy),
)
.await
};
let error = result.unwrap_err().to_string();
assert!(error.contains("attempt budget exhausted"));
assert!(!error.contains("output contract"));
}
upstream.assert_async().await;
}
#[tokio::test]
async fn owned_stream_repair_refusal_is_not_rewritten_as_provider_failure() {
let mut server = mockito::Server::new_async().await;
let first_attempt = server
.mock("POST", "/chat/completions")
.with_header("content-type", "text/event-stream")
.with_body(successful_stream("1"))
.expect(1)
.create_async()
.await;
let provider: Arc<dyn Provider> = Arc::new(OpenAiCompatible::new("local", server.url(), None));
let policy = Arc::new(RecordingPolicy::refusing(2));
let stream = ShimClient::new()
.stream_owned_with_policy(
provider,
"test",
&structured_request("local/test"),
&context(policy),
)
.await
.unwrap();
let chunks: Vec<_> = stream.collect().await;
assert_eq!(chunks.len(), 1);
let error = chunks[0].as_ref().unwrap_err().to_string();
assert!(error.contains("attempt budget exhausted"));
assert!(!error.contains("output contract"));
first_attempt.assert_async().await;
}
#[tokio::test]
async fn a_policy_coordinator_refusal_does_not_poison_provider_health() {
let mut server = mockito::Server::new_async().await;
let upstream = server
.mock("POST", "/chat/completions")
.expect(0)
.create_async()
.await;
let provider = OpenAiCompatible::new("local", server.url(), None);
let breaker = Arc::new(llmshim::breaker::ProviderBreaker::with_config(
llmshim::breaker::BreakerConfig {
window: std::time::Duration::from_secs(60),
trip_threshold: 1,
cooldown: std::time::Duration::from_secs(300),
},
));
let policy = Arc::new(RecordingPolicy::refusing_with_kind(
1,
AttemptPolicyRefusalKind::CoordinatorUnavailable,
));
let error = ShimClient::new()
.with_breaker(breaker.clone())
.completion_with_policy(&provider, "test", &request("local/test"), &context(policy))
.await
.unwrap_err();
assert!(error
.to_string()
.contains("attempt policy coordinator unavailable"));
assert!(breaker.admit("local").await);
upstream.assert_async().await;
}
#[tokio::test]
async fn local_repair_exhaustion_does_not_poison_provider_health() {
let mut server = mockito::Server::new_async().await;
let upstream = server
.mock("POST", "/chat/completions")
.with_body(response(json!("1")))
.expect(2)
.create_async()
.await;
let provider = OpenAiCompatible::new("local", server.url(), None);
let breaker = Arc::new(llmshim::breaker::ProviderBreaker::with_config(
llmshim::breaker::BreakerConfig {
window: std::time::Duration::from_secs(60),
trip_threshold: 1,
cooldown: std::time::Duration::from_secs(300),
},
));
let policy = Arc::new(RecordingPolicy::default());
let result = ShimClient::new()
.with_breaker(breaker.clone())
.completion_with_policy(
&provider,
"test",
&structured_request("local/test"),
&context(policy),
)
.await;
assert!(result.is_err());
assert!(breaker.admit("local").await);
upstream.assert_async().await;
}
#[tokio::test]
async fn fallback_acquires_the_actual_provider_for_each_model() {
let mut primary_server = mockito::Server::new_async().await;
let mut secondary_server = mockito::Server::new_async().await;
let primary = primary_server
.mock("POST", "/chat/completions")
.with_status(400)
.with_body("unsupported")
.expect(1)
.create_async()
.await;
let secondary = secondary_server
.mock("POST", "/chat/completions")
.with_body(response(json!("ok")))
.expect(1)
.create_async()
.await;
let router = llmshim::router::Router::new()
.register(
"primary",
Box::new(OpenAiCompatible::new("primary", primary_server.url(), None)),
)
.register(
"secondary",
Box::new(OpenAiCompatible::new(
"secondary",
secondary_server.url(),
None,
)),
);
let policy = Arc::new(RecordingPolicy::default());
let fallback = llmshim::FallbackConfig::new(vec![
"primary/first".to_owned(),
"secondary/second".to_owned(),
])
.max_retries(0);
let result = llmshim::completion_with_fallback_and_policy(
&router,
&request("ignored/model"),
&fallback,
None,
&context(policy.clone()),
)
.await
.unwrap();
assert_eq!(result["choices"][0]["message"]["content"], "ok");
primary.assert_async().await;
secondary.assert_async().await;
let acquired_targets: Vec<_> = policy
.events()
.into_iter()
.filter_map(|event| match event {
RecordedEvent::Acquired {
provider_name,
resolved_model,
..
} => Some((provider_name, resolved_model)),
_ => None,
})
.collect();
assert_eq!(
acquired_targets,
vec![
("primary".to_owned(), "first".to_owned()),
("secondary".to_owned(), "second".to_owned()),
]
);
}
#[tokio::test]
async fn a_terminal_policy_refusal_stops_the_fallback_chain() {
let mut primary_server = mockito::Server::new_async().await;
let mut secondary_server = mockito::Server::new_async().await;
let primary = primary_server
.mock("POST", "/chat/completions")
.expect(0)
.create_async()
.await;
let secondary = secondary_server
.mock("POST", "/chat/completions")
.expect(0)
.create_async()
.await;
let router = llmshim::router::Router::new()
.register(
"primary",
Box::new(OpenAiCompatible::new("primary", primary_server.url(), None)),
)
.register(
"secondary",
Box::new(OpenAiCompatible::new(
"secondary",
secondary_server.url(),
None,
)),
);
let policy = Arc::new(RecordingPolicy::refusing(1));
let fallback = llmshim::FallbackConfig::new(vec![
"primary/first".to_owned(),
"secondary/second".to_owned(),
]);
let error = llmshim::completion_with_fallback_and_policy(
&router,
&structured_request("ignored/model"),
&fallback,
None,
&context(policy),
)
.await
.unwrap_err();
assert!(error.to_string().contains("attempt budget exhausted"));
primary.assert_async().await;
secondary.assert_async().await;
}
#[tokio::test]
async fn an_upstream_body_matching_a_policy_message_keeps_upstream_provenance() {
let mut primary_server = mockito::Server::new_async().await;
let mut secondary_server = mockito::Server::new_async().await;
let primary = primary_server
.mock("POST", "/chat/completions")
.with_status(400)
.with_body("attempt budget exhausted")
.expect(1)
.create_async()
.await;
let secondary = secondary_server
.mock("POST", "/chat/completions")
.with_body(response(json!("ok")))
.expect(1)
.create_async()
.await;
let router = llmshim::router::Router::new()
.register(
"primary",
Box::new(OpenAiCompatible::new("primary", primary_server.url(), None)),
)
.register(
"secondary",
Box::new(OpenAiCompatible::new(
"secondary",
secondary_server.url(),
None,
)),
);
let policy = Arc::new(RecordingPolicy::default());
let fallback = llmshim::FallbackConfig::new(vec![
"primary/first".to_owned(),
"secondary/second".to_owned(),
])
.max_retries(0);
let result = llmshim::completion_with_fallback_and_policy(
&router,
&request("ignored/model"),
&fallback,
None,
&context(policy),
)
.await
.unwrap();
assert_eq!(result["choices"][0]["message"]["content"], "ok");
primary.assert_async().await;
secondary.assert_async().await;
}
#[tokio::test]
async fn fallback_skips_same_target_retries_after_local_request_build_failure() {
let mut invalid_server = mockito::Server::new_async().await;
let mut secondary_server = mockito::Server::new_async().await;
let invalid_upstream = invalid_server
.mock("POST", "/never")
.expect(0)
.create_async()
.await;
let secondary_upstream = secondary_server
.mock("POST", "/chat/completions")
.with_body(response(json!("ok")))
.expect(1)
.create_async()
.await;
let router = llmshim::router::Router::new()
.register(
"invalid-header",
Box::new(InvalidHeaderProvider {
endpoint: format!("{}/never", invalid_server.url()),
}),
)
.register(
"secondary",
Box::new(OpenAiCompatible::new(
"secondary",
secondary_server.url(),
None,
)),
);
let policy = Arc::new(RecordingPolicy::default());
let fallback = llmshim::FallbackConfig::new(vec![
"invalid-header/first".to_owned(),
"secondary/second".to_owned(),
])
.max_retries(3);
let result = llmshim::completion_with_fallback_and_policy(
&router,
&request("ignored/model"),
&fallback,
None,
&context(policy.clone()),
)
.await
.unwrap();
assert_eq!(result["choices"][0]["message"]["content"], "ok");
invalid_upstream.assert_async().await;
secondary_upstream.assert_async().await;
let acquired_providers: Vec<_> = policy
.events()
.into_iter()
.filter_map(|event| match event {
RecordedEvent::Acquired { provider_name, .. } => Some(provider_name),
_ => None,
})
.collect();
assert_eq!(acquired_providers, vec!["secondary"]);
}
#[tokio::test]
async fn fallback_moves_directly_after_local_repair_exhaustion() {
let mut primary_server = mockito::Server::new_async().await;
let mut secondary_server = mockito::Server::new_async().await;
let primary_upstream = primary_server
.mock("POST", "/chat/completions")
.with_body(response(json!("1")))
.expect(2)
.create_async()
.await;
let secondary_upstream = secondary_server
.mock("POST", "/chat/completions")
.with_body(response(json!("3")))
.expect(1)
.create_async()
.await;
let router = llmshim::router::Router::new()
.register(
"primary",
Box::new(OpenAiCompatible::new("primary", primary_server.url(), None)),
)
.register(
"secondary",
Box::new(OpenAiCompatible::new(
"secondary",
secondary_server.url(),
None,
)),
);
let policy = Arc::new(RecordingPolicy::default());
let fallback = llmshim::FallbackConfig::new(vec![
"primary/first".to_owned(),
"secondary/second".to_owned(),
])
.max_retries(3);
let result = llmshim::completion_with_fallback_and_policy(
&router,
&structured_request("ignored/model"),
&fallback,
None,
&context(policy.clone()),
)
.await
.unwrap();
assert_eq!(result["choices"][0]["message"]["content"], "3");
primary_upstream.assert_async().await;
secondary_upstream.assert_async().await;
let acquired_providers: Vec<_> = policy
.events()
.into_iter()
.filter_map(|event| match event {
RecordedEvent::Acquired { provider_name, .. } => Some(provider_name),
_ => None,
})
.collect();
assert_eq!(acquired_providers, vec!["primary", "primary", "secondary"]);
}
#[tokio::test]
async fn provider_limit_skips_same_provider_retries_and_models() {
let mut primary_server = mockito::Server::new_async().await;
let mut secondary_server = mockito::Server::new_async().await;
let primary_upstream = primary_server
.mock("POST", "/chat/completions")
.expect(0)
.create_async()
.await;
let secondary_upstream = secondary_server
.mock("POST", "/chat/completions")
.with_body(response(json!("ok")))
.expect(1)
.create_async()
.await;
let router = llmshim::router::Router::new()
.register(
"primary",
Box::new(OpenAiCompatible::new("primary", primary_server.url(), None)),
)
.register(
"secondary",
Box::new(OpenAiCompatible::new(
"secondary",
secondary_server.url(),
None,
)),
);
let policy = Arc::new(RecordingPolicy::refusing_with_kind(
1,
AttemptPolicyRefusalKind::ProviderLimit,
));
let fallback = llmshim::FallbackConfig::new(vec![
"primary/first".to_owned(),
"primary/second".to_owned(),
"secondary/third".to_owned(),
])
.max_retries(3);
let result = llmshim::completion_with_fallback_and_policy(
&router,
&request("ignored/model"),
&fallback,
None,
&context(policy.clone()),
)
.await
.unwrap();
assert_eq!(result["choices"][0]["message"]["content"], "ok");
primary_upstream.assert_async().await;
secondary_upstream.assert_async().await;
let acquired_providers: Vec<_> = policy
.events()
.into_iter()
.filter_map(|event| match event {
RecordedEvent::Acquired { provider_name, .. } => Some(provider_name),
_ => None,
})
.collect();
assert_eq!(acquired_providers, vec!["primary", "secondary"]);
let no_distinct_policy = Arc::new(RecordingPolicy::refusing_with_kind(
1,
AttemptPolicyRefusalKind::ProviderLimit,
));
let no_distinct_fallback = llmshim::FallbackConfig::new(vec![
"primary/first".to_owned(),
"primary/second".to_owned(),
])
.max_retries(3);
let error = llmshim::completion_with_fallback_and_policy(
&router,
&request("ignored/model"),
&no_distinct_fallback,
None,
&context(no_distinct_policy.clone()),
)
.await
.unwrap_err();
assert!(error
.to_string()
.contains("provider attempt limit exceeded"));
assert_eq!(
no_distinct_policy
.events()
.iter()
.filter(|event| matches!(event, RecordedEvent::Acquired { .. }))
.count(),
1
);
primary_upstream.assert_async().await;
}
#[tokio::test]
async fn failed_repair_observes_each_usage_before_returning_the_error() {
let mut server = mockito::Server::new_async().await;
let upstream = server
.mock("POST", "/chat/completions")
.with_body(response(json!("1")))
.expect(2)
.create_async()
.await;
let provider = OpenAiCompatible::new("openai", server.url(), None);
let policy = Arc::new(RecordingPolicy::default());
let result = ShimClient::new()
.completion_with_policy(
&provider,
"gpt-5.6-luna",
&structured_request("openai/gpt-5.6-luna"),
&context(policy.clone()),
)
.await;
assert!(result.is_err());
upstream.assert_async().await;
let events = policy.events();
let usage_events: Vec<_> = events
.iter()
.filter_map(|event| match event {
RecordedEvent::Usage { usage, .. } => Some(usage),
_ => None,
})
.collect();
assert_eq!(usage_events.len(), 2);
assert!(usage_events
.iter()
.all(|usage| usage["cost_usd"].as_f64().is_some_and(|cost| cost > 0.0)));
assert_eq!(
events
.iter()
.filter(|event| matches!(
event,
RecordedEvent::Finished {
outcome: AttemptOutcome::Completed {
accounting: AttemptAccounting::UsageObserved
},
..
}
))
.count(),
2
);
}
#[tokio::test]
async fn bounded_http_error_observes_reported_bill_before_preserving_provider_error() {
let mut server = mockito::Server::new_async().await;
let error_body = json!({
"error": {"message": "provider-specific failure"},
"usage": {"cost": 12_345_678.0}
})
.to_string();
let upstream = server
.mock("POST", "/chat/completions")
.with_status(400)
.with_body(&error_body)
.expect(1)
.create_async()
.await;
let provider = OpenRouter::new("test-key".into()).with_base_url(server.url());
let policy = Arc::new(RecordingPolicy::default());
let error = ShimClient::new()
.completion_with_policy(
&provider,
"x-ai/grok-4.7",
&request("openrouter/x-ai/grok-4.7"),
&context(policy.clone()),
)
.await
.unwrap_err();
assert!(matches!(
error,
llmshim::error::ShimError::ProviderError {
status: 400,
ref body,
..
} if body == &error_body
));
upstream.assert_async().await;
let events = policy.events();
let usage_index = events
.iter()
.position(|event| matches!(event, RecordedEvent::Usage { .. }))
.unwrap();
let finish_index = events
.iter()
.position(|event| {
matches!(
event,
RecordedEvent::Finished {
outcome: AttemptOutcome::HttpFailure {
status: 400,
accounting: AttemptAccounting::UsageObserved
},
..
}
)
})
.unwrap();
assert!(usage_index < finish_index);
let RecordedEvent::Usage { usage, .. } = &events[usage_index] else {
unreachable!()
};
assert_eq!(usage["cost_source"], "provider");
assert_eq!(usage["cost_usd"], 12_345_678.0);
}
#[tokio::test]
async fn retryable_http_error_does_not_retry_after_usage_observer_failure() {
let mut server = mockito::Server::new_async().await;
let upstream = server
.mock("POST", "/chat/completions")
.with_status(500)
.with_body(json!({"error": "retryable", "usage": {"cost": 1.0}}).to_string())
.expect(1)
.create_async()
.await;
let provider = OpenRouter::new("test-key".into()).with_base_url(server.url());
let policy = Arc::new(RecordingPolicy::failing(FailingEvent::Usage));
let error = ShimClient::new()
.completion_with_policy(
&provider,
"x-ai/grok-4.7",
&request("openrouter/x-ai/grok-4.7"),
&context(policy.clone()),
)
.await
.unwrap_err();
assert!(error.to_string().contains("coordinator unavailable"));
upstream.assert_async().await;
assert_eq!(
policy
.events()
.iter()
.filter(|event| matches!(event, RecordedEvent::Acquired { .. }))
.count(),
1
);
}
#[tokio::test]
async fn unary_usage_requires_present_input_and_output_counters_for_terminal_release() {
let mut server = mockito::Server::new_async().await;
let upstream = server
.mock("POST", "/chat/completions")
.with_body(
json!({
"id": "response",
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": "ok"},
"finish_reason": "stop"
}],
"usage": {"prompt_tokens": 7}
})
.to_string(),
)
.expect(1)
.create_async()
.await;
let provider = OpenRouter::new("test-key".into()).with_base_url(server.url());
let policy = Arc::new(RecordingPolicy::default());
ShimClient::new()
.completion_with_policy(
&provider,
"x-ai/grok-4.7",
&request("openrouter/x-ai/grok-4.7"),
&context(policy.clone()),
)
.await
.unwrap();
upstream.assert_async().await;
assert!(policy.events().iter().any(|event| matches!(
event,
RecordedEvent::Usage {
terminal: true,
counters_complete: false,
explicit_zero: false,
..
}
)));
}
#[tokio::test]
async fn unary_explicit_zero_counters_are_complete_but_provider_bill_needs_no_counters() {
for (usage, expected_complete, expected_zero, expected_source) in [
(
json!({"prompt_tokens": 0, "completion_tokens": 0}),
true,
true,
"catalog",
),
(json!({"cost": 0.0}), false, false, "provider"),
] {
let mut server = mockito::Server::new_async().await;
let upstream = server
.mock("POST", "/chat/completions")
.with_body(
json!({
"id": "response",
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": "ok"},
"finish_reason": "stop"
}],
"usage": usage
})
.to_string(),
)
.expect(1)
.create_async()
.await;
let provider = OpenRouter::new("test-key".into()).with_base_url(server.url());
let policy = Arc::new(RecordingPolicy::default());
ShimClient::new()
.completion_with_policy(
&provider,
"x-ai/grok-4.7",
&request("openrouter/x-ai/grok-4.7"),
&context(policy.clone()),
)
.await
.unwrap();
upstream.assert_async().await;
assert!(policy.events().iter().any(|event| matches!(
event,
RecordedEvent::Usage {
usage,
terminal: true,
counters_complete,
explicit_zero,
..
} if *counters_complete == expected_complete
&& *explicit_zero == expected_zero
&& usage["cost_source"] == expected_source
)));
}
}
#[tokio::test]
async fn chat_stream_partial_usage_followed_by_clean_finish_is_not_terminal_accounting() {
let mut server = mockito::Server::new_async().await;
let stream_body = format!(
"data: {}\n\ndata: {}\n\ndata: {}\n\ndata: [DONE]\n\n",
json!({"id":"r","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":null}]}),
json!({"id":"r","choices":[{"index":0,"delta":{},"finish_reason":null}],"usage":{"prompt_tokens":7}}),
json!({"id":"r","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]})
);
let upstream = server
.mock("POST", "/chat/completions")
.with_header("content-type", "text/event-stream")
.with_body(stream_body)
.expect(1)
.create_async()
.await;
let provider = OpenRouter::new("test-key".into()).with_base_url(server.url());
let policy = Arc::new(RecordingPolicy::default());
let chunks: Vec<_> = ShimClient::new()
.stream_with_policy(
&provider,
"x-ai/grok-4.7",
&request("openrouter/x-ai/grok-4.7"),
&context(policy.clone()),
)
.await
.unwrap()
.collect()
.await;
assert!(chunks.iter().all(Result::is_ok));
upstream.assert_async().await;
assert!(policy.events().iter().any(|event| matches!(
event,
RecordedEvent::Usage {
terminal: false,
counters_complete: false,
..
}
)));
assert!(!policy
.events()
.iter()
.any(|event| matches!(event, RecordedEvent::Usage { terminal: true, .. })));
}
#[tokio::test]
async fn chat_terminal_authority_survives_only_without_later_choice_activity() {
let complete_usage = json!({"prompt_tokens": 7, "completion_tokens": 3});
let unsafe_sequences = vec![
vec![
json!({"id":"r","choices":[],"usage":complete_usage}),
json!({"id":"r","choices":[{"index":0,"delta":{"content":"later"},"finish_reason":null}]}),
json!({"id":"r","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}),
],
vec![
json!({"id":"r","choices":[
{"index":0,"delta":{},"finish_reason":"stop"},
{"index":1,"delta":{},"finish_reason":null}
],"usage":{"prompt_tokens":7,"completion_tokens":3}}),
json!({"id":"r","choices":[{"index":1,"delta":{"content":"later"},"finish_reason":null}]}),
json!({"id":"r","choices":[{"index":1,"delta":{},"finish_reason":"stop"}]}),
],
vec![
json!({"id":"r","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":7,"completion_tokens":3}}),
json!({"id":"r","choices":[{"index":1,"delta":{"content":"new"},"finish_reason":null}]}),
json!({"id":"r","choices":[{"index":1,"delta":{},"finish_reason":"stop"}]}),
],
];
for sequence in unsafe_sequences {
let events = record_chat_stream(sequence).await;
assert_eq!(terminal_usage_count(&events), 0, "{events:?}");
assert!(events.iter().any(|event| matches!(
event,
RecordedEvent::Usage {
terminal: false,
counters_complete: true,
..
}
)));
}
let valid = record_chat_stream(vec![
json!({"id":"r","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":null}]}),
json!({"id":"r","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}),
json!({"id":"r","choices":[],"usage":{"prompt_tokens":7,"completion_tokens":3}}),
])
.await;
assert_eq!(terminal_usage_count(&valid), 1, "{valid:?}");
}
#[tokio::test]
async fn partial_provider_cost_precedes_lower_terminal_catalog_candidate() {
let events = record_chat_stream(vec![
json!({
"id":"r",
"choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":null}],
"usage":{"cost":1.0}
}),
json!({"id":"r","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}),
json!({
"id":"r",
"choices":[],
"usage":{"prompt_tokens":7,"completion_tokens":3}
}),
])
.await;
assert!(events.iter().any(|event| matches!(
event,
RecordedEvent::Usage {
usage,
terminal: false,
..
} if usage["cost_source"] == "provider" && usage["cost_usd"] == 1.0
)));
assert!(events.iter().any(|event| matches!(
event,
RecordedEvent::Usage {
usage,
terminal: true,
..
} if usage["cost_source"] == "catalog"
&& usage["cost_usd"].as_f64().is_some_and(|cost| cost < 1.0)
)));
}
#[tokio::test]
async fn terminal_provider_candidate_survives_catalog_trailer_until_later_activity() {
let corrected_chat = record_chat_stream(vec![
json!({"choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"cost":0.25}}),
json!({"choices":[],"usage":{"prompt_tokens":1_000_000,"completion_tokens":0}}),
])
.await;
assert!(corrected_chat.iter().any(|event| matches!(
event,
RecordedEvent::Usage {
usage,
terminal: true,
..
} if usage["cost_source"] == "provider" && usage["cost_usd"] == 0.25
)));
let invalidated_chat = record_chat_stream(vec![
json!({"choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"cost":0.25}}),
json!({"choices":[{"index":1,"delta":{"content":"later"},"finish_reason":null}]}),
json!({"choices":[{"index":1,"delta":{},"finish_reason":"stop"}]}),
json!({"choices":[],"usage":{"prompt_tokens":1_000_000,"completion_tokens":0}}),
])
.await;
assert!(invalidated_chat.iter().any(|event| matches!(
event,
RecordedEvent::Usage {
usage,
terminal: true,
..
} if usage["cost_source"] == "catalog"
)));
let corrected_gemini = record_gemini_stream(vec![
json!({"candidates":[{"index":0,"finishReason":"STOP","content":{"parts":[]}}],"usageMetadata":{"cost":0.25}}),
json!({"candidates":[],"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":3}}),
])
.await;
assert!(corrected_gemini.iter().any(|event| matches!(
event,
RecordedEvent::Usage {
usage,
terminal: true,
..
} if usage["cost_source"] == "provider" && usage["cost_usd"] == 0.25
)));
let invalidated_gemini = record_gemini_stream(vec![
json!({"candidates":[{"index":0,"finishReason":"STOP","content":{"parts":[]}}],"usageMetadata":{"cost":0.25}}),
json!({"candidates":[{"index":1,"content":{"parts":[{"text":"later"}]}}]}),
json!({"candidates":[{"index":1,"finishReason":"STOP","content":{"parts":[]}}]}),
json!({"candidates":[],"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":3}}),
])
.await;
assert!(invalidated_gemini.iter().any(|event| matches!(
event,
RecordedEvent::Usage {
usage,
terminal: true,
..
} if usage["cost_source"] == "catalog"
)));
}
#[tokio::test]
async fn public_terminal_usage_keeps_explicit_terminal_provider_correction() {
let mut server = mockito::Server::new_async().await;
let body = format!(
"data: {}\n\ndata: {}\n\ndata: [DONE]\n\n",
json!({"choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"cost":0.25}}),
json!({"choices":[],"usage":{"prompt_tokens":1_000_000,"completion_tokens":0}})
);
let upstream = server
.mock("POST", "/chat/completions")
.with_header("content-type", "text/event-stream")
.with_body(body)
.expect(1)
.create_async()
.await;
let provider = OpenRouter::new("test-key".into()).with_base_url(server.url());
let chunks: Vec<_> = ShimClient::new()
.stream(
&provider,
"x-ai/grok-4.7",
&request("openrouter/x-ai/grok-4.7"),
)
.await
.unwrap()
.collect::<Vec<_>>()
.await
.into_iter()
.collect::<Result<Vec<_>, _>>()
.unwrap();
upstream.assert_async().await;
assert!(chunks
.iter()
.filter_map(|chunk| serde_json::from_str::<Value>(chunk).ok())
.any(|chunk| {
chunk["usage"]["cost_source"] == "provider" && chunk["usage"]["cost_usd"] == 0.25
}));
}
#[tokio::test]
async fn public_terminal_usage_preserves_partial_provider_floor_separately() {
let mut server = mockito::Server::new_async().await;
let body = format!(
"data: {}\n\ndata: {}\n\ndata: {}\n\ndata: [DONE]\n\n",
json!({"choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":null}],"usage":{"cost":1.0}}),
json!({"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}),
json!({"choices":[],"usage":{"prompt_tokens":7,"completion_tokens":3}})
);
let upstream = server
.mock("POST", "/chat/completions")
.with_header("content-type", "text/event-stream")
.with_body(body)
.expect(1)
.create_async()
.await;
let provider = OpenRouter::new("test-key".into()).with_base_url(server.url());
let chunks: Vec<_> = ShimClient::new()
.stream(
&provider,
"x-ai/grok-4.7",
&request("openrouter/x-ai/grok-4.7"),
)
.await
.unwrap()
.collect::<Vec<_>>()
.await
.into_iter()
.collect::<Result<Vec<_>, _>>()
.unwrap();
upstream.assert_async().await;
let terminal_usage = chunks
.iter()
.filter_map(|chunk| serde_json::from_str::<Value>(chunk).ok())
.find_map(|chunk| {
chunk
.get("usage")
.filter(|usage| usage["provider_cost_floor_usd"].is_number())
.cloned()
})
.unwrap();
assert_eq!(terminal_usage["provider_cost_floor_usd"], 1.0);
assert_eq!(terminal_usage["cost_source"], "provider_floor");
assert_eq!(terminal_usage["cost_usd"], 1.0);
assert!(terminal_usage.get("cost").is_none());
}
#[tokio::test]
async fn gemini_terminal_authority_survives_only_without_later_candidate_activity() {
let unsafe_sequences = vec![
vec![
json!({"candidates":[],"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":3}}),
json!({"candidates":[{"index":0,"content":{"parts":[{"text":"later"}]}}]}),
json!({"candidates":[{"index":0,"finishReason":"STOP","content":{"parts":[]}}]}),
],
vec![
json!({"candidates":[
{"index":0,"finishReason":"STOP","content":{"parts":[]}},
{"index":1,"content":{"parts":[]}}
],"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":3}}),
json!({"candidates":[{"index":1,"content":{"parts":[{"text":"later"}]}}]}),
json!({"candidates":[{"index":1,"finishReason":"STOP","content":{"parts":[]}}]}),
],
vec![
json!({"candidates":[{"index":0,"finishReason":"STOP","content":{"parts":[]}}],"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":3}}),
json!({"candidates":[{"index":1,"content":{"parts":[{"text":"new"}]}}]}),
json!({"candidates":[{"index":1,"finishReason":"STOP","content":{"parts":[]}}]}),
],
];
for sequence in unsafe_sequences {
let events = record_gemini_stream(sequence).await;
assert_eq!(terminal_usage_count(&events), 0, "{events:?}");
}
let valid = record_gemini_stream(vec![
json!({"candidates":[{"index":0,"content":{"parts":[{"text":"ok"}]}}]}),
json!({"candidates":[{"index":0,"finishReason":"STOP","content":{"parts":[]}}]}),
json!({"candidates":[],"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":3}}),
])
.await;
assert_eq!(terminal_usage_count(&valid), 1, "{valid:?}");
}
#[tokio::test]
async fn anthropic_terminal_empty_usage_cannot_bless_earlier_partial_counters() {
let mut server = mockito::Server::new_async().await;
let stream_body = format!(
"data: {}\n\ndata: {}\n\ndata: {}\n\ndata: {}\n\ndata: {}\n\n",
json!({"type":"message_start","message":{"id":"m","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-6","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":7,"output_tokens":0}}}),
json!({"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}),
json!({"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"ok"}}),
json!({"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{}}),
json!({"type":"message_stop"})
);
let upstream = server
.mock("POST", "/messages")
.with_header("content-type", "text/event-stream")
.with_body(stream_body)
.expect(1)
.create_async()
.await;
let provider = Anthropic::new("test-key".into()).with_base_url(server.url());
let policy = Arc::new(RecordingPolicy::default());
let chunks: Vec<_> = ShimClient::new()
.stream_with_policy(
&provider,
"claude-sonnet-4-6",
&request("anthropic/claude-sonnet-4-6"),
&context(policy.clone()),
)
.await
.unwrap()
.collect()
.await;
assert!(chunks.iter().all(Result::is_ok), "{chunks:?}");
upstream.assert_async().await;
assert!(!policy.events().iter().any(|event| matches!(
event,
RecordedEvent::Usage {
terminal: true,
counters_complete: true,
..
}
)));
}
#[tokio::test]
async fn prepared_metadata_uses_resolved_and_native_targets_without_headers() {
let mut server = mockito::Server::new_async().await;
let upstream = server
.mock("POST", "/chat/completions")
.with_body(response(json!("ok")))
.expect(1)
.create_async()
.await;
let provider = OpenAiCompatible::new("resolved-provider", server.url(), Some("secret".into()));
let policy = Arc::new(RecordingPolicy::default());
let mut native_override_request = request("caller/model");
native_override_request["x-resolved-provider"] = json!({"model": "native-model"});
ShimClient::new()
.completion_with_policy(
&provider,
"resolved-model",
&native_override_request,
&context(policy.clone()),
)
.await
.unwrap();
upstream.assert_async().await;
let events = policy.events();
let acquired = events
.iter()
.find_map(|event| match event {
RecordedEvent::Acquired {
id,
kind,
provider_name,
resolved_model,
native_model,
wire,
account_fingerprint,
endpoint,
method,
native_body,
} => Some((
id,
kind,
provider_name,
resolved_model,
native_model,
wire,
account_fingerprint,
endpoint,
method,
native_body,
)),
_ => None,
})
.unwrap();
assert_ne!(*acquired.0, uuid::Uuid::nil());
assert_eq!(*acquired.1, AttemptKind::Completion);
assert_eq!(acquired.2, "resolved-provider");
assert_eq!(acquired.3, "resolved-model");
assert_eq!(acquired.4, "native-model");
assert_eq!(*acquired.5, llmshim::reasoning::WireFormat::OpenAiChat);
assert!(acquired
.6
.as_deref()
.is_some_and(|account| { account.starts_with("sha256:") && !account.contains("secret") }));
assert_eq!(acquired.7, &format!("{}/chat/completions", server.url()));
assert_eq!(acquired.8, "POST");
assert_eq!(acquired.9["model"], "native-model");
assert!(!format!("{acquired:?}").contains("secret"));
}
#[tokio::test]
async fn native_model_override_drives_attempt_pricing() {
let mut server = mockito::Server::new_async().await;
let upstream = server
.mock("POST", "/chat/completions")
.with_body(response(json!("ok")))
.expect(1)
.create_async()
.await;
let provider = OpenAiCompatible::new("openai", server.url(), None);
let policy = Arc::new(RecordingPolicy::default());
let mut overridden_request = request("openai/gpt-5.6-luna");
overridden_request["x-openai"] = json!({"model": "gpt-6-astra"});
let result = ShimClient::new()
.completion_with_policy(
&provider,
"gpt-5.6-luna",
&overridden_request,
&context(policy.clone()),
)
.await
.unwrap();
upstream.assert_async().await;
let usage = policy
.events()
.into_iter()
.find_map(|event| match event {
RecordedEvent::Usage { usage, .. } => Some(usage),
_ => None,
})
.unwrap();
let expected = llmshim::cost::cost_usd("openai", "gpt-6-astra", &usage).unwrap();
assert!((usage["cost_usd"].as_f64().unwrap() - expected).abs() < 1e-12);
assert!((result["usage"]["cost_usd"].as_f64().unwrap() - expected).abs() < 1e-12);
}
async fn assert_native_usage_precedes_invalid_response(
provider: &dyn Provider,
model: &str,
request_model: &str,
) {
let policy = Arc::new(RecordingPolicy::default());
let result = ShimClient::new()
.completion_with_policy(
provider,
model,
&request(request_model),
&context(policy.clone()),
)
.await;
assert!(result.is_err());
let events = policy.events();
let usage_index = events
.iter()
.position(|event| matches!(event, RecordedEvent::Usage { .. }))
.unwrap_or_else(|| {
panic!(
"{} native usage must be observed: {events:?}",
provider.name()
)
});
let invalid_index = events
.iter()
.position(|event| {
matches!(
event,
RecordedEvent::Finished {
outcome: AttemptOutcome::InvalidResponse {
accounting: AttemptAccounting::UsageObserved
},
..
}
)
})
.unwrap_or_else(|| panic!("{} invalid response missing: {events:?}", provider.name()));
assert!(usage_index < invalid_index);
}
#[tokio::test]
async fn native_adapter_usage_is_observed_before_terminal_validation_failure() {
let mut openai_server = mockito::Server::new_async().await;
let openai_mock = openai_server
.mock("POST", "/responses")
.with_body(
json!({
"status":"failed",
"output":[],
"usage":{"input_tokens":7,"output_tokens":3,"total_tokens":10}
})
.to_string(),
)
.expect(1)
.create_async()
.await;
assert_native_usage_precedes_invalid_response(
&OpenAi::new("unused".into()).with_base_url(openai_server.url()),
"gpt-5.6-luna",
"openai/gpt-5.6-luna",
)
.await;
openai_mock.assert_async().await;
let mut anthropic_server = mockito::Server::new_async().await;
let anthropic_mock = anthropic_server
.mock("POST", "/messages")
.with_body(
json!({
"stop_reason":"pause_turn",
"content":[],
"usage":{"input_tokens":7,"output_tokens":3}
})
.to_string(),
)
.expect(1)
.create_async()
.await;
assert_native_usage_precedes_invalid_response(
&Anthropic::new("unused".into()).with_base_url(anthropic_server.url()),
"claude-sonnet-5",
"anthropic/claude-sonnet-5",
)
.await;
anthropic_mock.assert_async().await;
let mut gemini_server = mockito::Server::new_async().await;
let gemini_mock = gemini_server
.mock("POST", "/models/gemini-3.8-flash:generateContent")
.match_query(mockito::Matcher::UrlEncoded("key".into(), "unused".into()))
.with_body(
json!({
"candidates":[{"finishReason":"OTHER","content":{"parts":[]}}],
"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":3,"totalTokenCount":10}
})
.to_string(),
)
.expect(1)
.create_async()
.await;
assert_native_usage_precedes_invalid_response(
&Gemini::new("unused".into()).with_base_url(gemini_server.url()),
"gemini-3.8-flash",
"gemini/gemini-3.8-flash",
)
.await;
gemini_mock.assert_async().await;
let mut xai_server = mockito::Server::new_async().await;
let xai_mock = xai_server
.mock("POST", "/responses")
.with_body(
json!({
"status":"failed",
"output":[],
"usage":{"input_tokens":7,"output_tokens":3,"total_tokens":10}
})
.to_string(),
)
.expect(1)
.create_async()
.await;
assert_native_usage_precedes_invalid_response(
&Xai::new("unused".into()).with_base_url(xai_server.url()),
"grok-4.7",
"xai/grok-4.7",
)
.await;
xai_mock.assert_async().await;
}
#[tokio::test]
async fn chatgpt_terminal_usage_and_account_binding_survive_collection_failure() {
let mut server = mockito::Server::new_async().await;
let auth_directory = tempfile::tempdir().unwrap();
let auth = ChatGptAuth::new(auth_directory.path().join("auth.json"));
std::fs::write(
auth.auth_path(),
json!({
"access_token":"test-access",
"refresh_token":"test-refresh",
"account_id":"private-account",
"expires_at":chrono::Utc::now().timestamp() + 3600
})
.to_string(),
)
.unwrap();
let terminal = json!({
"type":"response.completed",
"response":{
"id":"response",
"status":"failed",
"output":[],
"usage":{"input_tokens":7,"output_tokens":3,"total_tokens":10}
}
});
let upstream = server
.mock("POST", "/responses")
.match_header("chatgpt-account-id", "private-account")
.with_header("content-type", "text/event-stream")
.with_body(format!(
"event: response.completed\r\ndata:{terminal}\r\n\r\n"
))
.expect(1)
.create_async()
.await;
let provider = ChatGpt::new(auth).with_base_url(server.url());
let policy = Arc::new(RecordingPolicy::default());
let result = ShimClient::new()
.completion_with_policy(
&provider,
"gpt-6-astra",
&request("chatgpt/gpt-6-astra"),
&context(policy.clone()),
)
.await;
assert!(result.is_err());
upstream.assert_async().await;
let events = policy.events();
assert!(events.iter().any(|event| matches!(
event,
RecordedEvent::Usage { usage, .. } if usage["total_tokens"] == 10
)));
let account = events.iter().find_map(|event| match event {
RecordedEvent::Acquired {
account_fingerprint,
..
} => account_fingerprint.as_deref(),
_ => None,
});
assert!(account.is_some_and(|account| {
account.starts_with("sha256:") && !account.contains("private-account")
}));
assert!(!format!("{events:?}").contains("private-account"));
}
#[tokio::test]
async fn chatgpt_usage_precedes_terminal_retention_failure_on_the_public_client_path() {
let mut server = mockito::Server::new_async().await;
let auth_directory = tempfile::tempdir().unwrap();
let auth = ChatGptAuth::new(auth_directory.path().join("auth.json"));
std::fs::write(
auth.auth_path(),
json!({
"access_token":"test-access",
"refresh_token":"test-refresh",
"account_id":"test-account",
"expires_at":chrono::Utc::now().timestamp() + 3600
})
.to_string(),
)
.unwrap();
let terminal = json!({
"type":"response.completed",
"response":{
"id":"response",
"status":"completed",
"output":[{
"type":"message",
"role":"assistant",
"content":[{"type":"output_text","text":"x".repeat(4_096)}]
}],
"usage":{"input_tokens":7,"output_tokens":3,"total_tokens":10}
}
});
let upstream = server
.mock("POST", "/responses")
.with_header("content-type", "text/event-stream")
.with_body(format!("data:{terminal}\n\n"))
.expect(1)
.create_async()
.await;
let provider = ChatGpt::new(auth).with_base_url(server.url());
let policy = Arc::new(RecordingPolicy::default());
let client = ShimClient::new()
.with_stream_retention_limits(StreamRetentionLimits::new(512, 64, 4 * 1024, 64).unwrap())
.unwrap();
let error = client
.completion_with_policy(
&provider,
"gpt-6-astra",
&request("chatgpt/gpt-6-astra"),
&context(policy.clone()),
)
.await
.unwrap_err();
upstream.assert_async().await;
assert!(error
.to_string()
.contains("upstream stream retained state exceeds limit"));
let events = policy.events();
let usage_index = events
.iter()
.position(|event| {
matches!(
event,
RecordedEvent::Usage { usage, terminal: true, .. }
if usage["total_tokens"] == 10
)
})
.unwrap();
let failure_index = events
.iter()
.position(|event| {
matches!(
event,
RecordedEvent::Finished {
outcome: AttemptOutcome::InvalidResponse {
accounting: AttemptAccounting::UsageObserved
},
..
}
)
})
.unwrap();
assert!(usage_index < failure_index);
assert!(!format!("{events:?}").contains(&"x".repeat(128)));
}
async fn assert_native_stream_usage_precedes_failure(
provider: &dyn Provider,
model: &str,
request_model: &str,
) {
let policy = Arc::new(RecordingPolicy::default());
let chunks: Vec<_> = ShimClient::new()
.stream_with_policy(
provider,
model,
&request(request_model),
&context(policy.clone()),
)
.await
.unwrap_or_else(|error| panic!("{} stream did not open: {error:?}", provider.name()))
.collect()
.await;
assert!(chunks.iter().any(Result::is_err));
let events = policy.events();
let usage_index = events
.iter()
.position(|event| matches!(event, RecordedEvent::Usage { .. }))
.unwrap_or_else(|| {
panic!(
"{} native stream usage must be observed: {events:?}",
provider.name()
)
});
let failure_index = events
.iter()
.position(|event| {
matches!(
event,
RecordedEvent::Finished {
outcome: AttemptOutcome::StreamFailure {
accounting: AttemptAccounting::UsageObserved
},
..
}
)
})
.unwrap_or_else(|| panic!("{} stream failure missing: {events:?}", provider.name()));
assert!(usage_index < failure_index);
}
#[tokio::test]
async fn native_stream_usage_is_observed_before_normalization_failure() {
let responses_failure = json!({
"type":"response.failed",
"response":{
"status":"failed",
"usage":{"input_tokens":7,"output_tokens":3,"total_tokens":10}
}
});
let mut openai_server = mockito::Server::new_async().await;
let openai_mock = openai_server
.mock("POST", "/responses")
.with_header("content-type", "text/event-stream")
.with_body(format!("data: {responses_failure}\n\n"))
.expect(1)
.create_async()
.await;
assert_native_stream_usage_precedes_failure(
&OpenAi::new("unused".into()).with_base_url(openai_server.url()),
"gpt-5.6-luna",
"openai/gpt-5.6-luna",
)
.await;
openai_mock.assert_async().await;
let mut anthropic_server = mockito::Server::new_async().await;
let anthropic_mock = anthropic_server
.mock("POST", "/messages")
.with_header("content-type", "text/event-stream")
.with_body(format!(
"data: {}\n\n",
json!({
"type":"error",
"usage":{"input_tokens":7,"output_tokens":3},
"error":{"type":"overloaded_error"}
})
))
.expect(1)
.create_async()
.await;
assert_native_stream_usage_precedes_failure(
&Anthropic::new("unused".into()).with_base_url(anthropic_server.url()),
"claude-sonnet-5",
"anthropic/claude-sonnet-5",
)
.await;
anthropic_mock.assert_async().await;
let mut gemini_server = mockito::Server::new_async().await;
let gemini_mock = gemini_server
.mock("POST", "/models/gemini-3.8-flash:streamGenerateContent")
.match_query(mockito::Matcher::Exact("key=unused&alt=sse".into()))
.with_header("content-type", "text/event-stream")
.with_body(format!(
"data: {}\n\n",
json!({
"error":{"code":500},
"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":3,"totalTokenCount":10}
})
))
.expect(1)
.create_async()
.await;
assert_native_stream_usage_precedes_failure(
&Gemini::new("unused".into()).with_base_url(gemini_server.url()),
"gemini-3.8-flash",
"gemini/gemini-3.8-flash",
)
.await;
gemini_mock.assert_async().await;
let mut xai_server = mockito::Server::new_async().await;
let xai_mock = xai_server
.mock("POST", "/responses")
.with_header("content-type", "text/event-stream")
.with_body(format!("data: {responses_failure}\n\n"))
.expect(1)
.create_async()
.await;
assert_native_stream_usage_precedes_failure(
&Xai::new("unused".into()).with_base_url(xai_server.url()),
"grok-4.7",
"xai/grok-4.7",
)
.await;
xai_mock.assert_async().await;
}
#[tokio::test]
async fn usage_and_finish_observer_failures_stop_unary_work_and_retain_liability() {
for failing_event in [FailingEvent::Usage, FailingEvent::Finished] {
let mut server = mockito::Server::new_async().await;
let upstream = server
.mock("POST", "/chat/completions")
.with_body(response(json!("ok")))
.expect(1)
.create_async()
.await;
let provider = OpenAiCompatible::new("local", server.url(), None);
let policy = Arc::new(RecordingPolicy::failing(failing_event));
let error = ShimClient::new()
.completion_with_policy(
&provider,
"test",
&request("local/test"),
&context(policy.clone()),
)
.await
.unwrap_err();
assert!(error
.to_string()
.contains("attempt policy coordinator unavailable"));
upstream.assert_async().await;
assert!(policy
.events()
.iter()
.any(|event| matches!(event, RecordedEvent::Abandoned { .. })));
}
}
#[tokio::test]
async fn stream_observer_failure_is_caller_visible_before_completion() {
let mut server = mockito::Server::new_async().await;
let upstream = server
.mock("POST", "/chat/completions")
.with_header("content-type", "text/event-stream")
.with_body(successful_stream("ok"))
.expect(1)
.create_async()
.await;
let provider = OpenAiCompatible::new("local", server.url(), None);
let policy = Arc::new(RecordingPolicy::failing(FailingEvent::Usage));
let chunks: Vec<_> = ShimClient::new()
.stream_with_policy(
&provider,
"test",
&request("local/test"),
&context(policy.clone()),
)
.await
.unwrap()
.collect()
.await;
assert!(chunks
.iter()
.any(|chunk| chunk.as_ref().is_err_and(|error| {
error
.to_string()
.contains("attempt policy coordinator unavailable")
})));
assert!(policy
.events()
.iter()
.any(|event| matches!(event, RecordedEvent::Abandoned { .. })));
upstream.assert_async().await;
}
#[tokio::test]
async fn failed_drop_observation_is_conservative_and_does_not_panic() {
let mut server = mockito::Server::new_async().await;
let upstream = server
.mock("POST", "/chat/completions")
.with_header("content-type", "text/event-stream")
.with_body(successful_stream("unused"))
.expect(1)
.create_async()
.await;
let provider = OpenAiCompatible::new("local", server.url(), None);
let policy = Arc::new(RecordingPolicy::failing_abandonment());
let stream = ShimClient::new()
.stream_with_policy(
&provider,
"test",
&request("local/test"),
&context(policy.clone()),
)
.await
.unwrap();
drop(stream);
assert!(policy
.events()
.iter()
.any(|event| matches!(event, RecordedEvent::Abandoned { .. })));
upstream.assert_async().await;
}
#[tokio::test]
async fn cancelling_an_inflight_usage_observation_marks_the_attempt_abandoned() {
let mut server = mockito::Server::new_async().await;
let upstream = server
.mock("POST", "/chat/completions")
.with_header("content-type", "text/event-stream")
.with_body(successful_stream("ok"))
.expect(1)
.create_async()
.await;
let provider = OpenAiCompatible::new("local", server.url(), None);
let policy = Arc::new(CancellationPolicy::new());
let mut stream = ShimClient::new()
.stream_with_policy(
&provider,
"test",
&request("local/test"),
&DispatchPolicyContext::new(policy.clone()),
)
.await
.unwrap();
let task = tokio::spawn(async move { while stream.next().await.is_some() {} });
tokio::time::timeout(
std::time::Duration::from_secs(1),
policy.usage_started.notified(),
)
.await
.unwrap();
task.abort();
let _ = task.await;
assert!(policy.acquired.load(Ordering::SeqCst));
assert!(policy.abandoned.load(Ordering::SeqCst));
upstream.assert_async().await;
}
#[tokio::test]
async fn streams_report_usage_terminal_failure_and_abandonment() {
let mut server = mockito::Server::new_async().await;
let successful = server
.mock("POST", "/chat/completions")
.with_header("content-type", "text/event-stream")
.with_body(successful_stream("ok"))
.expect(1)
.create_async()
.await;
let provider = OpenAiCompatible::new("local", server.url(), None);
let successful_policy = Arc::new(RecordingPolicy::default());
let chunks: Vec<_> = ShimClient::new()
.stream_with_policy(
&provider,
"test",
&request("local/test"),
&context(successful_policy.clone()),
)
.await
.unwrap()
.collect()
.await;
assert!(chunks.iter().all(Result::is_ok));
successful.assert_async().await;
let successful_events = successful_policy.events();
assert!(successful_events.iter().any(|event| matches!(
event,
RecordedEvent::Usage { usage, .. } if usage["total_tokens"] == 10
)));
assert!(successful_events.iter().any(|event| matches!(
event,
RecordedEvent::Finished {
outcome: AttemptOutcome::Completed {
accounting: AttemptAccounting::UsageObserved
},
..
}
)));
successful.remove_async().await;
let failed = server
.mock("POST", "/chat/completions")
.with_header("content-type", "text/event-stream")
.with_body("data: {\"error\":{\"message\":\"failed\"}}\n\n")
.expect(1)
.create_async()
.await;
let failed_policy = Arc::new(RecordingPolicy::default());
let failed_chunks: Vec<_> = ShimClient::new()
.stream_with_policy(
&provider,
"test",
&request("local/test"),
&context(failed_policy.clone()),
)
.await
.unwrap()
.collect()
.await;
assert!(failed_chunks.iter().any(Result::is_err));
failed.assert_async().await;
assert!(failed_policy.events().iter().any(|event| matches!(
event,
RecordedEvent::Finished {
outcome: AttemptOutcome::StreamFailure {
accounting: AttemptAccounting::Unknown
},
..
}
)));
failed.remove_async().await;
let abandoned = server
.mock("POST", "/chat/completions")
.with_header("content-type", "text/event-stream")
.with_body(successful_stream("unused"))
.expect(1)
.create_async()
.await;
let abandoned_policy = Arc::new(RecordingPolicy::default());
let stream = ShimClient::new()
.stream_with_policy(
&provider,
"test",
&request("local/test"),
&context(abandoned_policy.clone()),
)
.await
.unwrap();
drop(stream);
abandoned.assert_async().await;
assert!(abandoned_policy.events().iter().any(|event| matches!(
event,
RecordedEvent::Abandoned {
outcome: AttemptOutcome::Abandoned {
kind: AttemptKind::Stream,
accounting: AttemptAccounting::Unknown
},
..
}
)));
}
#[tokio::test]
async fn cumulative_and_duplicate_stream_usage_share_one_attempt_identity() {
let mut server = mockito::Server::new_async().await;
let body = format!(
"data: {}\n\ndata: {}\n\ndata: {}\n\ndata: {}\n\ndata: {}\n\ndata: [DONE]\n\n",
json!({"id":"response","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":null}]}),
json!({"id":"response","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}),
json!({"choices":[],"usage":{"prompt_tokens":7,"completion_tokens":1,"total_tokens":8}}),
json!({"choices":[],"usage":{"prompt_tokens":7,"completion_tokens":1,"total_tokens":8}}),
json!({"choices":[],"usage":{"prompt_tokens":7,"completion_tokens":3,"total_tokens":10}}),
);
let upstream = server
.mock("POST", "/chat/completions")
.with_header("content-type", "text/event-stream")
.with_body(body)
.expect(1)
.create_async()
.await;
let provider = OpenAiCompatible::new("local", server.url(), None);
let policy = Arc::new(RecordingPolicy::default());
let chunks: Vec<_> = ShimClient::new()
.stream_with_policy(
&provider,
"test",
&request("local/test"),
&context(policy.clone()),
)
.await
.unwrap()
.collect()
.await;
assert!(chunks.iter().all(Result::is_ok));
upstream.assert_async().await;
let events = policy.events();
let usage_events: Vec<_> = events
.iter()
.filter_map(|event| match event {
RecordedEvent::Usage { id, usage, .. } => Some((*id, usage["total_tokens"].as_u64())),
_ => None,
})
.collect();
assert_eq!(usage_events.len(), 4);
assert_eq!(usage_events[0].1, Some(8));
assert_eq!(usage_events[1].1, Some(8));
assert_eq!(usage_events[2].1, Some(10));
assert_eq!(usage_events[3].1, Some(10));
assert_eq!(usage_events[0].0, usage_events[1].0);
assert_eq!(usage_events[1].0, usage_events[2].0);
assert_eq!(usage_events[2].0, usage_events[3].0);
assert!(events.iter().any(|event| matches!(
event,
RecordedEvent::Usage {
terminal: true,
usage,
..
} if usage["total_tokens"] == 10
)));
let finished_id = events.iter().find_map(|event| match event {
RecordedEvent::Finished {
id,
outcome:
AttemptOutcome::Completed {
accounting: AttemptAccounting::UsageObserved,
},
} => Some(*id),
_ => None,
});
assert_eq!(finished_id, Some(usage_events[0].0));
}
#[tokio::test]
async fn owned_buffered_stream_failed_repair_observes_both_attempts() {
let mut server = mockito::Server::new_async().await;
let invalid_streams = server
.mock("POST", "/chat/completions")
.with_header("content-type", "text/event-stream")
.with_body(successful_stream("1"))
.expect(2)
.create_async()
.await;
let router = llmshim::router::Router::new().register(
"local",
Box::new(OpenAiCompatible::new("local", server.url(), None)),
);
let policy = Arc::new(RecordingPolicy::default());
let stream = llmshim::stream_with_policy(
&router,
&structured_request("local/test"),
&context(policy.clone()),
)
.await
.unwrap();
let chunks: Vec<_> = stream.collect().await;
assert_eq!(chunks.len(), 1);
assert!(chunks[0].is_err());
invalid_streams.assert_async().await;
let events = policy.events();
assert_eq!(
events
.iter()
.filter(|event| matches!(event, RecordedEvent::Acquired { .. }))
.count(),
2
);
let attempts_with_usage: std::collections::HashSet<_> = events
.iter()
.filter_map(|event| match event {
RecordedEvent::Usage { id, .. } => Some(*id),
_ => None,
})
.collect();
assert_eq!(attempts_with_usage.len(), 2);
}
#[test]
fn policy_diagnostics_do_not_expose_native_request_content() {
struct NoopPolicy;
impl AttemptPolicy for NoopPolicy {
fn acquire<'a>(
&'a self,
_attempt: &'a PreparedAttempt<'a>,
) -> AttemptPolicyFuture<'a, Result<(), AttemptPolicyRefusal>> {
Box::pin(async { Ok(()) })
}
fn observe<'a>(
&'a self,
_attempt: &'a AttemptIdentity,
_event: AttemptEvent<'a>,
) -> AttemptPolicyFuture<'a, Result<(), AttemptPolicyError>> {
Box::pin(async { Ok(()) })
}
fn observe_abandoned(
&self,
_attempt: &AttemptIdentity,
_outcome: AttemptOutcome,
) -> Result<(), AttemptPolicyError> {
Ok(())
}
}
let context = DispatchPolicyContext::new(Arc::new(NoopPolicy));
assert_eq!(format!("{context:?}"), "DispatchPolicyContext { .. }");
}