test(rust): add Messages live provider harness (#45823)

* test(rust): add Messages live provider harness

* test(rust): support automatic Vertex tool selection
This commit is contained in:
yujonglee 2026-10-12 17:39:08 -07:00 • committed by GitHub
parent c2bc5aecff
commit 8db9037c90
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 440 additions and 0 deletions

View file

@ -4024,6 +4024,7 @@ dependencies = [
"litellm-cache-memory",
"litellm-cache-response",
"litellm-core-utils",
"litellm-framer",
"litellm-host",
"litellm-host-native",
"litellm-http",
@ -4107,9 +4108,11 @@ name = "litellm-inference-testing"
version = "0.1.0"
dependencies = [
"futures-util",
"litellm-auth",
"litellm-http",
"litellm-inference",
"litellm-secrets",
"tokio",
]
[[package]]

View file

@ -26,6 +26,7 @@ serde_json = { workspace = true, features = ["preserve_order"] }
tracing.workspace = true
[dev-dependencies]
litellm-framer = { workspace = true, features = ["sse"] }
litellm-cache-memory.workspace = true
litellm-host-native.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }

View file

@ -0,0 +1,5 @@
#[path = "live/support.rs"]
mod support;
#[path = "live/vertex_ai.rs"]
mod vertex_ai;

View file

@ -0,0 +1,265 @@
use std::{sync::OnceLock, time::Duration};
use futures_util::{FutureExt, TryStreamExt};
use litellm_host::interceptors::{
ExecutionFacts, Interceptors, RawResponse, RequestContext, ResultSource, WireRequest,
};
use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesRoute};
use litellm_inference_testing::live::{LiveResources, model, within_deadline};
use litellm_llms_types::formats::messages::MessagesResponse;
use serde_json::{Value, json};
pub(super) struct LiveCall {
provider: &'static str,
facts: OnceLock<ExecutionFacts>,
}
impl LiveCall {
pub(super) fn new(provider: &'static str) -> Self {
Self {
provider,
facts: OnceLock::new(),
}
}
pub(super) fn facts(&self) -> &ExecutionFacts {
self.facts
.get()
.expect("the route must deliver execution facts")
}
pub(super) fn assert_provider_result(&self) {
let facts = self.facts();
assert_eq!(facts.provider.provider, self.provider);
assert_eq!(facts.source, ResultSource::Provider);
println!(
"{}",
json!({"route": "messages", "provider": self.provider, "source": "provider"})
);
}
}
impl Interceptors<Error> for LiveCall {
async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> {
assert!(
self.facts.set(facts).is_ok(),
"result must be accepted exactly once"
);
Ok(())
}
async fn before_provider_request(
&self,
wire: WireRequest,
context: RequestContext,
) -> Result<WireRequest, Error> {
assert_eq!(context.custom_llm_provider, self.provider);
println!(
"{}",
json!({"route": "messages", "provider": self.provider, "model": context.model})
);
Ok(wire)
}
async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> {
Ok(())
}
}
#[rstest::fixture]
pub(super) fn route() -> MessagesRoute {
let resources = LiveResources::default();
MessagesRoute::new(resources.http, resources.auth, resources.secrets)
}
pub(super) async fn request(
route: &MessagesRoute,
call: MessagesCall,
host: &LiveCall,
) -> MessagesCallResponse {
route
.execute(call, host, None)
.boxed()
.await
.expect("live request failed")
}
pub(super) fn call(provider: &str, payload: Value) -> MessagesCall {
let Value::Object(fields) = payload else {
panic!("live payload must be an object")
};
let body = serde_json::from_value(Value::Object(
[("max_tokens".into(), json!(128))]
.into_iter()
.chain(fields)
.chain([("model".into(), json!(model("messages", provider)))])
.collect(),
))
.expect("live payload must be a valid Messages request");
MessagesCall {
body,
api_key: None,
api_base: None,
custom_llm_provider: Some(provider.into()),
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: Some(Duration::from_secs(60)),
shaping: Default::default(),
}
}
pub(super) fn complete(response: MessagesCallResponse) -> Box<MessagesResponse> {
match response {
MessagesCallResponse::Complete(message) => message,
MessagesCallResponse::Stream { .. } => panic!("a complete request returned a stream"),
}
}
pub(super) fn assert_text(message: &MessagesResponse) {
assert!(!message.id.is_empty());
assert_eq!(message.role, "assistant");
assert!(message.stop_reason.is_some());
assert!(
message
.content
.iter()
.any(|block| block["text"].as_str().is_some_and(|text| !text.is_empty()))
);
assert!(
message
.usage
.as_ref()
.and_then(|usage| usage["output_tokens"].as_u64())
.is_some_and(|count| count > 0)
);
}
pub(super) async fn stream_events(response: MessagesCallResponse) -> Vec<Value> {
let MessagesCallResponse::Stream { head, chunks } = response else {
panic!("stream request returned a complete response")
};
let content_types: Vec<_> = head
.headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("content-type"))
.map(|(_, value)| value.split(';').next().unwrap().trim())
.collect();
assert_eq!(content_types, ["text/event-stream"]);
let frames = litellm_framer::frames(chunks, litellm_framer::sse::SseCodec::default())
.try_collect::<Vec<_>>()
.await
.expect("live SSE framing failed");
let data_frames = match frames.last() {
Some(frame) if frame.data == "[DONE]" => &frames[..frames.len() - 1],
_ => frames.as_slice(),
};
let events: Vec<Value> = data_frames
.iter()
.map(|frame| serde_json::from_str(&frame.data).expect("SSE data must be JSON"))
.collect();
assert_eq!(
events.first().map(|event| event["type"].as_str()),
Some(Some("message_start"))
);
assert_eq!(
events.last().map(|event| event["type"].as_str()),
Some(Some("message_stop"))
);
assert!(!events.iter().any(|event| event["type"] == "error"));
assert!(events.iter().any(|event| {
event["usage"]["output_tokens"]
.as_u64()
.is_some_and(|count| count > 0)
}));
println!("{}", json!({"events": events}));
events
}
fn prompt() -> Value {
json!({"role": "user", "content": "Use the echo tool with value live-check, then tell me the result."})
}
fn tools() -> Value {
json!([{"name": "echo", "description": "Return a value", "input_schema": {
"type": "object", "properties": {"value": {"type": "string", "enum": ["live-check"]}}, "required": ["value"]
}}])
}
fn fragments(events: &[Value], index: u64, field: &str) -> String {
events
.iter()
.filter(|event| event["index"].as_u64() == Some(index))
.filter_map(|event| event["delta"][field].as_str())
.collect()
}
fn streamed_block(start: &Value, events: &[Value]) -> Value {
let index = start["index"]
.as_u64()
.expect("content block must have an index");
let block = &start["content_block"];
match block["type"].as_str() {
Some("tool_use") => {
let partial = fragments(events, index, "partial_json");
let input = match partial.is_empty() {
true => block["input"].clone(),
false => {
serde_json::from_str(&partial).expect("streamed tool arguments must be JSON")
}
};
json!({"type": "tool_use", "id": block["id"], "name": block["name"], "input": input})
}
Some("text") => {
json!({"type": "text", "text": format!("{}{}", block["text"].as_str().unwrap_or_default(), fragments(events, index, "text"))})
}
Some("thinking") => {
json!({"type": "thinking", "thinking": format!("{}{}", block["thinking"].as_str().unwrap_or_default(), fragments(events, index, "thinking")),
"signature": format!("{}{}", block["signature"].as_str().unwrap_or_default(), fragments(events, index, "signature"))})
}
_ => block.clone(),
}
}
pub(super) async fn tool_round_trip(
route: MessagesRoute,
provider: &'static str,
stream: bool,
tool_params: Value,
) {
within_deadline(async {
let first_host = LiveCall::new(provider);
let payload = json!({"messages": [prompt()], "tools": tools(), "stream": stream});
let fields = payload.as_object().expect("tool payload must be an object").clone();
let params = tool_params.as_object().expect("tool parameters must be an object").clone();
let first = request(&route, call(provider, Value::Object(fields.into_iter().chain(params).collect())), &first_host).await;
let content = match stream {
false => {
let message = complete(first);
assert_eq!(message.stop_reason.as_deref(), Some("tool_use"));
message.content
}
true => {
let events = stream_events(first).await;
assert!(events.iter().any(|event| event["delta"]["stop_reason"] == "tool_use"));
events.iter().filter(|event| event["type"] == "content_block_start")
.map(|event| streamed_block(event, &events)).collect()
}
};
let tool = content.iter().find(|block| block["type"] == "tool_use").expect("requested tool must be returned");
assert_eq!(tool["name"], "echo");
assert_eq!(tool["input"]["value"], "live-check");
let tool_id = tool["id"].as_str().filter(|id| !id.is_empty()).expect("tool must have an id");
first_host.assert_provider_result();
let second_host = LiveCall::new(provider);
let second = complete(request(&route, call(provider, json!({
"tools": tools(), "messages": [prompt(),
{"role": "assistant", "content": content},
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": tool_id, "content": "live-check"}]}]
})), &second_host).await);
assert_text(&second);
second_host.assert_provider_result();
println!("{}", serde_json::to_string(&second).unwrap());
})
.await;
}

View file

@ -0,0 +1,68 @@
use litellm_inference_messages::MessagesRoute;
use litellm_inference_testing::live::within_deadline;
use rstest::rstest;
use serde_json::{Value, json};
use super::support::{LiveCall, assert_text, call, complete, request, route, stream_events};
const PROVIDER: &str = "vertex_ai";
#[rstest]
#[case::basic(json!({"messages": [{"role": "user", "content": "Say hello."}]}))]
#[case::system_text(json!({
"system": "You are a helpful assistant.",
"messages": [{"role": "user", "content": "Say hello."}]
}))]
#[ignore = "calls a real provider, requires credentials and a live model"]
#[tokio::test]
async fn completion(route: MessagesRoute, #[case] payload: Value) {
within_deadline(async {
let host = LiveCall::new(PROVIDER);
let message = complete(request(&route, call(PROVIDER, payload), &host).await);
assert_text(&message);
host.assert_provider_result();
println!("{}", serde_json::to_string(&message).unwrap());
})
.await;
}
#[rstest]
#[ignore = "calls a real provider, requires credentials and a live model"]
#[tokio::test]
async fn streaming(route: MessagesRoute) {
within_deadline(async {
let host = LiveCall::new(PROVIDER);
let response = request(
&route,
call(
PROVIDER,
json!({"stream": true, "messages": [{"role": "user", "content": "Say hello."}]}),
),
&host,
)
.await;
let events = stream_events(response).await;
assert!(events.iter().any(|event| {
event["delta"]["text"]
.as_str()
.is_some_and(|text| !text.is_empty())
}));
host.assert_provider_result();
})
.await;
}
#[rstest]
#[case::complete(false)]
#[case::streaming(true)]
#[ignore = "calls a real provider, requires credentials and a live model"]
#[tokio::test]
async fn tool_round_trip(route: MessagesRoute, #[case] stream: bool) {
super::support::tool_round_trip(
route,
PROVIDER,
stream,
json!({"tool_choice": {"type": "auto"}}),
)
.await;
}

