package wecom import ( "context" "encoding/base64" "encoding/json" "fmt" "io" "net/http" "os" "path/filepath" "strings" "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/media" "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 wsSendMsgTimeout = 10 * time.Second wsRespondMsgTimeout = 10 * time.Second wsWelcomeMsgTimeout = 5 * time.Second // WeCom requires welcome reply within 5 seconds wsMaxReconnectWait = 60 * time.Second wsInitialReconnect = time.Second // WeCom requires finish=true within 6 minutes of the first stream frame. // wsStreamTickInterval controls how often we send an in-progress hint. // wsStreamMaxDuration is a safety margin below the 6-minute hard limit. wsStreamTickInterval = 30 * time.Second wsStreamMaxDuration = 5*time.Minute + 30*time.Second // 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 // wsStreamMaxContentBytes is the maximum UTF-8 byte length for the content field // of a single WeCom AI Bot stream / text / markdown frame. // Ref: https://developer.work.weixin.qq.com/document/path/101463 wsStreamMaxContentBytes = 20480 ) // wsImageHTTPClient is a shared HTTP client for downloading inbound images. // Reusing it enables connection pooling across multiple image downloads. var wsImageHTTPClient = &http.Client{Timeout: wsImageDownloadTimeout} // 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 // dedupe prevents duplicate message processing (WeCom may re-deliver). dedupe *MessageDeduplicator // 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. 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 ChatType uint32 StreamID string // our generated stream.id answerCh chan string // agent delivers its reply here via Send() ctx context.Context 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. 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"` } 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"` Stream *wsStreamContent `json:"stream,omitempty"` Text *wsTextContent `json:"text,omitempty"` Markdown *wsMarkdownContent `json:"markdown,omitempty"` Image *wsImageContent `json:"image,omitempty"` } type wsStreamContent struct { ID string `json:"id"` Finish bool `json:"finish"` Content string `json:"content,omitempty"` } // wsImageContent carries a base64-encoded image payload for outbound messages. type wsImageContent struct { Base64 string `json:"base64"` MD5 string `json:"md5"` } 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 { Content string `json:"content"` // 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"` File *struct { URL string `json:"url"` AESKey string `json:"aeskey,omitempty"` } `json:"file,omitempty"` Video *struct { URL string `json:"url"` AESKey string `json:"aeskey,omitempty"` } `json:"video,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.WithReasoningChannelID(cfg.ReasoningChannelID), ) return &WeComAIBotWSChannel{ BaseChannel: base, config: cfg, dedupe: NewMessageDeduplicator(wecomMaxProcessedMessages), reqStates: make(map[string]*wsReqState), 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 } // msg.ChatID carries the inbound req_id (set by dispatchWSAgentTask). // For cron-triggered messages, msg.ChatID is the real WeCom chat/user ID // and there will be no matching entry in reqStates; fall through to proactive push. task, route, ok := c.getReqState(msg.ChatID) if !ok { // No req_id record found — this is a cron/scheduler-originated message. // Send it as a proactive markdown push using the chat ID directly. logger.InfoCF("wecom_aibot", "Send: no req_id state, delivering via proactive push (cron/scheduler)", map[string]any{"chat_id": msg.ChatID}) if err := c.wsSendActivePush(msg.ChatID, 0, msg.Content); err != nil { logger.WarnCF("wecom_aibot", "Proactive push failed", map[string]any{"chat_id": msg.ChatID, "error": err.Error()}) return fmt.Errorf("websocket delivery failed: %w", channels.ErrSendFailed) } return nil } if task == nil { if time.Now().Before(route.ReadyAt) { // Keep using aibot_respond_msg within stream window; do not proactively // push unless wsStreamMaxDuration has elapsed. logger.WarnCF("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.Error()}) return fmt.Errorf("websocket delivery failed: %w", channels.ErrSendFailed) } 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 } // Non-blocking fast path: when answerCh has space, deliver without racing // against task.ctx.Done() (which fires when the task is canceled by a new // incoming message, but the response must still be sent). select { case task.answerCh <- msg.Content: return nil default: } // answerCh was full; block with cancellation guards. select { case task.answerCh <- msg.Content: case <-task.ctx.Done(): return nil case <-ctx.Done(): return ctx.Err() } return nil } // ---- Connection management ---- // wsBackoffResetDuration is the minimum duration a WebSocket connection must // stay up before we reset the reconnect backoff to its initial value. This // prevents a short burst of failures from causing long waits after later, // stable connection periods. const wsBackoffResetDuration = time.Minute // 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...") start := time.Now() if err := c.runConnection(); err != nil { elapsed := time.Since(start) // If the connection was stable for long enough, reset backoff so that // a previous burst of failures does not keep us at the maximum delay. if elapsed >= wsBackoffResetDuration { backoff = wsInitialReconnect } select { case <-c.ctx.Done(): return default: logger.WarnCF("wecom_aibot", "WebSocket connection lost, reconnecting", map[string]any{"error": err.Error(), "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. // It validates the server's pong response via sendAndWait; a failed pong // triggers a reconnection by closing the connection. func (c *WeComAIBotWSChannel) heartbeatLoop(conn *websocket.Conn) { ticker := time.NewTicker(wsHeartbeatInterval) defer ticker.Stop() for { select { case <-ticker.C: reqID := wsGenerateID() resp, err := c.sendAndWait(conn, reqID, wsCommand{ Cmd: "ping", Headers: wsHeaders{ReqID: reqID}, }, wsHeartbeatInterval) if err != nil { logger.WarnCF("wecom_aibot", "Heartbeat failed, closing connection", map[string]any{"error": err.Error()}) conn.Close() return } if resp.ErrCode != 0 { logger.WarnCF("wecom_aibot", "Heartbeat rejected", map[string]any{"errcode": resp.ErrCode, "errmsg": resp.ErrMsg}) conn.Close() return } logger.DebugCF("wecom_aibot", "Heartbeat pong received", 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.Error(), "raw": string(raw)}) continue } // Command responses have an empty Cmd field; forward to any waiting // sendAndWait() call, or silently drop if no one is waiting (e.g. // late responses after timeout). 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.Error()}) return } // Deduplicate by msgid (WeCom may re-deliver on network issues). if msg.MsgID != "" && !c.dedupe.MarkMessageProcessed(msg.MsgID) { logger.DebugCF("wecom_aibot", "Duplicate message ignored", map[string]any{"msgid": msg.MsgID}) 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) case "file": c.handleWSFileMessage(reqID, msg) case "video": c.handleWSVideoMessage(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.Error()}) return } // Deduplicate by msgid. if msg.MsgID != "" && !c.dedupe.MarkMessageProcessed(msg.MsgID) { logger.DebugCF("wecom_aibot", "Duplicate event ignored", map[string]any{"msgid": msg.MsgID}) 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 } c.dispatchWSAgentTask(reqID, msg, msg.Text.Content, nil) } // handleWSImageMessage downloads and stores the inbound image, then dispatches // it to the agent as a media-tagged message. func (c *WeComAIBotWSChannel) handleWSImageMessage(reqID string, msg WeComAIBotWSMessage) { if msg.Image == nil { logger.WarnC("wecom_aibot", "Image message missing image field") c.wsSendStreamFinish(reqID, wsGenerateID(), "Image message could not be processed.") return } c.wsHandleMediaMessage(reqID, msg, msg.Image.URL, msg.Image.AESKey, "image") } // wsHandleMediaMessage is a shared helper for image, file and video messages. // It downloads the resource, stores it in MediaStore, and dispatches to the agent. func (c *WeComAIBotWSChannel) wsHandleMediaMessage( reqID string, msg WeComAIBotWSMessage, resourceURL, aesKey, label string, ) { chatID := wsChatID(msg) ctx, cancel := context.WithTimeout(c.ctx, wsImageDownloadTimeout) defer cancel() ref, err := c.storeWSMedia(ctx, chatID, msg.MsgID, resourceURL, aesKey, wsLabelToDefaultExt(label)) if err != nil { logger.WarnCF("wecom_aibot", "Failed to download/store WS "+label, map[string]any{"error": err.Error(), "url": resourceURL}) c.wsSendStreamFinish(reqID, wsGenerateID(), strings.ToUpper(label[:1])+label[1:]+" message could not be processed.") return } c.dispatchWSAgentTask(reqID, msg, "["+label+"]", []string{ref}) } // handleWSMixedMessage handles mixed text+image messages. // All text parts are collected into the content string; all image parts are // downloaded and stored in MediaStore before dispatching to the agent. func (c *WeComAIBotWSChannel) handleWSMixedMessage(reqID string, msg WeComAIBotWSMessage) { if msg.Mixed == nil { logger.WarnC("wecom_aibot", "Mixed message has no content") c.wsSendStreamFinish(reqID, wsGenerateID(), "Mixed message type is not yet fully supported.") return } chatID := wsChatID(msg) ctx, cancel := context.WithTimeout(c.ctx, wsImageDownloadTimeout) defer cancel() var textParts []string var mediaRefs []string for _, item := range msg.Mixed.MsgItem { switch item.MsgType { case "text": if item.Text != nil && item.Text.Content != "" { textParts = append(textParts, item.Text.Content) } case "image": if item.Image != nil { ref, err := c.storeWSMedia(ctx, chatID, msg.MsgID+"-"+wsGenerateID(), item.Image.URL, item.Image.AESKey, ".jpg") if err != nil { logger.WarnCF("wecom_aibot", "Failed to download/store mixed image", map[string]any{"error": err.Error()}) } else { mediaRefs = append(mediaRefs, ref) } } default: logger.WarnCF("wecom_aibot", "Unsupported item type in mixed message", map[string]any{"msgtype": item.MsgType}) } } if len(textParts) == 0 && len(mediaRefs) == 0 { logger.WarnC("wecom_aibot", "Mixed message has no usable content") c.wsSendStreamFinish(reqID, wsGenerateID(), "Mixed message type is not yet fully supported.") return } content := strings.Join(textParts, "\n") if content == "" { content = "[images]" } c.dispatchWSAgentTask(reqID, msg, content, mediaRefs) } // dispatchWSAgentTask registers a new agent task, sends the opening stream frame, // and starts a goroutine that runs the agent and streams the reply back. // content is the text forwarded to the agent; mediaRefs are optional media // store references attached to the inbound message. func (c *WeComAIBotWSChannel) dispatchWSAgentTask( reqID string, msg WeComAIBotWSMessage, content string, mediaRefs []string, ) { userID := msg.From.UserID if userID == "" { userID = "unknown" } // actualChatID is the real WeCom chat/user ID used for peer identification. // reqID is used as the routing chatID so each turn is independently addressable. actualChatID := wsChatID(msg) streamID := wsGenerateID() chatType := wsChatTypeValue(msg.ChatType) taskCtx, taskCancel := context.WithCancel(c.ctx) task := &wsTask{ ReqID: reqID, ChatID: actualChatID, ChatType: chatType, StreamID: streamID, answerCh: make(chan string, 1), ctx: taskCtx, cancel: taskCancel, } // Each req_id is unique per WeCom turn; tasks run concurrently, no cancellation. 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}) // Send an empty stream opening frame (finish=false) immediately. c.wsSendStreamChunk(reqID, streamID, false, "") go func() { defer func() { taskCancel() c.clearReqTask(reqID, task) }() 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: actualChatID} metadata := map[string]string{ "channel": "wecom_aibot", "chat_id": actualChatID, "chat_type": msg.ChatType, "msg_type": msg.MsgType, "msgid": msg.MsgID, "aibotid": msg.AIBotID, "stream_id": streamID, } // Pass reqID as chatID: OutboundMessage.ChatID = reqID → Send() finds tasks[reqID]. c.HandleMessage(taskCtx, peer, reqID, userID, reqID, content, mediaRefs, metadata, sender) // Wait for the agent reply. While waiting, send periodic finish=false // hints so the user knows processing is still in progress. // WeCom requires finish=true within 6 minutes of the first stream frame; // wsStreamMaxDuration enforces that limit with a safety margin. waitHints := []string{ "⏳ Processing, please wait...", "⏳ Still processing, please wait...", "⏳ Almost there, please wait...", } ticker := time.NewTicker(wsStreamTickInterval) defer ticker.Stop() deadlineTimer := time.NewTimer(wsStreamMaxDuration) defer deadlineTimer.Stop() tickCount := 0 for { select { case answer := <-task.answerCh: // Split the answer into byte-bounded chunks and send as stream frames. // All but the last carry finish=false; the final frame closes the stream. chunks := splitWSContent(answer, wsStreamMaxContentBytes) for i, chunk := range chunks { c.wsSendStreamChunk(reqID, streamID, i == len(chunks)-1, chunk) } c.deleteReqState(reqID) return case <-ticker.C: hint := waitHints[tickCount%len(waitHints)] tickCount++ logger.DebugCF("wecom_aibot", "Sending stream progress hint", map[string]any{"chat_id": actualChatID, "tick": tickCount}) c.wsSendStreamChunk(reqID, streamID, false, hint) case <-deadlineTimer.C: logger.WarnCF("wecom_aibot", "Stream response deadline reached, closing stream; late reply will be pushed", map[string]any{"chat_id": actualChatID}) c.wsSendStreamFinish(reqID, streamID, "⏳ Processing is taking longer than expected, the response will be sent as a follow-up message.") return case <-taskCtx.Done(): // Give a short grace period so that a response queued in the bus // just before cancellation can still be delivered. This closes a // race where a rapid second message cancels this task after the // agent already published but before Send() wrote to answerCh. // // The connection is gone at this point, so we cannot use // wsSendStreamFinish. Try wsSendActivePush on the (possibly // already-restored) connection; if that also fails, leave the // route intact so Send() can push the reply once reconnected. select { case answer := <-task.answerCh: if err := c.wsSendActivePush(task.ChatID, task.ChatType, answer); err != nil { logger.WarnCF("wecom_aibot", "Grace-period push failed after task cancellation; reply may be lost", map[string]any{"req_id": reqID, "chat_id": task.ChatID, "error": err.Error()}) } else { c.deleteReqState(reqID) } case <-time.After(100 * time.Millisecond): } return } } }() } // handleWSVoiceMessage handles voice messages. // WeCom transcribes voice to text in the callback; if the transcription is // present it is dispatched as plain text to the agent. func (c *WeComAIBotWSChannel) handleWSVoiceMessage(reqID string, msg WeComAIBotWSMessage) { if msg.Voice != nil && msg.Voice.Content != "" { c.dispatchWSAgentTask(reqID, msg, msg.Voice.Content, nil) return } c.wsSendStreamFinish(reqID, wsGenerateID(), "Voice messages are not yet supported.") } // handleWSFileMessage handles file messages. func (c *WeComAIBotWSChannel) handleWSFileMessage(reqID string, msg WeComAIBotWSMessage) { if msg.File == nil { logger.WarnC("wecom_aibot", "File message missing file field") c.wsSendStreamFinish(reqID, wsGenerateID(), "File message could not be processed.") return } c.wsHandleMediaMessage(reqID, msg, msg.File.URL, msg.File.AESKey, "file") } // handleWSVideoMessage handles video messages. func (c *WeComAIBotWSChannel) handleWSVideoMessage(reqID string, msg WeComAIBotWSMessage) { if msg.Video == nil { logger.WarnC("wecom_aibot", "Video message missing video field") c.wsSendStreamFinish(reqID, wsGenerateID(), "Video message could not be processed.") return } c.wsHandleMediaMessage(reqID, msg, msg.Video.URL, msg.Video.AESKey, "video") } // ---- 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), }) cmd := wsCommand{ Cmd: "aibot_respond_msg", Headers: wsHeaders{ReqID: reqID}, Body: wsRespondMsgBody{ MsgType: "stream", Stream: &wsStreamContent{ ID: streamID, Finish: finish, Content: content, }, }, } if err := c.writeWSAndWait(cmd, wsRespondMsgTimeout); err != nil { logger.WarnCF("wecom_aibot", "Stream chunk ack failed", map[string]any{ "req_id": reqID, "stream_id": streamID, "finish": finish, "error": err, }) } } // wsSendStreamFinish sends the final aibot_respond_msg frame (finish=true, no images). 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}) cmd := wsCommand{ Cmd: "aibot_respond_welcome_msg", Headers: wsHeaders{ReqID: reqID}, Body: wsRespondMsgBody{ MsgType: "text", Text: &wsTextContent{Content: content}, }, } if err := c.writeWSAndWait(cmd, wsWelcomeMsgTimeout); err != nil { logger.WarnCF("wecom_aibot", "Welcome message ack failed", map[string]any{"req_id": reqID, "error": err.Error()}) } } // wsSendActivePush sends a proactive markdown message using aibot_send_msg. // Long content is automatically split into byte-bounded chunks (≤ wsStreamMaxContentBytes // each) and delivered as consecutive messages. // 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") } for _, chunk := range splitWSContent(content, wsStreamMaxContentBytes) { reqID := wsGenerateID() if err := c.writeWSAndWait(wsCommand{ Cmd: "aibot_send_msg", Headers: wsHeaders{ReqID: reqID}, Body: wsSendMsgBody{ ChatID: chatID, ChatType: chatType, MsgType: "markdown", Markdown: &wsMarkdownContent{Content: chunk}, }, }, wsSendMsgTimeout); err != nil { return err } } return nil } // writeWSAndWait writes cmd to the active connection and validates the command response. func (c *WeComAIBotWSChannel) writeWSAndWait(cmd wsCommand, timeout time.Duration) error { if cmd.Headers.ReqID == "" { return fmt.Errorf("req_id is empty") } c.connMu.Lock() conn := c.conn c.connMu.Unlock() if conn == nil { return fmt.Errorf("websocket not connected") } resp, err := c.sendAndWait(conn, cmd.Headers.ReqID, cmd, timeout) if err != nil { return err } if resp.ErrCode != 0 { return fmt.Errorf("%s rejected (errcode=%d): %s", cmd.Cmd, resp.ErrCode, resp.ErrMsg) } return nil } // cancelAllTasks cancels every pending agent task; called when the connection drops. // It also expires each task's stream window (ReadyAt = now) so that when the agent // eventually delivers its reply via Send(), the message is forwarded via // wsSendActivePush on the restored connection instead of being silently discarded. func (c *WeComAIBotWSChannel) cancelAllTasks() { c.reqStatesMu.Lock() defer c.reqStatesMu.Unlock() now := time.Now() for _, state := range c.reqStates { if state != nil && state.Task != nil { state.Task.cancel() state.Task = nil // Expire the stream window immediately so Send() uses wsSendActivePush. state.Route.ReadyAt = now } } } 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 } // wsChatID returns the effective chat ID from a WS message. // For group messages it is msg.ChatID; for single chats it falls back to the sender's UserID. func wsChatID(msg WeComAIBotWSMessage) string { if msg.ChatID != "" { return msg.ChatID } return msg.From.UserID } // 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) } // ---- Inbound media download helpers ---- // storeWSMedia downloads the resource at resourceURL (with optional AES-CBC // decryption) and stores it in the MediaStore. The file extension is inferred // from the HTTP Content-Type response header; defaultExt is used as a fallback // when the content type is absent or unrecognized. func (c *WeComAIBotWSChannel) storeWSMedia( ctx context.Context, chatID, msgID, resourceURL, aesKey, defaultExt string, ) (string, error) { store := c.GetMediaStore() if store == nil { return "", fmt.Errorf("no media store available") } const maxSize = 20 << 20 // 20 MB req, err := http.NewRequestWithContext(ctx, http.MethodGet, resourceURL, nil) if err != nil { return "", fmt.Errorf("create request: %w", err) } resp, err := wsImageHTTPClient.Do(req) if err != nil { return "", fmt.Errorf("download: %w", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { return "", fmt.Errorf("download HTTP %d", resp.StatusCode) } // Infer file extension from the Content-Type response header. ext := wsMediaExtFromContentType(resp.Header.Get("Content-Type")) if ext == "" { ext = defaultExt } // Buffer the media in memory, bounded to maxSize. data, err := io.ReadAll(io.LimitReader(resp.Body, int64(maxSize)+1)) if err != nil { return "", fmt.Errorf("read media: %w", err) } if len(data) > maxSize { return "", fmt.Errorf("media too large (> %d MB)", maxSize>>20) } // AES-CBC decryption if a key is present. if aesKey != "" { key, decErr := base64.StdEncoding.DecodeString(aesKey) if decErr != nil || len(key) != 32 { key, decErr = decodeWeComAESKey(aesKey) if decErr != nil { return "", fmt.Errorf("decode media AES key: %w", decErr) } } data, err = decryptAESCBC(key, data) if err != nil { return "", fmt.Errorf("decrypt media: %w", err) } } // Write to a temp file. The file is owned by the MediaStore and deleted by // store.ReleaseAll — no caller-side cleanup needed. mediaDir := filepath.Join(os.TempDir(), "picoclaw_media") if err = os.MkdirAll(mediaDir, 0o700); err != nil { return "", fmt.Errorf("mkdir: %w", err) } tmpFile, err := os.CreateTemp(mediaDir, msgID+"-*"+ext) if err != nil { return "", fmt.Errorf("create temp file: %w", err) } tmpPath := tmpFile.Name() _, writeErr := tmpFile.Write(data) closeErr := tmpFile.Close() if writeErr != nil { os.Remove(tmpPath) return "", fmt.Errorf("write media: %w", writeErr) } if closeErr != nil { os.Remove(tmpPath) return "", fmt.Errorf("close media: %w", closeErr) } scope := channels.BuildMediaScope("wecom_aibot", chatID, msgID) ref, err := store.Store(tmpPath, media.MediaMeta{ Filename: msgID + ext, Source: "wecom_aibot", }, scope) if err != nil { os.Remove(tmpPath) return "", fmt.Errorf("store: %w", err) } return ref, nil } // wsMediaExtFromContentType returns the lowercase file extension (with leading // dot) for the given Content-Type value, or "" when the type is unrecognized. func wsMediaExtFromContentType(contentType string) string { if contentType == "" { return "" } // Strip parameters (e.g. "image/jpeg; charset=utf-8" → "image/jpeg"). mt := strings.ToLower(strings.TrimSpace(strings.SplitN(contentType, ";", 2)[0])) switch mt { case "image/jpeg", "image/jpg": return ".jpg" case "image/png": return ".png" case "image/gif": return ".gif" case "image/webp": return ".webp" case "video/mp4": return ".mp4" case "video/mpeg", "video/x-mpeg": return ".mpeg" case "video/quicktime": return ".mov" case "video/webm": return ".webm" case "audio/mpeg", "audio/mp3": return ".mp3" case "audio/ogg": return ".ogg" case "audio/wav": return ".wav" case "application/pdf": return ".pdf" case "application/zip": return ".zip" case "application/x-rar-compressed", "application/vnd.rar": return ".rar" case "text/plain": return ".txt" case "application/msword": return ".doc" case "application/vnd.openxmlformats-officedocument.wordprocessingml.document": return ".docx" case "application/vnd.ms-excel": return ".xls" case "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": return ".xlsx" case "application/vnd.ms-powerpoint": return ".ppt" case "application/vnd.openxmlformats-officedocument.presentationml.presentation": return ".pptx" } return "" } // wsLabelToDefaultExt returns the default file extension for the given media label // used in wsHandleMediaMessage. It is the fallback when Content-Type detection fails. func wsLabelToDefaultExt(label string) string { switch label { case "image": return ".jpg" case "video": return ".mp4" default: // "file" and any future labels return ".bin" } } // ---- Content length helpers ---- // splitWSContent splits content into chunks each fitting within maxBytes UTF-8 // bytes, preserving code block integrity via channels.SplitMessage. // When SplitMessage still produces an oversized chunk (e.g. dense CJK content), // splitAtByteBoundary is applied as a last-resort byte-level fallback. func splitWSContent(content string, maxBytes int) []string { if len(content) <= maxBytes { return []string{content} } // SplitMessage works in runes. Use maxBytes as the rune limit: for pure ASCII // this is exact; for multibyte content the byte verification below catches // any chunk that still overflows. chunks := channels.SplitMessage(content, maxBytes) var result []string for _, chunk := range chunks { if len(chunk) <= maxBytes { result = append(result, chunk) } else { // Still too large in bytes (e.g. dense CJK); force-split at UTF-8 boundaries. result = append(result, splitAtByteBoundary(chunk, maxBytes)...) } } return result } // splitAtByteBoundary splits s into parts each ≤ maxBytes bytes by walking back // from the hard byte limit to find a valid UTF-8 rune start boundary. // This is a last-resort fallback; it does not try to preserve code blocks. func splitAtByteBoundary(s string, maxBytes int) []string { var parts []string for len(s) > maxBytes { end := maxBytes // Walk back past any UTF-8 continuation bytes (high two bits == 10). for end > 0 && s[end]>>6 == 0b10 { end-- } if end == 0 { end = maxBytes // shouldn't happen with valid UTF-8 } parts = append(parts, s[:end]) s = strings.TrimLeft(s[end:], " \t\n\r") } if s != "" { parts = append(parts, s) } return parts }