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