Skip to main content

switchyard_llm_client/
run.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Drive a libsy algorithm to completion, making the model calls it offloads.
5//!
6//! [`switchyard_libsy::Algorithm::run_stream`] is the whole libsy API: it yields a stream of
7//! steps and expects its consumer to serve every offloaded model call. [`run()`] is that
8//! consumer — it drives the stream with [`switchyard_libsy::drive`], hands routing-time calls to
9//! a [`RoutedLlmClient`], and serves the terminal routing outcome.
10//!
11//! libsy owns the stream mechanics; what this module adds is per-target request preparation,
12//! ordered candidate fallback, and the `libsy.client_call` span around each candidate. Each
13//! completion candidate exhausts its backend retry budget before fallback advances. Routing
14//! calls stop on the first candidate's failure unless the algorithm enables error recovery.
15//! A timeout stops completion calls.
16//!
17//! A Responses continuation can refer to state through `previous_response_id` or `conversation`.
18//! [`ClientRouter`] sends native Responses state back to its provider; for Chat or Anthropic
19//! targets it retains the canonical history needed to materialize the continuation.
20
21use 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
43/// Run one request to completion, serving every offloaded model call with `client`.
44///
45/// Returns the model selected by the algorithm and the final [`Response`]. `observer`, when
46/// present, receives metadata for the routing outcome, each completed routing or answer call,
47/// and the routing overhead.
48///
49/// `clients` resolves each offloaded call to the client for the target the algorithm
50/// selected — an algorithm may route among targets served by different providers, so this is
51/// a per-call lookup, not one client for the whole run. Use
52/// [`ClientRouter::single`](ClientRouter::single) when one client serves every target.
53///
54/// Routing calls are buffered; client failures stop the request unless the call enables recovery.
55/// Once routing completes, non-timeout failures may try the outcome's ordered fallback candidates.
56pub 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    // This says if we have an observer, put Some(..) in routing_observations.
67    // No observer means we don't want any routing_observations.
68    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
139/// Reports the served answer to the algorithm; a stream reports it before its final chunk
140/// so the client cannot send its next request first.
141async 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
194/// Run an algorithm to a routing decision without serving its terminal completion.
195///
196/// Routing-time calls are served normally. The returned request is prepared only for the selected
197/// target; use [`run`] to execute selected and fallback candidates.
198pub 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
255/// Emits completed routing calls after the outcome reveals whether one response became the answer.
256fn 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
278/// Serve one offloaded call and fulfill its promise.
279///
280/// LLM failures stop the run unless the call enables recovery. Decision results,
281/// including client failures, go back to the algorithm for its routing policy.
282async 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
315/// Try candidates in order until one succeeds or a failure stops fallback.
316async 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/// The stage of the Switchyard run in which a model call was issued.
370///
371/// A routing-phase call may still supply the final answer. This is deliberately
372/// distinct from the `routing`/`answer` role assigned after the outcome is known.
373#[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/// Call one candidate model and record its observation and span.
389#[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 is for span log
438    index: usize,
439    // count is for span log
440    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    // Resolved before the clock starts: picking the client is Switchyard's work, not
454    // the provider's, so it belongs in the routing overhead.
455    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        // Observe dispatch and live I/O once; buffering must not re-observe replay.
463        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
504// Routing logs need the final usage, but replay must keep signed provider events intact.
505async 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
531/// Whether a failed candidate is worth routing around.
532fn 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
565/// Route to `owner` without fallbacks and record why the algorithm did not run.
566fn 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
586/// The model that served the stored state and the request field that refers to it.
587struct 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    /// `None` means the provider owns the state and only routing must be pinned.
626    history: Option<Arc<MessageHistory>>,
627}
628
629#[derive(Clone)]
630struct CanonicalInput {
631    parent: Option<Arc<MessageHistory>>,
632    messages: Arc<[Message]>,
633}
634
635/// Stored Responses IDs and conversation IDs, with their provider or canonical state.
636///
637/// Keep earlier response IDs because a client can continue from any stored response.
638/// Reject new IDs at capacity rather than forgetting where an earlier response is stored.
639#[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
698/// Read the Responses `conversation` ID from a string or an object.
699fn 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/// Resolves a routed call's selected model to the client that serves it.
707///
708/// An algorithm routes among named targets; which provider each target lives on is the
709/// host's concern, and two targets in one run may sit on different providers. A router owns
710/// that mapping and resolves the client for each kind of call.
711///
712/// Cloning is cheap — the mapping is shared, so one router can serve every request.
713#[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    /// Provider-owned routing pins and materialized cross-format history.
725    state_owners: Mutex<StateOwners>,
726    /// Native Responses state only needs a routing pin when answer targets span clients.
727    track_provider_state: bool,
728}
729
730enum Routing {
731    /// One client serves every model.
732    Single(Arc<dyn RoutedLlmClient>),
733    /// Each model is served by the client configured for it.
734    ByModel(HashMap<ModelId, Arc<dyn RoutedLlmClient>>),
735}
736
737impl ClientRouter {
738    /// Build a router over `model name -> client`, for targets spread across providers.
739    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    /// Build a router with target prompts used by [`run`] and [`decide`].
744    ///
745    /// `routing_answer_target` identifies a routing-time call whose response may become the
746    /// answer. Pass `None` when routing only selects a later completion target. Prompt keys and
747    /// `routing_answer_target` use the resolved model IDs stored in `by_model`.
748    /// This constructor treats every model in `by_model` as a possible completion target.
749    /// Use [`Self::new_with_completion_targets`] to exclude classifier-only models from the
750    /// comparison of completion clients.
751    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    /// Build a router with an explicit list of models that can answer requests.
766    ///
767    /// `by_model` contains callable LLMs, including LLM classifiers. Only the model IDs in
768    /// `completion_targets` determine whether native Responses state needs a routing pin.
769    /// Cross-format state is recorded lazily when a Responses request uses a Chat or Anthropic
770    /// target. `target_prompts` and `routing_answer_target` have the same meaning as in
771    /// [`Self::new_with_target_prompts`].
772    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    /// Build a router with separate LLM and decision mappings, shared across clones.
788    /// LLM prompt and completion settings follow [`Self::new_with_completion_targets`].
789    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    /// A router that serves every model with one client — the single-provider case.
816    ///
817    /// [`TranslatingLlmClient`](crate::TranslatingLlmClient) already maps model names to
818    /// backends internally and rejects ones it does not know, so enumerating them here would
819    /// only duplicate that.
820    pub fn single(client: Arc<dyn RoutedLlmClient>) -> Self {
821        Self::single_with_decision_clients(client, HashMap::new())
822    }
823
824    /// Use one client for all LLM targets and explicit clients for decision targets.
825    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    /// Stop completion fallback on HTTP 429 for the given models.
843    ///
844    /// For translating clients, also use
845    /// [`TranslatingLlmClient::without_rate_limit_cooldown`](crate::TranslatingLlmClient::without_rate_limit_cooldown)
846    /// so later calls are not skipped with a generic cooldown error.
847    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    /// Resolves a target for model telemetry, including calls that fail before a response.
853    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    /// The client that serves `model`.
859    ///
860    /// Errors with [`LlmClientError::Configuration`] when the router maps models and has no
861    /// entry for this one, rather than silently sending the call to another provider.
862    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    /// Resolve a decision target without falling back to an LLM client.
879    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    /// Return the recorded model for the requested response or conversation ID.
892    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    /// Record provider-owned Responses state or canonical history for a cross-format response.
926    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    /// Check both IDs from a response or stream event before recording either. The response's
1045    /// `store` value overrides the request value; `false` skips only the response ID.
1046    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    /// Prepare a completion candidate with its configured target prompt.
1109    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    /// Prepare a routing call, adding a target prompt only when it generates a candidate answer.
1116    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    /// Rewrites the request for its selected model only.
1202    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        // Context overflow is retryable across candidates.
2569        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        // A provider-specific policy denial also permits another candidate.
2585        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        // Authentication failure is not retryable, so the second candidate is untouched.
2595        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                // Nothing listens on port 1, so the connection is refused.
2805                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        // Receiving a stream handle is a successful call and ends candidate selection.
2826        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        // A later in-stream failure surfaces during aggregation without trying another model.
2838        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    // Verbatim shape of an SGLang admission rejection on a streaming session:
2851    // HTTP 200, and the overflow arrives as the first SSE event.
2852    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    // A normal one-turn streamed answer, as the second candidate serves it.
2855    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    /// Runs a streaming request through the weak→strong candidate loop, serving
2858    /// `weak_sse` from the first candidate and [`STRONG_SSE`] from the second.
2859    /// The server is returned so response streams stay readable after the call.
2860    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        // Nothing has reached the caller when the first event arrives, so an overflow
2931        // there falls through to the next candidate exactly as its buffered twin does.
2932        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        // Only a classified overflow re-routes; any other in-band error event still
2944        // streams through to the consumer without trying another candidate.
2945        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        // Once content has streamed, a retry would repeat delivered output, so a later
2959        // overflow event keeps today's semantics and surfaces during aggregation.
2960        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        // A stream that ends before any event is not a classified failure; it still
2978        // aggregates to an empty response.
2979        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}