1use std::{
9 collections::{HashMap, HashSet},
10 pin::Pin,
11 sync::Arc,
12 time::{Duration, Instant},
13};
14
15use async_trait::async_trait;
16use futures::{Stream, StreamExt};
17use parking_lot::Mutex;
18use tracing::Instrument;
19
20use switchyard_protocol::{
28 Context, Decision, LlmClientError, Request, Response, RoutedLlmClient, Signals, Usage,
29};
30
31use super::driver::{DriverRequest, DriverStep, TypeErasedDriver};
32use crate::{DriverError, LibsyError, Result, observability};
33
34pub type StepStream = Pin<Box<dyn Stream<Item = Result<Step>> + Send>>;
38
39#[derive(Clone, Debug)]
41pub struct LlmCallObservation {
42 pub selected_model: String,
44 pub tier: Option<String>,
46 pub is_routed: bool,
48 pub is_success: bool,
50 pub duration: Duration,
52 pub usage: Option<Usage>,
54}
55
56#[derive(Clone, Debug, Eq, PartialEq)]
58pub struct MetricAttribute {
59 pub key: &'static str,
61 pub value: String,
63}
64
65impl MetricAttribute {
66 pub fn new(key: &'static str, value: impl Into<String>) -> Self {
68 Self {
69 key,
70 value: value.into(),
71 }
72 }
73}
74
75#[derive(Clone, Copy, Debug, PartialEq)]
77pub enum AlgorithmMetricValue {
78 Counter(u64),
80 Histogram(f64),
82}
83
84#[derive(Clone, Debug, PartialEq)]
86pub struct AlgorithmMetricObservation {
87 pub algorithm: String,
89 pub name: &'static str,
91 pub value: AlgorithmMetricValue,
93 pub attributes: Vec<MetricAttribute>,
95}
96
97#[derive(Clone, Debug)]
99pub enum RunObservation {
100 LlmCall(LlmCallObservation),
102 RoutingOverhead(Duration),
104 AlgorithmMetric(AlgorithmMetricObservation),
106}
107
108pub type RunObserver = Arc<dyn Fn(RunObservation) + Send + Sync>;
114
115#[derive(Clone)]
124pub struct RoutedRequest {
125 pub request: Request,
127 pub decision: Arc<dyn Decision>,
129 pub default_client: Option<Arc<dyn RoutedLlmClient>>,
133 pub ctx: Context,
137}
138
139pub struct CallLlmRequest {
147 inner: DriverRequest,
148 routed: RoutedRequest,
149}
150
151impl CallLlmRequest {
152 fn new(inner: DriverRequest) -> Self {
155 let routed = match inner.request::<RoutedRequest>() {
158 Ok(routed) => routed.clone(),
159 Err(_) => unreachable!("CallLlmRequest payload is always a RoutedRequest"),
160 };
161 Self { inner, routed }
162 }
163
164 pub fn get_routed(&self) -> &RoutedRequest {
168 &self.routed
169 }
170
171 pub fn get_request(&self) -> &Request {
173 &self.get_routed().request
174 }
175
176 pub fn get_decision(&self) -> &dyn Decision {
178 self.get_routed().decision.as_ref()
179 }
180
181 pub fn respond(self, result: Result<Response>) -> Result<()> {
185 self.inner.respond::<Response>(result)
186 }
187}
188
189#[derive(Clone)]
196pub struct Driver {
197 driver: TypeErasedDriver,
198 algorithm: Arc<str>,
199 routed_call: Arc<Mutex<Option<Duration>>>,
201 observer: Option<RunObserver>,
202}
203
204impl Driver {
205 pub(crate) fn new() -> Self {
208 Self::with_observer("", None)
209 }
210
211 fn with_observer(algorithm: &str, observer: Option<RunObserver>) -> Self {
212 Self {
213 driver: TypeErasedDriver::new(),
214 algorithm: Arc::from(algorithm),
215 routed_call: Arc::new(Mutex::new(None)),
216 observer,
217 }
218 }
219
220 pub(crate) fn routed_call_duration(&self) -> Option<Duration> {
222 *self.routed_call.lock()
223 }
224
225 pub(crate) fn observe_routing_overhead(&self, duration: Duration) {
228 if let Some(observer) = &self.observer {
229 observer(RunObservation::RoutingOverhead(duration));
230 }
231 }
232
233 pub fn record_counter(
237 &self,
238 name: &'static str,
239 delta: u64,
240 attributes: impl IntoIterator<Item = MetricAttribute>,
241 ) {
242 self.record_algorithm_metric(
243 name,
244 AlgorithmMetricValue::Counter(delta),
245 attributes.into_iter().collect(),
246 );
247 }
248
249 pub fn record_histogram(
252 &self,
253 name: &'static str,
254 sample: f64,
255 attributes: impl IntoIterator<Item = MetricAttribute>,
256 ) {
257 if !sample.is_finite() {
258 tracing::debug!(
259 metric = name,
260 sample,
261 "dropping non-finite histogram sample"
262 );
263 return;
264 }
265 self.record_algorithm_metric(
266 name,
267 AlgorithmMetricValue::Histogram(sample),
268 attributes.into_iter().collect(),
269 );
270 }
271
272 fn record_algorithm_metric(
273 &self,
274 name: &'static str,
275 value: AlgorithmMetricValue,
276 mut attributes: Vec<MetricAttribute>,
277 ) {
278 attributes.retain(|attribute| attribute.key != "algorithm");
279 let observation = AlgorithmMetricObservation {
280 algorithm: self.algorithm.to_string(),
281 name,
282 value,
283 attributes,
284 };
285 observability::record_algorithm_metric(&observation);
286 if let Some(observer) = &self.observer {
287 observer(RunObservation::AlgorithmMetric(observation));
288 }
289 }
290
291 #[tracing::instrument(
300 target = "libsy",
301 name = "libsy.llm_call",
302 skip_all,
303 fields(
304 algorithm = observability::algorithm_label(&routed.ctx),
305 selected_model = routed.decision.selected_model(),
306 openinference.span.kind = "CHAIN",
307 outcome = tracing::field::Empty,
308 error = tracing::field::Empty,
309 input_tokens = tracing::field::Empty,
310 output_tokens = tracing::field::Empty,
311 total_tokens = tracing::field::Empty,
312 reasoning_tokens = tracing::field::Empty,
313 )
314 )]
315 pub async fn call_llm(&self, routed: RoutedRequest) -> Result<Response> {
316 let algorithm = observability::algorithm_label(&routed.ctx).to_string();
317 let selected_model = routed.decision.selected_model().to_string();
318 let tier = routed.decision.routing_tier().map(str::to_string);
319 let is_routed = routed.decision.is_routed_call();
320 let started = Instant::now();
321 let result = self
322 .driver
323 .fulfill_request::<RoutedRequest, Response>(routed.ctx.clone(), routed)
324 .await;
325 let elapsed = started.elapsed();
326 observability::record_llm_call(
327 &algorithm,
328 &selected_model,
329 tier.as_deref(),
330 is_routed,
331 elapsed,
332 &result,
333 &tracing::Span::current(),
334 );
335 if let Some(observer) = &self.observer {
336 observer(RunObservation::LlmCall(LlmCallObservation {
337 selected_model,
338 tier,
339 is_routed,
340 is_success: result.is_ok(),
341 duration: elapsed,
342 usage: result
343 .as_ref()
344 .ok()
345 .and_then(|response| response.llm_response.as_agg())
346 .map(|response| response.usage.clone()),
347 }));
348 }
349 if is_routed && result.is_ok() {
352 *self.routed_call.lock() = Some(elapsed);
353 }
354 result
355 }
356
357 pub async fn call_llm_target(
363 &self,
364 ctx: Context,
365 target: &LlmTarget,
366 request: Request,
367 decision: Arc<dyn Decision>,
368 ) -> Result<Response> {
369 self.call_llm(RoutedRequest {
370 request,
371 decision,
372 default_client: target.llm_client.clone(),
373 ctx,
374 })
375 .await
376 }
377
378 pub async fn info(&self, ctx: Context, decision: Arc<dyn Decision>) -> Result<()> {
382 self.driver.info(ctx.clone(), decision.clone()).await?;
383 observability::record_decision(&ctx, decision.as_ref());
384 Ok(())
385 }
386
387 pub(crate) async fn finish(&self, ctx: Context, result: Result<Response>) -> Result<()> {
391 match result {
392 Ok(response) => self.driver.done(ctx, response).await,
393 Err(err) => self.driver.fail(ctx, err).await,
394 }
395 }
396
397 pub(crate) fn stream(&self) -> impl Stream<Item = Result<Step>> + use<> {
401 self.driver.stream().map(|item| match item? {
402 DriverStep::Request(req) => Ok(Step::CallLlm(Box::new(CallLlmRequest::new(req)))),
403 DriverStep::Info(payload) => payload
404 .downcast::<Arc<dyn Decision>>()
405 .map(|decision| Step::Decision(*decision))
406 .map_err(|_| {
407 DriverError::TypeMismatch {
408 expected: "Arc<dyn Decision>",
409 }
410 .into()
411 }),
412 DriverStep::Done(payload) => payload
413 .downcast::<Response>()
414 .map(Step::ReturnToAgent)
415 .map_err(|_| {
416 DriverError::TypeMismatch {
417 expected: "Response",
418 }
419 .into()
420 }),
421 })
422 }
423}
424
425impl Default for Driver {
426 fn default() -> Self {
427 Self::new()
428 }
429}
430
431pub enum Step {
433 CallLlm(Box<CallLlmRequest>),
437 Decision(Arc<dyn Decision>),
440 ReturnToAgent(Box<Response>),
442}
443
444struct AbortOnDrop(tokio::task::AbortHandle);
446
447impl Drop for AbortOnDrop {
448 fn drop(&mut self) {
449 self.0.abort();
450 }
451}
452
453#[derive(Clone)]
458pub struct LlmTarget {
459 pub semantic_name: String,
463 pub llm_client: Option<Arc<dyn RoutedLlmClient>>,
466}
467
468#[derive(Clone)]
472pub struct LlmTargetSet {
473 targets: Vec<LlmTarget>,
474}
475
476impl LlmTargetSet {
477 pub fn new(targets: Vec<LlmTarget>) -> Self {
479 Self { targets }
480 }
481
482 pub fn targets(&self) -> &[LlmTarget] {
484 &self.targets
485 }
486
487 pub fn get_target(&self, name: &str) -> Result<LlmTarget> {
489 self.targets
490 .iter()
491 .find(|t| t.semantic_name == name)
492 .cloned()
493 .ok_or_else(|| LibsyError::TargetNotFound {
494 target: name.to_string(),
495 })
496 }
497
498 pub fn resolve_target(&self, name: &str, ctx: &Context) -> Result<LlmTarget> {
501 let target = self.get_target(name)?;
502 if !ctx.is_excluded(&target.semantic_name) {
503 return Ok(target);
504 }
505 self.targets
506 .iter()
507 .find(|t| !ctx.is_excluded(&t.semantic_name))
508 .cloned()
509 .ok_or(LibsyError::AllTargetsExcluded)
510 }
511
512 pub fn count_tokens_client(&self) -> Option<Arc<dyn RoutedLlmClient>> {
516 self.targets.iter().find_map(|target| {
517 target
518 .llm_client
519 .as_ref()
520 .filter(|client| client.supports_count_tokens())
521 .cloned()
522 })
523 }
524}
525
526const MAX_EVICTION_SESSIONS: usize = 1_024;
529
530#[derive(Default)]
536pub(crate) struct SessionEvictions {
537 sessions: Mutex<HashMap<String, HashSet<String>>>,
538}
539
540impl SessionEvictions {
541 fn evicted_in(&self, session: Option<&str>) -> Vec<String> {
543 let Some(session) = session else {
544 return Vec::new();
545 };
546 self.sessions
547 .lock()
548 .get(session)
549 .map(|targets| targets.iter().cloned().collect())
550 .unwrap_or_default()
551 }
552
553 fn record(&self, session: Option<&str>, target: &str) {
556 let Some(session) = session else { return };
557 let mut sessions = self.sessions.lock();
558 if sessions.len() >= MAX_EVICTION_SESSIONS
559 && !sessions.contains_key(session)
560 && let Some(oldest) = sessions.keys().next().cloned()
561 {
562 sessions.remove(&oldest);
563 }
564 sessions
565 .entry(session.to_string())
566 .or_default()
567 .insert(target.to_string());
568 }
569}
570
571fn eligible_targets(targets: &LlmTargetSet, ctx: &Context) -> usize {
573 targets
574 .targets()
575 .iter()
576 .filter(|t| !ctx.is_excluded(&t.semantic_name))
577 .count()
578}
579
580pub(crate) fn exclude_evicted(
583 ctx: &mut Context,
584 targets: &LlmTargetSet,
585 evictions: &SessionEvictions,
586 session: Option<&str>,
587) {
588 for target in evictions.evicted_in(session) {
589 if eligible_targets(targets, ctx) <= 1 {
592 break;
593 }
594 ctx.exclude_target(target);
595 }
596}
597
598#[allow(clippy::too_many_arguments)]
606pub(crate) async fn call_llm_with_overflow_fallback(
607 mut ctx: Context,
608 driver: &Driver,
609 targets: &LlmTargetSet,
610 mut target: LlmTarget,
611 mut decision: Arc<dyn Decision>,
612 request: Request,
613 session: Option<&str>,
614 evictions: &SessionEvictions,
615 fallback_decision: impl Fn(&LlmTarget, &LlmTarget) -> Arc<dyn Decision>,
616) -> Result<Response> {
617 loop {
618 let result = driver
619 .call_llm_target(ctx.clone(), &target, request.clone(), decision.clone())
620 .await;
621 let Err(error) = result else { return result };
622 let LibsyError::ClientCall {
623 target: failed,
624 source: LlmClientError::ContextWindowExceeded { .. },
625 } = &error
626 else {
627 return Err(error);
628 };
629 if !ctx.exclude_target(failed) {
632 return Err(error);
633 }
634 evictions.record(session, failed);
635 let Ok(next) = targets.resolve_target(&target.semantic_name, &ctx) else {
636 return Err(error);
637 };
638 decision = fallback_decision(&target, &next);
639 target = next;
640 driver.info(ctx.clone(), decision.clone()).await?;
641 }
642}
643
644#[async_trait]
665pub trait Algorithm: Send + Sync + 'static {
666 fn name(&self) -> &str;
670
671 async fn create_run_task(
677 self: Arc<Self>,
678 ctx: Context,
679 driver: Driver,
680 request: Request,
681 ) -> Result<Response>;
682
683 #[allow(unused_variables)]
687 async fn process_signals(self: Arc<Self>, signals: Signals) -> Result<()> {
688 Ok(())
689 }
690
691 fn count_tokens_client(&self) -> Option<Arc<dyn RoutedLlmClient>> {
701 None
702 }
703
704 async fn count_tokens(&self, request: Request) -> Result<serde_json::Value> {
712 let client = self
713 .count_tokens_client()
714 .ok_or_else(|| LibsyError::AlgorithmError {
715 message: "no target supports count_tokens (needs an Anthropic upstream)"
716 .to_string(),
717 })?;
718 client
719 .count_tokens(request)
720 .await
721 .map_err(|source| LibsyError::client_call("count_tokens", source))
722 }
723
724 fn run_stream(
735 self: Arc<Self>,
736 ctx: Context,
737 request: Request,
738 observer: Option<RunObserver>,
739 ) -> StepStream {
740 let mut ctx = ctx;
743 ctx.values.insert(
744 observability::ALGORITHM_KEY.to_string(),
745 self.name().to_string(),
746 );
747 let driver = Driver::with_observer(self.name(), observer);
748 let task_driver = driver.clone();
749 let task_ctx = ctx.clone();
750 let stream = task_driver.stream();
751 let span = observability::run_span(self.name(), &request);
755 let observed_driver = task_driver.clone();
756 let handle = tokio::spawn(
757 async move {
758 observability::observe_run(
759 task_ctx.clone(),
760 observed_driver,
761 self.create_run_task(task_ctx, task_driver, request),
762 )
763 .await
764 }
765 .instrument(span),
766 );
767 let abort_guard = AbortOnDrop(handle.abort_handle());
769
770 let finish_driver = driver.clone();
771 let finish_ctx = ctx;
772 let tail: StepStream = Box::pin(
773 futures::stream::once(async move {
774 let result = match handle.await {
775 Ok(response) => response,
776 Err(source) => Err(LibsyError::AlgorithmTask { source }),
777 };
778 finish_driver.finish(finish_ctx, result).await
779 })
780 .filter_map(|finish_result| async move { finish_result.err().map(Err) }),
781 );
782
783 let stream: StepStream = Box::pin(stream);
784 Box::pin(futures::stream::select(stream, tail).map(move |step| {
785 let _keep_alive = &abort_guard;
787 step
788 }))
789 }
790
791 async fn run(
795 self: Arc<Self>,
796 ctx: Context,
797 request: Request,
798 ) -> Result<(Vec<Arc<dyn Decision>>, Response)> {
799 self.run_observed(ctx, request, None).await
800 }
801
802 async fn run_observed(
804 self: Arc<Self>,
805 ctx: Context,
806 request: Request,
807 observer: Option<RunObserver>,
808 ) -> Result<(Vec<Arc<dyn Decision>>, Response)> {
809 #[tracing::instrument(
815 target = "libsy",
816 name = "libsy.client_call",
817 skip_all,
818 fields(
819 algorithm = observability::algorithm_label(&call.get_routed().ctx),
820 switchyard.algorithm = observability::algorithm_label(&call.get_routed().ctx),
821 switchyard.routing.tier = tracing::field::Empty,
822 selected_model = call.get_decision().selected_model(),
823 otel.kind = "client",
824 otel.name = %format_args!("chat {}", call.get_decision().selected_model()),
825 openinference.span.kind = "LLM",
826 gen_ai.operation.name = "chat",
827 gen_ai.request.model = call.get_decision().selected_model(),
828 gen_ai.request.stream = tracing::field::Empty,
829 gen_ai.request.temperature = tracing::field::Empty,
830 gen_ai.request.top_p = tracing::field::Empty,
831 gen_ai.request.top_k = tracing::field::Empty,
832 gen_ai.request.max_tokens = tracing::field::Empty,
833 gen_ai.request.reasoning.level = tracing::field::Empty,
834 gen_ai.output.type = tracing::field::Empty,
835 gen_ai.conversation.id = tracing::field::Empty,
836 server.address = tracing::field::Empty,
837 server.port = tracing::field::Empty,
838 gen_ai.response.id = tracing::field::Empty,
839 gen_ai.response.model = tracing::field::Empty,
840 gen_ai.usage.input_tokens = tracing::field::Empty,
841 gen_ai.usage.output_tokens = tracing::field::Empty,
842 gen_ai.usage.cache_read.input_tokens = tracing::field::Empty,
843 gen_ai.usage.cache_creation.input_tokens = tracing::field::Empty,
844 gen_ai.usage.reasoning.output_tokens = tracing::field::Empty,
845 outcome = tracing::field::Empty,
846 otel.status_code = tracing::field::Empty,
847 error.type = tracing::field::Empty,
848 error = tracing::field::Empty,
849 )
850 )]
851 async fn serve(call: CallLlmRequest) -> Result<()> {
852 let span = tracing::Span::current();
853 observability::record_gen_ai_request(&span, &call.get_routed().request.llm_request);
854 if let Some(tier) = call.get_decision().routing_tier() {
855 span.record("switchyard.routing.tier", tier);
856 }
857 if let Some(session_id) = call
858 .get_routed()
859 .request
860 .metadata
861 .as_ref()
862 .and_then(|metadata| metadata.session_id.as_deref())
863 {
864 span.record("gen_ai.conversation.id", session_id);
865 }
866 let routed = call.get_routed().clone();
867 let target = routed.decision.selected_model().to_string();
868 let client =
869 routed
870 .default_client
871 .clone()
872 .ok_or_else(|| LibsyError::MissingClient {
873 target: target.clone(),
874 })?;
875 let result = client
876 .call(routed.ctx, routed.request, routed.decision)
877 .await
878 .map_err(|source| LibsyError::client_call(target, source));
879 let result = observability::observe_client_call(result);
880 call.respond(result)
881 }
882
883 let stream = self.run_stream(ctx, request, observer);
884 tokio::pin!(stream);
885
886 let mut trace: Vec<Arc<dyn Decision>> = Vec::new();
887 let mut in_flight = futures::stream::FuturesUnordered::new();
888 let mut final_response: Option<Response> = None;
889
890 loop {
891 tokio::select! {
892 Some(result) = in_flight.next() => match result {
893 Ok(()) => {}, Err(err) => return Err(err), },
896 step = stream.next() => {
897 match step {
898 None => break, Some(item) => match item? {
900 Step::CallLlm(call) => in_flight.push(serve(*call)),
901 Step::Decision(decision) => trace.push(decision),
902 Step::ReturnToAgent(response) => {
903 final_response = Some(*response);
904 break;
905 }
906 }
907 }
908 },
909 }
910 }
911 final_response
912 .map(|response| (trace, response))
913 .ok_or(LibsyError::MissingFinalResponse)
914 }
915}
916
917#[cfg(test)]
918mod tests {
919 use super::*;
920 use futures::StreamExt;
921 use switchyard_protocol::{
922 LlmResponse, LlmResponseChunk, completion_text, text_request, text_response,
923 };
924
925 #[derive(Debug, thiserror::Error)]
926 #[error("{0}")]
927 struct TestError(&'static str);
928
929 fn test_error(message: &'static str) -> LibsyError {
930 LibsyError::external("test", TestError(message))
931 }
932
933 struct EchoClient;
935
936 #[async_trait]
937 impl RoutedLlmClient for EchoClient {
938 async fn call(
939 &self,
940 _ctx: Context,
941 _request: Request,
942 decision: Arc<dyn Decision>,
943 ) -> std::result::Result<Response, LlmClientError> {
944 Ok(Response {
946 llm_response: LlmResponse::Agg(text_response(
947 None,
948 decision.selected_model().to_string(),
949 )),
950 metadata: None,
951 })
952 }
953 }
954
955 struct TestDecision {
958 model: String,
959 }
960
961 impl Decision for TestDecision {
962 fn selected_model(&self) -> &str {
963 &self.model
964 }
965 fn reasoning(&self) -> Option<&str> {
966 None
967 }
968 fn as_any(&self) -> &dyn std::any::Any {
969 self
970 }
971 }
972
973 struct TestAlgo {
974 target_set: LlmTargetSet,
975 }
976
977 #[async_trait]
978 impl Algorithm for TestAlgo {
979 fn name(&self) -> &str {
980 "test"
981 }
982
983 async fn create_run_task(
984 self: Arc<Self>,
985 ctx: Context,
986 driver: Driver,
987 request: Request,
988 ) -> Result<Response> {
989 let target = self
990 .target_set
991 .targets()
992 .first()
993 .ok_or(LibsyError::NoTargets)?
994 .clone();
995 let decision: Arc<dyn Decision> = Arc::new(TestDecision {
996 model: target.semantic_name.clone(),
997 });
998 driver.info(ctx.clone(), decision.clone()).await?;
999 driver
1000 .call_llm_target(ctx, &target, request, decision)
1001 .await
1002 }
1003 }
1004
1005 fn orch(target_set: LlmTargetSet) -> Arc<dyn Algorithm> {
1007 Arc::new(TestAlgo { target_set })
1008 }
1009
1010 fn request() -> Request {
1011 Request {
1012 llm_request: text_request(Some("auto".to_string()), "hi".to_string()),
1013 raw_request: None,
1014 metadata: None,
1015 }
1016 }
1017
1018 fn target_set(names: &[(&str, bool)]) -> LlmTargetSet {
1020 let targets = names
1021 .iter()
1022 .map(|(name, has_client)| LlmTarget {
1023 semantic_name: name.to_string(),
1024 llm_client: has_client.then(|| Arc::new(EchoClient) as Arc<dyn RoutedLlmClient>),
1025 })
1026 .collect();
1027 LlmTargetSet::new(targets)
1028 }
1029
1030 #[tokio::test]
1031 async fn observed_run_reports_one_successful_routed_call() -> Result<()> {
1032 let observations = Arc::new(Mutex::new(Vec::new()));
1033 let observed = observations.clone();
1034 let observer: RunObserver = Arc::new(move |observation| observed.lock().push(observation));
1035 let (_, response) = orch(target_set(&[("direct/model", true)]))
1036 .run_observed(Context::default(), request(), Some(observer))
1037 .await?;
1038 assert_eq!(
1039 response.llm_response.as_agg().map(completion_text),
1040 Some("direct/model".to_string())
1041 );
1042 let observations = observations.lock();
1043 assert_eq!(observations.len(), 2);
1044 let RunObservation::LlmCall(observation) = &observations[0] else {
1045 return Err(test_error("expected an LLM call observation"));
1046 };
1047 assert_eq!(observation.selected_model, "direct/model");
1048 assert!(observation.is_routed);
1049 assert!(observation.is_success);
1050 assert!(observation.usage.is_some());
1051 assert!(matches!(
1052 observations[1],
1053 RunObservation::RoutingOverhead(_)
1054 ));
1055 Ok(())
1056 }
1057
1058 #[test]
1059 fn target_lookup_returns_the_missing_target() {
1060 let error = target_set(&[]).get_target("missing").err();
1061 assert!(matches!(
1062 error,
1063 Some(LibsyError::TargetNotFound { target }) if target == "missing"
1064 ));
1065 }
1066
1067 struct StreamingClient {
1070 chunks: Vec<LlmResponseChunk>,
1071 }
1072
1073 #[async_trait]
1074 impl RoutedLlmClient for StreamingClient {
1075 async fn call(
1076 &self,
1077 _ctx: Context,
1078 _request: Request,
1079 _decision: Arc<dyn Decision>,
1080 ) -> std::result::Result<Response, LlmClientError> {
1081 let stream = futures::stream::iter(
1082 self.chunks
1083 .clone()
1084 .into_iter()
1085 .map(|chunk| Ok(chunk.into())),
1086 )
1087 .boxed();
1088 Ok(Response {
1089 llm_response: LlmResponse::Stream(stream),
1090 metadata: None,
1091 })
1092 }
1093 }
1094
1095 fn streaming_orch(chunks: Vec<LlmResponseChunk>) -> Arc<dyn Algorithm> {
1097 let target = LlmTarget {
1098 semantic_name: "stream/model".to_string(),
1099 llm_client: Some(Arc::new(StreamingClient { chunks }) as Arc<dyn RoutedLlmClient>),
1100 };
1101 orch(LlmTargetSet::new(vec![target]))
1102 }
1103
1104 #[tokio::test]
1105 async fn run_returns_a_streamed_response_the_caller_aggregates() -> Result<()> {
1106 let orch = streaming_orch(vec![
1109 LlmResponseChunk::MessageStart {
1110 id: Some("m1".to_string()),
1111 model: Some("stream/model".to_string()),
1112 },
1113 LlmResponseChunk::TextDelta {
1114 index: 0,
1115 text: "hel".to_string(),
1116 },
1117 LlmResponseChunk::TextDelta {
1118 index: 0,
1119 text: "lo".to_string(),
1120 },
1121 LlmResponseChunk::MessageStop {
1122 reason: Some("stop".to_string()),
1123 },
1124 ]);
1125 let (trace, response) = orch.run(Context::default(), request()).await?;
1126 let agg = response
1128 .llm_response
1129 .into_agg()
1130 .await
1131 .map_err(|error| LibsyError::external("aggregating response stream", error))?;
1132 assert_eq!(completion_text(&agg), "hello");
1133 assert_eq!(agg.model.as_deref(), Some("stream/model"));
1134 assert_eq!(trace.len(), 1);
1135 Ok(())
1136 }
1137
1138 #[tokio::test]
1139 async fn aggregating_a_streamed_response_propagates_a_mid_stream_error() -> Result<()> {
1140 let orch = streaming_orch(vec![
1143 LlmResponseChunk::TextDelta {
1144 index: 0,
1145 text: "partial".to_string(),
1146 },
1147 LlmResponseChunk::StreamError {
1148 message: "upstream exploded".to_string(),
1149 },
1150 ]);
1151 let (_, response) = orch.run(Context::default(), request()).await?;
1152 match response.llm_response.into_agg().await {
1153 Ok(_) => panic!("expected a mid-stream error, got an aggregate"),
1154 Err(err) => {
1155 assert!(err.to_string().contains("upstream exploded"));
1156 Ok(())
1157 }
1158 }
1159 }
1160
1161 #[tokio::test]
1162 async fn run_offloads_via_promise_then_returns_to_agent() -> Result<()> {
1163 let stream = orch(target_set(&[("offload/model", false)])).run_stream(
1166 Context::default(),
1167 request(),
1168 None,
1169 );
1170 tokio::pin!(stream);
1171
1172 let mut saw_call = false;
1173 let mut final_completion = None;
1174 while let Some(step) = stream.next().await {
1175 match step? {
1176 Step::CallLlm(call) => {
1177 saw_call = true;
1178 assert_eq!(call.get_decision().selected_model(), "offload/model");
1180 call.respond(Ok(Response {
1182 llm_response: LlmResponse::Agg(text_response(
1183 None,
1184 "fulfilled".to_string(),
1185 )),
1186 metadata: None,
1187 }))?;
1188 }
1189 Step::Decision(decision) => {
1190 assert_eq!(decision.selected_model(), "offload/model");
1191 }
1192 Step::ReturnToAgent(response) => {
1193 final_completion = Some(
1194 response
1195 .llm_response
1196 .as_agg()
1197 .map(completion_text)
1198 .unwrap_or_default(),
1199 );
1200 }
1201 }
1202 }
1203
1204 assert!(saw_call, "expected a CallLlm step before ReturnToAgent");
1205 assert_eq!(
1206 final_completion.ok_or_else(|| test_error("no ReturnToAgent step"))?,
1207 "fulfilled"
1208 );
1209 Ok(())
1210 }
1211
1212 #[tokio::test]
1213 async fn client_backed_target_offloads_with_a_default_client() -> Result<()> {
1214 let stream = orch(target_set(&[("direct/model", true)])).run_stream(
1217 Context::default(),
1218 request(),
1219 None,
1220 );
1221 tokio::pin!(stream);
1222
1223 let mut final_completion = None;
1224 while let Some(step) = stream.next().await {
1225 match step? {
1226 Step::CallLlm(call) => {
1227 let routed = call.get_routed().clone();
1228 let client = routed
1229 .default_client
1230 .clone()
1231 .ok_or_else(|| test_error("expected a default client"))?;
1232 let target = routed.decision.selected_model().to_string();
1233 let result = client
1234 .call(routed.ctx, routed.request, routed.decision)
1235 .await
1236 .map_err(|error| LibsyError::client_call(target, error));
1237 call.respond(result)?;
1238 }
1239 Step::Decision(_) => {}
1240 Step::ReturnToAgent(response) => {
1241 final_completion = Some(
1242 response
1243 .llm_response
1244 .as_agg()
1245 .map(completion_text)
1246 .unwrap_or_default(),
1247 );
1248 }
1249 }
1250 }
1251
1252 assert_eq!(
1254 final_completion.ok_or_else(|| test_error("no ReturnToAgent"))?,
1255 "direct/model"
1256 );
1257 Ok(())
1258 }
1259
1260 #[tokio::test]
1261 async fn run_returns_the_response_when_all_targets_have_clients() -> Result<()> {
1262 let (trace, response) = orch(target_set(&[("direct/model", true)]))
1265 .run(Context::default(), request())
1266 .await?;
1267 assert_eq!(
1269 response
1270 .llm_response
1271 .as_agg()
1272 .map(completion_text)
1273 .unwrap_or_default(),
1274 "direct/model"
1275 );
1276 assert_eq!(trace[0].selected_model(), "direct/model");
1277 Ok(())
1278 }
1279
1280 #[tokio::test]
1281 async fn run_errors_when_a_target_lacks_a_client() -> Result<()> {
1282 let error = orch(target_set(&[("offload/model", false)]))
1285 .run(Context::default(), request())
1286 .await
1287 .err()
1288 .ok_or_else(|| test_error("expected a missing-client error"))?;
1289 assert!(matches!(
1290 error,
1291 LibsyError::MissingClient { target } if target == "offload/model"
1292 ));
1293 Ok(())
1294 }
1295
1296 #[tokio::test(flavor = "multi_thread", worker_threads = 12)]
1297 async fn requests_are_processed_in_parallel() -> Result<()> {
1298 use std::time::Duration;
1299 use tokio::sync::Barrier;
1300
1301 const N: usize = 12;
1302
1303 struct BarrierClient {
1309 barrier: Arc<Barrier>,
1310 }
1311
1312 #[async_trait]
1313 impl RoutedLlmClient for BarrierClient {
1314 async fn call(
1315 &self,
1316 _ctx: Context,
1317 _request: Request,
1318 decision: Arc<dyn Decision>,
1319 ) -> std::result::Result<Response, LlmClientError> {
1320 self.barrier.wait().await;
1321 Ok(Response {
1322 llm_response: LlmResponse::Agg(text_response(
1323 None,
1324 decision.selected_model().to_string(),
1325 )),
1326 metadata: None,
1327 })
1328 }
1329 }
1330
1331 let barrier = Arc::new(Barrier::new(N));
1332 let targets = LlmTargetSet::new(vec![LlmTarget {
1333 semantic_name: "m".to_string(),
1334 llm_client: Some(Arc::new(BarrierClient {
1335 barrier: barrier.clone(),
1336 })),
1337 }]);
1338 let algo = orch(targets);
1340
1341 let mut handles = Vec::new();
1342 for _ in 0..N {
1343 let algo = algo.clone();
1344 handles.push(tokio::spawn(async move {
1345 algo.run(Context::default(), request())
1346 .await
1347 .map(|(_, response)| {
1348 response
1349 .llm_response
1350 .as_agg()
1351 .map(completion_text)
1352 .unwrap_or_default()
1353 })
1354 }));
1355 }
1356
1357 for handle in handles {
1358 let completion = tokio::time::timeout(Duration::from_secs(5), handle)
1360 .await
1361 .map_err(|error| LibsyError::external("waiting for test task", error))?
1362 .map_err(|source| LibsyError::AlgorithmTask { source })??;
1363 assert_eq!(completion, "m");
1364 }
1365 Ok(())
1366 }
1367
1368 #[tokio::test]
1369 async fn offload_error_propagates_back_to_the_algorithm() -> Result<()> {
1370 let stream = orch(target_set(&[("offload/model", false)])).run_stream(
1374 Context::default(),
1375 request(),
1376 None,
1377 );
1378 tokio::pin!(stream);
1379
1380 let mut saw_error = false;
1381 while let Some(step) = stream.next().await {
1382 match step {
1383 Ok(Step::CallLlm(call)) => {
1384 call.respond(Err(test_error("upstream model call failed")))?;
1385 }
1386 Ok(Step::Decision(_)) => {}
1387 Ok(Step::ReturnToAgent(..)) => {
1388 return Err(test_error(
1389 "expected the offload error to propagate, got a response",
1390 ));
1391 }
1392 Err(err) => {
1393 assert!(err.to_string().contains("upstream model call failed"));
1395 saw_error = true;
1396 }
1397 }
1398 }
1399
1400 assert!(saw_error, "expected an error step");
1401 Ok(())
1402 }
1403
1404 #[tokio::test]
1405 async fn dropping_the_stream_cancels_the_algorithm_task() -> Result<()> {
1406 use std::sync::atomic::{AtomicBool, Ordering};
1407 use std::time::Duration;
1408 use tokio::sync::mpsc;
1409
1410 struct DropGuard(Arc<AtomicBool>);
1413 impl Drop for DropGuard {
1414 fn drop(&mut self) {
1415 self.0.store(true, Ordering::SeqCst);
1416 }
1417 }
1418
1419 struct StuckAlgo {
1420 started: mpsc::UnboundedSender<()>,
1421 dropped: Arc<AtomicBool>,
1422 }
1423
1424 #[async_trait]
1425 impl Algorithm for StuckAlgo {
1426 fn name(&self) -> &str {
1427 "stuck"
1428 }
1429
1430 async fn create_run_task(
1431 self: Arc<Self>,
1432 _ctx: Context,
1433 _driver: Driver,
1434 _request: Request,
1435 ) -> Result<Response> {
1436 let _guard = DropGuard(self.dropped.clone());
1437 let _ = self.started.send(());
1438 std::future::pending::<()>().await;
1440 unreachable!()
1441 }
1442 }
1443
1444 let (started_tx, mut started_rx) = mpsc::unbounded_channel();
1445 let dropped = Arc::new(AtomicBool::new(false));
1446 let algo: Arc<dyn Algorithm> = Arc::new(StuckAlgo {
1447 started: started_tx,
1448 dropped: dropped.clone(),
1449 });
1450
1451 let stream = algo.run_stream(Context::default(), request(), None);
1452 started_rx
1453 .recv()
1454 .await
1455 .ok_or_else(|| test_error("task never started"))?;
1456 drop(stream);
1457 tokio::time::sleep(Duration::from_millis(100)).await;
1458
1459 assert!(
1460 dropped.load(Ordering::SeqCst),
1461 "algorithm task was NOT cancelled after dropping the stream"
1462 );
1463 Ok(())
1464 }
1465
1466 #[tokio::test]
1467 async fn create_run_task_panic_surfaces_as_a_stream_error() -> Result<()> {
1468 struct Panicky;
1471
1472 #[async_trait]
1473 impl Algorithm for Panicky {
1474 fn name(&self) -> &str {
1475 "panicky"
1476 }
1477
1478 async fn create_run_task(
1479 self: Arc<Self>,
1480 _ctx: Context,
1481 _driver: Driver,
1482 _request: Request,
1483 ) -> Result<Response> {
1484 panic!("boom");
1485 }
1486 }
1487
1488 let algo: Arc<dyn Algorithm> = Arc::new(Panicky);
1489 let stream = algo.run_stream(Context::default(), request(), None);
1490 tokio::pin!(stream);
1491
1492 let mut saw_error = false;
1493 while let Some(step) = stream.next().await {
1494 match step {
1495 Err(err) => {
1496 assert!(matches!(err, LibsyError::AlgorithmTask { .. }));
1497 saw_error = true;
1498 }
1499 Ok(_) => return Err(test_error("expected the panic to surface as an error step")),
1500 }
1501 }
1502
1503 assert!(saw_error, "expected an error step from the panicked task");
1504 Ok(())
1505 }
1506
1507 #[tokio::test]
1508 async fn run_returns_an_error_when_the_algorithm_task_panics() -> Result<()> {
1509 struct Panicky;
1512
1513 #[async_trait]
1514 impl Algorithm for Panicky {
1515 fn name(&self) -> &str {
1516 "panicky"
1517 }
1518
1519 async fn create_run_task(
1520 self: Arc<Self>,
1521 _ctx: Context,
1522 _driver: Driver,
1523 _request: Request,
1524 ) -> Result<Response> {
1525 panic!("boom");
1526 }
1527 }
1528
1529 let algo: Arc<dyn Algorithm> = Arc::new(Panicky);
1530 match algo.run(Context::default(), request()).await {
1531 Ok(_) => Err(test_error(
1532 "expected run to surface the algorithm panic as an error",
1533 )),
1534 Err(err) => {
1535 assert!(matches!(err, LibsyError::AlgorithmTask { .. }));
1536 Ok(())
1537 }
1538 }
1539 }
1540
1541 #[tokio::test]
1542 async fn cancelling_run_cancels_the_algorithm_task() -> Result<()> {
1543 use std::sync::atomic::{AtomicBool, Ordering};
1544 use std::time::Duration;
1545 use tokio::sync::mpsc;
1546
1547 struct DropGuard(Arc<AtomicBool>);
1550 impl Drop for DropGuard {
1551 fn drop(&mut self) {
1552 self.0.store(true, Ordering::SeqCst);
1553 }
1554 }
1555
1556 struct StuckAlgo {
1557 started: mpsc::UnboundedSender<()>,
1558 dropped: Arc<AtomicBool>,
1559 }
1560
1561 #[async_trait]
1562 impl Algorithm for StuckAlgo {
1563 fn name(&self) -> &str {
1564 "stuck"
1565 }
1566
1567 async fn create_run_task(
1568 self: Arc<Self>,
1569 _ctx: Context,
1570 _driver: Driver,
1571 _request: Request,
1572 ) -> Result<Response> {
1573 let _guard = DropGuard(self.dropped.clone());
1574 let _ = self.started.send(());
1575 std::future::pending::<()>().await;
1578 unreachable!()
1579 }
1580 }
1581
1582 let (started_tx, mut started_rx) = mpsc::unbounded_channel();
1583 let dropped = Arc::new(AtomicBool::new(false));
1584 let algo: Arc<dyn Algorithm> = Arc::new(StuckAlgo {
1585 started: started_tx,
1586 dropped: dropped.clone(),
1587 });
1588
1589 let run_task = tokio::spawn(async move { algo.run(Context::default(), request()).await });
1592 started_rx
1593 .recv()
1594 .await
1595 .ok_or_else(|| test_error("task never started"))?;
1596 run_task.abort();
1597 tokio::time::sleep(Duration::from_millis(100)).await;
1598
1599 assert!(
1600 dropped.load(Ordering::SeqCst),
1601 "algorithm task was NOT cancelled after cancelling run"
1602 );
1603 Ok(())
1604 }
1605
1606 struct LoserClient {
1612 started: Arc<tokio::sync::Notify>,
1613 delay: Option<std::time::Duration>,
1614 }
1615
1616 #[async_trait]
1617 impl RoutedLlmClient for LoserClient {
1618 async fn call(
1619 &self,
1620 _ctx: Context,
1621 _request: Request,
1622 decision: Arc<dyn Decision>,
1623 ) -> std::result::Result<Response, LlmClientError> {
1624 self.started.notify_one();
1625 match self.delay {
1626 Some(delay) => tokio::time::sleep(delay).await,
1627 None => std::future::pending::<()>().await,
1628 }
1629 Ok(Response {
1630 llm_response: LlmResponse::Agg(text_response(
1631 None,
1632 decision.selected_model().to_string(),
1633 )),
1634 metadata: None,
1635 })
1636 }
1637 }
1638
1639 struct GatedEchoClient {
1642 gate: Arc<tokio::sync::Notify>,
1643 }
1644
1645 #[async_trait]
1646 impl RoutedLlmClient for GatedEchoClient {
1647 async fn call(
1648 &self,
1649 _ctx: Context,
1650 _request: Request,
1651 decision: Arc<dyn Decision>,
1652 ) -> std::result::Result<Response, LlmClientError> {
1653 self.gate.notified().await;
1654 Ok(Response {
1655 llm_response: LlmResponse::Agg(text_response(
1656 None,
1657 decision.selected_model().to_string(),
1658 )),
1659 metadata: None,
1660 })
1661 }
1662 }
1663
1664 struct Hedge {
1667 winner: LlmTarget,
1668 loser: LlmTarget,
1669 }
1670
1671 #[async_trait]
1672 impl Algorithm for Hedge {
1673 fn name(&self) -> &str {
1674 "hedge"
1675 }
1676
1677 async fn create_run_task(
1678 self: Arc<Self>,
1679 ctx: Context,
1680 driver: Driver,
1681 request: Request,
1682 ) -> Result<Response> {
1683 let dec_w: Arc<dyn Decision> = Arc::new(TestDecision {
1684 model: self.winner.semantic_name.clone(),
1685 });
1686 let dec_l: Arc<dyn Decision> = Arc::new(TestDecision {
1687 model: self.loser.semantic_name.clone(),
1688 });
1689 let win = driver.call_llm_target(ctx.clone(), &self.winner, request.clone(), dec_w);
1690 let lose = driver.call_llm_target(ctx, &self.loser, request, dec_l);
1691 tokio::select! {
1693 res = win => res,
1694 res = lose => res,
1695 }
1696 }
1697 }
1698
1699 fn hedge(loser_delay: Option<std::time::Duration>) -> Arc<dyn Algorithm> {
1702 let started = Arc::new(tokio::sync::Notify::new());
1703 let winner = LlmTarget {
1704 semantic_name: "winner".to_string(),
1705 llm_client: Some(Arc::new(GatedEchoClient {
1706 gate: started.clone(),
1707 })),
1708 };
1709 let loser = LlmTarget {
1710 semantic_name: "loser".to_string(),
1711 llm_client: Some(Arc::new(LoserClient {
1712 started,
1713 delay: loser_delay,
1714 })),
1715 };
1716 Arc::new(Hedge { winner, loser })
1717 }
1718
1719 #[tokio::test]
1720 async fn run_returns_the_winner_without_a_late_loser_overwriting_it() -> Result<()> {
1721 let (_trace, response) = hedge(Some(std::time::Duration::from_millis(50)))
1724 .run(Context::default(), request())
1725 .await?;
1726 assert_eq!(
1727 response
1728 .llm_response
1729 .as_agg()
1730 .map(completion_text)
1731 .unwrap_or_default(),
1732 "winner"
1733 );
1734 Ok(())
1735 }
1736
1737 #[tokio::test]
1738 async fn run_returns_the_winner_without_hanging_on_a_pending_loser() -> Result<()> {
1739 let run = hedge(None).run(Context::default(), request());
1742 let (_trace, response) = tokio::time::timeout(std::time::Duration::from_secs(1), run)
1743 .await
1744 .map_err(|error| LibsyError::external("waiting for pending loser", error))??;
1745 assert_eq!(
1746 response
1747 .llm_response
1748 .as_agg()
1749 .map(completion_text)
1750 .unwrap_or_default(),
1751 "winner"
1752 );
1753 Ok(())
1754 }
1755
1756 #[tokio::test]
1757 async fn run_surfaces_a_terminal_error_with_many_calls_in_flight() -> Result<()> {
1758 use std::sync::atomic::{AtomicUsize, Ordering};
1759
1760 const N: usize = 10;
1763
1764 struct EnterThenPend {
1766 started: Arc<AtomicUsize>,
1767 all_started: Arc<tokio::sync::Notify>,
1768 n: usize,
1769 }
1770
1771 #[async_trait]
1772 impl RoutedLlmClient for EnterThenPend {
1773 async fn call(
1774 &self,
1775 _ctx: Context,
1776 _request: Request,
1777 _decision: Arc<dyn Decision>,
1778 ) -> std::result::Result<Response, LlmClientError> {
1779 if self.started.fetch_add(1, Ordering::SeqCst) + 1 == self.n {
1780 self.all_started.notify_one();
1781 }
1782 std::future::pending::<()>().await;
1783 unreachable!()
1784 }
1785 }
1786
1787 struct FanOutThenError {
1790 target: LlmTarget,
1791 all_started: Arc<tokio::sync::Notify>,
1792 n: usize,
1793 }
1794
1795 #[async_trait]
1796 impl Algorithm for FanOutThenError {
1797 fn name(&self) -> &str {
1798 "fan_out_then_error"
1799 }
1800
1801 async fn create_run_task(
1802 self: Arc<Self>,
1803 ctx: Context,
1804 driver: Driver,
1805 request: Request,
1806 ) -> Result<Response> {
1807 let offloads = futures::future::join_all((0..self.n).map(|i| {
1808 let decision: Arc<dyn Decision> = Arc::new(TestDecision {
1809 model: format!("m{i}"),
1810 });
1811 driver.call_llm_target(ctx.clone(), &self.target, request.clone(), decision)
1812 }));
1813 tokio::select! {
1814 _ = offloads => Err(test_error("offloads unexpectedly completed")),
1815 _ = self.all_started.notified() => {
1816 Err(test_error("terminal error while calls pending"))
1817 }
1818 }
1819 }
1820 }
1821
1822 let all_started = Arc::new(tokio::sync::Notify::new());
1823 let target = LlmTarget {
1824 semantic_name: "pending".to_string(),
1825 llm_client: Some(Arc::new(EnterThenPend {
1826 started: Arc::new(AtomicUsize::new(0)),
1827 all_started: all_started.clone(),
1828 n: N,
1829 })),
1830 };
1831 let algo: Arc<dyn Algorithm> = Arc::new(FanOutThenError {
1832 target,
1833 all_started,
1834 n: N,
1835 });
1836
1837 let run = algo.run(Context::default(), request());
1840 let result = tokio::time::timeout(std::time::Duration::from_millis(500), run)
1841 .await
1842 .map_err(|error| {
1843 LibsyError::external("waiting for terminal error with full call cap", error)
1844 })?;
1845 match result {
1846 Ok(_) => Err(test_error("expected the terminal error, got a response")),
1847 Err(err) => {
1848 assert!(
1849 err.to_string()
1850 .contains("terminal error while calls pending")
1851 );
1852 Ok(())
1853 }
1854 }
1855 }
1856}