feat(wecom): add WebSocket long-connection support for WeCom AI Bot
- Introduced WeComAIBotWSChannel to handle WebSocket connections. - Updated NewWeComAIBotChannel to prioritize WebSocket mode when BotID and Secret are provided. - Enhanced WeComAIBotConfig to include BotID and Secret for WebSocket mode. - Implemented message handling for text, image, voice, and mixed messages in WebSocket mode. - Added tests for WebSocket mode functionality and ensured backward compatibility with webhook mode. - Refactored existing code to improve clarity and maintainability.
This commit is contained in:
Zhang Rui
parent
7673b626b3
commit
6f5d1f336f
6 changed files with 1025 additions and 49 deletions
|
|
@ -204,6 +204,8 @@
|
|||
"wecom_aibot": {
|
||||
"_comment": "WeCom AI Bot (智能机器人) - Official WeCom AI Bot integration, supports proactive messaging and private chats.",
|
||||
"enabled": false,
|
||||
"bot_id": "YOUR_BOT_ID",
|
||||
"secret": "YOUR_SECRET",
|
||||
"token": "YOUR_TOKEN",
|
||||
"encoding_aes_key": "YOUR_43_CHAR_ENCODING_AES_KEY",
|
||||
"webhook_path": "/webhook/wecom-aibot",
|
||||
|
|
|
|||
|
|
@ -296,7 +296,9 @@ func (m *Manager) initChannels(channels *config.ChannelsConfig) error {
|
|||
m.initChannel("wecom", "WeCom")
|
||||
}
|
||||
|
||||
if channels.WeComAIBot.Enabled && channels.WeComAIBot.Token != "" {
|
||||
if m.config.Channels.WeComAIBot.Enabled &&
|
||||
((m.config.Channels.WeComAIBot.BotID != "" && m.config.Channels.WeComAIBot.Secret != "") ||
|
||||
m.config.Channels.WeComAIBot.Token != "") {
|
||||
m.initChannel("wecom_aibot", "WeCom AI Bot")
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -134,13 +134,25 @@ type WeComAIBotEncryptedResponse struct {
|
|||
Nonce string `json:"nonce"`
|
||||
}
|
||||
|
||||
// NewWeComAIBotChannel creates a new WeCom AI Bot channel instance
|
||||
// NewWeComAIBotChannel creates a WeCom AI Bot channel instance.
|
||||
// If cfg.BotID and cfg.Secret are both set, it returns a WeComAIBotWSChannel
|
||||
// using the WebSocket long-connection API.
|
||||
// Otherwise it returns the webhook-mode WeComAIBotChannel (requires Token +
|
||||
// EncodingAESKey).
|
||||
func NewWeComAIBotChannel(
|
||||
cfg config.WeComAIBotConfig,
|
||||
messageBus *bus.MessageBus,
|
||||
) (*WeComAIBotChannel, error) {
|
||||
) (channels.Channel, error) {
|
||||
// WebSocket long-connection mode takes priority when BotID + Secret are set.
|
||||
if cfg.BotID != "" && cfg.Secret != "" {
|
||||
logger.InfoC("wecom_aibot", "BotID and Secret provided, using WebSocket mode")
|
||||
return newWeComAIBotWSChannel(cfg, messageBus)
|
||||
}
|
||||
// Webhook (short-connection) mode.
|
||||
if cfg.Token == "" || cfg.EncodingAESKey == "" {
|
||||
return nil, fmt.Errorf("token and encoding_aes_key are required for WeCom AI Bot")
|
||||
return nil, fmt.Errorf(
|
||||
"WeCom AI Bot requires either (bot_id + secret) for WebSocket mode " +
|
||||
"or (token + encoding_aes_key) for webhook mode")
|
||||
}
|
||||
|
||||
base := channels.NewBaseChannel("wecom_aibot", cfg, messageBus, cfg.AllowFrom,
|
||||
|
|
@ -895,17 +907,80 @@ func (c *WeComAIBotChannel) encryptMessage(plaintext, receiveid string) (string,
|
|||
return base64.StdEncoding.EncodeToString(ciphertext), nil
|
||||
}
|
||||
|
||||
// generateStreamID generates a random stream ID
|
||||
func (c *WeComAIBotChannel) generateStreamID() string {
|
||||
// func (c *WeComAIBotChannel) downloadAndDecryptImage(
|
||||
// ctx context.Context,
|
||||
// imageURL string,
|
||||
// ) ([]byte, error) {
|
||||
// // Download image
|
||||
// req, err := http.NewRequestWithContext(ctx, http.MethodGet, imageURL, nil)
|
||||
// if err != nil {
|
||||
// return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
// }
|
||||
|
||||
// client := &http.Client{
|
||||
// Timeout: 15 * time.Second,
|
||||
// }
|
||||
|
||||
// resp, err := client.Do(req)
|
||||
// if err != nil {
|
||||
// return nil, fmt.Errorf("failed to download image: %w", err)
|
||||
// }
|
||||
// defer resp.Body.Close()
|
||||
|
||||
// if resp.StatusCode != http.StatusOK {
|
||||
// return nil, fmt.Errorf("download failed with status: %d", resp.StatusCode)
|
||||
// }
|
||||
|
||||
// // Limit image download to 20 MB to prevent memory exhaustion
|
||||
// const maxImageSize = 20 << 20 // 20 MB
|
||||
// encryptedData, err := io.ReadAll(io.LimitReader(resp.Body, maxImageSize+1))
|
||||
// if err != nil {
|
||||
// return nil, fmt.Errorf("failed to read image data: %w", err)
|
||||
// }
|
||||
// if len(encryptedData) > maxImageSize {
|
||||
// return nil, fmt.Errorf("image too large (exceeds %d MB)", maxImageSize>>20)
|
||||
// }
|
||||
|
||||
// logger.DebugCF("wecom_aibot", "Image downloaded", map[string]any{
|
||||
// "size": len(encryptedData),
|
||||
// })
|
||||
|
||||
// // Decode AES key
|
||||
// aesKey, err := decodeWeComAESKey(c.config.EncodingAESKey)
|
||||
// if err != nil {
|
||||
// return nil, err
|
||||
// }
|
||||
|
||||
// // Decrypt image (AES-CBC with IV = first 16 bytes of key, PKCS7 padding stripped)
|
||||
// decryptedData, err := decryptAESCBC(aesKey, encryptedData)
|
||||
// if err != nil {
|
||||
// return nil, fmt.Errorf("failed to decrypt image: %w", err)
|
||||
// }
|
||||
|
||||
// logger.DebugCF("wecom_aibot", "Image decrypted", map[string]any{
|
||||
// "size": len(decryptedData),
|
||||
// })
|
||||
|
||||
// return decryptedData, nil
|
||||
// }
|
||||
|
||||
// generateRandomID generates a cryptographically random alphanumeric ID of
|
||||
// length n. Used for stream IDs and WebSocket request IDs.
|
||||
func generateRandomID(n int) string {
|
||||
const letters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
|
||||
b := make([]byte, 10)
|
||||
b := make([]byte, n)
|
||||
for i := range b {
|
||||
n, _ := rand.Int(rand.Reader, big.NewInt(int64(len(letters))))
|
||||
b[i] = letters[n.Int64()]
|
||||
num, _ := rand.Int(rand.Reader, big.NewInt(int64(len(letters))))
|
||||
b[i] = letters[num.Int64()]
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// generateStreamID generates a random 10-character stream ID (webhook mode).
|
||||
func (c *WeComAIBotChannel) generateStreamID() string {
|
||||
return generateRandomID(10)
|
||||
}
|
||||
|
||||
// cleanupLoop periodically cleans up old streaming tasks
|
||||
func (c *WeComAIBotChannel) cleanupLoop() {
|
||||
ticker := time.NewTicker(5 * time.Minute)
|
||||
|
|
|
|||
|
|
@ -5,10 +5,13 @@ import (
|
|||
"testing"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/bus"
|
||||
"github.com/sipeed/picoclaw/pkg/channels"
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
)
|
||||
|
||||
func TestNewWeComAIBotChannel(t *testing.T) {
|
||||
// ---- Webhook mode tests ----
|
||||
|
||||
func TestNewWeComAIBotChannel_WebhookMode(t *testing.T) {
|
||||
t.Run("success with valid config", func(t *testing.T) {
|
||||
cfg := config.WeComAIBotConfig{
|
||||
Enabled: true,
|
||||
|
|
@ -22,14 +25,16 @@ func TestNewWeComAIBotChannel(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatalf("Expected no error, got %v", err)
|
||||
}
|
||||
|
||||
if ch == nil {
|
||||
t.Fatal("Expected channel to be created")
|
||||
}
|
||||
|
||||
if ch.Name() != "wecom_aibot" {
|
||||
t.Errorf("Expected name 'wecom_aibot', got '%s'", ch.Name())
|
||||
}
|
||||
// Webhook mode must implement WebhookHandler.
|
||||
if _, ok := ch.(channels.WebhookHandler); !ok {
|
||||
t.Error("Webhook mode channel should implement WebhookHandler")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("error with missing token", func(t *testing.T) {
|
||||
|
|
@ -37,10 +42,8 @@ func TestNewWeComAIBotChannel(t *testing.T) {
|
|||
Enabled: true,
|
||||
EncodingAESKey: "testkey1234567890123456789012345678901234567",
|
||||
}
|
||||
|
||||
messageBus := bus.NewMessageBus()
|
||||
_, err := NewWeComAIBotChannel(cfg, messageBus)
|
||||
|
||||
if err == nil {
|
||||
t.Fatal("Expected error for missing token, got nil")
|
||||
}
|
||||
|
|
@ -51,17 +54,15 @@ func TestNewWeComAIBotChannel(t *testing.T) {
|
|||
Enabled: true,
|
||||
Token: "test_token",
|
||||
}
|
||||
|
||||
messageBus := bus.NewMessageBus()
|
||||
_, err := NewWeComAIBotChannel(cfg, messageBus)
|
||||
|
||||
if err == nil {
|
||||
t.Fatal("Expected error for missing encoding key, got nil")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestWeComAIBotChannelStartStop(t *testing.T) {
|
||||
func TestWeComAIBotWebhookChannelStartStop(t *testing.T) {
|
||||
cfg := config.WeComAIBotConfig{
|
||||
Enabled: true,
|
||||
Token: "test_token",
|
||||
|
|
@ -76,22 +77,18 @@ func TestWeComAIBotChannelStartStop(t *testing.T) {
|
|||
|
||||
ctx := context.Background()
|
||||
|
||||
// Test Start
|
||||
if err := ch.Start(ctx); err != nil {
|
||||
t.Fatalf("Failed to start channel: %v", err)
|
||||
}
|
||||
|
||||
if !ch.IsRunning() {
|
||||
t.Error("Expected channel to be running")
|
||||
t.Error("Expected channel to be running after Start")
|
||||
}
|
||||
|
||||
// Test Stop
|
||||
if err := ch.Stop(ctx); err != nil {
|
||||
t.Fatalf("Failed to stop channel: %v", err)
|
||||
}
|
||||
|
||||
if ch.IsRunning() {
|
||||
t.Error("Expected channel to be stopped")
|
||||
t.Error("Expected channel to be stopped after Stop")
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -102,13 +99,16 @@ func TestWeComAIBotChannelWebhookPath(t *testing.T) {
|
|||
Token: "test_token",
|
||||
EncodingAESKey: "testkey1234567890123456789012345678901234567",
|
||||
}
|
||||
|
||||
messageBus := bus.NewMessageBus()
|
||||
ch, _ := NewWeComAIBotChannel(cfg, messageBus)
|
||||
|
||||
wh, ok := ch.(channels.WebhookHandler)
|
||||
if !ok {
|
||||
t.Fatal("Expected channel to implement WebhookHandler")
|
||||
}
|
||||
expectedPath := "/webhook/wecom-aibot"
|
||||
if ch.WebhookPath() != expectedPath {
|
||||
t.Errorf("Expected webhook path '%s', got '%s'", expectedPath, ch.WebhookPath())
|
||||
if wh.WebhookPath() != expectedPath {
|
||||
t.Errorf("Expected webhook path '%s', got '%s'", expectedPath, wh.WebhookPath())
|
||||
}
|
||||
})
|
||||
|
||||
|
|
@ -120,12 +120,15 @@ func TestWeComAIBotChannelWebhookPath(t *testing.T) {
|
|||
EncodingAESKey: "testkey1234567890123456789012345678901234567",
|
||||
WebhookPath: customPath,
|
||||
}
|
||||
|
||||
messageBus := bus.NewMessageBus()
|
||||
ch, _ := NewWeComAIBotChannel(cfg, messageBus)
|
||||
|
||||
if ch.WebhookPath() != customPath {
|
||||
t.Errorf("Expected webhook path '%s', got '%s'", customPath, ch.WebhookPath())
|
||||
wh, ok := ch.(channels.WebhookHandler)
|
||||
if !ok {
|
||||
t.Fatal("Expected channel to implement WebhookHandler")
|
||||
}
|
||||
if wh.WebhookPath() != customPath {
|
||||
t.Errorf("Expected webhook path '%s', got '%s'", customPath, wh.WebhookPath())
|
||||
}
|
||||
})
|
||||
}
|
||||
|
|
@ -136,19 +139,19 @@ func TestGenerateStreamID(t *testing.T) {
|
|||
Token: "test_token",
|
||||
EncodingAESKey: "testkey1234567890123456789012345678901234567",
|
||||
}
|
||||
|
||||
messageBus := bus.NewMessageBus()
|
||||
ch, _ := NewWeComAIBotChannel(cfg, messageBus)
|
||||
webhookCh, ok := ch.(*WeComAIBotChannel)
|
||||
if !ok {
|
||||
t.Fatal("Expected webhook mode channel")
|
||||
}
|
||||
|
||||
// Generate multiple IDs and check they are unique
|
||||
ids := make(map[string]bool)
|
||||
for i := 0; i < 100; i++ {
|
||||
id := ch.generateStreamID()
|
||||
|
||||
id := webhookCh.generateStreamID()
|
||||
if len(id) != 10 {
|
||||
t.Errorf("Expected stream ID length 10, got %d", len(id))
|
||||
}
|
||||
|
||||
if ids[id] {
|
||||
t.Errorf("Duplicate stream ID generated: %s", id)
|
||||
}
|
||||
|
|
@ -157,35 +160,33 @@ func TestGenerateStreamID(t *testing.T) {
|
|||
}
|
||||
|
||||
func TestEncryptDecrypt(t *testing.T) {
|
||||
// Use a valid 43-character base64 key (企业微信标准格式)
|
||||
cfg := config.WeComAIBotConfig{
|
||||
Enabled: true,
|
||||
Token: "test_token",
|
||||
EncodingAESKey: "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFG", // 43 characters
|
||||
}
|
||||
|
||||
messageBus := bus.NewMessageBus()
|
||||
ch, _ := NewWeComAIBotChannel(cfg, messageBus)
|
||||
webhookCh, ok := ch.(*WeComAIBotChannel)
|
||||
if !ok {
|
||||
t.Fatal("Expected webhook mode channel")
|
||||
}
|
||||
|
||||
plaintext := "Hello, World!"
|
||||
receiveid := ""
|
||||
|
||||
// Encrypt
|
||||
encrypted, err := ch.encryptMessage(plaintext, receiveid)
|
||||
encrypted, err := webhookCh.encryptMessage(plaintext, receiveid)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to encrypt message: %v", err)
|
||||
}
|
||||
|
||||
if encrypted == "" {
|
||||
t.Fatal("Encrypted message is empty")
|
||||
}
|
||||
|
||||
// Decrypt
|
||||
decrypted, err := decryptMessageWithVerify(encrypted, cfg.EncodingAESKey, receiveid)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to decrypt message: %v", err)
|
||||
}
|
||||
|
||||
if decrypted != plaintext {
|
||||
t.Errorf("Expected decrypted message '%s', got '%s'", plaintext, decrypted)
|
||||
}
|
||||
|
|
@ -198,13 +199,139 @@ func TestGenerateSignature(t *testing.T) {
|
|||
encrypt := "encrypted_msg"
|
||||
|
||||
signature := computeSignature(token, timestamp, nonce, encrypt)
|
||||
|
||||
if signature == "" {
|
||||
t.Error("Generated signature is empty")
|
||||
}
|
||||
|
||||
// Verify signature using verifySignature function
|
||||
if !verifySignature(token, signature, timestamp, nonce, encrypt) {
|
||||
t.Error("Generated signature does not verify correctly")
|
||||
}
|
||||
}
|
||||
|
||||
// ---- WebSocket long-connection mode tests ----
|
||||
|
||||
func TestNewWeComAIBotChannel_WSMode(t *testing.T) {
|
||||
t.Run("success with bot_id and secret", func(t *testing.T) {
|
||||
cfg := config.WeComAIBotConfig{
|
||||
Enabled: true,
|
||||
BotID: "test_bot_id",
|
||||
Secret: "test_secret",
|
||||
}
|
||||
messageBus := bus.NewMessageBus()
|
||||
ch, err := NewWeComAIBotChannel(cfg, messageBus)
|
||||
if err != nil {
|
||||
t.Fatalf("Expected no error, got %v", err)
|
||||
}
|
||||
if ch == nil {
|
||||
t.Fatal("Expected channel to be created")
|
||||
}
|
||||
if ch.Name() != "wecom_aibot" {
|
||||
t.Errorf("Expected name 'wecom_aibot', got '%s'", ch.Name())
|
||||
}
|
||||
// WebSocket mode must NOT implement WebhookHandler.
|
||||
if _, ok := ch.(channels.WebhookHandler); ok {
|
||||
t.Error("WebSocket mode channel should NOT implement WebhookHandler")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("ws mode takes priority over webhook fields", func(t *testing.T) {
|
||||
cfg := config.WeComAIBotConfig{
|
||||
Enabled: true,
|
||||
BotID: "test_bot_id",
|
||||
Secret: "test_secret",
|
||||
Token: "also_set",
|
||||
EncodingAESKey: "testkey1234567890123456789012345678901234567",
|
||||
}
|
||||
messageBus := bus.NewMessageBus()
|
||||
ch, err := NewWeComAIBotChannel(cfg, messageBus)
|
||||
if err != nil {
|
||||
t.Fatalf("Expected no error, got %v", err)
|
||||
}
|
||||
if _, ok := ch.(*WeComAIBotWSChannel); !ok {
|
||||
t.Error("Expected WebSocket mode channel when both BotID+Secret and Token+Key are set")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("error with missing bot_id", func(t *testing.T) {
|
||||
cfg := config.WeComAIBotConfig{
|
||||
Enabled: true,
|
||||
Secret: "test_secret",
|
||||
}
|
||||
messageBus := bus.NewMessageBus()
|
||||
_, err := NewWeComAIBotChannel(cfg, messageBus)
|
||||
// Missing bot_id alone means neither WS mode nor webhook mode is fully configured.
|
||||
if err == nil {
|
||||
t.Fatal("Expected error for missing bot_id, got nil")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("error with missing secret", func(t *testing.T) {
|
||||
cfg := config.WeComAIBotConfig{
|
||||
Enabled: true,
|
||||
BotID: "test_bot_id",
|
||||
}
|
||||
messageBus := bus.NewMessageBus()
|
||||
_, err := NewWeComAIBotChannel(cfg, messageBus)
|
||||
if err == nil {
|
||||
t.Fatal("Expected error for missing secret, got nil")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestWeComAIBotWSChannelStartStop(t *testing.T) {
|
||||
cfg := config.WeComAIBotConfig{
|
||||
Enabled: true,
|
||||
BotID: "test_bot_id",
|
||||
Secret: "test_secret",
|
||||
}
|
||||
messageBus := bus.NewMessageBus()
|
||||
ch, err := NewWeComAIBotChannel(cfg, messageBus)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create channel: %v", err)
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Start launches a background goroutine; it should not block or return an error.
|
||||
if err := ch.Start(ctx); err != nil {
|
||||
t.Fatalf("Failed to start channel: %v", err)
|
||||
}
|
||||
if !ch.IsRunning() {
|
||||
t.Error("Expected channel to be running after Start")
|
||||
}
|
||||
|
||||
// Stop should work regardless of whether the WebSocket actually connected.
|
||||
if err := ch.Stop(ctx); err != nil {
|
||||
t.Fatalf("Failed to stop channel: %v", err)
|
||||
}
|
||||
if ch.IsRunning() {
|
||||
t.Error("Expected channel to be stopped after Stop")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateRandomID(t *testing.T) {
|
||||
ids := make(map[string]bool)
|
||||
for i := 0; i < 200; i++ {
|
||||
id := generateRandomID(10)
|
||||
if len(id) != 10 {
|
||||
t.Errorf("Expected ID length 10, got %d", len(id))
|
||||
}
|
||||
if ids[id] {
|
||||
t.Errorf("Duplicate ID generated: %s", id)
|
||||
}
|
||||
ids[id] = true
|
||||
}
|
||||
}
|
||||
|
||||
func TestWSGenerateID(t *testing.T) {
|
||||
ids := make(map[string]bool)
|
||||
for i := 0; i < 200; i++ {
|
||||
id := wsGenerateID()
|
||||
if len(id) != 10 {
|
||||
t.Errorf("Expected ID length 10, got %d", len(id))
|
||||
}
|
||||
if ids[id] {
|
||||
t.Errorf("Duplicate wsGenerateID result: %s", id)
|
||||
}
|
||||
ids[id] = true
|
||||
}
|
||||
}
|
||||
|
|
|
|||
759
pkg/channels/wecom/aibot_ws.go
Normal file
759
pkg/channels/wecom/aibot_ws.go
Normal file
|
|
@ -0,0 +1,759 @@
|
|||
package wecom
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/bus"
|
||||
"github.com/sipeed/picoclaw/pkg/channels"
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
"github.com/sipeed/picoclaw/pkg/identity"
|
||||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
"github.com/sipeed/picoclaw/pkg/utils"
|
||||
)
|
||||
|
||||
// Long-connection WebSocket endpoint.
|
||||
// Ref: https://developer.work.weixin.qq.com/document/path/101463
|
||||
const (
|
||||
wsEndpoint = "wss://openws.work.weixin.qq.com"
|
||||
wsHeartbeatInterval = 30 * time.Second
|
||||
wsConnectTimeout = 15 * time.Second
|
||||
wsSubscribeTimeout = 10 * time.Second
|
||||
wsMaxReconnectWait = 60 * time.Second
|
||||
wsInitialReconnect = time.Second
|
||||
)
|
||||
|
||||
// WeComAIBotWSChannel implements channels.Channel for WeCom AI Bot using the
|
||||
// WebSocket long-connection API.
|
||||
// Unlike the webhook counterpart it does NOT implement WebhookHandler, so the
|
||||
// HTTP manager will not register any callback URL for it.
|
||||
type WeComAIBotWSChannel struct {
|
||||
*channels.BaseChannel
|
||||
config config.WeComAIBotConfig
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
|
||||
// conn is the active WebSocket connection; nil when disconnected.
|
||||
// All writes are serialized through connMu.
|
||||
conn *websocket.Conn
|
||||
connMu sync.Mutex
|
||||
|
||||
// tasks holds one live agent task per chatID.
|
||||
// A new message for the same chat cancels the previous task.
|
||||
tasks map[string]*wsTask
|
||||
tasksMu sync.Mutex
|
||||
|
||||
// reqPending correlates command req_ids with response channels.
|
||||
// Used only for subscribe/ping command-response pairs.
|
||||
reqPending map[string]chan wsEnvelope
|
||||
reqPendingMu sync.Mutex
|
||||
}
|
||||
|
||||
// wsTask tracks one in-progress agent reply for a single chat turn.
|
||||
type wsTask struct {
|
||||
ReqID string // req_id echoed in all replies for this turn
|
||||
ChatID string
|
||||
StreamID string // our generated stream.id
|
||||
CreatedTime time.Time
|
||||
answerCh chan string // agent delivers its reply here via Send()
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
}
|
||||
|
||||
// ---- WebSocket protocol types ----
|
||||
|
||||
// wsEnvelope is the generic JSON envelope for all WebSocket messages.
|
||||
type wsEnvelope struct {
|
||||
Cmd string `json:"cmd,omitempty"`
|
||||
Headers wsHeaders `json:"headers"`
|
||||
Body json.RawMessage `json:"body,omitempty"`
|
||||
ErrCode int `json:"errcode,omitempty"`
|
||||
ErrMsg string `json:"errmsg,omitempty"`
|
||||
}
|
||||
|
||||
type wsHeaders struct {
|
||||
ReqID string `json:"req_id"`
|
||||
}
|
||||
|
||||
// wsCommand is an outgoing request sent over the WebSocket.
|
||||
type wsCommand struct {
|
||||
Cmd string `json:"cmd"`
|
||||
Headers wsHeaders `json:"headers"`
|
||||
Body any `json:"body,omitempty"`
|
||||
}
|
||||
|
||||
// wsRespondMsgBody is the body for aibot_respond_msg / aibot_respond_welcome_msg.
|
||||
type wsRespondMsgBody struct {
|
||||
MsgType string `json:"msgtype"`
|
||||
Stream *wsStreamContent `json:"stream,omitempty"`
|
||||
Text *wsTextContent `json:"text,omitempty"`
|
||||
Markdown *wsMarkdownContent `json:"markdown,omitempty"`
|
||||
}
|
||||
|
||||
type wsStreamContent struct {
|
||||
ID string `json:"id"`
|
||||
Finish bool `json:"finish"`
|
||||
Content string `json:"content,omitempty"`
|
||||
}
|
||||
|
||||
type wsTextContent struct {
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
type wsMarkdownContent struct {
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
// WeComAIBotWSMessage is the decoded body of aibot_msg_callback /
|
||||
// aibot_event_callback in WebSocket long-connection mode.
|
||||
// The structure mirrors WeComAIBotMessage but includes extra fields
|
||||
// that only appear in long-connection callbacks (Voice, AESKey on Image/File).
|
||||
type WeComAIBotWSMessage struct {
|
||||
MsgID string `json:"msgid"`
|
||||
CreateTime int64 `json:"create_time,omitempty"`
|
||||
AIBotID string `json:"aibotid"`
|
||||
ChatID string `json:"chatid,omitempty"`
|
||||
ChatType string `json:"chattype,omitempty"` // "single" | "group"
|
||||
From struct {
|
||||
UserID string `json:"userid"`
|
||||
} `json:"from"`
|
||||
MsgType string `json:"msgtype"`
|
||||
Text *struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"text,omitempty"`
|
||||
Image *struct {
|
||||
URL string `json:"url"`
|
||||
AESKey string `json:"aeskey,omitempty"` // long-connection: per-resource decrypt key
|
||||
} `json:"image,omitempty"`
|
||||
Voice *struct {
|
||||
Text string `json:"text"` // WeCom transcribes voice to text in callbacks
|
||||
} `json:"voice,omitempty"`
|
||||
Mixed *struct {
|
||||
MsgItem []struct {
|
||||
MsgType string `json:"msgtype"`
|
||||
Text *struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"text,omitempty"`
|
||||
Image *struct {
|
||||
URL string `json:"url"`
|
||||
AESKey string `json:"aeskey,omitempty"`
|
||||
} `json:"image,omitempty"`
|
||||
} `json:"msg_item"`
|
||||
} `json:"mixed,omitempty"`
|
||||
Event *struct {
|
||||
EventType string `json:"eventtype"`
|
||||
} `json:"event,omitempty"`
|
||||
}
|
||||
|
||||
// ---- Constructor ----
|
||||
|
||||
// newWeComAIBotWSChannel creates a WeComAIBotWSChannel for WebSocket mode.
|
||||
func newWeComAIBotWSChannel(
|
||||
cfg config.WeComAIBotConfig,
|
||||
messageBus *bus.MessageBus,
|
||||
) (*WeComAIBotWSChannel, error) {
|
||||
if cfg.BotID == "" || cfg.Secret == "" {
|
||||
return nil, fmt.Errorf("bot_id and secret are required for WeCom AI Bot WebSocket mode")
|
||||
}
|
||||
|
||||
base := channels.NewBaseChannel("wecom_aibot", cfg, messageBus, cfg.AllowFrom,
|
||||
channels.WithMaxMessageLength(2048),
|
||||
channels.WithReasoningChannelID(cfg.ReasoningChannelID),
|
||||
)
|
||||
|
||||
return &WeComAIBotWSChannel{
|
||||
BaseChannel: base,
|
||||
config: cfg,
|
||||
tasks: make(map[string]*wsTask),
|
||||
reqPending: make(map[string]chan wsEnvelope),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ---- Channel interface ----
|
||||
|
||||
// Name implements channels.Channel.
|
||||
func (c *WeComAIBotWSChannel) Name() string { return "wecom_aibot" }
|
||||
|
||||
// Start connects to the WeCom WebSocket endpoint and begins message processing.
|
||||
func (c *WeComAIBotWSChannel) Start(ctx context.Context) error {
|
||||
logger.InfoC("wecom_aibot", "Starting WeCom AI Bot channel (WebSocket long-connection mode)...")
|
||||
c.ctx, c.cancel = context.WithCancel(ctx)
|
||||
c.SetRunning(true)
|
||||
go c.connectLoop()
|
||||
logger.InfoC("wecom_aibot", "WeCom AI Bot channel started (WebSocket mode)")
|
||||
return nil
|
||||
}
|
||||
|
||||
// Stop shuts down the channel and closes the WebSocket connection.
|
||||
func (c *WeComAIBotWSChannel) Stop(_ context.Context) error {
|
||||
logger.InfoC("wecom_aibot", "Stopping WeCom AI Bot channel (WebSocket mode)...")
|
||||
if c.cancel != nil {
|
||||
c.cancel()
|
||||
}
|
||||
c.connMu.Lock()
|
||||
if c.conn != nil {
|
||||
c.conn.Close()
|
||||
c.conn = nil
|
||||
}
|
||||
c.connMu.Unlock()
|
||||
c.SetRunning(false)
|
||||
logger.InfoC("wecom_aibot", "WeCom AI Bot channel stopped")
|
||||
return nil
|
||||
}
|
||||
|
||||
// Send delivers the agent reply for msg.ChatID.
|
||||
// The waiting task goroutine picks it up and writes the final stream response.
|
||||
func (c *WeComAIBotWSChannel) Send(ctx context.Context, msg bus.OutboundMessage) error {
|
||||
if !c.IsRunning() {
|
||||
return channels.ErrNotRunning
|
||||
}
|
||||
|
||||
c.tasksMu.Lock()
|
||||
task := c.tasks[msg.ChatID]
|
||||
c.tasksMu.Unlock()
|
||||
|
||||
if task == nil {
|
||||
logger.DebugCF("wecom_aibot", "Send: no active task for chat (may have finished or timed out)",
|
||||
map[string]any{"chat_id": msg.ChatID})
|
||||
return nil
|
||||
}
|
||||
|
||||
select {
|
||||
case task.answerCh <- msg.Content:
|
||||
case <-task.ctx.Done():
|
||||
return nil // task canceled (connection dropped)
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ---- Connection management ----
|
||||
|
||||
// connectLoop maintains the WebSocket connection, reconnecting on failure with
|
||||
// exponential backoff.
|
||||
func (c *WeComAIBotWSChannel) connectLoop() {
|
||||
backoff := wsInitialReconnect
|
||||
for {
|
||||
select {
|
||||
case <-c.ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
logger.InfoC("wecom_aibot", "Connecting to WeCom WebSocket endpoint...")
|
||||
if err := c.runConnection(); err != nil {
|
||||
select {
|
||||
case <-c.ctx.Done():
|
||||
return
|
||||
default:
|
||||
logger.WarnCF("wecom_aibot", "WebSocket connection lost, reconnecting",
|
||||
map[string]any{"error": err, "backoff": backoff.String()})
|
||||
select {
|
||||
case <-time.After(backoff):
|
||||
case <-c.ctx.Done():
|
||||
return
|
||||
}
|
||||
if backoff < wsMaxReconnectWait {
|
||||
backoff *= 2
|
||||
if backoff > wsMaxReconnectWait {
|
||||
backoff = wsMaxReconnectWait
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Clean exit (context canceled); stop reconnecting.
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// runConnection dials, subscribes, and runs the read/heartbeat loops until the
|
||||
// connection closes or the channel context is canceled.
|
||||
func (c *WeComAIBotWSChannel) runConnection() error {
|
||||
dialCtx, dialCancel := context.WithTimeout(c.ctx, wsConnectTimeout)
|
||||
conn, httpResp, err := websocket.DefaultDialer.DialContext(dialCtx, wsEndpoint, nil)
|
||||
dialCancel()
|
||||
if httpResp != nil {
|
||||
httpResp.Body.Close()
|
||||
}
|
||||
if err != nil {
|
||||
return fmt.Errorf("dial failed: %w", err)
|
||||
}
|
||||
|
||||
c.connMu.Lock()
|
||||
c.conn = conn
|
||||
c.connMu.Unlock()
|
||||
|
||||
defer func() {
|
||||
c.connMu.Lock()
|
||||
if c.conn == conn {
|
||||
c.conn = nil
|
||||
}
|
||||
c.connMu.Unlock()
|
||||
// Cancel any tasks that were started over this connection so their
|
||||
// agent goroutines do not keep running after the connection is gone.
|
||||
c.cancelAllTasks()
|
||||
}()
|
||||
|
||||
// ---- Read loop (must start BEFORE subscribing) ----
|
||||
// sendAndWait blocks waiting for the subscribe response on reqPending;
|
||||
// readLoop is the only goroutine that delivers messages to reqPending.
|
||||
// Starting readLoop first avoids a deadlock where sendAndWait times out
|
||||
// because no one reads the server's reply.
|
||||
readErrCh := make(chan error, 1)
|
||||
go func() { readErrCh <- c.readLoop(conn) }()
|
||||
|
||||
// ---- Subscribe ----
|
||||
reqID := wsGenerateID()
|
||||
resp, err := c.sendAndWait(conn, reqID, wsCommand{
|
||||
Cmd: "aibot_subscribe",
|
||||
Headers: wsHeaders{ReqID: reqID},
|
||||
Body: map[string]string{
|
||||
"bot_id": c.config.BotID,
|
||||
"secret": c.config.Secret,
|
||||
},
|
||||
}, wsSubscribeTimeout)
|
||||
if err != nil {
|
||||
conn.Close() // stop readLoop
|
||||
<-readErrCh
|
||||
return fmt.Errorf("subscribe failed: %w", err)
|
||||
}
|
||||
if resp.ErrCode != 0 {
|
||||
conn.Close()
|
||||
<-readErrCh
|
||||
return fmt.Errorf("subscribe rejected (errcode=%d): %s", resp.ErrCode, resp.ErrMsg)
|
||||
}
|
||||
|
||||
logger.InfoC("wecom_aibot", "WebSocket subscription successful")
|
||||
|
||||
// ---- Heartbeat goroutine ----
|
||||
hbDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(hbDone)
|
||||
c.heartbeatLoop(conn)
|
||||
}()
|
||||
|
||||
// Wait for the read loop to exit, then tear down the heartbeat.
|
||||
readErr := <-readErrCh
|
||||
conn.Close() // signal heartbeat to stop (idempotent)
|
||||
<-hbDone
|
||||
return readErr
|
||||
}
|
||||
|
||||
// sendAndWait registers a pending-response slot, sends cmd, and blocks until
|
||||
// the matching response arrives or the timeout/context fires.
|
||||
func (c *WeComAIBotWSChannel) sendAndWait(
|
||||
conn *websocket.Conn,
|
||||
reqID string,
|
||||
cmd wsCommand,
|
||||
timeout time.Duration,
|
||||
) (wsEnvelope, error) {
|
||||
ch := make(chan wsEnvelope, 1)
|
||||
c.reqPendingMu.Lock()
|
||||
c.reqPending[reqID] = ch
|
||||
c.reqPendingMu.Unlock()
|
||||
|
||||
cleanup := func() {
|
||||
c.reqPendingMu.Lock()
|
||||
delete(c.reqPending, reqID)
|
||||
c.reqPendingMu.Unlock()
|
||||
}
|
||||
|
||||
data, err := json.Marshal(cmd)
|
||||
if err != nil {
|
||||
cleanup()
|
||||
return wsEnvelope{}, fmt.Errorf("marshal command: %w", err)
|
||||
}
|
||||
c.connMu.Lock()
|
||||
err = conn.WriteMessage(websocket.TextMessage, data)
|
||||
c.connMu.Unlock()
|
||||
if err != nil {
|
||||
cleanup()
|
||||
return wsEnvelope{}, fmt.Errorf("write command: %w", err)
|
||||
}
|
||||
|
||||
timer := time.NewTimer(timeout)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case env := <-ch:
|
||||
return env, nil
|
||||
case <-timer.C:
|
||||
cleanup()
|
||||
return wsEnvelope{}, fmt.Errorf("timeout waiting for response (req_id=%s)", reqID)
|
||||
case <-c.ctx.Done():
|
||||
cleanup()
|
||||
return wsEnvelope{}, c.ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
// heartbeatLoop sends a ping every wsHeartbeatInterval until conn is closed.
|
||||
func (c *WeComAIBotWSChannel) heartbeatLoop(conn *websocket.Conn) {
|
||||
ticker := time.NewTicker(wsHeartbeatInterval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
reqID := wsGenerateID()
|
||||
data, _ := json.Marshal(wsCommand{
|
||||
Cmd: "ping",
|
||||
Headers: wsHeaders{ReqID: reqID},
|
||||
})
|
||||
c.connMu.Lock()
|
||||
err := conn.WriteMessage(websocket.TextMessage, data)
|
||||
c.connMu.Unlock()
|
||||
if err != nil {
|
||||
logger.WarnCF("wecom_aibot", "Heartbeat write failed", map[string]any{"error": err})
|
||||
return
|
||||
}
|
||||
logger.DebugCF("wecom_aibot", "Heartbeat sent", map[string]any{"req_id": reqID})
|
||||
case <-c.ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// readLoop reads WebSocket messages and dispatches them until the connection
|
||||
// closes or the channel is stopped.
|
||||
func (c *WeComAIBotWSChannel) readLoop(conn *websocket.Conn) error {
|
||||
for {
|
||||
_, raw, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
select {
|
||||
case <-c.ctx.Done():
|
||||
return nil // clean shutdown
|
||||
default:
|
||||
return fmt.Errorf("read error: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
var env wsEnvelope
|
||||
if err := json.Unmarshal(raw, &env); err != nil {
|
||||
logger.WarnCF("wecom_aibot", "Failed to parse WebSocket message",
|
||||
map[string]any{"error": err, "raw": string(raw)})
|
||||
continue
|
||||
}
|
||||
|
||||
// If there is a waiting sendAndWait() call for this req_id, forward
|
||||
// the envelope to it. Command responses have an empty Cmd field.
|
||||
if env.Cmd == "" && env.Headers.ReqID != "" {
|
||||
c.reqPendingMu.Lock()
|
||||
ch, ok := c.reqPending[env.Headers.ReqID]
|
||||
if ok {
|
||||
delete(c.reqPending, env.Headers.ReqID)
|
||||
}
|
||||
c.reqPendingMu.Unlock()
|
||||
if ok {
|
||||
ch <- env
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
// Dispatch to appropriate handler in a separate goroutine so the
|
||||
// read loop is never blocked by a slow agent.
|
||||
go c.handleEnvelope(env)
|
||||
}
|
||||
}
|
||||
|
||||
// ---- Message / event handlers ----
|
||||
|
||||
// handleEnvelope routes a WebSocket envelope to the right handler.
|
||||
func (c *WeComAIBotWSChannel) handleEnvelope(env wsEnvelope) {
|
||||
switch env.Cmd {
|
||||
case "aibot_msg_callback":
|
||||
c.handleMsgCallback(env)
|
||||
case "aibot_event_callback":
|
||||
c.handleEventCallback(env)
|
||||
default:
|
||||
logger.DebugCF("wecom_aibot", "Unhandled WebSocket command",
|
||||
map[string]any{"cmd": env.Cmd})
|
||||
}
|
||||
}
|
||||
|
||||
// handleMsgCallback processes aibot_msg_callback.
|
||||
func (c *WeComAIBotWSChannel) handleMsgCallback(env wsEnvelope) {
|
||||
var msg WeComAIBotWSMessage
|
||||
if err := json.Unmarshal(env.Body, &msg); err != nil {
|
||||
logger.WarnCF("wecom_aibot", "Failed to parse msg callback body",
|
||||
map[string]any{"error": err})
|
||||
return
|
||||
}
|
||||
|
||||
reqID := env.Headers.ReqID
|
||||
switch msg.MsgType {
|
||||
case "text":
|
||||
c.handleWSTextMessage(reqID, msg)
|
||||
case "image":
|
||||
c.handleWSImageMessage(reqID, msg)
|
||||
case "voice":
|
||||
c.handleWSVoiceMessage(reqID, msg)
|
||||
case "mixed":
|
||||
c.handleWSMixedMessage(reqID, msg)
|
||||
default:
|
||||
logger.WarnCF("wecom_aibot", "Unsupported message type",
|
||||
map[string]any{"msgtype": msg.MsgType})
|
||||
c.wsSendStreamFinish(reqID, wsGenerateID(),
|
||||
"Unsupported message type: "+msg.MsgType)
|
||||
}
|
||||
}
|
||||
|
||||
// handleEventCallback processes aibot_event_callback.
|
||||
func (c *WeComAIBotWSChannel) handleEventCallback(env wsEnvelope) {
|
||||
var msg WeComAIBotWSMessage
|
||||
if err := json.Unmarshal(env.Body, &msg); err != nil {
|
||||
logger.WarnCF("wecom_aibot", "Failed to parse event callback body",
|
||||
map[string]any{"error": err})
|
||||
return
|
||||
}
|
||||
|
||||
var eventType string
|
||||
if msg.Event != nil {
|
||||
eventType = msg.Event.EventType
|
||||
}
|
||||
logger.DebugCF("wecom_aibot", "Received event callback",
|
||||
map[string]any{"event_type": eventType})
|
||||
|
||||
switch eventType {
|
||||
case "enter_chat":
|
||||
if c.config.WelcomeMessage != "" {
|
||||
c.wsSendWelcomeMsg(env.Headers.ReqID, c.config.WelcomeMessage)
|
||||
}
|
||||
case "disconnected_event":
|
||||
// The server will close this connection after sending this event.
|
||||
// connectLoop will detect the closure and reconnect automatically.
|
||||
logger.WarnC("wecom_aibot",
|
||||
"Received disconnected_event: this connection is being replaced by a newer one")
|
||||
default:
|
||||
logger.DebugCF("wecom_aibot", "Unhandled event type",
|
||||
map[string]any{"event_type": eventType})
|
||||
}
|
||||
}
|
||||
|
||||
// handleWSTextMessage dispatches a plain-text message to the agent and streams
|
||||
// the reply back over the WebSocket connection.
|
||||
func (c *WeComAIBotWSChannel) handleWSTextMessage(reqID string, msg WeComAIBotWSMessage) {
|
||||
if msg.Text == nil {
|
||||
logger.ErrorC("wecom_aibot", "text message missing text field")
|
||||
return
|
||||
}
|
||||
|
||||
userID := msg.From.UserID
|
||||
if userID == "" {
|
||||
userID = "unknown"
|
||||
}
|
||||
chatID := msg.ChatID
|
||||
if chatID == "" {
|
||||
chatID = userID
|
||||
}
|
||||
|
||||
streamID := wsGenerateID()
|
||||
taskCtx, taskCancel := context.WithCancel(c.ctx)
|
||||
|
||||
task := &wsTask{
|
||||
ReqID: reqID,
|
||||
ChatID: chatID,
|
||||
StreamID: streamID,
|
||||
CreatedTime: time.Now(),
|
||||
answerCh: make(chan string, 1),
|
||||
ctx: taskCtx,
|
||||
cancel: taskCancel,
|
||||
}
|
||||
|
||||
c.tasksMu.Lock()
|
||||
// Cancel any previous task for this chat (user sent a new message mid-reply).
|
||||
if prev, ok := c.tasks[chatID]; ok {
|
||||
prev.cancel()
|
||||
}
|
||||
c.tasks[chatID] = task
|
||||
c.tasksMu.Unlock()
|
||||
|
||||
// Send an empty stream opening frame (finish=false) immediately so WeCom
|
||||
// shows the typing indicator while the agent is processing.
|
||||
c.wsSendStreamChunk(reqID, streamID, false, "")
|
||||
|
||||
go func() {
|
||||
defer func() {
|
||||
taskCancel()
|
||||
c.tasksMu.Lock()
|
||||
if c.tasks[chatID] == task {
|
||||
delete(c.tasks, chatID)
|
||||
}
|
||||
c.tasksMu.Unlock()
|
||||
}()
|
||||
|
||||
sender := bus.SenderInfo{
|
||||
Platform: "wecom_aibot",
|
||||
PlatformID: userID,
|
||||
CanonicalID: identity.BuildCanonicalID("wecom_aibot", userID),
|
||||
DisplayName: userID,
|
||||
}
|
||||
peerKind := "direct"
|
||||
if msg.ChatType == "group" {
|
||||
peerKind = "group"
|
||||
}
|
||||
peer := bus.Peer{Kind: peerKind, ID: chatID}
|
||||
metadata := map[string]string{
|
||||
"channel": "wecom_aibot",
|
||||
"chat_type": msg.ChatType,
|
||||
"msg_type": "text",
|
||||
"msgid": msg.MsgID,
|
||||
"aibotid": msg.AIBotID,
|
||||
"stream_id": streamID,
|
||||
}
|
||||
c.HandleMessage(taskCtx, peer, msg.MsgID, userID, chatID,
|
||||
msg.Text.Content, nil, metadata, sender)
|
||||
|
||||
// Wait for the agent reply and send it as the final stream frame.
|
||||
select {
|
||||
case answer := <-task.answerCh:
|
||||
c.wsSendStreamFinish(reqID, streamID, answer)
|
||||
case <-taskCtx.Done():
|
||||
// Connection dropped or task canceled; nothing to send.
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// handleWSImageMessage handles image messages.
|
||||
func (c *WeComAIBotWSChannel) handleWSImageMessage(reqID string, msg WeComAIBotWSMessage) {
|
||||
logger.WarnC("wecom_aibot", "Image messages not yet supported in WebSocket mode")
|
||||
content := "Image messages are not yet supported."
|
||||
if msg.Image != nil {
|
||||
content = fmt.Sprintf(
|
||||
"Image received (URL: %s), but image messages are not yet supported.",
|
||||
msg.Image.URL,
|
||||
)
|
||||
}
|
||||
c.wsSendStreamFinish(reqID, wsGenerateID(), content)
|
||||
}
|
||||
|
||||
// handleWSMixedMessage handles mixed text+image messages.
|
||||
// If the message contains a text part it is handled as a text message;
|
||||
// otherwise an unsupported-type notice is returned.
|
||||
func (c *WeComAIBotWSChannel) handleWSMixedMessage(reqID string, msg WeComAIBotWSMessage) {
|
||||
if msg.Mixed != nil {
|
||||
for _, item := range msg.Mixed.MsgItem {
|
||||
if item.MsgType == "text" && item.Text != nil {
|
||||
// Treat the text portion as a standalone text message.
|
||||
c.handleWSTextMessage(reqID, WeComAIBotWSMessage{
|
||||
MsgID: msg.MsgID,
|
||||
AIBotID: msg.AIBotID,
|
||||
ChatID: msg.ChatID,
|
||||
ChatType: msg.ChatType,
|
||||
From: msg.From,
|
||||
MsgType: "text",
|
||||
Text: item.Text,
|
||||
})
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
logger.WarnC("wecom_aibot", "Mixed message has no usable text part")
|
||||
c.wsSendStreamFinish(reqID, wsGenerateID(), "Mixed message type is not yet fully supported.")
|
||||
}
|
||||
|
||||
// handleWSVoiceMessage handles voice messages.
|
||||
// WeCom transcribes voice to text in the callback; if the transcription is
|
||||
// present it is forwarded as a text message.
|
||||
func (c *WeComAIBotWSChannel) handleWSVoiceMessage(reqID string, msg WeComAIBotWSMessage) {
|
||||
if msg.Voice != nil && msg.Voice.Text != "" {
|
||||
c.handleWSTextMessage(reqID, WeComAIBotWSMessage{
|
||||
MsgID: msg.MsgID,
|
||||
AIBotID: msg.AIBotID,
|
||||
ChatID: msg.ChatID,
|
||||
ChatType: msg.ChatType,
|
||||
From: msg.From,
|
||||
MsgType: "text",
|
||||
Text: &struct {
|
||||
Content string `json:"content"`
|
||||
}{Content: msg.Voice.Text},
|
||||
})
|
||||
return
|
||||
}
|
||||
c.wsSendStreamFinish(reqID, wsGenerateID(), "Voice messages are not yet supported.")
|
||||
}
|
||||
|
||||
// ---- WebSocket write helpers ----
|
||||
|
||||
// wsSendStreamChunk sends an aibot_respond_msg stream frame.
|
||||
func (c *WeComAIBotWSChannel) wsSendStreamChunk(reqID, streamID string, finish bool, content string) {
|
||||
logger.DebugCF("wecom_aibot", "Sending stream chunk", map[string]any{
|
||||
"stream_id": streamID,
|
||||
"finish": finish,
|
||||
"preview": utils.Truncate(content, 100),
|
||||
})
|
||||
c.writeWS(wsCommand{
|
||||
Cmd: "aibot_respond_msg",
|
||||
Headers: wsHeaders{ReqID: reqID},
|
||||
Body: wsRespondMsgBody{
|
||||
MsgType: "stream",
|
||||
Stream: &wsStreamContent{
|
||||
ID: streamID,
|
||||
Finish: finish,
|
||||
Content: content,
|
||||
},
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// wsSendStreamFinish sends the final aibot_respond_msg frame (finish=true).
|
||||
func (c *WeComAIBotWSChannel) wsSendStreamFinish(reqID, streamID, content string) {
|
||||
c.wsSendStreamChunk(reqID, streamID, true, content)
|
||||
}
|
||||
|
||||
// wsSendWelcomeMsg sends a text welcome message via aibot_respond_welcome_msg.
|
||||
func (c *WeComAIBotWSChannel) wsSendWelcomeMsg(reqID, content string) {
|
||||
logger.DebugCF("wecom_aibot", "Sending welcome message", map[string]any{"req_id": reqID})
|
||||
c.writeWS(wsCommand{
|
||||
Cmd: "aibot_respond_welcome_msg",
|
||||
Headers: wsHeaders{ReqID: reqID},
|
||||
Body: wsRespondMsgBody{
|
||||
MsgType: "text",
|
||||
Text: &wsTextContent{Content: content},
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// writeWS serializes cmd to JSON and writes it to the active WebSocket
|
||||
// connection. It is safe to call from multiple goroutines.
|
||||
func (c *WeComAIBotWSChannel) writeWS(cmd any) {
|
||||
data, err := json.Marshal(cmd)
|
||||
if err != nil {
|
||||
logger.ErrorCF("wecom_aibot", "Failed to marshal WebSocket command",
|
||||
map[string]any{"error": err})
|
||||
return
|
||||
}
|
||||
c.connMu.Lock()
|
||||
conn := c.conn
|
||||
if conn != nil {
|
||||
err = conn.WriteMessage(websocket.TextMessage, data)
|
||||
}
|
||||
c.connMu.Unlock()
|
||||
if conn == nil {
|
||||
logger.WarnC("wecom_aibot", "WebSocket connection unavailable, dropping outbound message")
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
logger.WarnCF("wecom_aibot", "WebSocket write failed", map[string]any{"error": err})
|
||||
}
|
||||
}
|
||||
|
||||
// cancelAllTasks cancels every pending agent task; called when the connection drops.
|
||||
func (c *WeComAIBotWSChannel) cancelAllTasks() {
|
||||
c.tasksMu.Lock()
|
||||
defer c.tasksMu.Unlock()
|
||||
for chatID, task := range c.tasks {
|
||||
task.cancel()
|
||||
delete(c.tasks, chatID)
|
||||
}
|
||||
}
|
||||
|
||||
// wsGenerateID generates a random 10-character alphanumeric ID.
|
||||
// It is package-level (not a method) so it can be shared by both channel modes.
|
||||
func wsGenerateID() string {
|
||||
return generateRandomID(10)
|
||||
}
|
||||
|
|
@ -475,12 +475,23 @@ type WeComAppConfig struct {
|
|||
|
||||
type WeComAIBotConfig struct {
|
||||
Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_ENABLED"`
|
||||
Token string `json:"token" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_TOKEN"`
|
||||
EncodingAESKey string `json:"encoding_aes_key" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_ENCODING_AES_KEY"`
|
||||
WebhookPath string `json:"webhook_path" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WEBHOOK_PATH"`
|
||||
|
||||
// WebSocket long-connection mode.
|
||||
// Set BotID + Secret to use this mode (no encryption, no webhook needed).
|
||||
// Ref: https://developer.work.weixin.qq.com/document/path/101463
|
||||
BotID string `json:"bot_id,omitempty" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_BOT_ID"`
|
||||
Secret string `json:"secret,omitempty" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_SECRET"`
|
||||
|
||||
// Webhook (short-connection) mode.
|
||||
// Set Token + EncodingAESKey to use this mode.
|
||||
Token string `json:"token,omitempty" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_TOKEN"`
|
||||
EncodingAESKey string `json:"encoding_aes_key,omitempty" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_ENCODING_AES_KEY"`
|
||||
WebhookPath string `json:"webhook_path,omitempty" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WEBHOOK_PATH"`
|
||||
|
||||
// Common fields
|
||||
AllowFrom FlexibleStringSlice `json:"allow_from" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_ALLOW_FROM"`
|
||||
ReplyTimeout int `json:"reply_timeout" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_REPLY_TIMEOUT"`
|
||||
MaxSteps int `json:"max_steps" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_MAX_STEPS"` // Maximum streaming steps
|
||||
MaxSteps int `json:"max_steps" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_MAX_STEPS"`
|
||||
WelcomeMessage string `json:"welcome_message" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WELCOME_MESSAGE"` // Sent on enter_chat event; empty = no welcome
|
||||
ReasoningChannelID string `json:"reasoning_channel_id" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_REASONING_CHANNEL_ID"`
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue