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