ai-usagebar 1.4.0

Omarchy/Waybar widgets + TUI for tracking multi-provider AI plan usage
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

//! Waybar widget binary. The library does all the work — this is just the
//! tokio bootstrap + clap parse.

use ai_usagebar::widget::cli::{AuthProvider, Cli, Command, NousAuthAction};
use ai_usagebar::widget::run::run;
use clap::Parser;

fn main() {
    let cli = Cli::parse();
    if let Some(Command::Account { action }) = &cli.command {
        std::process::exit(ai_usagebar::account::run(action));
    }
    if let Some(Command::Settings { action }) = &cli.command {
        std::process::exit(ai_usagebar::tui::settings::run_cli(action));
    }
    if let Some(Command::Auth { provider }) = &cli.command {
        let rt = match tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
        {
            Ok(rt) => rt,
            Err(_) => std::process::exit(1),
        };
        let code = match provider {
            AuthProvider::Nous { action } => rt.block_on(run_nous_auth(action)),
        };
        std::process::exit(code);
    }
    let rt = match tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
    {
        Ok(rt) => rt,
        Err(_) => {
            // Catastrophic — emit the always-valid ⚠ JSON and exit 0.
            println!(
                r#"{{"text":"⚠","tooltip":"failed to create tokio runtime","class":"critical"}}"#
            );
            std::process::exit(0);
        }
    };
    // An administrative report, not the widget: it needs the runtime but must
    // not go through the always-exit-0 Waybar contract — a script piping this
    // deserves a real exit code.
    if let Some(Command::Usage { json }) = &cli.command {
        std::process::exit(rt.block_on(ai_usagebar::report::run(*json)));
    }
    let code = rt.block_on(run(cli));
    std::process::exit(code);
}

async fn run_nous_auth(action: &NousAuthAction) -> i32 {
    let store = ai_usagebar::nous::credentials::CredentialStore::default();
    match action {
        NousAuthAction::Logout => match store.logout() {
            Ok(()) => {
                println!("Nous Research logout complete");
                0
            }
            Err(error) => {
                eprintln!("Nous Research logout failed: {error}");
                1
            }
        },
        NousAuthAction::Login => {
            let client = match reqwest::Client::builder()
                .timeout(ai_usagebar::vendor::HTTP_CLIENT_TIMEOUT)
                .redirect(ai_usagebar::vendor::same_origin_redirect_policy())
                .build()
            {
                Ok(client) => client,
                Err(_) => {
                    eprintln!("Nous Research login failed: could not initialize HTTP client");
                    return 1;
                }
            };
            let endpoints = ai_usagebar::nous::oauth::Endpoints::default();
            let device =
                match ai_usagebar::nous::oauth::request_device_code(&client, &endpoints).await {
                    Ok(device) => device,
                    Err(error) => {
                        eprintln!("Nous Research login failed: {error}");
                        return 1;
                    }
                };
            println!("Open {}", device.verification_uri);
            println!("Verification code: {}", device.user_code);
            let opener = ai_usagebar::nous::oauth::SystemBrowserOpener;
            if !ai_usagebar::nous::oauth::open_verification_url(
                &device.verification_uri_complete,
                &opener,
            ) {
                eprintln!("Browser opener unavailable; open the URL manually.");
            }
            let token =
                match ai_usagebar::nous::oauth::poll_for_token(&client, &endpoints.token, &device)
                    .await
                {
                    Ok(token) => token,
                    Err(error) => {
                        eprintln!("Nous Research login failed: {error}");
                        return 1;
                    }
                };
            let credential =
                match ai_usagebar::nous::oauth::credential_from_token(token, chrono::Utc::now()) {
                    Ok(credential) => credential,
                    Err(error) => {
                        eprintln!("Nous Research login failed: {error}");
                        return 1;
                    }
                };
            match ai_usagebar::nous::oauth::persist_credential(&store, credential) {
                Ok(()) => {
                    println!("Nous Research login complete");
                    0
                }
                Err(error) => {
                    eprintln!("Nous Research login failed: {error}");
                    1
                }
            }
        }
    }
}