Vectrum
Sign In
304 lines · 12.5 KB · Rust
//! The bar under the composer, like T3 Code's: how much of the plan's
//! limits is used, Claude's or Codex's as Settings → General → Usage bar
//! says, and how full the context window is and how fast the last response
//! came out. Clicking the limits shows them in full and switches whose
//! they are.

use std::time::{Duration, Instant, SystemTime};

use chrono::{DateTime, Local};
use gpui_kit::component::{h_flex, v_flex};
use gpui_kit::prelude::FluentBuilder as _;
use gpui_kit::*;

use super::{ChatView, Picker, format_tokens};
use crate::models::{Catalog, Provider};
use crate::settings::{self, UsageBar};
use crate::ui::composer;
use crate::ui::metal::{self, MetalButton, Palette};
use crate::ui::pickers;
use crate::usage::{self, CodexSources, Limit};

/// How old limits get before they're asked for again while the bar shows them.
const STALE: Duration = Duration::from_secs(180);

/// Codex's limits, asked of ChatGPT in the background.
#[derive(Default)]
pub struct CodexLimits {
    limits: Option<Vec<Limit>>,
    error: Option<String>,
    fetched: Option<Instant>,
    fetching: bool,
}

impl Global for CodexLimits {}

/// Keeps the limits the bar shows fresh: Codex's from ChatGPT, Claude's
/// from a health check of the Claude instance.
pub fn refresh_limits(force: bool, cx: &mut App) {
    match settings::get(|settings| settings.usage_bar) {
        UsageBar::Off => {}
        UsageBar::Claude => {
            let stale = |checked: Option<SystemTime>| {
                force || checked.and_then(|at| at.elapsed().ok()).is_none_or(|age| age >= STALE)
            };
            let id = Catalog::global(cx)
                .instances
                .iter()
                .find(|state| state.kind() == Provider::Claude && state.config.enabled)
                .filter(|state| stale(state.checked))
                .map(|state| state.config.id.clone());
            if let Some(id) = id {
                Catalog::refresh(&id, cx);
            }
        }
        UsageBar::Codex => {
            let state = cx.default_global::<CodexLimits>();
            if state.fetching || (!force && state.fetched.is_some_and(|at| at.elapsed() < STALE)) {
                return;
            }
            state.fetching = true;
            let sources = CodexSources::with_pi(pi_dirs(cx));
            let fetch = cx.background_executor().spawn(async move { usage::codex_limits(&sources) });
            cx.spawn(async move |cx| {
                let result = fetch.await;
                cx.update(|cx| {
                    let state = cx.default_global::<CodexLimits>();
                    state.fetching = false;
                    state.fetched = Some(Instant::now());
                    match result {
                        Ok(limits) => {
                            state.limits = Some(limits);
                            state.error = None;
                        }
                        Err(error) => state.error = Some(error),
                    }
                    cx.refresh_windows();
                })
            })
            .detach();
        }
    }
}

/// The Pi instances' home folders, where a Codex login may live.
fn pi_dirs(cx: &App) -> Vec<std::path::PathBuf> {
    let mut dirs: Vec<_> = Catalog::global(cx)
        .instances
        .iter()
        .filter(|state| state.kind() == Provider::Pi)
        .map(|state| match state.config.home.trim() {
            "" => settings::expand_home("~/.pi/agent"),
            home => settings::expand_home(home),
        })
        .collect();
    if dirs.is_empty() {
        dirs.push(settings::expand_home("~/.pi/agent"));
    }
    dirs.dedup();
    dirs
}

/// The limits to show, or why there are none.
enum Shown {
    Limits(Vec<Limit>),
    Loading,
    Missing(String),
}

