ai_usagebar/kimi/

lock.rs

//! Cross-process lock around a Kimi Code CLI credential refresh, speaking the
//! CLI's own lock protocol.
//!
//! Unlike every other lock in this codebase (which guards a file ai-usagebar
//! owns and therefore uses `flock` via `cache::acquire_lock_async`), this one
//! guards a file **another program** owns and refreshes on its own schedule.
//! kimi-code locks its credential store with `proper-lockfile`, whose contract
//! is a *directory*: `lock(target)` succeeds by `mkdir`ing `<target>.lock`, an
//! existing lock directory whose mtime is older than `stale` may be stolen, and
//! the holder keeps refreshing that mtime while it works. `flock` here would be
//! invisible to the CLI and both processes would rotate the same refresh token
//! at once, so this speaks the protocol the other side actually implements.
//!
//! Windows is deliberately unlocked: kimi-code disables its own lock there
//! (`if (process.platform === "win32") return undefined`), so taking one would
//! only be a lock against ourselves.

use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime};

use crate::error::{AppError, Result};

/// proper-lockfile's default `stale` for kimi-code's OAuth lock (5 s), after
/// which an abandoned lock directory may be taken over.
pub const STALE: Duration = Duration::from_secs(5);
/// How often the holder touches the lock directory. proper-lockfile refreshes
/// at `stale / 2`; matching that keeps a slow refresh from looking abandoned.
const HEARTBEAT: Duration = Duration::from_millis(2_500);
const POLL: Duration = Duration::from_millis(250);
/// Long enough to outlast a peer's refresh round-trip, short enough that a
/// widget tick never hangs on it.
pub const ACQUIRE_TIMEOUT: Duration = Duration::from_secs(20);

/// The directory `proper-lockfile` actually creates for `target`.
pub fn lock_dir_for(target: &Path) -> PathBuf {
    let mut dir = target.as_os_str().to_os_string();
    dir.push(".lock");
    PathBuf::from(dir)
}

/// Held lock. Dropping it releases the lock and stops the heartbeat.
#[derive(Debug)]
pub struct OauthLock {
    dir: PathBuf,
    owner: String,
    heartbeat: Option<tokio::task::JoinHandle<()>>,
}

impl Drop for OauthLock {
    fn drop(&mut self) {
        if let Some(handle) = self.heartbeat.take() {
            handle.abort();
        }
        // Only release what is still ours. A holder delayed past `STALE` can
        // resume after another process stole the lock and rebuilt the
        // directory; deleting unconditionally would destroy the replacement
        // (#495). A missing or foreign owner marker means "not ours" — the
        // cost of leaving a lock behind is one stale-interval steal by the
        // next acquirer, never data loss.
        if owns(&self.dir, &self.owner) {
            // `remove_dir_all`, not `remove_dir`: an aborted heartbeat can leave
            // its sentinel behind, and a lock directory we fail to remove would
            // block the CLI for `STALE` on every one of its own refreshes.
            let _ = std::fs::remove_dir_all(&self.dir);
        }
    }
}

/// Acquire the lock for `target` (kimi-code's `<home>/oauth/kimi-code`),
/// stealing it only once it has been abandoned for [`STALE`].
pub async fn acquire(target: &Path) -> Result<OauthLock> {
    acquire_with(target, ACQUIRE_TIMEOUT, STALE, POLL).await
}

