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