diff --git a/pkg/channels/wecom/aibot.go b/pkg/channels/wecom/aibot.go index 999f4f13b..2264b8492 100644 --- a/pkg/channels/wecom/aibot.go +++ b/pkg/channels/wecom/aibot.go @@ -158,6 +158,9 @@ func NewWeComAIBotChannel( "WeCom AI Bot requires either (bot_id + secret) for WebSocket 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, channels.WithMaxMessageLength(2048), @@ -709,7 +712,7 @@ func (c *WeComAIBotChannel) getStreamResponse(task *streamTask, timestamp, nonce default: if time.Now().After(task.Deadline) { // 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 closeStreamOnly = true logger.InfoCF( diff --git a/pkg/channels/wecom/aibot_test.go b/pkg/channels/wecom/aibot_test.go index 7c5ae67b1..01f3d3a41 100644 --- a/pkg/channels/wecom/aibot_test.go +++ b/pkg/channels/wecom/aibot_test.go @@ -2,6 +2,7 @@ package wecom import ( "context" + "encoding/json" "testing" "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) { cfg := config.WeComAIBotConfig{ 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 ---- func TestNewWeComAIBotChannel_WSMode(t *testing.T) { diff --git a/pkg/config/config.go b/pkg/config/config.go index d226bba51..a0b2eea81 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -247,7 +247,10 @@ type AgentDefaults struct { 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 { if d.MaxMediaSize > 0 { @@ -482,9 +485,10 @@ type WeComAIBotConfig struct { 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"` ReplyTimeout int `json:"reply_timeout" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_REPLY_TIMEOUT"` - MaxSteps int `json:"max_steps" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_MAX_STEPS"` - WelcomeMessage string `json:"welcome_message" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WELCOME_MESSAGE"` - ReasoningChannelID string `json:"reasoning_channel_id" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_REASONING_CHANNEL_ID"` + 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"` // Sent on enter_chat event; empty = no welcome + 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 { diff --git a/pkg/config/defaults.go b/pkg/config/defaults.go index 4038696c4..d44c73577 100644 --- a/pkg/config/defaults.go +++ b/pkg/config/defaults.go @@ -163,14 +163,15 @@ func DefaultConfig() *Config { ReplyTimeout: 5, }, WeComAIBot: WeComAIBotConfig{ - Enabled: false, - Token: "", - EncodingAESKey: "", - WebhookPath: "/webhook/wecom-aibot", - AllowFrom: FlexibleStringSlice{}, - ReplyTimeout: 5, - MaxSteps: 10, - WelcomeMessage: "Hello! I'm your AI assistant. How can I help you today?", + Enabled: false, + Token: "", + EncodingAESKey: "", + WebhookPath: "/webhook/wecom-aibot", + AllowFrom: FlexibleStringSlice{}, + ReplyTimeout: 5, + MaxSteps: 10, + WelcomeMessage: "Hello! I'm your AI assistant. How can I help you today?", + ProcessingMessage: DefaultWeComAIBotProcessingMessage, }, Pico: PicoConfig{ Enabled: false,