osdk_core/source/
select.rs

1//! Source selection: given a backend's default sources (plus user config),
2//! pick which one to use — by pin, priority order, or fastest-probe (auto).
3//!
4//! Probing and disk-cached results are the meat of M3; this module also exposes
5//! `active_source` used by every backend to resolve a single source now.
6
7use std::path::{Path, PathBuf};
8use std::time::{Duration, Instant};
9
10use crate::backend::{Backend, Ctx};
11use crate::error::{Error, Result};
12use crate::source::{
13    candidate_fingerprint, ProbeCache, ProbeResult, Selection, Source, SourceKind,
14};
15
16const SOURCE_CACHE_SCHEMA_VERSION: u32 = 2;
17
18#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
19struct VersionedProbeCache {
20    schema_version: u32,
21    candidate_fingerprint: String,
22    results: Vec<ProbeResult>,
23}
24
25impl VersionedProbeCache {
26    fn new(candidate_fingerprint: String, results: Vec<ProbeResult>) -> Self {
27        Self {
28            schema_version: SOURCE_CACHE_SCHEMA_VERSION,
29            candidate_fingerprint,
30            results,
31        }
32    }
33
34    fn is_compatible(&self, candidate_fingerprint: &str) -> bool {
35        self.schema_version == SOURCE_CACHE_SCHEMA_VERSION
36            && self.candidate_fingerprint == candidate_fingerprint
37    }
38
39    fn into_probe_cache(self) -> ProbeCache {
40        ProbeCache {
41            results: self.results,
42        }
43    }
44}
45
46/// Assemble the effective source list for a backend: defaults minus disabled,
47/// plus user custom sources, honoring per-tool config.
48pub fn effective_sources(ctx: &Ctx, backend: &dyn Backend) -> Vec<Source> {
49    let mut sources = backend.default_sources();
50    if let Some(tool_cfg) = ctx.config.tool_sources(backend.id()) {
51        if !tool_cfg.disable.is_empty() {
52            sources.retain(|s| !tool_cfg.disable.iter().any(|d| d == &s.id));
53        }
54        for custom in &tool_cfg.custom {
55            // custom overrides a builtin with the same id
56            sources.retain(|s| s.id != custom.id);
57            sources.push(custom.clone());
58        }
59    }
60    sources.retain(|s| s.enabled);
61    // Stable order by priority (lower first) for the `ordered` strategy.
62    sources.sort_by_key(|s| s.priority);
63    sources
64}
65
66/// Resolve the single source to use right now for `backend`.
67///
68/// - `pin` (config) always wins if the id exists.
69/// - `Selection::Auto` uses cached probe results when fresh, else probes.
70/// - `Selection::Ordered` returns the highest-priority enabled source.
71pub async fn active_source(ctx: &Ctx, backend: &dyn Backend) -> Result<Source> {
72    ranked_source_list(ctx, backend)
73        .await?
74        .into_iter()
75        .next()
76        .ok_or_else(|| Error::NoUsableSource {
77            tool: backend.id().to_string(),
78            tried: 0,
79        })
80}
81
82/// The full list of candidate sources, best-first. Used for download failover:
83/// callers try each in order until one succeeds.
84///
85/// A config pin (or one-shot `--source`) moves that source to the front.
86pub async fn ranked_source_list(ctx: &Ctx, backend: &dyn Backend) -> Result<Vec<Source>> {
87    ranked_source_candidates(ctx, backend, effective_sources(ctx, backend)).await
88}
89
90/// Rank an already-filtered effective source set with the same pin, cache,
91/// offline, and live-probe policy as [`ranked_source_list`]. Backends use this
92/// when they must narrow candidates before any network probe, for example to
93/// keep private package names away from public registries.
94pub async fn ranked_source_candidates(
95    ctx: &Ctx,
96    backend: &dyn Backend,
97    sources: Vec<Source>,
98) -> Result<Vec<Source>> {
99    if sources.is_empty() {
100        return Err(Error::NoUsableSource {
101            tool: backend.id().to_string(),
102            tried: 0,
103        });
104    }
105
106    // Explicit pin wins: put it first, keep the rest as fallbacks.
107    if let Some(tool_cfg) = ctx.config.tool_sources(backend.id()) {
108        if let Some(pin) = &tool_cfg.pin {
109            if let Some(idx) = sources.iter().position(|s| &s.id == pin) {
110                let mut ordered = sources.clone();
111                let pinned = ordered.remove(idx);
112                let mut out = vec![pinned];
113                out.extend(ordered);
114                return Ok(out);
115            }
116        }
117    }
118
119    if ctx.config.settings.offline {
120        if matches!(ctx.config.sources.selection, Selection::Auto) {
121            if let Some(cache) = load_cache(ctx, backend.id(), &sources) {
122                return Ok(order_sources_by_probe_results(sources, cache.results));
123            }
124        }
125        return Ok(sources);
126    }
127
128    match ctx.config.sources.selection {
129        Selection::Ordered | Selection::Pinned => Ok(sources),
130        Selection::Auto => {
131            let ranked_ids = ranked_sources(ctx, backend, &sources).await;
132            Ok(order_sources_by_ids(sources, &ranked_ids))
133        }
134    }
135}
136
137fn order_sources_by_probe_results(
138    sources: Vec<Source>,
139    mut results: Vec<ProbeResult>,
140) -> Vec<Source> {
141    results.sort_by(|left, right| right.score().total_cmp(&left.score()));
142    let ids = results
143        .into_iter()
144        .filter(|result| result.ok)
145        .map(|result| result.source_id)
146        .collect::<Vec<_>>();
147    order_sources_by_ids(sources, &ids)
148}
149
150fn order_sources_by_ids(sources: Vec<Source>, ids: &[String]) -> Vec<Source> {
151    // Append candidates without a successful probe so they remain available
152    // as fallbacks after the ranked sources.
153    let mut ordered = Vec::with_capacity(sources.len());
154    for id in ids {
155        if let Some(source) = sources.iter().find(|source| &source.id == id) {
156            ordered.push(source.clone());
157        }
158    }
159    for source in sources {
160        if !ordered.iter().any(|ranked| ranked.id == source.id) {
161            ordered.push(source);
162        }
163    }
164    ordered
165}
166
167/// Return source ids ranked best-first, using fresh cache or a live probe.
168async fn ranked_sources(ctx: &Ctx, backend: &dyn Backend, sources: &[Source]) -> Vec<String> {
169    // Try fresh cache first.
170    if let Some(cache) = load_cache(ctx, backend.id(), sources) {
171        let ttl = ctx.config.sources.cache_ttl_secs();
172        let now = crate::source::now_secs();
173        let fresh = cache
174            .results
175            .iter()
176            .all(|r| now.saturating_sub(r.measured_at) <= ttl)
177            && !cache.results.is_empty();
178        if fresh {
179            let mut results = cache.results.clone();
180            results.sort_by(|a, b| b.score().total_cmp(&a.score()));
181            return results
182                .into_iter()
183                .filter(|r| r.ok)
184                .map(|r| r.source_id)
185                .collect();
186        }
187    }
188
189    // Live probe.
190    let results = probe_all(ctx, backend, sources).await;
191    save_cache(ctx, backend.id(), sources, &results);
192    let mut ok: Vec<ProbeResult> = results.into_iter().filter(|r| r.ok).collect();
193    ok.sort_by(|a, b| b.score().total_cmp(&a.score()));
194    ok.into_iter().map(|r| r.source_id).collect()
195}
196
197/// Probe every source concurrently, returning results (failed ones included).
198pub async fn probe_all(ctx: &Ctx, backend: &dyn Backend, sources: &[Source]) -> Vec<ProbeResult> {
199    let timeout = Duration::from_millis(ctx.config.sources.probe_timeout_ms);
200    let mut handles = Vec::new();
201    for s in sources {
202        let url = backend.probe_url(ctx, s);
203        let client = ctx.client.clone();
204        let source = s.clone();
205        let to = timeout;
206        handles.push(tokio::spawn(async move {
207            match url {
208                Some(u) => probe_one(&client, &source, &u, to).await,
209                None => ProbeResult::failed(&source.id),
210            }
211        }));
212    }
213    let mut out = Vec::new();
214    for h in handles {
215        if let Ok(r) = h.await {
216            out.push(r)
217        }
218    }
219    out
220}
221
222/// Probe a single URL: measure time-to-first-byte and throughput over a bounded
223/// window, downloading at most ~1MB.
224async fn probe_one(
225    client: &reqwest::Client,
226    source: &Source,
227    url: &str,
228    timeout: Duration,
229) -> ProbeResult {
230    use futures_util::StreamExt;
231
232    let start = Instant::now();
233    let fut = async {
234        let resp = crate::http::get_source_response(client, source, url)
235            .await
236            .ok()?
237            .error_for_status()
238            .ok()?;
239        let ttfb = start.elapsed();
240        let mut stream = resp.bytes_stream();
241        let mut downloaded: u64 = 0;
242        let body_start = Instant::now();
243        while let Some(chunk) = stream.next().await {
244            match chunk {
245                Ok(c) => {
246                    downloaded += c.len() as u64;
247                    if downloaded >= 1_000_000 {
248                        break;
249                    }
250                }
251                Err(_) => break,
252            }
253        }
254        let secs = body_start.elapsed().as_secs_f64().max(0.001);
255        let throughput = downloaded as f64 / secs;
256        Some((ttfb, throughput, downloaded))
257    };
258
259    match tokio::time::timeout(timeout, fut).await {
260        Ok(Some((ttfb, throughput, downloaded))) if downloaded > 0 => ProbeResult {
261            source_id: source.id.clone(),
262            throughput,
263            ttfb_ms: ttfb.as_millis() as u64,
264            ok: true,
265            measured_at: crate::source::now_secs(),
266        },
267        _ => ProbeResult::failed(&source.id),
268    }
269}
270
271fn cache_path(ctx: &Ctx, tool: &str) -> PathBuf {
272    let mut path = ctx
273        .dirs
274        .sources_cache()
275        .join(crate::dirs::sanitize_tool_id(tool));
276    path.set_extension("json");
277    path
278}
279
280fn read_versioned_cache(path: &Path, candidate_fingerprint: &str) -> Option<ProbeCache> {
281    let bytes = std::fs::read(path).ok()?;
282    let cache: VersionedProbeCache = serde_json::from_slice(&bytes).ok()?;
283    cache
284        .is_compatible(candidate_fingerprint)
285        .then(|| cache.into_probe_cache())
286}
287
288fn load_cache(ctx: &Ctx, tool: &str, sources: &[Source]) -> Option<ProbeCache> {
289    let candidate_fingerprint = candidate_fingerprint(sources);
290    read_versioned_cache(&cache_path(ctx, tool), &candidate_fingerprint)
291}
292
293fn save_cache(ctx: &Ctx, tool: &str, sources: &[Source], results: &[ProbeResult]) {
294    let p = cache_path(ctx, tool);
295    if let Some(parent) = p.parent() {
296        let _ = std::fs::create_dir_all(parent);
297    }
298    let cache = VersionedProbeCache::new(candidate_fingerprint(sources), results.to_vec());
299    if let Ok(bytes) = serde_json::to_vec_pretty(&cache) {
300        let _ = std::fs::write(&p, bytes);
301    }
302}
303
304/// Force a refresh of the probe cache for a backend (used by `osdk source test`
305/// and `--refresh-sources`). Returns the fresh results.
306pub async fn refresh(ctx: &Ctx, backend: &dyn Backend) -> Result<Vec<ProbeResult>> {
307    if ctx.config.settings.offline {
308        return Err(Error::other("cannot refresh sources while offline"));
309    }
310    let sources = effective_sources(ctx, backend);
311    let results = probe_all(ctx, backend, &sources).await;
312    save_cache(ctx, backend.id(), &sources, &results);
313    Ok(results)
314}
315
316/// Human-readable kind label.
317pub fn kind_label(kind: SourceKind) -> &'static str {
318    match kind {
319        SourceKind::Official => "official",
320        SourceKind::Mirror => "mirror",
321        SourceKind::Custom => "custom",
322    }
323}
324
325#[cfg(test)]
326mod tests {
327    use super::*;
328    use std::collections::BTreeMap;
329    use std::io::{Read, Write};
330    use std::net::TcpListener;
331    use std::sync::Arc;
332
333    use async_trait::async_trait;
334
335    use crate::config::{Config, Settings};
336    use crate::dirs::Dirs;
337    use crate::platform::Platform;
338    use crate::store::Cas;
339    use crate::version::{ToolVersion, VersionInfo};
340
341    struct FixtureBackend {
342        id: &'static str,
343        sources: Vec<Source>,
344    }
345
346    #[async_trait]
347    impl Backend for FixtureBackend {
348        fn id(&self) -> &str {
349            self.id
350        }
351
352        fn default_sources(&self) -> Vec<Source> {
353            self.sources.clone()
354        }
355
356        fn probe_url(&self, _ctx: &Ctx, source: &Source) -> Option<String> {
357            Some(source.download_url.clone())
358        }
359
360        async fn list_remote_versions(&self, _ctx: &Ctx) -> Result<Vec<VersionInfo>> {
361            Ok(Vec::new())
362        }
363
364        async fn install(
365            &self,
366            _ctx: &crate::backend::InstallCtx<'_>,
367            _tv: &ToolVersion,
368        ) -> Result<()> {
369            Ok(())
370        }
371
372        fn bin_paths(&self, _ctx: &Ctx, _tv: &ToolVersion) -> Result<Vec<PathBuf>> {
373            Ok(Vec::new())
374        }
375
376        fn bin_names(&self, _ctx: &Ctx, _tv: &ToolVersion) -> Result<Vec<String>> {
377            Ok(Vec::new())
378        }
379    }
380
381    #[test]
382    fn cache_path_uses_nested_sanitized_tool_ids() {
383        let temporary = tempfile::tempdir().unwrap();
384        let ctx = test_ctx(temporary.path(), false);
385        assert_eq!(
386            cache_path(&ctx, "github:owner/repo"),
387            ctx.dirs
388                .sources_cache()
389                .join("github")
390                .join("owner")
391                .join("repo.json")
392        );
393        assert_eq!(
394            cache_path(&ctx, "../:owner\\\\repo"),
395            ctx.dirs.sources_cache().join("owner").join("repo.json")
396        );
397    }
398
399    #[test]
400    fn load_cache_treats_legacy_cache_as_stale() {
401        let temporary = tempfile::tempdir().unwrap();
402        let ctx = test_ctx(temporary.path(), false);
403        let sources = vec![Source::mirror("fixture", "https://mirror.example.test", 1)];
404        let legacy_path = cache_path(&ctx, "tool");
405        std::fs::write(
406            &legacy_path,
407            serde_json::to_vec_pretty(&ProbeCache {
408                results: vec![ProbeResult {
409                    source_id: "fixture".into(),
410                    throughput: 10.0,
411                    ttfb_ms: 1,
412                    ok: true,
413                    measured_at: crate::source::now_secs(),
414                }],
415            })
416            .unwrap(),
417        )
418        .unwrap();
419
420        assert!(load_cache(&ctx, "tool", &sources).is_none());
421    }
422
423    #[test]
424    fn load_cache_rejects_changed_candidate_sets() {
425        let temporary = tempfile::tempdir().unwrap();
426        let ctx = test_ctx(temporary.path(), false);
427        let base_sources = vec![
428            Source::mirror("fixture", "https://mirror.example.test", 1),
429            Source::mirror("extra", "https://extra.example.test", 2),
430        ];
431        let cache_path = cache_path(&ctx, "tool");
432        save_cache(
433            &ctx,
434            "tool",
435            &base_sources,
436            &[
437                ProbeResult {
438                    source_id: "fixture".into(),
439                    throughput: 10.0,
440                    ttfb_ms: 1,
441                    ok: true,
442                    measured_at: crate::source::now_secs(),
443                },
444                ProbeResult {
445                    source_id: "extra".into(),
446                    throughput: 5.0,
447                    ttfb_ms: 2,
448                    ok: true,
449                    measured_at: crate::source::now_secs(),
450                },
451            ],
452        );
453
454        let mut changed_download = base_sources.clone();
455        changed_download[0].download_url = "https://other.example.test".into();
456        assert!(load_cache(&ctx, "tool", &changed_download).is_none());
457
458        let mut changed_priority = base_sources.clone();
459        changed_priority[0].priority = 99;
460        assert!(load_cache(&ctx, "tool", &changed_priority).is_none());
461
462        let mut changed_enabled = base_sources.clone();
463        changed_enabled[0].enabled = false;
464        assert!(load_cache(&ctx, "tool", &changed_enabled).is_none());
465
466        let changed_add = vec![
467            base_sources[0].clone(),
468            base_sources[1].clone(),
469            Source::mirror("third", "https://third.example.test", 3),
470        ];
471        assert!(load_cache(&ctx, "tool", &changed_add).is_none());
472
473        let changed_remove = vec![base_sources[0].clone()];
474        assert!(load_cache(&ctx, "tool", &changed_remove).is_none());
475
476        assert!(cache_path.is_file());
477    }
478
479    #[tokio::test]
480    async fn ranked_sources_reprobes_when_candidates_change() {
481        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
482        let address = listener.local_addr().unwrap();
483        let server = std::thread::spawn(move || {
484            for _ in 0..3 {
485                let (mut stream, _) = listener.accept().unwrap();
486                let mut request = Vec::new();
487                let mut buffer = [0u8; 2048];
488                while !request.ends_with(b"\r\n\r\n") {
489                    let read = stream.read(&mut buffer).unwrap();
490                    if read == 0 {
491                        break;
492                    }
493                    request.extend_from_slice(&buffer[..read]);
494                }
495                stream
496                    .write_all(
497                        b"HTTP/1.1 200 OK\r\nContent-Length: 4\r\nConnection: close\r\n\r\ndata",
498                    )
499                    .unwrap();
500            }
501        });
502
503        let temporary = tempfile::tempdir().unwrap();
504        let ctx = test_ctx(temporary.path(), false);
505        let backend = FixtureBackend {
506            id: "tool",
507            sources: vec![Source::mirror("fixture", &format!("http://{address}"), 1)],
508        };
509
510        let first = ranked_sources(&ctx, &backend, &backend.sources).await;
511        assert_eq!(first, vec!["fixture"]);
512
513        let second = ranked_sources(&ctx, &backend, &backend.sources).await;
514        assert_eq!(second, vec!["fixture"]);
515
516        let changed_sources = vec![
517            Source::mirror("fixture", &format!("http://{address}"), 1),
518            Source::mirror("new", &format!("http://{address}"), 2),
519        ];
520        let mut third = ranked_sources(&ctx, &backend, &changed_sources).await;
521        third.sort();
522        assert_eq!(third, vec!["fixture", "new"]);
523
524        server.join().unwrap();
525    }
526
527    #[tokio::test]
528    async fn source_probe_applies_explicit_headers_without_forward_credentials() {
529        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
530        let address = listener.local_addr().unwrap();
531        let server = std::thread::spawn(move || {
532            let (mut stream, _) = listener.accept().unwrap();
533            let mut request = Vec::new();
534            let mut buffer = [0u8; 2048];
535            while !request.ends_with(b"\r\n\r\n") {
536                let read = stream.read(&mut buffer).unwrap();
537                if read == 0 {
538                    break;
539                }
540                request.extend_from_slice(&buffer[..read]);
541            }
542            let request = String::from_utf8(request).unwrap().to_ascii_lowercase();
543            assert!(request.contains("x-probe-key: source-secret"));
544            stream
545                .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 4\r\nConnection: close\r\n\r\ndata")
546                .unwrap();
547        });
548
549        let temporary = tempfile::tempdir().unwrap();
550        let ctx = test_ctx(temporary.path(), false);
551        let mut source = Source::mirror("fixture", &format!("http://{address}/"), 1);
552        source.forward_credentials = false;
553        source.headers = vec![("X-Probe-Key".into(), "source-secret".into())];
554        let backend = FixtureBackend {
555            id: "tool",
556            sources: vec![source.clone()],
557        };
558
559        let results = probe_all(&ctx, &backend, &[source]).await;
560        assert_eq!(results.len(), 1);
561        assert!(results[0].ok);
562        server.join().unwrap();
563    }
564
565    #[tokio::test]
566    async fn offline_ranked_source_list_reuses_compatible_probe_cache() {
567        let temporary = tempfile::tempdir().unwrap();
568        let online_ctx = test_ctx(temporary.path(), false);
569        let backend = FixtureBackend {
570            id: "github:owner/repo",
571            sources: vec![
572                Source::mirror("first", "https://fast.example.test", 20),
573                Source::mirror("second", "https://slow.example.test", 10),
574            ],
575        };
576        let effective = effective_sources(&online_ctx, &backend);
577        save_cache(
578            &online_ctx,
579            backend.id(),
580            &effective,
581            &[
582                ProbeResult {
583                    source_id: "first".into(),
584                    throughput: 100.0,
585                    ttfb_ms: 10,
586                    ok: true,
587                    measured_at: crate::source::now_secs(),
588                },
589                ProbeResult {
590                    source_id: "second".into(),
591                    throughput: 10.0,
592                    ttfb_ms: 100,
593                    ok: true,
594                    measured_at: crate::source::now_secs(),
595                },
596            ],
597        );
598
599        let ranked = ranked_sources(&online_ctx, &backend, &effective).await;
600        assert_eq!(ranked, vec!["first", "second"]);
601
602        let mut offline_ctx = test_ctx(temporary.path(), true);
603        offline_ctx.config.sources.selection = Selection::Auto;
604        let resolved = ranked_source_list(&offline_ctx, &backend).await.unwrap();
605        let order: Vec<_> = resolved.into_iter().map(|source| source.id).collect();
606        assert_eq!(order, vec!["first", "second"]);
607    }
608
609    fn test_ctx(root: &Path, offline: bool) -> Ctx {
610        let dirs = Dirs::resolve_from(|key| match key {
611            "OSDK_DATA_DIR" => Some(root.join("data").display().to_string()),
612            "OSDK_CACHE_DIR" => Some(root.join("cache").display().to_string()),
613            "OSDK_CONFIG_DIR" => Some(root.join("config").display().to_string()),
614            _ => None,
615        })
616        .unwrap();
617        dirs.ensure().unwrap();
618        let settings = Settings {
619            offline,
620            ..Default::default()
621        };
622        Ctx {
623            dirs: dirs.clone(),
624            platform: Platform::current(),
625            config: Config {
626                settings,
627                sources: Default::default(),
628                tools: BTreeMap::new(),
629                tool_configs: BTreeMap::new(),
630                global_tools: BTreeMap::new(),
631                global_tool_configs: BTreeMap::new(),
632                tool_origins: BTreeMap::new(),
633                aliases: BTreeMap::new(),
634                project_config_path: None,
635            },
636            client: reqwest::Client::new(),
637            cas: Arc::new(Cas::new(dirs.store.clone())),
638            show_progress: false,
639        }
640    }
641}
osdk_core/source/select.rs

osdk_core/source/
select.rs