View file

@ -7,7 +7,9 @@ repository.workspace = true
publish = false
[dependencies]
litellm-auth.workspace = true
futures-util.workspace = true
litellm-http.workspace = true
litellm-inference.workspace = true
litellm-secrets.workspace = true
tokio.workspace = true

View file

@ -0,0 +1,41 @@
# Live provider route tests
Each route crate owns `tests/live.rs`. Use named `rstest` provider cases for the behaviors that route supports. Shared setup lives in `litellm_inference_testing::live`, with production HTTP settings, authentication services, environment secrets and a deadline
These tests execute the Rust route against a real provider. They do not prove Python parity, proxy authorization, billing persistence or agent compatibility. Those belong in the corresponding boundary and gateway suites
All live cases use `#[ignore = "calls a real provider, requires credentials and a live model"]`. Normal Cargo runs report them as ignored. An explicitly selected case fails on missing credentials, unavailable models, provider errors or timeouts. Never return early and count missing coverage as a pass
Set `LITELLM_LIVE_<ROUTE>_<PROVIDER>_MODEL` to the latest model available to the test account. The model has no compiled default because vendor availability changes independently of LiteLLM. Credentials use the provider's existing configuration rather than a second test-only key convention. Export credentials in the shell before running tests. The harness does not load `.env` or modify process environment while tests run
## Messages
Completion checks text and usage. Streaming consumes SSE to completion and checks text, usage and the terminal event. Tool round trips request a tool call and send its returned ID back with a result. Each case checks provider execution facts and rejects cache results
Messages setup and assertions live in `inference-messages/tests/live/support.rs`. Provider modules under `tests/live/` select their payloads and tool policy. The shared helpers consume streams completely and reconstruct text, thinking signatures and tool arguments for the follow-up request
The initial provider case is `vertex_ai`. A registered ignored case is available coverage, not evidence that it passed against that provider
Run one provider's cases explicitly
```bash
LITELLM_LIVE_MESSAGES_VERTEX_AI_MODEL=<claude-model-on-vertex> \
cargo test --manifest-path litellm-rust/Cargo.toml \
-p litellm-inference-messages --test live vertex_ai \
-- --ignored --nocapture --test-threads=1
```
For Vertex AI, set `LITELLM_LIVE_MESSAGES_VERTEX_AI_MODEL` to a Claude model published in the project's location. Credentials, project and location resolve through the production Rust GCP authentication, from `GOOGLE_APPLICATION_CREDENTIALS` or `VERTEXAI_CREDENTIALS`, with `VERTEXAI_PROJECT` and `VERTEXAI_LOCATION`
Vertex AI cases live in `inference-messages/tests/live/vertex_ai.rs`. They cover plain and system messages, text streaming, and complete and streaming tool round trips. Tool cases use automatic selection because current Claude models reject forced `tool_choice`. Each case still requires a tool call, checks its arguments, and sends the tool result back
Tool cases require a model that supports forced `tool_choice`, except where a provider's section says otherwise. A provider rejection fails the selected case instead of skipping it
Select one behavior by its test name, for example `completion::case_1_basic`. To list every available cell without contacting any provider
```bash
cargo test --manifest-path litellm-rust/Cargo.toml \
-p litellm-inference-messages --test live -- --list
```
For another route, reuse `LiveResources`, `model` and `within_deadline`. Keep payload construction and response assertions in that route's test file. Add only supported provider behaviors, and document any gap instead of substituting mocks or skipping failures. Run cases serially to keep spending and rate limits predictable

