Skip to main content

mbx_cache_core/
agent.rs

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