feat: add CLI Proxy API quota segment with forward-compatible theme s… #64
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Reviewer's GuideImplements a new CLI Proxy API Quota statusline segment with configurable aliases/colors and caching, wires it through core, preview, and TUI editors, and enhances the theme system to share the segment definition across all presets with forward-compatible segment merging and improved text input handling. Sequence diagram for editing CLI Proxy API quota alias via TUIsequenceDiagram
actor User
participant App
participant CliProxyApiQuotaOptionsComponent as CliProxyOptions
participant NameInputComponent as NameInput
participant Config
participant PreviewComponent as Preview
User->>App: Navigate to CliProxyApiQuota segment
User->>App: Press Tab to Options field
User->>App: Press Enter
App->>CliProxyOptions: open()
App->>App: status_message = "Editing CPA Quota options"
loop Navigate options
User->>CliProxyOptions: Up/Down keys (via App event loop)
CliProxyOptions-->>App: move_selection(delta)
end
User->>App: Press Enter on Alias field
App->>CliProxyOptions: selected_field()
CliProxyOptions-->>App: Alias(CliProxyApiQuotaModelKind)
App->>Config: read current segment.options[alias_key]
App->>NameInput: open_with_value(title, placeholder, current_alias)
App->>App: text_input_target = CliProxyApiQuotaAlias(model)
User->>NameInput: Type new alias
User->>App: Press Enter
App->>NameInput: get_input()
NameInput-->>App: value
App->>App: apply_text_input(value)
App->>Config: update segment.options[alias_key] = value
App->>App: status_message = "Updated <model> alias"
App->>Preview: update_preview(&config)
App->>NameInput: close()
App->>App: text_input_target = None
User->>App: Press Esc
App->>CliProxyOptions: close()
Sequence diagram for rendering CLI Proxy API quota segment with cachingsequenceDiagram
participant StatusLineGenerator as StatusLineGen
participant CliProxyApiQuotaSegment as CPAQuotaSeg
participant Config
participant FileSystem as FS
StatusLineGen->>CPAQuotaSeg: collect(&InputData)
CPAQuotaSeg->>Config: Config::load()
Config-->>CPAQuotaSeg: Config (includes segments)
CPAQuotaSeg->>CPAQuotaSeg: find CliProxyApiQuota SegmentConfig
CPAQuotaSeg->>CPAQuotaSeg: read options (host,key,cache_duration,auth_type,separator)
CPAQuotaSeg->>FS: load_cache()
FS-->>CPAQuotaSeg: Option<CliProxyApiQuotaCache>
alt valid cache
CPAQuotaSeg->>CPAQuotaSeg: is_cache_valid(cache, cache_duration)
CPAQuotaSeg-->>CPAQuotaSeg: use cached quotas
else missing or expired cache
CPAQuotaSeg->>CPAQuotaSeg: fetch_all_quotas(host,key,auth_type)
CPAQuotaSeg-->>CPAQuotaSeg: quotas
alt quotas not empty
CPAQuotaSeg->>FS: save_cache(CliProxyApiQuotaCache)
else fetch failed
CPAQuotaSeg->>CPAQuotaSeg: fallback to cached quotas or empty
end
end
alt quotas empty
CPAQuotaSeg-->>StatusLineGen: None
else quotas present
CPAQuotaSeg->>CPAQuotaSeg: format_tracked_output(quotas, options, separator)
CPAQuotaSeg-->>StatusLineGen: Some(SegmentData{primary, metadata[raw_text=true]})
end
StatusLineGen->>StatusLineGen: render_segment(config, data - respects raw_text flag)
Class diagram for CLI Proxy API quota segment and optionsclassDiagram
direction LR
class CliProxyApiQuotaSegment {
+new() CliProxyApiQuotaSegment
+collect(input: InputData) Option~SegmentData~
+id() SegmentId
-normalize_model_text(text: &str) String
-tracked_model_for(model_id: &str, display_name: &str) Option~TrackedModel~
-tracked_model_for_quota(quota: &ModelQuota) Option~TrackedModel~
-get_alias(options: &HashMap~String, Value~, model: TrackedModel) String
-get_color(options: &HashMap~String, Value~, model: TrackedModel) AnsiColor
-apply_foreground_color(text: &str, color: &AnsiColor) String
-format_tracked_output(quotas: &Vec~ModelQuota~, options: &HashMap~String, Value~, separator: &str) String
-get_cache_path() Option~PathBuf~
-load_cache() Option~CliProxyApiQuotaCache~
-save_cache(cache: &CliProxyApiQuotaCache) void
-is_cache_valid(cache: &CliProxyApiQuotaCache, cache_duration: u64) bool
-get_auth_files(host: &str, key: &str) Option~Vec~AuthFile~~
-api_call(host: &str, key: &str, auth_index: &str, method: &str, url: &str, data: &str, extra_headers: Option~HashMap~&str,&str~~) Option~ApiCallResponse~
-get_antigravity_quota(host: &str, key: &str, auth_index: &str) Vec~ModelQuota~
-extract_project_from_name(name: &str) Option~String~
-get_gemini_cli_quota(host: &str, key: &str, auth_index: &str, project: &str) Vec~ModelQuota~
-fetch_all_quotas(host: &str, key: &str, auth_type_filter: &str) Vec~ModelQuota~
}
class Segment {
<<trait>>
+collect(input: InputData) Option~SegmentData~
+id() SegmentId
}
class SegmentData {
+primary: String
+secondary: String
+metadata: HashMap~String,String~
}
class TrackedModel {
<<enum>>
Opus
Gemini3Pro
Gemini3Flash
+alias_key() &str
+color_key() &str
+default_alias() &str
+default_color() AnsiColor
}
class CliProxyApiQuotaCache {
+quotas: Vec~ModelQuota~
+cached_at: String
}
class ModelQuota {
+model_id: String
+display_name: String
+remaining_fraction: f64
+auth_type: String
}
class AuthFile {
+auth_type: String
+auth_index: String
+label: Option~String~
+name: Option~String~
+disabled: Option~bool~
}
class AuthFilesResponse {
+files: Vec~AuthFile~
}
class ApiCallResponse {
+body: Option~String~
+error: Option~String~
}
class AntigravityModelsResponse {
+models: Option~HashMap~String,ModelInfo~>
}
class ModelInfo {
+display_name: Option~String~
+quota_info: Option~QuotaInfo~
}
class QuotaInfo {
+remaining_fraction: Option~f64~
}
class GeminiQuotaResponse {
+buckets: Option~Vec~GeminiBucket~~
}
class GeminiBucket {
+model_id: Option~String~
+remaining_fraction: Option~f64~
}
class CliProxyApiQuotaOptionsComponent {
+is_open: bool
-selected: usize
+new() CliProxyApiQuotaOptionsComponent
+open() void
+close() void
+move_selection(delta: i32) void
+selected_field() CliProxyApiQuotaOptionField
+render(f: Frame, area: Rect, config: Config, selected_segment: usize) void
-fields() &[CliProxyApiQuotaOptionField]
-get_alias(options: &HashMap~String,Value~, model: CliProxyApiQuotaModelKind) String
-get_color(options: &HashMap~String,Value~, model: CliProxyApiQuotaModelKind) Option~AnsiColor~
-color_to_desc(color: &Option~AnsiColor~) String
-to_ratatui_color(color: &AnsiColor) Color
}
class CliProxyApiQuotaModelKind {
<<enum>>
Opus
Gemini3Pro
Gemini3Flash
+display_name() &str
+alias_key() &str
+color_key() &str
+default_alias() &str
}
class CliProxyApiQuotaOptionField {
<<enum>>
Alias(CliProxyApiQuotaModelKind)
Color(CliProxyApiQuotaModelKind)
Separator
}
class App {
-config: Config
-cli_proxy_api_quota_options: CliProxyApiQuotaOptionsComponent
-text_input_target: Option~TextInputTarget~
-color_picker_target: Option~ColorPickerTarget~
+apply_text_input(value: String) void
+apply_selected_color(color: AnsiColor) void
+open_cli_proxy_api_quota_option_editor() void
}
class TextInputTarget {
<<enum>>
SaveThemeName
CliProxyApiQuotaAlias(CliProxyApiQuotaModelKind)
CliProxyApiQuotaSeparator
}
class ColorPickerTarget {
<<enum>>
CliProxyApiQuotaModelColor(CliProxyApiQuotaModelKind)
}
class ThemePresets {
+get_theme(theme_name: &str) Config
+builtin_theme(theme_name: &str) Option~Config~
+merge_missing_segments(config: Config, baseline: &Config) Config
+default_cli_proxy_api_quota_segment() SegmentConfig
}
class Config {
+segments: Vec~SegmentConfig~
+theme: String
+load() Result~Config,Error~
}
class SegmentConfig {
+id: SegmentId
+enabled: bool
+icon: IconConfig
+colors: ColorConfig
+styles: TextStyleConfig
+options: HashMap~String,Value~
}
class SegmentId {
<<enum>>
ContextWindow
Cost
Directory
Git
Model
Provider
Session
OutputStyle
Update
CliProxyApiQuota
}
Segment <|.. CliProxyApiQuotaSegment
CliProxyApiQuotaSegment --> SegmentData
CliProxyApiQuotaSegment --> CliProxyApiQuotaCache
CliProxyApiQuotaSegment --> ModelQuota
CliProxyApiQuotaSegment --> TrackedModel
CliProxyApiQuotaSegment --> AuthFilesResponse
CliProxyApiQuotaSegment --> ApiCallResponse
CliProxyApiQuotaSegment --> AntigravityModelsResponse
CliProxyApiQuotaSegment --> GeminiQuotaResponse
CliProxyApiQuotaOptionsComponent --> CliProxyApiQuotaModelKind
CliProxyApiQuotaOptionsComponent --> CliProxyApiQuotaOptionField
CliProxyApiQuotaOptionsComponent --> Config
CliProxyApiQuotaOptionsComponent --> SegmentConfig
App --> CliProxyApiQuotaOptionsComponent
App --> TextInputTarget
App --> ColorPickerTarget
App --> Config
ThemePresets --> SegmentConfig
ThemePresets --> Config
SegmentConfig --> SegmentId
Config --> SegmentConfig
CliProxyApiQuotaSegment --> SegmentId
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 2 issues, and left some high level feedback:
- The
TrackedModellogic incore/segments/cli_proxy_api_quota.rsandCliProxyApiQuotaModelKindplus alias/color key helpers in the TUI options component are essentially duplicated; consider extracting a single shared type/API for model kinds, alias/color keys, and defaults to avoid drift. - The
CliProxyApiQuotaSegment::collectimplementation loadsConfig::load()from disk and performs synchronous HTTP requests on each call, which may block statusline rendering; it would be safer to pass the segment options in from the caller and/or isolate the network/cache refresh on a slower background path with tighter timeouts. - ANSI foreground-color formatting logic for the quota output is duplicated between
CliProxyApiQuotaSegmentand the preview segment builder inPreviewComponent; extracting a small shared helper for applying foreground colors would reduce the chance of inconsistencies (e.g., reset codes) between preview and runtime rendering.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The `TrackedModel` logic in `core/segments/cli_proxy_api_quota.rs` and `CliProxyApiQuotaModelKind` plus alias/color key helpers in the TUI options component are essentially duplicated; consider extracting a single shared type/API for model kinds, alias/color keys, and defaults to avoid drift.
- The `CliProxyApiQuotaSegment::collect` implementation loads `Config::load()` from disk and performs synchronous HTTP requests on each call, which may block statusline rendering; it would be safer to pass the segment options in from the caller and/or isolate the network/cache refresh on a slower background path with tighter timeouts.
- ANSI foreground-color formatting logic for the quota output is duplicated between `CliProxyApiQuotaSegment` and the preview segment builder in `PreviewComponent`; extracting a small shared helper for applying foreground colors would reduce the chance of inconsistencies (e.g., reset codes) between preview and runtime rendering.
## Individual Comments
### Comment 1
<location> `src/core/segments/cli_proxy_api_quota.rs:499-508` </location>
<code_context>
+ fn collect(&self, _input: &InputData) -> Option<SegmentData> {
</code_context>
<issue_to_address>
**issue (performance):** Avoid unnecessary Config load and network work when the segment is disabled, and consider the impact of synchronous HTTP/FS calls with relatively long timeouts.
This `collect` implementation always reads the full `Config` and may perform HTTP requests (5s/10s timeouts) plus JSON (de)serialization on every call. A few concrete refinements:
1. Since you already locate `CliProxyApiQuota` via `find(|s| s.id == SegmentId::CliProxyApiQuota)`, you can early-return when the segment is missing or `!segment.enabled` to skip all IO/networking in that case.
2. On cache misses or expired cache, the synchronous HTTP calls can block prompt rendering for several seconds. You could:
* Lower the default timeouts or make them configurable.
* Introduce a short-lived negative cache/backoff after failures instead of retrying every invocation.
* Reuse config passed via `InputData` (if available) to avoid re-reading it from disk on each call.
</issue_to_address>
### Comment 2
<location> `src/ui/components/preview.rs:211-220` </location>
<code_context>
+ .unwrap_or(default)
+ };
+
+ let apply = |text: &str, c: &crate::config::AnsiColor| -> String {
+ let prefix = match c {
+ crate::config::AnsiColor::Color16 { c16 } => {
+ let code = if *c16 < 8 { 30 + c16 } else { 90 + (c16 - 8) };
+ format!("\x1b[{}m", code)
+ }
+ crate::config::AnsiColor::Color256 { c256 } => {
+ format!("\x1b[38;5;{}m", c256)
+ }
+ crate::config::AnsiColor::Rgb { r, g, b } => {
+ format!("\x1b[38;2;{};{};{}m", r, g, b)
+ }
+ };
</code_context>
<issue_to_address>
**suggestion:** Factor out duplicated ANSI color escape generation logic shared with CliProxyApiQuotaSegment.
This closure duplicates the foreground escape logic in `CliProxyApiQuotaSegment::apply_foreground_color`, which risks subtle divergence if color variants or reset behavior change. Please extract a shared helper (e.g., a small utility in `crate::config`) that converts `AnsiColor` to the correct escape sequence and use it in both places to keep behavior consistent and avoid duplication.
Suggested implementation:
```rust
let apply = |text: &str, c: &crate::config::AnsiColor| -> String {
```
```rust
let prefix = crate::config::ansi_color_escape_sequence(c);
format!("{prefix}{text}\x1b[0m")
};
```
To fully implement the refactor and eliminate duplication:
1. In the `crate::config` module (likely `src/config.rs` or similar), add a shared helper:
```rust
pub fn ansi_color_escape_sequence(color: &AnsiColor) -> String {
match color {
AnsiColor::Color16 { c16 } => {
let code = if *c16 < 8 { 30 + c16 } else { 90 + (c16 - 8) };
format!("\x1b[{}m", code)
}
AnsiColor::Color256 { c256 } => {
format!("\x1b[38;5;{}m", c256)
}
AnsiColor::Rgb { r, g, b } => {
format!("\x1b[38;2;{};{};{}m", r, g, b)
}
}
}
```
- Adjust the function name/signature if your config module structure differs.
- Ensure `AnsiColor` is in scope (`use crate::config::AnsiColor;` or similar).
2. In `CliProxyApiQuotaSegment::apply_foreground_color` (where the original duplicated logic lives), replace the inline `match` on `AnsiColor` with a call to `crate::config::ansi_color_escape_sequence(color)` and keep/reset behavior (`\x1b[0m`) consistent with how it's done in `preview.rs`.
3. If you prefer the helper to return only the prefix (without reset), adjust the `preview.rs` `apply` closure and `CliProxyApiQuotaSegment::apply_foreground_color` accordingly so both add the reset in a consistent place.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
…ystem - Add new CliProxyApiQuota segment to display quota usage from CLI Proxy API - Supports multiple auth types (antigravity, gemini-cli) - Shows quota percentage for Opus, Gemini 3 Pro, and Gemini 3 Flash models - Configurable via TUI with alias and color customization - Uses caching to minimize API calls - Enhance theme system with forward-compatible segment merging - Add merge_missing_segments() to auto-add new segments to existing configs - Centralize default segment config in presets.rs (reduces code duplication) - New segments are disabled by default to avoid breaking existing setups - Code quality improvements (addressing Sourcery AI feedback): - Unify TrackedModel type across segment and TUI components (eliminates duplication) - Add collect_with_options() to avoid loading Config from disk on each call - Extract apply_foreground_color() as shared ANSI color helper - Add open_with_value() method to NameInputComponent for pre-filled text inputs
rere43
force-pushed
the
feat/antigravity-quota
branch
from
9b3daaf to
60e1c9c
Compare
Author
|
close since not considering cliproxyapi |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
…ystem
Add new CliProxyApiQuota segment to display quota usage from CLI Proxy API
Enhance theme system with forward-compatible segment merging
Add open_with_value() method to NameInputComponent for pre-filled text inputs
Summary by Sourcery
Add a new CLI Proxy API quota statusline segment with configurable options and forward-compatible theme support.
New Features:
Enhancements: