From 9073fd047b240791df5b4e18bb6be2adb0b152f2 Mon Sep 17 00:00:00 2001 From: Zhang Rui Date: Tue, 10 Mar 2026 13:53:40 +0800 Subject: [PATCH] feat(wecom): refactor WeCom AI Bot to manage request states and late replies --- pkg/channels/wecom/aibot_ws.go | 187 ++++++++++++++++++++++++++++----- 1 file changed, 162 insertions(+), 25 deletions(-) diff --git a/pkg/channels/wecom/aibot_ws.go b/pkg/channels/wecom/aibot_ws.go index 3384e3017..8e5ca76fa 100644 --- a/pkg/channels/wecom/aibot_ws.go +++ b/pkg/channels/wecom/aibot_ws.go @@ -31,6 +31,7 @@ const ( wsHeartbeatInterval = 30 * time.Second wsConnectTimeout = 15 * time.Second wsSubscribeTimeout = 10 * time.Second + wsSendMsgTimeout = 10 * time.Second wsMaxReconnectWait = 60 * time.Second wsInitialReconnect = time.Second @@ -42,6 +43,9 @@ const ( // wsImageDownloadTimeout caps the time we spend downloading an inbound image. wsImageDownloadTimeout = 30 * time.Second + + // Keep req_id -> chat route for late fallback pushes after stream window closes. + wsLateReplyRouteTTL = 30 * time.Minute ) // WeComAIBotWSChannel implements channels.Channel for WeCom AI Bot using the @@ -59,11 +63,10 @@ type WeComAIBotWSChannel struct { conn *websocket.Conn connMu sync.Mutex - // tasks holds one live agent task per req_id (WeCom turn identifier). - // Each inbound message gets a unique req_id, so concurrent messages are - // handled independently without canceling each other. - tasks map[string]*wsTask - tasksMu sync.Mutex + // reqStates holds per-req_id runtime state. + // It unifies active task state and late-reply fallback routing. + reqStates map[string]*wsReqState + reqStatesMu sync.Mutex // reqPending correlates command req_ids with response channels. // Used only for subscribe/ping command-response pairs. @@ -75,6 +78,7 @@ type WeComAIBotWSChannel struct { type wsTask struct { ReqID string // req_id echoed in all replies for this turn ChatID string + ChatType uint32 StreamID string // our generated stream.id CreatedTime time.Time answerCh chan string // agent delivers its reply here via Send() @@ -83,6 +87,18 @@ type wsTask struct { cancel context.CancelFunc } +type wsReqState struct { + Task *wsTask + Route wsLateReplyRoute +} + +type wsLateReplyRoute struct { + ChatID string + ChatType uint32 + ReadyAt time.Time + ExpiresAt time.Time +} + // ---- WebSocket protocol types ---- // wsEnvelope is the generic JSON envelope for all WebSocket messages. @@ -105,6 +121,13 @@ type wsCommand struct { Body any `json:"body,omitempty"` } +type wsSendMsgBody struct { + ChatID string `json:"chatid"` + ChatType uint32 `json:"chat_type,omitempty"` + MsgType string `json:"msgtype"` + Markdown *wsMarkdownContent `json:"markdown,omitempty"` +} + // wsRespondMsgBody is the body for aibot_respond_msg / aibot_respond_welcome_msg. type wsRespondMsgBody struct { MsgType string `json:"msgtype"` @@ -194,7 +217,7 @@ func newWeComAIBotWSChannel( return &WeComAIBotWSChannel{ BaseChannel: base, config: cfg, - tasks: make(map[string]*wsTask), + reqStates: make(map[string]*wsReqState), reqPending: make(map[string]chan wsEnvelope), }, nil } @@ -239,13 +262,35 @@ func (c *WeComAIBotWSChannel) Send(ctx context.Context, msg bus.OutboundMessage) } // msg.ChatID carries the inbound req_id (set by dispatchWSAgentTask). - c.tasksMu.Lock() - task := c.tasks[msg.ChatID] - c.tasksMu.Unlock() + task, route, ok := c.getReqState(msg.ChatID) + if !ok { + logger.DebugCF("wecom_aibot", "Send: no active task/route for req_id (may be stale)", + map[string]any{"req_id": msg.ChatID}) + return nil + } if task == nil { - logger.DebugCF("wecom_aibot", "Send: no active task for req_id (may have finished or timed out)", - map[string]any{"req_id": msg.ChatID}) + if !ok { + logger.DebugCF("wecom_aibot", "Send: no active task/route for req_id (may be stale)", + map[string]any{"req_id": msg.ChatID}) + return nil + } + if time.Now().Before(route.ReadyAt) { + // Keep using aibot_respond_msg within stream window; do not proactively + // push unless wsStreamMaxDuration has elapsed. + logger.DebugCF("wecom_aibot", "Send: stream window still open, skip proactive push", + map[string]any{"req_id": msg.ChatID, "ready_at": route.ReadyAt.Format(time.RFC3339)}) + return nil + } + + if err := c.wsSendActivePush(route.ChatID, route.ChatType, msg.Content); err != nil { + logger.WarnCF("wecom_aibot", "Late reply proactive push failed", + map[string]any{"req_id": msg.ChatID, "chat_id": route.ChatID, "error": err}) + return err + } + logger.InfoCF("wecom_aibot", "Late reply delivered via proactive push", + map[string]any{"req_id": msg.ChatID, "chat_id": route.ChatID, "chat_type": route.ChatType}) + c.deleteReqState(msg.ChatID) return nil } @@ -683,11 +728,13 @@ func (c *WeComAIBotWSChannel) dispatchWSAgentTask( } streamID := wsGenerateID() + chatType := wsChatTypeValue(msg.ChatType) taskCtx, taskCancel := context.WithCancel(c.ctx) task := &wsTask{ ReqID: reqID, ChatID: actualChatID, + ChatType: chatType, StreamID: streamID, CreatedTime: time.Now(), answerCh: make(chan string, 1), @@ -695,11 +742,16 @@ func (c *WeComAIBotWSChannel) dispatchWSAgentTask( ctx: taskCtx, cancel: taskCancel, } - - c.tasksMu.Lock() // Each req_id is unique per WeCom turn; tasks run concurrently, no cancellation. - c.tasks[reqID] = task - c.tasksMu.Unlock() + c.setReqState(reqID, &wsReqState{ + Task: task, + Route: wsLateReplyRoute{ + ChatID: actualChatID, + ChatType: chatType, + ReadyAt: time.Now().Add(wsStreamMaxDuration), + ExpiresAt: time.Now().Add(wsLateReplyRouteTTL), + }, + }) logger.DebugCF("wecom_aibot", "Registered new agent task", map[string]any{"chat_id": actualChatID, "req_id": reqID, "stream_id": streamID}) @@ -710,11 +762,7 @@ func (c *WeComAIBotWSChannel) dispatchWSAgentTask( go func() { defer func() { taskCancel() - c.tasksMu.Lock() - if c.tasks[reqID] == task { - delete(c.tasks, reqID) - } - c.tasksMu.Unlock() + c.clearReqTask(reqID, task) }() sender := bus.SenderInfo{ @@ -760,6 +808,7 @@ func (c *WeComAIBotWSChannel) dispatchWSAgentTask( case answer := <-task.answerCh: // send final frame with finish=true; any media will come in subsequent frames (if at all) c.wsSendStreamFinish(reqID, streamID, answer) + c.deleteReqState(reqID) return case <-ticker.C: hint := waitHints[tickCount%len(waitHints)] @@ -781,6 +830,7 @@ func (c *WeComAIBotWSChannel) dispatchWSAgentTask( select { case answer := <-task.answerCh: c.wsSendStreamFinish(reqID, streamID, answer) + c.deleteReqState(reqID) case <-time.After(100 * time.Millisecond): } return @@ -851,6 +901,40 @@ func (c *WeComAIBotWSChannel) wsSendWelcomeMsg(reqID, content string) { }) } +// wsSendActivePush sends a proactive markdown message using aibot_send_msg. +// It is used as a fallback for late replies after stream response window expires. +func (c *WeComAIBotWSChannel) wsSendActivePush(chatID string, chatType uint32, content string) error { + if chatID == "" { + return fmt.Errorf("chatid is empty") + } + reqID := wsGenerateID() + + c.connMu.Lock() + conn := c.conn + c.connMu.Unlock() + if conn == nil { + return fmt.Errorf("websocket not connected") + } + + resp, err := c.sendAndWait(conn, reqID, wsCommand{ + Cmd: "aibot_send_msg", + Headers: wsHeaders{ReqID: reqID}, + Body: wsSendMsgBody{ + ChatID: chatID, + ChatType: chatType, + MsgType: "markdown", + Markdown: &wsMarkdownContent{Content: content}, + }, + }, wsSendMsgTimeout) + if err != nil { + return err + } + if resp.ErrCode != 0 { + return fmt.Errorf("aibot_send_msg rejected (errcode=%d): %s", resp.ErrCode, resp.ErrMsg) + } + return nil +} + // 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) { @@ -877,14 +961,67 @@ func (c *WeComAIBotWSChannel) writeWS(cmd any) { // 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) + c.reqStatesMu.Lock() + defer c.reqStatesMu.Unlock() + for _, state := range c.reqStates { + if state != nil && state.Task != nil { + state.Task.cancel() + state.Task = nil + } } } +func (c *WeComAIBotWSChannel) setReqState(reqID string, state *wsReqState) { + c.reqStatesMu.Lock() + defer c.reqStatesMu.Unlock() + now := time.Now() + for k, v := range c.reqStates { + if v == nil || now.After(v.Route.ExpiresAt) { + delete(c.reqStates, k) + } + } + c.reqStates[reqID] = state +} + +func (c *WeComAIBotWSChannel) getReqState(reqID string) (*wsTask, wsLateReplyRoute, bool) { + c.reqStatesMu.Lock() + defer c.reqStatesMu.Unlock() + state, ok := c.reqStates[reqID] + if !ok || state == nil { + return nil, wsLateReplyRoute{}, false + } + if time.Now().After(state.Route.ExpiresAt) { + delete(c.reqStates, reqID) + return nil, wsLateReplyRoute{}, false + } + return state.Task, state.Route, true +} + +func (c *WeComAIBotWSChannel) deleteReqState(reqID string) { + c.reqStatesMu.Lock() + delete(c.reqStates, reqID) + c.reqStatesMu.Unlock() +} + +func (c *WeComAIBotWSChannel) clearReqTask(reqID string, task *wsTask) { + c.reqStatesMu.Lock() + defer c.reqStatesMu.Unlock() + state, ok := c.reqStates[reqID] + if !ok || state == nil { + return + } + if state.Task == task { + state.Task = nil + } +} + +func wsChatTypeValue(chatType string) uint32 { + if chatType == "group" { + return 2 + } + return 1 +} + // 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 {