[streaming-quote-api] Stream per-solver results from the price competition (1/3)

crates/price-estimation/src/competition/mod.rs

@@ -1,9 +1,15 @@
 use {
     super::{QuoteVerificationMode, native::NativePriceEstimating},
-    crate::PriceEstimationError,
+    crate::{
+        PriceEstimateResult,
+        PriceEstimating,
+        PriceEstimationError,
+        Query,
+        StreamingPriceEstimating,
+    },
     futures::{
         future::{BoxFuture, FutureExt},
-        stream::{FuturesUnordered, StreamExt},
+        stream::{BoxStream, FuturesUnordered, StreamExt},
     },
     gas_price_estimation::GasPriceEstimating,
     model::order::OrderKind,
@@ -168,6 +174,21 @@ impl<T: Send + Sync + 'static> CompetitionEstimator<T> {
     }
 }
 
+impl StreamingPriceEstimating for CompetitionEstimator<Arc<dyn PriceEstimating>> {
+    /// Runs every estimator concurrently across all stages and yields each
+    /// result as it arrives. No ranking, no early return. The caller stops
+    /// by dropping the stream.
+    fn estimate_stream(&self, query: Arc<Query>) -> BoxStream<'_, PriceEstimateResult> {
+        let futures: FuturesUnordered<BoxFuture<'_, PriceEstimateResult>> = FuturesUnordered::new();
+        for stage in &self.stages {
+            for (_name, estimator) in stage {
+                futures.push(estimator.estimate(query.clone()));
+            }
+        }
+        futures.boxed()
+    }
+}
+
 struct Context<'a, ESTIMATOR, QUERY> {
     /// the estimator that is supposed to compute a price
     estimator: &'a ESTIMATOR,
@@ -253,11 +274,13 @@ mod tests {
             HEALTHY_PRICE_ESTIMATION_TIME,
             MockPriceEstimating,
             PriceEstimating,
+            PriceEstimationError,
             Query,
+            StreamingPriceEstimating,
         },
         alloy::primitives::{Address, U256},
         anyhow::anyhow,
-        futures::channel::oneshot::channel,
+        futures::{StreamExt, channel::oneshot::channel},
         model::order::OrderKind,
         number::nonzero::NonZeroU256,
         std::time::Duration,
@@ -598,6 +621,111 @@ mod tests {
         racing.estimate(query).await.unwrap();
     }
 
+    fn make_query() -> Arc<Query> {
+        Arc::new(Query {
+            verification: Default::default(),
+            sell_token: Address::with_last_byte(0),
+            buy_token: Address::with_last_byte(1),
+            in_amount: NonZeroU256::try_from(1).unwrap(),
+            kind: OrderKind::Sell,
+            block_dependent: false,
+            timeout: HEALTHY_PRICE_ESTIMATION_TIME,
+        })
+    }
+
+    #[tokio::test]
+    async fn estimate_stream_yields_all_results() {
+        let fast = {
+            let mut m = MockPriceEstimating::new();
+            m.expect_estimate().times(1).returning(|_| {
+                async {
+                    Ok(Estimate {
+                        out_amount: U256::from(1u64),
+                        gas: 1,
+                        ..Default::default()
+                    })
+                }
+                .boxed()
+            });
+            m
+        };
+        let slow = {
+            let mut m = MockPriceEstimating::new();
+            m.expect_estimate().times(1).returning(|_| {
+                async {
+                    sleep(Duration::from_millis(10)).await;
+                    Ok(Estimate {
+                        out_amount: U256::from(2u64),
+                        gas: 1,
+                        ..Default::default()
+                    })
+                }
+                .boxed()
+            });
+            m
+        };
+
+        let estimator: CompetitionEstimator<Arc<dyn PriceEstimating>> = CompetitionEstimator::new(
+            vec![vec![
+                ("fast".to_owned(), Arc::new(fast)),
+                ("slow".to_owned(), Arc::new(slow)),
+            ]],
+            PriceRanking::MaxOutAmount,
+        );
+
+        let results: Vec<_> = estimator.estimate_stream(make_query()).collect().await;
+
+        assert_eq!(results.len(), 2);
+        assert!(results.iter().all(|r| r.is_ok()));
+        // `fast` (out_amount 1) resolves immediately while `slow` sleeps, so a
+        // stream that yields as results arrive must emit `fast` first. A serial
+        // collect-then-yield implementation would fail this.
+        let amounts: Vec<_> = results
+            .iter()
+            .map(|r| r.as_ref().unwrap().out_amount)
+            .collect();
+        assert_eq!(amounts, vec![U256::from(1u64), U256::from(2u64)]);
+    }
+
+    #[tokio::test]
+    async fn estimate_stream_passes_through_errors() {
+        let ok = {
+            let mut m = MockPriceEstimating::new();
+            m.expect_estimate().times(1).returning(|_| {
+                async {
+                    Ok(Estimate {
+                        out_amount: U256::from(1u64),
+                        gas: 1,
+                        ..Default::default()
+                    })
+                }
+                .boxed()
+            });
+            m
+        };
+        let err = {
+            let mut m = MockPriceEstimating::new();
+            m.expect_estimate()
+                .times(1)
+                .returning(|_| async { Err(PriceEstimationError::NoLiquidity) }.boxed());
+            m
+        };
+
+        let estimator: CompetitionEstimator<Arc<dyn PriceEstimating>> = CompetitionEstimator::new(
+            vec![vec![
+                ("ok".to_owned(), Arc::new(ok)),
+                ("err".to_owned(), Arc::new(err)),
+            ]],
+            PriceRanking::MaxOutAmount,
+        );
+
+        let results: Vec<_> = estimator.estimate_stream(make_query()).collect().await;
+
+        assert_eq!(results.len(), 2);
+        assert_eq!(results.iter().filter(|r| r.is_ok()).count(), 1);
+        assert_eq!(results.iter().filter(|r| r.is_err()).count(), 1);
+    }
+
     #[test]
     fn custom_solver_errors_have_higher_priority_than_generic_errors() {
         let custom_errors = [

crates/price-estimation/src/lib.rs

@@ -6,7 +6,7 @@ use {
     crate::trade_finding::QuoteExecution,
     alloy::primitives::{Address, U256},
     anyhow::Result,
-    futures::future::BoxFuture,
+    futures::{future::BoxFuture, stream::BoxStream},
     model::order::OrderKind,
     number::nonzero::NonZeroU256,
     rate_limit::RateLimiter,
@@ -223,6 +223,13 @@ pub trait PriceEstimating: Send + Sync + 'static {
     fn estimate(&self, query: Arc<Query>) -> BoxFuture<'_, PriceEstimateResult>;
 }
 
+/// Like `PriceEstimating`, but yields every estimator's result as it
+/// completes instead of collapsing to the single best one.
+#[cfg_attr(any(test, feature = "test-util"), mockall::automock)]
+pub trait StreamingPriceEstimating: Send + Sync + 'static {
+    fn estimate_stream(&self, query: Arc<Query>) -> BoxStream<'_, PriceEstimateResult>;
+}
+
 pub const HEALTHY_PRICE_ESTIMATION_TIME: Duration = Duration::from_millis(5_000);
 
 pub async fn rate_limited<T>(