View file

@ -9,6 +9,8 @@ use litellm_secrets::{SecretValue, source::SecretSource};
use litellm_inference::resources::CoreResources;
pub mod live;
pub fn http_pool() -> HttpClientPool {
HttpClientPool::new(Arc::new(PublicDnsResolver))
}

View file

@ -0,0 +1,53 @@
use std::{future::Future, sync::Arc, time::Duration};
use litellm_http::{ClientVariant, HttpSettings, HttpSettingsLayer, Resolution};
use litellm_inference::resources::CoreResources;
use litellm_secrets::source::{EnvironmentSecrets, SecretSource};
pub const DEADLINE: Duration = Duration::from_secs(120);
pub struct LiveResources {
pub http: litellm_http::Client,
pub auth: Arc<litellm_auth::AuthServices>,
pub secrets: Arc<dyn SecretSource>,
}
impl Default for LiveResources {
fn default() -> Self {
let resources = CoreResources::new(Arc::new(super::http_pool()));
let settings =
HttpSettings::from_layers([HttpSettingsLayer::from_environment(&|key: &str| {
std::env::var(key).ok()
})]);
let resolution = Resolution::from(&settings);
let http = resources
.pool
.client(&resolution.config, ClientVariant::Provider)
.expect("live provider HTTP settings must be valid");
let secrets_http = resources
.pool
.client(&resolution.config, ClientVariant::NoRedirect)
.expect("live secret HTTP settings must be valid");
Self {
http,
auth: resources.auth,
secrets: Arc::new(EnvironmentSecrets::python_compatible(secrets_http)),
}
}
}
pub fn model(route: &str, provider: &str) -> String {
let name = format!("LITELLM_LIVE_{route}_{provider}_MODEL").to_ascii_uppercase();
std::env::var(&name)
.ok()
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| {
panic!("set {name} to an available current model before running this live case")
})
}
pub async fn within_deadline<T>(operation: impl Future<Output = T>) -> T {
tokio::time::timeout(DEADLINE, operation)
.await
.expect("live case exceeded its 120 second deadline")
}