From eaa011dae3bea6fe5ab2da3bf539a3e9f04a20ec Mon Sep 17 00:00:00 2001 From: Zezae Oh Date: Tue, 23 Jun 2026 14:53:33 +0900 Subject: [PATCH] feat(agent): honor per-agent tools.rate_limit_per_hour (#1263) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The agent config UI (tools-profile-section) and i18n already expose a per-agent "Rate Limit (per hour)" field saved into tools_config, but the backend ignored it — the tool rate limiter only ever used the global tools.rate_limit_per_hour. Wire the field through: ToolPolicySpec gains RateLimitPerHour; the agent loop threads it into the tool-execution context; the registry passes it to the limiter, which uses it in place of the global max when > 0 (0 inherits the global). --- internal/agent/loop_context.go | 3 +++ internal/config/config_channels.go | 3 +++ internal/tools/context_keys.go | 16 ++++++++++++++++ internal/tools/rate_limiter.go | 20 ++++++++++++++++---- internal/tools/rate_limiter_test.go | 22 ++++++++++++++++++++++ internal/tools/registry.go | 5 +++-- ui/web/src/types/agent.ts | 1 + 7 files changed, 64 insertions(+), 6 deletions(-) diff --git a/internal/agent/loop_context.go b/internal/agent/loop_context.go index e8804344..df46f1e4 100644 --- a/internal/agent/loop_context.go +++ b/internal/agent/loop_context.go @@ -113,6 +113,9 @@ func (l *Loop) injectContext(ctx context.Context, req *RunRequest) (contextSetup waitToolCfg = l.agentToolPolicy.Wait ctx = tools.WithWaitToolConfig(ctx, waitToolCfg) } + if l.agentToolPolicy != nil && l.agentToolPolicy.RateLimitPerHour > 0 { + ctx = tools.WithToolRateLimitOverride(ctx, l.agentToolPolicy.RateLimitPerHour) + } if l.sandboxCfg != nil { ctx = tools.WithSandboxConfig(ctx, l.sandboxCfg) } diff --git a/internal/config/config_channels.go b/internal/config/config_channels.go index c95ea267..5f6ff604 100644 --- a/internal/config/config_channels.go +++ b/internal/config/config_channels.go @@ -537,6 +537,9 @@ type ToolPolicySpec struct { ByProvider map[string]*ToolPolicySpec `json:"byProvider,omitempty"` Wait *WaitToolPolicy `json:"wait,omitempty"` ToolCallPrefix string `json:"toolCallPrefix,omitempty"` // prefix to strip from model's tool call names before registry lookup + // RateLimitPerHour overrides the global tools.rate_limit_per_hour for this + // agent (applied per session key). 0 = inherit the global limit. + RateLimitPerHour int `json:"rate_limit_per_hour,omitempty"` } // WaitToolPolicy configures per-agent safety bounds for the wait tool. diff --git a/internal/tools/context_keys.go b/internal/tools/context_keys.go index 923ac901..9af2b8e7 100644 --- a/internal/tools/context_keys.go +++ b/internal/tools/context_keys.go @@ -34,6 +34,10 @@ const ( ctxRunKind toolContextKey = "tool_run_kind" // "notification", "announce", "delegation" ) +// ctxRateLimitOverride carries a per-agent tool rate limit (calls/hour) that +// overrides the global tools.rate_limit_per_hour. 0 means "use the global". +const ctxRateLimitOverride toolContextKey = "tool_rate_limit_override" + // Well-known channel names used for routing and access control. const ( ChannelSystem = "system" @@ -160,6 +164,18 @@ func ToolSessionKeyFromCtx(ctx context.Context) string { return v } +// WithToolRateLimitOverride sets a per-agent tool rate limit (calls/hour) that +// overrides the global tools.rate_limit_per_hour for tools executed under ctx. +func WithToolRateLimitOverride(ctx context.Context, perHour int) context.Context { + return context.WithValue(ctx, ctxRateLimitOverride, perHour) +} + +// ToolRateLimitOverrideFromCtx returns the per-agent override, or 0 if unset. +func ToolRateLimitOverrideFromCtx(ctx context.Context) int { + v, _ := ctx.Value(ctxRateLimitOverride).(int) + return v +} + // WithRunKind injects the run classification (e.g. "notification") into context. func WithRunKind(ctx context.Context, kind string) context.Context { return context.WithValue(ctx, ctxRunKind, kind) diff --git a/internal/tools/rate_limiter.go b/internal/tools/rate_limiter.go index 0ef995ae..04a6bb13 100644 --- a/internal/tools/rate_limiter.go +++ b/internal/tools/rate_limiter.go @@ -28,12 +28,24 @@ func NewToolRateLimiter(maxPerHour int) *ToolRateLimiter { } } -// Allow checks if a tool execution is allowed for the given key. -// Returns nil if allowed, or an error describing the rate limit. +// Allow checks if a tool execution is allowed for the given key against the +// limiter's configured max. Returns nil if allowed, or an error. func (rl *ToolRateLimiter) Allow(key string) error { + return rl.AllowWithLimit(key, 0) +} + +// AllowWithLimit is Allow with a per-call max override (calls/hour). When +// maxOverride > 0 it replaces the configured max for this check — used for the +// per-agent tools.rate_limit_per_hour. maxOverride <= 0 uses the configured max. +func (rl *ToolRateLimiter) AllowWithLimit(key string, maxOverride int) error { rl.mu.Lock() defer rl.mu.Unlock() + max := rl.maxPerHr + if maxOverride > 0 { + max = maxOverride + } + now := time.Now() cutoff := now.Add(-rl.window) @@ -45,8 +57,8 @@ func (rl *ToolRateLimiter) Allow(key string) error { } entries = entries[start:] - if len(entries) >= rl.maxPerHr { - return fmt.Errorf("tool rate limit exceeded: %d actions/hour for key %s", rl.maxPerHr, key) + if len(entries) >= max { + return fmt.Errorf("tool rate limit exceeded: %d actions/hour for key %s", max, key) } // Record this action diff --git a/internal/tools/rate_limiter_test.go b/internal/tools/rate_limiter_test.go index 70cf20cc..e6ae2413 100644 --- a/internal/tools/rate_limiter_test.go +++ b/internal/tools/rate_limiter_test.go @@ -61,6 +61,28 @@ func TestToolRateLimiter_SeparateKeys(t *testing.T) { } } +func TestToolRateLimiter_AllowWithLimit_Override(t *testing.T) { + rl := NewToolRateLimiter(100) // global default 100 + + // A per-agent override of 2 caps this key at 2, regardless of the global 100. + if err := rl.AllowWithLimit("agentA", 2); err != nil { + t.Fatalf("call 1 should be allowed: %v", err) + } + if err := rl.AllowWithLimit("agentA", 2); err != nil { + t.Fatalf("call 2 should be allowed: %v", err) + } + if err := rl.AllowWithLimit("agentA", 2); err == nil { + t.Error("call 3 should be blocked by the override of 2") + } + + // maxOverride <= 0 falls back to the configured global (100), on its own key. + for i := range 3 { + if err := rl.AllowWithLimit("agentB", 0); err != nil { + t.Fatalf("agentB call %d should use global 100: %v", i, err) + } + } +} + func TestToolRateLimiter_WindowExpiry(t *testing.T) { rl := &ToolRateLimiter{ windows: make(map[string][]time.Time), diff --git a/internal/tools/registry.go b/internal/tools/registry.go index 46f13a14..c12a83f7 100644 --- a/internal/tools/registry.go +++ b/internal/tools/registry.go @@ -200,9 +200,10 @@ func (r *Registry) ExecuteWithContext(ctx context.Context, name string, args map ctx = WithToolAsyncCB(ctx, asyncCB) } - // Rate limit check (per session key) + // Rate limit check (per session key). A per-agent override + // (tools.rate_limit_per_hour, threaded via context) wins over the global max. if r.rateLimiter != nil && sessionKey != "" { - if err := r.rateLimiter.Allow(sessionKey); err != nil { + if err := r.rateLimiter.AllowWithLimit(sessionKey, ToolRateLimitOverrideFromCtx(ctx)); err != nil { return ErrorResult(err.Error()) } } diff --git a/ui/web/src/types/agent.ts b/ui/web/src/types/agent.ts index cf042d8d..8fdd2fc4 100644 --- a/ui/web/src/types/agent.ts +++ b/ui/web/src/types/agent.ts @@ -13,6 +13,7 @@ export interface ToolPolicyConfig { max_ms?: number; }; toolCallPrefix?: string; // prefix to strip from model's tool call names + rate_limit_per_hour?: number; // per-agent tool calls/hour (0 = use global) } export interface SubagentsConfig {