Report CLI limit reset status

2026-05-21 10:06:43 +02:00
parent 0788cffc88
commit 00a63df9a8
5 changed files with 327 additions and 79 deletions
--- 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

--- 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
--- 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
--- 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=
--- 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}
 }