From 3a8c4fbbd2a68e9b030edabd4c95f0e857a136de Mon Sep 17 00:00:00 2001 From: kevin Date: Mon, 3 Aug 2026 15:39:33 +0800 Subject: [PATCH 1/2] feat: add Kimi Code session monitoring and account quota MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Discover live `kimi` processes, map cwd to ~/.kimi-code sessions, and parse agents/main/wire.jsonl for tokens, model, tools, and context %. Account quota follows the Claude StatusLine pattern: a local companion script (abtop-usages.sh) reuses existing kimi login credentials and writes abtop-rate-limits.json; abtop only reads that file. The TUI auto-runs the companion when Kimi is enabled (cold-start sync ≤5s, then throttled background refresh). Disable with hidden_agents = ["kimi"]. Requires python3 for the companion script. --- AGENTS.md | 47 +- README.md | 36 +- src/app.rs | 26 +- src/collector/kimi.rs | 1081 +++++++++++++++++++++++++++++++++++ src/collector/mod.rs | 16 +- src/collector/rate_limit.rs | 111 +++- src/locale.rs | 2 + src/setup.rs | 586 ++++++++++++++++++- src/ui/quota.rs | 4 +- src/ui/sessions.rs | 1 + 10 files changed, 1837 insertions(+), 73 deletions(-) create mode 100644 src/collector/kimi.rs diff --git a/AGENTS.md b/AGENTS.md index d1a254c..227f808 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -2,7 +2,7 @@ AI agent monitor for your terminal. Like btop++, but for AI coding agents. -Supports Claude Code, Codex CLI, and OpenCode sessions. +Supports Claude Code, Codex CLI, OpenCode, and Kimi Code sessions. ## Language Policy @@ -20,7 +20,7 @@ English is mandatory for all project-facing work and communication. src/ ├── main.rs # Entry, terminal setup, event loop, --setup flag ├── app.rs # App state, tick logic, key handling, summary generation -├── setup.rs # StatusLine hook installation (abtop --setup) +├── setup.rs # abtop --setup: Claude StatusLine + Kimi usages companion ├── ui/ │ └── mod.rs # All panels in single file: header, context, quota, │ # tokens, projects, ports, sessions, footer @@ -29,8 +29,9 @@ src/ │ ├── claude.rs # Claude Code: session discovery, transcript parsing │ ├── codex.rs # Codex CLI: session discovery via ps+lsof, JSONL parsing │ ├── opencode.rs # OpenCode: session discovery via ps + SQLite DB parsing +│ ├── kimi.rs # Kimi Code: ps+cwd → sessions/wire.jsonl │ ├── process.rs # Child process tree (ps) + open ports (lsof) + git stats -│ └── rate_limit.rs # Rate limit file reading (~/.claude/abtop-rate-limits.json) +│ └── rate_limit.rs # Local rate-limit file reads (Claude/Kimi hook JSON) └── model/ ├── mod.rs # Re-exports └── session.rs # AgentSession, SessionStatus, RateLimitInfo, @@ -194,25 +195,37 @@ git -C {cwd} status --porcelain # added/modified file counts - Path: `~/.claude/projects/{encoded-path}/memory/` - Count files in directory + lines in `MEMORY.md` -### 9. Rate limit (Claude Code) +### 9. Rate limit (account quota) -NOT in transcript JSONL. Collected via StatusLine mechanism. +abtop **never** calls provider APIs for quota. It only reads local JSON files +written by companion hooks/scripts (`abtop --setup`). -`abtop --setup` automates this: creates a script at `~/.claude/abtop-statusline.sh` that writes rate limit JSON to `~/.claude/abtop-rate-limits.json`, and registers it in `~/.claude/settings.json`. +**Claude Code** — StatusLine mechanism: +- Installs `~/.claude/abtop-statusline.sh` and registers it in `~/.claude/settings.json` +- Writes `~/.claude/abtop-rate-limits.json` on each Claude response -File format read by abtop: +**Kimi Code** — companion script (no StatusLine equivalent): +- Auto-installs `~/.kimi-code/abtop-usages.sh` when credentials exist +- Cold start (file missing or ≥10 min old): run companion once with ≤5s wait so first quota paint is filled +- Warm refresh: background spawn (~2 min throttle) when file is merely stale +- Script does OAuth + `api.kimi.com/coding/v1/usages`; writes `~/.kimi-code/abtop-rate-limits.json` +- Out of the box after `kimi login` — no cron / manual `--setup` required (setup still works) + +**Codex** — still extracted from local rollout JSONL `token_count` events (no setup). + +Shared file format read by abtop: ```json { "source": "claude", - "five_hour": { "used_percentage": 35.0, "resets_at": 1774715000 }, + "five_hour": { "used_percentage": 35.0, "resets_at": 1774715000, "window_minutes": 300 }, "seven_day": { "used_percentage": 12.0, "resets_at": 1775320000 }, "updated_at": 1774714400 } ``` -- Rejects stale data (> 10 minutes old). -- `rate_limits` only present for Pro/Max subscribers. -- Account-level metric, shared across all sessions. +- UI dims stale data (> 10 minutes old). +- Account-level metric, shared across all sessions of that agent. - Show "—" when not configured or data unavailable. +- OpenCode: no quota row (no reliable local account-level source). ### 10. Other files - `~/.claude/stats-cache.json` — daily aggregates. Only updated on `/stats`, NOT real-time. @@ -386,8 +399,14 @@ Parsing/registry logic is unit-tested in `jump/mod.rs`; the thin `ps`/`osascript abtop reads transcripts, prompts, tool inputs, and memory files. These may contain secrets. - **`--once` output**: redact file contents from tool_use inputs. Show tool name + file path only, not content. - **TUI mode**: show tool name + first arg (file path), never show file contents or prompt text in session list. -- **No network**: abtop never sends data anywhere. All local reads. -- **Exception**: summary generation calls `claude --print` locally (no network by abtop itself, but claude may use its API). +- **Session discovery is local-only**: filesystem + `ps` / `lsof`. No abtop-owned API keys. +- **Network exceptions** (optional local helpers; abtop still never uploads transcripts): + - Summary generation shells out to `claude --print` (may use the user's Claude API). + - Kimi account quota: when Kimi is not hidden and `~/.kimi-code/credentials/kimi-code.json` + exists, the TUI may run `~/.kimi-code/abtop-usages.sh`, which calls Moonshot OAuth/usages + APIs using the user's existing Kimi login and writes only usage percentages to + `abtop-rate-limits.json`. Disable with `hidden_agents = ["kimi"]` or by not logging in. + Requires `python3` on PATH for the companion script. ## Gotchas @@ -407,3 +426,5 @@ abtop reads transcripts, prompts, tool inputs, and memory files. These may conta - **PID reuse in port cache**: invalidate cached ports when the set of tracked PIDs changes. - **Rate limit staleness**: reject rate limit data older than 10 minutes. - **`/clear` + multi-PID same cwd**: after `/clear`, Claude Code mints a new `sessionId` + `.jsonl` without rewriting `sessions/{PID}.json`. abtop overrides the stale sid by picking the newest transcript in the project dir, but this heuristic can't disambiguate ownership when two live `claude` PIDs share a cwd — so the override is disabled in that case and both sessions keep their original sid until exit. Use separate worktrees if live tracking is needed on both simultaneously. +- **Kimi multi-PID same cwd**: sessions are matched newest-first per workspace; two live `kimi` PIDs in one cwd claim the two newest session dirs. Prefer separate workspaces if both must be tracked precisely. +- **Kimi quota companion**: may refresh OAuth tokens in place under `credentials/kimi-code.json` (same file Kimi CLI uses). diff --git a/README.md b/README.md index 68116b7..8f04e84 100644 --- a/README.md +++ b/README.md @@ -2,8 +2,8 @@ **Like [btop](https://github.com/aristocratos/btop), but for your AI coding agents.** -See every Claude Code, Codex CLI, and OpenCode session at a glance — token usage, context window %, rate limits, child processes, open ports, and more. -Claude Code, Codex CLI, and OpenCode sessions are discovered from local process/file state, so multiple active profiles are supported across macOS, Linux, and Windows. +See every Claude Code, Codex CLI, OpenCode, and Kimi Code session at a glance — token usage, context window %, rate limits, child processes, open ports, and more. +Sessions are discovered from local process/file state, so multiple active profiles are supported across macOS, Linux, and Windows. ![demo](https://raw.githubusercontent.com/graykode/abtop/main/assets/demo.gif) @@ -14,7 +14,7 @@ Claude Code, Codex CLI, and OpenCode sessions are discovered from local process/ - Agent spawned a server and forgot to kill it? Orphan port detection. - Context window filling up? Per-session % bars with warnings. -All read-only. No API keys. No auth. +Session monitoring is local-only (filesystem + `ps` / `lsof`). abtop does not take its own API keys; Kimi account quota reuses your existing `kimi login` credentials via a local companion script. ## Install @@ -53,7 +53,7 @@ Pre-built binaries for all platforms are available on the [GitHub Releases](http abtop # Launch TUI abtop --once # Print snapshot and exit abtop --json # Print one JSON snapshot and exit (for scripts/tools) -abtop --setup # Install rate limit collection hook +abtop --setup # Install Claude StatusLine + Kimi usages companions abtop --theme dracula # Launch with a specific theme abtop --mouse # Enable mouse click/scroll navigation ``` @@ -75,21 +75,25 @@ tmux new -s work ## Supported Agents -| Feature | Claude Code | Codex CLI | OpenCode | -| ----------------- | :---------: | :-------: | :------: | -| Session Discovery | ✅ | ✅ | ✅ | -| Token Tracking | ✅ | ✅ | ✅ | -| Context Window % | ✅ | ✅ | ❌ | -| Status Detection | ✅ | ✅ | ✅ | -| Current Task | ✅ | ✅ | ❌ | -| Rate Limit | ✅ | ✅ | ❌ | -| Git Status | ✅ | ✅ | ✅ | -| Children / Ports | ✅ | ✅ | ✅ | -| Subagents | ✅ | ❌ | ❌ | -| Memory Status | ✅ | ❌ | ❌ | +| Feature | Claude Code | Codex CLI | OpenCode | Kimi Code | +| ----------------- | :---------: | :-------: | :------: | :-------: | +| Session Discovery | ✅ | ✅ | ✅ | ✅ | +| Token Tracking | ✅ | ✅ | ✅ | ✅ | +| Context Window % | ✅ | ✅ | ❌ | ✅ | +| Status Detection | ✅ | ✅ | ✅ | ✅ | +| Current Task | ✅ | ✅ | ❌ | ✅ | +| Rate Limit | ✅ | ✅ | ❌ | ✅ | +| Git Status | ✅ | ✅ | ✅ | ✅ | +| Children / Ports | ✅ | ✅ | ✅ | ✅ | +| Subagents | ✅ | ❌ | ❌ | names* | +| Memory Status | ✅ | ❌ | ❌ | ❌ | + +\*Kimi lists subagent names from session state; per-subagent tokens/status are not polled yet. OpenCode support reads the local SQLite database at `~/.local/share/opencode/opencode.db` (also the default location on Windows; `%LOCALAPPDATA%\opencode` and `%APPDATA%\opencode` are probed as fallbacks) and requires `sqlite3` in `PATH` (on Windows: `winget install SQLite.SQLite`). +Kimi Code sessions are discovered from live `kimi` processes and `~/.kimi-code/sessions/**/wire.jsonl` (override root with `KIMI_CODE_HOME`). Account quota reuses your existing `kimi login` file (`credentials/kimi-code.json`): a local companion script (`abtop-usages.sh`, needs `python3`) writes `abtop-rate-limits.json`; the TUI may run that script automatically (brief wait on cold start, then background refresh). Hide Kimi entirely with `hidden_agents = ["kimi"]`. + ## Themes 12 built-in themes, including 4 colorblind-friendly options (`high-contrast`, `protanopia`, `deuteranopia`, `tritanopia`). Press `t` to cycle at runtime, or launch with `--theme `. Your choice is saved to `~/.config/abtop/config.toml`. diff --git a/src/app.rs b/src/app.rs index c9de3cc..315c797 100644 --- a/src/app.rs +++ b/src/app.rs @@ -149,6 +149,8 @@ pub struct App { pub help_open: bool, /// View leader overlay (`v`) visibility. pub view_open: bool, + /// When false (`hidden_agents` includes "kimi"), skip the Kimi usages companion. + kimi_quota_enabled: bool, } impl App { @@ -172,6 +174,7 @@ impl App { let mut collector = MultiCollector::with_hidden_and_claude_config_dirs(hidden_agents, claude_config_dirs); collector.set_mcp_suppress(true); + let kimi_quota_enabled = !hidden_agents.iter().any(|h| h.eq_ignore_ascii_case("kimi")); Self { sessions: Vec::new(), selected: 0, @@ -215,6 +218,7 @@ impl App { agent_aggregate: AgentAggregate::default(), help_open: false, view_open: false, + kimi_quota_enabled, } } @@ -494,17 +498,27 @@ impl App { /// Full refresh used by the TUI: collect monitored data, then generate and /// retry session summaries. Equivalent to [`App::tick_no_summaries`] followed /// by [`App::drain_and_retry_summaries`]. + /// + /// Also may refresh Kimi quota via the local companion when credentials + /// exist (cold start: short blocking run; warm: background spawn). See + /// [`crate::setup::maybe_refresh_kimi_quota`]. pub fn tick(&mut self) { + // Before the local rate-limit read inside tick_no_summaries so a + // cold-start write is visible on this same tick. Kept out of + // tick_no_summaries so headless consumers never get network-side jobs. + // Skipped when `hidden_agents` includes "kimi". + if self.kimi_quota_enabled { + crate::setup::maybe_refresh_kimi_quota(); + } self.tick_no_summaries(); self.drain_and_retry_summaries(); } - /// Refresh all monitored data WITHOUT spawning background summary jobs. + /// Refresh all monitored data WITHOUT spawning background jobs. /// - /// `tick` additionally calls [`App::drain_and_retry_summaries`], which - /// shells out to `claude --print` to generate session titles. Headless - /// consumers (e.g. the web snapshot API) call this variant so they never - /// spawn subprocesses or consume the user's Claude quota. + /// Unlike [`App::tick`], this never shells out to `claude --print` (session + /// summaries) and never spawns the Kimi usages companion. Headless consumers + /// (e.g. the web snapshot API) call this so they only read local state. pub fn tick_no_summaries(&mut self) { self.collector.set_mcp_suppress(self.mcp_suppress_sessions); self.sessions = self.collector.collect(); @@ -539,6 +553,7 @@ impl App { if self.rate_limits.is_empty() || self.rate_limit_counter >= 5 { self.rate_limit_counter = 0; let extra_dirs = self.collector.all_config_dirs(); + // Local files only (Claude/Kimi hook JSON + Codex cache/live). self.rate_limits = read_rate_limits(&extra_dirs); // Merge live rate limits from agent collectors (e.g. Codex JSONL parsing) self.rate_limits.extend(self.collector.agent_rate_limits()); @@ -1002,6 +1017,7 @@ fn is_supported_agent_command(cmd: &str) -> bool { crate::collector::process::cmd_has_binary(cmd, "claude") || crate::collector::process::cmd_has_binary(cmd, "codex") || crate::collector::process::cmd_has_binary(cmd, "opencode") + || crate::collector::process::cmd_has_binary(cmd, "kimi") } fn is_killable_agent_command(cmd: &str) -> bool { diff --git a/src/collector/kimi.rs b/src/collector/kimi.rs new file mode 100644 index 0000000..23bbb73 --- /dev/null +++ b/src/collector/kimi.rs @@ -0,0 +1,1081 @@ +use super::{process, SharedProcessData}; +use crate::model::{AgentSession, ChildProcess, SessionStatus, SubAgent}; +use serde::Deserialize; +use serde_json::Value; +use std::collections::{HashMap, HashSet}; +use std::fs::{self, File}; +use std::io::{Read, Seek, SeekFrom}; +use std::path::{Path, PathBuf}; +use std::time::{SystemTime, UNIX_EPOCH}; + +/// Collector for Kimi Code (`kimi`) sessions. +/// +/// Discovery strategy: +/// 1. Find running `kimi` processes via shared `ps` data +/// 2. Map each PID's cwd to a workspace under `~/.kimi-code/sessions/wd_*` +/// 3. Pick the newest session dir for that workspace (by `state.json` mtime / +/// `updatedAt`), claiming one session per live PID when several share a cwd +/// 4. Parse `agents/main/wire.jsonl` for tokens, model, current tool, etc. +/// +/// Account quota is **not** fetched here. Like Claude, abtop only reads a local +/// file (`~/.kimi-code/abtop-rate-limits.json`). When Kimi credentials exist, +/// the TUI auto-installs/spawns `abtop-usages.sh` in the background to refresh +/// that file — network stays out of process. See `setup::maybe_refresh_kimi_quota`. +/// +/// Config root: `~/.kimi-code` (override with `KIMI_CODE_HOME`). +pub struct KimiCollector { + config_root: PathBuf, + /// Cached workspace id → root path (from workspaces.json). + workspace_roots: HashMap, + /// Cached wire parse state keyed by session id. + wire_cache: HashMap, + /// Version string from `~/.kimi-code/updates/latest.json` (best-effort). + version: String, +} + +#[derive(Debug, Clone, Default)] +struct WireState { + offset: u64, + partial: String, + total_input: u64, + total_output: u64, + total_cache_read: u64, + total_cache_create: u64, + turn_count: u32, + model: String, + effort: String, + /// Last step context size (inputOther + inputCacheRead). + last_context_tokens: u64, + last_activity_ms: u64, + first_prompt: String, + current_task: String, + /// tool.call without matching tool.result still open. + pending_tools: u32, + pending_since_ms: u64, + thinking_since_ms: u64, + token_history: Vec, + context_history: Vec, +} + +#[derive(Debug, Deserialize)] +struct StateFile { + #[serde(default, rename = "createdAt")] + created_at: String, + #[serde(default, rename = "updatedAt")] + updated_at: String, + #[serde(default)] + title: String, + #[serde(default, rename = "workDir")] + work_dir: String, + #[serde(default, rename = "lastPrompt")] + last_prompt: String, + #[serde(default)] + agents: HashMap, +} + +#[derive(Debug, Deserialize)] +struct AgentMeta { + #[serde(default, rename = "type")] + agent_type: String, +} + +#[derive(Debug, Deserialize)] +struct WorkspacesFile { + #[serde(default)] + workspaces: HashMap, +} + +#[derive(Debug, Deserialize)] +struct WorkspaceEntry { + #[serde(default)] + root: String, +} + +impl KimiCollector { + pub fn new() -> Self { + let config_root = std::env::var("KIMI_CODE_HOME") + .map(PathBuf::from) + .unwrap_or_else(|_| dirs::home_dir().unwrap_or_default().join(".kimi-code")); + Self { + config_root, + workspace_roots: HashMap::new(), + wire_cache: HashMap::new(), + version: String::new(), + } + } + + fn refresh_meta(&mut self) { + self.workspace_roots = load_workspace_roots(&self.config_root); + if self.version.is_empty() { + self.version = load_version(&self.config_root); + } + } + + fn collect_sessions(&mut self, shared: &SharedProcessData) -> Vec { + if !self.config_root.is_dir() { + return vec![]; + } + + // Refresh workspace map on slow ticks (or first run). + if shared.slow_tick || self.workspace_roots.is_empty() { + self.refresh_meta(); + } + + let kimi_pids = find_kimi_pids(&shared.process_info); + if kimi_pids.is_empty() { + return vec![]; + } + + // Group live PIDs by cwd. + let mut pids_by_cwd: HashMap> = HashMap::new(); + for pid in kimi_pids { + let Some(cwd) = get_process_cwd(pid) else { + continue; + }; + if cwd.len() < 2 { + continue; + } + pids_by_cwd.entry(cwd).or_default().push(pid); + } + + // Invert workspace map: root path → workspace id. + let mut root_to_wd: HashMap = HashMap::new(); + for (wd, root) in &self.workspace_roots { + root_to_wd.insert(root.clone(), wd.clone()); + } + + let now_ms = current_time_ms(); + let mut sessions = Vec::new(); + let mut live_session_ids: HashSet = HashSet::new(); + + for (cwd, mut pids) in pids_by_cwd { + pids.sort_unstable(); + let wd_id = root_to_wd + .get(&cwd) + .cloned() + .or_else(|| find_workspace_id_by_scan(&self.config_root, &cwd)); + + let Some(wd_id) = wd_id else { + // Live process but no known workspace — still show a stub row. + for pid in pids { + sessions.push(stub_session( + pid, + &cwd, + &shared.process_info, + &shared.children_map, + &shared.ports, + &self.config_root, + &self.version, + )); + } + continue; + }; + + let mut candidates = list_sessions_for_workspace(&self.config_root, &wd_id); + // Newest first so the first N map to the N live PIDs. + candidates.sort_by_key(|b| std::cmp::Reverse(b.updated_ms)); + + for (i, pid) in pids.into_iter().enumerate() { + let Some(meta) = candidates.get(i) else { + sessions.push(stub_session( + pid, + &cwd, + &shared.process_info, + &shared.children_map, + &shared.ports, + &self.config_root, + &self.version, + )); + continue; + }; + + live_session_ids.insert(meta.session_id.clone()); + let wire_path = meta + .session_dir + .join("agents") + .join("main") + .join("wire.jsonl"); + + let wire = self.parse_wire_incremental(&meta.session_id, &wire_path); + let state = read_state_file(&meta.session_dir.join("state.json")); + + let proc = shared.process_info.get(&pid); + let mem_mb = proc.map(|p| p.rss_kb / 1024).unwrap_or(0); + + let model = if !wire.model.is_empty() { + wire.model.clone() + } else { + "-".to_string() + }; + let context_window = context_window_for_kimi(&model); + let context_percent = if context_window > 0 && wire.last_context_tokens > 0 { + (wire.last_context_tokens as f64 / context_window as f64) * 100.0 + } else { + 0.0 + }; + + let activity_ms = wire + .last_activity_ms + .max(meta.updated_ms) + .max(wire_mtime_ms(&wire_path)); + let age_secs = now_ms.saturating_sub(activity_ms) / 1000; + let has_active_child = process::has_active_descendant( + pid, + &shared.children_map, + &shared.process_info, + 5.0, + ); + let cpu_active = proc.is_some_and(|p| p.cpu_pct > 1.0); + + let status = if wire.pending_tools > 0 || has_active_child { + SessionStatus::Executing + } else if age_secs < 30 || cpu_active || wire.thinking_since_ms > 0 { + SessionStatus::Thinking + } else { + SessionStatus::Waiting + }; + + let current_tasks = if !wire.current_task.is_empty() + && matches!(status, SessionStatus::Executing | SessionStatus::Thinking) + { + vec![wire.current_task.clone()] + } else if matches!(status, SessionStatus::Waiting) { + vec!["waiting for input".to_string()] + } else if matches!(status, SessionStatus::Thinking) { + vec!["thinking...".to_string()] + } else { + vec![] + }; + + let project_name = process::last_path_segment(&cwd).unwrap_or("?").to_string(); + + let started_at = + parse_iso_ms(state.as_ref().map(|s| s.created_at.as_str()).unwrap_or("")) + .unwrap_or(meta.updated_ms); + + let initial_prompt = { + let from_wire = wire.first_prompt.clone(); + let from_state = state + .as_ref() + .map(|s| { + if !s.title.is_empty() { + s.title.clone() + } else { + s.last_prompt.clone() + } + }) + .unwrap_or_default(); + if !from_wire.is_empty() { + from_wire + } else { + from_state + } + }; + + let subagents: Vec = state + .as_ref() + .map(|s| { + s.agents + .iter() + .filter(|(name, a)| *name != "main" && a.agent_type != "main") + .map(|(name, _)| SubAgent { + name: name.clone(), + // Kimi subagent wire files are not polled yet; + // surface presence only. + status: String::new(), + tokens: 0, + }) + .collect() + }) + .unwrap_or_default(); + + let children = collect_children(pid, shared); + + sessions.push(AgentSession { + agent_cli: "kimi", + pid, + session_id: meta.session_id.clone(), + cwd: cwd.clone(), + project_name, + started_at, + status, + model, + effort: wire.effort.clone(), + context_percent, + total_input_tokens: wire.total_input, + total_output_tokens: wire.total_output, + total_cache_read: wire.total_cache_read, + total_cache_create: wire.total_cache_create, + turn_count: wire.turn_count, + current_tasks, + mem_mb, + version: self.version.clone(), + git_branch: String::new(), + git_added: 0, + git_modified: 0, + token_history: wire.token_history.clone(), + context_history: wire.context_history.clone(), + compaction_count: 0, + context_window, + subagents, + mem_file_count: 0, + mem_line_count: 0, + children, + initial_prompt: super::redact_secrets(&super::sanitize_terminal_text( + &truncate_str(&initial_prompt, 200), + )), + first_assistant_text: String::new(), + chat_messages: vec![], + tool_calls: vec![], + pending_since_ms: wire.pending_since_ms, + thinking_since_ms: wire.thinking_since_ms, + file_accesses: vec![], + config_root: super::abbrev_path(&self.config_root), + }); + } + } + + // Drop wire cache entries for sessions no longer live. + self.wire_cache + .retain(|id, _| live_session_ids.contains(id)); + + sessions.sort_by_key(|s| std::cmp::Reverse(s.started_at)); + sessions + } + + fn parse_wire_incremental(&mut self, session_id: &str, path: &Path) -> WireState { + let mut state = self.wire_cache.get(session_id).cloned().unwrap_or_default(); + + let Ok(meta) = fs::metadata(path) else { + return state; + }; + let file_len = meta.len(); + + // File rotated / rewritten smaller — full rescan. + if file_len < state.offset { + state = WireState::default(); + } + + // Nothing new on disk (incomplete trailing line waits for more bytes). + if file_len == state.offset { + return state; + } + + let Ok(mut file) = File::open(path) else { + return state; + }; + let start_offset = state.offset; + if start_offset > 0 && file.seek(SeekFrom::Start(start_offset)).is_err() { + return state; + } + + let mut buf = String::new(); + if file.read_to_string(&mut buf).is_err() { + return state; + } + + // Prepend any incomplete trailing line from the previous read. + // Those bytes were already counted in `offset`; only `buf` is new. + let combined = if state.partial.is_empty() { + buf + } else { + let mut s = std::mem::take(&mut state.partial); + s.push_str(&buf); + s + }; + + for line in combined.split_inclusive('\n') { + if !line.ends_with('\n') { + state.partial = line.to_string(); + break; + } + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + if let Ok(value) = serde_json::from_str::(trimmed) { + apply_wire_event(&mut state, &value); + } + } + // All bytes through EOF were either parsed or buffered in `partial`. + state.offset = file_len; + + self.wire_cache + .insert(session_id.to_string(), state.clone()); + state + } +} + +impl Default for KimiCollector { + fn default() -> Self { + Self::new() + } +} + +impl super::AgentCollector for KimiCollector { + fn collect(&mut self, shared: &SharedProcessData) -> Vec { + self.collect_sessions(shared) + } +} + +// ── Wire event application ────────────────────────────────────────────────── + +fn apply_wire_event(state: &mut WireState, value: &Value) { + let Some(event_type) = value.get("type").and_then(|v| v.as_str()) else { + return; + }; + let time_ms = value.get("time").and_then(|v| v.as_u64()).unwrap_or(0); + if time_ms > state.last_activity_ms { + state.last_activity_ms = time_ms; + } + + match event_type { + "usage.record" => { + let scope = value + .get("usageScope") + .and_then(|v| v.as_str()) + .unwrap_or("turn"); + if scope != "turn" { + return; + } + if let Some(model) = value.get("model").and_then(|v| v.as_str()) { + if !model.is_empty() { + state.model = model.to_string(); + } + } + if let Some(usage) = value.get("usage") { + let input = usage_u64(usage, "inputOther"); + let output = usage_u64(usage, "output"); + let cache_read = usage_u64(usage, "inputCacheRead"); + let cache_create = usage_u64(usage, "inputCacheCreation"); + state.total_input = state.total_input.saturating_add(input); + state.total_output = state.total_output.saturating_add(output); + state.total_cache_read = state.total_cache_read.saturating_add(cache_read); + state.total_cache_create = state.total_cache_create.saturating_add(cache_create); + state.turn_count = state.turn_count.saturating_add(1); + let turn_tokens = input + output + cache_create; + state.token_history.push(turn_tokens); + if state.token_history.len() > 200 { + let drain = state.token_history.len() - 200; + state.token_history.drain(0..drain); + } + // Context usage excludes cache_create (same rationale as Claude #54). + state.last_context_tokens = input.saturating_add(cache_read); + state.context_history.push(state.last_context_tokens); + if state.context_history.len() > 200 { + let drain = state.context_history.len() - 200; + state.context_history.drain(0..drain); + } + } + // A usage record closes a thinking/generation phase. + state.thinking_since_ms = 0; + } + "llm.request" => { + if let Some(model) = value + .get("modelAlias") + .and_then(|v| v.as_str()) + .or_else(|| value.get("model").and_then(|v| v.as_str())) + { + if !model.is_empty() { + state.model = model.to_string(); + } + } + if let Some(effort) = value.get("thinkingEffort").and_then(|v| v.as_str()) { + state.effort = effort.to_string(); + } + if state.thinking_since_ms == 0 { + state.thinking_since_ms = time_ms; + } + } + "turn.prompt" => { + if state.first_prompt.is_empty() { + if let Some(text) = first_text_from_input(value.get("input")) { + state.first_prompt = text; + } + } + state.thinking_since_ms = time_ms; + } + "context.append_loop_event" => { + let Some(event) = value.get("event") else { + return; + }; + let Some(ev_type) = event.get("type").and_then(|v| v.as_str()) else { + return; + }; + match ev_type { + "tool.call" => { + state.pending_tools = state.pending_tools.saturating_add(1); + if state.pending_since_ms == 0 { + state.pending_since_ms = time_ms; + } + state.thinking_since_ms = 0; + state.current_task = format_tool_task(event); + } + "tool.result" => { + state.pending_tools = state.pending_tools.saturating_sub(1); + if state.pending_tools == 0 { + state.pending_since_ms = 0; + } + } + "step.end" => { + if let Some(usage) = event.get("usage") { + let input = usage_u64(usage, "inputOther"); + let cache_read = usage_u64(usage, "inputCacheRead"); + state.last_context_tokens = input.saturating_add(cache_read); + } + state.thinking_since_ms = 0; + } + "step.begin" if state.pending_tools == 0 && state.thinking_since_ms == 0 => { + state.thinking_since_ms = time_ms; + } + _ => {} + } + } + _ => {} + } +} + +fn usage_u64(usage: &Value, key: &str) -> u64 { + usage + .get(key) + .and_then(|v| v.as_u64().or_else(|| v.as_f64().map(|f| f as u64))) + .unwrap_or(0) +} + +fn first_text_from_input(input: Option<&Value>) -> Option { + let arr = input?.as_array()?; + for item in arr { + if item.get("type").and_then(|v| v.as_str()) == Some("text") { + if let Some(text) = item.get("text").and_then(|v| v.as_str()) { + let t = text.trim(); + if !t.is_empty() { + return Some(t.to_string()); + } + } + } + } + None +} + +fn format_tool_task(event: &Value) -> String { + let name = event.get("name").and_then(|v| v.as_str()).unwrap_or("tool"); + let args = event.get("args").cloned().unwrap_or(Value::Null); + let arg = match name { + "Bash" | "bash" => args + .get("command") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(), + "Read" | "Write" | "Edit" | "read" | "write" | "edit" => args + .get("file_path") + .or_else(|| args.get("path")) + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(), + "Grep" | "Glob" => args + .get("pattern") + .or_else(|| args.get("glob_pattern")) + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(), + "Agent" => args + .get("description") + .or_else(|| args.get("prompt")) + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(), + _ => { + // Prefer description when present (Kimi often provides a short one). + if let Some(d) = event.get("description").and_then(|v| v.as_str()) { + d.to_string() + } else { + String::new() + } + } + }; + let arg = truncate_str(&super::sanitize_terminal_text(&arg), 80); + if arg.is_empty() { + name.to_string() + } else { + format!("{name} {arg}") + } +} + +// ── Discovery helpers ─────────────────────────────────────────────────────── + +fn find_kimi_pids(process_info: &HashMap) -> Vec { + process_info + .iter() + .filter(|(_, info)| { + process::cmd_has_binary(&info.command, "kimi") && !info.command.contains("grep") + }) + .map(|(pid, _)| *pid) + .collect() +} + +#[derive(Debug, Clone)] +struct SessionMeta { + session_id: String, + session_dir: PathBuf, + updated_ms: u64, +} + +fn list_sessions_for_workspace(config_root: &Path, wd_id: &str) -> Vec { + let dir = config_root.join("sessions").join(wd_id); + let Ok(entries) = fs::read_dir(&dir) else { + return vec![]; + }; + let mut out = Vec::new(); + for entry in entries.flatten() { + let path = entry.path(); + if !path.is_dir() { + continue; + } + let name = entry.file_name().to_string_lossy().into_owned(); + if !name.starts_with("session_") { + continue; + } + let state_path = path.join("state.json"); + let updated_ms = fs::metadata(&state_path) + .and_then(|m| m.modified()) + .ok() + .and_then(|t| t.duration_since(UNIX_EPOCH).ok()) + .map(|d| d.as_millis() as u64) + .unwrap_or(0); + // Prefer state.updatedAt when parseable (more accurate than mtime). + let updated_ms = read_state_file(&state_path) + .and_then(|s| parse_iso_ms(&s.updated_at)) + .unwrap_or(updated_ms); + out.push(SessionMeta { + session_id: name, + session_dir: path, + updated_ms, + }); + } + out +} + +fn load_workspace_roots(config_root: &Path) -> HashMap { + let path = config_root.join("workspaces.json"); + let Ok(text) = fs::read_to_string(path) else { + return HashMap::new(); + }; + let Ok(parsed) = serde_json::from_str::(&text) else { + return HashMap::new(); + }; + parsed + .workspaces + .into_iter() + .filter(|(_, e)| !e.root.is_empty()) + .map(|(id, e)| (id, e.root)) + .collect() +} + +/// Fallback when workspaces.json is missing/stale: scan state.json files for workDir. +fn find_workspace_id_by_scan(config_root: &Path, cwd: &str) -> Option { + let sessions = config_root.join("sessions"); + let entries = fs::read_dir(sessions).ok()?; + for entry in entries.flatten() { + let path = entry.path(); + if !path.is_dir() { + continue; + } + let wd_id = entry.file_name().to_string_lossy().into_owned(); + if !wd_id.starts_with("wd_") { + continue; + } + // Check any session's state.json under this workspace. + if let Ok(subs) = fs::read_dir(&path) { + for sub in subs.flatten() { + let state_path = sub.path().join("state.json"); + if let Some(state) = read_state_file(&state_path) { + if paths_equal(&state.work_dir, cwd) { + return Some(wd_id); + } + } + } + } + } + None +} + +fn read_state_file(path: &Path) -> Option { + let text = fs::read_to_string(path).ok()?; + serde_json::from_str(&text).ok() +} + +fn load_version(config_root: &Path) -> String { + let path = config_root.join("updates").join("latest.json"); + let Ok(text) = fs::read_to_string(path) else { + return String::new(); + }; + let Ok(value) = serde_json::from_str::(&text) else { + return String::new(); + }; + value + .pointer("/lastSuccess/version") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string() +} + +fn wire_mtime_ms(path: &Path) -> u64 { + fs::metadata(path) + .and_then(|m| m.modified()) + .ok() + .and_then(|t| t.duration_since(UNIX_EPOCH).ok()) + .map(|d| d.as_millis() as u64) + .unwrap_or(0) +} + +fn stub_session( + pid: u32, + cwd: &str, + process_info: &HashMap, + children_map: &HashMap>, + ports: &HashMap>, + config_root: &Path, + version: &str, +) -> AgentSession { + let proc = process_info.get(&pid); + let mem_mb = proc.map(|p| p.rss_kb / 1024).unwrap_or(0); + let project_name = process::last_path_segment(cwd).unwrap_or("?").to_string(); + // Build a minimal SharedProcessData-like children walk. + let mut children = Vec::new(); + let mut stack: Vec = children_map.get(&pid).cloned().unwrap_or_default(); + let mut visited = HashSet::new(); + while let Some(cpid) = stack.pop() { + if !visited.insert(cpid) { + continue; + } + if let Some(cproc) = process_info.get(&cpid) { + let port = ports.get(&cpid).and_then(|v| v.first().copied()); + children.push(ChildProcess { + pid: cpid, + command: cproc.command.clone(), + mem_kb: cproc.rss_kb, + port, + }); + } + if let Some(gc) = children_map.get(&cpid) { + stack.extend(gc); + } + } + AgentSession { + agent_cli: "kimi", + pid, + session_id: format!("kimi-{pid}"), + cwd: cwd.to_string(), + project_name, + started_at: current_time_ms(), + status: SessionStatus::Unknown, + model: "-".to_string(), + effort: String::new(), + context_percent: 0.0, + total_input_tokens: 0, + total_output_tokens: 0, + total_cache_read: 0, + total_cache_create: 0, + turn_count: 0, + current_tasks: vec!["session metadata unavailable".to_string()], + mem_mb, + version: version.to_string(), + git_branch: String::new(), + git_added: 0, + git_modified: 0, + token_history: vec![], + context_history: vec![], + compaction_count: 0, + context_window: 200_000, + subagents: vec![], + mem_file_count: 0, + mem_line_count: 0, + children, + initial_prompt: String::new(), + first_assistant_text: String::new(), + chat_messages: vec![], + tool_calls: vec![], + pending_since_ms: 0, + thinking_since_ms: 0, + file_accesses: vec![], + config_root: super::abbrev_path(config_root), + } +} + +fn collect_children(pid: u32, shared: &SharedProcessData) -> Vec { + let mut children = Vec::new(); + let mut stack: Vec = shared.children_map.get(&pid).cloned().unwrap_or_default(); + let mut visited = HashSet::new(); + while let Some(cpid) = stack.pop() { + if !visited.insert(cpid) { + continue; + } + if let Some(cproc) = shared.process_info.get(&cpid) { + let port = shared.ports.get(&cpid).and_then(|v| v.first().copied()); + children.push(ChildProcess { + pid: cpid, + command: cproc.command.clone(), + mem_kb: cproc.rss_kb, + port, + }); + } + if let Some(gc) = shared.children_map.get(&cpid) { + stack.extend(gc); + } + } + children +} + +/// Context window for Kimi models. Values mirror `~/.kimi-code/config.toml`. +pub(crate) fn context_window_for_kimi(model: &str) -> u64 { + let m = model.to_ascii_lowercase(); + if m.contains("k3") && !m.contains("256") { + 1_048_576 + } else if m.contains("256") || m.contains("kimi-for-coding") || m.contains("k2") { + 262_144 + } else { + 200_000 + } +} + +fn parse_iso_ms(s: &str) -> Option { + if s.is_empty() { + return None; + } + // chrono parses RFC3339 with fractional seconds. + chrono::DateTime::parse_from_rfc3339(s) + .ok() + .map(|dt| dt.timestamp_millis() as u64) +} + +fn truncate_str(s: &str, max_chars: usize) -> String { + let mut out = String::new(); + for (i, ch) in s.chars().enumerate() { + if i >= max_chars { + out.push('…'); + break; + } + out.push(ch); + } + out +} + +fn current_time_ms() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64 +} + +#[cfg(target_os = "windows")] +fn paths_equal(a: &str, b: &str) -> bool { + let norm = |s: &str| { + s.replace('/', "\\") + .trim_end_matches('\\') + .to_ascii_lowercase() + }; + norm(a) == norm(b) +} + +#[cfg(not(target_os = "windows"))] +fn paths_equal(a: &str, b: &str) -> bool { + a == b +} + +/// Get the current working directory of a process. +#[cfg(target_os = "linux")] +fn get_process_cwd(pid: u32) -> Option { + fs::read_link(format!("/proc/{}/cwd", pid)) + .ok() + .map(|p| p.to_string_lossy().into_owned()) +} + +#[cfg(target_os = "windows")] +fn get_process_cwd(pid: u32) -> Option { + use sysinfo::{Pid, ProcessRefreshKind, ProcessesToUpdate, System, UpdateKind}; + let mut sys = System::new(); + let pid = Pid::from_u32(pid); + sys.refresh_processes_specifics( + ProcessesToUpdate::Some(&[pid]), + false, + ProcessRefreshKind::new().with_cwd(UpdateKind::Always), + ); + sys.process(pid) + .and_then(|p| p.cwd()) + .map(|p| p.to_string_lossy().into_owned()) +} + +#[cfg(all(not(target_os = "linux"), not(target_os = "windows")))] +fn get_process_cwd(pid: u32) -> Option { + use std::process::Command; + let output = Command::new("lsof") + .args(["-a", "-p", &pid.to_string(), "-d", "cwd", "-Fn"]) + .output() + .ok()?; + if !output.status.success() { + return None; + } + let stdout = String::from_utf8_lossy(&output.stdout); + stdout + .lines() + .find(|l| l.starts_with('n') && l.len() > 1) + .map(|l| l[1..].to_string()) +} + +// ── Tests ─────────────────────────────────────────────────────────────────── + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Write; + use tempfile::tempdir; + + #[test] + fn context_window_k3_is_1m() { + assert_eq!(context_window_for_kimi("kimi-code/k3"), 1_048_576); + assert_eq!(context_window_for_kimi("k3"), 1_048_576); + } + + #[test] + fn context_window_256k_variants() { + assert_eq!(context_window_for_kimi("kimi-code/k3-256k"), 262_144); + assert_eq!( + context_window_for_kimi("kimi-code/kimi-for-coding"), + 262_144 + ); + } + + #[test] + fn apply_wire_usage_and_tool() { + let mut state = WireState::default(); + let usage = serde_json::json!({ + "type": "usage.record", + "model": "kimi-code/k3", + "usage": { + "inputOther": 100, + "output": 50, + "inputCacheRead": 1000, + "inputCacheCreation": 0 + }, + "usageScope": "turn", + "time": 1000 + }); + apply_wire_event(&mut state, &usage); + assert_eq!(state.total_input, 100); + assert_eq!(state.total_output, 50); + assert_eq!(state.total_cache_read, 1000); + assert_eq!(state.turn_count, 1); + assert_eq!(state.last_context_tokens, 1100); + assert_eq!(state.model, "kimi-code/k3"); + + let tool = serde_json::json!({ + "type": "context.append_loop_event", + "event": { + "type": "tool.call", + "name": "Bash", + "args": {"command": "cargo test"} + }, + "time": 2000 + }); + apply_wire_event(&mut state, &tool); + assert_eq!(state.pending_tools, 1); + assert!(state.current_task.contains("Bash")); + assert!(state.current_task.contains("cargo test")); + + let result = serde_json::json!({ + "type": "context.append_loop_event", + "event": {"type": "tool.result", "toolCallId": "x"}, + "time": 3000 + }); + apply_wire_event(&mut state, &result); + assert_eq!(state.pending_tools, 0); + } + + #[test] + fn apply_wire_first_prompt() { + let mut state = WireState::default(); + let prompt = serde_json::json!({ + "type": "turn.prompt", + "input": [{"type": "text", "text": "fix the bug"}], + "time": 1 + }); + apply_wire_event(&mut state, &prompt); + assert_eq!(state.first_prompt, "fix the bug"); + } + + #[test] + fn incremental_wire_parse_across_chunks() { + let dir = tempdir().unwrap(); + let wire = dir.path().join("wire.jsonl"); + { + let mut f = File::create(&wire).unwrap(); + writeln!( + f, + r#"{{"type":"usage.record","model":"kimi-code/k3","usage":{{"inputOther":10,"output":5,"inputCacheRead":100,"inputCacheCreation":0}},"usageScope":"turn","time":1}}"# + ) + .unwrap(); + } + + let mut collector = KimiCollector { + config_root: dir.path().to_path_buf(), + workspace_roots: HashMap::new(), + wire_cache: HashMap::new(), + version: String::new(), + }; + let s1 = collector.parse_wire_incremental("sess1", &wire); + assert_eq!(s1.turn_count, 1); + assert_eq!(s1.total_input, 10); + + // Append another usage line + { + let mut f = fs::OpenOptions::new().append(true).open(&wire).unwrap(); + writeln!( + f, + r#"{{"type":"usage.record","model":"kimi-code/k3","usage":{{"inputOther":20,"output":7,"inputCacheRead":200,"inputCacheCreation":0}},"usageScope":"turn","time":2}}"# + ) + .unwrap(); + } + let s2 = collector.parse_wire_incremental("sess1", &wire); + assert_eq!(s2.turn_count, 2); + assert_eq!(s2.total_input, 30); + assert_eq!(s2.total_output, 12); + assert_eq!(s2.last_context_tokens, 220); + } + + #[test] + fn list_sessions_sorted_by_updated() { + let dir = tempdir().unwrap(); + let root = dir.path(); + let wd = root.join("sessions").join("wd_demo"); + fs::create_dir_all(wd.join("session_old")).unwrap(); + fs::create_dir_all(wd.join("session_new")).unwrap(); + fs::write( + wd.join("session_old").join("state.json"), + r#"{"createdAt":"2026-01-01T00:00:00.000Z","updatedAt":"2026-01-01T00:00:00.000Z","title":"old","workDir":"/tmp/demo"}"#, + ) + .unwrap(); + fs::write( + wd.join("session_new").join("state.json"), + r#"{"createdAt":"2026-02-01T00:00:00.000Z","updatedAt":"2026-02-01T00:00:00.000Z","title":"new","workDir":"/tmp/demo"}"#, + ) + .unwrap(); + + let mut sessions = list_sessions_for_workspace(root, "wd_demo"); + sessions.sort_by_key(|b| std::cmp::Reverse(b.updated_ms)); + assert_eq!(sessions.len(), 2); + assert_eq!(sessions[0].session_id, "session_new"); + assert_eq!(sessions[1].session_id, "session_old"); + } + + #[test] + fn load_workspace_roots_parses() { + let dir = tempdir().unwrap(); + fs::write( + dir.path().join("workspaces.json"), + r#"{"version":1,"workspaces":{"wd_x":{"root":"/home/u/proj","name":"proj"}}}"#, + ) + .unwrap(); + let map = load_workspace_roots(dir.path()); + assert_eq!(map.get("wd_x").map(String::as_str), Some("/home/u/proj")); + } +} diff --git a/src/collector/mod.rs b/src/collector/mod.rs index 930caf1..14f9435 100644 --- a/src/collector/mod.rs +++ b/src/collector/mod.rs @@ -1,5 +1,6 @@ pub mod claude; pub mod codex; +pub mod kimi; pub mod mcp; pub mod opencode; pub mod process; @@ -7,6 +8,7 @@ pub mod rate_limit; pub use claude::ClaudeCollector; pub use codex::CodexCollector; +pub use kimi::KimiCollector; pub use mcp::McpServer; pub use opencode::OpenCodeCollector; pub use rate_limit::read_rate_limits; @@ -335,6 +337,9 @@ impl MultiCollector { if !is_hidden("opencode") { collectors.push(Box::new(OpenCodeCollector::new())); } + if !is_hidden("kimi") { + collectors.push(Box::new(KimiCollector::new())); + } let codex_enabled = !is_hidden("codex"); Self { collectors, @@ -507,27 +512,27 @@ mod tests { #[test] fn with_hidden_empty_keeps_all_collectors() { let mc = MultiCollector::with_hidden(&[]); - assert_eq!(mc.collectors.len(), 3); + assert_eq!(mc.collectors.len(), 4); } #[test] fn with_hidden_codex_drops_codex_only() { let mc = MultiCollector::with_hidden(&["codex".to_string()]); - assert_eq!(mc.collectors.len(), 2); + assert_eq!(mc.collectors.len(), 3); } #[test] fn with_hidden_is_case_insensitive() { let mc = MultiCollector::with_hidden(&["CODEX".to_string()]); - assert_eq!(mc.collectors.len(), 2); + assert_eq!(mc.collectors.len(), 3); let mc = MultiCollector::with_hidden(&["Claude".to_string()]); - assert_eq!(mc.collectors.len(), 2); + assert_eq!(mc.collectors.len(), 3); } #[test] fn with_hidden_unknown_names_are_ignored() { let mc = MultiCollector::with_hidden(&["kiro".to_string(), "gemini".to_string()]); - assert_eq!(mc.collectors.len(), 3); + assert_eq!(mc.collectors.len(), 4); } #[test] @@ -536,6 +541,7 @@ mod tests { "claude".to_string(), "codex".to_string(), "opencode".to_string(), + "kimi".to_string(), ]); assert!(mc.collectors.is_empty()); } diff --git a/src/collector/rate_limit.rs b/src/collector/rate_limit.rs index e9722d5..8062b2c 100644 --- a/src/collector/rate_limit.rs +++ b/src/collector/rate_limit.rs @@ -2,8 +2,10 @@ use crate::model::RateLimitInfo; use serde::Deserialize; use std::path::{Path, PathBuf}; -/// File written by the StatusLine hook: ~/.claude/abtop-rate-limits.json -const CLAUDE_RATE_FILE: &str = "abtop-rate-limits.json"; +/// Shared filename written by companion hooks for Claude and Kimi. +/// Claude: `~/.claude/abtop-rate-limits.json` (StatusLine) +/// Kimi: `~/.kimi-code/abtop-rate-limits.json` (`abtop --setup` usages script) +const RATE_LIMIT_FILE: &str = "abtop-rate-limits.json"; /// Cached Codex rate limit: ~/.cache/abtop/codex-rate-limits.json const CODEX_CACHE_FILE: &str = "codex-rate-limits.json"; @@ -30,36 +32,68 @@ struct WindowInfo { window_minutes: Option, } -/// Read rate limit info from all known Claude config directories. -/// Checks the default ~/.claude, CLAUDE_CONFIG_DIR if set, and any -/// additional directories discovered from running Claude processes. +/// Read account-level rate limits from local hook files only. +/// +/// - Claude: `~/.claude/abtop-rate-limits.json` (+ CLAUDE_CONFIG_DIR / discovered dirs) +/// - Kimi: `~/.kimi-code/abtop-rate-limits.json` (+ KIMI_CODE_HOME) +/// +/// abtop never contacts provider APIs. Companion scripts write these files: +/// Claude StatusLine, and for Kimi an auto-spawned `abtop-usages.sh` when +/// credentials exist (also installable via `abtop --setup`). pub fn read_rate_limits(extra_dirs: &[PathBuf]) -> Vec { let mut results = Vec::new(); let mut seen = std::collections::HashSet::new(); - // Collect candidate directories: defaults + discovered - let mut dirs = Vec::new(); + // Claude config roots + let mut claude_dirs = Vec::new(); if let Some(home) = dirs::home_dir() { - dirs.push(home.join(".claude")); + claude_dirs.push(home.join(".claude")); } if let Ok(dir) = std::env::var("CLAUDE_CONFIG_DIR") { - dirs.push(PathBuf::from(dir)); + claude_dirs.push(PathBuf::from(dir)); } - dirs.extend_from_slice(extra_dirs); + claude_dirs.extend_from_slice(extra_dirs); - for dir in dirs { + for dir in claude_dirs { if !dir.is_dir() || !seen.insert(dir.clone()) { continue; } - let path = dir.join(CLAUDE_RATE_FILE); + let path = dir.join(RATE_LIMIT_FILE); if let Some(info) = read_rate_file(&path, "claude") { results.push(info); } } + // Kimi Code config root (local file written by abtop-usages.sh) + for dir in kimi_config_dirs() { + if !dir.is_dir() || !seen.insert(dir.clone()) { + continue; + } + let path = dir.join(RATE_LIMIT_FILE); + if let Some(info) = read_rate_file(&path, "kimi") { + results.push(info); + } + } + results } +fn kimi_config_dirs() -> Vec { + let mut dirs = Vec::new(); + if let Ok(dir) = std::env::var("KIMI_CODE_HOME") { + let p = PathBuf::from(dir); + if !p.as_os_str().is_empty() { + dirs.push(p); + } + } + if let Some(home) = dirs::home_dir() { + dirs.push(home.join(".kimi-code")); + // Legacy kimi-cli home, if a user points the companion script there. + dirs.push(home.join(".kimi")); + } + dirs +} + /// Read cached Codex rate limit (fallback when no live session provides one). /// Rate limits have their own `resets_at` expiry and the cache is refreshed /// whenever the next Codex session runs, so the reader keeps serving the last @@ -141,6 +175,16 @@ fn read_rate_file(path: &Path, default_source: &str) -> Option { file.source }; + // Default window lengths when the hook omits window_minutes: + // Claude/Codex-style 5h + 7d; Kimi short window is also 5h (300 min). + let default_short = 300; + let default_long = if source.eq_ignore_ascii_case("kimi") { + // Kimi's top-level `usage` window is account-period, not fixed 7d. + None + } else { + Some(10_080) + }; + Some(RateLimitInfo { source, five_hour_pct: file.five_hour.as_ref().map(|w| w.used_percentage), @@ -149,14 +193,53 @@ fn read_rate_file(path: &Path, default_source: &str) -> Option { .five_hour .as_ref() .and_then(|w| w.window_minutes) - .or(file.five_hour.as_ref().map(|_| 300)), + .or(file.five_hour.as_ref().map(|_| default_short)), seven_day_pct: file.seven_day.as_ref().map(|w| w.used_percentage), seven_day_resets_at: file.seven_day.as_ref().map(|w| w.resets_at), seven_day_window_minutes: file .seven_day .as_ref() .and_then(|w| w.window_minutes) - .or(file.seven_day.as_ref().map(|_| 10_080)), + .or(file.seven_day.as_ref().and(default_long)), updated_at: file.updated_at, }) } + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Write; + use tempfile::tempdir; + + #[test] + fn reads_kimi_hook_file_from_kimi_code_home() { + let dir = tempdir().unwrap(); + let kimi_home = dir.path().join(".kimi-code"); + std::fs::create_dir_all(&kimi_home).unwrap(); + let path = kimi_home.join(RATE_LIMIT_FILE); + let mut f = std::fs::File::create(&path).unwrap(); + write!( + f, + r#"{{"source":"kimi","five_hour":{{"used_percentage":34.0,"resets_at":1785746704,"window_minutes":300}},"seven_day":{{"used_percentage":31.0,"resets_at":1786067104}},"updated_at":1785739503}}"# + ) + .unwrap(); + + // Temporarily point home-like discovery via KIMI_CODE_HOME. + let prev = std::env::var_os("KIMI_CODE_HOME"); + std::env::set_var("KIMI_CODE_HOME", &kimi_home); + let results = read_rate_limits(&[]); + match prev { + Some(v) => std::env::set_var("KIMI_CODE_HOME", v), + None => std::env::remove_var("KIMI_CODE_HOME"), + } + + let kimi = results + .iter() + .find(|r| r.source.eq_ignore_ascii_case("kimi")) + .expect("kimi rate limit present"); + assert!((kimi.five_hour_pct.unwrap() - 34.0).abs() < 0.01); + assert_eq!(kimi.five_hour_window_minutes, Some(300)); + assert!((kimi.seven_day_pct.unwrap() - 31.0).abs() < 0.01); + assert!(kimi.seven_day_window_minutes.is_none()); + } +} diff --git a/src/locale.rs b/src/locale.rs index abf098f..113c616 100644 --- a/src/locale.rs +++ b/src/locale.rs @@ -202,6 +202,7 @@ static LOCALE_EN: LazyLock> = LazyLock::ne m.insert("quota.no_data", "no data"); m.insert("quota.abtop_setup", "abtop --setup"); m.insert("quota.run_codex", "run codex once"); + m.insert("quota.kimi_login", "kimi login"); m.insert("quota.total", "total"); m.insert("quota.in", "in"); @@ -447,6 +448,7 @@ static LOCALE_ZH: LazyLock> = LazyLock::ne m.insert("quota.no_data", "无数据"); m.insert("quota.abtop_setup", "abtop --setup"); m.insert("quota.run_codex", "运行一次 codex"); + m.insert("quota.kimi_login", "kimi login"); m.insert("quota.total", "总计"); m.insert("quota.in", "还有"); diff --git a/src/setup.rs b/src/setup.rs index 90f2d9a..42ca014 100644 --- a/src/setup.rs +++ b/src/setup.rs @@ -1,6 +1,10 @@ use serde_json::Value; use std::fs; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; +use std::process::{Child, Command, Stdio}; +use std::sync::Mutex; +use std::thread; +use std::time::{Duration, Instant, SystemTime}; const STATUSLINE_SCRIPT: &str = r#"#!/bin/bash # abtop StatusLine hook — writes rate limit data for abtop to read. @@ -32,6 +36,276 @@ with open(os.path.join(config_dir, 'abtop-rate-limits.json'), 'w') as f: " 2>/dev/null "#; +/// Companion script for Kimi Code account quota. +/// Network/auth stay outside abtop — this script writes the same local JSON +/// shape Claude's StatusLine hook produces (`abtop-rate-limits.json`). +/// +/// Logic mirrors TokenTracker's `fetchKimiLimits` (OAuth refresh + usages API). +/// Marker so we can detect an outdated installed script and rewrite it. +const KIMI_USAGES_SCRIPT_VERSION: &str = "abtop-kimi-usages-v1"; + +/// Min gap between background companion spawns (matches TokenTracker ~2 min). +const KIMI_REFRESH_INTERVAL: Duration = Duration::from_secs(120); + +/// Treat the local rate-limit file as fresh enough to skip a refresh spawn. +const KIMI_FILE_FRESH: Duration = Duration::from_secs(120); + +/// File missing or older than this → cold start: run the companion once +/// synchronously so the first quota paint is not empty. +/// Matches the quota panel's "stale" dim threshold (10 minutes). +const KIMI_FILE_ANCIENT: Duration = Duration::from_secs(600); + +/// Max wait for a cold-start companion run (network + token refresh). +const KIMI_COLD_START_TIMEOUT: Duration = Duration::from_secs(5); + +static KIMI_LAST_SPAWN: Mutex> = Mutex::new(None); + +const KIMI_USAGES_SCRIPT: &str = r#"#!/usr/bin/env python3 +"""abtop Kimi Code usage companion — writes local rate-limit JSON for abtop. + +# abtop-kimi-usages-v1 + +Installed automatically when abtop detects Kimi credentials, or via: + abtop --setup + +abtop may spawn this in the background; it only *reads* the JSON this writes. +Network/auth stay in this script — not in the abtop process. +""" +import json +import os +import sys +import time +import urllib.error +import urllib.parse +import urllib.request +from pathlib import Path + +CLIENT_ID = "17e5f671-d194-4dfb-9706-5516cb48c098" +TOKEN_URL = "https://auth.kimi.com/api/oauth/token" +USAGES_URL = "https://api.kimi.com/coding/v1/usages" +TIMEOUT = 8 + + +def kimi_home(): + explicit = os.environ.get("KIMI_CODE_HOME", "").strip() + if explicit: + return Path(explicit).expanduser() + home = Path.home() + code = home / ".kimi-code" + creds = code / "credentials" / "kimi-code.json" + if creds.is_file(): + try: + data = json.loads(creds.read_text()) + if data.get("access_token"): + return code + except Exception: + pass + legacy = home / ".kimi" + if (legacy / "credentials" / "kimi-code.json").is_file(): + return legacy + return code + + +def load_creds(path): + if not path.is_file(): + return None + try: + return json.loads(path.read_text()) + except Exception: + return None + + +def save_creds(path, creds): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(creds, indent=2) + "\n") + + +def expired(creds, now=None): + now = time.time() if now is None else now + exp = creds.get("expires_at") + try: + exp_f = float(exp) + except (TypeError, ValueError): + return False + if exp_f <= 0: + return False + return exp_f <= now + 30 + + +def refresh_token(creds, creds_path): + body = urllib.parse.urlencode( + { + "client_id": CLIENT_ID, + "grant_type": "refresh_token", + "refresh_token": creds.get("refresh_token") or "", + } + ).encode() + req = urllib.request.Request( + TOKEN_URL, + data=body, + method="POST", + headers={ + "Content-Type": "application/x-www-form-urlencoded", + "X-Msh-Platform": "kimi_cli", + }, + ) + with urllib.request.urlopen(req, timeout=TIMEOUT) as resp: + data = json.loads(resp.read().decode()) + if not data.get("access_token"): + raise RuntimeError("token refresh missing access_token") + expires_in = float(data.get("expires_in") or 900) + if expires_in <= 0: + expires_in = 900 + next_creds = { + "access_token": str(data["access_token"]), + "refresh_token": str(data.get("refresh_token") or creds.get("refresh_token") or ""), + "expires_at": time.time() + expires_in, + "scope": str(data.get("scope") or "kimi-code"), + "token_type": str(data.get("token_type") or "Bearer"), + "expires_in": expires_in, + } + save_creds(creds_path, next_creds) + return next_creds + + +def fetch_usages(access_token): + req = urllib.request.Request( + USAGES_URL, + headers={ + "Authorization": "Bearer " + access_token, + "Accept": "application/json", + }, + ) + with urllib.request.urlopen(req, timeout=TIMEOUT) as resp: + return json.loads(resp.read().decode()) + + +def num(value): + if value is None or value == "": + return None + try: + return float(value) + except (TypeError, ValueError): + return None + + +def parse_reset(value): + if not isinstance(value, str) or not value: + return None + # fromisoformat handles trailing Z poorly on some Python versions; normalize. + s = value.replace("Z", "+00:00") + try: + from datetime import datetime + + dt = datetime.fromisoformat(s) + return int(dt.timestamp()) + except Exception: + return None + + +def window_from_usage(data, window_minutes=None): + if not isinstance(data, dict): + return None + limit = num(data.get("limit")) + if limit is None or limit <= 0: + return None + used = num(data.get("used")) + if used is None: + remaining = num(data.get("remaining")) + if remaining is None: + return None + used = limit - remaining + pct = max(0.0, min(100.0, (used / limit) * 100.0)) + out = { + "used_percentage": pct, + "resets_at": parse_reset( + data.get("resetTime") or data.get("reset_at") or data.get("resetAt") + ) + or 0, + } + if window_minutes is not None: + out["window_minutes"] = window_minutes + return out + + +def window_minutes_from(entry): + if not isinstance(entry, dict): + return None + window = entry.get("window") + if not isinstance(window, dict): + return None + duration = num(window.get("duration")) + if duration is None or duration <= 0: + return None + unit = str(window.get("timeUnit") or "TIME_UNIT_MINUTE") + d = int(duration) + if unit in ("TIME_UNIT_SECOND", "SECOND", "seconds"): + return max(1, (d + 59) // 60) + if unit in ("TIME_UNIT_HOUR", "HOUR", "hours"): + return d * 60 + if unit in ("TIME_UNIT_DAY", "DAY", "days"): + return d * 24 * 60 + return d + + +def normalize(body): + limits = body.get("limits") if isinstance(body.get("limits"), list) else [] + short = None + if limits: + entry = limits[0] if isinstance(limits[0], dict) else None + if entry is not None: + detail = entry.get("detail") if isinstance(entry.get("detail"), dict) else entry + mins = window_minutes_from(entry) or 300 + short = window_from_usage(detail, mins) + long = window_from_usage(body.get("usage") if isinstance(body.get("usage"), dict) else None) + if short is None and long is None: + return None + out = {"source": "kimi", "updated_at": int(time.time())} + if short is not None: + out["five_hour"] = short + if long is not None: + out["seven_day"] = long + return out + + +def main(): + home = kimi_home() + creds_path = home / "credentials" / "kimi-code.json" + out_path = home / "abtop-rate-limits.json" + creds = load_creds(creds_path) + if not creds or not str(creds.get("access_token") or "").strip(): + print(f"abtop-usages: not logged in (missing {creds_path})", file=sys.stderr) + print(" run: kimi login", file=sys.stderr) + return 1 + try: + if expired(creds) and creds.get("refresh_token"): + creds = refresh_token(creds, creds_path) + try: + body = fetch_usages(str(creds["access_token"])) + except urllib.error.HTTPError as e: + if e.code == 401 and creds.get("refresh_token"): + creds = refresh_token(creds, creds_path) + body = fetch_usages(str(creds["access_token"])) + else: + raise + out = normalize(body) + if not out: + print("abtop-usages: could not parse usages response", file=sys.stderr) + return 1 + tmp = out_path.with_suffix(".tmp") + tmp.write_text(json.dumps(out) + "\n") + tmp.replace(out_path) + print(f"abtop-usages: wrote {out_path}") + return 0 + except Exception as e: + print(f"abtop-usages: {e}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) +"#; + fn claude_dir() -> PathBuf { std::env::var("CLAUDE_CONFIG_DIR") .ok() @@ -48,21 +322,217 @@ fn settings_path() -> PathBuf { claude_dir().join("settings.json") } +fn kimi_dir() -> PathBuf { + resolve_kimi_home().unwrap_or_else(|| { + std::env::var("KIMI_CODE_HOME") + .ok() + .map(PathBuf::from) + .filter(|p| !p.as_os_str().is_empty()) + .unwrap_or_else(|| dirs::home_dir().unwrap_or_default().join(".kimi-code")) + }) +} + +fn kimi_usages_script_path_for(home: &Path) -> PathBuf { + home.join("abtop-usages.sh") +} + +fn kimi_rate_file_for(home: &Path) -> PathBuf { + home.join("abtop-rate-limits.json") +} + +fn kimi_creds_path_for(home: &Path) -> PathBuf { + home.join("credentials").join("kimi-code.json") +} + +/// Prefer a Kimi home that already has login credentials. +fn resolve_kimi_home() -> Option { + let mut candidates = Vec::new(); + if let Ok(dir) = std::env::var("KIMI_CODE_HOME") { + let p = PathBuf::from(dir); + if !p.as_os_str().is_empty() { + candidates.push(p); + } + } + if let Some(home) = dirs::home_dir() { + candidates.push(home.join(".kimi-code")); + candidates.push(home.join(".kimi")); + } + for dir in &candidates { + if kimi_creds_path_for(dir).is_file() { + return Some(dir.clone()); + } + } + candidates.into_iter().next() +} + +fn kimi_has_credentials(home: &Path) -> bool { + let path = kimi_creds_path_for(home); + let Ok(text) = fs::read_to_string(path) else { + return false; + }; + // Cheap check — avoid parsing secrets into structured types here. + text.contains("access_token") && text.contains('"') +} + +/// Ensure the companion script exists (and is current), return its path. +pub fn ensure_kimi_usages_script() -> Option { + let home = resolve_kimi_home()?; + fs::create_dir_all(&home).ok()?; + let script = kimi_usages_script_path_for(&home); + let needs_write = match fs::read_to_string(&script) { + Ok(existing) => !existing.contains(KIMI_USAGES_SCRIPT_VERSION), + Err(_) => true, + }; + if needs_write { + fs::write(&script, KIMI_USAGES_SCRIPT).ok()?; + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let _ = fs::set_permissions(&script, fs::Permissions::from_mode(0o700)); + } + } + Some(script) +} + +fn rate_file_age(path: &Path) -> Option { + let meta = fs::metadata(path).ok()?; + let modified = meta.modified().ok()?; + SystemTime::now().duration_since(modified).ok() +} + +fn rate_file_is_fresh(path: &Path) -> bool { + rate_file_age(path).is_some_and(|d| d < KIMI_FILE_FRESH) +} + +/// Missing file, or older than [`KIMI_FILE_ANCIENT`] → worth a blocking cold start. +fn rate_file_is_cold(path: &Path) -> bool { + match rate_file_age(path) { + None => true, + Some(age) => age >= KIMI_FILE_ANCIENT, + } +} + +/// Run the companion and wait up to `timeout`, then kill if still running. +fn run_companion_sync(script: &Path, timeout: Duration) { + let mut child = match Command::new(script) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + { + Ok(c) => c, + Err(_) => return, + }; + wait_child_with_timeout(&mut child, timeout); +} + +fn wait_child_with_timeout(child: &mut Child, timeout: Duration) { + let start = Instant::now(); + loop { + match child.try_wait() { + Ok(Some(_)) => return, + Ok(None) => { + if start.elapsed() >= timeout { + let _ = child.kill(); + let _ = child.wait(); + return; + } + thread::sleep(Duration::from_millis(50)); + } + Err(_) => return, + } + } +} + +fn spawn_companion_background(script: &Path) { + let _ = Command::new(script) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn(); +} + +/// Out-of-the-box Kimi quota: if the user is logged in (`kimi login`), install +/// the companion script if needed and refresh when the local rate-limit file +/// is missing or stale. +/// +/// - **Cold start** (file missing or ≥10 min old): run the companion once +/// synchronously (≤5s) so the first quota paint is populated. +/// - **Warm refresh** (file present but >2 min old): fire-and-forget background +/// spawn so the TUI never blocks. +/// +/// abtop still does not perform HTTP itself — the child script does. +/// +/// Call from the full TUI [`crate::app::App::tick`] path only (not +/// `tick_no_summaries`). Background path is non-blocking; cold start may block +/// briefly once. +pub fn maybe_refresh_kimi_quota() { + let Some(home) = resolve_kimi_home() else { + return; + }; + if !kimi_has_credentials(&home) { + return; + } + + let rate_path = kimi_rate_file_for(&home); + if rate_file_is_fresh(&rate_path) { + return; + } + + let cold = rate_file_is_cold(&rate_path); + + { + let mut last = match KIMI_LAST_SPAWN.lock() { + Ok(g) => g, + Err(_) => return, + }; + if let Some(t) = *last { + if t.elapsed() < KIMI_REFRESH_INTERVAL { + return; + } + } + *last = Some(Instant::now()); + } + + let Some(script) = ensure_kimi_usages_script() else { + return; + }; + + if cold { + run_companion_sync(&script, KIMI_COLD_START_TIMEOUT); + } else { + spawn_companion_background(&script); + } +} + pub fn run_setup() { - println!("abtop --setup: configuring Claude Code StatusLine hook\n"); + println!("abtop --setup: installing local rate-limit companions\n"); + + setup_claude_statusline(); + setup_kimi_usages(); + + println!("\n done!"); + println!(" Claude: rate limits appear after the next Claude response (restart sessions)."); + println!(" Kimi: auto-refreshed while abtop runs if you are logged in (`kimi login`)."); + println!( + " companion script: {}", + kimi_usages_script_path_for(&kimi_dir()).display() + ); +} + +fn setup_claude_statusline() { + println!("[claude] StatusLine hook"); - // Ensure ~/.claude directory exists let dir = claude_dir(); if let Err(e) = fs::create_dir_all(&dir) { eprintln!(" ✗ failed to create {}: {}", dir.display(), e); - std::process::exit(1); + eprintln!(" skipping Claude setup"); + return; } - // Step 1: Write the statusline script let script = script_path(); match fs::write(&script, STATUSLINE_SCRIPT) { Ok(_) => { - // chmod +x #[cfg(unix)] { use std::os::unix::fs::PermissionsExt; @@ -72,18 +542,19 @@ pub fn run_setup() { } Err(e) => { eprintln!(" ✗ failed to write {}: {}", script.display(), e); - std::process::exit(1); + eprintln!(" skipping Claude setup"); + return; } } - // Step 2: Update settings.json let settings_file = settings_path(); let mut settings: Value = if settings_file.exists() { let content = match fs::read_to_string(&settings_file) { Ok(c) => c, Err(e) => { eprintln!(" ✗ cannot read {}: {}", settings_file.display(), e); - std::process::exit(1); + eprintln!(" skipping Claude settings update"); + return; } }; match serde_json::from_str(&content) { @@ -94,8 +565,8 @@ pub fn run_setup() { settings_file.display(), e ); - eprintln!(" fix the file manually before running --setup"); - std::process::exit(1); + eprintln!(" fix the file manually before re-running --setup"); + return; } } } else { @@ -104,7 +575,6 @@ pub fn run_setup() { let obj = settings.as_object_mut().unwrap(); - // Check if statusLine is already configured let expected_cmd = script.display().to_string(); if let Some(existing) = obj.get("statusLine") { if let Some(existing_obj) = existing.as_object() { @@ -112,15 +582,15 @@ pub fn run_setup() { let cmd_str = cmd.as_str().unwrap_or(""); if cmd_str != expected_cmd && !cmd_str.is_empty() { eprintln!(" ⚠ statusLine already configured: {}", cmd_str); - eprintln!(" to override, remove the existing statusLine key from:"); - eprintln!(" {}", settings_file.display()); - std::process::exit(1); + eprintln!(" leaving settings unchanged; script still written."); + eprintln!(" to use abtop's hook, point statusLine.command at:"); + eprintln!(" {}", expected_cmd); + return; } } } } - // Set statusLine config obj.insert( "statusLine".to_string(), serde_json::json!({ @@ -136,10 +606,88 @@ pub fn run_setup() { Ok(_) => println!(" ✓ updated {}", settings_file.display()), Err(e) => { eprintln!(" ✗ failed to update {}: {}", settings_file.display(), e); - std::process::exit(1); } } +} - println!("\n done! rate limit data will appear in abtop after the next Claude response."); - println!(" restart any running Claude Code sessions to activate."); +fn setup_kimi_usages() { + println!("\n[kimi] usages companion script"); + + let dir = kimi_dir(); + if let Err(e) = fs::create_dir_all(&dir) { + eprintln!(" ✗ failed to create {}: {}", dir.display(), e); + eprintln!(" skipping Kimi setup"); + return; + } + + match ensure_kimi_usages_script() { + Some(script) => { + println!(" ✓ wrote {}", script.display()); + println!( + " writes {} (abtop reads this file only)", + kimi_rate_file_for(&dir).display() + ); + if !kimi_has_credentials(&dir) { + println!(" ⚠ not logged in yet — run `kimi login`, then start abtop"); + return; + } + // Synchronous first fetch so the panel has data immediately. + run_companion_sync(&script, KIMI_COLD_START_TIMEOUT); + if kimi_rate_file_for(&dir).is_file() { + println!(" ✓ initial fetch succeeded"); + } else { + println!(" ⚠ initial fetch failed or timed out (check `kimi login` / network)"); + } + } + None => { + eprintln!(" ✗ failed to install Kimi companion script"); + eprintln!(" skipping Kimi setup"); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Write; + use tempfile::tempdir; + + #[test] + fn rate_file_missing_is_cold_not_fresh() { + let dir = tempdir().unwrap(); + let path = dir.path().join("missing.json"); + assert!(!rate_file_is_fresh(&path)); + assert!(rate_file_is_cold(&path)); + } + + #[test] + fn rate_file_just_written_is_fresh_not_cold() { + let dir = tempdir().unwrap(); + let path = dir.path().join("fresh.json"); + let mut f = fs::File::create(&path).unwrap(); + writeln!(f, "{{}}").unwrap(); + assert!(rate_file_is_fresh(&path)); + assert!(!rate_file_is_cold(&path)); + } + + #[cfg(unix)] + #[test] + fn wait_child_timeout_kills_slow_process() { + // A process that would hang if not killed (`sleep` is unix-portable). + let mut child = Command::new("sleep") + .arg("30") + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .expect("spawn sleep"); + let start = Instant::now(); + wait_child_with_timeout(&mut child, Duration::from_millis(200)); + assert!( + start.elapsed() < Duration::from_secs(2), + "timeout path should return quickly" + ); + // Child should be reaped. + assert!(child.try_wait().ok().flatten().is_some()); + } } diff --git a/src/ui/quota.rs b/src/ui/quota.rs index 1c65813..ca9ca7d 100644 --- a/src/ui/quota.rs +++ b/src/ui/quota.rs @@ -14,7 +14,7 @@ use super::{btop_block_active, fmt_tokens, grad_at, make_gradient, remaining_bar const STALE_SECS: u64 = 600; /// Fixed source order so columns stay stable across runs. -const SOURCES: &[&str] = &["claude", "codex"]; +const SOURCES: &[&str] = &["claude", "codex", "kimi"]; pub(crate) fn draw_quota_panel(f: &mut Frame, app: &App, area: Rect, theme: &Theme) { draw_quota_panel_active(f, app, area, theme, false); @@ -110,6 +110,8 @@ fn draw_source_column( let Some(rl) = rl else { let hint = if source.eq_ignore_ascii_case("claude") { t("quota.abtop_setup") + } else if source.eq_ignore_ascii_case("kimi") { + t("quota.kimi_login") } else { t("quota.run_codex") }; diff --git a/src/ui/sessions.rs b/src/ui/sessions.rs index 1a911eb..ad2e3f0 100644 --- a/src/ui/sessions.rs +++ b/src/ui/sessions.rs @@ -137,6 +137,7 @@ pub(crate) fn draw_sessions_panel_active( "claude" => ("*CC", Color::Rgb(217, 119, 87)), // #D97757 terracotta "codex" => (">CD", Color::Rgb(122, 157, 255)), // #7A9DFF periwinkle "opencode" => ("#OC", Color::Rgb(74, 222, 128)), // #4ADE80 emerald + "kimi" => ("~KM", Color::Rgb(168, 85, 247)), // #A855F7 violet (Moonshot/Kimi) other => { let fallback: String = other.chars().take(3).collect::().to_uppercase(); ( From 493115594a344572f52130b5c66f571ed4a98d14 Mon Sep 17 00:00:00 2001 From: kevin Date: Mon, 3 Aug 2026 17:21:27 +0800 Subject: [PATCH 2/2] fix: harden Kimi monitoring and quota --- .gitignore | 3 +- AGENTS.md | 70 ++-- README.md | 10 +- src/app.rs | 26 +- src/collector/kimi.rs | 254 +++++++++----- src/collector/mod.rs | 23 +- src/collector/opencode.rs | 70 +--- src/collector/process.rs | 59 ++++ src/collector/rate_limit.rs | 153 +++++---- src/kimi_usages.py | 262 +++++++++++++++ src/lib.rs | 2 +- src/locale.rs | 6 +- src/setup.rs | 651 +++++++++--------------------------- src/ui/quota.rs | 112 ++++--- 14 files changed, 879 insertions(+), 822 deletions(-) create mode 100644 src/kimi_usages.py diff --git a/.gitignore b/.gitignore index efcd773..b0b3177 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ /target +__pycache__/ .DS_Store .agent/ .claude/worktrees/ demo.tape -.idea/ \ No newline at end of file +.idea/ diff --git a/AGENTS.md b/AGENTS.md index 227f808..d765396 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -20,7 +20,8 @@ English is mandatory for all project-facing work and communication. src/ ├── main.rs # Entry, terminal setup, event loop, --setup flag ├── app.rs # App state, tick logic, key handling, summary generation -├── setup.rs # abtop --setup: Claude StatusLine + Kimi usages companion +├── setup.rs # Claude StatusLine + optional Kimi quota companion +├── kimi_usages.py # Explicitly installed Kimi quota helper ├── ui/ │ └── mod.rs # All panels in single file: header, context, quota, │ # tokens, projects, ports, sessions, footer @@ -31,7 +32,7 @@ src/ │ ├── opencode.rs # OpenCode: session discovery via ps + SQLite DB parsing │ ├── kimi.rs # Kimi Code: ps+cwd → sessions/wire.jsonl │ ├── process.rs # Child process tree (ps) + open ports (lsof) + git stats -│ └── rate_limit.rs # Local rate-limit file reads (Claude/Kimi hook JSON) +│ └── rate_limit.rs # Local Claude/Kimi quota file reads + Codex cache └── model/ ├── mod.rs # Re-exports └── session.rs # AgentSession, SessionStatus, RateLimitInfo, @@ -77,7 +78,7 @@ Panel rendering priority (top to bottom): Panel descriptions: - **¹context**: Left = token rate braille sparkline (200-point history). Right = per-session context % bars with yellow/red warning. -- **²quota**: Claude + Codex rate limit gauges side-by-side (5h and 7d windows with reset countdown). Quota is intentionally limited to Claude and Codex; do not add an OpenCode row unless OpenCode exposes a reliable account-level provider rate-limit source. +- **²quota**: Claude + Codex + Kimi rate limit gauges side-by-side. Kimi's long window is the account plan period rather than a fixed seven-day window. Do not add an OpenCode row unless OpenCode exposes a reliable account-level provider rate-limit source. - **³tokens**: Total token breakdown (in/out/cache) + per-turn sparkline for selected session. - **projects** (always visible): Per-project git branch + added/modified file counts. - **⁴ports**: Agent-spawned open ports + orphan ports (from dead sessions). Conflict detection. @@ -85,7 +86,7 @@ Panel descriptions: ## Data Sources -All read-only from local filesystem + `ps` + `lsof`. No API calls, no auth. +Session monitoring is read-only from the local filesystem plus `ps` and `lsof`. The explicitly installed Kimi quota companion is the only provider-API/auth exception. ### 1. Claude Code session discovery: process + config-root mapping @@ -172,13 +173,21 @@ Rate limits extracted from `token_count` events: - Discover running `opencode` processes via shared `ps` data. - Read recent sessions from OpenCode's SQLite DB through `sqlite3 -readonly -json`. - Match live PIDs to DB sessions by process cwd. OpenCode does not expose a PID/session mapping, so when multiple DB rows share one cwd, only live PIDs should be assigned and older rows should not be shown as live duplicates. -- OpenCode contributes session/token/project/port data, but not quota data. Quota remains Claude + Codex only. +- OpenCode contributes session/token/project/port data, but not quota data. -### 5. Subagents: `~/.claude/projects/{path}/{sessionId}/subagents/` +### 5. Kimi Code sessions: `~/.kimi-code/sessions/` +- Discover running `kimi` processes via shared `ps` data. +- Resolve PID → cwd locally (`/proc` on Linux, `sysinfo` on Windows, cached `lsof` on macOS). +- Map cwd to Kimi workspaces and select the newest matching live session directories. +- Parse `agents/main/wire.jsonl` incrementally for tokens, model, context, and current tools. +- `KIMI_CODE_HOME` overrides the default `~/.kimi-code` data root. +- Kimi contributes session/token/project/port data. Account quota comes only from the optional companion described below. + +### 6. Subagents: `~/.claude/projects/{path}/{sessionId}/subagents/` - `agent-{hash}.jsonl` — same JSONL format as main transcript - `agent-{hash}.meta.json` — `{ "agentType": "general-purpose", "description": "..." }` -### 6. Process tree: `ps` + `lsof` +### 7. Process tree: `ps` + `lsof` ```bash ps -eo pid,ppid,rss,%cpu,command # All processes lsof -i -P -n -sTCP:LISTEN # Open ports @@ -186,48 +195,39 @@ lsof -i -P -n -sTCP:LISTEN # Open ports - Build parent→children map from ppid - Map listening PID → parent agent PID → session -### 7. Git status per project +### 8. Git status per project ```bash git -C {cwd} status --porcelain # added/modified file counts ``` -### 8. Memory status +### 9. Memory status - Path: `~/.claude/projects/{encoded-path}/memory/` - Count files in directory + lines in `MEMORY.md` -### 9. Rate limit (account quota) +### 10. Rate limits (account quota) -abtop **never** calls provider APIs for quota. It only reads local JSON files -written by companion hooks/scripts (`abtop --setup`). +Quota is account-level and is not present in session transcript JSONL. -**Claude Code** — StatusLine mechanism: -- Installs `~/.claude/abtop-statusline.sh` and registers it in `~/.claude/settings.json` -- Writes `~/.claude/abtop-rate-limits.json` on each Claude response +**Claude Code** uses its StatusLine mechanism. `abtop --setup` creates `~/.claude/abtop-statusline.sh`, registers it in `~/.claude/settings.json`, and the hook writes `~/.claude/abtop-rate-limits.json` after Claude responses. -**Kimi Code** — companion script (no StatusLine equivalent): -- Auto-installs `~/.kimi-code/abtop-usages.sh` when credentials exist -- Cold start (file missing or ≥10 min old): run companion once with ≤5s wait so first quota paint is filled -- Warm refresh: background spawn (~2 min throttle) when file is merely stale -- Script does OAuth + `api.kimi.com/coding/v1/usages`; writes `~/.kimi-code/abtop-rate-limits.json` -- Out of the box after `kimi login` — no cron / manual `--setup` required (setup still works) +**Kimi Code** has no reliable local account-quota file. `abtop --setup` explicitly installs `~/.kimi-code/abtop-usages.sh` (or under `KIMI_CODE_HOME`). While the full TUI runs, abtop may execute that already-installed companion in the background at most once every two minutes. The companion uses the existing `kimi login` OAuth credentials, refreshes them when necessary, calls `https://api.kimi.com/coding/v1/usages`, and atomically writes `abtop-rate-limits.json`. It requires `python3`. The TUI never auto-installs the helper, never blocks on refresh, reaps every helper process, and never invokes it from `--once`, `--json`, or when Kimi is hidden. -**Codex** — still extracted from local rollout JSONL `token_count` events (no setup). +**Codex** quota comes from local `token_count` events and is cached under `~/.cache/abtop/`. -Shared file format read by abtop: +File format read by abtop: ```json { "source": "claude", - "five_hour": { "used_percentage": 35.0, "resets_at": 1774715000, "window_minutes": 300 }, + "five_hour": { "used_percentage": 35.0, "resets_at": 1774715000 }, "seven_day": { "used_percentage": 12.0, "resets_at": 1775320000 }, "updated_at": 1774714400 } ``` -- UI dims stale data (> 10 minutes old). -- Account-level metric, shared across all sessions of that agent. +- Claude hook data older than ten minutes is rejected. Kimi's last valid cache remains visible but is dimmed after ten minutes so transient refresh failures do not erase the gauge. +- Account-level metrics are shared across all sessions of that provider. - Show "—" when not configured or data unavailable. -- OpenCode: no quota row (no reliable local account-level source). -### 10. Other files +### 11. Other files - `~/.claude/stats-cache.json` — daily aggregates. Only updated on `/stats`, NOT real-time. - `~/.claude/history.jsonl` — prompt history with sessionId. @@ -400,13 +400,9 @@ abtop reads transcripts, prompts, tool inputs, and memory files. These may conta - **`--once` output**: redact file contents from tool_use inputs. Show tool name + file path only, not content. - **TUI mode**: show tool name + first arg (file path), never show file contents or prompt text in session list. - **Session discovery is local-only**: filesystem + `ps` / `lsof`. No abtop-owned API keys. -- **Network exceptions** (optional local helpers; abtop still never uploads transcripts): - - Summary generation shells out to `claude --print` (may use the user's Claude API). - - Kimi account quota: when Kimi is not hidden and `~/.kimi-code/credentials/kimi-code.json` - exists, the TUI may run `~/.kimi-code/abtop-usages.sh`, which calls Moonshot OAuth/usages - APIs using the user's existing Kimi login and writes only usage percentages to - `abtop-rate-limits.json`. Disable with `hidden_agents = ["kimi"]` or by not logging in. - Requires `python3` on PATH for the companion script. +- **Network exceptions**: + - Summary generation calls `claude --print` locally; Claude may use its API. + - After explicit `abtop --setup`, the Kimi quota companion uses the existing Kimi OAuth login and sends only an authenticated quota request to Kimi. It never receives transcript or prompt content from abtop. ## Gotchas @@ -424,7 +420,7 @@ abtop reads transcripts, prompts, tool inputs, and memory files. These may conta - **Undocumented internals**: all data sources are Claude Code/Codex implementation details, not stable APIs. Schema may change without notice. Defensive parsing with `serde(default)` everywhere. - **Terminal size**: minimum 80x24. Panels degrade gracefully when small (context panel hidden first). - **PID reuse in port cache**: invalidate cached ports when the set of tracked PIDs changes. -- **Rate limit staleness**: reject rate limit data older than 10 minutes. +- **Rate limit staleness**: reject Claude hook data older than 10 minutes; retain and dim stale Kimi cache data. - **`/clear` + multi-PID same cwd**: after `/clear`, Claude Code mints a new `sessionId` + `.jsonl` without rewriting `sessions/{PID}.json`. abtop overrides the stale sid by picking the newest transcript in the project dir, but this heuristic can't disambiguate ownership when two live `claude` PIDs share a cwd — so the override is disabled in that case and both sessions keep their original sid until exit. Use separate worktrees if live tracking is needed on both simultaneously. - **Kimi multi-PID same cwd**: sessions are matched newest-first per workspace; two live `kimi` PIDs in one cwd claim the two newest session dirs. Prefer separate workspaces if both must be tracked precisely. -- **Kimi quota companion**: may refresh OAuth tokens in place under `credentials/kimi-code.json` (same file Kimi CLI uses). +- **Kimi quota companion**: optional, installed only by `abtop --setup`, and may refresh the same OAuth credentials file used by Kimi Code. diff --git a/README.md b/README.md index 8f04e84..d55c556 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,7 @@ Sessions are discovered from local process/file state, so multiple active profil - Agent spawned a server and forgot to kill it? Orphan port detection. - Context window filling up? Per-session % bars with warnings. -Session monitoring is local-only (filesystem + `ps` / `lsof`). abtop does not take its own API keys; Kimi account quota reuses your existing `kimi login` credentials via a local companion script. +Session monitoring is read-only and local. Optional Kimi quota support is enabled explicitly with `abtop --setup`; its companion reuses the user's existing Kimi login and never reads session content. ## Install @@ -53,7 +53,7 @@ Pre-built binaries for all platforms are available on the [GitHub Releases](http abtop # Launch TUI abtop --once # Print snapshot and exit abtop --json # Print one JSON snapshot and exit (for scripts/tools) -abtop --setup # Install Claude StatusLine + Kimi usages companions +abtop --setup # Install Claude and Kimi quota companions abtop --theme dracula # Launch with a specific theme abtop --mouse # Enable mouse click/scroll navigation ``` @@ -82,17 +82,17 @@ tmux new -s work | Context Window % | ✅ | ✅ | ❌ | ✅ | | Status Detection | ✅ | ✅ | ✅ | ✅ | | Current Task | ✅ | ✅ | ❌ | ✅ | -| Rate Limit | ✅ | ✅ | ❌ | ✅ | +| Rate Limit | ✅ | ✅ | ❌ | ✅* | | Git Status | ✅ | ✅ | ✅ | ✅ | | Children / Ports | ✅ | ✅ | ✅ | ✅ | | Subagents | ✅ | ❌ | ❌ | names* | | Memory Status | ✅ | ❌ | ❌ | ❌ | -\*Kimi lists subagent names from session state; per-subagent tokens/status are not polled yet. +\*Kimi lists subagent names from session state; per-subagent tokens/status are not polled yet. Kimi quota requires `abtop --setup`, `python3`, a working network connection, and an existing `kimi login`. OpenCode support reads the local SQLite database at `~/.local/share/opencode/opencode.db` (also the default location on Windows; `%LOCALAPPDATA%\opencode` and `%APPDATA%\opencode` are probed as fallbacks) and requires `sqlite3` in `PATH` (on Windows: `winget install SQLite.SQLite`). -Kimi Code sessions are discovered from live `kimi` processes and `~/.kimi-code/sessions/**/wire.jsonl` (override root with `KIMI_CODE_HOME`). Account quota reuses your existing `kimi login` file (`credentials/kimi-code.json`): a local companion script (`abtop-usages.sh`, needs `python3`) writes `abtop-rate-limits.json`; the TUI may run that script automatically (brief wait on cold start, then background refresh). Hide Kimi entirely with `hidden_agents = ["kimi"]`. +Kimi Code sessions are discovered from live `kimi` processes and `~/.kimi-code/sessions/**/wire.jsonl` (override root with `KIMI_CODE_HOME`). Because Kimi does not write account quota locally, `abtop --setup` can explicitly install `~/.kimi-code/abtop-usages.sh`. The TUI then runs that companion in the background every two minutes; it reuses the existing Kimi OAuth login, calls Kimi's usage endpoint, and writes only quota percentages and reset times to `abtop-rate-limits.json`. Headless snapshots never run the companion. Hide Kimi sessions and quota refreshes with `hidden_agents = ["kimi"]`. ## Themes diff --git a/src/app.rs b/src/app.rs index 315c797..32e3419 100644 --- a/src/app.rs +++ b/src/app.rs @@ -149,7 +149,7 @@ pub struct App { pub help_open: bool, /// View leader overlay (`v`) visibility. pub view_open: bool, - /// When false (`hidden_agents` includes "kimi"), skip the Kimi usages companion. + /// Skip optional Kimi quota refreshes when Kimi is hidden. kimi_quota_enabled: bool, } @@ -498,15 +498,10 @@ impl App { /// Full refresh used by the TUI: collect monitored data, then generate and /// retry session summaries. Equivalent to [`App::tick_no_summaries`] followed /// by [`App::drain_and_retry_summaries`]. - /// - /// Also may refresh Kimi quota via the local companion when credentials - /// exist (cold start: short blocking run; warm: background spawn). See - /// [`crate::setup::maybe_refresh_kimi_quota`]. pub fn tick(&mut self) { - // Before the local rate-limit read inside tick_no_summaries so a - // cold-start write is visible on this same tick. Kept out of - // tick_no_summaries so headless consumers never get network-side jobs. - // Skipped when `hidden_agents` includes "kimi". + // Only runs a companion previously installed by the explicit + // `abtop --setup` flow. It is non-blocking and omitted from headless + // snapshots and hidden Kimi configurations. if self.kimi_quota_enabled { crate::setup::maybe_refresh_kimi_quota(); } @@ -514,11 +509,12 @@ impl App { self.drain_and_retry_summaries(); } - /// Refresh all monitored data WITHOUT spawning background jobs. + /// Refresh all monitored data WITHOUT spawning background summary jobs. /// - /// Unlike [`App::tick`], this never shells out to `claude --print` (session - /// summaries) and never spawns the Kimi usages companion. Headless consumers - /// (e.g. the web snapshot API) call this so they only read local state. + /// `tick` additionally calls [`App::drain_and_retry_summaries`], which + /// shells out to `claude --print` to generate session titles. Headless + /// consumers (e.g. the web snapshot API) call this variant so they never + /// spawn subprocesses or consume the user's Claude quota. pub fn tick_no_summaries(&mut self) { self.collector.set_mcp_suppress(self.mcp_suppress_sessions); self.sessions = self.collector.collect(); @@ -553,7 +549,6 @@ impl App { if self.rate_limits.is_empty() || self.rate_limit_counter >= 5 { self.rate_limit_counter = 0; let extra_dirs = self.collector.all_config_dirs(); - // Local files only (Claude/Kimi hook JSON + Codex cache/live). self.rate_limits = read_rate_limits(&extra_dirs); // Merge live rate limits from agent collectors (e.g. Codex JSONL parsing) self.rate_limits.extend(self.collector.agent_rate_limits()); @@ -1113,10 +1108,11 @@ mod tests { } #[test] - fn supported_agent_command_accepts_opencode() { + fn supported_agent_command_accepts_all_collectors() { assert!(is_supported_agent_command("/usr/local/bin/claude")); assert!(is_supported_agent_command("codex --resume abc")); assert!(is_supported_agent_command("/opt/homebrew/bin/opencode")); + assert!(is_supported_agent_command("/usr/local/bin/kimi")); assert!(!is_supported_agent_command("node server.js")); } diff --git a/src/collector/kimi.rs b/src/collector/kimi.rs index 23bbb73..d065585 100644 --- a/src/collector/kimi.rs +++ b/src/collector/kimi.rs @@ -17,16 +17,14 @@ use std::time::{SystemTime, UNIX_EPOCH}; /// `updatedAt`), claiming one session per live PID when several share a cwd /// 4. Parse `agents/main/wire.jsonl` for tokens, model, current tool, etc. /// -/// Account quota is **not** fetched here. Like Claude, abtop only reads a local -/// file (`~/.kimi-code/abtop-rate-limits.json`). When Kimi credentials exist, -/// the TUI auto-installs/spawns `abtop-usages.sh` in the background to refresh -/// that file — network stays out of process. See `setup::maybe_refresh_kimi_quota`. -/// /// Config root: `~/.kimi-code` (override with `KIMI_CODE_HOME`). pub struct KimiCollector { config_root: PathBuf, /// Cached workspace id → root path (from workspaces.json). workspace_roots: HashMap, + /// Cached PID → cwd mapping. On macOS this avoids running one `lsof` + /// process per Kimi session on every fast tick. + process_cwds: HashMap, /// Cached wire parse state keyed by session id. wire_cache: HashMap, /// Version string from `~/.kimi-code/updates/latest.json` (best-effort). @@ -94,11 +92,14 @@ struct WorkspaceEntry { impl KimiCollector { pub fn new() -> Self { let config_root = std::env::var("KIMI_CODE_HOME") + .ok() + .filter(|value| !value.trim().is_empty()) .map(PathBuf::from) - .unwrap_or_else(|_| dirs::home_dir().unwrap_or_default().join(".kimi-code")); + .unwrap_or_else(|| dirs::home_dir().unwrap_or_default().join(".kimi-code")); Self { config_root, workspace_roots: HashMap::new(), + process_cwds: HashMap::new(), wire_cache: HashMap::new(), version: String::new(), } @@ -113,6 +114,7 @@ impl KimiCollector { fn collect_sessions(&mut self, shared: &SharedProcessData) -> Vec { if !self.config_root.is_dir() { + self.process_cwds.clear(); return vec![]; } @@ -123,13 +125,16 @@ impl KimiCollector { let kimi_pids = find_kimi_pids(&shared.process_info); if kimi_pids.is_empty() { + self.process_cwds.clear(); return vec![]; } + self.refresh_process_cwds(&kimi_pids, shared.slow_tick); + // Group live PIDs by cwd. let mut pids_by_cwd: HashMap> = HashMap::new(); for pid in kimi_pids { - let Some(cwd) = get_process_cwd(pid) else { + let Some(cwd) = self.process_cwds.get(&pid).cloned() else { continue; }; if cwd.len() < 2 { @@ -236,7 +241,7 @@ impl KimiCollector { }; let current_tasks = if !wire.current_task.is_empty() - && matches!(status, SessionStatus::Executing | SessionStatus::Thinking) + && matches!(status, SessionStatus::Executing) { vec![wire.current_task.clone()] } else if matches!(status, SessionStatus::Waiting) { @@ -343,6 +348,24 @@ impl KimiCollector { sessions } + fn refresh_process_cwds(&mut self, pids: &[u32], force: bool) { + let live: HashSet = pids.iter().copied().collect(); + self.process_cwds.retain(|pid, _| live.contains(pid)); + + for &pid in pids { + if force || !self.process_cwds.contains_key(&pid) { + match process::get_process_cwd(pid) { + Some(cwd) => { + self.process_cwds.insert(pid, cwd); + } + None => { + self.process_cwds.remove(&pid); + } + } + } + } + } + fn parse_wire_incremental(&mut self, session_id: &str, path: &Path) -> WireState { let mut state = self.wire_cache.get(session_id).cloned().unwrap_or_default(); @@ -453,7 +476,10 @@ fn apply_wire_event(state: &mut WireState, value: &Value) { state.total_cache_read = state.total_cache_read.saturating_add(cache_read); state.total_cache_create = state.total_cache_create.saturating_add(cache_create); state.turn_count = state.turn_count.saturating_add(1); - let turn_tokens = input + output + cache_create; + let turn_tokens = input + .saturating_add(output) + .saturating_add(cache_read) + .saturating_add(cache_create); state.token_history.push(turn_tokens); if state.token_history.len() > 200 { let drain = state.token_history.len() - 200; @@ -515,6 +541,7 @@ fn apply_wire_event(state: &mut WireState, value: &Value) { state.pending_tools = state.pending_tools.saturating_sub(1); if state.pending_tools == 0 { state.pending_since_ms = 0; + state.current_task.clear(); } } "step.end" => { @@ -565,6 +592,9 @@ fn format_tool_task(event: &Value) -> String { .get("command") .and_then(|v| v.as_str()) .unwrap_or("") + .lines() + .next() + .unwrap_or("") .to_string(), "Read" | "Write" | "Edit" | "read" | "write" | "edit" => args .get("file_path") @@ -580,7 +610,6 @@ fn format_tool_task(event: &Value) -> String { .to_string(), "Agent" => args .get("description") - .or_else(|| args.get("prompt")) .and_then(|v| v.as_str()) .unwrap_or("") .to_string(), @@ -593,7 +622,9 @@ fn format_tool_task(event: &Value) -> String { } } }; - let arg = truncate_str(&super::sanitize_terminal_text(&arg), 80); + let arg = super::sanitize_terminal_text(&arg); + let arg = super::redact_secrets(&arg); + let arg = truncate_str(&arg, 40); if arg.is_empty() { name.to_string() } else { @@ -689,7 +720,7 @@ fn find_workspace_id_by_scan(config_root: &Path, cwd: &str) -> Option { for sub in subs.flatten() { let state_path = sub.path().join("state.json"); if let Some(state) = read_state_file(&state_path) { - if paths_equal(&state.work_dir, cwd) { + if process::paths_equal(&state.work_dir, cwd) { return Some(wd_id); } } @@ -740,27 +771,7 @@ fn stub_session( let proc = process_info.get(&pid); let mem_mb = proc.map(|p| p.rss_kb / 1024).unwrap_or(0); let project_name = process::last_path_segment(cwd).unwrap_or("?").to_string(); - // Build a minimal SharedProcessData-like children walk. - let mut children = Vec::new(); - let mut stack: Vec = children_map.get(&pid).cloned().unwrap_or_default(); - let mut visited = HashSet::new(); - while let Some(cpid) = stack.pop() { - if !visited.insert(cpid) { - continue; - } - if let Some(cproc) = process_info.get(&cpid) { - let port = ports.get(&cpid).and_then(|v| v.first().copied()); - children.push(ChildProcess { - pid: cpid, - command: cproc.command.clone(), - mem_kb: cproc.rss_kb, - port, - }); - } - if let Some(gc) = children_map.get(&cpid) { - stack.extend(gc); - } - } + let children = collect_children_from_maps(pid, process_info, children_map, ports); AgentSession { agent_cli: "kimi", pid, @@ -803,15 +814,29 @@ fn stub_session( } fn collect_children(pid: u32, shared: &SharedProcessData) -> Vec { + collect_children_from_maps( + pid, + &shared.process_info, + &shared.children_map, + &shared.ports, + ) +} + +fn collect_children_from_maps( + pid: u32, + process_info: &HashMap, + children_map: &HashMap>, + ports: &HashMap>, +) -> Vec { let mut children = Vec::new(); - let mut stack: Vec = shared.children_map.get(&pid).cloned().unwrap_or_default(); + let mut stack: Vec = children_map.get(&pid).cloned().unwrap_or_default(); let mut visited = HashSet::new(); while let Some(cpid) = stack.pop() { if !visited.insert(cpid) { continue; } - if let Some(cproc) = shared.process_info.get(&cpid) { - let port = shared.ports.get(&cpid).and_then(|v| v.first().copied()); + if let Some(cproc) = process_info.get(&cpid) { + let port = ports.get(&cpid).and_then(|v| v.first().copied()); children.push(ChildProcess { pid: cpid, command: cproc.command.clone(), @@ -819,7 +844,7 @@ fn collect_children(pid: u32, shared: &SharedProcessData) -> Vec { port, }); } - if let Some(gc) = shared.children_map.get(&cpid) { + if let Some(gc) = children_map.get(&cpid) { stack.extend(gc); } } @@ -867,61 +892,6 @@ fn current_time_ms() -> u64 { .as_millis() as u64 } -#[cfg(target_os = "windows")] -fn paths_equal(a: &str, b: &str) -> bool { - let norm = |s: &str| { - s.replace('/', "\\") - .trim_end_matches('\\') - .to_ascii_lowercase() - }; - norm(a) == norm(b) -} - -#[cfg(not(target_os = "windows"))] -fn paths_equal(a: &str, b: &str) -> bool { - a == b -} - -/// Get the current working directory of a process. -#[cfg(target_os = "linux")] -fn get_process_cwd(pid: u32) -> Option { - fs::read_link(format!("/proc/{}/cwd", pid)) - .ok() - .map(|p| p.to_string_lossy().into_owned()) -} - -#[cfg(target_os = "windows")] -fn get_process_cwd(pid: u32) -> Option { - use sysinfo::{Pid, ProcessRefreshKind, ProcessesToUpdate, System, UpdateKind}; - let mut sys = System::new(); - let pid = Pid::from_u32(pid); - sys.refresh_processes_specifics( - ProcessesToUpdate::Some(&[pid]), - false, - ProcessRefreshKind::new().with_cwd(UpdateKind::Always), - ); - sys.process(pid) - .and_then(|p| p.cwd()) - .map(|p| p.to_string_lossy().into_owned()) -} - -#[cfg(all(not(target_os = "linux"), not(target_os = "windows")))] -fn get_process_cwd(pid: u32) -> Option { - use std::process::Command; - let output = Command::new("lsof") - .args(["-a", "-p", &pid.to_string(), "-d", "cwd", "-Fn"]) - .output() - .ok()?; - if !output.status.success() { - return None; - } - let stdout = String::from_utf8_lossy(&output.stdout); - stdout - .lines() - .find(|l| l.starts_with('n') && l.len() > 1) - .map(|l| l[1..].to_string()) -} - // ── Tests ─────────────────────────────────────────────────────────────────── #[cfg(test)] @@ -965,6 +935,7 @@ mod tests { assert_eq!(state.total_output, 50); assert_eq!(state.total_cache_read, 1000); assert_eq!(state.turn_count, 1); + assert_eq!(state.token_history, vec![1150]); assert_eq!(state.last_context_tokens, 1100); assert_eq!(state.model, "kimi-code/k3"); @@ -989,6 +960,31 @@ mod tests { }); apply_wire_event(&mut state, &result); assert_eq!(state.pending_tools, 0); + assert!(state.current_task.is_empty()); + } + + #[test] + fn tool_task_redacts_secrets_and_uses_only_first_command_line() { + let tool = serde_json::json!({ + "name": "Bash", + "args": {"command": "curl -H 'Authorization: Bearer secret-token' example.com\necho leaked"} + }); + + let task = format_tool_task(&tool); + assert!(task.starts_with("Bash curl")); + assert!(task.contains("[REDACTED]")); + assert!(!task.contains("secret-token")); + assert!(!task.contains("echo leaked")); + } + + #[test] + fn agent_task_never_falls_back_to_prompt_text() { + let tool = serde_json::json!({ + "name": "Agent", + "args": {"prompt": "private prompt contents"} + }); + + assert_eq!(format_tool_task(&tool), "Agent"); } #[test] @@ -1019,6 +1015,7 @@ mod tests { let mut collector = KimiCollector { config_root: dir.path().to_path_buf(), workspace_roots: HashMap::new(), + process_cwds: HashMap::new(), wire_cache: HashMap::new(), version: String::new(), }; @@ -1078,4 +1075,77 @@ mod tests { let map = load_workspace_roots(dir.path()); assert_eq!(map.get("wd_x").map(String::as_str), Some("/home/u/proj")); } + + #[test] + fn collects_live_session_from_local_process_and_wire_state() { + let dir = tempdir().unwrap(); + let cwd = "/tmp/kimi-project"; + let session_dir = dir + .path() + .join("sessions") + .join("wd_demo") + .join("session_live"); + let agent_dir = session_dir.join("agents").join("main"); + fs::create_dir_all(&agent_dir).unwrap(); + fs::write( + dir.path().join("workspaces.json"), + format!(r#"{{"workspaces":{{"wd_demo":{{"root":"{cwd}"}}}}}}"#), + ) + .unwrap(); + fs::write( + session_dir.join("state.json"), + format!( + r#"{{"createdAt":"2026-08-03T00:00:00Z","updatedAt":"2026-08-03T00:00:01Z","title":"Local Kimi session","workDir":"{cwd}"}}"# + ), + ) + .unwrap(); + fs::write( + agent_dir.join("wire.jsonl"), + concat!( + "{\"type\":\"usage.record\",\"model\":\"kimi-code/k3\",\"usage\":{\"inputOther\":10,\"output\":5,\"inputCacheRead\":100,\"inputCacheCreation\":2},\"usageScope\":\"turn\",\"time\":1}\n", + "{\"type\":\"context.append_loop_event\",\"event\":{\"type\":\"tool.call\",\"name\":\"Read\",\"args\":{\"file_path\":\"src/main.rs\"}},\"time\":2}\n" + ), + ) + .unwrap(); + + let pid = 42; + let mut collector = KimiCollector { + config_root: dir.path().to_path_buf(), + workspace_roots: HashMap::new(), + process_cwds: HashMap::from([(pid, cwd.to_string())]), + wire_cache: HashMap::new(), + version: "0.6.0".to_string(), + }; + let shared = SharedProcessData { + process_info: HashMap::from([( + pid, + process::ProcInfo { + pid, + ppid: 1, + rss_kb: 2048, + cpu_pct: 0.0, + command: "/usr/local/bin/kimi".to_string(), + }, + )]), + children_map: HashMap::new(), + ports: HashMap::new(), + slow_tick: false, + mcp_server_pids: HashSet::new(), + mcp_owned_rollouts: HashSet::new(), + mcp_suppress: true, + desktop_rollout_fd_map: HashMap::new(), + }; + + let sessions = collector.collect_sessions(&shared); + assert_eq!(sessions.len(), 1); + let session = &sessions[0]; + assert_eq!(session.pid, pid); + assert_eq!(session.session_id, "session_live"); + assert_eq!(session.cwd, cwd); + assert_eq!(session.model, "kimi-code/k3"); + assert_eq!(session.status, SessionStatus::Executing); + assert_eq!(session.current_tasks, vec!["Read src/main.rs"]); + assert_eq!(session.total_tokens(), 117); + assert_eq!(session.token_history, vec![117]); + } } diff --git a/src/collector/mod.rs b/src/collector/mod.rs index 14f9435..3f59ddd 100644 --- a/src/collector/mod.rs +++ b/src/collector/mod.rs @@ -60,9 +60,10 @@ pub(crate) fn redact_secrets(s: &str) -> String { let mut result = s.to_string(); for pat in PATTERNS { while let Some(pos) = result.find(pat) { - let end = result[pos..] + let secret_start = pos + pat.len(); + let end = result[secret_start..] .find(char::is_whitespace) - .map(|i| pos + i) + .map(|i| secret_start + i) .unwrap_or(result.len()); result.replace_range(pos..end, "[REDACTED]"); } @@ -509,6 +510,18 @@ impl MultiCollector { mod tests { use super::*; + #[test] + fn redact_secrets_removes_bearer_token_value() { + let redacted = redact_secrets("Authorization: Bearer secret-token next"); + assert_eq!(redacted, "Authorization: [REDACTED] next"); + } + + #[test] + fn redact_secrets_removes_prefixed_token_value() { + let redacted = redact_secrets("token=sk-proj-secret next"); + assert_eq!(redacted, "token=[REDACTED] next"); + } + #[test] fn with_hidden_empty_keeps_all_collectors() { let mc = MultiCollector::with_hidden(&[]); @@ -521,6 +534,12 @@ mod tests { assert_eq!(mc.collectors.len(), 3); } + #[test] + fn with_hidden_kimi_drops_kimi_only() { + let mc = MultiCollector::with_hidden(&["kimi".to_string()]); + assert_eq!(mc.collectors.len(), 3); + } + #[test] fn with_hidden_is_case_insensitive() { let mc = MultiCollector::with_hidden(&["CODEX".to_string()]); diff --git a/src/collector/opencode.rs b/src/collector/opencode.rs index 608e3b6..ce38cab 100644 --- a/src/collector/opencode.rs +++ b/src/collector/opencode.rs @@ -1,4 +1,4 @@ -use super::{process, context_window_for_model}; +use super::{context_window_for_model, process}; use crate::model::{AgentSession, ChildProcess, SessionStatus}; use serde_json::Value; use std::collections::{HashMap, HashSet}; @@ -286,8 +286,8 @@ impl OpenCodeCollector { if claimed_pids.contains(&pid) { continue; } - if let Some(cwd) = get_process_cwd(pid) { - if paths_equal(&cwd, session_dir) { + if let Some(cwd) = process::get_process_cwd(pid) { + if process::paths_equal(&cwd, session_dir) { return Some(pid); } } @@ -464,24 +464,6 @@ fn truncate_field(s: &mut String, max_bytes: usize) { } } -/// Compare a process cwd with a DB session directory. -/// On Windows paths are case-insensitive and may mix `/` and `\`, so -/// normalize before comparing; elsewhere keep the exact comparison. -#[cfg(target_os = "windows")] -fn paths_equal(a: &str, b: &str) -> bool { - let norm = |s: &str| { - s.replace('/', "\\") - .trim_end_matches('\\') - .to_ascii_lowercase() - }; - norm(a) == norm(b) -} - -#[cfg(not(target_os = "windows"))] -fn paths_equal(a: &str, b: &str) -> bool { - a == b -} - /// On Windows, OpenCode builds (e.g. installed via npm) have been observed to /// keep the XDG-style `~/.local/share/opencode` layout, so prefer the same /// path as unix; fall back to probing `%LOCALAPPDATA%` / `%APPDATA%` in case @@ -505,52 +487,6 @@ fn windows_db_path(default: PathBuf) -> PathBuf { default } -/// Get the current working directory of a process. -/// Uses /proc on Linux, sysinfo (PEB) on Windows, lsof on macOS/other Unix. -#[cfg(target_os = "linux")] -fn get_process_cwd(pid: u32) -> Option { - std::fs::read_link(format!("/proc/{}/cwd", pid)) - .ok() - .map(|p| p.to_string_lossy().into_owned()) -} - -#[cfg(target_os = "windows")] -fn get_process_cwd(pid: u32) -> Option { - use sysinfo::{Pid, ProcessRefreshKind, ProcessesToUpdate, System, UpdateKind}; - // `lsof` does not exist on Windows; sysinfo reads the cwd from the - // process PEB. Refresh just this one PID — this runs only for the - // handful of opencode PIDs, once per tick. - let mut sys = System::new(); - let pid = Pid::from_u32(pid); - sys.refresh_processes_specifics( - ProcessesToUpdate::Some(&[pid]), - false, - ProcessRefreshKind::new().with_cwd(UpdateKind::Always), - ); - sys.process(pid) - .and_then(|p| p.cwd()) - .map(|p| p.to_string_lossy().into_owned()) -} - -#[cfg(all(not(target_os = "linux"), not(target_os = "windows")))] -fn get_process_cwd(pid: u32) -> Option { - // -a ANDs the selection terms; without it, lsof ORs `-p ` with - // `-d cwd` and returns cwd entries for unrelated processes too. - let output = Command::new("lsof") - .args(["-a", "-p", &pid.to_string(), "-d", "cwd", "-Fn"]) - .output() - .ok()?; - if !output.status.success() { - return None; - } - let stdout = String::from_utf8_lossy(&output.stdout); - // lsof -Fn output: lines starting with 'n' contain the path - stdout - .lines() - .find(|l| l.starts_with('n') && l.len() > 1) - .map(|l| l[1..].to_string()) -} - fn current_time_ms() -> u64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) diff --git a/src/collector/process.rs b/src/collector/process.rs index a700f8f..1d8504d 100644 --- a/src/collector/process.rs +++ b/src/collector/process.rs @@ -395,6 +395,65 @@ pub fn last_path_segment(s: &str) -> Option<&str> { segment } +/// Compare local filesystem paths using the host platform's path semantics. +#[cfg(target_os = "windows")] +pub fn paths_equal(a: &str, b: &str) -> bool { + let normalize = |value: &str| { + value + .replace('/', "\\") + .trim_end_matches('\\') + .to_ascii_lowercase() + }; + normalize(a) == normalize(b) +} + +#[cfg(not(target_os = "windows"))] +pub fn paths_equal(a: &str, b: &str) -> bool { + a == b +} + +/// Get the current working directory of a process. +/// Uses `/proc` on Linux, `sysinfo` on Windows, and `lsof` on other Unix hosts. +#[cfg(target_os = "linux")] +pub fn get_process_cwd(pid: u32) -> Option { + fs::read_link(format!("/proc/{pid}/cwd")) + .ok() + .map(|p| p.to_string_lossy().into_owned()) +} + +#[cfg(target_os = "windows")] +pub fn get_process_cwd(pid: u32) -> Option { + use sysinfo::{Pid, ProcessRefreshKind, ProcessesToUpdate, System, UpdateKind}; + + let mut sys = System::new(); + let pid = Pid::from_u32(pid); + sys.refresh_processes_specifics( + ProcessesToUpdate::Some(&[pid]), + false, + ProcessRefreshKind::new().with_cwd(UpdateKind::Always), + ); + sys.process(pid) + .and_then(|p| p.cwd()) + .map(|p| p.to_string_lossy().into_owned()) +} + +#[cfg(all(not(target_os = "linux"), not(target_os = "windows")))] +pub fn get_process_cwd(pid: u32) -> Option { + // `-a` ANDs the PID and cwd selectors; without it, lsof also returns cwd + // entries for unrelated processes. + let output = Command::new("lsof") + .args(["-a", "-p", &pid.to_string(), "-d", "cwd", "-Fn"]) + .output() + .ok()?; + if !output.status.success() { + return None; + } + String::from_utf8_lossy(&output.stdout) + .lines() + .find(|line| line.starts_with('n') && line.len() > 1) + .map(|line| line[1..].to_string()) +} + /// Check if a command string has a given binary name in executable position. /// Checks the first two argv tokens only (covers direct invocation and /// interpreter-wrapped scripts like `node /path/to/codex ...`). diff --git a/src/collector/rate_limit.rs b/src/collector/rate_limit.rs index 8062b2c..7a0bba3 100644 --- a/src/collector/rate_limit.rs +++ b/src/collector/rate_limit.rs @@ -2,18 +2,17 @@ use crate::model::RateLimitInfo; use serde::Deserialize; use std::path::{Path, PathBuf}; -/// Shared filename written by companion hooks for Claude and Kimi. -/// Claude: `~/.claude/abtop-rate-limits.json` (StatusLine) -/// Kimi: `~/.kimi-code/abtop-rate-limits.json` (`abtop --setup` usages script) +/// Shared filename written by the Claude hook and optional Kimi companion. const RATE_LIMIT_FILE: &str = "abtop-rate-limits.json"; /// Cached Codex rate limit: ~/.cache/abtop/codex-rate-limits.json const CODEX_CACHE_FILE: &str = "codex-rate-limits.json"; +/// Provider-written hook data older than this is no longer reliable. +const HOOK_STALE_SECS: u64 = 600; + #[derive(Debug, Deserialize)] struct RateLimitFile { - #[serde(default)] - source: String, #[serde(default)] five_hour: Option, #[serde(default)] @@ -32,45 +31,44 @@ struct WindowInfo { window_minutes: Option, } -/// Read account-level rate limits from local hook files only. +/// Read account-level rate limits from local cache files. /// -/// - Claude: `~/.claude/abtop-rate-limits.json` (+ CLAUDE_CONFIG_DIR / discovered dirs) -/// - Kimi: `~/.kimi-code/abtop-rate-limits.json` (+ KIMI_CODE_HOME) -/// -/// abtop never contacts provider APIs. Companion scripts write these files: -/// Claude StatusLine, and for Kimi an auto-spawned `abtop-usages.sh` when -/// credentials exist (also installable via `abtop --setup`). +/// Claude data comes from its StatusLine hook. Kimi data comes from the +/// optional companion installed by `abtop --setup`; this reader never performs +/// network or credential operations itself. pub fn read_rate_limits(extra_dirs: &[PathBuf]) -> Vec { let mut results = Vec::new(); let mut seen = std::collections::HashSet::new(); - // Claude config roots - let mut claude_dirs = Vec::new(); + // Collect candidate directories: defaults + discovered + let mut dirs = Vec::new(); if let Some(home) = dirs::home_dir() { - claude_dirs.push(home.join(".claude")); + dirs.push(home.join(".claude")); } if let Ok(dir) = std::env::var("CLAUDE_CONFIG_DIR") { - claude_dirs.push(PathBuf::from(dir)); + dirs.push(PathBuf::from(dir)); } - claude_dirs.extend_from_slice(extra_dirs); + dirs.extend_from_slice(extra_dirs); - for dir in claude_dirs { + for dir in dirs { if !dir.is_dir() || !seen.insert(dir.clone()) { continue; } let path = dir.join(RATE_LIMIT_FILE); - if let Some(info) = read_rate_file(&path, "claude") { + if let Some(info) = read_rate_file(&path, "claude", true) { results.push(info); } } - // Kimi Code config root (local file written by abtop-usages.sh) + // Keep the last valid Kimi value visible when a refresh fails. The quota + // panel dims cache data older than ten minutes, so users can distinguish a + // stale value without losing the last known account state entirely. for dir in kimi_config_dirs() { if !dir.is_dir() || !seen.insert(dir.clone()) { continue; } let path = dir.join(RATE_LIMIT_FILE); - if let Some(info) = read_rate_file(&path, "kimi") { + if let Some(info) = read_rate_file(&path, "kimi", false) { results.push(info); } } @@ -81,14 +79,13 @@ pub fn read_rate_limits(extra_dirs: &[PathBuf]) -> Vec { fn kimi_config_dirs() -> Vec { let mut dirs = Vec::new(); if let Ok(dir) = std::env::var("KIMI_CODE_HOME") { - let p = PathBuf::from(dir); - if !p.as_os_str().is_empty() { - dirs.push(p); + let path = PathBuf::from(dir); + if !path.as_os_str().is_empty() { + dirs.push(path); } } if let Some(home) = dirs::home_dir() { dirs.push(home.join(".kimi-code")); - // Legacy kimi-cli home, if a user points the companion script there. dirs.push(home.join(".kimi")); } dirs @@ -100,7 +97,7 @@ fn kimi_config_dirs() -> Vec { /// known value regardless of file age — the UI shows "N m ago" for staleness. pub fn read_codex_cache() -> Option { let path = codex_cache_path()?; - read_rate_file(&path, "codex") + read_rate_file(&path, "codex", false) } /// Write Codex rate limit to cache file (atomic: write temp + rename). @@ -160,26 +157,35 @@ fn codex_cache_path() -> Option { dirs::cache_dir().map(|d| d.join("abtop").join(CODEX_CACHE_FILE)) } -fn read_rate_file(path: &Path, default_source: &str) -> Option { +fn read_rate_file(path: &Path, default_source: &str, reject_stale: bool) -> Option { let content = std::fs::read_to_string(path).ok()?; let file: RateLimitFile = serde_json::from_str(&content).ok()?; + if reject_stale + && file.updated_at.is_some_and(|updated_at| { + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + now.saturating_sub(updated_at) > HOOK_STALE_SECS + }) + { + return None; + } + // Reject if both windows are absent if file.five_hour.is_none() && file.seven_day.is_none() { return None; } - let source = if file.source.is_empty() { - default_source.to_string() - } else { - file.source - }; + // The cache location determines the provider. Do not allow malformed or + // hand-edited JSON to impersonate another quota column via `source`. + let source = default_source.to_string(); - // Default window lengths when the hook omits window_minutes: - // Claude/Codex-style 5h + 7d; Kimi short window is also 5h (300 min). - let default_short = 300; - let default_long = if source.eq_ignore_ascii_case("kimi") { - // Kimi's top-level `usage` window is account-period, not fixed 7d. + // Kimi's long window is the account billing/subscription period, not a + // fixed seven-day window. Leave it unlengthed unless the companion reports + // a concrete duration so the UI can label it as the plan period. + let default_long_window = if source.eq_ignore_ascii_case("kimi") { None } else { Some(10_080) @@ -193,14 +199,14 @@ fn read_rate_file(path: &Path, default_source: &str) -> Option { .five_hour .as_ref() .and_then(|w| w.window_minutes) - .or(file.five_hour.as_ref().map(|_| default_short)), + .or(file.five_hour.as_ref().map(|_| 300)), seven_day_pct: file.seven_day.as_ref().map(|w| w.used_percentage), seven_day_resets_at: file.seven_day.as_ref().map(|w| w.resets_at), seven_day_window_minutes: file .seven_day .as_ref() .and_then(|w| w.window_minutes) - .or(file.seven_day.as_ref().and(default_long)), + .or(file.seven_day.as_ref().and(default_long_window)), updated_at: file.updated_at, }) } @@ -208,38 +214,55 @@ fn read_rate_file(path: &Path, default_source: &str) -> Option { #[cfg(test)] mod tests { use super::*; - use std::io::Write; use tempfile::tempdir; + fn write_rate_file(path: &Path, updated_at: u64) { + std::fs::write( + path, + format!( + r#"{{"source":"claude","five_hour":{{"used_percentage":25.0,"resets_at":0}},"updated_at":{updated_at}}}"# + ), + ) + .unwrap(); + } + #[test] - fn reads_kimi_hook_file_from_kimi_code_home() { + fn rejects_stale_hook_data() { let dir = tempdir().unwrap(); - let kimi_home = dir.path().join(".kimi-code"); - std::fs::create_dir_all(&kimi_home).unwrap(); - let path = kimi_home.join(RATE_LIMIT_FILE); - let mut f = std::fs::File::create(&path).unwrap(); - write!( - f, - r#"{{"source":"kimi","five_hour":{{"used_percentage":34.0,"resets_at":1785746704,"window_minutes":300}},"seven_day":{{"used_percentage":31.0,"resets_at":1786067104}},"updated_at":1785739503}}"# + let path = dir.path().join(RATE_LIMIT_FILE); + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + write_rate_file(&path, now.saturating_sub(HOOK_STALE_SECS + 1)); + + assert!(read_rate_file(&path, "claude", true).is_none()); + } + + #[test] + fn codex_cache_can_keep_stale_data() { + let dir = tempdir().unwrap(); + let path = dir.path().join(CODEX_CACHE_FILE); + write_rate_file(&path, 1); + + assert!(read_rate_file(&path, "codex", false).is_some()); + } + + #[test] + fn kimi_cache_keeps_stale_data_and_does_not_invent_a_weekly_window() { + let dir = tempdir().unwrap(); + let path = dir.path().join(RATE_LIMIT_FILE); + std::fs::write( + &path, + r#"{"source":"kimi","five_hour":{"used_percentage":40.0,"resets_at":10,"window_minutes":300},"seven_day":{"used_percentage":32.0,"resets_at":20},"updated_at":1}"#, ) .unwrap(); - // Temporarily point home-like discovery via KIMI_CODE_HOME. - let prev = std::env::var_os("KIMI_CODE_HOME"); - std::env::set_var("KIMI_CODE_HOME", &kimi_home); - let results = read_rate_limits(&[]); - match prev { - Some(v) => std::env::set_var("KIMI_CODE_HOME", v), - None => std::env::remove_var("KIMI_CODE_HOME"), - } - - let kimi = results - .iter() - .find(|r| r.source.eq_ignore_ascii_case("kimi")) - .expect("kimi rate limit present"); - assert!((kimi.five_hour_pct.unwrap() - 34.0).abs() < 0.01); - assert_eq!(kimi.five_hour_window_minutes, Some(300)); - assert!((kimi.seven_day_pct.unwrap() - 31.0).abs() < 0.01); - assert!(kimi.seven_day_window_minutes.is_none()); + let info = read_rate_file(&path, "kimi", false).expect("stale Kimi cache is retained"); + assert_eq!(info.five_hour_pct, Some(40.0)); + assert_eq!(info.five_hour_window_minutes, Some(300)); + assert_eq!(info.seven_day_pct, Some(32.0)); + assert_eq!(info.seven_day_window_minutes, None); + assert!(read_rate_file(&path, "kimi", true).is_none()); } } diff --git a/src/kimi_usages.py b/src/kimi_usages.py new file mode 100644 index 0000000..f8eba56 --- /dev/null +++ b/src/kimi_usages.py @@ -0,0 +1,262 @@ +#!/usr/bin/env python3 +"""Fetch Kimi Code account quota for abtop's local cache. + +# abtop-kimi-usages-v2 + +Installed explicitly by `abtop --setup`. The abtop process only reads the +resulting JSON file; this companion owns the optional network and OAuth work. +""" + +import json +import os +import stat +import sys +import tempfile +import time +import urllib.error +import urllib.parse +import urllib.request +from datetime import datetime +from pathlib import Path + +CLIENT_ID = "17e5f671-d194-4dfb-9706-5516cb48c098" +TOKEN_URL = "https://auth.kimi.com/api/oauth/token" +USAGES_URL = "https://api.kimi.com/coding/v1/usages" +TIMEOUT_SECONDS = 8 +MAX_RESPONSE_BYTES = 1024 * 1024 + + +def kimi_home(): + explicit = os.environ.get("KIMI_CODE_HOME", "").strip() + if explicit: + return Path(explicit).expanduser() + + home = Path.home() + for candidate in (home / ".kimi-code", home / ".kimi"): + if (candidate / "credentials" / "kimi-code.json").is_file(): + return candidate + return home / ".kimi-code" + + +def read_json(path): + try: + value = json.loads(path.read_text(encoding="utf-8")) + return value if isinstance(value, dict) else None + except (OSError, ValueError): + return None + + +def read_response(response): + payload = response.read(MAX_RESPONSE_BYTES + 1) + if len(payload) > MAX_RESPONSE_BYTES: + raise RuntimeError("response exceeded 1 MiB") + value = json.loads(payload.decode("utf-8")) + if not isinstance(value, dict): + raise RuntimeError("response was not a JSON object") + return value + + +def atomic_write_json(path, value, mode=None): + path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary = tempfile.mkstemp(prefix=path.name + ".", dir=path.parent) + temporary_path = Path(temporary) + try: + with os.fdopen(fd, "w", encoding="utf-8") as output: + json.dump(value, output, separators=(",", ":")) + output.write("\n") + output.flush() + os.fsync(output.fileno()) + if mode is not None: + os.chmod(temporary_path, mode) + os.replace(temporary_path, path) + finally: + try: + temporary_path.unlink() + except FileNotFoundError: + pass + + +def credentials_expired(credentials): + try: + expires_at = float(credentials.get("expires_at")) + except (TypeError, ValueError): + return False + return expires_at > 0 and expires_at <= time.time() + 30 + + +def refresh_credentials(credentials, path): + refresh_token = str(credentials.get("refresh_token") or "") + if not refresh_token: + raise RuntimeError("OAuth credentials expired; run `kimi login`") + + body = urllib.parse.urlencode( + { + "client_id": CLIENT_ID, + "grant_type": "refresh_token", + "refresh_token": refresh_token, + } + ).encode() + request = urllib.request.Request( + TOKEN_URL, + data=body, + method="POST", + headers={ + "Content-Type": "application/x-www-form-urlencoded", + "X-Msh-Platform": "kimi_cli", + }, + ) + with urllib.request.urlopen(request, timeout=TIMEOUT_SECONDS) as response: + refreshed = read_response(response) + + access_token = str(refreshed.get("access_token") or "") + if not access_token: + raise RuntimeError("OAuth refresh did not return an access token") + + try: + expires_in = max(1.0, float(refreshed.get("expires_in") or 900)) + except (TypeError, ValueError): + expires_in = 900.0 + + next_credentials = dict(credentials) + next_credentials.update(refreshed) + next_credentials["access_token"] = access_token + next_credentials["refresh_token"] = str( + refreshed.get("refresh_token") or refresh_token + ) + next_credentials["expires_at"] = time.time() + expires_in + + try: + existing_mode = stat.S_IMODE(path.stat().st_mode) + except OSError: + existing_mode = 0o600 + atomic_write_json(path, next_credentials, existing_mode) + return next_credentials + + +def fetch_usages(access_token): + request = urllib.request.Request( + USAGES_URL, + headers={ + "Authorization": "Bearer " + access_token, + "Accept": "application/json", + }, + ) + with urllib.request.urlopen(request, timeout=TIMEOUT_SECONDS) as response: + return read_response(response) + + +def number(value): + if value is None or value == "": + return None + try: + return float(value) + except (TypeError, ValueError): + return None + + +def reset_timestamp(value): + if not isinstance(value, str) or not value: + return None + try: + return int(datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp()) + except (TypeError, ValueError): + return None + + +def window_from_usage(value, window_minutes=None): + if not isinstance(value, dict): + return None + limit = number(value.get("limit")) + if limit is None or limit <= 0: + return None + used = number(value.get("used")) + if used is None: + remaining = number(value.get("remaining")) + if remaining is None: + return None + used = limit - remaining + + window = { + "used_percentage": max(0.0, min(100.0, used / limit * 100.0)), + "resets_at": reset_timestamp( + value.get("resetTime") or value.get("reset_at") or value.get("resetAt") + ) + or 0, + } + if window_minutes is not None: + window["window_minutes"] = window_minutes + return window + + +def window_minutes(value): + if not isinstance(value, dict) or not isinstance(value.get("window"), dict): + return None + duration = number(value["window"].get("duration")) + if duration is None or duration <= 0: + return None + unit = str(value["window"].get("timeUnit") or "TIME_UNIT_MINUTE") + duration = int(duration) + if unit in ("TIME_UNIT_SECOND", "SECOND", "seconds"): + return max(1, (duration + 59) // 60) + if unit in ("TIME_UNIT_HOUR", "HOUR", "hours"): + return duration * 60 + if unit in ("TIME_UNIT_DAY", "DAY", "days"): + return duration * 24 * 60 + return duration + + +def normalize_usages(body): + short = None + limits = body.get("limits") + if isinstance(limits, list) and limits: + entry = limits[0] if isinstance(limits[0], dict) else None + if entry is not None: + detail = entry.get("detail") + if not isinstance(detail, dict): + detail = entry + short = window_from_usage(detail, window_minutes(entry) or 300) + + long = window_from_usage(body.get("usage")) + if short is None and long is None: + return None + + result = {"source": "kimi", "updated_at": int(time.time())} + if short is not None: + result["five_hour"] = short + if long is not None: + result["seven_day"] = long + return result + + +def main(): + home = kimi_home() + credentials_path = home / "credentials" / "kimi-code.json" + output_path = home / "abtop-rate-limits.json" + credentials = read_json(credentials_path) + if not credentials or not str(credentials.get("access_token") or ""): + print("abtop-usages: Kimi is not logged in; run `kimi login`", file=sys.stderr) + return 1 + + try: + if credentials_expired(credentials): + credentials = refresh_credentials(credentials, credentials_path) + try: + body = fetch_usages(str(credentials["access_token"])) + except urllib.error.HTTPError as error: + if error.code != 401: + raise + credentials = refresh_credentials(credentials, credentials_path) + body = fetch_usages(str(credentials["access_token"])) + + result = normalize_usages(body) + if result is None: + raise RuntimeError("usage response contained no supported quota windows") + atomic_write_json(output_path, result, 0o600) + return 0 + except Exception as error: + print(f"abtop-usages: {error}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lib.rs b/src/lib.rs index b393efc..2524897 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -100,7 +100,7 @@ pub fn run() -> io::Result<()> { return run_update(); } - // --setup flag: configure StatusLine hook and exit + // --setup flag: configure optional quota companions and exit if std::env::args().any(|a| a == "--setup") { setup::run_setup(); return Ok(()); diff --git a/src/locale.rs b/src/locale.rs index 113c616..d5e9b3e 100644 --- a/src/locale.rs +++ b/src/locale.rs @@ -202,7 +202,8 @@ static LOCALE_EN: LazyLock> = LazyLock::ne m.insert("quota.no_data", "no data"); m.insert("quota.abtop_setup", "abtop --setup"); m.insert("quota.run_codex", "run codex once"); - m.insert("quota.kimi_login", "kimi login"); + m.insert("quota.kimi_setup", "abtop --setup"); + m.insert("quota.plan", "plan"); m.insert("quota.total", "total"); m.insert("quota.in", "in"); @@ -448,7 +449,8 @@ static LOCALE_ZH: LazyLock> = LazyLock::ne m.insert("quota.no_data", "无数据"); m.insert("quota.abtop_setup", "abtop --setup"); m.insert("quota.run_codex", "运行一次 codex"); - m.insert("quota.kimi_login", "kimi login"); + m.insert("quota.kimi_setup", "abtop --setup"); + m.insert("quota.plan", "plan"); m.insert("quota.total", "总计"); m.insert("quota.in", "还有"); diff --git a/src/setup.rs b/src/setup.rs index 42ca014..9bcba3f 100644 --- a/src/setup.rs +++ b/src/setup.rs @@ -36,276 +36,14 @@ with open(os.path.join(config_dir, 'abtop-rate-limits.json'), 'w') as f: " 2>/dev/null "#; -/// Companion script for Kimi Code account quota. -/// Network/auth stay outside abtop — this script writes the same local JSON -/// shape Claude's StatusLine hook produces (`abtop-rate-limits.json`). -/// -/// Logic mirrors TokenTracker's `fetchKimiLimits` (OAuth refresh + usages API). -/// Marker so we can detect an outdated installed script and rewrite it. -const KIMI_USAGES_SCRIPT_VERSION: &str = "abtop-kimi-usages-v1"; - -/// Min gap between background companion spawns (matches TokenTracker ~2 min). +const KIMI_USAGES_SCRIPT: &str = include_str!("kimi_usages.py"); +const KIMI_USAGES_SCRIPT_VERSION: &str = "abtop-kimi-usages-v2"; +const KIMI_USAGES_SCRIPT_MARKER: &str = "abtop-kimi-usages-v"; const KIMI_REFRESH_INTERVAL: Duration = Duration::from_secs(120); - -/// Treat the local rate-limit file as fresh enough to skip a refresh spawn. -const KIMI_FILE_FRESH: Duration = Duration::from_secs(120); - -/// File missing or older than this → cold start: run the companion once -/// synchronously so the first quota paint is not empty. -/// Matches the quota panel's "stale" dim threshold (10 minutes). -const KIMI_FILE_ANCIENT: Duration = Duration::from_secs(600); - -/// Max wait for a cold-start companion run (network + token refresh). -const KIMI_COLD_START_TIMEOUT: Duration = Duration::from_secs(5); +const KIMI_SETUP_TIMEOUT: Duration = Duration::from_secs(10); static KIMI_LAST_SPAWN: Mutex> = Mutex::new(None); -const KIMI_USAGES_SCRIPT: &str = r#"#!/usr/bin/env python3 -"""abtop Kimi Code usage companion — writes local rate-limit JSON for abtop. - -# abtop-kimi-usages-v1 - -Installed automatically when abtop detects Kimi credentials, or via: - abtop --setup - -abtop may spawn this in the background; it only *reads* the JSON this writes. -Network/auth stay in this script — not in the abtop process. -""" -import json -import os -import sys -import time -import urllib.error -import urllib.parse -import urllib.request -from pathlib import Path - -CLIENT_ID = "17e5f671-d194-4dfb-9706-5516cb48c098" -TOKEN_URL = "https://auth.kimi.com/api/oauth/token" -USAGES_URL = "https://api.kimi.com/coding/v1/usages" -TIMEOUT = 8 - - -def kimi_home(): - explicit = os.environ.get("KIMI_CODE_HOME", "").strip() - if explicit: - return Path(explicit).expanduser() - home = Path.home() - code = home / ".kimi-code" - creds = code / "credentials" / "kimi-code.json" - if creds.is_file(): - try: - data = json.loads(creds.read_text()) - if data.get("access_token"): - return code - except Exception: - pass - legacy = home / ".kimi" - if (legacy / "credentials" / "kimi-code.json").is_file(): - return legacy - return code - - -def load_creds(path): - if not path.is_file(): - return None - try: - return json.loads(path.read_text()) - except Exception: - return None - - -def save_creds(path, creds): - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(creds, indent=2) + "\n") - - -def expired(creds, now=None): - now = time.time() if now is None else now - exp = creds.get("expires_at") - try: - exp_f = float(exp) - except (TypeError, ValueError): - return False - if exp_f <= 0: - return False - return exp_f <= now + 30 - - -def refresh_token(creds, creds_path): - body = urllib.parse.urlencode( - { - "client_id": CLIENT_ID, - "grant_type": "refresh_token", - "refresh_token": creds.get("refresh_token") or "", - } - ).encode() - req = urllib.request.Request( - TOKEN_URL, - data=body, - method="POST", - headers={ - "Content-Type": "application/x-www-form-urlencoded", - "X-Msh-Platform": "kimi_cli", - }, - ) - with urllib.request.urlopen(req, timeout=TIMEOUT) as resp: - data = json.loads(resp.read().decode()) - if not data.get("access_token"): - raise RuntimeError("token refresh missing access_token") - expires_in = float(data.get("expires_in") or 900) - if expires_in <= 0: - expires_in = 900 - next_creds = { - "access_token": str(data["access_token"]), - "refresh_token": str(data.get("refresh_token") or creds.get("refresh_token") or ""), - "expires_at": time.time() + expires_in, - "scope": str(data.get("scope") or "kimi-code"), - "token_type": str(data.get("token_type") or "Bearer"), - "expires_in": expires_in, - } - save_creds(creds_path, next_creds) - return next_creds - - -def fetch_usages(access_token): - req = urllib.request.Request( - USAGES_URL, - headers={ - "Authorization": "Bearer " + access_token, - "Accept": "application/json", - }, - ) - with urllib.request.urlopen(req, timeout=TIMEOUT) as resp: - return json.loads(resp.read().decode()) - - -def num(value): - if value is None or value == "": - return None - try: - return float(value) - except (TypeError, ValueError): - return None - - -def parse_reset(value): - if not isinstance(value, str) or not value: - return None - # fromisoformat handles trailing Z poorly on some Python versions; normalize. - s = value.replace("Z", "+00:00") - try: - from datetime import datetime - - dt = datetime.fromisoformat(s) - return int(dt.timestamp()) - except Exception: - return None - - -def window_from_usage(data, window_minutes=None): - if not isinstance(data, dict): - return None - limit = num(data.get("limit")) - if limit is None or limit <= 0: - return None - used = num(data.get("used")) - if used is None: - remaining = num(data.get("remaining")) - if remaining is None: - return None - used = limit - remaining - pct = max(0.0, min(100.0, (used / limit) * 100.0)) - out = { - "used_percentage": pct, - "resets_at": parse_reset( - data.get("resetTime") or data.get("reset_at") or data.get("resetAt") - ) - or 0, - } - if window_minutes is not None: - out["window_minutes"] = window_minutes - return out - - -def window_minutes_from(entry): - if not isinstance(entry, dict): - return None - window = entry.get("window") - if not isinstance(window, dict): - return None - duration = num(window.get("duration")) - if duration is None or duration <= 0: - return None - unit = str(window.get("timeUnit") or "TIME_UNIT_MINUTE") - d = int(duration) - if unit in ("TIME_UNIT_SECOND", "SECOND", "seconds"): - return max(1, (d + 59) // 60) - if unit in ("TIME_UNIT_HOUR", "HOUR", "hours"): - return d * 60 - if unit in ("TIME_UNIT_DAY", "DAY", "days"): - return d * 24 * 60 - return d - - -def normalize(body): - limits = body.get("limits") if isinstance(body.get("limits"), list) else [] - short = None - if limits: - entry = limits[0] if isinstance(limits[0], dict) else None - if entry is not None: - detail = entry.get("detail") if isinstance(entry.get("detail"), dict) else entry - mins = window_minutes_from(entry) or 300 - short = window_from_usage(detail, mins) - long = window_from_usage(body.get("usage") if isinstance(body.get("usage"), dict) else None) - if short is None and long is None: - return None - out = {"source": "kimi", "updated_at": int(time.time())} - if short is not None: - out["five_hour"] = short - if long is not None: - out["seven_day"] = long - return out - - -def main(): - home = kimi_home() - creds_path = home / "credentials" / "kimi-code.json" - out_path = home / "abtop-rate-limits.json" - creds = load_creds(creds_path) - if not creds or not str(creds.get("access_token") or "").strip(): - print(f"abtop-usages: not logged in (missing {creds_path})", file=sys.stderr) - print(" run: kimi login", file=sys.stderr) - return 1 - try: - if expired(creds) and creds.get("refresh_token"): - creds = refresh_token(creds, creds_path) - try: - body = fetch_usages(str(creds["access_token"])) - except urllib.error.HTTPError as e: - if e.code == 401 and creds.get("refresh_token"): - creds = refresh_token(creds, creds_path) - body = fetch_usages(str(creds["access_token"])) - else: - raise - out = normalize(body) - if not out: - print("abtop-usages: could not parse usages response", file=sys.stderr) - return 1 - tmp = out_path.with_suffix(".tmp") - tmp.write_text(json.dumps(out) + "\n") - tmp.replace(out_path) - print(f"abtop-usages: wrote {out_path}") - return 0 - except Exception as e: - print(f"abtop-usages: {e}", file=sys.stderr) - return 1 - - -if __name__ == "__main__": - sys.exit(main()) -"#; - fn claude_dir() -> PathBuf { std::env::var("CLAUDE_CONFIG_DIR") .ok() @@ -322,217 +60,157 @@ fn settings_path() -> PathBuf { claude_dir().join("settings.json") } -fn kimi_dir() -> PathBuf { - resolve_kimi_home().unwrap_or_else(|| { - std::env::var("KIMI_CODE_HOME") - .ok() - .map(PathBuf::from) - .filter(|p| !p.as_os_str().is_empty()) - .unwrap_or_else(|| dirs::home_dir().unwrap_or_default().join(".kimi-code")) - }) -} - -fn kimi_usages_script_path_for(home: &Path) -> PathBuf { - home.join("abtop-usages.sh") -} - -fn kimi_rate_file_for(home: &Path) -> PathBuf { - home.join("abtop-rate-limits.json") -} - -fn kimi_creds_path_for(home: &Path) -> PathBuf { - home.join("credentials").join("kimi-code.json") -} - -/// Prefer a Kimi home that already has login credentials. fn resolve_kimi_home() -> Option { let mut candidates = Vec::new(); if let Ok(dir) = std::env::var("KIMI_CODE_HOME") { - let p = PathBuf::from(dir); - if !p.as_os_str().is_empty() { - candidates.push(p); + let path = PathBuf::from(dir); + if !path.as_os_str().is_empty() { + candidates.push(path); } } if let Some(home) = dirs::home_dir() { candidates.push(home.join(".kimi-code")); candidates.push(home.join(".kimi")); } - for dir in &candidates { - if kimi_creds_path_for(dir).is_file() { - return Some(dir.clone()); - } - } - candidates.into_iter().next() -} -fn kimi_has_credentials(home: &Path) -> bool { - let path = kimi_creds_path_for(home); - let Ok(text) = fs::read_to_string(path) else { - return false; - }; - // Cheap check — avoid parsing secrets into structured types here. - text.contains("access_token") && text.contains('"') + candidates + .iter() + .find(|home| home.join("credentials").join("kimi-code.json").is_file()) + .cloned() + .or_else(|| candidates.into_iter().next()) } -/// Ensure the companion script exists (and is current), return its path. -pub fn ensure_kimi_usages_script() -> Option { - let home = resolve_kimi_home()?; - fs::create_dir_all(&home).ok()?; - let script = kimi_usages_script_path_for(&home); - let needs_write = match fs::read_to_string(&script) { - Ok(existing) => !existing.contains(KIMI_USAGES_SCRIPT_VERSION), - Err(_) => true, - }; - if needs_write { - fs::write(&script, KIMI_USAGES_SCRIPT).ok()?; - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - let _ = fs::set_permissions(&script, fs::Permissions::from_mode(0o700)); - } - } - Some(script) +fn kimi_usages_script_path(home: &Path) -> PathBuf { + home.join("abtop-usages.sh") } -fn rate_file_age(path: &Path) -> Option { - let meta = fs::metadata(path).ok()?; - let modified = meta.modified().ok()?; - SystemTime::now().duration_since(modified).ok() +fn kimi_rate_file_path(home: &Path) -> PathBuf { + home.join("abtop-rate-limits.json") } -fn rate_file_is_fresh(path: &Path) -> bool { - rate_file_age(path).is_some_and(|d| d < KIMI_FILE_FRESH) +fn installed_kimi_usages_script() -> Option<(PathBuf, PathBuf)> { + let home = resolve_kimi_home()?; + let script = kimi_usages_script_path(&home); + let content = fs::read_to_string(&script).ok()?; + content + .contains(KIMI_USAGES_SCRIPT_MARKER) + .then_some((home, script)) } -/// Missing file, or older than [`KIMI_FILE_ANCIENT`] → worth a blocking cold start. -fn rate_file_is_cold(path: &Path) -> bool { - match rate_file_age(path) { - None => true, - Some(age) => age >= KIMI_FILE_ANCIENT, +fn install_kimi_usages_script(home: &Path) -> Result { + fs::create_dir_all(home)?; + let script = kimi_usages_script_path(home); + let current = fs::read_to_string(&script).unwrap_or_default(); + if !current.contains(KIMI_USAGES_SCRIPT_VERSION) { + fs::write(&script, KIMI_USAGES_SCRIPT)?; } + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + fs::set_permissions(&script, fs::Permissions::from_mode(0o700))?; + } + Ok(script) } -/// Run the companion and wait up to `timeout`, then kill if still running. -fn run_companion_sync(script: &Path, timeout: Duration) { - let mut child = match Command::new(script) +#[cfg(not(windows))] +fn spawn_kimi_companion(script: &Path) -> Option { + Command::new(script) .stdin(Stdio::null()) .stdout(Stdio::null()) .stderr(Stdio::null()) .spawn() - { - Ok(c) => c, - Err(_) => return, - }; - wait_child_with_timeout(&mut child, timeout); + .ok() } -fn wait_child_with_timeout(child: &mut Child, timeout: Duration) { - let start = Instant::now(); +#[cfg(windows)] +fn spawn_kimi_companion(script: &Path) -> Option { + ["python3", "python"].into_iter().find_map(|python| { + Command::new(python) + .arg(script) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .ok() + }) +} + +fn run_kimi_companion(script: &Path, timeout: Duration) -> bool { + let Some(mut child) = spawn_kimi_companion(script) else { + return false; + }; + let started = Instant::now(); loop { match child.try_wait() { - Ok(Some(_)) => return, - Ok(None) => { - if start.elapsed() >= timeout { - let _ = child.kill(); - let _ = child.wait(); - return; - } + Ok(Some(status)) => return status.success(), + Ok(None) if started.elapsed() < timeout => { thread::sleep(Duration::from_millis(50)); } - Err(_) => return, + Ok(None) => { + let _ = child.kill(); + let _ = child.wait(); + return false; + } + Err(_) => { + let _ = child.kill(); + let _ = child.wait(); + return false; + } } } } -fn spawn_companion_background(script: &Path) { - let _ = Command::new(script) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn(); +fn spawn_kimi_companion_background(script: &Path) -> bool { + let Some(mut child) = spawn_kimi_companion(script) else { + return false; + }; + thread::spawn(move || { + let _ = child.wait(); + }); + true +} + +fn file_age(path: &Path) -> Option { + let modified = fs::metadata(path).ok()?.modified().ok()?; + SystemTime::now().duration_since(modified).ok() } -/// Out-of-the-box Kimi quota: if the user is logged in (`kimi login`), install -/// the companion script if needed and refresh when the local rate-limit file -/// is missing or stale. -/// -/// - **Cold start** (file missing or ≥10 min old): run the companion once -/// synchronously (≤5s) so the first quota paint is populated. -/// - **Warm refresh** (file present but >2 min old): fire-and-forget background -/// spawn so the TUI never blocks. -/// -/// abtop still does not perform HTTP itself — the child script does. -/// -/// Call from the full TUI [`crate::app::App::tick`] path only (not -/// `tick_no_summaries`). Background path is non-blocking; cold start may block -/// briefly once. +/// Refresh the explicitly installed Kimi quota companion without blocking the +/// TUI. This never installs files and is not called by headless snapshot paths. pub fn maybe_refresh_kimi_quota() { - let Some(home) = resolve_kimi_home() else { + let Some((home, script)) = installed_kimi_usages_script() else { return; }; - if !kimi_has_credentials(&home) { - return; - } - - let rate_path = kimi_rate_file_for(&home); - if rate_file_is_fresh(&rate_path) { + if file_age(&kimi_rate_file_path(&home)).is_some_and(|age| age < KIMI_REFRESH_INTERVAL) { return; } - let cold = rate_file_is_cold(&rate_path); - - { - let mut last = match KIMI_LAST_SPAWN.lock() { - Ok(g) => g, - Err(_) => return, - }; - if let Some(t) = *last { - if t.elapsed() < KIMI_REFRESH_INTERVAL { - return; - } - } - *last = Some(Instant::now()); - } - - let Some(script) = ensure_kimi_usages_script() else { + let Ok(mut last_spawn) = KIMI_LAST_SPAWN.lock() else { return; }; - - if cold { - run_companion_sync(&script, KIMI_COLD_START_TIMEOUT); - } else { - spawn_companion_background(&script); + if last_spawn.is_some_and(|last| last.elapsed() < KIMI_REFRESH_INTERVAL) { + return; + } + if spawn_kimi_companion_background(&script) { + *last_spawn = Some(Instant::now()); } } -pub fn run_setup() { - println!("abtop --setup: installing local rate-limit companions\n"); - - setup_claude_statusline(); - setup_kimi_usages(); - - println!("\n done!"); - println!(" Claude: rate limits appear after the next Claude response (restart sessions)."); - println!(" Kimi: auto-refreshed while abtop runs if you are logged in (`kimi login`)."); - println!( - " companion script: {}", - kimi_usages_script_path_for(&kimi_dir()).display() - ); -} - -fn setup_claude_statusline() { +fn setup_claude_statusline() -> bool { println!("[claude] StatusLine hook"); + // Ensure ~/.claude directory exists let dir = claude_dir(); if let Err(e) = fs::create_dir_all(&dir) { eprintln!(" ✗ failed to create {}: {}", dir.display(), e); - eprintln!(" skipping Claude setup"); - return; + return false; } + // Step 1: Write the statusline script let script = script_path(); match fs::write(&script, STATUSLINE_SCRIPT) { Ok(_) => { + // chmod +x #[cfg(unix)] { use std::os::unix::fs::PermissionsExt; @@ -542,19 +220,18 @@ fn setup_claude_statusline() { } Err(e) => { eprintln!(" ✗ failed to write {}: {}", script.display(), e); - eprintln!(" skipping Claude setup"); - return; + return false; } } + // Step 2: Update settings.json let settings_file = settings_path(); let mut settings: Value = if settings_file.exists() { let content = match fs::read_to_string(&settings_file) { Ok(c) => c, Err(e) => { eprintln!(" ✗ cannot read {}: {}", settings_file.display(), e); - eprintln!(" skipping Claude settings update"); - return; + return false; } }; match serde_json::from_str(&content) { @@ -565,16 +242,20 @@ fn setup_claude_statusline() { settings_file.display(), e ); - eprintln!(" fix the file manually before re-running --setup"); - return; + eprintln!(" fix the file manually before running --setup"); + return false; } } } else { Value::Object(Default::default()) }; - let obj = settings.as_object_mut().unwrap(); + let Some(obj) = settings.as_object_mut() else { + eprintln!(" ✗ {} must contain a JSON object", settings_file.display()); + return false; + }; + // Check if statusLine is already configured let expected_cmd = script.display().to_string(); if let Some(existing) = obj.get("statusLine") { if let Some(existing_obj) = existing.as_object() { @@ -582,15 +263,15 @@ fn setup_claude_statusline() { let cmd_str = cmd.as_str().unwrap_or(""); if cmd_str != expected_cmd && !cmd_str.is_empty() { eprintln!(" ⚠ statusLine already configured: {}", cmd_str); - eprintln!(" leaving settings unchanged; script still written."); - eprintln!(" to use abtop's hook, point statusLine.command at:"); - eprintln!(" {}", expected_cmd); - return; + eprintln!(" to override, remove the existing statusLine key from:"); + eprintln!(" {}", settings_file.display()); + return false; } } } } + // Set statusLine config obj.insert( "statusLine".to_string(), serde_json::json!({ @@ -606,44 +287,48 @@ fn setup_claude_statusline() { Ok(_) => println!(" ✓ updated {}", settings_file.display()), Err(e) => { eprintln!(" ✗ failed to update {}: {}", settings_file.display(), e); + return false; } } + + true } -fn setup_kimi_usages() { - println!("\n[kimi] usages companion script"); +fn setup_kimi_usages() -> bool { + println!("\n[kimi] quota companion"); + let Some(home) = resolve_kimi_home() else { + eprintln!(" ✗ could not resolve the Kimi data directory"); + return false; + }; + let script = match install_kimi_usages_script(&home) { + Ok(script) => script, + Err(error) => { + eprintln!(" ✗ failed to install {}: {error}", home.display()); + return false; + } + }; + println!(" ✓ wrote {}", script.display()); - let dir = kimi_dir(); - if let Err(e) = fs::create_dir_all(&dir) { - eprintln!(" ✗ failed to create {}: {}", dir.display(), e); - eprintln!(" skipping Kimi setup"); - return; + if run_kimi_companion(&script, KIMI_SETUP_TIMEOUT) { + println!(" ✓ fetched initial quota data"); + } else { + println!(" ⚠ quota fetch unavailable; run `kimi login` and restart abtop"); } + true +} - match ensure_kimi_usages_script() { - Some(script) => { - println!(" ✓ wrote {}", script.display()); - println!( - " writes {} (abtop reads this file only)", - kimi_rate_file_for(&dir).display() - ); - if !kimi_has_credentials(&dir) { - println!(" ⚠ not logged in yet — run `kimi login`, then start abtop"); - return; - } - // Synchronous first fetch so the panel has data immediately. - run_companion_sync(&script, KIMI_COLD_START_TIMEOUT); - if kimi_rate_file_for(&dir).is_file() { - println!(" ✓ initial fetch succeeded"); - } else { - println!(" ⚠ initial fetch failed or timed out (check `kimi login` / network)"); - } - } - None => { - eprintln!(" ✗ failed to install Kimi companion script"); - eprintln!(" skipping Kimi setup"); - } +pub fn run_setup() { + println!("abtop --setup: configuring local quota companions\n"); + + let claude_ok = setup_claude_statusline(); + let kimi_ok = setup_kimi_usages(); + if !claude_ok && !kimi_ok { + std::process::exit(1); } + + println!("\n done!"); + println!(" Claude quota appears after the next Claude response."); + println!(" Kimi quota refreshes while the TUI runs after explicit setup."); } #[cfg(test)] @@ -653,41 +338,27 @@ mod tests { use tempfile::tempdir; #[test] - fn rate_file_missing_is_cold_not_fresh() { + fn file_age_distinguishes_missing_and_fresh_files() { let dir = tempdir().unwrap(); - let path = dir.path().join("missing.json"); - assert!(!rate_file_is_fresh(&path)); - assert!(rate_file_is_cold(&path)); - } + let path = dir.path().join("quota.json"); + assert_eq!(file_age(&path), None); - #[test] - fn rate_file_just_written_is_fresh_not_cold() { - let dir = tempdir().unwrap(); - let path = dir.path().join("fresh.json"); - let mut f = fs::File::create(&path).unwrap(); - writeln!(f, "{{}}").unwrap(); - assert!(rate_file_is_fresh(&path)); - assert!(!rate_file_is_cold(&path)); + let mut file = fs::File::create(&path).unwrap(); + writeln!(file, "{{}}").unwrap(); + assert!(file_age(&path).is_some_and(|age| age < Duration::from_secs(2))); } #[cfg(unix)] #[test] - fn wait_child_timeout_kills_slow_process() { - // A process that would hang if not killed (`sleep` is unix-portable). - let mut child = Command::new("sleep") - .arg("30") - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn() - .expect("spawn sleep"); - let start = Instant::now(); - wait_child_with_timeout(&mut child, Duration::from_millis(200)); - assert!( - start.elapsed() < Duration::from_secs(2), - "timeout path should return quickly" - ); - // Child should be reaped. - assert!(child.try_wait().ok().flatten().is_some()); + fn timed_companion_is_killed_and_reaped() { + let dir = tempdir().unwrap(); + let script = dir.path().join("slow.sh"); + fs::write(&script, "#!/bin/sh\nsleep 30\n").unwrap(); + use std::os::unix::fs::PermissionsExt; + fs::set_permissions(&script, fs::Permissions::from_mode(0o700)).unwrap(); + + let started = Instant::now(); + assert!(!run_kimi_companion(&script, Duration::from_millis(100))); + assert!(started.elapsed() < Duration::from_secs(2)); } } diff --git a/src/ui/quota.rs b/src/ui/quota.rs index ca9ca7d..05a6676 100644 --- a/src/ui/quota.rs +++ b/src/ui/quota.rs @@ -45,7 +45,7 @@ pub(crate) fn draw_quota_panel_active( let ticks_per_min = 30usize; let tokens_per_min: f64 = rates.iter().rev().take(ticks_per_min).sum(); - // Split into side-by-side columns: one per known source (CLAUDE | CODEX). + // Split into side-by-side columns: one per known quota source. // Columns are always rendered so the panel layout stays stable even when a // source has no data yet. let num_sources = SOURCES.len() as u16; @@ -105,13 +105,14 @@ fn draw_source_column( theme: &Theme, ) { let col_w_usize = area.width as usize; + let compact = area.width < 11; let bar_w = col_w_usize.saturating_sub(10).clamp(2, 8); let Some(rl) = rl else { let hint = if source.eq_ignore_ascii_case("claude") { t("quota.abtop_setup") } else if source.eq_ignore_ascii_case("kimi") { - t("quota.kimi_login") + t("quota.kimi_setup") } else { t("quota.run_codex") }; @@ -178,27 +179,35 @@ fn draw_source_column( }; let c = grad_at(cpu_grad, used_pct); let label_5h = format_window_label(rl.five_hour_window_minutes, t("quota.5h")); - let mut s = vec![styled_label( - format!(" {}", label_5h).as_str(), - theme.graph_text, - )]; - s.extend(remaining_bar(remaining, bar_w, cpu_grad, theme.meter_bg)); - s.push(Span::styled( - format!(" {:>3.0}%", remaining), - Style::default().fg(c), - )); - lines.push(Line::from(s)); - // Always reserve the row so both columns line up vertically; - // when there's nothing meaningful to show (stale source or the - // cached reset moment is past), render it blank. - lines.push(Line::from(Span::styled( - if reset.is_empty() { - String::new() - } else { - format!(" {}", reset) - }, - Style::default().fg(theme.graph_text), - ))); + if compact { + lines.push(Line::from(vec![ + styled_label(format!("{} ", label_5h).as_str(), theme.graph_text), + Span::styled(format!("{remaining:.0}%"), Style::default().fg(c)), + ])); + lines.push(Line::default()); + } else { + let mut s = vec![styled_label( + format!(" {}", label_5h).as_str(), + theme.graph_text, + )]; + s.extend(remaining_bar(remaining, bar_w, cpu_grad, theme.meter_bg)); + s.push(Span::styled( + format!(" {:>3.0}%", remaining), + Style::default().fg(c), + )); + lines.push(Line::from(s)); + // Always reserve the row so both columns line up vertically; + // when there's nothing meaningful to show (stale source or the + // cached reset moment is past), render it blank. + lines.push(Line::from(Span::styled( + if reset.is_empty() { + String::new() + } else { + format!(" {}", reset) + }, + Style::default().fg(theme.graph_text), + ))); + } } if let Some(used_pct) = rl.seven_day_pct { let remaining = (100.0 - used_pct).clamp(0.0, 100.0); @@ -210,28 +219,41 @@ fn draw_source_column( String::new() }; let c = grad_at(cpu_grad, used_pct); - let label_7d = format_window_label(rl.seven_day_window_minutes, t("quota.7d")); - let mut s = vec![styled_label( - format!(" {}", label_7d).as_str(), - theme.graph_text, - )]; - s.extend(remaining_bar(remaining, bar_w, cpu_grad, theme.meter_bg)); - s.push(Span::styled( - format!(" {:>3.0}%", remaining), - Style::default().fg(c), - )); - lines.push(Line::from(s)); - // Always reserve the row so both columns line up vertically; - // when there's nothing meaningful to show (stale source or the - // cached reset moment is past), render it blank. - lines.push(Line::from(Span::styled( - if reset.is_empty() { - String::new() - } else { - format!(" {}", reset) - }, - Style::default().fg(theme.graph_text), - ))); + let fallback = if source.eq_ignore_ascii_case("kimi") { + t("quota.plan") + } else { + t("quota.7d") + }; + let label_7d = format_window_label(rl.seven_day_window_minutes, fallback); + if compact { + lines.push(Line::from(vec![ + styled_label(format!("{} ", label_7d).as_str(), theme.graph_text), + Span::styled(format!("{remaining:.0}%"), Style::default().fg(c)), + ])); + lines.push(Line::default()); + } else { + let mut s = vec![styled_label( + format!(" {}", label_7d).as_str(), + theme.graph_text, + )]; + s.extend(remaining_bar(remaining, bar_w, cpu_grad, theme.meter_bg)); + s.push(Span::styled( + format!(" {:>3.0}%", remaining), + Style::default().fg(c), + )); + lines.push(Line::from(s)); + // Always reserve the row so both columns line up vertically; + // when there's nothing meaningful to show (stale source or the + // cached reset moment is past), render it blank. + lines.push(Line::from(Span::styled( + if reset.is_empty() { + String::new() + } else { + format!(" {}", reset) + }, + Style::default().fg(theme.graph_text), + ))); + } } f.render_widget(Paragraph::new(lines), area);