From 69605d7e5414786cfe004e93e392f8a90ac1638c Mon Sep 17 00:00:00 2001 From: Armel van Ravels Date: Fri, 9 Oct 2026 15:45:07 +0200 Subject: [PATCH] Initial commit --- .gitignore | 3 + Cargo.lock | 1125 +++++++++++++++++++++++++++++++++++++ Cargo.toml | 20 + README.md | 68 +++ fixtures/claude.json | 7 + fixtures/codex.json | 10 + fixtures/cursor-usage.txt | 12 + fixtures/cursor.txt | 1 + flake.lock | 27 + flake.nix | 30 + modules/home-manager.nix | 56 ++ src/activity.rs | 215 +++++++ src/cursor-guard.js | 91 +++ src/daemon.rs | 294 ++++++++++ src/discovery.rs | 227 ++++++++ src/formatting.rs | 161 ++++++ src/main.rs | 66 +++ src/model.rs | 355 ++++++++++++ src/parsers.rs | 442 +++++++++++++++ src/paths.rs | 36 ++ src/process.rs | 346 ++++++++++++ src/protocol.rs | 415 ++++++++++++++ src/providers.rs | 459 +++++++++++++++ src/refresh.rs | 312 ++++++++++ src/settings.rs | 88 +++ src/store.rs | 64 +++ src/waybar.rs | 258 +++++++++ 27 files changed, 5188 insertions(+) create mode 100644 .gitignore create mode 100644 Cargo.lock create mode 100644 Cargo.toml create mode 100644 README.md create mode 100644 fixtures/claude.json create mode 100644 fixtures/codex.json create mode 100644 fixtures/cursor-usage.txt create mode 100644 fixtures/cursor.txt create mode 100644 flake.lock create mode 100644 flake.nix create mode 100644 modules/home-manager.nix create mode 100644 src/activity.rs create mode 100644 src/cursor-guard.js create mode 100644 src/daemon.rs create mode 100644 src/discovery.rs create mode 100644 src/formatting.rs create mode 100644 src/main.rs create mode 100644 src/model.rs create mode 100644 src/parsers.rs create mode 100644 src/paths.rs create mode 100644 src/process.rs create mode 100644 src/protocol.rs create mode 100644 src/providers.rs create mode 100644 src/refresh.rs create mode 100644 src/settings.rs create mode 100644 src/store.rs create mode 100644 src/waybar.rs diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f717ddd --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +/target +/result +.direnv diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000..5987c11 --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,1125 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +dependencies = [ + "memchr", +] + +[[package]] +name = "ai-usage" +version = "0.1.0" +dependencies = [ + "chrono", + "futures-util", + "libc", + "regex", + "serde", + "serde_json", + "tokio", + "zbus", +] + +[[package]] +name = "android_system_properties" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae221649c9976a6f6c56ae1facf410f3ddb33cc661c4b7b61020a912d4237fbc" +dependencies = [ + "libc", +] + +[[package]] +name = "async-broadcast" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "435a87a52755b8f27fcf321ac4f04b2802e337c8c4872923137471ec39c37532" +dependencies = [ + "event-listener", + "event-listener-strategy", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-channel" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "924ed96dd52d1b75e9c1a3e6275715fd320f5f9439fb5a4a11fa51f4221158d2" +dependencies = [ + "concurrent-queue", + "event-listener-strategy", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-executor" +version = "1.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c96bf972d85afc50bf5ab8fe2d54d1586b4e0b46c97c50a0c9e71e2f7bcd812a" +dependencies = [ + "async-task", + "concurrent-queue", + "fastrand", + "futures-lite", + "pin-project-lite", + "slab", +] + +[[package]] +name = "async-io" +version = "2.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "456b8a8feb6f42d237746d4b3e9a178494627745c3c56c6ea55d92ba50d026fc" +dependencies = [ + "autocfg", + "cfg-if", + "concurrent-queue", + "futures-io", + "futures-lite", + "parking", + "polling", + "rustix", + "slab", + "windows-sys", +] + +[[package]] +name = "async-lock" +version = "3.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "290f7f2596bd5b78a9fec8088ccd89180d7f9f55b94b0576823bbbdc72ee8311" +dependencies = [ + "event-listener", + "event-listener-strategy", + "pin-project-lite", +] + +[[package]] +name = "async-process" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc50921ec0055cdd8a16de48773bfeec5c972598674347252c0399676be7da75" +dependencies = [ + "async-channel", + "async-io", + "async-lock", + "async-signal", + "async-task", + "blocking", + "cfg-if", + "event-listener", + "futures-lite", + "rustix", +] + +[[package]] +name = "async-recursion" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f8abc12baad266b1c8cec146854c195b5864b4221d4b2ca7296a7ae82d9e451" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.6", +] + +[[package]] +name = "async-signal" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52b5aaafa020cf5053a01f2a60e8ff5dccf550f0f77ec54a4e47285ac2bab485" +dependencies = [ + "async-io", + "async-lock", + "atomic-waker", + "cfg-if", + "futures-core", + "futures-io", + "rustix", + "signal-hook-registry", + "slab", + "windows-sys", +] + +[[package]] +name = "async-task" +version = "4.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b75356056920673b02621b35afd0f7dda9306d03c79a30f5c56c44cf256e3de" + +[[package]] +name = "async-trait" +version = "0.1.92" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.6", +] + +[[package]] +name = "atomic-waker" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "bitflags" +version = "2.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" + +[[package]] +name = "blocking" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a70e4329df6cb94385eed412ec92375c3cdd8a6e502493d1229b6414e4036dfa" +dependencies = [ + "async-channel", + "async-task", + "futures-io", + "futures-lite", + "piper", +] + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "bytes" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" + +[[package]] +name = "cc" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f74872d07caf508b30a21f6836e7d7016a2eaf7d9ff4f48deaa58cd8a0407630" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" + +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "iana-time-zone", + "num-traits", + "serde", + "windows-link", +] + +[[package]] +name = "concurrent-queue" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "crossbeam-utils" +version = "0.8.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6" + +[[package]] +name = "endi" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66b7e2430c6dff6a955451e2cfc438f09cea1965a9d6f87f7e3b90decc014099" + +[[package]] +name = "enumflags2" +version = "0.7.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1027f7680c853e056ebcec683615fb6fbbc07dbaa13b4d5d9442b146ded4ecef" +dependencies = [ + "enumflags2_derive", + "serde", +] + +[[package]] +name = "enumflags2_derive" +version = "0.7.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67c78a4d8fdf9953a5c9d458f9efe940fd97a0cab0941c075a813ac594733827" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "event-listener" +version = "5.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a23add41df1562121a9393cb065eab5146a1242410f23a644851e90cfd669d2" +dependencies = [ + "parking", + "pin-project-lite", +] + +[[package]] +name = "event-listener-strategy" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8be9f3dfaaffdae2972880079a491a1a8bb7cbed0b8dd7a347f668b4150a3b93" +dependencies = [ + "event-listener", + "pin-project-lite", +] + +[[package]] +name = "fastrand" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" + +[[package]] +name = "find-msvc-tools" +version = "0.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aedcfb3409746eddb02b9e19ebda1c3394f759a152e48ee875a0844d1b955484" + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-lite" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f78e10609fe0e0b3f4157ffab1876319b5b0db102a2c60dc4626306dc46b44ad" +dependencies = [ + "fastrand", + "futures-core", + "futures-io", + "parking", + "pin-project-lite", +] + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-core", + "futures-task", + "pin-project-lite", + "slab", +] + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "hermit-abi" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17592d60ebacc7d5e169f4663c5f84f9161cc90328abcfe8456f41e4dfcb284" + +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "js-sys" +version = "0.3.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7883d941dae510fb2d978fc3fe018c71c9e2892fd38854de3e8b92c2e5ad9cc5" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + +[[package]] +name = "libc" +version = "0.2.190" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce5d3ddc6d3fa000eb1536d85e147bfe31aacaba692ed6a876f95cb7c855be78" + +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + +[[package]] +name = "mio" +version = "1.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1788edb87fdc09c7e26304471e2f5be8cdefb1b6930d6e3985fc02ff53bf86ee" +dependencies = [ + "libc", + "wasi", + "windows-sys", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "ordered-stream" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9aa2b01e1d916879f73a53d01d1d6cee68adbb31d6d9177a8cfce093cced1d50" +dependencies = [ + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "parking" +version = "2.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "piper" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c835479a4443ded371d6c535cbfd8d31ad92c5d23ae9770a61bc155e4992a3c1" +dependencies = [ + "atomic-waker", + "fastrand", + "futures-io", +] + +[[package]] +name = "polling" +version = "3.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d0e4f59085d47d8241c88ead0f274e8a0cb551f3625263c05eb8dd897c34218" +dependencies = [ + "cfg-if", + "concurrent-queue", + "hermit-abi", + "pin-project-lite", + "rustix", + "windows-sys", +] + +[[package]] +name = "proc-macro-crate" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" +dependencies = [ + "toml_edit", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "rustix" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "891efababe418670775f199f0d233d84843c227a0949a883ce15b37c78d6629d" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys", +] + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.6", +] + +[[package]] +name = "serde_json" +version = "1.0.151" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_repr" +version = "0.1.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d3b1629de253c70a0508c3899572da79ca359fdab27c7920ff00406df418906" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.6", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "socket2" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8593e8e72159ed2257d083c7a454a85cbf854f37a0966d8d483aff8c8a3ebcee" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom", + "once_cell", + "rustix", + "windows-sys", +] + +[[package]] +name = "tokio" +version = "1.53.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e95f91fcc7a621e8b030f6aa23c71fe9838ae2fb4d8118b75602a328f5144044" +dependencies = [ + "bytes", + "libc", + "mio", + "pin-project-lite", + "signal-hook-registry", + "socket2", + "tokio-macros", + "windows-sys", +] + +[[package]] +name = "tokio-macros" +version = "2.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78773a2a397f451582ce068015985c33193cf6dea8b74d2a639fe457b2f07b0e" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.6", +] + +[[package]] +name = "toml_datetime" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b86d767906c6c42421dcba507eb9d203e779497710a47782a224bb871653053" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_edit" +version = "0.25.17+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3641d5bbb5349a79e1020a242d251efbc546ad8048d133958323ce9c40a9c9c" +dependencies = [ + "indexmap", + "toml_datetime", + "toml_parser", + "winnow", +] + +[[package]] +name = "toml_parser" +version = "1.1.5+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baa693a8032d7e1cada7d0041e96126df243179ff061456783ac7f12bda4744c" +dependencies = [ + "winnow", +] + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", +] + +[[package]] +name = "uds_windows" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2f6fb2847f6742cd76af783a2a2c49e9375d0a111c7bef6f71cd9e738c72d6e" +dependencies = [ + "memoffset", + "tempfile", + "windows-sys", +] + +[[package]] +name = "unicode-ident" +version = "1.0.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d245f478577f809a851594d02313b640fb437e0bb33866753cff937863096954" + +[[package]] +name = "uuid" +version = "1.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97277d36b9c3ace13e58fa6e753f8b0bbbf302a18dd193240d70f9e29681059a" +dependencies = [ + "js-sys", + "serde_core", + "wasm-bindgen", +] + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasm-bindgen" +version = "0.2.129" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9bb54f33acc68fd454578d9820b0bde1a1a3d17aa17bb7b6595806d02886d409" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.129" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e29d0c35b16e224a7eeb5cd2d25e3e1968fbd65604117b44d3b789d00ee8535" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.129" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f501a8bc3719dba86ef8ae4728879c08001bea749eb1333ac5b91e040e2a6b7" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn 3.0.6", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.129" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23f0c9c52aa7cd7d77769a4cfe2a9adb1b331f489a41d912ce14513d5ab995c6" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "winnow" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" +dependencies = [ + "memchr", +] + +[[package]] +name = "zbus" +version = "5.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5db4be7c075cb421e4b7ee645541604239bd243ba7c357511f4ff3a74b555907" +dependencies = [ + "async-broadcast", + "async-executor", + "async-io", + "async-lock", + "async-process", + "async-recursion", + "async-task", + "async-trait", + "blocking", + "enumflags2", + "event-listener", + "futures-core", + "futures-lite", + "hex", + "libc", + "ordered-stream", + "rustix", + "serde", + "serde_repr", + "tracing", + "uds_windows", + "uuid", + "windows-sys", + "winnow", + "zbus_macros", + "zbus_names", + "zvariant", +] + +[[package]] +name = "zbus_macros" +version = "5.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2990635d09ade6df1868f72f8cac69a876a90981e8bd3c40b1be413f8dc88f40" +dependencies = [ + "proc-macro-crate", + "proc-macro2", + "quote", + "syn 3.0.6", + "zbus_names", + "zvariant", + "zvariant_utils", +] + +[[package]] +name = "zbus_names" +version = "4.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8bf88b4a3ff53e883001e0e0115b297a9d53c31b9c1edd2bfdd853e3428624e" +dependencies = [ + "serde", + "winnow", + "zvariant", +] + +[[package]] +name = "zcheapstr" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1afec51604565183aeb5c54c20aeab286120d4e4460f7f76e3e8bb8c0d99473" +dependencies = [ + "serde", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" + +[[package]] +name = "zvariant" +version = "5.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c1d34c27cc6cdd1f458427519dd6b8612f7b7e3f7b9a0b2355d041dda9869147" +dependencies = [ + "endi", + "enumflags2", + "serde", + "winnow", + "zcheapstr", + "zvariant_derive", + "zvariant_utils", +] + +[[package]] +name = "zvariant_derive" +version = "5.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "864155e69b4352db0c7f374917bf45d1e0c8d17659c8b3dbf9795f3673f8c497" +dependencies = [ + "proc-macro-crate", + "proc-macro2", + "quote", + "syn 3.0.6", + "zvariant_utils", +] + +[[package]] +name = "zvariant_utils" +version = "4.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bad0294361a320b694a328460dc73add56c306150f5cb6bfafc44446120008a3" +dependencies = [ + "proc-macro2", + "quote", + "serde", + "syn 3.0.6", + "winnow", +] diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..c5b8339 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "ai-usage" +version = "0.1.0" +edition = "2021" +description = "AI usage monitor for Hyprland" +license = "MIT" + +[dependencies] +chrono = { version = "0.4", default-features = false, features = ["clock", "std", "serde"] } +futures-util = { version = "0.3", default-features = false, features = ["std"] } +libc = "0.2" +regex = "1" +serde = { version = "1", features = ["derive"] } +serde_json = "1" +tokio = { version = "1", features = ["fs", "io-util", "macros", "net", "rt-multi-thread", "signal", "sync", "time"] } +zbus = "5" + +[profile.release] +lto = "thin" +codegen-units = 1 diff --git a/README.md b/README.md new file mode 100644 index 0000000..0791681 --- /dev/null +++ b/README.md @@ -0,0 +1,68 @@ +# AI Usage Monitor voor Linux + +Een achtergrondmonitor voor Claude, Codex en Cursor op Hyprland. De daemon haalt usage op via de bestaande CLI-logins. Waybar leest alleen de lokale cache. Er is geen venster. + +De Mac-app in de naastliggende map blijft de bron van de providerregels. Deze Linux-app volgt dezelfde vensters: Claude en Codex tonen het 5-uurs- en weekvenster, Cursor toont **Cursor models**, **Other models** en de resetdatum. + +## Wat je ziet + +De balk toont resterende percentages, bijvoorbeeld `C 38% X 71% Cu 52%`. Een klik ververst meteen. De tooltip bevat resettijden. Kleuren volgen de resterende ruimte: groen boven 50%, oranje vanaf 20%, rood daaronder. Gegevens ouder dan 30 minuten zijn stale. Een mislukte ronde bewaart de laatste succesvolle meting. + +Verversen pauzeert als logind de sessie vergrendelt of het systeem in slaap gaat. Zonder een leesbare sessie blijft verversen uit. + +## Instellingen + +`$XDG_CONFIG_HOME/ai-usage/config.json` (meestal `~/.config/ai-usage/config.json`): + +```json +{ + "refreshMinutes": 5, + "enabledProviders": ["claude", "codex", "cursor"], + "executableOverrides": {}, + "showUsedPercentage": false +} +``` + +`refreshMinutes` is 5, 15, 30 of 60. `executableOverrides` kan `claude`, `codex` of `cursor` op een absoluut pad zetten. De cache staat in `$XDG_STATE_HOME/ai-usage/usage.json`. + +De daemon zoekt `claude`, `codex` en `cursor-agent` (daarna `agent`) op `PATH`, in `~/.local/bin` en in Cursor-extensies. Zonder Agent CLI leest Cursor alleen `cursorAuth/stripeSubscriptionStatus` uit `~/.config/Cursor/User/globalStorage/state.vscdb`. + +De tools gebruiken hun eigen login. Deze app leest geen tokens en logt geen CLI-uitvoer. Cursor draait met een netwerkfilter dat alleen usage-endpoints toestaat. Op Linux mag die CLI `secret-tool lookup` gebruiken om zijn bestaande login te lezen; andere subprocessen worden geblokkeerd. + +## Nix + +```nix +{ + inputs.ai-usage.url = "path:/home/avravels/ai-usage/ai-usage-linux"; + + outputs = { self, home-manager, ai-usage, ... }: { + # in je home-manager configuratie + }; +} +``` + +```nix +imports = [ inputs.ai-usage.homeManagerModules.default ]; + +programs.ai-usage = { + enable = true; + package = inputs.ai-usage.packages.x86_64-linux.default; +}; +``` + +Dit voegt het pakket toe, start `systemd --user` service `ai-usage` met de grafische sessie, en zet `custom/ai-usage` vooraan `modules-right` als Waybar al aan staat. Je eigen Nix-configuratie blijft onaangeroerd tot je deze module importeert. + +## Handmatig + +```sh +nix build +./result/bin/ai-usage daemon +./result/bin/ai-usage status +./result/bin/ai-usage waybar +``` + +Tests: + +```sh +cargo test +``` diff --git a/fixtures/claude.json b/fixtures/claude.json new file mode 100644 index 0000000..9211709 --- /dev/null +++ b/fixtures/claude.json @@ -0,0 +1,7 @@ +{ + "rate_limits_available": true, + "rate_limits": { + "five_hour": { "utilization": 62, "resets_at": "2026-10-02T12:22:00.000Z" }, + "seven_day": { "utilization": 38, "resets_at": "2026-10-04T07:00:00Z" } + } +} diff --git a/fixtures/codex.json b/fixtures/codex.json new file mode 100644 index 0000000..0b41e49 --- /dev/null +++ b/fixtures/codex.json @@ -0,0 +1,10 @@ +{ + "rateLimits": { "limitId": "codex", "primary": { "usedPercent": 29, "windowDurationMins": 300, "resetsAt": 1790941620 }, "secondary": null }, + "rateLimitsByLimitId": { + "codex": { + "primary": { "usedPercent": 29, "windowDurationMins": 300, "resetsAt": 1790941620 }, + "secondary": { "usedPercent": 56, "windowDurationMins": 10080, "resetsAt": 1791217800 } + }, + "other_model": { "primary": { "usedPercent": 99, "windowDurationMins": 300, "resetsAt": 1790941620 } } + } +} diff --git a/fixtures/cursor-usage.txt b/fixtures/cursor-usage.txt new file mode 100644 index 0000000..9741a27 --- /dev/null +++ b/fixtures/cursor-usage.txt @@ -0,0 +1,12 @@ +Usage • Example plan Resets Oct 30 +Monthly plan and on-demand usage + +Category Current Usage +Included 44% used █████░░░░░ + Auto 48% used █████░░░░░ + API 0% used ░░░░░░░░░░ +On-Demand Disabled —————————— + +On-demand usage is off + +View in dashboard: cursor.com/dashboard?tab=usage diff --git a/fixtures/cursor.txt b/fixtures/cursor.txt new file mode 100644 index 0000000..584c847 --- /dev/null +++ b/fixtures/cursor.txt @@ -0,0 +1 @@ +active diff --git a/flake.lock b/flake.lock new file mode 100644 index 0000000..ba6568a --- /dev/null +++ b/flake.lock @@ -0,0 +1,27 @@ +{ + "nodes": { + "nixpkgs": { + "locked": { + "lastModified": 1791426274, + "narHash": "sha256-MWhktfFtFE8Jl7bf9Pkp7TR5MSdYCjCswSo0Gkb4ywc=", + "owner": "NixOS", + "repo": "nixpkgs", + "rev": "7c8764b7c7b09b34f632464276218ef9090eaa11", + "type": "github" + }, + "original": { + "owner": "NixOS", + "ref": "nixos-26.05", + "repo": "nixpkgs", + "type": "github" + } + }, + "root": { + "inputs": { + "nixpkgs": "nixpkgs" + } + } + }, + "root": "root", + "version": 7 +} diff --git a/flake.nix b/flake.nix new file mode 100644 index 0000000..c4074f1 --- /dev/null +++ b/flake.nix @@ -0,0 +1,30 @@ +{ + description = "AI Usage Monitor for Hyprland"; + + inputs.nixpkgs.url = "github:NixOS/nixpkgs/nixos-26.05"; + + outputs = { self, nixpkgs }: + let + systems = [ "x86_64-linux" "aarch64-linux" ]; + each = nixpkgs.lib.genAttrs systems; + in { + packages = each (system: + let + pkgs = import nixpkgs { inherit system; }; + in { + ai-usage = pkgs.rustPlatform.buildRustPackage { + pname = "ai-usage"; + version = "0.1.1"; + src = ./.; + cargoLock.lockFile = ./Cargo.lock; + meta = { + description = "Waybar usage monitor for Claude, Codex and Cursor"; + mainProgram = "ai-usage"; + }; + }; + default = self.packages.${system}.ai-usage; + }); + + homeManagerModules.default = ./modules/home-manager.nix; + }; +} diff --git a/modules/home-manager.nix b/modules/home-manager.nix new file mode 100644 index 0000000..29288f6 --- /dev/null +++ b/modules/home-manager.nix @@ -0,0 +1,56 @@ +{ config, lib, pkgs, ... }: +let + cfg = config.programs.ai-usage; +in { + options.programs.ai-usage = { + enable = lib.mkEnableOption "AI Usage Monitor for Waybar"; + package = lib.mkOption { + type = lib.types.package; + description = "ai-usage package from this flake."; + }; + }; + + config = lib.mkIf cfg.enable (lib.mkMerge [ + { + home.packages = [ cfg.package ]; + + systemd.user.services.ai-usage = { + Unit = { + Description = "AI Usage Monitor"; + PartOf = [ "graphical-session.target" ]; + After = [ "graphical-session.target" ]; + }; + Service = { + ExecStart = "${lib.getExe cfg.package} daemon"; + Restart = "on-failure"; + RestartSec = 2; + Environment = [ + "DBUS_SESSION_BUS_ADDRESS=unix:path=%t/bus" + "SSL_CERT_FILE=/etc/ssl/certs/ca-certificates.crt" + "PATH=${lib.makeBinPath [ pkgs.sqlite ]}:%h/.local/bin:%h/.nix-profile/bin:/etc/profiles/per-user/%u/bin:/run/current-system/sw/bin" + ]; + }; + Install.WantedBy = [ "graphical-session.target" ]; + }; + } + (lib.mkIf config.programs.waybar.enable { + programs.waybar.settings.mainBar = { + "custom/ai-usage" = { + exec = "${lib.getExe cfg.package} waybar"; + return-type = "json"; + interval = 60; + signal = 1; + on-click = "${lib.getExe cfg.package} refresh"; + tooltip = true; + }; + modules-right = lib.mkBefore [ "custom/ai-usage" ]; + }; + programs.waybar.style = lib.mkAfter '' + #custom-ai-usage.ok { color: #a6e3a1; } + #custom-ai-usage.warning { color: #fab387; } + #custom-ai-usage.critical { color: #f38ba8; } + #custom-ai-usage.unknown { color: #cdd6f4; } + ''; + }) + ]); +} diff --git a/src/activity.rs b/src/activity.rs new file mode 100644 index 0000000..534002e --- /dev/null +++ b/src/activity.rs @@ -0,0 +1,215 @@ +use std::time::Duration; + +use futures_util::{Stream, StreamExt}; +use tokio::sync::watch; +use zbus::proxy::CacheProperties; +use zbus::{proxy, Connection}; + +/// Unknown session details fail closed. Lock and sleep stay separate so a wake +/// does not clear a lock. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Activity { + pub system_awake: bool, + pub screen_awake: Option, + pub session_active: Option, + pub unlocked: Option, +} + +impl Activity { + pub fn unknown() -> Self { + Self { + system_awake: true, + screen_awake: None, + session_active: None, + unlocked: None, + } + } + + pub fn is_active(self) -> bool { + self.system_awake && self.screen_awake == Some(true) && self.session_active == Some(true) && self.unlocked == Some(true) + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LogindSample { + pub sleeping: bool, + pub locked_hint: Option, + pub state: Option, +} + +pub fn activity_from(sample: Option<&LogindSample>) -> Activity { + let Some(sample) = sample else { + return Activity::unknown(); + }; + Activity { + system_awake: !sample.sleeping, + screen_awake: Some(true), + session_active: sample + .state + .as_deref() + .map(|state| state == "active" || state == "online"), + unlocked: sample.locked_hint.map(|locked| !locked), + } +} + +pub async fn watch(sender: watch::Sender) { + loop { + if let Err(error) = follow(&sender).await { + eprintln!("ai-usage: logind unavailable ({error}); updates paused"); + let _ = sender.send(Activity::unknown()); + tokio::time::sleep(Duration::from_secs(15)).await; + } + } +} + +async fn follow(sender: &watch::Sender) -> zbus::Result<()> { + let connection = Connection::system().await?; + let session = LoginSessionProxy::builder(&connection) + .cache_properties(CacheProperties::No) + .path("/org/freedesktop/login1/session/auto")? + .build() + .await?; + let mut sleeping = false; + publish(sender, &session, sleeping).await?; + + let manager = LoginManagerProxy::builder(&connection) + .cache_properties(CacheProperties::No) + .build() + .await + .ok(); + let mut sleep_signals = match manager { + Some(manager) => match manager.receive_prepare_for_sleep().await { + Ok(stream) => Some(stream), + Err(error) => { + eprintln!("ai-usage: sleep signal unavailable ({error}); polling the session"); + None + } + }, + None => None, + }; + let mut lock_signals = session.receive_lock().await.ok(); + let mut unlock_signals = session.receive_unlock().await.ok(); + let mut locked_changed = Some(session.receive_locked_hint_changed().await); + + loop { + tokio::select! { + signal = next_item(&mut sleep_signals) => { + match signal { + Some(signal) => { + if let Ok(args) = signal.args() { + sleeping = args.start; + publish(sender, &session, sleeping).await?; + } + } + None => sleep_signals = None, + } + } + signal = next_item(&mut lock_signals) => { + if signal.is_some() { + let _ = sender.send(locked_activity(sender.borrow().clone(), false, sleeping)); + } else { + lock_signals = None; + } + } + signal = next_item(&mut unlock_signals) => { + if signal.is_some() { + publish(sender, &session, sleeping).await?; + } else { + unlock_signals = None; + } + } + change = next_item(&mut locked_changed) => { + if change.is_some() { + publish(sender, &session, sleeping).await?; + } else { + locked_changed = None; + } + } + _ = tokio::time::sleep(Duration::from_secs(5)) => { + publish(sender, &session, sleeping).await?; + } + } + } +} + +async fn next_item(stream: &mut Option) -> Option +where + S: Stream + Unpin, +{ + match stream.as_mut() { + Some(stream) => stream.next().await, + None => std::future::pending().await, + } +} + +async fn publish(sender: &watch::Sender, session: &LoginSessionProxy<'_>, sleeping: bool) -> zbus::Result<()> { + let locked_hint = session.locked_hint().await?; + let state = session.state().await.unwrap_or_else(|_| "active".to_string()); + let _ = sender.send(activity_from(Some(&LogindSample { + sleeping, + locked_hint: Some(locked_hint), + state: Some(state), + }))); + Ok(()) +} + +fn locked_activity(mut activity: Activity, unlocked: bool, sleeping: bool) -> Activity { + activity.system_awake = !sleeping; + activity.unlocked = Some(unlocked); + activity +} + +#[proxy( + interface = "org.freedesktop.login1.Manager", + default_service = "org.freedesktop.login1", + default_path = "/org/freedesktop/login1" +)] +trait LoginManager { + #[zbus(signal)] + fn prepare_for_sleep(&self, start: bool) -> zbus::Result<()>; +} + +#[proxy( + interface = "org.freedesktop.login1.Session", + default_service = "org.freedesktop.login1" +)] +trait LoginSession { + #[zbus(property)] + fn locked_hint(&self) -> zbus::Result; + + #[zbus(property)] + fn state(&self) -> zbus::Result; + + #[zbus(signal)] + fn lock(&self) -> zbus::Result<()>; + + #[zbus(signal)] + fn unlock(&self) -> zbus::Result<()>; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn unknown_and_locked_sessions_pause() { + let active = activity_from(Some(&LogindSample { + sleeping: false, + locked_hint: Some(false), + state: Some("active".to_string()), + })); + assert!(active.is_active()); + let mut sleeping = LogindSample { + sleeping: true, + locked_hint: Some(false), + state: Some("active".to_string()), + }; + assert!(!activity_from(Some(&sleeping)).is_active()); + sleeping.sleeping = false; + sleeping.locked_hint = Some(true); + assert!(!activity_from(Some(&sleeping)).is_active()); + assert!(!activity_from(None).is_active()); + sleeping.locked_hint = None; + assert!(!activity_from(Some(&sleeping)).is_active()); + } +} diff --git a/src/cursor-guard.js b/src/cursor-guard.js new file mode 100644 index 0000000..48aaef7 --- /dev/null +++ b/src/cursor-guard.js @@ -0,0 +1,91 @@ +'use strict'; +const allowed = new Set([ + '/aiserver.v1.DashboardService/GetMe', + '/aiserver.v1.DashboardService/GetTeams', + '/aiserver.v1.DashboardService/GetUserPrivacyMode', + '/aiserver.v1.DashboardService/GetTeamAdminSettingsOrEmptyIfNotInTeam', + '/aiserver.v1.DashboardService/GetCurrentPeriodUsage', + '/aiserver.v1.DashboardService/GetHardLimit', + '/aiserver.v1.DashboardService/GetPlanInfo', + '/aiserver.v1.AiService/FetchServerConfig', + '/aiserver.v1.AiService/AvailableModels', + '/aiserver.v1.AiService/GetUsableModels', + '/aiserver.v1.AiService/GetDefaultModelForCli', + '/aiserver.v1.AiService/GetTeamSettings', + '/aiserver.v1.ServerConfigService/GetServerConfig', + '/aiserver.v1.AnalyticsService/BootstrapStatsig' +]); +function check(url) { + const value = new URL(url); + if (value.protocol === 'data:' || value.protocol === 'blob:') return; + if (value.protocol !== 'https:' || value.hostname !== 'api2.cursor.sh' || + (value.port && value.port !== '443') || value.username || value.password || + value.search || !allowed.has(value.pathname)) { + throw new Error('AI Usage Monitor restricted a non-usage CLI network request'); + } +} +const originalFetch = globalThis.fetch; +globalThis.fetch = function(input, options) { + try { check(typeof input === 'string' || input instanceof URL ? input : input.url); } + catch (error) { return Promise.reject(error); } + return originalFetch.call(this, input, { ...options, redirect: 'error' }); +}; +for (const name of ['http', 'https']) { + const module = require('node:' + name); + const request = module.request; + module.request = function(input, options, callback) { + let url; + if (typeof input === 'string' || input instanceof URL) { + const parsed = new URL(input); + if (options && typeof options === 'object') { + const host = options.hostname || options.host; + if (host) parsed.host = host; + if (options.port) parsed.port = String(options.port); + if (options.path) url = parsed.origin + options.path; + } + url ||= parsed.href; + } else { + url = (input.protocol || name + ':') + '//' + (input.hostname || input.host) + + (input.port ? ':' + input.port : '') + (input.path || '/'); + } + check(url); + return request.apply(this, arguments); + }; + module.get = function() { const req = module.request.apply(this, arguments); req.end(); return req; }; +} +const http2 = require('node:http2'); +const connect = http2.connect; +http2.connect = function(authority) { + const origin = new URL(authority); + if (origin.protocol !== 'https:' || origin.hostname !== 'api2.cursor.sh' || + (origin.port && origin.port !== '443') || origin.username || origin.password) { + throw new Error('AI Usage Monitor restricted a non-usage CLI connection'); + } + const session = connect.apply(this, arguments); + const request = session.request; + session.request = function(headers) { + if (headers[':authority'] && headers[':authority'] !== origin.host) { + throw new Error('AI Usage Monitor restricted a non-usage CLI connection'); + } + check(origin.origin + headers[':path']); + return request.apply(this, arguments); + }; + return session; +}; +const child = require('node:child_process'); +function baseName(file) { + if (typeof file !== 'string') return ''; + const parts = file.split('/'); + return parts[parts.length - 1]; +} +for (const name of ['spawn', 'spawnSync', 'exec', 'execSync', 'execFile', 'execFileSync', 'fork']) { + const original = child[name]; + child[name] = function(file, args, options) { + if (name === 'spawn' && baseName(file) === 'secret-tool' && Array.isArray(args) && + args[0] === 'lookup' && !(options && options.shell)) { + return original.apply(this, arguments); + } + throw new Error('AI Usage Monitor restricted a CLI subprocess'); + }; +} +require('node:module').syncBuiltinESMExports(); diff --git a/src/daemon.rs b/src/daemon.rs new file mode 100644 index 0000000..fd518ae --- /dev/null +++ b/src/daemon.rs @@ -0,0 +1,294 @@ +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +use chrono::Utc; +use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; +use tokio::net::UnixListener; +use tokio::sync::{watch, Notify}; + +use crate::activity::{self, Activity}; +use crate::paths; +use crate::providers; +use crate::refresh::{Engine, Reason}; +use crate::settings::Settings; +use crate::waybar; + +pub async fn serve() -> std::io::Result<()> { + let settings = Settings::load_default(); + let snapshot = crate::model::UsageSnapshot::read(&paths::snapshot_path()).unwrap_or_else(|_| { + eprintln!("ai-usage: saved usage snapshot is unreadable; starting empty"); + crate::model::UsageSnapshot::default() + }); + let engine = Engine::new(snapshot, settings); + let state = Arc::new(tokio::sync::Mutex::new(Daemon::new(engine))); + let notify = Arc::new(Notify::new()); + let (activity_tx, activity_rx) = watch::channel(Activity::unknown()); + let cancel = state.lock().await.cancel.clone(); + tokio::spawn(activity::watch(activity_tx)); + tokio::spawn(listen(state.clone(), notify.clone())); + let mut activity_rx = activity_rx; + let mut terminate = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?; + loop { + let activity = *activity_rx.borrow(); + { + let mut daemon = state.lock().await; + daemon.activity = activity; + if !activity.is_active() { + daemon.cancel.store(true, Ordering::Relaxed); + } + } + if activity.is_active() { + reload_settings(&state).await; + let manual = { + let mut daemon = state.lock().await; + let manual = daemon.manual; + daemon.manual = false; + manual + }; + refresh_due(&state, if manual { Reason::Manual } else { Reason::Timer }).await; + } + let delay = { + let daemon = state.lock().await; + if daemon.activity.is_active() { + daemon + .engine + .next_delay(Utc::now()) + .map(|duration| Duration::from_millis(duration.num_milliseconds().max(10) as u64)) + .unwrap_or(Duration::from_secs(60)) + } else { + Duration::from_secs(5) + } + }; + tokio::select! { + _ = tokio::time::sleep(delay) => {} + _ = notify.notified() => {} + _ = activity_rx.changed() => {} + _ = tokio::signal::ctrl_c() => break, + _ = terminate.recv() => break, + } + let _ = cancel; + } + Ok(()) +} + +struct Daemon { + engine: Engine, + activity: Activity, + manual: bool, + cancel: Arc, + config_stamp: Option, +} + +impl Daemon { + fn new(engine: Engine) -> Self { + Self { + engine, + activity: Activity::unknown(), + manual: true, + cancel: Arc::new(AtomicBool::new(false)), + config_stamp: None, + } + } +} + +async fn reload_settings(state: &Arc>) { + let path = paths::config_path(); + let stamp = std::fs::metadata(&path).and_then(|metadata| metadata.modified()).ok(); + let mut daemon = state.lock().await; + if stamp == daemon.config_stamp && daemon.config_stamp.is_some() { + return; + } + daemon.config_stamp = stamp; + let settings = Settings::load(&path); + daemon.engine.apply_settings(settings); + let snapshot = daemon.engine.snapshot.clone(); + drop(daemon); + if snapshot.write(&paths::snapshot_path()).is_ok() { + signal_waybar(); + } +} + +async fn refresh_due(state: &Arc>, reason: Reason) { + let (round, settings, cancel) = { + let mut daemon = state.lock().await; + if !daemon.activity.is_active() { + return; + } + let now = Utc::now(); + let Some(round) = daemon.engine.select(reason, now) else { + return; + }; + daemon.cancel = Arc::new(AtomicBool::new(false)); + (round, daemon.engine.settings.clone(), daemon.cancel.clone()) + }; + let mut joins = tokio::task::JoinSet::new(); + for provider in round.providers.clone() { + let cancel = cancel.clone(); + let override_path = settings.override_for(provider).map(str::to_string); + joins.spawn_blocking(move || { + let result = providers::fetch(provider, override_path.as_deref(), &cancel); + (provider, result) + }); + } + let mut results = std::collections::HashMap::new(); + while let Some(joined) = joins.join_next().await { + match joined { + Ok((provider, result)) => { + results.insert(provider, result); + } + Err(_) => {} + } + } + let cancelled = cancel.load(Ordering::Relaxed); + let mut daemon = state.lock().await; + if cancelled || !daemon.activity.is_active() { + daemon.engine.commit(round, results, true); + return; + } + if daemon.engine.commit(round, results, false) { + let snapshot = daemon.engine.snapshot.clone(); + drop(daemon); + if let Err(error) = snapshot.write(&paths::snapshot_path()) { + eprintln!("ai-usage: could not save usage snapshot ({error:?})"); + } else { + signal_waybar(); + } + eprintln!("ai-usage: provider refresh completed"); + } +} + +async fn listen(state: Arc>, notify: Arc) { + let path = paths::socket_path(); + if let Some(parent) = path.parent() { + let _ = std::fs::create_dir_all(parent); + } + let _ = std::fs::remove_file(&path); + let listener = match UnixListener::bind(&path) { + Ok(listener) => listener, + Err(error) => { + eprintln!("ai-usage: control socket unavailable ({error})"); + return; + } + }; + loop { + let Ok((stream, _)) = listener.accept().await else { + continue; + }; + let state = state.clone(); + let notify = notify.clone(); + tokio::spawn(async move { + let _ = handle(stream, state, notify).await; + }); + } +} + +async fn handle( + stream: tokio::net::UnixStream, + state: Arc>, + notify: Arc, +) -> std::io::Result<()> { + let (read, mut write) = stream.into_split(); + let mut line = String::new(); + BufReader::new(read).read_line(&mut line).await?; + match line.trim() { + "refresh" => { + { + let mut daemon = state.lock().await; + if !daemon.activity.is_active() { + drop(daemon); + write.write_all(b"paused\n").await?; + return Ok(()); + } + daemon.manual = true; + } + notify.notify_one(); + tokio::time::sleep(Duration::from_millis(50)).await; + let started = Utc::now(); + loop { + if !state.lock().await.manual && !refreshing_recent(&state, started).await { + break; + } + if Utc::now().signed_duration_since(started).num_seconds() > 45 { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + write.write_all(b"ok\n").await?; + } + "status" => { + let daemon = state.lock().await; + let text = waybar::status_text( + &daemon.engine.snapshot, + daemon.activity.is_active(), + Utc::now(), + &daemon.engine.settings.enabled_providers, + ); + write.write_all(text.as_bytes()).await?; + write.write_all(b"\n").await?; + } + _ => write.write_all(b"unknown\n").await?, + } + Ok(()) +} + +async fn refreshing_recent(state: &Arc>, started: chrono::DateTime) -> bool { + let daemon = state.lock().await; + daemon.engine.snapshot.written_at < started +} + +fn signal_waybar() { + let Ok(entries) = std::fs::read_dir("/proc") else { + return; + }; + let signal = libc::SIGRTMIN() + 1; + for entry in entries.flatten() { + let Ok(pid) = entry.file_name().to_string_lossy().parse::() else { + continue; + }; + let Ok(comm) = std::fs::read_to_string(format!("/proc/{pid}/comm")) else { + continue; + }; + if comm.trim() == "waybar" { + unsafe { + libc::kill(pid, signal); + } + } + } +} + +pub fn request_refresh() -> std::io::Result<()> { + use std::io::{Read, Write}; + use std::os::unix::net::UnixStream; + let mut stream = UnixStream::connect(paths::socket_path()).map_err(|_| { + std::io::Error::new(std::io::ErrorKind::NotFound, "ai-usage daemon is not running") + })?; + stream.set_read_timeout(Some(Duration::from_secs(50)))?; + stream.set_write_timeout(Some(Duration::from_secs(5)))?; + stream.write_all(b"refresh\n")?; + let mut response = String::new(); + stream.read_to_string(&mut response)?; + print!("{response}"); + if response.trim() == "paused" { + std::process::exit(0); + } + if response.trim() != "ok" { + return Err(std::io::Error::other("refresh failed")); + } + signal_waybar(); + Ok(()) +} + +pub fn request_status() -> std::io::Result { + use std::io::{Read, Write}; + use std::os::unix::net::UnixStream; + let mut stream = UnixStream::connect(paths::socket_path()).map_err(|_| { + std::io::Error::new(std::io::ErrorKind::NotFound, "ai-usage daemon is not running") + })?; + stream.set_read_timeout(Some(Duration::from_secs(5)))?; + stream.write_all(b"status\n")?; + let mut response = String::new(); + stream.read_to_string(&mut response)?; + Ok(response) +} + diff --git a/src/discovery.rs b/src/discovery.rs new file mode 100644 index 0000000..e83bdab --- /dev/null +++ b/src/discovery.rs @@ -0,0 +1,227 @@ +use std::cmp::Ordering; +use std::fs; +use std::os::unix::fs::PermissionsExt; +use std::path::{Path, PathBuf}; + +use crate::model::Provider; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ExecutableCandidate { + pub path: PathBuf, + pub origin: &'static str, +} + +pub fn candidates( + provider: Provider, + override_path: Option<&str>, + home: &Path, + path: &str, +) -> Vec { + if let Some(override_path) = override_path.map(str::trim).filter(|value| !value.is_empty()) { + return vec![ExecutableCandidate { + path: expand_tilde(override_path, home), + origin: "Manual override", + }]; + } + let binaries: &[&str] = match provider { + Provider::Cursor => &["cursor-agent", "agent"], + Provider::Claude => &["claude"], + Provider::Codex => &["codex"], + }; + let mut result = Vec::new(); + for directory in path.split(':').filter(|directory| directory.starts_with('/')) { + for binary in binaries { + result.push(ExecutableCandidate { + path: PathBuf::from(directory).join(binary), + origin: "PATH", + }); + } + } + let mut local = vec![ + PathBuf::from("/usr/local/bin"), + PathBuf::from("/opt/homebrew/bin"), + home.join(".local/bin"), + home.join(".claude/local"), + home.join(".npm-global/bin"), + home.join(".nix-profile/bin"), + PathBuf::from("/run/current-system/sw/bin"), + ]; + if let Ok(user) = std::env::var("USER") { + local.push(PathBuf::from(format!("/etc/profiles/per-user/{user}/bin"))); + } + for directory in local { + for binary in binaries { + result.push(ExecutableCandidate { + path: directory.join(binary), + origin: "Local installation", + }); + } + } + if provider != Provider::Cursor { + let extensions = home.join(".cursor/extensions"); + let prefix = match provider { + Provider::Claude => "anthropic.claude-code-", + Provider::Codex => "openai.chatgpt-", + Provider::Cursor => "", + }; + let mut folders = fs::read_dir(&extensions) + .ok() + .into_iter() + .flatten() + .filter_map(|entry| entry.ok()) + .map(|entry| entry.file_name().to_string_lossy().into_owned()) + .filter(|name| name.starts_with(prefix)) + .collect::>(); + folders.sort_by(|left, right| natural_cmp(right, left)); + let suffixes: &[&str] = match provider { + Provider::Claude => &["resources/native-binary/claude"], + Provider::Codex => &[ + "bin/linux-x86_64/codex", + "bin/linux-aarch64/codex", + ], + Provider::Cursor => &[], + }; + for folder in folders { + for suffix in suffixes { + result.push(ExecutableCandidate { + path: extensions.join(&folder).join(suffix), + origin: "Cursor extension", + }); + } + } + } + let mut seen = std::collections::HashSet::new(); + result + .into_iter() + .filter(|candidate| is_executable(&candidate.path)) + .filter(|candidate| { + let key = if provider == Provider::Cursor { + fs::canonicalize(&candidate.path).unwrap_or_else(|_| candidate.path.clone()) + } else { + candidate.path.clone() + }; + seen.insert(key) + }) + .collect() +} + +fn expand_tilde(path: &str, home: &Path) -> PathBuf { + if let Some(rest) = path.strip_prefix("~/") { + home.join(rest) + } else if path == "~" { + home.to_path_buf() + } else { + PathBuf::from(path) + } +} + +fn is_executable(path: &Path) -> bool { + let Ok(metadata) = fs::metadata(path) else { + return false; + }; + metadata.is_file() && metadata.permissions().mode() & 0o111 != 0 +} + +fn natural_cmp(left: &str, right: &str) -> Ordering { + let mut left_chars = left.chars().peekable(); + let mut right_chars = right.chars().peekable(); + loop { + match (left_chars.peek().copied(), right_chars.peek().copied()) { + (None, None) => return Ordering::Equal, + (None, Some(_)) => return Ordering::Less, + (Some(_), None) => return Ordering::Greater, + (Some(left_char), Some(right_char)) if left_char.is_ascii_digit() && right_char.is_ascii_digit() => { + let left_number = take_number(&mut left_chars); + let right_number = take_number(&mut right_chars); + match left_number.cmp(&right_number) { + Ordering::Equal => {} + other => return other, + } + } + (Some(left_char), Some(right_char)) => { + left_chars.next(); + right_chars.next(); + match left_char.cmp(&right_char) { + Ordering::Equal => {} + other => return other, + } + } + } + } +} + +fn take_number(chars: &mut std::iter::Peekable>) -> u64 { + let mut number = 0u64; + while let Some(char) = chars.peek().copied() { + if let Some(digit) = char.to_digit(10) { + number = number.saturating_mul(10).saturating_add(u64::from(digit)); + chars.next(); + } else { + break; + } + } + number +} + +#[cfg(test)] +mod tests { + use super::*; + use std::os::unix::fs::PermissionsExt; + + fn temp_home() -> PathBuf { + use std::sync::atomic::{AtomicU64, Ordering}; + static NEXT: AtomicU64 = AtomicU64::new(0); + let path = std::env::temp_dir().join(format!( + "ai-usage-home-{}-{}", + std::process::id(), + NEXT.fetch_add(1, Ordering::Relaxed) + )); + let _ = fs::remove_dir_all(&path); + fs::create_dir_all(&path).unwrap(); + path + } + + #[test] + fn cursor_agent_is_discovered_outside_path_and_editor_is_ignored() { + let home = temp_home(); + let binary = home.join(".local/bin/cursor-agent"); + fs::create_dir_all(binary.parent().unwrap()).unwrap(); + fs::write(&binary, []).unwrap(); + std::os::unix::fs::symlink(&binary, home.join(".local/bin/agent")).unwrap(); + let mode = fs::Permissions::from_mode(0o755); + fs::set_permissions(&binary, mode).unwrap(); + let found = candidates(Provider::Cursor, None, &home, ""); + assert!(found.iter().any(|candidate| candidate.path == binary)); + assert!(found.iter().all(|candidate| !candidate.path.ends_with("cursor"))); + assert!(found.iter().all(|candidate| candidate.path.file_name().unwrap() != "agent")); + let _ = fs::remove_dir_all(home); + } + + #[test] + fn missing_manual_override_is_kept() { + let found = candidates(Provider::Claude, Some("/does/not/exist"), Path::new("/tmp"), ""); + assert_eq!(found.len(), 1); + assert_eq!(found[0].origin, "Manual override"); + } + + #[test] + fn newest_extension_wins_when_there_is_no_path_binary() { + let home = temp_home(); + for version in ["2.1.9", "2.1.10"] { + let executable = home.join(format!( + ".cursor/extensions/anthropic.claude-code-{version}-linux-x64/resources/native-binary/claude" + )); + fs::create_dir_all(executable.parent().unwrap()).unwrap(); + fs::write(&executable, []).unwrap(); + fs::set_permissions(&executable, fs::Permissions::from_mode(0o755)).unwrap(); + } + let found = candidates(Provider::Claude, None, &home, ""); + let extensions: Vec<_> = found + .iter() + .filter(|candidate| candidate.origin == "Cursor extension") + .collect(); + assert_eq!(extensions.len(), 2); + assert!(extensions[0].path.display().to_string().contains("2.1.10")); + let _ = fs::remove_dir_all(home); + } +} diff --git a/src/formatting.rs b/src/formatting.rs new file mode 100644 index 0000000..ba63410 --- /dev/null +++ b/src/formatting.rs @@ -0,0 +1,161 @@ +use chrono::{DateTime, Datelike, Timelike, Utc}; + +use crate::model::UsageWindow; + +pub fn percentage(window: Option<&UsageWindow>, show_used: bool) -> String { + let Some(window) = window else { + return "—".to_string(); + }; + let value = if show_used { + window.used_percentage + } else { + window.remaining_percentage + }; + match value { + Some(value) => format!( + "{}% {}", + value.round() as i64, + if show_used { "used" } else { "left" } + ), + None => "—".to_string(), + } +} + +pub fn bar_number(window: Option<&UsageWindow>, show_used: bool) -> String { + let Some(window) = window else { + return "—".to_string(); + }; + let value = if show_used { + window.used_percentage + } else { + window.remaining_percentage + }; + match value { + Some(value) => format!("{}%", value.round() as i64), + None => "—".to_string(), + } +} + +pub fn reset(date: Option>, now: DateTime, compact: bool) -> String { + let Some(date) = date else { + return "Reset unknown".to_string(); + }; + let local_date = date.with_timezone(&chrono::Local); + let local_now = now.with_timezone(&chrono::Local); + let clock = format!("{:02}:{:02}", local_date.hour(), local_date.minute()); + if local_date.date_naive() == local_now.date_naive() { + return if compact { + clock + } else { + format!("Today {clock}") + }; + } + let weekday = weekday_name(local_date.weekday()); + if compact { + format!("{weekday} {clock}") + } else { + format!( + "{weekday} {} {} {clock}", + local_date.day(), + month_name(local_date.month()) + ) + } +} + +pub fn relative_reset(date: Option>, now: DateTime) -> String { + let Some(date) = date else { + return "Reset unknown".to_string(); + }; + let seconds = (date - now).num_seconds(); + if seconds <= 0 { + return "Awaiting reset update".to_string(); + } + let minutes = (seconds + 59) / 60; + if minutes >= 1440 { + return format!("Reset in {}d {}h", minutes / 1440, (minutes % 1440) / 60); + } + if minutes >= 60 { + return format!("Reset in {}h {}m", minutes / 60, minutes % 60); + } + format!("Reset in {minutes}m") +} + +pub fn clock(date: Option>) -> String { + let Some(date) = date else { + return "Never".to_string(); + }; + let local = date.with_timezone(&chrono::Local); + format!("{:02}:{:02}", local.hour(), local.minute()) +} + +pub fn parse_iso8601(value: &str) -> Option> { + DateTime::parse_from_rfc3339(value) + .ok() + .map(|parsed| parsed.with_timezone(&Utc)) +} + +fn weekday_name(weekday: chrono::Weekday) -> &'static str { + match weekday { + chrono::Weekday::Mon => "Mon", + chrono::Weekday::Tue => "Tue", + chrono::Weekday::Wed => "Wed", + chrono::Weekday::Thu => "Thu", + chrono::Weekday::Fri => "Fri", + chrono::Weekday::Sat => "Sat", + chrono::Weekday::Sun => "Sun", + } +} + +fn month_name(month: u32) -> &'static str { + match month { + 1 => "Jan", + 2 => "Feb", + 3 => "Mar", + 4 => "Apr", + 5 => "May", + 6 => "Jun", + 7 => "Jul", + 8 => "Aug", + 9 => "Sep", + 10 => "Oct", + 11 => "Nov", + _ => "Dec", + } +} + +pub fn level(remaining: Option) -> &'static str { + match remaining { + Some(remaining) if remaining > 50.0 => "ok", + Some(remaining) if remaining >= 20.0 => "warning", + Some(_) => "critical", + None => "unknown", + } +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::TimeZone; + + #[test] + fn relative_reset_rounds_up_and_rejects_the_past() { + let now = Utc.with_ymd_and_hms(2026, 10, 9, 12, 0, 0).unwrap(); + assert_eq!( + relative_reset(Some(now + chrono::Duration::seconds(90)), now), + "Reset in 2m" + ); + assert_eq!( + relative_reset(Some(now - chrono::Duration::seconds(1)), now), + "Awaiting reset update" + ); + assert_eq!(relative_reset(None, now), "Reset unknown"); + } + + #[test] + fn percentage_uses_the_selected_side() { + let window = UsageWindow::new(Some(62.0), None, None); + assert_eq!(percentage(Some(&window), false), "38% left"); + assert_eq!(bar_number(Some(&window), true), "62%"); + assert_eq!(percentage(None, false), "—"); + } +} diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..2295b3c --- /dev/null +++ b/src/main.rs @@ -0,0 +1,66 @@ +mod activity; +mod daemon; +mod discovery; +mod formatting; +mod model; +mod parsers; +mod paths; +mod process; +mod protocol; +mod providers; +mod refresh; +mod settings; +mod store; +mod waybar; + +use chrono::Utc; + +#[tokio::main] +async fn main() { + let command = std::env::args().nth(1); + match command.as_deref() { + Some("daemon") => { + if let Err(error) = daemon::serve().await { + eprintln!("ai-usage: {error}"); + std::process::exit(1); + } + } + Some("waybar") => { + let snapshot = model::UsageSnapshot::read(&paths::snapshot_path()).unwrap_or_default(); + let settings = settings::Settings::load(&paths::config_path()); + println!( + "{}", + waybar::json(&waybar::render(&snapshot, Utc::now(), &settings.enabled_providers)) + ); + } + Some("refresh") => { + if let Err(error) = daemon::request_refresh() { + eprintln!("ai-usage: {error}"); + std::process::exit(1); + } + } + Some("status") => match daemon::request_status() { + Ok(text) => print!("{text}"), + Err(error) => { + let snapshot = model::UsageSnapshot::read(&paths::snapshot_path()).unwrap_or_default(); + let settings = settings::Settings::load(&paths::config_path()); + println!( + "{}", + waybar::status_text(&snapshot, false, Utc::now(), &settings.enabled_providers) + ); + eprintln!("ai-usage: {error}"); + std::process::exit(1); + } + }, + _ => { + eprintln!( + "Usage: ai-usage \n\n\ + daemon refresh Claude, Codex and Cursor in the background\n\ + waybar print one Waybar JSON line from the saved snapshot\n\ + refresh ask the daemon to refresh now\n\ + status print the current snapshot" + ); + std::process::exit(2); + } + } +} diff --git a/src/model.rs b/src/model.rs new file mode 100644 index 0000000..db13980 --- /dev/null +++ b/src/model.rs @@ -0,0 +1,355 @@ +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; + +pub const STALE_AFTER_SECS: i64 = 30 * 60; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum Provider { + Claude, + Codex, + Cursor, +} + +impl Provider { + pub const ALL: [Provider; 3] = [Provider::Claude, Provider::Codex, Provider::Cursor]; + + pub fn as_str(self) -> &'static str { + match self { + Provider::Claude => "claude", + Provider::Codex => "codex", + Provider::Cursor => "cursor", + } + } + + pub fn display_name(self) -> &'static str { + match self { + Provider::Claude => "Claude", + Provider::Codex => "Codex", + Provider::Cursor => "Cursor", + } + } + + pub fn bar_label(self) -> &'static str { + match self { + Provider::Claude => "C", + Provider::Codex => "X", + Provider::Cursor => "Cu", + } + } +} + +fn clamp(value: Option) -> Option { + let value = value?; + if !value.is_finite() { + return None; + } + Some(value.clamp(0.0, 100.0)) +} + +#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct UsageWindow { + pub used_percentage: Option, + pub remaining_percentage: Option, + #[serde(default, with = "optional_rfc3339")] + pub reset_at: Option>, +} + +impl UsageWindow { + pub fn new( + used_percentage: Option, + remaining_percentage: Option, + reset_at: Option>, + ) -> Self { + let used = clamp(used_percentage); + let remaining = clamp(remaining_percentage); + Self { + used_percentage: used.or_else(|| remaining.map(|value| 100.0 - value)), + remaining_percentage: remaining.or_else(|| used.map(|value| 100.0 - value)), + reset_at, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct UsageMetric { + pub name: String, + pub value: String, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct BillingUsagePool { + pub name: String, + pub window: UsageWindow, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ProviderUsage { + pub provider: Provider, + #[serde(with = "rfc3339")] + pub fetched_at: DateTime, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub short_window: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub weekly_window: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub billing_window: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub billing_reset_label: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub billing_pools: Option>, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub additional_usage: Vec, +} + +impl ProviderUsage { + pub fn display_billing_pools(&self) -> Vec { + if let Some(pools) = &self.billing_pools { + return pools.clone(); + } + if self.provider != Provider::Cursor { + return Vec::new(); + } + [("Auto", "Cursor models"), ("API", "Other models")] + .into_iter() + .filter_map(|(old_name, name)| { + let metric = self.additional_usage.iter().find(|metric| { + metric.name == old_name || metric.name == name + })?; + let value = metric.value.strip_suffix("% used")?; + if !regex_percent(value) { + return None; + } + let used = value.parse::().ok()?; + if !used.is_finite() { + return None; + } + Some(BillingUsagePool { + name: name.to_string(), + window: UsageWindow::new(Some(used), None, None), + }) + }) + .collect() + } + + pub fn display_additional_usage(&self) -> Vec { + if self.provider != Provider::Cursor { + return self.additional_usage.clone(); + } + let pools: Vec = self + .display_billing_pools() + .into_iter() + .map(|pool| pool.name) + .collect(); + self.additional_usage + .iter() + .filter(|metric| { + let name = match metric.name.as_str() { + "Auto" => "Cursor models", + "API" => "Other models", + other => other, + }; + !pools.iter().any(|pool| pool == name) + }) + .cloned() + .collect() + } +} + +fn regex_percent(value: &str) -> bool { + let mut seen_dot = false; + let mut digits = false; + for char in value.chars() { + if char.is_ascii_digit() { + digits = true; + continue; + } + if char == '.' && !seen_dot && digits { + seen_dot = true; + digits = false; + continue; + } + return false; + } + digits +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum FetchStatus { + Available, + Disabled, + NotSignedIn, + Unsupported, + Failed, + Pending, +} + +impl FetchStatus { + pub fn label(self) -> &'static str { + match self { + FetchStatus::Available => "Available", + FetchStatus::Disabled => "Disabled", + FetchStatus::NotSignedIn => "Sign in with the provider", + FetchStatus::Unsupported => "Usage unavailable", + FetchStatus::Failed => "Update failed", + FetchStatus::Pending => "Awaiting first update", + } + } +} + +#[derive(Debug, Clone, PartialEq)] +pub struct ProviderFailure { + pub status: FetchStatus, + pub message: String, +} + +impl ProviderFailure { + pub fn new(status: FetchStatus, message: impl Into) -> Self { + Self { + status, + message: message.into(), + } + } +} + +impl std::fmt::Display for ProviderFailure { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(&self.message) + } +} + +impl std::error::Error for ProviderFailure {} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ProviderState { + pub provider: Provider, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub last_successful_usage: Option, + #[serde(default, with = "optional_rfc3339")] + pub last_attempt_at: Option>, + pub status: FetchStatus, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub message: Option, +} + +impl ProviderState { + pub fn pending(provider: Provider) -> Self { + Self { + provider, + last_successful_usage: None, + last_attempt_at: None, + status: FetchStatus::Pending, + message: None, + } + } + + pub fn is_stale(&self, at: DateTime) -> bool { + let Some(usage) = &self.last_successful_usage else { + return false; + }; + (at - usage.fetched_at).num_seconds() > STALE_AFTER_SECS + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct UsageSnapshot { + pub schema_version: u32, + #[serde(with = "rfc3339")] + pub written_at: DateTime, + pub show_used_percentage: bool, + pub is_demo: bool, + pub providers: Vec, +} + +impl Default for UsageSnapshot { + fn default() -> Self { + Self { + schema_version: 1, + written_at: Utc::now(), + show_used_percentage: false, + is_demo: false, + providers: Provider::ALL + .into_iter() + .map(ProviderState::pending) + .collect(), + } + } +} + +impl UsageSnapshot { + pub fn state(&self, provider: Provider) -> ProviderState { + self.providers + .iter() + .find(|state| state.provider == provider) + .cloned() + .unwrap_or_else(|| ProviderState::pending(provider)) + } + + pub fn state_mut(&mut self, provider: Provider) -> &mut ProviderState { + if let Some(index) = self + .providers + .iter() + .position(|state| state.provider == provider) + { + return &mut self.providers[index]; + } + self.providers.push(ProviderState::pending(provider)); + self.providers.last_mut().expect("provider just inserted") + } +} + +mod rfc3339 { + use chrono::{DateTime, Utc}; + use serde::{Deserialize, Deserializer, Serializer}; + + pub fn serialize(value: &DateTime, serializer: S) -> Result + where + S: Serializer, + { + serializer.serialize_str(&value.to_rfc3339_opts(chrono::SecondsFormat::Secs, true)) + } + + pub fn deserialize<'de, D>(deserializer: D) -> Result, D::Error> + where + D: Deserializer<'de>, + { + let text = String::deserialize(deserializer)?; + DateTime::parse_from_rfc3339(&text) + .map(|value| value.with_timezone(&Utc)) + .map_err(serde::de::Error::custom) + } +} + +mod optional_rfc3339 { + use chrono::{DateTime, Utc}; + use serde::{Deserialize, Deserializer, Serializer}; + + pub fn serialize(value: &Option>, serializer: S) -> Result + where + S: Serializer, + { + match value { + Some(value) => serializer.serialize_some(&value.to_rfc3339_opts(chrono::SecondsFormat::Secs, true)), + None => serializer.serialize_none(), + } + } + + pub fn deserialize<'de, D>(deserializer: D) -> Result>, D::Error> + where + D: Deserializer<'de>, + { + let value = Option::::deserialize(deserializer)?; + value + .map(|text| { + DateTime::parse_from_rfc3339(&text) + .map(|parsed| parsed.with_timezone(&Utc)) + .map_err(serde::de::Error::custom) + }) + .transpose() + } +} diff --git a/src/parsers.rs b/src/parsers.rs new file mode 100644 index 0000000..a5eb156 --- /dev/null +++ b/src/parsers.rs @@ -0,0 +1,442 @@ +use chrono::{DateTime, TimeZone, Utc}; +use serde::Deserialize; + +use crate::formatting::parse_iso8601; +use crate::model::{ + BillingUsagePool, FetchStatus, Provider, ProviderFailure, ProviderUsage, UsageMetric, UsageWindow, +}; + +pub fn parse_claude(data: &[u8], fetched_at: DateTime) -> Result { + let response: ClaudeResponse = serde_json::from_slice(data).map_err(|_| { + ProviderFailure::new( + FetchStatus::Unsupported, + "Claude's usage protocol changed. Check the detected CLI version.", + ) + })?; + if response.rate_limits_available == Some(false) { + return Err(ProviderFailure::new( + FetchStatus::NotSignedIn, + "Subscription limits are unavailable. Check Claude's existing subscription login.", + )); + } + let short = normalize_claude(response.rate_limits.as_ref().and_then(|limits| limits.five_hour.as_ref()))?; + let week = normalize_claude(response.rate_limits.as_ref().and_then(|limits| limits.seven_day.as_ref()))?; + if short.as_ref().and_then(|window| window.used_percentage).is_none() + && week.as_ref().and_then(|window| window.used_percentage).is_none() + { + return Err(ProviderFailure::new( + FetchStatus::Failed, + "Claude returned no subscription usage. Open Claude Code and check /usage, then refresh. Last successful usage is retained.", + )); + } + Ok(ProviderUsage { + provider: Provider::Claude, + fetched_at, + short_window: short, + weekly_window: week, + billing_window: None, + billing_reset_label: None, + billing_pools: None, + additional_usage: Vec::new(), + }) +} + +fn normalize_claude(window: Option<&ClaudeWindow>) -> Result, ProviderFailure> { + let Some(window) = window else { + return Ok(None); + }; + let reset_at = match &window.resets_at { + Some(text) => Some(parse_iso8601(text).ok_or_else(|| { + ProviderFailure::new(FetchStatus::Failed, "Claude returned an unreadable reset timestamp.") + })?), + None => None, + }; + Ok(Some(UsageWindow::new(window.utilization, None, reset_at))) +} + +#[derive(Deserialize)] +struct ClaudeResponse { + rate_limits_available: Option, + rate_limits: Option, +} + +#[derive(Deserialize)] +struct ClaudeLimits { + five_hour: Option, + seven_day: Option, +} + +#[derive(Deserialize)] +struct ClaudeWindow { + utilization: Option, + resets_at: Option, +} + +pub fn parse_codex(data: &[u8], fetched_at: DateTime) -> Result { + let response: CodexResponse = serde_json::from_slice(data).map_err(|_| { + ProviderFailure::new( + FetchStatus::Unsupported, + "Codex's usage protocol changed. Check the detected CLI version.", + ) + })?; + let bucket = response + .rate_limits_by_limit_id + .as_ref() + .and_then(|buckets| buckets.get("codex")) + .or_else(|| { + response.rate_limits.as_ref().filter(|bucket| { + bucket.limit_id.as_deref().is_none_or(|id| id == "codex") + }) + }); + let Some(bucket) = bucket else { + return Err(ProviderFailure::new( + FetchStatus::Unsupported, + "Codex returned no subscription rate-limit bucket.", + )); + }; + let mut short = None; + let mut week = None; + let mut metrics = Vec::new(); + for window in [bucket.primary.as_ref(), bucket.secondary.as_ref()] + .into_iter() + .flatten() + { + let normalized = UsageWindow::new( + Some(window.used_percent), + None, + window.resets_at.map(|seconds| Utc.timestamp_opt(seconds as i64, 0).single().unwrap_or_else(Utc::now)), + ); + match window.window_duration_mins { + Some(300) => short = Some(normalized), + Some(10080) => week = Some(normalized), + Some(minutes) => metrics.push(UsageMetric { + name: format!("{minutes}m window"), + value: crate::formatting::percentage(Some(&normalized), false), + }), + None => metrics.push(UsageMetric { + name: "Unspecified window".to_string(), + value: crate::formatting::percentage(Some(&normalized), false), + }), + } + } + if short.is_none() && week.is_none() && metrics.is_empty() { + return Err(ProviderFailure::new( + FetchStatus::Failed, + "Codex returned no usage windows.", + )); + } + Ok(ProviderUsage { + provider: Provider::Codex, + fetched_at, + short_window: short, + weekly_window: week, + billing_window: None, + billing_reset_label: None, + billing_pools: None, + additional_usage: metrics, + }) +} + +#[derive(Deserialize)] +struct CodexResponse { + #[serde(rename = "rateLimits")] + rate_limits: Option, + #[serde(rename = "rateLimitsByLimitId")] + rate_limits_by_limit_id: Option>, +} + +#[derive(Deserialize)] +struct CodexBucket { + #[serde(rename = "limitId")] + limit_id: Option, + primary: Option, + secondary: Option, +} + +#[derive(Deserialize)] +struct CodexWindow { + #[serde(rename = "usedPercent")] + used_percent: f64, + #[serde(rename = "windowDurationMins")] + window_duration_mins: Option, + #[serde(rename = "resetsAt")] + resets_at: Option, +} + +pub fn terminal_text(data: &[u8]) -> String { + let text = String::from_utf8_lossy(data); + let without_osc = OSC_RE.replace_all(&text, ""); + let without_csi = CSI_RE.replace_all(&without_osc, ""); + without_csi.replace('\r', "") +} + +pub fn parse_cursor(data: &[u8], fetched_at: DateTime) -> Result { + let text = terminal_text(data); + let Some(header) = text.rfind("Monthly plan and on-demand usage") else { + return Err(ProviderFailure::new( + FetchStatus::Unsupported, + "Cursor did not return a supported billing-usage screen.", + )); + }; + let before = &text[..header]; + let screen = &text[header + "Monthly plan and on-demand usage".len()..]; + let included = percent("Included", screen)?; + let auto = percent("Auto", screen)?; + let api = percent("API", screen)?; + let on_demand = capture(ON_DEMAND, screen).ok_or_else(|| { + ProviderFailure::new( + FetchStatus::Unsupported, + "Cursor's usage screen is incomplete or changed. Check the Agent CLI version.", + ) + })?; + if included.is_none() || auto.is_none() || api.is_none() { + return Err(ProviderFailure::new( + FetchStatus::Unsupported, + "Cursor's usage screen is incomplete or changed. Check the Agent CLI version.", + )); + } + let reset_source = before + .lines() + .rev() + .take(3) + .collect::>() + .into_iter() + .rev() + .collect::>() + .join("\n"); + let reset = capture(RESET, &reset_source); + Ok(ProviderUsage { + provider: Provider::Cursor, + fetched_at, + short_window: None, + weekly_window: None, + billing_window: Some(UsageWindow::new(included, None, None)), + billing_reset_label: reset, + billing_pools: Some(vec![ + BillingUsagePool { + name: "Cursor models".to_string(), + window: UsageWindow::new(auto, None, None), + }, + BillingUsagePool { + name: "Other models".to_string(), + window: UsageWindow::new(api, None, None), + }, + ]), + additional_usage: vec![UsageMetric { + name: "On-demand".to_string(), + value: on_demand, + }], + }) +} + +fn percent(label: &str, screen: &str) -> Result, ProviderFailure> { + let pattern = format!(r"(?m)^\s*{label}\s+([0-9]+(?:\.[0-9]+)?)% used\b"); + let regex = regex::Regex::new(&pattern).map_err(|_| { + ProviderFailure::new(FetchStatus::Failed, "Cursor usage could not be read.") + })?; + let Some(found) = regex.captures(screen).and_then(|caps| caps.get(1)) else { + return Ok(None); + }; + let number = found.as_str().parse::().ok().filter(|value| value.is_finite()); + Ok(number) +} + +fn capture(pattern: &str, value: &str) -> Option { + regex::Regex::new(pattern) + .ok()? + .captures(value)? + .get(1) + .map(|item| item.as_str().to_string()) +} + +const ON_DEMAND: &str = r"(?m)^[ \t]*On-Demand[ \t]+(Disabled|Unavailable|\$[0-9,]+(?:\.[0-9]{1,2})?(?: / \$[0-9,]+(?:\.[0-9]{1,2})?)?)(?:[ \t]+[█░—]+)?[ \t]*\n"; +const RESET: &str = r"Resets ((?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec) (?:[1-9]|[12][0-9]|3[01]))\b"; + +macro_rules! lazy_static_regex { + ($name:ident, $pattern:expr) => { + static $name: std::sync::LazyLock = + std::sync::LazyLock::new(|| regex::Regex::new($pattern).expect("static regex")); + }; +} +pub fn subscription_status(text: &str) -> Option { + let status = text.trim(); + const KNOWN: &[&str] = &[ + "active", + "trialing", + "past_due", + "canceled", + "unpaid", + "incomplete", + "incomplete_expired", + "paused", + ]; + if !KNOWN.contains(&status) { + return None; + } + Some(UsageMetric { + name: "Subscription".to_string(), + value: status.replace('_', " "), + }) +} + +lazy_static_regex!(OSC_RE, r"\x1B\][\s\S]*?(?:\x07|\x1B\\)"); +lazy_static_regex!(CSI_RE, r"\x1B\[[0-?]*[ -/]*[@-~]"); + +#[cfg(test)] +mod tests { + use super::*; + + fn fixture(name: &str) -> Vec { + match name { + "claude" => include_bytes!("../fixtures/claude.json").to_vec(), + "codex" => include_bytes!("../fixtures/codex.json").to_vec(), + "cursor" => include_bytes!("../fixtures/cursor.txt").to_vec(), + "cursor-usage" => include_bytes!("../fixtures/cursor-usage.txt").to_vec(), + _ => panic!("unknown fixture"), + } + } + + #[test] + fn claude_windows_and_fractional_timestamp() { + let usage = parse_claude(&fixture("claude"), Utc::now()).unwrap(); + assert_eq!(usage.short_window.unwrap().remaining_percentage, Some(38.0)); + assert_eq!(usage.weekly_window.unwrap().remaining_percentage, Some(62.0)); + assert_eq!( + usage.short_window.unwrap().reset_at, + parse_iso8601("2026-10-02T12:22:00Z") + ); + } + + #[test] + fn claude_missing_limits_do_not_fabricate_zero() { + let error = parse_claude(br#"{"rate_limits_available":true,"rate_limits":null}"#, Utc::now()).unwrap_err(); + assert_eq!(error.status, FetchStatus::Failed); + } + + #[test] + fn claude_non_subscription_auth() { + let error = parse_claude(br#"{"rate_limits_available":false}"#, Utc::now()).unwrap_err(); + assert_eq!(error.status, FetchStatus::NotSignedIn); + } + + #[test] + fn claude_invalid_reset_is_an_error() { + let json = br#"{"rate_limits":{"five_hour":{"utilization":42,"resets_at":"tomorrow"}}}"#; + assert!(parse_claude(json, Utc::now()).is_err()); + } + + #[test] + fn claude_missing_reset_remains_unknown() { + let json = br#"{"rate_limits":{"five_hour":{"utilization":42}}}"#; + let usage = parse_claude(json, Utc::now()).unwrap(); + assert_eq!(usage.short_window.unwrap().remaining_percentage, Some(58.0)); + assert!(usage.short_window.unwrap().reset_at.is_none()); + assert!(usage.weekly_window.is_none()); + } + + #[test] + fn codex_uses_matching_bucket_and_actual_durations() { + let usage = parse_codex(&fixture("codex"), Utc::now()).unwrap(); + assert_eq!(usage.short_window.as_ref().unwrap().remaining_percentage, Some(71.0)); + assert_eq!(usage.weekly_window.as_ref().unwrap().remaining_percentage, Some(44.0)); + assert_eq!( + usage.short_window.unwrap().reset_at.unwrap().timestamp(), + 1_790_941_620 + ); + } + + #[test] + fn codex_reversed_windows() { + let json = br#"{"rateLimits":{"primary":{"usedPercent":42,"windowDurationMins":10080},"secondary":{"usedPercent":7,"windowDurationMins":300}}}"#; + let usage = parse_codex(json, Utc::now()).unwrap(); + assert_eq!(usage.short_window.unwrap().used_percentage, Some(7.0)); + assert_eq!(usage.weekly_window.unwrap().used_percentage, Some(42.0)); + } + + #[test] + fn codex_unknown_window_is_not_called_five_hours() { + let json = br#"{"rateLimits":{"primary":{"usedPercent":25,"windowDurationMins":15}}}"#; + let usage = parse_codex(json, Utc::now()).unwrap(); + assert!(usage.short_window.is_none()); + assert_eq!(usage.additional_usage[0].name, "15m window"); + } + + #[test] + fn codex_missing_duration_is_not_guessed() { + let json = br#"{"rateLimits":{"primary":{"usedPercent":25}}}"#; + let usage = parse_codex(json, Utc::now()).unwrap(); + assert!(usage.short_window.is_none()); + assert_eq!(usage.additional_usage[0].name, "Unspecified window"); + } + + #[test] + fn codex_wrong_bucket_is_unsupported() { + let json = br#"{"rateLimits":{"limitId":"other","primary":{"usedPercent":25,"windowDurationMins":300}}}"#; + let error = parse_codex(json, Utc::now()).unwrap_err(); + assert_eq!(error.status, FetchStatus::Unsupported); + } + + #[test] + fn malformed_data_is_rejected() { + assert!(parse_claude(b"broken", Utc::now()).is_err()); + assert!(parse_codex(b"broken", Utc::now()).is_err()); + } + + #[test] + fn cursor_status_is_not_an_allowance() { + let text = String::from_utf8(fixture("cursor")).unwrap(); + let metric = subscription_status(&text).unwrap(); + assert_eq!(metric.name, "Subscription"); + assert_eq!(metric.value, "active"); + assert!(subscription_status("unexpected raw value").is_none()); + } + + #[test] + fn cursor_billing_usage_keeps_reset_label_and_independent_categories() { + let usage = parse_cursor(&fixture("cursor-usage"), Utc::now()).unwrap(); + assert_eq!(usage.billing_window.unwrap().remaining_percentage, Some(56.0)); + assert!(usage.billing_window.unwrap().reset_at.is_none()); + assert!(usage.short_window.is_none()); + assert!(usage.weekly_window.is_none()); + assert_eq!(usage.billing_reset_label.as_deref(), Some("Oct 30")); + let pools = usage.billing_pools.unwrap(); + assert_eq!(pools[0].name, "Cursor models"); + assert_eq!(pools[0].window.used_percentage, Some(48.0)); + assert_eq!(pools[1].name, "Other models"); + assert_eq!(pools[1].window.used_percentage, Some(0.0)); + assert_eq!(usage.additional_usage[0].value, "Disabled"); + } + + #[test] + fn cursor_ansi_and_private_metadata_are_not_persisted() { + let screen = fixture("cursor-usage"); + let mut text = b"\x1b]0;private@example.invalid\x07\x1b[2J\x1b[36m".to_vec(); + text.extend_from_slice(&screen); + text.extend_from_slice(b"\x1b[0m"); + let usage = parse_cursor(&text, Utc::now()).unwrap(); + assert_eq!(usage.billing_window.unwrap().used_percentage, Some(44.0)); + let encoded = serde_json::to_string(&usage).unwrap(); + assert!(!encoded.contains("private")); + } + + #[test] + fn cursor_incomplete_changed_and_enterprise_screens_rejected() { + let screen = String::from_utf8(fixture("cursor-usage")).unwrap(); + let changed = screen.replace("Included", "New label"); + let private_row = screen.replace("On-Demand Disabled", "On-Demand private-account"); + for invalid in ["active", "Personal spend for the current billing period", changed.as_str(), private_row.as_str()] { + assert!(parse_cursor(invalid.as_bytes(), Utc::now()).is_err(), "{invalid}"); + } + } + + #[test] + fn cursor_waits_for_full_on_demand_row_and_accepts_spend_limits() { + let data = String::from_utf8(fixture("cursor-usage")).unwrap(); + let prefix = data.find("On-Demand").unwrap(); + let partial = format!("{}On-Demand $0.75 ", &data[..prefix]); + assert!(parse_cursor(partial.as_bytes(), Utc::now()).is_err()); + let complete = format!("{partial}/ $40.00 █░░░░\n"); + let usage = parse_cursor(complete.as_bytes(), Utc::now()).unwrap(); + assert_eq!(usage.additional_usage.last().unwrap().value, "$0.75 / $40.00"); + } +} diff --git a/src/paths.rs b/src/paths.rs new file mode 100644 index 0000000..b323349 --- /dev/null +++ b/src/paths.rs @@ -0,0 +1,36 @@ +use std::path::PathBuf; + +pub fn home() -> PathBuf { + std::env::var_os("HOME") + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from("/")) +} + +pub fn config_dir() -> PathBuf { + std::env::var_os("XDG_CONFIG_HOME") + .map(PathBuf::from) + .unwrap_or_else(|| home().join(".config")) + .join("ai-usage") +} + +pub fn config_path() -> PathBuf { + config_dir().join("config.json") +} + +pub fn state_dir() -> PathBuf { + std::env::var_os("XDG_STATE_HOME") + .map(PathBuf::from) + .unwrap_or_else(|| home().join(".local/state")) + .join("ai-usage") +} + +pub fn snapshot_path() -> PathBuf { + state_dir().join("usage.json") +} + +pub fn socket_path() -> PathBuf { + let runtime = std::env::var_os("XDG_RUNTIME_DIR") + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from(format!("/run/user/{}", unsafe { libc::getuid() }))); + runtime.join("ai-usage.sock") +} diff --git a/src/process.rs b/src/process.rs new file mode 100644 index 0000000..661cfa6 --- /dev/null +++ b/src/process.rs @@ -0,0 +1,346 @@ +use std::io::Write; +use std::os::fd::{AsRawFd, FromRawFd, OwnedFd}; +use std::os::unix::process::CommandExt; +use std::path::Path; +use std::process::{Command, Stdio}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::{Duration, Instant}; + +#[derive(Debug)] +pub struct ProcessResult { + pub stdout: Vec, + pub stderr: Vec, + pub exit_status: i32, +} + +#[derive(Debug, PartialEq, Eq)] +pub enum ProcessError { + Launch, + TimedOut, + Cancelled, + OutputLimit, + Io, +} + +#[derive(Debug, PartialEq, Eq)] +pub enum StreamAction { + None, + Send(Vec), + Finish, +} + +pub trait ProcessHooks { + fn on_output(&mut self, chunk: &[u8], at: Duration) -> StreamAction; + fn on_poll(&mut self, at: Duration) -> StreamAction; +} + +pub struct CommandSpec<'a> { + pub executable: &'a Path, + pub arguments: &'a [String], + pub environment: &'a [(String, String)], + pub working_directory: Option<&'a Path>, + pub timeout: Duration, + pub terminal: bool, + pub initial_input: &'a [u8], + pub cancel: &'a AtomicBool, +} + +struct Reader { + fd: OwnedFd, + open: bool, + is_stdout: bool, +} + +pub fn run(spec: CommandSpec<'_>, hooks: Option<&mut dyn ProcessHooks>) -> Result { + if spec.cancel.load(Ordering::Relaxed) { + return Err(ProcessError::Cancelled); + } + let mut command = Command::new(spec.executable); + command.args(spec.arguments); + command.env_clear(); + for (key, value) in spec.environment { + command.env(key, value); + } + if let Some(directory) = spec.working_directory { + command.current_dir(directory); + } + unsafe { + command.pre_exec(|| { + libc::setpgid(0, 0); + libc::signal(libc::SIGPIPE, libc::SIG_DFL); + Ok(()) + }); + } + + let mut pty_write = None; + let mut readers = Vec::new(); + let mut child_stdin = None; + if spec.terminal { + let (master, slave) = open_pty()?; + let stdout = duplicate(&slave)?; + let stderr = duplicate(&slave)?; + command.stdin(Stdio::from(slave)); + command.stdout(Stdio::from(stdout)); + command.stderr(Stdio::from(stderr)); + set_nonblocking(master.as_raw_fd())?; + readers.push(Reader { + fd: master, + open: true, + is_stdout: true, + }); + pty_write = Some(readers[0].fd.try_clone().map_err(|_| ProcessError::Io)?); + } else { + command.stdin(Stdio::piped()).stdout(Stdio::piped()).stderr(Stdio::piped()); + } + + let mut child = command.spawn().map_err(|_| ProcessError::Launch)?; + let pid = child.id() as i32; + unsafe { + libc::setpgid(pid, pid); + } + let _guard = KillGuard { pid }; + + if !spec.terminal { + child_stdin = child.stdin.take(); + if let Some(stdout) = child.stdout.take() { + let fd = OwnedFd::from(stdout); + set_nonblocking(fd.as_raw_fd())?; + readers.push(Reader { fd, open: true, is_stdout: true }); + } + if let Some(stderr) = child.stderr.take() { + let fd = OwnedFd::from(stderr); + set_nonblocking(fd.as_raw_fd())?; + readers.push(Reader { fd, open: true, is_stdout: false }); + } + } + + let origin = Instant::now(); + if !spec.initial_input.is_empty() { + write_all(pty_write.as_ref(), child_stdin.as_mut(), spec.initial_input)?; + } + let keep_stdin = hooks.is_some() || spec.terminal; + if !keep_stdin { + child_stdin.take(); + } + + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + let mut child_exited = false; + let mut drained_after_exit = false; + let mut status_code = 0; + let mut hooks = hooks; + loop { + if spec.cancel.load(Ordering::Relaxed) { + return Err(ProcessError::Cancelled); + } + if origin.elapsed() >= spec.timeout { + return Err(ProcessError::TimedOut); + } + for reader in &mut readers { + if !reader.open { + continue; + } + loop { + if spec.cancel.load(Ordering::Relaxed) { + return Err(ProcessError::Cancelled); + } + let mut buffer = [0u8; 16384]; + let count = unsafe { libc::read(reader.fd.as_raw_fd(), buffer.as_mut_ptr().cast(), buffer.len()) }; + if count > 0 { + let chunk = &buffer[..count as usize]; + if reader.is_stdout { + stdout.extend_from_slice(chunk); + if let Some(hooks) = hooks.as_mut() { + match hooks.on_output(chunk, origin.elapsed()) { + StreamAction::None => {} + StreamAction::Send(input) => { + write_all(pty_write.as_ref(), child_stdin.as_mut(), &input)?; + } + StreamAction::Finish => { + return Ok(ProcessResult { stdout, stderr, exit_status: 0 }); + } + } + } + } else { + stderr.extend_from_slice(chunk); + } + if stdout.len() + stderr.len() > 4 * 1024 * 1024 { + return Err(ProcessError::OutputLimit); + } + continue; + } + if count == 0 { + reader.open = false; + break; + } + let error = std::io::Error::last_os_error(); + match error.raw_os_error() { + Some(libc::EINTR) => continue, + Some(libc::EAGAIN) => break, + Some(libc::EIO) if spec.terminal => { + reader.open = false; + break; + } + _ => return Err(ProcessError::Io), + } + } + } + if let Some(hooks) = hooks.as_mut() { + match hooks.on_poll(origin.elapsed()) { + StreamAction::None => {} + StreamAction::Send(input) => write_all(pty_write.as_ref(), child_stdin.as_mut(), &input)?, + StreamAction::Finish => return Ok(ProcessResult { stdout, stderr, exit_status: 0 }), + } + } + if !child_exited { + match child.try_wait() { + Ok(Some(status)) => { + child_exited = true; + status_code = status.code().unwrap_or(-1); + unsafe { + libc::kill(-pid, libc::SIGTERM); + libc::kill(-pid, libc::SIGKILL); + } + } + Ok(None) => {} + Err(_) => return Err(ProcessError::Io), + } + } + let output_open = readers.iter().any(|reader| reader.open); + if child_exited && !output_open { + return Ok(ProcessResult { + stdout, + stderr, + exit_status: status_code, + }); + } + if child_exited && spec.terminal { + if drained_after_exit { + return Ok(ProcessResult { + stdout, + stderr, + exit_status: status_code, + }); + } + drained_after_exit = true; + continue; + } + let wait_ms = if hooks.is_some() { 100 } else { 250 }; + let remaining = spec.timeout.saturating_sub(origin.elapsed()).as_millis() as i32; + poll_readers(&readers, wait_ms.min(remaining).max(0)); + } +} + +fn write_all( + master: Option<&OwnedFd>, + mut stdin: Option<&mut std::process::ChildStdin>, + data: &[u8], +) -> Result<(), ProcessError> { + if data.len() > 4096 { + return Err(ProcessError::Io); + } + let mut offset = 0; + while offset < data.len() { + let written = if let Some(master) = master { + unsafe { libc::write(master.as_raw_fd(), data[offset..].as_ptr().cast(), data.len() - offset) } + } else if let Some(stdin) = stdin.as_mut() { + match stdin.write(&data[offset..]) { + Ok(count) => count as isize, + Err(error) if error.kind() == std::io::ErrorKind::Interrupted => continue, + Err(_) => return Err(ProcessError::Io), + } + } else { + return Err(ProcessError::Io); + }; + if written < 0 { + let error = std::io::Error::last_os_error(); + if error.raw_os_error() == Some(libc::EINTR) { + continue; + } + return Err(ProcessError::Io); + } + offset += written as usize; + } + Ok(()) +} + +fn poll_readers(readers: &[Reader], millis: i32) { + let mut fds: Vec = readers + .iter() + .filter(|reader| reader.open) + .map(|reader| libc::pollfd { + fd: reader.fd.as_raw_fd(), + events: libc::POLLIN, + revents: 0, + }) + .collect(); + if fds.is_empty() { + std::thread::sleep(Duration::from_millis(millis.max(0) as u64)); + return; + } + unsafe { + libc::poll(fds.as_mut_ptr(), fds.len() as libc::nfds_t, millis); + } +} + +fn open_pty() -> Result<(OwnedFd, OwnedFd), ProcessError> { + let mut master = -1; + let mut slave = -1; + let mut size = libc::winsize { + ws_row: 50, + ws_col: 120, + ws_xpixel: 0, + ws_ypixel: 0, + }; + let result = unsafe { libc::openpty(&mut master, &mut slave, std::ptr::null_mut(), std::ptr::null(), &mut size) }; + if result != 0 { + return Err(ProcessError::Io); + } + unsafe { Ok((OwnedFd::from_raw_fd(master), OwnedFd::from_raw_fd(slave))) } +} + +fn duplicate(fd: &OwnedFd) -> Result { + let duplicated = unsafe { libc::fcntl(fd.as_raw_fd(), libc::F_DUPFD_CLOEXEC, 0) }; + if duplicated < 0 { + Err(ProcessError::Io) + } else { + Ok(unsafe { OwnedFd::from_raw_fd(duplicated) }) + } +} + +fn set_nonblocking(fd: i32) -> Result<(), ProcessError> { + let flags = unsafe { libc::fcntl(fd, libc::F_GETFL) }; + if flags < 0 || unsafe { libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) } < 0 { + return Err(ProcessError::Io); + } + Ok(()) +} + +struct KillGuard { + pid: i32, +} + +impl Drop for KillGuard { + fn drop(&mut self) { + unsafe { + libc::kill(-self.pid, libc::SIGTERM); + } + let deadline = Instant::now() + Duration::from_millis(500); + let mut status = 0; + let mut reaped = false; + while Instant::now() < deadline { + let result = unsafe { libc::waitpid(self.pid, &mut status, libc::WNOHANG) }; + if result == self.pid { + reaped = true; + break; + } + std::thread::sleep(Duration::from_millis(10)); + } + unsafe { + libc::kill(-self.pid, libc::SIGKILL); + if !reaped { + libc::waitpid(self.pid, &mut status, 0); + } + } + } +} diff --git a/src/protocol.rs b/src/protocol.rs new file mode 100644 index 0000000..956d9b7 --- /dev/null +++ b/src/protocol.rs @@ -0,0 +1,415 @@ +use std::time::Duration; + +use serde_json::{json, Value}; + +use crate::model::{FetchStatus, ProviderFailure, ProviderUsage}; +use crate::parsers::{self, parse_cursor}; +use crate::process::StreamAction; + +pub struct UsageProtocolSession { + kind: ProtocolKind, + pending: Vec, + usage_sent: bool, + pub usage_data: Option>, + pub failure: Option, +} + +#[derive(Clone, Copy, PartialEq, Eq)] +pub enum ProtocolKind { + Claude, + Codex, +} + +impl UsageProtocolSession { + pub fn new(kind: ProtocolKind) -> Self { + Self { + kind, + pending: Vec::new(), + usage_sent: false, + usage_data: None, + failure: None, + } + } + + pub fn initialize() -> Vec { + line(&json!({ + "type": "control_request", + "request_id": "ai-usage-init", + "request": { "subtype": "initialize" } + })) + } + + pub fn codex_initialize() -> Vec { + line(&json!({ + "id": 1, + "method": "initialize", + "params": { + "clientInfo": { + "name": "ai_usage_monitor", + "title": "AI Usage Monitor", + "version": "1.0" + } + } + })) + } + + pub fn consume(&mut self, chunk: &[u8]) -> StreamAction { + self.pending.extend_from_slice(chunk); + while let Some(index) = self.pending.iter().position(|byte| *byte == b'\n') { + let raw = self.pending.drain(..=index).collect::>(); + let bytes = &raw[..raw.len().saturating_sub(1)]; + let Ok(message) = serde_json::from_slice::(bytes) else { + continue; + }; + match self.kind { + ProtocolKind::Codex => { + let Some(id) = message.get("id").and_then(Value::as_i64) else { + continue; + }; + if id != 1 && id != 2 { + continue; + } + if message.get("error").is_some() { + self.failure = Some(ProviderFailure::new( + FetchStatus::Failed, + "Codex rejected the usage request. Check its existing subscription login.", + )); + return StreamAction::Finish; + } + if id == 1 && !self.usage_sent { + self.usage_sent = true; + let mut next = line_value(&json!({"method": "initialized", "params": {}})); + next.extend(line_value(&json!({ + "id": 2, + "method": "account/rateLimits/read", + "params": { "excludeResetCreditDetails": true } + }))); + return StreamAction::Send(next); + } + if id == 2 { + if let Some(result) = message.get("result").and_then(Value::as_object) { + let mut allowed = serde_json::Map::new(); + for key in ["rateLimits", "rateLimitsByLimitId"] { + if let Some(value) = result.get(key) { + allowed.insert(key.to_string(), value.clone()); + } + } + self.usage_data = serde_json::to_vec(&Value::Object(allowed)).ok(); + return StreamAction::Finish; + } + } + } + ProtocolKind::Claude => { + if message.get("type").and_then(Value::as_str) != Some("control_response") { + continue; + } + let Some(response) = message.get("response").and_then(Value::as_object) else { + continue; + }; + let Some(id) = response.get("request_id").and_then(Value::as_str) else { + continue; + }; + if id != "ai-usage-init" && id != "ai-usage-read" { + continue; + } + if response.get("subtype").and_then(Value::as_str) != Some("success") { + self.failure = Some(ProviderFailure::new( + FetchStatus::Unsupported, + "Claude's get_usage request is unsupported or failed. Open Claude to check its existing login.", + )); + return StreamAction::Finish; + } + if id == "ai-usage-init" && !self.usage_sent { + self.usage_sent = true; + return StreamAction::Send(line(&json!({ + "type": "control_request", + "request_id": "ai-usage-read", + "request": { "subtype": "get_usage", "skip_behaviors": true } + }))); + } + if id == "ai-usage-read" { + if let Some(result) = response.get("response").and_then(Value::as_object) { + let mut allowed = serde_json::Map::new(); + for key in ["rate_limits_available", "rate_limits"] { + if let Some(value) = result.get(key) { + allowed.insert(key.to_string(), value.clone()); + } + } + self.usage_data = serde_json::to_vec(&Value::Object(allowed)).ok(); + return StreamAction::Finish; + } + } + } + } + } + StreamAction::None + } +} + +fn line(value: &Value) -> Vec { + line_value(value) +} + +fn line_value(value: &Value) -> Vec { + let mut data = serde_json::to_vec(value).unwrap_or_default(); + data.push(b'\n'); + data +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum CursorStage { + Startup, + Trusting, + SelectingUsage, + Submitted, +} + +pub struct CursorUsageSession { + stage: CursorStage, + transcript: Vec, + pending_input: Option>, + last_output_at: Duration, + workspace: String, + workspace_canonical: String, + pub usage: Option, + pub failure: Option, +} + +impl CursorUsageSession { + pub fn new(workspace: impl Into) -> Self { + let workspace = workspace.into(); + let workspace_canonical = std::fs::canonicalize(&workspace) + .map(|path| path.display().to_string()) + .unwrap_or_else(|_| workspace.clone()); + Self { + stage: CursorStage::Startup, + transcript: Vec::new(), + pending_input: None, + last_output_at: Duration::ZERO, + workspace, + workspace_canonical, + usage: None, + failure: None, + } + } + + pub fn consume(&mut self, chunk: &[u8], at: Duration) -> StreamAction { + self.last_output_at = at; + self.transcript.extend_from_slice(chunk); + if self.transcript.len() > 1024 * 1024 { + self.failure = Some(ProviderFailure::new( + FetchStatus::Failed, + "Cursor produced too much terminal output.", + )); + return StreamAction::Finish; + } + if self.pending_input.is_some() { + return StreamAction::None; + } + let text = parsers::terminal_text(&self.transcript); + if text.contains("Not logged in") + || text.contains("Authentication required") + || text.contains("Sign in to Cursor") + { + self.failure = Some(ProviderFailure::new( + FetchStatus::NotSignedIn, + "Sign in using Cursor Agent CLI (agent login), then refresh.", + )); + return StreamAction::Finish; + } + match self.stage { + CursorStage::Startup | CursorStage::Trusting => { + if self.stage == CursorStage::Startup && text.contains("Workspace Trust Required") { + if !text.contains("Trust this workspace") { + return StreamAction::None; + } + if !text.contains(&self.workspace) && !text.contains(&self.workspace_canonical) { + self.failure = Some(ProviderFailure::new( + FetchStatus::Failed, + "Cursor requested trust for an unexpected workspace.", + )); + return StreamAction::Finish; + } + self.stage = CursorStage::Trusting; + self.transcript.clear(); + self.pending_input = Some(b"a".to_vec()); + return StreamAction::None; + } + if text.contains("Plan, search, build anything") { + self.stage = CursorStage::SelectingUsage; + self.transcript.clear(); + self.pending_input = Some(b"/usage".to_vec()); + return StreamAction::None; + } + } + CursorStage::SelectingUsage => { + if text.contains("/usage") && text.contains("Show plan and on-demand usage") { + self.stage = CursorStage::Submitted; + self.transcript.clear(); + self.pending_input = Some(b"\r".to_vec()); + return StreamAction::None; + } + } + CursorStage::Submitted => { + if text.contains("Failed to load usage data") { + self.failure = Some(ProviderFailure::new( + FetchStatus::Failed, + "Cursor could not load usage. Check its existing login and network connection.", + )); + return StreamAction::Finish; + } + if text.contains("Usage details are not available") + || text.contains("Personal spend for the current billing period") + { + self.failure = Some(ProviderFailure::new( + FetchStatus::Unsupported, + "Cursor did not expose an included allowance for this plan in the CLI.", + )); + return StreamAction::Finish; + } + if let Ok(parsed) = parse_cursor(&self.transcript, chrono::Utc::now()) { + self.usage = Some(parsed); + return StreamAction::Finish; + } + } + } + StreamAction::None + } + + pub fn poll(&mut self, at: Duration) -> StreamAction { + let Some(input) = &self.pending_input else { + return StreamAction::None; + }; + if at.saturating_sub(self.last_output_at) < Duration::from_millis(300) { + return StreamAction::None; + } + let input = input.clone(); + self.pending_input = None; + StreamAction::Send(input) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::parsers::parse_claude; + use crate::parsers::parse_codex; + + #[test] + fn chunked_codex_protocol_ignores_notifications_and_metadata() { + let mut session = UsageProtocolSession::new(ProtocolKind::Codex); + let notification = line(&json!({"method": "unrelated/updated", "params": {"private": "ignored"}})); + assert!(matches!(session.consume(¬ification), StreamAction::None)); + let init_reply = line(&json!({"id": 1, "result": {}})); + assert!(matches!(session.consume(&init_reply[..4]), StreamAction::None)); + let StreamAction::Send(command) = session.consume(&init_reply[4..]) else { + panic!("no usage request"); + }; + let methods: Vec = command + .split(|byte| *byte == b'\n') + .filter(|line| !line.is_empty()) + .map(|line| { + serde_json::from_slice::(line) + .unwrap() + .get("method") + .and_then(Value::as_str) + .unwrap() + .to_string() + }) + .collect(); + assert_eq!(methods, ["initialized", "account/rateLimits/read"]); + let payload: Value = serde_json::from_slice(include_bytes!("../fixtures/codex.json")).unwrap(); + assert!(matches!( + session.consume(&line(&json!({"id": 2, "result": payload}))), + StreamAction::Finish + )); + let usage = parse_codex(session.usage_data.as_ref().unwrap(), chrono::Utc::now()).unwrap(); + assert_eq!(usage.short_window.unwrap().remaining_percentage, Some(71.0)); + } + + #[test] + fn claude_protocol_discards_session_and_behavior_data() { + let mut session = UsageProtocolSession::new(ProtocolKind::Claude); + session.consume(&line(&json!({ + "type": "control_response", + "response": {"request_id": "ai-usage-init", "subtype": "success", "response": {}} + }))); + let mut payload: Value = serde_json::from_slice(include_bytes!("../fixtures/claude.json")).unwrap(); + payload["session"] = json!({"private": "never save"}); + payload["behaviors"] = json!({"private": "never save"}); + session.consume(&line(&json!({ + "type": "control_response", + "response": {"request_id": "ai-usage-read", "subtype": "success", "response": payload} + }))); + let data = session.usage_data.unwrap(); + let text = String::from_utf8(data.clone()).unwrap(); + assert!(!text.contains("never save")); + assert_eq!( + parse_claude(&data, chrono::Utc::now()).unwrap().short_window.unwrap().remaining_percentage, + Some(38.0) + ); + } + + #[test] + fn remote_errors_are_not_copied_to_state() { + let mut session = UsageProtocolSession::new(ProtocolKind::Codex); + session.consume(&line(&json!({"id": 2, "error": {"message": "private-account-or-token"}}))); + let message = session.failure.unwrap().message; + assert!(!message.contains("private-account-or-token")); + } + + #[test] + fn cursor_session_waits_for_menu_and_chunked_complete_usage() { + let mut session = CursorUsageSession::new("/tmp/monitor-fixture"); + assert!(matches!( + session.consume(b"Workspace Trust Required", Duration::ZERO), + StreamAction::None + )); + session.consume(b"\n/tmp/monitor-fixture\nTrust this workspace", Duration::ZERO); + assert!(matches!(session.poll(Duration::from_millis(100)), StreamAction::None)); + match session.poll(Duration::from_secs(1)) { + StreamAction::Send(trust) => assert_eq!(trust, b"a"), + other => panic!("no trust key: {other:?}"), + } + session.consume(b"Plan, search, build anything", Duration::from_secs(2)); + match session.poll(Duration::from_secs(3)) { + StreamAction::Send(command) => assert_eq!(command, b"/usage"), + other => panic!("no usage command: {other:?}"), + } + assert!(matches!( + session.consume(b"-> /usage", Duration::from_secs(4)), + StreamAction::None + )); + session.consume(b" Show plan and on-demand usage", Duration::from_secs(4)); + match session.poll(Duration::from_secs(5)) { + StreamAction::Send(enter) => assert_eq!(enter, b"\r"), + other => panic!("did not accept usage command: {other:?}"), + } + let data = include_bytes!("../fixtures/cursor-usage.txt"); + for byte in data.iter().take(data.len() / 2).copied() { + session.consume(&[byte], Duration::from_secs(6)); + } + assert!(session.usage.is_none()); + assert!(matches!( + session.consume(&data[data.len() / 2..], Duration::from_secs(6)), + StreamAction::Finish + )); + assert_eq!(session.usage.unwrap().billing_window.unwrap().remaining_percentage, Some(56.0)); + } + + #[test] + fn cursor_session_does_not_trust_unexpected_workspace() { + let mut session = CursorUsageSession::new("/tmp/monitor-fixture"); + assert!(matches!( + session.consume( + b"Workspace Trust Required\n/other/project\nTrust this workspace", + Duration::ZERO + ), + StreamAction::Finish + )); + assert!(session.usage.is_none()); + let failure = session.failure.unwrap(); + assert_eq!(failure.status, FetchStatus::Failed); + assert!(!failure.message.contains("/other/project")); + } +} diff --git a/src/providers.rs b/src/providers.rs new file mode 100644 index 0000000..94aa9f8 --- /dev/null +++ b/src/providers.rs @@ -0,0 +1,459 @@ +use std::fs; +use std::os::unix::fs::DirBuilderExt; +use std::path::{Path, PathBuf}; +use std::sync::atomic::AtomicBool; +use std::time::Duration; + +use chrono::Utc; +use serde_json::Value; + +use crate::discovery; +use crate::model::{FetchStatus, Provider, ProviderFailure, ProviderUsage, UsageMetric}; +use crate::parsers::{self, parse_claude, parse_codex}; +use crate::process::{self, CommandSpec, ProcessError, ProcessHooks, StreamAction}; +use crate::protocol::{CursorUsageSession, ProtocolKind, UsageProtocolSession}; + +const GUARD: &str = include_str!("cursor-guard.js"); + +const ALLOWED_ENV: &[&str] = &[ + "HOME", + "PATH", + "LANG", + "LC_ALL", + "LC_CTYPE", + "TMPDIR", + "USER", + "LOGNAME", + "CODEX_HOME", + "CLAUDE_CONFIG_DIR", + "SSL_CERT_FILE", + "SSL_CERT_DIR", + "NIX_SSL_CERT_FILE", + "XDG_RUNTIME_DIR", + "XDG_CONFIG_HOME", + "XDG_DATA_HOME", + "XDG_CACHE_HOME", + "DBUS_SESSION_BUS_ADDRESS", +]; + +pub fn fetch( + provider: Provider, + override_path: Option<&str>, + cancel: &AtomicBool, +) -> Result { + let home = home_dir(); + let path = std::env::var("PATH").unwrap_or_default(); + let executable = discovery::candidates(provider, override_path, &home, &path) + .into_iter() + .next(); + match provider { + Provider::Claude => fetch_claude(executable.as_ref().map(|item| item.path.as_path()), cancel), + Provider::Codex => fetch_codex(executable.as_ref().map(|item| item.path.as_path()), cancel), + Provider::Cursor => fetch_cursor(executable.as_ref().map(|item| item.path.as_path()), &home, cancel), + } +} + +fn fetch_claude(executable: Option<&Path>, cancel: &AtomicBool) -> Result { + let Some(executable) = executable else { + return Err(ProviderFailure::new( + FetchStatus::Unsupported, + "Claude executable not found. Set its path in Settings.", + )); + }; + let mut session = UsageProtocolSession::new(ProtocolKind::Claude); + let mut environment = local_environment(); + for name in [ + "DISABLE_TELEMETRY", + "DISABLE_ERROR_REPORTING", + "DISABLE_GROWTHBOOK", + "DISABLE_AUTOUPDATER", + "DISABLE_UPDATES", + "DISABLE_FEEDBACK_COMMAND", + "DISABLE_BUG_COMMAND", + "CLAUDE_CODE_DISABLE_FEEDBACK_SURVEY", + "CLAUDE_CODE_DISABLE_OFFICIAL_MARKETPLACE_AUTOINSTALL", + ] { + environment.push((name.to_string(), "1".to_string())); + } + environment.push(("CLAUDE_CODE_ENABLE_TELEMETRY".to_string(), "0".to_string())); + environment.push(("OTEL_METRICS_EXPORTER".to_string(), "none".to_string())); + environment.push(("OTEL_LOGS_EXPORTER".to_string(), "none".to_string())); + environment.push(("OTEL_TRACES_EXPORTER".to_string(), "none".to_string())); + let arguments = strings(&[ + "--print", + "--input-format", + "stream-json", + "--output-format", + "stream-json", + "--verbose", + "--no-session-persistence", + "--safe-mode", + "--setting-sources", + "", + "--strict-mcp-config", + "--mcp-config", + "{\"mcpServers\":{}}", + "--tools", + "", + ]); + let initial = UsageProtocolSession::initialize(); + let directory = std::env::temp_dir(); + let mut hooks = ProtocolHooks { session: &mut session }; + let result = process::run( + CommandSpec { + executable, + arguments: &arguments, + environment: &environment, + working_directory: Some(&directory), + timeout: Duration::from_secs(30), + terminal: false, + initial_input: &initial, + cancel, + }, + Some(&mut hooks), + ); + map_process(result)?; + if let Some(failure) = session.failure { + return Err(failure); + } + let data = session.usage_data.ok_or_else(|| { + ProviderFailure::new( + FetchStatus::Unsupported, + "Claude did not return a structured usage response.", + ) + })?; + parse_claude(&data, Utc::now()) +} + +fn fetch_codex(executable: Option<&Path>, cancel: &AtomicBool) -> Result { + let Some(executable) = executable else { + return Err(ProviderFailure::new( + FetchStatus::Unsupported, + "Codex executable not found. Set its path in Settings.", + )); + }; + let mut session = UsageProtocolSession::new(ProtocolKind::Codex); + let environment = local_environment(); + let arguments = strings(&[ + "-c", + "analytics.enabled=false", + "-c", + "otel.exporter=\"none\"", + "-c", + "otel.trace_exporter=\"none\"", + "app-server", + "--listen", + "stdio://", + ]); + let initial = UsageProtocolSession::codex_initialize(); + let directory = std::env::temp_dir(); + let mut hooks = ProtocolHooks { session: &mut session }; + let result = process::run( + CommandSpec { + executable, + arguments: &arguments, + environment: &environment, + working_directory: Some(&directory), + timeout: Duration::from_secs(30), + terminal: false, + initial_input: &initial, + cancel, + }, + Some(&mut hooks), + ); + map_process(result)?; + if let Some(failure) = session.failure { + return Err(failure); + } + let data = session.usage_data.ok_or_else(|| { + ProviderFailure::new(FetchStatus::Unsupported, "Codex did not return a rate-limit response.") + })?; + parse_codex(&data, Utc::now()) +} + +fn fetch_cursor(executable: Option<&Path>, home: &Path, cancel: &AtomicBool) -> Result { + if let Some(executable) = executable { + return fetch_cursor_cli(executable, cancel); + } + fetch_subscription_status(home, cancel) +} + +fn fetch_cursor_cli(executable: &Path, cancel: &AtomicBool) -> Result { + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|duration| duration.as_nanos()) + .unwrap_or(0); + let root = std::env::temp_dir().join(format!("ai-usage-cursor-{}-{nanos}", std::process::id())); + let workspace = root.join("workspace"); + let config = root.join("config"); + fs::DirBuilder::new() + .recursive(true) + .mode(0o700) + .create(&workspace) + .and_then(|_| fs::create_dir_all(&config)) + .map_err(|_| { + ProviderFailure::new( + FetchStatus::Failed, + "Cursor's temporary usage workspace could not be created.", + ) + })?; + let _cleanup = TempDir(root.clone()); + let guard = root.join("usage-only.cjs"); + fs::write(&guard, GUARD).map_err(|_| { + ProviderFailure::new( + FetchStatus::Failed, + "Cursor's usage-only process configuration could not be created.", + ) + })?; + let mut environment = local_environment(); + environment.push(("NODE_OPTIONS".to_string(), format!("--require \"{}\"", guard.display()))); + environment.push(("CURSOR_CONFIG_DIR".to_string(), config.display().to_string())); + environment.push(("CURSOR_DATA_DIR".to_string(), root.join("data").display().to_string())); + environment.push(( + "NODE_COMPILE_CACHE".to_string(), + root.join("compile-cache").display().to_string(), + )); + environment.push(("TERM".to_string(), "xterm-256color".to_string())); + environment.push(("LANG".to_string(), "en_US.UTF-8".to_string())); + environment.push(("LC_ALL".to_string(), "en_US.UTF-8".to_string())); + environment.push(("NO_COLOR".to_string(), "1".to_string())); + environment.push(("NO_OPEN_BROWSER".to_string(), "1".to_string())); + environment.push(("OTEL_TRACES_SAMPLER".to_string(), "always_off".to_string())); + let common = [ + "--disable-auto-update", + "--disable-project-configs", + "--disable-indexing", + "--disable-codebase-ref", + "--exclude-workspace-context", + ]; + let mut status_environment = environment.clone(); + status_environment.retain(|(key, _)| key != "CURSOR_CONFIG_DIR"); + let status_arguments = strings(&{ + let mut args = common.to_vec(); + args.extend(["status", "--format", "json"]); + args + }); + let status = process::run( + CommandSpec { + executable, + arguments: &status_arguments, + environment: &status_environment, + working_directory: Some(&workspace), + timeout: Duration::from_secs(15), + terminal: false, + initial_input: &[], + cancel, + }, + None, + ) + .map_err(process_failure)?; + let login: Value = serde_json::from_slice(&status.stdout).map_err(|_| { + ProviderFailure::new( + FetchStatus::Unsupported, + "Cursor Agent CLI did not return its authentication status. Check its path and version.", + ) + })?; + let signed_in = login.get("isAuthenticated").and_then(Value::as_bool).ok_or_else(|| { + ProviderFailure::new( + FetchStatus::Unsupported, + "Cursor Agent CLI did not return its authentication status. Check its path and version.", + ) + })?; + if status.exit_status != 0 || !signed_in { + return Err(ProviderFailure::new( + FetchStatus::NotSignedIn, + "Sign in using Cursor Agent CLI (agent login), then refresh.", + )); + } + let workspace_text = workspace.display().to_string(); + let mut argument_strings: Vec = common.iter().map(|item| (*item).to_string()).collect(); + argument_strings.push("--workspace".to_string()); + argument_strings.push(workspace_text.clone()); + if let Some(team) = login + .get("userInfo") + .and_then(|info| info.get("teamId")) + .and_then(|value| value.as_i64().or_else(|| value.as_str().and_then(|text| text.parse().ok()))) + { + if team > 0 { + argument_strings.push("--header".to_string()); + argument_strings.push(format!("x-cursor-team-id: {team}")); + } + } + let mut session = CursorUsageSession::new(workspace_text); + let mut hooks = CursorHooks { session: &mut session }; + let result = process::run( + CommandSpec { + executable, + arguments: &argument_strings, + environment: &environment, + working_directory: Some(&workspace), + timeout: Duration::from_secs(30), + terminal: true, + initial_input: &[], + cancel, + }, + Some(&mut hooks), + ); + map_process(result)?; + if let Some(failure) = session.failure { + return Err(failure); + } + session.usage.ok_or_else(|| { + ProviderFailure::new( + FetchStatus::Unsupported, + "Cursor Agent CLI did not return a recognized /usage screen. Check its version.", + ) + }) +} + +fn fetch_subscription_status(home: &Path, cancel: &AtomicBool) -> Result { + let database = home.join(".config/Cursor/User/globalStorage/state.vscdb"); + if !database.exists() { + return Err(ProviderFailure::new( + FetchStatus::Unsupported, + "Cursor's local database is unavailable; no reliable usage source was found.", + )); + } + let sqlite = which("sqlite3").ok_or_else(|| { + ProviderFailure::new( + FetchStatus::Unsupported, + "Cursor's local database is unavailable; no reliable usage source was found.", + ) + })?; + let arguments = strings(&[ + "-readonly", + "-cmd", + ".timeout 1000", + &database.display().to_string(), + "SELECT value FROM ItemTable WHERE key = 'cursorAuth/stripeSubscriptionStatus';", + ]); + let result = process::run( + CommandSpec { + executable: &sqlite, + arguments: &arguments, + environment: &local_environment(), + working_directory: None, + timeout: Duration::from_secs(5), + terminal: false, + initial_input: &[], + cancel, + }, + None, + ) + .map_err(process_failure)?; + let text = String::from_utf8_lossy(&result.stdout); + let metric = (result.exit_status == 0) + .then(|| parsers::subscription_status(&text)) + .flatten() + .ok_or_else(|| { + ProviderFailure::new( + FetchStatus::Unsupported, + "Cursor exposes no recognized local subscription status or usage allowance.", + ) + })?; + Ok(ProviderUsage { + provider: Provider::Cursor, + fetched_at: Utc::now(), + short_window: None, + weekly_window: None, + billing_window: None, + billing_reset_label: None, + billing_pools: None, + additional_usage: vec![ + metric, + UsageMetric { + name: "Usage".to_string(), + value: "Install Cursor Agent CLI".to_string(), + }, + ], + }) +} + +struct ProtocolHooks<'a> { + session: &'a mut UsageProtocolSession, +} + +impl ProcessHooks for ProtocolHooks<'_> { + fn on_output(&mut self, chunk: &[u8], _at: Duration) -> StreamAction { + self.session.consume(chunk) + } + fn on_poll(&mut self, _at: Duration) -> StreamAction { + StreamAction::None + } +} + +struct CursorHooks<'a> { + session: &'a mut CursorUsageSession, +} + +impl ProcessHooks for CursorHooks<'_> { + fn on_output(&mut self, chunk: &[u8], at: Duration) -> StreamAction { + self.session.consume(chunk, at) + } + fn on_poll(&mut self, at: Duration) -> StreamAction { + self.session.poll(at) + } +} + +struct TempDir(PathBuf); + +impl Drop for TempDir { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.0); + } +} + +fn map_process(result: Result) -> Result<(), ProviderFailure> { + result.map(|_| ()).map_err(process_failure) +} + +fn process_failure(error: ProcessError) -> ProviderFailure { + match error { + ProcessError::TimedOut => ProviderFailure::new( + FetchStatus::Failed, + "The provider timed out. Last successful usage is retained.", + ), + ProcessError::Cancelled => ProviderFailure::new(FetchStatus::Failed, "Update cancelled."), + ProcessError::Launch => ProviderFailure::new( + FetchStatus::Failed, + "The provider executable could not start. Check its path.", + ), + ProcessError::OutputLimit | ProcessError::Io => ProviderFailure::new( + FetchStatus::Failed, + "The provider process failed. Check its path and existing login.", + ), + } +} + +fn local_environment() -> Vec<(String, String)> { + let mut values: Vec<(String, String)> = std::env::vars() + .filter(|(key, _)| ALLOWED_ENV.contains(&key.as_str())) + .collect(); + values.sort(); + if !values.iter().any(|(key, _)| key == "SSL_CERT_FILE") { + let bundle = Path::new("/etc/ssl/certs/ca-certificates.crt"); + if bundle.exists() { + values.push(("SSL_CERT_FILE".to_string(), bundle.display().to_string())); + } + } + values +} + +fn strings(values: &[&str]) -> Vec { + values.iter().map(|value| (*value).to_string()).collect() +} + +fn which(name: &str) -> Option { + std::env::var_os("PATH").and_then(|path| { + std::env::split_paths(&path).find_map(|directory| { + let candidate = directory.join(name); + candidate.is_file().then_some(candidate) + }) + }) +} + +fn home_dir() -> PathBuf { + std::env::var_os("HOME") + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from("/")) +} diff --git a/src/refresh.rs b/src/refresh.rs new file mode 100644 index 0000000..1607fb9 --- /dev/null +++ b/src/refresh.rs @@ -0,0 +1,312 @@ +use std::collections::{HashMap, HashSet}; + +use chrono::{DateTime, Utc}; + +use crate::model::{FetchStatus, Provider, ProviderFailure, ProviderUsage, UsageSnapshot}; +use crate::settings::Settings; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Reason { + Startup, + Timer, + Manual, + Resume, + Settings, +} + +pub struct Engine { + pub snapshot: UsageSnapshot, + pub settings: Settings, + generation: u64, + forced: HashSet, +} + +pub struct Round { + generation: u64, + pub providers: Vec, + attempted: DateTime, + forced: HashSet, +} + +impl Engine { + pub fn new(snapshot: UsageSnapshot, settings: Settings) -> Self { + let mut engine = Self { + snapshot, + settings, + generation: 0, + forced: HashSet::new(), + }; + engine.sync_visibility(); + engine + } + + pub fn apply_settings(&mut self, updated: Settings) { + let old = self.settings.clone(); + if old == updated { + return; + } + self.settings = updated; + for provider in Provider::ALL { + let path_changed = old.override_for(provider) != self.settings.override_for(provider); + let newly_enabled = self.settings.enabled(provider) && !old.enabled(provider); + if self.settings.enabled(provider) && (path_changed || newly_enabled) { + self.forced.insert(provider); + } + if !self.settings.enabled(provider) { + self.forced.remove(&provider); + } + } + self.sync_visibility(); + let selection_changed = old.enabled_providers != self.settings.enabled_providers + || Provider::ALL.iter().any(|provider| old.override_for(*provider) != self.settings.override_for(*provider)); + if selection_changed { + self.generation = self.generation.wrapping_add(1); + } + } + + pub fn due(&self, now: DateTime) -> Vec { + Provider::ALL + .into_iter() + .filter(|provider| self.settings.enabled(*provider)) + .filter(|provider| { + let elapsed = now.signed_duration_since(self.last_attempt(*provider)); + elapsed.num_seconds() >= i64::from(self.settings.refresh_minutes) * 60 + }) + .collect() + } + + pub fn select(&mut self, reason: Reason, now: DateTime) -> Option { + let mut selected: Vec = if reason == Reason::Manual { + self.settings.enabled_providers.clone() + } else { + let mut providers = self.due(now); + for provider in &self.forced { + if self.settings.enabled(*provider) && !providers.contains(provider) { + providers.push(*provider); + } + } + providers + }; + selected.retain(|provider| self.settings.enabled(*provider)); + if selected.is_empty() { + return None; + } + let forced = self + .forced + .iter() + .copied() + .filter(|provider| selected.contains(provider)) + .collect::>(); + for provider in &selected { + self.forced.remove(provider); + } + let generation = self.generation; + Some(Round { + generation, + providers: selected, + attempted: now, + forced, + }) + } + + pub fn commit( + &mut self, + round: Round, + results: HashMap>, + cancelled: bool, + ) -> bool { + if cancelled || round.generation != self.generation { + for provider in round.forced { + if self.settings.enabled(provider) { + self.forced.insert(provider); + } + } + return false; + } + for provider in round.providers { + if !self.settings.enabled(provider) { + continue; + } + let previous = self.snapshot.state(provider).last_successful_usage.clone(); + let state = self.snapshot.state_mut(provider); + state.last_attempt_at = Some(round.attempted); + match results.get(&provider) { + Some(Ok(usage)) => { + state.last_successful_usage = Some(usage.clone()); + state.status = FetchStatus::Available; + state.message = None; + } + Some(Err(failure)) if failure.message == "Update cancelled." => { + state.last_successful_usage = previous; + } + Some(Err(failure)) => { + state.last_successful_usage = previous; + state.status = failure.status; + state.message = Some(failure.message.clone()); + } + None => { + state.last_successful_usage = previous; + state.status = FetchStatus::Failed; + state.message = Some("Provider update failed. Check the app diagnostics.".to_string()); + } + } + } + self.snapshot.written_at = Utc::now(); + self.snapshot.show_used_percentage = self.settings.show_used_percentage; + self.snapshot.is_demo = false; + true + } + + pub fn next_delay(&self, now: DateTime) -> Option { + self.settings + .enabled_providers + .iter() + .map(|provider| { + self.last_attempt(*provider) + chrono::Duration::seconds(i64::from(self.settings.refresh_minutes) * 60) + }) + .min() + .map(|next| { + let delay = next.signed_duration_since(now); + if delay > chrono::Duration::milliseconds(10) { + delay + } else { + chrono::Duration::milliseconds(10) + } + }) + } + + fn sync_visibility(&mut self) { + for provider in Provider::ALL { + if !self.settings.enabled(provider) { + let state = self.snapshot.state_mut(provider); + state.status = FetchStatus::Disabled; + state.message = None; + } else if self.snapshot.state(provider).status == FetchStatus::Disabled { + let state = self.snapshot.state_mut(provider); + state.status = if state.last_successful_usage.is_none() { + FetchStatus::Pending + } else { + FetchStatus::Available + }; + } + } + self.snapshot.show_used_percentage = self.settings.show_used_percentage; + } + + fn last_attempt(&self, provider: Provider) -> DateTime { + let state = self.snapshot.state(provider); + state + .last_attempt_at + .or(state.last_successful_usage.as_ref().map(|usage| usage.fetched_at)) + .unwrap_or_else(|| DateTime::from_timestamp(0, 0).expect("epoch")) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::model::{FetchStatus, ProviderState, ProviderUsage, UsageWindow}; + use crate::settings::Settings; + + fn at(seconds: i64) -> DateTime { + DateTime::from_timestamp(seconds, 0).unwrap() + } + + fn cached(when: DateTime) -> UsageSnapshot { + UsageSnapshot { + schema_version: 1, + written_at: when, + show_used_percentage: false, + is_demo: false, + providers: Provider::ALL + .into_iter() + .map(|provider| ProviderState { + provider, + last_successful_usage: Some(ProviderUsage { + provider, + fetched_at: when, + short_window: Some(UsageWindow::new(Some(12.0), None, None)), + weekly_window: None, + billing_window: None, + billing_reset_label: None, + billing_pools: None, + additional_usage: Vec::new(), + }), + last_attempt_at: Some(when), + status: FetchStatus::Available, + message: None, + }) + .collect(), + } + } + + #[test] + fn recent_cache_is_not_due_until_the_interval() { + let snapshot = cached(at(1_800_000_000)); + let mut settings = Settings::default(); + settings.enabled_providers = vec![Provider::Codex]; + let engine = Engine::new(snapshot, settings); + assert!(engine.due(at(1_800_000_000 + 60)).is_empty()); + assert_eq!(engine.due(at(1_800_000_000 + 5 * 60)), vec![Provider::Codex]); + } + + #[test] + fn manual_selects_every_enabled_provider() { + let mut engine = Engine::new(cached(at(1_800_000_000)), Settings::default()); + let round = engine.select(Reason::Manual, at(1_800_000_100)).unwrap(); + assert_eq!(round.providers.len(), 3); + } + + #[test] + fn failure_keeps_the_previous_success() { + let when = at(1_800_000_000); + let mut engine = Engine::new(cached(when), Settings::default()); + let round = engine.select(Reason::Manual, when).unwrap(); + let mut results = HashMap::new(); + results.insert( + Provider::Claude, + Err(ProviderFailure::new(FetchStatus::Failed, "Fixture failure")), + ); + assert!(engine.commit(round, results, false)); + let state = engine.snapshot.state(Provider::Claude); + assert_eq!(state.status, FetchStatus::Failed); + assert_eq!(state.message.as_deref(), Some("Fixture failure")); + assert_eq!(state.last_successful_usage.unwrap().fetched_at, when); + } + + #[test] + fn cancelled_round_does_not_change_the_cache() { + let when = at(1_800_000_000); + let mut engine = Engine::new(cached(when), Settings::default()); + let before = engine.snapshot.clone(); + let round = engine.select(Reason::Timer, when + chrono::Duration::hours(1)).unwrap(); + let mut results = HashMap::new(); + results.insert( + Provider::Codex, + Ok(ProviderUsage { + provider: Provider::Codex, + fetched_at: when + chrono::Duration::hours(1), + short_window: Some(UsageWindow::new(Some(90.0), None, None)), + weekly_window: None, + billing_window: None, + billing_reset_label: None, + billing_pools: None, + additional_usage: Vec::new(), + }), + ); + assert!(!engine.commit(round, results, true)); + assert_eq!(engine.snapshot.state(Provider::Codex).last_successful_usage, before.state(Provider::Codex).last_successful_usage); + } + + #[test] + fn enabling_a_provider_forces_a_refresh() { + let mut settings = Settings::default(); + settings.enabled_providers = vec![Provider::Codex]; + let mut engine = Engine::new(cached(Utc::now()), settings); + let mut updated = Settings::default(); + updated.enabled_providers = vec![Provider::Codex, Provider::Claude]; + engine.apply_settings(updated); + assert!(engine.forced.contains(&Provider::Claude)); + assert_eq!(engine.snapshot.state(Provider::Cursor).status, FetchStatus::Disabled); + } +} diff --git a/src/settings.rs b/src/settings.rs new file mode 100644 index 0000000..c489568 --- /dev/null +++ b/src/settings.rs @@ -0,0 +1,88 @@ +use std::collections::HashMap; +use std::fs; +use std::path::Path; + +use serde::{Deserialize, Serialize}; + +use crate::model::Provider; +use crate::paths; + +pub const INTERVALS: [u32; 4] = [5, 15, 30, 60]; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct Settings { + pub refresh_minutes: u32, + pub enabled_providers: Vec, + #[serde(default)] + pub executable_overrides: HashMap, + pub show_used_percentage: bool, +} + +impl Default for Settings { + fn default() -> Self { + Self { + refresh_minutes: 5, + enabled_providers: Provider::ALL.to_vec(), + executable_overrides: HashMap::new(), + show_used_percentage: false, + } + } +} + +impl Settings { + pub fn normalize(mut self) -> Self { + if !INTERVALS.contains(&self.refresh_minutes) { + self.refresh_minutes = 5; + } + self.enabled_providers.sort_by_key(|provider| match provider { + Provider::Claude => 0, + Provider::Codex => 1, + Provider::Cursor => 2, + }); + self.enabled_providers.dedup(); + self.executable_overrides + .retain(|_, value| !value.trim().is_empty()); + self + } + + pub fn enabled(&self, provider: Provider) -> bool { + self.enabled_providers.contains(&provider) + } + + pub fn override_for(&self, provider: Provider) -> Option<&str> { + self.executable_overrides + .get(&provider) + .map(String::as_str) + .filter(|value| !value.trim().is_empty()) + } + + pub fn load(path: &Path) -> Self { + let Ok(data) = fs::read(path) else { + return Self::default(); + }; + serde_json::from_slice::(&data) + .map(Self::normalize) + .unwrap_or_else(|_| Self::default()) + } + + pub fn save(&self, path: &Path) -> std::io::Result<()> { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent)?; + } + let data = serde_json::to_vec_pretty(self).unwrap_or_default(); + let temporary = path.with_extension("json.tmp"); + fs::write(&temporary, data)?; + fs::rename(temporary, path) + } + + pub fn load_default() -> Self { + let path = paths::config_path(); + if !path.exists() { + let settings = Self::default(); + let _ = settings.save(&path); + return settings; + } + Self::load(&path) + } +} diff --git a/src/store.rs b/src/store.rs new file mode 100644 index 0000000..9c25bad --- /dev/null +++ b/src/store.rs @@ -0,0 +1,64 @@ +use std::fs; +use std::path::Path; + +use crate::model::{Provider, UsageSnapshot}; + +#[derive(Debug)] +pub struct StoreError; + +impl UsageSnapshot { + pub fn read(path: &Path) -> Result { + if !path.exists() { + return Ok(Self::default()); + } + let data = fs::read(path).map_err(|_| StoreError)?; + let snapshot: UsageSnapshot = serde_json::from_slice(&data).map_err(|_| StoreError)?; + if snapshot.schema_version != 1 { + return Err(StoreError); + } + let mut seen = std::collections::HashSet::new(); + for state in &snapshot.providers { + if !seen.insert(state.provider) { + return Err(StoreError); + } + } + let mut snapshot = snapshot; + for provider in Provider::ALL { + if !snapshot.providers.iter().any(|state| state.provider == provider) { + snapshot.providers.push(crate::model::ProviderState::pending(provider)); + } + } + Ok(snapshot) + } + + pub fn write(&self, path: &Path) -> Result<(), StoreError> { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent).map_err(|_| StoreError)?; + } + let data = serde_json::to_vec_pretty(self).map_err(|_| StoreError)?; + let temporary = path.with_extension("json.tmp"); + fs::write(&temporary, &data).map_err(|_| StoreError)?; + fs::rename(&temporary, path).map_err(|_| StoreError) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::Utc; + + #[test] + fn snapshot_round_trip() { + let directory = std::env::temp_dir().join(format!("ai-usage-store-{}", std::process::id())); + let _ = fs::remove_dir_all(&directory); + fs::create_dir_all(&directory).unwrap(); + let path = directory.join("usage.json"); + let mut snapshot = UsageSnapshot::default(); + snapshot.written_at = Utc::now(); + snapshot.write(&path).unwrap(); + let loaded = UsageSnapshot::read(&path).unwrap(); + assert_eq!(loaded.schema_version, 1); + assert_eq!(loaded.providers.len(), 3); + let _ = fs::remove_dir_all(directory); + } +} diff --git a/src/waybar.rs b/src/waybar.rs new file mode 100644 index 0000000..2a4945f --- /dev/null +++ b/src/waybar.rs @@ -0,0 +1,258 @@ +use chrono::{DateTime, Utc}; + +use crate::formatting::{self, bar_number, level}; +use crate::model::{Provider, ProviderState, UsageSnapshot, UsageWindow}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct WaybarView { + pub text: String, + pub tooltip: String, + pub class: String, +} + +pub fn render(snapshot: &UsageSnapshot, now: DateTime, enabled: &[Provider]) -> WaybarView { + let mut parts = Vec::new(); + let mut tooltip = Vec::new(); + let mut worst = "unknown"; + let mut severity = level_rank("unknown"); + let mut any = false; + if snapshot.is_demo { + tooltip.push("DEMO".to_string()); + } + for provider in Provider::ALL { + if !enabled.contains(&provider) { + continue; + } + let state = snapshot.state(provider); + if state.status == crate::model::FetchStatus::Disabled { + continue; + } + any = true; + let primary = displayed_window(provider, &state); + parts.push(format!( + "{} {}", + provider.bar_label(), + bar_number(primary.as_ref(), snapshot.show_used_percentage) + )); + let remaining = primary.as_ref().and_then(|window| window.remaining_percentage); + let mut provider_level = level(remaining); + if state.is_stale(now) && provider_level == "ok" { + provider_level = "warning"; + } + if level_rank(provider_level) > severity { + severity = level_rank(provider_level); + worst = provider_level; + } + tooltip.push(provider_block(provider, &state, snapshot.show_used_percentage, now)); + } + if !any { + return WaybarView { + text: "AI usage off".to_string(), + tooltip: "Every provider is disabled.".to_string(), + class: "unknown".to_string(), + }; + } + if tooltip.iter().all(|line| line.contains("Awaiting first update")) && parts.iter().all(|part| part.ends_with("—")) { + worst = "unknown"; + } + WaybarView { + text: parts.join(" "), + tooltip: tooltip.join("\n"), + class: worst.to_string(), + } +} + +pub fn json(view: &WaybarView) -> String { + serde_json::json!({ + "text": view.text, + "tooltip": view.tooltip, + "class": view.class, + }) + .to_string() +} + +fn level_rank(class: &str) -> u8 { + match class { + "critical" => 3, + "warning" => 2, + "ok" => 1, + _ => 0, + } +} + +fn displayed_window(provider: Provider, state: &ProviderState) -> Option { + let usage = state.last_successful_usage.as_ref()?; + match provider { + Provider::Cursor => usage + .display_billing_pools() + .into_iter() + .next() + .map(|pool| pool.window) + .or_else(|| usage.billing_window.clone()), + _ => usage + .short_window + .clone() + .or_else(|| usage.weekly_window.clone()), + } +} + +fn provider_block(provider: Provider, state: &ProviderState, show_used: bool, now: DateTime) -> String { + let mut lines = vec![provider.display_name().to_string()]; + if state.is_stale(now) { + lines.push("Stale".to_string()); + } + if let Some(message) = &state.message { + lines.push(message.clone()); + } + let Some(usage) = &state.last_successful_usage else { + lines.push(state.status.label().to_string()); + return lines.join("\n"); + }; + if provider == Provider::Cursor { + for pool in usage.display_billing_pools() { + lines.push(format!( + "{} {}", + pool.name, + formatting::percentage(Some(&pool.window), show_used) + )); + } + if usage.display_billing_pools().is_empty() { + if let Some(window) = &usage.billing_window { + lines.push(format!("Included {}", formatting::percentage(Some(window), show_used))); + } + } + lines.push( + usage + .billing_reset_label + .as_ref() + .map(|label| format!("Resets {label}")) + .unwrap_or_else(|| "Reset unknown".to_string()), + ); + for metric in usage.display_additional_usage() { + lines.push(format!("{}: {}", metric.name, metric.value)); + } + } else { + lines.push(window_line("5h", usage.short_window.as_ref(), show_used, now)); + lines.push(window_line("Week", usage.weekly_window.as_ref(), show_used, now)); + for metric in &usage.additional_usage { + lines.push(format!("{}: {}", metric.name, metric.value)); + } + } + lines.push(format!("Updated {}", formatting::clock(Some(usage.fetched_at)))); + lines.join("\n") +} + +fn window_line(label: &str, window: Option<&UsageWindow>, show_used: bool, now: DateTime) -> String { + let percent = formatting::percentage(window, show_used); + let reset = window + .and_then(|window| window.reset_at) + .map(|date| formatting::reset(Some(date), now, true)) + .unwrap_or_else(|| "—".to_string()); + let relative = formatting::relative_reset(window.and_then(|window| window.reset_at), now); + format!("{label} {percent} {reset} {relative}") +} + +pub fn status_text(snapshot: &UsageSnapshot, active: bool, now: DateTime, enabled: &[Provider]) -> String { + let view = render(snapshot, now, enabled); + format!( + "{}\n{}\n{}", + if active { "Updates active" } else { "Updates paused" }, + view.text, + view.tooltip + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::model::{BillingUsagePool, FetchStatus, ProviderState, ProviderUsage}; + + #[test] + fn bar_uses_remaining_percentages_and_the_worst_level() { + let now = Utc::now(); + let snapshot = UsageSnapshot { + schema_version: 1, + written_at: now, + show_used_percentage: false, + is_demo: false, + providers: vec![ + state(Provider::Claude, Some(usage_window(Provider::Claude, 62.0, now))), + state(Provider::Codex, Some(usage_window(Provider::Codex, 29.0, now))), + cursor_state(now), + ], + }; + let view = render(&snapshot, now, &Provider::ALL); + assert_eq!(view.text, "C 38% X 71% Cu 52%"); + assert_eq!(view.class, "warning"); + assert!(view.tooltip.contains("Cursor models")); + assert!(view.tooltip.contains("Resets Oct 30")); + } + + #[test] + fn bar_hides_providers_that_are_not_enabled() { + let now = Utc::now(); + let snapshot = UsageSnapshot { + schema_version: 1, + written_at: now, + show_used_percentage: false, + is_demo: false, + providers: vec![ + state(Provider::Claude, Some(usage_window(Provider::Claude, 62.0, now))), + state(Provider::Codex, Some(usage_window(Provider::Codex, 29.0, now))), + cursor_state(now), + ], + }; + let view = render(&snapshot, now, &[Provider::Cursor]); + assert_eq!(view.text, "Cu 52%"); + assert!(!view.tooltip.contains("Claude")); + assert!(view.tooltip.contains("Cursor")); + } + + fn state(provider: Provider, usage: Option) -> ProviderState { + ProviderState { + provider, + last_successful_usage: usage, + last_attempt_at: Some(Utc::now()), + status: FetchStatus::Available, + message: None, + } + } + + fn usage_window(provider: Provider, used: f64, now: chrono::DateTime) -> ProviderUsage { + ProviderUsage { + provider, + fetched_at: now, + short_window: Some(UsageWindow::new(Some(used), None, None)), + weekly_window: None, + billing_window: None, + billing_reset_label: None, + billing_pools: None, + additional_usage: Vec::new(), + } + } + + fn cursor_state(now: chrono::DateTime) -> ProviderState { + state( + Provider::Cursor, + Some(ProviderUsage { + provider: Provider::Cursor, + fetched_at: now, + short_window: None, + weekly_window: None, + billing_window: Some(UsageWindow::new(Some(44.0), None, None)), + billing_reset_label: Some("Oct 30".to_string()), + billing_pools: Some(vec![ + BillingUsagePool { + name: "Cursor models".to_string(), + window: UsageWindow::new(Some(48.0), None, None), + }, + BillingUsagePool { + name: "Other models".to_string(), + window: UsageWindow::new(Some(0.0), None, None), + }, + ]), + additional_usage: vec![], + }), + ) + } +}