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