Skip to main content

kache_core/
lib.rs

1#[cfg(feature = "planning")]
2use std::collections::{HashMap, HashSet};
3
4#[cfg(feature = "planning")]
5use anyhow::Result;
6#[cfg(feature = "planning")]
7use async_trait::async_trait;
8use serde::{Deserialize, Serialize};
9
10pub mod timeline;
11
12#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
13pub struct BuildIntent {
14    #[serde(default)]
15    pub crate_names: Vec<String>,
16    #[serde(default)]
17    pub namespace: Option<String>,
18    #[serde(default)]
19    pub cargo_lock_deps: Vec<(String, String)>,
20    /// Rank-0 prefetch key: lock digest + target + profile. Absent on older
21    /// clients and when no lockfile was visible at session start.
22    #[serde(default)]
23    pub identity_key: Option<String>,
24}
25
26/// Which source produced a candidate, i.e. how much to trust it
27/// (kunobi-ninja/kache#617).
28///
29/// `Unknown` is the `#[serde(other)]` arm so a newer planner naming a source
30/// this build has never heard of degrades to "untrusted" instead of failing
31/// the whole plan.
32#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
33#[cfg_attr(kani, derive(kani::Arbitrary))]
34#[serde(rename_all = "snake_case")]
35pub enum CandidateSource {
36    /// Recorded action list for this lockfile + target + profile.
37    Manifest,
38    /// Exact lockfile-shard match: this build's dependency set produced it.
39    Shard,
40    /// This machine built this crate before.
41    History,
42    /// A crate NAME matched something in the remote listing. A crate name is
43    /// not a build identity, so this is a guess.
44    KeyCache,
45    #[default]
46    #[serde(other)]
47    Unknown,
48}
49
50impl CandidateSource {
51    /// Confidence rank, lower is better. Only the ORDER matters; these are not
52    /// probabilities and must not be presented as any.
53    pub fn confidence_rank(self) -> u8 {
54        match self {
55            CandidateSource::Manifest => 0,
56            CandidateSource::Shard => 1,
57            CandidateSource::History => 2,
58            CandidateSource::KeyCache => 3,
59            CandidateSource::Unknown => 4,
60        }
61    }
62}
63
64#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
65pub struct PrefetchCandidate {
66    pub cache_key: String,
67    pub crate_name: String,
68    /// What a miss would cost to rebuild. `None` = unknown, which is NOT the
69    /// same as zero: an un-backfilled store row reads 0, and treating that as
70    /// "free to fetch and worthless to have" would bury it (#617).
71    #[serde(default, skip_serializing_if = "Option::is_none")]
72    pub compile_time_ms: Option<u64>,
73    /// Stored artifact size. An admission and ranking ESTIMATE, never a
74    /// promise about compressed transfer bytes: anything enforced has to be
75    /// counted on the wire.
76    #[serde(default, skip_serializing_if = "Option::is_none")]
77    pub size_bytes: Option<u64>,
78    #[serde(default)]
79    pub source: CandidateSource,
80    /// Position in the build's dependency order, i.e. roughly when the build
81    /// will ask for it. `None` = not in the intent's crate list.
82    #[serde(default, skip_serializing_if = "Option::is_none")]
83    pub demand_index: Option<u32>,
84}
85
86impl PrefetchCandidate {
87    /// A candidate with no metadata, as the pre-#617 wire produced.
88    pub fn new(cache_key: String, crate_name: String) -> Self {
89        Self {
90            cache_key,
91            crate_name,
92            compile_time_ms: None,
93            size_bytes: None,
94            source: CandidateSource::Unknown,
95            demand_index: None,
96        }
97    }
98
99    pub fn with_source(mut self, source: CandidateSource) -> Self {
100        self.source = source;
101        self
102    }
103}
104
105#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
106#[serde(rename_all = "snake_case")]
107pub enum PrefetchDisposition {
108    Execute,
109    UseFallback,
110    DoNothing,
111}
112
113#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
114pub struct PrefetchPlan {
115    #[serde(default)]
116    pub plan_id: Option<String>,
117    #[serde(default)]
118    pub planner: Option<String>,
119    pub disposition: PrefetchDisposition,
120    #[serde(default)]
121    pub candidates: Vec<PrefetchCandidate>,
122}
123
124/// How many candidates each source may contribute to one plan
125/// (kunobi-ninja/kache#616).
126///
127/// These bound plan COMPOSITION, which is a different job from the daemon's
128/// key/byte/time budgets: those bound resource use and are the trust boundary,
129/// these stop one low-confidence source from crowding out better candidates
130/// before the budget is even reached. `0` disables a cap.
131#[derive(Debug, Clone, Copy, PartialEq, Eq)]
132pub struct PlanLimits {
133    /// Key-cache variants per crate. A crate name is not a build identity, so
134    /// of `n` variants at most one can be the right one; taking many is paying
135    /// `n` downloads for at most one hit.
136    pub key_cache_per_crate: usize,
137    /// Key-cache candidates across the whole plan, so a build with hundreds of
138    /// unresolved crates cannot fill the plan with guesses.
139    pub key_cache_total: usize,
140    /// History entries per crate. Unlike shards, where one crate legitimately
141    /// has several compile units, extra history rows for one crate are older
142    /// variants.
143    pub history_per_crate: usize,
144}
145
146impl Default for PlanLimits {
147    fn default() -> Self {
148        Self {
149            key_cache_per_crate: 2,
150            key_cache_total: 64,
151            history_per_crate: 2,
152        }
153    }
154}
155
156/// Urgency bucket width, in dependency-order positions (#617).
157///
158/// Demand order is bucketed rather than used exactly because crate-name order
159/// from a lockfile (or a previous graph walk) only approximates when cargo
160/// will actually ask (cargo reorders for parallelism, build scripts, proc
161/// macros, features). Treating position 40 and 45 as meaningfully different
162/// is false precision; 40 versus 400 is real. Roughly the prefetch
163/// concurrency, so one window is about one wave of downloads. Identity
164/// manifests stamp demand_index from last-compile order when they have it.
165#[cfg(feature = "planning")]
166const URGENCY_BUCKET: u32 = 16;
167
168/// Sort key for dispatch order, lowest first (#617).
169///
170/// Lexicographic, deliberately, rather than a weighted score: a weighted sum
171/// needs coefficients, and nothing can calibrate them until #618 makes
172/// "arrived before it was demanded" measurable. Every element here is an
173/// ordering, not a magnitude.
174///
175/// 1. Urgency bucket. Prefetch races the build, so a high-value artifact
176///    needed at minute eight loses to a medium-value one needed at second
177///    five. Candidates with no demand index sort last.
178/// 2. Confidence. Within one wave, prefer the source most likely to be right.
179/// 3. Value, descending. Expensive rebuilds first, so the limited slots buy
180///    the most avoided work. Unknown cost sorts after known cost rather than
181///    being scored as zero.
182///
183/// Callers must apply this as a STABLE sort: equal keys keep source order,
184/// which is the planner's confidence-merge order.
185#[cfg(feature = "planning")]
186pub fn dispatch_sort_key(
187    candidate: &PrefetchCandidate,
188) -> (u32, u8, std::cmp::Reverse<Option<u64>>) {
189    let bucket = candidate
190        .demand_index
191        .map(|index| index / URGENCY_BUCKET)
192        .unwrap_or(u32::MAX);
193    (
194        bucket,
195        candidate.source.confidence_rank(),
196        // `Option` orders `None` before every `Some`, so `Reverse` puts every
197        // known cost (including zero) before unknown while retaining
198        // descending order among known costs.
199        std::cmp::Reverse(candidate.compile_time_ms),
200    )
201}
202
203#[cfg(feature = "planning")]
204fn sort_candidates_for_dispatch(candidates: &mut [PrefetchCandidate]) {
205    candidates.sort_by_key(dispatch_sort_key);
206}
207
208/// Truncate `items` to `limit`, returning how many were dropped. `0` disables.
209///
210/// Shared by every composition cap so "0 means unlimited" is defined once
211/// rather than re-derived at each call site.
212#[cfg(feature = "planning")]
213fn cap_to(items: &mut Vec<PrefetchCandidate>, limit: usize) -> usize {
214    if limit == 0 || items.len() <= limit {
215        return 0;
216    }
217    let dropped = items.len() - limit;
218    items.truncate(limit);
219    dropped
220}
221
222/// What a plan left out, so a truncated plan is distinguishable from one that
223/// had nothing more to offer (#616).
224#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
225pub struct PlanComposition {
226    pub from_identity: usize,
227    pub from_shards: usize,
228    pub from_history: usize,
229    pub from_key_cache: usize,
230    /// Dropped by [`PlanLimits`], by source.
231    pub dropped_history_per_crate: usize,
232    pub dropped_key_cache_per_crate: usize,
233    pub dropped_key_cache_total: usize,
234}
235
236impl PlanComposition {
237    pub fn dropped_total(&self) -> usize {
238        self.dropped_history_per_crate
239            + self.dropped_key_cache_per_crate
240            + self.dropped_key_cache_total
241    }
242}
243
244#[cfg(feature = "planning")]
245#[async_trait]
246pub trait PlannerDataSource {
247    async fn shard_candidates(
248        &self,
249        namespace: &str,
250        deps: &[(String, String)],
251    ) -> Result<Vec<PrefetchCandidate>>;
252
253    async fn history_candidates(&self, crate_names: &[String]) -> Result<Vec<PrefetchCandidate>>;
254
255    async fn key_cache_keys_for_crate(&self, crate_name: &str) -> Result<Vec<String>>;
256
257    /// Rank-0: recorded actions for this lock+target+profile. A miss is
258    /// `Ok(vec![])`, not fatal — shards and history still run.
259    async fn identity_candidates(&self, identity_key: &str) -> Result<Vec<PrefetchCandidate>>;
260}
261
262#[cfg(feature = "planning")]
263pub async fn build_prefetch_plan<T>(
264    source: &T,
265    intent: &BuildIntent,
266    planner_name: &str,
267) -> Result<PrefetchPlan>
268where
269    T: PlannerDataSource + Sync + ?Sized,
270{
271    build_prefetch_plan_with_limits(source, intent, planner_name, PlanLimits::default())
272        .await
273        .map(|(plan, _composition)| plan)
274}
275
276/// [`build_prefetch_plan`] with explicit composition limits, also returning
277/// what the limits dropped so a caller can report it (#616).
278#[cfg(feature = "planning")]
279pub async fn build_prefetch_plan_with_limits<T>(
280    source: &T,
281    intent: &BuildIntent,
282    planner_name: &str,
283    limits: PlanLimits,
284) -> Result<(PrefetchPlan, PlanComposition)>
285where
286    T: PlannerDataSource + Sync + ?Sized,
287{
288    let crate_order = crate_query_order(intent);
289    let demand_index = demand_index_map(&crate_order);
290    let mut seen = HashSet::new();
291    let mut resolved_crates = HashSet::new();
292    let mut candidates = Vec::new();
293    let mut composition = PlanComposition::default();
294
295    // Sources are merged in descending order of confidence, each one filling
296    // only the crates the ones before it left unresolved.
297    if let Some(identity_key) = intent
298        .identity_key
299        .as_deref()
300        .map(str::trim)
301        .filter(|key| !key.is_empty())
302    {
303        match source.identity_candidates(identity_key).await {
304            Ok(identity_candidates) => {
305                for (index, candidate) in identity_candidates.into_iter().enumerate() {
306                    resolved_crates.insert(candidate.crate_name.clone());
307                    if seen.insert(candidate.cache_key.clone()) {
308                        let mut candidate = candidate.with_source(CandidateSource::Manifest);
309                        if candidate.demand_index.is_none() {
310                            candidate.demand_index = u32::try_from(index).ok();
311                        }
312                        candidates.push(candidate);
313                    }
314                }
315                composition.from_identity = candidates.len();
316            }
317            Err(_) => {
318                // Not fatal: shards and history still run.
319            }
320        }
321    }
322    let identity_resolved = resolved_crates.clone();
323
324    // Shard lookups used to RETURN as soon as they produced anything
325    // (kunobi-ninja/kache#614), but a shard hit is exact per bucket: one
326    // dependency bump invalidates one of `NUM_SHARDS` buckets while the rest
327    // still match, so a single matching shard was enough to short-circuit
328    // history and key-cache recovery for every crate in every bucket that
329    // missed.
330    if let Some(namespace) = intent.namespace.as_deref()
331        && !intent.cargo_lock_deps.is_empty()
332    {
333        // A failed shard lookup is not fatal: the sources below still run.
334        if let Ok(shard_candidates) = source
335            .shard_candidates(namespace, &intent.cargo_lock_deps)
336            .await
337        {
338            for candidate in order_candidates_by_crate_order(shard_candidates, intent) {
339                if identity_resolved.contains(&candidate.crate_name) {
340                    continue;
341                }
342                resolved_crates.insert(candidate.crate_name.clone());
343                if seen.insert(candidate.cache_key.clone()) {
344                    composition.from_shards += 1;
345                    candidates.push(candidate.with_source(CandidateSource::Shard));
346                }
347            }
348        }
349    }
350
351    let unresolved = |resolved: &HashSet<String>| -> Vec<String> {
352        crate_order
353            .iter()
354            .filter(|name| !resolved.contains(*name))
355            .cloned()
356            .collect()
357    };
358
359    let history_query = unresolved(&resolved_crates);
360    if !history_query.is_empty() {
361        // Group by crate so the per-crate cap keeps the FIRST entries, which
362        // both data sources return most-recently-used first.
363        let mut by_crate: HashMap<String, Vec<PrefetchCandidate>> = HashMap::new();
364        for candidate in order_candidates_by_crate_order(
365            source.history_candidates(&history_query).await?,
366            intent,
367        ) {
368            by_crate
369                .entry(candidate.crate_name.clone())
370                .or_default()
371                .push(candidate.with_source(CandidateSource::History));
372        }
373
374        for crate_name in &history_query {
375            let Some(mut for_crate) = by_crate.remove(crate_name) else {
376                continue;
377            };
378            composition.dropped_history_per_crate +=
379                cap_to(&mut for_crate, limits.history_per_crate);
380            for candidate in for_crate {
381                resolved_crates.insert(candidate.crate_name.clone());
382                if seen.insert(candidate.cache_key.clone()) {
383                    composition.from_history += 1;
384                    candidates.push(candidate);
385                }
386            }
387        }
388    }
389
390    // The key cache is the weakest source: it maps a crate NAME to every cache
391    // key in the bucket, with no target, toolchain, profile, or feature
392    // dimension, so of `n` variants at most one can be right. Capping is not a
393    // fix for that (dimensioning the remote layout is, separately) but it stops
394    // the guesses crowding out better candidates (#616).
395    for crate_name in unresolved(&resolved_crates) {
396        if limits.key_cache_total > 0 && composition.from_key_cache >= limits.key_cache_total {
397            // Whatever this crate would have offered is dropped wholesale; count
398            // it so the truncation is visible rather than inferred.
399            composition.dropped_key_cache_total += source
400                .key_cache_keys_for_crate(&crate_name)
401                .await?
402                .into_iter()
403                .filter(|key| !seen.contains(key))
404                .count();
405            continue;
406        }
407
408        let mut for_crate: Vec<PrefetchCandidate> = source
409            .key_cache_keys_for_crate(&crate_name)
410            .await?
411            .into_iter()
412            .map(|cache_key| {
413                PrefetchCandidate::new(cache_key, crate_name.clone())
414                    .with_source(CandidateSource::KeyCache)
415            })
416            .collect();
417
418        composition.dropped_key_cache_per_crate +=
419            cap_to(&mut for_crate, limits.key_cache_per_crate);
420
421        for candidate in for_crate {
422            if limits.key_cache_total > 0 && composition.from_key_cache >= limits.key_cache_total {
423                composition.dropped_key_cache_total += 1;
424                continue;
425            }
426            if seen.insert(candidate.cache_key.clone()) {
427                composition.from_key_cache += 1;
428                candidates.push(candidate);
429            }
430        }
431    }
432
433    // Stamp demand position so the daemon can rank without the intent: it only
434    // receives the plan, and `PrefetchRequest::from_plan` drops everything else.
435    // Identity manifests already carry last-compile order; leave those alone.
436    for candidate in &mut candidates {
437        if candidate.demand_index.is_none() {
438            candidate.demand_index = demand_index.get(&candidate.crate_name).copied();
439        }
440    }
441
442    // Dispatch order (#617). Stable, so equal keys keep the confidence-merge
443    // order the sources were appended in.
444    sort_candidates_for_dispatch(&mut candidates);
445
446    Ok((execute_plan(planner_name, candidates), composition))
447}
448
449#[cfg(feature = "planning")]
450fn execute_plan(planner_name: &str, candidates: Vec<PrefetchCandidate>) -> PrefetchPlan {
451    let planner = planner_name.trim();
452    PrefetchPlan {
453        plan_id: None,
454        planner: Some(if planner.is_empty() {
455            "planner".to_string()
456        } else {
457            planner.to_string()
458        }),
459        disposition: PrefetchDisposition::Execute,
460        candidates,
461    }
462}
463
464/// crate name -> position in dependency order, for [`dispatch_sort_key`].
465#[cfg(feature = "planning")]
466fn demand_index_map(crate_order: &[String]) -> HashMap<String, u32> {
467    crate_order
468        .iter()
469        .enumerate()
470        .map(|(index, name)| (name.clone(), index as u32))
471        .collect()
472}
473
474#[cfg(feature = "planning")]
475fn crate_query_order(intent: &BuildIntent) -> Vec<String> {
476    let mut seen = HashSet::new();
477    intent
478        .crate_names
479        .iter()
480        .filter(|crate_name| seen.insert((*crate_name).clone()))
481        .cloned()
482        .collect()
483}
484
485#[cfg(feature = "planning")]
486fn order_candidates_by_crate_order(
487    mut candidates: Vec<PrefetchCandidate>,
488    intent: &BuildIntent,
489) -> Vec<PrefetchCandidate> {
490    if intent.crate_names.is_empty() {
491        return candidates;
492    }
493
494    let mut priority = HashMap::new();
495    for (index, crate_name) in crate_query_order(intent).iter().enumerate() {
496        priority.entry(crate_name.clone()).or_insert(index);
497    }
498
499    let mut indexed_candidates = candidates.drain(..).enumerate().collect::<Vec<_>>();
500    indexed_candidates.sort_by_key(|(index, candidate)| {
501        (
502            priority
503                .get(&candidate.crate_name)
504                .copied()
505                .unwrap_or(usize::MAX),
506            *index,
507        )
508    });
509    indexed_candidates
510        .into_iter()
511        .map(|(_, candidate)| candidate)
512        .collect()
513}
514
515/// Bounded model-checking harnesses for the planner's dispatch contract.
516///
517/// These mirror the three lexicographic priorities documented by
518/// [`dispatch_sort_key`] and cover every primitive input value, rather than a
519/// sampled set. Kani injects its support crate when `cargo kani` sets
520/// `cfg(kani)`, so normal builds carry no additional dependency.
521#[cfg(all(kani, feature = "planning"))]
522mod kani_proofs {
523    use super::*;
524
525    fn candidate(
526        source: CandidateSource,
527        compile_time_ms: Option<u64>,
528        demand_index: Option<u32>,
529    ) -> PrefetchCandidate {
530        PrefetchCandidate {
531            cache_key: String::new(),
532            crate_name: String::new(),
533            compile_time_ms,
534            size_bytes: None,
535            source,
536            demand_index,
537        }
538    }
539
540    #[kani::proof]
541    fn known_demand_always_precedes_unknown_demand() {
542        let demand_index: u32 = kani::any();
543        let known = candidate(kani::any(), kani::any(), Some(demand_index));
544        let unknown = candidate(kani::any(), kani::any(), None);
545        kani::cover!();
546
547        assert!(dispatch_sort_key(&known) < dispatch_sort_key(&unknown));
548    }
549
550    #[kani::proof]
551    fn earlier_urgency_bucket_always_wins() {
552        let earlier: u32 = kani::any();
553        let later: u32 = kani::any();
554        kani::assume(earlier / URGENCY_BUCKET < later / URGENCY_BUCKET);
555        let earlier_candidate = candidate(kani::any(), kani::any(), Some(earlier));
556        let later_candidate = candidate(kani::any(), kani::any(), Some(later));
557        kani::cover!();
558
559        assert!(dispatch_sort_key(&earlier_candidate) < dispatch_sort_key(&later_candidate));
560    }
561
562    #[kani::proof]
563    fn confidence_always_wins_within_a_bucket() {
564        let better_source: CandidateSource = kani::any();
565        let worse_source: CandidateSource = kani::any();
566        kani::assume(better_source.confidence_rank() < worse_source.confidence_rank());
567        let better_demand: u32 = kani::any();
568        let worse_demand: u32 = kani::any();
569        kani::assume(better_demand / URGENCY_BUCKET == worse_demand / URGENCY_BUCKET);
570        let better = candidate(better_source, kani::any(), Some(better_demand));
571        let worse = candidate(worse_source, kani::any(), Some(worse_demand));
572        kani::cover!();
573
574        assert!(dispatch_sort_key(&better) < dispatch_sort_key(&worse));
575    }
576
577    #[kani::proof]
578    fn higher_compile_cost_wins_after_urgency_and_confidence() {
579        let cheaper: u64 = kani::any();
580        let costlier: u64 = kani::any();
581        kani::assume(cheaper < costlier);
582        let cheaper_demand: u32 = kani::any();
583        let costlier_demand: u32 = kani::any();
584        kani::assume(cheaper_demand / URGENCY_BUCKET == costlier_demand / URGENCY_BUCKET);
585        let source: CandidateSource = kani::any();
586        let cheaper_candidate = candidate(source, Some(cheaper), Some(cheaper_demand));
587        let costlier_candidate = candidate(source, Some(costlier), Some(costlier_demand));
588        kani::cover!();
589
590        assert!(dispatch_sort_key(&costlier_candidate) < dispatch_sort_key(&cheaper_candidate));
591    }
592
593    #[kani::proof]
594    fn known_compile_cost_including_zero_precedes_unknown_cost() {
595        let known_cost: u64 = kani::any();
596        let demand_index: u32 = kani::any();
597        let source: CandidateSource = kani::any();
598        let known = candidate(source, Some(known_cost), Some(demand_index));
599        let unknown = candidate(source, None, Some(demand_index));
600        kani::cover!(known_cost == 0);
601
602        assert!(dispatch_sort_key(&known) < dispatch_sort_key(&unknown));
603    }
604
605    #[kani::proof]
606    fn dispatch_sort_preserves_equal_key_order() {
607        let source: CandidateSource = kani::any();
608        let compile_time_ms: Option<u64> = kani::any();
609        let demand_index: Option<u32> = kani::any();
610        let mut first = candidate(source, compile_time_ms, demand_index);
611        first.size_bytes = Some(1);
612        let mut second = candidate(source, compile_time_ms, demand_index);
613        second.size_bytes = Some(2);
614        let mut candidates = [first, second];
615
616        sort_candidates_for_dispatch(&mut candidates);
617        kani::cover!();
618
619        assert_eq!(candidates[0].size_bytes, Some(1));
620        assert_eq!(candidates[1].size_bytes, Some(2));
621    }
622}
623
624#[cfg(test)]
625mod tests {
626    use super::*;
627
628    #[cfg(feature = "planning")]
629    use std::collections::HashMap;
630
631    #[cfg(feature = "planning")]
632    use anyhow::anyhow;
633
634    #[cfg(feature = "planning")]
635    use proptest::prelude::*;
636
637    #[cfg(feature = "planning")]
638    #[derive(Default)]
639    struct FakePlannerDataSource {
640        identity_candidates: Vec<PrefetchCandidate>,
641        shard_candidates: Vec<PrefetchCandidate>,
642        shard_error: bool,
643        history_candidates: Vec<PrefetchCandidate>,
644        history_by_crate: HashMap<String, String>,
645        key_cache: HashMap<String, Vec<String>>,
646    }
647
648    #[cfg(feature = "planning")]
649    #[async_trait]
650    impl PlannerDataSource for FakePlannerDataSource {
651        async fn shard_candidates(
652            &self,
653            _namespace: &str,
654            _deps: &[(String, String)],
655        ) -> Result<Vec<PrefetchCandidate>> {
656            if self.shard_error {
657                Err(anyhow!("shard lookup failed"))
658            } else {
659                Ok(self.shard_candidates.clone())
660            }
661        }
662
663        async fn history_candidates(
664            &self,
665            crate_names: &[String],
666        ) -> Result<Vec<PrefetchCandidate>> {
667            if !self.history_candidates.is_empty() {
668                return Ok(self.history_candidates.clone());
669            }
670
671            Ok(crate_names
672                .iter()
673                .filter_map(|crate_name| {
674                    self.history_by_crate.get(crate_name).map(|cache_key| {
675                        PrefetchCandidate::new(cache_key.clone(), crate_name.clone())
676                    })
677                })
678                .collect())
679        }
680
681        async fn key_cache_keys_for_crate(&self, crate_name: &str) -> Result<Vec<String>> {
682            Ok(self.key_cache.get(crate_name).cloned().unwrap_or_default())
683        }
684
685        async fn identity_candidates(&self, _identity_key: &str) -> Result<Vec<PrefetchCandidate>> {
686            Ok(self.identity_candidates.clone())
687        }
688    }
689
690    #[test]
691    fn test_build_intent_serde_roundtrip() {
692        let intent = BuildIntent {
693            crate_names: vec!["serde".into(), "tokio".into()],
694            namespace: Some("x86_64/hash/release".into()),
695            cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
696            identity_key: None,
697        };
698
699        let json = serde_json::to_string(&intent).unwrap();
700        let parsed: BuildIntent = serde_json::from_str(&json).unwrap();
701        assert_eq!(parsed, intent);
702    }
703
704    #[test]
705    fn test_build_intent_defaults_missing_fields() {
706        let parsed: BuildIntent = serde_json::from_str(r#"{"crate_names":["serde"]}"#).unwrap();
707        assert_eq!(parsed.crate_names, vec!["serde"]);
708        assert!(parsed.namespace.is_none());
709        assert!(parsed.cargo_lock_deps.is_empty());
710        assert!(parsed.identity_key.is_none());
711    }
712
713    #[test]
714    fn test_prefetch_plan_serde_roundtrip() {
715        let plan = PrefetchPlan {
716            plan_id: Some("plan-1".into()),
717            planner: Some("local".into()),
718            disposition: PrefetchDisposition::Execute,
719            candidates: vec![PrefetchCandidate::new("abc".into(), "serde".into())],
720        };
721
722        let json = serde_json::to_string(&plan).unwrap();
723        let parsed: PrefetchPlan = serde_json::from_str(&json).unwrap();
724        assert_eq!(parsed, plan);
725    }
726
727    #[test]
728    fn test_prefetch_plan_missing_disposition_is_rejected() {
729        let err = serde_json::from_str::<PrefetchPlan>(
730            r#"{"planner":"legacy","candidates":[{"cache_key":"abc","crate_name":"serde"}]}"#,
731        )
732        .unwrap_err();
733        assert!(err.to_string().contains("missing field"));
734    }
735
736    #[test]
737    fn test_prefetch_plan_do_nothing_roundtrip() {
738        let plan = PrefetchPlan {
739            plan_id: Some("plan-2".into()),
740            planner: Some("remote".into()),
741            disposition: PrefetchDisposition::DoNothing,
742            candidates: vec![],
743        };
744
745        let json = serde_json::to_string(&plan).unwrap();
746        let parsed: PrefetchPlan = serde_json::from_str(&json).unwrap();
747        assert_eq!(parsed, plan);
748    }
749
750    #[cfg(feature = "planning")]
751    #[tokio::test]
752    async fn test_build_prefetch_plan_prefers_shard_candidates() {
753        let source = FakePlannerDataSource {
754            shard_candidates: vec![PrefetchCandidate::new("from-shard".into(), "serde".into())],
755            ..Default::default()
756        };
757        let intent = BuildIntent {
758            crate_names: vec!["serde".into()],
759            namespace: Some("linux/hash/release".into()),
760            cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
761            identity_key: None,
762        };
763
764        let (plan, composition) =
765            build_prefetch_plan_with_limits(&source, &intent, "fallback", PlanLimits::default())
766                .await
767                .unwrap();
768
769        assert_eq!(plan.disposition, PrefetchDisposition::Execute);
770        assert_eq!(plan.planner.as_deref(), Some("fallback"));
771        assert_eq!(plan.candidates.len(), 1);
772        assert_eq!(plan.candidates[0].cache_key, "from-shard");
773        assert_eq!(composition.from_shards, 1);
774    }
775
776    #[cfg(feature = "planning")]
777    #[tokio::test]
778    async fn test_build_prefetch_plan_prefers_identity_manifest_over_shards() {
779        let source = FakePlannerDataSource {
780            identity_candidates: vec![PrefetchCandidate::new(
781                "from-identity".into(),
782                "serde".into(),
783            )],
784            shard_candidates: vec![PrefetchCandidate::new("from-shard".into(), "serde".into())],
785            ..Default::default()
786        };
787        let intent = BuildIntent {
788            crate_names: vec!["serde".into()],
789            namespace: Some("linux/hash/release".into()),
790            cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
791            identity_key: Some("id/abcd/x86_64-unknown-linux-gnu/release".into()),
792        };
793
794        let plan = build_prefetch_plan(&source, &intent, "fallback")
795            .await
796            .unwrap();
797
798        assert_eq!(plan.candidates.len(), 1);
799        assert_eq!(plan.candidates[0].cache_key, "from-identity");
800        assert_eq!(plan.candidates[0].source, CandidateSource::Manifest);
801        assert_eq!(plan.candidates[0].demand_index, Some(0));
802    }
803
804    #[cfg(feature = "planning")]
805    #[tokio::test]
806    async fn test_build_prefetch_plan_falls_back_to_history_and_key_cache() {
807        let mut source = FakePlannerDataSource {
808            shard_error: true,
809            history_candidates: vec![PrefetchCandidate::new("history-key".into(), "serde".into())],
810            ..Default::default()
811        };
812        source.key_cache.insert(
813            "tokio".into(),
814            vec!["tokio-key".into(), "history-key".into()],
815        );
816
817        let intent = BuildIntent {
818            crate_names: vec!["serde".into(), "tokio".into()],
819            namespace: Some("linux/hash/debug".into()),
820            cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
821            identity_key: None,
822        };
823
824        let plan = build_prefetch_plan(&source, &intent, "fallback")
825            .await
826            .unwrap();
827
828        assert_eq!(plan.disposition, PrefetchDisposition::Execute);
829        assert_eq!(plan.candidates.len(), 2);
830        assert_eq!(plan.candidates[0].cache_key, "history-key");
831        assert_eq!(plan.candidates[1].cache_key, "tokio-key");
832    }
833
834    #[cfg(feature = "planning")]
835    #[tokio::test]
836    async fn test_build_prefetch_plan_orders_shard_candidates_by_crate_order() {
837        let source = FakePlannerDataSource {
838            shard_candidates: vec![
839                PrefetchCandidate::new("app-key".into(), "app".into()),
840                PrefetchCandidate::new("dep-key".into(), "dep".into()),
841                PrefetchCandidate::new("middle-key".into(), "middle".into()),
842            ],
843            ..Default::default()
844        };
845        let intent = BuildIntent {
846            crate_names: vec!["dep".into(), "middle".into(), "app".into()],
847            namespace: Some("linux/hash/debug".into()),
848            cargo_lock_deps: vec![("dep".into(), "1.0.0".into())],
849            identity_key: None,
850        };
851
852        let plan = build_prefetch_plan(&source, &intent, "fallback")
853            .await
854            .unwrap();
855
856        let keys = plan
857            .candidates
858            .iter()
859            .map(|candidate| candidate.cache_key.as_str())
860            .collect::<Vec<_>>();
861        assert_eq!(keys, vec!["dep-key", "middle-key", "app-key"]);
862    }
863
864    /// A PARTIAL shard hit must not suppress the lower-confidence sources
865    /// (kunobi-ninja/kache#614).
866    ///
867    /// Shard matching is exact per bucket, so a dependency bump invalidates
868    /// one bucket while the rest still match. The planner used to return as
869    /// soon as shards produced anything, so the crates in the missed buckets
870    /// were dropped from the plan even though history and the key cache could
871    /// resolve them.
872    #[cfg(feature = "planning")]
873    #[tokio::test]
874    async fn test_build_prefetch_plan_fills_crates_a_partial_shard_hit_missed() {
875        let mut source = FakePlannerDataSource {
876            // Only `dep` is in a bucket that still matches.
877            shard_candidates: vec![PrefetchCandidate::new("dep-shard-key".into(), "dep".into())],
878            history_by_crate: HashMap::from([("middle".into(), "middle-history-key".into())]),
879            ..Default::default()
880        };
881        source
882            .key_cache
883            .insert("app".into(), vec!["app-key-cache-key".into()]);
884
885        let intent = BuildIntent {
886            crate_names: vec!["dep".into(), "middle".into(), "app".into()],
887            namespace: Some("linux/hash/debug".into()),
888            cargo_lock_deps: vec![("dep".into(), "1.0.0".into())],
889            identity_key: None,
890        };
891
892        let plan = build_prefetch_plan(&source, &intent, "fallback")
893            .await
894            .unwrap();
895
896        let keys = plan
897            .candidates
898            .iter()
899            .map(|candidate| candidate.cache_key.as_str())
900            .collect::<Vec<_>>();
901        assert_eq!(
902            keys,
903            vec!["dep-shard-key", "middle-history-key", "app-key-cache-key"],
904            "each source should fill the crates the higher-confidence ones left unresolved"
905        );
906    }
907
908    /// A crate the shards already resolved is not re-queried from the
909    /// lower-confidence sources (#614): shard keys are exact, history and the
910    /// key cache are not, so they only fill gaps.
911    #[cfg(feature = "planning")]
912    #[tokio::test]
913    async fn test_build_prefetch_plan_does_not_requery_shard_resolved_crates() {
914        let mut source = FakePlannerDataSource {
915            shard_candidates: vec![PrefetchCandidate::new(
916                "serde-shard-key".into(),
917                "serde".into(),
918            )],
919            history_by_crate: HashMap::from([("serde".into(), "serde-stale-history-key".into())]),
920            ..Default::default()
921        };
922        source
923            .key_cache
924            .insert("serde".into(), vec!["serde-stale-key-cache-key".into()]);
925
926        let intent = BuildIntent {
927            crate_names: vec!["serde".into()],
928            namespace: Some("linux/hash/debug".into()),
929            cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
930            identity_key: None,
931        };
932
933        let plan = build_prefetch_plan(&source, &intent, "fallback")
934            .await
935            .unwrap();
936
937        let keys = plan
938            .candidates
939            .iter()
940            .map(|candidate| candidate.cache_key.as_str())
941            .collect::<Vec<_>>();
942        assert_eq!(keys, vec!["serde-shard-key"]);
943    }
944
945    #[cfg(feature = "planning")]
946    #[tokio::test]
947    async fn test_build_prefetch_plan_queries_history_by_crate_order() {
948        let source = FakePlannerDataSource {
949            history_by_crate: HashMap::from([
950                ("app".into(), "app-key".into()),
951                ("dep".into(), "dep-key".into()),
952                ("middle".into(), "middle-key".into()),
953            ]),
954            ..Default::default()
955        };
956        let intent = BuildIntent {
957            crate_names: vec!["dep".into(), "middle".into(), "app".into()],
958            namespace: None,
959            cargo_lock_deps: vec![],
960            identity_key: None,
961        };
962
963        let plan = build_prefetch_plan(&source, &intent, "fallback")
964            .await
965            .unwrap();
966
967        let keys = plan
968            .candidates
969            .iter()
970            .map(|candidate| candidate.cache_key.as_str())
971            .collect::<Vec<_>>();
972        assert_eq!(keys, vec!["dep-key", "middle-key", "app-key"]);
973    }
974
975    #[cfg(feature = "planning")]
976    proptest! {
977        #[test]
978        fn property_crate_query_order_preserves_first_occurrence(
979            crate_ids in proptest::collection::vec(any::<u8>(), 1..64),
980        ) {
981            let mut crate_names = crate_ids
982                .into_iter()
983                .map(|id| format!("crate-{id}"))
984                .collect::<Vec<_>>();
985            crate_names.push(crate_names[0].clone());
986            let intent = BuildIntent {
987                crate_names: crate_names.clone(),
988                ..Default::default()
989            };
990
991            let mut expected = Vec::new();
992            for crate_name in crate_names {
993                if !expected.contains(&crate_name) {
994                    expected.push(crate_name);
995                }
996            }
997
998            prop_assert_eq!(crate_query_order(&intent), expected);
999        }
1000
1001        #[test]
1002        fn property_candidate_ordering_is_a_stable_priority_permutation(
1003            requested_ids in proptest::collection::vec(0u8..8, 1..24),
1004            candidate_ids in proptest::collection::vec(0u8..16, 0..48),
1005        ) {
1006            let requested_names = requested_ids
1007                .into_iter()
1008                .map(|id| format!("crate-{id}"))
1009                .collect::<Vec<_>>();
1010            let intent = BuildIntent {
1011                crate_names: requested_names.clone(),
1012                ..Default::default()
1013            };
1014
1015            // Interleave two unknown and two known candidates so even minimal
1016            // shrunk inputs exercise both priority and stable tie ordering.
1017            let mut candidate_names = Vec::with_capacity(candidate_ids.len() + 4);
1018            candidate_names.push("not-requested".to_string());
1019            candidate_names.push(requested_names[0].clone());
1020            candidate_names.extend(
1021                candidate_ids
1022                    .into_iter()
1023                    .map(|id| format!("crate-{id}")),
1024            );
1025            candidate_names.push("not-requested".to_string());
1026            candidate_names.push(requested_names[0].clone());
1027
1028            let candidates = candidate_names
1029                .into_iter()
1030                .enumerate()
1031                .map(|(index, crate_name)| {
1032                    PrefetchCandidate::new(format!("key-{index}"), crate_name)
1033                })
1034                .collect::<Vec<_>>();
1035
1036            let mut expected = candidates.clone();
1037            expected.sort_by_key(|candidate| {
1038                requested_names
1039                    .iter()
1040                    .position(|name| name == &candidate.crate_name)
1041                    .unwrap_or(usize::MAX)
1042            });
1043
1044            prop_assert_eq!(
1045                order_candidates_by_crate_order(candidates, &intent),
1046                expected
1047            );
1048        }
1049    }
1050    // ── Composition caps and ranking (#616, #617) ────────────────────────
1051
1052    #[cfg(feature = "planning")]
1053    fn candidate(
1054        key: &str,
1055        crate_name: &str,
1056        source: CandidateSource,
1057        compile_time_ms: Option<u64>,
1058        demand_index: Option<u32>,
1059    ) -> PrefetchCandidate {
1060        PrefetchCandidate {
1061            cache_key: key.into(),
1062            crate_name: crate_name.into(),
1063            compile_time_ms,
1064            size_bytes: None,
1065            source,
1066            demand_index,
1067        }
1068    }
1069
1070    /// Confidence ordering is the whole point of the rank; assert the ORDER,
1071    /// not the literal numbers.
1072    #[test]
1073    fn test_confidence_rank_orders_sources() {
1074        assert!(
1075            CandidateSource::Manifest.confidence_rank() < CandidateSource::Shard.confidence_rank()
1076        );
1077        assert!(
1078            CandidateSource::Shard.confidence_rank() < CandidateSource::History.confidence_rank()
1079        );
1080        assert!(
1081            CandidateSource::History.confidence_rank()
1082                < CandidateSource::KeyCache.confidence_rank()
1083        );
1084        assert!(
1085            CandidateSource::KeyCache.confidence_rank()
1086                < CandidateSource::Unknown.confidence_rank()
1087        );
1088    }
1089
1090    /// An unrecognised source name degrades to `Unknown` instead of failing the
1091    /// whole plan (forward compatibility with a newer planner).
1092    #[test]
1093    fn test_unknown_candidate_source_deserializes() {
1094        let candidate: PrefetchCandidate =
1095            serde_json::from_str(r#"{"cache_key":"k","crate_name":"c","source":"telepathy"}"#)
1096                .unwrap();
1097        assert_eq!(candidate.source, CandidateSource::Unknown);
1098    }
1099
1100    /// A pre-#617 candidate has no metadata and must not be rejected, and the
1101    /// missing fields must read as unknown rather than zero.
1102    #[test]
1103    fn test_legacy_candidate_wire_still_parses() {
1104        let candidate: PrefetchCandidate =
1105            serde_json::from_str(r#"{"cache_key":"k","crate_name":"c"}"#).unwrap();
1106        assert_eq!(candidate.compile_time_ms, None);
1107        assert_eq!(candidate.size_bytes, None);
1108        assert_eq!(candidate.demand_index, None);
1109        assert_eq!(candidate.source, CandidateSource::Unknown);
1110    }
1111
1112    /// Urgency dominates value: prefetch races the build, so a costly artifact
1113    /// needed late loses to a cheaper one needed now (#617).
1114    #[cfg(feature = "planning")]
1115    #[test]
1116    fn test_dispatch_order_prefers_urgency_over_value() {
1117        let urgent_cheap = candidate("a", "a", CandidateSource::Shard, Some(10), Some(0));
1118        let late_expensive = candidate("b", "b", CandidateSource::Shard, Some(100_000), Some(500));
1119        assert!(dispatch_sort_key(&urgent_cheap) < dispatch_sort_key(&late_expensive));
1120    }
1121
1122    /// Within one urgency window, value decides.
1123    #[cfg(feature = "planning")]
1124    #[test]
1125    fn test_dispatch_order_prefers_value_within_a_bucket() {
1126        let cheap = candidate("a", "a", CandidateSource::Shard, Some(10), Some(0));
1127        let costly = candidate("b", "b", CandidateSource::Shard, Some(5_000), Some(1));
1128        assert!(
1129            dispatch_sort_key(&costly) < dispatch_sort_key(&cheap),
1130            "same bucket, same source: the expensive rebuild goes first"
1131        );
1132    }
1133
1134    /// Positions inside one bucket are treated as equally urgent: lockfile
1135    /// order only approximates demand time, so finer distinctions are false
1136    /// precision.
1137    #[cfg(feature = "planning")]
1138    #[test]
1139    fn test_dispatch_order_buckets_nearby_positions_together() {
1140        let first = candidate("a", "a", CandidateSource::Shard, Some(10), Some(0));
1141        let nearby = candidate("b", "b", CandidateSource::Shard, Some(10), Some(15));
1142        let next_bucket = candidate("c", "c", CandidateSource::Shard, Some(10), Some(16));
1143        assert_eq!(
1144            dispatch_sort_key(&first),
1145            dispatch_sort_key(&nearby),
1146            "positions 0 and 15 share a bucket"
1147        );
1148        assert!(
1149            dispatch_sort_key(&nearby) < dispatch_sort_key(&next_bucket),
1150            "position 16 starts the next bucket"
1151        );
1152    }
1153
1154    /// Within a bucket, confidence beats value: a shard match is worth more
1155    /// than a bigger number from a guess.
1156    #[cfg(feature = "planning")]
1157    #[test]
1158    fn test_dispatch_order_prefers_confidence_within_a_bucket() {
1159        let shard_cheap = candidate("a", "a", CandidateSource::Shard, Some(1), Some(0));
1160        let guess_costly = candidate("b", "b", CandidateSource::KeyCache, Some(99_999), Some(0));
1161        assert!(dispatch_sort_key(&shard_cheap) < dispatch_sort_key(&guess_costly));
1162    }
1163
1164    /// Unknown cost sorts after known cost, but is not dropped and is not
1165    /// treated as zero-value against a *different* bucket.
1166    #[cfg(feature = "planning")]
1167    #[test]
1168    fn test_dispatch_order_places_unknown_cost_last_within_its_bucket() {
1169        let known = candidate("a", "a", CandidateSource::Shard, Some(1), Some(0));
1170        let known_zero = candidate("zero", "zero", CandidateSource::Shard, Some(0), Some(0));
1171        let unknown = candidate("b", "b", CandidateSource::Shard, None, Some(0));
1172        assert!(dispatch_sort_key(&known) < dispatch_sort_key(&unknown));
1173        assert!(dispatch_sort_key(&known_zero) < dispatch_sort_key(&unknown));
1174
1175        // ...but an unknown-cost candidate needed NOW still beats a known-cost
1176        // one needed much later.
1177        let late_known = candidate("c", "c", CandidateSource::Shard, Some(50_000), Some(999));
1178        assert!(dispatch_sort_key(&unknown) < dispatch_sort_key(&late_known));
1179    }
1180
1181    /// A candidate outside the intent's crate list sorts last rather than first.
1182    #[cfg(feature = "planning")]
1183    #[test]
1184    fn test_dispatch_order_places_unknown_demand_last() {
1185        let known = candidate("a", "a", CandidateSource::Shard, Some(1), Some(100_000));
1186        let no_demand = candidate("b", "b", CandidateSource::Shard, Some(50_000), None);
1187        assert!(dispatch_sort_key(&known) < dispatch_sort_key(&no_demand));
1188    }
1189
1190    /// The production dispatch path applies the key with a stable sort.
1191    #[cfg(feature = "planning")]
1192    #[test]
1193    fn test_dispatch_sort_orders_candidates_and_preserves_ties() {
1194        let late = candidate("late", "late", CandidateSource::Shard, Some(1), Some(32));
1195        let first_tie = candidate("first", "first", CandidateSource::Shard, Some(10), Some(0));
1196        let second_tie = candidate(
1197            "second",
1198            "second",
1199            CandidateSource::Shard,
1200            Some(10),
1201            Some(0),
1202        );
1203        let mut candidates = vec![late, first_tie, second_tie];
1204
1205        sort_candidates_for_dispatch(&mut candidates);
1206
1207        let keys = candidates
1208            .iter()
1209            .map(|candidate| candidate.cache_key.as_str())
1210            .collect::<Vec<_>>();
1211        assert_eq!(keys, vec!["first", "second", "late"]);
1212    }
1213
1214    /// The per-crate key-cache cap keeps the first N and reports the rest.
1215    #[cfg(feature = "planning")]
1216    #[tokio::test]
1217    async fn test_key_cache_per_crate_cap() {
1218        let mut source = FakePlannerDataSource::default();
1219        source.key_cache.insert(
1220            "serde".into(),
1221            vec!["k1".into(), "k2".into(), "k3".into(), "k4".into()],
1222        );
1223        let intent = BuildIntent {
1224            crate_names: vec!["serde".into()],
1225            ..Default::default()
1226        };
1227
1228        let (plan, composition) = build_prefetch_plan_with_limits(
1229            &source,
1230            &intent,
1231            "fallback",
1232            PlanLimits {
1233                key_cache_per_crate: 2,
1234                ..PlanLimits::default()
1235            },
1236        )
1237        .await
1238        .unwrap();
1239
1240        assert_eq!(plan.candidates.len(), 2, "capped to two variants");
1241        assert_eq!(composition.from_key_cache, 2);
1242        assert_eq!(composition.dropped_key_cache_per_crate, 2);
1243        assert!(composition.dropped_total() > 0, "truncation is visible");
1244    }
1245
1246    /// The plan-wide key-cache cap stops one weak source filling the plan, and
1247    /// counts what whole crates it skipped.
1248    #[cfg(feature = "planning")]
1249    #[tokio::test]
1250    async fn test_key_cache_total_cap_counts_skipped_crates() {
1251        let mut source = FakePlannerDataSource::default();
1252        for name in ["a", "b", "c"] {
1253            source
1254                .key_cache
1255                .insert(name.into(), vec![format!("{name}-k1")]);
1256        }
1257        let intent = BuildIntent {
1258            crate_names: vec!["a".into(), "b".into(), "c".into()],
1259            ..Default::default()
1260        };
1261
1262        let (plan, composition) = build_prefetch_plan_with_limits(
1263            &source,
1264            &intent,
1265            "fallback",
1266            PlanLimits {
1267                key_cache_total: 1,
1268                ..PlanLimits::default()
1269            },
1270        )
1271        .await
1272        .unwrap();
1273
1274        assert_eq!(plan.candidates.len(), 1);
1275        assert_eq!(composition.from_key_cache, 1);
1276        assert_eq!(
1277            composition.dropped_key_cache_total, 2,
1278            "the two skipped crates are counted, not silently missing"
1279        );
1280    }
1281
1282    /// `0` disables a cap rather than dropping everything.
1283    #[cfg(feature = "planning")]
1284    #[tokio::test]
1285    async fn test_zero_limit_disables_the_cap() {
1286        let mut source = FakePlannerDataSource::default();
1287        source
1288            .key_cache
1289            .insert("serde".into(), vec!["k1".into(), "k2".into(), "k3".into()]);
1290        let intent = BuildIntent {
1291            crate_names: vec!["serde".into()],
1292            ..Default::default()
1293        };
1294
1295        let (plan, composition) = build_prefetch_plan_with_limits(
1296            &source,
1297            &intent,
1298            "fallback",
1299            PlanLimits {
1300                key_cache_per_crate: 0,
1301                key_cache_total: 0,
1302                history_per_crate: 0,
1303            },
1304        )
1305        .await
1306        .unwrap();
1307
1308        assert_eq!(plan.candidates.len(), 3);
1309        assert_eq!(composition.dropped_total(), 0);
1310    }
1311
1312    /// Candidates are stamped with their source and demand position, so the
1313    /// daemon can rank without the intent (it only receives the plan).
1314    #[cfg(feature = "planning")]
1315    #[tokio::test]
1316    async fn test_plan_stamps_source_and_demand_index() {
1317        let source = FakePlannerDataSource {
1318            history_by_crate: HashMap::from([
1319                ("dep".into(), "dep-key".into()),
1320                ("app".into(), "app-key".into()),
1321            ]),
1322            ..Default::default()
1323        };
1324        let intent = BuildIntent {
1325            crate_names: vec!["dep".into(), "app".into()],
1326            ..Default::default()
1327        };
1328
1329        let (plan, _) =
1330            build_prefetch_plan_with_limits(&source, &intent, "fallback", PlanLimits::default())
1331                .await
1332                .unwrap();
1333
1334        let dep = plan
1335            .candidates
1336            .iter()
1337            .find(|c| c.crate_name == "dep")
1338            .expect("dep planned");
1339        let app = plan
1340            .candidates
1341            .iter()
1342            .find(|c| c.crate_name == "app")
1343            .expect("app planned");
1344        assert_eq!(dep.source, CandidateSource::History);
1345        assert_eq!(dep.demand_index, Some(0));
1346        assert_eq!(app.demand_index, Some(1), "position follows crate order");
1347    }
1348
1349    /// `cap_to` computes the DROPPED count, not a ratio. 5 items capped to 2
1350    /// drops 3; a divide would say 2. (Mutation-driven: `-` vs `/` agree on
1351    /// 4-and-2, so the earlier cap test could not tell them apart.)
1352    #[cfg(feature = "planning")]
1353    #[test]
1354    fn test_cap_to_reports_the_dropped_count() {
1355        let mut items: Vec<PrefetchCandidate> = (0..5)
1356            .map(|i| PrefetchCandidate::new(format!("k{i}"), "c".into()))
1357            .collect();
1358        assert_eq!(cap_to(&mut items, 2), 3);
1359        assert_eq!(items.len(), 2);
1360        assert_eq!(items[0].cache_key, "k0", "keeps the FIRST N");
1361
1362        let mut exact: Vec<PrefetchCandidate> =
1363            vec![PrefetchCandidate::new("a".into(), "c".into())];
1364        assert_eq!(cap_to(&mut exact, 1), 0, "exactly at the limit drops none");
1365        assert_eq!(cap_to(&mut exact, 0), 0, "0 disables the cap");
1366        assert_eq!(exact.len(), 1);
1367    }
1368
1369    /// `dropped_total` sums every drop class. Each field is distinct so a
1370    /// swapped operator or a missing term changes the result.
1371    #[test]
1372    fn test_dropped_total_sums_every_class() {
1373        let composition = PlanComposition {
1374            dropped_history_per_crate: 1,
1375            dropped_key_cache_per_crate: 2,
1376            dropped_key_cache_total: 4,
1377            ..PlanComposition::default()
1378        };
1379        assert_eq!(composition.dropped_total(), 7);
1380        assert_eq!(PlanComposition::default().dropped_total(), 0);
1381    }
1382
1383    /// The per-crate history cap keeps the most-recent entries and reports the
1384    /// rest. History is ordered most-recently-used first by both data sources,
1385    /// so the cap keeps the freshest variants.
1386    #[cfg(feature = "planning")]
1387    #[tokio::test]
1388    async fn test_history_per_crate_cap() {
1389        let source = FakePlannerDataSource {
1390            history_candidates: vec![
1391                PrefetchCandidate::new("serde-newest".into(), "serde".into()),
1392                PrefetchCandidate::new("serde-older".into(), "serde".into()),
1393                PrefetchCandidate::new("serde-oldest".into(), "serde".into()),
1394            ],
1395            ..Default::default()
1396        };
1397        let intent = BuildIntent {
1398            crate_names: vec!["serde".into()],
1399            ..Default::default()
1400        };
1401
1402        let (plan, composition) = build_prefetch_plan_with_limits(
1403            &source,
1404            &intent,
1405            "fallback",
1406            PlanLimits {
1407                history_per_crate: 1,
1408                ..PlanLimits::default()
1409            },
1410        )
1411        .await
1412        .unwrap();
1413
1414        assert_eq!(composition.from_history, 1);
1415        assert_eq!(composition.dropped_history_per_crate, 2);
1416        assert_eq!(
1417            plan.candidates
1418                .iter()
1419                .map(|c| c.cache_key.as_str())
1420                .collect::<Vec<_>>(),
1421            vec!["serde-newest"]
1422        );
1423    }
1424
1425    /// Once the plan-wide key-cache budget is spent, a later crate is skipped
1426    /// WHOLESALE, and everything it would have offered is counted, not just the
1427    /// per-crate-capped subset.
1428    #[cfg(feature = "planning")]
1429    #[tokio::test]
1430    async fn test_key_cache_total_skips_a_whole_crate_and_counts_all_of_it() {
1431        let mut source = FakePlannerDataSource::default();
1432        source.key_cache.insert("a".into(), vec!["a1".into()]);
1433        source.key_cache.insert(
1434            "b".into(),
1435            vec![
1436                "b1".into(),
1437                "b2".into(),
1438                "b3".into(),
1439                "b4".into(),
1440                "b5".into(),
1441            ],
1442        );
1443        let intent = BuildIntent {
1444            crate_names: vec!["a".into(), "b".into()],
1445            ..Default::default()
1446        };
1447
1448        let (_plan, composition) = build_prefetch_plan_with_limits(
1449            &source,
1450            &intent,
1451            "fallback",
1452            PlanLimits {
1453                key_cache_per_crate: 2,
1454                key_cache_total: 1,
1455                history_per_crate: 2,
1456            },
1457        )
1458        .await
1459        .unwrap();
1460
1461        assert_eq!(composition.from_key_cache, 1, "only crate `a` fits");
1462        assert_eq!(
1463            composition.dropped_key_cache_total, 5,
1464            "all five of `b`'s variants are counted, not the two a per-crate \
1465             cap would have admitted"
1466        );
1467        assert_eq!(
1468            composition.dropped_key_cache_per_crate, 0,
1469            "`b` was skipped before the per-crate cap applied"
1470        );
1471    }
1472
1473    /// The budget can also run out PART WAY through one crate's variants: the
1474    /// remainder is dropped and counted individually.
1475    #[cfg(feature = "planning")]
1476    #[tokio::test]
1477    async fn test_key_cache_total_can_run_out_mid_crate() {
1478        let mut source = FakePlannerDataSource::default();
1479        source
1480            .key_cache
1481            .insert("serde".into(), vec!["k1".into(), "k2".into(), "k3".into()]);
1482        let intent = BuildIntent {
1483            crate_names: vec!["serde".into()],
1484            ..Default::default()
1485        };
1486
1487        let (_plan, composition) = build_prefetch_plan_with_limits(
1488            &source,
1489            &intent,
1490            "fallback",
1491            PlanLimits {
1492                key_cache_per_crate: 2,
1493                key_cache_total: 1,
1494                history_per_crate: 2,
1495            },
1496        )
1497        .await
1498        .unwrap();
1499
1500        assert_eq!(composition.from_key_cache, 1);
1501        assert_eq!(
1502            composition.dropped_key_cache_per_crate, 1,
1503            "k3 dropped by the per-crate cap"
1504        );
1505        assert_eq!(
1506            composition.dropped_key_cache_total, 1,
1507            "k2 dropped by the plan-wide cap, mid-crate"
1508        );
1509    }
1510}