ai_usagebar/kimi/
lock.rs

1//! Cross-process lock around a Kimi Code CLI credential refresh, speaking the
2//! CLI's own lock protocol.
3//!
4//! Unlike every other lock in this codebase (which guards a file ai-usagebar
5//! owns and therefore uses `flock` via `cache::acquire_lock_async`), this one
6//! guards a file **another program** owns and refreshes on its own schedule.
7//! kimi-code locks its credential store with `proper-lockfile`, whose contract
8//! is a *directory*: `lock(target)` succeeds by `mkdir`ing `<target>.lock`, an
9//! existing lock directory whose mtime is older than `stale` may be stolen, and
10//! the holder keeps refreshing that mtime while it works. `flock` here would be
11//! invisible to the CLI and both processes would rotate the same refresh token
12//! at once, so this speaks the protocol the other side actually implements.
13//!
14//! Windows is deliberately unlocked: kimi-code disables its own lock there
15//! (`if (process.platform === "win32") return undefined`), so taking one would
16//! only be a lock against ourselves.
17
18use std::path::{Path, PathBuf};
19use std::time::{Duration, SystemTime};
20
21use crate::error::{AppError, Result};
22
23/// proper-lockfile's default `stale` for kimi-code's OAuth lock (5 s), after
24/// which an abandoned lock directory may be taken over.
25pub const STALE: Duration = Duration::from_secs(5);
26/// How often the holder touches the lock directory. proper-lockfile refreshes
27/// at `stale / 2`; matching that keeps a slow refresh from looking abandoned.
28const HEARTBEAT: Duration = Duration::from_millis(2_500);
29const POLL: Duration = Duration::from_millis(250);
30/// Long enough to outlast a peer's refresh round-trip, short enough that a
31/// widget tick never hangs on it.
32pub const ACQUIRE_TIMEOUT: Duration = Duration::from_secs(20);
33
34/// The directory `proper-lockfile` actually creates for `target`.
35pub fn lock_dir_for(target: &Path) -> PathBuf {
36    let mut dir = target.as_os_str().to_os_string();
37    dir.push(".lock");
38    PathBuf::from(dir)
39}
40
41/// Held lock. Dropping it releases the lock and stops the heartbeat.
42#[derive(Debug)]
43pub struct OauthLock {
44    dir: PathBuf,
45    heartbeat: Option<tokio::task::JoinHandle<()>>,
46}
47
48impl Drop for OauthLock {
49    fn drop(&mut self) {
50        if let Some(handle) = self.heartbeat.take() {
51            handle.abort();
52        }
53        // `remove_dir_all`, not `remove_dir`: an aborted heartbeat can leave
54        // its sentinel behind, and a lock directory we fail to remove would
55        // block the CLI for `STALE` on every one of its own refreshes.
56        let _ = std::fs::remove_dir_all(&self.dir);
57    }
58}
59
60/// Acquire the lock for `target` (kimi-code's `<home>/oauth/kimi-code`),
61/// stealing it only once it has been abandoned for [`STALE`].
62pub async fn acquire(target: &Path) -> Result<OauthLock> {
63    acquire_with(target, ACQUIRE_TIMEOUT, STALE, POLL).await
64}
65
66/// Seam for tests: they pass a tiny `stale`/`poll` so the steal path does not
67/// need a five-second wall-clock wait.
68pub async fn acquire_with(
69    target: &Path,
70    timeout: Duration,
71    stale: Duration,
72    poll: Duration,
73) -> Result<OauthLock> {
74    let dir = lock_dir_for(target);
75    if let Some(parent) = dir.parent() {
76        std::fs::create_dir_all(parent).map_err(|e| AppError::io_at(parent, e))?;
77    }
78
79    let deadline = tokio::time::Instant::now() + timeout;
80    loop {
81        match std::fs::create_dir(&dir) {
82            Ok(()) => return Ok(hold(dir)),
83            Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {
84                if is_stale(&dir, stale) {
85                    // Best-effort steal: losing the race just means another
86                    // process took it over first, and the next loop waits.
87                    let _ = std::fs::remove_dir_all(&dir);
88                }
89            }
90            Err(e) => return Err(AppError::io_at(&dir, e)),
91        }
92        if tokio::time::Instant::now() >= deadline {
93            return Err(AppError::Transport(format!(
94                "timed out waiting for the Kimi Code CLI credential lock at {}",
95                crate::display::sanitize_untrusted_path(&dir)
96            )));
97        }
98        tokio::time::sleep(poll).await;
99    }
100}
101
102fn hold(dir: PathBuf) -> OauthLock {
103    let beat_dir = dir.clone();
104    let heartbeat = tokio::runtime::Handle::try_current().ok().map(|_| {
105        tokio::spawn(async move {
106            loop {
107                tokio::time::sleep(HEARTBEAT).await;
108                touch(&beat_dir);
109            }
110        })
111    });
112    OauthLock { dir, heartbeat }
113}
114
115/// Bump the lock directory's mtime, which is what proper-lockfile reads to
116/// decide whether a lock was abandoned. Adding and removing an entry is how a
117/// POSIX directory's mtime moves — there is no `utimes` in `std`, and pulling a
118/// crate in for one syscall is not worth a dependency.
119fn touch(dir: &Path) {
120    let sentinel = dir.join(".ai-usagebar-heartbeat");
121    if std::fs::write(&sentinel, b"").is_ok() {
122        let _ = std::fs::remove_file(&sentinel);
123    }
124}
125
126fn is_stale(dir: &Path, stale: Duration) -> bool {
127    let Ok(modified) = std::fs::metadata(dir).and_then(|m| m.modified()) else {
128        // No mtime to trust: leave the lock alone rather than steal blindly.
129        return false;
130    };
131    SystemTime::now()
132        .duration_since(modified)
133        .is_ok_and(|age| age > stale)
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139    use tempfile::TempDir;
140
141    fn target(td: &TempDir) -> PathBuf {
142        td.path().join("oauth").join("kimi-code")
143    }
144
145    #[test]
146    fn lock_dir_matches_proper_lockfile_naming() {
147        assert_eq!(
148            lock_dir_for(Path::new("/home/u/.kimi-code/oauth/kimi-code")),
149            PathBuf::from("/home/u/.kimi-code/oauth/kimi-code.lock")
150        );
151    }
152
153    #[tokio::test]
154    async fn acquire_creates_the_lock_directory_and_drop_removes_it() {
155        let td = TempDir::new().unwrap();
156        let target = target(&td);
157        let dir = lock_dir_for(&target);
158        {
159            let _lock = acquire_with(&target, Duration::from_secs(1), STALE, POLL)
160                .await
161                .unwrap();
162            assert!(dir.is_dir());
163        }
164        assert!(!dir.exists());
165    }
166
167    #[tokio::test]
168    async fn a_fresh_foreign_lock_is_waited_out_not_stolen() {
169        let td = TempDir::new().unwrap();
170        let target = target(&td);
171        let dir = lock_dir_for(&target);
172        std::fs::create_dir_all(&dir).unwrap();
173
174        let err = acquire_with(
175            &target,
176            Duration::from_millis(60),
177            Duration::from_secs(3_600),
178            Duration::from_millis(10),
179        )
180        .await
181        .unwrap_err();
182        assert!(matches!(err, AppError::Transport(_)), "{err:?}");
183        assert!(dir.is_dir(), "a live peer's lock must survive");
184    }
185
186    #[tokio::test]
187    async fn an_abandoned_lock_is_stolen_once_stale() {
188        let td = TempDir::new().unwrap();
189        let target = target(&td);
190        let dir = lock_dir_for(&target);
191        std::fs::create_dir_all(&dir).unwrap();
192
193        let lock = acquire_with(
194            &target,
195            Duration::from_secs(2),
196            Duration::ZERO,
197            Duration::from_millis(10),
198        )
199        .await
200        .unwrap();
201        assert!(dir.is_dir());
202        drop(lock);
203        assert!(!dir.exists());
204    }
205
206    #[tokio::test]
207    async fn a_leftover_heartbeat_sentinel_does_not_block_release() {
208        let td = TempDir::new().unwrap();
209        let target = target(&td);
210        let dir = lock_dir_for(&target);
211        {
212            let _lock = acquire_with(&target, Duration::from_secs(1), STALE, POLL)
213                .await
214                .unwrap();
215            std::fs::write(dir.join(".ai-usagebar-heartbeat"), b"").unwrap();
216        }
217        assert!(!dir.exists());
218    }
219
220    #[test]
221    fn touch_moves_the_directory_mtime_forward() {
222        let td = TempDir::new().unwrap();
223        let dir = td.path().join("kimi-code.lock");
224        std::fs::create_dir(&dir).unwrap();
225        let before = std::fs::metadata(&dir).unwrap().modified().unwrap();
226        // Filesystem mtime granularity can be coarse; assert the touch is not
227        // *older*, and that the sentinel never outlives the call.
228        touch(&dir);
229        let after = std::fs::metadata(&dir).unwrap().modified().unwrap();
230        assert!(after >= before);
231        assert!(!dir.join(".ai-usagebar-heartbeat").exists());
232    }
233}
ai_usagebar/kimi/lock.rs

ai_usagebar/kimi/
lock.rs
