byre/internal/agents/usage.go

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
}