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