opencode/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(qq): add WebSocket reconnection, API retry, and rate limiting

Browse source
This commit is contained in:
kent.liu 2026-03-19 17:04:55 +08:00
parent 828971d549
commit 760f13b1c6
3 changed files with 488 additions and 48 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

291
pkg/channels/qq/qq.go
View file

@ -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) {
@ -94,6 +113,8 @@ func NewQQChannel(cfg config.QQConfig, messageBus *bus.MessageBus) (*QQChannel,
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,7 +249,24 @@ func (c *QQChannel) Send(ctx context.Context, msg bus.OutboundMessage) error {
}
}
// Route to group or C2C.
// 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(),
})
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)
@ -243,12 +274,29 @@ func (c *QQChannel) Send(ctx context.Context, msg bus.OutboundMessage) error {
_, err = c.api.PostC2CMessage(ctx, msg.ChatID, msgToCreate)
}
if err != nil {
logger.ErrorCF("qq", "Failed to send message", map[string]any{
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<<uint(attempt))
if backoff > retryMaxDelay {
backoff = retryMaxDelay
}
return backoff
}

86
pkg/channels/qq/qq_ratelimit.go Normal file
View file

@ -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)
}

139
pkg/channels/qq/qq_ratelimit_test.go Normal file
View file

@ -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")
}
}