/// Seam for tests: they pass a tiny `stale`/`poll` so the steal path does not
/// need a five-second wall-clock wait.
pub async fn acquire_with(
    target: &Path,
    timeout: Duration,
    stale: Duration,
    poll: Duration,
) -> Result<OauthLock> {
    let dir = lock_dir_for(target);
    if let Some(parent) = dir.parent() {
        std::fs::create_dir_all(parent).map_err(|e| AppError::io_at(parent, e))?;
    }

    let deadline = tokio::time::Instant::now() + timeout;
    loop {
        match std::fs::create_dir(&dir) {
            Ok(()) => return Ok(hold(dir)),
            Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {
                if is_stale(&dir, stale) {
                    // Best-effort steal: losing the race just means another
                    // process took it over first, and the next loop waits.
                    let _ = std::fs::remove_dir_all(&dir);
                }
            }
            Err(e) => return Err(AppError::io_at(&dir, e)),
        }
        if tokio::time::Instant::now() >= deadline {
            return Err(AppError::Transport(format!(
                "timed out waiting for the Kimi Code CLI credential lock at {}",
                crate::display::sanitize_untrusted_path(&dir)
            )));
        }
        tokio::time::sleep(poll).await;
    }
}

fn hold(dir: PathBuf) -> OauthLock {
    // Identity for the ownership check (#495): unique per acquisition, so a
    // holder that resumes after its lock was stolen can tell the replacement
    // is not its own. Process id + nanos + a process-local counter keeps it
    // collision-free without a new dependency.
    static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
    let owner = format!(
        "{}-{}-{}",
        std::process::id(),
        SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .map(|d| d.as_nanos())
            .unwrap_or(0),
        COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
    );
    // Best-effort marker: if the write fails, this holder simply never
    // deletes on release (missing marker = not ours), never a foreign one.
    let _ = std::fs::write(dir.join(OWNER_FILE), owner.as_bytes());

    let beat_dir = dir.clone();
    let beat_owner = owner.clone();
    let heartbeat = tokio::runtime::Handle::try_current().ok().map(|_| {
        tokio::spawn(async move {
            loop {
                tokio::time::sleep(HEARTBEAT).await;
                if !touch(&beat_dir, &beat_owner) {
                    // The lock was stolen while we slept: stop touching a
                    // directory that belongs to someone else (#495).
                    break;
                }
            }
        })
    });
    OauthLock {
        dir,
        owner,
        heartbeat,
    }
}

/// The acquisition-time identity marker inside a held lock directory.
const OWNER_FILE: &str = ".ai-usagebar-owner";

/// Whether `dir` is still held under `owner`'s identity.
fn owns(dir: &Path, owner: &str) -> bool {
    std::fs::read(dir.join(OWNER_FILE)).is_ok_and(|bytes| bytes == owner.as_bytes())
}

/// Bump the lock directory's mtime, which is what proper-lockfile reads to
/// decide whether a lock was abandoned. Adding and removing an entry is how a
/// POSIX directory's mtime moves — there is no `utimes` in `std`, and pulling a
/// crate in for one syscall is not worth a dependency.
///
/// Returns false when the directory no longer belongs to `owner`, so the
/// heartbeat stops touching a stolen lock (#495).
fn touch(dir: &Path, owner: &str) -> bool {
    if !owns(dir, owner) {
        return false;
    }
    let sentinel = dir.join(".ai-usagebar-heartbeat");
    if std::fs::write(&sentinel, b"").is_ok() {
        let _ = std::fs::remove_file(&sentinel);
    }
    true
}

