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 d1a254c..d765396 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,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 # StatusLine hook installation (abtop --setup) +├── 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 @@ -29,8 +30,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 Claude/Kimi quota file reads + Codex cache └── model/ ├── mod.rs # Re-exports └── session.rs # AgentSession, SessionStatus, RateLimitInfo, @@ -76,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. @@ -84,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 @@ -171,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 @@ -185,20 +195,24 @@ 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 (Claude Code) +### 10. Rate limits (account quota) + +Quota is account-level and is not present in session transcript JSONL. + +**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. -NOT in transcript JSONL. Collected via StatusLine mechanism. +**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. -`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`. +**Codex** quota comes from local `token_count` events and is cached under `~/.cache/abtop/`. File format read by abtop: ```json @@ -209,12 +223,11 @@ File format read by abtop: "updated_at": 1774714400 } ``` -- Rejects stale data (> 10 minutes old). -- `rate_limits` only present for Pro/Max subscribers. -- Account-level metric, shared across all sessions. +- 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. -### 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. @@ -386,8 +399,10 @@ 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**: + - 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 @@ -405,5 +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**: 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 68116b7..d55c556 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 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 rate limit collection hook +abtop --setup # Install Claude and Kimi quota 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. 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`). 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 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..32e3419 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, + /// Skip optional Kimi quota refreshes when Kimi is hidden. + 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, } } @@ -495,6 +499,12 @@ impl App { /// retry session summaries. Equivalent to [`App::tick_no_summaries`] followed /// by [`App::drain_and_retry_summaries`]. pub fn tick(&mut self) { + // 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(); + } self.tick_no_summaries(); self.drain_and_retry_summaries(); } @@ -1002,6 +1012,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 { @@ -1097,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 new file mode 100644 index 0000000..d065585 --- /dev/null +++ b/src/collector/kimi.rs @@ -0,0 +1,1151 @@ +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. +/// +/// 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). + 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") + .ok() + .filter(|value| !value.trim().is_empty()) + .map(PathBuf::from) + .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(), + } + } + + 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() { + self.process_cwds.clear(); + 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() { + 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) = self.process_cwds.get(&pid).cloned() 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) + { + 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 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(); + + 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 + .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; + 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; + state.current_task.clear(); + } + } + "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("") + .lines() + .next() + .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") + .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 = 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 { + 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 process::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(); + let children = collect_children_from_maps(pid, process_info, children_map, ports); + 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 { + 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 = 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); + } + } + 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 +} + +// ── 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.token_history, vec![1150]); + 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); + 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] + 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(), + process_cwds: 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")); + } + + #[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 930caf1..3f59ddd 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; @@ -58,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]"); } @@ -335,6 +338,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, @@ -504,30 +510,48 @@ 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(&[]); - 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_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()]); - 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 +560,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/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 e9722d5..7a0bba3 100644 --- a/src/collector/rate_limit.rs +++ b/src/collector/rate_limit.rs @@ -2,16 +2,17 @@ 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 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)] @@ -30,9 +31,11 @@ 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 cache files. +/// +/// 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(); @@ -51,8 +54,21 @@ pub fn read_rate_limits(extra_dirs: &[PathBuf]) -> Vec { if !dir.is_dir() || !seen.insert(dir.clone()) { continue; } - let path = dir.join(CLAUDE_RATE_FILE); - if let Some(info) = read_rate_file(&path, "claude") { + let path = dir.join(RATE_LIMIT_FILE); + if let Some(info) = read_rate_file(&path, "claude", true) { + results.push(info); + } + } + + // 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", false) { results.push(info); } } @@ -60,13 +76,28 @@ pub fn read_rate_limits(extra_dirs: &[PathBuf]) -> Vec { results } +fn kimi_config_dirs() -> Vec { + let mut dirs = Vec::new(); + if let Ok(dir) = std::env::var("KIMI_CODE_HOME") { + 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")); + 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 /// 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). @@ -126,19 +157,38 @@ 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() + // 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(); + + // 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 { - file.source + Some(10_080) }; Some(RateLimitInfo { @@ -156,7 +206,63 @@ fn read_rate_file(path: &Path, default_source: &str) -> Option { .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_window)), updated_at: file.updated_at, }) } + +#[cfg(test)] +mod tests { + use super::*; + 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 rejects_stale_hook_data() { + let dir = tempdir().unwrap(); + 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(); + + 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 abf098f..d5e9b3e 100644 --- a/src/locale.rs +++ b/src/locale.rs @@ -202,6 +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_setup", "abtop --setup"); + m.insert("quota.plan", "plan"); m.insert("quota.total", "total"); m.insert("quota.in", "in"); @@ -447,6 +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_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 90f2d9a..9bcba3f 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,14 @@ with open(os.path.join(config_dir, 'abtop-rate-limits.json'), 'w') as f: " 2>/dev/null "#; +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); +const KIMI_SETUP_TIMEOUT: Duration = Duration::from_secs(10); + +static KIMI_LAST_SPAWN: Mutex> = Mutex::new(None); + fn claude_dir() -> PathBuf { std::env::var("CLAUDE_CONFIG_DIR") .ok() @@ -48,14 +60,150 @@ fn settings_path() -> PathBuf { claude_dir().join("settings.json") } -pub fn run_setup() { - println!("abtop --setup: configuring Claude Code StatusLine hook\n"); +fn resolve_kimi_home() -> Option { + let mut candidates = Vec::new(); + if let Ok(dir) = std::env::var("KIMI_CODE_HOME") { + 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")); + } + + candidates + .iter() + .find(|home| home.join("credentials").join("kimi-code.json").is_file()) + .cloned() + .or_else(|| candidates.into_iter().next()) +} + +fn kimi_usages_script_path(home: &Path) -> PathBuf { + home.join("abtop-usages.sh") +} + +fn kimi_rate_file_path(home: &Path) -> PathBuf { + home.join("abtop-rate-limits.json") +} + +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)) +} + +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) +} + +#[cfg(not(windows))] +fn spawn_kimi_companion(script: &Path) -> Option { + Command::new(script) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .ok() +} + +#[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(status)) => return status.success(), + Ok(None) if started.elapsed() < timeout => { + thread::sleep(Duration::from_millis(50)); + } + Ok(None) => { + let _ = child.kill(); + let _ = child.wait(); + return false; + } + Err(_) => { + let _ = child.kill(); + let _ = child.wait(); + return false; + } + } + } +} + +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() +} + +/// 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, script)) = installed_kimi_usages_script() else { + return; + }; + if file_age(&kimi_rate_file_path(&home)).is_some_and(|age| age < KIMI_REFRESH_INTERVAL) { + return; + } + + let Ok(mut last_spawn) = KIMI_LAST_SPAWN.lock() else { + return; + }; + if last_spawn.is_some_and(|last| last.elapsed() < KIMI_REFRESH_INTERVAL) { + return; + } + if spawn_kimi_companion_background(&script) { + *last_spawn = Some(Instant::now()); + } +} + +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); - std::process::exit(1); + return false; } // Step 1: Write the statusline script @@ -72,7 +220,7 @@ pub fn run_setup() { } Err(e) => { eprintln!(" ✗ failed to write {}: {}", script.display(), e); - std::process::exit(1); + return false; } } @@ -83,7 +231,7 @@ pub fn run_setup() { Ok(c) => c, Err(e) => { eprintln!(" ✗ cannot read {}: {}", settings_file.display(), e); - std::process::exit(1); + return false; } }; match serde_json::from_str(&content) { @@ -95,14 +243,17 @@ pub fn run_setup() { e ); eprintln!(" fix the file manually before running --setup"); - std::process::exit(1); + 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(); @@ -114,7 +265,7 @@ pub fn run_setup() { eprintln!(" ⚠ statusLine already configured: {}", cmd_str); eprintln!(" to override, remove the existing statusLine key from:"); eprintln!(" {}", settings_file.display()); - std::process::exit(1); + return false; } } } @@ -136,10 +287,78 @@ pub fn run_setup() { Ok(_) => println!(" ✓ updated {}", settings_file.display()), Err(e) => { eprintln!(" ✗ failed to update {}: {}", settings_file.display(), e); - std::process::exit(1); + return false; } } - println!("\n done! rate limit data will appear in abtop after the next Claude response."); - println!(" restart any running Claude Code sessions to activate."); + true +} + +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()); + + 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 +} + +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)] +mod tests { + use super::*; + use std::io::Write; + use tempfile::tempdir; + + #[test] + fn file_age_distinguishes_missing_and_fresh_files() { + let dir = tempdir().unwrap(); + let path = dir.path().join("quota.json"); + assert_eq!(file_age(&path), None); + + 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 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 1c65813..05a6676 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); @@ -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,11 +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_setup") } else { t("quota.run_codex") }; @@ -176,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); @@ -208,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); 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(); (