use std::path::PathBuf;
use tokio::sync::mpsc;
use crate::models::adapters::driver::{StreamProtocol, drive_stream};
use crate::models::error::{BackendError, ModelError};
use crate::models::stream::StreamEvent;
use crate::models::types::{FinishReason, ModelResponse};
#[derive(Debug, PartialEq, Eq)]
struct Outcome {
streamed_text: String,
streamed_reasoning: String,
streamed_tool_calls: Vec<(String, String)>,
content: String,
thinking: Option<String>,
response_tool_calls: Vec<(String, String)>,
stop_reason: Option<FinishReason>,
usage: Option<(usize, usize)>,
}
#[derive(Debug, PartialEq, Eq)]
enum Failure {
ProviderError,
StreamCut,
}
type ScenarioResult = std::result::Result<Outcome, Failure>;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Delivery {
Whole,
OneByteAtATime,
}
async fn run<P: StreamProtocol>(protocol: P, body: &str, delivery: Delivery) -> ScenarioResult {
let bytes = body.as_bytes();
let chunks: Vec<std::result::Result<Vec<u8>, String>> = match delivery {
Delivery::Whole => vec![Ok(bytes.to_vec())],
Delivery::OneByteAtATime => bytes.iter().map(|b| Ok(vec![*b])).collect(),
};
let (tx, mut rx) = mpsc::channel::<StreamEvent>(4096);
let result = drive_stream(futures::stream::iter(chunks), protocol, Some(&tx)).await;
drop(tx);
let mut events = Vec::new();
while let Some(event) = rx.recv().await {
events.push(event);
}
match result {
Ok(response) => Ok(normalize(&events, response)),
Err(e) => Err(classify(&e)),
}
}
fn classify(err: &ModelError) -> Failure {
if matches!(err, ModelError::Backend(BackendError::ProviderError { .. })) {
return Failure::ProviderError;
}
assert!(
matches!(err, ModelError::StreamError(_)),
"scenario failed in an unexpected way: {err:?}"
);
Failure::StreamCut
}
fn normalize(events: &[StreamEvent], response: ModelResponse) -> Outcome {
let mut streamed_text = String::new();
let mut streamed_reasoning = String::new();
let mut streamed_tool_calls = Vec::new();
let mut off_contract: Vec<String> = Vec::new();
for event in events {
match event {
StreamEvent::Text(t) => streamed_text.push_str(t),
StreamEvent::Reasoning(c) => streamed_reasoning.push_str(&c.text),
StreamEvent::ToolCall(tc) => streamed_tool_calls
.push((tc.function.name.clone(), tc.function.arguments.to_string())),
StreamEvent::Status(_) | StreamEvent::Done { .. } => {
off_contract.push(format!("{event:?}"));
},
}
}
assert!(
off_contract.is_empty(),
"adapter emitted {off_contract:?} from inside the stream"
);
Outcome {
streamed_text,
streamed_reasoning,
streamed_tool_calls,
content: response.content,
thinking: response.thinking,
response_tool_calls: response
.tool_calls
.unwrap_or_default()
.into_iter()
.map(|tc| (tc.function.name, tc.function.arguments.to_string()))
.collect(),
stop_reason: response.stop_reason,
usage: response
.usage
.map(|u| (u.input_total_tokens(), u.output_total_tokens())),
}
}
fn fixture(provider: &str, name: &str) -> String {
let path: PathBuf = [
env!("CARGO_MANIFEST_DIR"),
"tests",
"fixtures",
"streams",
provider,
name,
]
.iter()
.collect();
assert!(path.is_file(), "missing fixture {}", path.display());
std::fs::read_to_string(&path).expect("fixture reads")
}
mod protocols {
use super::super::anthropic::AnthropicStream;
use super::super::gemini::GeminiStream;
use super::super::meta::MetaStream;
use super::super::ollama::OllamaStream;
use super::super::openai_compat::OpenAICompatStream;
pub(super) fn anthropic() -> AnthropicStream {
AnthropicStream::new("claude-test".to_string())
}
pub(super) fn gemini() -> GeminiStream {
GeminiStream::new("gemini-test".to_string())
}
pub(super) fn ollama() -> OllamaStream {
OllamaStream::new("ollama-test".to_string())
}
pub(super) fn meta() -> MetaStream {
MetaStream::new("muse-spark-test".to_string())
}
pub(super) fn openai_compat() -> OpenAICompatStream {
let profile = crate::models::lookup_provider("deepinfra").expect("registry has deepinfra");
OpenAICompatStream::new(profile, "openai-test".to_string())
}
}
macro_rules! scenario {
($build:expr, $provider:literal, $file:literal) => {{
let body = fixture($provider, $file);
let whole = run($build(), &body, Delivery::Whole).await;
let dribbled = run($build(), &body, Delivery::OneByteAtATime).await;
assert_eq!(
whole, dribbled,
"{}/{}: chunk boundaries changed the result",
$provider, $file
);
whole
}};
}
macro_rules! scenario_body {
($build:expr, $provider:literal, $body:expr) => {{
let result = run($build(), &$body, Delivery::Whole).await;
(result, $provider)
}};
}
macro_rules! all_providers {
($sse:literal, $ndjson:literal) => {
vec![
scenario!(protocols::anthropic, "anthropic", $sse),
scenario!(protocols::gemini, "gemini", $sse),
scenario!(protocols::openai_compat, "openai_compat", $sse),
scenario!(protocols::meta, "meta", $sse),
scenario!(protocols::ollama, "ollama", $ndjson),
]
};
}
#[tokio::test]
async fn text_deltas_concatenate_everywhere() {
for outcome in all_providers!("text.sse", "text.ndjson") {
let outcome = outcome.expect("text scenario succeeds");
assert_eq!(outcome.streamed_text, "Hello, world");
assert_eq!(outcome.content, "Hello, world");
assert_eq!(outcome.streamed_reasoning, "");
assert_eq!(outcome.stop_reason, Some(FinishReason::Stop));
assert_eq!(outcome.usage, Some((12, 7)));
}
}
#[tokio::test]
async fn reasoning_splits_from_content_everywhere() {
for outcome in all_providers!("reasoning.sse", "reasoning.ndjson") {
let outcome = outcome.expect("reasoning scenario succeeds");
assert_eq!(outcome.streamed_reasoning, "weighing options");
assert_eq!(outcome.thinking.as_deref(), Some("weighing options"));
assert_eq!(outcome.streamed_text, "the answer");
assert_eq!(outcome.content, "the answer");
assert_eq!(outcome.usage, Some((20, 11)));
}
}
#[tokio::test]
async fn tool_calls_reassemble_everywhere() {
for outcome in all_providers!("tool_call.sse", "tool_call.ndjson") {
let outcome = outcome.expect("tool_call scenario succeeds");
let expected = vec![("read_file".to_string(), r#"{"path":"a.txt"}"#.to_string())];
assert_eq!(outcome.streamed_tool_calls, expected);
assert_eq!(outcome.response_tool_calls, expected);
assert_eq!(outcome.usage, Some((30, 15)));
}
}
#[tokio::test]
async fn a_real_truncation_survives_as_length_everywhere() {
for outcome in all_providers!("truncation.sse", "truncation.ndjson") {
let outcome = outcome.expect("truncation is a completion, not a failure");
assert_eq!(outcome.stop_reason, Some(FinishReason::Length));
assert_eq!(outcome.content, "cut off");
}
}
#[tokio::test]
async fn a_mid_stream_error_frame_is_typed_everywhere() {
for outcome in all_providers!("error_frame.sse", "error_frame.ndjson") {
assert_eq!(outcome, Err(Failure::ProviderError));
}
}
#[tokio::test]
async fn a_body_cut_before_the_terminal_frame_is_an_error_everywhere() {
for outcome in all_providers!("abnormal_close.sse", "abnormal_close.ndjson") {
assert_eq!(outcome, Err(Failure::StreamCut));
}
}
#[tokio::test]
async fn ollama_keeps_the_frame_its_body_closed_on() {
let outcome = scenario!(protocols::ollama, "ollama", "no_trailing_newline.ndjson")
.expect("the final object still counts");
assert_eq!(outcome.content, "Hello");
assert_eq!(outcome.stop_reason, Some(FinishReason::Stop));
assert_eq!(outcome.usage, Some((12, 7)));
}
fn oversized(provider: &str, name: &str, needle: &str) -> String {
let filler = "w".repeat(crate::constants::MAX_RESPONSE_CHARS + 16);
let body = fixture(provider, name);
assert!(body.contains(needle), "{provider}/{name} lacks {needle:?}");
body.replace(needle, &filler)
}
async fn oversized_everywhere(needle: &str) -> Vec<(ScenarioResult, &'static str)> {
vec![
scenario_body!(
protocols::anthropic,
"anthropic",
oversized("anthropic", "reasoning.sse", needle)
),
scenario_body!(
protocols::gemini,
"gemini",
oversized("gemini", "reasoning.sse", needle)
),
scenario_body!(
protocols::openai_compat,
"openai_compat",
oversized("openai_compat", "reasoning.sse", needle)
),
scenario_body!(
protocols::meta,
"meta",
oversized("meta", "reasoning.sse", needle)
),
scenario_body!(
protocols::ollama,
"ollama",
oversized("ollama", "reasoning.ndjson", needle)
),
]
}
#[tokio::test]
async fn a_reasoning_trace_past_the_cap_leaves_the_answer_everywhere() {
for (result, provider) in oversized_everywhere("weighing ").await {
let outcome = result
.map_err(|e| format!("{provider}: {e:?}"))
.expect("the stream decoded");
assert_eq!(
outcome.content, "the answer",
"{provider}: the answer after an oversized trace was dropped"
);
let thinking = outcome.thinking.unwrap_or_default();
assert!(
thinking.ends_with(super::accumulator::TRUNCATION_MARKER),
"{provider}: the trace was not marked as truncated"
);
}
}
#[tokio::test]
async fn content_past_the_cap_is_marked_everywhere() {
let cap = crate::constants::MAX_RESPONSE_CHARS;
for (result, provider) in oversized_everywhere("the answer").await {
let outcome = result
.map_err(|e| format!("{provider}: {e:?}"))
.expect("the stream decoded");
assert!(
outcome
.content
.ends_with(super::accumulator::TRUNCATION_MARKER),
"{provider}: oversized content was not marked"
);
assert!(
outcome.content.len() <= cap + super::accumulator::TRUNCATION_MARKER.len(),
"{provider}: content kept growing past the cap ({} bytes)",
outcome.content.len()
);
assert_eq!(
outcome.thinking.as_deref(),
Some("weighing options"),
"{provider}: the trace was disturbed by the content cap"
);
}
}
fn strip_usage(body: &str) -> String {
fn scrub(value: &mut serde_json::Value) {
match value {
serde_json::Value::Object(map) => {
for key in ["usage", "usageMetadata", "prompt_eval_count", "eval_count"] {
map.remove(key);
}
map.values_mut().for_each(scrub);
},
serde_json::Value::Array(items) => items.iter_mut().for_each(scrub),
_ => {},
}
}
body.lines()
.map(|line| {
let (prefix, json) = match line.strip_prefix("data: ") {
Some(rest) => ("data: ", rest),
None => ("", line),
};
match serde_json::from_str::<serde_json::Value>(json) {
Ok(mut value) => {
scrub(&mut value);
format!("{prefix}{value}")
},
Err(_) => line.to_string(),
}
})
.collect::<Vec<_>>()
.join("\n")
+ "\n"
}
#[tokio::test]
async fn no_usage_frame_means_no_usage_everywhere() {
let cases = vec![
scenario_body!(
protocols::anthropic,
"anthropic",
strip_usage(&fixture("anthropic", "text.sse"))
),
scenario_body!(
protocols::gemini,
"gemini",
strip_usage(&fixture("gemini", "text.sse"))
),
scenario_body!(
protocols::openai_compat,
"openai_compat",
strip_usage(&fixture("openai_compat", "text.sse"))
),
scenario_body!(
protocols::meta,
"meta",
strip_usage(&fixture("meta", "text.sse"))
),
scenario_body!(
protocols::ollama,
"ollama",
strip_usage(&fixture("ollama", "text.ndjson"))
),
];
for (result, provider) in cases {
let outcome = result
.map_err(|e| format!("{provider}: {e:?}"))
.expect("the stream decoded");
assert_eq!(outcome.usage, None, "{provider}: fabricated a usage record");
}
}