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::{info, warn};
11use std::collections::{BTreeMap, BTreeSet};
12use std::fs;
13use std::path::{Path, PathBuf};
14use std::sync::atomic::{AtomicU64, Ordering};
15use std::sync::{Arc, Mutex, Weak};
16use std::time::{Duration, Instant, SystemTime};
17
18const ACTION_PROMISE_WAIT: Duration = Duration::from_secs(60 * 60);
22use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};
23
24mod file_digest;
25mod manifest;
26mod prefetch;
27mod stats;
28mod wire;
29
30#[cfg(test)]
31pub(crate) use prefetch::select_prefetch_actions;
32
33pub use file_digest::{
34 FileDigestCache, FileDigestResolution, FileDigestScope, FileIdentity, FileObjectIdentity,
35 FileSnapshot, NoFileDigestCache, RecordedFileDigest, digest_file,
36};
37pub use manifest::{is_task_identity, task_manifest_actions};
38use manifest::{
39 merge_remote_task_manifest, merge_task_manifests, task_manifest_dir, validate_task_identity,
40 validate_task_manifest,
41};
42pub use stats::{AgentStats, CompilerStats};
43use wire::MAX_REQUEST_BYTES;
44pub use wire::{
45 AGENT_PROTOCOL_VERSION, ActionDiagnostic, AgentEvent, AgentEventObserver, AgentRequest,
46 AgentResponse, RestoreStats,
47};
48
49const MAX_EXECUTABLE_IDENTITIES: usize = 64;
50const MAX_EXECUTABLE_IDENTITY_SIZE: usize = 64 * 1024;
51const MAX_EXECUTABLE_IDENTITY_BYTES: usize = 256 * 1024;
52const TASK_ACTION_MANIFEST_VERSION: u8 = 1;
53const MAX_TASK_ACTION_PREDICTIONS: usize = 16 * 1024;
54const MAX_WARNING_BYTES: usize = 4 * 1024;
56const MAX_WARNINGS: usize = 128;
62const ACTION_DIAGNOSTIC_PREFIX: &str = "@mbx-action-diagnostic\t";
63const MAX_FILE_DIGEST_BATCH: usize = 16 * 1024;
65const MAX_FILE_DIGEST_ENTRIES: usize = 1024 * 1024;
70const MAX_CONCURRENT_FILE_DIGESTS: usize = 8;
73const MAX_REMOTE_TRANSFERS: usize = 64;
74const MAX_PREFETCH_TRANSFERS: usize = 48;
75const MAX_CONCURRENT_BLOB_PACKS: usize = 8;
81const MAX_PREFETCH_ACTIONS: usize = 1024;
92const MAX_PREFETCH_ACTION_WAVE: usize = 32;
97const MAX_PREFETCH_BATCH_LOOKUPS: usize = 1;
103const PREFETCH_ACTION_BATCH_DELAY: Duration = Duration::from_millis(5);
104const MAX_PREFETCH_DIRECTORY_OBJECTS: usize = 100_000;
105const MAX_PREFETCH_OBJECTS_PER_WAVE: usize = 100_000;
106const DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES: u64 = 5 * 1024 * 1024 * 1024;
107
108pub type TaskFallbacks = Arc<dyn Fn() -> Vec<String> + Send + Sync>;
111type FileDigestFlightKey = (FileDigestScope, FileIdentity);
112type FileDigestFlights = BTreeMap<FileDigestFlightKey, Weak<FileDigestFlight>>;
113
114struct FileDigestFlight {
115 lock: tokio::sync::Mutex<()>,
116 resolution: Mutex<Option<FileDigestResolution>>,
117}
118
119const NEWEST_MANIFEST_CANDIDATES: usize = 8;
131
132pub struct AgentRemoteCache {
134 pub client: RemoteCacheClient,
136 pub mode: RemoteCacheMode,
138 pub staging_dir: PathBuf,
140}
141
142#[derive(Default)]
143struct AtomicAgentStats {
144 lookups: AtomicU64,
145 unconsulted: AtomicU64,
146 hits: AtomicU64,
147 stores: AtomicU64,
148 stored_bytes: AtomicU64,
149 verifications: AtomicU64,
150 divergences: AtomicU64,
151 downloaded_bytes: AtomicU64,
152 uploaded_bytes: AtomicU64,
153 background_uploads: AtomicU64,
154 background_upload_failures: AtomicU64,
155 remote_blob_pack_uploads: AtomicU64,
156 remote_blob_pack_upload_blobs: AtomicU64,
157 upload_drain_duration_ns: AtomicU64,
158 prefetched_actions: AtomicU64,
159 predictions_loaded: AtomicU64,
160 remote_failures: AtomicU64,
161 remote_manifest_lookups: AtomicU64,
162 remote_manifest_lookup_duration_ns: AtomicU64,
163 remote_action_lookups: AtomicU64,
164 remote_action_lookup_duration_ns: AtomicU64,
165 remote_blob_requests: AtomicU64,
166 remote_blob_pack_requests: AtomicU64,
167 remote_blob_pack_blobs: AtomicU64,
168 remote_blob_transfer_duration_ns: AtomicU64,
169 local_cas_write_duration_ns: AtomicU64,
170 prefetch_runs: AtomicU64,
171 prefetch_duration_ns: AtomicU64,
172 materialization_duration_ns: AtomicU64,
173 bypasses: Mutex<BTreeMap<String, u64>>,
174 avoided_compiler_duration_ns: AtomicU64,
175 compiler: Mutex<BTreeMap<String, CompilerStats>>,
176 slow_compilations: Mutex<BTreeMap<String, u64>>,
177 restored_output_files: AtomicU64,
178 restored_output_bytes: AtomicU64,
179 reflinked_output_files: AtomicU64,
180 reflinked_output_bytes: AtomicU64,
181 copied_output_files: AtomicU64,
182 copied_output_bytes: AtomicU64,
183 reused_output_files: AtomicU64,
184 reused_output_bytes: AtomicU64,
185}
186
187struct AtomicDurationTimer<'a> {
188 started: Instant,
189 target: &'a AtomicU64,
190}
191
192impl<'a> AtomicDurationTimer<'a> {
193 fn start(target: &'a AtomicU64) -> Self {
194 Self {
195 started: Instant::now(),
196 target,
197 }
198 }
199}
200
201impl Drop for AtomicDurationTimer<'_> {
202 fn drop(&mut self) {
203 atomic_saturating_add(self.target, duration_ns(self.started));
204 }
205}
206
207fn duration_ns(started: Instant) -> u64 {
208 started.elapsed().as_nanos().try_into().unwrap_or(u64::MAX)
209}
210
211fn validate_crate_name(crate_name: Option<&str>) -> Result<()> {
212 if let Some(crate_name) = crate_name
213 && (crate_name.len() > 256 || crate_name.contains(['\0', '\n', '\r']))
214 {
215 bail!("invalid compiler crate name");
216 }
217 Ok(())
218}
219
220#[derive(serde::Deserialize)]
221struct ActionDiagnosticEnvelope {
222 outcome: String,
223 crate_name: Option<String>,
224 diagnostic: ActionDiagnostic,
225}
226
227fn parse_action_diagnostic(
228 message: &str,
229) -> Option<Result<(String, Option<String>, ActionDiagnostic)>> {
230 let payload = message.strip_prefix(ACTION_DIAGNOSTIC_PREFIX)?;
231 Some((|| {
232 let envelope: ActionDiagnosticEnvelope = serde_json::from_str(payload)?;
233 if !matches!(envelope.outcome.as_str(), "hit" | "miss") {
234 bail!("invalid action diagnostic outcome");
235 }
236 validate_crate_name(envelope.crate_name.as_deref())?;
237 Ok((envelope.outcome, envelope.crate_name, envelope.diagnostic))
238 })())
239}
240
241fn atomic_saturating_add(target: &AtomicU64, value: u64) {
242 let _ = target.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
243 Some(current.saturating_add(value))
244 });
245}
246
247fn queue_prefetch_digest(
248 verified: &BTreeMap<CacheDigest, PathBuf>,
249 pending: &mut BTreeMap<CacheDigest, ()>,
250 digest: CacheDigest,
251) {
252 if verified.contains_key(&digest) || pending.contains_key(&digest) {
253 return;
254 }
255 pending.insert(digest, ());
256}
257
258fn queue_prefetch_directory(
259 seen: &BTreeMap<CacheDigest, ()>,
260 pending: &mut BTreeMap<CacheDigest, ()>,
261 digest: CacheDigest,
262 limit: usize,
263) -> bool {
264 if seen.contains_key(&digest) || pending.contains_key(&digest) {
265 return true;
266 }
267 if seen.len().saturating_add(pending.len()) >= limit {
268 return false;
269 }
270 pending.insert(digest, ());
271 true
272}
273
274pub type TaskLoader = dyn for<'a> Fn(&'a CacheAgent, &'a str) -> BoxFuture<'a, ()> + Send + Sync;
276
277#[derive(Clone)]
283pub struct CacheAgent {
284 cas: LocalCas,
285 actions: LocalActionCache,
286 verified_blobs: Arc<Mutex<BTreeMap<CacheDigest, VerifiedBlob>>>,
287 version: Arc<str>,
288 write_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
289 action_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
290 stats: Arc<AtomicAgentStats>,
291 observer: Option<Arc<dyn AgentEventObserver>>,
292 observer_emission: Arc<Mutex<()>>,
293 task_loader: Option<Arc<TaskLoader>>,
296 executable_identities: Arc<Mutex<BTreeMap<ExecutableIdentityKey, Vec<u8>>>>,
297 manifest_dir: Arc<PathBuf>,
298 task_actions: Arc<Mutex<BTreeMap<String, TaskActionState>>>,
299 task_fallbacks: Arc<Mutex<BTreeMap<String, TaskFallbacks>>>,
301 next_task_run: Arc<AtomicU64>,
302 manifest_write_lock: Arc<Mutex<()>>,
303 remote: Option<Arc<RemoteCacheClient>>,
304 remote_mode: RemoteCacheMode,
305 remote_staging_dir: Arc<PathBuf>,
306 remote_download_limit: u64,
307 remote_download_bytes: Arc<AtomicU64>,
308 pending_remote_actions: Arc<Mutex<BTreeMap<CacheDigest, RemoteActionResult>>>,
309 remote_transfers: Arc<tokio::sync::Semaphore>,
310 prefetch_transfers: Arc<tokio::sync::Semaphore>,
311 prefetch_tasks: Arc<Mutex<Vec<tokio::task::JoinHandle<()>>>>,
312 warnings: Arc<Mutex<BTreeSet<String>>>,
315 file_digests: Arc<Mutex<BTreeMap<(FileDigestScope, PathBuf), RecordedFileDigest>>>,
318 file_digest_locks: Arc<Mutex<FileDigestFlights>>,
320 file_digest_permits: Arc<tokio::sync::Semaphore>,
321 #[cfg(test)]
322 file_digest_reads: Arc<AtomicU64>,
323 uploads: Option<UploadQueue>,
325}
326
327fn validate_file_digest_record(entry: &RecordedFileDigest) -> Result<()> {
330 if !entry.file.path.is_absolute() {
331 bail!("file-digest records need absolute paths");
332 }
333 entry.digest.validate()?;
334 if entry.file.len != entry.digest.size {
335 bail!("file-digest record length does not match its digest");
336 }
337 Ok(())
338}
339
340struct AgentUploadSink {
342 stats: Arc<AtomicAgentStats>,
343}
344
345impl UploadSink for AgentUploadSink {
346 fn record_blob_uploaded(&self, bytes: u64) {
347 self.stats
348 .background_uploads
349 .fetch_add(1, Ordering::Relaxed);
350 self.stats
351 .uploaded_bytes
352 .fetch_add(bytes, Ordering::Relaxed);
353 }
354
355 fn record_action_uploaded(&self) {
356 self.stats
357 .background_uploads
358 .fetch_add(1, Ordering::Relaxed);
359 }
360
361 fn record_blob_pack_uploaded(&self, blobs: u64) {
362 self.stats
363 .remote_blob_pack_uploads
364 .fetch_add(1, Ordering::Relaxed);
365 self.stats
366 .remote_blob_pack_upload_blobs
367 .fetch_add(blobs, Ordering::Relaxed);
368 }
369
370 fn record_upload_failure(&self) {
371 self.stats
372 .background_upload_failures
373 .fetch_add(1, Ordering::Relaxed);
374 self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
375 }
376}
377
378#[derive(Debug, Clone)]
384struct VerifiedBlob {
385 path: PathBuf,
386 len: u64,
387 modified: SystemTime,
388}
389
390impl VerifiedBlob {
391 fn describe(path: &Path) -> Option<Self> {
394 let metadata = std::fs::metadata(path).ok()?;
395 Some(Self {
396 path: path.to_path_buf(),
397 len: metadata.len(),
398 modified: metadata.modified().ok()?,
399 })
400 }
401
402 fn is_unchanged(&self) -> bool {
404 let Ok(metadata) = std::fs::metadata(&self.path) else {
405 return false;
406 };
407 metadata.len() == self.len && metadata.modified().is_ok_and(|now| now == self.modified)
408 }
409}
410
411#[derive(Debug, Clone, Default)]
412struct TaskActionState {
413 manifest: String,
414 baseline_loaded: bool,
415 predictions: BTreeMap<CacheDigest, ActionPrediction>,
416 pending_predictions: BTreeMap<CacheDigest, ActionPrediction>,
417 prefetched_adapters: BTreeSet<String>,
418 remote_etag: Option<String>,
419}
420
421fn activate_prediction_adapter(
423 state: &mut TaskActionState,
424 invocation: &CacheDigest,
425) -> (Option<ActionPrediction>, Option<Vec<ActionPrediction>>) {
426 let prediction = state.predictions.get(invocation).cloned();
427 let prefetch = prediction.as_ref().and_then(|prediction| {
428 if !state.prefetched_adapters.insert(prediction.adapter.clone()) {
429 return None;
430 }
431 Some(
432 state
433 .predictions
434 .values()
435 .filter(|candidate| candidate.adapter == prediction.adapter)
436 .cloned()
437 .collect(),
438 )
439 });
440 (prediction, prefetch)
441}
442
443struct PrefetchedAction {
444 adapter: String,
445 result: RemoteActionResult,
446}
447
448struct RemoteDownloadReservation {
449 counter: Arc<AtomicU64>,
450 reserved: u64,
451 committed: bool,
452}
453
454impl RemoteDownloadReservation {
455 fn bytes(&self) -> u64 {
456 self.reserved
457 }
458
459 fn commit(mut self, bytes: u64) {
460 debug_assert!(bytes <= self.reserved);
461 self.counter
462 .fetch_sub(self.reserved.saturating_sub(bytes), Ordering::AcqRel);
463 self.committed = true;
464 }
465}
466
467impl Drop for RemoteDownloadReservation {
468 fn drop(&mut self) {
469 if !self.committed {
470 self.counter.fetch_sub(self.reserved, Ordering::AcqRel);
471 }
472 }
473}
474
475#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
476struct ExecutableIdentityKey {
477 executable: PathBuf,
478 environment: BTreeMap<String, Option<String>>,
479}
480
481impl CacheAgent {
482 pub fn new(cache_dir: impl Into<PathBuf>, version: impl Into<Arc<str>>) -> Self {
484 Self::build(cache_dir.into(), version.into(), None, 0)
485 }
486
487 pub fn new_remote(
489 cache_dir: impl Into<PathBuf>,
490 version: impl Into<Arc<str>>,
491 remote: AgentRemoteCache,
492 ) -> Self {
493 Self::build(
494 cache_dir.into(),
495 version.into(),
496 Some(remote),
497 DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES,
498 )
499 }
500
501 pub fn new_remote_with_download_limit(
503 cache_dir: impl Into<PathBuf>,
504 version: impl Into<Arc<str>>,
505 remote: AgentRemoteCache,
506 max_remote_download_bytes: u64,
507 ) -> Self {
508 Self::build(
509 cache_dir.into(),
510 version.into(),
511 Some(remote),
512 max_remote_download_bytes,
513 )
514 }
515
516 fn build(
517 cache_dir: PathBuf,
518 version: Arc<str>,
519 remote: Option<AgentRemoteCache>,
520 remote_download_limit: u64,
521 ) -> Self {
522 let remote_mode = remote
523 .as_ref()
524 .map_or(RemoteCacheMode::ReadOnly, |remote| remote.mode);
525 let remote_staging_dir = remote.as_ref().map_or_else(
526 || cache_dir.join("remote"),
527 |remote| remote.staging_dir.clone(),
528 );
529 let remote = remote.map(|remote| Arc::new(remote.client));
530 let stats = Arc::new(AtomicAgentStats::default());
531 let remote_transfers = Arc::new(tokio::sync::Semaphore::new(MAX_REMOTE_TRANSFERS));
532 let uploads = remote
535 .clone()
536 .filter(|_| remote_mode.writes())
537 .map(|client| {
538 UploadQueue::new(
539 client,
540 Arc::new(AgentUploadSink {
541 stats: stats.clone(),
542 }),
543 remote_transfers.clone(),
544 )
545 });
546 Self {
547 cas: LocalCas::new(cache_dir.clone()),
548 actions: LocalActionCache::new(cache_dir.clone()),
549 verified_blobs: Arc::new(Mutex::new(BTreeMap::new())),
550 version,
551 write_locks: Arc::new(Mutex::new(BTreeMap::new())),
552 action_locks: Arc::new(Mutex::new(BTreeMap::new())),
553 stats,
554 observer: None,
555 observer_emission: Arc::new(Mutex::new(())),
556 task_loader: None,
557 executable_identities: Arc::new(Mutex::new(BTreeMap::new())),
558 manifest_dir: Arc::new(task_manifest_dir(&cache_dir)),
559 task_actions: Arc::new(Mutex::new(BTreeMap::new())),
560 task_fallbacks: Arc::new(Mutex::new(BTreeMap::new())),
561 next_task_run: Arc::new(AtomicU64::new(0)),
562 manifest_write_lock: Arc::new(Mutex::new(())),
563 remote,
564 remote_mode,
565 remote_staging_dir: Arc::new(remote_staging_dir),
566 remote_download_limit,
567 remote_download_bytes: Arc::new(AtomicU64::new(0)),
568 pending_remote_actions: Arc::new(Mutex::new(BTreeMap::new())),
569 remote_transfers,
570 prefetch_transfers: Arc::new(tokio::sync::Semaphore::new(MAX_PREFETCH_TRANSFERS)),
571 prefetch_tasks: Arc::new(Mutex::new(Vec::new())),
572 warnings: Arc::new(Mutex::new(BTreeSet::new())),
573 file_digests: Arc::new(Mutex::new(BTreeMap::new())),
574 file_digest_locks: Arc::new(Mutex::new(BTreeMap::new())),
575 file_digest_permits: Arc::new(tokio::sync::Semaphore::new(MAX_CONCURRENT_FILE_DIGESTS)),
576 #[cfg(test)]
577 file_digest_reads: Arc::new(AtomicU64::new(0)),
578 uploads,
579 }
580 }
581
582 #[must_use]
584 pub fn with_observer(mut self, observer: Arc<dyn AgentEventObserver>) -> Self {
585 self.observer = Some(observer);
586 self
587 }
588
589 pub fn with_task_loader(mut self, loader: Arc<TaskLoader>) -> Self {
596 self.task_loader = Some(loader);
597 self
598 }
599
600 async fn ensure_task_loaded(&self, task: &str) {
603 let Some(loader) = &self.task_loader else {
604 return;
605 };
606 if self.task_actions.lock().unwrap().contains_key(task) {
607 return;
608 }
609 loader(self, task).await;
610 }
611
612 fn emit(&self, event: impl FnOnce() -> AgentEvent) {
613 if let Some(observer) = &self.observer {
614 let _emission = self.observer_emission.lock().unwrap();
615 observer.event(event());
616 }
617 }
618
619 fn emit_action(&self, diagnostic: Option<AgentEvent>, action: AgentEvent) {
620 if let Some(observer) = &self.observer {
621 let _emission = self.observer_emission.lock().unwrap();
622 if let Some(diagnostic) = diagnostic {
623 observer.event(diagnostic);
624 }
625 observer.event(action);
626 }
627 }
628
629 fn reserve_remote_download(&self, bytes: u64) -> Result<RemoteDownloadReservation> {
630 let mut current = self.remote_download_bytes.load(Ordering::Acquire);
631 loop {
632 let next = current
633 .checked_add(bytes)
634 .ok_or_else(|| eyre::eyre!("remote cache download budget overflowed"))?;
635 if next > self.remote_download_limit {
636 bail!(
637 "remote cache download budget exceeded: {} bytes requested with {} of {} bytes already used",
638 bytes,
639 current,
640 self.remote_download_limit
641 );
642 }
643 match self.remote_download_bytes.compare_exchange_weak(
644 current,
645 next,
646 Ordering::AcqRel,
647 Ordering::Acquire,
648 ) {
649 Ok(_) => {
650 return Ok(RemoteDownloadReservation {
651 counter: self.remote_download_bytes.clone(),
652 reserved: bytes,
653 committed: false,
654 });
655 }
656 Err(observed) => current = observed,
657 }
658 }
659 }
660
661 fn reserve_remote_download_up_to(&self, requested: u64) -> Result<RemoteDownloadReservation> {
662 let mut current = self.remote_download_bytes.load(Ordering::Acquire);
663 loop {
664 let reserved = requested.min(self.remote_download_limit.saturating_sub(current));
665 let next = current + reserved;
666 match self.remote_download_bytes.compare_exchange_weak(
667 current,
668 next,
669 Ordering::AcqRel,
670 Ordering::Acquire,
671 ) {
672 Ok(_) => {
673 return Ok(RemoteDownloadReservation {
674 counter: self.remote_download_bytes.clone(),
675 reserved,
676 committed: false,
677 });
678 }
679 Err(observed) => current = observed,
680 }
681 }
682 }
683
684 pub async fn begin_task(&self, task: &str) -> Result<String> {
686 self.begin_task_with_remote_errors(task, false, None, true)
687 .await
688 }
689
690 pub async fn begin_task_on_prediction(&self, task: &str) -> Result<String> {
693 self.begin_task_with_remote_errors(task, false, None, false)
694 .await
695 }
696
697 pub async fn begin_session_task(&self, task: &str) -> Result<()> {
703 self.begin_task_with_remote_errors(task, false, Some(task.to_string()), false)
704 .await?;
705 Ok(())
706 }
707
708 pub fn register_task_fallbacks(
719 &self,
720 task: &str,
721 fallbacks: impl Fn() -> Vec<String> + Send + Sync + 'static,
722 ) -> Result<()> {
723 validate_task_identity(task)?;
724 self.task_fallbacks
725 .lock()
726 .unwrap()
727 .insert(task.to_string(), Arc::new(fallbacks));
728 Ok(())
729 }
730
731 pub async fn prefetch_task(&self, task: &str) -> Result<String> {
733 let run = self
734 .begin_task_with_remote_errors(task, true, None, true)
735 .await?;
736 self.wait_for_prefetches().await;
737 Ok(run)
738 }
739
740 async fn begin_task_with_remote_errors(
741 &self,
742 task: &str,
743 strict: bool,
744 run: Option<String>,
745 eager_prefetch: bool,
746 ) -> Result<String> {
747 validate_task_identity(task)?;
748 let early_manifest = {
754 let _write_guard = self.manifest_write_lock.lock().unwrap();
755 let _file_guard = self.lock_task_manifest(task)?;
756 self.load_task_manifest(task)?
757 };
758 let early_actions: BTreeSet<_> = early_manifest
759 .iter()
760 .flat_map(|manifest| manifest.predictions.iter())
761 .map(|prediction| prediction.action.clone())
762 .collect();
763 if eager_prefetch {
764 self.spawn_prefetch_predictions(
765 early_manifest
766 .as_ref()
767 .map(|manifest| manifest.predictions.clone())
768 .unwrap_or_default(),
769 );
770 }
771 let (remote_manifest, mut remote_etag) = if self.remote_mode.reads() {
772 match self.get_remote_task_manifest(task).await {
773 Ok(Some((manifest, etag))) => (Some(manifest), Some(etag)),
774 Ok(None) => (None, None),
775 Err(error) => {
776 if strict {
777 return Err(error).wrap_err_with(|| {
778 format!("remote task action manifest lookup failed for {task}")
779 });
780 }
781 self.note_remote_failure();
782 warn!("remote task action manifest lookup failed for {task}: {error}");
783 (None, None)
784 }
785 }
786 } else {
787 (None, None)
788 };
789 let manifest = if self.remote.is_some() && self.remote_mode.reads() {
793 let _write_guard = self.manifest_write_lock.lock().unwrap();
794 let _file_guard = self.lock_task_manifest(task)?;
795 let local_manifest = self.load_task_manifest(task)?;
796 let manifest = match (remote_manifest, local_manifest) {
797 (Some(remote), Some(local)) => {
798 let (manifest, merged) = merge_remote_task_manifest(task, remote, local);
799 if !merged {
800 remote_etag = None;
801 }
802 Some(manifest)
803 }
804 (Some(remote), None) => Some(remote),
805 (None, local) => local,
806 };
807 if let Some(manifest) = &manifest {
808 self.persist_task_manifest(manifest)?;
809 }
810 manifest
811 } else {
812 early_manifest
813 };
814 let (manifest, remote_etag) = match manifest {
815 Some(manifest) => (Some(manifest), remote_etag),
816 None => (self.inherit_task_manifest(task, strict).await?, None),
819 };
820 let mut state = if let Some(manifest) = manifest {
821 TaskActionState {
822 manifest: task.to_string(),
823 baseline_loaded: true,
824 predictions: manifest
825 .predictions
826 .into_iter()
827 .map(|prediction| (prediction.invocation.clone(), prediction))
828 .collect(),
829 pending_predictions: BTreeMap::new(),
830 prefetched_adapters: BTreeSet::new(),
831 remote_etag,
832 }
833 } else {
834 TaskActionState {
835 manifest: task.to_string(),
836 baseline_loaded: true,
837 remote_etag,
838 ..TaskActionState::default()
839 }
840 };
841 let run = run.unwrap_or_else(|| {
842 let sequence = self.next_task_run.fetch_add(1, Ordering::Relaxed);
843 CacheDigest::blake3(format!("{task}\0{}\0{sequence}", std::process::id()).as_bytes())
844 .hash
845 });
846 self.stats.predictions_loaded.fetch_max(
850 state.predictions.len().try_into().unwrap_or(u64::MAX),
851 Ordering::Relaxed,
852 );
853 if eager_prefetch {
854 state.prefetched_adapters.extend(
855 state
856 .predictions
857 .values()
858 .map(|prediction| prediction.adapter.clone()),
859 );
860 }
861 let predictions = state
865 .predictions
866 .values()
867 .filter(|prediction| !early_actions.contains(&prediction.action))
868 .cloned()
869 .collect();
870 self.task_actions.lock().unwrap().insert(run.clone(), state);
871 if eager_prefetch {
872 self.spawn_prefetch_predictions(predictions);
873 }
874 Ok(run)
875 }
876
877 async fn inherit_task_manifest(
885 &self,
886 task: &str,
887 strict: bool,
888 ) -> Result<Option<TaskActionManifest>> {
889 let fallbacks = self.task_fallbacks.lock().unwrap().get(task).cloned();
890 let Some(fallbacks) = fallbacks else {
891 return Ok(None);
892 };
893 let named = tokio::task::spawn_blocking(move || fallbacks()).await?;
894 let mut tried = BTreeSet::from([task.to_string()]);
895 let candidates = named.into_iter().map(|identity| (identity, true)).chain(
896 self.newest_task_identities()
897 .into_iter()
898 .map(|identity| (identity, false)),
899 );
900 for (identity, named) in candidates {
901 if !tried.insert(identity.clone()) || validate_task_identity(&identity).is_err() {
902 continue;
903 }
904 let mut found = self.load_task_manifest(&identity)?;
905 if found.is_none() && named && self.remote_mode.reads() {
906 found = match self.get_remote_task_manifest(&identity).await {
907 Ok(manifest) => manifest.map(|(manifest, _)| manifest),
908 Err(error) => {
909 if strict {
910 return Err(error).wrap_err_with(|| {
911 format!("remote task action manifest lookup failed for {identity}")
912 });
913 }
914 self.note_remote_failure();
915 warn!("remote task action manifest lookup failed for {identity}: {error}");
916 None
917 }
918 };
919 }
920 let Some(mut manifest) = found else {
921 continue;
922 };
923 if manifest.predictions.is_empty()
924 || (!named && manifest.predictions.len() > MAX_TASK_ACTION_PREDICTIONS / 2)
925 {
926 continue;
927 }
928 info!(
929 "no predictions were recorded for {task}; inheriting {} from {identity}",
930 manifest.predictions.len()
931 );
932 manifest.task = task.to_string();
933 validate_task_manifest(&manifest, task)?;
934 let _write_guard = self.manifest_write_lock.lock().unwrap();
935 let _file_guard = self.lock_task_manifest(task)?;
936 if let Some(recorded) = self.load_task_manifest(task)? {
940 return Ok(Some(recorded));
941 }
942 self.persist_task_manifest(&manifest)?;
943 return Ok(Some(manifest));
944 }
945 Ok(None)
946 }
947
948 fn newest_task_identities(&self) -> Vec<String> {
950 let Ok(entries) = fs::read_dir(self.manifest_dir.as_path()) else {
951 return Vec::new();
952 };
953 let mut manifests: Vec<(std::time::SystemTime, String)> = entries
954 .filter_map(|entry| {
955 let entry = entry.ok()?;
956 let name = entry.file_name();
957 let identity = name.to_str()?.strip_suffix(".json")?.to_string();
958 validate_task_identity(&identity).ok()?;
959 let modified = entry.metadata().ok()?.modified().ok()?;
960 Some((modified, identity))
961 })
962 .collect();
963 manifests.sort_by(|left, right| right.cmp(left));
964 manifests
965 .into_iter()
966 .take(NEWEST_MANIFEST_CANDIDATES)
967 .map(|(_, identity)| identity)
968 .collect()
969 }
970
971 pub async fn cancel_prefetches(&self) {
973 let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
974 for task in &tasks {
975 task.abort();
976 }
977 for task in tasks {
978 if let Err(error) = task.await
979 && !error.is_cancelled()
980 {
981 warn!("remote action prefetch task failed: {error}");
982 }
983 }
984 }
985
986 pub async fn wait_for_uploads(&self) {
993 let Some(uploads) = &self.uploads else {
994 return;
995 };
996 let _timer = AtomicDurationTimer::start(&self.stats.upload_drain_duration_ns);
997 uploads.drain().await;
998 }
999
1000 async fn wait_for_prefetches(&self) {
1001 let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
1002 for task in tasks {
1003 if let Err(error) = task.await {
1004 warn!("remote action prefetch task failed: {error}");
1005 }
1006 }
1007 }
1008
1009 pub async fn commit_task(&self, run: &str) -> Result<()> {
1011 self.commit_task_actions(run).await.map(|_| ())
1012 }
1013
1014 pub async fn commit_task_actions(&self, run: &str) -> Result<Vec<ActionPrediction>> {
1020 validate_task_identity(run)?;
1021 let state = {
1022 let mut runs = self.task_actions.lock().unwrap();
1023 let state = runs
1024 .get(run)
1025 .ok_or_else(|| eyre::eyre!("task action manifest baseline was not loaded"))?;
1026 if !state.baseline_loaded {
1027 bail!("task action manifest baseline was not loaded");
1028 }
1029 if state.pending_predictions.is_empty()
1035 && (self.remote.is_none() || !self.remote_mode.writes())
1036 {
1037 runs.remove(run);
1038 return Ok(Vec::new());
1039 }
1040 state.clone()
1041 };
1042 let task = state.manifest;
1043 validate_task_identity(&task)?;
1044 let completed = state
1045 .pending_predictions
1046 .values()
1047 .cloned()
1048 .collect::<Vec<_>>();
1049 let (manifest, introduced) = {
1050 let _write_guard = self.manifest_write_lock.lock().unwrap();
1051 let _file_guard = self.lock_task_manifest(&task)?;
1052 let mut predictions = self
1053 .load_task_manifest(&task)?
1054 .map(|manifest| {
1055 manifest
1056 .predictions
1057 .into_iter()
1058 .map(|prediction| (prediction.invocation.clone(), prediction))
1059 .collect::<BTreeMap<_, _>>()
1060 })
1061 .unwrap_or_default();
1062 let introduced: BTreeSet<CacheDigest> = state
1067 .pending_predictions
1068 .keys()
1069 .filter(|invocation| !predictions.contains_key(*invocation))
1070 .cloned()
1071 .collect();
1072 predictions.extend(state.pending_predictions);
1077 let manifest = TaskActionManifest {
1078 version: TASK_ACTION_MANIFEST_VERSION,
1079 task: task.clone(),
1080 predictions: predictions.into_values().collect(),
1081 };
1082 validate_task_manifest(&manifest, &task)?;
1083 self.persist_task_manifest(&manifest)?;
1084 (manifest, introduced)
1085 };
1086 self.task_actions.lock().unwrap().remove(run);
1087 if self.remote_mode.writes() {
1088 let mut manifest = manifest;
1092 if let Some(uploads) = &self.uploads {
1093 let actions: Vec<CacheDigest> = manifest
1094 .predictions
1095 .iter()
1096 .map(|prediction| prediction.action.clone())
1097 .collect();
1098 let unpublished = uploads.wait_for_actions(&actions).await;
1099 let withheld = manifest
1100 .predictions
1101 .iter()
1102 .filter(|prediction| {
1103 introduced.contains(&prediction.invocation)
1104 && unpublished.contains(&prediction.action)
1105 })
1106 .count();
1107 if withheld > 0 {
1108 warn!(
1111 "{withheld} of {} predicted actions were not published, so the remote task action manifest omits them",
1112 manifest.predictions.len()
1113 );
1114 manifest.predictions.retain(|prediction| {
1115 !(introduced.contains(&prediction.invocation)
1116 && unpublished.contains(&prediction.action))
1117 });
1118 }
1119 }
1120 match self
1121 .put_remote_task_manifest(&task, manifest, state.remote_etag)
1122 .await
1123 {
1124 Ok(remote_manifest) => {
1125 let _write_guard = self.manifest_write_lock.lock().unwrap();
1126 let reconciliation = (|| {
1127 let _file_guard = self.lock_task_manifest(&task)?;
1128 let manifest = match self.load_task_manifest(&task)? {
1129 Some(local) => {
1130 merge_remote_task_manifest(&task, remote_manifest, local).0
1131 }
1132 None => remote_manifest,
1133 };
1134 self.persist_task_manifest(&manifest)
1135 })();
1136 if let Err(error) = reconciliation {
1137 warn!(
1138 "remote task action manifest reconciliation failed for {task}: {error}"
1139 );
1140 }
1141 }
1142 Err(error) => {
1143 self.note_remote_failure();
1144 warn!("remote task action manifest upload failed for {task}: {error}");
1145 }
1146 }
1147 }
1148 Ok(completed)
1149 }
1150
1151 fn task_manifest_path(&self, task: &str) -> PathBuf {
1152 self.manifest_dir.join(format!("{task}.json"))
1153 }
1154
1155 fn task_manifest_lock_path(&self, task: &str) -> PathBuf {
1156 self.manifest_dir.join("locks").join(format!("{task}.lock"))
1157 }
1158
1159 fn lock_task_manifest(&self, task: &str) -> Result<fslock::LockFile> {
1160 let path = self.task_manifest_lock_path(task);
1161 fs::create_dir_all(path.parent().expect("task manifest lock has a parent"))?;
1162 let mut lock = fslock::LockFile::open(&path)?;
1163 lock.lock()?;
1164 Ok(lock)
1165 }
1166
1167 fn load_task_manifest(&self, task: &str) -> Result<Option<TaskActionManifest>> {
1168 match fs::read(self.task_manifest_path(task)) {
1169 Ok(contents) => Ok(Some(self.parse_task_manifest(task, &contents, false)?)),
1170 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
1171 Err(error) => Err(error.into()),
1172 }
1173 }
1174
1175 fn parse_task_manifest(
1176 &self,
1177 task: &str,
1178 contents: &[u8],
1179 require_canonical: bool,
1180 ) -> Result<TaskActionManifest> {
1181 let manifest: TaskActionManifest = serde_json::from_slice(contents)?;
1182 validate_task_manifest(&manifest, task)?;
1183 if require_canonical && canonical_json(&manifest)? != contents {
1184 bail!("task action manifest is not canonical JSON");
1185 }
1186 Ok(manifest)
1187 }
1188
1189 fn task_manifest_selector(task: &str) -> Result<(Vec<u8>, CacheDigest)> {
1190 TaskActionManifest::selector(task)
1191 }
1192
1193 fn persist_task_manifest(&self, manifest: &TaskActionManifest) -> Result<()> {
1194 let bytes = canonical_json(manifest)?;
1195 fs::create_dir_all(self.manifest_dir.as_path())?;
1196 let mut temporary = tempfile::NamedTempFile::new_in(self.manifest_dir.as_path())?;
1197 std::io::Write::write_all(temporary.as_file_mut(), &bytes)?;
1198 temporary.as_file_mut().sync_all()?;
1199 temporary
1200 .persist(self.task_manifest_path(&manifest.task))
1201 .map_err(|error| error.error)?;
1202 Ok(())
1203 }
1204
1205 async fn get_remote_task_manifest(
1206 &self,
1207 task: &str,
1208 ) -> Result<Option<(TaskActionManifest, String)>> {
1209 let Some(remote) = &self.remote else {
1210 return Ok(None);
1211 };
1212 let (_, selector) = Self::task_manifest_selector(task)?;
1213 let _permit = self.remote_transfers.acquire().await?;
1214 self.stats
1215 .remote_manifest_lookups
1216 .fetch_add(1, Ordering::Relaxed);
1217 let _timer = AtomicDurationTimer::start(&self.stats.remote_manifest_lookup_duration_ns);
1218 let Some(remote_manifest) = remote.get_action_manifest(&selector).await? else {
1219 return Ok(None);
1220 };
1221 let manifest = self.parse_task_manifest(task, &remote_manifest.bytes, true)?;
1222 Ok(Some((manifest, remote_manifest.etag)))
1223 }
1224
1225 async fn put_remote_task_manifest(
1226 &self,
1227 task: &str,
1228 mut manifest: TaskActionManifest,
1229 mut expected_etag: Option<String>,
1230 ) -> Result<TaskActionManifest> {
1231 let Some(remote) = &self.remote else {
1232 return Ok(manifest);
1233 };
1234 let (_, selector) = Self::task_manifest_selector(task)?;
1235 for _ in 0..4 {
1236 let bytes = canonical_json(&manifest)?;
1237 let outcome = {
1238 let _permit = self.remote_transfers.acquire().await?;
1239 remote
1240 .put_action_manifest(&selector, &bytes, expected_etag.as_deref())
1241 .await?
1242 };
1243 match outcome {
1244 ManifestPutOutcome::Stored => return Ok(manifest),
1245 ManifestPutOutcome::PreconditionFailed => {
1246 let Some((remote_manifest, etag)) = self.get_remote_task_manifest(task).await?
1247 else {
1248 expected_etag = None;
1249 continue;
1250 };
1251 manifest = merge_task_manifests(task, Some(remote_manifest), manifest)?;
1252 expected_etag = Some(etag);
1253 }
1254 }
1255 }
1256 bail!("remote task action manifest changed too frequently")
1257 }
1258
1259 pub fn stats(&self) -> AgentStats {
1261 AgentStats {
1262 session_duration_ns: 0,
1263 lookups: self.stats.lookups.load(Ordering::Relaxed),
1264 unconsulted: self.stats.unconsulted.load(Ordering::Relaxed),
1265 hits: self.stats.hits.load(Ordering::Relaxed),
1266 stores: self.stats.stores.load(Ordering::Relaxed),
1267 stored_bytes: self.stats.stored_bytes.load(Ordering::Relaxed),
1268 verifications: self.stats.verifications.load(Ordering::Relaxed),
1269 divergences: self.stats.divergences.load(Ordering::Relaxed),
1270 downloaded_bytes: self.stats.downloaded_bytes.load(Ordering::Relaxed),
1271 uploaded_bytes: self.stats.uploaded_bytes.load(Ordering::Relaxed),
1272 background_uploads: self.stats.background_uploads.load(Ordering::Relaxed),
1273 background_upload_failures: self
1274 .stats
1275 .background_upload_failures
1276 .load(Ordering::Relaxed),
1277 remote_blob_pack_uploads: self.stats.remote_blob_pack_uploads.load(Ordering::Relaxed),
1278 remote_blob_pack_upload_blobs: self
1279 .stats
1280 .remote_blob_pack_upload_blobs
1281 .load(Ordering::Relaxed),
1282 upload_drain_duration_ns: self.stats.upload_drain_duration_ns.load(Ordering::Relaxed),
1283 prefetched_actions: self.stats.prefetched_actions.load(Ordering::Relaxed),
1284 predictions_loaded: self.stats.predictions_loaded.load(Ordering::Relaxed),
1285 bypasses: self.stats.bypasses.lock().unwrap().clone(),
1286 avoided_compiler_duration_ns: self
1287 .stats
1288 .avoided_compiler_duration_ns
1289 .load(Ordering::Relaxed),
1290 compiler: self.stats.compiler.lock().unwrap().clone(),
1291 slow_compilations: self.stats.slow_compilations.lock().unwrap().clone(),
1292 remote_failures: self.stats.remote_failures.load(Ordering::Relaxed),
1293 remote_manifest_lookups: self.stats.remote_manifest_lookups.load(Ordering::Relaxed),
1294 remote_manifest_lookup_duration_ns: self
1295 .stats
1296 .remote_manifest_lookup_duration_ns
1297 .load(Ordering::Relaxed),
1298 remote_action_lookups: self.stats.remote_action_lookups.load(Ordering::Relaxed),
1299 remote_action_lookup_duration_ns: self
1300 .stats
1301 .remote_action_lookup_duration_ns
1302 .load(Ordering::Relaxed),
1303 remote_blob_requests: self.stats.remote_blob_requests.load(Ordering::Relaxed),
1304 remote_blob_pack_requests: self.stats.remote_blob_pack_requests.load(Ordering::Relaxed),
1305 remote_blob_pack_blobs: self.stats.remote_blob_pack_blobs.load(Ordering::Relaxed),
1306 remote_blob_transfer_duration_ns: self
1307 .stats
1308 .remote_blob_transfer_duration_ns
1309 .load(Ordering::Relaxed),
1310 local_cas_write_duration_ns: self
1311 .stats
1312 .local_cas_write_duration_ns
1313 .load(Ordering::Relaxed),
1314 prefetch_runs: self.stats.prefetch_runs.load(Ordering::Relaxed),
1315 prefetch_duration_ns: self.stats.prefetch_duration_ns.load(Ordering::Relaxed),
1316 materialization_duration_ns: self
1317 .stats
1318 .materialization_duration_ns
1319 .load(Ordering::Relaxed),
1320 restored_output_files: self.stats.restored_output_files.load(Ordering::Relaxed),
1321 restored_output_bytes: self.stats.restored_output_bytes.load(Ordering::Relaxed),
1322 reflinked_output_files: self.stats.reflinked_output_files.load(Ordering::Relaxed),
1323 reflinked_output_bytes: self.stats.reflinked_output_bytes.load(Ordering::Relaxed),
1324 copied_output_files: self.stats.copied_output_files.load(Ordering::Relaxed),
1325 copied_output_bytes: self.stats.copied_output_bytes.load(Ordering::Relaxed),
1326 reused_output_files: self.stats.reused_output_files.load(Ordering::Relaxed),
1327 reused_output_bytes: self.stats.reused_output_bytes.load(Ordering::Relaxed),
1328 }
1329 }
1330
1331 pub async fn handle_requests(
1338 &self,
1339 requests: impl IntoIterator<Item = AgentRequest>,
1340 ) -> Vec<AgentResponse> {
1341 let mut connection = ConnectionUploads::default();
1342 let mut responses = Vec::new();
1343 for request in requests {
1344 responses.push(self.respond_on(request, &mut connection).await);
1345 }
1346 responses
1347 }
1348
1349 fn write_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1350 Self::digest_lock(&self.write_locks, digest)
1351 }
1352
1353 fn action_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1354 Self::digest_lock(&self.action_locks, digest)
1355 }
1356
1357 fn digest_lock(
1358 locks: &Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>,
1359 digest: &CacheDigest,
1360 ) -> Arc<tokio::sync::Mutex<()>> {
1361 let mut locks = locks.lock().unwrap();
1362 locks.retain(|_, lock| lock.strong_count() > 0);
1363 if let Some(lock) = locks.get(digest).and_then(Weak::upgrade) {
1364 return lock;
1365 }
1366 let lock = Arc::new(tokio::sync::Mutex::new(()));
1367 locks.insert(digest.clone(), Arc::downgrade(&lock));
1368 lock
1369 }
1370
1371 #[cfg(test)]
1376 async fn respond(&self, request: AgentRequest) -> AgentResponse {
1377 self.respond_on(request, &mut ConnectionUploads::default())
1378 .await
1379 }
1380
1381 async fn respond_on(
1382 &self,
1383 request: AgentRequest,
1384 connection: &mut ConnectionUploads,
1385 ) -> AgentResponse {
1386 match &request {
1387 AgentRequest::FindActionPrediction { task, .. }
1388 | AgentRequest::RecordActionPrediction { task, .. } => {
1389 self.ensure_task_loaded(task).await;
1390 }
1391 _ => {}
1392 }
1393 let result = match request {
1394 AgentRequest::BeginTask { task } => self
1395 .begin_task(&task)
1396 .await
1397 .map(|run| AgentResponse::TaskBegun { run }),
1398 AgentRequest::CommitTask { run } => self
1399 .commit_task(&run)
1400 .await
1401 .map(|()| AgentResponse::TaskCommitted),
1402 AgentRequest::FindBlob { digest } => self.find_blob(&digest).await,
1403 AgentRequest::FindBlobs { digests } => self.find_blobs(digests).await,
1404 AgentRequest::StoreBlob { digest, source } => {
1405 self.store_blob(&digest, &source, connection).await
1406 }
1407 AgentRequest::FindActionResult { action } => {
1408 self.stats.lookups.fetch_add(1, Ordering::Relaxed);
1409 self.find_action_result(&action).await
1410 }
1411 AgentRequest::RecordActionHit {
1412 action,
1413 restore,
1414 crate_name,
1415 } => {
1416 let diagnostic = connection
1417 .take_action_diagnostic("hit", crate_name.as_deref())
1418 .map(|diagnostic| AgentEvent::ActionDiagnostic {
1419 outcome: "hit".into(),
1420 crate_name: crate_name.clone(),
1421 diagnostic,
1422 });
1423 self.record_action_hit(&action, restore, crate_name, diagnostic)
1424 }
1425 AgentRequest::RecordBypass { kind } => {
1426 *self
1427 .stats
1428 .bypasses
1429 .lock()
1430 .unwrap()
1431 .entry(kind.clone())
1432 .or_insert(0) += 1;
1433 self.emit(|| AgentEvent::Bypass { kind });
1434 Ok(AgentResponse::BypassRecorded)
1435 }
1436 AgentRequest::RecordUnconsulted => {
1437 self.stats.unconsulted.fetch_add(1, Ordering::Relaxed);
1438 self.emit(|| AgentEvent::Unconsulted);
1439 Ok(AgentResponse::UnconsultedRecorded)
1440 }
1441 AgentRequest::RecordWarning { message } => match parse_action_diagnostic(&message) {
1442 Some(Ok((outcome, crate_name, diagnostic))) => {
1443 connection.record_action_diagnostic(outcome, crate_name, diagnostic);
1444 Ok(AgentResponse::WarningRecorded)
1445 }
1446 Some(Err(error)) => Err(error),
1447 None => self.record_warning(message),
1448 },
1449 AgentRequest::FindFileDigests { scope, files } => self.find_file_digests(scope, files),
1450 AgentRequest::JoinActionPromise {
1451 adapter,
1452 invocation,
1453 } => self.join_action_promise(&adapter, &invocation).await,
1454 AgentRequest::CompleteActionPromise { claim, prediction } => {
1455 self.complete_action_promise(&claim, &prediction).await
1456 }
1457 AgentRequest::ResolveFileDigests { scope, files } => {
1458 self.resolve_file_digests(scope, files).await
1459 }
1460 AgentRequest::RecordFileDigests { scope, entries } => {
1461 self.record_file_digests(scope, entries)
1462 }
1463 AgentRequest::RecordCompilerInvocation {
1464 outcome,
1465 crate_name,
1466 duration_ns,
1467 } => {
1468 let diagnostic = connection
1469 .take_action_diagnostic(&outcome, crate_name.as_deref())
1470 .map(|diagnostic| AgentEvent::ActionDiagnostic {
1471 outcome: outcome.clone(),
1472 crate_name: crate_name.clone(),
1473 diagnostic,
1474 });
1475 self.record_compiler_invocation(
1476 &outcome,
1477 crate_name.as_deref(),
1478 duration_ns,
1479 diagnostic,
1480 )
1481 }
1482 AgentRequest::RecordActionVerification { matched, restore } => {
1483 self.record_materialization(restore);
1484 self.stats.verifications.fetch_add(1, Ordering::Relaxed);
1485 if !matched {
1486 self.stats.divergences.fetch_add(1, Ordering::Relaxed);
1487 }
1488 self.emit(|| AgentEvent::Verification { matched, restore });
1489 Ok(AgentResponse::ActionVerificationRecorded)
1490 }
1491 AgentRequest::StoreActionResult { result } => {
1492 self.store_action_result(&result, connection).await
1493 }
1494 AgentRequest::FindActionPrediction { task, invocation } => {
1495 self.find_action_prediction(&task, &invocation)
1496 }
1497 AgentRequest::RecordActionPrediction { task, prediction } => {
1498 self.record_action_prediction(&task, prediction)
1499 }
1500 AgentRequest::FindExecutableIdentity {
1501 executable,
1502 environment,
1503 } => self.find_executable_identity(executable, environment),
1504 AgentRequest::StoreExecutableIdentity {
1505 executable,
1506 environment,
1507 stdout,
1508 } => self.store_executable_identity(executable, environment, stdout),
1509 AgentRequest::Hello { .. } => {
1510 Err(eyre::eyre!("hello is only valid as the first request"))
1511 }
1512 };
1513 result.unwrap_or_else(|error| AgentResponse::Error {
1514 message: error.to_string(),
1515 })
1516 }
1517
1518 async fn find_blob(&self, digest: &CacheDigest) -> Result<AgentResponse> {
1519 if let Some(path) = self.find_verified_blob(digest)? {
1520 return Ok(AgentResponse::Blob { path: Some(path) });
1521 }
1522 if !self.remote_mode.reads() {
1523 return Ok(AgentResponse::Blob { path: None });
1524 }
1525 let Some(remote) = &self.remote else {
1526 return Ok(AgentResponse::Blob { path: None });
1527 };
1528 match self.fetch_remote_blob(remote, digest).await {
1529 Ok(path) => Ok(AgentResponse::Blob { path: Some(path) }),
1530 Err(error) => {
1531 warn!(
1532 "remote cache blob lookup failed for {}: {error}",
1533 digest.hash
1534 );
1535 Ok(AgentResponse::Blob { path: None })
1536 }
1537 }
1538 }
1539
1540 async fn find_blobs(&self, digests: Vec<CacheDigest>) -> Result<AgentResponse> {
1541 let mut paths = BTreeMap::new();
1542 let mut missing = Vec::new();
1543 for digest in &digests {
1544 match self.find_verified_blob(digest)? {
1545 Some(path) => {
1546 paths.insert(digest.clone(), path);
1547 }
1548 None => {
1549 missing.push(digest.clone());
1550 }
1551 }
1552 }
1553
1554 if !missing.is_empty()
1555 && self.remote_mode.reads()
1556 && let Some(remote) = &self.remote
1557 {
1558 paths.extend(self.fetch_remote_blobs(remote, missing, None).await);
1559 }
1560
1561 Ok(AgentResponse::Blobs {
1562 paths: digests
1563 .into_iter()
1564 .map(|digest| paths.get(&digest).cloned())
1565 .collect(),
1566 })
1567 }
1568
1569 async fn store_blob(
1570 &self,
1571 digest: &CacheDigest,
1572 source: &Path,
1573 connection: &mut ConnectionUploads,
1574 ) -> Result<AgentResponse> {
1575 let path = {
1576 let lock = self.write_lock(digest);
1577 let _guard = lock.lock().await;
1578 if let Some(path) = self.find_verified_blob(digest)? {
1579 path
1580 } else {
1581 let path = self.cas.store_file(digest, source)?;
1582 self.remember_verified_blob(digest, &path);
1583 self.stats.stores.fetch_add(1, Ordering::Relaxed);
1584 self.stats
1585 .stored_bytes
1586 .fetch_add(digest.size, Ordering::Relaxed);
1587 path
1588 }
1589 };
1590 if let Some(uploads) = &self.uploads {
1593 uploads.queue_blob(digest, path.clone(), connection);
1594 }
1595 Ok(AgentResponse::Stored { path })
1596 }
1597
1598 fn find_verified_blob(&self, digest: &CacheDigest) -> Result<Option<PathBuf>> {
1599 let remembered = self.verified_blobs.lock().unwrap().get(digest).cloned();
1600 if let Some(remembered) = remembered {
1601 if remembered.is_unchanged() {
1612 return Ok(Some(remembered.path));
1613 }
1614 self.verified_blobs.lock().unwrap().remove(digest);
1615 }
1616 let path = self.cas.find(digest)?;
1617 if let Some(path) = &path {
1618 self.remember_verified_blob(digest, path);
1619 }
1620 Ok(path)
1621 }
1622
1623 fn remember_verified_blob(&self, digest: &CacheDigest, path: &Path) {
1624 let Some(verified) = VerifiedBlob::describe(path) else {
1625 return;
1629 };
1630 self.verified_blobs
1631 .lock()
1632 .unwrap()
1633 .insert(digest.clone(), verified);
1634 }
1635
1636 async fn find_action_result(&self, action: &CacheDigest) -> Result<AgentResponse> {
1637 if let Some(result) = self.actions.find(action)? {
1638 return Ok(AgentResponse::ActionResult {
1639 result: Some(result),
1640 });
1641 }
1642 if !self.remote_mode.reads() {
1643 return Ok(AgentResponse::ActionResult { result: None });
1644 }
1645 let Some(remote) = &self.remote else {
1646 return Ok(AgentResponse::ActionResult { result: None });
1647 };
1648 let lock = self.action_lock(action);
1649 let _guard = lock.lock().await;
1650 if let Some(result) = self.actions.find(action)? {
1651 return Ok(AgentResponse::ActionResult {
1652 result: Some(result),
1653 });
1654 }
1655 if let Some(result) = self
1656 .pending_remote_actions
1657 .lock()
1658 .unwrap()
1659 .get(action)
1660 .cloned()
1661 {
1662 return Ok(AgentResponse::ActionResult {
1663 result: Some(result),
1664 });
1665 }
1666 let _permit = self.remote_transfers.acquire().await?;
1667 match self.get_remote_action_result(remote, action).await {
1668 Ok(Some(result)) => {
1669 self.pending_remote_actions
1670 .lock()
1671 .unwrap()
1672 .insert(action.clone(), result.clone());
1673 Ok(AgentResponse::ActionResult {
1674 result: Some(result),
1675 })
1676 }
1677 Ok(None) => Ok(AgentResponse::ActionResult { result: None }),
1678 Err(error) => {
1679 self.note_remote_failure();
1680 warn!(
1681 "remote cache action lookup failed for {}: {error}",
1682 action.hash
1683 );
1684 Ok(AgentResponse::ActionResult { result: None })
1685 }
1686 }
1687 }
1688
1689 async fn store_action_result(
1690 &self,
1691 result: &RemoteActionResult,
1692 connection: &ConnectionUploads,
1693 ) -> Result<AgentResponse> {
1694 let path = self.actions.store(result)?;
1695 if let Some(uploads) = &self.uploads {
1696 uploads.queue_action_result(result, connection);
1697 }
1698 Ok(AgentResponse::ActionStored { path })
1699 }
1700
1701 async fn join_action_promise(
1702 &self,
1703 adapter: &str,
1704 invocation: &CacheDigest,
1705 ) -> Result<AgentResponse> {
1706 if !self.remote_mode.reads() || !self.remote_mode.writes() {
1707 return Ok(AgentResponse::ActionPromise {
1708 claim: None,
1709 prediction: None,
1710 });
1711 }
1712 let Some(remote) = &self.remote else {
1713 return Ok(AgentResponse::ActionPromise {
1714 claim: None,
1715 prediction: None,
1716 });
1717 };
1718 let deadline = Instant::now() + ACTION_PROMISE_WAIT;
1719 loop {
1720 let _permit = self.remote_transfers.acquire().await?;
1721 let state = remote.join_action_promise(invocation, adapter).await;
1722 drop(_permit);
1723 match state {
1724 Ok(Some(ActionPromiseState::Claimed { claim })) => {
1725 return Ok(AgentResponse::ActionPromise {
1726 claim: Some(claim),
1727 prediction: None,
1728 });
1729 }
1730 Ok(Some(ActionPromiseState::Complete { prediction })) => {
1731 return Ok(AgentResponse::ActionPromise {
1732 claim: None,
1733 prediction: Some(prediction),
1734 });
1735 }
1736 Ok(Some(ActionPromiseState::Pending { retry_after_ms }))
1737 if Instant::now() < deadline =>
1738 {
1739 tokio::time::sleep(Duration::from_millis(retry_after_ms.clamp(10, 5_000)))
1740 .await;
1741 }
1742 Ok(Some(ActionPromiseState::Pending { .. }) | None) => {
1743 return Ok(AgentResponse::ActionPromise {
1744 claim: None,
1745 prediction: None,
1746 });
1747 }
1748 Ok(Some(_)) => {
1749 return Ok(AgentResponse::ActionPromise {
1750 claim: None,
1751 prediction: None,
1752 });
1753 }
1754 Err(error) => {
1755 self.note_remote_failure();
1756 warn!(
1757 "remote cache action promise failed for {}: {error}",
1758 invocation.hash
1759 );
1760 return Ok(AgentResponse::ActionPromise {
1761 claim: None,
1762 prediction: None,
1763 });
1764 }
1765 }
1766 }
1767 }
1768
1769 async fn complete_action_promise(
1770 &self,
1771 claim: &str,
1772 prediction: &ActionPrediction,
1773 ) -> Result<AgentResponse> {
1774 prediction.validate()?;
1775 let Some(remote) = &self.remote else {
1776 return Ok(AgentResponse::ActionPromiseCompleted);
1777 };
1778 let Some(uploads) = &self.uploads else {
1779 return Ok(AgentResponse::ActionPromiseCompleted);
1780 };
1781 if uploads
1782 .wait_for_actions(std::slice::from_ref(&prediction.action))
1783 .await
1784 .contains(&prediction.action)
1785 {
1786 return Ok(AgentResponse::ActionPromiseCompleted);
1789 }
1790 let completion = ActionPromiseCompletion {
1791 claim: claim.to_string(),
1792 prediction: prediction.clone(),
1793 };
1794 let _permit = self.remote_transfers.acquire().await?;
1795 if let Err(error) = remote
1796 .complete_action_promise(&prediction.invocation, &completion)
1797 .await
1798 {
1799 self.note_remote_failure();
1800 warn!(
1801 "remote cache action promise completion failed for {}: {error}",
1802 prediction.invocation.hash
1803 );
1804 }
1805 Ok(AgentResponse::ActionPromiseCompleted)
1806 }
1807
1808 fn note_remote_failure(&self) {
1810 self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
1811 }
1812
1813 async fn get_remote_action_result(
1814 &self,
1815 remote: &RemoteCacheClient,
1816 action: &CacheDigest,
1817 ) -> Result<Option<RemoteActionResult>> {
1818 self.stats
1819 .remote_action_lookups
1820 .fetch_add(1, Ordering::Relaxed);
1821 let _timer = AtomicDurationTimer::start(&self.stats.remote_action_lookup_duration_ns);
1822 remote.get_action_result(action).await
1823 }
1824
1825 fn record_action_hit(
1826 &self,
1827 action: &CacheDigest,
1828 restore: RestoreStats,
1829 crate_name: Option<String>,
1830 diagnostic: Option<AgentEvent>,
1831 ) -> Result<AgentResponse> {
1832 validate_crate_name(crate_name.as_deref())?;
1833 if self.actions.find(action)?.is_none() {
1834 let pending = self.pending_remote_actions.lock().unwrap().remove(action);
1835 if let Some(result) = pending {
1836 self.actions.store(&result)?;
1837 } else {
1838 bail!("cannot record a hit for a missing action result");
1839 }
1840 }
1841 self.record_restore(restore);
1842 self.stats.hits.fetch_add(1, Ordering::Relaxed);
1843 self.emit_action(
1844 diagnostic,
1845 AgentEvent::ActionHit {
1846 crate_name,
1847 restore,
1848 },
1849 );
1850 Ok(AgentResponse::ActionHitRecorded)
1851 }
1852
1853 fn record_restore(&self, restore: RestoreStats) {
1854 self.record_materialization(restore);
1855 atomic_saturating_add(
1856 &self.stats.avoided_compiler_duration_ns,
1857 restore.avoided_compiler_duration_ns,
1858 );
1859 atomic_saturating_add(&self.stats.restored_output_files, restore.output_files);
1860 atomic_saturating_add(&self.stats.restored_output_bytes, restore.output_bytes);
1861 atomic_saturating_add(
1862 &self.stats.reflinked_output_files,
1863 restore.reflinked_output_files,
1864 );
1865 atomic_saturating_add(
1866 &self.stats.reflinked_output_bytes,
1867 restore.reflinked_output_bytes,
1868 );
1869 atomic_saturating_add(&self.stats.copied_output_files, restore.copied_output_files);
1870 atomic_saturating_add(&self.stats.copied_output_bytes, restore.copied_output_bytes);
1871 atomic_saturating_add(&self.stats.reused_output_files, restore.reused_output_files);
1872 atomic_saturating_add(&self.stats.reused_output_bytes, restore.reused_output_bytes);
1873 }
1874
1875 fn record_compiler_invocation(
1876 &self,
1877 outcome: &str,
1878 crate_name: Option<&str>,
1879 duration_ns: u64,
1880 diagnostic: Option<AgentEvent>,
1881 ) -> Result<AgentResponse> {
1882 if !matches!(
1883 outcome,
1884 "miss" | "unconsulted" | "bypass" | "verification" | "incremental"
1885 ) {
1886 bail!("invalid compiler invocation outcome");
1887 }
1888 validate_crate_name(crate_name)?;
1889 let mut compiler = self.stats.compiler.lock().unwrap();
1890 let stats = compiler.entry(outcome.to_string()).or_default();
1891 stats.invocations = stats.invocations.saturating_add(1);
1892 stats.duration_ns = stats.duration_ns.saturating_add(duration_ns);
1893 drop(compiler);
1894 if outcome != "verification"
1895 && let Some(crate_name) = crate_name.filter(|name| !name.is_empty())
1896 {
1897 let mut slow = self.stats.slow_compilations.lock().unwrap();
1898 let duration = slow.entry(crate_name.to_string()).or_default();
1899 *duration = duration.saturating_add(duration_ns);
1900 }
1901 self.emit_action(
1902 diagnostic,
1903 AgentEvent::CompilerInvocation {
1904 outcome: outcome.to_string(),
1905 crate_name: crate_name.map(str::to_string),
1906 duration_ns,
1907 },
1908 );
1909 Ok(AgentResponse::CompilerInvocationRecorded)
1910 }
1911
1912 fn record_materialization(&self, restore: RestoreStats) {
1913 atomic_saturating_add(&self.stats.materialization_duration_ns, restore.duration_ns);
1914 }
1915
1916 fn record_warning(&self, message: String) -> Result<AgentResponse> {
1925 if message.is_empty()
1926 || message.len() > MAX_WARNING_BYTES
1927 || message.contains(['\n', '\r', '\0'])
1928 {
1929 bail!("invalid shim warning");
1930 }
1931 let mut warnings = self.warnings.lock().unwrap();
1932 if !warnings.contains(&message) && warnings.len() < MAX_WARNINGS {
1933 eprintln!("mbx[warning]: {message}");
1934 self.emit(|| AgentEvent::Warning {
1935 message: message.clone(),
1936 });
1937 warnings.insert(message);
1938 }
1939 Ok(AgentResponse::WarningRecorded)
1940 }
1941
1942 fn find_file_digests(
1949 &self,
1950 scope: FileDigestScope,
1951 files: Vec<FileIdentity>,
1952 ) -> Result<AgentResponse> {
1953 if files.len() > MAX_FILE_DIGEST_BATCH {
1954 bail!("too many file-digest lookups in one request");
1955 }
1956 let ledger = self.file_digests.lock().unwrap();
1957 let digests = files
1958 .into_iter()
1959 .map(|file| {
1960 let recorded = ledger.get(&(scope, file.path.clone()))?;
1961 (recorded.file == file).then(|| recorded.digest.clone())
1962 })
1963 .collect();
1964 Ok(AgentResponse::FileDigests { digests })
1965 }
1966
1967 async fn resolve_file_digests(
1970 &self,
1971 scope: FileDigestScope,
1972 files: Vec<FileIdentity>,
1973 ) -> Result<AgentResponse> {
1974 if files.len() > MAX_FILE_DIGEST_BATCH {
1975 bail!("too many file-digest resolutions in one request");
1976 }
1977 for file in &files {
1978 if !file.path.is_absolute() {
1979 bail!("file-digest resolutions need absolute paths");
1980 }
1981 }
1982 let resolutions = stream::iter(
1983 files
1984 .into_iter()
1985 .map(|file| async move { self.resolve_file_digest(scope, file).await }),
1986 )
1987 .buffered(MAX_CONCURRENT_FILE_DIGESTS)
1988 .collect()
1989 .await;
1990 Ok(AgentResponse::FileDigestsResolved { resolutions })
1991 }
1992
1993 async fn resolve_file_digest(
1994 &self,
1995 scope: FileDigestScope,
1996 file: FileIdentity,
1997 ) -> FileDigestResolution {
1998 let lock = {
1999 let key = (scope, file.clone());
2000 let mut locks = self.file_digest_locks.lock().unwrap();
2001 locks.retain(|_, lock| lock.strong_count() > 0);
2002 if let Some(lock) = locks.get(&key).and_then(Weak::upgrade) {
2003 lock
2004 } else {
2005 let lock = Arc::new(FileDigestFlight {
2006 lock: tokio::sync::Mutex::new(()),
2007 resolution: Mutex::new(None),
2008 });
2009 locks.insert(key, Arc::downgrade(&lock));
2010 lock
2011 }
2012 };
2013 let _flight = lock.lock.lock().await;
2014 if let Some(resolution) = lock.resolution.lock().unwrap().clone() {
2015 return resolution;
2016 }
2017 if let Some(digest) = self
2018 .file_digests
2019 .lock()
2020 .unwrap()
2021 .get(&(scope, file.path.clone()))
2022 .filter(|recorded| recorded.file == file)
2023 .map(|recorded| recorded.digest.clone())
2024 {
2025 return FileDigestResolution::Digest(digest);
2026 }
2027 let Ok(_permit) = self.file_digest_permits.acquire().await else {
2028 return FileDigestResolution::Unresolved;
2029 };
2030 #[cfg(test)]
2031 self.file_digest_reads.fetch_add(1, Ordering::Relaxed);
2032 let path = file.path.clone();
2033 let resolved = tokio::task::spawn_blocking(move || {
2034 let resolution = digest_file(scope, &path)?;
2035 let current = std::fs::metadata(&path)
2036 .and_then(|metadata| FileIdentity::for_digest_cache(&path, &metadata));
2037 Ok::<_, std::io::Error>((resolution, current?))
2038 })
2039 .await;
2040 let Ok(Ok((resolution, current))) = resolved else {
2041 return FileDigestResolution::Unresolved;
2042 };
2043 let resolution = match resolution {
2044 FileDigestResolution::Digest(digest)
2045 if current.as_ref() == Some(&file) && digest.size == file.len =>
2046 {
2047 let _ = self.record_file_digests(
2048 scope,
2049 vec![RecordedFileDigest {
2050 file: file.clone(),
2051 digest: digest.clone(),
2052 }],
2053 );
2054 FileDigestResolution::Digest(digest)
2055 }
2056 FileDigestResolution::EmbeddedTimestampMacro if current.as_ref() == Some(&file) => {
2057 FileDigestResolution::EmbeddedTimestampMacro
2058 }
2059 FileDigestResolution::Digest(_)
2060 | FileDigestResolution::EmbeddedTimestampMacro
2061 | FileDigestResolution::Unresolved => FileDigestResolution::Unresolved,
2062 };
2063 *lock.resolution.lock().unwrap() = Some(resolution.clone());
2064 resolution
2065 }
2066
2067 fn record_file_digests(
2073 &self,
2074 scope: FileDigestScope,
2075 entries: Vec<RecordedFileDigest>,
2076 ) -> Result<AgentResponse> {
2077 if entries.len() > MAX_FILE_DIGEST_BATCH {
2078 bail!("too many file-digest records in one request");
2079 }
2080 for entry in &entries {
2081 validate_file_digest_record(entry)?;
2082 }
2083 let mut ledger = self.file_digests.lock().unwrap();
2084 for entry in entries {
2085 if ledger.len() >= MAX_FILE_DIGEST_ENTRIES
2086 && !ledger.contains_key(&(scope, entry.file.path.clone()))
2087 {
2088 break;
2089 }
2090 ledger.insert((scope, entry.file.path.clone()), entry);
2091 }
2092 Ok(AgentResponse::FileDigestsRecorded)
2093 }
2094
2095 pub fn seed_file_digests(&self, entries: Vec<(FileDigestScope, RecordedFileDigest)>) -> usize {
2104 self.seed_file_digests_with(|| entries)
2105 }
2106
2107 pub fn seed_file_digests_with(
2113 &self,
2114 load: impl FnOnce() -> Vec<(FileDigestScope, RecordedFileDigest)>,
2115 ) -> usize {
2116 let mut ledger = self.file_digests.lock().unwrap();
2117 let entries = load();
2118 let mut seeded = 0;
2119 for (scope, entry) in entries {
2120 if ledger.len() >= MAX_FILE_DIGEST_ENTRIES {
2121 break;
2122 }
2123 if validate_file_digest_record(&entry).is_err() {
2124 continue;
2125 }
2126 if let std::collections::btree_map::Entry::Vacant(slot) =
2127 ledger.entry((scope, entry.file.path.clone()))
2128 {
2129 slot.insert(entry);
2130 seeded += 1;
2131 }
2132 }
2133 seeded
2134 }
2135
2136 pub fn file_digests(&self) -> Vec<(FileDigestScope, RecordedFileDigest)> {
2139 self.file_digests
2140 .lock()
2141 .unwrap()
2142 .iter()
2143 .map(|((scope, _), entry)| (*scope, entry.clone()))
2144 .collect()
2145 }
2146
2147 fn find_action_prediction(
2148 &self,
2149 task: &str,
2150 invocation: &CacheDigest,
2151 ) -> Result<AgentResponse> {
2152 validate_task_identity(task)?;
2153 invocation.validate()?;
2154 let (prediction, prefetch) = {
2155 let mut tasks = self.task_actions.lock().unwrap();
2156 let Some(state) = tasks.get_mut(task) else {
2157 return Ok(AgentResponse::ActionPrediction { prediction: None });
2158 };
2159 activate_prediction_adapter(state, invocation)
2160 };
2161 if let Some(prefetch) = prefetch {
2162 self.spawn_prefetch_predictions(prefetch);
2163 }
2164 Ok(AgentResponse::ActionPrediction { prediction })
2165 }
2166
2167 fn record_action_prediction(
2168 &self,
2169 task: &str,
2170 prediction: ActionPrediction,
2171 ) -> Result<AgentResponse> {
2172 validate_task_identity(task)?;
2173 prediction.validate()?;
2174 let mut tasks = self.task_actions.lock().unwrap();
2175 let state = tasks.entry(task.to_string()).or_default();
2176 if !state.predictions.contains_key(&prediction.invocation)
2177 && state.predictions.len() >= MAX_TASK_ACTION_PREDICTIONS
2178 {
2179 bail!("task action manifest contains too many predictions");
2180 }
2181 state
2182 .predictions
2183 .insert(prediction.invocation.clone(), prediction.clone());
2184 state
2185 .pending_predictions
2186 .insert(prediction.invocation.clone(), prediction);
2187 Ok(AgentResponse::ActionPredictionRecorded)
2188 }
2189
2190 fn executable_identity_key(
2191 &self,
2192 executable: PathBuf,
2193 environment: BTreeMap<String, Option<String>>,
2194 ) -> Result<ExecutableIdentityKey> {
2195 if !environment.keys().all(|name| {
2200 matches!(
2201 name.as_str(),
2202 "MBX_FUSE_LD"
2203 | "RUSTUP_HOME"
2204 | "RUSTUP_TOOLCHAIN"
2205 | "SDKROOT"
2206 | "MACOSX_DEPLOYMENT_TARGET"
2207 | "LIB"
2208 | "UCRTVersion"
2209 | "UniversalCRTSdkDir"
2210 | "VCToolsInstallDir"
2211 | "VCToolsVersion"
2212 | "WindowsSdkDir"
2213 | "WindowsSDKVersion"
2214 )
2215 }) {
2216 bail!("executable identity contains an unsupported environment variable");
2217 }
2218 Ok(ExecutableIdentityKey {
2219 executable,
2220 environment,
2221 })
2222 }
2223
2224 fn find_executable_identity(
2225 &self,
2226 executable: PathBuf,
2227 environment: BTreeMap<String, Option<String>>,
2228 ) -> Result<AgentResponse> {
2229 let key = self.executable_identity_key(executable, environment)?;
2230 let stdout = self
2231 .executable_identities
2232 .lock()
2233 .unwrap()
2234 .get(&key)
2235 .cloned();
2236 Ok(AgentResponse::ExecutableIdentity { stdout })
2237 }
2238
2239 fn store_executable_identity(
2240 &self,
2241 executable: PathBuf,
2242 environment: BTreeMap<String, Option<String>>,
2243 stdout: Vec<u8>,
2244 ) -> Result<AgentResponse> {
2245 if stdout.len() > MAX_EXECUTABLE_IDENTITY_SIZE {
2246 bail!("executable identity exceeds {MAX_EXECUTABLE_IDENTITY_SIZE} bytes");
2247 }
2248 let key = self.executable_identity_key(executable, environment)?;
2249 let mut identities = self.executable_identities.lock().unwrap();
2250 let is_new = !identities.contains_key(&key);
2251 let previous_size = identities.get(&key).map_or(0, Vec::len);
2252 if is_new && identities.len() >= MAX_EXECUTABLE_IDENTITIES {
2253 bail!("executable identity cache contains too many entries");
2254 }
2255 let retained_bytes = identities.values().map(Vec::len).sum::<usize>();
2256 if retained_bytes - previous_size + stdout.len() > MAX_EXECUTABLE_IDENTITY_BYTES {
2257 bail!("executable identity cache contains too many bytes");
2258 }
2259 identities.insert(key, stdout.clone());
2260 Ok(AgentResponse::ExecutableIdentity {
2261 stdout: Some(stdout),
2262 })
2263 }
2264
2265 pub async fn handle_connection<S>(&self, stream: S) -> Result<()>
2267 where
2268 S: AsyncRead + AsyncWrite + Unpin,
2269 {
2270 let (reader, mut writer) = tokio::io::split(stream);
2271 let mut reader = BufReader::new(reader);
2272 let hello = read_request(&mut reader)
2273 .await?
2274 .ok_or_else(|| eyre::eyre!("connection closed before the agent handshake"))?;
2275 let request: AgentRequest = serde_json::from_str(&hello)?;
2276 match request {
2277 AgentRequest::Hello {
2278 protocol,
2279 client_version,
2280 } if protocol == AGENT_PROTOCOL_VERSION && client_version == self.version.as_ref() => {}
2281 AgentRequest::Hello { protocol, .. } if protocol != AGENT_PROTOCOL_VERSION => {
2282 send_response(
2283 &mut writer,
2284 &AgentResponse::Error {
2285 message: format!(
2286 "unsupported agent protocol {protocol}; expected {AGENT_PROTOCOL_VERSION}"
2287 ),
2288 },
2289 )
2290 .await?;
2291 return Ok(());
2292 }
2293 AgentRequest::Hello { client_version, .. } => {
2294 send_response(
2295 &mut writer,
2296 &AgentResponse::Error {
2297 message: format!(
2298 "cache client {client_version} does not match agent {}",
2299 self.version
2300 ),
2301 },
2302 )
2303 .await?;
2304 return Ok(());
2305 }
2306 _ => bail!("the first agent request must be hello"),
2307 }
2308 send_response(
2309 &mut writer,
2310 &AgentResponse::Hello {
2311 protocol: AGENT_PROTOCOL_VERSION,
2312 agent_version: self.version.to_string(),
2313 },
2314 )
2315 .await?;
2316
2317 let mut connection = ConnectionUploads::default();
2321 while let Some(line) = read_request(&mut reader).await? {
2322 let response = match serde_json::from_str(&line) {
2323 Ok(request) => self.respond_on(request, &mut connection).await,
2324 Err(error) => AgentResponse::Error {
2325 message: format!("invalid agent request: {error}"),
2326 },
2327 };
2328 send_response(&mut writer, &response).await?;
2329 }
2330 Ok(())
2331 }
2332}
2333
2334async fn read_request<R>(reader: &mut R) -> Result<Option<String>>
2340where
2341 R: AsyncBufRead + Unpin,
2342{
2343 let mut line = Vec::new();
2344 loop {
2345 let available = reader.fill_buf().await?;
2346 if available.is_empty() {
2347 break;
2348 }
2349 let (consumed, complete) = match available.iter().position(|byte| *byte == b'\n') {
2350 Some(index) => (index, true),
2351 None => (available.len(), false),
2352 };
2353 if line.len() + consumed > MAX_REQUEST_BYTES {
2354 bail!("agent request exceeded {MAX_REQUEST_BYTES} bytes");
2355 }
2356 line.extend_from_slice(&available[..consumed]);
2357 reader.consume(consumed + usize::from(complete));
2359 if complete {
2360 return Ok(Some(String::from_utf8(line)?));
2361 }
2362 }
2363 if line.is_empty() {
2364 Ok(None)
2365 } else {
2366 Ok(Some(String::from_utf8(line)?))
2367 }
2368}
2369
2370async fn send_response(
2371 writer: &mut (impl AsyncWrite + Unpin),
2372 response: &AgentResponse,
2373) -> Result<()> {
2374 let mut encoded = serde_json::to_vec(response)?;
2375 encoded.push(b'\n');
2376 writer.write_all(&encoded).await?;
2377 writer.flush().await?;
2378 Ok(())
2379}
2380
2381#[cfg(test)]
2382#[path = "agent_tests.rs"]
2383mod tests;