fn shown_limits(source: UsageBar, instance: &str, cx: &App) -> Shown {
    match source {
        UsageBar::Off => Shown::Missing(String::new()),
        UsageBar::Claude => {
            let claude = Catalog::global(cx).instances.iter().filter(|state| state.kind() == Provider::Claude);
            // The thread's own Claude account, else the first that reports limits.
            let state = claude.clone().find(|state| state.config.id == instance && state.usage.is_some()).or_else(|| {
                claude.clone().find(|state| state.usage.is_some())
            });
            match state.and_then(|state| state.usage.as_ref()) {
                Some(reply) => Shown::Limits(usage::limits_from_claude(reply)),
                None if claude.clone().any(|state| state.checked.is_none() && state.config.enabled) => Shown::Loading,
                None => Shown::Missing("Claude reports no plan limits for this login.".into()),
            }
        }
        UsageBar::Codex => match cx.try_global::<CodexLimits>() {
            Some(CodexLimits { limits: Some(limits), .. }) => Shown::Limits(limits.clone()),
            Some(CodexLimits { error: Some(error), .. }) => Shown::Missing(error.clone()),
            _ => Shown::Loading,
        },
    }
}

fn source_name(source: UsageBar) -> &'static str {
    match source {
        UsageBar::Codex => "Codex",
        _ => "Claude",
    }
}

fn limit_color(limit: &Limit, p: &Palette) -> Hsla {
    match limit.severity.as_str() {
        "critical" => p.red,
        "warning" => p.amber,
        _ => p.ink3,
    }
}

/// `5h ▬▬ 12%`.
fn limit_chip(limit: &Limit, p: &Palette) -> Div {
    let fraction = (limit.percent / 100.).clamp(0., 1.) as f32;
    h_flex()
        .gap(px(6.))
        .child(div().text_color(p.ink3).child(limit.short.clone().unwrap_or_default()))
        .child(
            div()
                .w(px(44.))
                .h(px(4.))
                .rounded_full()
                .bg(p.ink3.opacity(0.22))
                .child(div().h_full().rounded_full().bg(limit_color(limit, p)).w(relative(fraction))),
        )
        .child(div().text_color(p.ink2).child(format!("{:.0}%", limit.percent)))
}

/// `Resets in 2h 14m`.
fn resets_text(at: DateTime<Local>, now: DateTime<Local>) -> String {
    let left = at.signed_duration_since(now);
    if left.num_seconds() <= 0 {
        "Reset".into()
    } else if left.num_hours() >= 24 {
        format!("Resets in {}d {}h", left.num_days(), left.num_hours() % 24)
    } else if left.num_minutes() >= 60 {
        format!("Resets in {}h {}m", left.num_hours(), left.num_minutes() % 60)
    } else {
        format!("Resets in {}m", left.num_minutes().max(1))
    }
}

impl ChatView {
    /// The bar under the composer, unless Settings → General → Usage bar
    /// turns it off.
    pub(super) fn render_usage_bar(&self, cx: &mut Context<Self>) -> Option<AnyElement> {
        let source = settings::get(|settings| settings.usage_bar);
        if source == UsageBar::Off {
            return None;
        }
        let p = metal::palette(cx);
        let session = self.session.read(cx);
        let (tokens, window, speed) = (session.context_tokens, session.context_window, session.tokens_per_sec);

        let leading = match shown_limits(source, &session.instance, cx) {
            Shown::Limits(limits) if limits.iter().any(|limit| limit.short.is_some()) => h_flex()
                .gap(px(14.))
                .children(limits.iter().filter(|limit| limit.short.is_some()).map(|limit| limit_chip(limit, p))),
            Shown::Limits(_) | Shown::Missing(_) => {
                h_flex().text_color(p.ink3).child(format!("{} limits unavailable", source_name(source)))
            }
            Shown::Loading => h_flex().text_color(p.ink3).child(format!("{} limits…", source_name(source))),
        };
        let trigger = MetalButton::new("usage-limits")
            .ghost()
            .mini()
            .w_auto()
            .h(px(22.))
            .px(px(6.))
            .text_size(px(12.))
            .leading(leading)
            .tooltip(format!("{} plan limits · click for details", source_name(source)));
        let limits = self.picker_anchored(Picker::Usage, trigger, Anchor::BottomLeft, cx, |this, cx| {
            this.render_usage_menu(cx)
        });

        let figure = |value: String, unit: &'static str| {
            h_flex().gap(px(4.)).child(div().text_color(p.ink2).child(value)).child(div().text_color(p.ink3).child(unit))
        };
        let context = match (tokens, window) {
            (Some(tokens), Some(window)) => {
                Some(figure(format!("{:.0}%", tokens as f64 * 100. / window.max(1) as f64), "context"))
            }
            (Some(tokens), None) => Some(figure(format_tokens(tokens), "context")),
            _ => None,
        };
        let speed = speed.map(|speed| figure(format!("{speed:.0}"), "tok/s"));
        let dot = (context.is_some() && speed.is_some()).then(|| div().text_color(p.ink3.opacity(0.6)).child("·"));

