Initial commit
This commit is contained in:
Armel van Ravels committed
commit
69605d7e54
27 files changed
+5188
No files matched your search
@@ -0,0 +1,3 @@
|
||||
/target
|
||||
/result
|
||||
.direnv
|
||||
Generated
+1125
File diff suppressed because it is too large.
Load diff
+20
@@ -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
|
||||
@@ -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
|
||||
```
|
||||
@@ -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" }
|
||||
}
|
||||
}
|
||||
@@ -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 } }
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
@@ -0,0 +1 @@
|
||||
active
|
||||
Generated
+27
@@ -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
|
||||
}
|
||||
@@ -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;
|
||||
};
|
||||
}
|
||||
@@ -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; }
|
||||
'';
|
||||
})
|
||||
]);
|
||||
}
|
||||
+215
@@ -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<bool>,
|
||||
pub session_active: Option<bool>,
|
||||
pub unlocked: Option<bool>,
|
||||
}
|
||||
|
||||
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<bool>,
|
||||
pub state: Option<String>,
|
||||
}
|
||||
|
||||
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<Activity>) {
|
||||
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<Activity>) -> 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<S>(stream: &mut Option<S>) -> Option<S::Item>
|
||||
where
|
||||
S: Stream + Unpin,
|
||||
{
|
||||
match stream.as_mut() {
|
||||
Some(stream) => stream.next().await,
|
||||
None => std::future::pending().await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn publish(sender: &watch::Sender<Activity>, 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<bool>;
|
||||
|
||||
#[zbus(property)]
|
||||
fn state(&self) -> zbus::Result<String>;
|
||||
|
||||
#[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());
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
+294
@@ -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<AtomicBool>,
|
||||
config_stamp: Option<std::time::SystemTime>,
|
||||
}
|
||||
|
||||
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<tokio::sync::Mutex<Daemon>>) {
|
||||
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<tokio::sync::Mutex<Daemon>>, 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<tokio::sync::Mutex<Daemon>>, notify: Arc<Notify>) {
|
||||
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<tokio::sync::Mutex<Daemon>>,
|
||||
notify: Arc<Notify>,
|
||||
) -> 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<tokio::sync::Mutex<Daemon>>, started: chrono::DateTime<Utc>) -> 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::<i32>() 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<String> {
|
||||
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)
|
||||
}
|
||||
|
||||
@@ -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<ExecutableCandidate> {
|
||||
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::<Vec<_>>();
|
||||
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<std::str::Chars<'_>>) -> 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);
|
||||
}
|
||||
}
|
||||
@@ -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<DateTime<Utc>>, now: DateTime<Utc>, 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<DateTime<Utc>>, now: DateTime<Utc>) -> 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<DateTime<Utc>>) -> 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<Utc>> {
|
||||
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<f64>) -> &'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), "—");
|
||||
}
|
||||
}
|
||||
+66
@@ -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 <daemon|waybar|refresh|status>\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);
|
||||
}
|
||||
}
|
||||
}
|
||||
+355
@@ -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<f64>) -> Option<f64> {
|
||||
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<f64>,
|
||||
pub remaining_percentage: Option<f64>,
|
||||
#[serde(default, with = "optional_rfc3339")]
|
||||
pub reset_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
impl UsageWindow {
|
||||
pub fn new(
|
||||
used_percentage: Option<f64>,
|
||||
remaining_percentage: Option<f64>,
|
||||
reset_at: Option<DateTime<Utc>>,
|
||||
) -> 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<Utc>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub short_window: Option<UsageWindow>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub weekly_window: Option<UsageWindow>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub billing_window: Option<UsageWindow>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub billing_reset_label: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub billing_pools: Option<Vec<BillingUsagePool>>,
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub additional_usage: Vec<UsageMetric>,
|
||||
}
|
||||
|
||||
impl ProviderUsage {
|
||||
pub fn display_billing_pools(&self) -> Vec<BillingUsagePool> {
|
||||
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::<f64>().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<UsageMetric> {
|
||||
if self.provider != Provider::Cursor {
|
||||
return self.additional_usage.clone();
|
||||
}
|
||||
let pools: Vec<String> = 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<String>) -> 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<ProviderUsage>,
|
||||
#[serde(default, with = "optional_rfc3339")]
|
||||
pub last_attempt_at: Option<DateTime<Utc>>,
|
||||
pub status: FetchStatus,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub message: Option<String>,
|
||||
}
|
||||
|
||||
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<Utc>) -> 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<Utc>,
|
||||
pub show_used_percentage: bool,
|
||||
pub is_demo: bool,
|
||||
pub providers: Vec<ProviderState>,
|
||||
}
|
||||
|
||||
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<S>(value: &DateTime<Utc>, serializer: S) -> Result<S::Ok, S::Error>
|
||||
where
|
||||
S: Serializer,
|
||||
{
|
||||
serializer.serialize_str(&value.to_rfc3339_opts(chrono::SecondsFormat::Secs, true))
|
||||
}
|
||||
|
||||
pub fn deserialize<'de, D>(deserializer: D) -> Result<DateTime<Utc>, 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<S>(value: &Option<DateTime<Utc>>, serializer: S) -> Result<S::Ok, S::Error>
|
||||
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<Option<DateTime<Utc>>, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
{
|
||||
let value = Option::<String>::deserialize(deserializer)?;
|
||||
value
|
||||
.map(|text| {
|
||||
DateTime::parse_from_rfc3339(&text)
|
||||
.map(|parsed| parsed.with_timezone(&Utc))
|
||||
.map_err(serde::de::Error::custom)
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
}
|
||||
+442
@@ -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<Utc>) -> Result<ProviderUsage, ProviderFailure> {
|
||||
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<Option<UsageWindow>, 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<bool>,
|
||||
rate_limits: Option<ClaudeLimits>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ClaudeLimits {
|
||||
five_hour: Option<ClaudeWindow>,
|
||||
seven_day: Option<ClaudeWindow>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ClaudeWindow {
|
||||
utilization: Option<f64>,
|
||||
resets_at: Option<String>,
|
||||
}
|
||||
|
||||
pub fn parse_codex(data: &[u8], fetched_at: DateTime<Utc>) -> Result<ProviderUsage, ProviderFailure> {
|
||||
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<CodexBucket>,
|
||||
#[serde(rename = "rateLimitsByLimitId")]
|
||||
rate_limits_by_limit_id: Option<std::collections::BTreeMap<String, CodexBucket>>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct CodexBucket {
|
||||
#[serde(rename = "limitId")]
|
||||
limit_id: Option<String>,
|
||||
primary: Option<CodexWindow>,
|
||||
secondary: Option<CodexWindow>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct CodexWindow {
|
||||
#[serde(rename = "usedPercent")]
|
||||
used_percent: f64,
|
||||
#[serde(rename = "windowDurationMins")]
|
||||
window_duration_mins: Option<i64>,
|
||||
#[serde(rename = "resetsAt")]
|
||||
resets_at: Option<f64>,
|
||||
}
|
||||
|
||||
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<Utc>) -> Result<ProviderUsage, ProviderFailure> {
|
||||
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::<Vec<_>>()
|
||||
.into_iter()
|
||||
.rev()
|
||||
.collect::<Vec<_>>()
|
||||
.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<Option<f64>, 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::<f64>().ok().filter(|value| value.is_finite());
|
||||
Ok(number)
|
||||
}
|
||||
|
||||
fn capture(pattern: &str, value: &str) -> Option<String> {
|
||||
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<regex::Regex> =
|
||||
std::sync::LazyLock::new(|| regex::Regex::new($pattern).expect("static regex"));
|
||||
};
|
||||
}
|
||||
pub fn subscription_status(text: &str) -> Option<UsageMetric> {
|
||||
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<u8> {
|
||||
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");
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
+346
@@ -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<u8>,
|
||||
pub stderr: Vec<u8>,
|
||||
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<u8>),
|
||||
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<ProcessResult, ProcessError> {
|
||||
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<libc::pollfd> = 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<OwnedFd, ProcessError> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
+415
@@ -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<u8>,
|
||||
usage_sent: bool,
|
||||
pub usage_data: Option<Vec<u8>>,
|
||||
pub failure: Option<ProviderFailure>,
|
||||
}
|
||||
|
||||
#[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<u8> {
|
||||
line(&json!({
|
||||
"type": "control_request",
|
||||
"request_id": "ai-usage-init",
|
||||
"request": { "subtype": "initialize" }
|
||||
}))
|
||||
}
|
||||
|
||||
pub fn codex_initialize() -> Vec<u8> {
|
||||
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::<Vec<_>>();
|
||||
let bytes = &raw[..raw.len().saturating_sub(1)];
|
||||
let Ok(message) = serde_json::from_slice::<Value>(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<u8> {
|
||||
line_value(value)
|
||||
}
|
||||
|
||||
fn line_value(value: &Value) -> Vec<u8> {
|
||||
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<u8>,
|
||||
pending_input: Option<Vec<u8>>,
|
||||
last_output_at: Duration,
|
||||
workspace: String,
|
||||
workspace_canonical: String,
|
||||
pub usage: Option<ProviderUsage>,
|
||||
pub failure: Option<ProviderFailure>,
|
||||
}
|
||||
|
||||
impl CursorUsageSession {
|
||||
pub fn new(workspace: impl Into<String>) -> 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<String> = command
|
||||
.split(|byte| *byte == b'\n')
|
||||
.filter(|line| !line.is_empty())
|
||||
.map(|line| {
|
||||
serde_json::from_slice::<Value>(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"));
|
||||
}
|
||||
}
|
||||
@@ -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<ProviderUsage, ProviderFailure> {
|
||||
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<ProviderUsage, ProviderFailure> {
|
||||
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<ProviderUsage, ProviderFailure> {
|
||||
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<ProviderUsage, ProviderFailure> {
|
||||
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<ProviderUsage, ProviderFailure> {
|
||||
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<String> = 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<ProviderUsage, ProviderFailure> {
|
||||
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<process::ProcessResult, ProcessError>) -> 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<String> {
|
||||
values.iter().map(|value| (*value).to_string()).collect()
|
||||
}
|
||||
|
||||
fn which(name: &str) -> Option<PathBuf> {
|
||||
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("/"))
|
||||
}
|
||||
+312
@@ -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<Provider>,
|
||||
}
|
||||
|
||||
pub struct Round {
|
||||
generation: u64,
|
||||
pub providers: Vec<Provider>,
|
||||
attempted: DateTime<Utc>,
|
||||
forced: HashSet<Provider>,
|
||||
}
|
||||
|
||||
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<Utc>) -> Vec<Provider> {
|
||||
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<Utc>) -> Option<Round> {
|
||||
let mut selected: Vec<Provider> = 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::<HashSet<_>>();
|
||||
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<Provider, Result<ProviderUsage, ProviderFailure>>,
|
||||
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<Utc>) -> Option<chrono::Duration> {
|
||||
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<Utc> {
|
||||
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<Utc> {
|
||||
DateTime::from_timestamp(seconds, 0).unwrap()
|
||||
}
|
||||
|
||||
fn cached(when: DateTime<Utc>) -> 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);
|
||||
}
|
||||
}
|
||||
@@ -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<Provider>,
|
||||
#[serde(default)]
|
||||
pub executable_overrides: HashMap<Provider, String>,
|
||||
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::<Settings>(&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)
|
||||
}
|
||||
}
|
||||
@@ -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<Self, StoreError> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
+258
@@ -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<Utc>, 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<UsageWindow> {
|
||||
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<Utc>) -> 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<Utc>) -> 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<Utc>, 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<ProviderUsage>) -> 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<Utc>) -> 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<Utc>) -> 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![],
|
||||
}),
|
||||
)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user