Add usage commands, a vendor-neutral usage format and a knowledge base

2026-10-19 15:08:57 -07:00
parent 912225cbd1
commit 8a395d4d6e
17 changed files with 1010 additions and 403 deletions
@@ -0,0 +1,205 @@
+package agents
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"log/slog"
+	"maps"
+	"os"
+	"os/exec"
+	"slices"
+	"strings"
+	"syscall"
+	"time"
+
+	"g.fc.run/ai/byre/internal/api"
+)
+
+const usageTimeout = 10 * time.Second
+
+func (m *Manager) loadUsage() {
+	cmds, err := m.store.UsageCommands(context.Background())
+	if err != nil {
+		slog.Error("load usage commands", "err", err)
+		return
+	}
+	m.mu.Lock()
+	m.usageCommands = cmds
+	m.mu.Unlock()
+	go m.readCodexRateLimit()
+	for _, v := range []api.AgentSessionVendor{api.Claude, api.Codex, api.Antigravity} {
+		go m.readUsage(v)
+	}
+}
+
+func (m *Manager) UsageCommands() api.UsageCommands {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return api.UsageCommands{
+		Claude:      m.usageCommands[api.Claude],
+		Codex:       m.usageCommands[api.Codex],
+		Antigravity: m.usageCommands[api.Antigravity],
+	}
+}
+
+func (m *Manager) SetUsageCommands(ctx context.Context, actor string, c api.UsageCommands) error {
+	cmds := map[api.AgentSessionVendor]string{
+		api.Claude:      strings.TrimSpace(c.Claude),
+		api.Codex:       strings.TrimSpace(c.Codex),
+		api.Antigravity: strings.TrimSpace(c.Antigravity),
+	}
+	err := m.store.SetUsageCommands(ctx, actor, cmds)
+	if err != nil {
+		return err
+	}
+
+	changed := []api.AgentSessionVendor{}
+	m.mu.Lock()
+	for v, cmd := range cmds {
+		if m.usageCommands[v] != cmd {
+			changed = append(changed, v)
+			delete(m.rateLimits, v)
+		}
+	}
+	m.usageCommands = cmds
+	m.mu.Unlock()
+
+	m.store.Broadcast(api.Event{Kind: "admin.health", Actor: "byre", Payload: map[string]any{}})
+	for _, v := range changed {
+		if v == api.Codex && cmds[v] == "" {
+			go m.readCodexRateLimit()
+		}
+		go m.readUsage(v)
+	}
+	return nil
+}
+
+func (m *Manager) readUsage(vendor api.AgentSessionVendor) {
+	mu := m.usageMu[vendor]
+	mu.Lock()
+	defer mu.Unlock()
+
+	m.mu.Lock()
+	cmdline := m.usageCommands[vendor]
+	m.mu.Unlock()
+	if cmdline == "" {
+		if vendor == api.Antigravity {
+			m.readAgyRateLimit()
+		}
+		return
+	}
+
+	u, err := runUsageCommand(cmdline)
+	m.mu.Lock()
+	if m.usageCommands[vendor] != cmdline {
+		m.mu.Unlock()
+		return
+	}
+	if err != nil {
+		delete(m.rateLimits, vendor)
+	} else {
+		m.rateLimits[vendor] = u
+	}
+	m.mu.Unlock()
+	if err != nil {
+		slog.Error("usage command", "vendor", vendor, "err", err)
+	}
+	m.store.Broadcast(api.Event{Kind: "admin.health", Actor: "byre", Payload: map[string]any{}})
+}
+
+func (m *Manager) setVendorUsage(vendor api.AgentSessionVendor, u api.Usage) {
+	m.mu.Lock()
+	if m.usageCommands[vendor] != "" {
+		m.mu.Unlock()
+		return
+	}
+	m.rateLimits[vendor] = u
+	m.mu.Unlock()
+	m.store.Broadcast(api.Event{Kind: "admin.health", Actor: "byre", Payload: map[string]any{}})
+}
+
+func runUsageCommand(cmdline string) (api.Usage, error) {
+	shell := os.Getenv("SHELL")
+	if shell == "" {
+		return api.Usage{}, errors.New("SHELL is not set in byre's environment")
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), usageTimeout)
+	defer cancel()
+	cmd := exec.CommandContext(ctx, shell, "-l", "-c", cmdline)
+	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
+	cmd.Cancel = func() error {
+		return syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
+	}
+	cmd.WaitDelay = time.Second
+	out, err := cmd.Output()
+	if err != nil {
+		return api.Usage{}, fmt.Errorf("%s: %w", cmdline, err)
+	}
+	return parseUsage(out)
+}
+
+func parseUsage(out []byte) (api.Usage, error) {
+	in := struct {
+		UpdatedAt *time.Time `json:"updated_at"`
+		Rejected  *string    `json:"rejected"`
+		Windows   []struct {
+			api.UsageWindow
+			ResetsInSeconds *float64 `json:"resets_in_seconds"`
+		} `json:"windows"`
+	}{}
+	err := json.Unmarshal(out, &in)
+	if err != nil {
+		return api.Usage{}, fmt.Errorf("usage output: %w", err)
+	}
+
+	u := api.Usage{UpdatedAt: time.Now(), Rejected: in.Rejected, Windows: []api.UsageWindow{}}
+	if in.UpdatedAt != nil {
+		u.UpdatedAt = *in.UpdatedAt
+	}
+	for _, w := range in.Windows {
+		if w.ResetsAt == nil && w.ResetsInSeconds != nil {
+			w.ResetsAt = new(u.UpdatedAt.Add(time.Duration(*w.ResetsInSeconds * float64(time.Second))))
+		}
+		u.Windows = append(u.Windows, w.UsageWindow)
+	}
+	return u, nil
+}
+
+func percentWindow(label string, fraction float64, resetsAt int64, windowSeconds int64) api.UsageWindow {
+	w := api.UsageWindow{Label: label, Used: fraction * 100, Limit: new(100.0), Unit: "percent"}
+	if resetsAt != 0 {
+		w.ResetsAt = new(time.Unix(resetsAt, 0))
+	}
+	if windowSeconds != 0 {
+		w.WindowSeconds = &windowSeconds
+	}
+	return w
+}
+
+func claudeUsage(info map[string]any) api.Usage {
+	u := api.Usage{UpdatedAt: time.Now(), Windows: []api.UsageWindow{}}
+	if info["status"] == "rejected" {
+		reason, _ := info["rateLimitType"].(string)
+		u.Rejected = new(reason)
+	}
+	windows, _ := info["unifiedWindows"].(map[string]any)
+	for _, name := range slices.Sorted(maps.Keys(windows)) {
+		w, ok := windows[name].(map[string]any)
+		if !ok {
+			continue
+		}
+		used, _ := w["utilization"].(float64)
+		resets, _ := w["resetsAt"].(float64)
+		seconds := int64(0)
+		switch {
+		case strings.HasPrefix(name, "five_hour"):
+			seconds = 5 * 60 * 60
+		case strings.HasPrefix(name, "seven_day"):
+			seconds = 7 * 24 * 60 * 60
+		}
+		u.Windows = append(u.Windows, percentWindow(name, used, int64(resets), seconds))
+	}
+	return u
+}