From 00a63df9a80f402b8f824063adff7df9d6445104 Mon Sep 17 00:00:00 2001 From: Gemma Date: Thu, 21 May 2026 10:06:43 +0200 Subject: [PATCH] Report CLI limit reset status --- README.md | 31 ++-- docker-compose.yml | 2 + go.mod | 5 +- go.sum | 2 + main.go | 366 ++++++++++++++++++++++++++++++++++++++------- 5 files changed, 327 insertions(+), 79 deletions(-) diff --git a/README.md b/README.md index a5b220a..5ba6a9a 100644 --- a/README.md +++ b/README.md @@ -4,31 +4,20 @@ Small Go service that publishes Claude Code and Codex CLI status to Home Assista ## What it reports -The CLIs do not expose a stable standalone "remaining quota" API. This service therefore performs a small real CLI probe and classifies the result: +This publishes the limit/reset style status shown by the interactive CLIs: -- `ok`: CLI authenticated and a tiny probe completed. -- `limited`: output matched rate/usage/quota-limit errors. -- `auth_required`: login/auth is missing or invalid. -- `timeout` / `error`: probe failed for another reason. +- Claude Code: captures `/usage` in a PTY and parses current-session/current-week percent used and reset text. +- Codex CLI: runs a tiny probe, captures `/status` in a PTY, and parses Codex's recorded `rate_limits` object for 5h/week percent used and reset timestamps. -It also publishes any usage object returned by the CLI probe: +It reports: -- Claude: JSON result `usage` from `claude -p ... --output-format json`. -- Codex: JSONL `usage` from `codex exec --json` turn completion. +- `status`: `ok`, `limited`, `auth_required`, `timeout`, or `error`. +- `primary`: short window limit, normally 5h/session, with `used_percent`, `window_minutes`, `resets_at`/`reset_text` where exposed. +- `secondary`: weekly/all-model limit, with percent and reset data where exposed. +- `raw_report`: cleaned slash-command output for debugging/comparison with the CLI. +- `plan_type`: provider plan where exposed. -It now also reports rolling local token usage for each CLI: - -- `session_5h.used_tokens` and `session_5h.percent_used` -- `week.used_tokens` and `week.percent_used` - -The percent fields are computed against env-configured token limits: - -- `CLAUDE_SESSION_5H_TOKEN_LIMIT` -- `CLAUDE_WEEK_TOKEN_LIMIT` -- `CODEX_SESSION_5H_TOKEN_LIMIT` -- `CODEX_WEEK_TOKEN_LIMIT` - -Set a limit to `0` to publish raw used tokens without a percentage. +Token usage is not the main signal. Any local token counters remain as optional attributes only for diagnostics/backward compatibility. ## MQTT topics diff --git a/docker-compose.yml b/docker-compose.yml index 002fc44..8f8b23d 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -4,9 +4,11 @@ services: image: ai-limits-mqtt:latest container_name: ai-limits-mqtt restart: unless-stopped + working_dir: /root/ai-limits-mqtt env_file: - .env volumes: + - /root/ai-limits-mqtt:/root/ai-limits-mqtt:ro - /root/.claude:/root/.claude - /root/.claude.json:/root/.claude.json - /root/.codex:/root/.codex diff --git a/go.mod b/go.mod index 5c46bb3..8dad59b 100644 --- a/go.mod +++ b/go.mod @@ -2,7 +2,10 @@ module ai-limits-mqtt go 1.23 -require github.com/eclipse/paho.mqtt.golang v1.5.0 +require ( + github.com/creack/pty v1.1.24 + github.com/eclipse/paho.mqtt.golang v1.5.0 +) require ( github.com/gorilla/websocket v1.5.3 // indirect diff --git a/go.sum b/go.sum index b555dd0..1e1f8b5 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,5 @@ +github.com/creack/pty v1.1.24 h1:bJrF4RRfyJnbTJqzRLHzcGaZK1NeM5kTC9jGgovnR1s= +github.com/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE= github.com/eclipse/paho.mqtt.golang v1.5.0 h1:EH+bUVJNgttidWFkLLVKaQPGmkTUfQQqjOsyvMGvD6o= github.com/eclipse/paho.mqtt.golang v1.5.0/go.mod h1:du/2qNQVqJf/Sqs4MEL77kR8QTqANF7XU7Fk0aOTAgk= github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= diff --git a/main.go b/main.go index 8e2b5ce..b3753f5 100644 --- a/main.go +++ b/main.go @@ -14,11 +14,13 @@ import ( "os" "os/exec" "path/filepath" + "regexp" "strconv" "strings" "sync" "time" + "github.com/creack/pty" mqtt "github.com/eclipse/paho.mqtt.golang" ) @@ -39,6 +41,9 @@ type Config struct { CodexProbePrompt string ClaudeExtraArgs []string CodexExtraArgs []string + RunSlashCommand bool + ClaudeSlashTimeout time.Duration + CodexSlashTimeout time.Duration ClaudeHistoryDir string CodexStateDB string ClaudeSession5hLimit int64 @@ -54,6 +59,13 @@ type LimitWindow struct { PercentUsed *float64 `json:"percent_used,omitempty"` } +type RateLimitWindowStatus struct { + UsedPercent *float64 `json:"used_percent,omitempty"` + WindowMinutes int64 `json:"window_minutes,omitempty"` + ResetsAt int64 `json:"resets_at,omitempty"` + ResetText string `json:"reset_text,omitempty"` +} + type ProviderStatus struct { Provider string `json:"provider"` Available bool `json:"available"` @@ -66,6 +78,10 @@ type ProviderStatus struct { ExitCode int `json:"exit_code"` Version string `json:"version,omitempty"` Auth string `json:"auth,omitempty"` + PlanType string `json:"plan_type,omitempty"` + RawReport string `json:"raw_report,omitempty"` + Primary RateLimitWindowStatus `json:"primary,omitempty"` + Secondary RateLimitWindowStatus `json:"secondary,omitempty"` Usage map[string]interface{} `json:"usage,omitempty"` Session5h LimitWindow `json:"session_5h"` Week LimitWindow `json:"week"` @@ -119,6 +135,9 @@ func loadConfig() Config { CodexProbePrompt: env("CODEX_PROBE_PROMPT", "Reply with OK only."), ClaudeExtraArgs: splitArgs(env("CLAUDE_EXTRA_ARGS", "")), CodexExtraArgs: splitArgs(env("CODEX_EXTRA_ARGS", "")), + RunSlashCommand: envBool("RUN_SLASH_COMMAND", true), + ClaudeSlashTimeout: envDuration("CLAUDE_SLASH_TIMEOUT", 35*time.Second), + CodexSlashTimeout: envDuration("CODEX_SLASH_TIMEOUT", 30*time.Second), ClaudeHistoryDir: env("CLAUDE_HISTORY_DIR", "/root/.claude/projects"), CodexStateDB: env("CODEX_STATE_DB", "/root/.codex/state_5.sqlite"), ClaudeSession5hLimit: envInt64("CLAUDE_SESSION_5H_TOKEN_LIMIT", 0), @@ -196,40 +215,30 @@ func probeClaude(parent context.Context, cfg Config) ProviderStatus { st.Message = compact(auth, 500) return st } - if !cfg.RunProbeCommand { + + if cfg.RunSlashCommand { + started := time.Now() + report, err := runSlash(parent, cfg.ClaudeSlashTimeout, []string{cfg.ClaudeCommand}, "/usage", 5*time.Second, 24*time.Second) + st.LatencyMS = time.Since(started).Milliseconds() + if err != nil && strings.TrimSpace(report) == "" { + st.Status = classifyFailure("", err.Error()) + st.Message = compact(err.Error(), 500) + return st + } + clean := cleanTerminal(report) + parseClaudeUsageReport(clean, &st) + st.RawReport = compact(clean, 1800) st.Available = true - st.Status = "auth_ok" + st.Status = "ok" + if st.Message == "" { + st.Message = "Claude /usage captured" + } return st } - args := []string{"-p", cfg.ClaudeProbePrompt, "--output-format", "json", "--max-turns", "1", "--no-session-persistence"} - args = append(args, cfg.ClaudeExtraArgs...) - started := time.Now() - out, exitCode, err := runCommand(parent, cfg.ProbeTimeout, cfg.ClaudeCommand, args...) - st.LatencyMS = time.Since(started).Milliseconds() - st.ExitCode = exitCode - if err != nil && strings.TrimSpace(out) == "" { - st.Status = classifyFailure(out, err.Error()) - st.Message = compact(err.Error(), 500) - st.Limited = st.Status == "limited" - return st - } - var doc map[string]interface{} - if json.Unmarshal([]byte(lastJSON(out)), &doc) == nil { - if usage, ok := doc["usage"].(map[string]interface{}); ok { - st.Usage = usage - } - if isErr, _ := doc["is_error"].(bool); !isErr && stringField(doc, "subtype") == "success" { - st.Available = true - st.Status = "ok" - st.Message = stringField(doc, "result") - return st - } - } - st.Status = classifyFailure(out, "") - st.Limited = st.Status == "limited" - st.Message = compact(out, 800) - st.LimitResetText = extractReset(out) + st.Available = true + st.Status = "auth_ok" + st.Message = "Claude auth OK" return st } @@ -243,41 +252,284 @@ func probeCodex(parent context.Context, cfg Config) ProviderStatus { st.Message = compact(auth, 500) return st } - if !cfg.RunProbeCommand { - st.Available = true - st.Status = "auth_ok" - return st + + started := time.Now() + if cfg.RunProbeCommand { + args := []string{"exec", "--skip-git-repo-check", "--ephemeral", "--json", cfg.CodexProbePrompt} + args = append(args, cfg.CodexExtraArgs...) + out, exitCode, err := runCommand(parent, cfg.ProbeTimeout, cfg.CodexCommand, args...) + st.LatencyMS = time.Since(started).Milliseconds() + st.ExitCode = exitCode + if err != nil && strings.TrimSpace(out) == "" { + st.Status = classifyFailure(out, err.Error()) + st.Message = compact(err.Error(), 500) + st.Limited = st.Status == "limited" + return st + } + if usage := parseCodexUsage(out); usage != nil { + st.Usage = usage + } + if exitCode != 0 { + st.Status = classifyFailure(out, "") + st.Limited = st.Status == "limited" + st.Message = compact(out, 800) + st.LimitResetText = extractReset(out) + if limits, ok := latestCodexRateLimits(cfg); ok { + applyCodexRateLimits(&st, limits) + } + return st + } } - args := []string{"exec", "--skip-git-repo-check", "--ephemeral", "--json", cfg.CodexProbePrompt} - args = append(args, cfg.CodexExtraArgs...) - started := time.Now() - out, exitCode, err := runCommand(parent, cfg.ProbeTimeout, cfg.CodexCommand, args...) - st.LatencyMS = time.Since(started).Milliseconds() - st.ExitCode = exitCode - if err != nil && strings.TrimSpace(out) == "" { - st.Status = classifyFailure(out, err.Error()) - st.Message = compact(err.Error(), 500) - st.Limited = st.Status == "limited" - return st + if cfg.RunSlashCommand { + report, err := runSlash(parent, cfg.CodexSlashTimeout, []string{cfg.CodexCommand}, "/status", 5*time.Second, 20*time.Second) + if err == nil || strings.TrimSpace(report) != "" { + st.RawReport = compact(cleanTerminal(report), 1800) + } } - usage := parseCodexUsage(out) - if usage != nil { - st.Usage = usage + if limits, ok := latestCodexRateLimits(cfg); ok { + applyCodexRateLimits(&st, limits) } - if exitCode == 0 && strings.Contains(out, "turn.completed") { - st.Available = true - st.Status = "ok" - st.Message = "OK" - return st + st.Available = true + st.Status = "ok" + if st.Message == "" { + st.Message = "Codex status captured" } - st.Status = classifyFailure(out, "") - st.Limited = st.Status == "limited" - st.Message = compact(out, 800) - st.LimitResetText = extractReset(out) return st } +func runSlash(parent context.Context, timeout time.Duration, argv []string, slash string, sendAfter, stopAfter time.Duration) (string, error) { + ctx, cancel := context.WithTimeout(parent, timeout) + defer cancel() + cmd := exec.CommandContext(ctx, argv[0], argv[1:]...) + cmd.Env = append(os.Environ(), "TERM=xterm-256color", "COLUMNS=140", "LINES=48") + ptmx, err := pty.StartWithSize(cmd, &pty.Winsize{Rows: 48, Cols: 140}) + if err != nil { + return "", err + } + defer func() { _ = ptmx.Close() }() + + var buf bytes.Buffer + deadline := time.Now().Add(timeout) + sentTrust := false + sent := false + stopped := false + start := time.Now() + tmp := make([]byte, 8192) + for time.Now().Before(deadline) { + _ = ptmx.SetReadDeadline(time.Now().Add(200 * time.Millisecond)) + n, _ := ptmx.Read(tmp) + if n > 0 { + buf.Write(tmp[:n]) + } + if !sentTrust && time.Since(start) >= 2*time.Second { + _, _ = ptmx.Write([]byte("\r")) + sentTrust = true + } + if !sent && time.Since(start) >= sendAfter { + _, _ = ptmx.Write([]byte(slash + "\r")) + sent = true + } + if sent && !stopped && time.Since(start) >= stopAfter { + _, _ = ptmx.Write([]byte{3}) + stopped = true + } + if stopped && ctx.Err() != nil { + break + } + if stopped && time.Since(start) > stopAfter+2*time.Second { + break + } + } + _ = cmd.Process.Kill() + _ = cmd.Wait() + if ctx.Err() != nil { + return buf.String(), ctx.Err() + } + return buf.String(), nil +} + +func cleanTerminal(s string) string { + s = regexp.MustCompile(`\x1b\][^\x07]*(?:\x07|\x1b\\)`).ReplaceAllString(s, "") + s = regexp.MustCompile(`\x1b\[[0-?]*[ -/]*[@-~]`).ReplaceAllString(s, "") + s = strings.ReplaceAll(s, "\r", "\n") + s = strings.ReplaceAll(s, "\u00a0", " ") + // Collapse runs of blank lines/spaces but keep line boundaries readable. + lines := strings.Split(s, "\n") + out := make([]string, 0, len(lines)) + for _, line := range lines { + line = strings.TrimSpace(regexp.MustCompile(`[ \t]+`).ReplaceAllString(line, " ")) + if line != "" { + out = append(out, line) + } + } + return strings.Join(out, "\n") +} + +func parseClaudeUsageReport(report string, st *ProviderStatus) { + lower := strings.ToLower(report) + st.PlanType = "subscription" + if strings.Contains(lower, "claude max") { + st.PlanType = "max" + } + if strings.Contains(lower, "usage credits are off") { + st.Auth = strings.TrimSpace(st.Auth + " usage_credits=off") + } + parseClaudeWindow := func(label string) RateLimitWindowStatus { + compactLower := strings.ReplaceAll(lower, " ", "") + compactLabel := strings.ReplaceAll(strings.ToLower(label), " ", "") + idx := strings.Index(compactLower, compactLabel) + if idx < 0 { + return RateLimitWindowStatus{} + } + // The compact index is close enough for slicing because removing spaces only shifts backwards. + if idx > len(report) { + idx = len(report) / 2 + } + chunk := report[idx:] + if len(chunk) > 700 { + chunk = chunk[:700] + } + var w RateLimitWindowStatus + if m := regexp.MustCompile(`(?i)(\d+(?:\.\d+)?)\s*%\s*used`).FindStringSubmatch(chunk); len(m) == 2 { + if f, err := strconv.ParseFloat(m[1], 64); err == nil { + w.UsedPercent = &f + } + } + if m := regexp.MustCompile(`(?i)resets\s*([^\n]+)`).FindStringSubmatch(chunk); len(m) == 2 { + w.ResetText = strings.TrimSpace(m[1]) + } + return w + } + st.Primary = parseClaudeWindow("Current session") + st.Primary.WindowMinutes = 300 + st.Secondary = parseClaudeWindow("Current week") + st.Secondary.WindowMinutes = 10080 + parts := []string{} + if st.Primary.UsedPercent != nil { + parts = append(parts, fmt.Sprintf("session %.1f%% used resets %s", *st.Primary.UsedPercent, st.Primary.ResetText)) + } + if st.Secondary.UsedPercent != nil { + parts = append(parts, fmt.Sprintf("week %.1f%% used resets %s", *st.Secondary.UsedPercent, st.Secondary.ResetText)) + } + if len(parts) > 0 { + st.Message = strings.Join(parts, "; ") + } +} + +func latestCodexRateLimits(cfg Config) (map[string]interface{}, bool) { + root := env("CODEX_SESSIONS_DIR", "/root/.codex/sessions") + var newest string + var newestMod time.Time + _ = filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error { + if err != nil || d.IsDir() || !strings.HasSuffix(path, ".jsonl") { + return nil + } + info, err := d.Info() + if err != nil { + return nil + } + if info.ModTime().After(newestMod) { + newestMod = info.ModTime() + newest = path + } + return nil + }) + if newest == "" { + return nil, false + } + f, err := os.Open(newest) + if err != nil { + return nil, false + } + defer f.Close() + scanner := bufio.NewScanner(f) + scanner.Buffer(make([]byte, 0, 64*1024), 8*1024*1024) + var last map[string]interface{} + for scanner.Scan() { + line := scanner.Bytes() + if !bytes.Contains(line, []byte("rate_limits")) { + continue + } + var ev map[string]interface{} + if json.Unmarshal(line, &ev) != nil { + continue + } + payload, _ := ev["payload"].(map[string]interface{}) + if rl, ok := payload["rate_limits"].(map[string]interface{}); ok { + last = rl + } + if rl, ok := ev["rate_limits"].(map[string]interface{}); ok { + last = rl + } + } + return last, last != nil +} + +func applyCodexRateLimits(st *ProviderStatus, rl map[string]interface{}) { + st.PlanType = stringFromAny(rl["plan_type"]) + if s := stringFromAny(rl["limit_id"]); s != "" { + st.Auth = strings.TrimSpace(st.Auth + " limit_id=" + s) + } + if s := stringFromAny(rl["rate_limit_reached_type"]); s != "" { + st.Limited = true + st.Status = "limited" + st.Message = "Codex rate limit reached: " + s + } + st.Primary = parseRateLimitMap(rl["primary"]) + st.Secondary = parseRateLimitMap(rl["secondary"]) + parts := []string{} + if st.Primary.UsedPercent != nil { + parts = append(parts, fmt.Sprintf("5h %.1f%% used resets %s", *st.Primary.UsedPercent, st.Primary.ResetText)) + } + if st.Secondary.UsedPercent != nil { + parts = append(parts, fmt.Sprintf("week %.1f%% used resets %s", *st.Secondary.UsedPercent, st.Secondary.ResetText)) + } + if len(parts) > 0 && st.Message == "" { + st.Message = strings.Join(parts, "; ") + } +} + +func parseRateLimitMap(v interface{}) RateLimitWindowStatus { + m, _ := v.(map[string]interface{}) + if m == nil { + return RateLimitWindowStatus{} + } + w := RateLimitWindowStatus{} + if f, ok := floatFromAny(m["used_percent"]); ok { + w.UsedPercent = &f + } + if f, ok := floatFromAny(m["window_minutes"]); ok { + w.WindowMinutes = int64(f) + } + if f, ok := floatFromAny(m["resets_at"]); ok { + w.ResetsAt = int64(f) + w.ResetText = time.Unix(w.ResetsAt, 0).Local().Format("Jan 2 15:04 MST") + } + return w +} + +func floatFromAny(v interface{}) (float64, bool) { + switch x := v.(type) { + case float64: + return x, true + case int64: + return float64(x), true + case int: + return float64(x), true + case json.Number: + f, err := x.Float64() + return f, err == nil + } + return 0, false +} +func stringFromAny(v interface{}) string { + if s, ok := v.(string); ok { + return s + } + return "" +} + func baseStatus(provider string) ProviderStatus { return ProviderStatus{Provider: provider, Status: "unknown", CheckedAt: time.Now().UTC().Format(time.RFC3339), ExitCode: -1} }