1use std::collections::HashSet;
2use std::path::{Path, PathBuf};
3use std::process::Stdio;
4
5use serde::{Deserialize, Serialize};
6
7use super::cursor::CursorProbeResult;
8use super::probe_refresh::ProbeCacheBranch;
9use crate::error::MarsError;
10
11const SCHEMA_VERSION: u32 = 1;
12const DEFAULT_TTL_SECS: u64 = 60;
13
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct ProbeCacheEntry {
16 pub schema_version: u32,
17 pub fetched_at: u64,
18 pub last_attempt_at: u64,
19 pub last_error: Option<String>,
20 pub result: Option<CursorProbeResult>,
21}
22
23#[derive(Debug, Clone)]
24pub enum CachedCursorProbeOutcome {
25 Hit(CursorProbeResult),
26 Stale(CursorProbeResult),
27 Miss(CursorProbeResult),
28 Failed(CursorProbeResult),
29 Unavailable,
30}
31
32impl CachedCursorProbeOutcome {
33 pub fn result(&self) -> Option<&CursorProbeResult> {
34 match self {
35 Self::Hit(r) | Self::Stale(r) | Self::Miss(r) | Self::Failed(r) => Some(r),
36 Self::Unavailable => None,
37 }
38 }
39
40 pub fn cache_status(&self) -> &'static str {
41 match self {
42 Self::Hit(_) => "hit",
43 Self::Stale(_) => "stale",
44 Self::Miss(_) => "miss",
45 Self::Failed(_) => "failed",
46 Self::Unavailable => "skipped",
47 }
48 }
49}
50pub fn read_cached_probe_result_usable() -> Option<CursorProbeResult> {
54 let entry = read_cache_tolerant()?;
55 if !is_usable(&entry) {
56 return None;
57 }
58 entry.result
59}
60
61fn cache_dir() -> Result<PathBuf, MarsError> {
62 let root = crate::platform::cache::global_cache_root()?;
63 Ok(root.join("availability"))
64}
65
66fn cache_path() -> Result<PathBuf, MarsError> {
67 Ok(cache_dir()?.join("cursor-probe.json"))
68}
69
70fn lock_path() -> Result<PathBuf, MarsError> {
71 Ok(cache_dir()?.join(".cursor-probe.lock"))
72}
73
74fn ttl_secs() -> u64 {
75 std::env::var("MARS_PROBE_CACHE_TTL_SECS")
76 .ok()
77 .and_then(|v| v.parse::<u64>().ok())
78 .unwrap_or(DEFAULT_TTL_SECS)
79}
80
81fn now_unix_secs() -> u64 {
82 std::time::SystemTime::now()
83 .duration_since(std::time::UNIX_EPOCH)
84 .unwrap_or_default()
85 .as_secs()
86}
87
88fn is_fresh(entry: &ProbeCacheEntry) -> bool {
89 let ttl = ttl_secs();
90 let now = now_unix_secs();
91 if entry.fetched_at > now {
92 return false;
93 }
94 (now - entry.fetched_at) < ttl
95}
96
97fn is_usable(entry: &ProbeCacheEntry) -> bool {
98 entry.result.as_ref().is_some_and(|r| r.model_probe_success)
99}
100
101fn read_cache_tolerant() -> Option<ProbeCacheEntry> {
102 read_cache_tolerant_at(&cache_path().ok()?)
103}
104
105fn read_cache_tolerant_at(path: &Path) -> Option<ProbeCacheEntry> {
106 let content = std::fs::read_to_string(path).ok()?;
107 let entry: ProbeCacheEntry = serde_json::from_str(&content).ok()?;
108 if entry.schema_version != SCHEMA_VERSION {
109 return None;
110 }
111 Some(entry)
112}
113
114fn write_cache(entry: &ProbeCacheEntry) -> Result<(), MarsError> {
115 write_cache_at(&cache_path()?, entry)
116}
117
118fn write_cache_at(path: &Path, entry: &ProbeCacheEntry) -> Result<(), MarsError> {
119 let json = serde_json::to_string_pretty(entry)
120 .map_err(|e| MarsError::Internal(format!("probe cache serialize: {e}")))?;
121 crate::fs::atomic_write(path, json.as_bytes())
122}
123
124struct FileLock {
125 _file: std::fs::File,
126}
127
128fn try_lock() -> Option<FileLock> {
129 lock_at(&lock_path().ok()?, true)
130}
131
132fn blocking_lock() -> Option<FileLock> {
133 lock_at(&lock_path().ok()?, false)
134}
135
136fn lock_at(path: &Path, nonblocking: bool) -> Option<FileLock> {
137 if let Some(parent) = path.parent() {
138 std::fs::create_dir_all(parent).ok()?;
139 }
140 let file = std::fs::OpenOptions::new()
141 .create(true)
142 .write(true)
143 .truncate(false)
144 .open(path)
145 .ok()?;
146
147 #[cfg(unix)]
148 {
149 use std::os::unix::io::AsRawFd;
150 let flags = if nonblocking {
151 libc::LOCK_EX | libc::LOCK_NB
152 } else {
153 libc::LOCK_EX
154 };
155 let ret = unsafe { libc::flock(file.as_raw_fd(), flags) };
156 if ret != 0 {
157 return None;
158 }
159 }
160
161 #[cfg(windows)]
162 {
163 use std::os::windows::io::AsRawHandle;
164 use windows_sys::Win32::Foundation::HANDLE;
165 use windows_sys::Win32::Storage::FileSystem::{
166 LOCKFILE_EXCLUSIVE_LOCK, LOCKFILE_FAIL_IMMEDIATELY, LockFileEx,
167 };
168 let handle = file.as_raw_handle() as HANDLE;
169 let mut overlapped = unsafe { std::mem::zeroed() };
170 let flags = if nonblocking {
171 LOCKFILE_EXCLUSIVE_LOCK | LOCKFILE_FAIL_IMMEDIATELY
172 } else {
173 LOCKFILE_EXCLUSIVE_LOCK
174 };
175 let ret = unsafe { LockFileEx(handle, flags, 0, 1, 0, &mut overlapped) };
176 if ret == 0 {
177 return None;
178 }
179 }
180
181 Some(FileLock { _file: file })
182}
183
184pub fn probe_cached(
185 installed: &HashSet<String>,
186 mars_offline: bool,
187 probe_refresh: super::ProbeRefreshMode,
188) -> CachedCursorProbeOutcome {
189 if !super::should_probe_cursor(installed, mars_offline) {
190 return CachedCursorProbeOutcome::Unavailable;
191 }
192
193 probe_cached_impl(
194 mars_offline,
195 probe_refresh,
196 &cache_path().ok(),
197 super::cursor::probe,
198 || spawn_detached_refresh().map_err(|_| ()),
199 )
200}
201
202fn probe_cached_impl<F, S>(
203 mars_offline: bool,
204 probe_refresh: super::ProbeRefreshMode,
205 path: &Option<PathBuf>,
206 probe: F,
207 spawn_refresh: S,
208) -> CachedCursorProbeOutcome
209where
210 F: Fn() -> CursorProbeResult,
211 S: Fn() -> Result<(), ()>,
212{
213 let cached = path.as_deref().and_then(read_cache_tolerant_at);
214 match super::probe_refresh::resolve_probe_cache_branch(
215 cached,
216 mars_offline,
217 probe_refresh,
218 |entry| entry.result.as_ref().filter(|_| is_usable(entry)),
219 is_fresh,
220 || trigger_background_refresh_with(spawn_refresh),
221 ) {
222 ProbeCacheBranch::Hit(result) => CachedCursorProbeOutcome::Hit(result),
223 ProbeCacheBranch::Stale(result) => CachedCursorProbeOutcome::Stale(result),
224 ProbeCacheBranch::Unavailable => CachedCursorProbeOutcome::Unavailable,
225 ProbeCacheBranch::SynchronousProbe => synchronous_probe_with(path, probe),
226 }
227}
228
229fn trigger_background_refresh_with<S>(spawn_refresh: S)
230where
231 S: Fn() -> Result<(), ()>,
232{
233 let Some(lock) = try_lock() else { return };
234 if let Some(entry) = read_cache_tolerant()
235 && is_fresh(&entry)
236 && is_usable(&entry)
237 {
238 drop(lock);
239 return;
240 }
241 let _ = spawn_refresh();
242 drop(lock);
243}
244
245fn synchronous_probe_with<F>(path: &Option<PathBuf>, probe: F) -> CachedCursorProbeOutcome
246where
247 F: Fn() -> CursorProbeResult,
248{
249 let lock = blocking_lock();
250
251 if lock.is_some()
252 && let Some(path) = path
253 && let Some(entry) = read_cache_tolerant_at(path)
254 && is_usable(&entry)
255 {
256 if is_fresh(&entry) {
257 return CachedCursorProbeOutcome::Hit(entry.result.unwrap());
258 }
259 let probe_result = probe();
260 if probe_result.model_probe_success {
261 write_probe_attempt(path, probe_result.clone());
262 return CachedCursorProbeOutcome::Miss(probe_result);
263 } else {
264 write_failed_attempt(path, &entry, &probe_result);
265 return CachedCursorProbeOutcome::Stale(entry.result.unwrap());
266 }
267 }
268
269 let probe_result = probe();
270 if let Some(path) = path {
271 write_probe_attempt(path, probe_result.clone());
272 }
273 drop(lock);
274
275 if probe_result.model_probe_success {
276 CachedCursorProbeOutcome::Miss(probe_result)
277 } else {
278 CachedCursorProbeOutcome::Failed(probe_result)
279 }
280}
281
282fn write_probe_attempt(path: &Path, probe_result: CursorProbeResult) {
283 let now = now_unix_secs();
284 let entry = ProbeCacheEntry {
285 schema_version: SCHEMA_VERSION,
286 fetched_at: now,
287 last_attempt_at: now,
288 last_error: if probe_result.model_probe_success {
289 None
290 } else {
291 probe_result.error.clone()
292 },
293 result: Some(probe_result),
294 };
295
296 if let Err(e) = write_cache_at(path, &entry) {
297 eprintln!("debug: probe cache write failed: {e}");
298 }
299}
300
301fn write_failed_attempt(path: &Path, existing: &ProbeCacheEntry, failed_probe: &CursorProbeResult) {
302 let now = now_unix_secs();
303 let entry = ProbeCacheEntry {
304 schema_version: SCHEMA_VERSION,
305 fetched_at: existing.fetched_at,
306 last_attempt_at: now,
307 last_error: failed_probe.error.clone(),
308 result: existing.result.clone(),
309 };
310
311 if let Err(e) = write_cache_at(path, &entry) {
312 eprintln!("debug: probe cache write failed: {e}");
313 }
314}
315
316fn spawn_detached_refresh() -> std::io::Result<()> {
317 let mars_bin = std::env::current_exe()?;
318 let mut cmd = std::process::Command::new(mars_bin);
319 cmd.args(["models", "__refresh-probe", "--target", "cursor"]);
320 cmd.stdin(Stdio::null());
321 cmd.stdout(Stdio::null());
322 cmd.stderr(Stdio::null());
323
324 #[cfg(unix)]
325 {
326 use std::os::unix::process::CommandExt;
327 unsafe {
328 cmd.pre_exec(|| {
329 libc::setsid();
330 Ok(())
331 });
332 }
333 }
334
335 #[cfg(windows)]
336 {
337 use std::os::windows::process::CommandExt;
338 cmd.creation_flags(0x00000008);
339 }
340
341 cmd.spawn()?;
342 Ok(())
343}
344
345pub fn run_refresh_probe_command() -> Result<i32, MarsError> {
346 let Some(_lock) = blocking_lock() else {
347 return Ok(0);
348 };
349
350 if let Some(entry) = read_cache_tolerant()
351 && is_fresh(&entry)
352 && is_usable(&entry)
353 {
354 return Ok(0);
355 }
356
357 let probe_result = super::cursor::probe();
358 let now = now_unix_secs();
359 let existing = read_cache_tolerant();
360
361 let entry = if probe_result.model_probe_success {
362 ProbeCacheEntry {
363 schema_version: SCHEMA_VERSION,
364 fetched_at: now,
365 last_attempt_at: now,
366 last_error: None,
367 result: Some(probe_result),
368 }
369 } else {
370 ProbeCacheEntry {
371 schema_version: SCHEMA_VERSION,
372 fetched_at: existing.as_ref().map(|e| e.fetched_at).unwrap_or(0),
373 last_attempt_at: now,
374 last_error: probe_result.error.clone(),
375 result: existing.and_then(|e| e.result),
376 }
377 };
378 let _ = write_cache(&entry);
379
380 Ok(0)
381}
382
383#[cfg(test)]
384mod tests {
385 use super::*;
386 use std::cell::Cell;
387 use tempfile::TempDir;
388
389 fn ok_result() -> CursorProbeResult {
390 CursorProbeResult {
391 slugs: vec!["gpt-5.5-high".to_string()],
392 model_probe_success: true,
393 error: None,
394 }
395 }
396
397 fn fail_result() -> CursorProbeResult {
398 CursorProbeResult {
399 model_probe_success: false,
400 error: Some("boom".to_string()),
401 ..CursorProbeResult::default()
402 }
403 }
404
405 fn entry(fetched_at: u64, result: Option<CursorProbeResult>) -> ProbeCacheEntry {
406 ProbeCacheEntry {
407 schema_version: SCHEMA_VERSION,
408 fetched_at,
409 last_attempt_at: fetched_at,
410 last_error: None,
411 result,
412 }
413 }
414
415 fn cache_file(temp: &TempDir) -> PathBuf {
416 temp.path().join("availability").join("cursor-probe.json")
417 }
418
419 fn write_entry(path: &Path, entry: &ProbeCacheEntry) {
420 write_cache_at(path, entry).unwrap();
421 }
422
423 #[test]
424 fn fresh_hit_returns_cached_result() {
425 let temp = TempDir::new().unwrap();
426 let path = cache_file(&temp);
427 write_entry(&path, &entry(now_unix_secs(), Some(ok_result())));
428
429 let outcome = probe_cached_impl(
430 false,
431 crate::models::probes::ProbeRefreshMode::Background,
432 &Some(path),
433 fail_result,
434 || Ok(()),
435 );
436 assert!(matches!(outcome, CachedCursorProbeOutcome::Hit(_)));
437 assert_eq!(outcome.result().unwrap().slugs[0], "gpt-5.5-high");
438 }
439
440 #[test]
441 fn stale_entry_returns_stale_outcome() {
442 let temp = TempDir::new().unwrap();
443 let path = cache_file(&temp);
444 write_entry(&path, &entry(1, Some(ok_result())));
445
446 let outcome = probe_cached_impl(
447 false,
448 crate::models::probes::ProbeRefreshMode::Background,
449 &Some(path),
450 fail_result,
451 || Ok(()),
452 );
453 assert!(matches!(outcome, CachedCursorProbeOutcome::Stale(_)));
454 }
455
456 #[test]
457 fn stale_cache_preserved_on_failed_probe() {
458 let temp = TempDir::new().unwrap();
459 let path = cache_file(&temp);
460 write_entry(&path, &entry(1, Some(ok_result())));
461
462 let outcome = probe_cached_impl(
463 false,
464 crate::models::probes::ProbeRefreshMode::Background,
465 &Some(path.clone()),
466 fail_result,
467 || Ok(()),
468 );
469
470 assert!(matches!(outcome, CachedCursorProbeOutcome::Stale(_)));
471
472 let on_disk = read_cache_tolerant_at(&path).unwrap();
473 assert!(on_disk.result.as_ref().unwrap().model_probe_success);
474 assert_eq!(on_disk.fetched_at, 1);
475 }
476
477 #[test]
478 fn missing_cache_runs_synchronous_probe() {
479 let temp = TempDir::new().unwrap();
480 let path = cache_file(&temp);
481 let called = Cell::new(false);
482 let outcome = probe_cached_impl(
483 false,
484 crate::models::probes::ProbeRefreshMode::Background,
485 &Some(path.clone()),
486 || {
487 called.set(true);
488 ok_result()
489 },
490 || Ok(()),
491 );
492
493 assert!(called.get());
494 assert!(matches!(outcome, CachedCursorProbeOutcome::Miss(_)));
495 assert!(read_cache_tolerant_at(&path).is_some());
496 }
497
498 #[test]
499 fn invalid_json_is_cache_miss() {
500 let temp = TempDir::new().unwrap();
501 let path = cache_file(&temp);
502 std::fs::create_dir_all(path.parent().unwrap()).unwrap();
503 std::fs::write(&path, "not json").unwrap();
504
505 let outcome = probe_cached_impl(
506 false,
507 crate::models::probes::ProbeRefreshMode::Background,
508 &Some(path),
509 ok_result,
510 || Ok(()),
511 );
512 assert!(matches!(outcome, CachedCursorProbeOutcome::Miss(_)));
513 }
514
515 #[test]
516 fn incompatible_schema_is_cache_miss() {
517 let temp = TempDir::new().unwrap();
518 let path = cache_file(&temp);
519 let mut old = entry(now_unix_secs(), Some(ok_result()));
520 old.schema_version = 999;
521 write_entry(&path, &old);
522
523 let outcome = probe_cached_impl(
524 false,
525 crate::models::probes::ProbeRefreshMode::Background,
526 &Some(path),
527 ok_result,
528 || Ok(()),
529 );
530 assert!(matches!(outcome, CachedCursorProbeOutcome::Miss(_)));
531 }
532
533 #[test]
534 fn future_fetched_at_is_stale() {
535 let future = entry(now_unix_secs() + 3600, Some(ok_result()));
536 assert!(!is_fresh(&future));
537 }
538
539 #[test]
540 fn ttl_override_controls_freshness() {
541 let _guard = EnvGuard::set("MARS_PROBE_CACHE_TTL_SECS", "9999");
542 let recent = entry(now_unix_secs().saturating_sub(10), Some(ok_result()));
543 assert!(is_fresh(&recent));
544 }
545
546 #[test]
547 fn write_failure_degrades_gracefully() {
548 let temp = TempDir::new().unwrap();
549 let path = temp.path().join("availability");
550 std::fs::write(&path, "file blocks directory").unwrap();
551 let blocked = path.join("cursor-probe.json");
552
553 let outcome = probe_cached_impl(
554 false,
555 crate::models::probes::ProbeRefreshMode::Background,
556 &Some(blocked),
557 ok_result,
558 || Ok(()),
559 );
560 assert!(matches!(outcome, CachedCursorProbeOutcome::Miss(_)));
561 }
562
563 #[test]
564 fn cold_probe_failure_returns_failed_outcome_with_error_status() {
565 let temp = TempDir::new().unwrap();
566 let path = cache_file(&temp);
567
568 let outcome = probe_cached_impl(
569 false,
570 crate::models::probes::ProbeRefreshMode::Background,
571 &Some(path),
572 fail_result,
573 || Ok(()),
574 );
575
576 assert!(matches!(outcome, CachedCursorProbeOutcome::Failed(_)));
577 assert_eq!(outcome.cache_status(), "failed");
578 assert!(!outcome.result().unwrap().model_probe_success);
579 assert_eq!(outcome.result().unwrap().error.as_deref(), Some("boom"));
580 }
581
582 struct EnvGuard {
583 key: &'static str,
584 prev: Option<std::ffi::OsString>,
585 }
586
587 impl EnvGuard {
588 fn set(key: &'static str, value: &str) -> Self {
589 let prev = std::env::var_os(key);
590 unsafe { std::env::set_var(key, value) };
591 Self { key, prev }
592 }
593 }
594
595 impl Drop for EnvGuard {
596 fn drop(&mut self) {
597 if let Some(prev) = &self.prev {
598 unsafe { std::env::set_var(self.key, prev) };
599 } else {
600 unsafe { std::env::remove_var(self.key) };
601 }
602 }
603 }
604}