feat(aegisctl): rate-limiter status/reset/gc + auto-GC on startup by Lincyaw · Pull Request #69 · OperationsPAI/AegisLab

Lincyaw · 2026-04-18T12:10:00Z

Closes OperationsPAI/aegis#21

Implements OperationsPAI/aegis#21. Adds 3 new backend endpoints for rate-limiter inspection and remediation, plus an auto-GC pass at consumer startup that releases tokens held by terminal-state tasks. Prevents the restart_service bucket from getting stuck at HELD=CAP after a task crash. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Copilot

Pull request overview

Adds an admin surface (API + aegisctl) for inspecting and remediating Redis token-bucket rate limiters, and runs a GC pass on consumer startup to prevent leaked tokens from blocking future work (Ops issue #21).

Changes:

Introduces /api/v2/rate-limiters endpoints to list buckets/holders, reset a bucket, and GC leaked tokens.
Adds aegisctl rate-limiter {status|reset|gc} commands and documents them.
Runs a startup GC pass in consumer and both modes; adds unit tests around GC behavior.

Reviewed changes

Copilot reviewed 9 out of 10 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
src/service/producer/rate_limiter.go	Implements list/reset/GC logic against Redis sets + DB task state lookup
src/service/producer/rate_limiter_test.go	Adds regression/unit tests for GC + terminal-state detection
src/handlers/v2/rate_limiters.go	Exposes new v2 handlers with Swagger annotations
src/router/v2.go	Wires new `/api/v2/rate-limiters` routes with JWT + admin guards
src/main.go	Runs rate-limiter GC on startup for consumer/both modes
src/dto/rate_limiter.go	Adds DTOs for list + GC responses
src/cmd/aegisctl/cmd/rate_limiter.go	Adds `aegisctl rate-limiter` subcommands
docs/aegisctl-cli-spec.md	Documents new CLI commands and auto-GC behavior
src/go.mod / src/go.sum	Adds test dependencies (miniredis + sqlite driver)

Copilot · 2026-04-18T12:15:21Z

+	var task database.Task
+	err := db.WithContext(ctx).Select("state").Where("id = ?", taskID).First(&task).Error
+	if err != nil {
+		if strings.Contains(err.Error(), "record not found") {
+			return 0, false, nil
+		}
+		return 0, false, err
+	}


lookupTaskStateWith detects missing tasks by checking strings.Contains(err.Error(), "record not found"), which is brittle and can break GC/list behavior (tokens won’t be released for tasks that are actually missing). Use errors.Is(err, gorm.ErrRecordNotFound) (as used elsewhere in the repo) to reliably detect the not-found case.

Copilot · 2026-04-18T12:15:21Z

+func knownBuckets() map[string]int {
+	return map[string]int{
+		consts.RestartPedestalTokenBucket: consts.MaxConcurrentRestartPedestal,
+		consts.BuildContainerTokenBucket:  consts.MaxConcurrentBuildContainer,
+		consts.AlgoExecutionTokenBucket:   consts.MaxConcurrentAlgoExecution,
+	}


knownBuckets() hard-codes capacities using consts.MaxConcurrent* defaults, but the consumer rate limiter capacity is configurable via consts.MaxTokensKey* (see service/consumer/rate_limiter.go). This means the status endpoint can report an incorrect Capacity when config overrides are used; consider reading the configured value (with fallback to the default) so the admin output matches runtime behavior.

Copilot · 2026-04-18T12:15:21Z

+// isTerminalState mirrors service/producer/task.go:isTaskTerminal.
+// Issue #21 spells these "Success / Failed / -1"; codebase uses
+// TaskCompleted (3), TaskError (-1), TaskCancelled (-2).
+func isTerminalState(state consts.TaskState) bool {
+	return state == consts.TaskCompleted || state == consts.TaskError || state == consts.TaskCancelled
+}


isTerminalState duplicates isTaskTerminal in service/producer/task.go. Since both are in the same package, reusing the existing helper would prevent future divergence (e.g., if terminal states change).

Copilot · 2026-04-18T12:15:22Z

+// GCRateLimiters releases tokens held by terminal-state tasks.
+func GCRateLimiters(ctx context.Context) (released int, touchedBuckets int, err error) {
+	return gcRateLimitersWith(ctx, client.GetRedisClient(), database.DB, knownBuckets())
+}
+
+// gcRateLimitersWith is the testable core.
+func gcRateLimitersWith(ctx context.Context, redisCli *redis.Client, db *gorm.DB, buckets map[string]int) (released int, touchedBuckets int, err error) {
+	for key := range buckets {


gcRateLimitersWith accepts buckets map[string]int but never uses the map values; this makes the API confusing and suggests capacity is relevant when it isn't. Consider changing the parameter to a slice/set of keys, or actually use the capacity values (e.g., for reporting) if they’re intended to matter.

Suggested change

// GCRateLimiters releases tokens held by terminal-state tasks.

func GCRateLimiters(ctx context.Context) (released int, touchedBuckets int, err error) {

return gcRateLimitersWith(ctx, client.GetRedisClient(), database.DB, knownBuckets())

}

// gcRateLimitersWith is the testable core.

func gcRateLimitersWith(ctx context.Context, redisCli *redis.Client, db *gorm.DB, buckets map[string]int) (released int, touchedBuckets int, err error) {

for key := range buckets {

func bucketKeys(buckets map[string]int) []string {

keys := make([]string, 0, len(buckets))

for key := range buckets {

keys = append(keys, key)

}

return keys

}

// GCRateLimiters releases tokens held by terminal-state tasks.

func GCRateLimiters(ctx context.Context) (released int, touchedBuckets int, err error) {

return gcRateLimitersWith(ctx, client.GetRedisClient(), database.DB, bucketKeys(knownBuckets()))

}

// gcRateLimitersWith is the testable core.

func gcRateLimitersWith(ctx context.Context, redisCli *redis.Client, db *gorm.DB, buckets []string) (released int, touchedBuckets int, err error) {

for _, key := range buckets {

Copilot · 2026-04-18T12:15:22Z

+// ListRateLimiters returns each token_bucket:* bucket with its holders.
+func ListRateLimiters(ctx context.Context) (*dto.RateLimiterListResp, error) {
+	redisCli := client.GetRedisClient()
+	bucketCaps := knownBuckets()
+
+	iter := redisCli.Scan(ctx, 0, tokenBucketKeyPrefix+"*", 0).Iterator()
+	for iter.Next(ctx) {
+		key := iter.Val()
+		if _, ok := bucketCaps[key]; !ok {
+			bucketCaps[key] = 0
+		}
+	}
+	if err := iter.Err(); err != nil {
+		return nil, fmt.Errorf("scan token buckets: %w", err)
+	}
+
+	items := make([]dto.RateLimiterItem, 0, len(bucketCaps))
+	for key, capacity := range bucketCaps {
+		holders, err := redisCli.SMembers(ctx, key).Result()
+		if err != nil && err != redis.Nil {
+			return nil, fmt.Errorf("smembers %s: %w", key, err)
+		}
+		holderItems := make([]dto.RateLimiterHolder, 0, len(holders))
+		for _, taskID := range holders {
+			state, found, err := lookupTaskState(ctx, taskID)
+			if err != nil {
+				logrus.WithError(err).WithField("task_id", taskID).
+					Warn("lookup task state for rate-limiter holder")
+			}
+			stateName := "Unknown"
+			terminal := false
+			if found {
+				stateName = consts.GetTaskStateName(state)
+				terminal = isTerminalState(state)
+			} else {
+				terminal = true
+			}
+			holderItems = append(holderItems, dto.RateLimiterHolder{
+				TaskID: taskID, TaskState: stateName, IsTerminal: terminal,
+			})
+		}
+		items = append(items, dto.RateLimiterItem{
+			Bucket:   strings.TrimPrefix(key, tokenBucketKeyPrefix),
+			Key:      key,
+			Capacity: capacity,
+			Held:     len(holders),
+			Holders:  holderItems,
+		})
+	}
+	return &dto.RateLimiterListResp{Items: items}, nil
+}
+
+// ResetRateLimiter deletes the given bucket key from Redis.
+func ResetRateLimiter(ctx context.Context, bucket string) error {
+	key := resolveBucketKey(bucket)
+	if _, ok := knownBuckets()[key]; !ok {
+		return fmt.Errorf("%w: unknown bucket %q", consts.ErrBadRequest, bucket)
+	}
+	redisCli := client.GetRedisClient()
+	n, err := redisCli.Del(ctx, key).Result()
+	if err != nil {
+		return fmt.Errorf("del %s: %w", key, err)
+	}
+	if n == 0 {
+		return fmt.Errorf("%w: bucket %q not present in redis", consts.ErrNotFound, bucket)
+	}
+	logrus.WithField("bucket", key).Warn("rate-limiter bucket reset")
+	return nil
+}


This file adds new admin behavior (ListRateLimiters + ResetRateLimiter) but tests currently only cover GC/terminal-state logic. Adding unit tests for listing (including unknown buckets) and reset (unknown bucket, missing key, successful delete) would help prevent regressions in these new API behaviors.

Copilot · 2026-04-18T12:15:22Z

+func newTestDB(t *testing.T) *gorm.DB {
+	t.Helper()
+	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
+	require.NoError(t, err)
+	require.NoError(t, db.AutoMigrate(&database.Task{}))
+	return db
+}
+
+func newTestRedis(t *testing.T) *redis.Client {
+	t.Helper()
+	mr, err := miniredis.Run()
+	require.NoError(t, err)
+	t.Cleanup(mr.Close)
+	return redis.NewClient(&redis.Options{Addr: mr.Addr()})
+}


Test helpers create a *redis.Client (and a GORM sqlite DB) without closing them; this can leak goroutines/connections across tests. Consider adding t.Cleanup(func(){ _ = rdb.Close() }) and closing the underlying *sql.DB from db.DB() in newTestDB as well.

Copilot · 2026-04-18T12:15:22Z

+// ListRateLimiters returns each token_bucket:* bucket with its holders.
+func ListRateLimiters(ctx context.Context) (*dto.RateLimiterListResp, error) {
+	redisCli := client.GetRedisClient()
+	bucketCaps := knownBuckets()
+
+	iter := redisCli.Scan(ctx, 0, tokenBucketKeyPrefix+"*", 0).Iterator()
+	for iter.Next(ctx) {
+		key := iter.Val()
+		if _, ok := bucketCaps[key]; !ok {
+			bucketCaps[key] = 0
+		}
+	}
+	if err := iter.Err(); err != nil {
+		return nil, fmt.Errorf("scan token buckets: %w", err)
+	}
+


ListRateLimiters intentionally discovers and returns unknown token_bucket:* keys (setting Capacity to 0), but ResetRateLimiter/GCRateLimiters only operate on knownBuckets(). This makes the API inconsistent: an operator can see an unknown bucket in status but can’t reset/GC it via the admin endpoints. Consider either (a) filtering the list to known buckets only, or (b) allowing reset/GC for any token_bucket:* key (with strict prefix validation) so the admin tooling can remediate what it displays.

Suggested change

// ListRateLimiters returns each token_bucket:* bucket with its holders.

func ListRateLimiters(ctx context.Context) (*dto.RateLimiterListResp, error) {

redisCli := client.GetRedisClient()

bucketCaps := knownBuckets()

iter := redisCli.Scan(ctx, 0, tokenBucketKeyPrefix+"*", 0).Iterator()

for iter.Next(ctx) {

key := iter.Val()

if _, ok := bucketCaps[key]; !ok {

bucketCaps[key] = 0

}

}

if err := iter.Err(); err != nil {

return nil, fmt.Errorf("scan token buckets: %w", err)

}

// ListRateLimiters returns each configured token bucket with its holders.

func ListRateLimiters(ctx context.Context) (*dto.RateLimiterListResp, error) {

redisCli := client.GetRedisClient()

bucketCaps := knownBuckets()

Copilot AI review requested due to automatic review settings April 18, 2026 12:10

Lincyaw had a problem deploying to test-server April 18, 2026 12:10 — with GitHub Actions Failure

Copilot started reviewing on behalf of Lincyaw April 18, 2026 12:10 View session

Copilot AI reviewed Apr 18, 2026

View reviewed changes

Lincyaw closed this Apr 19, 2026

Lincyaw deleted the issue-21/rate-limiter-gc branch April 19, 2026 03:47

Navigation Menu

feat(aegisctl): rate-limiter status/reset/gc + auto-GC on startup - #69

feat(aegisctl): rate-limiter status/reset/gc + auto-GC on startup#69
Lincyaw wants to merge 1 commit into
mainfrom
issue-21/rate-limiter-gc

Lincyaw commented Apr 18, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Apr 18, 2026

Uh oh!

Copilot AI Apr 18, 2026

Uh oh!

Copilot AI Apr 18, 2026

Uh oh!

Copilot AI Apr 18, 2026

Uh oh!

Copilot AI Apr 18, 2026

Uh oh!

Copilot AI Apr 18, 2026

Uh oh!

Copilot AI Apr 18, 2026

Uh oh!

This branch had an error being deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

Lincyaw commented Apr 18, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Apr 18, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 18, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 18, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 18, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 18, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 18, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 18, 2026

Choose a reason for hiding this comment

Uh oh!

This branch had an error being deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants