1use std::collections::{HashMap, HashSet};
22use std::sync::Arc;
23use std::time::Instant;
24
25use futures::{StreamExt, TryStreamExt, stream};
26use http::StatusCode;
27use parking_lot::Mutex;
28use serde_json::{Value, json};
29use switchyard_libsy::{
30 Algorithm, AnswerSummary, Call, CallDecision, LibsyError, OutcomeMetadata, Result,
31 RoutingOutcome, RuntimeModels, drive,
32};
33use switchyard_protocol::{
34 AggLlmResponse, LlmClientError, LlmResponse, LlmResponseChunk, LlmResponseStream, Message,
35 ModelId, Request, Response, ResponseAccumulator, RoutedDecisionClient, RoutedLlmClient,
36 RoutingFallbackReason, Usage, WireFormat,
37};
38use switchyard_translation::prepare_request_for_target;
39
40use crate::observation::{ModelCallObservation, RunObservation, RunObserver};
41use crate::{metrics, observability};
42
43pub async fn run(
57 algorithm: Arc<dyn Algorithm>,
58 clients: ClientRouter,
59 request: Request,
60 models: Arc<RuntimeModels>,
61 observer: Option<RunObserver>,
62) -> Result<(ModelId, Response)> {
63 let algorithm_name = algorithm.name().to_string();
64 let run_started = Instant::now();
65 let routing_clients = clients.clone();
66 let routing_observations = observer.as_ref().map(|_| Arc::new(Mutex::new(Vec::new())));
69 let outcome = match clients.stored_state_owner(&request) {
70 Some(owner) => Ok(continue_on(owner, &algorithm_name, request)),
71 None => {
72 drive(Arc::clone(&algorithm), request, models, {
73 let routing_observations = routing_observations.clone();
74 move |call| serve(routing_clients.clone(), call, routing_observations.clone())
75 })
76 .await
77 }
78 };
79 let answered_model = outcome
80 .as_ref()
81 .ok()
82 .and_then(|outcome| outcome.response.as_ref())
83 .and_then(Response::served_model);
84 emit_routing_observations(&observer, &routing_observations, answered_model);
85 let outcome = outcome?;
86 let overhead = run_started.elapsed();
87 metrics::record_routing_overhead(&algorithm_name, overhead);
88
89 let selected_model_id = outcome.selected_model_id()?.clone();
90 if let Some(observer) = &observer
91 && let Some(metadata) = outcome.metadata
92 {
93 observer(RunObservation::Outcome(metadata));
94 }
95 let result = if let Some(response) = outcome.response {
96 let result = clients.remember_state_owner(&outcome.request, response);
97 let served_model = result
98 .as_ref()
99 .ok()
100 .and_then(Response::served_model)
101 .unwrap_or(&selected_model_id);
102 let served_model = clients.provider_model(served_model).clone();
103 metrics::observe_routed_request(&algorithm_name, &served_model, None, result)
104 } else {
105 let observe = |observation| {
106 if let Some(observer) = &observer {
107 observer(RunObservation::AnswerCall(observation));
108 }
109 };
110 call_first_available(
111 &clients,
112 &algorithm_name,
113 &outcome.request,
114 outcome.fallback_request.as_ref(),
115 &outcome.candidate_requests,
116 &outcome.selected_model_ids,
117 &observe,
118 )
119 .await
120 .and_then(|(request, response)| clients.remember_state_owner(request, response))
121 };
122 if let Some(observer) = &observer {
123 observer(RunObservation::RoutingOverhead(overhead));
124 }
125 let mut response = result?;
126 if algorithm.tracks_answers()
127 && let Some(session) = session_id(&outcome.request)
128 {
129 response = report_answer(algorithm, session, response).await;
130 }
131 Ok((selected_model_id, response))
132}
133
134fn session_id(request: &Request) -> Option<String> {
135 let session = request.metadata.as_ref()?.session_id.as_deref()?;
136 (!session.is_empty()).then(|| session.to_string())
137}
138
139async fn report_answer(
142 algorithm: Arc<dyn Algorithm>,
143 session: String,
144 mut response: Response,
145) -> Response {
146 let stream = match response.llm_response {
147 LlmResponse::Agg(aggregate) => {
148 algorithm
149 .observe_answer(&session, AnswerSummary::of(&aggregate))
150 .await;
151 response.llm_response = LlmResponse::Agg(aggregate);
152 return response;
153 }
154 LlmResponse::Stream(stream) => stream,
155 };
156 algorithm
157 .observe_answer(&session, AnswerSummary { final_reply: false })
158 .await;
159 let report = (algorithm, session, Some(ResponseAccumulator::new()));
160 response.llm_response = LlmResponse::Stream(Box::pin(stream::unfold(
161 (stream, report),
162 |(mut stream, (algorithm, session, mut accumulator))| async move {
163 let item = stream.next().await?;
164 let failed = match &item {
165 Err(_) => true,
166 Ok(event) => event
167 .normalized()
168 .iter()
169 .any(|chunk| matches!(chunk, LlmResponseChunk::StreamError { .. })),
170 };
171 if failed {
172 accumulator = None;
173 algorithm
174 .observe_answer(&session, AnswerSummary { final_reply: false })
175 .await;
176 }
177 let mut stopped = false;
178 if let (Ok(event), Some(answer)) = (&item, accumulator.as_mut()) {
179 for chunk in event.normalized() {
180 stopped |= matches!(chunk, LlmResponseChunk::MessageStop { .. });
181 answer.push(chunk.clone());
182 }
183 }
184 if stopped && let Some(answer) = accumulator.take() {
185 let answer = AnswerSummary::of(&answer.finish());
186 algorithm.observe_answer(&session, answer).await;
187 }
188 Some((item, (stream, (algorithm, session, accumulator))))
189 },
190 )));
191 response
192}
193
194pub async fn decide(
199 algorithm: Arc<dyn Algorithm>,
200 clients: ClientRouter,
201 request: Request,
202 models: Arc<RuntimeModels>,
203) -> Result<RoutingOutcome> {
204 let routing_clients = clients.clone();
205 let mut outcome = match clients.stored_state_owner(&request) {
206 Some(owner) => continue_on(owner, algorithm.name(), request),
207 None => {
208 drive(algorithm, request, models, move |call| {
209 serve(routing_clients.clone(), call, None)
210 })
211 .await?
212 }
213 };
214 let selected_model_id = outcome.selected_model_id()?.clone();
215 outcome.request = clients.prepare_completion_request(outcome.request, &selected_model_id);
216 outcome.response = outcome
217 .response
218 .map(|response| clients.remember_state_owner(&outcome.request, response))
219 .transpose()?;
220 Ok(outcome)
221}
222
223#[tracing::instrument(
224 name = "libsy.decision_client_call",
225 skip_all,
226 fields(algorithm = %call.algorithm, selected_model = %call.model, outcome = tracing::field::Empty),
227)]
228async fn serve_decision(
229 clients: &ClientRouter,
230 call: CallDecision,
231 observations: &Option<Arc<Mutex<Vec<RunObservation>>>>,
232) -> Result<()> {
233 let client = clients.route_decision(&call.model);
234 let model = call.model.clone();
235 let request = call.request.clone();
236 call.respond(async {
237 let started = Instant::now();
238 let result = async { client?.call(request).await }.await;
239 tracing::Span::current().record("outcome", if result.is_ok() { "ok" } else { "error" });
240 if let Some(observations) = observations {
241 observations
242 .lock()
243 .push(RunObservation::DecisionCall(ModelCallObservation {
244 selected_model: model.clone(),
245 is_success: result.is_ok(),
246 duration: started.elapsed(),
247 usage: result.as_ref().ok().map(|response| response.usage.clone()),
248 }));
249 }
250 result.map_err(|error| LibsyError::client_call(model, error))
251 })
252 .await
253}
254
255fn emit_routing_observations(
257 observer: &Option<RunObserver>,
258 observations: &Option<Arc<Mutex<Vec<RunObservation>>>>,
259 answered_model: Option<&ModelId>,
260) {
261 let (Some(observer), Some(observations)) = (observer, observations) else {
262 return;
263 };
264 let mut answer_observed = false;
265 for observation in observations.lock().drain(..) {
266 observer(match observation {
267 RunObservation::LlmCall(call)
268 if !answer_observed && answered_model == Some(&call.selected_model) =>
269 {
270 answer_observed = true;
271 RunObservation::AnswerCall(call)
272 }
273 observation => observation,
274 });
275 }
276}
277
278async fn serve(
283 clients: ClientRouter,
284 call: Call,
285 observations: Option<Arc<Mutex<Vec<RunObservation>>>>,
286) -> Result<()> {
287 let call = match call {
288 Call::Model(call) => *call,
289 Call::Decision(call) => return serve_decision(&clients, *call, &observations).await,
290 };
291 let observe = |observation| {
292 if let Some(observations) = &observations {
293 observations
294 .lock()
295 .push(RunObservation::LlmCall(observation));
296 }
297 };
298 let target = call.models.first().cloned().ok_or(LibsyError::NoTargets)?;
299 let request = clients.prepare_routing_request(call.request.clone(), &target);
300 let algorithm = call.algorithm.clone();
301 let candidates = call.models.len();
302 call.respond(call_one(
303 &clients,
304 &target,
305 request,
306 &algorithm,
307 &observe,
308 0,
309 candidates,
310 CallPhase::Routing,
311 ))
312 .await
313}
314
315async fn call_first_available<'a>(
317 clients: &ClientRouter,
318 algorithm: &str,
319 request: &'a Request,
320 fallback_request: Option<&'a Request>,
321 candidate_requests: &'a HashMap<ModelId, Request>,
322 models: &[ModelId],
323 observe: &(dyn Fn(ModelCallObservation) + Send + Sync),
324) -> Result<(&'a Request, Response)> {
325 for (index, target) in models.iter().enumerate() {
326 let request = candidate_requests
327 .get(target)
328 .unwrap_or(match fallback_request {
329 Some(fallback) if index > 0 => fallback,
330 _ => request,
331 });
332 let prepared = clients.prepare_completion_request(request.clone(), target);
333 match call_one(
334 clients,
335 target,
336 prepared,
337 algorithm,
338 observe,
339 index,
340 models.len(),
341 CallPhase::Completion,
342 )
343 .await
344 {
345 Ok(response) => return Ok((request, response)),
346 Err(error) if index + 1 == models.len() => return Err(error),
347 Err(error)
348 if clients.rate_limit_fallback_disabled.contains(target)
349 && matches!(&error, LibsyError::ClientCall {
350 source: LlmClientError::UpstreamHttp { status, .. }, ..
351 } if *status == StatusCode::TOO_MANY_REQUESTS) =>
352 {
353 return Err(error);
354 }
355 Err(error) => match fallback_reason(&error, fallback_request.is_none()) {
356 Some(reason) => tracing::info!(
357 from = %target,
358 to = %models[index + 1],
359 reason = reason.as_str(),
360 "model call failed; trying next candidate"
361 ),
362 None => return Err(error),
363 },
364 }
365 }
366 Err(LibsyError::NoTargets)
367}
368
369#[derive(Clone, Copy)]
374enum CallPhase {
375 Routing,
376 Completion,
377}
378
379impl CallPhase {
380 const fn as_str(self) -> &'static str {
381 match self {
382 Self::Routing => "routing",
383 Self::Completion => "completion",
384 }
385 }
386}
387
388#[allow(clippy::too_many_arguments)]
390#[tracing::instrument(
391 target = "libsy",
392 name = "libsy.client_call",
393 skip_all,
394 fields(
395 algorithm = algorithm,
396 switchyard.algorithm = algorithm,
397 switchyard.candidate = index + 1,
398 switchyard.candidate_count = count,
399 switchyard.call_phase = phase.as_str(),
400 selected_model = %clients.provider_model(model_id),
401 selected_target = %model_id,
402 otel.kind = "client",
403 otel.name = %format_args!("chat {}", clients.provider_model(model_id)),
404 openinference.span.kind = "LLM",
405 gen_ai.operation.name = "chat",
406 gen_ai.request.model = %clients.provider_model(model_id),
407 gen_ai.request.stream = tracing::field::Empty,
408 gen_ai.request.temperature = tracing::field::Empty,
409 gen_ai.request.top_p = tracing::field::Empty,
410 gen_ai.request.top_k = tracing::field::Empty,
411 gen_ai.request.max_tokens = tracing::field::Empty,
412 gen_ai.request.reasoning.level = tracing::field::Empty,
413 gen_ai.output.type = tracing::field::Empty,
414 gen_ai.conversation.id = tracing::field::Empty,
415 server.address = tracing::field::Empty,
416 server.port = tracing::field::Empty,
417 gen_ai.response.id = tracing::field::Empty,
418 gen_ai.response.model = tracing::field::Empty,
419 gen_ai.usage.input_tokens = tracing::field::Empty,
420 gen_ai.usage.output_tokens = tracing::field::Empty,
421 gen_ai.usage.cache_read.input_tokens = tracing::field::Empty,
422 gen_ai.usage.cache_creation.input_tokens = tracing::field::Empty,
423 gen_ai.usage.cache_creation_1h.input_tokens = tracing::field::Empty,
424 gen_ai.usage.reasoning.output_tokens = tracing::field::Empty,
425 outcome = tracing::field::Empty,
426 otel.status_code = tracing::field::Empty,
427 error.type = tracing::field::Empty,
428 error = tracing::field::Empty,
429 )
430)]
431async fn call_one(
432 clients: &ClientRouter,
433 model_id: &ModelId,
434 request: Request,
435 algorithm: &str,
436 observe: &(dyn Fn(ModelCallObservation) + Send + Sync),
437 index: usize,
439 count: usize,
441 phase: CallPhase,
442) -> Result<Response> {
443 let is_routing = matches!(phase, CallPhase::Routing);
444 let span = tracing::Span::current();
445 observability::record_gen_ai_request(&span, &request.llm_request);
446 if let Some(session_id) = request
447 .metadata
448 .as_ref()
449 .and_then(|metadata| metadata.session_id.as_deref())
450 {
451 span.record("gen_ai.conversation.id", session_id);
452 }
453 let client = clients.route(model_id);
456 let started = Instant::now();
457 let mut routing_usage = None;
458 let result = async {
459 let response = async { client?.call(request).await }
460 .await
461 .map_err(|source| LibsyError::client_call(model_id.clone(), source));
462 let mut response = observability::observe_client_call(response)?;
464 if is_routing && let LlmResponse::Stream(chunks) = response.llm_response {
465 let (chunks, usage) = buffer_routing_stream(chunks)
466 .await
467 .map_err(|source| LibsyError::client_call(model_id.clone(), source))?;
468 routing_usage = usage;
469 response.llm_response = LlmResponse::Stream(chunks);
470 }
471 Ok(response)
472 }
473 .await;
474 let duration = started.elapsed();
475
476 let result = result.map(|mut response| {
477 response.set_served_model(model_id);
478 response
479 });
480 let result = if is_routing {
481 result
482 } else {
483 metrics::observe_routed_request(
484 algorithm,
485 clients.provider_model(model_id),
486 Some(duration),
487 result,
488 )
489 };
490 observe(ModelCallObservation {
491 selected_model: model_id.clone(),
492 is_success: result.is_ok(),
493 duration,
494 usage: result
495 .as_ref()
496 .ok()
497 .and_then(|response| response.llm_response.as_agg())
498 .map(|response| response.usage.clone())
499 .or(routing_usage),
500 });
501 result
502}
503
504async fn buffer_routing_stream(
506 mut chunks: LlmResponseStream,
507) -> std::result::Result<(LlmResponseStream, Option<Usage>), LlmClientError> {
508 let mut events = Vec::new();
509 let mut usage = None;
510 while let Some(event) = chunks.try_next().await? {
511 for chunk in event.normalized() {
512 match chunk {
513 LlmResponseChunk::Usage(latest) => usage = Some(latest.clone()),
514 LlmResponseChunk::DecodeError { message } => {
515 return Err(LlmClientError::ResponseTranslation(message.clone()));
516 }
517 LlmResponseChunk::StreamError { message } => {
518 return Err(LlmClientError::UpstreamHttp {
519 status: StatusCode::BAD_GATEWAY,
520 body: message.clone(),
521 });
522 }
523 _ => {}
524 }
525 }
526 events.push(event);
527 }
528 Ok((stream::iter(events.into_iter().map(Ok)).boxed(), usage))
529}
530
531fn fallback_reason(
533 error: &LibsyError,
534 allow_policy_fallback: bool,
535) -> Option<RoutingFallbackReason> {
536 let LibsyError::ClientCall { source, .. } = error else {
537 return None;
538 };
539 match source {
540 LlmClientError::ContextWindowExceeded { .. } => Some(RoutingFallbackReason::ContextWindow),
541 LlmClientError::Transport { .. } | LlmClientError::TemporarilyUnavailable => {
542 Some(RoutingFallbackReason::Unavailable)
543 }
544 LlmClientError::UpstreamHttp { status, body } => {
545 let policy_denial = serde_json::from_str::<Value>(body).is_ok_and(|value| {
546 value["error"]["code"].as_str() == Some("content_policy_violation")
547 });
548 if policy_denial && !allow_policy_fallback {
549 return None;
550 }
551 ((*status == StatusCode::BAD_REQUEST && policy_denial)
552 || matches!(
553 *status,
554 StatusCode::FORBIDDEN
555 | StatusCode::REQUEST_TIMEOUT
556 | StatusCode::TOO_MANY_REQUESTS
557 )
558 || status.is_server_error())
559 .then_some(RoutingFallbackReason::Unavailable)
560 }
561 _ => None,
562 }
563}
564
565fn continue_on(owner: StateOwner, algorithm: &str, mut request: Request) -> RoutingOutcome {
567 tracing::debug!(
568 target = %owner.model,
569 trigger = owner.trigger,
570 "routing Responses continuation to the model that holds its state"
571 );
572 if let Some(history) = owner.history {
573 let mut messages = Vec::with_capacity(history.len + request.llm_request.messages.len());
574 history.extend(&mut messages);
575 messages.append(&mut request.llm_request.messages);
576 request.llm_request.messages = messages;
577 }
578 let mut outcome = RoutingOutcome::route_to(owner.model, Vec::new(), request);
579 outcome.metadata = Some(OutcomeMetadata::new(
580 algorithm.to_string(),
581 Some(json!({"source": "stored_state", "trigger": owner.trigger})),
582 ));
583 outcome
584}
585
586struct StateOwner {
588 model: ModelId,
589 trigger: &'static str,
590 history: Option<Arc<MessageHistory>>,
591}
592
593struct MessageHistory {
594 parent: Option<Arc<Self>>,
595 segment: Arc<[Message]>,
596 len: usize,
597}
598
599impl MessageHistory {
600 fn new(parent: Option<Arc<Self>>, segment: Arc<[Message]>) -> Self {
601 let len = parent.as_ref().map_or(0, |history| history.len) + segment.len();
602 Self {
603 parent,
604 segment,
605 len,
606 }
607 }
608
609 fn extend(&self, messages: &mut Vec<Message>) {
610 let mut segments = Vec::new();
611 let mut history = Some(self);
612 while let Some(node) = history {
613 segments.push(node.segment.as_ref());
614 history = node.parent.as_deref();
615 }
616 for segment in segments.into_iter().rev() {
617 messages.extend_from_slice(segment);
618 }
619 }
620}
621
622#[derive(Clone)]
623struct StoredState {
624 model: ModelId,
625 history: Option<Arc<MessageHistory>>,
627}
628
629#[derive(Clone)]
630struct CanonicalInput {
631 parent: Option<Arc<MessageHistory>>,
632 messages: Arc<[Message]>,
633}
634
635#[derive(Default)]
640struct StateOwners {
641 by_id: HashMap<String, StoredState>,
642}
643
644const MAX_STATE_OWNERS: usize = 65_536;
645
646impl StateOwners {
647 fn owner(&self, id: &str) -> Option<&StoredState> {
648 self.by_id.get(id)
649 }
650
651 fn remember(
652 &mut self,
653 response_id: Option<&str>,
654 conversation_id: Option<&str>,
655 model: &ModelId,
656 history: Option<Arc<MessageHistory>>,
657 ) -> std::result::Result<(), LlmClientError> {
658 let materialized = history.is_some();
659 let ids = [
660 response_id,
661 conversation_id.filter(|id| Some(*id) != response_id),
662 ];
663 let mut new_ids = 0;
664 for id in ids.into_iter().flatten() {
665 match self.by_id.get(id) {
666 Some(owner) if owner.model != *model && !materialized => {
667 return Err(LlmClientError::ResponseStateConflict);
668 }
669 None => new_ids += 1,
670 _ => {}
671 }
672 }
673 if self.by_id.len() + new_ids > MAX_STATE_OWNERS {
674 return Err(LlmClientError::ResponseStateLimitExceeded {
675 limit: MAX_STATE_OWNERS,
676 });
677 }
678 let state = StoredState {
679 model: model.clone(),
680 history,
681 };
682 if let Some(id) = response_id {
683 if materialized {
684 self.by_id.insert(id.to_owned(), state.clone());
685 } else {
686 self.by_id
687 .entry(id.to_owned())
688 .or_insert_with(|| state.clone());
689 }
690 }
691 if let Some(id) = conversation_id {
692 self.by_id.entry(id.to_owned()).or_insert(state);
693 }
694 Ok(())
695 }
696}
697
698fn conversation_id(fields: &serde_json::Map<String, Value>) -> Option<&str> {
700 let conversation = fields.get("conversation")?;
701 conversation
702 .as_str()
703 .or_else(|| conversation.get("id").and_then(Value::as_str))
704}
705
706#[derive(Clone)]
714pub struct ClientRouter {
715 inner: Arc<ClientRouting>,
716 rate_limit_fallback_disabled: Arc<HashSet<ModelId>>,
717}
718
719struct ClientRouting {
720 routing: Routing,
721 decision_clients: HashMap<ModelId, Arc<dyn RoutedDecisionClient>>,
722 target_prompts: HashMap<ModelId, String>,
723 routing_answer_target: Option<ModelId>,
724 state_owners: Mutex<StateOwners>,
726 track_provider_state: bool,
728}
729
730enum Routing {
731 Single(Arc<dyn RoutedLlmClient>),
733 ByModel(HashMap<ModelId, Arc<dyn RoutedLlmClient>>),
735}
736
737impl ClientRouter {
738 pub fn new(by_model: HashMap<ModelId, Arc<dyn RoutedLlmClient>>) -> Self {
740 Self::new_with_target_prompts(by_model, HashMap::new(), None)
741 }
742
743 pub fn new_with_target_prompts(
752 by_model: HashMap<ModelId, Arc<dyn RoutedLlmClient>>,
753 target_prompts: HashMap<ModelId, String>,
754 routing_answer_target: Option<ModelId>,
755 ) -> Self {
756 let completion_targets = by_model.keys().cloned().collect::<Vec<_>>();
757 Self::new_with_completion_targets(
758 by_model,
759 target_prompts,
760 routing_answer_target,
761 &completion_targets,
762 )
763 }
764
765 pub fn new_with_completion_targets(
773 by_model: HashMap<ModelId, Arc<dyn RoutedLlmClient>>,
774 target_prompts: HashMap<ModelId, String>,
775 routing_answer_target: Option<ModelId>,
776 completion_targets: &[ModelId],
777 ) -> Self {
778 Self::new_with_decision_clients(
779 by_model,
780 HashMap::new(),
781 target_prompts,
782 routing_answer_target,
783 completion_targets,
784 )
785 }
786
787 pub fn new_with_decision_clients(
790 by_model: HashMap<ModelId, Arc<dyn RoutedLlmClient>>,
791 decision_clients: HashMap<ModelId, Arc<dyn RoutedDecisionClient>>,
792 target_prompts: HashMap<ModelId, String>,
793 routing_answer_target: Option<ModelId>,
794 completion_targets: &[ModelId],
795 ) -> Self {
796 let mut clients = completion_targets
797 .iter()
798 .filter_map(|model| by_model.get(model));
799 let spans_providers = clients
800 .next()
801 .is_some_and(|first| clients.any(|client| !Arc::ptr_eq(first, client)));
802 Self {
803 rate_limit_fallback_disabled: Arc::default(),
804 inner: Arc::new(ClientRouting {
805 routing: Routing::ByModel(by_model),
806 decision_clients,
807 target_prompts,
808 routing_answer_target,
809 state_owners: Mutex::default(),
810 track_provider_state: spans_providers,
811 }),
812 }
813 }
814
815 pub fn single(client: Arc<dyn RoutedLlmClient>) -> Self {
821 Self::single_with_decision_clients(client, HashMap::new())
822 }
823
824 pub fn single_with_decision_clients(
826 client: Arc<dyn RoutedLlmClient>,
827 decision_clients: HashMap<ModelId, Arc<dyn RoutedDecisionClient>>,
828 ) -> Self {
829 Self {
830 rate_limit_fallback_disabled: Arc::default(),
831 inner: Arc::new(ClientRouting {
832 routing: Routing::Single(client),
833 decision_clients,
834 target_prompts: HashMap::new(),
835 routing_answer_target: None,
836 state_owners: Mutex::default(),
837 track_provider_state: false,
838 }),
839 }
840 }
841
842 pub fn without_rate_limit_fallback(mut self, models: HashSet<ModelId>) -> Self {
848 self.rate_limit_fallback_disabled = Arc::new(models);
849 self
850 }
851
852 fn provider_model<'a>(&'a self, target: &'a ModelId) -> &'a ModelId {
854 self.route(target)
855 .map_or(target, |client| client.provider_model(target))
856 }
857
858 pub fn route(
863 &self,
864 model: &ModelId,
865 ) -> std::result::Result<&Arc<dyn RoutedLlmClient>, LlmClientError> {
866 match &self.inner.routing {
867 Routing::Single(client) => Ok(client),
868 Routing::ByModel(by_model) => {
869 by_model
870 .get(model)
871 .ok_or_else(|| LlmClientError::Configuration {
872 message: format!("no llm client is configured for model {model:?}"),
873 })
874 }
875 }
876 }
877
878 fn route_decision(
880 &self,
881 model: &ModelId,
882 ) -> std::result::Result<&Arc<dyn RoutedDecisionClient>, LlmClientError> {
883 self.inner
884 .decision_clients
885 .get(model)
886 .ok_or_else(|| LlmClientError::Configuration {
887 message: format!("no decision client is configured for model {model:?}"),
888 })
889 }
890
891 fn stored_state_owner(&self, request: &Request) -> Option<StateOwner> {
893 let fields = &request.llm_request.extensions.fields;
894 let (trigger, id) = match fields.get("previous_response_id").and_then(Value::as_str) {
895 Some(id) => ("previous_response_id", id),
896 None => ("conversation", conversation_id(fields)?),
897 };
898 let state = self.inner.state_owners.lock().owner(id)?.clone();
899 Some(StateOwner {
900 model: state.model,
901 trigger,
902 history: state.history,
903 })
904 }
905
906 fn canonical_input(&self, request: &Request) -> CanonicalInput {
907 let parent = request
908 .llm_request
909 .extensions
910 .fields
911 .get("previous_response_id")
912 .and_then(Value::as_str)
913 .and_then(|id| self.inner.state_owners.lock().owner(id)?.history.clone());
914 let parent_len = parent.as_ref().map_or(0, |history| history.len);
915 let (parent, messages) = match request.llm_request.messages.get(parent_len..) {
916 Some(messages) => (parent, messages.to_vec()),
917 None => (None, request.llm_request.messages.clone()),
918 };
919 CanonicalInput {
920 parent,
921 messages: Arc::from(messages),
922 }
923 }
924
925 fn remember_state_owner(&self, request: &Request, mut response: Response) -> Result<Response> {
927 let Some(model) = response.served_model().cloned() else {
928 return Ok(response);
929 };
930 let fields = &request.llm_request.extensions.fields;
931 let store = fields.get("store").and_then(Value::as_bool) != Some(false);
932 let conversation = conversation_id(fields).map(str::to_owned);
933 let responses_format = WireFormat::OpenAiResponses.into();
934 let responses_request = request
935 .llm_request
936 .preservation
937 .requests
938 .contains_key(&responses_format)
939 || fields.contains_key("previous_response_id")
940 || fields.contains_key("conversation");
941 if !responses_request && !self.inner.track_provider_state {
942 return Ok(response);
943 }
944 let canonical_input = responses_request.then(|| self.canonical_input(request));
945 response.llm_response = match response.llm_response {
946 LlmResponse::Agg(agg) => {
947 if let Some(body) = agg.preservation.responses.get(&responses_format) {
948 if self.inner.track_provider_state {
949 self.remember_response(body, &model, store, conversation.as_deref())
950 .map_err(|error| LibsyError::client_call(model.clone(), error))?;
951 }
952 } else if let Some(input) = &canonical_input {
953 self.remember_canonical_response(&agg, &model, store, input)
954 .map_err(|error| LibsyError::client_call(model.clone(), error))?;
955 }
956 LlmResponse::Agg(agg)
957 }
958 LlmResponse::Stream(stream) => {
959 let router = self.clone();
960 LlmResponse::Stream(Box::pin(futures_util::stream::try_unfold(
961 (
962 stream,
963 router,
964 model,
965 conversation,
966 canonical_input,
967 ResponseAccumulator::new(),
968 false,
969 true,
970 ),
971 move |(
972 mut stream,
973 router,
974 model,
975 conversation,
976 canonical_input,
977 mut accumulator,
978 mut native_responses,
979 mut valid,
980 )| async move {
981 let Some(event) = stream.next().await else {
982 if let Some(input) = &canonical_input
983 && !native_responses
984 && valid
985 {
986 router.remember_canonical_response(
987 &accumulator.finish(),
988 &model,
989 store,
990 input,
991 )?;
992 }
993 return Ok(None);
994 };
995 let event = event?;
996 if let Some(preserved) = event.preservation()
997 && preserved.source().as_str() == WireFormat::OpenAiResponses.as_str()
998 {
999 native_responses = true;
1000 if router.inner.track_provider_state
1001 && let Some(body) = preserved.raw().get("response")
1002 {
1003 router.remember_response(
1004 body,
1005 &model,
1006 store,
1007 conversation.as_deref(),
1008 )?;
1009 }
1010 }
1011 if responses_request {
1012 for chunk in event.normalized() {
1013 if matches!(
1014 chunk,
1015 LlmResponseChunk::DecodeError { .. }
1016 | LlmResponseChunk::StreamError { .. }
1017 ) {
1018 valid = false;
1019 } else {
1020 accumulator.push(chunk.clone());
1021 }
1022 }
1023 }
1024 Ok(Some((
1025 event,
1026 (
1027 stream,
1028 router,
1029 model,
1030 conversation,
1031 canonical_input,
1032 accumulator,
1033 native_responses,
1034 valid,
1035 ),
1036 )))
1037 },
1038 )))
1039 }
1040 };
1041 Ok(response)
1042 }
1043
1044 fn remember_response(
1047 &self,
1048 body: &Value,
1049 model: &ModelId,
1050 store: bool,
1051 conversation: Option<&str>,
1052 ) -> std::result::Result<(), LlmClientError> {
1053 let mut owners = self.inner.state_owners.lock();
1054 let response_id = body
1055 .get("id")
1056 .and_then(Value::as_str)
1057 .filter(|_| body.get("store").and_then(Value::as_bool).unwrap_or(store));
1058 let conversation = body.as_object().and_then(conversation_id).or(conversation);
1059 if let Err(error) = owners.remember(response_id, conversation, model, None) {
1060 if matches!(error, LlmClientError::ResponseStateLimitExceeded { .. }) {
1061 tracing::warn!(%error, "Responses state owner capacity reached; state was not retained");
1062 } else {
1063 tracing::error!(%error, "could not record Responses state owner");
1064 return Err(error);
1065 }
1066 }
1067 Ok(())
1068 }
1069
1070 fn remember_canonical_response(
1071 &self,
1072 response: &AggLlmResponse,
1073 model: &ModelId,
1074 store: bool,
1075 input: &CanonicalInput,
1076 ) -> std::result::Result<(), LlmClientError> {
1077 let response_id = response.id.as_deref().filter(|_| store);
1078 if response_id.is_none() {
1079 return Ok(());
1080 }
1081 let mut segment = input.messages.to_vec();
1082 segment.extend(
1083 response
1084 .outputs
1085 .iter()
1086 .map(|output| Message::new(output.role, output.content.clone())),
1087 );
1088 let history = Arc::new(MessageHistory::new(
1089 input.parent.clone(),
1090 Arc::from(segment),
1091 ));
1092 let result =
1093 self.inner
1094 .state_owners
1095 .lock()
1096 .remember(response_id, None, model, Some(history));
1097 if let Err(error) = result {
1098 if matches!(error, LlmClientError::ResponseStateLimitExceeded { .. }) {
1099 tracing::warn!(%error, "cross-format Responses state capacity reached; history was not retained");
1100 return Ok(());
1101 }
1102 tracing::error!(%error, "could not record cross-format Responses state");
1103 return Err(error);
1104 }
1105 Ok(())
1106 }
1107
1108 fn prepare_completion_request(&self, mut request: Request, target: &ModelId) -> Request {
1110 let prompt = self.inner.target_prompts.get(target).map(String::as_str);
1111 prepare_request_for_target(&mut request.llm_request, target, prompt);
1112 request
1113 }
1114
1115 fn prepare_routing_request(&self, mut request: Request, target: &ModelId) -> Request {
1117 let prompt = if self.inner.routing_answer_target.as_ref() == Some(target) {
1118 self.inner.target_prompts.get(target).map(String::as_str)
1119 } else {
1120 None
1121 };
1122 prepare_request_for_target(&mut request.llm_request, target, prompt);
1123 request
1124 }
1125}
1126
1127impl FromIterator<(ModelId, Arc<dyn RoutedLlmClient>)> for ClientRouter {
1128 fn from_iter<I: IntoIterator<Item = (ModelId, Arc<dyn RoutedLlmClient>)>>(iter: I) -> Self {
1129 Self::new(iter.into_iter().collect())
1130 }
1131}
1132
1133#[cfg(test)]
1134mod tests {
1135 use std::collections::{BTreeMap, BTreeSet};
1136
1137 use super::*;
1138 use async_trait::async_trait;
1139 use futures::StreamExt;
1140 use http::StatusCode;
1141 use switchyard_libsy::{Driver, RoutingOutcome};
1142 use switchyard_protocol::{
1143 ContentBlock, LlmResponse, LlmResponseChunk, LlmResponseStreamEvent, Message, Role,
1144 ToolResult, completion_text, text_request, text_response,
1145 };
1146 use wiremock::matchers::method;
1147 use wiremock::{Mock, MockServer, ResponseTemplate};
1148
1149 use crate::{Backend, HttpBackendConfig, ModelConfig, TranslatingLlmClient};
1150 use switchyard_protocol::Category;
1151
1152 struct CandidateAlgorithm {}
1153
1154 struct AnsweredAlgorithm {
1155 model: ModelId,
1156 }
1157
1158 #[async_trait]
1159 impl Algorithm for CandidateAlgorithm {
1160 fn name(&self) -> &str {
1161 "candidate_test"
1162 }
1163
1164 async fn route(
1165 self: Arc<Self>,
1166 driver: Driver,
1167 request: Request,
1168 ) -> Result<RoutingOutcome> {
1169 let models = driver.models_for(&Category::Any);
1170 let selected_model = models.first().cloned().ok_or(LibsyError::NoTargets)?;
1171 Ok(RoutingOutcome::route_to(
1172 selected_model,
1173 models.iter().skip(1).cloned().collect(),
1174 request,
1175 ))
1176 }
1177 }
1178
1179 #[async_trait]
1180 impl Algorithm for AnsweredAlgorithm {
1181 fn name(&self) -> &str {
1182 "answered_test"
1183 }
1184
1185 async fn route(
1186 self: Arc<Self>,
1187 driver: Driver,
1188 request: Request,
1189 ) -> Result<RoutingOutcome> {
1190 let response = driver
1191 .call_model(request.clone(), vec![self.model.clone()])
1192 .await?;
1193 Ok(RoutingOutcome::answered(
1194 self.model.clone(),
1195 request,
1196 response,
1197 ))
1198 }
1199 }
1200
1201 struct RewritingAlgorithm {}
1203
1204 #[async_trait]
1205 impl Algorithm for RewritingAlgorithm {
1206 fn name(&self) -> &str {
1207 "rewriting_test"
1208 }
1209
1210 async fn route(
1211 self: Arc<Self>,
1212 _driver: Driver,
1213 request: Request,
1214 ) -> Result<RoutingOutcome> {
1215 let mut rewritten = request.clone();
1216 rewritten
1217 .llm_request
1218 .messages
1219 .push(Message::text(Role::User, "for weak only"));
1220 let mut outcome =
1221 RoutingOutcome::route_to("weak".into(), vec!["strong".into()], rewritten);
1222 outcome.fallback_request = Some(request);
1223 Ok(outcome)
1224 }
1225 }
1226
1227 #[derive(Clone, Copy)]
1228 enum FirstOutcome {
1229 ContextWindow,
1230 ContentPolicy(StatusCode),
1231 Unauthorized,
1232 StreamSuccess,
1233 MidStreamError,
1234 }
1235
1236 struct CandidateClient {
1237 calls: Mutex<Vec<ModelId>>,
1238 requests: Mutex<Vec<Request>>,
1239 first: FirstOutcome,
1240 }
1241
1242 #[async_trait]
1243 impl RoutedLlmClient for CandidateClient {
1244 async fn call(&self, request: Request) -> std::result::Result<Response, LlmClientError> {
1245 let model = request.model_id().unwrap_or_default();
1246 self.calls.lock().push(model.clone());
1247 self.requests.lock().push(request);
1248 if model == "weak" {
1249 return match self.first {
1250 FirstOutcome::ContextWindow => Err(LlmClientError::ContextWindowExceeded {
1251 model,
1252 message: "too long".to_string(),
1253 }),
1254 FirstOutcome::ContentPolicy(status) => Err(LlmClientError::UpstreamHttp {
1255 status,
1256 body: r#"{"error":{"code":"content_policy_violation","message":"request blocked by content policy","type":"invalid_request_error"},"metadata":{"documentation_section":"context window"}}"#.to_string(),
1257 }),
1258 FirstOutcome::Unauthorized => Err(LlmClientError::UpstreamHttp {
1259 status: StatusCode::UNAUTHORIZED,
1260 body: "unauthorized".to_string(),
1261 }),
1262 FirstOutcome::StreamSuccess => Ok(stream_response(vec![
1263 LlmResponseChunk::TextDelta {
1264 index: 0,
1265 text: "streamed".to_string(),
1266 },
1267 LlmResponseChunk::MessageStop {
1268 reason: Some("stop".to_string()),
1269 },
1270 ])),
1271 FirstOutcome::MidStreamError => Ok(stream_response(vec![
1272 LlmResponseChunk::TextDelta {
1273 index: 0,
1274 text: "partial".to_string(),
1275 },
1276 LlmResponseChunk::StreamError {
1277 message: "stream failed".to_string(),
1278 },
1279 LlmResponseChunk::MessageStop { reason: Some("stop".into()) },
1280 ])),
1281 };
1282 }
1283 Ok(Response {
1284 llm_response: LlmResponse::Agg(AggLlmResponse {
1285 id: Some("candidate_response".into()),
1286 ..text_response(Some(model.to_string()), model)
1287 }),
1288 metadata: None,
1289 upstream_headers: http::HeaderMap::new(),
1290 })
1291 }
1292 }
1293
1294 fn stream_response(chunks: Vec<LlmResponseChunk>) -> Response {
1295 Response {
1296 llm_response: LlmResponse::Stream(
1297 futures::stream::iter(
1298 chunks
1299 .into_iter()
1300 .map(|chunk| Ok(LlmResponseStreamEvent::from(chunk))),
1301 )
1302 .boxed(),
1303 ),
1304 metadata: None,
1305 upstream_headers: http::HeaderMap::new(),
1306 }
1307 }
1308
1309 fn request() -> Request {
1310 Request {
1311 llm_request: text_request(Some("auto".to_string()), "hello".to_string()),
1312 raw_request: None,
1313 metadata: None,
1314 }
1315 }
1316
1317 #[tokio::test]
1318 async fn system_one_serves_decisions_and_returns_failures_to_the_classifier()
1319 -> std::result::Result<(), Box<dyn std::error::Error>> {
1320 use std::time::Duration;
1321 use switchyard_libsy::{
1322 CapabilityJudgeConfig, DecisionJudgeConfig, LlmClassifierConfig, LlmTaskClassifier,
1323 TaskClassifierConfig,
1324 };
1325 use switchyard_protocol::{
1326 BooleanEstimate, DecisionRequest, DecisionValue, Probability, ScoreValue,
1327 };
1328 use wiremock::matchers::{body_partial_json, header, path};
1329
1330 let server = MockServer::start().await;
1331 let client = Arc::new(crate::SystemOneClient::new(
1332 format!("{}/v1/systemone", server.uri()).parse()?,
1333 "test-key".into(),
1334 Duration::from_secs(2),
1335 )?);
1336 assert!(matches!(
1337 crate::SystemOneClient::new(
1338 format!("{}/v1/systemone", server.uri()).parse()?,
1339 "invalid\nkey".into(),
1340 Duration::from_secs(2),
1341 ),
1342 Err(LlmClientError::Configuration { .. })
1343 ));
1344 let mixed: DecisionRequest = serde_json::from_value(json!({
1345 "model": "judge", "input": [Message::text(Role::User, "A simple task")],
1346 "questions": {
1347 "boolean": {"instructions": "Is it simple?", "kind": {"type": "boolean", "data": {
1348 "true_description": ["Simple"], "false_description": null
1349 }}},
1350 "route": {"instructions": "Choose", "kind": {"type": "choice", "data": {
1351 "options": [{"id": "capable", "description": "A wins"}, {"id": "efficient", "description": "B wins"}]
1352 }}},
1353 "score": {"instructions": "Rate difficulty", "kind": {"type": "score", "data": {
1354 "levels": ["Easy", {"description": "Hard"}]
1355 }}}
1356 }
1357 }))?;
1358 let mut body = json!({
1359 "model": "jev-1.13.0", "usage": {"input_tokens": 42, "output_tokens": 3},
1360 "answers": {
1361 "boolean": {"type": "noul", "noul": 0.9},
1362 "route": {"type": "choice", "choice": "efficient", "confidence": 0.7,
1363 "probabilities": {"capable": 0.2, "efficient": 0.8}},
1364 "score": {"type": "score", "score": 0.6, "confidence": 0.3,
1365 "legend": {"0": "Easy", "1": {"description": "Hard"}},
1366 "probabilities": {"1": 0.6, "0": 0.4}}
1367 }
1368 });
1369 Mock::given(method("POST"))
1370 .and(path("/v1/systemone"))
1371 .and(header("authorization", "Bearer test-key"))
1372 .and(body_partial_json(json!({
1373 "model": "judge", "state": mixed.input,
1374 "questions": {
1375 "boolean": {"type": "noul", "criteria": {"true": ["Simple"]}},
1376 "route": {"type": "choice", "instructions": "Choose",
1377 "criteria": {"capable": "A wins", "efficient": "B wins"}},
1378 "score": {"type": "score", "criteria": ["Easy", {"description": "Hard"}]}
1379 }
1380 })))
1381 .respond_with(ResponseTemplate::new(200).set_body_json(&body))
1382 .expect(1)
1383 .mount(&server)
1384 .await;
1385 let response = client.call(mixed.clone()).await?;
1386 assert_eq!(response.model.as_deref(), Some("jev-1.13.0"));
1387 assert_eq!(response.id, None);
1388 assert_eq!(response.usage.input_tokens, Some(42));
1389 assert_eq!(response.usage.total_tokens, None);
1390 assert_eq!(
1391 response.answers["boolean"].value,
1392 DecisionValue::Boolean(BooleanEstimate::ProbabilityTrue(Probability(0.9)))
1393 );
1394 assert_eq!(
1395 response.answers["score"].value,
1396 DecisionValue::Score {
1397 value: ScoreValue(0.6),
1398 probabilities: Some(vec![Probability(0.4), Probability(0.6)]),
1399 }
1400 );
1401 assert_eq!(
1402 response.answers["route"].provider_confidence.map(|c| c.0),
1403 Some(0.7)
1404 );
1405 server.verify().await;
1406
1407 let mut partial = body.clone();
1408 partial["answers"]
1409 .as_object_mut()
1410 .unwrap()
1411 .remove("boolean");
1412 let mut invalid_responses = vec![json!({"answers": {}}), partial];
1413 for (id, answer) in [
1414 ("unknown-question", json!({"type": "noul", "noul": 0.9})),
1415 ("boolean", json!({"type": "score", "score": 0.0})),
1416 (
1417 "route",
1418 json!({"type": "choice", "choice": "unknown-option",
1419 "probabilities": {"capable": 0.2, "efficient": 0.8}}),
1420 ),
1421 ("score", json!({"type": "score", "score": -0.1})),
1422 (
1423 "score",
1424 json!({"type": "score", "score": 1.1,
1425 "probabilities": {"0": 0.4, "1": 0.6}}),
1426 ),
1427 ] {
1428 let mut invalid = body.clone();
1429 invalid["answers"][id] = answer;
1430 if id == "unknown-question" {
1431 invalid["answers"]
1432 .as_object_mut()
1433 .unwrap()
1434 .remove("boolean");
1435 }
1436 invalid_responses.push(invalid);
1437 }
1438 for template in invalid_responses
1439 .into_iter()
1440 .map(|body| ResponseTemplate::new(200).set_body_json(body))
1441 .chain([ResponseTemplate::new(200).set_body_string("invalid JSON")])
1442 {
1443 server.reset().await;
1444 Mock::given(method("POST"))
1445 .respond_with(template)
1446 .expect(1)
1447 .mount(&server)
1448 .await;
1449 assert!(matches!(
1450 client.call(mixed.clone()).await,
1451 Err(LlmClientError::ResponseTranslation(_) | LlmClientError::InvalidResponse { .. })
1452 ));
1453 server.verify().await;
1454 }
1455
1456 body["answers"].as_object_mut().unwrap().remove("boolean");
1457 body["answers"].as_object_mut().unwrap().remove("score");
1458 let models = Arc::new(RuntimeModels::new(HashMap::from([
1459 (Category::Judge, vec!["judge".into()]),
1460 (Category::Capable, vec!["capable".into()]),
1461 (Category::Efficient, vec!["efficient".into()]),
1462 (Category::Any, vec!["efficient".into(), "capable".into()]),
1463 ])));
1464 let algorithm = || -> Result<Arc<dyn Algorithm>> {
1465 Ok(Arc::new(LlmTaskClassifier::new(
1466 LlmClassifierConfig::Capability {
1467 config: TaskClassifierConfig {
1468 judge: CapabilityJudgeConfig::Decision(DecisionJudgeConfig {
1469 cutoff: Some(0.4),
1470 ..DecisionJudgeConfig::default()
1471 }),
1472 ..TaskClassifierConfig::default()
1473 },
1474 },
1475 )?))
1476 };
1477 for (template, expected, is_success) in [
1478 (
1479 ResponseTemplate::new(200).set_body_json(&body),
1480 "efficient",
1481 true,
1482 ),
1483 (
1484 ResponseTemplate::new(503).set_body_string("unavailable"),
1485 "capable",
1486 false,
1487 ),
1488 ] {
1489 server.reset().await;
1490 Mock::given(method("POST"))
1491 .and(path("/v1/systemone"))
1492 .respond_with(template)
1493 .expect(2)
1494 .mount(&server)
1495 .await;
1496 let llm = Arc::new(CandidateClient {
1497 calls: Mutex::default(),
1498 requests: Mutex::default(),
1499 first: FirstOutcome::Unauthorized,
1500 });
1501 let clients = ClientRouter::single_with_decision_clients(
1502 llm.clone(),
1503 HashMap::from([(
1504 ModelId::from("judge"),
1505 client.clone() as Arc<dyn RoutedDecisionClient>,
1506 )]),
1507 );
1508 let outcome = decide(algorithm()?, clients.clone(), request(), models.clone()).await?;
1509 assert_eq!(outcome.selected_model_id()?, expected);
1510 assert!(llm.calls.lock().is_empty());
1511 let events = Arc::new(Mutex::new(Vec::new()));
1512 let captured = events.clone();
1513 let (selected, _) = run(
1514 algorithm()?,
1515 clients,
1516 request(),
1517 models.clone(),
1518 Some(Arc::new(move |event| captured.lock().push(event))),
1519 )
1520 .await?;
1521 assert_eq!(selected, expected);
1522 assert_eq!(&*llm.calls.lock(), &[ModelId::from(expected)]);
1523 {
1524 let events = events.lock();
1525 let RunObservation::DecisionCall(call) = &events[0] else {
1526 panic!("missing decision observation")
1527 };
1528 assert_eq!(call.selected_model, "judge");
1529 assert_eq!(call.is_success, is_success);
1530 assert_eq!(
1531 call.usage.as_ref().and_then(|u| u.input_tokens),
1532 is_success.then_some(42)
1533 );
1534 assert!(
1535 events
1536 .iter()
1537 .any(|event| matches!(event, RunObservation::AnswerCall(_)))
1538 );
1539 }
1540 server.verify().await;
1541 }
1542 Ok(())
1543 }
1544
1545 fn instruction_text(request: &Request) -> Vec<&str> {
1546 request
1547 .llm_request
1548 .instructions
1549 .iter()
1550 .flat_map(|instruction| &instruction.content)
1551 .filter_map(|block| match block {
1552 ContentBlock::Text { text } => Some(text.as_str()),
1553 _ => None,
1554 })
1555 .collect()
1556 }
1557
1558 async fn run_candidates(
1559 first: FirstOutcome,
1560 ) -> (Arc<CandidateClient>, Result<(ModelId, Response)>) {
1561 let client = Arc::new(CandidateClient {
1562 calls: Mutex::new(Vec::new()),
1563 requests: Mutex::new(Vec::new()),
1564 first,
1565 });
1566 let algorithm = Arc::new(CandidateAlgorithm {});
1567 let models = to_category_map(&["weak", "strong"]);
1568 let result = run(
1569 algorithm,
1570 ClientRouter::single(client.clone()),
1571 request(),
1572 models,
1573 None,
1574 )
1575 .await;
1576 (client, result)
1577 }
1578
1579 fn to_category_map(names: &[&str]) -> Arc<RuntimeModels> {
1580 Arc::new(RuntimeModels::new(
1581 [(
1582 Category::Any,
1583 names.iter().map(|name| ModelId::from(*name)).collect(),
1584 )]
1585 .into(),
1586 ))
1587 }
1588
1589 #[tokio::test]
1590 async fn answered_outcome_does_not_make_a_second_model_call() -> Result<()> {
1591 let client = Arc::new(CandidateClient {
1592 calls: Mutex::new(Vec::new()),
1593 requests: Mutex::new(Vec::new()),
1594 first: FirstOutcome::StreamSuccess,
1595 });
1596 let observations = Arc::new(Mutex::new(Vec::new()));
1597 let observed = Arc::clone(&observations);
1598 let observer: RunObserver = Arc::new(move |event| observed.lock().push(event));
1599
1600 let (selected, response) = run(
1601 Arc::new(AnsweredAlgorithm {
1602 model: "weak".into(),
1603 }),
1604 ClientRouter::single(client.clone()),
1605 request(),
1606 Arc::new(RuntimeModels::default()),
1607 Some(observer),
1608 )
1609 .await?;
1610
1611 assert_eq!(selected, "weak");
1612 assert_eq!(response.served_model().map(ModelId::as_str), Some("weak"));
1613 assert_eq!(&*client.calls.lock(), &[ModelId::from("weak")]);
1614 let observations = observations.lock();
1615 assert!(matches!(observations[0], RunObservation::AnswerCall(_)));
1616 assert!(matches!(
1617 observations[1],
1618 RunObservation::Outcome(ref metadata) if metadata.algorithm == "answered_test"
1619 ));
1620 assert!(matches!(
1621 observations[2],
1622 RunObservation::RoutingOverhead(_)
1623 ));
1624 assert_eq!(observations.len(), 3);
1625 Ok(())
1626 }
1627
1628 #[tokio::test]
1629 async fn served_answer_is_reported_before_the_client_sees_it_end() -> Result<()> {
1630 #[derive(Default)]
1631 struct ReportingAlgorithm {
1632 answers: Mutex<Vec<(String, bool)>>,
1633 tracking: bool,
1634 }
1635
1636 #[async_trait]
1637 impl Algorithm for ReportingAlgorithm {
1638 fn name(&self) -> &str {
1639 "reporting_test"
1640 }
1641
1642 async fn route(
1643 self: Arc<Self>,
1644 _driver: Driver,
1645 request: Request,
1646 ) -> Result<RoutingOutcome> {
1647 Ok(RoutingOutcome::route_to("weak".into(), Vec::new(), request))
1648 }
1649
1650 fn tracks_answers(&self) -> bool {
1651 self.tracking
1652 }
1653
1654 async fn observe_answer(&self, session_id: &str, answer: AnswerSummary) {
1655 self.answers
1656 .lock()
1657 .push((session_id.to_string(), answer.final_reply));
1658 }
1659 }
1660
1661 assert!(!CandidateAlgorithm {}.tracks_answers());
1662 for (tracking, first, final_reply, cancel) in [
1663 (false, FirstOutcome::StreamSuccess, true, false),
1664 (true, FirstOutcome::StreamSuccess, true, false),
1665 (true, FirstOutcome::MidStreamError, false, false),
1666 (true, FirstOutcome::StreamSuccess, false, true),
1667 ] {
1668 let algorithm = Arc::new(ReportingAlgorithm {
1669 tracking,
1670 ..Default::default()
1671 });
1672 let client = Arc::new(CandidateClient {
1673 calls: Mutex::new(Vec::new()),
1674 requests: Mutex::new(Vec::new()),
1675 first,
1676 });
1677 let request = Request {
1678 metadata: Some(switchyard_protocol::Metadata {
1679 session_id: Some("session-1".to_string()),
1680 ..Default::default()
1681 }),
1682 ..request()
1683 };
1684 let (_, response) = run(
1685 algorithm.clone(),
1686 ClientRouter::single(client),
1687 request,
1688 to_category_map(&["weak"]),
1689 None,
1690 )
1691 .await?;
1692 let LlmResponse::Stream(mut stream) = response.llm_response else {
1693 panic!("expected a streamed answer");
1694 };
1695 if !tracking {
1696 assert!(algorithm.answers.lock().is_empty());
1697 while let Some(item) = stream.next().await {
1698 assert!(item.is_ok());
1699 }
1700 assert!(algorithm.answers.lock().is_empty());
1701 continue;
1702 }
1703 if cancel {
1704 stream.next().await;
1705 drop(stream);
1706 assert_eq!(
1707 *algorithm.answers.lock(),
1708 [("session-1".to_string(), false)]
1709 );
1710 continue;
1711 }
1712 while let Some(item) = stream.next().await {
1713 let stops = item.as_ref().is_ok_and(|event| {
1714 event
1715 .normalized()
1716 .iter()
1717 .any(|chunk| matches!(chunk, LlmResponseChunk::MessageStop { .. }))
1718 });
1719 if stops {
1720 assert_eq!(algorithm.answers.lock().len(), 2);
1721 }
1722 }
1723 assert_eq!(
1724 *algorithm.answers.lock(),
1725 if final_reply {
1726 vec![
1727 ("session-1".to_string(), false),
1728 ("session-1".to_string(), true),
1729 ]
1730 } else {
1731 vec![
1732 ("session-1".to_string(), false),
1733 ("session-1".to_string(), false),
1734 ]
1735 }
1736 );
1737 }
1738 Ok(())
1739 }
1740
1741 #[test]
1742 fn answer_observation_keeps_call_order() {
1743 let pending = Some(Arc::new(Mutex::new(
1744 ["answer", "judge"]
1745 .map(|model| ModelCallObservation {
1746 selected_model: model.into(),
1747 is_success: true,
1748 duration: std::time::Duration::ZERO,
1749 usage: None,
1750 })
1751 .map(RunObservation::LlmCall)
1752 .into(),
1753 )));
1754 let emitted = Arc::new(Mutex::new(Vec::new()));
1755 let captured = Arc::clone(&emitted);
1756 let observer: RunObserver = Arc::new(move |event| captured.lock().push(event));
1757 let answer = ModelId::from("answer");
1758
1759 emit_routing_observations(&Some(observer), &pending, Some(&answer));
1760
1761 assert!(matches!(
1762 &emitted.lock()[..],
1763 [RunObservation::AnswerCall(answer), RunObservation::LlmCall(judge)]
1764 if answer.selected_model == "answer" && judge.selected_model == "judge"
1765 ));
1766 }
1767
1768 #[tokio::test]
1769 async fn responses_state_tracking_uses_upstream_storage_and_format() -> Result<()> {
1770 for (responses, store, conversation) in [
1771 (true, true, None),
1772 (true, false, None),
1773 (true, false, Some(json!("conv_plain"))),
1774 (true, false, Some(json!({"id": "conv_object"}))),
1775 (false, true, None),
1776 ] {
1777 for stream in [false, true] {
1778 let server = MockServer::start().await;
1779 Mock::given(method("POST"))
1780 .respond_with(move |request: &wiremock::Request| {
1781 let body: Value = serde_json::from_slice(&request.body).expect("request JSON");
1782 let response = if responses {
1783 json!({
1784 "id": "resp_seed", "object": "response", "model": "weak",
1785 "status": "completed", "output": [], "store": body["store"],
1786 "conversation": body["conversation"],
1787 })
1788 } else if stream {
1789 json!({"id": "resp_seed", "model": "weak", "choices": [{"index": 0, "delta": {"content": "ok"}, "finish_reason": "stop"}]})
1790 } else {
1791 json!({"id": "resp_seed", "model": "weak", "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}]})
1792 };
1793 if stream {
1794 let event = if responses { json!({"type": "response.completed", "response": response}) } else { response };
1795 ResponseTemplate::new(200)
1796 .insert_header("content-type", "text/event-stream")
1797 .set_body_string(if responses {
1798 format!("event: response.completed\ndata: {event}\n\n")
1799 } else {
1800 format!("data: {event}\n\ndata: [DONE]\n\n")
1801 })
1802 } else {
1803 ResponseTemplate::new(200).set_body_json(response)
1804 }
1805 })
1806 .mount(&server).await;
1807 let mut clients = HashMap::new();
1808 for model in ["weak", "strong"] {
1809 let config = HttpBackendConfig {
1810 base_url: server.uri(),
1811 api_key: None,
1812 forward_auth: false,
1813 extra_headers: BTreeMap::new(),
1814 extra_body: BTreeMap::from([("store".to_string(), json!(store))]),
1815 omit_body_fields: BTreeSet::new(),
1816 reasoning_effort: None,
1817 max_retries: 0,
1818 failure_cooldown: std::time::Duration::ZERO,
1819 timeout: None,
1820 };
1821 let backend = if responses {
1822 Backend::OpenAiResponses(config)
1823 } else {
1824 Backend::OpenAiChat(config)
1825 };
1826 let client: Arc<dyn RoutedLlmClient> = Arc::new(
1827 TranslatingLlmClient::new(&[ModelConfig::new(model, backend, None)])
1828 .map_err(|error| LibsyError::external("building test client", error))?,
1829 );
1830 clients.insert(ModelId::from(model), client);
1831 }
1832 let clients = ClientRouter::new(clients);
1833 let mut seed = request();
1834 seed.llm_request.stream = stream;
1835 if let Some(conversation) = &conversation {
1836 seed.llm_request
1837 .extensions
1838 .fields
1839 .insert("conversation".to_string(), conversation.clone());
1840 }
1841 let (_, response) = run(
1842 Arc::new(switchyard_libsy::Passthrough),
1843 clients.clone(),
1844 seed,
1845 to_category_map(&["weak"]),
1846 None,
1847 )
1848 .await?;
1849 response
1850 .llm_response
1851 .into_agg()
1852 .await
1853 .map_err(|error| LibsyError::client_call("weak", error))?;
1854 let mut continuations = vec![(
1855 "previous_response_id",
1856 json!("resp_seed"),
1857 responses && store,
1858 )];
1859 if let Some(conversation) = conversation.clone() {
1860 continuations.push(("conversation", conversation, true));
1861 }
1862 for (field, id, pinned) in continuations {
1863 let mut follow = request();
1864 follow
1865 .llm_request
1866 .extensions
1867 .fields
1868 .insert(field.to_string(), id);
1869 let outcome = decide(
1870 Arc::new(switchyard_libsy::Passthrough),
1871 clients.clone(),
1872 follow,
1873 to_category_map(&["strong"]),
1874 )
1875 .await?;
1876 assert_eq!(
1877 outcome.selected_model_id()?.as_str(),
1878 if pinned { "weak" } else { "strong" },
1879 "responses={responses}, store={store}, stream={stream}, field={field}"
1880 );
1881 }
1882 let requests = server.received_requests().await.expect("request recording");
1883 assert_eq!(requests.len(), 1);
1884 assert_eq!(
1885 serde_json::from_slice::<Value>(&requests[0].body).expect("request JSON")["store"],
1886 store
1887 );
1888 }
1889 }
1890 Ok(())
1891 }
1892
1893 #[tokio::test]
1894 async fn responses_stored_tool_continuation_materializes_for_anthropic() -> Result<()> {
1895 let server = MockServer::start().await;
1896 Mock::given(method("POST"))
1897 .respond_with(|request: &wiremock::Request| {
1898 let body: Value = serde_json::from_slice(&request.body).expect("request JSON");
1899 let blocks = body["messages"]
1900 .as_array()
1901 .into_iter()
1902 .flatten()
1903 .filter_map(|message| message["content"].as_array())
1904 .flatten()
1905 .collect::<Vec<_>>();
1906 let has_result = blocks.iter().any(|block| block["type"] == "tool_result");
1907 let has_call = blocks.iter().any(|block| block["type"] == "tool_use");
1908 let has_thanks = blocks
1909 .iter()
1910 .any(|block| block["type"] == "text" && block["text"] == "thanks")
1911 || body["messages"].as_array().is_some_and(|messages| {
1912 messages
1913 .iter()
1914 .any(|message| message["content"] == "thanks")
1915 });
1916 if has_result && !has_call {
1917 return ResponseTemplate::new(400).set_body_json(json!({
1918 "type": "error",
1919 "error": {
1920 "type": "invalid_request_error",
1921 "message": "tool_result has no corresponding tool_use"
1922 }
1923 }));
1924 }
1925 let response = if has_thanks {
1926 json!({
1927 "id": "msg_third", "type": "message", "role": "assistant",
1928 "model": "weak", "content": [{"type": "text", "text": "done"}],
1929 "stop_reason": "end_turn", "usage": {"input_tokens": 6, "output_tokens": 1}
1930 })
1931 } else if has_result {
1932 json!({
1933 "id": "msg_follow", "type": "message", "role": "assistant",
1934 "model": "weak", "content": [{"type": "text", "text": "sunny"}],
1935 "stop_reason": "end_turn", "usage": {"input_tokens": 4, "output_tokens": 1}
1936 })
1937 } else {
1938 json!({
1939 "id": "msg_seed", "type": "message", "role": "assistant",
1940 "model": "weak",
1941 "content": [{
1942 "type": "tool_use", "id": "toolu_weather", "name": "get_weather",
1943 "input": {"city": "Paris"}
1944 }],
1945 "stop_reason": "tool_use", "usage": {"input_tokens": 3, "output_tokens": 1}
1946 })
1947 };
1948 ResponseTemplate::new(200).set_body_json(response)
1949 })
1950 .mount(&server)
1951 .await;
1952
1953 let client: Arc<dyn RoutedLlmClient> = Arc::new(
1954 TranslatingLlmClient::new(&[ModelConfig::new(
1955 "weak",
1956 Backend::Anthropic(HttpBackendConfig {
1957 base_url: server.uri(),
1958 api_key: None,
1959 forward_auth: false,
1960 extra_headers: BTreeMap::new(),
1961 extra_body: BTreeMap::new(),
1962 omit_body_fields: BTreeSet::new(),
1963 reasoning_effort: None,
1964 max_retries: 0,
1965 failure_cooldown: std::time::Duration::ZERO,
1966 timeout: None,
1967 }),
1968 None,
1969 )])
1970 .map_err(|error| LibsyError::external("building test client", error))?,
1971 );
1972 let clients = ClientRouter::new(HashMap::from([(ModelId::from("weak"), client)]));
1973 let models = to_category_map(&["weak"]);
1974
1975 let seed = Request {
1976 llm_request: switchyard_translation::decode_request(
1977 WireFormat::OpenAiResponses,
1978 &json!({
1979 "model": "route",
1980 "input": "Call get_weather for Paris",
1981 "tools": [{
1982 "type": "function", "name": "get_weather",
1983 "parameters": {"type": "object"}
1984 }],
1985 "store": true
1986 }),
1987 )
1988 .map_err(|error| LibsyError::external("decoding seed request", error))?,
1989 raw_request: None,
1990 metadata: None,
1991 };
1992 let (_, response) = run(
1993 Arc::new(switchyard_libsy::Passthrough),
1994 clients.clone(),
1995 seed,
1996 models.clone(),
1997 None,
1998 )
1999 .await?;
2000 assert_eq!(
2001 response
2002 .llm_response
2003 .as_agg()
2004 .and_then(|agg| agg.id.as_deref()),
2005 Some("msg_seed")
2006 );
2007
2008 let follow = Request {
2009 llm_request: switchyard_translation::decode_request(
2010 WireFormat::OpenAiResponses,
2011 &json!({
2012 "model": "route",
2013 "previous_response_id": "msg_seed",
2014 "input": [{
2015 "type": "function_call_output",
2016 "call_id": "toolu_weather",
2017 "output": "Paris is sunny"
2018 }],
2019 "tools": [{
2020 "type": "function", "name": "get_weather",
2021 "parameters": {"type": "object"}
2022 }],
2023 "store": true
2024 }),
2025 )
2026 .map_err(|error| LibsyError::external("decoding continuation request", error))?,
2027 raw_request: None,
2028 metadata: None,
2029 };
2030 let (_, response) = run(
2031 Arc::new(switchyard_libsy::Passthrough),
2032 clients.clone(),
2033 follow,
2034 models.clone(),
2035 None,
2036 )
2037 .await?;
2038 assert_eq!(
2039 completion_text(
2040 response
2041 .llm_response
2042 .as_agg()
2043 .expect("buffered Anthropic response"),
2044 ),
2045 "sunny"
2046 );
2047
2048 let third = Request {
2049 llm_request: switchyard_translation::decode_request(
2050 WireFormat::OpenAiResponses,
2051 &json!({
2052 "model": "route",
2053 "previous_response_id": "msg_follow",
2054 "input": "thanks",
2055 "store": true
2056 }),
2057 )
2058 .map_err(|error| LibsyError::external("decoding chained request", error))?,
2059 raw_request: None,
2060 metadata: None,
2061 };
2062 let history = clients
2063 .stored_state_owner(&third)
2064 .and_then(|owner| owner.history)
2065 .expect("chained canonical history");
2066 assert_eq!(history.len, 4);
2067 assert_eq!(history.segment.len(), 2);
2068 assert!(history.parent.is_some());
2069 let (_, response) = run(
2070 Arc::new(switchyard_libsy::Passthrough),
2071 clients,
2072 third,
2073 models,
2074 None,
2075 )
2076 .await?;
2077 assert_eq!(
2078 completion_text(
2079 response
2080 .llm_response
2081 .as_agg()
2082 .expect("buffered Anthropic response"),
2083 ),
2084 "done"
2085 );
2086
2087 let requests = server.received_requests().await.expect("request recording");
2088 assert_eq!(requests.len(), 3);
2089 let follow: Value = serde_json::from_slice(&requests[1].body)
2090 .map_err(|error| LibsyError::external("decoding captured request", error))?;
2091 assert_eq!(follow["messages"][1]["content"][0]["type"], "tool_use");
2092 assert_eq!(follow["messages"][2]["content"][0]["type"], "tool_result");
2093 let third: Value = serde_json::from_slice(&requests[2].body)
2094 .map_err(|error| LibsyError::external("decoding captured request", error))?;
2095 assert_eq!(third["messages"][1]["content"][0]["type"], "tool_use");
2096 assert_eq!(third["messages"][2]["content"][0]["type"], "tool_result");
2097 assert_eq!(third["messages"][3]["content"], "sunny");
2098 assert_eq!(third["messages"][4]["content"][0]["text"], "thanks");
2099 Ok(())
2100 }
2101
2102 #[tokio::test]
2103 async fn streamed_cross_format_state_is_recorded_only_after_completion() -> Result<()> {
2104 let client = Arc::new(CandidateClient {
2105 calls: Mutex::new(Vec::new()),
2106 requests: Mutex::new(Vec::new()),
2107 first: FirstOutcome::StreamSuccess,
2108 });
2109 let clients = ClientRouter::new(HashMap::from([(
2110 ModelId::from("weak"),
2111 client as Arc<dyn RoutedLlmClient>,
2112 )]));
2113 let mut seed = request();
2114 seed.llm_request.messages = vec![Message::text(Role::User, "inspect")];
2115 seed.llm_request
2116 .preservation
2117 .requests
2118 .insert(WireFormat::OpenAiResponses.into(), json!({}));
2119 let mut response = stream_response(vec![
2120 LlmResponseChunk::MessageStart {
2121 id: Some("msg_stream".to_string()),
2122 model: Some("weak".to_string()),
2123 },
2124 LlmResponseChunk::ToolCallDelta {
2125 index: 0,
2126 id: Some("toolu_stream".to_string()),
2127 name: Some("lookup".to_string()),
2128 arguments_delta: Some("{}".to_string()),
2129 },
2130 LlmResponseChunk::MessageStop {
2131 reason: Some("tool_use".to_string()),
2132 },
2133 ]);
2134 response.set_served_model(&ModelId::from("weak"));
2135 let response = clients.remember_state_owner(&seed, response)?;
2136
2137 let mut follow = request();
2138 follow
2139 .llm_request
2140 .extensions
2141 .fields
2142 .insert("previous_response_id".to_string(), json!("msg_stream"));
2143 follow.llm_request.messages = vec![Message::new(
2144 Role::Tool,
2145 vec![ContentBlock::ToolResult(ToolResult {
2146 tool_call_id: "toolu_stream".to_string(),
2147 content: vec![ContentBlock::Text {
2148 text: "ready".to_string(),
2149 }],
2150 is_error: None,
2151 })],
2152 )];
2153 assert!(clients.stored_state_owner(&follow).is_none());
2154 response
2155 .llm_response
2156 .into_agg()
2157 .await
2158 .map_err(|error| LibsyError::client_call("weak", error))?;
2159
2160 let outcome = continue_on(
2161 clients
2162 .stored_state_owner(&follow)
2163 .expect("completed stream state"),
2164 "passthrough",
2165 follow,
2166 );
2167 assert_eq!(outcome.request.llm_request.messages.len(), 3);
2168 assert!(matches!(
2169 outcome.request.llm_request.messages[1].content.as_slice(),
2170 [ContentBlock::ToolCall(call)] if call.id == "toolu_stream"
2171 ));
2172 assert!(matches!(
2173 outcome.request.llm_request.messages[2].content.as_slice(),
2174 [ContentBlock::ToolResult(result)] if result.tool_call_id == "toolu_stream"
2175 ));
2176 Ok(())
2177 }
2178
2179 #[test]
2180 fn cross_format_response_with_store_false_is_not_retained() -> Result<()> {
2181 let client = Arc::new(CandidateClient {
2182 calls: Mutex::new(Vec::new()),
2183 requests: Mutex::new(Vec::new()),
2184 first: FirstOutcome::StreamSuccess,
2185 });
2186 let clients = ClientRouter::new(HashMap::from([(
2187 ModelId::from("weak"),
2188 client as Arc<dyn RoutedLlmClient>,
2189 )]));
2190 let mut seed = request();
2191 seed.llm_request
2192 .extensions
2193 .fields
2194 .insert("store".to_string(), json!(false));
2195 seed.llm_request
2196 .preservation
2197 .requests
2198 .insert(WireFormat::OpenAiResponses.into(), json!({}));
2199 let mut agg = text_response(Some("weak".to_string()), "done".to_string());
2200 agg.id = Some("msg_not_stored".to_string());
2201 let mut response = Response {
2202 llm_response: LlmResponse::Agg(agg),
2203 metadata: None,
2204 upstream_headers: http::HeaderMap::new(),
2205 };
2206 response.set_served_model(&ModelId::from("weak"));
2207 clients.remember_state_owner(&seed, response)?;
2208
2209 let mut follow = request();
2210 follow
2211 .llm_request
2212 .extensions
2213 .fields
2214 .insert("previous_response_id".to_string(), json!("msg_not_stored"));
2215 assert!(clients.stored_state_owner(&follow).is_none());
2216 Ok(())
2217 }
2218
2219 #[tokio::test]
2220 async fn state_capacity_does_not_fail_completed_responses() -> Result<()> {
2221 let client = Arc::new(CandidateClient {
2222 calls: Mutex::new(Vec::new()),
2223 requests: Mutex::new(Vec::new()),
2224 first: FirstOutcome::StreamSuccess,
2225 });
2226 let clients = ClientRouter::new(HashMap::from([(
2227 ModelId::from("weak"),
2228 client as Arc<dyn RoutedLlmClient>,
2229 )]));
2230 let model = ModelId::from("weak");
2231 clients.inner.state_owners.lock().by_id = (0..MAX_STATE_OWNERS)
2232 .map(|id| {
2233 (
2234 format!("existing_{id}"),
2235 StoredState {
2236 model: model.clone(),
2237 history: None,
2238 },
2239 )
2240 })
2241 .collect();
2242
2243 let mut request = request();
2244 request
2245 .llm_request
2246 .preservation
2247 .requests
2248 .insert(WireFormat::OpenAiResponses.into(), json!({}));
2249 let mut agg = text_response(Some("weak".to_string()), "done".to_string());
2250 agg.id = Some("overflow_buffered".to_string());
2251 let mut buffered = Response {
2252 llm_response: LlmResponse::Agg(agg),
2253 metadata: None,
2254 upstream_headers: http::HeaderMap::new(),
2255 };
2256 buffered.set_served_model(&model);
2257 clients.remember_state_owner(&request, buffered)?;
2258
2259 let mut streamed = stream_response(vec![
2260 LlmResponseChunk::MessageStart {
2261 id: Some("overflow_streamed".to_string()),
2262 model: Some("weak".to_string()),
2263 },
2264 LlmResponseChunk::TextDelta {
2265 index: 0,
2266 text: "done".to_string(),
2267 },
2268 LlmResponseChunk::MessageStop {
2269 reason: Some("stop".to_string()),
2270 },
2271 ]);
2272 streamed.set_served_model(&model);
2273 clients
2274 .remember_state_owner(&request, streamed)?
2275 .llm_response
2276 .into_agg()
2277 .await
2278 .map_err(|error| LibsyError::client_call("weak", error))?;
2279
2280 assert_eq!(
2281 clients.inner.state_owners.lock().by_id.len(),
2282 MAX_STATE_OWNERS
2283 );
2284 Ok(())
2285 }
2286
2287 #[test]
2288 fn state_owners_preserve_existing_records_at_capacity()
2289 -> std::result::Result<(), LlmClientError> {
2290 let mut owners = StateOwners::default();
2291 let model = ModelId::from("model/a");
2292 for id in 0..MAX_STATE_OWNERS - 1 {
2293 owners.remember(Some(&format!("resp_{id}")), None, &model, None)?;
2294 }
2295 assert!(matches!(
2296 owners.remember(Some("new_response"), Some("new_conversation"), &model, None,),
2297 Err(LlmClientError::ResponseStateLimitExceeded {
2298 limit: MAX_STATE_OWNERS
2299 })
2300 ));
2301 assert!(owners.owner("new_response").is_none());
2302 assert!(owners.owner("new_conversation").is_none());
2303 owners.remember(Some("last_id"), Some("last_id"), &model, None)?;
2304 owners.remember(Some("last_id"), Some("resp_0"), &model, None)?;
2305 assert!(matches!(
2306 owners.remember(Some("over_limit"), None, &model, None),
2307 Err(LlmClientError::ResponseStateLimitExceeded { .. })
2308 ));
2309 assert!(matches!(
2310 owners.remember(Some("resp_0"), None, &ModelId::from("model/b"), None,),
2311 Err(LlmClientError::ResponseStateConflict)
2312 ));
2313 assert_eq!(
2314 owners.owner("resp_0").map(|state| &state.model),
2315 Some(&model)
2316 );
2317 assert_eq!(owners.by_id.len(), MAX_STATE_OWNERS);
2318 Ok(())
2319 }
2320
2321 #[test]
2322 fn concurrent_state_owners_do_not_exceed_capacity() {
2323 let model = ModelId::from("model/a");
2324 let owners = Arc::new(Mutex::new(StateOwners {
2325 by_id: (0..MAX_STATE_OWNERS - 16)
2326 .map(|id| {
2327 (
2328 format!("existing_{id}"),
2329 StoredState {
2330 model: model.clone(),
2331 history: None,
2332 },
2333 )
2334 })
2335 .collect(),
2336 }));
2337 let barrier = Arc::new(std::sync::Barrier::new(32));
2338 let accepted = std::thread::scope(|scope| {
2339 let handles = (0..32)
2340 .map(|id| {
2341 let owners = owners.clone();
2342 let barrier = barrier.clone();
2343 let model = model.clone();
2344 scope.spawn(move || {
2345 barrier.wait();
2346 let result = owners.lock().remember(
2347 Some(&format!("resp_{id}")),
2348 Some(&format!("conv_{id}")),
2349 &model,
2350 None,
2351 );
2352 match result {
2353 Ok(()) => Some(id),
2354 Err(LlmClientError::ResponseStateLimitExceeded { .. }) => None,
2355 other => panic!("unexpected state tracking result: {other:?}"),
2356 }
2357 })
2358 })
2359 .collect::<Vec<_>>();
2360 handles
2361 .into_iter()
2362 .filter_map(|handle| handle.join().expect("request thread panicked"))
2363 .collect::<Vec<_>>()
2364 });
2365 let owners = owners.lock();
2366 assert_eq!(accepted.len(), 8);
2367 assert_eq!(owners.by_id.len(), MAX_STATE_OWNERS);
2368 assert_eq!(
2369 owners.owner("existing_0").map(|state| &state.model),
2370 Some(&model)
2371 );
2372 for id in accepted {
2373 assert_eq!(
2374 owners
2375 .owner(&format!("resp_{id}"))
2376 .map(|state| &state.model),
2377 Some(&model)
2378 );
2379 assert_eq!(
2380 owners
2381 .owner(&format!("conv_{id}"))
2382 .map(|state| &state.model),
2383 Some(&model)
2384 );
2385 }
2386 }
2387
2388 #[tokio::test]
2389 async fn each_fallback_candidate_receives_only_its_own_prompt() -> Result<()> {
2390 let client = Arc::new(CandidateClient {
2391 calls: Mutex::new(Vec::new()),
2392 requests: Mutex::new(Vec::new()),
2393 first: FirstOutcome::ContextWindow,
2394 });
2395 let routed_client: Arc<dyn RoutedLlmClient> = client.clone();
2396 let clients = ClientRouter::new_with_target_prompts(
2397 HashMap::from([
2398 (ModelId::from("weak"), Arc::clone(&routed_client)),
2399 (ModelId::from("strong"), routed_client),
2400 ]),
2401 HashMap::from([
2402 ("weak".into(), "weak prompt".to_string()),
2403 ("strong".into(), "strong prompt".to_string()),
2404 ]),
2405 None,
2406 );
2407
2408 run(
2409 Arc::new(CandidateAlgorithm {}),
2410 clients,
2411 request(),
2412 to_category_map(&["weak", "strong"]),
2413 None,
2414 )
2415 .await?;
2416
2417 let calls = client.requests.lock();
2418 assert_eq!(calls.len(), 2);
2419 assert_eq!(instruction_text(&calls[0]), ["weak prompt"]);
2420 assert_eq!(instruction_text(&calls[1]), ["strong prompt"]);
2421 Ok(())
2422 }
2423
2424 #[tokio::test]
2425 async fn decision_prompts_a_routing_response_target() -> Result<()> {
2426 let client = Arc::new(CandidateClient {
2427 calls: Mutex::new(Vec::new()),
2428 requests: Mutex::new(Vec::new()),
2429 first: FirstOutcome::StreamSuccess,
2430 });
2431 let routed_client: Arc<dyn RoutedLlmClient> = client.clone();
2432 let clients = ClientRouter::new_with_target_prompts(
2433 HashMap::from([(ModelId::from("answer"), routed_client)]),
2434 HashMap::from([("answer".into(), "answer prompt".to_string())]),
2435 Some("answer".into()),
2436 );
2437
2438 decide(
2439 Arc::new(AnsweredAlgorithm {
2440 model: "answer".into(),
2441 }),
2442 clients,
2443 request(),
2444 Arc::new(RuntimeModels::default()),
2445 )
2446 .await?;
2447
2448 let calls = client.requests.lock();
2449 assert_eq!(calls.len(), 1);
2450 assert_eq!(instruction_text(&calls[0]), ["answer prompt"]);
2451 Ok(())
2452 }
2453
2454 #[tokio::test]
2455 async fn decision_prepares_only_the_selected_request() -> Result<()> {
2456 let client = Arc::new(CandidateClient {
2457 calls: Mutex::new(Vec::new()),
2458 requests: Mutex::new(Vec::new()),
2459 first: FirstOutcome::StreamSuccess,
2460 });
2461 let routed_client: Arc<dyn RoutedLlmClient> = client;
2462 let clients = ClientRouter::new_with_target_prompts(
2463 HashMap::from([
2464 (ModelId::from("weak"), Arc::clone(&routed_client)),
2465 (ModelId::from("strong"), routed_client),
2466 ]),
2467 HashMap::from([
2468 (ModelId::from("weak"), "weak prompt".to_string()),
2469 (ModelId::from("strong"), "strong prompt".to_string()),
2470 ]),
2471 None,
2472 );
2473
2474 let outcome = decide(
2475 Arc::new(CandidateAlgorithm {}),
2476 clients,
2477 request(),
2478 to_category_map(&["weak", "strong"]),
2479 )
2480 .await?;
2481
2482 assert_eq!(
2483 outcome.selected_model_ids,
2484 [ModelId::from("weak"), ModelId::from("strong")]
2485 );
2486 assert_eq!(instruction_text(&outcome.request), ["weak prompt"]);
2487 Ok(())
2488 }
2489
2490 #[test]
2491 fn fallback_only_accepts_context_policy_and_unavailable_failures() {
2492 let error = |source| LibsyError::client_call("target", source);
2493 assert_eq!(
2494 fallback_reason(
2495 &error(LlmClientError::ContextWindowExceeded {
2496 model: "target".into(),
2497 message: "too long".to_string(),
2498 }),
2499 true
2500 ),
2501 Some(RoutingFallbackReason::ContextWindow)
2502 );
2503 for status in [
2504 StatusCode::FORBIDDEN,
2505 StatusCode::REQUEST_TIMEOUT,
2506 StatusCode::TOO_MANY_REQUESTS,
2507 StatusCode::INTERNAL_SERVER_ERROR,
2508 StatusCode::from_u16(599).expect("599 is a valid status"),
2509 ] {
2510 assert_eq!(
2511 fallback_reason(
2512 &error(LlmClientError::UpstreamHttp {
2513 status,
2514 body: "failed".to_string(),
2515 }),
2516 true
2517 ),
2518 Some(RoutingFallbackReason::Unavailable)
2519 );
2520 }
2521 for (body, expected) in [
2522 (
2523 r#"{"error":{"code":"content_policy_violation"}}"#,
2524 Some(RoutingFallbackReason::Unavailable),
2525 ),
2526 (
2527 r#"{"error":{"code":"invalid_request_error"},"metadata":"content_policy_violation"}"#,
2528 None,
2529 ),
2530 (r#"{"error":{"message":"content_policy_violation"}}"#, None),
2531 ("content_policy_violation", None),
2532 ] {
2533 assert_eq!(
2534 fallback_reason(
2535 &error(LlmClientError::UpstreamHttp {
2536 status: StatusCode::BAD_REQUEST,
2537 body: body.to_string(),
2538 }),
2539 true
2540 ),
2541 expected,
2542 "{body}"
2543 );
2544 }
2545 for status in [
2546 StatusCode::BAD_REQUEST,
2547 StatusCode::UNAUTHORIZED,
2548 StatusCode::NOT_FOUND,
2549 StatusCode::CONFLICT,
2550 StatusCode::from_u16(499).expect("499 is a valid status"),
2551 StatusCode::from_u16(600).expect("600 is a valid status"),
2552 ] {
2553 assert_eq!(
2554 fallback_reason(
2555 &error(LlmClientError::UpstreamHttp {
2556 status,
2557 body: "failed".to_string(),
2558 }),
2559 true
2560 ),
2561 None
2562 );
2563 }
2564 }
2565
2566 #[tokio::test]
2567 async fn candidate_failures_follow_the_fallback_policy() -> Result<()> {
2568 let (client, result) = run_candidates(FirstOutcome::ContextWindow).await;
2570 let (_, response) = result?;
2571 assert_eq!(
2572 &*client.calls.lock(),
2573 &[ModelId::from("weak"), "strong".into()]
2574 );
2575 assert_eq!(
2576 response
2577 .llm_response
2578 .as_agg()
2579 .map(|response| response.model.as_deref()),
2580 Some(Some("strong"))
2581 );
2582 assert_eq!(response.served_model().map(ModelId::as_str), Some("strong"));
2583
2584 let (client, result) =
2586 run_candidates(FirstOutcome::ContentPolicy(StatusCode::BAD_REQUEST)).await;
2587 let (_, response) = result?;
2588 assert_eq!(
2589 &*client.calls.lock(),
2590 &[ModelId::from("weak"), "strong".into()]
2591 );
2592 assert_eq!(response.served_model().map(ModelId::as_str), Some("strong"));
2593
2594 let (client, result) = run_candidates(FirstOutcome::Unauthorized).await;
2596 assert!(matches!(
2597 result,
2598 Err(LibsyError::ClientCall {
2599 source: LlmClientError::UpstreamHttp {
2600 status: StatusCode::UNAUTHORIZED,
2601 ..
2602 },
2603 ..
2604 })
2605 ));
2606 assert_eq!(&*client.calls.lock(), &[ModelId::from("weak")]);
2607 Ok(())
2608 }
2609
2610 #[tokio::test]
2611 async fn fallback_candidates_receive_the_fallback_request() -> Result<()> {
2612 let client = Arc::new(CandidateClient {
2613 calls: Mutex::new(Vec::new()),
2614 requests: Mutex::new(Vec::new()),
2615 first: FirstOutcome::ContextWindow,
2616 });
2617 let clients = ClientRouter::single(client.clone());
2618 let mut seed = request();
2619 seed.llm_request
2620 .preservation
2621 .requests
2622 .insert(WireFormat::OpenAiResponses.into(), json!({}));
2623 run(
2624 Arc::new(RewritingAlgorithm {}),
2625 clients.clone(),
2626 seed,
2627 to_category_map(&["weak", "strong"]),
2628 None,
2629 )
2630 .await?;
2631 let sent: Vec<_> = client
2632 .requests
2633 .lock()
2634 .iter()
2635 .map(|request| request.llm_request.messages.clone())
2636 .collect();
2637 assert_eq!(sent.len(), 2);
2638 assert_eq!(sent[0].len(), request().llm_request.messages.len() + 1);
2639 assert_eq!(sent[1], request().llm_request.messages);
2640 let mut follow = request();
2641 follow
2642 .llm_request
2643 .extensions
2644 .fields
2645 .insert("previous_response_id".into(), json!("candidate_response"));
2646 follow.llm_request.messages = vec![Message::text(Role::User, "continue")];
2647 let owner = clients
2648 .stored_state_owner(&follow)
2649 .expect("fallback Responses state");
2650 let continued = continue_on(owner, "test", follow);
2651 let mut expected = sent[1].clone();
2652 expected.push(Message::text(Role::Assistant, "strong"));
2653 expected.push(Message::text(Role::User, "continue"));
2654 assert_eq!(continued.request.llm_request.messages, expected);
2655 Ok(())
2656 }
2657
2658 #[tokio::test]
2659 async fn fallback_request_does_not_bypass_content_policy() -> Result<()> {
2660 for status in [StatusCode::BAD_REQUEST, StatusCode::FORBIDDEN] {
2661 let client = Arc::new(CandidateClient {
2662 calls: Mutex::new(Vec::new()),
2663 requests: Mutex::new(Vec::new()),
2664 first: FirstOutcome::ContentPolicy(status),
2665 });
2666 let result = run(
2667 Arc::new(RewritingAlgorithm {}),
2668 ClientRouter::single(client.clone()),
2669 request(),
2670 to_category_map(&["weak", "strong"]),
2671 None,
2672 )
2673 .await;
2674 assert!(
2675 matches!(result, Err(LibsyError::ClientCall { source: LlmClientError::UpstreamHttp { status: actual, .. }, .. }) if actual == status)
2676 );
2677 assert_eq!(&*client.calls.lock(), &[ModelId::from("weak")]);
2678 }
2679 Ok(())
2680 }
2681
2682 #[tokio::test]
2683 async fn retry_exhaustion_falls_back_and_optional_cooldown_skips_failed_backend() -> Result<()>
2684 {
2685 let server = MockServer::start().await;
2686 let calls = Arc::new(Mutex::new(Vec::new()));
2687 let observed_calls = Arc::clone(&calls);
2688 Mock::given(method("POST"))
2689 .respond_with(move |request: &wiremock::Request| {
2690 let body: serde_json::Value =
2691 serde_json::from_slice(&request.body).unwrap_or(serde_json::Value::Null);
2692 let model = body["model"].as_str().unwrap_or_default().to_string();
2693 observed_calls.lock().push(model.clone());
2694 if model == "weak" {
2695 ResponseTemplate::new(503)
2696 .insert_header("retry-after", "0")
2697 .set_body_string("unavailable")
2698 } else {
2699 ResponseTemplate::new(200).set_body_json(serde_json::json!({
2700 "id": "answer",
2701 "model": "strong",
2702 "choices": [{
2703 "index": 0,
2704 "message": {"role": "assistant", "content": "ok"},
2705 "finish_reason": "stop"
2706 }],
2707 "usage": {}
2708 }))
2709 }
2710 })
2711 .mount(&server)
2712 .await;
2713
2714 let client_with_cooldown = |cooldown| {
2715 let backend = || {
2716 Backend::OpenAiChat(HttpBackendConfig {
2717 base_url: format!("{}/v1", server.uri()),
2718 api_key: None,
2719 forward_auth: false,
2720 extra_headers: BTreeMap::new(),
2721 extra_body: BTreeMap::new(),
2722 omit_body_fields: BTreeSet::new(),
2723 reasoning_effort: None,
2724 max_retries: 2,
2725 failure_cooldown: cooldown,
2726 timeout: None,
2727 })
2728 };
2729 TranslatingLlmClient::new(&[
2730 ModelConfig::new("weak", backend(), None),
2731 ModelConfig::new("strong", backend(), None),
2732 ])
2733 .map(Arc::new)
2734 .map_err(|error| LibsyError::external("building test client", error))
2735 };
2736 let call = async |client: Arc<TranslatingLlmClient>| {
2737 run(
2738 Arc::new(CandidateAlgorithm {}),
2739 ClientRouter::single(client),
2740 request(),
2741 to_category_map(&["weak", "strong"]),
2742 None,
2743 )
2744 .await?;
2745 Ok::<_, LibsyError>(std::mem::take(&mut *calls.lock()))
2746 };
2747
2748 let client = client_with_cooldown(std::time::Duration::ZERO)?;
2749 assert_eq!(
2750 call(client.clone()).await?,
2751 ["weak", "weak", "weak", "strong"]
2752 );
2753 assert_eq!(call(client).await?, ["weak", "weak", "weak", "strong"]);
2754
2755 let cooldown = std::time::Duration::from_secs(60);
2756 let client = client_with_cooldown(cooldown)?;
2757 assert_eq!(
2758 call(client.clone()).await?,
2759 ["weak", "weak", "weak", "strong"]
2760 );
2761 assert_eq!(call(client.clone()).await?, ["strong"]);
2762
2763 tokio::time::pause();
2764 tokio::time::advance(cooldown).await;
2765 tokio::time::resume();
2766
2767 assert_eq!(call(client).await?, ["weak", "weak", "weak", "strong"]);
2768 Ok(())
2769 }
2770
2771 #[tokio::test]
2772 async fn refused_connection_falls_back_to_the_next_candidate() -> Result<()> {
2773 let server = MockServer::start().await;
2774 Mock::given(method("POST"))
2775 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
2776 "id": "answer",
2777 "model": "strong",
2778 "choices": [{
2779 "index": 0,
2780 "message": {"role": "assistant", "content": "ok"},
2781 "finish_reason": "stop"
2782 }],
2783 "usage": {}
2784 })))
2785 .mount(&server)
2786 .await;
2787
2788 let backend = |base_url: String| {
2789 Backend::OpenAiChat(HttpBackendConfig {
2790 base_url,
2791 api_key: None,
2792 forward_auth: false,
2793 extra_headers: BTreeMap::new(),
2794 extra_body: BTreeMap::new(),
2795 omit_body_fields: BTreeSet::new(),
2796 reasoning_effort: None,
2797 max_retries: 0,
2798 failure_cooldown: std::time::Duration::ZERO,
2799 timeout: None,
2800 })
2801 };
2802 let client = Arc::new(
2803 TranslatingLlmClient::new(&[
2804 ModelConfig::new("weak", backend("http://127.0.0.1:1/v1".to_string()), None),
2806 ModelConfig::new("strong", backend(format!("{}/v1", server.uri())), None),
2807 ])
2808 .map_err(|error| LibsyError::external("building test client", error))?,
2809 );
2810 let (_, response) = run(
2811 Arc::new(CandidateAlgorithm {}),
2812 ClientRouter::single(client),
2813 request(),
2814 to_category_map(&["weak", "strong"]),
2815 None,
2816 )
2817 .await?;
2818
2819 assert_eq!(response.served_model().map(ModelId::as_str), Some("strong"));
2820 Ok(())
2821 }
2822
2823 #[tokio::test]
2824 async fn streams_are_outside_the_candidate_fallback_boundary() -> Result<()> {
2825 let (client, result) = run_candidates(FirstOutcome::StreamSuccess).await;
2827 let (_, response) = result?;
2828 assert_eq!(response.served_model().map(ModelId::as_str), Some("weak"));
2829 let aggregate = response
2830 .llm_response
2831 .into_agg()
2832 .await
2833 .map_err(|source| LibsyError::client_call("weak", source))?;
2834 assert_eq!(&*client.calls.lock(), &[ModelId::from("weak")]);
2835 assert_eq!(completion_text(&aggregate), "streamed");
2836
2837 let (client, result) = run_candidates(FirstOutcome::MidStreamError).await;
2839 let (_, response) = result?;
2840 let error = response
2841 .llm_response
2842 .into_agg()
2843 .await
2844 .expect_err("the stream should fail during aggregation");
2845 assert!(error.to_string().contains("stream failed"));
2846 assert_eq!(&*client.calls.lock(), &[ModelId::from("weak")]);
2847 Ok(())
2848 }
2849
2850 const OVERFLOW_SSE: &str = "data: {\"error\":{\"code\":400,\"message\":\"Input length (600013 tokens) exceeds the maximum allowed length (536826 tokens). Use a shorter input or enable --allow-auto-truncate.\",\"object\":\"error\",\"param\":null,\"type\":\"BAD_REQUEST\"},\"model\":\"weak\"}\n\ndata: [DONE]\n\n";
2853
2854 const STRONG_SSE: &str = "data: {\"id\":\"answer\",\"model\":\"strong\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"ok\"},\"finish_reason\":null}]}\n\ndata: {\"id\":\"answer\",\"model\":\"strong\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n";
2856
2857 async fn run_streaming_candidates(
2861 weak_sse: &'static str,
2862 ) -> (
2863 MockServer,
2864 Arc<Mutex<Vec<String>>>,
2865 Result<(ModelId, Response)>,
2866 ) {
2867 let server = MockServer::start().await;
2868 let calls = Arc::new(Mutex::new(Vec::new()));
2869 let observed_calls = Arc::clone(&calls);
2870 Mock::given(method("POST"))
2871 .respond_with(move |request: &wiremock::Request| {
2872 let body: serde_json::Value =
2873 serde_json::from_slice(&request.body).unwrap_or(serde_json::Value::Null);
2874 let model = body["model"].as_str().unwrap_or_default().to_string();
2875 observed_calls.lock().push(model.clone());
2876 let sse = if model == "weak" {
2877 weak_sse
2878 } else {
2879 STRONG_SSE
2880 };
2881 ResponseTemplate::new(200)
2882 .insert_header("content-type", "text/event-stream")
2883 .set_body_string(sse)
2884 })
2885 .mount(&server)
2886 .await;
2887
2888 let backend = || {
2889 Backend::OpenAiChat(HttpBackendConfig {
2890 base_url: format!("{}/v1", server.uri()),
2891 api_key: None,
2892 forward_auth: false,
2893 extra_headers: BTreeMap::new(),
2894 extra_body: BTreeMap::new(),
2895 omit_body_fields: BTreeSet::new(),
2896 reasoning_effort: None,
2897 max_retries: 0,
2898 failure_cooldown: std::time::Duration::ZERO,
2899 timeout: None,
2900 })
2901 };
2902 let client = Arc::new(
2903 TranslatingLlmClient::new(&[
2904 ModelConfig::new("weak", backend(), None),
2905 ModelConfig::new("strong", backend(), None),
2906 ])
2907 .expect("building test client"),
2908 );
2909 let algorithm = Arc::new(CandidateAlgorithm {});
2910 let mut llm_request = text_request(Some("auto".to_string()), "hello".to_string());
2911 llm_request.stream = true;
2912 let request = Request {
2913 llm_request,
2914 raw_request: None,
2915 metadata: None,
2916 };
2917 let result = run(
2918 algorithm,
2919 ClientRouter::single(client),
2920 request,
2921 to_category_map(&["weak", "strong"]),
2922 None,
2923 )
2924 .await;
2925 (server, calls, result)
2926 }
2927
2928 #[tokio::test]
2929 async fn first_stream_event_overflow_falls_back_to_the_next_candidate() -> Result<()> {
2930 let (_server, calls, result) = run_streaming_candidates(OVERFLOW_SSE).await;
2933 let (_, response) = result?;
2934 assert_eq!(&*calls.lock(), &["weak", "strong"]);
2935 assert_eq!(response.served_model().map(ModelId::as_str), Some("strong"));
2936 let aggregate = response
2937 .llm_response
2938 .into_agg()
2939 .await
2940 .map_err(|source| LibsyError::client_call("strong", source))?;
2941 assert_eq!(completion_text(&aggregate), "ok");
2942
2943 let (_server, calls, result) = run_streaming_candidates(
2946 "data: {\"error\":{\"code\":500,\"message\":\"engine exploded\",\"object\":\"error\"},\"model\":\"weak\"}\n\ndata: [DONE]\n\n",
2947 )
2948 .await;
2949 let (_, response) = result?;
2950 assert_eq!(&*calls.lock(), &["weak"]);
2951 let error = response
2952 .llm_response
2953 .into_agg()
2954 .await
2955 .expect_err("the in-band error should surface to the consumer");
2956 assert!(error.to_string().contains("engine exploded"));
2957
2958 let (_server, calls, result) = run_streaming_candidates(
2961 "data: {\"id\":\"turn\",\"model\":\"weak\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"partial\"},\"finish_reason\":null}]}\n\ndata: {\"error\":{\"code\":400,\"message\":\"Input length (600013 tokens) exceeds the maximum allowed length (536826 tokens).\",\"object\":\"error\",\"type\":\"BAD_REQUEST\"},\"model\":\"weak\"}\n\ndata: [DONE]\n\n",
2962 )
2963 .await;
2964 let (_, response) = result?;
2965 assert_eq!(&*calls.lock(), &["weak"]);
2966 let error = response
2967 .llm_response
2968 .into_agg()
2969 .await
2970 .expect_err("the mid-stream overflow should surface to the consumer");
2971 assert!(
2972 error
2973 .to_string()
2974 .contains("exceeds the maximum allowed length")
2975 );
2976
2977 let (_server, calls, result) = run_streaming_candidates("data: [DONE]\n\n").await;
2980 let (_, response) = result?;
2981 assert_eq!(&*calls.lock(), &["weak"]);
2982 let aggregate = response
2983 .llm_response
2984 .into_agg()
2985 .await
2986 .map_err(|source| LibsyError::client_call("weak", source))?;
2987 assert_eq!(completion_text(&aggregate), "");
2988 Ok(())
2989 }
2990}