1use std::{
9 collections::{HashMap, HashSet},
10 future::Future,
11 pin::Pin,
12 sync::Arc,
13 time::Instant,
14};
15
16use async_trait::async_trait;
17use futures::{Stream, StreamExt};
18use parking_lot::Mutex;
19use tokio::sync::{mpsc, oneshot};
20use tokio_stream::wrappers::ReceiverStream;
21use tracing::Instrument;
22
23use switchyard_protocol::{
31 Context, Decision, LlmClientError, Request, Response, RoutingFallbackReason, Signals,
32};
33
34use crate::{DriverError, LibsyError, Result, observability};
35
36pub type StepStream = Pin<Box<dyn Stream<Item = Result<Step>> + Send>>;
40
41#[derive(Clone)]
50pub struct RoutedRequest {
51 pub request: Request,
53 pub decision: Arc<Decision>,
55 pub ctx: Context,
58}
59
60pub struct CallLlmRequest {
68 routed: RoutedRequest,
69 reply: oneshot::Sender<Result<Response>>,
70}
71
72impl CallLlmRequest {
73 pub fn get_routed(&self) -> &RoutedRequest {
76 &self.routed
77 }
78
79 pub fn get_request(&self) -> &Request {
81 &self.get_routed().request
82 }
83
84 pub fn get_decision(&self) -> &Decision {
86 self.get_routed().decision.as_ref()
87 }
88
89 pub fn respond(self, result: Result<Response>) -> Result<()> {
93 self.reply
94 .send(result)
95 .map_err(|_| DriverError::ResponseDropped.into())
96 }
97}
98
99#[derive(Clone)]
106pub struct Driver {
107 step_tx: mpsc::Sender<Result<Step>>,
108}
109
110impl Driver {
111 pub(crate) fn new() -> (Self, mpsc::Receiver<Result<Step>>) {
114 let (step_tx, step_rx) = mpsc::channel(1);
119 (Self { step_tx }, step_rx)
120 }
121
122 #[tracing::instrument(
131 target = "libsy",
132 name = "libsy.llm_call",
133 skip_all,
134 fields(
135 algorithm = observability::algorithm_label(&routed.ctx),
136 selected_model = routed.decision.selected_model_id(),
137 openinference.span.kind = "CHAIN",
138 outcome = tracing::field::Empty,
139 error = tracing::field::Empty,
140 input_tokens = tracing::field::Empty,
141 output_tokens = tracing::field::Empty,
142 total_tokens = tracing::field::Empty,
143 reasoning_tokens = tracing::field::Empty,
144 )
145 )]
146 pub async fn call_llm(&self, routed: RoutedRequest) -> Result<Response> {
147 let algorithm = observability::algorithm_label(&routed.ctx).to_string();
148 let decision = Arc::clone(&routed.decision);
149 let is_answer_call = decision.is_answer_call();
150 let started = Instant::now();
151 let (reply, response) = oneshot::channel::<Result<Response>>();
152 let call = CallLlmRequest { routed, reply };
153 let result = async {
154 self.step_tx
155 .send(Ok(Step::CallLlm(Box::new(call))))
156 .await
157 .map_err(|_| DriverError::StreamClosed)?;
158 response
159 .await
160 .map_err(|_| LibsyError::from(DriverError::ResponseDropped))?
161 }
162 .await;
163 let elapsed = started.elapsed();
164 observability::record_llm_call(
165 &algorithm,
166 decision.selected_model_id(),
167 is_answer_call,
168 elapsed,
169 &result,
170 &tracing::Span::current(),
171 );
172 result
173 }
174
175 pub async fn info(&self, ctx: Context, decision: Arc<Decision>) -> Result<()> {
179 self.step_tx
180 .send(Ok(Step::Decision(decision.clone())))
181 .await
182 .map_err(|_| DriverError::StreamClosed)?;
183 observability::record_decision(&ctx, decision.as_ref());
184 Ok(())
185 }
186
187 pub(crate) async fn finish(&self, result: Result<Response>) -> Result<()> {
191 let step = result.map(|response| Step::ReturnToAgent(Box::new(response)));
192 self.step_tx
193 .send(step)
194 .await
195 .map_err(|_| DriverError::StreamClosed.into())
196 }
197}
198
199pub enum Step {
201 CallLlm(Box<CallLlmRequest>),
204 Decision(Arc<Decision>),
207 ReturnToAgent(Box<Response>),
209}
210
211pub async fn drive<F, Fut>(
224 algorithm: Arc<dyn Algorithm>,
225 ctx: Context,
226 request: Request,
227 serve: F,
228) -> Result<(Vec<Arc<Decision>>, Response)>
229where
230 F: Fn(CallLlmRequest) -> Fut,
231 Fut: Future<Output = Result<()>>,
232{
233 let stream = algorithm.run_stream(ctx, request);
234 tokio::pin!(stream);
235
236 let mut trace: Vec<Arc<Decision>> = Vec::new();
237 let mut in_flight = futures::stream::FuturesUnordered::new();
238 let mut final_response: Option<Response> = None;
239
240 loop {
241 tokio::select! {
242 Some(result) = in_flight.next() => match result {
243 Ok(()) => {}, Err(err) => return Err(err), },
246 step = stream.next() => {
247 match step {
248 None => break, Some(item) => match item? {
250 Step::CallLlm(call) => in_flight.push(serve(*call)),
251 Step::Decision(decision) => trace.push(decision),
252 Step::ReturnToAgent(response) => {
253 final_response = Some(*response);
254 break;
255 }
256 }
257 }
258 },
259 }
260 }
261 final_response
262 .map(|response| (trace, response))
263 .ok_or(LibsyError::MissingFinalResponse)
264}
265
266struct AbortOnDrop(tokio::task::AbortHandle);
268
269impl Drop for AbortOnDrop {
270 fn drop(&mut self) {
271 self.0.abort();
272 }
273}
274
275#[derive(Clone)]
279pub struct LlmTarget {
280 pub semantic_name: String,
284}
285
286#[derive(Clone)]
290pub struct LlmTargetSet {
291 targets: Vec<LlmTarget>,
292}
293
294impl LlmTargetSet {
295 pub fn new(targets: Vec<LlmTarget>) -> Self {
297 Self { targets }
298 }
299
300 pub fn targets(&self) -> &[LlmTarget] {
302 &self.targets
303 }
304
305 pub fn get_target(&self, name: &str) -> Result<LlmTarget> {
307 self.targets
308 .iter()
309 .find(|t| t.semantic_name == name)
310 .cloned()
311 .ok_or_else(|| LibsyError::TargetNotFound {
312 target: name.to_string(),
313 })
314 }
315
316 pub fn resolve_target(&self, name: &str, ctx: &Context) -> Result<LlmTarget> {
319 let target = self.get_target(name)?;
320 if !ctx.is_excluded(&target.semantic_name) {
321 return Ok(target);
322 }
323 self.targets
324 .iter()
325 .find(|t| !ctx.is_excluded(&t.semantic_name))
326 .cloned()
327 .ok_or(LibsyError::AllTargetsExcluded)
328 }
329}
330
331#[derive(Clone, Hash, PartialEq, Eq)]
335pub(crate) enum RoutingIdentity {
336 Session(String),
338 Subagent { session: String, agent: String },
340}
341
342impl RoutingIdentity {
343 pub(crate) fn from_request(request: &Request) -> Option<Self> {
348 let metadata = request.metadata.as_ref()?;
349 let session = metadata.session_id.as_deref().filter(|id| !id.is_empty())?;
350 if metadata.is_subagent {
351 let agent = metadata.agent_id.as_deref().filter(|id| !id.is_empty())?;
352 Some(Self::Subagent {
353 session: session.to_string(),
354 agent: agent.to_string(),
355 })
356 } else {
357 Some(Self::Session(session.to_string()))
358 }
359 }
360
361 fn session(&self) -> &str {
363 match self {
364 Self::Session(session) | Self::Subagent { session, .. } => session,
365 }
366 }
367}
368
369const MAX_EVICTION_IDENTITIES: usize = 1_024;
372
373#[derive(Default)]
379pub(crate) struct SessionEvictions {
380 by_identity: Mutex<HashMap<RoutingIdentity, HashSet<String>>>,
381}
382
383impl SessionEvictions {
384 pub(crate) fn remove_session(&self, session: &str) {
386 self.by_identity
387 .lock()
388 .retain(|identity, _| identity.session() != session);
389 }
390
391 fn evicted_for(&self, identity: Option<&RoutingIdentity>) -> Vec<String> {
393 let Some(identity) = identity else {
394 return Vec::new();
395 };
396 self.by_identity
397 .lock()
398 .get(identity)
399 .map(|targets| targets.iter().cloned().collect())
400 .unwrap_or_default()
401 }
402
403 fn record(&self, identity: Option<&RoutingIdentity>, target: &str) {
406 let Some(identity) = identity else { return };
407 let mut histories = self.by_identity.lock();
408 if histories.len() >= MAX_EVICTION_IDENTITIES
409 && !histories.contains_key(identity)
410 && let Some(oldest) = histories.keys().next().cloned()
411 {
412 histories.remove(&oldest);
413 }
414 histories
415 .entry(identity.clone())
416 .or_default()
417 .insert(target.to_string());
418 }
419}
420
421fn eligible_targets(targets: &LlmTargetSet, ctx: &Context) -> usize {
423 targets
424 .targets()
425 .iter()
426 .filter(|t| !ctx.is_excluded(&t.semantic_name))
427 .count()
428}
429
430pub(crate) fn exclude_evicted(
433 ctx: &mut Context,
434 targets: &LlmTargetSet,
435 evictions: &SessionEvictions,
436 identity: Option<&RoutingIdentity>,
437) {
438 for target in evictions.evicted_for(identity) {
439 if eligible_targets(targets, ctx) <= 1 {
442 break;
443 }
444 ctx.exclude_target(target);
445 }
446}
447
448fn classify_fallback(error: &LibsyError) -> Option<(&str, RoutingFallbackReason)> {
450 let LibsyError::ClientCall { target, source } = error else {
451 return None;
452 };
453 let reason = match source {
454 LlmClientError::ContextWindowExceeded { .. } => RoutingFallbackReason::ContextWindow,
455 LlmClientError::Transport { .. } | LlmClientError::Timeout { .. } => {
456 RoutingFallbackReason::Unavailable
457 }
458 LlmClientError::UpstreamHttp { status, .. }
459 if matches!(*status, 403 | 408 | 429) || (500..=599).contains(status) =>
460 {
461 RoutingFallbackReason::Unavailable
462 }
463 _ => return None,
464 };
465 Some((target, reason))
466}
467
468#[allow(clippy::too_many_arguments)]
476pub(crate) async fn call_llm_with_fallback(
477 mut ctx: Context,
478 driver: &Driver,
479 targets: &LlmTargetSet,
480 mut target: LlmTarget,
481 mut decision: Arc<Decision>,
482 request: Request,
483 identity: Option<&RoutingIdentity>,
484 evictions: &SessionEvictions,
485 target_unavailable: impl Fn(&Request, &str),
486 fallback_decision: impl Fn(&LlmTarget, &LlmTarget, RoutingFallbackReason) -> Arc<Decision>,
487) -> Result<Response> {
488 loop {
489 let result = driver
490 .call_llm(RoutedRequest {
491 request: request.clone(),
492 decision: decision.clone(),
493 ctx: ctx.clone(),
494 })
495 .await;
496 let Err(error) = result else { return result };
497 let Some((failed, reason)) = classify_fallback(&error) else {
498 return Err(error);
499 };
500 if !ctx.exclude_target(failed) {
503 return Err(error);
504 }
505 match reason {
506 RoutingFallbackReason::ContextWindow => evictions.record(identity, failed),
507 RoutingFallbackReason::Unavailable => target_unavailable(&request, failed),
508 }
509 let Ok(next) = targets.resolve_target(&target.semantic_name, &ctx) else {
510 return Err(error);
511 };
512 decision = fallback_decision(&target, &next, reason);
513 target = next;
514 driver.info(ctx.clone(), decision.clone()).await?;
515 }
516}
517
518#[async_trait]
540pub trait Algorithm: Send + Sync + 'static {
541 fn name(&self) -> &str;
545
546 async fn create_run_task(
552 self: Arc<Self>,
553 ctx: Context,
554 driver: Driver,
555 request: Request,
556 ) -> Result<Response>;
557
558 #[allow(unused_variables)]
562 async fn process_signals(self: Arc<Self>, signals: Signals) -> Result<()> {
563 Ok(())
564 }
565
566 fn run_stream(self: Arc<Self>, ctx: Context, request: Request) -> StepStream {
575 let mut ctx = ctx;
578 ctx.values.insert(
579 observability::ALGORITHM_KEY.to_string(),
580 self.name().to_string(),
581 );
582 let (driver, step_rx) = Driver::new();
583 let task_driver = driver.clone();
584 let task_ctx = ctx.clone();
585 let stream = ReceiverStream::new(step_rx);
586 let span = observability::run_span(self.name(), &request);
590 let handle = tokio::spawn(
591 async move {
592 observability::observe_run(
593 task_ctx.clone(),
594 self.create_run_task(task_ctx, task_driver, request),
595 )
596 .await
597 }
598 .instrument(span),
599 );
600 let abort_guard = AbortOnDrop(handle.abort_handle());
602
603 let finish_driver = driver.clone();
604 let tail: StepStream = Box::pin(
605 futures::stream::once(async move {
606 let result = match handle.await {
607 Ok(response) => response,
608 Err(source) => Err(LibsyError::AlgorithmTask { source }),
609 };
610 finish_driver.finish(result).await
611 })
612 .filter_map(|finish_result| async move { finish_result.err().map(Err) }),
613 );
614
615 let stream: StepStream = Box::pin(stream);
616 Box::pin(futures::stream::select(stream, tail).map(move |step| {
617 let _keep_alive = &abort_guard;
619 step
620 }))
621 }
622}
623
624#[cfg(test)]
625mod tests {
626 use super::*;
627 use crate::core::testing::{Serve, ServeResult, echo, reply, test_drive};
628 use futures::StreamExt;
629 use switchyard_protocol::{
630 LlmResponse, LlmResponseChunk, completion_text, text_request, text_response,
631 };
632
633 #[derive(Debug, thiserror::Error)]
634 #[error("{0}")]
635 struct TestError(&'static str);
636
637 fn test_error(message: &'static str) -> LibsyError {
638 LibsyError::external("test", TestError(message))
639 }
640
641 fn classified_client_error(source: LlmClientError) -> Option<RoutingFallbackReason> {
642 classify_fallback(&LibsyError::client_call("target", source)).map(|(_, reason)| reason)
643 }
644
645 #[test]
646 fn route_fallback_only_accepts_context_and_unavailable_failures() {
647 assert_eq!(
648 classified_client_error(LlmClientError::ContextWindowExceeded {
649 model: "target".to_string(),
650 message: "too long".to_string(),
651 }),
652 Some(RoutingFallbackReason::ContextWindow)
653 );
654 for source in [
655 LlmClientError::Transport {
656 source: Box::new(std::io::Error::other("connection failed")),
657 },
658 LlmClientError::Timeout {
659 source: Box::new(std::io::Error::other("request timed out")),
660 },
661 ] {
662 assert_eq!(
663 classified_client_error(source),
664 Some(RoutingFallbackReason::Unavailable)
665 );
666 }
667 for (status, expected) in [
668 (400, None),
669 (401, None),
670 (403, Some(RoutingFallbackReason::Unavailable)),
671 (404, None),
672 (408, Some(RoutingFallbackReason::Unavailable)),
673 (409, None),
674 (429, Some(RoutingFallbackReason::Unavailable)),
675 (499, None),
676 (500, Some(RoutingFallbackReason::Unavailable)),
677 (599, Some(RoutingFallbackReason::Unavailable)),
678 (600, None),
679 ] {
680 assert_eq!(
681 classified_client_error(LlmClientError::UpstreamHttp {
682 status,
683 body: "failed".to_string(),
684 }),
685 expected
686 );
687 }
688 assert_eq!(
689 classified_client_error(LlmClientError::InvalidResponse {
690 source: Box::new(std::io::Error::other("invalid response")),
691 }),
692 None
693 );
694 }
695
696 fn test_decision(selected_model_id: String) -> Arc<Decision> {
698 Arc::new(Decision::new(selected_model_id, None, true))
699 }
700
701 struct TestAlgo {
704 target_set: LlmTargetSet,
705 }
706
707 #[async_trait]
708 impl Algorithm for TestAlgo {
709 fn name(&self) -> &str {
710 "test"
711 }
712
713 async fn create_run_task(
714 self: Arc<Self>,
715 ctx: Context,
716 driver: Driver,
717 request: Request,
718 ) -> Result<Response> {
719 let target = self
720 .target_set
721 .targets()
722 .first()
723 .ok_or(LibsyError::NoTargets)?
724 .clone();
725 let decision = test_decision(target.semantic_name.clone());
726 driver.info(ctx.clone(), decision.clone()).await?;
727 driver
728 .call_llm(RoutedRequest {
729 request,
730 decision,
731 ctx,
732 })
733 .await
734 }
735 }
736
737 fn orch(target_set: LlmTargetSet) -> Arc<dyn Algorithm> {
739 Arc::new(TestAlgo { target_set })
740 }
741
742 fn request() -> Request {
743 Request {
744 llm_request: text_request(Some("auto".to_string()), "hi".to_string()),
745 raw_request: None,
746 metadata: None,
747 }
748 }
749
750 fn target_set(names: &[&str]) -> LlmTargetSet {
751 let targets = names
752 .iter()
753 .map(|name| LlmTarget {
754 semantic_name: name.to_string(),
755 })
756 .collect();
757 LlmTargetSet::new(targets)
758 }
759
760 fn routed(model: &str) -> RoutedRequest {
761 RoutedRequest {
762 request: request(),
763 decision: test_decision(model.to_string()),
764 ctx: Context::default(),
765 }
766 }
767
768 #[tokio::test]
769 async fn typed_driver_preserves_call_and_stream_boundaries() -> Result<()> {
770 tokio::time::timeout(std::time::Duration::from_secs(1), async {
771 let (driver, mut step_rx) = Driver::new();
774 let first_driver = driver.clone();
775 let mut first =
776 tokio::spawn(async move { first_driver.call_llm(routed("first")).await });
777 let second = tokio::spawn(async move { driver.call_llm(routed("second")).await });
778
779 let mut calls = HashMap::new();
780 for _ in 0..2 {
781 let step = step_rx.recv().await.ok_or(DriverError::StreamClosed)??;
782 let Step::CallLlm(call) = step else {
783 return Err(test_error("expected a CallLlm step"));
784 };
785 calls.insert(call.get_decision().selected_model_id().to_string(), call);
786 }
787 assert!(
788 tokio::time::timeout(std::time::Duration::from_millis(20), &mut first)
789 .await
790 .is_err(),
791 "call completed before the host responded"
792 );
793 calls
794 .remove("second")
795 .ok_or_else(|| test_error("missing second call"))?
796 .respond(Ok(reply("second response")))?;
797 calls
798 .remove("first")
799 .ok_or_else(|| test_error("missing first call"))?
800 .respond(Ok(reply("first response")))?;
801
802 let first_response = first
803 .await
804 .map_err(|source| LibsyError::AlgorithmTask { source })??;
805 let second_response = second
806 .await
807 .map_err(|source| LibsyError::AlgorithmTask { source })??;
808 assert_eq!(
809 first_response.llm_response.as_agg().map(completion_text),
810 Some("first response".to_string())
811 );
812 assert_eq!(
813 second_response.llm_response.as_agg().map(completion_text),
814 Some("second response".to_string())
815 );
816
817 let (driver, mut step_rx) = Driver::new();
819 let producer = tokio::spawn(async move { driver.call_llm(routed("dropped")).await });
820 let step = step_rx.recv().await.ok_or(DriverError::StreamClosed)??;
821 let Step::CallLlm(call) = step else {
822 return Err(test_error("expected a CallLlm step"));
823 };
824 drop(call);
825 let result = producer
826 .await
827 .map_err(|source| LibsyError::AlgorithmTask { source })?;
828 assert!(matches!(
829 result,
830 Err(LibsyError::Driver(DriverError::ResponseDropped))
831 ));
832
833 let (driver, step_rx) = Driver::new();
835 drop(step_rx);
836 let decision = test_decision("closed".to_string());
837 let result = driver.info(Context::default(), decision).await;
838 assert!(matches!(
839 result,
840 Err(LibsyError::Driver(DriverError::StreamClosed))
841 ));
842 Ok(())
843 })
844 .await
845 .map_err(|error| LibsyError::external("waiting for typed driver boundaries", error))?
846 }
847
848 #[test]
849 fn target_lookup_returns_the_missing_target() {
850 let error = target_set(&[]).get_target("missing").err();
851 assert!(matches!(
852 error,
853 Some(LibsyError::TargetNotFound { target }) if target == "missing"
854 ));
855 }
856
857 fn streaming_orch(chunks: Vec<LlmResponseChunk>) -> (Arc<dyn Algorithm>, impl Serve) {
860 let algo = orch(target_set(&["stream/model"]));
861 let serve = move |_decision: Arc<Decision>, _request: Request| {
862 let chunks = chunks.clone();
863 async move {
864 let stream =
865 futures::stream::iter(chunks.into_iter().map(|chunk| Ok(chunk.into()))).boxed();
866 Ok(Response {
867 llm_response: LlmResponse::Stream(stream),
868 metadata: None,
869 })
870 }
871 };
872 (algo, serve)
873 }
874
875 #[tokio::test]
876 async fn run_returns_a_streamed_response_the_caller_aggregates() -> Result<()> {
877 let (orch, serve) = streaming_orch(vec![
880 LlmResponseChunk::MessageStart {
881 id: Some("m1".to_string()),
882 model: Some("stream/model".to_string()),
883 },
884 LlmResponseChunk::TextDelta {
885 index: 0,
886 text: "hel".to_string(),
887 },
888 LlmResponseChunk::TextDelta {
889 index: 0,
890 text: "lo".to_string(),
891 },
892 LlmResponseChunk::MessageStop {
893 reason: Some("stop".to_string()),
894 },
895 ]);
896 let (trace, response) = test_drive(orch, Context::default(), request(), serve).await?;
897 let agg = response
899 .llm_response
900 .into_agg()
901 .await
902 .map_err(|error| LibsyError::external("aggregating response stream", error))?;
903 assert_eq!(completion_text(&agg), "hello");
904 assert_eq!(agg.model.as_deref(), Some("stream/model"));
905 assert_eq!(trace.len(), 1);
906 Ok(())
907 }
908
909 #[tokio::test]
910 async fn aggregating_a_streamed_response_propagates_a_mid_stream_error() -> Result<()> {
911 let (orch, serve) = streaming_orch(vec![
914 LlmResponseChunk::TextDelta {
915 index: 0,
916 text: "partial".to_string(),
917 },
918 LlmResponseChunk::StreamError {
919 message: "upstream exploded".to_string(),
920 },
921 ]);
922 let (_, response) = test_drive(orch, Context::default(), request(), serve).await?;
923 match response.llm_response.into_agg().await {
924 Ok(_) => panic!("expected a mid-stream error, got an aggregate"),
925 Err(err) => {
926 assert!(err.to_string().contains("upstream exploded"));
927 Ok(())
928 }
929 }
930 }
931
932 #[tokio::test]
933 async fn run_offloads_via_promise_then_returns_to_agent() -> Result<()> {
934 let stream = orch(target_set(&["offload/model"])).run_stream(Context::default(), request());
937 tokio::pin!(stream);
938
939 let mut saw_call = false;
940 let mut final_completion = None;
941 while let Some(step) = stream.next().await {
942 match step? {
943 Step::CallLlm(call) => {
944 saw_call = true;
945 assert_eq!(call.get_decision().selected_model_id(), "offload/model");
947 call.respond(Ok(Response {
949 llm_response: LlmResponse::Agg(text_response(
950 None,
951 "fulfilled".to_string(),
952 )),
953 metadata: None,
954 }))?;
955 }
956 Step::Decision(decision) => {
957 assert_eq!(decision.selected_model_id(), "offload/model");
958 }
959 Step::ReturnToAgent(response) => {
960 final_completion = Some(
961 response
962 .llm_response
963 .as_agg()
964 .map(completion_text)
965 .unwrap_or_default(),
966 );
967 }
968 }
969 }
970
971 assert!(saw_call, "expected a CallLlm step before ReturnToAgent");
972 assert_eq!(
973 final_completion.ok_or_else(|| test_error("no ReturnToAgent step"))?,
974 "fulfilled"
975 );
976 Ok(())
977 }
978
979 #[tokio::test]
980 async fn a_driven_run_returns_the_trace_and_the_final_response() -> Result<()> {
981 let (trace, response) = test_drive(
982 orch(target_set(&["direct/model"])),
983 Context::default(),
984 request(),
985 echo(),
986 )
987 .await?;
988 assert_eq!(
990 response
991 .llm_response
992 .as_agg()
993 .map(completion_text)
994 .unwrap_or_default(),
995 "direct/model"
996 );
997 assert_eq!(trace[0].selected_model_id(), "direct/model");
998 Ok(())
999 }
1000
1001 #[tokio::test(flavor = "multi_thread", worker_threads = 12)]
1002 async fn requests_are_processed_in_parallel() -> Result<()> {
1003 use std::time::Duration;
1004 use tokio::sync::Barrier;
1005
1006 const N: usize = 12;
1007
1008 let barrier = Arc::new(Barrier::new(N));
1013 let algo = orch(target_set(&["m"]));
1015
1016 let mut handles = Vec::new();
1017 for _ in 0..N {
1018 let algo = algo.clone();
1019 let barrier = barrier.clone();
1020 let serve = move |decision: Arc<Decision>, _request: Request| {
1021 let barrier = barrier.clone();
1022 async move {
1023 barrier.wait().await;
1024 Ok(reply(decision.selected_model_id()))
1025 }
1026 };
1027 handles.push(tokio::spawn(async move {
1028 test_drive(algo, Context::default(), request(), serve)
1029 .await
1030 .map(|(_, response)| {
1031 response
1032 .llm_response
1033 .as_agg()
1034 .map(completion_text)
1035 .unwrap_or_default()
1036 })
1037 }));
1038 }
1039
1040 for handle in handles {
1041 let completion = tokio::time::timeout(Duration::from_secs(5), handle)
1043 .await
1044 .map_err(|error| LibsyError::external("waiting for test task", error))?
1045 .map_err(|source| LibsyError::AlgorithmTask { source })??;
1046 assert_eq!(completion, "m");
1047 }
1048 Ok(())
1049 }
1050
1051 #[tokio::test]
1052 async fn offload_error_propagates_back_to_the_algorithm() -> Result<()> {
1053 let stream = orch(target_set(&["offload/model"])).run_stream(Context::default(), request());
1057 tokio::pin!(stream);
1058
1059 let mut saw_error = false;
1060 while let Some(step) = stream.next().await {
1061 match step {
1062 Ok(Step::CallLlm(call)) => {
1063 call.respond(Err(test_error("upstream model call failed")))?;
1064 }
1065 Ok(Step::Decision(_)) => {}
1066 Ok(Step::ReturnToAgent(..)) => {
1067 return Err(test_error(
1068 "expected the offload error to propagate, got a response",
1069 ));
1070 }
1071 Err(err) => {
1072 assert!(err.to_string().contains("upstream model call failed"));
1074 saw_error = true;
1075 }
1076 }
1077 }
1078
1079 assert!(saw_error, "expected an error step");
1080 Ok(())
1081 }
1082
1083 #[tokio::test]
1084 async fn dropping_the_stream_cancels_the_algorithm_task() -> Result<()> {
1085 use std::sync::atomic::{AtomicBool, Ordering};
1086 use std::time::Duration;
1087 use tokio::sync::mpsc;
1088
1089 struct DropGuard(Arc<AtomicBool>);
1092 impl Drop for DropGuard {
1093 fn drop(&mut self) {
1094 self.0.store(true, Ordering::SeqCst);
1095 }
1096 }
1097
1098 struct StuckAlgo {
1099 started: mpsc::UnboundedSender<()>,
1100 dropped: Arc<AtomicBool>,
1101 }
1102
1103 #[async_trait]
1104 impl Algorithm for StuckAlgo {
1105 fn name(&self) -> &str {
1106 "stuck"
1107 }
1108
1109 async fn create_run_task(
1110 self: Arc<Self>,
1111 _ctx: Context,
1112 _driver: Driver,
1113 _request: Request,
1114 ) -> Result<Response> {
1115 let _guard = DropGuard(self.dropped.clone());
1116 let _ = self.started.send(());
1117 std::future::pending::<()>().await;
1119 unreachable!()
1120 }
1121 }
1122
1123 let (started_tx, mut started_rx) = mpsc::unbounded_channel();
1124 let dropped = Arc::new(AtomicBool::new(false));
1125 let algo: Arc<dyn Algorithm> = Arc::new(StuckAlgo {
1126 started: started_tx,
1127 dropped: dropped.clone(),
1128 });
1129
1130 let stream = algo.run_stream(Context::default(), request());
1131 started_rx
1132 .recv()
1133 .await
1134 .ok_or_else(|| test_error("task never started"))?;
1135 drop(stream);
1136 tokio::time::sleep(Duration::from_millis(100)).await;
1137
1138 assert!(
1139 dropped.load(Ordering::SeqCst),
1140 "algorithm task was NOT cancelled after dropping the stream"
1141 );
1142 Ok(())
1143 }
1144
1145 #[tokio::test]
1146 async fn create_run_task_panic_surfaces_as_a_stream_error() -> Result<()> {
1147 struct Panicky;
1150
1151 #[async_trait]
1152 impl Algorithm for Panicky {
1153 fn name(&self) -> &str {
1154 "panicky"
1155 }
1156
1157 async fn create_run_task(
1158 self: Arc<Self>,
1159 _ctx: Context,
1160 _driver: Driver,
1161 _request: Request,
1162 ) -> Result<Response> {
1163 panic!("boom");
1164 }
1165 }
1166
1167 let algo: Arc<dyn Algorithm> = Arc::new(Panicky);
1168 let stream = algo.run_stream(Context::default(), request());
1169 tokio::pin!(stream);
1170
1171 let mut saw_error = false;
1172 while let Some(step) = stream.next().await {
1173 match step {
1174 Err(err) => {
1175 assert!(matches!(err, LibsyError::AlgorithmTask { .. }));
1176 saw_error = true;
1177 }
1178 Ok(_) => return Err(test_error("expected the panic to surface as an error step")),
1179 }
1180 }
1181
1182 assert!(saw_error, "expected an error step from the panicked task");
1183 Ok(())
1184 }
1185
1186 #[tokio::test]
1187 async fn run_returns_an_error_when_the_algorithm_task_panics() -> Result<()> {
1188 struct Panicky;
1191
1192 #[async_trait]
1193 impl Algorithm for Panicky {
1194 fn name(&self) -> &str {
1195 "panicky"
1196 }
1197
1198 async fn create_run_task(
1199 self: Arc<Self>,
1200 _ctx: Context,
1201 _driver: Driver,
1202 _request: Request,
1203 ) -> Result<Response> {
1204 panic!("boom");
1205 }
1206 }
1207
1208 let algo: Arc<dyn Algorithm> = Arc::new(Panicky);
1209 match test_drive(algo, Context::default(), request(), echo()).await {
1210 Ok(_) => Err(test_error(
1211 "expected the run to surface the algorithm panic as an error",
1212 )),
1213 Err(err) => {
1214 assert!(matches!(err, LibsyError::AlgorithmTask { .. }));
1215 Ok(())
1216 }
1217 }
1218 }
1219
1220 #[tokio::test]
1221 async fn cancelling_run_cancels_the_algorithm_task() -> Result<()> {
1222 use std::sync::atomic::{AtomicBool, Ordering};
1223 use std::time::Duration;
1224 use tokio::sync::mpsc;
1225
1226 struct DropGuard(Arc<AtomicBool>);
1229 impl Drop for DropGuard {
1230 fn drop(&mut self) {
1231 self.0.store(true, Ordering::SeqCst);
1232 }
1233 }
1234
1235 struct StuckAlgo {
1236 started: mpsc::UnboundedSender<()>,
1237 dropped: Arc<AtomicBool>,
1238 }
1239
1240 #[async_trait]
1241 impl Algorithm for StuckAlgo {
1242 fn name(&self) -> &str {
1243 "stuck"
1244 }
1245
1246 async fn create_run_task(
1247 self: Arc<Self>,
1248 _ctx: Context,
1249 _driver: Driver,
1250 _request: Request,
1251 ) -> Result<Response> {
1252 let _guard = DropGuard(self.dropped.clone());
1253 let _ = self.started.send(());
1254 std::future::pending::<()>().await;
1257 unreachable!()
1258 }
1259 }
1260
1261 let (started_tx, mut started_rx) = mpsc::unbounded_channel();
1262 let dropped = Arc::new(AtomicBool::new(false));
1263 let algo: Arc<dyn Algorithm> = Arc::new(StuckAlgo {
1264 started: started_tx,
1265 dropped: dropped.clone(),
1266 });
1267
1268 let run_task =
1271 tokio::spawn(
1272 async move { test_drive(algo, Context::default(), request(), echo()).await },
1273 );
1274 started_rx
1275 .recv()
1276 .await
1277 .ok_or_else(|| test_error("task never started"))?;
1278 run_task.abort();
1279 tokio::time::sleep(Duration::from_millis(100)).await;
1280
1281 assert!(
1282 dropped.load(Ordering::SeqCst),
1283 "algorithm task was NOT cancelled after cancelling run"
1284 );
1285 Ok(())
1286 }
1287
1288 struct Hedge {
1293 winner: LlmTarget,
1294 loser: LlmTarget,
1295 }
1296
1297 #[async_trait]
1298 impl Algorithm for Hedge {
1299 fn name(&self) -> &str {
1300 "hedge"
1301 }
1302
1303 async fn create_run_task(
1304 self: Arc<Self>,
1305 ctx: Context,
1306 driver: Driver,
1307 request: Request,
1308 ) -> Result<Response> {
1309 let dec_w = test_decision(self.winner.semantic_name.clone());
1310 let dec_l = test_decision(self.loser.semantic_name.clone());
1311 let win = driver.call_llm(RoutedRequest {
1312 request: request.clone(),
1313 decision: dec_w,
1314 ctx: ctx.clone(),
1315 });
1316 let lose = driver.call_llm(RoutedRequest {
1317 request,
1318 decision: dec_l,
1319 ctx,
1320 });
1321 tokio::select! {
1323 res = win => res,
1324 res = lose => res,
1325 }
1326 }
1327 }
1328
1329 fn hedge(loser_delay: Option<std::time::Duration>) -> (Arc<dyn Algorithm>, impl Serve) {
1333 let started = Arc::new(tokio::sync::Notify::new());
1334 let algo = Arc::new(Hedge {
1335 winner: LlmTarget {
1336 semantic_name: "winner".to_string(),
1337 },
1338 loser: LlmTarget {
1339 semantic_name: "loser".to_string(),
1340 },
1341 });
1342 let serve = move |decision: Arc<Decision>, _request: Request| {
1343 let started = started.clone();
1344 async move {
1345 if decision.selected_model_id() == "loser" {
1346 started.notify_one();
1347 match loser_delay {
1348 Some(delay) => tokio::time::sleep(delay).await,
1349 None => std::future::pending::<()>().await,
1350 }
1351 } else {
1352 started.notified().await;
1353 }
1354 Ok(reply(decision.selected_model_id()))
1355 }
1356 };
1357 (algo, serve)
1358 }
1359
1360 #[tokio::test]
1361 async fn run_returns_the_winner_without_a_late_loser_overwriting_it() -> Result<()> {
1362 let (algo, serve) = hedge(Some(std::time::Duration::from_millis(50)));
1365 let (_trace, response) = test_drive(algo, Context::default(), request(), serve).await?;
1366 assert_eq!(
1367 response
1368 .llm_response
1369 .as_agg()
1370 .map(completion_text)
1371 .unwrap_or_default(),
1372 "winner"
1373 );
1374 Ok(())
1375 }
1376
1377 #[tokio::test]
1378 async fn run_returns_the_winner_without_hanging_on_a_pending_loser() -> Result<()> {
1379 let (algo, serve) = hedge(None);
1382 let run = test_drive(algo, Context::default(), request(), serve);
1383 let (_trace, response) = tokio::time::timeout(std::time::Duration::from_secs(1), run)
1384 .await
1385 .map_err(|error| LibsyError::external("waiting for pending loser", error))??;
1386 assert_eq!(
1387 response
1388 .llm_response
1389 .as_agg()
1390 .map(completion_text)
1391 .unwrap_or_default(),
1392 "winner"
1393 );
1394 Ok(())
1395 }
1396
1397 #[tokio::test]
1398 async fn run_surfaces_a_terminal_error_with_many_calls_in_flight() -> Result<()> {
1399 use std::sync::atomic::{AtomicUsize, Ordering};
1400
1401 const N: usize = 10;
1404
1405 struct FanOutThenError {
1408 all_started: Arc<tokio::sync::Notify>,
1409 n: usize,
1410 }
1411
1412 #[async_trait]
1413 impl Algorithm for FanOutThenError {
1414 fn name(&self) -> &str {
1415 "fan_out_then_error"
1416 }
1417
1418 async fn create_run_task(
1419 self: Arc<Self>,
1420 ctx: Context,
1421 driver: Driver,
1422 request: Request,
1423 ) -> Result<Response> {
1424 let offloads = futures::future::join_all((0..self.n).map(|i| {
1425 let decision = test_decision(format!("m{i}"));
1426 driver.call_llm(RoutedRequest {
1427 request: request.clone(),
1428 decision,
1429 ctx: ctx.clone(),
1430 })
1431 }));
1432 tokio::select! {
1433 _ = offloads => Err(test_error("offloads unexpectedly completed")),
1434 _ = self.all_started.notified() => {
1435 Err(test_error("terminal error while calls pending"))
1436 }
1437 }
1438 }
1439 }
1440
1441 let all_started = Arc::new(tokio::sync::Notify::new());
1442 let algo: Arc<dyn Algorithm> = Arc::new(FanOutThenError {
1443 all_started: all_started.clone(),
1444 n: N,
1445 });
1446
1447 let started = Arc::new(AtomicUsize::new(0));
1449 let serve = move |_decision: Arc<Decision>, _request: Request| {
1450 let started = started.clone();
1451 let all_started = all_started.clone();
1452 async move {
1453 if started.fetch_add(1, Ordering::SeqCst) + 1 == N {
1454 all_started.notify_one();
1455 }
1456 std::future::pending::<ServeResult>().await
1457 }
1458 };
1459
1460 let run = test_drive(algo, Context::default(), request(), serve);
1463 let result = tokio::time::timeout(std::time::Duration::from_millis(500), run)
1464 .await
1465 .map_err(|error| {
1466 LibsyError::external("waiting for terminal error with full call cap", error)
1467 })?;
1468 match result {
1469 Ok(_) => Err(test_error("expected the terminal error, got a response")),
1470 Err(err) => {
1471 assert!(
1472 err.to_string()
1473 .contains("terminal error while calls pending")
1474 );
1475 Ok(())
1476 }
1477 }
1478 }
1479}