Skip to main content

paygress/
discovery.rs

1// Discovery Client
2//
3// Used by end users to discover available providers on Nostr
4// and interact with them for spawning workloads.
5
6use anyhow::{Context, Result};
7use serde::{Deserialize, Serialize};
8use tracing::{info, warn};
9
10use crate::nostr::{
11    NostrRelaySubscriber, RelayConfig, ProviderOfferContent, ProviderInfo, 
12    ProviderFilter, PodSpec,
13};
14
15/// Discovery client for finding providers
16pub struct DiscoveryClient {
17    nostr: NostrRelaySubscriber,
18}
19
20impl DiscoveryClient {
21    /// Create a new discovery client
22    pub async fn new(relays: Vec<String>) -> Result<Self> {
23        let config = RelayConfig {
24            relays,
25            private_key: None, // Read-only client doesn't need a key
26        };
27        
28        let nostr = NostrRelaySubscriber::new(config).await?;
29        
30        Ok(Self { nostr })
31    }
32
33    /// Create with a private key (for sending spawn requests)
34    pub async fn new_with_key(relays: Vec<String>, private_key: String) -> Result<Self> {
35        let config = RelayConfig {
36            relays,
37            private_key: Some(private_key),
38        };
39        
40        let nostr = NostrRelaySubscriber::new(config).await?;
41        
42        Ok(Self { nostr })
43    }
44
45    /// Get the client's public key (npub)
46    pub fn get_npub(&self) -> String {
47        self.nostr.get_service_public_key()
48    }
49
50    /// List all available providers
51    pub async fn list_providers(&self, filter: Option<ProviderFilter>) -> Result<Vec<ProviderInfo>> {
52        let offers = self.nostr.query_providers().await?;
53        
54        let mut providers = Vec::new();
55        
56        // Optimisation: Fetch all heartbeats in parallel (batch query)
57        let provider_npubs: Vec<String> = offers.iter().map(|o| o.provider_npub.clone()).collect();
58        let heartbeats = self.nostr.get_latest_heartbeats_multi(provider_npubs).await?;
59        
60        for offer in offers {
61            // Check if provider is online (has recent heartbeat)
62            let (is_online, last_seen) = match heartbeats.get(&offer.provider_npub) {
63                Some(hb) => {
64                    let now = std::time::SystemTime::now()
65                        .duration_since(std::time::UNIX_EPOCH)
66                        .map(|d| d.as_secs())
67                        .unwrap_or(0);
68                    // Consider online if heartbeat within last 2 minutes
69                    (now - hb.timestamp < 120, hb.timestamp)
70                }
71                None => (false, 0),
72            };
73
74
75
76            let provider = ProviderInfo {
77                npub: offer.provider_npub.clone(),
78                hostname: offer.hostname,
79                location: offer.location,
80                capabilities: offer.capabilities,
81                specs: offer.specs,
82                whitelisted_mints: offer.whitelisted_mints,
83                uptime_percent: offer.uptime_percent,
84                total_jobs_completed: offer.total_jobs_completed,
85                last_seen,
86                is_online,
87            };
88
89            // Apply filters
90            if let Some(ref f) = filter {
91                if let Some(ref cap) = f.capability {
92                    if !provider.capabilities.contains(cap) {
93                        continue;
94                    }
95                }
96                if let Some(min_uptime) = f.min_uptime {
97                    if provider.uptime_percent < min_uptime {
98                        continue;
99                    }
100                }
101                if let Some(min_mem) = f.min_memory_mb {
102                    if !provider.specs.iter().any(|s| s.memory_mb >= min_mem) {
103                        continue;
104                    }
105                }
106                if let Some(min_cpu) = f.min_cpu {
107                    if !provider.specs.iter().any(|s| s.cpu_millicores >= min_cpu) {
108                        continue;
109                    }
110                }
111            }
112
113            providers.push(provider);
114        }
115
116        info!("Found {} providers matching filter", providers.len());
117        Ok(providers)
118    }
119
120    /// Get details of a specific provider (supports exact match or prefix of at least 8 chars)
121    /// Accepts both hex pubkeys and bech32 npub format.
122    pub async fn get_provider(&self, npub: &str) -> Result<Option<ProviderInfo>> {
123        let providers = self.list_providers(None).await?;
124
125        // Normalize input to hex for comparison (provider npubs are stored as hex)
126        let lookup_hex = match nostr_sdk::PublicKey::parse(npub) {
127            Ok(pk) => pk.to_hex(),
128            Err(_) => npub.to_string(),
129        };
130
131        // precise match first
132        if let Some(p) = providers.iter().find(|p| p.npub == lookup_hex) {
133            return Ok(Some(p.clone()));
134        }
135
136        // try prefix match if long enough
137        if lookup_hex.len() >= 8 {
138            let matches: Vec<&ProviderInfo> = providers.iter()
139                .filter(|p| p.npub.starts_with(&lookup_hex))
140                .collect();
141
142            if matches.len() == 1 {
143                return Ok(Some(matches[0].clone()));
144            }
145        }
146
147        Ok(None)
148    }
149
150    /// Check if a provider is online
151    pub async fn is_provider_online(&self, npub: &str) -> bool {
152        match self.get_provider(npub).await {
153            Ok(Some(p)) => p.is_online,
154            _ => false,
155        }
156    }
157
158    /// Get uptime percentage for a provider
159    pub async fn get_uptime(&self, npub: &str, days: u32) -> Result<f32> {
160        // Resolve full npub
161        let full_npub = if let Ok(Some(p)) = self.get_provider(npub).await {
162            p.npub
163        } else {
164            npub.to_string()
165        };
166        self.nostr.calculate_uptime(&full_npub, days).await
167    }
168
169    /// Get the underlying Nostr client (for sending messages)
170    pub fn nostr(&self) -> &NostrRelaySubscriber {
171        &self.nostr
172    }
173
174    /// Sort providers by various criteria
175    pub fn sort_providers(providers: &mut [ProviderInfo], sort_by: &str) {
176        match sort_by {
177            "price" => {
178                providers.sort_by(|a, b| {
179                    let a_rate = a.specs.first().map(|s| s.rate_msats_per_sec).unwrap_or(u64::MAX);
180                    let b_rate = b.specs.first().map(|s| s.rate_msats_per_sec).unwrap_or(u64::MAX);
181                    a_rate.cmp(&b_rate)
182                });
183            }
184            "uptime" => {
185                providers.sort_by(|a, b| {
186                    b.uptime_percent.partial_cmp(&a.uptime_percent).unwrap_or(std::cmp::Ordering::Equal)
187                });
188            }
189            "capacity" => {
190                providers.sort_by(|a, b| {
191                    let a_mem = a.specs.iter().map(|s| s.memory_mb).max().unwrap_or(0);
192                    let b_mem = b.specs.iter().map(|s| s.memory_mb).max().unwrap_or(0);
193                    b_mem.cmp(&a_mem)
194                });
195            }
196            "jobs" => {
197                providers.sort_by(|a, b| b.total_jobs_completed.cmp(&a.total_jobs_completed));
198            }
199            _ => {} // No sorting
200        }
201    }
202
203    /// Format provider list for display
204    pub fn format_provider_table(providers: &[ProviderInfo]) -> String {
205        use std::fmt::Write;
206        
207        let mut output = String::new();
208        
209        writeln!(&mut output, "┌──────────────────────────────────────────────────────────────────────────────────────────────────────┐").unwrap();
210        writeln!(&mut output, "│ {:^16} │ {:^18} │ {:^10} │ {:^8} │ {:^8} │ {:^10} │ {:^6} │", 
211            "ID", "PROVIDER", "LOCATION", "UPTIME", "CHEAPEST", "LXC/VM", "ONLINE").unwrap();
212        writeln!(&mut output, "├──────────────────────────────────────────────────────────────────────────────────────────────────────┤").unwrap();
213        
214        for p in providers {
215            let id = truncate_str(&p.npub, 16);
216            let location = p.location.as_deref().unwrap_or("Unknown");
217            let cheapest = p.specs.iter()
218                .map(|s| s.rate_msats_per_sec)
219                .min()
220                .map(|r| format!("{}m/s", r))
221                .unwrap_or_else(|| "-".to_string());
222            let capabilities = p.capabilities.join("/");
223            let online = if p.is_online { "✓" } else { "✗" };
224            
225            writeln!(&mut output, "│ {:16} │ {:18} │ {:^10} │ {:>6.1}% │ {:>8} │ {:^10} │ {:^6} │",
226                id,
227                truncate_str(&p.hostname, 18),
228                truncate_str(location, 10),
229                p.uptime_percent,
230                cheapest,
231                capabilities,
232                online
233            ).unwrap();
234        }
235        
236        writeln!(&mut output, "└──────────────────────────────────────────────────────────────────────────────────────────────────────┘").unwrap();
237        
238        output
239    }
240
241    /// Format single provider details
242    pub fn format_provider_details(provider: &ProviderInfo) -> String {
243        use std::fmt::Write;
244        
245        let mut output = String::new();
246        
247        writeln!(&mut output, "┌────────────────────────────────────────────────────────────┐").unwrap();
248        writeln!(&mut output, "│ Provider: {}",  provider.hostname).unwrap();
249        writeln!(&mut output, "├────────────────────────────────────────────────────────────┤").unwrap();
250        writeln!(&mut output, "│ NPUB:       {}", truncate_str(&provider.npub, 45)).unwrap();
251        writeln!(&mut output, "│ Location:   {}", provider.location.as_deref().unwrap_or("Unknown")).unwrap();
252        writeln!(&mut output, "│ Uptime:     {:.1}%", provider.uptime_percent).unwrap();
253        writeln!(&mut output, "│ Jobs Done:  {}", provider.total_jobs_completed).unwrap();
254        writeln!(&mut output, "│ Status:     {}", if provider.is_online { "🟢 Online" } else { "🔴 Offline" }).unwrap();
255        writeln!(&mut output, "│ Supports:   {}", provider.capabilities.join(", ")).unwrap();
256        writeln!(&mut output, "├────────────────────────────────────────────────────────────┤").unwrap();
257        writeln!(&mut output, "│ Available Tiers:").unwrap();
258        
259        for spec in &provider.specs {
260            writeln!(&mut output, "│   • {} ({}) - {} msat/sec",
261                spec.name, spec.id, spec.rate_msats_per_sec).unwrap();
262            writeln!(&mut output, "│     {} vCPU, {} MB RAM",
263                spec.cpu_millicores / 1000, spec.memory_mb).unwrap();
264        }
265        
266        writeln!(&mut output, "├────────────────────────────────────────────────────────────┤").unwrap();
267        writeln!(&mut output, "│ Accepted Mints:").unwrap();
268        for mint in &provider.whitelisted_mints {
269            writeln!(&mut output, "│   • {}", mint).unwrap();
270        }
271        writeln!(&mut output, "└────────────────────────────────────────────────────────────┘").unwrap();
272        
273        output
274    }
275}
276
277/// Helper to truncate strings for display
278fn truncate_str(s: &str, max_len: usize) -> &str {
279    if s.len() <= max_len {
280        s
281    } else {
282        &s[..max_len - 2]
283    }
284}
285
286#[cfg(test)]
287mod tests {
288    use super::*;
289
290    #[test]
291    fn test_format_provider_table() {
292        let providers = vec![
293            ProviderInfo {
294                npub: "npub123".to_string(),
295                hostname: "Test Provider".to_string(),
296                location: Some("US-East".to_string()),
297                capabilities: vec!["lxc".to_string()],
298                specs: vec![PodSpec {
299                    id: "basic".to_string(),
300                    name: "Basic".to_string(),
301                    description: "Test".to_string(),
302                    cpu_millicores: 1000,
303                    memory_mb: 1024,
304                    rate_msats_per_sec: 50,
305                }],
306                whitelisted_mints: vec![],
307                uptime_percent: 99.5,
308                total_jobs_completed: 10,
309                last_seen: 0,
310                is_online: true,
311            }
312        ];
313
314        let table = DiscoveryClient::format_provider_table(&providers);
315        assert!(table.contains("Test Provider"));
316        assert!(table.contains("99.5%"));
317    }
318}