feat(wecom_aibot): make processing message configurable

This commit is contained in:
Hoshina 2026-03-19 19:34:54 +08:00
parent 9a25fad20a
commit d301445734
4 changed files with 116 additions and 13 deletions

View file

@ -158,6 +158,9 @@ func NewWeComAIBotChannel(
"WeCom AI Bot requires either (bot_id + secret) for WebSocket mode " + "WeCom AI Bot requires either (bot_id + secret) for WebSocket mode " +
"or (token + encoding_aes_key) for webhook mode") "or (token + encoding_aes_key) for webhook mode")
} }
if cfg.ProcessingMessage == "" {
cfg.ProcessingMessage = config.DefaultWeComAIBotProcessingMessage
}
base := channels.NewBaseChannel("wecom_aibot", cfg, messageBus, cfg.AllowFrom, base := channels.NewBaseChannel("wecom_aibot", cfg, messageBus, cfg.AllowFrom,
channels.WithMaxMessageLength(2048), channels.WithMaxMessageLength(2048),
@ -709,7 +712,7 @@ func (c *WeComAIBotChannel) getStreamResponse(task *streamTask, timestamp, nonce
default: default:
if time.Now().After(task.Deadline) { if time.Now().After(task.Deadline) {
// Deadline reached: close the stream with a notice, then wait for agent via response_url. // Deadline reached: close the stream with a notice, then wait for agent via response_url.
content = "⏳ Processing, please wait. The results will be sent shortly." content = c.config.ProcessingMessage
finish = true finish = true
closeStreamOnly = true closeStreamOnly = true
logger.InfoCF( logger.InfoCF(

View file

@ -2,6 +2,7 @@ package wecom
import ( import (
"context" "context"
"encoding/json"
"testing" "testing"
"time" "time"
@ -134,6 +135,79 @@ func TestWeComAIBotChannelWebhookPath(t *testing.T) {
}) })
} }
func TestWeComAIBotChannelGetStreamResponseProcessingMessage(t *testing.T) {
validAESKey := "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFG"
t.Run("uses default processing message", func(t *testing.T) {
cfg := config.WeComAIBotConfig{
Enabled: true,
Token: "test_token",
EncodingAESKey: validAESKey,
}
messageBus := bus.NewMessageBus()
ch, err := NewWeComAIBotChannel(cfg, messageBus)
if err != nil {
t.Fatalf("Failed to create channel: %v", err)
}
task := &streamTask{
StreamID: "stream-default",
ChatID: "chat-default",
Deadline: time.Now().Add(-time.Second),
}
ch.streamTasks[task.StreamID] = task
ch.chatTasks[task.ChatID] = []*streamTask{task}
resp := decodeStreamResponse(t, ch, ch.getStreamResponse(task, "1234567890", "nonce"))
if !resp.Stream.Finish {
t.Fatal("Expected finished stream response after deadline")
}
if resp.Stream.Content != config.DefaultWeComAIBotProcessingMessage {
t.Fatalf("Expected default processing message %q, got %q",
config.DefaultWeComAIBotProcessingMessage, resp.Stream.Content)
}
if !task.StreamClosed {
t.Fatal("Expected task stream to be marked closed")
}
if _, ok := ch.streamTasks[task.StreamID]; ok {
t.Fatal("Expected closed stream task to be removed from streamTasks")
}
if len(ch.chatTasks[task.ChatID]) != 1 {
t.Fatalf("Expected task to remain queued for response_url delivery, got %d entries",
len(ch.chatTasks[task.ChatID]))
}
})
t.Run("uses custom processing message", func(t *testing.T) {
cfg := config.WeComAIBotConfig{
Enabled: true,
Token: "test_token",
EncodingAESKey: validAESKey,
ProcessingMessage: "Please wait a moment. The result will be delivered in a follow-up message.",
}
messageBus := bus.NewMessageBus()
ch, err := NewWeComAIBotChannel(cfg, messageBus)
if err != nil {
t.Fatalf("Failed to create channel: %v", err)
}
task := &streamTask{
StreamID: "stream-custom",
ChatID: "chat-custom",
Deadline: time.Now().Add(-time.Second),
}
resp := decodeStreamResponse(t, ch, ch.getStreamResponse(task, "1234567890", "nonce"))
if resp.Stream.Content != cfg.ProcessingMessage {
t.Fatalf("Expected custom processing message %q, got %q", cfg.ProcessingMessage, resp.Stream.Content)
}
})
}
func TestGenerateStreamID(t *testing.T) { func TestGenerateStreamID(t *testing.T) {
cfg := config.WeComAIBotConfig{ cfg := config.WeComAIBotConfig{
Enabled: true, Enabled: true,
@ -208,6 +282,27 @@ func TestGenerateSignature(t *testing.T) {
} }
} }
func decodeStreamResponse(t *testing.T, ch *WeComAIBotChannel, encryptedResponse string) WeComAIBotStreamResponse {
t.Helper()
var wrapped WeComAIBotEncryptedResponse
if err := json.Unmarshal([]byte(encryptedResponse), &wrapped); err != nil {
t.Fatalf("Failed to unmarshal encrypted response: %v", err)
}
plaintext, err := decryptMessageWithVerify(wrapped.Encrypt, ch.config.EncodingAESKey, "")
if err != nil {
t.Fatalf("Failed to decrypt response: %v", err)
}
var resp WeComAIBotStreamResponse
if err := json.Unmarshal([]byte(plaintext), &resp); err != nil {
t.Fatalf("Failed to unmarshal decrypted response: %v", err)
}
return resp
}
// ---- WebSocket long-connection mode tests ---- // ---- WebSocket long-connection mode tests ----
func TestNewWeComAIBotChannel_WSMode(t *testing.T) { func TestNewWeComAIBotChannel_WSMode(t *testing.T) {

View file

@ -247,7 +247,10 @@ type AgentDefaults struct {
ToolFeedback ToolFeedbackConfig `json:"tool_feedback,omitempty"` ToolFeedback ToolFeedbackConfig `json:"tool_feedback,omitempty"`
} }
const DefaultMaxMediaSize = 20 * 1024 * 1024 // 20 MB const (
DefaultMaxMediaSize = 20 * 1024 * 1024 // 20 MB
DefaultWeComAIBotProcessingMessage = "⏳ Processing, please wait. The results will be sent shortly."
)
func (d *AgentDefaults) GetMaxMediaSize() int { func (d *AgentDefaults) GetMaxMediaSize() int {
if d.MaxMediaSize > 0 { if d.MaxMediaSize > 0 {
@ -482,9 +485,10 @@ type WeComAIBotConfig struct {
WebhookPath string `json:"webhook_path,omitempty" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WEBHOOK_PATH"` WebhookPath string `json:"webhook_path,omitempty" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WEBHOOK_PATH"`
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"` MaxSteps int `json:"max_steps" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_MAX_STEPS"` // Maximum streaming steps
WelcomeMessage string `json:"welcome_message" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WELCOME_MESSAGE"` 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"` ProcessingMessage string `json:"processing_message,omitempty" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_PROCESSING_MESSAGE"`
ReasoningChannelID string `json:"reasoning_channel_id" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_REASONING_CHANNEL_ID"`
} }
type PicoConfig struct { type PicoConfig struct {

View file

@ -163,14 +163,15 @@ func DefaultConfig() *Config {
ReplyTimeout: 5, ReplyTimeout: 5,
}, },
WeComAIBot: WeComAIBotConfig{ WeComAIBot: WeComAIBotConfig{
Enabled: false, Enabled: false,
Token: "", Token: "",
EncodingAESKey: "", EncodingAESKey: "",
WebhookPath: "/webhook/wecom-aibot", WebhookPath: "/webhook/wecom-aibot",
AllowFrom: FlexibleStringSlice{}, AllowFrom: FlexibleStringSlice{},
ReplyTimeout: 5, ReplyTimeout: 5,
MaxSteps: 10, MaxSteps: 10,
WelcomeMessage: "Hello! I'm your AI assistant. How can I help you today?", WelcomeMessage: "Hello! I'm your AI assistant. How can I help you today?",
ProcessingMessage: DefaultWeComAIBotProcessingMessage,
}, },
Pico: PicoConfig{ Pico: PicoConfig{
Enabled: false, Enabled: false,