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            AgentRequest::FindActionResult { action } => {
1954                self.stats.lookups.fetch_add(1, Ordering::Relaxed);
1955                self.find_action_result(&action).await
1956            }
1957            AgentRequest::RecordActionHit {
1958                action,
1959                restore,
1960                crate_name,
1961            } => {
1962                let diagnostic = connection
1963                    .take_action_diagnostic("hit", crate_name.as_deref())
1964                    .map(|diagnostic| AgentEvent::ActionDiagnostic {
1965                        outcome: "hit".into(),
1966                        crate_name: crate_name.clone(),
1967                        diagnostic,
1968                    });
1969                self.record_action_hit(&action, restore, crate_name, diagnostic)
1970            }
1971            AgentRequest::RecordBypass { kind } => {
1972                *self
1973                    .stats
1974                    .bypasses
1975                    .lock()
1976                    .unwrap()
1977                    .entry(kind.clone())
1978                    .or_insert(0) += 1;
1979                // Held until the connection's next compiler run, which is the
1980                // compilation this bypass describes when the shim reports the
1981                // two together. A reason nothing claims is emitted on its own.
1982                if let Some(unclaimed) = connection.hold_bypass(kind) {
1983                    self.emit(|| AgentEvent::Bypass { kind: unclaimed });
1984                }
1985                Ok(AgentResponse::BypassRecorded)
1986            }
1987            AgentRequest::RecordUnconsulted => {
1988                self.stats.unconsulted.fetch_add(1, Ordering::Relaxed);
1989                self.emit(|| AgentEvent::Unconsulted);
1990                Ok(AgentResponse::UnconsultedRecorded)
1991            }
1992            AgentRequest::RecordWarning { message } => match parse_action_diagnostic(&message) {
1993                Some(Ok((outcome, crate_name, diagnostic))) => {
1994                    connection.record_action_diagnostic(outcome, crate_name, diagnostic);
1995                    Ok(AgentResponse::WarningRecorded)
1996                }
1997                Some(Err(error)) => Err(error),
1998                None => self.record_diagnostic(DiagnosticLevel::Warning, message),
1999            },
2000            AgentRequest::RecordDebug { target, message } => (|| {
2001                if target.is_empty()
2002                    || target.len() > 256
2003                    || target.contains(['\n', '\r', '\0'])
2004                    || message.is_empty()
2005                    || message.len() > MAX_DIAGNOSTIC_BYTES
2006                    || message.contains(['\n', '\r', '\0'])
2007                {
2008                    bail!("invalid shim debug record");
2009                }
2010                if target == "mbx::unit-outcome" {
2011                    let (unit, outcome): (String, String) = serde_json::from_str(&message)?;
2012                    // Recorded as what the cache did, which is the column the
2013                    // TUI shows. Whether a compilation kept incremental state
2014                    // is reported by the summary, not per unit.
2015                    let (outcome, _) = lookup_outcome(&outcome);
2016                    if unit.len() > 256
2017                        || !unit.contains(':')
2018                        || !matches!(
2019                            outcome,
2020                            "hit" | "miss" | "bypass" | "unconsulted" | "verification"
2021                        )
2022                    {
2023                        bail!("invalid unit outcome");
2024                    }
2025                    self.stats
2026                        .unit_outcomes
2027                        .lock()
2028                        .unwrap()
2029                        .entry(unit)
2030                        .or_default()
2031                        .insert(outcome.to_string());
2032                } else {
2033                    debug!(target: &target, "{message}");
2034                }
2035                Ok(AgentResponse::DebugRecorded)
2036            })(),
2037            AgentRequest::RecordError { message } => {
2038                self.record_diagnostic(DiagnosticLevel::Error, message)
2039            }
2040            AgentRequest::FindFileDigests { scope, files } => self.find_file_digests(scope, files),
2041            AgentRequest::JoinActionPromise {
2042                adapter,
2043                invocation,
2044            } => self.join_action_promise(&adapter, &invocation).await,
2045            AgentRequest::CompleteActionPromise { claim, prediction } => {
2046                self.complete_action_promise(&claim, &prediction).await
2047            }
2048            AgentRequest::ResolveFileDigests { scope, files } => {
2049                self.resolve_file_digests(scope, files).await
2050            }
2051            AgentRequest::RecordFileDigests { scope, entries } => {
2052                self.record_file_digests(scope, entries)
2053            }
2054            AgentRequest::RecordWrapperTiming { timing } => (|| {
2055                if timing.spans.len() > 512
2056                    || timing.phases_ns.len() > 32
2057                    || timing.adapter.len() > 64
2058                    || timing.unit.as_ref().is_some_and(|s| s.len() > 1024)
2059                    || timing.unit_id.as_ref().is_some_and(|s| s.len() > 256)
2060                    || timing.dependencies.len() > 4096
2061                    || timing.dependencies.iter().any(|s| s.len() > 256)
2062                    || timing.phases_ns.keys().any(|s| s.len() > 64)
2063                    || timing.spans.iter().any(|s| s.name.len() > 64)
2064                {
2065                    bail!("wrapper timing exceeds telemetry limits");
2066                }
2067                {
2068                    let mut phases = self.stats.wrapper_phases_ns.lock().unwrap();
2069                    for (phase, duration) in &timing.phases_ns {
2070                        let total = phases.entry(phase.clone()).or_default();
2071                        *total = total.saturating_add(*duration);
2072                    }
2073                }
2074                self.emit(|| AgentEvent::WrapperTiming { timing });
2075                Ok(AgentResponse::WrapperTimingRecorded)
2076            })(),
2077            AgentRequest::RecordCompilerInvocation {
2078                outcome,
2079                crate_name,
2080                duration_ns,
2081            } => {
2082                // Normalized before the diagnostic is claimed: it was stored
2083                // under the lookup outcome, and an unnormalized name here would
2084                // never find it.
2085                let (outcome, incremental) = lookup_outcome(&outcome);
2086                let outcome = outcome.to_string();
2087                let bypass = if outcome == "bypass" {
2088                    connection.take_bypass()
2089                } else {
2090                    self.release_bypass(connection);
2091                    None
2092                };
2093                let diagnostic = connection
2094                    .take_action_diagnostic(&outcome, crate_name.as_deref())
2095                    .map(|diagnostic| AgentEvent::ActionDiagnostic {
2096                        outcome: outcome.clone(),
2097                        crate_name: crate_name.clone(),
2098                        diagnostic,
2099                    });
2100                self.record_compiler_invocation(
2101                    &outcome,
2102                    incremental,
2103                    crate_name.as_deref(),
2104                    duration_ns,
2105                    diagnostic,
2106                    bypass,
2107                )
2108            }
2109            AgentRequest::RecordActionVerification { matched, restore } => {
2110                self.record_materialization(restore);
2111                self.stats.verifications.fetch_add(1, Ordering::Relaxed);
2112                if !matched {
2113                    self.stats.divergences.fetch_add(1, Ordering::Relaxed);
2114                }
2115                self.emit(|| AgentEvent::Verification { matched, restore });
2116                Ok(AgentResponse::ActionVerificationRecorded)
2117            }
2118            AgentRequest::StoreActionResult { result } => {
2119                self.store_action_result(&result, connection).await
2120            }
2121            AgentRequest::FindActionPrediction { task, invocation } => {
2122                self.find_action_prediction(&task, &invocation)
2123            }
2124            AgentRequest::RecordActionPrediction { task, prediction } => {
2125                self.record_action_prediction(&task, prediction)
2126            }
2127            AgentRequest::FindExecutableIdentity {
2128                executable,
2129                environment,
2130            } => self.find_executable_identity(executable, environment),
2131            AgentRequest::StoreExecutableIdentity {
2132                executable,
2133                environment,
2134                stdout,
2135                pins,
2136            } => self.store_executable_identity(executable, environment, stdout, pins),
2137            AgentRequest::Hello { .. } => {
2138                Err(eyre::eyre!("hello is only valid as the first request"))
2139            }
2140        };
2141        result.unwrap_or_else(|error| AgentResponse::Error {
2142            message: error.to_string(),
2143        })
2144    }
2145
2146    async fn find_blob(&self, digest: &CacheDigest) -> Result<AgentResponse> {
2147        if let Some(path) = self.find_verified_blob(digest)? {
2148            return Ok(AgentResponse::Blob { path: Some(path) });
2149        }
2150        if !self.remote_mode.reads() {
2151            return Ok(AgentResponse::Blob { path: None });
2152        }
2153        let Some(remote) = &self.remote else {
2154            return Ok(AgentResponse::Blob { path: None });
2155        };
2156        match self.fetch_remote_blob(remote, digest).await {
2157            Ok(path) => Ok(AgentResponse::Blob { path: Some(path) }),
2158            Err(error) => {
2159                warn!(
2160                    "remote cache blob lookup failed for {}: {error}",
2161                    digest.hash
2162                );
2163                Ok(AgentResponse::Blob { path: None })
2164            }
2165        }
2166    }
2167
2168    async fn find_blobs(&self, digests: Vec<CacheDigest>) -> Result<AgentResponse> {
2169        let mut paths = BTreeMap::new();
2170        let mut missing = Vec::new();
2171        for digest in &digests {
2172            match self.find_verified_blob(digest)? {
2173                Some(path) => {
2174                    paths.insert(digest.clone(), path);
2175                }
2176                None => {
2177                    missing.push(digest.clone());
2178                }
2179            }
2180        }
2181
2182        if !missing.is_empty()
2183            && self.remote_mode.reads()
2184            && let Some(remote) = &self.remote
2185        {
2186            paths.extend(self.fetch_remote_blobs(remote, missing, None).await);
2187        }
2188
2189        Ok(AgentResponse::Blobs {
2190            paths: digests
2191                .into_iter()
2192                .map(|digest| paths.get(&digest).cloned())
2193                .collect(),
2194        })
2195    }
2196
2197    async fn store_blob(
2198        &self,
2199        digest: &CacheDigest,
2200        source: &Path,
2201        connection: &mut ConnectionUploads,
2202    ) -> Result<AgentResponse> {
2203        let path = {
2204            let lock = self.write_lock(digest);
2205            let _guard = lock.lock().await;
2206            if let Some(path) = self.find_verified_blob(digest)? {
2207                path
2208            } else {
2209                let timer = AtomicDurationTimer::start(&self.stats.local_cas_write_duration_ns);
2210                let path = self.cas.store_file(digest, source)?;
2211                drop(timer);
2212                self.remember_verified_blob(digest, &path);
2213                self.stats.stores.fetch_add(1, Ordering::Relaxed);
2214                self.stats
2215                    .stored_bytes
2216                    .fetch_add(digest.size, Ordering::Relaxed);
2217                path
2218            }
2219        };
2220        // The object is durable locally, so the build has what it needs and the
2221        // remote publication can happen after this request returns.
2222        if let Some(uploads) = &self.uploads {
2223            uploads.queue_blob(digest, path.clone(), connection);
2224        }
2225        Ok(AgentResponse::Stored { path })
2226    }
2227
2228    fn find_verified_blob(&self, digest: &CacheDigest) -> Result<Option<PathBuf>> {
2229        let remembered = self.verified_blobs.lock().unwrap().get(digest).cloned();
2230        if let Some(remembered) = remembered {
2231            // The contents behind this digest were hashed in full when the
2232            // session first reached for them. Hashing again on every later
2233            // lookup would re-read each dependency's rlib once per crate that
2234            // links it, so what is rechecked here is that the file has not
2235            // been written since: an overwrite moves the modification time,
2236            // and truncation or eviction changes the length or removes it.
2237            // Only a replacement that reproduces both -- which no writer in
2238            // the store does, since blobs are published by rename under a
2239            // content-derived name -- would go unnoticed until `mbx cache
2240            // verify` reads the store back in full.
2241            if remembered.is_unchanged() {
2242                return Ok(Some(remembered.path));
2243            }
2244            self.verified_blobs.lock().unwrap().remove(digest);
2245        }
2246        // Blobs are published by rename without an fsync, so a crash can leave
2247        // one torn under a valid name. It is not restored from, and it is not
2248        // an error either: the compilation runs again and its publication
2249        // replaces the bad bytes, the same way a blob that was never stored
2250        // gets there. Reporting it as a failure would surface a warning for a
2251        // build that is behaving correctly.
2252        let path = match self.cas.find(digest) {
2253            Ok(path) => path,
2254            Err(error) => {
2255                warn!(
2256                    "local CAS blob {} is not restorable and will be republished: {error}",
2257                    digest.hash
2258                );
2259                return Ok(None);
2260            }
2261        };
2262        if let Some(path) = &path {
2263            self.remember_verified_blob(digest, path);
2264        }
2265        Ok(path)
2266    }
2267
2268    fn remember_verified_blob(&self, digest: &CacheDigest, path: &Path) {
2269        let Some(verified) = VerifiedBlob::describe(path) else {
2270            // Without an identity to compare against there is nothing to
2271            // shortcut safely, so the digest stays unremembered and every
2272            // lookup re-reads it.
2273            return;
2274        };
2275        self.verified_blobs
2276            .lock()
2277            .unwrap()
2278            .insert(digest.clone(), verified);
2279    }
2280
2281    async fn find_action_result(&self, action: &CacheDigest) -> Result<AgentResponse> {
2282        if let Some(result) = self.actions.find(action)? {
2283            return Ok(AgentResponse::ActionResult {
2284                result: Some(result),
2285            });
2286        }
2287        if !self.remote_mode.reads() {
2288            return Ok(AgentResponse::ActionResult { result: None });
2289        }
2290        let Some(remote) = &self.remote else {
2291            return Ok(AgentResponse::ActionResult { result: None });
2292        };
2293        let lock = self.action_lock(action);
2294        let _guard = lock.lock().await;
2295        if let Some(result) = self.actions.find(action)? {
2296            return Ok(AgentResponse::ActionResult {
2297                result: Some(result),
2298            });
2299        }
2300        if let Some(result) = self
2301            .pending_remote_actions
2302            .lock()
2303            .unwrap()
2304            .get(action)
2305            .cloned()
2306        {
2307            return Ok(AgentResponse::ActionResult {
2308                result: Some(result),
2309            });
2310        }
2311        let _permit = self.remote_transfers.acquire().await?;
2312        match self.get_remote_action_result(remote, action).await {
2313            Ok(Some(result)) => {
2314                self.pending_remote_actions
2315                    .lock()
2316                    .unwrap()
2317                    .insert(action.clone(), result.clone());
2318                Ok(AgentResponse::ActionResult {
2319                    result: Some(result),
2320                })
2321            }
2322            Ok(None) => Ok(AgentResponse::ActionResult { result: None }),
2323            Err(error) => {
2324                self.note_remote_failure();
2325                warn!(
2326                    "remote cache action lookup failed for {}: {error}",
2327                    action.hash
2328                );
2329                Ok(AgentResponse::ActionResult { result: None })
2330            }
2331        }
2332    }
2333
2334    async fn store_action_result(
2335        &self,
2336        result: &RemoteActionResult,
2337        connection: &ConnectionUploads,
2338    ) -> Result<AgentResponse> {
2339        let path = self.actions.store(result)?;
2340        if let Some(uploads) = &self.uploads {
2341            uploads.queue_action_result(result, connection);
2342        }
2343        Ok(AgentResponse::ActionStored { path })
2344    }
2345
2346    async fn join_action_promise(
2347        &self,
2348        adapter: &str,
2349        invocation: &CacheDigest,
2350    ) -> Result<AgentResponse> {
2351        if !self.remote_mode.reads() || !self.remote_mode.writes() {
2352            return Ok(AgentResponse::ActionPromise {
2353                claim: None,
2354                prediction: None,
2355            });
2356        }
2357        let Some(remote) = &self.remote else {
2358            return Ok(AgentResponse::ActionPromise {
2359                claim: None,
2360                prediction: None,
2361            });
2362        };
2363        let deadline = Instant::now() + ACTION_PROMISE_WAIT;
2364        loop {
2365            let _permit = self.remote_transfers.acquire().await?;
2366            let state = remote.join_action_promise(invocation, adapter).await;
2367            drop(_permit);
2368            match state {
2369                Ok(Some(ActionPromiseState::Claimed { claim })) => {
2370                    return Ok(AgentResponse::ActionPromise {
2371                        claim: Some(claim),
2372                        prediction: None,
2373                    });
2374                }
2375                Ok(Some(ActionPromiseState::Complete { prediction })) => {
2376                    return Ok(AgentResponse::ActionPromise {
2377                        claim: None,
2378                        prediction: Some(prediction),
2379                    });
2380                }
2381                Ok(Some(ActionPromiseState::Pending { retry_after_ms }))
2382                    if Instant::now() < deadline =>
2383                {
2384                    tokio::time::sleep(Duration::from_millis(retry_after_ms.clamp(10, 5_000)))
2385                        .await;
2386                }
2387                Ok(Some(ActionPromiseState::Pending { .. }) | None) => {
2388                    return Ok(AgentResponse::ActionPromise {
2389                        claim: None,
2390                        prediction: None,
2391                    });
2392                }
2393                Ok(Some(_)) => {
2394                    return Ok(AgentResponse::ActionPromise {
2395                        claim: None,
2396                        prediction: None,
2397                    });
2398                }
2399                Err(error) => {
2400                    self.note_remote_failure();
2401                    warn!(
2402                        "remote cache action promise failed for {}: {error}",
2403                        invocation.hash
2404                    );
2405                    return Ok(AgentResponse::ActionPromise {
2406                        claim: None,
2407                        prediction: None,
2408                    });
2409                }
2410            }
2411        }
2412    }
2413
2414    async fn complete_action_promise(
2415        &self,
2416        claim: &str,
2417        prediction: &ActionPrediction,
2418    ) -> Result<AgentResponse> {
2419        prediction.validate()?;
2420        let Some(remote) = &self.remote else {
2421            return Ok(AgentResponse::ActionPromiseCompleted);
2422        };
2423        let Some(uploads) = &self.uploads else {
2424            return Ok(AgentResponse::ActionPromiseCompleted);
2425        };
2426        if uploads
2427            .wait_for_actions(std::slice::from_ref(&prediction.action))
2428            .await
2429            .contains(&prediction.action)
2430        {
2431            // Never promise an action result the server does not hold. The
2432            // server lease expires and another runner gets to repair it.
2433            return Ok(AgentResponse::ActionPromiseCompleted);
2434        }
2435        let completion = ActionPromiseCompletion {
2436            claim: claim.to_string(),
2437            prediction: prediction.clone(),
2438        };
2439        let _permit = self.remote_transfers.acquire().await?;
2440        if let Err(error) = remote
2441            .complete_action_promise(&prediction.invocation, &completion)
2442            .await
2443        {
2444            self.note_remote_failure();
2445            warn!(
2446                "remote cache action promise completion failed for {}: {error}",
2447                prediction.invocation.hash
2448            );
2449        }
2450        Ok(AgentResponse::ActionPromiseCompleted)
2451    }
2452
2453    /// Record that a remote operation failed and the build carried on without it.
2454    fn note_remote_failure(&self) {
2455        self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
2456    }
2457
2458    async fn get_remote_action_result(
2459        &self,
2460        remote: &RemoteCacheClient,
2461        action: &CacheDigest,
2462    ) -> Result<Option<RemoteActionResult>> {
2463        self.stats
2464            .remote_action_lookups
2465            .fetch_add(1, Ordering::Relaxed);
2466        let _timer = AtomicDurationTimer::start(&self.stats.remote_action_lookup_duration_ns);
2467        remote.get_action_result(action).await
2468    }
2469
2470    fn record_action_hit(
2471        &self,
2472        action: &CacheDigest,
2473        restore: RestoreStats,
2474        crate_name: Option<String>,
2475        diagnostic: Option<AgentEvent>,
2476    ) -> Result<AgentResponse> {
2477        validate_crate_name(crate_name.as_deref())?;
2478        if self.actions.find(action)?.is_none() {
2479            let pending = self.pending_remote_actions.lock().unwrap().remove(action);
2480            if let Some(result) = pending {
2481                self.actions.store(&result)?;
2482            } else {
2483                bail!("cannot record a hit for a missing action result");
2484            }
2485        }
2486        self.record_restore(restore);
2487        self.stats.hits.fetch_add(1, Ordering::Relaxed);
2488        self.emit_action(
2489            diagnostic,
2490            AgentEvent::ActionHit {
2491                crate_name,
2492                restore,
2493            },
2494        );
2495        Ok(AgentResponse::ActionHitRecorded)
2496    }
2497
2498    fn record_restore(&self, restore: RestoreStats) {
2499        self.record_materialization(restore);
2500        atomic_saturating_add(
2501            &self.stats.avoided_compiler_duration_ns,
2502            restore.avoided_compiler_duration_ns,
2503        );
2504        atomic_saturating_add(&self.stats.restored_output_files, restore.output_files);
2505        atomic_saturating_add(&self.stats.restored_output_bytes, restore.output_bytes);
2506        atomic_saturating_add(
2507            &self.stats.reflinked_output_files,
2508            restore.reflinked_output_files,
2509        );
2510        atomic_saturating_add(
2511            &self.stats.reflinked_output_bytes,
2512            restore.reflinked_output_bytes,
2513        );
2514        atomic_saturating_add(
2515            &self.stats.hardlinked_output_files,
2516            restore.hardlinked_output_files,
2517        );
2518        atomic_saturating_add(
2519            &self.stats.hardlinked_output_bytes,
2520            restore.hardlinked_output_bytes,
2521        );
2522        atomic_saturating_add(&self.stats.copied_output_files, restore.copied_output_files);
2523        atomic_saturating_add(&self.stats.copied_output_bytes, restore.copied_output_bytes);
2524        atomic_saturating_add(&self.stats.reused_output_files, restore.reused_output_files);
2525        atomic_saturating_add(&self.stats.reused_output_bytes, restore.reused_output_bytes);
2526    }
2527
2528    fn record_compiler_invocation(
2529        &self,
2530        outcome: &str,
2531        incremental: bool,
2532        crate_name: Option<&str>,
2533        duration_ns: u64,
2534        diagnostic: Option<AgentEvent>,
2535        bypass: Option<String>,
2536    ) -> Result<AgentResponse> {
2537        if !matches!(outcome, "miss" | "unconsulted" | "bypass" | "verification") {
2538            bail!("invalid compiler invocation outcome");
2539        }
2540        validate_crate_name(crate_name)?;
2541        if incremental {
2542            self.stats
2543                .incremental_compilations
2544                .fetch_add(1, Ordering::Relaxed);
2545        }
2546        let mut compiler = self.stats.compiler.lock().unwrap();
2547        let stats = compiler.entry(outcome.to_string()).or_default();
2548        stats.invocations = stats.invocations.saturating_add(1);
2549        stats.duration_ns = stats.duration_ns.saturating_add(duration_ns);
2550        drop(compiler);
2551        if outcome != "verification"
2552            && let Some(crate_name) = crate_name.filter(|name| !name.is_empty())
2553        {
2554            let mut slow = self.stats.slow_compilations.lock().unwrap();
2555            let duration = slow.entry(crate_name.to_string()).or_default();
2556            *duration = duration.saturating_add(duration_ns);
2557        }
2558        let action = match bypass {
2559            Some(kind) => AgentEvent::BypassedCompilation {
2560                kind,
2561                crate_name: crate_name.map(str::to_string),
2562                duration_ns,
2563            },
2564            None => AgentEvent::CompilerInvocation {
2565                outcome: outcome.to_string(),
2566                crate_name: crate_name.map(str::to_string),
2567                duration_ns,
2568            },
2569        };
2570        self.emit_action(diagnostic, action);
2571        Ok(AgentResponse::CompilerInvocationRecorded)
2572    }
2573
2574    fn record_materialization(&self, restore: RestoreStats) {
2575        atomic_saturating_add(&self.stats.materialization_duration_ns, restore.duration_ns);
2576    }
2577
2578    /// Surface a shim diagnostic on this process's stderr.
2579    ///
2580    /// The agent lives in the process that owns the session, so printing here
2581    /// reaches the terminal running the build rather than the stderr of the
2582    /// compilation the shim stands in for -- which build scripts read as part
2583    /// of the compiler's answer. Deduplicated because one cause tends to fire
2584    /// once per compilation, and capped so a message unique per compilation
2585    /// cannot scroll the build away. The level is part of the deduplication
2586    /// key, so a fatal reason is never hidden by an earlier warning with the
2587    /// same text.
2588    fn record_diagnostic(&self, level: DiagnosticLevel, message: String) -> Result<AgentResponse> {
2589        if message.is_empty()
2590            || message.len() > MAX_DIAGNOSTIC_BYTES
2591            || message.contains(['\n', '\r', '\0'])
2592        {
2593            bail!("invalid shim {}", level.label());
2594        }
2595        let mut diagnostics = self.diagnostics.lock().unwrap();
2596        let key = (level, message.clone());
2597        if !diagnostics.contains(&key) {
2598            // An acknowledgement promises the diagnostic was surfaced. A
2599            // fatal shim error must be able to fall back locally when this
2600            // channel can no longer display it.
2601            if diagnostics.len() >= MAX_DIAGNOSTICS {
2602                bail!("shim diagnostic limit reached");
2603            }
2604            eprintln!("mbx[{}]: {message}", level.label());
2605            self.emit(|| match level {
2606                DiagnosticLevel::Warning => AgentEvent::Warning {
2607                    message: message.clone(),
2608                },
2609                DiagnosticLevel::Error => AgentEvent::Error {
2610                    message: message.clone(),
2611                },
2612            });
2613            diagnostics.insert(key);
2614        }
2615        Ok(match level {
2616            DiagnosticLevel::Warning => AgentResponse::WarningRecorded,
2617            DiagnosticLevel::Error => AgentResponse::ErrorRecorded,
2618        })
2619    }
2620
2621    /// Answer file-digest lookups from the session ledger.
2622    ///
2623    /// A recorded digest is returned only when the requested identity matches
2624    /// the recorded one exactly. The caller statted the file to build the
2625    /// identity, so this is the [`VerifiedBlob`] freshness check with the
2626    /// stat moved to the side that was already making it.
2627    fn find_file_digests(
2628        &self,
2629        scope: FileDigestScope,
2630        files: Vec<FileIdentity>,
2631    ) -> Result<AgentResponse> {
2632        if files.len() > MAX_FILE_DIGEST_BATCH {
2633            bail!("too many file-digest lookups in one request");
2634        }
2635        let ledger = self.file_digests.lock().unwrap();
2636        let digests = files
2637            .into_iter()
2638            .map(|file| {
2639                let recorded = ledger.get(&(scope, file.path.clone()))?;
2640                (recorded.file == file).then(|| recorded.digest.clone())
2641            })
2642            .collect();
2643        Ok(AgentResponse::FileDigests { digests })
2644    }
2645
2646    /// Resolve ledger misses inside the agent so concurrent shims that name the
2647    /// same NFS object wait for and reuse one read instead of stampeding it.
2648    async fn resolve_file_digests(
2649        &self,
2650        scope: FileDigestScope,
2651        files: Vec<FileIdentity>,
2652    ) -> Result<AgentResponse> {
2653        if files.len() > MAX_FILE_DIGEST_BATCH {
2654            bail!("too many file-digest resolutions in one request");
2655        }
2656        for file in &files {
2657            if !file.path.is_absolute() {
2658                bail!("file-digest resolutions need absolute paths");
2659            }
2660        }
2661        let resolutions = stream::iter(
2662            files
2663                .into_iter()
2664                .map(|file| async move { self.resolve_file_digest(scope, file).await }),
2665        )
2666        .buffered(MAX_CONCURRENT_FILE_DIGESTS)
2667        .collect()
2668        .await;
2669        Ok(AgentResponse::FileDigestsResolved { resolutions })
2670    }
2671
2672    async fn resolve_file_digest(
2673        &self,
2674        scope: FileDigestScope,
2675        file: FileIdentity,
2676    ) -> FileDigestResolution {
2677        let lock = {
2678            let key = (scope, file.clone());
2679            let mut locks = self.file_digest_locks.lock().unwrap();
2680            locks.retain(|_, lock| lock.strong_count() > 0);
2681            if let Some(lock) = locks.get(&key).and_then(Weak::upgrade) {
2682                lock
2683            } else {
2684                let lock = Arc::new(FileDigestFlight {
2685                    lock: tokio::sync::Mutex::new(()),
2686                    resolution: Mutex::new(None),
2687                });
2688                locks.insert(key, Arc::downgrade(&lock));
2689                lock
2690            }
2691        };
2692        let _flight = lock.lock.lock().await;
2693        if let Some(resolution) = lock.resolution.lock().unwrap().clone() {
2694            return resolution;
2695        }
2696        if let Some(digest) = self
2697            .file_digests
2698            .lock()
2699            .unwrap()
2700            .get(&(scope, file.path.clone()))
2701            .filter(|recorded| recorded.file == file)
2702            .map(|recorded| recorded.digest.clone())
2703            .filter(|digest| digest.size == file.len && file.still_describes().unwrap_or(false))
2704        {
2705            return FileDigestResolution::Digest(digest);
2706        }
2707        let Ok(_permit) = self.file_digest_permits.acquire().await else {
2708            return FileDigestResolution::Unresolved;
2709        };
2710        #[cfg(test)]
2711        self.file_digest_reads.fetch_add(1, Ordering::Relaxed);
2712        let path = file.path.clone();
2713        let resolved =
2714            tokio::task::spawn_blocking(move || digest_file_validated(scope, &path)).await;
2715        let Ok(Ok(Some(observed))) = resolved else {
2716            return FileDigestResolution::Unresolved;
2717        };
2718        if observed.cache_identity.as_ref() != Some(&file) {
2719            return FileDigestResolution::Unresolved;
2720        }
2721        let resolution = match observed.resolution {
2722            FileDigestResolution::Digest(digest) if digest.size == file.len => {
2723                let _ = self.record_file_digests(
2724                    scope,
2725                    vec![RecordedFileDigest {
2726                        file: file.clone(),
2727                        digest: digest.clone(),
2728                    }],
2729                );
2730                FileDigestResolution::Digest(digest)
2731            }
2732            FileDigestResolution::EmbeddedTimestampMacro => {
2733                FileDigestResolution::EmbeddedTimestampMacro
2734            }
2735            FileDigestResolution::Digest(_) | FileDigestResolution::Unresolved => {
2736                FileDigestResolution::Unresolved
2737            }
2738        };
2739        *lock.resolution.lock().unwrap() = Some(resolution.clone());
2740        resolution
2741    }
2742
2743    /// Record digests of files a shim read in full, for later reuse.
2744    ///
2745    /// Capped rather than evicted: a session that outgrows the cap loses the
2746    /// shortcut for the overflow and nothing else, and no workload observed so
2747    /// far comes near it.
2748    fn record_file_digests(
2749        &self,
2750        scope: FileDigestScope,
2751        entries: Vec<RecordedFileDigest>,
2752    ) -> Result<AgentResponse> {
2753        if entries.len() > MAX_FILE_DIGEST_BATCH {
2754            bail!("too many file-digest records in one request");
2755        }
2756        for entry in &entries {
2757            validate_file_digest_record(entry)?;
2758        }
2759        let mut ledger = self.file_digests.lock().unwrap();
2760        for entry in entries {
2761            if ledger.len() >= MAX_FILE_DIGEST_ENTRIES
2762                && !ledger.contains_key(&(scope, entry.file.path.clone()))
2763            {
2764                break;
2765            }
2766            ledger.insert((scope, entry.file.path.clone()), entry);
2767        }
2768        Ok(AgentResponse::FileDigestsRecorded)
2769    }
2770
2771    /// Start the file-digest ledger from entries an earlier session left
2772    /// behind, so a file nothing has touched since is not read again by the
2773    /// first compilation of this session that names it.
2774    ///
2775    /// Each entry stands only while its recorded identity still matches the
2776    /// disk, the same rule a lookup applies to what this session recorded, so
2777    /// a stale seed costs a hash and nothing else. Entries this session has
2778    /// already recorded are kept over seeded ones. Returns how many were taken.
2779    pub fn seed_file_digests(&self, entries: Vec<(FileDigestScope, RecordedFileDigest)>) -> usize {
2780        self.seed_file_digests_with(|| entries)
2781    }
2782
2783    /// Seed the ledger from entries produced under its lock.
2784    ///
2785    /// A lookup that arrives while `load` is still reading waits for it rather
2786    /// than missing, so a caller can start the read in the background and let
2787    /// whatever else the build is doing overlap it.
2788    pub fn seed_file_digests_with(
2789        &self,
2790        load: impl FnOnce() -> Vec<(FileDigestScope, RecordedFileDigest)>,
2791    ) -> usize {
2792        let mut ledger = self.file_digests.lock().unwrap();
2793        let entries = load();
2794        let mut seeded = 0;
2795        for (scope, entry) in entries {
2796            if ledger.len() >= MAX_FILE_DIGEST_ENTRIES {
2797                break;
2798            }
2799            if validate_file_digest_record(&entry).is_err() {
2800                continue;
2801            }
2802            if let std::collections::btree_map::Entry::Vacant(slot) =
2803                ledger.entry((scope, entry.file.path.clone()))
2804            {
2805                slot.insert(entry);
2806                seeded += 1;
2807            }
2808        }
2809        seeded
2810    }
2811
2812    /// Everything the file-digest ledger holds, for a later session to start
2813    /// from.
2814    pub fn file_digests(&self) -> Vec<(FileDigestScope, RecordedFileDigest)> {
2815        self.file_digests
2816            .lock()
2817            .unwrap()
2818            .iter()
2819            .map(|((scope, _), entry)| (*scope, entry.clone()))
2820            .collect()
2821    }
2822
2823    fn find_action_prediction(
2824        &self,
2825        task: &str,
2826        invocation: &CacheDigest,
2827    ) -> Result<AgentResponse> {
2828        validate_task_identity(task)?;
2829        invocation.validate()?;
2830        let (prediction, prefetch) = {
2831            let mut tasks = self.task_actions.lock().unwrap();
2832            let Some(state) = tasks.get_mut(task) else {
2833                return Ok(AgentResponse::ActionPrediction { prediction: None });
2834            };
2835            activate_prediction_adapter(state, invocation)
2836        };
2837        if let Some(prefetch) = prefetch {
2838            self.spawn_prefetch_predictions(prefetch);
2839        }
2840        Ok(AgentResponse::ActionPrediction { prediction })
2841    }
2842
2843    fn record_action_prediction(
2844        &self,
2845        task: &str,
2846        prediction: ActionPrediction,
2847    ) -> Result<AgentResponse> {
2848        validate_task_identity(task)?;
2849        prediction.validate()?;
2850        let mut tasks = self.task_actions.lock().unwrap();
2851        let state = tasks.entry(task.to_string()).or_default();
2852        if !state
2853            .pending_predictions
2854            .contains_key(&prediction.invocation)
2855            && state.pending_predictions.len() >= MAX_TASK_ACTION_PREDICTIONS
2856        {
2857            bail!("this task run contains too many action predictions");
2858        }
2859        let current = state.predictions.get(&prediction.invocation);
2860        if current != Some(&prediction) {
2861            match state.changed_predictions.get(&prediction.invocation) {
2862                // Back to what the manifest held: no longer a change.
2863                Some(baseline) if baseline.as_ref() == Some(&prediction) => {
2864                    state.changed_predictions.remove(&prediction.invocation);
2865                }
2866                // Changed again; the baseline it left is already recorded.
2867                Some(_) => {}
2868                None => {
2869                    state
2870                        .changed_predictions
2871                        .insert(prediction.invocation.clone(), current.cloned());
2872                }
2873            }
2874        }
2875        state
2876            .predictions
2877            .insert(prediction.invocation.clone(), prediction.clone());
2878        state
2879            .pending_predictions
2880            .insert(prediction.invocation.clone(), prediction);
2881        Ok(AgentResponse::ActionPredictionRecorded)
2882    }
2883
2884    fn executable_identity_key(
2885        &self,
2886        executable: PathBuf,
2887        environment: BTreeMap<String, Option<String>>,
2888    ) -> Result<ExecutableIdentityKey> {
2889        // Restricted to the variables that actually select what an identity
2890        // probe reports: the toolchain rustup resolves, the SDK a linker
2891        // driver builds against, a `-fuse-ld` linker selection, and the
2892        // search path a driver finds its linker on. Anything else would let
2893        // one key stand for two different compilers.
2894        if !environment.keys().all(|name| {
2895            matches!(
2896                name.as_str(),
2897                "COMPILER_PATH"
2898                    | "GCC_EXEC_PREFIX"
2899                    | "LIBRARY_PATH"
2900                    | "MBX_FUSE_LD"
2901                    | "PATH"
2902                    | "RUSTUP_HOME"
2903                    | "RUSTUP_TOOLCHAIN"
2904                    | "SDKROOT"
2905                    | "MACOSX_DEPLOYMENT_TARGET"
2906                    | "LIB"
2907                    | "UCRTVersion"
2908                    | "UniversalCRTSdkDir"
2909                    | "VCToolsInstallDir"
2910                    | "VCToolsVersion"
2911                    | "WindowsSdkDir"
2912                    | "WindowsSDKVersion"
2913            )
2914        }) {
2915            bail!("executable identity contains an unsupported environment variable");
2916        }
2917        Ok(ExecutableIdentityKey {
2918            executable,
2919            environment,
2920        })
2921    }
2922
2923    fn find_executable_identity(
2924        &self,
2925        executable: PathBuf,
2926        environment: BTreeMap<String, Option<String>>,
2927    ) -> Result<AgentResponse> {
2928        let key = self.executable_identity_key(executable, environment)?;
2929        let remembered = self
2930            .executable_identities
2931            .lock()
2932            .unwrap()
2933            .get(&key)
2934            .cloned();
2935        let stdout = match remembered {
2936            Some(stdout) => Some(stdout),
2937            None => {
2938                let persisted = self.load_persisted_identity(&key);
2939                if let Some(stdout) = &persisted {
2940                    // Memoized like a fresh probe: the pins were checked once
2941                    // for this session, which is as often as a probe runs.
2942                    let _ = self.remember_executable_identity(&key, stdout.clone());
2943                }
2944                persisted
2945            }
2946        };
2947        Ok(AgentResponse::ExecutableIdentity { stdout })
2948    }
2949
2950    fn store_executable_identity(
2951        &self,
2952        executable: PathBuf,
2953        environment: BTreeMap<String, Option<String>>,
2954        stdout: Vec<u8>,
2955        pins: Vec<PinnedFile>,
2956    ) -> Result<AgentResponse> {
2957        if stdout.len() > MAX_EXECUTABLE_IDENTITY_SIZE {
2958            bail!("executable identity exceeds {MAX_EXECUTABLE_IDENTITY_SIZE} bytes");
2959        }
2960        let key = self.executable_identity_key(executable, environment)?;
2961        self.remember_executable_identity(&key, stdout.clone())?;
2962        // Best-effort, like the ledger: a session that cannot write it probes
2963        // again next time, which is what it would have done anyway. A pin
2964        // that no longer holds means a file moved under the probe, and what
2965        // it printed describes neither the old file nor the new one.
2966        if pins.is_empty() {
2967        } else if !pins.iter().all(PinnedFile::holds) {
2968            debug!("executable identity was not persisted: a pinned file changed during the probe");
2969        } else if let Err(error) = self.persist_executable_identity(&key, pins, &stdout) {
2970            debug!("executable identity was not persisted: {error:#}");
2971        }
2972        Ok(AgentResponse::ExecutableIdentity {
2973            stdout: Some(stdout),
2974        })
2975    }
2976
2977    fn remember_executable_identity(
2978        &self,
2979        key: &ExecutableIdentityKey,
2980        stdout: Vec<u8>,
2981    ) -> Result<()> {
2982        let mut identities = self.executable_identities.lock().unwrap();
2983        let is_new = !identities.contains_key(key);
2984        let previous_size = identities.get(key).map_or(0, Vec::len);
2985        if is_new && identities.len() >= MAX_EXECUTABLE_IDENTITIES {
2986            bail!("executable identity cache contains too many entries");
2987        }
2988        let retained_bytes = identities.values().map(Vec::len).sum::<usize>();
2989        if retained_bytes - previous_size + stdout.len() > MAX_EXECUTABLE_IDENTITY_BYTES {
2990            bail!("executable identity cache contains too many bytes");
2991        }
2992        identities.insert(key.clone(), stdout);
2993        Ok(())
2994    }
2995
2996    fn persisted_identity_path(&self, key: &ExecutableIdentityKey) -> Result<PathBuf> {
2997        let selector = serde_json::to_vec(&(&key.executable, &key.environment))?;
2998        let digest = CacheDigest::blake3(&selector);
2999        Ok(self.identity_dir.join(format!("{}.json", digest.hash)))
3000    }
3001
3002    /// What an earlier session recorded for `key`, if every file it was
3003    /// pinned to is still as that session saw it.
3004    fn load_persisted_identity(&self, key: &ExecutableIdentityKey) -> Option<Vec<u8>> {
3005        let path = self.persisted_identity_path(key).ok()?;
3006        let bytes = fs::read(path).ok()?;
3007        let persisted: PersistedExecutableIdentity = serde_json::from_slice(&bytes).ok()?;
3008        if persisted.version != PERSISTED_EXECUTABLE_IDENTITY_VERSION
3009            || persisted.executable != key.executable
3010            || persisted.environment != key.environment
3011            || persisted.pins.is_empty()
3012            || persisted.stdout.len() > MAX_EXECUTABLE_IDENTITY_SIZE
3013        {
3014            return None;
3015        }
3016        persisted
3017            .pins
3018            .iter()
3019            .all(PinnedFile::holds)
3020            .then(|| persisted.stdout.into_bytes())
3021    }
3022
3023    fn persist_executable_identity(
3024        &self,
3025        key: &ExecutableIdentityKey,
3026        pins: Vec<PinnedFile>,
3027        stdout: &[u8],
3028    ) -> Result<()> {
3029        // An identity that is not text is kept for this session only.
3030        let Ok(stdout) = std::str::from_utf8(stdout) else {
3031            return Ok(());
3032        };
3033        let persisted = PersistedExecutableIdentity {
3034            version: PERSISTED_EXECUTABLE_IDENTITY_VERSION,
3035            executable: key.executable.clone(),
3036            environment: key.environment.clone(),
3037            pins,
3038            stdout: stdout.to_owned(),
3039        };
3040        let path = self.persisted_identity_path(key)?;
3041        fs::create_dir_all(self.identity_dir.as_path())?;
3042        // Not synced: a torn record reads as no record, and costs one probe.
3043        let mut temporary = tempfile::NamedTempFile::new_in(self.identity_dir.as_path())?;
3044        std::io::Write::write_all(temporary.as_file_mut(), &serde_json::to_vec(&persisted)?)?;
3045        temporary.persist(path).map_err(|error| error.error)?;
3046        Ok(())
3047    }
3048
3049    /// Serve newline-delimited protocol requests on an authenticated session stream.
3050    pub async fn handle_connection<S>(&self, stream: S) -> Result<()>
3051    where
3052        S: AsyncRead + AsyncWrite + Unpin,
3053    {
3054        let (reader, mut writer) = tokio::io::split(stream);
3055        let mut reader = BufReader::new(reader);
3056        let hello = read_request(&mut reader)
3057            .await?
3058            .ok_or_else(|| eyre::eyre!("connection closed before the agent handshake"))?;
3059        let request: AgentRequest = serde_json::from_str(&hello)?;
3060        match request {
3061            AgentRequest::Hello {
3062                protocol,
3063                client_version,
3064            } if protocol == AGENT_PROTOCOL_VERSION && client_version == self.version.as_ref() => {}
3065            AgentRequest::Hello { protocol, .. } if protocol != AGENT_PROTOCOL_VERSION => {
3066                send_response(
3067                    &mut writer,
3068                    &AgentResponse::Error {
3069                        message: format!(
3070                            "unsupported agent protocol {protocol}; expected {AGENT_PROTOCOL_VERSION}"
3071                        ),
3072                    },
3073                )
3074                .await?;
3075                return Ok(());
3076            }
3077            AgentRequest::Hello { client_version, .. } => {
3078                send_response(
3079                    &mut writer,
3080                    &AgentResponse::Error {
3081                        message: format!(
3082                            "cache client {client_version} does not match agent {}",
3083                            self.version
3084                        ),
3085                    },
3086                )
3087                .await?;
3088                return Ok(());
3089            }
3090            _ => bail!("the first agent request must be hello"),
3091        }
3092        send_response(
3093            &mut writer,
3094            &AgentResponse::Hello {
3095                protocol: AGENT_PROTOCOL_VERSION,
3096                agent_version: self.version.to_string(),
3097            },
3098        )
3099        .await?;
3100
3101        // Tickets accumulate for the life of the connection because a shim
3102        // publishes a compilation's blobs and the action result naming them over
3103        // one connection, in that order.
3104        let mut connection = ConnectionUploads::default();
3105        let served = async {
3106            while let Some(line) = read_request(&mut reader).await? {
3107                let response = match serde_json::from_str(&line) {
3108                    Ok(request) => self.respond_on(request, &mut connection).await,
3109                    Err(error) => AgentResponse::Error {
3110                        message: format!("invalid agent request: {error}"),
3111                    },
3112                };
3113                send_response(&mut writer, &response).await?;
3114            }
3115            Ok(())
3116        }
3117        .await;
3118        // However the connection ended, a reason it reported still counts.
3119        self.release_bypass(&mut connection);
3120        served
3121    }
3122}
3123
3124/// Read one newline-delimited request, refusing one that grows past the cap.
3125///
3126/// Any process running as this user can open the session socket, so a request
3127/// that never terminates its line must not be able to grow the agent's memory
3128/// without bound.
3129async fn read_request<R>(reader: &mut R) -> Result<Option<String>>
3130where
3131    R: AsyncBufRead + Unpin,
3132{
3133    let mut line = Vec::new();
3134    loop {
3135        let available = reader.fill_buf().await?;
3136        if available.is_empty() {
3137            break;
3138        }
3139        let (consumed, complete) = match available.iter().position(|byte| *byte == b'\n') {
3140            Some(index) => (index, true),
3141            None => (available.len(), false),
3142        };
3143        if line.len() + consumed > MAX_REQUEST_BYTES {
3144            bail!("agent request exceeded {MAX_REQUEST_BYTES} bytes");
3145        }
3146        line.extend_from_slice(&available[..consumed]);
3147        // The newline itself is consumed but never kept.
3148        reader.consume(consumed + usize::from(complete));
3149        if complete {
3150            return Ok(Some(String::from_utf8(line)?));
3151        }
3152    }
3153    if line.is_empty() {
3154        Ok(None)
3155    } else {
3156        Ok(Some(String::from_utf8(line)?))
3157    }
3158}
3159
3160async fn send_response(
3161    writer: &mut (impl AsyncWrite + Unpin),
3162    response: &AgentResponse,
3163) -> Result<()> {
3164    let mut encoded = serde_json::to_vec(response)?;
3165    encoded.push(b'\n');
3166    writer.write_all(&encoded).await?;
3167    writer.flush().await?;
3168    Ok(())
3169}
3170
3171#[cfg(test)]
3172#[path = "agent_tests.rs"]
3173mod tests;