Skip to main content

tycho_simulation/price_level_stream/
stream.rs

1use std::{
2    collections::{hash_map::Entry, HashMap, HashSet},
3    time::Duration,
4};
5
6use chrono::Utc;
7use num_bigint::BigUint;
8use tokio_stream::{Stream, StreamExt};
9use tycho_common::{
10    models::{token::Token, Chain},
11    simulation::protocol_sim::ProtocolSim,
12    Bytes,
13};
14
15use super::{
16    config::{
17        default_denied_pamms, default_served_pamms, PriceLevelStreamConfig,
18        DEFAULT_AUTO_DETECTED_GAS_COST,
19    },
20    fallback_router::fetch_fallback_router_venues,
21    state::{PriceLevelStreamQuote, PriceLevelStreamState},
22    titan::{
23        self, ConnectionSettings, TitanPairLevels, TitanPammLevels, TitanPriceLevel,
24        TitanPriceLevelMessage, TITAN_PRICE_LEVEL_URL,
25    },
26};
27use crate::protocol::models::{ProtocolComponent, Update};
28
29/// Static attribute under which each emitted component carries its pAMM venue address.
30pub const PAMM_ADDRESS_ATTRIBUTE: &str = "pamm_address";
31
32/// Builds a stream of [`Update`]s from the Titan pAMM price level WebSocket.
33///
34/// A new builder serves no pAMMs: register the known venues via
35/// [`with_known_pamms`](Self::with_known_pamms), individual ones via
36/// [`add_pamm`](Self::add_pamm), or opt into serving unknown streamed venues via
37/// [`auto_detect`](Self::auto_detect); [`with_tokens`](Self::with_tokens) provides the token
38/// metadata pairs are interpreted with.
39///
40/// One component is emitted per (pAMM, token pair), identified by the concatenation
41/// `pamm ++ token0 ++ token1` (tokens sorted ascending), under the protocol system
42/// `pricelevelstream:{pamm}` — or `propammfallback:{pamm}` for venues on the PropAMMRouter
43/// whitelist, unless [`without_fallback_router`](Self::without_fallback_router) turns that off.
44/// The venue address is exposed through the [`PAMM_ADDRESS_ATTRIBUTE`] static attribute for
45/// downstream encoding.
46pub struct PriceLevelStreamBuilder {
47    registry: HashMap<Bytes, PriceLevelStreamConfig>,
48    denied: HashSet<Bytes>,
49    tokens: HashMap<Bytes, Token>,
50    url: Option<String>,
51    auto_detect: bool,
52    auto_detected_gas_cost: Option<BigUint>,
53    connection: ConnectionSettings,
54    /// Whether [`build`](Self::build) reads the PropAMMRouter whitelist and serves the venues on
55    /// it under the `propammfallback:` family.
56    fallback_router: bool,
57}
58
59impl Default for PriceLevelStreamBuilder {
60    fn default() -> Self {
61        Self {
62            registry: HashMap::new(),
63            denied: HashSet::new(),
64            tokens: HashMap::new(),
65            url: None,
66            auto_detect: false,
67            auto_detected_gas_cost: None,
68            connection: ConnectionSettings::default(),
69            fallback_router: true,
70        }
71    }
72}
73
74impl PriceLevelStreamBuilder {
75    pub fn new() -> Self {
76        Self::default()
77    }
78
79    /// Enables serving pAMMs that are not registered via
80    /// [`with_known_pamms`](Self::with_known_pamms) or [`add_pamm`](Self::add_pamm)
81    /// (disabled by default).
82    ///
83    /// When enabled, any unknown streamed venue — except denied ones (see
84    /// [`deny_pamm`](Self::deny_pamm)) — is served under its full lowercase hex address
85    /// as the name, with the default gas cost. A venue's protocol system therefore changes from
86    /// the address form (`pricelevelstream:{0xaddress}`) to a name (`pricelevelstream:{name}`)
87    /// once it gets registered — via [`add_pamm`](Self::add_pamm) or a release's
88    /// [`default_served_pamms`] recognizing it; the name-independent identifiers — the component id
89    /// and the [`PAMM_ADDRESS_ATTRIBUTE`] — stay stable across such renames.
90    pub fn auto_detect(mut self, enabled: bool) -> Self {
91        self.auto_detect = enabled;
92        self
93    }
94
95    /// Overrides the per-swap gas cost that auto-detected pAMMs (see
96    /// [`auto_detect`](Self::auto_detect)) are served with. Defaults to the maximum over the
97    /// known venue profiles, as the conservative choice. Registered venues are unaffected —
98    /// their gas cost comes from their [`PriceLevelStreamConfig`].
99    pub fn auto_detected_gas_cost(mut self, gas_cost: BigUint) -> Self {
100        self.auto_detected_gas_cost = Some(gas_cost);
101        self
102    }
103
104    /// Overrides the stream endpoint, e.g. to connect to a closer Titan region than the default
105    /// (see <https://docs.titanbuilder.xyz/propamms/takers>).
106    pub fn endpoint(mut self, url: impl Into<String>) -> Self {
107        self.url = Some(url.into());
108        self
109    }
110
111    /// Overrides how long a single connection attempt may take before it is aborted and retried
112    /// (default: 10s), so a hung TCP/TLS handshake cannot block the stream forever.
113    pub fn connect_timeout(mut self, timeout: Duration) -> Self {
114        self.connection.connect_timeout = timeout;
115        self
116    }
117
118    /// Overrides the longest gap between Titan messages tolerated before the connection is
119    /// treated as dead and re-established (default: 30s). Titan pushes several updates per
120    /// second, so a multi-second silence means a stalled or half-open connection.
121    pub fn read_idle_timeout(mut self, timeout: Duration) -> Self {
122        self.connection.read_idle_timeout = timeout;
123        self
124    }
125
126    /// Overrides the cap on the exponential reconnect backoff of `2^attempt` seconds
127    /// (default: 32s).
128    pub fn max_backoff(mut self, max_backoff: Duration) -> Self {
129        self.connection.max_backoff = max_backoff;
130        self
131    }
132
133    /// Registers a pAMM to be served under the given configuration, overriding any default,
134    /// denied, or auto-detected one for the same address.
135    ///
136    /// Between [`add_pamm`](Self::add_pamm) and [`deny_pamm`](Self::deny_pamm) for the same
137    /// address, the later call wins; the defaults applied by
138    /// [`with_known_pamms`](Self::with_known_pamms) never override either, in any call order.
139    pub fn add_pamm(mut self, config: PriceLevelStreamConfig) -> Self {
140        self.denied.remove(&config.address);
141        self.registry
142            .insert(config.address.clone(), config);
143        self
144    }
145
146    /// Excludes a venue from being served: drops its current registration (default or explicit)
147    /// and blocks auto-detecting it.
148    ///
149    /// Between [`add_pamm`](Self::add_pamm) and [`deny_pamm`](Self::deny_pamm) for the same
150    /// address, the later call wins; the defaults applied by
151    /// [`with_known_pamms`](Self::with_known_pamms) never override either, in any call order —
152    /// so denying a venue from the default set works whether the denial comes before or after
153    /// [`with_known_pamms`](Self::with_known_pamms).
154    pub fn deny_pamm(mut self, address: Bytes) -> Self {
155        self.registry.remove(&address);
156        self.denied.insert(address);
157        self
158    }
159
160    /// Applies what is known about the streamed venues: registers the known-good ones
161    /// ([`default_served_pamms`]) to be served and denies the known-bad ones
162    /// ([`default_denied_pamms`]) — venues that stream quotes but whose swaps are not executable.
163    ///
164    /// These defaults never override an explicit [`add_pamm`](Self::add_pamm) or
165    /// [`deny_pamm`](Self::deny_pamm) for the same address, regardless of call order.
166    pub fn with_known_pamms(mut self) -> Self {
167        for config in default_served_pamms() {
168            if self.denied.contains(&config.address) {
169                continue;
170            }
171            self.registry
172                .entry(config.address.clone())
173                .or_insert(config);
174        }
175        for address in default_denied_pamms() {
176            if self.registry.contains_key(&address) {
177                continue;
178            }
179            self.denied.insert(address);
180        }
181        self
182    }
183
184    /// Provides the token metadata used to build components and interpret amounts. Pairs whose
185    /// tokens are missing here are skipped.
186    pub fn with_tokens(mut self, tokens: HashMap<Bytes, Token>) -> Self {
187        self.tokens = tokens;
188        self
189    }
190
191    /// Keeps every venue on the direct `pricelevelstream:{name}` path, so swaps execute on the
192    /// venues themselves and a stale maker quote reverts the route.
193    ///
194    /// By default [`build`](Self::build) emits venues on Titan's PropAMMRouter whitelist under
195    /// `propammfallback:{name}` instead, so tycho-execution routes their swaps through the
196    /// router. Opt out when the direct call is what you want to measure or execute, or to skip
197    /// the whitelist read at startup.
198    pub fn without_fallback_router(mut self) -> Self {
199        self.fallback_router = false;
200        self
201    }
202
203    /// Consumes the builder and opens the stream.
204    ///
205    /// Venues on Titan's PropAMMRouter whitelist are served under `propammfallback:{name}`, so
206    /// tycho-execution routes their swaps through the router. The router falls back to a
207    /// single-hop Uniswap V3 pool when the venue reverts — which a stale maker quote does in any
208    /// simulation against a mined block. Only whitelisted venues may use the family: the router
209    /// reverts `UnknownVenue` for others, so every swap would execute on the Uniswap V3 fallback
210    /// at a worse price than the venue gives.
211    ///
212    /// Reading that whitelist needs a node at `RPC_URL` (from the environment, falling back to
213    /// `.env`), and degrades instead of failing: without the variable, or when the read fails, a
214    /// warning is logged and every venue stays on the direct `pricelevelstream:` path.
215    /// [`without_fallback_router`](Self::without_fallback_router) skips the read and takes the
216    /// direct path unconditionally.
217    ///
218    /// The whitelist is read once, on the first poll, and never re-read — it is governance-gated
219    /// and changes rarely, and renaming a running component's protocol system would churn every
220    /// consumer's component set. Restart the stream to pick up a whitelist change.
221    ///
222    /// The connection is established lazily on first poll and maintained (with reconnects) for as
223    /// long as the stream is polled; it never terminates on its own, and dropping the stream
224    /// closes the connection. Frames that contain no served pAMM produce no update.
225    ///
226    /// Each streamed frame is a complete snapshot of everything Titan currently streams, so
227    /// every update carries the full set of the frame's pair states, with `new_pairs` /
228    /// `removed_pairs` derived by diffing against the previous frame — a pair (or a whole
229    /// venue) the stream stops serving is removed. Frames older than an already processed one
230    /// are skipped, so updates never move backwards in block number. Pairs whose tokens are
231    /// missing from the provided token metadata are skipped.
232    pub fn build(self) -> impl Stream<Item = Update> + Send {
233        let Self {
234            registry,
235            denied,
236            tokens,
237            url,
238            auto_detect,
239            auto_detected_gas_cost,
240            connection,
241            fallback_router,
242        } = self;
243        if registry.is_empty() && !auto_detect {
244            tracing::warn!(
245                "No pAMMs registered and auto-detection is off; the stream will never produce \
246                 an update"
247            );
248        }
249        if tokens.is_empty() {
250            tracing::warn!(
251                "No token metadata provided; every streamed pair will be skipped and the stream \
252                 will never produce an update"
253            );
254        }
255        let url = url.unwrap_or_else(|| TITAN_PRICE_LEVEL_URL.to_string());
256        let auto_detected_gas_cost =
257            auto_detected_gas_cost.unwrap_or_else(|| BigUint::from(DEFAULT_AUTO_DETECTED_GAS_COST));
258
259        futures::FutureExt::flatten_stream(async move {
260            let router_venues = if fallback_router {
261                fetch_router_venues(rpc_url_from_env()).await
262            } else {
263                HashSet::new()
264            };
265            let mut tracker = SnapshotTracker::new(
266                registry,
267                denied,
268                tokens,
269                auto_detect,
270                auto_detected_gas_cost,
271                router_venues,
272            );
273
274            titan::messages(url, connection).filter_map(move |message| tracker.process(message))
275        })
276    }
277}
278
279/// The node URL the whitelist is read from: `RPC_URL` from the environment, falling back to
280/// `.env`.
281fn rpc_url_from_env() -> Option<String> {
282    std::env::var("RPC_URL")
283        .ok()
284        .or_else(|| {
285            dotenv::dotenv().ok()?;
286            std::env::var("RPC_URL").ok()
287        })
288}
289
290/// The PropAMMRouter's whitelisted venues, or an empty set when they cannot be read — no node
291/// URL, or a failed call. Both cases warn and leave every venue on the direct path, so a
292/// misconfigured deployment loses the Uniswap V3 fallback instead of losing the stream.
293async fn fetch_router_venues(rpc_url: Option<String>) -> HashSet<Bytes> {
294    let Some(rpc_url) = rpc_url else {
295        tracing::warn!(
296            "RPC_URL is not set; pAMM swaps execute on the venues directly, without the \
297             PropAMMRouter's Uniswap V3 fallback"
298        );
299        return HashSet::new();
300    };
301
302    match fetch_fallback_router_venues(&rpc_url).await {
303        Ok(venues) => venues.into_iter().collect(),
304        Err(e) => {
305            tracing::warn!(
306                error = %e,
307                "Could not read the PropAMMRouter venue whitelist; pAMM swaps execute on the \
308                 venues directly, without the Uniswap V3 fallback"
309            );
310            HashSet::new()
311        }
312    }
313}
314
315/// Turns Titan frames into [`Update`]s, tracking the previously emitted components so pair
316/// additions and removals can be diffed against the last snapshot.
317struct SnapshotTracker {
318    registry: HashMap<Bytes, PriceLevelStreamConfig>,
319    /// Venues excluded from auto-detection. The builder keeps this disjoint from the registry:
320    /// denying removes any registration and registering removes any denial.
321    denied: HashSet<Bytes>,
322    tokens: HashMap<Bytes, Token>,
323    /// Whether frames from pAMMs absent from the registry get an address-named configuration
324    /// synthesized (and cached in the registry) instead of being skipped.
325    auto_detect: bool,
326    /// The per-swap gas cost synthesized auto-detected configurations are served with.
327    auto_detected_gas_cost: BigUint,
328    /// Venues whose components are emitted under the `propammfallback:` family, so their swaps
329    /// execute through Titan's PropAMMRouter instead of the venue directly.
330    router_venues: HashSet<Bytes>,
331    /// Components of the last emitted snapshot, across all pAMMs. A frame is a complete
332    /// snapshot of everything Titan currently streams, so removals are diffed globally: a
333    /// known component a frame does not re-emit is gone — including when its venue vanishes
334    /// from the stream entirely.
335    components: HashMap<String, ProtocolComponent>,
336    /// The newest block number processed so far. Frames targeting an older block (e.g.
337    /// delivered around a reconnect) are stale and skipped wholesale — processing one would
338    /// emit superseded states and churn the global diff.
339    newest_block: u64,
340}
341
342impl SnapshotTracker {
343    fn new(
344        registry: HashMap<Bytes, PriceLevelStreamConfig>,
345        denied: HashSet<Bytes>,
346        tokens: HashMap<Bytes, Token>,
347        auto_detect: bool,
348        auto_detected_gas_cost: BigUint,
349        router_venues: HashSet<Bytes>,
350    ) -> Self {
351        Self {
352            registry,
353            denied,
354            tokens,
355            auto_detect,
356            auto_detected_gas_cost,
357            router_venues,
358            components: HashMap::new(),
359            newest_block: 0,
360        }
361    }
362
363    /// Processes one frame into an [`Update`], or `None` if the frame targets an older block
364    /// than an already processed one or contains nothing relevant (no registered pAMM with at
365    /// least one known pair or a pair removal).
366    fn process(&mut self, message: TitanPriceLevelMessage) -> Option<Update> {
367        if message.block_number < self.newest_block {
368            tracing::warn!(
369                block_number = message.block_number,
370                newest_block = self.newest_block,
371                "Skipping out-of-order price level frame"
372            );
373            return None;
374        }
375        self.newest_block = message.block_number;
376
377        let mut states: HashMap<String, Box<dyn ProtocolSim>> = HashMap::new();
378        let mut new_pairs = HashMap::new();
379        // The frame is a complete snapshot: every known component is presumed gone until the
380        // frame re-emits it below.
381        let mut previous = std::mem::take(&mut self.components);
382
383        for TitanPammLevels { pamm, pairs } in message.pamms {
384            let config = match self.registry.entry(pamm.clone()) {
385                Entry::Occupied(entry) => &*entry.into_mut(),
386                Entry::Vacant(entry) => {
387                    if !self.auto_detect {
388                        tracing::debug!(%pamm, "Skipping unregistered pAMM");
389                        continue;
390                    }
391                    if self.denied.contains(&pamm) {
392                        tracing::debug!(%pamm, "Skipping denied pAMM");
393                        continue;
394                    }
395                    tracing::info!(%pamm, "Serving auto-detected pAMM");
396                    &*entry.insert(PriceLevelStreamConfig::auto_detected(
397                        pamm.clone(),
398                        self.auto_detected_gas_cost.clone(),
399                    ))
400                }
401            };
402
403            // Merge the frame's per-direction ladders into one entry per unordered token pair.
404            let mut merged_pairs: HashMap<(Bytes, Bytes), (Vec<_>, Vec<_>)> = HashMap::new();
405            for TitanPairLevels { token_in, token_out, order_book } in pairs {
406                if !self.tokens.contains_key(&token_in) || !self.tokens.contains_key(&token_out) {
407                    tracing::debug!(%token_in, %token_out, "Skipping pair with unknown token");
408                    continue;
409                }
410                let sells_token0 = token_in < token_out;
411                let key = if sells_token0 {
412                    (token_in.clone(), token_out.clone())
413                } else {
414                    (token_out.clone(), token_in.clone())
415                };
416                let quotes = order_book
417                    .into_iter()
418                    .map(|TitanPriceLevel { amount_in, amount_out }| {
419                        PriceLevelStreamQuote::new(amount_in, amount_out)
420                    })
421                    .collect();
422                let entry = merged_pairs.entry(key).or_default();
423                if sells_token0 {
424                    entry.0 = quotes;
425                } else {
426                    entry.1 = quotes;
427                }
428            }
429
430            for ((token0, token1), (quotes_0_to_1, quotes_1_to_0)) in merged_pairs {
431                let id = component_id(&config.address, &token0, &token1);
432                let id_string = id.to_string();
433                let component = previous
434                    .remove(&id_string)
435                    .unwrap_or_else(|| {
436                        let via_router = self
437                            .router_venues
438                            .contains(&config.address);
439                        let component =
440                            build_component(&self.tokens, config, id, &token0, &token1, via_router);
441                        new_pairs.insert(id_string.clone(), component.clone());
442                        component
443                    });
444
445                let state = PriceLevelStreamState::new(
446                    token0,
447                    token1,
448                    quotes_0_to_1,
449                    quotes_1_to_0,
450                    config.gas_cost.clone(),
451                );
452
453                states.insert(id_string.clone(), Box::new(state));
454                self.components
455                    .insert(id_string, component);
456            }
457        }
458
459        // Every re-emitted pair was moved back into `self.components` above — whatever remains
460        // is gone: the pair, or its whole venue, is no longer streamed.
461        let removed_pairs = previous;
462
463        if states.is_empty() && new_pairs.is_empty() && removed_pairs.is_empty() {
464            return None;
465        }
466
467        Some(
468            // Quotes target the block currently being built, hence partial. Sync states stay
469            // empty (like the RFQ path) because no full block header is available.
470            Update::new(message.block_number, states, new_pairs)
471                .set_is_partial(true)
472                .set_removed_pairs(removed_pairs),
473        )
474    }
475}
476
477fn build_component(
478    tokens: &HashMap<Bytes, Token>,
479    config: &PriceLevelStreamConfig,
480    id: Bytes,
481    token0: &Bytes,
482    token1: &Bytes,
483    via_router: bool,
484) -> ProtocolComponent {
485    let protocol_system =
486        if via_router { config.fallback_protocol_system() } else { config.protocol_system() };
487    ProtocolComponent::new(
488        id,
489        protocol_system.clone(),
490        protocol_system,
491        // Titan builds Ethereum L1 blocks; the stream carries no other chains.
492        Chain::Ethereum,
493        vec![tokens[token0].clone(), tokens[token1].clone()],
494        vec![config.address.clone()],
495        HashMap::from([(PAMM_ADDRESS_ATTRIBUTE.to_string(), config.address.clone())]),
496        Bytes::default(),
497        Utc::now().naive_utc(),
498    )
499}
500
501/// The component identity of a (pAMM, pair) combination: `pamm ++ token0 ++ token1`.
502fn component_id(pamm: &Bytes, token0: &Bytes, token1: &Bytes) -> Bytes {
503    Bytes::from([pamm.as_ref(), token0.as_ref(), token1.as_ref()].concat())
504}
505
506#[cfg(test)]
507mod tests {
508    use std::str::FromStr;
509
510    use num_bigint::BigUint;
511
512    use super::*;
513
514    const PAMM: &str = "0x5979458912f80b96d30d4220af8e2e4925a33320";
515    const WBTC: &str = "0x2260fac5e5542a773aa44fbcfedf7c193bc2c599";
516    const USDC: &str = "0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48";
517    const WETH: &str = "0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2";
518
519    fn token(address: &str, symbol: &str, decimals: u32) -> Token {
520        Token::new(
521            &Bytes::from_str(address).unwrap(),
522            symbol,
523            decimals,
524            0,
525            &[Some(10_000)],
526            Chain::Ethereum,
527            100,
528        )
529    }
530
531    fn tokens() -> HashMap<Bytes, Token> {
532        [token(WBTC, "WBTC", 8), token(USDC, "USDC", 6), token(WETH, "WETH", 18)]
533            .into_iter()
534            .map(|token| (token.address.clone(), token))
535            .collect()
536    }
537
538    fn tracker() -> SnapshotTracker {
539        let config = PriceLevelStreamConfig::new(
540            "fermiswap",
541            Bytes::from_str(PAMM).unwrap(),
542            BigUint::from(120_000u64),
543        );
544        SnapshotTracker::new(
545            HashMap::from([(config.address.clone(), config)]),
546            HashSet::new(),
547            tokens(),
548            false,
549            BigUint::from(DEFAULT_AUTO_DETECTED_GAS_COST),
550            HashSet::new(),
551        )
552    }
553
554    fn level(amount_in: u64, amount_out: u64) -> TitanPriceLevel {
555        TitanPriceLevel {
556            amount_in: BigUint::from(amount_in),
557            amount_out: BigUint::from(amount_out),
558        }
559    }
560
561    fn pair_levels(
562        token_in: &str,
563        token_out: &str,
564        order_book: Vec<TitanPriceLevel>,
565    ) -> TitanPairLevels {
566        TitanPairLevels {
567            token_in: Bytes::from_str(token_in).unwrap(),
568            token_out: Bytes::from_str(token_out).unwrap(),
569            order_book,
570        }
571    }
572
573    fn message(block_number: u64, pairs: Vec<TitanPairLevels>) -> TitanPriceLevelMessage {
574        TitanPriceLevelMessage {
575            block_number,
576            pamms: vec![TitanPammLevels { pamm: Bytes::from_str(PAMM).unwrap(), pairs }],
577        }
578    }
579
580    fn wbtc_usdc_pairs() -> Vec<TitanPairLevels> {
581        vec![
582            pair_levels(WBTC, USDC, vec![level(100_000_000, 100_000_000_000)]),
583            pair_levels(USDC, WBTC, vec![level(100_000_000_000, 99_000_000)]),
584        ]
585    }
586
587    fn expected_id() -> String {
588        // pamm ++ token0 ++ token1 with WBTC < USDC.
589        format!("{PAMM}{}{}", &WBTC[2..], &USDC[2..])
590    }
591
592    #[test]
593    fn first_snapshot_emits_new_pair_with_both_directions() {
594        let mut tracker = tracker();
595        let Update {
596            block_number_or_timestamp,
597            is_partial,
598            sync_states,
599            states,
600            new_pairs,
601            removed_pairs,
602        } = tracker
603            .process(message(100, wbtc_usdc_pairs()))
604            .expect("update expected");
605
606        assert_eq!(block_number_or_timestamp, 100);
607        assert!(is_partial);
608        assert!(sync_states.is_empty());
609        assert!(removed_pairs.is_empty());
610
611        let id = expected_id();
612        let component = &new_pairs[&id];
613        assert_eq!(component.protocol_system, "pricelevelstream:fermiswap");
614        assert_eq!(
615            component.static_attributes[PAMM_ADDRESS_ATTRIBUTE],
616            Bytes::from_str(PAMM).unwrap()
617        );
618
619        let PriceLevelStreamState { token0, token1, quotes_0_to_1, quotes_1_to_0, gas_cost } =
620            states[&id]
621                .as_any()
622                .downcast_ref::<PriceLevelStreamState>()
623                .expect("price level state");
624        assert_eq!(token0, &Bytes::from_str(WBTC).unwrap());
625        assert_eq!(token1, &Bytes::from_str(USDC).unwrap());
626        assert_eq!(quotes_0_to_1.len(), 1);
627        assert_eq!(quotes_1_to_0.len(), 1);
628        assert_eq!(quotes_0_to_1[0].amount_in, BigUint::from(100_000_000u64));
629        assert_eq!(gas_cost, &BigUint::from(120_000u64));
630    }
631
632    #[test]
633    fn repeated_snapshot_is_not_a_new_pair() {
634        let mut tracker = tracker();
635        tracker
636            .process(message(100, wbtc_usdc_pairs()))
637            .expect("update expected");
638        let update = tracker
639            .process(message(101, wbtc_usdc_pairs()))
640            .expect("update expected");
641
642        assert!(update.new_pairs.is_empty());
643        assert!(update.removed_pairs.is_empty());
644        assert!(update
645            .states
646            .contains_key(&expected_id()));
647    }
648
649    #[test]
650    fn dropped_pair_is_removed() {
651        let mut tracker = tracker();
652        tracker
653            .process(message(100, wbtc_usdc_pairs()))
654            .expect("update expected");
655        let weth_usdc =
656            vec![pair_levels(WETH, USDC, vec![level(1_000_000_000_000_000_000, 3_000_000_000)])];
657        let update = tracker
658            .process(message(101, weth_usdc))
659            .expect("update expected");
660
661        assert_eq!(update.removed_pairs.len(), 1);
662        assert!(update
663            .removed_pairs
664            .contains_key(&expected_id()));
665        assert_eq!(update.new_pairs.len(), 1);
666        assert_eq!(update.states.len(), 1);
667    }
668
669    #[test]
670    fn out_of_order_frame_is_skipped() {
671        let mut tracker = tracker();
672        tracker
673            .process(message(101, wbtc_usdc_pairs()))
674            .expect("update expected");
675
676        // A frame for an older block is stale: no update, and the caches stay untouched even
677        // though the frame's snapshot differs completely.
678        let stale =
679            vec![pair_levels(WETH, USDC, vec![level(1_000_000_000_000_000_000, 3_000_000_000)])];
680        assert!(tracker
681            .process(message(100, stale))
682            .is_none());
683
684        // The next current frame diffs against the pre-stale state: nothing was added or
685        // removed in between.
686        let update = tracker
687            .process(message(102, wbtc_usdc_pairs()))
688            .expect("update expected");
689        assert!(update.new_pairs.is_empty());
690        assert!(update.removed_pairs.is_empty());
691    }
692
693    #[test]
694    fn vanished_pamm_has_its_pairs_removed() {
695        let mut tracker = tracker();
696        tracker
697            .process(message(100, wbtc_usdc_pairs()))
698            .expect("update expected");
699
700        // The next frame no longer contains the pAMM at all: a complete snapshot without a
701        // venue means the venue is gone, pairs and all.
702        let update = tracker
703            .process(TitanPriceLevelMessage { block_number: 101, pamms: vec![] })
704            .expect("update expected");
705        assert!(update.states.is_empty());
706        assert!(update.new_pairs.is_empty());
707        assert_eq!(update.removed_pairs.len(), 1);
708        assert!(update
709            .removed_pairs
710            .contains_key(&expected_id()));
711
712        // Nothing served and nothing changed: no update.
713        assert!(tracker
714            .process(TitanPriceLevelMessage { block_number: 102, pamms: vec![] })
715            .is_none());
716
717        // A venue that reappears is a new pair again.
718        let update = tracker
719            .process(message(103, wbtc_usdc_pairs()))
720            .expect("update expected");
721        assert!(update
722            .new_pairs
723            .contains_key(&expected_id()));
724    }
725
726    #[test]
727    fn unregistered_pamm_produces_no_update_without_auto_detection() {
728        let mut tracker = SnapshotTracker::new(
729            HashMap::new(),
730            HashSet::new(),
731            tokens(),
732            false,
733            BigUint::from(DEFAULT_AUTO_DETECTED_GAS_COST),
734            HashSet::new(),
735        );
736        assert!(tracker
737            .process(message(100, wbtc_usdc_pairs()))
738            .is_none());
739    }
740
741    #[test]
742    fn denied_pamm_is_not_auto_detected() {
743        let denied = HashSet::from([Bytes::from_str(PAMM).unwrap()]);
744        let mut tracker = SnapshotTracker::new(
745            HashMap::new(),
746            denied,
747            tokens(),
748            true,
749            BigUint::from(DEFAULT_AUTO_DETECTED_GAS_COST),
750            HashSet::new(),
751        );
752        assert!(tracker
753            .process(message(100, wbtc_usdc_pairs()))
754            .is_none());
755    }
756
757    #[test]
758    fn explicit_add_and_deny_are_last_wins() {
759        let address = Bytes::from_str(PAMM).unwrap();
760        let custom =
761            || PriceLevelStreamConfig::new("custom", Bytes::from_str(PAMM).unwrap(), 1u64.into());
762
763        let builder = PriceLevelStreamBuilder::new()
764            .add_pamm(custom())
765            .deny_pamm(address.clone());
766        assert!(!builder.registry.contains_key(&address));
767        assert!(builder.denied.contains(&address));
768
769        let builder = PriceLevelStreamBuilder::new()
770            .deny_pamm(address.clone())
771            .add_pamm(custom());
772        assert_eq!(builder.registry[&address].protocol, "custom");
773        assert!(builder.denied.is_empty());
774    }
775
776    #[test]
777    fn defaults_never_override_explicit_calls() {
778        // Denying a venue from the default set works in either call order.
779        let fermiswap_router = Bytes::from_str(PAMM).unwrap();
780        for builder in [
781            PriceLevelStreamBuilder::new()
782                .deny_pamm(fermiswap_router.clone())
783                .with_known_pamms(),
784            PriceLevelStreamBuilder::new()
785                .with_known_pamms()
786                .deny_pamm(fermiswap_router.clone()),
787        ] {
788            assert!(!builder
789                .registry
790                .contains_key(&fermiswap_router));
791            assert!(builder
792                .denied
793                .contains(&fermiswap_router));
794            // The other defaults are unaffected.
795            assert!(!builder.registry.is_empty());
796        }
797
798        // Registering a venue from the default deny set works in either call order.
799        let denied_venue = default_denied_pamms().remove(0);
800        let custom = || PriceLevelStreamConfig::new("custom", denied_venue.clone(), 1u64.into());
801        for builder in [
802            PriceLevelStreamBuilder::new()
803                .add_pamm(custom())
804                .with_known_pamms(),
805            PriceLevelStreamBuilder::new()
806                .with_known_pamms()
807                .add_pamm(custom()),
808        ] {
809            assert_eq!(builder.registry[&denied_venue].protocol, "custom");
810            assert!(!builder.denied.contains(&denied_venue));
811        }
812    }
813
814    #[test]
815    fn auto_detected_pamm_is_served_under_its_address() {
816        let mut tracker = SnapshotTracker::new(
817            HashMap::new(),
818            HashSet::new(),
819            tokens(),
820            true,
821            BigUint::from(DEFAULT_AUTO_DETECTED_GAS_COST),
822            HashSet::new(),
823        );
824        let update = tracker
825            .process(message(100, wbtc_usdc_pairs()))
826            .expect("update expected");
827
828        let component = &update.new_pairs[&expected_id()];
829        assert_eq!(component.protocol_system, format!("pricelevelstream:{PAMM}"));
830        let state = update.states[&expected_id()]
831            .as_any()
832            .downcast_ref::<PriceLevelStreamState>()
833            .expect("price level state");
834        assert_eq!(state.gas_cost, BigUint::from(DEFAULT_AUTO_DETECTED_GAS_COST));
835
836        // The synthesized config is cached: the next snapshot is not a new pair again.
837        let update = tracker
838            .process(message(101, wbtc_usdc_pairs()))
839            .expect("update expected");
840        assert!(update.new_pairs.is_empty());
841    }
842
843    #[test]
844    fn auto_detected_gas_cost_override_applies() {
845        let mut tracker = SnapshotTracker::new(
846            HashMap::new(),
847            HashSet::new(),
848            tokens(),
849            true,
850            BigUint::from(42_000u64),
851            HashSet::new(),
852        );
853        let update = tracker
854            .process(message(100, wbtc_usdc_pairs()))
855            .expect("update expected");
856
857        let state = update.states[&expected_id()]
858            .as_any()
859            .downcast_ref::<PriceLevelStreamState>()
860            .expect("price level state");
861        assert_eq!(state.gas_cost, BigUint::from(42_000u64));
862    }
863
864    #[test]
865    fn with_known_pamms_registers_known_venues() {
866        // PAMM is the FermiSwap router, one of the default venues.
867        let fermiswap_router = Bytes::from_str(PAMM).unwrap();
868
869        let builder = PriceLevelStreamBuilder::new();
870        assert!(builder.registry.is_empty());
871        assert!(builder.denied.is_empty());
872
873        let builder = builder.with_known_pamms();
874        assert_eq!(builder.registry[&fermiswap_router].protocol, "fermiswap");
875        // The known-bad venues get denied alongside, and never overlap the served defaults.
876        assert!(!builder.denied.is_empty());
877        assert!(builder.denied.is_disjoint(
878            &builder
879                .registry
880                .keys()
881                .cloned()
882                .collect()
883        ));
884
885        // An `add_pamm` entry wins over the default for the same address, in either call order.
886        let custom =
887            || PriceLevelStreamConfig::new("custom", fermiswap_router.clone(), BigUint::from(1u64));
888        for builder in [
889            PriceLevelStreamBuilder::new()
890                .add_pamm(custom())
891                .with_known_pamms(),
892            PriceLevelStreamBuilder::new()
893                .with_known_pamms()
894                .add_pamm(custom()),
895        ] {
896            assert_eq!(builder.registry[&fermiswap_router].protocol, "custom");
897            assert_eq!(builder.registry[&fermiswap_router].gas_cost, BigUint::from(1u64));
898        }
899    }
900
901    /// A venue on the router's whitelist is emitted under `propammfallback:{name}`, so its swaps
902    /// execute through Titan's PropAMMRouter; identity and attributes stay the same.
903    #[test]
904    fn whitelisted_venue_is_served_under_the_fallback_family() {
905        let config = PriceLevelStreamConfig::new(
906            "fermiswap",
907            Bytes::from_str(PAMM).unwrap(),
908            BigUint::from(120_000u64),
909        );
910        let mut tracker = SnapshotTracker::new(
911            HashMap::from([(config.address.clone(), config)]),
912            HashSet::new(),
913            tokens(),
914            false,
915            BigUint::from(DEFAULT_AUTO_DETECTED_GAS_COST),
916            HashSet::from([Bytes::from_str(PAMM).unwrap()]),
917        );
918
919        let update = tracker
920            .process(message(100, wbtc_usdc_pairs()))
921            .expect("update expected");
922
923        let component = &update.new_pairs[&expected_id()];
924        assert_eq!(component.protocol_system, "propammfallback:fermiswap");
925        assert_eq!(
926            component.static_attributes[PAMM_ADDRESS_ATTRIBUTE],
927            Bytes::from_str(PAMM).unwrap()
928        );
929    }
930
931    /// The whitelist check is by address, so it also covers auto-detected, address-named venues.
932    #[test]
933    fn auto_detected_whitelisted_venue_is_served_under_the_fallback_family() {
934        let mut tracker = SnapshotTracker::new(
935            HashMap::new(),
936            HashSet::new(),
937            tokens(),
938            true,
939            BigUint::from(DEFAULT_AUTO_DETECTED_GAS_COST),
940            HashSet::from([Bytes::from_str(PAMM).unwrap()]),
941        );
942
943        let update = tracker
944            .process(message(100, wbtc_usdc_pairs()))
945            .expect("update expected");
946
947        let component = &update.new_pairs[&expected_id()];
948        assert_eq!(component.protocol_system, format!("propammfallback:{PAMM}"));
949    }
950
951    /// A venue absent from the whitelist keeps the direct `pricelevelstream:{name}` family — the
952    /// router reverts `UnknownVenue` for it, which would send every swap to the Uniswap V3
953    /// fallback.
954    #[test]
955    fn unwhitelisted_venue_keeps_the_direct_family() {
956        let other_venue =
957            Bytes::from_str("0x71e790dd841c8a9061487cb3e78c288e75ce0b3d").expect("valid address");
958        let config = PriceLevelStreamConfig::new(
959            "fermiswap",
960            Bytes::from_str(PAMM).unwrap(),
961            BigUint::from(120_000u64),
962        );
963        let mut tracker = SnapshotTracker::new(
964            HashMap::from([(config.address.clone(), config)]),
965            HashSet::new(),
966            tokens(),
967            false,
968            BigUint::from(DEFAULT_AUTO_DETECTED_GAS_COST),
969            HashSet::from([other_venue]),
970        );
971
972        let update = tracker
973            .process(message(100, wbtc_usdc_pairs()))
974            .expect("update expected");
975
976        assert_eq!(update.new_pairs[&expected_id()].protocol_system, "pricelevelstream:fermiswap");
977    }
978
979    /// The PropAMMRouter path is the default; `without_fallback_router` is the way off it.
980    #[test]
981    fn fallback_router_is_on_unless_opted_out() {
982        assert!(PriceLevelStreamBuilder::new().fallback_router);
983        assert!(
984            !PriceLevelStreamBuilder::new()
985                .without_fallback_router()
986                .fallback_router
987        );
988    }
989
990    /// Without a node URL there is nothing to read the whitelist from: every venue stays on the
991    /// direct path instead of the build failing.
992    #[tokio::test]
993    async fn missing_rpc_url_leaves_every_venue_on_the_direct_path() {
994        assert!(fetch_router_venues(None)
995            .await
996            .is_empty());
997    }
998
999    /// A failed whitelist read degrades the same way as a missing node URL.
1000    #[tokio::test]
1001    async fn failed_whitelist_read_leaves_every_venue_on_the_direct_path() {
1002        assert!(fetch_router_venues(Some("not a url".to_string()))
1003            .await
1004            .is_empty());
1005    }
1006
1007    /// The families this stream emits are the ones tycho-execution resolves an encoder for. A
1008    /// drift between the two makes every route through a pAMM fail to encode.
1009    #[test]
1010    fn families_match_the_execution_side_prefixes() {
1011        use tycho_execution::encoding::evm::{PRICE_LEVEL_STREAM_PREFIX, PROPAMM_FALLBACK_PREFIX};
1012
1013        use super::super::config::{PRICE_LEVEL_STREAM_FAMILY, PROPAMM_FALLBACK_FAMILY};
1014
1015        assert_eq!(format!("{PRICE_LEVEL_STREAM_FAMILY}:"), PRICE_LEVEL_STREAM_PREFIX);
1016        assert_eq!(format!("{PROPAMM_FALLBACK_FAMILY}:"), PROPAMM_FALLBACK_PREFIX);
1017    }
1018
1019    #[test]
1020    fn unknown_tokens_are_skipped() {
1021        let mut tracker = tracker();
1022        let unknown = vec![pair_levels(
1023            "0x1111111111111111111111111111111111111111",
1024            USDC,
1025            vec![level(1, 1)],
1026        )];
1027        assert!(tracker
1028            .process(message(100, unknown))
1029            .is_none());
1030    }
1031}