mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-11 03:13:24 +00:00
feat(agent): honor per-agent tools.rate_limit_per_hour (#1263)
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).
This commit is contained in:
1 parent
c0531e1e07
commit
eaa011dae3
7 files changed
+64
-6
No files matched your search
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in new issue
Block a user