1use std::{
9 collections::{HashMap, HashSet},
10 future::Future,
11 pin::Pin,
12 sync::Arc,
13 time::{Duration, Instant},
14};
15
16use async_trait::async_trait;
17use futures::{Stream, StreamExt};
18use parking_lot::Mutex;
19use tracing::Instrument;
20
21use switchyard_protocol::{
29 Context, Decision, LlmClientError, Request, Response, RoutingFallbackReason, Signals, Usage,
30};
31
32use super::driver::{DriverRequest, DriverStep, TypeErasedDriver};
33use crate::{DriverError, LibsyError, Result, observability};
34
35pub type StepStream = Pin<Box<dyn Stream<Item = Result<Step>> + Send>>;
39
40#[derive(Clone, Debug)]
42pub struct LlmCallObservation {
43 pub selected_model: String,
45 pub tier: Option<String>,
47 pub is_routed: bool,
49 pub is_success: bool,
51 pub duration: Duration,
53 pub usage: Option<Usage>,
55}
56
57#[derive(Clone, Debug)]
59pub enum RunObservation {
60 LlmCall(LlmCallObservation),
62 RoutingOverhead(Duration),
64}
65
66pub type RunObserver = Arc<dyn Fn(RunObservation) + Send + Sync>;
72
73#[derive(Clone)]
82pub struct RoutedRequest {
83 pub request: Request,
85 pub decision: Arc<dyn Decision>,
87 pub ctx: Context,
90}
91
92pub struct CallLlmRequest {
101 inner: DriverRequest,
102 routed: RoutedRequest,
103}
104
105impl CallLlmRequest {
106 fn new(inner: DriverRequest) -> Self {
109 let routed = match inner.request::<RoutedRequest>() {
112 Ok(routed) => routed.clone(),
113 Err(_) => unreachable!("CallLlmRequest payload is always a RoutedRequest"),
114 };
115 Self { inner, routed }
116 }
117
118 pub fn get_routed(&self) -> &RoutedRequest {
121 &self.routed
122 }
123
124 pub fn get_request(&self) -> &Request {
126 &self.get_routed().request
127 }
128
129 pub fn get_decision(&self) -> &dyn Decision {
131 self.get_routed().decision.as_ref()
132 }
133
134 pub fn respond(self, result: Result<Response>) -> Result<()> {
138 self.inner.respond::<Response>(result)
139 }
140}
141
142#[derive(Clone)]
149pub struct Driver {
150 driver: TypeErasedDriver,
151 routed_call: Arc<Mutex<Option<Duration>>>,
153 observer: Option<RunObserver>,
154}
155
156impl Driver {
157 pub(crate) fn new() -> Self {
160 Self::with_observer(None)
161 }
162
163 fn with_observer(observer: Option<RunObserver>) -> Self {
164 Self {
165 driver: TypeErasedDriver::new(),
166 routed_call: Arc::new(Mutex::new(None)),
167 observer,
168 }
169 }
170
171 pub(crate) fn routed_call_duration(&self) -> Option<Duration> {
173 *self.routed_call.lock()
174 }
175
176 pub(crate) fn observe_routing_overhead(&self, duration: Duration) {
179 if let Some(observer) = &self.observer {
180 observer(RunObservation::RoutingOverhead(duration));
181 }
182 }
183
184 #[tracing::instrument(
193 target = "libsy",
194 name = "libsy.llm_call",
195 skip_all,
196 fields(
197 algorithm = observability::algorithm_label(&routed.ctx),
198 selected_model = routed.decision.selected_model(),
199 openinference.span.kind = "CHAIN",
200 outcome = tracing::field::Empty,
201 error = tracing::field::Empty,
202 input_tokens = tracing::field::Empty,
203 output_tokens = tracing::field::Empty,
204 total_tokens = tracing::field::Empty,
205 reasoning_tokens = tracing::field::Empty,
206 )
207 )]
208 pub async fn call_llm(&self, routed: RoutedRequest) -> Result<Response> {
209 let algorithm = observability::algorithm_label(&routed.ctx).to_string();
210 let selected_model = routed.decision.selected_model().to_string();
211 let tier = routed.decision.routing_tier().map(str::to_string);
212 let is_routed = routed.decision.is_routed_call();
213 let started = Instant::now();
214 let result = self
215 .driver
216 .fulfill_request::<RoutedRequest, Response>(routed.ctx.clone(), routed)
217 .await;
218 let elapsed = started.elapsed();
219 observability::record_llm_call(
220 &algorithm,
221 &selected_model,
222 tier.as_deref(),
223 is_routed,
224 elapsed,
225 &result,
226 &tracing::Span::current(),
227 );
228 if let Some(observer) = &self.observer {
229 observer(RunObservation::LlmCall(LlmCallObservation {
230 selected_model,
231 tier,
232 is_routed,
233 is_success: result.is_ok(),
234 duration: elapsed,
235 usage: result
236 .as_ref()
237 .ok()
238 .and_then(|response| response.llm_response.as_agg())
239 .map(|response| response.usage.clone()),
240 }));
241 }
242 if is_routed && result.is_ok() {
245 *self.routed_call.lock() = Some(elapsed);
246 }
247 result
248 }
249
250 pub async fn info(&self, ctx: Context, decision: Arc<dyn Decision>) -> Result<()> {
254 self.driver.info(ctx.clone(), decision.clone()).await?;
255 observability::record_decision(&ctx, decision.as_ref());
256 Ok(())
257 }
258
259 pub(crate) async fn finish(&self, ctx: Context, result: Result<Response>) -> Result<()> {
263 match result {
264 Ok(response) => self.driver.done(ctx, response).await,
265 Err(err) => self.driver.fail(ctx, err).await,
266 }
267 }
268
269 pub(crate) fn stream(&self) -> impl Stream<Item = Result<Step>> + use<> {
273 self.driver.stream().map(|item| match item? {
274 DriverStep::Request(req) => Ok(Step::CallLlm(Box::new(CallLlmRequest::new(req)))),
275 DriverStep::Info(payload) => payload
276 .downcast::<Arc<dyn Decision>>()
277 .map(|decision| Step::Decision(*decision))
278 .map_err(|_| {
279 DriverError::TypeMismatch {
280 expected: "Arc<dyn Decision>",
281 }
282 .into()
283 }),
284 DriverStep::Done(payload) => payload
285 .downcast::<Response>()
286 .map(Step::ReturnToAgent)
287 .map_err(|_| {
288 DriverError::TypeMismatch {
289 expected: "Response",
290 }
291 .into()
292 }),
293 })
294 }
295}
296
297impl Default for Driver {
298 fn default() -> Self {
299 Self::new()
300 }
301}
302
303pub enum Step {
305 CallLlm(Box<CallLlmRequest>),
308 Decision(Arc<dyn Decision>),
311 ReturnToAgent(Box<Response>),
313}
314
315pub async fn drive<F, Fut>(
330 algorithm: Arc<dyn Algorithm>,
331 ctx: Context,
332 request: Request,
333 observer: Option<RunObserver>,
334 serve: F,
335) -> Result<(Vec<Arc<dyn Decision>>, Response)>
336where
337 F: Fn(CallLlmRequest) -> Fut,
338 Fut: Future<Output = Result<()>>,
339{
340 let stream = algorithm.run_stream(ctx, request, observer);
341 tokio::pin!(stream);
342
343 let mut trace: Vec<Arc<dyn Decision>> = Vec::new();
344 let mut in_flight = futures::stream::FuturesUnordered::new();
345 let mut final_response: Option<Response> = None;
346
347 loop {
348 tokio::select! {
349 Some(result) = in_flight.next() => match result {
350 Ok(()) => {}, Err(err) => return Err(err), },
353 step = stream.next() => {
354 match step {
355 None => break, Some(item) => match item? {
357 Step::CallLlm(call) => in_flight.push(serve(*call)),
358 Step::Decision(decision) => trace.push(decision),
359 Step::ReturnToAgent(response) => {
360 final_response = Some(*response);
361 break;
362 }
363 }
364 }
365 },
366 }
367 }
368 final_response
369 .map(|response| (trace, response))
370 .ok_or(LibsyError::MissingFinalResponse)
371}
372
373struct AbortOnDrop(tokio::task::AbortHandle);
375
376impl Drop for AbortOnDrop {
377 fn drop(&mut self) {
378 self.0.abort();
379 }
380}
381
382#[derive(Clone)]
386pub struct LlmTarget {
387 pub semantic_name: String,
391}
392
393#[derive(Clone)]
397pub struct LlmTargetSet {
398 targets: Vec<LlmTarget>,
399}
400
401impl LlmTargetSet {
402 pub fn new(targets: Vec<LlmTarget>) -> Self {
404 Self { targets }
405 }
406
407 pub fn targets(&self) -> &[LlmTarget] {
409 &self.targets
410 }
411
412 pub fn get_target(&self, name: &str) -> Result<LlmTarget> {
414 self.targets
415 .iter()
416 .find(|t| t.semantic_name == name)
417 .cloned()
418 .ok_or_else(|| LibsyError::TargetNotFound {
419 target: name.to_string(),
420 })
421 }
422
423 pub fn resolve_target(&self, name: &str, ctx: &Context) -> Result<LlmTarget> {
426 let target = self.get_target(name)?;
427 if !ctx.is_excluded(&target.semantic_name) {
428 return Ok(target);
429 }
430 self.targets
431 .iter()
432 .find(|t| !ctx.is_excluded(&t.semantic_name))
433 .cloned()
434 .ok_or(LibsyError::AllTargetsExcluded)
435 }
436}
437
438#[derive(Clone, Hash, PartialEq, Eq)]
442pub(crate) enum RoutingIdentity {
443 Session(String),
445 Subagent { session: String, agent: String },
447}
448
449impl RoutingIdentity {
450 pub(crate) fn from_request(request: &Request) -> Option<Self> {
455 let metadata = request.metadata.as_ref()?;
456 let session = metadata.session_id.as_deref().filter(|id| !id.is_empty())?;
457 if metadata.is_subagent {
458 let agent = metadata.agent_id.as_deref().filter(|id| !id.is_empty())?;
459 Some(Self::Subagent {
460 session: session.to_string(),
461 agent: agent.to_string(),
462 })
463 } else {
464 Some(Self::Session(session.to_string()))
465 }
466 }
467
468 fn session(&self) -> &str {
470 match self {
471 Self::Session(session) | Self::Subagent { session, .. } => session,
472 }
473 }
474}
475
476const MAX_EVICTION_IDENTITIES: usize = 1_024;
479
480#[derive(Default)]
486pub(crate) struct SessionEvictions {
487 by_identity: Mutex<HashMap<RoutingIdentity, HashSet<String>>>,
488}
489
490impl SessionEvictions {
491 pub(crate) fn remove_session(&self, session: &str) {
493 self.by_identity
494 .lock()
495 .retain(|identity, _| identity.session() != session);
496 }
497
498 fn evicted_for(&self, identity: Option<&RoutingIdentity>) -> Vec<String> {
500 let Some(identity) = identity else {
501 return Vec::new();
502 };
503 self.by_identity
504 .lock()
505 .get(identity)
506 .map(|targets| targets.iter().cloned().collect())
507 .unwrap_or_default()
508 }
509
510 fn record(&self, identity: Option<&RoutingIdentity>, target: &str) {
513 let Some(identity) = identity else { return };
514 let mut histories = self.by_identity.lock();
515 if histories.len() >= MAX_EVICTION_IDENTITIES
516 && !histories.contains_key(identity)
517 && let Some(oldest) = histories.keys().next().cloned()
518 {
519 histories.remove(&oldest);
520 }
521 histories
522 .entry(identity.clone())
523 .or_default()
524 .insert(target.to_string());
525 }
526}
527
528fn eligible_targets(targets: &LlmTargetSet, ctx: &Context) -> usize {
530 targets
531 .targets()
532 .iter()
533 .filter(|t| !ctx.is_excluded(&t.semantic_name))
534 .count()
535}
536
537pub(crate) fn exclude_evicted(
540 ctx: &mut Context,
541 targets: &LlmTargetSet,
542 evictions: &SessionEvictions,
543 identity: Option<&RoutingIdentity>,
544) {
545 for target in evictions.evicted_for(identity) {
546 if eligible_targets(targets, ctx) <= 1 {
549 break;
550 }
551 ctx.exclude_target(target);
552 }
553}
554
555fn classify_fallback(error: &LibsyError) -> Option<(&str, RoutingFallbackReason)> {
557 let LibsyError::ClientCall { target, source } = error else {
558 return None;
559 };
560 let reason = match source {
561 LlmClientError::ContextWindowExceeded { .. } => RoutingFallbackReason::ContextWindow,
562 LlmClientError::Transport { .. } | LlmClientError::Timeout { .. } => {
563 RoutingFallbackReason::Unavailable
564 }
565 LlmClientError::UpstreamHttp { status, .. }
566 if matches!(*status, 403 | 408 | 429) || (500..=599).contains(status) =>
567 {
568 RoutingFallbackReason::Unavailable
569 }
570 _ => return None,
571 };
572 Some((target, reason))
573}
574
575#[allow(clippy::too_many_arguments)]
583pub(crate) async fn call_llm_with_fallback(
584 mut ctx: Context,
585 driver: &Driver,
586 targets: &LlmTargetSet,
587 mut target: LlmTarget,
588 mut decision: Arc<dyn Decision>,
589 request: Request,
590 identity: Option<&RoutingIdentity>,
591 evictions: &SessionEvictions,
592 target_unavailable: impl Fn(&Request, &str),
593 fallback_decision: impl Fn(&LlmTarget, &LlmTarget, RoutingFallbackReason) -> Arc<dyn Decision>,
594) -> Result<Response> {
595 loop {
596 let result = driver
597 .call_llm(RoutedRequest {
598 request: request.clone(),
599 decision: decision.clone(),
600 ctx: ctx.clone(),
601 })
602 .await;
603 let Err(error) = result else { return result };
604 let Some((failed, reason)) = classify_fallback(&error) else {
605 return Err(error);
606 };
607 if !ctx.exclude_target(failed) {
610 return Err(error);
611 }
612 match reason {
613 RoutingFallbackReason::ContextWindow => evictions.record(identity, failed),
614 RoutingFallbackReason::Unavailable => target_unavailable(&request, failed),
615 }
616 let Ok(next) = targets.resolve_target(&target.semantic_name, &ctx) else {
617 return Err(error);
618 };
619 decision = fallback_decision(&target, &next, reason);
620 target = next;
621 driver.info(ctx.clone(), decision.clone()).await?;
622 }
623}
624
625#[async_trait]
647pub trait Algorithm: Send + Sync + 'static {
648 fn name(&self) -> &str;
652
653 async fn create_run_task(
659 self: Arc<Self>,
660 ctx: Context,
661 driver: Driver,
662 request: Request,
663 ) -> Result<Response>;
664
665 #[allow(unused_variables)]
669 async fn process_signals(self: Arc<Self>, signals: Signals) -> Result<()> {
670 Ok(())
671 }
672
673 fn run_stream(
683 self: Arc<Self>,
684 ctx: Context,
685 request: Request,
686 observer: Option<RunObserver>,
687 ) -> StepStream {
688 let mut ctx = ctx;
691 ctx.values.insert(
692 observability::ALGORITHM_KEY.to_string(),
693 self.name().to_string(),
694 );
695 let driver = Driver::with_observer(observer);
696 let task_driver = driver.clone();
697 let task_ctx = ctx.clone();
698 let stream = task_driver.stream();
699 let span = observability::run_span(self.name(), &request);
703 let observed_driver = task_driver.clone();
704 let handle = tokio::spawn(
705 async move {
706 observability::observe_run(
707 task_ctx.clone(),
708 observed_driver,
709 self.create_run_task(task_ctx, task_driver, request),
710 )
711 .await
712 }
713 .instrument(span),
714 );
715 let abort_guard = AbortOnDrop(handle.abort_handle());
717
718 let finish_driver = driver.clone();
719 let finish_ctx = ctx;
720 let tail: StepStream = Box::pin(
721 futures::stream::once(async move {
722 let result = match handle.await {
723 Ok(response) => response,
724 Err(source) => Err(LibsyError::AlgorithmTask { source }),
725 };
726 finish_driver.finish(finish_ctx, result).await
727 })
728 .filter_map(|finish_result| async move { finish_result.err().map(Err) }),
729 );
730
731 let stream: StepStream = Box::pin(stream);
732 Box::pin(futures::stream::select(stream, tail).map(move |step| {
733 let _keep_alive = &abort_guard;
735 step
736 }))
737 }
738}
739
740#[cfg(test)]
741mod tests {
742 use super::*;
743 use crate::core::testing::{Serve, ServeResult, drive, drive_observed, echo, reply};
744 use futures::StreamExt;
745 use switchyard_protocol::{
746 LlmResponse, LlmResponseChunk, completion_text, text_request, text_response,
747 };
748
749 #[derive(Debug, thiserror::Error)]
750 #[error("{0}")]
751 struct TestError(&'static str);
752
753 fn test_error(message: &'static str) -> LibsyError {
754 LibsyError::external("test", TestError(message))
755 }
756
757 fn classified_client_error(source: LlmClientError) -> Option<RoutingFallbackReason> {
758 classify_fallback(&LibsyError::client_call("target", source)).map(|(_, reason)| reason)
759 }
760
761 #[test]
762 fn route_fallback_only_accepts_context_and_unavailable_failures() {
763 assert_eq!(
764 classified_client_error(LlmClientError::ContextWindowExceeded {
765 model: "target".to_string(),
766 message: "too long".to_string(),
767 }),
768 Some(RoutingFallbackReason::ContextWindow)
769 );
770 for source in [
771 LlmClientError::Transport {
772 source: Box::new(std::io::Error::other("connection failed")),
773 },
774 LlmClientError::Timeout {
775 source: Box::new(std::io::Error::other("request timed out")),
776 },
777 ] {
778 assert_eq!(
779 classified_client_error(source),
780 Some(RoutingFallbackReason::Unavailable)
781 );
782 }
783 for (status, expected) in [
784 (400, None),
785 (401, None),
786 (403, Some(RoutingFallbackReason::Unavailable)),
787 (404, None),
788 (408, Some(RoutingFallbackReason::Unavailable)),
789 (409, None),
790 (429, Some(RoutingFallbackReason::Unavailable)),
791 (499, None),
792 (500, Some(RoutingFallbackReason::Unavailable)),
793 (599, Some(RoutingFallbackReason::Unavailable)),
794 (600, None),
795 ] {
796 assert_eq!(
797 classified_client_error(LlmClientError::UpstreamHttp {
798 status,
799 body: "failed".to_string(),
800 }),
801 expected
802 );
803 }
804 assert_eq!(
805 classified_client_error(LlmClientError::InvalidResponse {
806 source: Box::new(std::io::Error::other("invalid response")),
807 }),
808 None
809 );
810 }
811
812 struct TestDecision {
815 model: String,
816 }
817
818 impl Decision for TestDecision {
819 fn selected_model(&self) -> &str {
820 &self.model
821 }
822 fn reasoning(&self) -> Option<&str> {
823 None
824 }
825 fn as_any(&self) -> &dyn std::any::Any {
826 self
827 }
828 }
829
830 struct TestAlgo {
831 target_set: LlmTargetSet,
832 }
833
834 #[async_trait]
835 impl Algorithm for TestAlgo {
836 fn name(&self) -> &str {
837 "test"
838 }
839
840 async fn create_run_task(
841 self: Arc<Self>,
842 ctx: Context,
843 driver: Driver,
844 request: Request,
845 ) -> Result<Response> {
846 let target = self
847 .target_set
848 .targets()
849 .first()
850 .ok_or(LibsyError::NoTargets)?
851 .clone();
852 let decision: Arc<dyn Decision> = Arc::new(TestDecision {
853 model: target.semantic_name.clone(),
854 });
855 driver.info(ctx.clone(), decision.clone()).await?;
856 driver
857 .call_llm(RoutedRequest {
858 request,
859 decision,
860 ctx,
861 })
862 .await
863 }
864 }
865
866 fn orch(target_set: LlmTargetSet) -> Arc<dyn Algorithm> {
868 Arc::new(TestAlgo { target_set })
869 }
870
871 fn request() -> Request {
872 Request {
873 llm_request: text_request(Some("auto".to_string()), "hi".to_string()),
874 raw_request: None,
875 metadata: None,
876 }
877 }
878
879 fn target_set(names: &[&str]) -> LlmTargetSet {
880 let targets = names
881 .iter()
882 .map(|name| LlmTarget {
883 semantic_name: name.to_string(),
884 })
885 .collect();
886 LlmTargetSet::new(targets)
887 }
888
889 #[tokio::test]
890 async fn observed_run_reports_one_successful_routed_call() -> Result<()> {
891 let observations = Arc::new(Mutex::new(Vec::new()));
892 let observed = observations.clone();
893 let observer: RunObserver = Arc::new(move |observation| observed.lock().push(observation));
894 let (_, response) = drive_observed(
895 orch(target_set(&["direct/model"])),
896 Context::default(),
897 request(),
898 Some(observer),
899 echo(),
900 )
901 .await?;
902 assert_eq!(
903 response.llm_response.as_agg().map(completion_text),
904 Some("direct/model".to_string())
905 );
906 let observations = observations.lock();
907 assert_eq!(observations.len(), 2);
908 let RunObservation::LlmCall(observation) = &observations[0] else {
909 return Err(test_error("expected an LLM call observation"));
910 };
911 assert_eq!(observation.selected_model, "direct/model");
912 assert!(observation.is_routed);
913 assert!(observation.is_success);
914 assert!(observation.usage.is_some());
915 assert!(matches!(
916 observations[1],
917 RunObservation::RoutingOverhead(_)
918 ));
919 Ok(())
920 }
921
922 #[test]
923 fn target_lookup_returns_the_missing_target() {
924 let error = target_set(&[]).get_target("missing").err();
925 assert!(matches!(
926 error,
927 Some(LibsyError::TargetNotFound { target }) if target == "missing"
928 ));
929 }
930
931 fn streaming_orch(chunks: Vec<LlmResponseChunk>) -> (Arc<dyn Algorithm>, impl Serve) {
934 let algo = orch(target_set(&["stream/model"]));
935 let serve = move |_decision: Arc<dyn Decision>, _request: Request| {
936 let chunks = chunks.clone();
937 async move {
938 let stream =
939 futures::stream::iter(chunks.into_iter().map(|chunk| Ok(chunk.into()))).boxed();
940 Ok(Response {
941 llm_response: LlmResponse::Stream(stream),
942 metadata: None,
943 })
944 }
945 };
946 (algo, serve)
947 }
948
949 #[tokio::test]
950 async fn run_returns_a_streamed_response_the_caller_aggregates() -> Result<()> {
951 let (orch, serve) = streaming_orch(vec![
954 LlmResponseChunk::MessageStart {
955 id: Some("m1".to_string()),
956 model: Some("stream/model".to_string()),
957 },
958 LlmResponseChunk::TextDelta {
959 index: 0,
960 text: "hel".to_string(),
961 },
962 LlmResponseChunk::TextDelta {
963 index: 0,
964 text: "lo".to_string(),
965 },
966 LlmResponseChunk::MessageStop {
967 reason: Some("stop".to_string()),
968 },
969 ]);
970 let (trace, response) = drive(orch, Context::default(), request(), serve).await?;
971 let agg = response
973 .llm_response
974 .into_agg()
975 .await
976 .map_err(|error| LibsyError::external("aggregating response stream", error))?;
977 assert_eq!(completion_text(&agg), "hello");
978 assert_eq!(agg.model.as_deref(), Some("stream/model"));
979 assert_eq!(trace.len(), 1);
980 Ok(())
981 }
982
983 #[tokio::test]
984 async fn aggregating_a_streamed_response_propagates_a_mid_stream_error() -> Result<()> {
985 let (orch, serve) = streaming_orch(vec![
988 LlmResponseChunk::TextDelta {
989 index: 0,
990 text: "partial".to_string(),
991 },
992 LlmResponseChunk::StreamError {
993 message: "upstream exploded".to_string(),
994 },
995 ]);
996 let (_, response) = drive(orch, Context::default(), request(), serve).await?;
997 match response.llm_response.into_agg().await {
998 Ok(_) => panic!("expected a mid-stream error, got an aggregate"),
999 Err(err) => {
1000 assert!(err.to_string().contains("upstream exploded"));
1001 Ok(())
1002 }
1003 }
1004 }
1005
1006 #[tokio::test]
1007 async fn run_offloads_via_promise_then_returns_to_agent() -> Result<()> {
1008 let stream =
1011 orch(target_set(&["offload/model"])).run_stream(Context::default(), request(), None);
1012 tokio::pin!(stream);
1013
1014 let mut saw_call = false;
1015 let mut final_completion = None;
1016 while let Some(step) = stream.next().await {
1017 match step? {
1018 Step::CallLlm(call) => {
1019 saw_call = true;
1020 assert_eq!(call.get_decision().selected_model(), "offload/model");
1022 call.respond(Ok(Response {
1024 llm_response: LlmResponse::Agg(text_response(
1025 None,
1026 "fulfilled".to_string(),
1027 )),
1028 metadata: None,
1029 }))?;
1030 }
1031 Step::Decision(decision) => {
1032 assert_eq!(decision.selected_model(), "offload/model");
1033 }
1034 Step::ReturnToAgent(response) => {
1035 final_completion = Some(
1036 response
1037 .llm_response
1038 .as_agg()
1039 .map(completion_text)
1040 .unwrap_or_default(),
1041 );
1042 }
1043 }
1044 }
1045
1046 assert!(saw_call, "expected a CallLlm step before ReturnToAgent");
1047 assert_eq!(
1048 final_completion.ok_or_else(|| test_error("no ReturnToAgent step"))?,
1049 "fulfilled"
1050 );
1051 Ok(())
1052 }
1053
1054 #[tokio::test]
1055 async fn a_driven_run_returns_the_trace_and_the_final_response() -> Result<()> {
1056 let (trace, response) = drive(
1057 orch(target_set(&["direct/model"])),
1058 Context::default(),
1059 request(),
1060 echo(),
1061 )
1062 .await?;
1063 assert_eq!(
1065 response
1066 .llm_response
1067 .as_agg()
1068 .map(completion_text)
1069 .unwrap_or_default(),
1070 "direct/model"
1071 );
1072 assert_eq!(trace[0].selected_model(), "direct/model");
1073 Ok(())
1074 }
1075
1076 #[tokio::test(flavor = "multi_thread", worker_threads = 12)]
1077 async fn requests_are_processed_in_parallel() -> Result<()> {
1078 use std::time::Duration;
1079 use tokio::sync::Barrier;
1080
1081 const N: usize = 12;
1082
1083 let barrier = Arc::new(Barrier::new(N));
1088 let algo = orch(target_set(&["m"]));
1090
1091 let mut handles = Vec::new();
1092 for _ in 0..N {
1093 let algo = algo.clone();
1094 let barrier = barrier.clone();
1095 let serve = move |decision: Arc<dyn Decision>, _request: Request| {
1096 let barrier = barrier.clone();
1097 async move {
1098 barrier.wait().await;
1099 Ok(reply(decision.selected_model()))
1100 }
1101 };
1102 handles.push(tokio::spawn(async move {
1103 drive(algo, Context::default(), request(), serve)
1104 .await
1105 .map(|(_, response)| {
1106 response
1107 .llm_response
1108 .as_agg()
1109 .map(completion_text)
1110 .unwrap_or_default()
1111 })
1112 }));
1113 }
1114
1115 for handle in handles {
1116 let completion = tokio::time::timeout(Duration::from_secs(5), handle)
1118 .await
1119 .map_err(|error| LibsyError::external("waiting for test task", error))?
1120 .map_err(|source| LibsyError::AlgorithmTask { source })??;
1121 assert_eq!(completion, "m");
1122 }
1123 Ok(())
1124 }
1125
1126 #[tokio::test]
1127 async fn offload_error_propagates_back_to_the_algorithm() -> Result<()> {
1128 let stream =
1132 orch(target_set(&["offload/model"])).run_stream(Context::default(), request(), None);
1133 tokio::pin!(stream);
1134
1135 let mut saw_error = false;
1136 while let Some(step) = stream.next().await {
1137 match step {
1138 Ok(Step::CallLlm(call)) => {
1139 call.respond(Err(test_error("upstream model call failed")))?;
1140 }
1141 Ok(Step::Decision(_)) => {}
1142 Ok(Step::ReturnToAgent(..)) => {
1143 return Err(test_error(
1144 "expected the offload error to propagate, got a response",
1145 ));
1146 }
1147 Err(err) => {
1148 assert!(err.to_string().contains("upstream model call failed"));
1150 saw_error = true;
1151 }
1152 }
1153 }
1154
1155 assert!(saw_error, "expected an error step");
1156 Ok(())
1157 }
1158
1159 #[tokio::test]
1160 async fn dropping_the_stream_cancels_the_algorithm_task() -> Result<()> {
1161 use std::sync::atomic::{AtomicBool, Ordering};
1162 use std::time::Duration;
1163 use tokio::sync::mpsc;
1164
1165 struct DropGuard(Arc<AtomicBool>);
1168 impl Drop for DropGuard {
1169 fn drop(&mut self) {
1170 self.0.store(true, Ordering::SeqCst);
1171 }
1172 }
1173
1174 struct StuckAlgo {
1175 started: mpsc::UnboundedSender<()>,
1176 dropped: Arc<AtomicBool>,
1177 }
1178
1179 #[async_trait]
1180 impl Algorithm for StuckAlgo {
1181 fn name(&self) -> &str {
1182 "stuck"
1183 }
1184
1185 async fn create_run_task(
1186 self: Arc<Self>,
1187 _ctx: Context,
1188 _driver: Driver,
1189 _request: Request,
1190 ) -> Result<Response> {
1191 let _guard = DropGuard(self.dropped.clone());
1192 let _ = self.started.send(());
1193 std::future::pending::<()>().await;
1195 unreachable!()
1196 }
1197 }
1198
1199 let (started_tx, mut started_rx) = mpsc::unbounded_channel();
1200 let dropped = Arc::new(AtomicBool::new(false));
1201 let algo: Arc<dyn Algorithm> = Arc::new(StuckAlgo {
1202 started: started_tx,
1203 dropped: dropped.clone(),
1204 });
1205
1206 let stream = algo.run_stream(Context::default(), request(), None);
1207 started_rx
1208 .recv()
1209 .await
1210 .ok_or_else(|| test_error("task never started"))?;
1211 drop(stream);
1212 tokio::time::sleep(Duration::from_millis(100)).await;
1213
1214 assert!(
1215 dropped.load(Ordering::SeqCst),
1216 "algorithm task was NOT cancelled after dropping the stream"
1217 );
1218 Ok(())
1219 }
1220
1221 #[tokio::test]
1222 async fn create_run_task_panic_surfaces_as_a_stream_error() -> Result<()> {
1223 struct Panicky;
1226
1227 #[async_trait]
1228 impl Algorithm for Panicky {
1229 fn name(&self) -> &str {
1230 "panicky"
1231 }
1232
1233 async fn create_run_task(
1234 self: Arc<Self>,
1235 _ctx: Context,
1236 _driver: Driver,
1237 _request: Request,
1238 ) -> Result<Response> {
1239 panic!("boom");
1240 }
1241 }
1242
1243 let algo: Arc<dyn Algorithm> = Arc::new(Panicky);
1244 let stream = algo.run_stream(Context::default(), request(), None);
1245 tokio::pin!(stream);
1246
1247 let mut saw_error = false;
1248 while let Some(step) = stream.next().await {
1249 match step {
1250 Err(err) => {
1251 assert!(matches!(err, LibsyError::AlgorithmTask { .. }));
1252 saw_error = true;
1253 }
1254 Ok(_) => return Err(test_error("expected the panic to surface as an error step")),
1255 }
1256 }
1257
1258 assert!(saw_error, "expected an error step from the panicked task");
1259 Ok(())
1260 }
1261
1262 #[tokio::test]
1263 async fn run_returns_an_error_when_the_algorithm_task_panics() -> Result<()> {
1264 struct Panicky;
1267
1268 #[async_trait]
1269 impl Algorithm for Panicky {
1270 fn name(&self) -> &str {
1271 "panicky"
1272 }
1273
1274 async fn create_run_task(
1275 self: Arc<Self>,
1276 _ctx: Context,
1277 _driver: Driver,
1278 _request: Request,
1279 ) -> Result<Response> {
1280 panic!("boom");
1281 }
1282 }
1283
1284 let algo: Arc<dyn Algorithm> = Arc::new(Panicky);
1285 match drive(algo, Context::default(), request(), echo()).await {
1286 Ok(_) => Err(test_error(
1287 "expected the run to surface the algorithm panic as an error",
1288 )),
1289 Err(err) => {
1290 assert!(matches!(err, LibsyError::AlgorithmTask { .. }));
1291 Ok(())
1292 }
1293 }
1294 }
1295
1296 #[tokio::test]
1297 async fn cancelling_run_cancels_the_algorithm_task() -> Result<()> {
1298 use std::sync::atomic::{AtomicBool, Ordering};
1299 use std::time::Duration;
1300 use tokio::sync::mpsc;
1301
1302 struct DropGuard(Arc<AtomicBool>);
1305 impl Drop for DropGuard {
1306 fn drop(&mut self) {
1307 self.0.store(true, Ordering::SeqCst);
1308 }
1309 }
1310
1311 struct StuckAlgo {
1312 started: mpsc::UnboundedSender<()>,
1313 dropped: Arc<AtomicBool>,
1314 }
1315
1316 #[async_trait]
1317 impl Algorithm for StuckAlgo {
1318 fn name(&self) -> &str {
1319 "stuck"
1320 }
1321
1322 async fn create_run_task(
1323 self: Arc<Self>,
1324 _ctx: Context,
1325 _driver: Driver,
1326 _request: Request,
1327 ) -> Result<Response> {
1328 let _guard = DropGuard(self.dropped.clone());
1329 let _ = self.started.send(());
1330 std::future::pending::<()>().await;
1333 unreachable!()
1334 }
1335 }
1336
1337 let (started_tx, mut started_rx) = mpsc::unbounded_channel();
1338 let dropped = Arc::new(AtomicBool::new(false));
1339 let algo: Arc<dyn Algorithm> = Arc::new(StuckAlgo {
1340 started: started_tx,
1341 dropped: dropped.clone(),
1342 });
1343
1344 let run_task =
1347 tokio::spawn(async move { drive(algo, Context::default(), request(), echo()).await });
1348 started_rx
1349 .recv()
1350 .await
1351 .ok_or_else(|| test_error("task never started"))?;
1352 run_task.abort();
1353 tokio::time::sleep(Duration::from_millis(100)).await;
1354
1355 assert!(
1356 dropped.load(Ordering::SeqCst),
1357 "algorithm task was NOT cancelled after cancelling run"
1358 );
1359 Ok(())
1360 }
1361
1362 struct Hedge {
1367 winner: LlmTarget,
1368 loser: LlmTarget,
1369 }
1370
1371 #[async_trait]
1372 impl Algorithm for Hedge {
1373 fn name(&self) -> &str {
1374 "hedge"
1375 }
1376
1377 async fn create_run_task(
1378 self: Arc<Self>,
1379 ctx: Context,
1380 driver: Driver,
1381 request: Request,
1382 ) -> Result<Response> {
1383 let dec_w: Arc<dyn Decision> = Arc::new(TestDecision {
1384 model: self.winner.semantic_name.clone(),
1385 });
1386 let dec_l: Arc<dyn Decision> = Arc::new(TestDecision {
1387 model: self.loser.semantic_name.clone(),
1388 });
1389 let win = driver.call_llm(RoutedRequest {
1390 request: request.clone(),
1391 decision: dec_w,
1392 ctx: ctx.clone(),
1393 });
1394 let lose = driver.call_llm(RoutedRequest {
1395 request,
1396 decision: dec_l,
1397 ctx,
1398 });
1399 tokio::select! {
1401 res = win => res,
1402 res = lose => res,
1403 }
1404 }
1405 }
1406
1407 fn hedge(loser_delay: Option<std::time::Duration>) -> (Arc<dyn Algorithm>, impl Serve) {
1411 let started = Arc::new(tokio::sync::Notify::new());
1412 let algo = Arc::new(Hedge {
1413 winner: LlmTarget {
1414 semantic_name: "winner".to_string(),
1415 },
1416 loser: LlmTarget {
1417 semantic_name: "loser".to_string(),
1418 },
1419 });
1420 let serve = move |decision: Arc<dyn Decision>, _request: Request| {
1421 let started = started.clone();
1422 async move {
1423 if decision.selected_model() == "loser" {
1424 started.notify_one();
1425 match loser_delay {
1426 Some(delay) => tokio::time::sleep(delay).await,
1427 None => std::future::pending::<()>().await,
1428 }
1429 } else {
1430 started.notified().await;
1431 }
1432 Ok(reply(decision.selected_model()))
1433 }
1434 };
1435 (algo, serve)
1436 }
1437
1438 #[tokio::test]
1439 async fn run_returns_the_winner_without_a_late_loser_overwriting_it() -> Result<()> {
1440 let (algo, serve) = hedge(Some(std::time::Duration::from_millis(50)));
1443 let (_trace, response) = drive(algo, Context::default(), request(), serve).await?;
1444 assert_eq!(
1445 response
1446 .llm_response
1447 .as_agg()
1448 .map(completion_text)
1449 .unwrap_or_default(),
1450 "winner"
1451 );
1452 Ok(())
1453 }
1454
1455 #[tokio::test]
1456 async fn run_returns_the_winner_without_hanging_on_a_pending_loser() -> Result<()> {
1457 let (algo, serve) = hedge(None);
1460 let run = drive(algo, Context::default(), request(), serve);
1461 let (_trace, response) = tokio::time::timeout(std::time::Duration::from_secs(1), run)
1462 .await
1463 .map_err(|error| LibsyError::external("waiting for pending loser", error))??;
1464 assert_eq!(
1465 response
1466 .llm_response
1467 .as_agg()
1468 .map(completion_text)
1469 .unwrap_or_default(),
1470 "winner"
1471 );
1472 Ok(())
1473 }
1474
1475 #[tokio::test]
1476 async fn run_surfaces_a_terminal_error_with_many_calls_in_flight() -> Result<()> {
1477 use std::sync::atomic::{AtomicUsize, Ordering};
1478
1479 const N: usize = 10;
1482
1483 struct FanOutThenError {
1486 all_started: Arc<tokio::sync::Notify>,
1487 n: usize,
1488 }
1489
1490 #[async_trait]
1491 impl Algorithm for FanOutThenError {
1492 fn name(&self) -> &str {
1493 "fan_out_then_error"
1494 }
1495
1496 async fn create_run_task(
1497 self: Arc<Self>,
1498 ctx: Context,
1499 driver: Driver,
1500 request: Request,
1501 ) -> Result<Response> {
1502 let offloads = futures::future::join_all((0..self.n).map(|i| {
1503 let decision: Arc<dyn Decision> = Arc::new(TestDecision {
1504 model: format!("m{i}"),
1505 });
1506 driver.call_llm(RoutedRequest {
1507 request: request.clone(),
1508 decision,
1509 ctx: ctx.clone(),
1510 })
1511 }));
1512 tokio::select! {
1513 _ = offloads => Err(test_error("offloads unexpectedly completed")),
1514 _ = self.all_started.notified() => {
1515 Err(test_error("terminal error while calls pending"))
1516 }
1517 }
1518 }
1519 }
1520
1521 let all_started = Arc::new(tokio::sync::Notify::new());
1522 let algo: Arc<dyn Algorithm> = Arc::new(FanOutThenError {
1523 all_started: all_started.clone(),
1524 n: N,
1525 });
1526
1527 let started = Arc::new(AtomicUsize::new(0));
1529 let serve = move |_decision: Arc<dyn Decision>, _request: Request| {
1530 let started = started.clone();
1531 let all_started = all_started.clone();
1532 async move {
1533 if started.fetch_add(1, Ordering::SeqCst) + 1 == N {
1534 all_started.notify_one();
1535 }
1536 std::future::pending::<ServeResult>().await
1537 }
1538 };
1539
1540 let run = drive(algo, Context::default(), request(), serve);
1543 let result = tokio::time::timeout(std::time::Duration::from_millis(500), run)
1544 .await
1545 .map_err(|error| {
1546 LibsyError::external("waiting for terminal error with full call cap", error)
1547 })?;
1548 match result {
1549 Ok(_) => Err(test_error("expected the terminal error, got a response")),
1550 Err(err) => {
1551 assert!(
1552 err.to_string()
1553 .contains("terminal error while calls pending")
1554 );
1555 Ok(())
1556 }
1557 }
1558 }
1559}