use async_trait::async_trait;
use std::sync::Arc;
use crate::error::AiMuxError;
use crate::generate::{GenerateTextOptions, GenerateTextResult, generate_text};
use crate::language_model::LanguageModel;
use crate::language_model_message::LanguageModelPrompt;
use crate::message::{MessageContent, ModelMessage, ModelPrompt};
use crate::options::{CallOptions, ToolChoice};
use crate::recording::Recording;
use crate::result::{GenerateContent, GenerateResult, StreamResult};
use crate::stream_part::StreamPart;
use crate::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage};
pub trait ReplayMatcher: Send + Sync {
fn r#match<'a>(
&self,
options: &CallOptions,
recordings: &'a [Recording],
) -> Result<&'a Recording, AiMuxError>;
}
pub struct ExactMatcher {
provider: String,
model_id: String,
}
impl ExactMatcher {
pub fn new(provider: impl Into<String>, model_id: impl Into<String>) -> Self {
Self {
provider: provider.into(),
model_id: model_id.into(),
}
}
}
impl ReplayMatcher for ExactMatcher {
fn r#match<'a>(
&self,
options: &CallOptions,
recordings: &'a [Recording],
) -> Result<&'a Recording, AiMuxError> {
let needle = canonical_call_key(options);
recordings
.iter()
.find(|r| {
r.provider.provider == self.provider
&& r.provider.model_id == self.model_id
&& canonical_keys_match(&canonical_recording_key(r), &needle)
})
.ok_or_else(|| {
AiMuxError::InvalidArgument("mock replay: no exact matching recording".into())
})
}
}
fn canonical_call_key(opts: &CallOptions) -> serde_json::Value {
serde_json::json!({
"prompt": serde_json::to_value(&opts.prompt).unwrap_or_default(),
"temperature": opts.temperature,
"max_output_tokens": opts.max_output_tokens,
"seed": opts.seed,
"response_format": serde_json::to_value(&opts.response_format).unwrap_or_default(),
"tools": serde_json::to_value(&opts.tools).unwrap_or_default(),
"tool_choice": serde_json::to_value(&opts.tool_choice).unwrap_or_default(),
"headers": serde_json::to_value(&opts.headers).unwrap_or_default(),
"provider_options": serde_json::to_value(&opts.provider_options).unwrap_or_default(),
"body_overrides": serde_json::to_value(&opts.body_overrides).unwrap_or_default(),
})
}
fn canonical_recording_key(rec: &Recording) -> serde_json::Value {
let o = &rec.input.options;
serde_json::json!({
"prompt": serde_json::to_value(&rec.input.prompt).unwrap_or_default(),
"temperature": o.get("temperature").cloned().unwrap_or_default(),
"max_output_tokens": o.get("max_output_tokens").cloned().unwrap_or_default(),
"seed": o.get("seed").cloned().unwrap_or_default(),
"response_format": o.get("response_format").cloned().unwrap_or_default(),
"tools": o.get("tools").cloned().unwrap_or_default(),
"tool_choice": o
.get("tool_choice")
.cloned()
.unwrap_or_else(|| serde_json::to_value(ToolChoice::default()).unwrap_or_default()),
"headers": o.get("headers").cloned().unwrap_or_default(),
"provider_options": o.get("provider_options").cloned().unwrap_or_default(),
"body_overrides": o.get("body_overrides").cloned().unwrap_or_default(),
})
}
fn redaction_aware_eq(rec: &serde_json::Value, req: &serde_json::Value) -> bool {
if rec.as_str() == Some("[REDACTED]") {
return !req.is_null();
}
if rec.is_null() || req.is_null() {
return rec.is_null() && req.is_null();
}
match (rec, req) {
(serde_json::Value::Object(a), serde_json::Value::Object(b)) => {
if a.len() != b.len() {
return false; }
a.iter().all(|(k, rv)| {
b.get(k)
.map(|qv| redaction_aware_eq(rv, qv))
.unwrap_or(false)
})
}
(serde_json::Value::Array(a), serde_json::Value::Array(b)) => {
a.len() == b.len() && a.iter().zip(b).all(|(x, y)| redaction_aware_eq(x, y))
}
(x, y) => x == y,
}
}
fn canonical_keys_match(rec_key: &serde_json::Value, call_key: &serde_json::Value) -> bool {
const REDACTED_FIELDS: [&str; 3] = ["headers", "provider_options", "body_overrides"];
let (ro, co) = match (rec_key.as_object(), call_key.as_object()) {
(Some(a), Some(b)) => (a, b),
_ => return rec_key == call_key,
};
if ro.len() != co.len() {
return false;
}
ro.iter().all(|(k, rv)| {
co.get(k)
.map(|qv| {
if REDACTED_FIELDS.contains(&k.as_str()) {
redaction_aware_eq(rv, qv)
} else {
rv == qv
}
})
.unwrap_or(false)
})
}
pub struct ScoreMatcher {
provider: String,
model_id: String,
}
impl ScoreMatcher {
pub fn new(provider: impl Into<String>, model_id: impl Into<String>) -> Self {
Self {
provider: provider.into(),
model_id: model_id.into(),
}
}
}
impl ReplayMatcher for ScoreMatcher {
fn r#match<'a>(
&self,
options: &CallOptions,
recordings: &'a [Recording],
) -> Result<&'a Recording, AiMuxError> {
let mut best: Option<(&Recording, u64)> = None;
for r in recordings {
if r.provider.provider != self.provider || r.provider.model_id != self.model_id {
continue;
}
let score = match_score(options, r);
if score == 0 {
continue; }
if best.map(|(_, s)| score > s).unwrap_or(true) {
best = Some((r, score));
}
}
best.map(|(r, _)| r).ok_or_else(|| {
AiMuxError::InvalidArgument("mock replay: no scoring match for this input".into())
})
}
}
fn message_text(m: &crate::language_model_message::LanguageModelPromptMessage) -> String {
match m.content.first() {
Some(crate::content::ContentPart::Text { text, .. }) => text.clone(),
_ => serde_json::to_string(&m.content).unwrap_or_default(),
}
}
fn match_score(options: &CallOptions, rec: &Recording) -> u64 {
let common = options
.prompt
.iter()
.zip(rec.input.prompt.iter())
.take_while(|(a, b)| a.role == b.role && a.content == b.content)
.count();
let min_len = options.prompt.len().min(rec.input.prompt.len());
let lcp: u64 = if min_len == 0 {
0
} else {
let lcp_idx = common.min(min_len - 1);
match (options.prompt.get(lcp_idx), rec.input.prompt.get(lcp_idx)) {
(Some(a), Some(b)) => message_text(a)
.bytes()
.zip(message_text(b).bytes())
.take_while(|(x, y)| x == y)
.count() as u64,
_ => 0,
}
};
let prompt_relevance = (common as u64) * 100 + lcp;
if prompt_relevance == 0 {
return 0; }
let mut score = prompt_relevance;
if options.temperature
== rec
.input
.options
.get("temperature")
.and_then(serde_json::Value::as_f64)
{
score += 1;
}
score
}
pub struct PrefixMatcher {
provider: String,
model_id: String,
}
impl PrefixMatcher {
pub fn new(provider: impl Into<String>, model_id: impl Into<String>) -> Self {
Self {
provider: provider.into(),
model_id: model_id.into(),
}
}
}
impl ReplayMatcher for PrefixMatcher {
fn r#match<'a>(
&self,
options: &CallOptions,
recordings: &'a [Recording],
) -> Result<&'a Recording, AiMuxError> {
let mut best: Option<(&Recording, usize)> = None;
for r in recordings {
if r.provider.provider != self.provider || r.provider.model_id != self.model_id {
continue;
}
let common = options
.prompt
.iter()
.zip(r.input.prompt.iter())
.take_while(|(a, b)| a.role == b.role && a.content == b.content)
.count();
if common < r.input.prompt.len() {
continue; }
if common > 0 && best.map(|(_, c)| common > c).unwrap_or(true) {
best = Some((r, common));
}
}
best.map(|(r, _)| r).ok_or_else(|| {
AiMuxError::InvalidArgument("mock replay: no prefix match for this input".into())
})
}
}
pub struct MockReplayModel {
provider: String,
model_id: String,
recordings: Vec<Recording>,
matcher: Arc<dyn ReplayMatcher>,
}
impl MockReplayModel {
pub fn new(
provider: impl Into<String>,
model_id: impl Into<String>,
recordings: Vec<Recording>,
) -> Self {
let provider = provider.into();
let model_id = model_id.into();
let matcher = Arc::new(ScoreMatcher::new(provider.clone(), model_id.clone()));
Self {
provider,
model_id,
recordings,
matcher,
}
}
pub fn with_matcher(
provider: impl Into<String>,
model_id: impl Into<String>,
recordings: Vec<Recording>,
matcher: Arc<dyn ReplayMatcher>,
) -> Self {
Self {
provider: provider.into(),
model_id: model_id.into(),
recordings,
matcher,
}
}
pub fn from_jsonl(path: &str) -> Result<Self, AiMuxError> {
let content = std::fs::read_to_string(path)
.map_err(|e| AiMuxError::InvalidArgument(format!("mock replay: {e}")))?;
let mut recordings = Vec::new();
for (idx, line) in content.lines().enumerate() {
if line.trim().is_empty() {
continue;
}
let rec: Recording = serde_json::from_str(line)
.map_err(|e| AiMuxError::JsonParse(format!("mock replay line {}: {e}", idx + 1)))?;
recordings.push(rec);
}
if recordings.is_empty() {
return Err(AiMuxError::InvalidArgument(
"mock replay: empty recording file".into(),
));
}
let provider = recordings[0].provider.provider.clone();
let model_id = recordings[0].provider.model_id.clone();
Ok(Self::new(provider, model_id, recordings))
}
#[must_use]
pub fn recordings(&self) -> &[Recording] {
&self.recordings
}
}
#[async_trait]
impl LanguageModel for MockReplayModel {
fn provider(&self) -> &str {
&self.provider
}
fn model_id(&self) -> &str {
&self.model_id
}
async fn do_generate(&self, options: &CallOptions) -> Result<GenerateResult, AiMuxError> {
let rec = self.matcher.r#match(options, &self.recordings)?;
rebuild_generate_result(rec)
}
async fn do_stream(&self, options: &CallOptions) -> Result<StreamResult, AiMuxError> {
let rec = self.matcher.r#match(options, &self.recordings)?;
rebuild_stream_result(rec)
}
}
fn last_response(rec: &Recording) -> Result<&crate::recording::ResponseRecord, AiMuxError> {
for e in rec.exchanges.iter().rev() {
if !e.finalized || e.error.is_some() {
continue;
}
if let Some(resp) = e.response.as_ref()
&& (200..300).contains(&resp.status)
{
return Ok(resp);
}
}
Err(AiMuxError::InvalidArgument(
"mock replay: recording has no successful (2xx, finalized, error-free) response".into(),
))
}
fn parse_finish(raw: Option<&str>) -> (FinishReasonUnified, Option<String>) {
let unified = match raw {
Some("stop") => FinishReasonUnified::Stop,
Some("length") => FinishReasonUnified::Length,
Some("content_filter") => FinishReasonUnified::ContentFilter,
Some("tool_calls") => FinishReasonUnified::ToolCalls,
Some("error") => FinishReasonUnified::Error,
_ => FinishReasonUnified::Other,
};
(unified, raw.map(std::string::ToString::to_string))
}
fn u32_from_json(v: &serde_json::Value, field: &str) -> Result<Option<u32>, AiMuxError> {
v.as_u64()
.map(|n| {
u32::try_from(n).map_err(|_| {
AiMuxError::JsonParse(format!("mock replay: usage '{field}' overflows u32: {n}"))
})
})
.transpose()
}
fn parse_usage(v: &serde_json::Value) -> Result<Usage, AiMuxError> {
let u = &v["usage"];
if u.is_null() {
return Ok(Usage::default());
}
let prompt_details = &u["prompt_tokens_details"];
let completion_details = &u["completion_tokens_details"];
Ok(Usage {
input_tokens: crate::types::TokenUsage {
total: u32_from_json(&u["prompt_tokens"], "prompt_tokens")?,
no_cache: None,
cache_read: u32_from_json(
&prompt_details["cached_tokens"],
"prompt_tokens_details.cached_tokens",
)?,
cache_write: None,
text: None,
reasoning: None,
},
output_tokens: crate::types::TokenUsage {
total: u32_from_json(&u["completion_tokens"], "completion_tokens")?,
no_cache: None,
cache_read: None,
cache_write: None,
text: None,
reasoning: u32_from_json(
&completion_details["reasoning_tokens"],
"completion_tokens_details.reasoning_tokens",
)?,
},
raw: Some(u.clone()),
})
}
fn parse_tool_arguments(args: &str) -> serde_json::Value {
serde_json::from_str(args).unwrap_or_else(|_| serde_json::Value::String(args.to_string()))
}
fn rebuild_generate_result(rec: &Recording) -> Result<GenerateResult, AiMuxError> {
let resp = last_response(rec)?;
let body = resp
.body
.as_deref()
.ok_or_else(|| AiMuxError::InvalidArgument("mock replay: response has no body".into()))?;
let v: serde_json::Value = serde_json::from_str(body)
.map_err(|e| AiMuxError::JsonParse(format!("mock replay: {e}")))?;
let message = v
.get("choices")
.and_then(|c| c.get(0))
.and_then(|choice| choice.get("message"))
.ok_or_else(|| {
AiMuxError::UnsupportedFunctionality(
"mock replay: response is not an OpenAI chat.completions body (no choices[0].message)"
.into(),
)
})?;
let mut content = Vec::new();
if let Some(text) = message.get("content").and_then(|c| c.as_str())
&& !text.is_empty()
{
content.push(GenerateContent::Text {
text: text.to_string(),
provider_metadata: None,
});
}
if let Some(tool_calls) = message.get("tool_calls").and_then(|t| t.as_array()) {
for tc in tool_calls {
let tool_call_id = tc
.get("id")
.and_then(|x| x.as_str())
.unwrap_or_default()
.to_string();
let function = tc.get("function").cloned().unwrap_or_default();
let tool_name = function
.get("name")
.and_then(|x| x.as_str())
.unwrap_or_default()
.to_string();
let args_str = function
.get("arguments")
.and_then(|x| x.as_str())
.unwrap_or("");
let input = parse_tool_arguments(args_str);
content.push(GenerateContent::ToolCall {
tool_call_id,
tool_name,
input,
provider_executed: None,
dynamic: None,
thought_signature: None,
provider_metadata: None,
});
}
}
let raw_finish = v["choices"][0]["finish_reason"]
.as_str()
.map(std::string::ToString::to_string);
let (unified, raw) = parse_finish(raw_finish.as_deref());
Ok(GenerateResult {
content,
finish_reason: FinishReason { unified, raw },
usage: parse_usage(&v)?,
warnings: Vec::new(),
provider_metadata: None,
response: ResponseMetadata {
id: v["id"].as_str().map(std::string::ToString::to_string),
timestamp: None,
model_id: v["model"].as_str().map(std::string::ToString::to_string),
},
request_body: None,
response_headers: None,
})
}
struct ToolCallAccumulator {
id: String,
name: String,
arguments: String,
}
fn rebuild_stream_result(rec: &Recording) -> Result<StreamResult, AiMuxError> {
let resp = last_response(rec)?;
let body = resp
.body
.as_deref()
.ok_or_else(|| AiMuxError::InvalidArgument("mock replay: response has no body".into()))?;
let mut parts: Vec<Result<StreamPart, AiMuxError>> = Vec::new();
parts.push(Ok(StreamPart::StreamStart { warnings: vec![] }));
let text_id = "0".to_string();
let reasoning_id = "reasoning-0".to_string();
let mut text_started = false;
let mut reasoning_started = false;
let mut response_meta_emitted = false;
let mut final_usage = Usage::default();
let mut final_finish: Option<FinishReason> = None;
let mut tool_calls: std::collections::HashMap<usize, ToolCallAccumulator> =
std::collections::HashMap::new();
let mut tool_order: Vec<usize> = Vec::new();
let mut saw_openai = false;
for block in body.split("\n\n") {
let line = block.trim();
let data = line.strip_prefix("data:").unwrap_or(line).trim();
if data.is_empty() || data == "[DONE]" {
continue;
}
let Ok(v) = serde_json::from_str::<serde_json::Value>(data) else {
continue;
};
let choices_arr = v.get("choices").and_then(|c| c.as_array());
if choices_arr.is_some() {
saw_openai = true;
}
if !response_meta_emitted
&& (v.get("id").and_then(|x| x.as_str()).is_some()
|| v.get("model").and_then(|x| x.as_str()).is_some())
{
response_meta_emitted = true;
parts.push(Ok(StreamPart::ResponseMetadata {
id: v["id"].as_str().map(std::string::ToString::to_string),
timestamp: None,
model_id: v["model"].as_str().map(std::string::ToString::to_string),
}));
}
if v.get("usage").is_some_and(|u| !u.is_null()) {
final_usage = parse_usage(&v)?;
}
if let Some(choices) = choices_arr {
for choice in choices {
let delta = &choice["delta"];
let reasoning_delta = delta
.get("reasoning_content")
.and_then(|x| x.as_str())
.or_else(|| delta.get("reasoning").and_then(|x| x.as_str()));
if let Some(reasoning) = reasoning_delta
&& !reasoning.is_empty()
{
if !reasoning_started {
reasoning_started = true;
parts.push(Ok(StreamPart::ReasoningStart {
id: reasoning_id.clone(),
provider_metadata: None,
}));
}
parts.push(Ok(StreamPart::ReasoningDelta {
id: reasoning_id.clone(),
delta: reasoning.to_string(),
provider_metadata: None,
}));
}
if let Some(content) = delta.get("content").and_then(|x| x.as_str())
&& !content.is_empty()
{
if reasoning_started {
parts.push(Ok(StreamPart::ReasoningEnd {
id: reasoning_id.clone(),
provider_metadata: None,
}));
reasoning_started = false;
}
if !text_started {
text_started = true;
parts.push(Ok(StreamPart::TextStart {
id: text_id.clone(),
provider_metadata: None,
}));
}
parts.push(Ok(StreamPart::TextDelta {
id: text_id.clone(),
delta: content.to_string(),
provider_metadata: None,
}));
}
if let Some(tcs) = delta.get("tool_calls").and_then(|t| t.as_array()) {
if reasoning_started {
parts.push(Ok(StreamPart::ReasoningEnd {
id: reasoning_id.clone(),
provider_metadata: None,
}));
reasoning_started = false;
}
for dtc in tcs {
let idx = dtc
.get("index")
.and_then(serde_json::Value::as_u64)
.map(|n| n as usize)
.unwrap_or(0);
let func = dtc.get("function").cloned().unwrap_or_default();
let is_new = !tool_calls.contains_key(&idx);
if is_new {
let id = dtc
.get("id")
.and_then(|x| x.as_str())
.unwrap_or_default()
.to_string();
let name = func
.get("name")
.and_then(|x| x.as_str())
.unwrap_or_default()
.to_string();
tool_calls.insert(
idx,
ToolCallAccumulator {
id: id.clone(),
name: name.clone(),
arguments: String::new(),
},
);
tool_order.push(idx);
parts.push(Ok(StreamPart::ToolInputStart {
id,
tool_name: name,
provider_executed: None,
dynamic: None,
title: None,
provider_metadata: None,
}));
}
if let Some(args) = func.get("arguments").and_then(|x| x.as_str())
&& (!is_new || !args.is_empty())
&& let Some(acc) = tool_calls.get_mut(&idx)
{
acc.arguments.push_str(args);
parts.push(Ok(StreamPart::ToolInputDelta {
id: acc.id.clone(),
delta: args.to_string(),
provider_metadata: None,
}));
}
}
}
if let Some(fr) = choice.get("finish_reason").and_then(|x| x.as_str()) {
if reasoning_started {
parts.push(Ok(StreamPart::ReasoningEnd {
id: reasoning_id.clone(),
provider_metadata: None,
}));
reasoning_started = false;
}
if text_started {
parts.push(Ok(StreamPart::TextEnd {
id: text_id.clone(),
provider_metadata: None,
}));
text_started = false;
}
for &i in &tool_order {
if let Some(acc) = tool_calls.get(&i) {
parts.push(Ok(StreamPart::ToolInputEnd {
id: acc.id.clone(),
provider_metadata: None,
}));
let input = parse_tool_arguments(&acc.arguments);
parts.push(Ok(StreamPart::ToolCall {
tool_call_id: acc.id.clone(),
tool_name: acc.name.clone(),
input,
provider_executed: None,
dynamic: None,
thought_signature: None,
provider_metadata: None,
}));
}
}
tool_calls.clear();
tool_order.clear();
let (unified, raw) = parse_finish(Some(fr));
final_finish = Some(FinishReason { unified, raw });
}
}
}
}
if reasoning_started {
parts.push(Ok(StreamPart::ReasoningEnd {
id: reasoning_id,
provider_metadata: None,
}));
}
if text_started {
parts.push(Ok(StreamPart::TextEnd {
id: text_id,
provider_metadata: None,
}));
}
for &i in &tool_order {
if let Some(acc) = tool_calls.get(&i) {
parts.push(Ok(StreamPart::ToolInputEnd {
id: acc.id.clone(),
provider_metadata: None,
}));
let input = parse_tool_arguments(&acc.arguments);
parts.push(Ok(StreamPart::ToolCall {
tool_call_id: acc.id.clone(),
tool_name: acc.name.clone(),
input,
provider_executed: None,
dynamic: None,
thought_signature: None,
provider_metadata: None,
}));
}
}
if !saw_openai {
return Err(AiMuxError::UnsupportedFunctionality(
"mock replay: response is not an OpenAI chat.completions SSE stream".into(),
));
}
parts.push(Ok(StreamPart::Finish {
finish_reason: final_finish.unwrap_or(FinishReason {
unified: FinishReasonUnified::Stop,
raw: None,
}),
usage: final_usage,
provider_metadata: None,
}));
Ok(StreamResult {
stream: Box::pin(futures::stream::iter(parts)),
request_body: None,
response_headers: None,
})
}
#[derive(Debug, Clone, Default)]
pub struct ReplayOverrides {
pub prompt: Option<LanguageModelPrompt>,
pub temperature: Option<f64>,
pub max_output_tokens: Option<u32>,
}
pub async fn replay_with_model(
recording: &Recording,
model: &dyn LanguageModel,
overrides: Option<&ReplayOverrides>,
) -> Result<GenerateTextResult, AiMuxError> {
let mut call_options: CallOptions = serde_json::from_value(recording.input.options.clone())
.map_err(|e| AiMuxError::JsonParse(format!("mock replay: input options invalid: {e}")))?;
if let Some(o) = overrides {
if let Some(prompt) = &o.prompt {
call_options.prompt = prompt.clone();
}
if o.temperature.is_some() {
call_options.temperature = o.temperature;
}
if o.max_output_tokens.is_some() {
call_options.max_output_tokens = o.max_output_tokens;
}
}
let prompt = model_prompt_from_lm(&call_options.prompt);
let options = generate_options_from_call_options(call_options);
generate_text(model, prompt, options).await
}
fn model_prompt_from_lm(prompt: &LanguageModelPrompt) -> ModelPrompt {
ModelPrompt::Messages(
prompt
.iter()
.map(|m| ModelMessage {
role: m.role,
content: MessageContent::Parts(m.content.clone()),
})
.collect(),
)
}
fn generate_options_from_call_options(o: CallOptions) -> GenerateTextOptions {
GenerateTextOptions {
max_output_tokens: o.max_output_tokens,
temperature: o.temperature,
stop_sequences: o.stop_sequences,
top_p: o.top_p,
top_k: o.top_k,
presence_penalty: o.presence_penalty,
frequency_penalty: o.frequency_penalty,
response_format: o.response_format,
seed: o.seed,
tools: o.tools,
tool_choice: if o.tool_choice == crate::tool::ToolChoice::Auto {
None
} else {
Some(o.tool_choice)
},
headers: o.headers,
provider_options: o.provider_options,
reasoning: o.reasoning,
instructions: None,
body_overrides: o.body_overrides,
max_retries: o.max_retries,
timeout: o.timeout,
session_id: o.session_id,
abort_signal: None,
include_raw_chunks: o.include_raw_chunks,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::generate::GenerateTextOptions;
use crate::recording::{
HttpExchange, HttpRecord, InputRecord, ProviderRecord, ResponseRecord, TimingRecord,
};
use futures::StreamExt;
fn sample_options(text: &str, temperature: Option<f64>) -> CallOptions {
let prompt = vec![crate::language_model_message::LanguageModelPromptMessage {
role: crate::message::Role::User,
content: vec![crate::content::ContentPart::Text {
text: text.to_string(),
provider_options: None,
}],
provider_options: None,
}];
GenerateTextOptions {
temperature,
..Default::default()
}
.into_call_options(prompt)
}
fn openai_recording(trace_id: &str, prompt_text: &str, reply: &str, finish: &str) -> Recording {
let input_prompt: crate::language_model_message::LanguageModelPrompt =
vec![crate::language_model_message::LanguageModelPromptMessage {
role: crate::message::Role::User,
content: vec![crate::content::ContentPart::Text {
text: prompt_text.to_string(),
provider_options: None,
}],
provider_options: None,
}];
let body = serde_json::json!({
"id": "chatcmpl-mock",
"object": "chat.completion",
"model": "gpt-4o",
"choices": [{
"index": 0,
"message": { "role": "assistant", "content": reply },
"finish_reason": finish
}],
"usage": { "prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8 }
});
Recording {
schema: crate::recording::RECORDING_SCHEMA,
call_id: trace_id.to_string(),
recorded_at: "2026-08-06T00:00:00Z".to_string(),
input: InputRecord {
prompt: input_prompt,
options: serde_json::json!({ "temperature": 0.7 }),
},
provider: ProviderRecord {
provider: "openai".into(),
model_id: "gpt-4o".into(),
base_url: None,
api_key_source: "none".into(),
profile: None,
provider_options: None,
},
exchanges: vec![HttpExchange {
attempt: 0,
request: HttpRecord {
method: "post".into(),
url: "https://api.openai.com/v1/chat/completions".into(),
headers: vec![],
body: Some("{}".into()),
},
response: Some(ResponseRecord {
status: 200,
headers: vec![],
body: Some(body.to_string()),
stream_chunks: None,
ttfb_ms: None,
}),
timing: TimingRecord {
latency_ms: 10,
ttfb_ms: None,
},
error: None,
finalized: true,
}],
outcome: crate::recording::OutcomeRecord {
status: crate::recording::OutcomeStatus::Success,
finish_reason: Some("stop".into()),
error: None,
usage: None,
},
complete: true,
transport_closed: true,
session_id: None,
step: None,
}
}
#[test]
fn exact_matcher_hits_only_identical_prompt() {
let recs = [openai_recording("t1", "ping", "pong", "stop")];
let matcher = ExactMatcher::new("openai", "gpt-4o");
assert!(
matcher
.r#match(&sample_options("ping", Some(0.7)), &recs)
.is_ok()
);
assert!(
matcher
.r#match(&sample_options("pong", Some(0.7)), &recs)
.is_err()
);
}
#[test]
fn exact_matcher_different_temperature_misses() {
let recs = [openai_recording("t1", "ping", "pong", "stop")]; let matcher = ExactMatcher::new("openai", "gpt-4o");
assert!(
matcher
.r#match(&sample_options("ping", Some(0.1)), &recs)
.is_err()
);
assert!(
matcher
.r#match(&sample_options("ping", Some(0.7)), &recs)
.is_ok()
);
}
#[test]
fn exact_matcher_different_options_misses() {
let mut rec = openai_recording("t1", "ping", "pong", "stop");
let mut call = sample_options("ping", Some(0.7));
call.max_output_tokens = Some(128);
rec.input.options = serde_json::to_value(&call).unwrap();
let recs = [rec];
let matcher = ExactMatcher::new("openai", "gpt-4o");
let mut req = sample_options("ping", Some(0.7));
req.max_output_tokens = Some(256);
assert!(matcher.r#match(&req, &recs).is_err());
let mut req2 = sample_options("ping", Some(0.7));
req2.max_output_tokens = Some(128);
assert!(matcher.r#match(&req2, &recs).is_ok());
}
#[test]
fn exact_matcher_different_headers_misses() {
let mut rec = openai_recording("t1", "ping", "pong", "stop");
let mut call = sample_options("ping", Some(0.7));
call.headers = Some([("x-custom".into(), "a".into())].into_iter().collect());
rec.input.options = serde_json::to_value(&call).unwrap();
let recs = [rec];
let matcher = ExactMatcher::new("openai", "gpt-4o");
let mut req = sample_options("ping", Some(0.7));
req.headers = Some([("x-custom".into(), "b".into())].into_iter().collect());
assert!(matcher.r#match(&req, &recs).is_err());
let mut req2 = sample_options("ping", Some(0.7));
req2.headers = Some([("x-custom".into(), "a".into())].into_iter().collect());
assert!(matcher.r#match(&req2, &recs).is_ok());
let mut req3 = sample_options("ping", Some(0.7));
req3.headers = Some(
[
("x-custom".into(), "a".into()),
("x-extra".into(), "z".into()),
]
.into_iter()
.collect(),
);
assert!(matcher.r#match(&req3, &recs).is_err());
}
#[test]
fn exact_matcher_different_provider_options_misses() {
let mut rec = openai_recording("t1", "ping", "pong", "stop");
let mut call = sample_options("ping", Some(0.7));
call.provider_options = Some(
[("openai".into(), serde_json::json!({ "foo": 1 }))]
.into_iter()
.collect(),
);
rec.input.options = serde_json::to_value(&call).unwrap();
let recs = [rec];
let matcher = ExactMatcher::new("openai", "gpt-4o");
let mut req = sample_options("ping", Some(0.7));
req.provider_options = Some(
[("openai".into(), serde_json::json!({ "foo": 2 }))]
.into_iter()
.collect(),
);
assert!(matcher.r#match(&req, &recs).is_err());
let mut req2 = sample_options("ping", Some(0.7));
req2.provider_options = Some(
[("openai".into(), serde_json::json!({ "foo": 1 }))]
.into_iter()
.collect(),
);
assert!(matcher.r#match(&req2, &recs).is_ok());
}
#[test]
fn exact_matcher_redacted_values_match_any() {
let mut rec = openai_recording("t1", "ping", "pong", "stop");
let mut call = sample_options("ping", Some(0.7));
call.headers = Some(
[
("authorization".into(), "sk-secret".into()),
("content-type".into(), "application/json".into()),
]
.into_iter()
.collect(),
);
let mut opts_val = serde_json::to_value(&call).unwrap();
if let Some(h) = opts_val.get_mut("headers") {
*h = crate::recording::redact_json(h.clone());
}
rec.input.options = opts_val;
assert_eq!(
rec.input.options["headers"]["authorization"],
serde_json::json!("[REDACTED]")
);
let recs = [rec];
let matcher = ExactMatcher::new("openai", "gpt-4o");
let mut req = sample_options("ping", Some(0.7));
req.headers = Some(
[
("authorization".into(), "sk-other".into()),
("content-type".into(), "application/json".into()),
]
.into_iter()
.collect(),
);
assert!(matcher.r#match(&req, &recs).is_ok());
let mut req2 = sample_options("ping", Some(0.7));
req2.headers = Some(
[
("authorization".into(), "sk-other".into()),
("content-type".into(), "text/plain".into()),
]
.into_iter()
.collect(),
);
assert!(matcher.r#match(&req2, &recs).is_err());
let mut req3 = sample_options("ping", Some(0.7));
req3.headers = Some(
[("content-type".into(), "application/json".into())]
.into_iter()
.collect(),
);
assert!(matcher.r#match(&req3, &recs).is_err());
let req4 = sample_options("ping", Some(0.7));
assert!(matcher.r#match(&req4, &recs).is_err());
}
#[test]
fn score_matcher_prefers_longer_prefix_and_ties_to_first() {
let recs = [
openai_recording("t1", "hello", "a", "stop"),
openai_recording("t2", "hello world", "b", "stop"),
];
let matcher = ScoreMatcher::new("openai", "gpt-4o");
let hit = matcher
.r#match(&sample_options("hello world!", None), &recs)
.unwrap();
assert_eq!(hit.call_id, "t2");
let hit = matcher
.r#match(&sample_options("hello", None), &recs)
.unwrap();
assert_eq!(hit.call_id, "t1");
let hit = matcher
.r#match(&sample_options("hello", Some(0.7)), &recs)
.unwrap();
assert_eq!(hit.call_id, "t1");
}
#[test]
fn match_score_empty_prompt_no_underflow() {
let rec = openai_recording("t1", "hello", "a", "stop");
let empty_req = GenerateTextOptions::default().into_call_options(vec![]);
assert_eq!(match_score(&empty_req, &rec), 0);
let mut rec_empty = rec.clone();
rec_empty.input.prompt = vec![];
assert_eq!(match_score(&sample_options("hello", None), &rec_empty), 0);
assert_eq!(match_score(&empty_req, &rec_empty), 0);
}
#[test]
fn score_matcher_unrelated_prompt_with_matching_temperature_misses() {
let mut rec = openai_recording("t1", "alpha", "a", "stop");
rec.input.options = serde_json::json!({}); let recs = [rec];
let matcher = ScoreMatcher::new("openai", "gpt-4o");
assert!(
matcher
.r#match(&sample_options("zzzzz", None), &recs)
.is_err()
);
}
#[test]
fn score_matcher_related_prompt_with_matching_temperature_hits() {
let recs = [openai_recording("t1", "hello", "a", "stop")]; let matcher = ScoreMatcher::new("openai", "gpt-4o");
let hit = matcher
.r#match(&sample_options("hello", Some(0.7)), &recs)
.unwrap();
assert_eq!(hit.call_id, "t1");
}
#[test]
fn prefix_matcher_matches_message_prefix() {
let rec = openai_recording("t1", "hello", "a", "stop"); let matcher = PrefixMatcher::new("openai", "gpt-4o");
let prompt = vec![
crate::language_model_message::LanguageModelPromptMessage {
role: crate::message::Role::User,
content: vec![crate::content::ContentPart::Text {
text: "hello".into(),
provider_options: None,
}],
provider_options: None,
},
crate::language_model_message::LanguageModelPromptMessage {
role: crate::message::Role::User,
content: vec![crate::content::ContentPart::Text {
text: "world".into(),
provider_options: None,
}],
provider_options: None,
},
];
let opts = GenerateTextOptions::default().into_call_options(prompt);
let recs = [rec.clone()];
let hit = matcher.r#match(&opts, &recs).unwrap();
assert_eq!(hit.call_id, "t1");
let opts2 = GenerateTextOptions::default().into_call_options(vec![
crate::language_model_message::LanguageModelPromptMessage {
role: crate::message::Role::User,
content: vec![crate::content::ContentPart::Text {
text: "hi".into(),
provider_options: None,
}],
provider_options: None,
},
]);
assert!(matcher.r#match(&opts2, &[rec]).is_err());
}
#[test]
fn prefix_matcher_prefers_longest_prefix() {
let mut rec_a = openai_recording("ta", "hello", "a", "stop");
let mut rec_b = openai_recording("tb", "hello world", "b", "stop");
let mk = |text: &str| crate::language_model_message::LanguageModelPromptMessage {
role: crate::message::Role::User,
content: vec![crate::content::ContentPart::Text {
text: text.into(),
provider_options: None,
}],
provider_options: None,
};
rec_a.input.prompt = vec![mk("hello")];
rec_b.input.prompt = vec![mk("hello"), mk("world")];
let matcher = PrefixMatcher::new("openai", "gpt-4o");
let opts = GenerateTextOptions::default().into_call_options(vec![
mk("hello"),
mk("world"),
mk("!"),
]);
let recs = [rec_a, rec_b];
let hit = matcher.r#match(&opts, &recs).unwrap();
assert_eq!(hit.call_id, "tb");
}
#[test]
fn unmatched_returns_clear_error() {
let recs = [openai_recording("t1", "ping", "pong", "stop")];
let model = MockReplayModel::new("openai", "gpt-4o", vec![recs[0].clone()]);
let rt = tokio::runtime::Runtime::new().unwrap();
let err = rt.block_on(async {
model
.do_generate(&sample_options("nope", None))
.await
.unwrap_err()
});
assert!(err.to_string().contains("mock replay"), "{err}");
}
#[test]
fn mock_model_rebuilds_generate_result_from_openai_body() {
let rec = openai_recording("t1", "ping", "pong", "stop");
let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
let rt = tokio::runtime::Runtime::new().unwrap();
let result = rt.block_on(async {
model
.do_generate(&sample_options("ping", None))
.await
.unwrap()
});
let Some(GenerateContent::Text { text, .. }) = result.content.first() else {
panic!("expected text content");
};
assert_eq!(text, "pong");
assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop);
assert_eq!(result.usage.input_tokens.total, Some(5));
assert_eq!(result.response.model_id.as_deref(), Some("gpt-4o"));
}
#[test]
fn mock_model_usage_overflow_errors() {
let mut rec = openai_recording("t1", "ping", "pong", "stop");
let body = serde_json::json!({
"id": "chatcmpl-mock",
"model": "gpt-4o",
"choices": [{
"index": 0,
"message": { "role": "assistant", "content": "pong" },
"finish_reason": "stop"
}],
"usage": { "prompt_tokens": 5_000_000_000_u64, "completion_tokens": 3, "total_tokens": 8 }
});
rec.exchanges[0].response.as_mut().unwrap().body = Some(body.to_string());
let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
let rt = tokio::runtime::Runtime::new().unwrap();
let err = rt.block_on(async {
model
.do_generate(&sample_options("ping", None))
.await
.unwrap_err()
});
assert!(matches!(err, AiMuxError::JsonParse(_)), "{err}");
assert!(err.to_string().contains("overflows u32"), "{err}");
}
#[test]
fn mock_model_rebuilds_stream_from_sse() {
let sse = "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\",\"content\":\"hel\"}}]}\n\n\
data: {\"choices\":[{\"delta\":{\"content\":\"lo\"}}]}\n\n\
data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n\
data: [DONE]\n\n";
let mut rec = openai_recording("t1", "ping", "x", "stop");
rec.exchanges[0].response.as_mut().unwrap().body = Some(sse.to_string());
rec.exchanges[0].response.as_mut().unwrap().stream_chunks = Some(4);
let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
let rt = tokio::runtime::Runtime::new().unwrap();
let parts: Vec<StreamPart> = rt.block_on(async {
model
.do_stream(&sample_options("ping", None))
.await
.unwrap()
.stream
.collect::<Vec<_>>()
.await
.into_iter()
.map(|p| p.unwrap())
.collect()
});
assert!(matches!(parts[0], StreamPart::StreamStart { .. }));
let deltas: Vec<String> = parts
.iter()
.filter_map(|p| match p {
StreamPart::TextDelta { delta, .. } => Some(delta.clone()),
_ => None,
})
.collect();
assert_eq!(deltas, vec!["hel", "lo"]);
assert!(parts.iter().any(|p| matches!(p, StreamPart::Finish { .. })));
}
#[test]
fn non_openai_body_returns_unsupported() {
let mut rec = openai_recording("t1", "ping", "x", "stop");
rec.exchanges[0].response.as_mut().unwrap().body = Some("{\"foo\":\"bar\"}".to_string()); let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
let rt = tokio::runtime::Runtime::new().unwrap();
let err = rt.block_on(async {
model
.do_generate(&sample_options("ping", None))
.await
.unwrap_err()
});
assert!(
matches!(err, AiMuxError::UnsupportedFunctionality(_)),
"{err}"
);
}
#[test]
fn non_openai_stream_returns_unsupported() {
let mut rec = openai_recording("t1", "ping", "x", "stop");
rec.exchanges[0].response.as_mut().unwrap().body =
Some("event: message_start\ndata: {\"type\":\"message_start\"}\n\n".to_string());
rec.exchanges[0].response.as_mut().unwrap().stream_chunks = Some(1);
let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
let rt = tokio::runtime::Runtime::new().unwrap();
let err = rt.block_on(async {
model
.do_stream(&sample_options("ping", None))
.await
.unwrap_err()
});
assert!(
matches!(err, AiMuxError::UnsupportedFunctionality(_)),
"{err}"
);
}
#[test]
fn mock_model_rebuilds_tool_calls_from_openai_body() {
let mut rec = openai_recording("t1", "ping", "x", "tool_calls");
let body = serde_json::json!({
"id": "chatcmpl-mock",
"model": "gpt-4o",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": null,
"tool_calls": [{
"id": "call_abc",
"type": "function",
"function": { "name": "get_weather", "arguments": "{\"city\":\"SF\"}" }
}]
},
"finish_reason": "tool_calls"
}],
"usage": { "prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8 }
});
rec.exchanges[0].response.as_mut().unwrap().body = Some(body.to_string());
let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
let rt = tokio::runtime::Runtime::new().unwrap();
let result = rt.block_on(async {
model
.do_generate(&sample_options("ping", None))
.await
.unwrap()
});
assert_eq!(result.content.len(), 1);
let Some(GenerateContent::ToolCall {
tool_call_id,
tool_name,
input,
..
}) = result.content.first()
else {
panic!("expected ToolCall, got {:?}", result.content);
};
assert_eq!(tool_call_id, "call_abc");
assert_eq!(tool_name, "get_weather");
assert_eq!(input, &serde_json::json!({ "city": "SF" }));
assert_eq!(result.finish_reason.unified, FinishReasonUnified::ToolCalls);
}
#[test]
fn mock_model_tool_call_bad_json_arguments_falls_back_to_string() {
let mut rec = openai_recording("t1", "ping", "x", "tool_calls");
let body = serde_json::json!({
"id": "chatcmpl-mock",
"model": "gpt-4o",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": null,
"tool_calls": [{
"id": "call_1",
"type": "function",
"function": { "name": "f", "arguments": "not-json{" }
}]
},
"finish_reason": "tool_calls"
}]
});
rec.exchanges[0].response.as_mut().unwrap().body = Some(body.to_string());
let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
let rt = tokio::runtime::Runtime::new().unwrap();
let result = rt.block_on(async {
model
.do_generate(&sample_options("ping", None))
.await
.unwrap()
});
let Some(GenerateContent::ToolCall { input, .. }) = result.content.first() else {
panic!("expected ToolCall");
};
assert_eq!(input, &serde_json::json!("not-json{"));
}
#[test]
fn mock_model_rebuilds_stream_tool_calls() {
let chunk1 = serde_json::json!({
"id": "chatcmpl-1", "model": "gpt-4o",
"choices": [{"index": 0, "delta": {"role": "assistant", "tool_calls": [
{"index": 0, "id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": ""}}
]}}]
});
let chunk2 = serde_json::json!({
"choices": [{"index": 0, "delta": {"tool_calls": [
{"index": 0, "function": {"arguments": "{\"city\":\"SF\"}"}}
]}}]
});
let chunk3 = serde_json::json!({"choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]});
let sse = format!("data: {chunk1}\n\ndata: {chunk2}\n\ndata: {chunk3}\n\ndata: [DONE]\n\n");
let mut rec = openai_recording("t1", "ping", "x", "stop");
rec.exchanges[0].response.as_mut().unwrap().body = Some(sse);
rec.exchanges[0].response.as_mut().unwrap().stream_chunks = Some(4);
let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
let rt = tokio::runtime::Runtime::new().unwrap();
let parts: Vec<StreamPart> = rt.block_on(async {
model
.do_stream(&sample_options("ping", None))
.await
.unwrap()
.stream
.collect::<Vec<_>>()
.await
.into_iter()
.map(|p| p.unwrap())
.collect()
});
assert!(parts.iter().any(|p| matches!(
p,
StreamPart::ResponseMetadata { id, model_id, .. }
if id.as_deref() == Some("chatcmpl-1") && model_id.as_deref() == Some("gpt-4o")
)));
assert!(parts.iter().any(|p| matches!(
p,
StreamPart::ToolInputStart { tool_name, .. } if tool_name == "get_weather"
)));
let deltas: Vec<String> = parts
.iter()
.filter_map(|p| match p {
StreamPart::ToolInputDelta { delta, .. } => Some(delta.clone()),
_ => None,
})
.collect();
assert_eq!(deltas, vec!["{\"city\":\"SF\"}"]);
assert!(
parts
.iter()
.any(|p| matches!(p, StreamPart::ToolInputEnd { .. }))
);
let calls: Vec<_> = parts
.iter()
.filter_map(|p| match p {
StreamPart::ToolCall {
tool_call_id,
tool_name,
input,
..
} => Some((tool_call_id.clone(), tool_name.clone(), input.clone())),
_ => None,
})
.collect();
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].0, "call_1");
assert_eq!(calls[0].1, "get_weather");
assert_eq!(calls[0].2, serde_json::json!({"city":"SF"}));
assert!(parts.iter().any(|p| matches!(
p,
StreamPart::Finish { finish_reason, .. }
if finish_reason.unified == FinishReasonUnified::ToolCalls
)));
}
#[test]
fn mock_model_rebuilds_stream_metadata_and_usage() {
let chunk1 = serde_json::json!({
"id": "chatcmpl-2", "model": "gpt-4o",
"choices": [{"index": 0, "delta": {"content": "hi"}}]
});
let chunk2 = serde_json::json!({
"id": "chatcmpl-2", "choices": [],
"usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9}
});
let sse = format!("data: {chunk1}\n\ndata: {chunk2}\n\ndata: [DONE]\n\n");
let mut rec = openai_recording("t1", "ping", "x", "stop");
rec.exchanges[0].response.as_mut().unwrap().body = Some(sse);
rec.exchanges[0].response.as_mut().unwrap().stream_chunks = Some(3);
let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
let rt = tokio::runtime::Runtime::new().unwrap();
let parts: Vec<StreamPart> = rt.block_on(async {
model
.do_stream(&sample_options("ping", None))
.await
.unwrap()
.stream
.collect::<Vec<_>>()
.await
.into_iter()
.map(|p| p.unwrap())
.collect()
});
assert!(parts.iter().any(|p| matches!(
p,
StreamPart::ResponseMetadata { id, model_id, .. }
if id.as_deref() == Some("chatcmpl-2") && model_id.as_deref() == Some("gpt-4o")
)));
let deltas: Vec<String> = parts
.iter()
.filter_map(|p| match p {
StreamPart::TextDelta { delta, .. } => Some(delta.clone()),
_ => None,
})
.collect();
assert_eq!(deltas, vec!["hi"]);
let finish = parts
.iter()
.find_map(|p| match p {
StreamPart::Finish { usage, .. } => Some(usage.clone()),
_ => None,
})
.unwrap();
assert_eq!(finish.input_tokens.total, Some(7));
assert_eq!(finish.output_tokens.total, Some(2));
}
#[test]
fn last_response_skips_failed_attempt() {
let mut rec = openai_recording("t1", "ping", "pong", "stop");
rec.exchanges[0].response.as_mut().unwrap().status = 500;
rec.exchanges[0].attempt = 0;
let success_body = serde_json::json!({
"id": "chatcmpl-ok", "model": "gpt-4o",
"choices": [{"index":0,"message":{"role":"assistant","content":"pong"},"finish_reason":"stop"}],
"usage": {"prompt_tokens":5,"completion_tokens":3,"total_tokens":8}
});
rec.exchanges.push(HttpExchange {
attempt: 1,
request: HttpRecord {
method: "post".into(),
url: "https://api.openai.com/v1/chat/completions".into(),
headers: vec![],
body: Some("{}".into()),
},
response: Some(ResponseRecord {
status: 200,
headers: vec![],
body: Some(success_body.to_string()),
stream_chunks: None,
ttfb_ms: None,
}),
timing: TimingRecord {
latency_ms: 10,
ttfb_ms: None,
},
error: None,
finalized: true,
});
let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
let rt = tokio::runtime::Runtime::new().unwrap();
let result = rt.block_on(async {
model
.do_generate(&sample_options("ping", None))
.await
.unwrap()
});
let Some(GenerateContent::Text { text, .. }) = result.content.first() else {
panic!("expected text");
};
assert_eq!(text, "pong");
assert_eq!(result.response.id.as_deref(), Some("chatcmpl-ok"));
}
#[test]
fn last_response_no_successful_response_errors() {
let mut rec = openai_recording("t1", "ping", "pong", "stop");
rec.exchanges[0].response.as_mut().unwrap().status = 500;
let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
let rt = tokio::runtime::Runtime::new().unwrap();
let err = rt.block_on(async {
model
.do_generate(&sample_options("ping", None))
.await
.unwrap_err()
});
assert!(err.to_string().contains("no successful"), "{err}");
}
#[test]
fn from_jsonl_loads_recordings() {
let dir = std::env::temp_dir().join(format!("aimux-replay-jsonl-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
let path = dir.join("rec.jsonl");
let rec = openai_recording("t1", "ping", "pong", "stop");
std::fs::write(&path, serde_json::to_string(&rec).unwrap() + "\n").unwrap();
let model = MockReplayModel::from_jsonl(path.to_str().unwrap()).unwrap();
assert_eq!(model.provider(), "openai");
assert_eq!(model.model_id(), "gpt-4o");
assert_eq!(model.recordings().len(), 1);
let _ = std::fs::remove_dir_all(&dir);
}
#[derive(Clone)]
struct EchoModel {
provider: &'static str,
model_id: &'static str,
}
impl EchoModel {
fn new() -> Self {
Self {
provider: "openai",
model_id: "gpt-4o",
}
}
}
#[async_trait]
impl LanguageModel for EchoModel {
fn provider(&self) -> &str {
self.provider
}
fn model_id(&self) -> &str {
self.model_id
}
async fn do_generate(&self, options: &CallOptions) -> Result<GenerateResult, AiMuxError> {
let text = options
.prompt
.iter()
.filter_map(|m| match m.content.first() {
Some(crate::content::ContentPart::Text { text, .. }) => Some(text.clone()),
_ => None,
})
.collect::<Vec<_>>()
.join(" ");
Ok(GenerateResult {
content: vec![GenerateContent::Text {
text,
provider_metadata: None,
}],
finish_reason: FinishReason {
unified: FinishReasonUnified::Stop,
raw: Some("stop".into()),
},
usage: Usage::default(),
warnings: vec![],
provider_metadata: None,
response: ResponseMetadata::default(),
request_body: None,
response_headers: None,
})
}
async fn do_stream(&self, _options: &CallOptions) -> Result<StreamResult, AiMuxError> {
unreachable!("not used in tests")
}
}
fn optioned_recording(prompt_text: &str) -> Recording {
let mut rec = openai_recording("t1", prompt_text, "pong", "stop");
let mut call = sample_options(prompt_text, Some(0.7));
call.max_output_tokens = Some(128);
rec.input.options = serde_json::to_value(&call).unwrap();
rec
}
#[test]
fn replay_with_model_rebuilds_input_and_resends() {
let rec = optioned_recording("hello");
let model = EchoModel::new();
let rt = tokio::runtime::Runtime::new().unwrap();
let result = rt.block_on(async { replay_with_model(&rec, &model, None).await.unwrap() });
assert_eq!(result.text, "hello");
}
#[test]
fn replay_with_model_applies_overrides() {
let rec = optioned_recording("hello");
let model = EchoModel::new();
let overrides = ReplayOverrides {
prompt: Some(vec![
crate::language_model_message::LanguageModelPromptMessage {
role: crate::message::Role::User,
content: vec![crate::content::ContentPart::Text {
text: "overridden".into(),
provider_options: None,
}],
provider_options: None,
},
]),
..Default::default()
};
let rt = tokio::runtime::Runtime::new().unwrap();
let result = rt.block_on(async {
replay_with_model(&rec, &model, Some(&overrides))
.await
.unwrap()
});
assert_eq!(result.text, "overridden");
}
#[test]
fn replay_with_model_bad_input_options_errors() {
let mut rec = openai_recording("t1", "hello", "pong", "stop");
rec.input.options = serde_json::json!({ "not": "call-options" });
let model = EchoModel::new();
let rt = tokio::runtime::Runtime::new().unwrap();
let err = rt.block_on(async { replay_with_model(&rec, &model, None).await.unwrap_err() });
assert!(matches!(err, AiMuxError::JsonParse(_)), "{err}");
}
}