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:
parent
7673b626b3
commit
6f5d1f336f
6 changed files with 1025 additions and 49 deletions
|
|
@ -204,6 +204,8 @@
|
||||||
"wecom_aibot": {
|
"wecom_aibot": {
|
||||||
"_comment": "WeCom AI Bot (智能机器人) - Official WeCom AI Bot integration, supports proactive messaging and private chats.",
|
"_comment": "WeCom AI Bot (智能机器人) - Official WeCom AI Bot integration, supports proactive messaging and private chats.",
|
||||||
"enabled": false,
|
"enabled": false,
|
||||||
|
"bot_id": "YOUR_BOT_ID",
|
||||||
|
"secret": "YOUR_SECRET",
|
||||||
"token": "YOUR_TOKEN",
|
"token": "YOUR_TOKEN",
|
||||||
"encoding_aes_key": "YOUR_43_CHAR_ENCODING_AES_KEY",
|
"encoding_aes_key": "YOUR_43_CHAR_ENCODING_AES_KEY",
|
||||||
"webhook_path": "/webhook/wecom-aibot",
|
"webhook_path": "/webhook/wecom-aibot",
|
||||||
|
|
|
||||||
|
|
@ -296,7 +296,9 @@ func (m *Manager) initChannels(channels *config.ChannelsConfig) error {
|
||||||
m.initChannel("wecom", "WeCom")
|
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")
|
m.initChannel("wecom_aibot", "WeCom AI Bot")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -134,13 +134,25 @@ type WeComAIBotEncryptedResponse struct {
|
||||||
Nonce string `json:"nonce"`
|
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(
|
func NewWeComAIBotChannel(
|
||||||
cfg config.WeComAIBotConfig,
|
cfg config.WeComAIBotConfig,
|
||||||
messageBus *bus.MessageBus,
|
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 == "" {
|
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,
|
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
|
return base64.StdEncoding.EncodeToString(ciphertext), nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// generateStreamID generates a random stream ID
|
// func (c *WeComAIBotChannel) downloadAndDecryptImage(
|
||||||
func (c *WeComAIBotChannel) generateStreamID() string {
|
// 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"
|
const letters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
|
||||||
b := make([]byte, 10)
|
b := make([]byte, n)
|
||||||
for i := range b {
|
for i := range b {
|
||||||
n, _ := rand.Int(rand.Reader, big.NewInt(int64(len(letters))))
|
num, _ := rand.Int(rand.Reader, big.NewInt(int64(len(letters))))
|
||||||
b[i] = letters[n.Int64()]
|
b[i] = letters[num.Int64()]
|
||||||
}
|
}
|
||||||
return string(b)
|
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
|
// cleanupLoop periodically cleans up old streaming tasks
|
||||||
func (c *WeComAIBotChannel) cleanupLoop() {
|
func (c *WeComAIBotChannel) cleanupLoop() {
|
||||||
ticker := time.NewTicker(5 * time.Minute)
|
ticker := time.NewTicker(5 * time.Minute)
|
||||||
|
|
|
||||||
|
|
@ -5,10 +5,13 @@ import (
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
"github.com/sipeed/picoclaw/pkg/bus"
|
"github.com/sipeed/picoclaw/pkg/bus"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/channels"
|
||||||
"github.com/sipeed/picoclaw/pkg/config"
|
"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) {
|
t.Run("success with valid config", func(t *testing.T) {
|
||||||
cfg := config.WeComAIBotConfig{
|
cfg := config.WeComAIBotConfig{
|
||||||
Enabled: true,
|
Enabled: true,
|
||||||
|
|
@ -22,14 +25,16 @@ func TestNewWeComAIBotChannel(t *testing.T) {
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("Expected no error, got %v", err)
|
t.Fatalf("Expected no error, got %v", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
if ch == nil {
|
if ch == nil {
|
||||||
t.Fatal("Expected channel to be created")
|
t.Fatal("Expected channel to be created")
|
||||||
}
|
}
|
||||||
|
|
||||||
if ch.Name() != "wecom_aibot" {
|
if ch.Name() != "wecom_aibot" {
|
||||||
t.Errorf("Expected name 'wecom_aibot', got '%s'", ch.Name())
|
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) {
|
t.Run("error with missing token", func(t *testing.T) {
|
||||||
|
|
@ -37,10 +42,8 @@ func TestNewWeComAIBotChannel(t *testing.T) {
|
||||||
Enabled: true,
|
Enabled: true,
|
||||||
EncodingAESKey: "testkey1234567890123456789012345678901234567",
|
EncodingAESKey: "testkey1234567890123456789012345678901234567",
|
||||||
}
|
}
|
||||||
|
|
||||||
messageBus := bus.NewMessageBus()
|
messageBus := bus.NewMessageBus()
|
||||||
_, err := NewWeComAIBotChannel(cfg, messageBus)
|
_, err := NewWeComAIBotChannel(cfg, messageBus)
|
||||||
|
|
||||||
if err == nil {
|
if err == nil {
|
||||||
t.Fatal("Expected error for missing token, got nil")
|
t.Fatal("Expected error for missing token, got nil")
|
||||||
}
|
}
|
||||||
|
|
@ -51,17 +54,15 @@ func TestNewWeComAIBotChannel(t *testing.T) {
|
||||||
Enabled: true,
|
Enabled: true,
|
||||||
Token: "test_token",
|
Token: "test_token",
|
||||||
}
|
}
|
||||||
|
|
||||||
messageBus := bus.NewMessageBus()
|
messageBus := bus.NewMessageBus()
|
||||||
_, err := NewWeComAIBotChannel(cfg, messageBus)
|
_, err := NewWeComAIBotChannel(cfg, messageBus)
|
||||||
|
|
||||||
if err == nil {
|
if err == nil {
|
||||||
t.Fatal("Expected error for missing encoding key, got nil")
|
t.Fatal("Expected error for missing encoding key, got nil")
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestWeComAIBotChannelStartStop(t *testing.T) {
|
func TestWeComAIBotWebhookChannelStartStop(t *testing.T) {
|
||||||
cfg := config.WeComAIBotConfig{
|
cfg := config.WeComAIBotConfig{
|
||||||
Enabled: true,
|
Enabled: true,
|
||||||
Token: "test_token",
|
Token: "test_token",
|
||||||
|
|
@ -76,22 +77,18 @@ func TestWeComAIBotChannelStartStop(t *testing.T) {
|
||||||
|
|
||||||
ctx := context.Background()
|
ctx := context.Background()
|
||||||
|
|
||||||
// Test Start
|
|
||||||
if err := ch.Start(ctx); err != nil {
|
if err := ch.Start(ctx); err != nil {
|
||||||
t.Fatalf("Failed to start channel: %v", err)
|
t.Fatalf("Failed to start channel: %v", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
if !ch.IsRunning() {
|
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 {
|
if err := ch.Stop(ctx); err != nil {
|
||||||
t.Fatalf("Failed to stop channel: %v", err)
|
t.Fatalf("Failed to stop channel: %v", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
if ch.IsRunning() {
|
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",
|
Token: "test_token",
|
||||||
EncodingAESKey: "testkey1234567890123456789012345678901234567",
|
EncodingAESKey: "testkey1234567890123456789012345678901234567",
|
||||||
}
|
}
|
||||||
|
|
||||||
messageBus := bus.NewMessageBus()
|
messageBus := bus.NewMessageBus()
|
||||||
ch, _ := NewWeComAIBotChannel(cfg, messageBus)
|
ch, _ := NewWeComAIBotChannel(cfg, messageBus)
|
||||||
|
|
||||||
|
wh, ok := ch.(channels.WebhookHandler)
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("Expected channel to implement WebhookHandler")
|
||||||
|
}
|
||||||
expectedPath := "/webhook/wecom-aibot"
|
expectedPath := "/webhook/wecom-aibot"
|
||||||
if ch.WebhookPath() != expectedPath {
|
if wh.WebhookPath() != expectedPath {
|
||||||
t.Errorf("Expected webhook path '%s', got '%s'", expectedPath, ch.WebhookPath())
|
t.Errorf("Expected webhook path '%s', got '%s'", expectedPath, wh.WebhookPath())
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
|
|
||||||
|
|
@ -120,12 +120,15 @@ func TestWeComAIBotChannelWebhookPath(t *testing.T) {
|
||||||
EncodingAESKey: "testkey1234567890123456789012345678901234567",
|
EncodingAESKey: "testkey1234567890123456789012345678901234567",
|
||||||
WebhookPath: customPath,
|
WebhookPath: customPath,
|
||||||
}
|
}
|
||||||
|
|
||||||
messageBus := bus.NewMessageBus()
|
messageBus := bus.NewMessageBus()
|
||||||
ch, _ := NewWeComAIBotChannel(cfg, messageBus)
|
ch, _ := NewWeComAIBotChannel(cfg, messageBus)
|
||||||
|
|
||||||
if ch.WebhookPath() != customPath {
|
wh, ok := ch.(channels.WebhookHandler)
|
||||||
t.Errorf("Expected webhook path '%s', got '%s'", customPath, ch.WebhookPath())
|
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",
|
Token: "test_token",
|
||||||
EncodingAESKey: "testkey1234567890123456789012345678901234567",
|
EncodingAESKey: "testkey1234567890123456789012345678901234567",
|
||||||
}
|
}
|
||||||
|
|
||||||
messageBus := bus.NewMessageBus()
|
messageBus := bus.NewMessageBus()
|
||||||
ch, _ := NewWeComAIBotChannel(cfg, messageBus)
|
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)
|
ids := make(map[string]bool)
|
||||||
for i := 0; i < 100; i++ {
|
for i := 0; i < 100; i++ {
|
||||||
id := ch.generateStreamID()
|
id := webhookCh.generateStreamID()
|
||||||
|
|
||||||
if len(id) != 10 {
|
if len(id) != 10 {
|
||||||
t.Errorf("Expected stream ID length 10, got %d", len(id))
|
t.Errorf("Expected stream ID length 10, got %d", len(id))
|
||||||
}
|
}
|
||||||
|
|
||||||
if ids[id] {
|
if ids[id] {
|
||||||
t.Errorf("Duplicate stream ID generated: %s", id)
|
t.Errorf("Duplicate stream ID generated: %s", id)
|
||||||
}
|
}
|
||||||
|
|
@ -157,35 +160,33 @@ func TestGenerateStreamID(t *testing.T) {
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestEncryptDecrypt(t *testing.T) {
|
func TestEncryptDecrypt(t *testing.T) {
|
||||||
// Use a valid 43-character base64 key (企业微信标准格式)
|
|
||||||
cfg := config.WeComAIBotConfig{
|
cfg := config.WeComAIBotConfig{
|
||||||
Enabled: true,
|
Enabled: true,
|
||||||
Token: "test_token",
|
Token: "test_token",
|
||||||
EncodingAESKey: "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFG", // 43 characters
|
EncodingAESKey: "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFG", // 43 characters
|
||||||
}
|
}
|
||||||
|
|
||||||
messageBus := bus.NewMessageBus()
|
messageBus := bus.NewMessageBus()
|
||||||
ch, _ := NewWeComAIBotChannel(cfg, messageBus)
|
ch, _ := NewWeComAIBotChannel(cfg, messageBus)
|
||||||
|
webhookCh, ok := ch.(*WeComAIBotChannel)
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("Expected webhook mode channel")
|
||||||
|
}
|
||||||
|
|
||||||
plaintext := "Hello, World!"
|
plaintext := "Hello, World!"
|
||||||
receiveid := ""
|
receiveid := ""
|
||||||
|
|
||||||
// Encrypt
|
encrypted, err := webhookCh.encryptMessage(plaintext, receiveid)
|
||||||
encrypted, err := ch.encryptMessage(plaintext, receiveid)
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("Failed to encrypt message: %v", err)
|
t.Fatalf("Failed to encrypt message: %v", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
if encrypted == "" {
|
if encrypted == "" {
|
||||||
t.Fatal("Encrypted message is empty")
|
t.Fatal("Encrypted message is empty")
|
||||||
}
|
}
|
||||||
|
|
||||||
// Decrypt
|
|
||||||
decrypted, err := decryptMessageWithVerify(encrypted, cfg.EncodingAESKey, receiveid)
|
decrypted, err := decryptMessageWithVerify(encrypted, cfg.EncodingAESKey, receiveid)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("Failed to decrypt message: %v", err)
|
t.Fatalf("Failed to decrypt message: %v", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
if decrypted != plaintext {
|
if decrypted != plaintext {
|
||||||
t.Errorf("Expected decrypted message '%s', got '%s'", plaintext, decrypted)
|
t.Errorf("Expected decrypted message '%s', got '%s'", plaintext, decrypted)
|
||||||
}
|
}
|
||||||
|
|
@ -198,13 +199,139 @@ func TestGenerateSignature(t *testing.T) {
|
||||||
encrypt := "encrypted_msg"
|
encrypt := "encrypted_msg"
|
||||||
|
|
||||||
signature := computeSignature(token, timestamp, nonce, encrypt)
|
signature := computeSignature(token, timestamp, nonce, encrypt)
|
||||||
|
|
||||||
if signature == "" {
|
if signature == "" {
|
||||||
t.Error("Generated signature is empty")
|
t.Error("Generated signature is empty")
|
||||||
}
|
}
|
||||||
|
|
||||||
// Verify signature using verifySignature function
|
|
||||||
if !verifySignature(token, signature, timestamp, nonce, encrypt) {
|
if !verifySignature(token, signature, timestamp, nonce, encrypt) {
|
||||||
t.Error("Generated signature does not verify correctly")
|
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)
|
||||||
|
}
|
||||||
|
|
@ -474,13 +474,24 @@ type WeComAppConfig struct {
|
||||||
}
|
}
|
||||||
|
|
||||||
type WeComAIBotConfig struct {
|
type WeComAIBotConfig struct {
|
||||||
Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_ENABLED"`
|
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"`
|
// WebSocket long-connection mode.
|
||||||
WebhookPath string `json:"webhook_path" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WEBHOOK_PATH"`
|
// 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"`
|
AllowFrom FlexibleStringSlice `json:"allow_from" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_ALLOW_FROM"`
|
||||||
ReplyTimeout int `json:"reply_timeout" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_REPLY_TIMEOUT"`
|
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
|
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"`
|
ReasoningChannelID string `json:"reasoning_channel_id" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_REASONING_CHANNEL_ID"`
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue