byre/internal/agents/ratelimit_internal_test.go

package agents

import (
	"context"
	"path/filepath"
	"testing"
	"time"

	"g.fc.run/misc/byre/internal/api"
	"g.fc.run/misc/byre/internal/claude"
	"g.fc.run/misc/byre/internal/store"
)

func TestRateLimitedTurn(t *testing.T) {
	ctx := context.Background()
	st, err := store.Open(filepath.Join(t.TempDir(), "byre.db"), func(api.Event) {})
	if err != nil {
		t.Fatalf("open store: %v", err)
	}
	t.Cleanup(func() { st.Close() })

	now := time.Now().UTC()
	_, err = st.EnsureUser(ctx, "alice", "Alice")
	if err != nil {
		t.Fatalf("ensure user: %v", err)
	}
	p := api.Project{Id: "p1", Name: "demo", RepoPath: "/nonexistent", DefaultBranch: "main", Remote: "none", CreatedBy: "alice", CreatedAt: now}
	err = st.CreateProject(ctx, p)
	if err != nil {
		t.Fatalf("create project: %v", err)
	}
	task := api.Task{Id: "t1", ProjectId: "p1", Title: "t", Branch: "b", WorktreePath: "/nonexistent", BaseRef: "x", Status: api.TaskStatusActive, CreatedBy: "alice", CreatedAt: now}
	err = st.CreateTask(ctx, task)
	if err != nil {
		t.Fatalf("create task: %v", err)
	}
	sess := api.AgentSession{Id: "s1", TaskId: "t1", Vendor: api.Claude, VendorSessionId: "s1", State: api.AgentSessionStateRunning, Role: api.Implementer, CreatedBy: "alice", CreatedAt: now}
	err = st.CreateSession(ctx, sess, "p1")
	if err != nil {
		t.Fatalf("create session: %v", err)
	}

	m := New(st, DefaultLimits, t.TempDir())
	l := &live{id: "s1", dir: t.TempDir(), vendor: api.Claude, taskID: "t1", projectID: "p1"}

	feed := func(typ, raw string) {
		t.Helper()
		err := m.handle(ctx, l, claude.Event{Type: typ, Raw: []byte(raw)})
		if err != nil {
			t.Fatalf("handle %s: %v", raw, err)
		}
	}

	feed("rate_limit_event", `{"type":"rate_limit_event","rate_limit_info":{"status":"allowed","rateLimitType":"five_hour","resetsAt":1791972000}}`)
	usage, err := m.Usage(ctx)
	if err != nil {
		t.Fatalf("usage: %v", err)
	}
	if usage.RateLimits["claude"]["status"] != "allowed" {
		t.Fatalf("rate limit status not recorded: %+v", usage.RateLimits)
	}

	feed("rate_limit_event", `{"type":"rate_limit_event","rate_limit_info":{"status":"rejected","rateLimitType":"five_hour","resetsAt":1791972000}}`)
	feed("result", `{"type":"result","subtype":"success","is_error":true,"result":"API Error: rate limited","num_turns":1}`)

	items, err := st.ListInbox(ctx, store.InboxFilter{SessionID: "s1", State: api.Open})
	if err != nil {
		t.Fatalf("list inbox: %v", err)
	}
	if len(items) != 1 || items[0].Kind != api.RateLimited || items[0].Priority != api.Blocked {
		t.Fatalf("want one blocked rate_limited item, got %+v", items)
	}
	limit, ok := items[0].Payload["rate_limit"].(map[string]any)
	if !ok || limit["status"] != "rejected" || items[0].Payload["vendor"] != "claude" {
		t.Fatalf("unexpected payload %+v", items[0].Payload)
	}
}