Skip to main content

mbx_cache_core/
agent.rs

1use crate::uploads::{ConnectionUploads, UploadQueue, UploadSink};
2use crate::{
3    ActionPrediction, ActionPromiseCompletion, ActionPromiseState, CacheDigest, CacheDirectory,
4    LocalActionCache, LocalCas, MAX_STAGED_BLOB_PACK_BYTES, ManifestPutOutcome, RemoteActionResult,
5    RemoteCacheClient, RemoteCacheMode, RustcMetadata, TaskActionManifest, blob_pack_chunk,
6    canonical_json,
7};
8use eyre::{Context, Result, bail};
9use futures_util::{FutureExt, StreamExt, future::BoxFuture, stream};
10use log::{debug, warn};
11use serde::{Deserialize, Serialize};
12use std::collections::{BTreeMap, BTreeSet};
13use std::fs;
14use std::path::{Path, PathBuf};
15use std::sync::atomic::{AtomicU64, Ordering};
16use std::sync::{Arc, Mutex, Weak};
17use std::time::{Duration, Instant, SystemTime};
18
19/// Fleet claims are leases rather than correctness locks. This only bounds how
20/// long one shim waits through repeated pending responses before degrading to
21/// an ordinary compilation; the server independently expires abandoned claims.
22const ACTION_PROMISE_WAIT: Duration = Duration::from_secs(60 * 60);
23use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};
24
25mod file_digest;
26mod manifest;
27mod prefetch;
28mod stats;
29mod wire;
30
31#[cfg(test)]
32pub(crate) use prefetch::select_prefetch_actions;
33
34pub use file_digest::{
35    FileDigestCache, FileDigestResolution, FileDigestScope, FileIdentity, FileObjectIdentity,
36    FileSnapshot, NoFileDigestCache, RecordedFileDigest, ValidatedFileDigest, digest_file,
37    digest_file_validated,
38};
39pub use manifest::{
40    is_task_identity, task_manifest_actions,
41    update_task_predictions_in_order as merge_task_action_predictions,
42};
43use manifest::{
44    merge_task_manifests, task_manifest_dir, update_task_predictions, validate_task_identity,
45    validate_task_manifest,
46};
47pub use stats::{AgentStats, CompilerStats};
48use wire::MAX_REQUEST_BYTES;
49pub use wire::{
50    AGENT_PROTOCOL_VERSION, ActionDiagnostic, AgentEvent, AgentEventObserver, AgentRequest,
51    AgentResponse, PinnedFile, PinnedState, RestoreStats, WrapperSpan, WrapperTiming,
52};
53
54const MAX_EXECUTABLE_IDENTITIES: usize = 64;
55const MAX_EXECUTABLE_IDENTITY_SIZE: usize = 64 * 1024;
56const MAX_EXECUTABLE_IDENTITY_BYTES: usize = 256 * 1024;
57const TASK_ACTION_MANIFEST_VERSION: u8 = 1;
58const MAX_TASK_ACTION_PREDICTIONS: usize = 16 * 1024;
59const TASK_ACTION_LRU_REFRESH_MARGIN: usize = MAX_TASK_ACTION_PREDICTIONS / 16;
60/// Longest shim diagnostic the agent accepts, in bytes.
61const MAX_DIAGNOSTIC_BYTES: usize = 4 * 1024;
62/// Most distinct shim diagnostics one session surfaces before going quiet.
63///
64/// A failure mode that repeats across a build tends to repeat with the same
65/// message, which deduplication already collapses; the cap only guards
66/// against a message that embeds something unique per compilation.
67const MAX_DIAGNOSTICS: usize = 128;
68const ACTION_DIAGNOSTIC_PREFIX: &str = "@mbx-action-diagnostic\t";
69
70#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
71enum DiagnosticLevel {
72    Warning,
73    Error,
74}
75
76impl DiagnosticLevel {
77    fn label(self) -> &'static str {
78        match self {
79            Self::Warning => "warning",
80            Self::Error => "error",
81        }
82    }
83}
84
85/// Most file identities one digest lookup or record may carry.
86const MAX_FILE_DIGEST_BATCH: usize = 16 * 1024;
87/// Most recorded file digests one session retains across every scope.
88///
89/// Roughly two orders of magnitude above the largest workspace measured; the
90/// cap exists so a pathological build bounds the agent instead of growing it.
91const MAX_FILE_DIGEST_ENTRIES: usize = 1024 * 1024;
92/// Distinct cold file reads allowed at once. Identical reads coalesce before
93/// this limit, while the limit prevents a wide cold build from flooding NFS.
94const MAX_CONCURRENT_FILE_DIGESTS: usize = 8;
95const MAX_REMOTE_TRANSFERS: usize = 64;
96const MAX_PREFETCH_TRANSFERS: usize = 48;
97/// Most blob packs downloaded at once.
98///
99/// Packs are individually bounded to keep their staging footprint predictable.
100/// Allowing a few independent streams prevents one large Azure object stream
101/// from serializing an entire workspace restore.
102const MAX_CONCURRENT_BLOB_PACKS: usize = 8;
103/// Most predicted actions whose complete output closures are downloaded
104/// speculatively for one task.
105///
106/// Manifests deliberately retain predictions across compatible builds, so a
107/// large workspace can describe far more actions than the next invocation
108/// will request. Keep the most expensive recorded actions warm and let
109/// foreground lookups fetch only unusually large tails on demand. The remote
110/// download-byte budget remains the primary bound on speculative transfer. A
111/// 1,024-action ceiling covers the measured large-workspace manifests while
112/// retaining a guard against pathological manifests and stale prediction sets.
113const MAX_PREFETCH_ACTIONS: usize = 1024;
114// Resolve speculative actions progressively so the most valuable predictions
115// become usable first and cancellation at the end of a build abandons a small
116// tail instead of one workspace-sized transfer wave. Action-result lookup is
117// still batched independently; this only bounds each output-closure download.
118const MAX_PREFETCH_ACTION_WAVE: usize = 32;
119/// Batched action lookups issued at once.
120///
121/// One request already asks about hundreds of actions. Serial batches keep a
122/// fleet of concurrent CI jobs from multiplying metadata pressure while blob
123/// transfers from earlier answers are also underway.
124const MAX_PREFETCH_BATCH_LOOKUPS: usize = 1;
125const PREFETCH_ACTION_BATCH_DELAY: Duration = Duration::from_millis(5);
126const MAX_PREFETCH_DIRECTORY_OBJECTS: usize = 100_000;
127const MAX_PREFETCH_OBJECTS_PER_WAVE: usize = 100_000;
128const DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES: u64 = 5 * 1024 * 1024 * 1024;
129
130/// Names the identities a task may inherit predictions from, consulted only
131/// once nothing has been recorded under its own.
132pub type TaskFallbacks = Arc<dyn Fn() -> Vec<String> + Send + Sync>;
133type FileDigestFlightKey = (FileDigestScope, FileIdentity);
134type FileDigestFlights = BTreeMap<FileDigestFlightKey, Weak<FileDigestFlight>>;
135
136struct FileDigestFlight {
137    lock: tokio::sync::Mutex<()>,
138    resolution: Mutex<Option<FileDigestResolution>>,
139}
140
141/// How many of the store's most recently written manifests a task with
142/// fallbacks tries after the named identities yield nothing.
143///
144/// A runner that restored a cache bundle holds the manifests of the builds
145/// that produced it and no version-control history to name them by, so what
146/// was written last is the best remaining guess at the lockfile before this
147/// one. Each candidate costs one parse, and a manifest from an unrelated
148/// workspace only fails to match. One is skipped when it fills more than half
149/// the prediction limit: an inherited manifest is kept under the new identity,
150/// and a foreign one that large would leave this workspace's own recordings
151/// no room.
152const NEWEST_MANIFEST_CANDIDATES: usize = 8;
153
154/// Remote action-cache access owned by one task session.
155pub struct AgentRemoteCache {
156    /// Remote protocol client used by the agent.
157    pub client: RemoteCacheClient,
158    /// Permitted remote read/write operations.
159    pub mode: RemoteCacheMode,
160    /// Directory used for verified downloads before CAS ingestion.
161    pub staging_dir: PathBuf,
162}
163
164#[derive(Default)]
165struct AtomicAgentStats {
166    wrapper_phases_ns: Mutex<BTreeMap<String, u64>>,
167    lookups: AtomicU64,
168    unconsulted: AtomicU64,
169    hits: AtomicU64,
170    stores: AtomicU64,
171    stored_bytes: AtomicU64,
172    verifications: AtomicU64,
173    divergences: AtomicU64,
174    downloaded_bytes: AtomicU64,
175    uploaded_bytes: AtomicU64,
176    background_uploads: AtomicU64,
177    background_upload_failures: AtomicU64,
178    remote_blob_pack_uploads: AtomicU64,
179    remote_blob_pack_upload_blobs: AtomicU64,
180    upload_drain_duration_ns: AtomicU64,
181    prefetched_actions: AtomicU64,
182    predictions_loaded: AtomicU64,
183    remote_failures: AtomicU64,
184    remote_manifest_lookups: AtomicU64,
185    remote_manifest_lookup_duration_ns: AtomicU64,
186    remote_action_lookups: AtomicU64,
187    remote_action_lookup_duration_ns: AtomicU64,
188    remote_blob_requests: AtomicU64,
189    remote_blob_pack_requests: AtomicU64,
190    remote_blob_pack_blobs: AtomicU64,
191    remote_blob_transfer_duration_ns: AtomicU64,
192    local_cas_write_duration_ns: AtomicU64,
193    prefetch_runs: AtomicU64,
194    prefetch_duration_ns: AtomicU64,
195    materialization_duration_ns: AtomicU64,
196    bypasses: Mutex<BTreeMap<String, u64>>,
197    avoided_compiler_duration_ns: AtomicU64,
198    compiler: Mutex<BTreeMap<String, CompilerStats>>,
199    unit_outcomes: Mutex<BTreeMap<String, BTreeSet<String>>>,
200    slow_compilations: Mutex<BTreeMap<String, u64>>,
201    restored_output_files: AtomicU64,
202    restored_output_bytes: AtomicU64,
203    reflinked_output_files: AtomicU64,
204    reflinked_output_bytes: AtomicU64,
205    copied_output_files: AtomicU64,
206    copied_output_bytes: AtomicU64,
207    reused_output_files: AtomicU64,
208    reused_output_bytes: AtomicU64,
209}
210
211struct AtomicDurationTimer<'a> {
212    started: Instant,
213    target: &'a AtomicU64,
214}
215
216impl<'a> AtomicDurationTimer<'a> {
217    fn start(target: &'a AtomicU64) -> Self {
218        Self {
219            started: Instant::now(),
220            target,
221        }
222    }
223}
224
225impl Drop for AtomicDurationTimer<'_> {
226    fn drop(&mut self) {
227        atomic_saturating_add(self.target, duration_ns(self.started));
228    }
229}
230
231fn duration_ns(started: Instant) -> u64 {
232    started.elapsed().as_nanos().try_into().unwrap_or(u64::MAX)
233}
234
235fn validate_crate_name(crate_name: Option<&str>) -> Result<()> {
236    if let Some(crate_name) = crate_name
237        && (crate_name.len() > 256 || crate_name.contains(['\0', '\n', '\r']))
238    {
239        bail!("invalid compiler crate name");
240    }
241    Ok(())
242}
243
244#[derive(serde::Deserialize)]
245struct ActionDiagnosticEnvelope {
246    outcome: String,
247    crate_name: Option<String>,
248    diagnostic: ActionDiagnostic,
249}
250
251fn parse_action_diagnostic(
252    message: &str,
253) -> Option<Result<(String, Option<String>, ActionDiagnostic)>> {
254    let payload = message.strip_prefix(ACTION_DIAGNOSTIC_PREFIX)?;
255    Some((|| {
256        let envelope: ActionDiagnosticEnvelope = serde_json::from_str(payload)?;
257        if !matches!(envelope.outcome.as_str(), "hit" | "miss") {
258            bail!("invalid action diagnostic outcome");
259        }
260        validate_crate_name(envelope.crate_name.as_deref())?;
261        Ok((envelope.outcome, envelope.crate_name, envelope.diagnostic))
262    })())
263}
264
265fn atomic_saturating_add(target: &AtomicU64, value: u64) {
266    let _ = target.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
267        Some(current.saturating_add(value))
268    });
269}
270
271fn queue_prefetch_digest(
272    verified: &BTreeMap<CacheDigest, PathBuf>,
273    pending: &mut BTreeMap<CacheDigest, ()>,
274    digest: CacheDigest,
275) {
276    if verified.contains_key(&digest) || pending.contains_key(&digest) {
277        return;
278    }
279    pending.insert(digest, ());
280}
281
282fn queue_prefetch_directory(
283    seen: &BTreeMap<CacheDigest, ()>,
284    pending: &mut BTreeMap<CacheDigest, ()>,
285    digest: CacheDigest,
286    limit: usize,
287) -> bool {
288    if seen.contains_key(&digest) || pending.contains_key(&digest) {
289        return true;
290    }
291    if seen.len().saturating_add(pending.len()) >= limit {
292        return false;
293    }
294    pending.insert(digest, ());
295    true
296}
297
298/// See [`CacheAgent::with_task_loader`].
299pub type TaskLoader = dyn for<'a> Fn(&'a CacheAgent, &'a str) -> BoxFuture<'a, ()> + Send + Sync;
300
301/// Shared state for an agent hosted by the process that owns a build session.
302///
303/// Transport listeners deliberately live in the embedder so the session
304/// lifecycle owns them. This type only contains ecosystem-independent CAS and
305/// protocol logic.
306#[derive(Clone)]
307pub struct CacheAgent {
308    cas: LocalCas,
309    actions: LocalActionCache,
310    verified_blobs: Arc<Mutex<BTreeMap<CacheDigest, VerifiedBlob>>>,
311    version: Arc<str>,
312    write_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
313    action_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
314    stats: Arc<AtomicAgentStats>,
315    observer: Option<Arc<dyn AgentEventObserver>>,
316    observer_emission: Arc<Mutex<()>>,
317    /// Loads a task the first time a request names it, so a connection that
318    /// only reports a bypass never waits for a manifest it will not read.
319    task_loader: Option<Arc<TaskLoader>>,
320    executable_identities: Arc<Mutex<BTreeMap<ExecutableIdentityKey, Vec<u8>>>>,
321    /// Where identities pinned to their executables outlive the session.
322    identity_dir: Arc<PathBuf>,
323    manifest_dir: Arc<PathBuf>,
324    task_actions: Arc<Mutex<BTreeMap<String, TaskActionState>>>,
325    /// Where a task may inherit predictions from, by task identity.
326    task_fallbacks: Arc<Mutex<BTreeMap<String, TaskFallbacks>>>,
327    next_task_run: Arc<AtomicU64>,
328    manifest_write_lock: Arc<Mutex<()>>,
329    remote: Option<Arc<RemoteCacheClient>>,
330    remote_mode: RemoteCacheMode,
331    remote_staging_dir: Arc<PathBuf>,
332    remote_download_limit: u64,
333    remote_download_bytes: Arc<AtomicU64>,
334    pending_remote_actions: Arc<Mutex<BTreeMap<CacheDigest, RemoteActionResult>>>,
335    remote_transfers: Arc<tokio::sync::Semaphore>,
336    prefetch_transfers: Arc<tokio::sync::Semaphore>,
337    prefetch_tasks: Arc<Mutex<Vec<tokio::task::JoinHandle<()>>>>,
338    /// Distinct shim diagnostics already surfaced, so a diagnostic that fires
339    /// once per compilation is printed once per session. The severity is part
340    /// of each key: a fatal reason must still be surfaced after the same text
341    /// appeared as a warning.
342    diagnostics: Arc<Mutex<BTreeSet<(DiagnosticLevel, String)>>>,
343    /// Digests of files shims hashed or wrote this session, keyed by scope and
344    /// path, each entry standing while its recorded identity matches the disk.
345    file_digests: Arc<Mutex<BTreeMap<(FileDigestScope, PathBuf), RecordedFileDigest>>>,
346    /// Per-identity flights that make concurrent cold lookups share one read.
347    file_digest_locks: Arc<Mutex<FileDigestFlights>>,
348    file_digest_permits: Arc<tokio::sync::Semaphore>,
349    #[cfg(test)]
350    file_digest_reads: Arc<AtomicU64>,
351    /// Deferred remote publication, present only when the session may write.
352    uploads: Option<UploadQueue>,
353}
354
355/// What every file-digest record must satisfy before the ledger will answer
356/// with it, whether a shim sent it or an earlier session left it behind.
357fn validate_file_digest_record(entry: &RecordedFileDigest) -> Result<()> {
358    if !entry.file.path.is_absolute() {
359        bail!("file-digest records need absolute paths");
360    }
361    entry.digest.validate()?;
362    if entry.file.len != entry.digest.size {
363        bail!("file-digest record length does not match its digest");
364    }
365    Ok(())
366}
367
368/// Records background upload activity against a session's statistics.
369struct AgentUploadSink {
370    stats: Arc<AtomicAgentStats>,
371}
372
373impl UploadSink for AgentUploadSink {
374    fn record_blob_uploaded(&self, bytes: u64) {
375        self.stats
376            .background_uploads
377            .fetch_add(1, Ordering::Relaxed);
378        self.stats
379            .uploaded_bytes
380            .fetch_add(bytes, Ordering::Relaxed);
381    }
382
383    fn record_action_uploaded(&self) {
384        self.stats
385            .background_uploads
386            .fetch_add(1, Ordering::Relaxed);
387    }
388
389    fn record_blob_pack_uploaded(&self, blobs: u64) {
390        self.stats
391            .remote_blob_pack_uploads
392            .fetch_add(1, Ordering::Relaxed);
393        self.stats
394            .remote_blob_pack_upload_blobs
395            .fetch_add(blobs, Ordering::Relaxed);
396    }
397
398    fn record_upload_failure(&self) {
399        self.stats
400            .background_upload_failures
401            .fetch_add(1, Ordering::Relaxed);
402        self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
403    }
404}
405
406/// A CAS blob whose contents this session hashed, and the file identity that
407/// says the hash still describes what is on disk.
408///
409/// Length alone would miss an overwrite that keeps the size; the modification
410/// time is what makes a rewrite visible without reading the bytes again.
411#[derive(Debug, Clone)]
412struct VerifiedBlob {
413    path: PathBuf,
414    len: u64,
415    modified: SystemTime,
416}
417
418impl VerifiedBlob {
419    /// Describe a file that was just verified, or nothing when the filesystem
420    /// will not report a modification time to compare against later.
421    fn describe(path: &Path) -> Option<Self> {
422        let metadata = std::fs::metadata(path).ok()?;
423        Some(Self {
424            path: path.to_path_buf(),
425            len: metadata.len(),
426            modified: metadata.modified().ok()?,
427        })
428    }
429
430    /// Whether the file still has the identity it had when it was verified.
431    fn is_unchanged(&self) -> bool {
432        let Ok(metadata) = std::fs::metadata(&self.path) else {
433            return false;
434        };
435        metadata.len() == self.len && metadata.modified().is_ok_and(|now| now == self.modified)
436    }
437}
438
439#[derive(Debug, Clone, Default)]
440struct TaskActionState {
441    manifest: String,
442    baseline_loaded: bool,
443    predictions: BTreeMap<CacheDigest, ActionPrediction>,
444    pending_predictions: BTreeMap<CacheDigest, ActionPrediction>,
445    /// The pending predictions that differ from what the manifest already
446    /// held when this run loaded it, each with what it held: nothing, or
447    /// the baseline's own prediction. A hit re-records its prediction
448    /// byte-for-byte so the receipt can name it, and that is no reason to
449    /// rewrite the manifest; nor is a prediction that moved and then moved
450    /// back, which is why the baseline value is kept rather than a flag.
451    changed_predictions: BTreeMap<CacheDigest, Option<ActionPrediction>>,
452    /// Values observed in the remote manifest when this run began. `None`
453    /// means provenance is unknown because the lookup failed or the baseline
454    /// was inherited locally; it must not make every inherited value look
455    /// unpublished.
456    remote_predictions: Option<BTreeMap<CacheDigest, ActionPrediction>>,
457    /// The ordered remote snapshot associated with `remote_etag`. A
458    /// read-write publisher merges this back immediately before its PUT so a
459    /// bounded begin-time view cannot retract remote-only history.
460    remote_manifest: Option<TaskActionManifest>,
461    /// The task identity has a useful local baseline but no known remote copy.
462    /// Even an otherwise hit-only run must publish an inherited or missing
463    /// remote manifest once.
464    remote_publication_needed: bool,
465    prefetched_adapters: BTreeSet<String>,
466    remote_etag: Option<String>,
467}
468
469/// Match an invocation and, on an adapter's first match, select its prefetch wave.
470fn activate_prediction_adapter(
471    state: &mut TaskActionState,
472    invocation: &CacheDigest,
473) -> (Option<ActionPrediction>, Option<Vec<ActionPrediction>>) {
474    let prediction = state.predictions.get(invocation).cloned();
475    let prefetch = prediction.as_ref().and_then(|prediction| {
476        if !state.prefetched_adapters.insert(prediction.adapter.clone()) {
477            return None;
478        }
479        Some(
480            state
481                .predictions
482                .values()
483                .filter(|candidate| candidate.adapter == prediction.adapter)
484                .cloned()
485                .collect(),
486        )
487    });
488    (prediction, prefetch)
489}
490
491/// Protect this run's complete working set and then as many of the newest
492/// concurrent changes as the bounded manifest can carry.
493fn required_local_predictions(
494    predictions: &[ActionPrediction],
495    updates: &BTreeMap<CacheDigest, ActionPrediction>,
496    concurrent_updates: &BTreeMap<CacheDigest, ActionPrediction>,
497) -> BTreeSet<CacheDigest> {
498    let mut required: BTreeSet<_> = updates.keys().cloned().collect();
499    for prediction in predictions.iter().rev() {
500        if required.len() == MAX_TASK_ACTION_PREDICTIONS {
501            break;
502        }
503        if concurrent_updates.contains_key(&prediction.invocation) {
504            required.insert(prediction.invocation.clone());
505        }
506    }
507    required
508}
509
510/// Find disk values written after this run began, using the saved pre-recording
511/// value for invocations this run also changed.
512fn concurrent_task_predictions(
513    predictions: &[ActionPrediction],
514    state: &TaskActionState,
515) -> BTreeMap<CacheDigest, ActionPrediction> {
516    predictions
517        .iter()
518        .filter(|prediction| {
519            let baseline = match state.changed_predictions.get(&prediction.invocation) {
520                Some(baseline) => baseline.as_ref(),
521                None => state.predictions.get(&prediction.invocation),
522            };
523            baseline != Some(*prediction)
524        })
525        .map(|prediction| (prediction.invocation.clone(), prediction.clone()))
526        .collect()
527}
528
529/// Turn hit re-recordings into LRU refreshes of the value current at commit
530/// time. A run that actually introduced or changed a prediction remains
531/// authoritative; a hit on its unchanged baseline must not clobber a newer
532/// concurrent value for the same invocation.
533fn effective_task_updates(
534    predictions: &[ActionPrediction],
535    state: &TaskActionState,
536) -> BTreeMap<CacheDigest, ActionPrediction> {
537    let current: BTreeMap<_, _> = predictions
538        .iter()
539        .map(|prediction| (prediction.invocation.clone(), prediction))
540        .collect();
541    state
542        .pending_predictions
543        .iter()
544        .map(|(invocation, update)| {
545            let unchanged_hit = !state.changed_predictions.contains_key(invocation);
546            let effective = if unchanged_hit {
547                current.get(invocation).copied().unwrap_or(update).clone()
548            } else {
549                update.clone()
550            };
551            (invocation.clone(), effective)
552        })
553        .collect()
554}
555
556fn should_refresh_task_prediction_order(prediction_count: usize, has_recorded: bool) -> bool {
557    has_recorded && prediction_count >= MAX_TASK_ACTION_PREDICTIONS - TASK_ACTION_LRU_REFRESH_MARGIN
558}
559
560/// Choose the last remotely observed value for every prediction this run
561/// replaces. A local-only value is not proof that its action was published.
562fn replaced_task_predictions(
563    remote: &BTreeMap<CacheDigest, ActionPrediction>,
564    candidates: &BTreeSet<CacheDigest>,
565) -> Vec<ActionPrediction> {
566    candidates
567        .iter()
568        .filter_map(|invocation| remote.get(invocation).cloned())
569        .collect()
570}
571
572/// Decide which local values may outrank a remote manifest during publication.
573///
574/// Values changed by this run and concurrent local writers are authoritative
575/// once their actions are publishable. An unchanged hit is only an LRU hint:
576/// it must not replace a newer value learned remotely after this run began.
577fn remote_publication_precedence(
578    state: &TaskActionState,
579    completed: &[ActionPrediction],
580    concurrent_updates: &BTreeMap<CacheDigest, ActionPrediction>,
581    proven_uploads: &BTreeSet<CacheDigest>,
582    published: &BTreeMap<CacheDigest, &ActionPrediction>,
583    withheld: &BTreeSet<CacheDigest>,
584) -> (BTreeSet<CacheDigest>, BTreeSet<CacheDigest>) {
585    let publication_base_invocations: BTreeSet<_> = state
586        .remote_manifest
587        .iter()
588        .flat_map(|manifest| manifest.predictions.iter())
589        .map(|prediction| prediction.invocation.clone())
590        .collect();
591    let unchanged_hits: BTreeSet<_> = completed
592        .iter()
593        .filter(|prediction| {
594            !state
595                .changed_predictions
596                .contains_key(&prediction.invocation)
597        })
598        .map(|prediction| prediction.invocation.clone())
599        .collect();
600    let required = completed
601        .iter()
602        .filter(|prediction| !unchanged_hits.contains(&prediction.invocation))
603        .chain(concurrent_updates.values().filter(|prediction| {
604            state.remote_predictions.is_some() || proven_uploads.contains(&prediction.action)
605        }))
606        .filter(|prediction| {
607            !withheld.contains(&prediction.invocation)
608                && published.get(&prediction.invocation).copied() == Some(*prediction)
609        })
610        .map(|prediction| prediction.invocation.clone())
611        .collect();
612    let fallbacks = withheld
613        .iter()
614        .filter(|invocation| published.contains_key(*invocation))
615        .cloned()
616        .chain(
617            unchanged_hits
618                .into_iter()
619                .filter(|invocation| published.contains_key(invocation)),
620        )
621        .chain(
622            concurrent_updates
623                .iter()
624                .filter(|(invocation, prediction)| {
625                    state.remote_predictions.is_none()
626                        && !proven_uploads.contains(&prediction.action)
627                        && published.contains_key(*invocation)
628                })
629                .map(|(invocation, _)| invocation.clone()),
630        )
631        .chain(
632            published
633                .keys()
634                .filter(|invocation| {
635                    state.remote_predictions.is_none()
636                        && state.remote_manifest.is_some()
637                        && !publication_base_invocations.contains(*invocation)
638                })
639                .cloned(),
640        )
641        .collect();
642    (required, fallbacks)
643}
644
645/// Combine baseline provenance with this queue's observed upload outcomes.
646/// A held action is proven only when it is absent from the failed set.
647fn unpublished_task_actions(
648    mut unproven: BTreeSet<CacheDigest>,
649    held: &BTreeSet<CacheDigest>,
650    failed: BTreeSet<CacheDigest>,
651) -> BTreeSet<CacheDigest> {
652    unproven.retain(|action| !held.contains(action));
653    unproven.extend(failed);
654    unproven
655}
656
657/// Find final values the task's remote baseline does not prove were published.
658fn unproven_task_predictions<'a>(
659    predictions: impl IntoIterator<Item = &'a ActionPrediction>,
660    remote: &BTreeMap<CacheDigest, ActionPrediction>,
661) -> BTreeMap<CacheDigest, CacheDigest> {
662    predictions
663        .into_iter()
664        .filter(|prediction| remote.get(&prediction.invocation) != Some(*prediction))
665        .map(|prediction| (prediction.invocation.clone(), prediction.action.clone()))
666        .collect()
667}
668
669/// Remove new or replacement predictions whose actions did not publish,
670/// restoring only the last value observed in the remote manifest.
671fn withhold_unpublished_updates(
672    manifest: &mut TaskActionManifest,
673    candidates: &BTreeSet<CacheDigest>,
674    replaced: &[ActionPrediction],
675    unpublished: &BTreeSet<CacheDigest>,
676) -> BTreeSet<CacheDigest> {
677    let withheld: BTreeSet<_> = manifest
678        .predictions
679        .iter()
680        .filter(|prediction| {
681            candidates.contains(&prediction.invocation) && unpublished.contains(&prediction.action)
682        })
683        .map(|prediction| prediction.invocation.clone())
684        .collect();
685    manifest
686        .predictions
687        .retain(|prediction| !withheld.contains(&prediction.invocation));
688    let mut retained: BTreeSet<_> = manifest
689        .predictions
690        .iter()
691        .map(|prediction| prediction.invocation.clone())
692        .collect();
693    for previous in replaced.iter().rev() {
694        if withheld.contains(&previous.invocation) && retained.insert(previous.invocation.clone()) {
695            // This invocation was used by the current run. Keep its last
696            // published value at the newest end of the durable LRU even though
697            // the replacement could not be advertised.
698            manifest.predictions.push(previous.clone());
699        }
700    }
701    withheld
702}
703
704struct PrefetchedAction {
705    adapter: String,
706    result: RemoteActionResult,
707}
708
709struct RemoteDownloadReservation {
710    counter: Arc<AtomicU64>,
711    reserved: u64,
712    committed: bool,
713}
714
715impl RemoteDownloadReservation {
716    fn bytes(&self) -> u64 {
717        self.reserved
718    }
719
720    fn commit(mut self, bytes: u64) {
721        debug_assert!(bytes <= self.reserved);
722        self.counter
723            .fetch_sub(self.reserved.saturating_sub(bytes), Ordering::AcqRel);
724        self.committed = true;
725    }
726}
727
728impl Drop for RemoteDownloadReservation {
729    fn drop(&mut self) {
730        if !self.committed {
731            self.counter.fetch_sub(self.reserved, Ordering::AcqRel);
732        }
733    }
734}
735
736#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
737struct ExecutableIdentityKey {
738    executable: PathBuf,
739    environment: BTreeMap<String, Option<String>>,
740}
741
742// 2: pins carry the state the shim saw, and cover the search directories.
743const PERSISTED_EXECUTABLE_IDENTITY_VERSION: u8 = 2;
744
745/// An executable identity kept across sessions.
746///
747/// It stands while every pinned file is as the probing shim saw it. The
748/// probe reads those files and nothing else, so an unchanged set of them
749/// would print the same answer again.
750#[derive(Debug, Serialize, Deserialize)]
751struct PersistedExecutableIdentity {
752    version: u8,
753    executable: PathBuf,
754    environment: BTreeMap<String, Option<String>>,
755    pins: Vec<PinnedFile>,
756    stdout: String,
757}
758
759impl CacheAgent {
760    /// Create an agent backed by the cache rooted at `cache_dir`.
761    pub fn new(cache_dir: impl Into<PathBuf>, version: impl Into<Arc<str>>) -> Self {
762        Self::build(cache_dir.into(), version.into(), None, 0)
763    }
764
765    /// Create an agent with local-first access to a remote action cache.
766    pub fn new_remote(
767        cache_dir: impl Into<PathBuf>,
768        version: impl Into<Arc<str>>,
769        remote: AgentRemoteCache,
770    ) -> Self {
771        Self::build(
772            cache_dir.into(),
773            version.into(),
774            Some(remote),
775            DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES,
776        )
777    }
778
779    /// Create a remote agent with a cumulative download budget for the session.
780    pub fn new_remote_with_download_limit(
781        cache_dir: impl Into<PathBuf>,
782        version: impl Into<Arc<str>>,
783        remote: AgentRemoteCache,
784        max_remote_download_bytes: u64,
785    ) -> Self {
786        Self::build(
787            cache_dir.into(),
788            version.into(),
789            Some(remote),
790            max_remote_download_bytes,
791        )
792    }
793
794    fn build(
795        cache_dir: PathBuf,
796        version: Arc<str>,
797        remote: Option<AgentRemoteCache>,
798        remote_download_limit: u64,
799    ) -> Self {
800        let remote_mode = remote
801            .as_ref()
802            .map_or(RemoteCacheMode::ReadOnly, |remote| remote.mode);
803        let remote_staging_dir = remote.as_ref().map_or_else(
804            || cache_dir.join("remote"),
805            |remote| remote.staging_dir.clone(),
806        );
807        let remote = remote.map(|remote| Arc::new(remote.client));
808        let stats = Arc::new(AtomicAgentStats::default());
809        let remote_transfers = Arc::new(tokio::sync::Semaphore::new(MAX_REMOTE_TRANSFERS));
810        // A session that cannot publish never queues an upload, so it does not
811        // need somewhere to queue one.
812        let uploads = remote
813            .clone()
814            .filter(|_| remote_mode.writes())
815            .map(|client| {
816                UploadQueue::new(
817                    client,
818                    Arc::new(AgentUploadSink {
819                        stats: stats.clone(),
820                    }),
821                    remote_transfers.clone(),
822                )
823            });
824        Self {
825            cas: LocalCas::new(cache_dir.clone()),
826            actions: LocalActionCache::new(cache_dir.clone()),
827            verified_blobs: Arc::new(Mutex::new(BTreeMap::new())),
828            version,
829            write_locks: Arc::new(Mutex::new(BTreeMap::new())),
830            action_locks: Arc::new(Mutex::new(BTreeMap::new())),
831            stats,
832            observer: None,
833            observer_emission: Arc::new(Mutex::new(())),
834            task_loader: None,
835            executable_identities: Arc::new(Mutex::new(BTreeMap::new())),
836            identity_dir: Arc::new(cache_dir.join("executable-identities").join("v1")),
837            manifest_dir: Arc::new(task_manifest_dir(&cache_dir)),
838            task_actions: Arc::new(Mutex::new(BTreeMap::new())),
839            task_fallbacks: Arc::new(Mutex::new(BTreeMap::new())),
840            next_task_run: Arc::new(AtomicU64::new(0)),
841            manifest_write_lock: Arc::new(Mutex::new(())),
842            remote,
843            remote_mode,
844            remote_staging_dir: Arc::new(remote_staging_dir),
845            remote_download_limit,
846            remote_download_bytes: Arc::new(AtomicU64::new(0)),
847            pending_remote_actions: Arc::new(Mutex::new(BTreeMap::new())),
848            remote_transfers,
849            prefetch_transfers: Arc::new(tokio::sync::Semaphore::new(MAX_PREFETCH_TRANSFERS)),
850            prefetch_tasks: Arc::new(Mutex::new(Vec::new())),
851            diagnostics: Arc::new(Mutex::new(BTreeSet::new())),
852            file_digests: Arc::new(Mutex::new(BTreeMap::new())),
853            file_digest_locks: Arc::new(Mutex::new(BTreeMap::new())),
854            file_digest_permits: Arc::new(tokio::sync::Semaphore::new(MAX_CONCURRENT_FILE_DIGESTS)),
855            #[cfg(test)]
856            file_digest_reads: Arc::new(AtomicU64::new(0)),
857            uploads,
858        }
859    }
860
861    /// Report each accounted cache decision to `observer` as it happens.
862    #[must_use]
863    pub fn with_observer(mut self, observer: Arc<dyn AgentEventObserver>) -> Self {
864        self.observer = Some(observer);
865        self
866    }
867
868    /// Load tasks on demand.
869    ///
870    /// The loader is called with the agent and the task's identity before the
871    /// first request that names a task this agent has not begun, and is
872    /// expected to begin it; a second call for the same task must be a no-op.
873    /// Requests that name no task never wait on it.
874    pub fn with_task_loader(mut self, loader: Arc<TaskLoader>) -> Self {
875        self.task_loader = Some(loader);
876        self
877    }
878
879    /// Make sure `task` has been begun before a request reads or records
880    /// under it.
881    async fn ensure_task_loaded(&self, task: &str) {
882        let Some(loader) = &self.task_loader else {
883            return;
884        };
885        if self.task_actions.lock().unwrap().contains_key(task) {
886            return;
887        }
888        loader(self, task).await;
889    }
890
891    fn emit(&self, event: impl FnOnce() -> AgentEvent) {
892        if let Some(observer) = &self.observer {
893            let _emission = self.observer_emission.lock().unwrap();
894            observer.event(event());
895        }
896    }
897
898    fn emit_action(&self, diagnostic: Option<AgentEvent>, action: AgentEvent) {
899        if let Some(observer) = &self.observer {
900            let _emission = self.observer_emission.lock().unwrap();
901            if let Some(diagnostic) = diagnostic {
902                observer.event(diagnostic);
903            }
904            observer.event(action);
905        }
906    }
907
908    fn reserve_remote_download(&self, bytes: u64) -> Result<RemoteDownloadReservation> {
909        let mut current = self.remote_download_bytes.load(Ordering::Acquire);
910        loop {
911            let next = current
912                .checked_add(bytes)
913                .ok_or_else(|| eyre::eyre!("remote cache download budget overflowed"))?;
914            if next > self.remote_download_limit {
915                bail!(
916                    "remote cache download budget exceeded: {} bytes requested with {} of {} bytes already used",
917                    bytes,
918                    current,
919                    self.remote_download_limit
920                );
921            }
922            match self.remote_download_bytes.compare_exchange_weak(
923                current,
924                next,
925                Ordering::AcqRel,
926                Ordering::Acquire,
927            ) {
928                Ok(_) => {
929                    return Ok(RemoteDownloadReservation {
930                        counter: self.remote_download_bytes.clone(),
931                        reserved: bytes,
932                        committed: false,
933                    });
934                }
935                Err(observed) => current = observed,
936            }
937        }
938    }
939
940    fn reserve_remote_download_up_to(&self, requested: u64) -> Result<RemoteDownloadReservation> {
941        let mut current = self.remote_download_bytes.load(Ordering::Acquire);
942        loop {
943            let reserved = requested.min(self.remote_download_limit.saturating_sub(current));
944            let next = current + reserved;
945            match self.remote_download_bytes.compare_exchange_weak(
946                current,
947                next,
948                Ordering::AcqRel,
949                Ordering::Acquire,
950            ) {
951                Ok(_) => {
952                    return Ok(RemoteDownloadReservation {
953                        counter: self.remote_download_bytes.clone(),
954                        reserved,
955                        committed: false,
956                    });
957                }
958                Err(observed) => current = observed,
959            }
960        }
961    }
962
963    /// Load the last committed action manifest for a task into this session.
964    pub async fn begin_task(&self, task: &str) -> Result<String> {
965        self.begin_task_with_remote_errors(task, false, None, true)
966            .await
967    }
968
969    /// Load a task but defer each adapter's speculative downloads until its
970    /// first prediction matches the current build.
971    pub async fn begin_task_on_prediction(&self, task: &str) -> Result<String> {
972        self.begin_task_with_remote_errors(task, false, None, false)
973            .await
974    }
975
976    /// Load a task using its identity as the run key.
977    ///
978    /// A task-scoped process has only one run for an identity, so it can defer
979    /// this work until its first client connects while putting the identity in
980    /// the client's environment ahead of time.
981    pub async fn begin_session_task(&self, task: &str) -> Result<()> {
982        self.begin_task_with_remote_errors(task, false, Some(task.to_string()), false)
983            .await?;
984        Ok(())
985    }
986
987    /// Name where a task's predictions may come from when nothing has been
988    /// recorded under its own identity.
989    ///
990    /// The identities are produced on demand rather than taken up front,
991    /// because finding them can cost a few version-control commands and most
992    /// builds never need them. The first manifest found, locally and then on
993    /// the remote, is adopted under the task's own identity, so the commands
994    /// that follow, tests and lints included, start from it too, and a trusted
995    /// build publishes it there. When none of the named identities has one,
996    /// the store's newest manifests are tried.
997    pub fn register_task_fallbacks(
998        &self,
999        task: &str,
1000        fallbacks: impl Fn() -> Vec<String> + Send + Sync + 'static,
1001    ) -> Result<()> {
1002        validate_task_identity(task)?;
1003        self.task_fallbacks
1004            .lock()
1005            .unwrap()
1006            .insert(task.to_string(), Arc::new(fallbacks));
1007        Ok(())
1008    }
1009
1010    /// Load a task and finish its prefetch, surfacing remote lookup failures.
1011    pub async fn prefetch_task(&self, task: &str) -> Result<String> {
1012        let run = self
1013            .begin_task_with_remote_errors(task, true, None, true)
1014            .await?;
1015        self.wait_for_prefetches().await;
1016        Ok(run)
1017    }
1018
1019    async fn begin_task_with_remote_errors(
1020        &self,
1021        task: &str,
1022        strict: bool,
1023        run: Option<String>,
1024        eager_prefetch: bool,
1025    ) -> Result<String> {
1026        validate_task_identity(task)?;
1027        // The local manifest is already enough to start useful work. Do not
1028        // leave those predictions idle while a high-latency remote answers the
1029        // manifest lookup; the authoritative snapshot is loaded again below
1030        // before it is merged, so another process can still update it while
1031        // this request is in flight.
1032        let early_manifest = {
1033            let _write_guard = self.manifest_write_lock.lock().unwrap();
1034            let _file_guard = self.lock_task_manifest(task)?;
1035            self.load_task_manifest(task)?
1036        };
1037        let early_actions: BTreeSet<_> = early_manifest
1038            .iter()
1039            .flat_map(|manifest| manifest.predictions.iter())
1040            .map(|prediction| prediction.action.clone())
1041            .collect();
1042        if eager_prefetch {
1043            self.spawn_prefetch_predictions(
1044                early_manifest
1045                    .as_ref()
1046                    .map(|manifest| manifest.predictions.clone())
1047                    .unwrap_or_default(),
1048            );
1049        }
1050        let (remote_manifest, mut remote_etag, mut remote_predictions, remote_missing) =
1051            if self.remote.is_some() && self.remote_mode.reads() {
1052                match self.get_remote_task_manifest(task).await {
1053                    Ok(Some((manifest, etag))) => {
1054                        let predictions = manifest
1055                            .predictions
1056                            .iter()
1057                            .map(|prediction| (prediction.invocation.clone(), prediction.clone()))
1058                            .collect();
1059                        (Some(manifest), Some(etag), Some(predictions), false)
1060                    }
1061                    Ok(None) => (None, None, Some(BTreeMap::new()), true),
1062                    Err(error) => {
1063                        if strict {
1064                            return Err(error).wrap_err_with(|| {
1065                                format!("remote task action manifest lookup failed for {task}")
1066                            });
1067                        }
1068                        self.note_remote_failure();
1069                        warn!("remote task action manifest lookup failed for {task}: {error}");
1070                        (None, None, None, false)
1071                    }
1072                }
1073            } else {
1074                (None, None, None, false)
1075            };
1076        let remote_publication_base = remote_manifest.clone();
1077        // Without a remote there is nothing to reconcile: the manifest just
1078        // read is the whole truth, and writing it back byte for byte would
1079        // serialize every prediction for nothing.
1080        let manifest = if self.remote.is_some() && self.remote_mode.reads() {
1081            let _write_guard = self.manifest_write_lock.lock().unwrap();
1082            let _file_guard = self.lock_task_manifest(task)?;
1083            let local_manifest = self.load_task_manifest(task)?;
1084            let local_baseline = local_manifest.clone();
1085            let manifest = match (remote_manifest, local_manifest) {
1086                (Some(remote), Some(local)) => {
1087                    let required = local
1088                        .predictions
1089                        .iter()
1090                        .map(|prediction| prediction.invocation.clone())
1091                        .collect();
1092                    match merge_task_manifests(
1093                        task,
1094                        remote,
1095                        local.clone(),
1096                        &required,
1097                        &BTreeSet::new(),
1098                    ) {
1099                        Ok(manifest) => Some(manifest),
1100                        Err(error) => {
1101                            warn!("remote task action manifest merge failed for {task}: {error}");
1102                            remote_etag = None;
1103                            Some(local)
1104                        }
1105                    }
1106                }
1107                (Some(remote), None) => Some(remote),
1108                (None, local) => local,
1109            };
1110            if manifest != local_baseline
1111                && let Some(manifest) = &manifest
1112            {
1113                self.persist_task_manifest(manifest)?;
1114            }
1115            manifest
1116        } else {
1117            early_manifest
1118        };
1119        let (manifest, remote_etag, inherited) = match manifest {
1120            Some(manifest) => (Some(manifest), remote_etag, false),
1121            // Inherited from another identity, so the remote's copy of that
1122            // one is no precondition for publishing under this one.
1123            None => {
1124                let (manifest, inherited_provenance, inherited) =
1125                    self.inherit_task_manifest(task, strict).await?;
1126                if inherited {
1127                    remote_predictions = inherited_provenance;
1128                }
1129                (manifest, None, inherited)
1130            }
1131        };
1132        let mut state = if let Some(manifest) = manifest {
1133            let write_only_baseline = self.remote_mode.writes()
1134                && !self.remote_mode.reads()
1135                && !manifest.predictions.is_empty();
1136            TaskActionState {
1137                manifest: task.to_string(),
1138                baseline_loaded: true,
1139                predictions: manifest
1140                    .predictions
1141                    .into_iter()
1142                    .map(|prediction| (prediction.invocation.clone(), prediction))
1143                    .collect(),
1144                pending_predictions: BTreeMap::new(),
1145                changed_predictions: BTreeMap::new(),
1146                remote_predictions,
1147                remote_manifest: remote_publication_base,
1148                remote_publication_needed: self.remote.is_some()
1149                    && self.remote_mode.writes()
1150                    && (remote_missing || inherited || write_only_baseline),
1151                prefetched_adapters: BTreeSet::new(),
1152                remote_etag,
1153            }
1154        } else {
1155            TaskActionState {
1156                manifest: task.to_string(),
1157                baseline_loaded: true,
1158                remote_predictions,
1159                remote_manifest: remote_publication_base,
1160                remote_etag,
1161                ..TaskActionState::default()
1162            }
1163        };
1164        let run = run.unwrap_or_else(|| {
1165            let sequence = self.next_task_run.fetch_add(1, Ordering::Relaxed);
1166            CacheDigest::blake3(format!("{task}\0{}\0{sequence}", std::process::id()).as_bytes())
1167                .hash
1168        });
1169        // The largest baseline rather than a sum: beginning the same task again
1170        // in one session reloads the same manifest, and counting it twice would
1171        // overstate what there was to match.
1172        self.stats.predictions_loaded.fetch_max(
1173            state.predictions.len().try_into().unwrap_or(u64::MAX),
1174            Ordering::Relaxed,
1175        );
1176        if eager_prefetch {
1177            state.prefetched_adapters.extend(
1178                state
1179                    .predictions
1180                    .values()
1181                    .map(|prediction| prediction.adapter.clone()),
1182            );
1183        }
1184        // The early local wave already owns these actions. Only launch a
1185        // second wave for predictions learned from the remote or from a local
1186        // writer that committed while its lookup was in flight.
1187        let predictions = state
1188            .predictions
1189            .values()
1190            .filter(|prediction| !early_actions.contains(&prediction.action))
1191            .cloned()
1192            .collect();
1193        self.task_actions.lock().unwrap().insert(run.clone(), state);
1194        if eager_prefetch {
1195            self.spawn_prefetch_predictions(predictions);
1196        }
1197        Ok(run)
1198    }
1199
1200    /// Adopt the manifest of a fallback identity for a task that has none.
1201    ///
1202    /// The named identities are tried in the order given, each in the local
1203    /// store and then on the remote, and then the store's newest manifests,
1204    /// locally only. The first one found is persisted under `task` so the
1205    /// next command sees a complete baseline rather than only what this one
1206    /// recorded, and a later session finds it without asking again.
1207    async fn inherit_task_manifest(
1208        &self,
1209        task: &str,
1210        strict: bool,
1211    ) -> Result<(
1212        Option<TaskActionManifest>,
1213        Option<BTreeMap<CacheDigest, ActionPrediction>>,
1214        bool,
1215    )> {
1216        let fallbacks = self.task_fallbacks.lock().unwrap().get(task).cloned();
1217        let Some(fallbacks) = fallbacks else {
1218            return Ok((None, None, false));
1219        };
1220        let named = tokio::task::spawn_blocking(move || fallbacks()).await?;
1221        let mut tried = BTreeSet::from([task.to_string()]);
1222        let candidates = named.into_iter().map(|identity| (identity, true)).chain(
1223            self.newest_task_identities()
1224                .into_iter()
1225                .map(|identity| (identity, false)),
1226        );
1227        for (identity, named) in candidates {
1228            if !tried.insert(identity.clone()) || validate_task_identity(&identity).is_err() {
1229                continue;
1230            }
1231            let mut found = self.load_task_manifest(&identity)?;
1232            let mut provenance = None;
1233            if self.remote_mode.reads() && (found.is_some() || named) {
1234                match self.get_remote_task_manifest(&identity).await {
1235                    Ok(Some((remote, _))) => {
1236                        provenance = Some(
1237                            remote
1238                                .predictions
1239                                .iter()
1240                                .map(|prediction| {
1241                                    (prediction.invocation.clone(), prediction.clone())
1242                                })
1243                                .collect(),
1244                        );
1245                        if found.is_none() {
1246                            found = Some(remote);
1247                        }
1248                    }
1249                    Ok(None) => provenance = Some(BTreeMap::new()),
1250                    Err(error) => {
1251                        // A local fallback is already useful without remote
1252                        // provenance. Keep that baseline with unknown
1253                        // provenance rather than making strict prefetch fail
1254                        // because its optional proof lookup was unavailable.
1255                        if strict && found.is_none() {
1256                            return Err(error).wrap_err_with(|| {
1257                                format!("remote task action manifest lookup failed for {identity}")
1258                            });
1259                        }
1260                        self.note_remote_failure();
1261                        warn!("remote task action manifest lookup failed for {identity}: {error}");
1262                    }
1263                }
1264            }
1265            let Some(mut manifest) = found else {
1266                continue;
1267            };
1268            if manifest.predictions.is_empty()
1269                || (!named && manifest.predictions.len() > MAX_TASK_ACTION_PREDICTIONS / 2)
1270            {
1271                continue;
1272            }
1273            debug!(
1274                "no predictions were recorded for {task}; inheriting {} from {identity}",
1275                manifest.predictions.len()
1276            );
1277            manifest.task = task.to_string();
1278            validate_task_manifest(&manifest, task)?;
1279            let _write_guard = self.manifest_write_lock.lock().unwrap();
1280            let _file_guard = self.lock_task_manifest(task)?;
1281            // Another process may have recorded this identity while the
1282            // fallbacks were being found. What it wrote describes this
1283            // lockfile and wins over an inheritance.
1284            if let Some(recorded) = self.load_task_manifest(task)? {
1285                return Ok((Some(recorded), None, false));
1286            }
1287            self.persist_task_manifest(&manifest)?;
1288            return Ok((Some(manifest), provenance, true));
1289        }
1290        Ok((None, None, false))
1291    }
1292
1293    /// The identities of the store's most recently written manifests.
1294    fn newest_task_identities(&self) -> Vec<String> {
1295        let Ok(entries) = fs::read_dir(self.manifest_dir.as_path()) else {
1296            return Vec::new();
1297        };
1298        let mut manifests: Vec<(std::time::SystemTime, String)> = entries
1299            .filter_map(|entry| {
1300                let entry = entry.ok()?;
1301                let name = entry.file_name();
1302                let identity = name.to_str()?.strip_suffix(".json")?.to_string();
1303                validate_task_identity(&identity).ok()?;
1304                let modified = entry.metadata().ok()?.modified().ok()?;
1305                Some((modified, identity))
1306            })
1307            .collect();
1308        manifests.sort_by(|left, right| right.cmp(left));
1309        manifests
1310            .into_iter()
1311            .take(NEWEST_MANIFEST_CANDIDATES)
1312            .map(|(_, identity)| identity)
1313            .collect()
1314    }
1315
1316    /// Cancel speculative downloads before the owning session exits.
1317    pub async fn cancel_prefetches(&self) {
1318        let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
1319        for task in &tasks {
1320            task.abort();
1321        }
1322        for task in tasks {
1323            if let Err(error) = task.await
1324                && !error.is_cancelled()
1325            {
1326                warn!("remote action prefetch task failed: {error}");
1327            }
1328        }
1329    }
1330
1331    /// Publish everything a build queued, before the session stops.
1332    ///
1333    /// Store requests return once an object is durable locally, so at this point
1334    /// the remote cache may still be behind the local one. Uploads run on the
1335    /// session's runtime and are abandoned if it goes away, so a session that
1336    /// wants them published has to wait here.
1337    pub async fn wait_for_uploads(&self) {
1338        let Some(uploads) = &self.uploads else {
1339            return;
1340        };
1341        let _timer = AtomicDurationTimer::start(&self.stats.upload_drain_duration_ns);
1342        uploads.drain().await;
1343    }
1344
1345    async fn wait_for_prefetches(&self) {
1346        let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
1347        for task in tasks {
1348            if let Err(error) = task.await {
1349                warn!("remote action prefetch task failed: {error}");
1350            }
1351        }
1352    }
1353
1354    /// Atomically publish the completed actions collected by a task run.
1355    pub async fn commit_task(&self, run: &str) -> Result<()> {
1356        self.commit_task_actions(run).await.map(|_| ())
1357    }
1358
1359    /// Publish a task run and return exactly the predictions completed by it.
1360    ///
1361    /// The persisted task manifest also carries predictions inherited from
1362    /// earlier runs. Callers that need a receipt for this one run must not
1363    /// mistake that cumulative manifest for the work the run completed.
1364    pub async fn commit_task_actions(&self, run: &str) -> Result<Vec<ActionPrediction>> {
1365        validate_task_identity(run)?;
1366        let mut state = {
1367            let mut runs = self.task_actions.lock().unwrap();
1368            let state = runs
1369                .get(run)
1370                .ok_or_else(|| eyre::eyre!("task action manifest baseline was not loaded"))?;
1371            if !state.baseline_loaded {
1372                bail!("task action manifest baseline was not loaded");
1373            }
1374            // A hit-only run leaves a comfortably bounded manifest as it found
1375            // it. Refresh durable LRU order only near the cap, where it can
1376            // affect the next pruning decision, or when this task identity
1377            // still needs its first remote publication.
1378            debug!(
1379                "committing {} changed of {} recorded predictions for {}",
1380                state.changed_predictions.len(),
1381                state.pending_predictions.len(),
1382                state.manifest
1383            );
1384            let refresh_hit_order = should_refresh_task_prediction_order(
1385                state.predictions.len(),
1386                !state.pending_predictions.is_empty(),
1387            );
1388            if state.changed_predictions.is_empty()
1389                && !refresh_hit_order
1390                && !state.remote_publication_needed
1391            {
1392                let completed = state.pending_predictions.values().cloned().collect();
1393                runs.remove(run);
1394                return Ok(completed);
1395            }
1396            state.clone()
1397        };
1398        let task = state.manifest.clone();
1399        validate_task_identity(&task)?;
1400        let write_only = self.remote_mode.writes() && !self.remote_mode.reads();
1401        // Write-only mode does not use remote predictions for cache hits, but
1402        // it still needs the current manifest as a publication precondition
1403        // and provenance baseline. The conditional update path may read the
1404        // same metadata after a conflict, so doing it up front does not weaken
1405        // the mode's promise not to satisfy build lookups from the remote.
1406        let write_only_remote = if write_only {
1407            match self.get_remote_task_manifest(&task).await {
1408                Ok(Some((manifest, etag))) => {
1409                    let remote_predictions: BTreeMap<_, _> = manifest
1410                        .predictions
1411                        .iter()
1412                        .map(|prediction| (prediction.invocation.clone(), prediction.clone()))
1413                        .collect();
1414                    let has_missing_fallback = state
1415                        .predictions
1416                        .keys()
1417                        .any(|invocation| !remote_predictions.contains_key(invocation));
1418                    state.remote_predictions = if has_missing_fallback {
1419                        None
1420                    } else {
1421                        Some(remote_predictions)
1422                    };
1423                    state.remote_etag = Some(etag);
1424                    state.remote_manifest = Some(manifest.clone());
1425                    // Same-valued or conflicting remote entries already cover
1426                    // this unknown-provenance baseline. Missing invocations do
1427                    // not: publish those as fallbacks without letting a stale
1428                    // local value overwrite a value the remote already knows.
1429                    state.remote_publication_needed = has_missing_fallback;
1430                    Some(manifest)
1431                }
1432                Ok(None) => {
1433                    state.remote_predictions = if state.predictions.is_empty() {
1434                        Some(BTreeMap::new())
1435                    } else {
1436                        None
1437                    };
1438                    state.remote_publication_needed = !state.predictions.is_empty();
1439                    None
1440                }
1441                Err(error) => {
1442                    state.remote_predictions = None;
1443                    state.remote_publication_needed = false;
1444                    self.note_remote_failure();
1445                    warn!("remote task action manifest lookup failed for {task}: {error}");
1446                    None
1447                }
1448            }
1449        } else {
1450            None
1451        };
1452        let completed = state
1453            .pending_predictions
1454            .values()
1455            .cloned()
1456            .collect::<Vec<_>>();
1457        if write_only
1458            && state.changed_predictions.is_empty()
1459            && !should_refresh_task_prediction_order(
1460                state.predictions.len(),
1461                !state.pending_predictions.is_empty(),
1462            )
1463            && !state.remote_publication_needed
1464        {
1465            self.task_actions.lock().unwrap().remove(run);
1466            return Ok(completed);
1467        }
1468        let (manifest, concurrent_updates) = {
1469            let _write_guard = self.manifest_write_lock.lock().unwrap();
1470            let _file_guard = self.lock_task_manifest(&task)?;
1471            let current_manifest = self.load_task_manifest(&task)?;
1472            let predictions = current_manifest
1473                .as_ref()
1474                .map(|manifest| manifest.predictions.clone())
1475                .unwrap_or_default();
1476            let concurrent_updates = concurrent_task_predictions(&predictions, &state);
1477            let effective_updates = effective_task_updates(&predictions, &state);
1478            let required =
1479                required_local_predictions(&predictions, &effective_updates, &concurrent_updates);
1480            let predictions = update_task_predictions(predictions, &effective_updates, &required)?;
1481            let manifest = TaskActionManifest {
1482                version: TASK_ACTION_MANIFEST_VERSION,
1483                task: task.clone(),
1484                predictions,
1485            };
1486            validate_task_manifest(&manifest, &task)?;
1487            if current_manifest.as_ref() != Some(&manifest) {
1488                self.persist_task_manifest(&manifest)?;
1489            }
1490            (manifest, concurrent_updates)
1491        };
1492        self.task_actions.lock().unwrap().remove(run);
1493        if self.remote_mode.writes() {
1494            // A manifest advertises the actions it predicts, so it must not
1495            // reach the remote cache before the results a reader would then go
1496            // looking for -- nor name a result that never got there at all.
1497            let mut manifest = manifest;
1498            let (held, failed) = if let Some(uploads) = &self.uploads {
1499                let actions: Vec<CacheDigest> = manifest
1500                    .predictions
1501                    .iter()
1502                    .map(|prediction| prediction.action.clone())
1503                    .collect();
1504                uploads.wait_for_action_outcomes(&actions).await
1505            } else {
1506                (BTreeSet::new(), BTreeSet::new())
1507            };
1508            let proven_uploads = held.difference(&failed).cloned().collect();
1509            let (withhold_candidates, replaced, unpublished) =
1510                if let Some(remote_predictions) = &state.remote_predictions {
1511                    let unproven_predictions =
1512                        unproven_task_predictions(&manifest.predictions, remote_predictions);
1513                    let unproven = unproven_predictions.values().cloned().collect();
1514                    let candidates = unproven_predictions.keys().cloned().collect();
1515                    let replaced = replaced_task_predictions(remote_predictions, &candidates);
1516                    // A successful upload by this queue proves an otherwise
1517                    // local-only value is now safe to advertise.
1518                    let unpublished = unpublished_task_actions(unproven, &held, failed);
1519                    (candidates, replaced, unpublished)
1520                } else {
1521                    // Unknown provenance must not turn an inherited baseline
1522                    // into thousands of false negatives. Still withhold values
1523                    // this run actually changed unless its own upload proves
1524                    // them publishable.
1525                    let candidates: BTreeSet<_> = state
1526                        .changed_predictions
1527                        .keys()
1528                        .filter(|invocation| {
1529                            manifest
1530                                .predictions
1531                                .iter()
1532                                .any(|prediction| &prediction.invocation == *invocation)
1533                        })
1534                        .cloned()
1535                        .collect();
1536                    let unproven = manifest
1537                        .predictions
1538                        .iter()
1539                        .filter(|prediction| candidates.contains(&prediction.invocation))
1540                        .map(|prediction| prediction.action.clone())
1541                        .collect();
1542                    let unpublished = unpublished_task_actions(unproven, &held, failed);
1543                    (candidates, Vec::new(), unpublished)
1544                };
1545            let prediction_count = manifest.predictions.len();
1546            let withheld = withhold_unpublished_updates(
1547                &mut manifest,
1548                &withhold_candidates,
1549                &replaced,
1550                &unpublished,
1551            );
1552            if !withheld.is_empty() {
1553                // The local manifest keeps them: this checkout can still use
1554                // what it built, and a later session can publish it.
1555                warn!(
1556                    "{} of {} predicted actions were not published, so the remote task action manifest omits them",
1557                    withheld.len(),
1558                    prediction_count
1559                );
1560            }
1561            let published: BTreeMap<_, _> = manifest
1562                .predictions
1563                .iter()
1564                .map(|prediction| (prediction.invocation.clone(), prediction))
1565                .collect();
1566            let (required, fallbacks) = remote_publication_precedence(
1567                &state,
1568                &completed,
1569                &concurrent_updates,
1570                &proven_uploads,
1571                &published,
1572                &withheld,
1573            );
1574            let publication_base = write_only_remote.or(state.remote_manifest);
1575            if let Some(remote) = publication_base {
1576                manifest = merge_task_manifests(&task, remote, manifest, &required, &fallbacks)?;
1577            }
1578            match self
1579                .put_remote_task_manifest(&task, manifest, state.remote_etag, &required, &fallbacks)
1580                .await
1581            {
1582                Ok(remote_manifest) => {
1583                    let _write_guard = self.manifest_write_lock.lock().unwrap();
1584                    let reconciliation = (|| {
1585                        let _file_guard = self.lock_task_manifest(&task)?;
1586                        let manifest = match self.load_task_manifest(&task)? {
1587                            Some(local) => {
1588                                let required = local
1589                                    .predictions
1590                                    .iter()
1591                                    .map(|prediction| prediction.invocation.clone())
1592                                    .collect();
1593                                match merge_task_manifests(
1594                                    &task,
1595                                    remote_manifest,
1596                                    local.clone(),
1597                                    &required,
1598                                    &BTreeSet::new(),
1599                                ) {
1600                                    Ok(manifest) => manifest,
1601                                    Err(error) => {
1602                                        warn!(
1603                                            "remote task action manifest merge failed for {task}: {error}"
1604                                        );
1605                                        local
1606                                    }
1607                                }
1608                            }
1609                            None => remote_manifest,
1610                        };
1611                        self.persist_task_manifest(&manifest)
1612                    })();
1613                    if let Err(error) = reconciliation {
1614                        warn!(
1615                            "remote task action manifest reconciliation failed for {task}: {error}"
1616                        );
1617                    }
1618                }
1619                Err(error) => {
1620                    self.note_remote_failure();
1621                    warn!("remote task action manifest upload failed for {task}: {error}");
1622                }
1623            }
1624        }
1625        Ok(completed)
1626    }
1627
1628    fn task_manifest_path(&self, task: &str) -> PathBuf {
1629        self.manifest_dir.join(format!("{task}.json"))
1630    }
1631
1632    fn task_manifest_lock_path(&self, task: &str) -> PathBuf {
1633        self.manifest_dir.join("locks").join(format!("{task}.lock"))
1634    }
1635
1636    fn lock_task_manifest(&self, task: &str) -> Result<fslock::LockFile> {
1637        let path = self.task_manifest_lock_path(task);
1638        fs::create_dir_all(path.parent().expect("task manifest lock has a parent"))?;
1639        let mut lock = fslock::LockFile::open(&path)?;
1640        lock.lock()?;
1641        Ok(lock)
1642    }
1643
1644    fn load_task_manifest(&self, task: &str) -> Result<Option<TaskActionManifest>> {
1645        match fs::read(self.task_manifest_path(task)) {
1646            Ok(contents) => Ok(Some(self.parse_task_manifest(task, &contents, false)?)),
1647            Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
1648            Err(error) => Err(error.into()),
1649        }
1650    }
1651
1652    fn parse_task_manifest(
1653        &self,
1654        task: &str,
1655        contents: &[u8],
1656        require_canonical: bool,
1657    ) -> Result<TaskActionManifest> {
1658        let manifest: TaskActionManifest = serde_json::from_slice(contents)?;
1659        validate_task_manifest(&manifest, task)?;
1660        if require_canonical && canonical_json(&manifest)? != contents {
1661            bail!("task action manifest is not canonical JSON");
1662        }
1663        Ok(manifest)
1664    }
1665
1666    fn task_manifest_selector(task: &str) -> Result<(Vec<u8>, CacheDigest)> {
1667        TaskActionManifest::selector(task)
1668    }
1669
1670    fn persist_task_manifest(&self, manifest: &TaskActionManifest) -> Result<()> {
1671        let bytes = canonical_json(manifest)?;
1672        fs::create_dir_all(self.manifest_dir.as_path())?;
1673        let mut temporary = tempfile::NamedTempFile::new_in(self.manifest_dir.as_path())?;
1674        std::io::Write::write_all(temporary.as_file_mut(), &bytes)?;
1675        temporary.as_file_mut().sync_all()?;
1676        temporary
1677            .persist(self.task_manifest_path(&manifest.task))
1678            .map_err(|error| error.error)?;
1679        Ok(())
1680    }
1681
1682    async fn get_remote_task_manifest(
1683        &self,
1684        task: &str,
1685    ) -> Result<Option<(TaskActionManifest, String)>> {
1686        let Some(remote) = &self.remote else {
1687            return Ok(None);
1688        };
1689        let (_, selector) = Self::task_manifest_selector(task)?;
1690        let _permit = self.remote_transfers.acquire().await?;
1691        self.stats
1692            .remote_manifest_lookups
1693            .fetch_add(1, Ordering::Relaxed);
1694        let _timer = AtomicDurationTimer::start(&self.stats.remote_manifest_lookup_duration_ns);
1695        let Some(remote_manifest) = remote.get_action_manifest(&selector).await? else {
1696            return Ok(None);
1697        };
1698        let manifest = self.parse_task_manifest(task, &remote_manifest.bytes, true)?;
1699        Ok(Some((manifest, remote_manifest.etag)))
1700    }
1701
1702    async fn put_remote_task_manifest(
1703        &self,
1704        task: &str,
1705        mut manifest: TaskActionManifest,
1706        mut expected_etag: Option<String>,
1707        required: &BTreeSet<CacheDigest>,
1708        fallbacks: &BTreeSet<CacheDigest>,
1709    ) -> Result<TaskActionManifest> {
1710        let Some(remote) = &self.remote else {
1711            return Ok(manifest);
1712        };
1713        let (_, selector) = Self::task_manifest_selector(task)?;
1714        for _ in 0..4 {
1715            let bytes = canonical_json(&manifest)?;
1716            let outcome = {
1717                let _permit = self.remote_transfers.acquire().await?;
1718                remote
1719                    .put_action_manifest(&selector, &bytes, expected_etag.as_deref())
1720                    .await?
1721            };
1722            match outcome {
1723                ManifestPutOutcome::Stored => return Ok(manifest),
1724                ManifestPutOutcome::PreconditionFailed => {
1725                    let Some((remote_manifest, etag)) = self.get_remote_task_manifest(task).await?
1726                    else {
1727                        expected_etag = None;
1728                        continue;
1729                    };
1730                    manifest =
1731                        merge_task_manifests(task, remote_manifest, manifest, required, fallbacks)?;
1732                    expected_etag = Some(etag);
1733                }
1734            }
1735        }
1736        bail!("remote task action manifest changed too frequently")
1737    }
1738
1739    /// Return a snapshot of this session's cache activity.
1740    pub fn stats(&self) -> AgentStats {
1741        AgentStats {
1742            session_duration_ns: 0,
1743            lookups: self.stats.lookups.load(Ordering::Relaxed),
1744            unconsulted: self.stats.unconsulted.load(Ordering::Relaxed),
1745            hits: self.stats.hits.load(Ordering::Relaxed),
1746            stores: self.stats.stores.load(Ordering::Relaxed),
1747            stored_bytes: self.stats.stored_bytes.load(Ordering::Relaxed),
1748            verifications: self.stats.verifications.load(Ordering::Relaxed),
1749            divergences: self.stats.divergences.load(Ordering::Relaxed),
1750            downloaded_bytes: self.stats.downloaded_bytes.load(Ordering::Relaxed),
1751            uploaded_bytes: self.stats.uploaded_bytes.load(Ordering::Relaxed),
1752            background_uploads: self.stats.background_uploads.load(Ordering::Relaxed),
1753            background_upload_failures: self
1754                .stats
1755                .background_upload_failures
1756                .load(Ordering::Relaxed),
1757            remote_blob_pack_uploads: self.stats.remote_blob_pack_uploads.load(Ordering::Relaxed),
1758            remote_blob_pack_upload_blobs: self
1759                .stats
1760                .remote_blob_pack_upload_blobs
1761                .load(Ordering::Relaxed),
1762            upload_drain_duration_ns: self.stats.upload_drain_duration_ns.load(Ordering::Relaxed),
1763            prefetched_actions: self.stats.prefetched_actions.load(Ordering::Relaxed),
1764            predictions_loaded: self.stats.predictions_loaded.load(Ordering::Relaxed),
1765            bypasses: self.stats.bypasses.lock().unwrap().clone(),
1766            avoided_compiler_duration_ns: self
1767                .stats
1768                .avoided_compiler_duration_ns
1769                .load(Ordering::Relaxed),
1770            wrapper_phases_ns: self.stats.wrapper_phases_ns.lock().unwrap().clone(),
1771            compiler: self.stats.compiler.lock().unwrap().clone(),
1772            slow_compilations: self.stats.slow_compilations.lock().unwrap().clone(),
1773            unit_outcomes: self.stats.unit_outcomes.lock().unwrap().clone(),
1774            remote_failures: self.stats.remote_failures.load(Ordering::Relaxed),
1775            remote_manifest_lookups: self.stats.remote_manifest_lookups.load(Ordering::Relaxed),
1776            remote_manifest_lookup_duration_ns: self
1777                .stats
1778                .remote_manifest_lookup_duration_ns
1779                .load(Ordering::Relaxed),
1780            remote_action_lookups: self.stats.remote_action_lookups.load(Ordering::Relaxed),
1781            remote_action_lookup_duration_ns: self
1782                .stats
1783                .remote_action_lookup_duration_ns
1784                .load(Ordering::Relaxed),
1785            remote_blob_requests: self.stats.remote_blob_requests.load(Ordering::Relaxed),
1786            remote_blob_pack_requests: self.stats.remote_blob_pack_requests.load(Ordering::Relaxed),
1787            remote_blob_pack_blobs: self.stats.remote_blob_pack_blobs.load(Ordering::Relaxed),
1788            remote_blob_transfer_duration_ns: self
1789                .stats
1790                .remote_blob_transfer_duration_ns
1791                .load(Ordering::Relaxed),
1792            local_cas_write_duration_ns: self
1793                .stats
1794                .local_cas_write_duration_ns
1795                .load(Ordering::Relaxed),
1796            prefetch_runs: self.stats.prefetch_runs.load(Ordering::Relaxed),
1797            prefetch_duration_ns: self.stats.prefetch_duration_ns.load(Ordering::Relaxed),
1798            materialization_duration_ns: self
1799                .stats
1800                .materialization_duration_ns
1801                .load(Ordering::Relaxed),
1802            restored_output_files: self.stats.restored_output_files.load(Ordering::Relaxed),
1803            restored_output_bytes: self.stats.restored_output_bytes.load(Ordering::Relaxed),
1804            reflinked_output_files: self.stats.reflinked_output_files.load(Ordering::Relaxed),
1805            reflinked_output_bytes: self.stats.reflinked_output_bytes.load(Ordering::Relaxed),
1806            copied_output_files: self.stats.copied_output_files.load(Ordering::Relaxed),
1807            copied_output_bytes: self.stats.copied_output_bytes.load(Ordering::Relaxed),
1808            reused_output_files: self.stats.reused_output_files.load(Ordering::Relaxed),
1809            reused_output_bytes: self.stats.reused_output_bytes.load(Ordering::Relaxed),
1810        }
1811    }
1812
1813    /// Handle requests without a transport connection.
1814    ///
1815    /// Persistent wrappers use this entry point when Cargo invokes mbx outside
1816    /// an orchestrated session. It intentionally has the same response
1817    /// semantics as [`Self::handle_connection`], while leaving framing and the
1818    /// version handshake to callers that actually cross a process boundary.
1819    pub async fn handle_requests(
1820        &self,
1821        requests: impl IntoIterator<Item = AgentRequest>,
1822    ) -> Vec<AgentResponse> {
1823        let mut connection = ConnectionUploads::default();
1824        let mut responses = Vec::new();
1825        for request in requests {
1826            responses.push(self.respond_on(request, &mut connection).await);
1827        }
1828        responses
1829    }
1830
1831    fn write_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1832        Self::digest_lock(&self.write_locks, digest)
1833    }
1834
1835    fn action_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1836        Self::digest_lock(&self.action_locks, digest)
1837    }
1838
1839    fn digest_lock(
1840        locks: &Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>,
1841        digest: &CacheDigest,
1842    ) -> Arc<tokio::sync::Mutex<()>> {
1843        let mut locks = locks.lock().unwrap();
1844        locks.retain(|_, lock| lock.strong_count() > 0);
1845        if let Some(lock) = locks.get(digest).and_then(Weak::upgrade) {
1846            return lock;
1847        }
1848        let lock = Arc::new(tokio::sync::Mutex::new(()));
1849        locks.insert(digest.clone(), Arc::downgrade(&lock));
1850        lock
1851    }
1852
1853    /// Answer one request outside any connection.
1854    ///
1855    /// An upload queued this way has no sibling requests to order against, so
1856    /// the queue falls back to treating every blob it holds as a prerequisite.
1857    #[cfg(test)]
1858    async fn respond(&self, request: AgentRequest) -> AgentResponse {
1859        self.respond_on(request, &mut ConnectionUploads::default())
1860            .await
1861    }
1862
1863    async fn respond_on(
1864        &self,
1865        request: AgentRequest,
1866        connection: &mut ConnectionUploads,
1867    ) -> AgentResponse {
1868        match &request {
1869            AgentRequest::FindActionPrediction { task, .. }
1870            | AgentRequest::RecordActionPrediction { task, .. } => {
1871                self.ensure_task_loaded(task).await;
1872            }
1873            _ => {}
1874        }
1875        let result = match request {
1876            AgentRequest::BeginTask { task } => self
1877                .begin_task(&task)
1878                .await
1879                .map(|run| AgentResponse::TaskBegun { run }),
1880            AgentRequest::CommitTask { run } => self
1881                .commit_task(&run)
1882                .await
1883                .map(|()| AgentResponse::TaskCommitted),
1884            AgentRequest::FindBlob { digest } => self.find_blob(&digest).await,
1885            AgentRequest::FindBlobs { digests } => self.find_blobs(digests).await,
1886            AgentRequest::StoreBlob { digest, source } => {
1887                self.store_blob(&digest, &source, connection).await
1888            }
1889            AgentRequest::FindActionResult { action } => {
1890                self.stats.lookups.fetch_add(1, Ordering::Relaxed);
1891                self.find_action_result(&action).await
1892            }
1893            AgentRequest::RecordActionHit {
1894                action,
1895                restore,
1896                crate_name,
1897            } => {
1898                let diagnostic = connection
1899                    .take_action_diagnostic("hit", crate_name.as_deref())
1900                    .map(|diagnostic| AgentEvent::ActionDiagnostic {
1901                        outcome: "hit".into(),
1902                        crate_name: crate_name.clone(),
1903                        diagnostic,
1904                    });
1905                self.record_action_hit(&action, restore, crate_name, diagnostic)
1906            }
1907            AgentRequest::RecordBypass { kind } => {
1908                *self
1909                    .stats
1910                    .bypasses
1911                    .lock()
1912                    .unwrap()
1913                    .entry(kind.clone())
1914                    .or_insert(0) += 1;
1915                self.emit(|| AgentEvent::Bypass { kind });
1916                Ok(AgentResponse::BypassRecorded)
1917            }
1918            AgentRequest::RecordUnconsulted => {
1919                self.stats.unconsulted.fetch_add(1, Ordering::Relaxed);
1920                self.emit(|| AgentEvent::Unconsulted);
1921                Ok(AgentResponse::UnconsultedRecorded)
1922            }
1923            AgentRequest::RecordWarning { message } => match parse_action_diagnostic(&message) {
1924                Some(Ok((outcome, crate_name, diagnostic))) => {
1925                    connection.record_action_diagnostic(outcome, crate_name, diagnostic);
1926                    Ok(AgentResponse::WarningRecorded)
1927                }
1928                Some(Err(error)) => Err(error),
1929                None => self.record_diagnostic(DiagnosticLevel::Warning, message),
1930            },
1931            AgentRequest::RecordDebug { target, message } => (|| {
1932                if target.is_empty()
1933                    || target.len() > 256
1934                    || target.contains(['\n', '\r', '\0'])
1935                    || message.is_empty()
1936                    || message.len() > MAX_DIAGNOSTIC_BYTES
1937                    || message.contains(['\n', '\r', '\0'])
1938                {
1939                    bail!("invalid shim debug record");
1940                }
1941                if target == "mbx::unit-outcome" {
1942                    let (unit, outcome): (String, String) = serde_json::from_str(&message)?;
1943                    if unit.len() > 256
1944                        || !unit.contains(':')
1945                        || !matches!(
1946                            outcome.as_str(),
1947                            "hit"
1948                                | "miss"
1949                                | "bypass"
1950                                | "unconsulted"
1951                                | "incremental"
1952                                | "verification"
1953                        )
1954                    {
1955                        bail!("invalid unit outcome");
1956                    }
1957                    self.stats
1958                        .unit_outcomes
1959                        .lock()
1960                        .unwrap()
1961                        .entry(unit)
1962                        .or_default()
1963                        .insert(outcome);
1964                } else {
1965                    debug!(target: &target, "{message}");
1966                }
1967                Ok(AgentResponse::DebugRecorded)
1968            })(),
1969            AgentRequest::RecordError { message } => {
1970                self.record_diagnostic(DiagnosticLevel::Error, message)
1971            }
1972            AgentRequest::FindFileDigests { scope, files } => self.find_file_digests(scope, files),
1973            AgentRequest::JoinActionPromise {
1974                adapter,
1975                invocation,
1976            } => self.join_action_promise(&adapter, &invocation).await,
1977            AgentRequest::CompleteActionPromise { claim, prediction } => {
1978                self.complete_action_promise(&claim, &prediction).await
1979            }
1980            AgentRequest::ResolveFileDigests { scope, files } => {
1981                self.resolve_file_digests(scope, files).await
1982            }
1983            AgentRequest::RecordFileDigests { scope, entries } => {
1984                self.record_file_digests(scope, entries)
1985            }
1986            AgentRequest::RecordWrapperTiming { timing } => (|| {
1987                if timing.spans.len() > 512
1988                    || timing.phases_ns.len() > 32
1989                    || timing.adapter.len() > 64
1990                    || timing.unit.as_ref().is_some_and(|s| s.len() > 1024)
1991                    || timing.phases_ns.keys().any(|s| s.len() > 64)
1992                    || timing.spans.iter().any(|s| s.name.len() > 64)
1993                {
1994                    bail!("wrapper timing exceeds telemetry limits");
1995                }
1996                {
1997                    let mut phases = self.stats.wrapper_phases_ns.lock().unwrap();
1998                    for (phase, duration) in &timing.phases_ns {
1999                        let total = phases.entry(phase.clone()).or_default();
2000                        *total = total.saturating_add(*duration);
2001                    }
2002                }
2003                self.emit(|| AgentEvent::WrapperTiming { timing });
2004                Ok(AgentResponse::WrapperTimingRecorded)
2005            })(),
2006            AgentRequest::RecordCompilerInvocation {
2007                outcome,
2008                crate_name,
2009                duration_ns,
2010            } => {
2011                let diagnostic = connection
2012                    .take_action_diagnostic(&outcome, crate_name.as_deref())
2013                    .map(|diagnostic| AgentEvent::ActionDiagnostic {
2014                        outcome: outcome.clone(),
2015                        crate_name: crate_name.clone(),
2016                        diagnostic,
2017                    });
2018                self.record_compiler_invocation(
2019                    &outcome,
2020                    crate_name.as_deref(),
2021                    duration_ns,
2022                    diagnostic,
2023                )
2024            }
2025            AgentRequest::RecordActionVerification { matched, restore } => {
2026                self.record_materialization(restore);
2027                self.stats.verifications.fetch_add(1, Ordering::Relaxed);
2028                if !matched {
2029                    self.stats.divergences.fetch_add(1, Ordering::Relaxed);
2030                }
2031                self.emit(|| AgentEvent::Verification { matched, restore });
2032                Ok(AgentResponse::ActionVerificationRecorded)
2033            }
2034            AgentRequest::StoreActionResult { result } => {
2035                self.store_action_result(&result, connection).await
2036            }
2037            AgentRequest::FindActionPrediction { task, invocation } => {
2038                self.find_action_prediction(&task, &invocation)
2039            }
2040            AgentRequest::RecordActionPrediction { task, prediction } => {
2041                self.record_action_prediction(&task, prediction)
2042            }
2043            AgentRequest::FindExecutableIdentity {
2044                executable,
2045                environment,
2046            } => self.find_executable_identity(executable, environment),
2047            AgentRequest::StoreExecutableIdentity {
2048                executable,
2049                environment,
2050                stdout,
2051                pins,
2052            } => self.store_executable_identity(executable, environment, stdout, pins),
2053            AgentRequest::Hello { .. } => {
2054                Err(eyre::eyre!("hello is only valid as the first request"))
2055            }
2056        };
2057        result.unwrap_or_else(|error| AgentResponse::Error {
2058            message: error.to_string(),
2059        })
2060    }
2061
2062    async fn find_blob(&self, digest: &CacheDigest) -> Result<AgentResponse> {
2063        if let Some(path) = self.find_verified_blob(digest)? {
2064            return Ok(AgentResponse::Blob { path: Some(path) });
2065        }
2066        if !self.remote_mode.reads() {
2067            return Ok(AgentResponse::Blob { path: None });
2068        }
2069        let Some(remote) = &self.remote else {
2070            return Ok(AgentResponse::Blob { path: None });
2071        };
2072        match self.fetch_remote_blob(remote, digest).await {
2073            Ok(path) => Ok(AgentResponse::Blob { path: Some(path) }),
2074            Err(error) => {
2075                warn!(
2076                    "remote cache blob lookup failed for {}: {error}",
2077                    digest.hash
2078                );
2079                Ok(AgentResponse::Blob { path: None })
2080            }
2081        }
2082    }
2083
2084    async fn find_blobs(&self, digests: Vec<CacheDigest>) -> Result<AgentResponse> {
2085        let mut paths = BTreeMap::new();
2086        let mut missing = Vec::new();
2087        for digest in &digests {
2088            match self.find_verified_blob(digest)? {
2089                Some(path) => {
2090                    paths.insert(digest.clone(), path);
2091                }
2092                None => {
2093                    missing.push(digest.clone());
2094                }
2095            }
2096        }
2097
2098        if !missing.is_empty()
2099            && self.remote_mode.reads()
2100            && let Some(remote) = &self.remote
2101        {
2102            paths.extend(self.fetch_remote_blobs(remote, missing, None).await);
2103        }
2104
2105        Ok(AgentResponse::Blobs {
2106            paths: digests
2107                .into_iter()
2108                .map(|digest| paths.get(&digest).cloned())
2109                .collect(),
2110        })
2111    }
2112
2113    async fn store_blob(
2114        &self,
2115        digest: &CacheDigest,
2116        source: &Path,
2117        connection: &mut ConnectionUploads,
2118    ) -> Result<AgentResponse> {
2119        let path = {
2120            let lock = self.write_lock(digest);
2121            let _guard = lock.lock().await;
2122            if let Some(path) = self.find_verified_blob(digest)? {
2123                path
2124            } else {
2125                let path = self.cas.store_file(digest, source)?;
2126                self.remember_verified_blob(digest, &path);
2127                self.stats.stores.fetch_add(1, Ordering::Relaxed);
2128                self.stats
2129                    .stored_bytes
2130                    .fetch_add(digest.size, Ordering::Relaxed);
2131                path
2132            }
2133        };
2134        // The object is durable locally, so the build has what it needs and the
2135        // remote publication can happen after this request returns.
2136        if let Some(uploads) = &self.uploads {
2137            uploads.queue_blob(digest, path.clone(), connection);
2138        }
2139        Ok(AgentResponse::Stored { path })
2140    }
2141
2142    fn find_verified_blob(&self, digest: &CacheDigest) -> Result<Option<PathBuf>> {
2143        let remembered = self.verified_blobs.lock().unwrap().get(digest).cloned();
2144        if let Some(remembered) = remembered {
2145            // The contents behind this digest were hashed in full when the
2146            // session first reached for them. Hashing again on every later
2147            // lookup would re-read each dependency's rlib once per crate that
2148            // links it, so what is rechecked here is that the file has not
2149            // been written since: an overwrite moves the modification time,
2150            // and truncation or eviction changes the length or removes it.
2151            // Only a replacement that reproduces both -- which no writer in
2152            // the store does, since blobs are published by rename under a
2153            // content-derived name -- would go unnoticed until `mbx cache
2154            // verify` reads the store back in full.
2155            if remembered.is_unchanged() {
2156                return Ok(Some(remembered.path));
2157            }
2158            self.verified_blobs.lock().unwrap().remove(digest);
2159        }
2160        let path = self.cas.find(digest)?;
2161        if let Some(path) = &path {
2162            self.remember_verified_blob(digest, path);
2163        }
2164        Ok(path)
2165    }
2166
2167    fn remember_verified_blob(&self, digest: &CacheDigest, path: &Path) {
2168        let Some(verified) = VerifiedBlob::describe(path) else {
2169            // Without an identity to compare against there is nothing to
2170            // shortcut safely, so the digest stays unremembered and every
2171            // lookup re-reads it.
2172            return;
2173        };
2174        self.verified_blobs
2175            .lock()
2176            .unwrap()
2177            .insert(digest.clone(), verified);
2178    }
2179
2180    async fn find_action_result(&self, action: &CacheDigest) -> Result<AgentResponse> {
2181        if let Some(result) = self.actions.find(action)? {
2182            return Ok(AgentResponse::ActionResult {
2183                result: Some(result),
2184            });
2185        }
2186        if !self.remote_mode.reads() {
2187            return Ok(AgentResponse::ActionResult { result: None });
2188        }
2189        let Some(remote) = &self.remote else {
2190            return Ok(AgentResponse::ActionResult { result: None });
2191        };
2192        let lock = self.action_lock(action);
2193        let _guard = lock.lock().await;
2194        if let Some(result) = self.actions.find(action)? {
2195            return Ok(AgentResponse::ActionResult {
2196                result: Some(result),
2197            });
2198        }
2199        if let Some(result) = self
2200            .pending_remote_actions
2201            .lock()
2202            .unwrap()
2203            .get(action)
2204            .cloned()
2205        {
2206            return Ok(AgentResponse::ActionResult {
2207                result: Some(result),
2208            });
2209        }
2210        let _permit = self.remote_transfers.acquire().await?;
2211        match self.get_remote_action_result(remote, action).await {
2212            Ok(Some(result)) => {
2213                self.pending_remote_actions
2214                    .lock()
2215                    .unwrap()
2216                    .insert(action.clone(), result.clone());
2217                Ok(AgentResponse::ActionResult {
2218                    result: Some(result),
2219                })
2220            }
2221            Ok(None) => Ok(AgentResponse::ActionResult { result: None }),
2222            Err(error) => {
2223                self.note_remote_failure();
2224                warn!(
2225                    "remote cache action lookup failed for {}: {error}",
2226                    action.hash
2227                );
2228                Ok(AgentResponse::ActionResult { result: None })
2229            }
2230        }
2231    }
2232
2233    async fn store_action_result(
2234        &self,
2235        result: &RemoteActionResult,
2236        connection: &ConnectionUploads,
2237    ) -> Result<AgentResponse> {
2238        let path = self.actions.store(result)?;
2239        if let Some(uploads) = &self.uploads {
2240            uploads.queue_action_result(result, connection);
2241        }
2242        Ok(AgentResponse::ActionStored { path })
2243    }
2244
2245    async fn join_action_promise(
2246        &self,
2247        adapter: &str,
2248        invocation: &CacheDigest,
2249    ) -> Result<AgentResponse> {
2250        if !self.remote_mode.reads() || !self.remote_mode.writes() {
2251            return Ok(AgentResponse::ActionPromise {
2252                claim: None,
2253                prediction: None,
2254            });
2255        }
2256        let Some(remote) = &self.remote else {
2257            return Ok(AgentResponse::ActionPromise {
2258                claim: None,
2259                prediction: None,
2260            });
2261        };
2262        let deadline = Instant::now() + ACTION_PROMISE_WAIT;
2263        loop {
2264            let _permit = self.remote_transfers.acquire().await?;
2265            let state = remote.join_action_promise(invocation, adapter).await;
2266            drop(_permit);
2267            match state {
2268                Ok(Some(ActionPromiseState::Claimed { claim })) => {
2269                    return Ok(AgentResponse::ActionPromise {
2270                        claim: Some(claim),
2271                        prediction: None,
2272                    });
2273                }
2274                Ok(Some(ActionPromiseState::Complete { prediction })) => {
2275                    return Ok(AgentResponse::ActionPromise {
2276                        claim: None,
2277                        prediction: Some(prediction),
2278                    });
2279                }
2280                Ok(Some(ActionPromiseState::Pending { retry_after_ms }))
2281                    if Instant::now() < deadline =>
2282                {
2283                    tokio::time::sleep(Duration::from_millis(retry_after_ms.clamp(10, 5_000)))
2284                        .await;
2285                }
2286                Ok(Some(ActionPromiseState::Pending { .. }) | None) => {
2287                    return Ok(AgentResponse::ActionPromise {
2288                        claim: None,
2289                        prediction: None,
2290                    });
2291                }
2292                Ok(Some(_)) => {
2293                    return Ok(AgentResponse::ActionPromise {
2294                        claim: None,
2295                        prediction: None,
2296                    });
2297                }
2298                Err(error) => {
2299                    self.note_remote_failure();
2300                    warn!(
2301                        "remote cache action promise failed for {}: {error}",
2302                        invocation.hash
2303                    );
2304                    return Ok(AgentResponse::ActionPromise {
2305                        claim: None,
2306                        prediction: None,
2307                    });
2308                }
2309            }
2310        }
2311    }
2312
2313    async fn complete_action_promise(
2314        &self,
2315        claim: &str,
2316        prediction: &ActionPrediction,
2317    ) -> Result<AgentResponse> {
2318        prediction.validate()?;
2319        let Some(remote) = &self.remote else {
2320            return Ok(AgentResponse::ActionPromiseCompleted);
2321        };
2322        let Some(uploads) = &self.uploads else {
2323            return Ok(AgentResponse::ActionPromiseCompleted);
2324        };
2325        if uploads
2326            .wait_for_actions(std::slice::from_ref(&prediction.action))
2327            .await
2328            .contains(&prediction.action)
2329        {
2330            // Never promise an action result the server does not hold. The
2331            // server lease expires and another runner gets to repair it.
2332            return Ok(AgentResponse::ActionPromiseCompleted);
2333        }
2334        let completion = ActionPromiseCompletion {
2335            claim: claim.to_string(),
2336            prediction: prediction.clone(),
2337        };
2338        let _permit = self.remote_transfers.acquire().await?;
2339        if let Err(error) = remote
2340            .complete_action_promise(&prediction.invocation, &completion)
2341            .await
2342        {
2343            self.note_remote_failure();
2344            warn!(
2345                "remote cache action promise completion failed for {}: {error}",
2346                prediction.invocation.hash
2347            );
2348        }
2349        Ok(AgentResponse::ActionPromiseCompleted)
2350    }
2351
2352    /// Record that a remote operation failed and the build carried on without it.
2353    fn note_remote_failure(&self) {
2354        self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
2355    }
2356
2357    async fn get_remote_action_result(
2358        &self,
2359        remote: &RemoteCacheClient,
2360        action: &CacheDigest,
2361    ) -> Result<Option<RemoteActionResult>> {
2362        self.stats
2363            .remote_action_lookups
2364            .fetch_add(1, Ordering::Relaxed);
2365        let _timer = AtomicDurationTimer::start(&self.stats.remote_action_lookup_duration_ns);
2366        remote.get_action_result(action).await
2367    }
2368
2369    fn record_action_hit(
2370        &self,
2371        action: &CacheDigest,
2372        restore: RestoreStats,
2373        crate_name: Option<String>,
2374        diagnostic: Option<AgentEvent>,
2375    ) -> Result<AgentResponse> {
2376        validate_crate_name(crate_name.as_deref())?;
2377        if self.actions.find(action)?.is_none() {
2378            let pending = self.pending_remote_actions.lock().unwrap().remove(action);
2379            if let Some(result) = pending {
2380                self.actions.store(&result)?;
2381            } else {
2382                bail!("cannot record a hit for a missing action result");
2383            }
2384        }
2385        self.record_restore(restore);
2386        self.stats.hits.fetch_add(1, Ordering::Relaxed);
2387        self.emit_action(
2388            diagnostic,
2389            AgentEvent::ActionHit {
2390                crate_name,
2391                restore,
2392            },
2393        );
2394        Ok(AgentResponse::ActionHitRecorded)
2395    }
2396
2397    fn record_restore(&self, restore: RestoreStats) {
2398        self.record_materialization(restore);
2399        atomic_saturating_add(
2400            &self.stats.avoided_compiler_duration_ns,
2401            restore.avoided_compiler_duration_ns,
2402        );
2403        atomic_saturating_add(&self.stats.restored_output_files, restore.output_files);
2404        atomic_saturating_add(&self.stats.restored_output_bytes, restore.output_bytes);
2405        atomic_saturating_add(
2406            &self.stats.reflinked_output_files,
2407            restore.reflinked_output_files,
2408        );
2409        atomic_saturating_add(
2410            &self.stats.reflinked_output_bytes,
2411            restore.reflinked_output_bytes,
2412        );
2413        atomic_saturating_add(&self.stats.copied_output_files, restore.copied_output_files);
2414        atomic_saturating_add(&self.stats.copied_output_bytes, restore.copied_output_bytes);
2415        atomic_saturating_add(&self.stats.reused_output_files, restore.reused_output_files);
2416        atomic_saturating_add(&self.stats.reused_output_bytes, restore.reused_output_bytes);
2417    }
2418
2419    fn record_compiler_invocation(
2420        &self,
2421        outcome: &str,
2422        crate_name: Option<&str>,
2423        duration_ns: u64,
2424        diagnostic: Option<AgentEvent>,
2425    ) -> Result<AgentResponse> {
2426        if !matches!(
2427            outcome,
2428            "miss" | "unconsulted" | "bypass" | "verification" | "incremental"
2429        ) {
2430            bail!("invalid compiler invocation outcome");
2431        }
2432        validate_crate_name(crate_name)?;
2433        let mut compiler = self.stats.compiler.lock().unwrap();
2434        let stats = compiler.entry(outcome.to_string()).or_default();
2435        stats.invocations = stats.invocations.saturating_add(1);
2436        stats.duration_ns = stats.duration_ns.saturating_add(duration_ns);
2437        drop(compiler);
2438        if outcome != "verification"
2439            && let Some(crate_name) = crate_name.filter(|name| !name.is_empty())
2440        {
2441            let mut slow = self.stats.slow_compilations.lock().unwrap();
2442            let duration = slow.entry(crate_name.to_string()).or_default();
2443            *duration = duration.saturating_add(duration_ns);
2444        }
2445        self.emit_action(
2446            diagnostic,
2447            AgentEvent::CompilerInvocation {
2448                outcome: outcome.to_string(),
2449                crate_name: crate_name.map(str::to_string),
2450                duration_ns,
2451            },
2452        );
2453        Ok(AgentResponse::CompilerInvocationRecorded)
2454    }
2455
2456    fn record_materialization(&self, restore: RestoreStats) {
2457        atomic_saturating_add(&self.stats.materialization_duration_ns, restore.duration_ns);
2458    }
2459
2460    /// Surface a shim diagnostic on this process's stderr.
2461    ///
2462    /// The agent lives in the process that owns the session, so printing here
2463    /// reaches the terminal running the build rather than the stderr of the
2464    /// compilation the shim stands in for -- which build scripts read as part
2465    /// of the compiler's answer. Deduplicated because one cause tends to fire
2466    /// once per compilation, and capped so a message unique per compilation
2467    /// cannot scroll the build away. The level is part of the deduplication
2468    /// key, so a fatal reason is never hidden by an earlier warning with the
2469    /// same text.
2470    fn record_diagnostic(&self, level: DiagnosticLevel, message: String) -> Result<AgentResponse> {
2471        if message.is_empty()
2472            || message.len() > MAX_DIAGNOSTIC_BYTES
2473            || message.contains(['\n', '\r', '\0'])
2474        {
2475            bail!("invalid shim {}", level.label());
2476        }
2477        let mut diagnostics = self.diagnostics.lock().unwrap();
2478        let key = (level, message.clone());
2479        if !diagnostics.contains(&key) {
2480            // An acknowledgement promises the diagnostic was surfaced. A
2481            // fatal shim error must be able to fall back locally when this
2482            // channel can no longer display it.
2483            if diagnostics.len() >= MAX_DIAGNOSTICS {
2484                bail!("shim diagnostic limit reached");
2485            }
2486            eprintln!("mbx[{}]: {message}", level.label());
2487            self.emit(|| match level {
2488                DiagnosticLevel::Warning => AgentEvent::Warning {
2489                    message: message.clone(),
2490                },
2491                DiagnosticLevel::Error => AgentEvent::Error {
2492                    message: message.clone(),
2493                },
2494            });
2495            diagnostics.insert(key);
2496        }
2497        Ok(match level {
2498            DiagnosticLevel::Warning => AgentResponse::WarningRecorded,
2499            DiagnosticLevel::Error => AgentResponse::ErrorRecorded,
2500        })
2501    }
2502
2503    /// Answer file-digest lookups from the session ledger.
2504    ///
2505    /// A recorded digest is returned only when the requested identity matches
2506    /// the recorded one exactly. The caller statted the file to build the
2507    /// identity, so this is the [`VerifiedBlob`] freshness check with the
2508    /// stat moved to the side that was already making it.
2509    fn find_file_digests(
2510        &self,
2511        scope: FileDigestScope,
2512        files: Vec<FileIdentity>,
2513    ) -> Result<AgentResponse> {
2514        if files.len() > MAX_FILE_DIGEST_BATCH {
2515            bail!("too many file-digest lookups in one request");
2516        }
2517        let ledger = self.file_digests.lock().unwrap();
2518        let digests = files
2519            .into_iter()
2520            .map(|file| {
2521                let recorded = ledger.get(&(scope, file.path.clone()))?;
2522                (recorded.file == file).then(|| recorded.digest.clone())
2523            })
2524            .collect();
2525        Ok(AgentResponse::FileDigests { digests })
2526    }
2527
2528    /// Resolve ledger misses inside the agent so concurrent shims that name the
2529    /// same NFS object wait for and reuse one read instead of stampeding it.
2530    async fn resolve_file_digests(
2531        &self,
2532        scope: FileDigestScope,
2533        files: Vec<FileIdentity>,
2534    ) -> Result<AgentResponse> {
2535        if files.len() > MAX_FILE_DIGEST_BATCH {
2536            bail!("too many file-digest resolutions in one request");
2537        }
2538        for file in &files {
2539            if !file.path.is_absolute() {
2540                bail!("file-digest resolutions need absolute paths");
2541            }
2542        }
2543        let resolutions = stream::iter(
2544            files
2545                .into_iter()
2546                .map(|file| async move { self.resolve_file_digest(scope, file).await }),
2547        )
2548        .buffered(MAX_CONCURRENT_FILE_DIGESTS)
2549        .collect()
2550        .await;
2551        Ok(AgentResponse::FileDigestsResolved { resolutions })
2552    }
2553
2554    async fn resolve_file_digest(
2555        &self,
2556        scope: FileDigestScope,
2557        file: FileIdentity,
2558    ) -> FileDigestResolution {
2559        let lock = {
2560            let key = (scope, file.clone());
2561            let mut locks = self.file_digest_locks.lock().unwrap();
2562            locks.retain(|_, lock| lock.strong_count() > 0);
2563            if let Some(lock) = locks.get(&key).and_then(Weak::upgrade) {
2564                lock
2565            } else {
2566                let lock = Arc::new(FileDigestFlight {
2567                    lock: tokio::sync::Mutex::new(()),
2568                    resolution: Mutex::new(None),
2569                });
2570                locks.insert(key, Arc::downgrade(&lock));
2571                lock
2572            }
2573        };
2574        let _flight = lock.lock.lock().await;
2575        if let Some(resolution) = lock.resolution.lock().unwrap().clone() {
2576            return resolution;
2577        }
2578        if let Some(digest) = self
2579            .file_digests
2580            .lock()
2581            .unwrap()
2582            .get(&(scope, file.path.clone()))
2583            .filter(|recorded| recorded.file == file)
2584            .map(|recorded| recorded.digest.clone())
2585            .filter(|digest| digest.size == file.len && file.still_describes().unwrap_or(false))
2586        {
2587            return FileDigestResolution::Digest(digest);
2588        }
2589        let Ok(_permit) = self.file_digest_permits.acquire().await else {
2590            return FileDigestResolution::Unresolved;
2591        };
2592        #[cfg(test)]
2593        self.file_digest_reads.fetch_add(1, Ordering::Relaxed);
2594        let path = file.path.clone();
2595        let resolved =
2596            tokio::task::spawn_blocking(move || digest_file_validated(scope, &path)).await;
2597        let Ok(Ok(Some(observed))) = resolved else {
2598            return FileDigestResolution::Unresolved;
2599        };
2600        if observed.cache_identity.as_ref() != Some(&file) {
2601            return FileDigestResolution::Unresolved;
2602        }
2603        let resolution = match observed.resolution {
2604            FileDigestResolution::Digest(digest) if digest.size == file.len => {
2605                let _ = self.record_file_digests(
2606                    scope,
2607                    vec![RecordedFileDigest {
2608                        file: file.clone(),
2609                        digest: digest.clone(),
2610                    }],
2611                );
2612                FileDigestResolution::Digest(digest)
2613            }
2614            FileDigestResolution::EmbeddedTimestampMacro => {
2615                FileDigestResolution::EmbeddedTimestampMacro
2616            }
2617            FileDigestResolution::Digest(_) | FileDigestResolution::Unresolved => {
2618                FileDigestResolution::Unresolved
2619            }
2620        };
2621        *lock.resolution.lock().unwrap() = Some(resolution.clone());
2622        resolution
2623    }
2624
2625    /// Record digests of files a shim read in full, for later reuse.
2626    ///
2627    /// Capped rather than evicted: a session that outgrows the cap loses the
2628    /// shortcut for the overflow and nothing else, and no workload observed so
2629    /// far comes near it.
2630    fn record_file_digests(
2631        &self,
2632        scope: FileDigestScope,
2633        entries: Vec<RecordedFileDigest>,
2634    ) -> Result<AgentResponse> {
2635        if entries.len() > MAX_FILE_DIGEST_BATCH {
2636            bail!("too many file-digest records in one request");
2637        }
2638        for entry in &entries {
2639            validate_file_digest_record(entry)?;
2640        }
2641        let mut ledger = self.file_digests.lock().unwrap();
2642        for entry in entries {
2643            if ledger.len() >= MAX_FILE_DIGEST_ENTRIES
2644                && !ledger.contains_key(&(scope, entry.file.path.clone()))
2645            {
2646                break;
2647            }
2648            ledger.insert((scope, entry.file.path.clone()), entry);
2649        }
2650        Ok(AgentResponse::FileDigestsRecorded)
2651    }
2652
2653    /// Start the file-digest ledger from entries an earlier session left
2654    /// behind, so a file nothing has touched since is not read again by the
2655    /// first compilation of this session that names it.
2656    ///
2657    /// Each entry stands only while its recorded identity still matches the
2658    /// disk, the same rule a lookup applies to what this session recorded, so
2659    /// a stale seed costs a hash and nothing else. Entries this session has
2660    /// already recorded are kept over seeded ones. Returns how many were taken.
2661    pub fn seed_file_digests(&self, entries: Vec<(FileDigestScope, RecordedFileDigest)>) -> usize {
2662        self.seed_file_digests_with(|| entries)
2663    }
2664
2665    /// Seed the ledger from entries produced under its lock.
2666    ///
2667    /// A lookup that arrives while `load` is still reading waits for it rather
2668    /// than missing, so a caller can start the read in the background and let
2669    /// whatever else the build is doing overlap it.
2670    pub fn seed_file_digests_with(
2671        &self,
2672        load: impl FnOnce() -> Vec<(FileDigestScope, RecordedFileDigest)>,
2673    ) -> usize {
2674        let mut ledger = self.file_digests.lock().unwrap();
2675        let entries = load();
2676        let mut seeded = 0;
2677        for (scope, entry) in entries {
2678            if ledger.len() >= MAX_FILE_DIGEST_ENTRIES {
2679                break;
2680            }
2681            if validate_file_digest_record(&entry).is_err() {
2682                continue;
2683            }
2684            if let std::collections::btree_map::Entry::Vacant(slot) =
2685                ledger.entry((scope, entry.file.path.clone()))
2686            {
2687                slot.insert(entry);
2688                seeded += 1;
2689            }
2690        }
2691        seeded
2692    }
2693
2694    /// Everything the file-digest ledger holds, for a later session to start
2695    /// from.
2696    pub fn file_digests(&self) -> Vec<(FileDigestScope, RecordedFileDigest)> {
2697        self.file_digests
2698            .lock()
2699            .unwrap()
2700            .iter()
2701            .map(|((scope, _), entry)| (*scope, entry.clone()))
2702            .collect()
2703    }
2704
2705    fn find_action_prediction(
2706        &self,
2707        task: &str,
2708        invocation: &CacheDigest,
2709    ) -> Result<AgentResponse> {
2710        validate_task_identity(task)?;
2711        invocation.validate()?;
2712        let (prediction, prefetch) = {
2713            let mut tasks = self.task_actions.lock().unwrap();
2714            let Some(state) = tasks.get_mut(task) else {
2715                return Ok(AgentResponse::ActionPrediction { prediction: None });
2716            };
2717            activate_prediction_adapter(state, invocation)
2718        };
2719        if let Some(prefetch) = prefetch {
2720            self.spawn_prefetch_predictions(prefetch);
2721        }
2722        Ok(AgentResponse::ActionPrediction { prediction })
2723    }
2724
2725    fn record_action_prediction(
2726        &self,
2727        task: &str,
2728        prediction: ActionPrediction,
2729    ) -> Result<AgentResponse> {
2730        validate_task_identity(task)?;
2731        prediction.validate()?;
2732        let mut tasks = self.task_actions.lock().unwrap();
2733        let state = tasks.entry(task.to_string()).or_default();
2734        if !state
2735            .pending_predictions
2736            .contains_key(&prediction.invocation)
2737            && state.pending_predictions.len() >= MAX_TASK_ACTION_PREDICTIONS
2738        {
2739            bail!("this task run contains too many action predictions");
2740        }
2741        let current = state.predictions.get(&prediction.invocation);
2742        if current != Some(&prediction) {
2743            match state.changed_predictions.get(&prediction.invocation) {
2744                // Back to what the manifest held: no longer a change.
2745                Some(baseline) if baseline.as_ref() == Some(&prediction) => {
2746                    state.changed_predictions.remove(&prediction.invocation);
2747                }
2748                // Changed again; the baseline it left is already recorded.
2749                Some(_) => {}
2750                None => {
2751                    state
2752                        .changed_predictions
2753                        .insert(prediction.invocation.clone(), current.cloned());
2754                }
2755            }
2756        }
2757        state
2758            .predictions
2759            .insert(prediction.invocation.clone(), prediction.clone());
2760        state
2761            .pending_predictions
2762            .insert(prediction.invocation.clone(), prediction);
2763        Ok(AgentResponse::ActionPredictionRecorded)
2764    }
2765
2766    fn executable_identity_key(
2767        &self,
2768        executable: PathBuf,
2769        environment: BTreeMap<String, Option<String>>,
2770    ) -> Result<ExecutableIdentityKey> {
2771        // Restricted to the variables that actually select what an identity
2772        // probe reports: the toolchain rustup resolves, the SDK a linker
2773        // driver builds against, a `-fuse-ld` linker selection, and the
2774        // search path a driver finds its linker on. Anything else would let
2775        // one key stand for two different compilers.
2776        if !environment.keys().all(|name| {
2777            matches!(
2778                name.as_str(),
2779                "COMPILER_PATH"
2780                    | "GCC_EXEC_PREFIX"
2781                    | "LIBRARY_PATH"
2782                    | "MBX_FUSE_LD"
2783                    | "PATH"
2784                    | "RUSTUP_HOME"
2785                    | "RUSTUP_TOOLCHAIN"
2786                    | "SDKROOT"
2787                    | "MACOSX_DEPLOYMENT_TARGET"
2788                    | "LIB"
2789                    | "UCRTVersion"
2790                    | "UniversalCRTSdkDir"
2791                    | "VCToolsInstallDir"
2792                    | "VCToolsVersion"
2793                    | "WindowsSdkDir"
2794                    | "WindowsSDKVersion"
2795            )
2796        }) {
2797            bail!("executable identity contains an unsupported environment variable");
2798        }
2799        Ok(ExecutableIdentityKey {
2800            executable,
2801            environment,
2802        })
2803    }
2804
2805    fn find_executable_identity(
2806        &self,
2807        executable: PathBuf,
2808        environment: BTreeMap<String, Option<String>>,
2809    ) -> Result<AgentResponse> {
2810        let key = self.executable_identity_key(executable, environment)?;
2811        let remembered = self
2812            .executable_identities
2813            .lock()
2814            .unwrap()
2815            .get(&key)
2816            .cloned();
2817        let stdout = match remembered {
2818            Some(stdout) => Some(stdout),
2819            None => {
2820                let persisted = self.load_persisted_identity(&key);
2821                if let Some(stdout) = &persisted {
2822                    // Memoized like a fresh probe: the pins were checked once
2823                    // for this session, which is as often as a probe runs.
2824                    let _ = self.remember_executable_identity(&key, stdout.clone());
2825                }
2826                persisted
2827            }
2828        };
2829        Ok(AgentResponse::ExecutableIdentity { stdout })
2830    }
2831
2832    fn store_executable_identity(
2833        &self,
2834        executable: PathBuf,
2835        environment: BTreeMap<String, Option<String>>,
2836        stdout: Vec<u8>,
2837        pins: Vec<PinnedFile>,
2838    ) -> Result<AgentResponse> {
2839        if stdout.len() > MAX_EXECUTABLE_IDENTITY_SIZE {
2840            bail!("executable identity exceeds {MAX_EXECUTABLE_IDENTITY_SIZE} bytes");
2841        }
2842        let key = self.executable_identity_key(executable, environment)?;
2843        self.remember_executable_identity(&key, stdout.clone())?;
2844        // Best-effort, like the ledger: a session that cannot write it probes
2845        // again next time, which is what it would have done anyway. A pin
2846        // that no longer holds means a file moved under the probe, and what
2847        // it printed describes neither the old file nor the new one.
2848        if pins.is_empty() {
2849        } else if !pins.iter().all(PinnedFile::holds) {
2850            debug!("executable identity was not persisted: a pinned file changed during the probe");
2851        } else if let Err(error) = self.persist_executable_identity(&key, pins, &stdout) {
2852            debug!("executable identity was not persisted: {error:#}");
2853        }
2854        Ok(AgentResponse::ExecutableIdentity {
2855            stdout: Some(stdout),
2856        })
2857    }
2858
2859    fn remember_executable_identity(
2860        &self,
2861        key: &ExecutableIdentityKey,
2862        stdout: Vec<u8>,
2863    ) -> Result<()> {
2864        let mut identities = self.executable_identities.lock().unwrap();
2865        let is_new = !identities.contains_key(key);
2866        let previous_size = identities.get(key).map_or(0, Vec::len);
2867        if is_new && identities.len() >= MAX_EXECUTABLE_IDENTITIES {
2868            bail!("executable identity cache contains too many entries");
2869        }
2870        let retained_bytes = identities.values().map(Vec::len).sum::<usize>();
2871        if retained_bytes - previous_size + stdout.len() > MAX_EXECUTABLE_IDENTITY_BYTES {
2872            bail!("executable identity cache contains too many bytes");
2873        }
2874        identities.insert(key.clone(), stdout);
2875        Ok(())
2876    }
2877
2878    fn persisted_identity_path(&self, key: &ExecutableIdentityKey) -> Result<PathBuf> {
2879        let selector = serde_json::to_vec(&(&key.executable, &key.environment))?;
2880        let digest = CacheDigest::blake3(&selector);
2881        Ok(self.identity_dir.join(format!("{}.json", digest.hash)))
2882    }
2883
2884    /// What an earlier session recorded for `key`, if every file it was
2885    /// pinned to is still as that session saw it.
2886    fn load_persisted_identity(&self, key: &ExecutableIdentityKey) -> Option<Vec<u8>> {
2887        let path = self.persisted_identity_path(key).ok()?;
2888        let bytes = fs::read(path).ok()?;
2889        let persisted: PersistedExecutableIdentity = serde_json::from_slice(&bytes).ok()?;
2890        if persisted.version != PERSISTED_EXECUTABLE_IDENTITY_VERSION
2891            || persisted.executable != key.executable
2892            || persisted.environment != key.environment
2893            || persisted.pins.is_empty()
2894            || persisted.stdout.len() > MAX_EXECUTABLE_IDENTITY_SIZE
2895        {
2896            return None;
2897        }
2898        persisted
2899            .pins
2900            .iter()
2901            .all(PinnedFile::holds)
2902            .then(|| persisted.stdout.into_bytes())
2903    }
2904
2905    fn persist_executable_identity(
2906        &self,
2907        key: &ExecutableIdentityKey,
2908        pins: Vec<PinnedFile>,
2909        stdout: &[u8],
2910    ) -> Result<()> {
2911        // An identity that is not text is kept for this session only.
2912        let Ok(stdout) = std::str::from_utf8(stdout) else {
2913            return Ok(());
2914        };
2915        let persisted = PersistedExecutableIdentity {
2916            version: PERSISTED_EXECUTABLE_IDENTITY_VERSION,
2917            executable: key.executable.clone(),
2918            environment: key.environment.clone(),
2919            pins,
2920            stdout: stdout.to_owned(),
2921        };
2922        let path = self.persisted_identity_path(key)?;
2923        fs::create_dir_all(self.identity_dir.as_path())?;
2924        // Not synced: a torn record reads as no record, and costs one probe.
2925        let mut temporary = tempfile::NamedTempFile::new_in(self.identity_dir.as_path())?;
2926        std::io::Write::write_all(temporary.as_file_mut(), &serde_json::to_vec(&persisted)?)?;
2927        temporary.persist(path).map_err(|error| error.error)?;
2928        Ok(())
2929    }
2930
2931    /// Serve newline-delimited protocol requests on an authenticated session stream.
2932    pub async fn handle_connection<S>(&self, stream: S) -> Result<()>
2933    where
2934        S: AsyncRead + AsyncWrite + Unpin,
2935    {
2936        let (reader, mut writer) = tokio::io::split(stream);
2937        let mut reader = BufReader::new(reader);
2938        let hello = read_request(&mut reader)
2939            .await?
2940            .ok_or_else(|| eyre::eyre!("connection closed before the agent handshake"))?;
2941        let request: AgentRequest = serde_json::from_str(&hello)?;
2942        match request {
2943            AgentRequest::Hello {
2944                protocol,
2945                client_version,
2946            } if protocol == AGENT_PROTOCOL_VERSION && client_version == self.version.as_ref() => {}
2947            AgentRequest::Hello { protocol, .. } if protocol != AGENT_PROTOCOL_VERSION => {
2948                send_response(
2949                    &mut writer,
2950                    &AgentResponse::Error {
2951                        message: format!(
2952                            "unsupported agent protocol {protocol}; expected {AGENT_PROTOCOL_VERSION}"
2953                        ),
2954                    },
2955                )
2956                .await?;
2957                return Ok(());
2958            }
2959            AgentRequest::Hello { client_version, .. } => {
2960                send_response(
2961                    &mut writer,
2962                    &AgentResponse::Error {
2963                        message: format!(
2964                            "cache client {client_version} does not match agent {}",
2965                            self.version
2966                        ),
2967                    },
2968                )
2969                .await?;
2970                return Ok(());
2971            }
2972            _ => bail!("the first agent request must be hello"),
2973        }
2974        send_response(
2975            &mut writer,
2976            &AgentResponse::Hello {
2977                protocol: AGENT_PROTOCOL_VERSION,
2978                agent_version: self.version.to_string(),
2979            },
2980        )
2981        .await?;
2982
2983        // Tickets accumulate for the life of the connection because a shim
2984        // publishes a compilation's blobs and the action result naming them over
2985        // one connection, in that order.
2986        let mut connection = ConnectionUploads::default();
2987        while let Some(line) = read_request(&mut reader).await? {
2988            let response = match serde_json::from_str(&line) {
2989                Ok(request) => self.respond_on(request, &mut connection).await,
2990                Err(error) => AgentResponse::Error {
2991                    message: format!("invalid agent request: {error}"),
2992                },
2993            };
2994            send_response(&mut writer, &response).await?;
2995        }
2996        Ok(())
2997    }
2998}
2999
3000/// Read one newline-delimited request, refusing one that grows past the cap.
3001///
3002/// Any process running as this user can open the session socket, so a request
3003/// that never terminates its line must not be able to grow the agent's memory
3004/// without bound.
3005async fn read_request<R>(reader: &mut R) -> Result<Option<String>>
3006where
3007    R: AsyncBufRead + Unpin,
3008{
3009    let mut line = Vec::new();
3010    loop {
3011        let available = reader.fill_buf().await?;
3012        if available.is_empty() {
3013            break;
3014        }
3015        let (consumed, complete) = match available.iter().position(|byte| *byte == b'\n') {
3016            Some(index) => (index, true),
3017            None => (available.len(), false),
3018        };
3019        if line.len() + consumed > MAX_REQUEST_BYTES {
3020            bail!("agent request exceeded {MAX_REQUEST_BYTES} bytes");
3021        }
3022        line.extend_from_slice(&available[..consumed]);
3023        // The newline itself is consumed but never kept.
3024        reader.consume(consumed + usize::from(complete));
3025        if complete {
3026            return Ok(Some(String::from_utf8(line)?));
3027        }
3028    }
3029    if line.is_empty() {
3030        Ok(None)
3031    } else {
3032        Ok(Some(String::from_utf8(line)?))
3033    }
3034}
3035
3036async fn send_response(
3037    writer: &mut (impl AsyncWrite + Unpin),
3038    response: &AgentResponse,
3039) -> Result<()> {
3040    let mut encoded = serde_json::to_vec(response)?;
3041    encoded.push(b'\n');
3042    writer.write_all(&encoded).await?;
3043    writer.flush().await?;
3044    Ok(())
3045}
3046
3047#[cfg(test)]
3048#[path = "agent_tests.rs"]
3049mod tests;