From d7bed298019891b8afb8c2e6564a1a529273ca40 Mon Sep 17 00:00:00 2001 From: Max Qian <64824374+AstroAir@users.noreply.github.com> Date: Fri, 27 Mar 2026 23:59:11 +0800 Subject: [PATCH] feat(core): add outgoing message rate limiting per platform --- cmd/cc-connect/main.go | 28 +++++++ config.example.toml | 18 ++++ config/config.go | 17 +++- core/engine.go | 47 ++++++++--- core/outgoing_ratelimit.go | 132 +++++++++++++++++++++++++++++ core/outgoing_ratelimit_test.go | 143 ++++++++++++++++++++++++++++++++ 6 files changed, 374 insertions(+), 11 deletions(-) create mode 100644 core/outgoing_ratelimit.go create mode 100644 core/outgoing_ratelimit_test.go diff --git a/cmd/cc-connect/main.go b/cmd/cc-connect/main.go index 2b6e79a22..4de694b08 100644 --- a/cmd/cc-connect/main.go +++ b/cmd/cc-connect/main.go @@ -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) }) diff --git a/config.example.toml b/config.example.toml index e01fe6d1b..ceb62659d 100644 --- a/config.example.toml +++ b/config.example.toml @@ -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 设置 # ============================================================================= diff --git a/config/config.go b/config/config.go index af4deeb7f..c91f69a48 100644 --- a/config/config.go +++ b/config/config.go @@ -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" diff --git a/core/engine.go b/core/engine.go index 19466dfd2..9c6d030df 100644 --- a/core/engine.go +++ b/core/engine.go @@ -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) diff --git a/core/outgoing_ratelimit.go b/core/outgoing_ratelimit.go new file mode 100644 index 000000000..209d93d5b --- /dev/null +++ b/core/outgoing_ratelimit.go @@ -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. + } + } +} diff --git a/core/outgoing_ratelimit_test.go b/core/outgoing_ratelimit_test.go new file mode 100644 index 000000000..c4f174774 --- /dev/null +++ b/core/outgoing_ratelimit_test.go @@ -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") + } +}