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