From 760f13b1c62122e9b5bfd9720bf060e69d1efd5a Mon Sep 17 00:00:00 2001 From: "kent.liu" Date: Thu, 19 Mar 2026 17:04:55 +0800 Subject: [PATCH] feat(qq): add WebSocket reconnection, API retry, and rate limiting --- pkg/channels/qq/qq.go | 311 ++++++++++++++++++++++----- pkg/channels/qq/qq_ratelimit.go | 86 ++++++++ pkg/channels/qq/qq_ratelimit_test.go | 139 ++++++++++++ 3 files changed, 488 insertions(+), 48 deletions(-) create mode 100644 pkg/channels/qq/qq_ratelimit.go create mode 100644 pkg/channels/qq/qq_ratelimit_test.go diff --git a/pkg/channels/qq/qq.go b/pkg/channels/qq/qq.go index 1a48369f8..a47128a63 100644 --- a/pkg/channels/qq/qq.go +++ b/pkg/channels/qq/qq.go @@ -41,6 +41,16 @@ const ( typingResend = 8 * time.Second typingSeconds = 10 bytesPerMiB = 1024 * 1024 + + // Reconnection constants + reconnectInitial = 5 * time.Second + reconnectMax = 5 * time.Minute + reconnectMultiplier = 2.0 + + // Retry constants + maxRetries = 3 + retryInitialDelay = 500 * time.Millisecond + retryMaxDelay = 10 * time.Second ) type qqAPI interface { @@ -80,6 +90,15 @@ type QQChannel struct { // done is closed on Stop to shut down the dedup janitor. done chan struct{} stopOnce sync.Once + + // Reconnection state + reconnectMu sync.Mutex + reconnecting bool + stopReconnect chan struct{} + + // Rate limiting + groupRateLimiter *rateLimiter + directRateLimiter *rateLimiter } func NewQQChannel(cfg config.QQConfig, messageBus *bus.MessageBus) (*QQChannel, error) { @@ -90,10 +109,12 @@ func NewQQChannel(cfg config.QQConfig, messageBus *bus.MessageBus) (*QQChannel, ) return &QQChannel{ - BaseChannel: base, - config: cfg, - dedup: make(map[string]time.Time), - done: make(chan struct{}), + BaseChannel: base, + config: cfg, + dedup: make(map[string]time.Time), + done: make(chan struct{}), + groupRateLimiter: newRateLimiter(500 * time.Millisecond), // 20 msg/min with headroom + directRateLimiter: newRateLimiter(200 * time.Millisecond), // 5 msg/sec with headroom }, nil } @@ -105,9 +126,13 @@ func (c *QQChannel) Start(ctx context.Context) error { botgo.SetLogger(newBotGoLogger("botgo")) logger.InfoC("qq", "Starting QQ bot (WebSocket mode)") - // Reinitialize shutdown signal for clean restart. + // Reinitialize shutdown signals for clean restart. c.done = make(chan struct{}) c.stopOnce = sync.Once{} + c.stopReconnect = make(chan struct{}) + c.reconnectMu.Lock() + c.reconnecting = false + c.reconnectMu.Unlock() // create token source credentials := &token.QQBotCredentials{ @@ -127,44 +152,12 @@ func (c *QQChannel) Start(ctx context.Context) error { // initialize OpenAPI client c.api = botgo.NewOpenAPI(c.config.AppID, c.tokenSource).WithTimeout(5 * time.Second) - // register event handlers - intent := event.RegisterHandlers( - c.handleC2CMessage(), - c.handleGroupATMessage(), - ) - - // get WebSocket endpoint - wsInfo, err := c.api.WS(c.ctx, nil, "") - if err != nil { - return fmt.Errorf("failed to get websocket info: %w", err) - } - - logger.InfoCF("qq", "Got WebSocket info", map[string]any{ - "shards": wsInfo.Shards, - }) - - // create and save sessionManager - c.sessionManager = botgo.NewSessionManager() - - // start WebSocket connection in goroutine to avoid blocking - go func() { - if err := c.sessionManager.Start(wsInfo, c.tokenSource, &intent); err != nil { - logger.ErrorCF("qq", "WebSocket session error", map[string]any{ - "error": err.Error(), - }) - c.SetRunning(false) - } - }() + // start session with reconnection support + go c.startSession() // start dedup janitor goroutine go c.dedupJanitor() - // Pre-register reasoning_channel_id as group chat if configured, - // so outbound-only destinations are routed correctly. - if c.config.ReasoningChannelID != "" { - c.chatType.Store(c.config.ReasoningChannelID, "group") - } - c.SetRunning(true) logger.InfoC("qq", "QQ bot started successfully") @@ -175,6 +168,16 @@ func (c *QQChannel) Stop(ctx context.Context) error { logger.InfoC("qq", "Stopping QQ bot") c.SetRunning(false) + // Signal reconnection loop to stop + c.reconnectMu.Lock() + c.reconnecting = false + c.reconnectMu.Unlock() + + // Close stop channel to terminate reconnect goroutine + if c.stopReconnect != nil { + close(c.stopReconnect) + } + // Signal the dedup janitor to stop (idempotent). c.stopOnce.Do(func() { close(c.done) }) @@ -207,6 +210,17 @@ func (c *QQChannel) Send(ctx context.Context, msg bus.OutboundMessage) error { chatKind := c.getChatKind(msg.ChatID) + // Apply rate limiting before sending + if chatKind == "group" { + if err := c.groupRateLimiter.waitWithContext(msg.ChatID, ctx); err != nil { + return fmt.Errorf("qq send: %w", err) + } + } else { + if err := c.directRateLimiter.waitWithContext(msg.ChatID, ctx); err != nil { + return fmt.Errorf("qq send: %w", err) + } + } + // Build message with content. msgToCreate := &dto.MessageToCreate{ Content: msg.Content, @@ -235,20 +249,54 @@ func (c *QQChannel) Send(ctx context.Context, msg bus.OutboundMessage) error { } } - // Route to group or C2C. - var err error - if chatKind == "group" { - _, err = c.api.PostGroupMessage(ctx, msg.ChatID, msgToCreate) - } else { - _, err = c.api.PostC2CMessage(ctx, msg.ChatID, msgToCreate) - } + // Route to group or C2C with retry. + var lastErr error + for attempt := 0; attempt < maxRetries; attempt++ { + if attempt > 0 { + backoff := calculateBackoff(attempt - 1) + logger.InfoCF("qq", "Retrying send", map[string]any{ + "attempt": attempt, + "chat_id": msg.ChatID, + "backoff": backoff.String(), + }) - if err != nil { - logger.ErrorCF("qq", "Failed to send message", map[string]any{ + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(backoff): + } + } + + var err error + if chatKind == "group" { + _, err = c.api.PostGroupMessage(ctx, msg.ChatID, msgToCreate) + } else { + _, err = c.api.PostC2CMessage(ctx, msg.ChatID, msgToCreate) + } + + if err == nil { + return nil + } + + lastErr = err + logger.WarnCF("qq", "Send attempt failed", map[string]any{ + "attempt": attempt + 1, "chat_id": msg.ChatID, "chat_kind": chatKind, "error": err.Error(), }) + + if !isRetryableError(err) { + break + } + } + + if lastErr != nil { + logger.ErrorCF("qq", "Failed to send message after retries", map[string]any{ + "chat_id": msg.ChatID, + "chat_kind": chatKind, + "error": lastErr.Error(), + }) return fmt.Errorf("qq send: %w", channels.ErrTemporary) } @@ -915,3 +963,170 @@ func sanitizeURLs(text string) string { return scheme + domain + path }) } + +// startSession starts the WebSocket session with reconnection support. +func (c *QQChannel) startSession() { + go func() { + for { + select { + case <-c.ctx.Done(): + return + case <-c.stopReconnect: + return + default: + c.runSession() + // Session ended unexpectedly, attempt reconnect + if c.IsRunning() { + c.reconnect() + } + } + } + }() +} + +// runSession runs a single WebSocket session. +func (c *QQChannel) runSession() { + // Get WebSocket endpoint + wsInfo, err := c.api.WS(c.ctx, nil, "") + if err != nil { + logger.ErrorCF("qq", "Failed to get websocket info", map[string]any{ + "error": err.Error(), + }) + return + } + + logger.InfoCF("qq", "Got WebSocket info", map[string]any{ + "shards": wsInfo.Shards, + }) + + // Create and start session + c.sessionManager = botgo.NewSessionManager() + intent := event.RegisterHandlers( + c.handleC2CMessage(), + c.handleGroupATMessage(), + ) + + if err := c.sessionManager.Start(wsInfo, c.tokenSource, &intent); err != nil { + logger.ErrorCF("qq", "WebSocket session error", map[string]any{ + "error": err.Error(), + }) + c.SetRunning(false) + } +} + +// reconnect attempts to reconnect with exponential backoff. +func (c *QQChannel) reconnect() { + c.reconnectMu.Lock() + if c.reconnecting { + c.reconnectMu.Unlock() + return + } + c.reconnecting = true + c.reconnectMu.Unlock() + + defer func() { + c.reconnectMu.Lock() + c.reconnecting = false + c.reconnectMu.Unlock() + }() + + backoff := reconnectInitial + for { + select { + case <-c.ctx.Done(): + return + case <-c.stopReconnect: + return + default: + } + + logger.InfoCF("qq", "Attempting to reconnect", map[string]any{ + "backoff": backoff.String(), + }) + + // Reset internal state for new session + c.resetSessionState() + + c.runSession() + + if c.IsRunning() { + logger.InfoC("qq", "Reconnected successfully") + return + } + + logger.WarnCF("qq", "Reconnect failed, retrying", map[string]any{ + "backoff": backoff.String(), + }) + + select { + case <-c.ctx.Done(): + return + case <-c.stopReconnect: + return + case <-time.After(backoff): + if backoff < reconnectMax { + backoff = time.Duration(float64(backoff) * reconnectMultiplier) + if backoff > reconnectMax { + backoff = reconnectMax + } + } + } + } +} + +// resetSessionState clears transient state that may be invalid after reconnection. +func (c *QQChannel) resetSessionState() { + // Clear transient state that may be invalid after reconnection + c.chatType.Range(func(key, value interface{}) bool { + c.chatType.Delete(key) + return true + }) + c.lastMsgID.Range(func(key, value interface{}) bool { + c.lastMsgID.Delete(key) + return true + }) + c.msgSeqCounters.Range(func(key, value interface{}) bool { + c.msgSeqCounters.Delete(key) + return true + }) + + // Clear rate limiters to allow immediate sending after reconnection + if c.groupRateLimiter != nil { + c.groupRateLimiter.clearAll() + } + if c.directRateLimiter != nil { + c.directRateLimiter.clearAll() + } + + // Re-register reasoning channel if configured + if c.config.ReasoningChannelID != "" { + c.chatType.Store(c.config.ReasoningChannelID, "group") + } +} + +// isRetryableError returns true if the error should be retried. +func isRetryableError(err error) bool { + if err == nil { + return false + } + errStr := err.Error() + // Common transient error patterns + return strings.Contains(errStr, "timeout") || + strings.Contains(errStr, "context deadline") || + strings.Contains(errStr, "connection reset") || + strings.Contains(errStr, "connection refused") || + strings.Contains(errStr, "temporary failure") || + strings.Contains(errStr, "429") || + strings.Contains(errStr, "500") || + strings.Contains(errStr, "502") || + strings.Contains(errStr, "503") +} + +// calculateBackoff returns the next backoff duration with exponential increase. +func calculateBackoff(attempt int) time.Duration { + backoff := retryInitialDelay * time.Duration(1< retryMaxDelay { + backoff = retryMaxDelay + } + return backoff +} diff --git a/pkg/channels/qq/qq_ratelimit.go b/pkg/channels/qq/qq_ratelimit.go new file mode 100644 index 000000000..4a83432d1 --- /dev/null +++ b/pkg/channels/qq/qq_ratelimit.go @@ -0,0 +1,86 @@ +package qq + +import ( + "context" + "sync" + "time" +) + +// rateLimiter implements a per-chat token bucket rate limiter. +type rateLimiter struct { + mu sync.Mutex + lastSend map[string]time.Time + interval time.Duration +} + +// newRateLimiter creates a rate limiter with the specified minimum interval between sends. +func newRateLimiter(interval time.Duration) *rateLimiter { + return &rateLimiter{ + lastSend: make(map[string]time.Time), + interval: interval, + } +} + +// wait waits if necessary until enough time has passed since the last send to this chat. +func (r *rateLimiter) wait(chatID string) { + r.mu.Lock() + defer r.mu.Unlock() + + now := time.Now() + if lastSend, ok := r.lastSend[chatID]; ok { + elapsed := now.Sub(lastSend) + if elapsed < r.interval { + time.Sleep(r.interval - elapsed) + now = time.Now() + } + } + + r.lastSend[chatID] = now +} + +// waitWithContext waits with context cancellation support. +func (r *rateLimiter) waitWithContext(chatID string, ctx context.Context) error { + r.mu.Lock() + now := time.Now() + var waitDuration time.Duration + + if lastSend, ok := r.lastSend[chatID]; ok { + elapsed := now.Sub(lastSend) + if elapsed < r.interval { + waitDuration = r.interval - elapsed + } + } + r.mu.Unlock() + + if waitDuration > 0 { + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(waitDuration): + } + + r.mu.Lock() + r.lastSend[chatID] = time.Now() + r.mu.Unlock() + } else { + r.mu.Lock() + r.lastSend[chatID] = now + r.mu.Unlock() + } + + return nil +} + +// clear removes the rate limit entry for a chat (e.g., after reconnection). +func (r *rateLimiter) clear(chatID string) { + r.mu.Lock() + defer r.mu.Unlock() + delete(r.lastSend, chatID) +} + +// clearAll removes all rate limit entries. +func (r *rateLimiter) clearAll() { + r.mu.Lock() + defer r.mu.Unlock() + r.lastSend = make(map[string]time.Time) +} diff --git a/pkg/channels/qq/qq_ratelimit_test.go b/pkg/channels/qq/qq_ratelimit_test.go new file mode 100644 index 000000000..c91d97af2 --- /dev/null +++ b/pkg/channels/qq/qq_ratelimit_test.go @@ -0,0 +1,139 @@ +package qq + +import ( + "context" + "sync" + "testing" + "time" +) + +func TestRateLimiter_Wait(t *testing.T) { + limiter := newRateLimiter(50 * time.Millisecond) + + // First call should not wait + start := time.Now() + limiter.wait("chat-1") + elapsed := time.Since(start) + if elapsed > 10*time.Millisecond { + t.Errorf("first wait should be immediate, got %v", elapsed) + } + + // Second call to same chat should wait + start = time.Now() + limiter.wait("chat-1") + elapsed = time.Since(start) + if elapsed < 40*time.Millisecond { + t.Errorf("second wait should be at least 40ms, got %v", elapsed) + } + + // Different chat should not wait + start = time.Now() + limiter.wait("chat-2") + elapsed = time.Since(start) + if elapsed > 10*time.Millisecond { + t.Errorf("different chat should not wait, got %v", elapsed) + } +} + +func TestRateLimiter_Clear(t *testing.T) { + limiter := newRateLimiter(100 * time.Millisecond) + + limiter.wait("chat-1") + + limiter.clear("chat-1") + + start := time.Now() + limiter.wait("chat-1") + elapsed := time.Since(start) + if elapsed > 10*time.Millisecond { + t.Errorf("after clear, wait should be immediate, got %v", elapsed) + } +} + +func TestRateLimiter_ClearAll(t *testing.T) { + limiter := newRateLimiter(100 * time.Millisecond) + + limiter.wait("chat-1") + limiter.wait("chat-2") + + limiter.clearAll() + + start := time.Now() + limiter.wait("chat-1") + elapsed := time.Since(start) + if elapsed > 10*time.Millisecond { + t.Errorf("after clearAll, wait should be immediate, got %v", elapsed) + } +} + +func TestRateLimiter_WaitWithContext_Canceled(t *testing.T) { + limiter := newRateLimiter(100 * time.Millisecond) + + // First call to record a send time + limiter.wait("chat-1") + + ctx, cancel := context.WithCancel(context.Background()) + + // Cancel immediately + cancel() + + start := time.Now() + err := limiter.waitWithContext("chat-1", ctx) + elapsed := time.Since(start) + + if err == nil { + t.Error("expected context cancellation error") + } + if elapsed >= 50*time.Millisecond { + t.Errorf("should have returned early due to cancellation, got %v", elapsed) + } +} + +func TestRateLimiter_WaitWithContext_Success(t *testing.T) { + limiter := newRateLimiter(50 * time.Millisecond) + + ctx := context.Background() + + start := time.Now() + err := limiter.waitWithContext("chat-1", ctx) + elapsed := time.Since(start) + + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if elapsed > 10*time.Millisecond { + t.Errorf("first wait should be immediate, got %v", elapsed) + } +} + +func TestRateLimiter_Concurrent(t *testing.T) { + limiter := newRateLimiter(10 * time.Millisecond) + + var wg sync.WaitGroup + chatIDs := []string{"chat-1", "chat-2", "chat-3"} + + // Run 10 iterations for each chat concurrently + for i := 0; i < 10; i++ { + for _, chatID := range chatIDs { + wg.Add(1) + go func(id string) { + defer wg.Done() + limiter.wait(id) + }(chatID) + } + } + + // All goroutines should complete within reasonable time + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + // Success + case <-time.After(5 * time.Second): + t.Error("concurrent rate limiter test timed out") + } +}