fn is_stale(dir: &Path, stale: Duration) -> bool {
    let Ok(modified) = std::fs::metadata(dir).and_then(|m| m.modified()) else {
        // No mtime to trust: leave the lock alone rather than steal blindly.
        return false;
    };
    SystemTime::now()
        .duration_since(modified)
        .is_ok_and(|age| age > stale)
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;

    fn target(td: &TempDir) -> PathBuf {
        td.path().join("oauth").join("kimi-code")
    }

    #[test]
    fn lock_dir_matches_proper_lockfile_naming() {
        assert_eq!(
            lock_dir_for(Path::new("/home/u/.kimi-code/oauth/kimi-code")),
            PathBuf::from("/home/u/.kimi-code/oauth/kimi-code.lock")
        );
    }

    #[tokio::test]
    async fn acquire_creates_the_lock_directory_and_drop_removes_it() {
        let td = TempDir::new().unwrap();
        let target = target(&td);
        let dir = lock_dir_for(&target);
        {
            let _lock = acquire_with(&target, Duration::from_secs(1), STALE, POLL)
                .await
                .unwrap();
            assert!(dir.is_dir());
        }
        assert!(!dir.exists());
    }

    #[tokio::test]
    async fn a_fresh_foreign_lock_is_waited_out_not_stolen() {
        let td = TempDir::new().unwrap();
        let target = target(&td);
        let dir = lock_dir_for(&target);
        std::fs::create_dir_all(&dir).unwrap();

        let err = acquire_with(
            &target,
            Duration::from_millis(60),
            Duration::from_secs(3_600),
            Duration::from_millis(10),
        )
        .await
        .unwrap_err();
        assert!(matches!(err, AppError::Transport(_)), "{err:?}");
        assert!(dir.is_dir(), "a live peer's lock must survive");
    }

    #[tokio::test]
    async fn an_abandoned_lock_is_stolen_once_stale() {
        let td = TempDir::new().unwrap();
        let target = target(&td);
        let dir = lock_dir_for(&target);
        std::fs::create_dir_all(&dir).unwrap();

        let lock = acquire_with(
            &target,
            Duration::from_secs(2),
            Duration::ZERO,
            Duration::from_millis(10),
        )
        .await
        .unwrap();
        assert!(dir.is_dir());
        drop(lock);
        assert!(!dir.exists());
    }

    #[tokio::test]
    async fn a_leftover_heartbeat_sentinel_does_not_block_release() {
        let td = TempDir::new().unwrap();
        let target = target(&td);
        let dir = lock_dir_for(&target);
        {
            let _lock = acquire_with(&target, Duration::from_secs(1), STALE, POLL)
                .await
                .unwrap();
            std::fs::write(dir.join(".ai-usagebar-heartbeat"), b"").unwrap();
        }
        assert!(!dir.exists());
    }

    #[test]
    fn touch_moves_the_directory_mtime_forward() {
        let td = TempDir::new().unwrap();
        let dir = td.path().join("kimi-code.lock");
        std::fs::create_dir(&dir).unwrap();
        std::fs::write(dir.join(OWNER_FILE), b"owner").unwrap();
        let before = std::fs::metadata(&dir).unwrap().modified().unwrap();
        // Filesystem mtime granularity can be coarse; assert the touch is not
        // *older*, and that the sentinel never outlives the call.
        assert!(touch(&dir, "owner"));
        let after = std::fs::metadata(&dir).unwrap().modified().unwrap();
        assert!(after >= before);
        assert!(!dir.join(".ai-usagebar-heartbeat").exists());
        // A foreign owner's directory is never touched (#495).
        assert!(!touch(&dir, "someone-else"));
    }

    /// Regression for #495: a holder delayed past `STALE` resumes after
    /// another process stole the lock; dropping it must preserve the
    /// replacement, and its heartbeat must stop touching the foreign
    /// directory. Reproduction contributed with the issue.
    #[cfg(unix)]
    #[tokio::test]
    async fn dropping_previous_holder_preserves_replacement_lock() {
        let td = TempDir::new().unwrap();
        let target = target(&td);
        let dir = lock_dir_for(&target);
        let poll = Duration::from_millis(1);
        let first = acquire_with(&target, Duration::from_secs(1), STALE, poll)
            .await
            .unwrap();

        std::fs::File::open(&dir)
            .unwrap()
            .set_times(
                std::fs::FileTimes::new()
                    .set_modified(SystemTime::now() - STALE - Duration::from_secs(1)),
            )
            .unwrap();

        let replacement = acquire_with(&target, Duration::from_secs(1), STALE, poll)
            .await
            .unwrap();
        assert!(dir.is_dir());

        // Give the original holder a full heartbeat interval to notice the
        // replacement before releasing it. This does not restore ownership.
        tokio::time::sleep(Duration::from_millis(2_600)).await;
        drop(first);

        assert!(dir.is_dir(), "old holder removed replacement holder's lock");
        drop(replacement);
    }
}