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