1use async_trait::async_trait;
16use std::sync::Arc;
17
18use crate::error::AiMuxError;
19use crate::generate::{GenerateTextOptions, GenerateTextResult, generate_text};
20use crate::language_model::LanguageModel;
21use crate::language_model_message::LanguageModelPrompt;
22use crate::message::{MessageContent, ModelMessage, ModelPrompt};
23use crate::options::{CallOptions, ToolChoice};
24use crate::recording::Recording;
25use crate::result::{GenerateContent, GenerateResult, StreamResult};
26use crate::stream_part::StreamPart;
27use crate::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage};
28
29pub trait ReplayMatcher: Send + Sync {
33 fn r#match<'a>(
40 &self,
41 options: &CallOptions,
42 recordings: &'a [Recording],
43 ) -> Result<&'a Recording, AiMuxError>;
44}
45
46pub struct ExactMatcher {
56 provider: String,
57 model_id: String,
58}
59
60impl ExactMatcher {
61 pub fn new(provider: impl Into<String>, model_id: impl Into<String>) -> Self {
62 Self {
63 provider: provider.into(),
64 model_id: model_id.into(),
65 }
66 }
67}
68
69impl ReplayMatcher for ExactMatcher {
70 fn r#match<'a>(
71 &self,
72 options: &CallOptions,
73 recordings: &'a [Recording],
74 ) -> Result<&'a Recording, AiMuxError> {
75 let needle = canonical_call_key(options);
76 recordings
77 .iter()
78 .find(|r| {
79 r.provider.provider == self.provider
80 && r.provider.model_id == self.model_id
81 && canonical_keys_match(&canonical_recording_key(r), &needle)
82 })
83 .ok_or_else(|| {
84 AiMuxError::InvalidArgument("mock replay: no exact matching recording".into())
85 })
86 }
87}
88
89fn canonical_call_key(opts: &CallOptions) -> serde_json::Value {
97 serde_json::json!({
98 "prompt": serde_json::to_value(&opts.prompt).unwrap_or_default(),
99 "temperature": opts.temperature,
100 "max_output_tokens": opts.max_output_tokens,
101 "seed": opts.seed,
102 "response_format": serde_json::to_value(&opts.response_format).unwrap_or_default(),
103 "tools": serde_json::to_value(&opts.tools).unwrap_or_default(),
104 "tool_choice": serde_json::to_value(&opts.tool_choice).unwrap_or_default(),
105 "headers": serde_json::to_value(&opts.headers).unwrap_or_default(),
106 "provider_options": serde_json::to_value(&opts.provider_options).unwrap_or_default(),
107 "body_overrides": serde_json::to_value(&opts.body_overrides).unwrap_or_default(),
108 })
109}
110
111fn canonical_recording_key(rec: &Recording) -> serde_json::Value {
114 let o = &rec.input.options;
115 serde_json::json!({
116 "prompt": serde_json::to_value(&rec.input.prompt).unwrap_or_default(),
117 "temperature": o.get("temperature").cloned().unwrap_or_default(),
118 "max_output_tokens": o.get("max_output_tokens").cloned().unwrap_or_default(),
119 "seed": o.get("seed").cloned().unwrap_or_default(),
120 "response_format": o.get("response_format").cloned().unwrap_or_default(),
121 "tools": o.get("tools").cloned().unwrap_or_default(),
122 "tool_choice": o
123 .get("tool_choice")
124 .cloned()
125 .unwrap_or_else(|| serde_json::to_value(ToolChoice::default()).unwrap_or_default()),
126 "headers": o.get("headers").cloned().unwrap_or_default(),
127 "provider_options": o.get("provider_options").cloned().unwrap_or_default(),
128 "body_overrides": o.get("body_overrides").cloned().unwrap_or_default(),
129 })
130}
131
132fn redaction_aware_eq(rec: &serde_json::Value, req: &serde_json::Value) -> bool {
147 if rec.as_str() == Some("[REDACTED]") {
149 return !req.is_null();
150 }
151 if rec.is_null() || req.is_null() {
153 return rec.is_null() && req.is_null();
154 }
155 match (rec, req) {
157 (serde_json::Value::Object(a), serde_json::Value::Object(b)) => {
158 if a.len() != b.len() {
159 return false; }
161 a.iter().all(|(k, rv)| {
162 b.get(k)
163 .map(|qv| redaction_aware_eq(rv, qv))
164 .unwrap_or(false)
165 })
166 }
167 (serde_json::Value::Array(a), serde_json::Value::Array(b)) => {
168 a.len() == b.len() && a.iter().zip(b).all(|(x, y)| redaction_aware_eq(x, y))
169 }
170 (x, y) => x == y,
171 }
172}
173
174fn canonical_keys_match(rec_key: &serde_json::Value, call_key: &serde_json::Value) -> bool {
182 const REDACTED_FIELDS: [&str; 3] = ["headers", "provider_options", "body_overrides"];
183 let (ro, co) = match (rec_key.as_object(), call_key.as_object()) {
184 (Some(a), Some(b)) => (a, b),
185 _ => return rec_key == call_key,
186 };
187 if ro.len() != co.len() {
188 return false;
189 }
190 ro.iter().all(|(k, rv)| {
191 co.get(k)
192 .map(|qv| {
193 if REDACTED_FIELDS.contains(&k.as_str()) {
194 redaction_aware_eq(rv, qv)
195 } else {
196 rv == qv
197 }
198 })
199 .unwrap_or(false)
200 })
201}
202
203pub struct ScoreMatcher {
210 provider: String,
211 model_id: String,
212}
213
214impl ScoreMatcher {
215 pub fn new(provider: impl Into<String>, model_id: impl Into<String>) -> Self {
216 Self {
217 provider: provider.into(),
218 model_id: model_id.into(),
219 }
220 }
221}
222
223impl ReplayMatcher for ScoreMatcher {
224 fn r#match<'a>(
225 &self,
226 options: &CallOptions,
227 recordings: &'a [Recording],
228 ) -> Result<&'a Recording, AiMuxError> {
229 let mut best: Option<(&Recording, u64)> = None;
230 for r in recordings {
231 if r.provider.provider != self.provider || r.provider.model_id != self.model_id {
232 continue;
233 }
234 let score = match_score(options, r);
235 if score == 0 {
236 continue; }
238 if best.map(|(_, s)| score > s).unwrap_or(true) {
239 best = Some((r, score));
240 }
241 }
242 best.map(|(r, _)| r).ok_or_else(|| {
243 AiMuxError::InvalidArgument("mock replay: no scoring match for this input".into())
244 })
245 }
246}
247
248fn message_text(m: &crate::language_model_message::LanguageModelPromptMessage) -> String {
250 match m.content.first() {
251 Some(crate::content::ContentPart::Text { text, .. }) => text.clone(),
252 _ => serde_json::to_string(&m.content).unwrap_or_default(),
253 }
254}
255
256fn match_score(options: &CallOptions, rec: &Recording) -> u64 {
263 let common = options
265 .prompt
266 .iter()
267 .zip(rec.input.prompt.iter())
268 .take_while(|(a, b)| a.role == b.role && a.content == b.content)
269 .count();
270 let min_len = options.prompt.len().min(rec.input.prompt.len());
274 let lcp: u64 = if min_len == 0 {
275 0
276 } else {
277 let lcp_idx = common.min(min_len - 1);
278 match (options.prompt.get(lcp_idx), rec.input.prompt.get(lcp_idx)) {
279 (Some(a), Some(b)) => message_text(a)
280 .bytes()
281 .zip(message_text(b).bytes())
282 .take_while(|(x, y)| x == y)
283 .count() as u64,
284 _ => 0,
285 }
286 };
287 let prompt_relevance = (common as u64) * 100 + lcp;
288 if prompt_relevance == 0 {
289 return 0; }
291 let mut score = prompt_relevance;
293 if options.temperature
294 == rec
295 .input
296 .options
297 .get("temperature")
298 .and_then(serde_json::Value::as_f64)
299 {
300 score += 1;
301 }
302 score
303}
304
305pub struct PrefixMatcher {
311 provider: String,
312 model_id: String,
313}
314
315impl PrefixMatcher {
316 pub fn new(provider: impl Into<String>, model_id: impl Into<String>) -> Self {
317 Self {
318 provider: provider.into(),
319 model_id: model_id.into(),
320 }
321 }
322}
323
324impl ReplayMatcher for PrefixMatcher {
325 fn r#match<'a>(
326 &self,
327 options: &CallOptions,
328 recordings: &'a [Recording],
329 ) -> Result<&'a Recording, AiMuxError> {
330 let mut best: Option<(&Recording, usize)> = None;
331 for r in recordings {
332 if r.provider.provider != self.provider || r.provider.model_id != self.model_id {
333 continue;
334 }
335 let common = options
337 .prompt
338 .iter()
339 .zip(r.input.prompt.iter())
340 .take_while(|(a, b)| a.role == b.role && a.content == b.content)
341 .count();
342 if common < r.input.prompt.len() {
343 continue; }
345 if common > 0 && best.map(|(_, c)| common > c).unwrap_or(true) {
346 best = Some((r, common));
347 }
348 }
349 best.map(|(r, _)| r).ok_or_else(|| {
350 AiMuxError::InvalidArgument("mock replay: no prefix match for this input".into())
351 })
352 }
353}
354
355pub struct MockReplayModel {
359 provider: String,
360 model_id: String,
361 recordings: Vec<Recording>,
362 matcher: Arc<dyn ReplayMatcher>,
363}
364
365impl MockReplayModel {
366 pub fn new(
368 provider: impl Into<String>,
369 model_id: impl Into<String>,
370 recordings: Vec<Recording>,
371 ) -> Self {
372 let provider = provider.into();
373 let model_id = model_id.into();
374 let matcher = Arc::new(ScoreMatcher::new(provider.clone(), model_id.clone()));
375 Self {
376 provider,
377 model_id,
378 recordings,
379 matcher,
380 }
381 }
382
383 pub fn with_matcher(
385 provider: impl Into<String>,
386 model_id: impl Into<String>,
387 recordings: Vec<Recording>,
388 matcher: Arc<dyn ReplayMatcher>,
389 ) -> Self {
390 Self {
391 provider: provider.into(),
392 model_id: model_id.into(),
393 recordings,
394 matcher,
395 }
396 }
397
398 pub fn from_jsonl(path: &str) -> Result<Self, AiMuxError> {
406 let content = std::fs::read_to_string(path)
407 .map_err(|e| AiMuxError::InvalidArgument(format!("mock replay: {e}")))?;
408 let mut recordings = Vec::new();
409 for (idx, line) in content.lines().enumerate() {
410 if line.trim().is_empty() {
411 continue;
412 }
413 let rec: Recording = serde_json::from_str(line)
414 .map_err(|e| AiMuxError::JsonParse(format!("mock replay line {}: {e}", idx + 1)))?;
415 recordings.push(rec);
416 }
417 if recordings.is_empty() {
418 return Err(AiMuxError::InvalidArgument(
419 "mock replay: empty recording file".into(),
420 ));
421 }
422 let provider = recordings[0].provider.provider.clone();
423 let model_id = recordings[0].provider.model_id.clone();
424 Ok(Self::new(provider, model_id, recordings))
425 }
426
427 #[must_use]
429 pub fn recordings(&self) -> &[Recording] {
430 &self.recordings
431 }
432}
433
434#[async_trait]
435impl LanguageModel for MockReplayModel {
436 fn provider(&self) -> &str {
437 &self.provider
438 }
439
440 fn model_id(&self) -> &str {
441 &self.model_id
442 }
443
444 async fn do_generate(&self, options: &CallOptions) -> Result<GenerateResult, AiMuxError> {
445 let rec = self.matcher.r#match(options, &self.recordings)?;
446 rebuild_generate_result(rec)
447 }
448
449 async fn do_stream(&self, options: &CallOptions) -> Result<StreamResult, AiMuxError> {
450 let rec = self.matcher.r#match(options, &self.recordings)?;
451 rebuild_stream_result(rec)
452 }
453}
454
455fn last_response(rec: &Recording) -> Result<&crate::recording::ResponseRecord, AiMuxError> {
464 for e in rec.exchanges.iter().rev() {
465 if !e.finalized || e.error.is_some() {
466 continue;
467 }
468 if let Some(resp) = e.response.as_ref()
469 && (200..300).contains(&resp.status)
470 {
471 return Ok(resp);
472 }
473 }
474 Err(AiMuxError::InvalidArgument(
475 "mock replay: recording has no successful (2xx, finalized, error-free) response".into(),
476 ))
477}
478
479fn parse_finish(raw: Option<&str>) -> (FinishReasonUnified, Option<String>) {
481 let unified = match raw {
482 Some("stop") => FinishReasonUnified::Stop,
483 Some("length") => FinishReasonUnified::Length,
484 Some("content_filter") => FinishReasonUnified::ContentFilter,
485 Some("tool_calls") => FinishReasonUnified::ToolCalls,
486 Some("error") => FinishReasonUnified::Error,
487 _ => FinishReasonUnified::Other,
488 };
489 (unified, raw.map(std::string::ToString::to_string))
490}
491
492fn u32_from_json(v: &serde_json::Value, field: &str) -> Result<Option<u32>, AiMuxError> {
495 v.as_u64()
496 .map(|n| {
497 u32::try_from(n).map_err(|_| {
498 AiMuxError::JsonParse(format!("mock replay: usage '{field}' overflows u32: {n}"))
499 })
500 })
501 .transpose()
502}
503
504fn parse_usage(v: &serde_json::Value) -> Result<Usage, AiMuxError> {
506 let u = &v["usage"];
507 if u.is_null() {
508 return Ok(Usage::default());
509 }
510 let prompt_details = &u["prompt_tokens_details"];
511 let completion_details = &u["completion_tokens_details"];
512 Ok(Usage {
513 input_tokens: crate::types::TokenUsage {
514 total: u32_from_json(&u["prompt_tokens"], "prompt_tokens")?,
515 no_cache: None,
516 cache_read: u32_from_json(
517 &prompt_details["cached_tokens"],
518 "prompt_tokens_details.cached_tokens",
519 )?,
520 cache_write: None,
521 text: None,
522 reasoning: None,
523 },
524 output_tokens: crate::types::TokenUsage {
525 total: u32_from_json(&u["completion_tokens"], "completion_tokens")?,
526 no_cache: None,
527 cache_read: None,
528 cache_write: None,
529 text: None,
530 reasoning: u32_from_json(
531 &completion_details["reasoning_tokens"],
532 "completion_tokens_details.reasoning_tokens",
533 )?,
534 },
535 raw: Some(u.clone()),
536 })
537}
538
539fn parse_tool_arguments(args: &str) -> serde_json::Value {
546 serde_json::from_str(args).unwrap_or_else(|_| serde_json::Value::String(args.to_string()))
547}
548
549fn rebuild_generate_result(rec: &Recording) -> Result<GenerateResult, AiMuxError> {
556 let resp = last_response(rec)?;
557 let body = resp
558 .body
559 .as_deref()
560 .ok_or_else(|| AiMuxError::InvalidArgument("mock replay: response has no body".into()))?;
561 let v: serde_json::Value = serde_json::from_str(body)
562 .map_err(|e| AiMuxError::JsonParse(format!("mock replay: {e}")))?;
563
564 let message = v
566 .get("choices")
567 .and_then(|c| c.get(0))
568 .and_then(|choice| choice.get("message"))
569 .ok_or_else(|| {
570 AiMuxError::UnsupportedFunctionality(
571 "mock replay: response is not an OpenAI chat.completions body (no choices[0].message)"
572 .into(),
573 )
574 })?;
575
576 let mut content = Vec::new();
579 if let Some(text) = message.get("content").and_then(|c| c.as_str())
580 && !text.is_empty()
581 {
582 content.push(GenerateContent::Text {
583 text: text.to_string(),
584 provider_metadata: None,
585 });
586 }
587 if let Some(tool_calls) = message.get("tool_calls").and_then(|t| t.as_array()) {
588 for tc in tool_calls {
589 let tool_call_id = tc
590 .get("id")
591 .and_then(|x| x.as_str())
592 .unwrap_or_default()
593 .to_string();
594 let function = tc.get("function").cloned().unwrap_or_default();
595 let tool_name = function
596 .get("name")
597 .and_then(|x| x.as_str())
598 .unwrap_or_default()
599 .to_string();
600 let args_str = function
601 .get("arguments")
602 .and_then(|x| x.as_str())
603 .unwrap_or("");
604 let input = parse_tool_arguments(args_str);
605 content.push(GenerateContent::ToolCall {
606 tool_call_id,
607 tool_name,
608 input,
609 provider_executed: None,
610 dynamic: None,
611 thought_signature: None,
612 provider_metadata: None,
613 });
614 }
615 }
616
617 let raw_finish = v["choices"][0]["finish_reason"]
618 .as_str()
619 .map(std::string::ToString::to_string);
620 let (unified, raw) = parse_finish(raw_finish.as_deref());
621
622 Ok(GenerateResult {
623 content,
624 finish_reason: FinishReason { unified, raw },
625 usage: parse_usage(&v)?,
626 warnings: Vec::new(),
627 provider_metadata: None,
628 response: ResponseMetadata {
629 id: v["id"].as_str().map(std::string::ToString::to_string),
630 timestamp: None,
631 model_id: v["model"].as_str().map(std::string::ToString::to_string),
632 },
633 request_body: None,
634 response_headers: None,
635 })
636}
637
638struct ToolCallAccumulator {
640 id: String,
641 name: String,
642 arguments: String,
643}
644
645fn rebuild_stream_result(rec: &Recording) -> Result<StreamResult, AiMuxError> {
659 let resp = last_response(rec)?;
660 let body = resp
661 .body
662 .as_deref()
663 .ok_or_else(|| AiMuxError::InvalidArgument("mock replay: response has no body".into()))?;
664
665 let mut parts: Vec<Result<StreamPart, AiMuxError>> = Vec::new();
666 parts.push(Ok(StreamPart::StreamStart { warnings: vec![] }));
667
668 let text_id = "0".to_string();
669 let reasoning_id = "reasoning-0".to_string();
670 let mut text_started = false;
671 let mut reasoning_started = false;
672 let mut response_meta_emitted = false;
673 let mut final_usage = Usage::default();
674 let mut final_finish: Option<FinishReason> = None;
675 let mut tool_calls: std::collections::HashMap<usize, ToolCallAccumulator> =
677 std::collections::HashMap::new();
678 let mut tool_order: Vec<usize> = Vec::new();
679 let mut saw_openai = false;
680 for block in body.split("\n\n") {
681 let line = block.trim();
682 let data = line.strip_prefix("data:").unwrap_or(line).trim();
683 if data.is_empty() || data == "[DONE]" {
684 continue;
685 }
686 let Ok(v) = serde_json::from_str::<serde_json::Value>(data) else {
687 continue;
688 };
689
690 let choices_arr = v.get("choices").and_then(|c| c.as_array());
692 if choices_arr.is_some() {
693 saw_openai = true;
694 }
695
696 if !response_meta_emitted
698 && (v.get("id").and_then(|x| x.as_str()).is_some()
699 || v.get("model").and_then(|x| x.as_str()).is_some())
700 {
701 response_meta_emitted = true;
702 parts.push(Ok(StreamPart::ResponseMetadata {
703 id: v["id"].as_str().map(std::string::ToString::to_string),
704 timestamp: None,
705 model_id: v["model"].as_str().map(std::string::ToString::to_string),
706 }));
707 }
708
709 if v.get("usage").is_some_and(|u| !u.is_null()) {
712 final_usage = parse_usage(&v)?;
713 }
714
715 if let Some(choices) = choices_arr {
716 for choice in choices {
717 let delta = &choice["delta"];
718
719 let reasoning_delta = delta
721 .get("reasoning_content")
722 .and_then(|x| x.as_str())
723 .or_else(|| delta.get("reasoning").and_then(|x| x.as_str()));
724 if let Some(reasoning) = reasoning_delta
725 && !reasoning.is_empty()
726 {
727 if !reasoning_started {
728 reasoning_started = true;
729 parts.push(Ok(StreamPart::ReasoningStart {
730 id: reasoning_id.clone(),
731 provider_metadata: None,
732 }));
733 }
734 parts.push(Ok(StreamPart::ReasoningDelta {
735 id: reasoning_id.clone(),
736 delta: reasoning.to_string(),
737 provider_metadata: None,
738 }));
739 }
740
741 if let Some(content) = delta.get("content").and_then(|x| x.as_str())
743 && !content.is_empty()
744 {
745 if reasoning_started {
746 parts.push(Ok(StreamPart::ReasoningEnd {
747 id: reasoning_id.clone(),
748 provider_metadata: None,
749 }));
750 reasoning_started = false;
751 }
752 if !text_started {
753 text_started = true;
754 parts.push(Ok(StreamPart::TextStart {
755 id: text_id.clone(),
756 provider_metadata: None,
757 }));
758 }
759 parts.push(Ok(StreamPart::TextDelta {
760 id: text_id.clone(),
761 delta: content.to_string(),
762 provider_metadata: None,
763 }));
764 }
765
766 if let Some(tcs) = delta.get("tool_calls").and_then(|t| t.as_array()) {
768 if reasoning_started {
769 parts.push(Ok(StreamPart::ReasoningEnd {
770 id: reasoning_id.clone(),
771 provider_metadata: None,
772 }));
773 reasoning_started = false;
774 }
775 for dtc in tcs {
776 let idx = dtc
777 .get("index")
778 .and_then(serde_json::Value::as_u64)
779 .map(|n| n as usize)
780 .unwrap_or(0);
781 let func = dtc.get("function").cloned().unwrap_or_default();
782 let is_new = !tool_calls.contains_key(&idx);
783 if is_new {
784 let id = dtc
785 .get("id")
786 .and_then(|x| x.as_str())
787 .unwrap_or_default()
788 .to_string();
789 let name = func
790 .get("name")
791 .and_then(|x| x.as_str())
792 .unwrap_or_default()
793 .to_string();
794 tool_calls.insert(
795 idx,
796 ToolCallAccumulator {
797 id: id.clone(),
798 name: name.clone(),
799 arguments: String::new(),
800 },
801 );
802 tool_order.push(idx);
803 parts.push(Ok(StreamPart::ToolInputStart {
804 id,
805 tool_name: name,
806 provider_executed: None,
807 dynamic: None,
808 title: None,
809 provider_metadata: None,
810 }));
811 }
812 if let Some(args) = func.get("arguments").and_then(|x| x.as_str())
815 && (!is_new || !args.is_empty())
816 && let Some(acc) = tool_calls.get_mut(&idx)
817 {
818 acc.arguments.push_str(args);
819 parts.push(Ok(StreamPart::ToolInputDelta {
820 id: acc.id.clone(),
821 delta: args.to_string(),
822 provider_metadata: None,
823 }));
824 }
825 }
826 }
827
828 if let Some(fr) = choice.get("finish_reason").and_then(|x| x.as_str()) {
830 if reasoning_started {
831 parts.push(Ok(StreamPart::ReasoningEnd {
832 id: reasoning_id.clone(),
833 provider_metadata: None,
834 }));
835 reasoning_started = false;
836 }
837 if text_started {
838 parts.push(Ok(StreamPart::TextEnd {
839 id: text_id.clone(),
840 provider_metadata: None,
841 }));
842 text_started = false;
843 }
844 for &i in &tool_order {
845 if let Some(acc) = tool_calls.get(&i) {
846 parts.push(Ok(StreamPart::ToolInputEnd {
847 id: acc.id.clone(),
848 provider_metadata: None,
849 }));
850 let input = parse_tool_arguments(&acc.arguments);
851 parts.push(Ok(StreamPart::ToolCall {
852 tool_call_id: acc.id.clone(),
853 tool_name: acc.name.clone(),
854 input,
855 provider_executed: None,
856 dynamic: None,
857 thought_signature: None,
858 provider_metadata: None,
859 }));
860 }
861 }
862 tool_calls.clear();
863 tool_order.clear();
864 let (unified, raw) = parse_finish(Some(fr));
865 final_finish = Some(FinishReason { unified, raw });
866 }
867 }
868 }
869 }
870
871 if reasoning_started {
873 parts.push(Ok(StreamPart::ReasoningEnd {
874 id: reasoning_id,
875 provider_metadata: None,
876 }));
877 }
878 if text_started {
879 parts.push(Ok(StreamPart::TextEnd {
880 id: text_id,
881 provider_metadata: None,
882 }));
883 }
884 for &i in &tool_order {
885 if let Some(acc) = tool_calls.get(&i) {
886 parts.push(Ok(StreamPart::ToolInputEnd {
887 id: acc.id.clone(),
888 provider_metadata: None,
889 }));
890 let input = parse_tool_arguments(&acc.arguments);
891 parts.push(Ok(StreamPart::ToolCall {
892 tool_call_id: acc.id.clone(),
893 tool_name: acc.name.clone(),
894 input,
895 provider_executed: None,
896 dynamic: None,
897 thought_signature: None,
898 provider_metadata: None,
899 }));
900 }
901 }
902
903 if !saw_openai {
905 return Err(AiMuxError::UnsupportedFunctionality(
906 "mock replay: response is not an OpenAI chat.completions SSE stream".into(),
907 ));
908 }
909
910 parts.push(Ok(StreamPart::Finish {
911 finish_reason: final_finish.unwrap_or(FinishReason {
912 unified: FinishReasonUnified::Stop,
913 raw: None,
914 }),
915 usage: final_usage,
916 provider_metadata: None,
917 }));
918
919 Ok(StreamResult {
920 stream: Box::pin(futures::stream::iter(parts)),
921 request_body: None,
922 response_headers: None,
923 })
924}
925
926#[derive(Debug, Clone, Default)]
933pub struct ReplayOverrides {
934 pub prompt: Option<LanguageModelPrompt>,
936 pub temperature: Option<f64>,
938 pub max_output_tokens: Option<u32>,
940}
941
942pub async fn replay_with_model(
962 recording: &Recording,
963 model: &dyn LanguageModel,
964 overrides: Option<&ReplayOverrides>,
965) -> Result<GenerateTextResult, AiMuxError> {
966 let mut call_options: CallOptions = serde_json::from_value(recording.input.options.clone())
968 .map_err(|e| AiMuxError::JsonParse(format!("mock replay: input options invalid: {e}")))?;
969
970 if let Some(o) = overrides {
972 if let Some(prompt) = &o.prompt {
973 call_options.prompt = prompt.clone();
974 }
975 if o.temperature.is_some() {
976 call_options.temperature = o.temperature;
977 }
978 if o.max_output_tokens.is_some() {
979 call_options.max_output_tokens = o.max_output_tokens;
980 }
981 }
982
983 let prompt = model_prompt_from_lm(&call_options.prompt);
985 let options = generate_options_from_call_options(call_options);
986 generate_text(model, prompt, options).await
987}
988
989fn model_prompt_from_lm(prompt: &LanguageModelPrompt) -> ModelPrompt {
994 ModelPrompt::Messages(
995 prompt
996 .iter()
997 .map(|m| ModelMessage {
998 role: m.role,
999 content: MessageContent::Parts(m.content.clone()),
1000 })
1001 .collect(),
1002 )
1003}
1004
1005fn generate_options_from_call_options(o: CallOptions) -> GenerateTextOptions {
1010 GenerateTextOptions {
1011 max_output_tokens: o.max_output_tokens,
1012 temperature: o.temperature,
1013 stop_sequences: o.stop_sequences,
1014 top_p: o.top_p,
1015 top_k: o.top_k,
1016 presence_penalty: o.presence_penalty,
1017 frequency_penalty: o.frequency_penalty,
1018 response_format: o.response_format,
1019 seed: o.seed,
1020 tools: o.tools,
1021 tool_choice: if o.tool_choice == crate::tool::ToolChoice::Auto {
1022 None
1023 } else {
1024 Some(o.tool_choice)
1025 },
1026 headers: o.headers,
1027 provider_options: o.provider_options,
1028 reasoning: o.reasoning,
1029 instructions: None,
1030 body_overrides: o.body_overrides,
1031 max_retries: o.max_retries,
1032 timeout: o.timeout,
1033 session_id: o.session_id,
1034 abort_signal: None,
1035 include_raw_chunks: o.include_raw_chunks,
1036 }
1037}
1038
1039#[cfg(test)]
1042mod tests {
1043 use super::*;
1044 use crate::generate::GenerateTextOptions;
1045 use crate::recording::{
1046 HttpExchange, HttpRecord, InputRecord, ProviderRecord, ResponseRecord, TimingRecord,
1047 };
1048 use futures::StreamExt;
1049
1050 fn sample_options(text: &str, temperature: Option<f64>) -> CallOptions {
1051 let prompt = vec![crate::language_model_message::LanguageModelPromptMessage {
1052 role: crate::message::Role::User,
1053 content: vec![crate::content::ContentPart::Text {
1054 text: text.to_string(),
1055 provider_options: None,
1056 }],
1057 provider_options: None,
1058 }];
1059 GenerateTextOptions {
1060 temperature,
1061 ..Default::default()
1062 }
1063 .into_call_options(prompt)
1064 }
1065
1066 fn openai_recording(trace_id: &str, prompt_text: &str, reply: &str, finish: &str) -> Recording {
1068 let input_prompt: crate::language_model_message::LanguageModelPrompt =
1069 vec![crate::language_model_message::LanguageModelPromptMessage {
1070 role: crate::message::Role::User,
1071 content: vec![crate::content::ContentPart::Text {
1072 text: prompt_text.to_string(),
1073 provider_options: None,
1074 }],
1075 provider_options: None,
1076 }];
1077 let body = serde_json::json!({
1078 "id": "chatcmpl-mock",
1079 "object": "chat.completion",
1080 "model": "gpt-4o",
1081 "choices": [{
1082 "index": 0,
1083 "message": { "role": "assistant", "content": reply },
1084 "finish_reason": finish
1085 }],
1086 "usage": { "prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8 }
1087 });
1088 Recording {
1089 schema: crate::recording::RECORDING_SCHEMA,
1090 call_id: trace_id.to_string(),
1091 recorded_at: "2026-08-06T00:00:00Z".to_string(),
1092 input: InputRecord {
1093 prompt: input_prompt,
1094 options: serde_json::json!({ "temperature": 0.7 }),
1095 },
1096 provider: ProviderRecord {
1097 provider: "openai".into(),
1098 model_id: "gpt-4o".into(),
1099 base_url: None,
1100 api_key_source: "none".into(),
1101 profile: None,
1102 provider_options: None,
1103 },
1104 exchanges: vec![HttpExchange {
1105 attempt: 0,
1106 request: HttpRecord {
1107 method: "post".into(),
1108 url: "https://api.openai.com/v1/chat/completions".into(),
1109 headers: vec![],
1110 body: Some("{}".into()),
1111 },
1112 response: Some(ResponseRecord {
1113 status: 200,
1114 headers: vec![],
1115 body: Some(body.to_string()),
1116 stream_chunks: None,
1117 ttfb_ms: None,
1118 }),
1119 timing: TimingRecord {
1120 latency_ms: 10,
1121 ttfb_ms: None,
1122 },
1123 error: None,
1124 finalized: true,
1125 }],
1126 outcome: crate::recording::OutcomeRecord {
1127 status: crate::recording::OutcomeStatus::Success,
1128 finish_reason: Some("stop".into()),
1129 error: None,
1130 usage: None,
1131 },
1132 complete: true,
1133 transport_closed: true,
1134 session_id: None,
1135 step: None,
1136 }
1137 }
1138
1139 #[test]
1140 fn exact_matcher_hits_only_identical_prompt() {
1141 let recs = [openai_recording("t1", "ping", "pong", "stop")];
1144 let matcher = ExactMatcher::new("openai", "gpt-4o");
1145 assert!(
1146 matcher
1147 .r#match(&sample_options("ping", Some(0.7)), &recs)
1148 .is_ok()
1149 );
1150 assert!(
1152 matcher
1153 .r#match(&sample_options("pong", Some(0.7)), &recs)
1154 .is_err()
1155 );
1156 }
1157
1158 #[test]
1159 fn exact_matcher_different_temperature_misses() {
1160 let recs = [openai_recording("t1", "ping", "pong", "stop")]; let matcher = ExactMatcher::new("openai", "gpt-4o");
1163 assert!(
1164 matcher
1165 .r#match(&sample_options("ping", Some(0.1)), &recs)
1166 .is_err()
1167 );
1168 assert!(
1170 matcher
1171 .r#match(&sample_options("ping", Some(0.7)), &recs)
1172 .is_ok()
1173 );
1174 }
1175
1176 #[test]
1177 fn exact_matcher_different_options_misses() {
1178 let mut rec = openai_recording("t1", "ping", "pong", "stop");
1180 let mut call = sample_options("ping", Some(0.7));
1183 call.max_output_tokens = Some(128);
1184 rec.input.options = serde_json::to_value(&call).unwrap();
1185 let recs = [rec];
1186 let matcher = ExactMatcher::new("openai", "gpt-4o");
1187
1188 let mut req = sample_options("ping", Some(0.7));
1190 req.max_output_tokens = Some(256);
1191 assert!(matcher.r#match(&req, &recs).is_err());
1192
1193 let mut req2 = sample_options("ping", Some(0.7));
1195 req2.max_output_tokens = Some(128);
1196 assert!(matcher.r#match(&req2, &recs).is_ok());
1197 }
1198
1199 #[test]
1200 fn exact_matcher_different_headers_misses() {
1201 let mut rec = openai_recording("t1", "ping", "pong", "stop");
1203 let mut call = sample_options("ping", Some(0.7));
1204 call.headers = Some([("x-custom".into(), "a".into())].into_iter().collect());
1205 rec.input.options = serde_json::to_value(&call).unwrap();
1206 let recs = [rec];
1207 let matcher = ExactMatcher::new("openai", "gpt-4o");
1208
1209 let mut req = sample_options("ping", Some(0.7));
1211 req.headers = Some([("x-custom".into(), "b".into())].into_iter().collect());
1212 assert!(matcher.r#match(&req, &recs).is_err());
1213
1214 let mut req2 = sample_options("ping", Some(0.7));
1216 req2.headers = Some([("x-custom".into(), "a".into())].into_iter().collect());
1217 assert!(matcher.r#match(&req2, &recs).is_ok());
1218
1219 let mut req3 = sample_options("ping", Some(0.7));
1221 req3.headers = Some(
1222 [
1223 ("x-custom".into(), "a".into()),
1224 ("x-extra".into(), "z".into()),
1225 ]
1226 .into_iter()
1227 .collect(),
1228 );
1229 assert!(matcher.r#match(&req3, &recs).is_err());
1230 }
1231
1232 #[test]
1233 fn exact_matcher_different_provider_options_misses() {
1234 let mut rec = openai_recording("t1", "ping", "pong", "stop");
1236 let mut call = sample_options("ping", Some(0.7));
1237 call.provider_options = Some(
1238 [("openai".into(), serde_json::json!({ "foo": 1 }))]
1239 .into_iter()
1240 .collect(),
1241 );
1242 rec.input.options = serde_json::to_value(&call).unwrap();
1243 let recs = [rec];
1244 let matcher = ExactMatcher::new("openai", "gpt-4o");
1245
1246 let mut req = sample_options("ping", Some(0.7));
1248 req.provider_options = Some(
1249 [("openai".into(), serde_json::json!({ "foo": 2 }))]
1250 .into_iter()
1251 .collect(),
1252 );
1253 assert!(matcher.r#match(&req, &recs).is_err());
1254
1255 let mut req2 = sample_options("ping", Some(0.7));
1257 req2.provider_options = Some(
1258 [("openai".into(), serde_json::json!({ "foo": 1 }))]
1259 .into_iter()
1260 .collect(),
1261 );
1262 assert!(matcher.r#match(&req2, &recs).is_ok());
1263 }
1264
1265 #[test]
1266 fn exact_matcher_redacted_values_match_any() {
1267 let mut rec = openai_recording("t1", "ping", "pong", "stop");
1270 let mut call = sample_options("ping", Some(0.7));
1271 call.headers = Some(
1272 [
1273 ("authorization".into(), "sk-secret".into()),
1274 ("content-type".into(), "application/json".into()),
1275 ]
1276 .into_iter()
1277 .collect(),
1278 );
1279 let mut opts_val = serde_json::to_value(&call).unwrap();
1283 if let Some(h) = opts_val.get_mut("headers") {
1284 *h = crate::recording::redact_json(h.clone());
1285 }
1286 rec.input.options = opts_val;
1287 assert_eq!(
1289 rec.input.options["headers"]["authorization"],
1290 serde_json::json!("[REDACTED]")
1291 );
1292 let recs = [rec];
1293 let matcher = ExactMatcher::new("openai", "gpt-4o");
1294
1295 let mut req = sample_options("ping", Some(0.7));
1297 req.headers = Some(
1298 [
1299 ("authorization".into(), "sk-other".into()),
1300 ("content-type".into(), "application/json".into()),
1301 ]
1302 .into_iter()
1303 .collect(),
1304 );
1305 assert!(matcher.r#match(&req, &recs).is_ok());
1306
1307 let mut req2 = sample_options("ping", Some(0.7));
1309 req2.headers = Some(
1310 [
1311 ("authorization".into(), "sk-other".into()),
1312 ("content-type".into(), "text/plain".into()),
1313 ]
1314 .into_iter()
1315 .collect(),
1316 );
1317 assert!(matcher.r#match(&req2, &recs).is_err());
1318
1319 let mut req3 = sample_options("ping", Some(0.7));
1321 req3.headers = Some(
1322 [("content-type".into(), "application/json".into())]
1323 .into_iter()
1324 .collect(),
1325 );
1326 assert!(matcher.r#match(&req3, &recs).is_err());
1327
1328 let req4 = sample_options("ping", Some(0.7));
1330 assert!(matcher.r#match(&req4, &recs).is_err());
1331 }
1332
1333 #[test]
1334 fn score_matcher_prefers_longer_prefix_and_ties_to_first() {
1335 let recs = [
1336 openai_recording("t1", "hello", "a", "stop"),
1337 openai_recording("t2", "hello world", "b", "stop"),
1338 ];
1339 let matcher = ScoreMatcher::new("openai", "gpt-4o");
1340 let hit = matcher
1342 .r#match(&sample_options("hello world!", None), &recs)
1343 .unwrap();
1344 assert_eq!(hit.call_id, "t2");
1345 let hit = matcher
1347 .r#match(&sample_options("hello", None), &recs)
1348 .unwrap();
1349 assert_eq!(hit.call_id, "t1");
1350 let hit = matcher
1352 .r#match(&sample_options("hello", Some(0.7)), &recs)
1353 .unwrap();
1354 assert_eq!(hit.call_id, "t1");
1355 }
1356
1357 #[test]
1358 fn match_score_empty_prompt_no_underflow() {
1359 let rec = openai_recording("t1", "hello", "a", "stop");
1361 let empty_req = GenerateTextOptions::default().into_call_options(vec![]);
1362 assert_eq!(match_score(&empty_req, &rec), 0);
1364 let mut rec_empty = rec.clone();
1366 rec_empty.input.prompt = vec![];
1367 assert_eq!(match_score(&sample_options("hello", None), &rec_empty), 0);
1368 assert_eq!(match_score(&empty_req, &rec_empty), 0);
1370 }
1371
1372 #[test]
1373 fn score_matcher_unrelated_prompt_with_matching_temperature_misses() {
1374 let mut rec = openai_recording("t1", "alpha", "a", "stop");
1377 rec.input.options = serde_json::json!({}); let recs = [rec];
1379 let matcher = ScoreMatcher::new("openai", "gpt-4o");
1380 assert!(
1382 matcher
1383 .r#match(&sample_options("zzzzz", None), &recs)
1384 .is_err()
1385 );
1386 }
1387
1388 #[test]
1389 fn score_matcher_related_prompt_with_matching_temperature_hits() {
1390 let recs = [openai_recording("t1", "hello", "a", "stop")]; let matcher = ScoreMatcher::new("openai", "gpt-4o");
1393 let hit = matcher
1394 .r#match(&sample_options("hello", Some(0.7)), &recs)
1395 .unwrap();
1396 assert_eq!(hit.call_id, "t1");
1397 }
1398
1399 #[test]
1400 fn prefix_matcher_matches_message_prefix() {
1401 let rec = openai_recording("t1", "hello", "a", "stop"); let matcher = PrefixMatcher::new("openai", "gpt-4o");
1403 let prompt = vec![
1405 crate::language_model_message::LanguageModelPromptMessage {
1406 role: crate::message::Role::User,
1407 content: vec![crate::content::ContentPart::Text {
1408 text: "hello".into(),
1409 provider_options: None,
1410 }],
1411 provider_options: None,
1412 },
1413 crate::language_model_message::LanguageModelPromptMessage {
1414 role: crate::message::Role::User,
1415 content: vec![crate::content::ContentPart::Text {
1416 text: "world".into(),
1417 provider_options: None,
1418 }],
1419 provider_options: None,
1420 },
1421 ];
1422 let opts = GenerateTextOptions::default().into_call_options(prompt);
1423 let recs = [rec.clone()];
1424 let hit = matcher.r#match(&opts, &recs).unwrap();
1425 assert_eq!(hit.call_id, "t1");
1426
1427 let opts2 = GenerateTextOptions::default().into_call_options(vec![
1429 crate::language_model_message::LanguageModelPromptMessage {
1430 role: crate::message::Role::User,
1431 content: vec![crate::content::ContentPart::Text {
1432 text: "hi".into(),
1433 provider_options: None,
1434 }],
1435 provider_options: None,
1436 },
1437 ]);
1438 assert!(matcher.r#match(&opts2, &[rec]).is_err());
1439 }
1440
1441 #[test]
1442 fn prefix_matcher_prefers_longest_prefix() {
1443 let mut rec_a = openai_recording("ta", "hello", "a", "stop");
1445 let mut rec_b = openai_recording("tb", "hello world", "b", "stop");
1446 let mk = |text: &str| crate::language_model_message::LanguageModelPromptMessage {
1447 role: crate::message::Role::User,
1448 content: vec![crate::content::ContentPart::Text {
1449 text: text.into(),
1450 provider_options: None,
1451 }],
1452 provider_options: None,
1453 };
1454 rec_a.input.prompt = vec![mk("hello")];
1455 rec_b.input.prompt = vec![mk("hello"), mk("world")];
1456
1457 let matcher = PrefixMatcher::new("openai", "gpt-4o");
1458 let opts = GenerateTextOptions::default().into_call_options(vec![
1460 mk("hello"),
1461 mk("world"),
1462 mk("!"),
1463 ]);
1464 let recs = [rec_a, rec_b];
1465 let hit = matcher.r#match(&opts, &recs).unwrap();
1466 assert_eq!(hit.call_id, "tb");
1467 }
1468
1469 #[test]
1470 fn unmatched_returns_clear_error() {
1471 let recs = [openai_recording("t1", "ping", "pong", "stop")];
1472 let model = MockReplayModel::new("openai", "gpt-4o", vec![recs[0].clone()]);
1473 let rt = tokio::runtime::Runtime::new().unwrap();
1474 let err = rt.block_on(async {
1475 model
1476 .do_generate(&sample_options("nope", None))
1477 .await
1478 .unwrap_err()
1479 });
1480 assert!(err.to_string().contains("mock replay"), "{err}");
1481 }
1482
1483 #[test]
1484 fn mock_model_rebuilds_generate_result_from_openai_body() {
1485 let rec = openai_recording("t1", "ping", "pong", "stop");
1486 let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
1487 let rt = tokio::runtime::Runtime::new().unwrap();
1488 let result = rt.block_on(async {
1489 model
1490 .do_generate(&sample_options("ping", None))
1491 .await
1492 .unwrap()
1493 });
1494 let Some(GenerateContent::Text { text, .. }) = result.content.first() else {
1495 panic!("expected text content");
1496 };
1497 assert_eq!(text, "pong");
1498 assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop);
1499 assert_eq!(result.usage.input_tokens.total, Some(5));
1500 assert_eq!(result.response.model_id.as_deref(), Some("gpt-4o"));
1501 }
1502
1503 #[test]
1504 fn mock_model_usage_overflow_errors() {
1505 let mut rec = openai_recording("t1", "ping", "pong", "stop");
1507 let body = serde_json::json!({
1508 "id": "chatcmpl-mock",
1509 "model": "gpt-4o",
1510 "choices": [{
1511 "index": 0,
1512 "message": { "role": "assistant", "content": "pong" },
1513 "finish_reason": "stop"
1514 }],
1515 "usage": { "prompt_tokens": 5_000_000_000_u64, "completion_tokens": 3, "total_tokens": 8 }
1516 });
1517 rec.exchanges[0].response.as_mut().unwrap().body = Some(body.to_string());
1518 let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
1519 let rt = tokio::runtime::Runtime::new().unwrap();
1520 let err = rt.block_on(async {
1521 model
1522 .do_generate(&sample_options("ping", None))
1523 .await
1524 .unwrap_err()
1525 });
1526 assert!(matches!(err, AiMuxError::JsonParse(_)), "{err}");
1527 assert!(err.to_string().contains("overflows u32"), "{err}");
1528 }
1529
1530 #[test]
1531 fn mock_model_rebuilds_stream_from_sse() {
1532 let sse = "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\",\"content\":\"hel\"}}]}\n\n\
1533 data: {\"choices\":[{\"delta\":{\"content\":\"lo\"}}]}\n\n\
1534 data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n\
1535 data: [DONE]\n\n";
1536 let mut rec = openai_recording("t1", "ping", "x", "stop");
1537 rec.exchanges[0].response.as_mut().unwrap().body = Some(sse.to_string());
1538 rec.exchanges[0].response.as_mut().unwrap().stream_chunks = Some(4);
1539 let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
1540 let rt = tokio::runtime::Runtime::new().unwrap();
1541 let parts: Vec<StreamPart> = rt.block_on(async {
1542 model
1543 .do_stream(&sample_options("ping", None))
1544 .await
1545 .unwrap()
1546 .stream
1547 .collect::<Vec<_>>()
1548 .await
1549 .into_iter()
1550 .map(|p| p.unwrap())
1551 .collect()
1552 });
1553 assert!(matches!(parts[0], StreamPart::StreamStart { .. }));
1554 let deltas: Vec<String> = parts
1555 .iter()
1556 .filter_map(|p| match p {
1557 StreamPart::TextDelta { delta, .. } => Some(delta.clone()),
1558 _ => None,
1559 })
1560 .collect();
1561 assert_eq!(deltas, vec!["hel", "lo"]);
1562 assert!(parts.iter().any(|p| matches!(p, StreamPart::Finish { .. })));
1563 }
1564
1565 #[test]
1566 fn non_openai_body_returns_unsupported() {
1567 let mut rec = openai_recording("t1", "ping", "x", "stop");
1570 rec.exchanges[0].response.as_mut().unwrap().body = Some("{\"foo\":\"bar\"}".to_string()); let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
1572 let rt = tokio::runtime::Runtime::new().unwrap();
1573 let err = rt.block_on(async {
1574 model
1575 .do_generate(&sample_options("ping", None))
1576 .await
1577 .unwrap_err()
1578 });
1579 assert!(
1580 matches!(err, AiMuxError::UnsupportedFunctionality(_)),
1581 "{err}"
1582 );
1583 }
1584
1585 #[test]
1586 fn non_openai_stream_returns_unsupported() {
1587 let mut rec = openai_recording("t1", "ping", "x", "stop");
1589 rec.exchanges[0].response.as_mut().unwrap().body =
1590 Some("event: message_start\ndata: {\"type\":\"message_start\"}\n\n".to_string());
1591 rec.exchanges[0].response.as_mut().unwrap().stream_chunks = Some(1);
1592 let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
1593 let rt = tokio::runtime::Runtime::new().unwrap();
1594 let err = rt.block_on(async {
1595 model
1596 .do_stream(&sample_options("ping", None))
1597 .await
1598 .unwrap_err()
1599 });
1600 assert!(
1601 matches!(err, AiMuxError::UnsupportedFunctionality(_)),
1602 "{err}"
1603 );
1604 }
1605
1606 #[test]
1607 fn mock_model_rebuilds_tool_calls_from_openai_body() {
1608 let mut rec = openai_recording("t1", "ping", "x", "tool_calls");
1610 let body = serde_json::json!({
1611 "id": "chatcmpl-mock",
1612 "model": "gpt-4o",
1613 "choices": [{
1614 "index": 0,
1615 "message": {
1616 "role": "assistant",
1617 "content": null,
1618 "tool_calls": [{
1619 "id": "call_abc",
1620 "type": "function",
1621 "function": { "name": "get_weather", "arguments": "{\"city\":\"SF\"}" }
1622 }]
1623 },
1624 "finish_reason": "tool_calls"
1625 }],
1626 "usage": { "prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8 }
1627 });
1628 rec.exchanges[0].response.as_mut().unwrap().body = Some(body.to_string());
1629 let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
1630 let rt = tokio::runtime::Runtime::new().unwrap();
1631 let result = rt.block_on(async {
1632 model
1633 .do_generate(&sample_options("ping", None))
1634 .await
1635 .unwrap()
1636 });
1637 assert_eq!(result.content.len(), 1);
1638 let Some(GenerateContent::ToolCall {
1639 tool_call_id,
1640 tool_name,
1641 input,
1642 ..
1643 }) = result.content.first()
1644 else {
1645 panic!("expected ToolCall, got {:?}", result.content);
1646 };
1647 assert_eq!(tool_call_id, "call_abc");
1648 assert_eq!(tool_name, "get_weather");
1649 assert_eq!(input, &serde_json::json!({ "city": "SF" }));
1650 assert_eq!(result.finish_reason.unified, FinishReasonUnified::ToolCalls);
1651 }
1652
1653 #[test]
1654 fn mock_model_tool_call_bad_json_arguments_falls_back_to_string() {
1655 let mut rec = openai_recording("t1", "ping", "x", "tool_calls");
1657 let body = serde_json::json!({
1658 "id": "chatcmpl-mock",
1659 "model": "gpt-4o",
1660 "choices": [{
1661 "index": 0,
1662 "message": {
1663 "role": "assistant",
1664 "content": null,
1665 "tool_calls": [{
1666 "id": "call_1",
1667 "type": "function",
1668 "function": { "name": "f", "arguments": "not-json{" }
1669 }]
1670 },
1671 "finish_reason": "tool_calls"
1672 }]
1673 });
1674 rec.exchanges[0].response.as_mut().unwrap().body = Some(body.to_string());
1675 let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
1676 let rt = tokio::runtime::Runtime::new().unwrap();
1677 let result = rt.block_on(async {
1678 model
1679 .do_generate(&sample_options("ping", None))
1680 .await
1681 .unwrap()
1682 });
1683 let Some(GenerateContent::ToolCall { input, .. }) = result.content.first() else {
1684 panic!("expected ToolCall");
1685 };
1686 assert_eq!(input, &serde_json::json!("not-json{"));
1687 }
1688
1689 #[test]
1690 fn mock_model_rebuilds_stream_tool_calls() {
1691 let chunk1 = serde_json::json!({
1693 "id": "chatcmpl-1", "model": "gpt-4o",
1694 "choices": [{"index": 0, "delta": {"role": "assistant", "tool_calls": [
1695 {"index": 0, "id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": ""}}
1696 ]}}]
1697 });
1698 let chunk2 = serde_json::json!({
1699 "choices": [{"index": 0, "delta": {"tool_calls": [
1700 {"index": 0, "function": {"arguments": "{\"city\":\"SF\"}"}}
1701 ]}}]
1702 });
1703 let chunk3 = serde_json::json!({"choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]});
1704 let sse = format!("data: {chunk1}\n\ndata: {chunk2}\n\ndata: {chunk3}\n\ndata: [DONE]\n\n");
1705 let mut rec = openai_recording("t1", "ping", "x", "stop");
1706 rec.exchanges[0].response.as_mut().unwrap().body = Some(sse);
1707 rec.exchanges[0].response.as_mut().unwrap().stream_chunks = Some(4);
1708 let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
1709 let rt = tokio::runtime::Runtime::new().unwrap();
1710 let parts: Vec<StreamPart> = rt.block_on(async {
1711 model
1712 .do_stream(&sample_options("ping", None))
1713 .await
1714 .unwrap()
1715 .stream
1716 .collect::<Vec<_>>()
1717 .await
1718 .into_iter()
1719 .map(|p| p.unwrap())
1720 .collect()
1721 });
1722 assert!(parts.iter().any(|p| matches!(
1724 p,
1725 StreamPart::ResponseMetadata { id, model_id, .. }
1726 if id.as_deref() == Some("chatcmpl-1") && model_id.as_deref() == Some("gpt-4o")
1727 )));
1728 assert!(parts.iter().any(|p| matches!(
1730 p,
1731 StreamPart::ToolInputStart { tool_name, .. } if tool_name == "get_weather"
1732 )));
1733 let deltas: Vec<String> = parts
1735 .iter()
1736 .filter_map(|p| match p {
1737 StreamPart::ToolInputDelta { delta, .. } => Some(delta.clone()),
1738 _ => None,
1739 })
1740 .collect();
1741 assert_eq!(deltas, vec!["{\"city\":\"SF\"}"]);
1742 assert!(
1744 parts
1745 .iter()
1746 .any(|p| matches!(p, StreamPart::ToolInputEnd { .. }))
1747 );
1748 let calls: Vec<_> = parts
1749 .iter()
1750 .filter_map(|p| match p {
1751 StreamPart::ToolCall {
1752 tool_call_id,
1753 tool_name,
1754 input,
1755 ..
1756 } => Some((tool_call_id.clone(), tool_name.clone(), input.clone())),
1757 _ => None,
1758 })
1759 .collect();
1760 assert_eq!(calls.len(), 1);
1761 assert_eq!(calls[0].0, "call_1");
1762 assert_eq!(calls[0].1, "get_weather");
1763 assert_eq!(calls[0].2, serde_json::json!({"city":"SF"}));
1764 assert!(parts.iter().any(|p| matches!(
1766 p,
1767 StreamPart::Finish { finish_reason, .. }
1768 if finish_reason.unified == FinishReasonUnified::ToolCalls
1769 )));
1770 }
1771
1772 #[test]
1773 fn mock_model_rebuilds_stream_metadata_and_usage() {
1774 let chunk1 = serde_json::json!({
1777 "id": "chatcmpl-2", "model": "gpt-4o",
1778 "choices": [{"index": 0, "delta": {"content": "hi"}}]
1779 });
1780 let chunk2 = serde_json::json!({
1781 "id": "chatcmpl-2", "choices": [],
1782 "usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9}
1783 });
1784 let sse = format!("data: {chunk1}\n\ndata: {chunk2}\n\ndata: [DONE]\n\n");
1785 let mut rec = openai_recording("t1", "ping", "x", "stop");
1786 rec.exchanges[0].response.as_mut().unwrap().body = Some(sse);
1787 rec.exchanges[0].response.as_mut().unwrap().stream_chunks = Some(3);
1788 let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
1789 let rt = tokio::runtime::Runtime::new().unwrap();
1790 let parts: Vec<StreamPart> = rt.block_on(async {
1791 model
1792 .do_stream(&sample_options("ping", None))
1793 .await
1794 .unwrap()
1795 .stream
1796 .collect::<Vec<_>>()
1797 .await
1798 .into_iter()
1799 .map(|p| p.unwrap())
1800 .collect()
1801 });
1802 assert!(parts.iter().any(|p| matches!(
1803 p,
1804 StreamPart::ResponseMetadata { id, model_id, .. }
1805 if id.as_deref() == Some("chatcmpl-2") && model_id.as_deref() == Some("gpt-4o")
1806 )));
1807 let deltas: Vec<String> = parts
1808 .iter()
1809 .filter_map(|p| match p {
1810 StreamPart::TextDelta { delta, .. } => Some(delta.clone()),
1811 _ => None,
1812 })
1813 .collect();
1814 assert_eq!(deltas, vec!["hi"]);
1815 let finish = parts
1816 .iter()
1817 .find_map(|p| match p {
1818 StreamPart::Finish { usage, .. } => Some(usage.clone()),
1819 _ => None,
1820 })
1821 .unwrap();
1822 assert_eq!(finish.input_tokens.total, Some(7));
1823 assert_eq!(finish.output_tokens.total, Some(2));
1824 }
1825
1826 #[test]
1827 fn last_response_skips_failed_attempt() {
1828 let mut rec = openai_recording("t1", "ping", "pong", "stop");
1830 rec.exchanges[0].response.as_mut().unwrap().status = 500;
1831 rec.exchanges[0].attempt = 0;
1832 let success_body = serde_json::json!({
1833 "id": "chatcmpl-ok", "model": "gpt-4o",
1834 "choices": [{"index":0,"message":{"role":"assistant","content":"pong"},"finish_reason":"stop"}],
1835 "usage": {"prompt_tokens":5,"completion_tokens":3,"total_tokens":8}
1836 });
1837 rec.exchanges.push(HttpExchange {
1838 attempt: 1,
1839 request: HttpRecord {
1840 method: "post".into(),
1841 url: "https://api.openai.com/v1/chat/completions".into(),
1842 headers: vec![],
1843 body: Some("{}".into()),
1844 },
1845 response: Some(ResponseRecord {
1846 status: 200,
1847 headers: vec![],
1848 body: Some(success_body.to_string()),
1849 stream_chunks: None,
1850 ttfb_ms: None,
1851 }),
1852 timing: TimingRecord {
1853 latency_ms: 10,
1854 ttfb_ms: None,
1855 },
1856 error: None,
1857 finalized: true,
1858 });
1859 let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
1860 let rt = tokio::runtime::Runtime::new().unwrap();
1861 let result = rt.block_on(async {
1862 model
1863 .do_generate(&sample_options("ping", None))
1864 .await
1865 .unwrap()
1866 });
1867 let Some(GenerateContent::Text { text, .. }) = result.content.first() else {
1868 panic!("expected text");
1869 };
1870 assert_eq!(text, "pong");
1871 assert_eq!(result.response.id.as_deref(), Some("chatcmpl-ok"));
1873 }
1874
1875 #[test]
1876 fn last_response_no_successful_response_errors() {
1877 let mut rec = openai_recording("t1", "ping", "pong", "stop");
1879 rec.exchanges[0].response.as_mut().unwrap().status = 500;
1880 let model = MockReplayModel::new("openai", "gpt-4o", vec![rec]);
1881 let rt = tokio::runtime::Runtime::new().unwrap();
1882 let err = rt.block_on(async {
1883 model
1884 .do_generate(&sample_options("ping", None))
1885 .await
1886 .unwrap_err()
1887 });
1888 assert!(err.to_string().contains("no successful"), "{err}");
1889 }
1890
1891 #[test]
1892 fn from_jsonl_loads_recordings() {
1893 let dir = std::env::temp_dir().join(format!("aimux-replay-jsonl-{}", std::process::id()));
1894 let _ = std::fs::remove_dir_all(&dir);
1895 std::fs::create_dir_all(&dir).unwrap();
1896 let path = dir.join("rec.jsonl");
1897 let rec = openai_recording("t1", "ping", "pong", "stop");
1898 std::fs::write(&path, serde_json::to_string(&rec).unwrap() + "\n").unwrap();
1899
1900 let model = MockReplayModel::from_jsonl(path.to_str().unwrap()).unwrap();
1901 assert_eq!(model.provider(), "openai");
1902 assert_eq!(model.model_id(), "gpt-4o");
1903 assert_eq!(model.recordings().len(), 1);
1904 let _ = std::fs::remove_dir_all(&dir);
1905 }
1906
1907 #[derive(Clone)]
1911 struct EchoModel {
1912 provider: &'static str,
1913 model_id: &'static str,
1914 }
1915
1916 impl EchoModel {
1917 fn new() -> Self {
1918 Self {
1919 provider: "openai",
1920 model_id: "gpt-4o",
1921 }
1922 }
1923 }
1924
1925 #[async_trait]
1926 impl LanguageModel for EchoModel {
1927 fn provider(&self) -> &str {
1928 self.provider
1929 }
1930 fn model_id(&self) -> &str {
1931 self.model_id
1932 }
1933 async fn do_generate(&self, options: &CallOptions) -> Result<GenerateResult, AiMuxError> {
1934 let text = options
1935 .prompt
1936 .iter()
1937 .filter_map(|m| match m.content.first() {
1938 Some(crate::content::ContentPart::Text { text, .. }) => Some(text.clone()),
1939 _ => None,
1940 })
1941 .collect::<Vec<_>>()
1942 .join(" ");
1943 Ok(GenerateResult {
1944 content: vec![GenerateContent::Text {
1945 text,
1946 provider_metadata: None,
1947 }],
1948 finish_reason: FinishReason {
1949 unified: FinishReasonUnified::Stop,
1950 raw: Some("stop".into()),
1951 },
1952 usage: Usage::default(),
1953 warnings: vec![],
1954 provider_metadata: None,
1955 response: ResponseMetadata::default(),
1956 request_body: None,
1957 response_headers: None,
1958 })
1959 }
1960 async fn do_stream(&self, _options: &CallOptions) -> Result<StreamResult, AiMuxError> {
1961 unreachable!("not used in tests")
1962 }
1963 }
1964
1965 fn optioned_recording(prompt_text: &str) -> Recording {
1968 let mut rec = openai_recording("t1", prompt_text, "pong", "stop");
1969 let mut call = sample_options(prompt_text, Some(0.7));
1970 call.max_output_tokens = Some(128);
1971 rec.input.options = serde_json::to_value(&call).unwrap();
1972 rec
1973 }
1974
1975 #[test]
1976 fn replay_with_model_rebuilds_input_and_resends() {
1977 let rec = optioned_recording("hello");
1978 let model = EchoModel::new();
1979 let rt = tokio::runtime::Runtime::new().unwrap();
1980 let result = rt.block_on(async { replay_with_model(&rec, &model, None).await.unwrap() });
1981 assert_eq!(result.text, "hello");
1982 }
1983
1984 #[test]
1985 fn replay_with_model_applies_overrides() {
1986 let rec = optioned_recording("hello");
1987 let model = EchoModel::new();
1988 let overrides = ReplayOverrides {
1989 prompt: Some(vec![
1990 crate::language_model_message::LanguageModelPromptMessage {
1991 role: crate::message::Role::User,
1992 content: vec![crate::content::ContentPart::Text {
1993 text: "overridden".into(),
1994 provider_options: None,
1995 }],
1996 provider_options: None,
1997 },
1998 ]),
1999 ..Default::default()
2000 };
2001 let rt = tokio::runtime::Runtime::new().unwrap();
2002 let result = rt.block_on(async {
2003 replay_with_model(&rec, &model, Some(&overrides))
2004 .await
2005 .unwrap()
2006 });
2007 assert_eq!(result.text, "overridden");
2008 }
2009
2010 #[test]
2011 fn replay_with_model_bad_input_options_errors() {
2012 let mut rec = openai_recording("t1", "hello", "pong", "stop");
2013 rec.input.options = serde_json::json!({ "not": "call-options" });
2014 let model = EchoModel::new();
2015 let rt = tokio::runtime::Runtime::new().unwrap();
2016 let err = rt.block_on(async { replay_with_model(&rec, &model, None).await.unwrap_err() });
2017 assert!(matches!(err, AiMuxError::JsonParse(_)), "{err}");
2018 }
2019}