78 lines
2.7 KiB
Go
78 lines
2.7 KiB
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", 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)
|
|
}
|
|
}
|