1use std::{
8 collections::HashMap, future::Future, panic::AssertUnwindSafe, pin::Pin, sync::Arc,
9 time::Instant,
10};
11
12use async_trait::async_trait;
13use futures::{FutureExt, Stream, StreamExt};
14use parking_lot::Mutex;
15use serde_json::Value;
16use tokio::sync::{mpsc, oneshot};
17use tokio_stream::wrappers::ReceiverStream;
18use tracing::Instrument;
19
20use switchyard_protocol::{
28 Category, DecisionRequest, DecisionResponse, ModelId, Request, Response,
29};
30
31use crate::{DriverError, LibsyError, Result, State, observability};
32
33pub type StepStream = Pin<Box<dyn Stream<Item = Result<Step>> + Send>>;
37
38#[derive(Clone, Debug, Default)]
49pub struct RuntimeModels {
50 by_category: HashMap<Category, Vec<ModelId>>,
52 subagent: Option<HashMap<Category, Vec<ModelId>>>,
53}
54
55impl RuntimeModels {
56 pub fn new(by_category: HashMap<Category, Vec<ModelId>>) -> Self {
58 Self {
59 by_category,
60 subagent: None,
61 }
62 }
63
64 pub fn with_subagent(mut self, models: HashMap<Category, Vec<ModelId>>) -> Self {
66 self.subagent = Some(models);
67 self
68 }
69
70 pub fn models_for(&self, category: &Category) -> &[ModelId] {
72 self.by_category.get(category).map_or(&[], Vec::as_slice)
73 }
74
75 pub fn subagent_models_for(&self, category: &Category) -> &[ModelId] {
77 self.subagent
78 .as_ref()
79 .and_then(|models| models.get(category))
80 .map_or(&[], Vec::as_slice)
81 }
82}
83
84impl From<HashMap<Category, Vec<ModelId>>> for RuntimeModels {
85 fn from(by_category: HashMap<Category, Vec<ModelId>>) -> Self {
86 Self::new(by_category)
87 }
88}
89
90#[derive(Clone, Copy)]
92enum Scope {
93 Parent,
94 Subagent,
95}
96
97pub struct CallModel {
107 pub algorithm: String,
110 pub request: Request,
112 pub models: Vec<ModelId>,
114 pub recover_errors: bool,
116 reply: Option<oneshot::Sender<Result<Response>>>,
118 started: Instant,
119 span: tracing::Span,
121}
122
123impl CallModel {
124 pub async fn respond(mut self, work: impl Future<Output = Result<Response>>) -> Result<()> {
130 let reply = self.reply.as_mut().ok_or(DriverError::ResponseDropped)?;
131 let result = tokio::select! {
132 biased;
133 _ = reply.closed() => return Ok(()),
134 result = work => result,
135 };
136 if self.reply.as_ref().is_some_and(oneshot::Sender::is_closed) {
137 return Ok(());
138 }
139 if !self.recover_errors
140 && let Err(error) = result
141 {
142 return self.fail(error);
143 }
144 self.record(result.is_ok());
145 if let Ok(response) = &result {
146 observability::record_llm_response(response, &self.span);
147 }
148 self.reply
149 .take()
150 .ok_or(DriverError::ResponseDropped)?
151 .send(result)
152 .ok();
153 Ok(())
154 }
155
156 pub fn fail(mut self, error: LibsyError) -> Result<()> {
159 self.reply = None;
160 self.record(false);
161 Err(error)
162 }
163
164 fn record(&self, is_ok: bool) {
165 observability::record_llm_call(
166 &self.algorithm,
167 self.models
168 .first()
169 .map(ModelId::as_str)
170 .unwrap_or("NoTargets"),
171 self.started.elapsed(),
172 is_ok,
173 );
174 self.span
175 .record("outcome", if is_ok { "ok" } else { "error" });
176 }
177}
178
179impl Drop for CallModel {
180 fn drop(&mut self) {
181 if self.reply.is_some() {
182 self.record(false);
183 }
184 }
185}
186
187pub struct CallDecision {
190 pub algorithm: String,
192 pub request: DecisionRequest,
194 pub model: ModelId,
196 reply: Option<oneshot::Sender<Result<DecisionResponse>>>,
197 started: Instant,
198 span: tracing::Span,
200}
201
202impl CallDecision {
203 pub async fn respond(
207 mut self,
208 work: impl Future<Output = Result<DecisionResponse>>,
209 ) -> Result<()> {
210 let reply = self.reply.as_mut().ok_or(DriverError::ResponseDropped)?;
211 let result = tokio::select! {
212 biased;
213 _ = reply.closed() => return Ok(()),
214 result = work => result,
215 };
216 self.record(result.is_ok());
217 if let Ok(response) = &result {
218 observability::record_decision_response(response, &self.span);
219 }
220 self.reply
221 .take()
222 .ok_or(DriverError::ResponseDropped)?
223 .send(result)
224 .ok();
225 Ok(())
226 }
227
228 pub fn fail(mut self, error: LibsyError) -> Result<()> {
230 self.reply = None;
231 self.record(false);
232 Err(error)
233 }
234
235 fn record(&self, is_ok: bool) {
236 observability::record_decision_call(
237 &self.algorithm,
238 &self.model,
239 self.started.elapsed(),
240 is_ok,
241 );
242 self.span
243 .record("outcome", if is_ok { "ok" } else { "error" });
244 }
245}
246
247impl Drop for CallDecision {
248 fn drop(&mut self) {
249 if self.reply.is_some() {
250 self.record(false);
251 }
252 }
253}
254
255pub struct RoutingOutcome {
257 pub selected_model_ids: Vec<ModelId>,
259 pub request: Request,
261 pub fallback_request: Option<Request>,
265 pub response: Option<Response>,
267 pub metadata: Option<crate::OutcomeMetadata>,
272 pub state: State,
274}
275
276impl RoutingOutcome {
277 pub fn selected_model_id(&self) -> Result<&ModelId> {
280 self.selected_model_ids.first().ok_or(LibsyError::NoTargets)
281 }
282
283 pub fn route_to(
287 selected_model_id: ModelId,
288 fallback_models: Vec<ModelId>,
289 mut request: Request,
290 ) -> Self {
291 request.llm_request.model = Some(selected_model_id.to_string());
292 let mut selected_model_ids = Vec::with_capacity(1 + fallback_models.len());
293 selected_model_ids.push(selected_model_id);
294 selected_model_ids.extend(fallback_models);
295 Self {
296 selected_model_ids,
297 request,
298 fallback_request: None,
299 response: None,
300 metadata: None,
301 state: State::default(),
302 }
303 }
304
305 pub fn answered(selected_model_id: ModelId, mut request: Request, response: Response) -> Self {
308 request.llm_request.model = Some(selected_model_id.to_string());
309 Self {
310 selected_model_ids: vec![selected_model_id],
311 request,
312 fallback_request: None,
313 response: Some(response),
314 metadata: None,
315 state: State::default(),
316 }
317 }
318}
319
320#[derive(Clone)]
322pub struct Driver {
323 step_tx: mpsc::Sender<Result<Step>>,
324
325 algorithm: String,
327
328 evidence: Arc<Mutex<Option<Value>>>,
330
331 state: Arc<Mutex<State>>,
333
334 models: Arc<RuntimeModels>,
336
337 scope: Scope,
339}
340
341impl Driver {
342 pub(crate) fn new(
345 algorithm: &str,
346 models: Arc<RuntimeModels>,
347 ) -> (Self, mpsc::Receiver<Result<Step>>) {
348 let (step_tx, step_rx) = mpsc::channel(1);
353 (
354 Self {
355 step_tx,
356 algorithm: algorithm.to_string(),
357 evidence: Arc::new(Mutex::new(None)),
358 state: Arc::new(Mutex::new(State::default())),
359 models,
360 scope: Scope::Parent,
361 },
362 step_rx,
363 )
364 }
365
366 pub(crate) fn take_state(&self) -> State {
367 std::mem::take(&mut *self.state.lock())
368 }
369
370 pub(crate) fn set_state(&self, state: State) {
371 *self.state.lock() = state;
372 }
373
374 pub(crate) fn set_evidence(&self, evidence: Value) {
376 *self.evidence.lock() = Some(evidence);
377 }
378
379 pub(crate) fn set_evidence_if_empty(&self, evidence: Value) {
381 let mut current = self.evidence.lock();
382 if current.is_none() {
383 *current = Some(evidence);
384 }
385 }
386
387 pub async fn call_model(&self, request: Request, models: Vec<ModelId>) -> Result<Response> {
397 self.call_model_with_error_recovery(request, models, false)
398 .await
399 }
400
401 #[tracing::instrument(
403 target = "libsy",
404 name = "libsy.llm_call",
405 skip_all,
406 fields(
407 algorithm = self.algorithm,
408 selected_model = %models.first().map(ModelId::as_str).unwrap_or("NoTargets"),
409 openinference.span.kind = "CHAIN",
410 outcome = tracing::field::Empty,
411 input_tokens = tracing::field::Empty,
412 output_tokens = tracing::field::Empty,
413 total_tokens = tracing::field::Empty,
414 reasoning_tokens = tracing::field::Empty,
415 )
416 )]
417 pub(crate) async fn call_model_with_error_recovery(
418 &self,
419 mut request: Request,
420 models: Vec<ModelId>,
421 recover_errors: bool,
422 ) -> Result<Response> {
423 let Some(selected_model_id) = models.first() else {
424 return Err(LibsyError::NoTargets);
425 };
426 request.llm_request.model = Some(selected_model_id.to_string());
427 let started = Instant::now();
428 let (reply, response) = oneshot::channel::<Result<Response>>();
429 let call = CallModel {
430 algorithm: self.algorithm.clone(),
431 request,
432 models,
433 recover_errors,
434 reply: Some(reply),
435 started,
436 span: tracing::Span::current(),
437 };
438 self.step_tx
439 .send(Ok(Step::CallModel(Box::new(call))))
440 .await
441 .map_err(|_| DriverError::StreamClosed)?;
442 response
443 .await
444 .map_err(|_| LibsyError::from(DriverError::ResponseDropped))?
445 }
446
447 #[tracing::instrument(
449 target = "libsy",
450 name = "libsy.decision_call",
451 skip_all,
452 fields(
453 algorithm = self.algorithm,
454 selected_model = %model,
455 openinference.span.kind = "CHAIN",
456 outcome = tracing::field::Empty,
457 input_tokens = tracing::field::Empty,
458 output_tokens = tracing::field::Empty,
459 total_tokens = tracing::field::Empty,
460 reasoning_tokens = tracing::field::Empty,
461 gen_ai.response.id = tracing::field::Empty,
462 gen_ai.response.model = tracing::field::Empty,
463 ),
464 )]
465 pub async fn call_decision(
466 &self,
467 mut request: DecisionRequest,
468 model: ModelId,
469 ) -> Result<DecisionResponse> {
470 request.model = Some(model.clone());
471 let started = Instant::now();
472 let (reply, response) = oneshot::channel();
473 let call = CallDecision {
474 algorithm: self.algorithm.clone(),
475 request,
476 model,
477 reply: Some(reply),
478 started,
479 span: tracing::Span::current(),
480 };
481 self.step_tx
482 .send(Ok(Step::CallDecision(Box::new(call))))
483 .await
484 .map_err(|_| DriverError::StreamClosed)?;
485 response
486 .await
487 .map_err(|_| LibsyError::from(DriverError::ResponseDropped))?
488 }
489
490 pub fn models_for(&self, category: &Category) -> &[ModelId] {
492 match self.scope {
493 Scope::Parent => self.models.models_for(category),
494 Scope::Subagent => self.models.subagent_models_for(category),
495 }
496 }
497
498 pub fn first_model_for(&self, category: &Category) -> Result<&ModelId> {
500 self.models_for(category)
501 .first()
502 .ok_or_else(|| LibsyError::AlgorithmError {
503 message: format!("no models available for category {}", category.as_str()),
504 })
505 }
506
507 pub fn for_subagent(&self) -> Result<Self> {
510 if self.models.subagent.is_none() {
511 return Err(LibsyError::AlgorithmError {
512 message: "delegated work has no sub-agent models".to_string(),
513 });
514 }
515 Ok(Self {
516 scope: Scope::Subagent,
517 ..self.clone()
518 })
519 }
520
521 pub(crate) async fn finish(&self, result: Result<RoutingOutcome>) -> Result<()> {
525 let result = result.map(|mut outcome| {
526 outcome.state = self.take_state();
527 let metadata = outcome.metadata.get_or_insert_with(|| {
528 crate::OutcomeMetadata::new(self.algorithm.clone(), self.evidence.lock().take())
529 });
530 observability::record_outcome(metadata, &outcome.selected_model_ids);
531 outcome
532 });
533 let selected_model = result
534 .as_ref()
535 .ok()
536 .and_then(|outcome| outcome.selected_model_id().ok().cloned());
537 let step = result.map(|outcome| Step::Done(Box::new(outcome)));
538 self.step_tx
539 .send(step)
540 .await
541 .map_err(|_| DriverError::StreamClosed)?;
542 if let Some(selected_model) = selected_model {
543 observability::record_decision(&self.algorithm, &selected_model);
544 }
545 Ok(())
546 }
547}
548
549pub enum Step {
552 CallModel(Box<CallModel>),
555 CallDecision(Box<CallDecision>),
557 Done(Box<RoutingOutcome>),
559}
560
561pub enum Call {
564 Model(Box<CallModel>),
566 Decision(Box<CallDecision>),
568}
569
570pub async fn drive<F, Fut>(
580 algorithm: Arc<dyn Algorithm>,
581 request: Request,
582 models: Arc<RuntimeModels>,
583 serve: F,
584) -> Result<RoutingOutcome>
585where
586 F: Fn(Call) -> Fut,
587 Fut: Future<Output = Result<()>>,
588{
589 drive_with_state(algorithm, request, models, State::default(), serve).await
590}
591
592pub(crate) async fn drive_with_state<F, Fut>(
594 algorithm: Arc<dyn Algorithm>,
595 request: Request,
596 models: Arc<RuntimeModels>,
597 state: State,
598 serve: F,
599) -> Result<RoutingOutcome>
600where
601 F: Fn(Call) -> Fut,
602 Fut: Future<Output = Result<()>>,
603{
604 let stream = algorithm.run_stream(request, models, state);
605 tokio::pin!(stream);
606
607 let mut in_flight = futures::stream::FuturesUnordered::new();
608 let mut final_outcome: Option<RoutingOutcome> = None;
609
610 loop {
611 tokio::select! {
612 Some(result) = in_flight.next() => match result {
613 Ok(()) => {}, Err(err) => return Err(err), },
616 step = stream.next() => {
617 match step {
618 None => break, Some(item) => match item? {
620 Step::CallModel(call) => in_flight.push(serve(Call::Model(call))),
621 Step::CallDecision(call) => {
622 in_flight.push(serve(Call::Decision(call)));
623 }
624 Step::Done(outcome) => {
625 final_outcome = Some(*outcome);
626 break;
627 }
628 }
629 }
630 },
631 }
632 }
633 final_outcome.ok_or(LibsyError::MissingFinalResponse)
634}
635
636fn panic_message(payload: &(dyn std::any::Any + Send)) -> String {
638 payload
639 .downcast_ref::<&'static str>()
640 .map(|message| (*message).to_string())
641 .or_else(|| payload.downcast_ref::<String>().cloned())
642 .unwrap_or_else(|| "unknown panic payload".to_string())
643}
644
645struct AbortOnDrop(tokio::task::AbortHandle);
647
648impl Drop for AbortOnDrop {
649 fn drop(&mut self) {
650 self.0.abort();
651 }
652}
653
654pub(crate) fn ensure_model_is_target(targets: &[ModelId], name: &ModelId) -> Result<()> {
659 targets
660 .iter()
661 .any(|target| target == name)
662 .then_some(())
663 .ok_or_else(|| LibsyError::TargetNotFound {
664 target: name.clone(),
665 })
666}
667
668#[derive(Clone, Hash, PartialEq, Eq)]
671pub(crate) enum RoutingIdentity {
672 Session(String),
674 Subagent { session: String, agent: String },
676}
677
678impl RoutingIdentity {
679 pub(crate) fn from_request(request: &Request) -> Option<Self> {
684 let metadata = request.metadata.as_ref()?;
685 let session = metadata.session_id.as_deref().filter(|id| !id.is_empty())?;
686 if metadata.is_subagent {
687 let agent = metadata.agent_id.as_deref().filter(|id| !id.is_empty())?;
688 Some(Self::Subagent {
689 session: session.to_string(),
690 agent: agent.to_string(),
691 })
692 } else {
693 Some(Self::Session(session.to_string()))
694 }
695 }
696}
697
698#[async_trait]
733pub trait Algorithm: Send + Sync + 'static {
734 fn name(&self) -> &str;
738
739 async fn route(self: Arc<Self>, driver: Driver, request: Request) -> Result<RoutingOutcome>;
742
743 fn run_stream(
751 self: Arc<Self>,
752 request: Request,
753 models: Arc<RuntimeModels>,
754 state: State,
755 ) -> StepStream {
756 let (driver, step_rx) = Driver::new(self.name(), models);
757 driver.set_state(state);
758 let span = observability::run_span(self.name(), &request);
759 let handle = tokio::spawn(
760 async move {
761 let algorithm = self.name().to_string();
762 let route = AssertUnwindSafe(self.route(driver.clone(), request)).catch_unwind();
764 let result = observability::observe_run(&algorithm, async move {
765 route.await.unwrap_or_else(|payload| {
766 Err(LibsyError::AlgorithmError {
767 message: format!(
768 "algorithm task panicked: {}",
769 panic_message(payload.as_ref())
770 ),
771 })
772 })
773 })
774 .await;
775
776 let _ = driver.finish(result).await;
777 }
778 .instrument(span),
779 );
780 let abort_guard = AbortOnDrop(handle.abort_handle());
782 Box::pin(ReceiverStream::new(step_rx).map(move |step| {
783 let _keep_alive = &abort_guard;
785 step
786 }))
787 }
788}
789
790#[cfg(test)]
791mod tests {
792 use std::collections::HashMap;
793
794 use super::*;
795 use crate::core::testing::{Serve, ServeResult, echo, reply, serve_decision, test_drive};
796 use futures::StreamExt;
797 use switchyard_protocol::{
798 LlmResponse, LlmResponseChunk, completion_text, text_request, text_response,
799 };
800
801 #[derive(Debug, thiserror::Error)]
802 #[error("{0}")]
803 struct TestError(&'static str);
804
805 fn test_error(message: &'static str) -> LibsyError {
806 LibsyError::external("test", TestError(message))
807 }
808
809 struct TestAlgo {
812 target_set: Vec<ModelId>,
813 }
814
815 #[async_trait]
816 impl Algorithm for TestAlgo {
817 fn name(&self) -> &str {
818 "test"
819 }
820
821 async fn route(
822 self: Arc<Self>,
823 driver: Driver,
824 request: Request,
825 ) -> Result<RoutingOutcome> {
826 let target = self
827 .target_set
828 .first()
829 .ok_or(LibsyError::NoTargets)?
830 .clone();
831 let response = driver
832 .call_model(request.clone(), vec![target.clone()])
833 .await?;
834 driver.set_evidence(serde_json::json!({"source": "test"}));
835 driver.set_evidence_if_empty(serde_json::json!({"source": "ignored"}));
836 Ok(RoutingOutcome::answered(target, request, response))
837 }
838 }
839
840 fn orch(target_set: Vec<ModelId>) -> Arc<dyn Algorithm> {
842 Arc::new(TestAlgo { target_set })
843 }
844
845 fn request() -> Request {
846 Request {
847 llm_request: text_request(Some("auto".to_string()), "hi".to_string()),
848 raw_request: None,
849 metadata: None,
850 }
851 }
852
853 #[test]
854 fn routing_outcome_constructors_stamp_selection_and_preserve_payloads() {
855 let outcome = RoutingOutcome::route_to(
856 "selected".into(),
857 target_set(&["fallback-one", "fallback-two"]),
858 request(),
859 );
860
861 assert_eq!(
862 outcome.selected_model_ids,
863 target_set(&["selected", "fallback-one", "fallback-two"])
864 );
865 assert_eq!(outcome.request.model_id().as_deref(), Some("selected"));
866 assert!(outcome.response.is_none());
867 assert!(outcome.metadata.is_none());
868
869 let outcome = RoutingOutcome::route_to("only".into(), Vec::new(), request());
870 assert_eq!(outcome.selected_model_ids, target_set(&["only"]));
871
872 let outcome = RoutingOutcome::answered(
873 "answered".into(),
874 request(),
875 Response {
876 llm_response: LlmResponse::Agg(text_response(None, "existing")),
877 metadata: None,
878 upstream_headers: http::HeaderMap::new(),
879 },
880 );
881
882 assert_eq!(outcome.selected_model_ids, target_set(&["answered"]));
883 assert_eq!(outcome.request.model_id().as_deref(), Some("answered"));
884 assert_eq!(
885 outcome
886 .response
887 .as_ref()
888 .and_then(|response| response.llm_response.as_agg())
889 .map(completion_text),
890 Some("existing".to_string())
891 );
892 }
893
894 fn target_set(names: &[&str]) -> Vec<ModelId> {
895 names.iter().map(|name| ModelId::from(*name)).collect()
896 }
897
898 #[tokio::test]
899 async fn typed_driver_preserves_call_and_stream_boundaries() -> Result<()> {
900 tokio::time::timeout(std::time::Duration::from_secs(1), async {
901 let (driver, mut step_rx) = Driver::new("test", Arc::new(RuntimeModels::default()));
904 let first_driver = driver.clone();
905 let mut first = tokio::spawn(async move {
906 first_driver
907 .call_model(request(), vec![ModelId::from("first")])
908 .await
909 });
910 let second = tokio::spawn(async move {
911 driver
912 .call_model(request(), vec![ModelId::from("second")])
913 .await
914 });
915
916 let mut calls = HashMap::new();
917 for _ in 0..2 {
918 let step = step_rx.recv().await.ok_or(DriverError::StreamClosed)??;
919 let Step::CallModel(call) = step else {
920 return Err(test_error("expected a CallModel step"));
921 };
922 let selected_model = call
923 .models
924 .first()
925 .ok_or_else(|| test_error("model call has no candidates"))?
926 .to_string();
927 calls.insert(selected_model, call);
928 }
929 assert!(
930 tokio::time::timeout(std::time::Duration::from_millis(20), &mut first)
931 .await
932 .is_err(),
933 "call completed before the host responded"
934 );
935 calls
936 .remove("second")
937 .ok_or_else(|| test_error("missing second call"))?
938 .respond(std::future::ready(Ok(reply("second response")))).await?;
939 calls
940 .remove("first")
941 .ok_or_else(|| test_error("missing first call"))?
942 .respond(std::future::ready(Ok(reply("first response")))).await?;
943
944 let first_response = first
945 .await
946 .map_err(|source| LibsyError::external("joining a test task", source))??;
947 let second_response = second
948 .await
949 .map_err(|source| LibsyError::external("joining a test task", source))??;
950 assert_eq!(
951 first_response.llm_response.as_agg().map(completion_text),
952 Some("first response".to_string())
953 );
954 assert_eq!(
955 second_response.llm_response.as_agg().map(completion_text),
956 Some("second response".to_string())
957 );
958
959 let (driver, mut step_rx) = Driver::new("test", Arc::new(RuntimeModels::default()));
961 let producer = tokio::spawn(async move {
962 driver
963 .call_model(request(), vec![ModelId::from("dropped")])
964 .await
965 });
966 let step = step_rx.recv().await.ok_or(DriverError::StreamClosed)??;
967 let Step::CallModel(call) = step else {
968 return Err(test_error("expected a CallModel step"));
969 };
970 drop(call);
971 let result = producer
972 .await
973 .map_err(|source| LibsyError::external("joining a test task", source))?;
974 assert!(matches!(
975 result,
976 Err(LibsyError::Driver(DriverError::ResponseDropped))
977 ));
978
979 let (driver, step_rx) = Driver::new("test", Arc::new(RuntimeModels::default()));
981 drop(step_rx);
982 let result = driver
983 .call_model(request(), vec![ModelId::from("closed")])
984 .await;
985 assert!(matches!(
986 result,
987 Err(LibsyError::Driver(DriverError::StreamClosed))
988 ));
989
990 fn decision_request() -> DecisionRequest {
991 DecisionRequest {
992 model: Some("overwritten".into()),
993 context: serde_json::json!({"task": "choose a route"}),
994 questions: Default::default(),
995 }
996 }
997
998 fn decision_response() -> DecisionResponse {
999 DecisionResponse {
1000 id: Some("decision-1".to_string()),
1001 model: Some("provider-model".into()),
1002 answers: [(
1003 "p_solve".to_string(),
1004 switchyard_protocol::DecisionAnswer {
1005 value: switchyard_protocol::DecisionValue::Boolean(
1006 switchyard_protocol::BooleanEstimate::ProbabilityTrue(
1007 switchyard_protocol::Probability(0.8),
1008 ),
1009 ),
1010 provider_confidence: None,
1011 },
1012 )]
1013 .into(),
1014 usage: Default::default(),
1015 }
1016 }
1017
1018 struct MixedCalls(&'static str);
1019
1020 #[async_trait]
1021 impl Algorithm for MixedCalls {
1022 fn name(&self) -> &str {
1023 "mixed"
1024 }
1025
1026 async fn route(
1027 self: Arc<Self>,
1028 driver: Driver,
1029 request: Request,
1030 ) -> Result<RoutingOutcome> {
1031 let (llm, decision) = tokio::join!(
1032 driver.call_model(request.clone(), vec!["llm".into()]),
1033 driver.call_decision(decision_request(), "decision".into()),
1034 );
1035 assert_eq!(
1036 llm?.llm_response.as_agg().map(completion_text),
1037 Some("llm reply".into())
1038 );
1039 match self.0 {
1040 "mock" => assert_eq!(decision?, DecisionResponse {
1041 id: None,
1042 model: Some("decision".into()),
1043 answers: Default::default(),
1044 usage: Default::default(),
1045 }),
1046 "reply" => assert_eq!(decision?, decision_response()),
1047 "error" => assert!(matches!(
1048 decision,
1049 Err(LibsyError::AlgorithmError { message }) if message == "provider failed"
1050 )),
1051 "drop" => assert!(matches!(
1052 decision,
1053 Err(LibsyError::Driver(DriverError::ResponseDropped))
1054 )),
1055 "abort" => return std::future::pending().await,
1056 _ => unreachable!(),
1057 }
1058 Ok(RoutingOutcome::route_to("answer".into(), vec![], request))
1059 }
1060 }
1061
1062 for mode in ["mock", "reply", "error", "drop", "abort"] {
1063 let barrier = Arc::new(tokio::sync::Barrier::new(2));
1065 let outcome = drive(
1066 Arc::new(MixedCalls(mode)),
1067 request(),
1068 Arc::new(RuntimeModels::default()),
1069 move |call| {
1070 let barrier = barrier.clone();
1071 async move {
1072 barrier.wait().await;
1073 let call = match call {
1074 Call::Model(call) => return call.respond(std::future::ready(Ok(reply("llm reply")))).await,
1075 Call::Decision(call) => *call,
1076 };
1077 assert_eq!(call.algorithm, "mixed");
1078 assert_eq!(call.model, "decision");
1079 assert_eq!(call.request.model, Some("decision".into()));
1080 assert_eq!(call.request.context, decision_request().context);
1081 match mode {
1082 "mock" => serve_decision(call).await,
1083 "reply" => call.respond(std::future::ready(Ok(decision_response()))).await,
1084 "error" => call.respond(std::future::ready(Err(LibsyError::AlgorithmError {
1085 message: "provider failed".into(),
1086 }))).await,
1087 "drop" => {
1088 drop(call);
1089 Ok(())
1090 }
1091 "abort" => call.fail(test_error("host aborted")),
1092 _ => unreachable!(),
1093 }
1094 }
1095 },
1096 )
1097 .await;
1098 if mode == "abort" {
1099 assert!(matches!(outcome, Err(LibsyError::External { .. })));
1100 } else {
1101 assert_eq!(outcome?.selected_model_id()?, "answer");
1102 }
1103 }
1104
1105 let (driver, mut steps) = Driver::new("test", Arc::new(RuntimeModels::default()));
1106 let mut pending = Box::pin(driver.call_decision(decision_request(), "decision".into()));
1107 assert!(futures::poll!(&mut pending).is_pending());
1108 let Some(Ok(Step::CallDecision(call))) = steps.recv().await else {
1109 return Err(test_error("expected a decision call"));
1110 };
1111 drop(pending);
1112 call.respond(async { panic!("cancelled work must not be polled") }).await?;
1113 drop(steps);
1114 assert!(matches!(
1115 driver.call_decision(decision_request(), "decision".into()).await,
1116 Err(LibsyError::Driver(DriverError::StreamClosed))
1117 ));
1118 Ok(())
1119 })
1120 .await
1121 .map_err(|error| LibsyError::external("waiting for typed driver boundaries", error))?
1122 }
1123
1124 #[test]
1125 fn target_lookup_returns_the_missing_target() {
1126 let error = ensure_model_is_target(&target_set(&[]), &ModelId::from("missing")).err();
1127 assert!(matches!(
1128 error,
1129 Some(LibsyError::TargetNotFound { target }) if target == "missing"
1130 ));
1131 }
1132
1133 fn streaming_orch(chunks: Vec<LlmResponseChunk>) -> (Arc<dyn Algorithm>, impl Serve) {
1136 let algo = orch(target_set(&["stream/model"]));
1137 let serve = move |_target: ModelId, _request: Request| {
1138 let chunks = chunks.clone();
1139 async move {
1140 let stream =
1141 futures::stream::iter(chunks.into_iter().map(|chunk| Ok(chunk.into()))).boxed();
1142 Ok(Response {
1143 llm_response: LlmResponse::Stream(stream),
1144 metadata: None,
1145 upstream_headers: http::HeaderMap::new(),
1146 })
1147 }
1148 };
1149 (algo, serve)
1150 }
1151
1152 #[tokio::test]
1153 async fn run_returns_a_streamed_response_the_caller_aggregates() -> Result<()> {
1154 let (orch, serve) = streaming_orch(vec![
1157 LlmResponseChunk::MessageStart {
1158 id: Some("m1".to_string()),
1159 model: Some("stream/model".to_string()),
1160 },
1161 LlmResponseChunk::TextDelta {
1162 index: 0,
1163 text: "hel".to_string(),
1164 },
1165 LlmResponseChunk::TextDelta {
1166 index: 0,
1167 text: "lo".to_string(),
1168 },
1169 LlmResponseChunk::MessageStop {
1170 reason: Some("stop".to_string()),
1171 },
1172 ]);
1173 let (selected_model, response) = test_drive(orch, request(), serve).await?;
1174 let agg = response
1176 .llm_response
1177 .into_agg()
1178 .await
1179 .map_err(|error| LibsyError::external("aggregating response stream", error))?;
1180 assert_eq!(completion_text(&agg), "hello");
1181 assert_eq!(agg.model.as_deref(), Some("stream/model"));
1182 assert_eq!(selected_model, "stream/model");
1183 Ok(())
1184 }
1185
1186 #[tokio::test]
1187 async fn aggregating_a_streamed_response_propagates_a_mid_stream_error() -> Result<()> {
1188 let (orch, serve) = streaming_orch(vec![
1191 LlmResponseChunk::TextDelta {
1192 index: 0,
1193 text: "partial".to_string(),
1194 },
1195 LlmResponseChunk::StreamError {
1196 message: "upstream exploded".to_string(),
1197 },
1198 ]);
1199 let (_, response) = test_drive(orch, request(), serve).await?;
1200 match response.llm_response.into_agg().await {
1201 Ok(_) => panic!("expected a mid-stream error, got an aggregate"),
1202 Err(err) => {
1203 assert!(err.to_string().contains("upstream exploded"));
1204 Ok(())
1205 }
1206 }
1207 }
1208
1209 #[tokio::test]
1210 async fn run_offloads_via_promise_then_finishes() -> Result<()> {
1211 let stream = orch(target_set(&["offload/model"])).run_stream(
1214 request(),
1215 Arc::new(RuntimeModels::default()),
1216 State::default(),
1217 );
1218 tokio::pin!(stream);
1219
1220 let mut saw_call = false;
1221 let mut final_completion = None;
1222 while let Some(step) = stream.next().await {
1223 match step? {
1224 Step::CallDecision(_) => return Err(test_error("unexpected decision call")),
1225 Step::CallModel(call) => {
1226 saw_call = true;
1227 assert_eq!(call.models, vec![ModelId::from("offload/model")]);
1228 call.respond(std::future::ready(Ok(Response {
1230 llm_response: LlmResponse::Agg(text_response(
1231 None,
1232 "fulfilled".to_string(),
1233 )),
1234 metadata: None,
1235 upstream_headers: http::HeaderMap::new(),
1236 })))
1237 .await?;
1238 }
1239 Step::Done(outcome) => {
1240 let metadata = outcome
1241 .metadata
1242 .as_ref()
1243 .expect("run_stream should attach outcome metadata");
1244 assert_eq!(metadata.algorithm, "test");
1245 assert_eq!(
1246 uuid::Uuid::parse_str(metadata.outcome_id())
1247 .expect("outcome id should be a UUID")
1248 .get_version_num(),
1249 7
1250 );
1251 assert_eq!(
1252 metadata.evidence,
1253 Some(serde_json::json!({"source": "test"}))
1254 );
1255 let response = outcome
1256 .response
1257 .ok_or_else(|| test_error("expected an answered outcome"))?;
1258 final_completion = Some(
1259 response
1260 .llm_response
1261 .as_agg()
1262 .map(completion_text)
1263 .unwrap_or_default(),
1264 );
1265 }
1266 }
1267 }
1268
1269 assert!(saw_call, "expected a CallModel step before Done");
1270 assert_eq!(
1271 final_completion.ok_or_else(|| test_error("no Done step"))?,
1272 "fulfilled"
1273 );
1274 Ok(())
1275 }
1276
1277 #[tokio::test(flavor = "multi_thread", worker_threads = 12)]
1278 async fn requests_are_processed_in_parallel() -> Result<()> {
1279 use std::time::Duration;
1280 use tokio::sync::Barrier;
1281
1282 const N: usize = 12;
1283
1284 let barrier = Arc::new(Barrier::new(N));
1289 let algo = orch(target_set(&["m"]));
1291
1292 let mut handles = Vec::new();
1293 for _ in 0..N {
1294 let algo = algo.clone();
1295 let barrier = barrier.clone();
1296 let serve = move |target: ModelId, _request: Request| {
1297 let barrier = barrier.clone();
1298 async move {
1299 barrier.wait().await;
1300 Ok(reply(target))
1301 }
1302 };
1303 handles.push(tokio::spawn(async move {
1304 test_drive(algo, request(), serve)
1305 .await
1306 .map(|(_, response)| {
1307 response
1308 .llm_response
1309 .as_agg()
1310 .map(completion_text)
1311 .unwrap_or_default()
1312 })
1313 }));
1314 }
1315
1316 for handle in handles {
1317 let completion = tokio::time::timeout(Duration::from_secs(5), handle)
1319 .await
1320 .map_err(|error| LibsyError::external("waiting for test task", error))?
1321 .map_err(|source| LibsyError::external("joining a test task", source))??;
1322 assert_eq!(completion, "m");
1323 }
1324 Ok(())
1325 }
1326
1327 #[tokio::test]
1328 async fn offload_error_propagates_back_to_the_algorithm() -> Result<()> {
1329 let stream = orch(target_set(&["offload/model"])).run_stream(
1333 request(),
1334 Arc::new(RuntimeModels::default()),
1335 State::default(),
1336 );
1337 tokio::pin!(stream);
1338
1339 let mut saw_error = false;
1340 while let Some(step) = stream.next().await {
1341 match step {
1342 Ok(Step::CallDecision(_)) => return Err(test_error("unexpected decision call")),
1343 Ok(Step::CallModel(mut call)) => {
1344 call.recover_errors = true;
1345 call.respond(std::future::ready(Err(test_error(
1346 "upstream model call failed",
1347 ))))
1348 .await?;
1349 }
1350 Ok(Step::Done(..)) => {
1351 return Err(test_error(
1352 "expected the offload error to propagate, got a response",
1353 ));
1354 }
1355 Err(err) => {
1356 assert!(err.to_string().contains("upstream model call failed"));
1358 saw_error = true;
1359 }
1360 }
1361 }
1362
1363 assert!(saw_error, "expected an error step");
1364 Ok(())
1365 }
1366
1367 #[tokio::test]
1368 async fn dropping_the_stream_cancels_the_algorithm_task() -> Result<()> {
1369 use std::sync::atomic::{AtomicBool, Ordering};
1370 use std::time::Duration;
1371 use tokio::sync::mpsc;
1372
1373 struct DropGuard(Arc<AtomicBool>);
1376 impl Drop for DropGuard {
1377 fn drop(&mut self) {
1378 self.0.store(true, Ordering::SeqCst);
1379 }
1380 }
1381
1382 struct StuckAlgo {
1383 started: mpsc::UnboundedSender<()>,
1384 dropped: Arc<AtomicBool>,
1385 }
1386
1387 #[async_trait]
1388 impl Algorithm for StuckAlgo {
1389 fn name(&self) -> &str {
1390 "stuck"
1391 }
1392
1393 async fn route(
1394 self: Arc<Self>,
1395 _driver: Driver,
1396 _request: Request,
1397 ) -> Result<RoutingOutcome> {
1398 let _guard = DropGuard(self.dropped.clone());
1399 let _ = self.started.send(());
1400 std::future::pending::<()>().await;
1402 unreachable!()
1403 }
1404 }
1405
1406 let (started_tx, mut started_rx) = mpsc::unbounded_channel();
1407 let dropped = Arc::new(AtomicBool::new(false));
1408 let algo: Arc<dyn Algorithm> = Arc::new(StuckAlgo {
1409 started: started_tx,
1410 dropped: dropped.clone(),
1411 });
1412
1413 let stream = algo.run_stream(
1414 request(),
1415 Arc::new(RuntimeModels::default()),
1416 State::default(),
1417 );
1418 started_rx
1419 .recv()
1420 .await
1421 .ok_or_else(|| test_error("task never started"))?;
1422 drop(stream);
1423 tokio::time::sleep(Duration::from_millis(100)).await;
1424
1425 assert!(
1426 dropped.load(Ordering::SeqCst),
1427 "algorithm task was NOT cancelled after dropping the stream"
1428 );
1429 Ok(())
1430 }
1431
1432 #[tokio::test]
1433 async fn route_panic_surfaces_as_a_stream_error() -> Result<()> {
1434 struct Panicky;
1437
1438 #[async_trait]
1439 impl Algorithm for Panicky {
1440 fn name(&self) -> &str {
1441 "panicky"
1442 }
1443
1444 async fn route(
1445 self: Arc<Self>,
1446 _driver: Driver,
1447 _request: Request,
1448 ) -> Result<RoutingOutcome> {
1449 panic!("boom");
1450 }
1451 }
1452
1453 let algo: Arc<dyn Algorithm> = Arc::new(Panicky);
1454 let stream = algo.run_stream(
1455 request(),
1456 Arc::new(RuntimeModels::default()),
1457 State::default(),
1458 );
1459 tokio::pin!(stream);
1460
1461 let mut saw_error = false;
1462 while let Some(step) = stream.next().await {
1463 match step {
1464 Err(err) => {
1465 assert!(err.to_string().contains("algorithm task panicked: boom"));
1467 saw_error = true;
1468 }
1469 Ok(_) => return Err(test_error("expected the panic to surface as an error step")),
1470 }
1471 }
1472
1473 assert!(saw_error, "expected an error step from the panicked task");
1474 Ok(())
1475 }
1476
1477 #[tokio::test]
1481 async fn a_panic_with_a_leaked_driver_clone_still_terminates_the_run() -> Result<()> {
1482 struct LeakyPanic;
1483
1484 #[async_trait]
1485 impl Algorithm for LeakyPanic {
1486 fn name(&self) -> &str {
1487 "leaky_panic"
1488 }
1489
1490 async fn route(
1491 self: Arc<Self>,
1492 driver: Driver,
1493 _request: Request,
1494 ) -> Result<RoutingOutcome> {
1495 tokio::spawn(async move {
1496 let _keep_alive = driver;
1498 std::future::pending::<()>().await;
1499 });
1500 tokio::task::yield_now().await;
1501 panic!("boom");
1502 }
1503 }
1504
1505 let algo: Arc<dyn Algorithm> = Arc::new(LeakyPanic);
1506 let result = tokio::time::timeout(
1508 std::time::Duration::from_secs(1),
1509 test_drive(algo, request(), echo()),
1510 )
1511 .await
1512 .map_err(|error| LibsyError::external("waiting for the panicked run to end", error))?;
1513
1514 match result {
1515 Ok(_) => Err(test_error(
1516 "expected the panic to end the run with an error",
1517 )),
1518 Err(err) => {
1519 assert!(err.to_string().contains("algorithm task panicked: boom"));
1520 Ok(())
1521 }
1522 }
1523 }
1524
1525 #[tokio::test]
1526 async fn cancelling_run_cancels_the_algorithm_task() -> Result<()> {
1527 use std::sync::atomic::{AtomicBool, Ordering};
1528 use std::time::Duration;
1529 use tokio::sync::mpsc;
1530
1531 struct DropGuard(Arc<AtomicBool>);
1534 impl Drop for DropGuard {
1535 fn drop(&mut self) {
1536 self.0.store(true, Ordering::SeqCst);
1537 }
1538 }
1539
1540 struct StuckAlgo {
1541 started: mpsc::UnboundedSender<()>,
1542 dropped: Arc<AtomicBool>,
1543 }
1544
1545 #[async_trait]
1546 impl Algorithm for StuckAlgo {
1547 fn name(&self) -> &str {
1548 "stuck"
1549 }
1550
1551 async fn route(
1552 self: Arc<Self>,
1553 _driver: Driver,
1554 _request: Request,
1555 ) -> Result<RoutingOutcome> {
1556 let _guard = DropGuard(self.dropped.clone());
1557 let _ = self.started.send(());
1558 std::future::pending::<()>().await;
1561 unreachable!()
1562 }
1563 }
1564
1565 let (started_tx, mut started_rx) = mpsc::unbounded_channel();
1566 let dropped = Arc::new(AtomicBool::new(false));
1567 let algo: Arc<dyn Algorithm> = Arc::new(StuckAlgo {
1568 started: started_tx,
1569 dropped: dropped.clone(),
1570 });
1571
1572 let run_task = tokio::spawn(async move { test_drive(algo, request(), echo()).await });
1575 started_rx
1576 .recv()
1577 .await
1578 .ok_or_else(|| test_error("task never started"))?;
1579 run_task.abort();
1580 tokio::time::sleep(Duration::from_millis(100)).await;
1581
1582 assert!(
1583 dropped.load(Ordering::SeqCst),
1584 "algorithm task was NOT cancelled after cancelling run"
1585 );
1586 Ok(())
1587 }
1588
1589 struct Hedge {
1594 winner: String,
1595 loser: String,
1596 }
1597
1598 #[async_trait]
1599 impl Algorithm for Hedge {
1600 fn name(&self) -> &str {
1601 "hedge"
1602 }
1603
1604 async fn route(
1605 self: Arc<Self>,
1606 driver: Driver,
1607 request: Request,
1608 ) -> Result<RoutingOutcome> {
1609 let outcome_request = request.clone();
1610 let win = driver.call_model(request.clone(), vec![self.winner.clone().into()]);
1611 let lose = driver.call_model(request, vec![self.loser.clone().into()]);
1612 tokio::select! {
1614 res = win => Ok(RoutingOutcome::answered(
1615 self.winner.clone().into(),
1616 outcome_request,
1617 res?,
1618 )),
1619 res = lose => Ok(RoutingOutcome::answered(
1620 self.loser.clone().into(),
1621 outcome_request,
1622 res?,
1623 )),
1624 }
1625 }
1626 }
1627
1628 fn hedge(loser_delay: Option<std::time::Duration>) -> (Arc<dyn Algorithm>, impl Serve) {
1632 let started = Arc::new(tokio::sync::Notify::new());
1633 let algo = Arc::new(Hedge {
1634 winner: "winner".to_string(),
1635 loser: "loser".to_string(),
1636 });
1637 let serve = move |target: ModelId, _request: Request| {
1638 let started = started.clone();
1639 async move {
1640 if target == "loser" {
1641 started.notify_one();
1642 match loser_delay {
1643 Some(delay) => tokio::time::sleep(delay).await,
1644 None => std::future::pending::<()>().await,
1645 }
1646 } else {
1647 started.notified().await;
1648 }
1649 Ok(reply(target))
1650 }
1651 };
1652 (algo, serve)
1653 }
1654
1655 #[tokio::test]
1656 async fn run_returns_the_winner_without_a_late_loser_overwriting_it() -> Result<()> {
1657 let (algo, serve) = hedge(Some(std::time::Duration::from_millis(50)));
1660 let (_, response) = test_drive(algo, request(), serve).await?;
1661 assert_eq!(
1662 response
1663 .llm_response
1664 .as_agg()
1665 .map(completion_text)
1666 .unwrap_or_default(),
1667 "winner"
1668 );
1669 Ok(())
1670 }
1671
1672 #[tokio::test]
1673 async fn run_returns_the_winner_without_hanging_on_a_pending_loser() -> Result<()> {
1674 let (algo, serve) = hedge(None);
1677 let run = test_drive(algo, request(), serve);
1678 let (_, response) = tokio::time::timeout(std::time::Duration::from_secs(1), run)
1679 .await
1680 .map_err(|error| LibsyError::external("waiting for pending loser", error))??;
1681 assert_eq!(
1682 response
1683 .llm_response
1684 .as_agg()
1685 .map(completion_text)
1686 .unwrap_or_default(),
1687 "winner"
1688 );
1689 Ok(())
1690 }
1691
1692 #[tokio::test]
1693 async fn run_surfaces_a_terminal_error_with_many_calls_in_flight() -> Result<()> {
1694 use std::sync::atomic::{AtomicUsize, Ordering};
1695
1696 const N: usize = 10;
1699
1700 struct FanOutThenError {
1703 all_started: Arc<tokio::sync::Notify>,
1704 n: usize,
1705 }
1706
1707 #[async_trait]
1708 impl Algorithm for FanOutThenError {
1709 fn name(&self) -> &str {
1710 "fan_out_then_error"
1711 }
1712
1713 async fn route(
1714 self: Arc<Self>,
1715 driver: Driver,
1716 request: Request,
1717 ) -> Result<RoutingOutcome> {
1718 let offloads = futures::future::join_all(
1719 (0..self.n)
1720 .map(|i| driver.call_model(request.clone(), vec![format!("m{i}").into()])),
1721 );
1722 tokio::select! {
1723 _ = offloads => Err(test_error("offloads unexpectedly completed")),
1724 _ = self.all_started.notified() => {
1725 Err(test_error("terminal error while calls pending"))
1726 }
1727 }
1728 }
1729 }
1730
1731 let all_started = Arc::new(tokio::sync::Notify::new());
1732 let algo: Arc<dyn Algorithm> = Arc::new(FanOutThenError {
1733 all_started: all_started.clone(),
1734 n: N,
1735 });
1736
1737 let started = Arc::new(AtomicUsize::new(0));
1739 let serve = move |_target: ModelId, _request: Request| {
1740 let started = started.clone();
1741 let all_started = all_started.clone();
1742 async move {
1743 if started.fetch_add(1, Ordering::SeqCst) + 1 == N {
1744 all_started.notify_one();
1745 }
1746 std::future::pending::<ServeResult>().await
1747 }
1748 };
1749
1750 let run = test_drive(algo, request(), serve);
1753 let result = tokio::time::timeout(std::time::Duration::from_millis(500), run)
1754 .await
1755 .map_err(|error| {
1756 LibsyError::external("waiting for terminal error with full call cap", error)
1757 })?;
1758 match result {
1759 Ok(_) => Err(test_error("expected the terminal error, got a response")),
1760 Err(err) => {
1761 assert!(
1762 err.to_string()
1763 .contains("terminal error while calls pending")
1764 );
1765 Ok(())
1766 }
1767 }
1768 }
1769}