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::Result;
7use tracing::info;
8
9use crate::nostr::{NostrRelaySubscriber, ProviderFilter, ProviderInfo, RelayConfig};
10
11/// Discovery client for finding providers
12pub struct DiscoveryClient {
13    nostr: NostrRelaySubscriber,
14}
15
16impl DiscoveryClient {
17    /// Create a new discovery client
18    pub async fn new(relays: Vec<String>) -> Result<Self> {
19        let config = RelayConfig {
20            relays,
21            private_key: None, // Read-only client doesn't need a key
22        };
23
24        let nostr = NostrRelaySubscriber::new(config).await?;
25
26        Ok(Self { nostr })
27    }
28
29    /// Create with a private key (for sending spawn requests)
30    pub async fn new_with_key(relays: Vec<String>, private_key: String) -> Result<Self> {
31        let config = RelayConfig {
32            relays,
33            private_key: Some(private_key),
34        };
35
36        let nostr = NostrRelaySubscriber::new(config).await?;
37
38        Ok(Self { nostr })
39    }
40
41    /// Get the client's public key (npub)
42    pub fn get_npub(&self) -> String {
43        self.nostr.get_service_public_key()
44    }
45
46    /// List all available providers
47    pub async fn list_providers(
48        &self,
49        filter: Option<ProviderFilter>,
50    ) -> Result<Vec<ProviderInfo>> {
51        let offers = self.nostr.query_providers().await?;
52
53        let mut providers = Vec::new();
54
55        // Optimisation: Fetch all heartbeats in parallel (batch query)
56        let provider_npubs: Vec<String> = offers.iter().map(|o| o.provider_npub.clone()).collect();
57        let heartbeats = self
58            .nostr
59            .get_latest_heartbeats_multi(provider_npubs)
60            .await?;
61
62        for offer in offers {
63            // Check if provider is online (has recent heartbeat)
64            let (is_online, last_seen) = match heartbeats.get(&offer.provider_npub) {
65                Some(hb) => {
66                    let now = std::time::SystemTime::now()
67                        .duration_since(std::time::UNIX_EPOCH)
68                        .map(|d| d.as_secs())
69                        .unwrap_or(0);
70                    // Consider online if heartbeat within last 2 minutes
71                    (now - hb.timestamp < 120, hb.timestamp)
72                }
73                None => (false, 0),
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                isolation_level: offer.isolation_level,
88            };
89
90            // Apply filters
91            if let Some(ref f) = filter {
92                if let Some(ref cap) = f.capability {
93                    if !provider.capabilities.contains(cap) {
94                        continue;
95                    }
96                }
97                if let Some(min_uptime) = f.min_uptime {
98                    if provider.uptime_percent < min_uptime {
99                        continue;
100                    }
101                }
102                if let Some(min_mem) = f.min_memory_mb {
103                    if !provider.specs.iter().any(|s| s.memory_mb >= min_mem) {
104                        continue;
105                    }
106                }
107                if let Some(min_cpu) = f.min_cpu {
108                    if !provider.specs.iter().any(|s| s.cpu_millicores >= min_cpu) {
109                        continue;
110                    }
111                }
112                if let Some(min_iso) = f.isolation_level {
113                    if !provider.isolation_level.meets(min_iso) {
114                        continue;
115                    }
116                }
117            }
118
119            providers.push(provider);
120        }
121
122        info!("Found {} providers matching filter", providers.len());
123        Ok(providers)
124    }
125
126    /// Get details of a specific provider (supports exact match or prefix of at least 8 chars)
127    /// Accepts both hex pubkeys and bech32 npub format.
128    pub async fn get_provider(&self, npub: &str) -> Result<Option<ProviderInfo>> {
129        let providers = self.list_providers(None).await?;
130
131        // Normalize input to hex for comparison (provider npubs are stored as hex)
132        let lookup_hex = match nostr_sdk::PublicKey::parse(npub) {
133            Ok(pk) => pk.to_hex(),
134            Err(_) => npub.to_string(),
135        };
136
137        // precise match first
138        if let Some(p) = providers.iter().find(|p| p.npub == lookup_hex) {
139            return Ok(Some(p.clone()));
140        }
141
142        // try prefix match if long enough
143        if lookup_hex.len() >= 8 {
144            let matches: Vec<&ProviderInfo> = providers
145                .iter()
146                .filter(|p| p.npub.starts_with(&lookup_hex))
147                .collect();
148
149            if matches.len() == 1 {
150                return Ok(Some(matches[0].clone()));
151            }
152        }
153
154        Ok(None)
155    }
156
157    /// Check if a provider is online
158    pub async fn is_provider_online(&self, npub: &str) -> bool {
159        match self.get_provider(npub).await {
160            Ok(Some(p)) => p.is_online,
161            _ => false,
162        }
163    }
164
165    /// Get uptime percentage for a provider
166    pub async fn get_uptime(&self, npub: &str, days: u32) -> Result<f32> {
167        // Resolve full npub
168        let full_npub = if let Ok(Some(p)) = self.get_provider(npub).await {
169            p.npub
170        } else {
171            npub.to_string()
172        };
173        self.nostr.calculate_uptime(&full_npub, days).await
174    }
175
176    /// Get the underlying Nostr client (for sending messages)
177    pub fn nostr(&self) -> &NostrRelaySubscriber {
178        &self.nostr
179    }
180
181    /// Sort providers by various criteria
182    pub fn sort_providers(providers: &mut [ProviderInfo], sort_by: &str) {
183        match sort_by {
184            "price" => {
185                providers.sort_by(|a, b| {
186                    let a_rate = a
187                        .specs
188                        .first()
189                        .map(|s| s.rate_msats_per_sec)
190                        .unwrap_or(u64::MAX);
191                    let b_rate = b
192                        .specs
193                        .first()
194                        .map(|s| s.rate_msats_per_sec)
195                        .unwrap_or(u64::MAX);
196                    a_rate.cmp(&b_rate)
197                });
198            }
199            "uptime" => {
200                providers.sort_by(|a, b| {
201                    b.uptime_percent
202                        .partial_cmp(&a.uptime_percent)
203                        .unwrap_or(std::cmp::Ordering::Equal)
204                });
205            }
206            "capacity" => {
207                providers.sort_by(|a, b| {
208                    let a_mem = a.specs.iter().map(|s| s.memory_mb).max().unwrap_or(0);
209                    let b_mem = b.specs.iter().map(|s| s.memory_mb).max().unwrap_or(0);
210                    b_mem.cmp(&a_mem)
211                });
212            }
213            "jobs" => {
214                providers.sort_by(|a, b| b.total_jobs_completed.cmp(&a.total_jobs_completed));
215            }
216            _ => {} // No sorting
217        }
218    }
219
220    /// Format provider list for display
221    pub fn format_provider_table(providers: &[ProviderInfo]) -> String {
222        use std::fmt::Write;
223
224        let mut output = String::new();
225
226        // Replaced the historically-uniform `LXC/VM` column with
227        // `TIER` (the offer's isolation level) — every provider
228        // today reports `lxc/vm`, so the old column was decorative.
229        // `TIER` is the only column that distinguishes a Docker
230        // provider from a per-VM KVM provider in the listing.
231        writeln!(&mut output, "┌──────────────────────────────────────────────────────────────────────────────────────────────────────┐").unwrap();
232        writeln!(
233            &mut output,
234            "│ {:^16} │ {:^18} │ {:^10} │ {:^8} │ {:^8} │ {:^10} │ {:^6} │",
235            "ID", "PROVIDER", "LOCATION", "UPTIME", "CHEAPEST", "TIER", "ONLINE"
236        )
237        .unwrap();
238        writeln!(&mut output, "├──────────────────────────────────────────────────────────────────────────────────────────────────────┤").unwrap();
239
240        for p in providers {
241            let id = truncate_str(&p.npub, 16);
242            let location = p.location.as_deref().unwrap_or("Unknown");
243            let cheapest = p
244                .specs
245                .iter()
246                .map(|s| s.rate_msats_per_sec)
247                .min()
248                .map(|r| format!("{}m/s", r))
249                .unwrap_or_else(|| "-".to_string());
250            // Compact tier label that fits the 10-char column.
251            // `attested-research-tier` is too long; abbreviate.
252            let tier = match p.isolation_level {
253                crate::nostr::IsolationLevel::SharedKernel => "shared",
254                crate::nostr::IsolationLevel::DedicatedHost => "dedicated",
255                crate::nostr::IsolationLevel::AttestedResearchTier => "attested",
256            };
257            let online = if p.is_online { "✓" } else { "✗" };
258
259            writeln!(
260                &mut output,
261                "│ {:16} │ {:18} │ {:^10} │ {:>6.1}% │ {:>8} │ {:^10} │ {:^6} │",
262                id,
263                truncate_str(&p.hostname, 18),
264                truncate_str(location, 10),
265                p.uptime_percent,
266                cheapest,
267                tier,
268                online
269            )
270            .unwrap();
271        }
272
273        writeln!(&mut output, "└──────────────────────────────────────────────────────────────────────────────────────────────────────┘").unwrap();
274
275        output
276    }
277
278    /// Format single provider details
279    pub fn format_provider_details(provider: &ProviderInfo) -> String {
280        use std::fmt::Write;
281
282        let mut output = String::new();
283
284        writeln!(
285            &mut output,
286            "┌────────────────────────────────────────────────────────────┐"
287        )
288        .unwrap();
289        writeln!(&mut output, "│ Provider: {}", provider.hostname).unwrap();
290        writeln!(
291            &mut output,
292            "├────────────────────────────────────────────────────────────┤"
293        )
294        .unwrap();
295        writeln!(
296            &mut output,
297            "│ NPUB:       {}",
298            truncate_str(&provider.npub, 45)
299        )
300        .unwrap();
301        writeln!(
302            &mut output,
303            "│ Location:   {}",
304            provider.location.as_deref().unwrap_or("Unknown")
305        )
306        .unwrap();
307        writeln!(&mut output, "│ Uptime:     {:.1}%", provider.uptime_percent).unwrap();
308        writeln!(
309            &mut output,
310            "│ Jobs Done:  {}",
311            provider.total_jobs_completed
312        )
313        .unwrap();
314        writeln!(
315            &mut output,
316            "│ Status:     {}",
317            if provider.is_online {
318                "🟢 Online"
319            } else {
320                "🔴 Offline"
321            }
322        )
323        .unwrap();
324        writeln!(
325            &mut output,
326            "│ Supports:   {}",
327            provider.capabilities.join(", ")
328        )
329        .unwrap();
330        // Full slug here (vs the abbreviated form in the table).
331        // Annotated so a reader who's only just discovering the
332        // tier system understands what each label means without
333        // bouncing to the docs.
334        let iso_annotation = match provider.isolation_level {
335            crate::nostr::IsolationLevel::SharedKernel => " (containers; co-tenant boundary only)",
336            crate::nostr::IsolationLevel::DedicatedHost => {
337                " (per-VM; no co-tenants, but operator can read guest)"
338            }
339            crate::nostr::IsolationLevel::AttestedResearchTier => {
340                " (SEV-SNP / TDX; operator cannot read guest memory)"
341            }
342        };
343        writeln!(
344            &mut output,
345            "│ Isolation:  {}{}",
346            provider.isolation_level.slug(),
347            iso_annotation
348        )
349        .unwrap();
350        writeln!(
351            &mut output,
352            "├────────────────────────────────────────────────────────────┤"
353        )
354        .unwrap();
355        writeln!(&mut output, "│ Available Tiers:").unwrap();
356
357        for spec in &provider.specs {
358            writeln!(
359                &mut output,
360                "│   • {} ({}) - {} msat/sec",
361                spec.name, spec.id, spec.rate_msats_per_sec
362            )
363            .unwrap();
364            writeln!(
365                &mut output,
366                "│     {} vCPU, {} MB RAM",
367                spec.cpu_millicores / 1000,
368                spec.memory_mb
369            )
370            .unwrap();
371        }
372
373        writeln!(
374            &mut output,
375            "├────────────────────────────────────────────────────────────┤"
376        )
377        .unwrap();
378        writeln!(&mut output, "│ Accepted Mints:").unwrap();
379        for mint in &provider.whitelisted_mints {
380            writeln!(&mut output, "│   • {}", mint).unwrap();
381        }
382        writeln!(
383            &mut output,
384            "└────────────────────────────────────────────────────────────┘"
385        )
386        .unwrap();
387
388        output
389    }
390}
391
392/// Helper to truncate strings for display
393fn truncate_str(s: &str, max_len: usize) -> &str {
394    if s.len() <= max_len {
395        s
396    } else {
397        &s[..max_len - 2]
398    }
399}
400
401#[cfg(test)]
402mod tests {
403    use super::*;
404    use crate::nostr::PodSpec;
405
406    #[test]
407    fn test_format_provider_table() {
408        let providers = vec![ProviderInfo {
409            npub: "npub123".to_string(),
410            hostname: "Test Provider".to_string(),
411            location: Some("US-East".to_string()),
412            capabilities: vec!["lxc".to_string()],
413            specs: vec![PodSpec {
414                id: "basic".to_string(),
415                name: "Basic".to_string(),
416                description: "Test".to_string(),
417                cpu_millicores: 1000,
418                memory_mb: 1024,
419                rate_msats_per_sec: 50,
420            }],
421            whitelisted_mints: vec![],
422            uptime_percent: 99.5,
423            total_jobs_completed: 10,
424            last_seen: 0,
425            is_online: true,
426            isolation_level: crate::nostr::IsolationLevel::SharedKernel,
427        }];
428
429        let table = DiscoveryClient::format_provider_table(&providers);
430        assert!(table.contains("Test Provider"));
431        assert!(table.contains("99.5%"));
432    }
433}