Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions cmd/gateway.go
Original file line number Diff line number Diff line change
Expand Up @@ -192,6 +192,18 @@ func runGateway() {
slog.Info("system_configs applied to in-memory config", "keys", len(sysConfigs))
}
}

// Re-apply tool rate limiter using DB-overlaid config. setupToolRegistry
// initialised the limiter from the JSON5 default before ApplySystemConfigs
// ran, so DB-driven changes to tools.rate_limit_per_hour were lost. Replace
// the limiter object now that cfg reflects the DB value. Safe: server has
// not started, no in-flight tool calls.
if cfg.Tools.RateLimitPerHour > 0 {
toolsReg.SetRateLimiter(tools.NewToolRateLimiter(cfg.Tools.RateLimitPerHour))
slog.Info("tool rate limiting reapplied from system_configs", "per_hour", cfg.Tools.RateLimitPerHour)
} else {
toolsReg.SetRateLimiter(nil)
}
setupMemoryEmbeddings(pgStores, providerRegistry)

// Resolve background provider for consolidation + vault enrichment.
Expand Down
36 changes: 36 additions & 0 deletions internal/tools/registry_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,42 @@ func TestRegistry_ExecuteWithContext_RateLimiting(t *testing.T) {
}
}

// SetRateLimiter must be re-callable so that startup code can swap the limiter
// after DB-overlaid config is applied (cmd/gateway.go re-applies once
// system_configs has overlaid the JSON5 default).
func TestRegistry_SetRateLimiter_ReplacesPriorLimiter(t *testing.T) {
reg := NewRegistry()
reg.SetRateLimiter(NewToolRateLimiter(1)) // simulate JSON5 default
reg.Register(&mockTool{name: "tool"})

// Replace with higher limit (simulates DB overlay = 5)
reg.SetRateLimiter(NewToolRateLimiter(5))

for i := range 5 {
result := reg.ExecuteWithContext(context.Background(), "tool", nil,
"", "", "", "session-replace", nil)
if result.IsError {
t.Errorf("call %d should succeed under new 5/h limit: %s", i, result.ForLLM)
}
}
result := reg.ExecuteWithContext(context.Background(), "tool", nil,
"", "", "", "session-replace", nil)
if !result.IsError {
t.Error("6th call should hit the 5/h limit")
}

// Disable rate limiting via nil — verifies the gateway path that disables
// the limiter when cfg.Tools.RateLimitPerHour <= 0.
reg.SetRateLimiter(nil)
for i := range 10 {
result := reg.ExecuteWithContext(context.Background(), "tool", nil,
"", "", "", "session-replace", nil)
if result.IsError {
t.Errorf("call %d after nil limiter should be unbounded", i)
}
}
}

func TestRegistry_ExecuteWithContext_NoRateLimitWithoutSessionKey(t *testing.T) {
reg := NewRegistry()
reg.SetRateLimiter(NewToolRateLimiter(1))
Expand Down
Loading