feat(core): add outgoing message rate limiting per platform

This commit is contained in:
Max Qian authored and GitHub committed 2026-03-27 23:59:11 +08:00
1 parent 1225b7ee18
commit d7bed29801
6 files changed
+374 -11

No files matched your search

+28
View File
@@ -316,6 +316,34 @@ func main() {
})
}
}
// Wire outgoing rate limiting
{
var maxPS float64
if cfg.OutgoingRateLimit.MaxPerSecond != nil {
maxPS = *cfg.OutgoingRateLimit.MaxPerSecond
}
var burst int
if cfg.OutgoingRateLimit.Burst != nil {
burst = *cfg.OutgoingRateLimit.Burst
}
defaults := core.OutgoingRateLimitCfg{MaxPerSecond: maxPS, Burst: burst}
overrides := make(map[string]core.OutgoingRateLimitCfg)
for name, pc := range cfg.OutgoingRateLimit.Platforms {
var mps float64
if pc.MaxPerSecond != nil {
mps = *pc.MaxPerSecond
}
var b int
if pc.Burst != nil {
b = *pc.Burst
}
overrides[name] = core.OutgoingRateLimitCfg{MaxPerSecond: mps, Burst: b}
}
if maxPS > 0 || len(overrides) > 0 {
engine.SetOutgoingRateLimitCfg(defaults, overrides)
}
}
engine.SetDisplaySaveFunc(func(thinkingMaxLen, toolMaxLen *int) error {
return config.SaveDisplayConfig(thinkingMaxLen, toolMaxLen)
})
+18
View File
@@ -84,6 +84,24 @@ level = "info" # debug, info, warn, error
# max_messages = 20 # Max messages per window; 0 = disabled (default: 20) / 窗口内最大消息数;0 = 禁用(默认 20)
# window_secs = 60 # Window size in seconds (default: 60) / 窗口时间秒数(默认 60)
# =============================================================================
# Outgoing Rate Limit / 出站速率限制
# Throttle how fast messages are sent TO platforms.
# Prevents account bans on platforms with strict API rate limits (e.g. WeChat Work).
# 限制向平台发送消息的速率,防止因频率过高被封号(如企业微信)。
# =============================================================================
# [outgoing_rate_limit]
# max_per_second = 0 # Global default msgs/sec; 0 = unlimited (default) / 全局默认每秒消息数;0 = 无限制
# burst = 3 # Max burst size; default = ceil(max_per_second) / 最大突发数量
# Per-platform overrides / 每平台覆盖配置:
# [outgoing_rate_limit.platforms.wecom]
# max_per_second = 1
# [outgoing_rate_limit.platforms.telegram]
# max_per_second = 25
# [outgoing_rate_limit.platforms.feishu]
# max_per_second = 5
# =============================================================================
# Relay Settings / Relay 设置
# =============================================================================
+16 -1
View File
@@ -30,7 +30,8 @@ type Config struct {
TTS TTSConfig `toml:"tts"`
Display DisplayConfig `toml:"display"`
StreamPreview StreamPreviewConfig `toml:"stream_preview"` // real-time streaming preview
RateLimit RateLimitConfig `toml:"rate_limit"` // per-session rate limiting
RateLimit RateLimitConfig `toml:"rate_limit"` // per-session rate limiting
OutgoingRateLimit OutgoingRateLimitConfig `toml:"outgoing_rate_limit"` // outgoing message throttling
Relay RelayConfig `toml:"relay"` // bot-to-bot relay behavior
Quiet *bool `toml:"quiet,omitempty"` // global default for quiet mode; project-level overrides this
Cron CronConfig `toml:"cron"`
@@ -91,6 +92,20 @@ type RateLimitConfig struct {
WindowSecs *int `toml:"window_secs"` // window size in seconds; default 60
}
// OutgoingRateLimitConfig controls how fast messages are sent TO platforms.
// Prevents account bans on platforms with strict API rate limits (e.g. WeChat Work).
type OutgoingRateLimitConfig struct {
MaxPerSecond *float64 `toml:"max_per_second"` // messages per second; 0 = unlimited (default)
Burst *int `toml:"burst"` // max burst size; default = ceil(max_per_second)
Platforms map[string]OutgoingRateLimitPlatConfig `toml:"platforms"` // per-platform overrides keyed by platform type name
}
// OutgoingRateLimitPlatConfig is a per-platform override for outgoing rate limiting.
type OutgoingRateLimitPlatConfig struct {
MaxPerSecond *float64 `toml:"max_per_second"`
Burst *int `toml:"burst"`
}
// UsersConfig controls per-user role assignments and policies within a project.
type UsersConfig struct {
DefaultRole string `toml:"default_role,omitempty"` // role for unmatched users; default "member"
+37 -10
View File
@@ -176,6 +176,7 @@ type Engine struct {
userRolesMu sync.RWMutex // protects userRoles, disabledCmds, and adminFrom
rateLimiter *RateLimiter
outgoingRL *OutgoingRateLimiter
streamPreview StreamPreviewCfg
relayManager *RelayManager
eventIdleTimeout time.Duration
@@ -667,6 +668,11 @@ func (e *Engine) SetRateLimitCfg(cfg RateLimitCfg) {
e.rateLimiter = NewRateLimiter(cfg.MaxMessages, cfg.Window)
}
// SetOutgoingRateLimitCfg configures per-platform outgoing message throttling.
func (e *Engine) SetOutgoingRateLimitCfg(defaults OutgoingRateLimitCfg, overrides map[string]OutgoingRateLimitCfg) {
e.outgoingRL = NewOutgoingRateLimiter(defaults, overrides)
}
// checkRateLimit returns true if the message is allowed, false if rate-limited.
// It checks per-user role-based limits first, then falls back to the global limiter.
func (e *Engine) checkRateLimit(msg *Message) bool {
@@ -2449,10 +2455,7 @@ func (e *Engine) processInteractiveEvents(state *interactiveState, session *Sess
unsent := strings.Join(textParts[segmentStart:], "")
if unsent != "" {
for _, chunk := range splitMessage(unsent, maxPlatformMessageLen) {
if err := p.Send(e.ctx, replyCtx, chunk); err != nil {
slog.Error("failed to send reply", "error", err, "msg_id", msgID)
return
}
e.send(p, replyCtx, chunk)
}
}
}
@@ -2464,10 +2467,7 @@ func (e *Engine) processInteractiveEvents(state *interactiveState, session *Sess
} else {
slog.Debug("EventResult: sending via p.Send (preview inactive or failed)", "response_len", len(fullResponse), "chunks", len(splitMessage(fullResponse, maxPlatformMessageLen)))
for _, chunk := range splitMessage(fullResponse, maxPlatformMessageLen) {
if err := p.Send(e.ctx, replyCtx, chunk); err != nil {
slog.Error("failed to send reply", "error", err, "msg_id", msgID)
return
}
e.send(p, replyCtx, chunk)
}
}
@@ -5922,8 +5922,14 @@ func (e *Engine) sendAskQuestionPrompt(p Platform, replyCtx any, questions []Use
e.send(p, replyCtx, sb.String())
}
// send wraps p.Send with error logging and slow-operation warnings.
// send wraps p.Send with error logging, slow-operation warnings, and outgoing rate limiting.
func (e *Engine) send(p Platform, replyCtx any, content string) {
if e.outgoingRL != nil {
if err := e.outgoingRL.Wait(e.ctx, p.Name()); err != nil {
slog.Warn("outgoing rate limit: context cancelled", "platform", p.Name(), "error", err)
return
}
}
start := time.Now()
if err := p.Send(e.ctx, replyCtx, content); err != nil {
slog.Error("platform send failed", "platform", p.Name(), "error", err, "content_len", len(content))
@@ -5955,8 +5961,14 @@ func drainEvents(ch <-chan Event) {
}
}
// reply wraps p.Reply with error logging and slow-operation warnings.
// reply wraps p.Reply with error logging, slow-operation warnings, and outgoing rate limiting.
func (e *Engine) reply(p Platform, replyCtx any, content string) {
if e.outgoingRL != nil {
if err := e.outgoingRL.Wait(e.ctx, p.Name()); err != nil {
slog.Warn("outgoing rate limit: context cancelled", "platform", p.Name(), "error", err)
return
}
}
start := time.Now()
if err := p.Reply(e.ctx, replyCtx, content); err != nil {
slog.Error("platform reply failed", "platform", p.Name(), "error", err, "content_len", len(content))
@@ -5969,6 +5981,11 @@ func (e *Engine) reply(p Platform, replyCtx any, content string) {
// replyWithButtons sends a reply with inline buttons if the platform supports it,
// otherwise falls back to plain text reply.
func (e *Engine) replyWithButtons(p Platform, replyCtx any, content string, buttons [][]ButtonOption) {
if e.outgoingRL != nil {
if err := e.outgoingRL.Wait(e.ctx, p.Name()); err != nil {
return
}
}
if bs, ok := p.(InlineButtonSender); ok {
if err := bs.SendWithButtons(e.ctx, replyCtx, content, buttons); err == nil {
return
@@ -5989,6 +6006,11 @@ func (e *Engine) replyWithCard(p Platform, replyCtx any, card *Card) {
slog.Error("replyWithCard: nil card", "platform", p.Name())
return
}
if e.outgoingRL != nil {
if err := e.outgoingRL.Wait(e.ctx, p.Name()); err != nil {
return
}
}
if cs, ok := p.(CardSender); ok {
if err := cs.ReplyCard(e.ctx, replyCtx, card); err != nil {
slog.Error("card reply failed", "platform", p.Name(), "error", err)
@@ -6004,6 +6026,11 @@ func (e *Engine) sendWithCard(p Platform, replyCtx any, card *Card) {
slog.Error("sendWithCard: nil card", "platform", p.Name())
return
}
if e.outgoingRL != nil {
if err := e.outgoingRL.Wait(e.ctx, p.Name()); err != nil {
return
}
}
if cs, ok := p.(CardSender); ok {
if err := cs.SendCard(e.ctx, replyCtx, card); err != nil {
slog.Error("card send failed", "platform", p.Name(), "error", err)
+132
View File
@@ -0,0 +1,132 @@
package core
import (
"context"
"math"
"sync"
"time"
)
// OutgoingRateLimitCfg holds the resolved rate-limit parameters for outgoing messages.
type OutgoingRateLimitCfg struct {
MaxPerSecond float64 // messages per second; 0 = disabled / unlimited
Burst int // max burst size; 0 = use ceil(MaxPerSecond)
}
func (c OutgoingRateLimitCfg) disabled() bool {
return c.MaxPerSecond <= 0
}
func (c OutgoingRateLimitCfg) effectiveBurst() int {
if c.Burst > 0 {
return c.Burst
}
return int(math.Ceil(c.MaxPerSecond))
}
// OutgoingRateLimiter throttles outgoing messages sent to platforms using a
// per-platform token bucket. It never drops messages; callers block until
// the rate budget allows the send.
type OutgoingRateLimiter struct {
mu sync.Mutex
buckets map[string]*tokenBucket // key = platform name
defaults OutgoingRateLimitCfg
overrides map[string]OutgoingRateLimitCfg // per-platform overrides
}
type tokenBucket struct {
tokens float64
maxTokens float64
refillRate float64 // tokens per second
lastRefill time.Time
}
// NewOutgoingRateLimiter creates a rate limiter with global defaults and
// optional per-platform overrides. If defaults.MaxPerSecond <= 0 and no
// overrides are set, the limiter is effectively disabled.
func NewOutgoingRateLimiter(defaults OutgoingRateLimitCfg, overrides map[string]OutgoingRateLimitCfg) *OutgoingRateLimiter {
if overrides == nil {
overrides = make(map[string]OutgoingRateLimitCfg)
}
return &OutgoingRateLimiter{
buckets: make(map[string]*tokenBucket),
defaults: defaults,
overrides: overrides,
}
}
// cfgFor returns the effective config for a platform.
func (orl *OutgoingRateLimiter) cfgFor(platform string) OutgoingRateLimitCfg {
if c, ok := orl.overrides[platform]; ok {
return c
}
return orl.defaults
}
// bucketFor returns (or lazily creates) the token bucket for a platform.
// Must be called with orl.mu held.
func (orl *OutgoingRateLimiter) bucketFor(platform string) *tokenBucket {
if b, ok := orl.buckets[platform]; ok {
return b
}
cfg := orl.cfgFor(platform)
burst := cfg.effectiveBurst()
b := &tokenBucket{
tokens: float64(burst), // start full
maxTokens: float64(burst),
refillRate: cfg.MaxPerSecond,
lastRefill: time.Now(),
}
orl.buckets[platform] = b
return b
}
// refill adds tokens based on elapsed time since last refill.
func (b *tokenBucket) refill() {
now := time.Now()
elapsed := now.Sub(b.lastRefill).Seconds()
if elapsed <= 0 {
return
}
b.tokens += elapsed * b.refillRate
if b.tokens > b.maxTokens {
b.tokens = b.maxTokens
}
b.lastRefill = now
}
// Wait blocks until the rate limiter allows one message to be sent to the
// given platform. Returns nil on success, or the context error if ctx is
// cancelled while waiting.
func (orl *OutgoingRateLimiter) Wait(ctx context.Context, platform string) error {
cfg := orl.cfgFor(platform)
if cfg.disabled() {
return nil
}
for {
orl.mu.Lock()
b := orl.bucketFor(platform)
b.refill()
if b.tokens >= 1 {
b.tokens--
orl.mu.Unlock()
return nil
}
// Calculate wait time until 1 token is available.
deficit := 1.0 - b.tokens
waitSecs := deficit / b.refillRate
orl.mu.Unlock()
timer := time.NewTimer(time.Duration(waitSecs * float64(time.Second)))
select {
case <-ctx.Done():
timer.Stop()
return ctx.Err()
case <-timer.C:
// Loop back to try consuming a token.
}
}
}
+143
View File
@@ -0,0 +1,143 @@
package core
import (
"context"
"sync"
"testing"
"time"
)
func TestOutgoingRateLimiter_Disabled(t *testing.T) {
orl := NewOutgoingRateLimiter(OutgoingRateLimitCfg{MaxPerSecond: 0}, nil)
// Should return immediately when disabled.
start := time.Now()
for i := 0; i < 100; i++ {
if err := orl.Wait(context.Background(), "test"); err != nil {
t.Fatalf("unexpected error: %v", err)
}
}
if elapsed := time.Since(start); elapsed > 50*time.Millisecond {
t.Fatalf("disabled limiter should be instant, took %v", elapsed)
}
}
func TestOutgoingRateLimiter_BurstThenThrottle(t *testing.T) {
// 2 msgs/sec with burst=2: first 2 should be instant, 3rd should wait ~500ms.
orl := NewOutgoingRateLimiter(OutgoingRateLimitCfg{MaxPerSecond: 2, Burst: 2}, nil)
ctx := context.Background()
start := time.Now()
// Consume the burst
for i := 0; i < 2; i++ {
if err := orl.Wait(ctx, "p"); err != nil {
t.Fatalf("burst msg %d: %v", i, err)
}
}
burstElapsed := time.Since(start)
if burstElapsed > 50*time.Millisecond {
t.Fatalf("burst should be instant, took %v", burstElapsed)
}
// Third message should throttle
before := time.Now()
if err := orl.Wait(ctx, "p"); err != nil {
t.Fatalf("throttled msg: %v", err)
}
waited := time.Since(before)
if waited < 300*time.Millisecond || waited > 800*time.Millisecond {
t.Fatalf("expected ~500ms wait, got %v", waited)
}
}
func TestOutgoingRateLimiter_PerPlatformOverride(t *testing.T) {
defaults := OutgoingRateLimitCfg{MaxPerSecond: 100} // fast default
overrides := map[string]OutgoingRateLimitCfg{
"slow": {MaxPerSecond: 2, Burst: 1},
}
orl := NewOutgoingRateLimiter(defaults, overrides)
ctx := context.Background()
// "fast" platform should be instant (burst=100)
start := time.Now()
for i := 0; i < 10; i++ {
_ = orl.Wait(ctx, "fast")
}
if time.Since(start) > 50*time.Millisecond {
t.Fatal("fast platform should not throttle within burst")
}
// "slow" platform: burst=1, second call should wait ~500ms
_ = orl.Wait(ctx, "slow")
before := time.Now()
_ = orl.Wait(ctx, "slow")
waited := time.Since(before)
if waited < 300*time.Millisecond {
t.Fatalf("slow platform should throttle, waited only %v", waited)
}
}
func TestOutgoingRateLimiter_ContextCancellation(t *testing.T) {
orl := NewOutgoingRateLimiter(OutgoingRateLimitCfg{MaxPerSecond: 1, Burst: 1}, nil)
ctx, cancel := context.WithCancel(context.Background())
// Consume the single burst token
_ = orl.Wait(ctx, "p")
// Cancel context immediately
cancel()
err := orl.Wait(ctx, "p")
if err == nil {
t.Fatal("expected context error, got nil")
}
}
func TestOutgoingRateLimiter_ConcurrentAccess(t *testing.T) {
orl := NewOutgoingRateLimiter(OutgoingRateLimitCfg{MaxPerSecond: 1000, Burst: 100}, nil)
ctx := context.Background()
var wg sync.WaitGroup
for i := 0; i < 10; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for j := 0; j < 10; j++ {
_ = orl.Wait(ctx, "p")
}
}()
}
wg.Wait()
}
func TestOutgoingRateLimiter_DefaultBurst(t *testing.T) {
// When Burst is 0, effectiveBurst = ceil(MaxPerSecond)
cfg := OutgoingRateLimitCfg{MaxPerSecond: 1.5, Burst: 0}
if cfg.effectiveBurst() != 2 {
t.Fatalf("expected burst=2 for rate 1.5, got %d", cfg.effectiveBurst())
}
cfg2 := OutgoingRateLimitCfg{MaxPerSecond: 5, Burst: 0}
if cfg2.effectiveBurst() != 5 {
t.Fatalf("expected burst=5 for rate 5, got %d", cfg2.effectiveBurst())
}
}
func TestOutgoingRateLimiter_DisabledPlatformOverride(t *testing.T) {
// Global rate is set but a specific platform is disabled via MaxPerSecond=0
defaults := OutgoingRateLimitCfg{MaxPerSecond: 1, Burst: 1}
overrides := map[string]OutgoingRateLimitCfg{
"unlimited": {MaxPerSecond: 0},
}
orl := NewOutgoingRateLimiter(defaults, overrides)
ctx := context.Background()
// "unlimited" platform should be instant even after consuming tokens
start := time.Now()
for i := 0; i < 50; i++ {
if err := orl.Wait(ctx, "unlimited"); err != nil {
t.Fatalf("unexpected error: %v", err)
}
}
if time.Since(start) > 50*time.Millisecond {
t.Fatal("unlimited override platform should not throttle")
}
}