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