use std::{collections::HashMap, sync::Arc, time::Duration};
use alloy::primitives::{Address, U256};
use futures::StreamExt;
use serde::Deserialize;
use tokio::{
sync::watch,
time::{sleep, timeout},
};
use tokio_tungstenite::{connect_async, tungstenite::Message};
use tracing::{info, warn};
use super::{FailurePolicy, OverrideSnapshot, StateOverrideProvider};
type VenueOverrides = HashMap<Address, HashMap<U256, U256>>;
#[derive(Debug, Deserialize)]
struct TitanMessage {
#[serde(default)]
slot: Option<u64>,
#[serde(rename = "blockNumber", default)]
block_number: Option<u64>,
#[serde(default)]
timestamp: Option<u64>,
#[serde(flatten)]
overrides: HashMap<String, serde_json::Value>,
}
#[derive(Debug, Deserialize)]
struct AddressEntry {
#[serde(rename = "stateOverride", default)]
state_override: HashMap<Address, AccountState>,
}
#[derive(Debug, Deserialize)]
struct AccountState {
#[serde(rename = "stateDiff", default)]
state_diff: HashMap<U256, U256>,
}
const TITAN_URL: &str = "wss://eu.rpc.titanbuilder.xyz/ws/pamm_quote_stream";
const TITAN_URL_ENV: &str = "TITAN_PAMM_STREAM_URL";
const FERMISWAP_PROTOCOL_SYSTEM: &str = "vm:fermiswap";
const KIPSELI_PROTOCOL_SYSTEM: &str = "vm:kipseli";
const BOPAMM_PROTOCOL_SYSTEM: &str = "vm:bopamm";
const FERMISWAP_VENUES: &[&str] =
&["0xb1076fe3ab5e28005c7c323bac5ac06a680d452e", "0x5979458912f80b96d30d4220af8e2e4925a33320"];
const KIPSELI_VENUES: &[&str] =
&["0x5cdbe59400cc2efdcc2b54acca4a99fe00dd588c", "0x71e790dd841c8a9061487cb3e78c288e75ce0b3d"];
const BOPAMM_VENUES: &[&str] = &["0x160141a205f5ddcf096ba3f48b7ed21eb52c62ea"];
const READ_IDLE_TIMEOUT: Duration = Duration::from_secs(30);
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
const MAX_BACKOFF_EXP: u32 = 5;
const TITAN_QUOTE_TTL_SECS: u64 = 12;
const BEACON_GENESIS_TIMESTAMP: u64 = 1_606_824_023;
const SECONDS_PER_SLOT: u64 = 12;
const MAX_LANE_TIMESTAMP_LEAD_SECS: u64 = 2 * SECONDS_PER_SLOT;
fn known_pamm_protocols() -> &'static [&'static str] {
&[FERMISWAP_PROTOCOL_SYSTEM, KIPSELI_PROTOCOL_SYSTEM, BOPAMM_PROTOCOL_SYSTEM]
}
pub fn default_providers(
protocols: impl IntoIterator<Item = String>,
) -> HashMap<String, Arc<dyn StateOverrideProvider>> {
let needed: Vec<&'static str> = protocols
.into_iter()
.filter_map(|protocol| {
known_pamm_protocols()
.iter()
.copied()
.find(|&known| known == protocol.as_str())
})
.collect();
if needed.is_empty() {
return HashMap::new();
}
let url = std::env::var(TITAN_URL_ENV).unwrap_or_else(|_| TITAN_URL.to_string());
let provider: Arc<dyn StateOverrideProvider> = Arc::new(TitanProvider::spawn(url, &needed));
needed
.into_iter()
.map(|protocol| (protocol.to_string(), provider.clone()))
.collect()
}
pub struct TitanProvider {
receivers: HashMap<String, watch::Receiver<OverrideSnapshot>>,
}
impl TitanProvider {
pub fn spawn(url: String, protocols: &[&str]) -> Self {
let mut receivers = HashMap::new();
let mut senders = Vec::new();
for &protocol in protocols {
let Some(venues) = Self::venues_for_protocol(protocol) else {
warn!("Unknown protocol: {protocol}");
continue;
};
let (tx, rx) = watch::channel(OverrideSnapshot::default());
receivers.insert(protocol.to_string(), rx);
senders.push((venues, tx));
}
tokio::spawn(Self::run(url, senders));
Self { receivers }
}
async fn run(url: String, senders: Vec<(Vec<Address>, watch::Sender<OverrideSnapshot>)>) {
let mut attempt: u32 = 0;
loop {
if senders
.iter()
.all(|(_, tx)| tx.is_closed())
{
warn!("All Titan override receivers dropped; stopping quote-stream task");
return;
}
match timeout(CONNECT_TIMEOUT, connect_async(url.as_str())).await {
Ok(Ok((mut ws_stream, _))) => {
info!(%url, "Connected to Titan pAMM quote stream");
loop {
let message = match timeout(READ_IDLE_TIMEOUT, ws_stream.next()).await {
Ok(Some(message)) => message,
Ok(None) => {
warn!("Titan quote stream ended; reconnecting");
break;
}
Err(_elapsed) => {
warn!(
idle_secs = READ_IDLE_TIMEOUT.as_secs(),
"No Titan message within idle timeout; reconnecting"
);
break;
}
};
match message {
Ok(Message::Text(text)) => {
attempt = 0;
match serde_json::from_str::<TitanMessage>(text.as_str()) {
Ok(message) => {
for (venues, sender) in &senders {
match Self::extract_any_venue(&message, venues) {
Ok(Some(snapshot)) => {
let _ = sender.send(snapshot);
}
Ok(None) => {}
Err(e) => {
warn!(?venues, error = %e, "Failed to extract Titan venue snapshot");
}
}
}
}
Err(e) => warn!(error = %e, "Failed to parse Titan message"),
}
if senders
.iter()
.all(|(_, tx)| tx.is_closed())
{
warn!("All Titan override receivers dropped; stopping quote-stream task");
return;
}
}
Ok(Message::Binary(bytes)) => {
warn!(len = bytes.len(), "Ignoring unexpected binary Titan frame");
}
Ok(Message::Ping(_)) | Ok(Message::Pong(_)) => {}
Ok(Message::Frame(_)) => {}
Ok(Message::Close(frame)) => {
warn!(?frame, "Titan quote stream closed by server; reconnecting");
break;
}
Err(e) => {
warn!(error = %e, "Titan quote stream read error; reconnecting");
break;
}
}
}
}
Ok(Err(e)) => {
warn!(error = %e, "Failed to connect to Titan quote stream; retrying");
}
Err(_elapsed) => {
warn!(
timeout_secs = CONNECT_TIMEOUT.as_secs(),
"Titan connect timed out; retrying"
);
}
}
attempt = attempt.saturating_add(1);
let backoff = Duration::from_secs(2u64.pow(attempt.min(MAX_BACKOFF_EXP)));
warn!(seconds = backoff.as_secs(), attempt, "Backing off before reconnecting to Titan");
sleep(backoff).await;
}
}
fn venues_for_protocol(protocol_system: &str) -> Option<Vec<Address>> {
let raw = match protocol_system {
FERMISWAP_PROTOCOL_SYSTEM => FERMISWAP_VENUES,
KIPSELI_PROTOCOL_SYSTEM => KIPSELI_VENUES,
BOPAMM_PROTOCOL_SYSTEM => BOPAMM_VENUES,
_ => return None,
};
let venues = raw
.iter()
.map(|venue| {
venue
.parse()
.unwrap_or_else(|e| panic!("hardcoded venue address {venue} must parse: {e}"))
})
.collect();
Some(venues)
}
fn extract_any_venue(
message: &TitanMessage,
venues: &[Address],
) -> Result<Option<OverrideSnapshot>, String> {
for venue in venues {
if let Some(snapshot) = Self::extract_venue(message, venue)? {
return Ok(Some(snapshot));
}
}
Ok(None)
}
fn extract_venue(
message: &TitanMessage,
venue: &Address,
) -> Result<Option<OverrideSnapshot>, String> {
let venue_key = format!("{venue:#x}");
let Some(raw) = message.overrides.get(&venue_key) else {
return Ok(None);
};
let entry: AddressEntry = serde_json::from_value(raw.clone())
.map_err(|e| format!("invalid venue override: {e}"))?;
let block_number = message.block_number;
let mut storage: VenueOverrides = HashMap::new();
for (account, account_override) in entry.state_override {
if !account_override.state_diff.is_empty() {
storage.insert(account, account_override.state_diff);
}
}
let block_timestamp = Self::pending_block_timestamp(message, &storage);
let expires_at = message
.timestamp
.map(|ns| ns / 1_000_000_000 + TITAN_QUOTE_TTL_SECS);
Ok(Some(OverrideSnapshot {
block_number,
block_timestamp,
storage: Arc::new(storage),
expires_at,
failure_policy: FailurePolicy::FallbackToIndexedState,
}))
}
fn pending_block_timestamp(message: &TitanMessage, storage: &VenueOverrides) -> Option<u64> {
let wall_clock = message
.timestamp
.map(|ns| ns / 1_000_000_000);
if let Some(slot_ts) = Self::slot_timestamp(message.slot, wall_clock) {
return Some(slot_ts);
}
match (Self::max_lane_timestamp(storage), wall_clock) {
(Some(lane), Some(now))
if (now..=now + MAX_LANE_TIMESTAMP_LEAD_SECS).contains(&lane) =>
{
Some(lane)
}
(Some(_), wall_clock) | (None, wall_clock) => wall_clock,
}
}
fn slot_timestamp(slot: Option<u64>, wall_clock: Option<u64>) -> Option<u64> {
let derived = slot?
.checked_mul(SECONDS_PER_SLOT)
.and_then(|offset| BEACON_GENESIS_TIMESTAMP.checked_add(offset))?;
match wall_clock {
Some(now)
if !(now.saturating_sub(SECONDS_PER_SLOT)..=
now.saturating_add(MAX_LANE_TIMESTAMP_LEAD_SECS))
.contains(&derived) =>
{
None
}
Some(_) | None => Some(derived),
}
}
fn max_lane_timestamp(overrides: &VenueOverrides) -> Option<u64> {
const MIN_PLAUSIBLE: u64 = 1_000_000_000; const MAX_PLAUSIBLE: u64 = u32::MAX as u64; overrides
.values()
.flat_map(|slots| slots.values())
.filter_map(|value| {
let candidate = u64::try_from(*value >> 224).ok()?;
(MIN_PLAUSIBLE..=MAX_PLAUSIBLE)
.contains(&candidate)
.then_some(candidate)
})
.max()
}
}
impl StateOverrideProvider for TitanProvider {
fn subscribe(&self, protocol_system: &str) -> Option<watch::Receiver<OverrideSnapshot>> {
self.receivers
.get(protocol_system)
.cloned()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn known_pamm_protocols_contains_fermiswap() {
assert!(known_pamm_protocols().contains(&FERMISWAP_PROTOCOL_SYSTEM));
}
#[test]
fn venues_for_protocol_resolves_all_known_aliases() {
let cases = [
(FERMISWAP_PROTOCOL_SYSTEM, FERMISWAP_VENUES),
(KIPSELI_PROTOCOL_SYSTEM, KIPSELI_VENUES),
(BOPAMM_PROTOCOL_SYSTEM, BOPAMM_VENUES),
];
for (protocol, expected) in cases {
let venues = TitanProvider::venues_for_protocol(protocol).expect("venues must resolve");
let expected: Vec<Address> = expected
.iter()
.map(|raw| raw.parse().unwrap())
.collect();
assert_eq!(venues, expected, "venue mismatch for {protocol}");
assert!(!venues.is_empty(), "every protocol needs at least one venue");
}
}
#[test]
fn venues_for_protocol_returns_none_for_unknown() {
assert!(TitanProvider::venues_for_protocol("vm:unknown").is_none());
}
const SAMPLE_MESSAGE: &str = r#"{
"slot": 14285824,
"blockNumber": 25051224,
"timestamp": 1778253913749564761,
"0xb1076fe3ab5e28005c7c323bac5ac06a680d452e": {
"stateOverride": {
"0x1038c87766e36d1925889e6f26d10e0012d50fed": {
"balance": "0x0",
"nonce": "0x1",
"stateDiff": {
"0x156b1d71de08fed89d0fce38008e2b9d03a8998077e394b20597ef3d148f5ebc": "0x000000000000000000000000000000000000000000000000000000017e405801"
}
}
}
}
}"#;
#[test]
fn extract_venue_reads_state_diff() {
let venue = FERMISWAP_VENUES[0]
.parse::<Address>()
.unwrap();
let message = serde_json::from_str::<TitanMessage>(SAMPLE_MESSAGE).expect("valid JSON");
let snapshot = TitanProvider::extract_venue(&message, &venue)
.expect("extract must succeed")
.expect("venue is present");
assert_eq!(snapshot.block_number, Some(25051224));
assert_eq!(snapshot.block_timestamp, Some(1778253911));
assert_eq!(snapshot.expires_at, Some(1778253913 + TITAN_QUOTE_TTL_SECS));
assert_eq!(snapshot.failure_policy, FailurePolicy::FallbackToIndexedState);
let account = "0x1038c87766e36d1925889e6f26d10e0012d50fed"
.parse::<Address>()
.unwrap();
let slot = "0x156b1d71de08fed89d0fce38008e2b9d03a8998077e394b20597ef3d148f5ebc"
.parse::<U256>()
.unwrap();
let value = "0x000000000000000000000000000000000000000000000000000000017e405801"
.parse::<U256>()
.unwrap();
assert_eq!(
snapshot
.storage
.get(&account)
.and_then(|s| s.get(&slot)),
Some(&value)
);
}
#[test]
fn extract_venue_returns_none_for_absent_venue() {
let venue = KIPSELI_VENUES[0]
.parse::<Address>()
.unwrap();
let message = serde_json::from_str::<TitanMessage>(SAMPLE_MESSAGE).expect("valid JSON");
assert!(TitanProvider::extract_venue(&message, &venue)
.unwrap()
.is_none());
}
#[test]
fn extract_any_venue_falls_back_to_a_later_alias() {
let message = serde_json::from_str::<TitanMessage>(SAMPLE_MESSAGE).expect("valid JSON");
let aliases = [
FERMISWAP_VENUES[1]
.parse::<Address>()
.unwrap(),
FERMISWAP_VENUES[0]
.parse::<Address>()
.unwrap(),
];
let snapshot = TitanProvider::extract_any_venue(&message, &aliases)
.expect("extract must succeed")
.expect("oracle alias is present in the frame");
assert_eq!(snapshot.block_number, Some(25051224));
}
#[test]
fn extract_any_venue_returns_none_when_no_alias_matches() {
let message = serde_json::from_str::<TitanMessage>(SAMPLE_MESSAGE).expect("valid JSON");
let aliases: Vec<Address> = KIPSELI_VENUES
.iter()
.map(|raw| raw.parse().unwrap())
.collect();
assert!(TitanProvider::extract_any_venue(&message, &aliases)
.unwrap()
.is_none());
}
fn message(slot: Option<u64>, timestamp: Option<u64>) -> TitanMessage {
TitanMessage { slot, block_number: Some(25445022), timestamp, overrides: HashMap::new() }
}
fn lane_overrides(lane_ts: u64) -> VenueOverrides {
let account = "0xda7afeed01fe625cf15d187a19f94b45f00b8c5f"
.parse::<Address>()
.unwrap();
HashMap::from([(account, HashMap::from([(U256::from(1u64), U256::from(lane_ts) << 224)]))])
}
const WALL_CLOCK_NS: u64 = 1782997641189454379;
const PENDING_BLOCK_TS: u64 = 1782997643;
#[test]
fn pending_block_timestamp_derives_from_beacon_slot() {
let message = message(Some(14681135), Some(WALL_CLOCK_NS));
assert_eq!(
TitanProvider::pending_block_timestamp(&message, &lane_overrides(1_800_000_000)),
Some(PENDING_BLOCK_TS)
);
}
#[test]
fn pending_block_timestamp_rejects_implausible_slot() {
let message = message(Some(1), Some(WALL_CLOCK_NS));
assert_eq!(
TitanProvider::pending_block_timestamp(&message, &HashMap::new()),
Some(1782997641)
);
}
#[test]
fn pending_block_timestamp_survives_overflowing_slot() {
let message = message(Some(u64::MAX), Some(WALL_CLOCK_NS));
assert_eq!(
TitanProvider::pending_block_timestamp(&message, &HashMap::new()),
Some(1782997641)
);
}
#[test]
fn pending_block_timestamp_ignores_lane_without_wall_clock() {
let message = message(None, None);
assert!(TitanProvider::pending_block_timestamp(
&message,
&lane_overrides(PENDING_BLOCK_TS)
)
.is_none());
}
#[test]
fn pending_block_timestamp_backs_up_to_lane_timestamps() {
let message = message(None, Some(WALL_CLOCK_NS));
assert_eq!(
TitanProvider::pending_block_timestamp(&message, &lane_overrides(PENDING_BLOCK_TS)),
Some(PENDING_BLOCK_TS)
);
}
#[test]
fn pending_block_timestamp_rejects_implausible_lane_timestamps() {
let message = message(None, Some(WALL_CLOCK_NS));
assert_eq!(
TitanProvider::pending_block_timestamp(&message, &lane_overrides(1_557_751_927)),
Some(1782997641)
);
}
#[test]
fn pending_block_timestamp_falls_back_to_wall_clock() {
let message = message(None, Some(WALL_CLOCK_NS));
assert_eq!(
TitanProvider::pending_block_timestamp(&message, &HashMap::new()),
Some(1782997641)
);
}
#[test]
fn pending_block_timestamp_none_without_any_signal() {
let message = message(None, None);
assert!(TitanProvider::pending_block_timestamp(&message, &HashMap::new()).is_none());
}
}