        Some(
            h_flex()
                .w_full()
                .justify_between()
                .gap(px(12.))
                .pt(px(6.))
                .pl(px(14.))
                .pr(px(20.))
                .text_size(px(12.))
                .child(limits)
                .child(h_flex().flex_none().gap(px(8.)).children(context).children(dot).children(speed))
                .into_any_element(),
        )
    }

    /// Each limit with when it resets, and whose limits the bar shows.
    fn render_usage_menu(&mut self, cx: &mut Context<Self>) -> AnyElement {
        let p = metal::palette(cx);
        let source = settings::get(|settings| settings.usage_bar);
        let shown = shown_limits(source, &self.session.read(cx).instance, cx);
        let now = Local::now();
        let mut menu = pickers::panel(p).w(px(300.)).child(composer::menu_label(format!("{} plan limits", source_name(source)), p));
        match shown {
            Shown::Limits(limits) if !limits.is_empty() => {
                for limit in &limits {
                    menu = menu.child(
                        v_flex()
                            .px(px(8.))
                            .py(px(6.))
                            .gap(px(5.))
                            .text_size(px(12.5))
                            .child(
                                h_flex()
                                    .justify_between()
                                    .child(div().text_color(p.ink).child(limit.label.clone()))
                                    .child(div().text_color(p.ink2).child(format!("{:.0}% used", limit.percent))),
                            )
                            .child(
                                div().h(px(5.)).w_full().rounded_full().bg(p.ink3.opacity(0.18)).child(
                                    div()
                                        .h_full()
                                        .rounded_full()
                                        .bg(limit_color(limit, p))
                                        .w(relative((limit.percent / 100.).clamp(0., 1.) as f32)),
                                ),
                            )
                            .when_some(limit.resets_at, |this, at| {
                                this.child(div().text_size(px(11.5)).text_color(p.ink3).child(resets_text(at, now)))
                            }),
                    );
                }
            }
            Shown::Loading => menu = menu.child(note("Checking…", p)),
            Shown::Limits(_) => menu = menu.child(note("This login reports no limits.", p)),
            Shown::Missing(why) => menu = menu.child(note(why, p)),
        }
        menu = menu.child(composer::menu_separator(p)).child(composer::menu_label("Show limits from", p));
        for (option, label) in [(UsageBar::Claude, "Claude"), (UsageBar::Codex, "Codex")] {
            menu = menu.child(composer::menu_item(SharedString::from(format!("usage-source-{label}")), option == source, true, p).child(label).on_click(
                cx.listener(move |this, _, window, cx| {
                    settings::update(|settings| settings.usage_bar = option);
                    refresh_limits(false, cx);
                    this.set_picker(Picker::Usage, false, window, cx);
                }),
            ));
        }
        menu.child(composer::menu_separator(p))
            .child(composer::menu_item("usage-bar-hide", false, true, p).child("Hide usage bar").on_click(cx.listener(
                |this, _, window, cx| {
                    settings::update(|settings| settings.usage_bar = UsageBar::Off);
                    this.set_picker(Picker::Usage, false, window, cx);
                },
            )))
            .into_any_element()
    }
}

fn note(text: impl Into<SharedString>, p: &Palette) -> Div {
    div().px(px(8.)).py(px(6.)).text_size(px(12.)).text_color(p.ink3).child(text.into())
}