feat(wecom): enhance WeCom AI Bot with image handling and media support

pkg/channels/wecom/aibot_ws.go

@@ -2,8 +2,15 @@ package wecom
 
 import (
 	"context"
+	"crypto/md5"
+	"encoding/base64"
 	"encoding/json"
 	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strings"
 	"sync"
 	"time"
 
@@ -14,6 +21,7 @@ import (
 	"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"
 )
 
@@ -32,6 +40,13 @@ const (
 	// 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
+
+	// wsMediaWaitTimeout is how long the goroutine waits for images (delivered via
+	// SendMedia → mediaCh) AFTER sending the text finish frame, before giving up.
+	wsMediaWaitTimeout = 500 * time.Millisecond
 )
 
 // WeComAIBotWSChannel implements channels.Channel for WeCom AI Bot using the
@@ -58,6 +73,23 @@ type WeComAIBotWSChannel struct {
 	// Used only for subscribe/ping command-response pairs.
 	reqPending   map[string]chan wsEnvelope
 	reqPendingMu sync.Mutex
+
+	// lastReqIDs records the most recent req_id per chat so that SendMedia can
+	// send images even after the stream has finished.
+	lastReqIDs   map[string]wsReqIDRecord
+	lastReqIDsMu sync.Mutex
+
+	// tasksByMsgID allows Send() to route a response to the exact task that
+	// originated from a given inbound message ID, even if a newer task has
+	// since replaced it in the tasks map (concurrent-message race prevention).
+	// Protected by tasksMu.
+	tasksByMsgID map[string]*wsTask
+}
+
+// wsReqIDRecord stores a req_id and its expiry time.
+type wsReqIDRecord struct {
+	ReqID     string
+	ExpiresAt time.Time
 }
 
 // wsTask tracks one in-progress agent reply for a single chat turn.
@@ -67,6 +99,7 @@ type wsTask struct {
 	StreamID    string // our generated stream.id
 	CreatedTime time.Time
 	answerCh    chan string          // agent delivers its reply here via Send()
+	mediaCh     chan []bus.MediaPart // agent's media attachments (buffered: 1)
 	ctx         context.Context
 	cancel      context.CancelFunc
 }
@@ -99,6 +132,7 @@ type wsRespondMsgBody struct {
 	Stream   *wsStreamContent   `json:"stream,omitempty"`
 	Text     *wsTextContent     `json:"text,omitempty"`
 	Markdown *wsMarkdownContent `json:"markdown,omitempty"`
+	Image    *wsImageContent    `json:"image,omitempty"`
 }
 
 type wsStreamContent struct {
@@ -107,6 +141,12 @@ type wsStreamContent struct {
 	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"`
 }
@@ -176,7 +216,9 @@ func newWeComAIBotWSChannel(
 		BaseChannel:  base,
 		config:       cfg,
 		tasks:        make(map[string]*wsTask),
+		tasksByMsgID: make(map[string]*wsTask),
 		reqPending:   make(map[string]chan wsEnvelope),
+		lastReqIDs:   make(map[string]wsReqIDRecord),
 	}, nil
 }
 
@@ -220,25 +262,82 @@ func (c *WeComAIBotWSChannel) Send(ctx context.Context, msg bus.OutboundMessage)
 	}
 
 	c.tasksMu.Lock()
-	task := c.tasks[msg.ChatID]
+	var task *wsTask
+	// Prefer exact per-message routing to avoid stale responses going to the
+	// wrong (newer) task when two messages arrive in rapid succession.
+	if msg.MessageID != "" {
+		task = c.tasksByMsgID[msg.MessageID]
+	}
+	if task == nil {
+		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})
+			map[string]any{"chat_id": msg.ChatID, "message_id": msg.MessageID})
 		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 // task canceled (connection dropped)
+		return nil
 	case <-ctx.Done():
 		return ctx.Err()
 	}
 	return nil
 }
 
+// SendMedia implements channels.MediaSender.
+// if there is an active task for this chat, media is sent to the task's mediaCh and will be delivered after the text reply.
+func (c *WeComAIBotWSChannel) SendMedia(ctx context.Context, msg bus.OutboundMediaMessage) error {
+	if !c.IsRunning() {
+		return channels.ErrNotRunning
+	}
+	c.tasksMu.Lock()
+	task := c.tasks[msg.ChatID]
+	c.tasksMu.Unlock()
+	logger.InfoCF("wecom_aibot", "SendMedia called",
+		map[string]any{"chat_id": msg.ChatID, "parts": len(msg.Parts), "has_task": task != nil})
+	if task != nil {
+		select {
+		case task.mediaCh <- msg.Parts:
+			return nil
+		case <-task.ctx.Done():
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+			logger.DebugCF("wecom_aibot", "SendMedia: mediaCh full or task done", map[string]any{"chat_id": msg.ChatID})
+		}
+	}
+	reqID := c.getLastReqID(msg.ChatID)
+	if reqID == "" {
+		logger.WarnCF("wecom_aibot", "SendMedia: no active req_id for chat, dropping media",
+			map[string]any{"chat_id": msg.ChatID})
+		return nil
+	}
+	store := c.GetMediaStore()
+	if store == nil {
+		return nil
+	}
+	for _, part := range msg.Parts {
+		if part.Type == "image" {
+			c.sendWSStandaloneImage(reqID, part, store)
+		}
+	}
+	return nil
+}
+
 // ---- Connection management ----
 
 // connectLoop maintains the WebSocket connection, reconnecting on failure with
@@ -547,7 +646,101 @@ func (c *WeComAIBotWSChannel) handleWSTextMessage(reqID string, msg WeComAIBotWS
 		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
+	}
+
+	chatID := msg.ChatID
+	if chatID == "" {
+		chatID = msg.From.UserID
+	}
+
+	ctx, cancel := context.WithTimeout(c.ctx, wsImageDownloadTimeout)
+	defer cancel()
+
+	var mediaRefs []string
+	ref, err := c.storeWSImage(ctx, chatID, msg.MsgID, msg.Image.URL, msg.Image.AESKey)
+	if err != nil {
+		logger.WarnCF("wecom_aibot", "Failed to download/store WS image",
+			map[string]any{"error": err, "url": msg.Image.URL})
+	} else {
+		mediaRefs = append(mediaRefs, ref)
+	}
+
+	c.dispatchWSAgentTask(reqID, msg, "[image]", mediaRefs)
+}
+
+// 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 := msg.ChatID
+	if chatID == "" {
+		chatID = msg.From.UserID
+	}
+
+	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.storeWSImage(ctx, chatID,
+					msg.MsgID+"-"+wsGenerateID(), item.Image.URL, item.Image.AESKey)
+				if err != nil {
+					logger.WarnCF("wecom_aibot", "Failed to download/store mixed image",
+						map[string]any{"error": err})
+				} else {
+					mediaRefs = append(mediaRefs, ref)
+				}
+			}
+		}
+	}
+
+	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"
@@ -566,6 +759,7 @@ func (c *WeComAIBotWSChannel) handleWSTextMessage(reqID string, msg WeComAIBotWS
 		StreamID:    streamID,
 		CreatedTime: time.Now(),
 		answerCh:    make(chan string, 1),
+		mediaCh:     make(chan []bus.MediaPart, 1),
 		ctx:         taskCtx,
 		cancel:      taskCancel,
 	}
@@ -576,10 +770,16 @@ func (c *WeComAIBotWSChannel) handleWSTextMessage(reqID string, msg WeComAIBotWS
 		prev.cancel()
 	}
 	c.tasks[chatID] = task
+	// Also register by inbound message ID for precise response routing.
+	if msg.MsgID != "" {
+		c.tasksByMsgID[msg.MsgID] = task
+	}
 	c.tasksMu.Unlock()
 
-	// Send an empty stream opening frame (finish=false) immediately so WeCom
+	// Record this reqID so SendMedia can route images even after the stream closes.
-	// shows the typing indicator while the agent is processing.
+	c.recordLastReqID(chatID, reqID)
+
+	// Send an empty stream opening frame (finish=false) immediately.
 	c.wsSendStreamChunk(reqID, streamID, false, "")
 
 	go func() {
@@ -589,6 +789,9 @@ func (c *WeComAIBotWSChannel) handleWSTextMessage(reqID string, msg WeComAIBotWS
 			if c.tasks[chatID] == task {
 				delete(c.tasks, chatID)
 			}
+			if msg.MsgID != "" && c.tasksByMsgID[msg.MsgID] == task {
+				delete(c.tasksByMsgID, msg.MsgID)
+			}
 			c.tasksMu.Unlock()
 		}()
 
@@ -606,14 +809,13 @@ func (c *WeComAIBotWSChannel) handleWSTextMessage(reqID string, msg WeComAIBotWS
 		metadata := map[string]string{
 			"channel":   "wecom_aibot",
 			"chat_type": msg.ChatType,
-			"msg_type":  "text",
+			"msg_type":  msg.MsgType,
 			"msgid":     msg.MsgID,
 			"aibotid":   msg.AIBotID,
 			"stream_id": streamID,
 		}
-		// PublishInbound is non-blocking; the agent will call Send() when done.
 		c.HandleMessage(taskCtx, peer, msg.MsgID, userID, chatID,
-			msg.Text.Content, nil, metadata, sender)
+			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.
@@ -632,72 +834,54 @@ func (c *WeComAIBotWSChannel) handleWSTextMessage(reqID string, msg WeComAIBotWS
 		for {
 			select {
 			case answer := <-task.answerCh:
-				// Agent replied — deliver the real answer as the final frame.
+				// 1. send final frame with finish=true; any media will come in subsequent frames (if at all)
 				c.wsSendStreamFinish(reqID, streamID, answer)
+				// 2. wait briefly for any media to arrive via SendMedia → mediaCh,
+				// and send each image in its own frame (WeCom does not support multiple images in one frame,
+				// and the agent may have sent them at different times)
+				var mediaParts []bus.MediaPart
+				select {
+				case mediaParts = <-task.mediaCh:
+					logger.InfoCF("wecom_aibot", "Sending queued media images",
+						map[string]any{"chat_id": chatID, "count": len(mediaParts)})
+				case <-time.After(wsMediaWaitTimeout):
+				}
+				if store := c.GetMediaStore(); store != nil {
+					for _, part := range mediaParts {
+						if part.Type == "image" {
+							c.sendWSStandaloneImage(reqID, part, store)
+						}
+					}
+				}
 				return
 			case <-ticker.C:
-				// Send an in-progress hint (finish=false). WeCom replaces the
-				// displayed content with the latest non-finish frame, so the
-				// user just sees the most recent hint, not an accumulation.
 				hint := waitHints[tickCount%len(waitHints)]
 				tickCount++
 				logger.DebugCF("wecom_aibot", "Sending stream progress hint",
 					map[string]any{"chat_id": chatID, "tick": tickCount})
 				c.wsSendStreamChunk(reqID, streamID, false, hint)
 			case <-deadlineTimer.C:
-				// Hard deadline reached before the agent replied. Close the
-				// stream with a timeout notice (still within the 6-minute window).
 				logger.WarnCF("wecom_aibot", "Stream deadline reached without agent reply",
 					map[string]any{"chat_id": chatID, "stream_id": streamID})
 				c.wsSendStreamFinish(reqID, streamID,
 					"⏳ Processing is taking longer than expected. Please resend your message to try again.")
 				return
 			case <-taskCtx.Done():
-				// Connection dropped or task canceled; nothing to send.
+				// 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.
+				select {
+				case answer := <-task.answerCh:
+					c.wsSendStreamFinish(reqID, streamID, answer)
+				case <-time.After(100 * time.Millisecond):
+				}
 				return
 			}
 		}
 	}()
 }
 
-// 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.
@@ -742,7 +926,7 @@ func (c *WeComAIBotWSChannel) wsSendStreamChunk(reqID, streamID string, finish b
 	})
 }
 
-// wsSendStreamFinish sends the final aibot_respond_msg frame (finish=true).
+// 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)
 }
@@ -799,3 +983,136 @@ func (c *WeComAIBotWSChannel) cancelAllTasks() {
 func wsGenerateID() string {
 	return generateRandomID(10)
 }
+
+// ---- req_id tracking (for SendMedia after stream close) ----
+
+// recordLastReqID stores the req_id for chatID with a TTL slightly beyond the
+// stream max duration so that SendMedia can find it after finish=true is sent.
+func (c *WeComAIBotWSChannel) recordLastReqID(chatID, reqID string) {
+	c.lastReqIDsMu.Lock()
+	c.lastReqIDs[chatID] = wsReqIDRecord{
+		ReqID:     reqID,
+		ExpiresAt: time.Now().Add(wsStreamMaxDuration + 2*time.Minute),
+	}
+	c.lastReqIDsMu.Unlock()
+}
+
+// getLastReqID returns the most recent req_id for chatID, or "" if none/expired.
+func (c *WeComAIBotWSChannel) getLastReqID(chatID string) string {
+	c.lastReqIDsMu.Lock()
+	r, ok := c.lastReqIDs[chatID]
+	c.lastReqIDsMu.Unlock()
+	if !ok || time.Now().After(r.ExpiresAt) {
+		return ""
+	}
+	return r.ReqID
+}
+
+// ---- Image send helpers ----
+
+// sendWSStandaloneImage reads a media part from store and sends it as a
+// standalone aibot_respond_msg with msgtype "image".
+func (c *WeComAIBotWSChannel) sendWSStandaloneImage(reqID string, part bus.MediaPart, store media.MediaStore) {
+	localPath, err := store.Resolve(part.Ref)
+	if err != nil {
+		logger.ErrorCF("wecom_aibot", "SendMedia: resolve failed", map[string]any{"ref": part.Ref, "error": err})
+		return
+	}
+	data, err := os.ReadFile(localPath)
+	if err != nil {
+		logger.ErrorCF("wecom_aibot", "SendMedia: read file failed", map[string]any{"path": localPath, "error": err})
+		return
+	}
+	hash := md5.Sum(data)
+	logger.InfoCF("wecom_aibot", "Sending standalone image",
+		map[string]any{"req_id": reqID, "path": localPath, "bytes": len(data)})
+	c.writeWS(wsCommand{
+		Cmd:     "aibot_respond_msg",
+		Headers: wsHeaders{ReqID: reqID},
+		Body: wsRespondMsgBody{
+			MsgType: "image",
+			Image: &wsImageContent{
+				Base64: base64.StdEncoding.EncodeToString(data),
+				MD5:    fmt.Sprintf("%x", hash),
+			},
+		},
+	})
+}
+
+// ---- Inbound image download helpers ----
+
+// downloadWSImage fetches and optionally decrypts a WeCom WS image resource.
+// aesKey is the per-resource AES key provided in the callback (may be empty).
+func (c *WeComAIBotWSChannel) downloadWSImage(ctx context.Context, imageURL, aesKey string) ([]byte, error) {
+	const maxSize = 20 << 20 // 20 MB
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, imageURL, nil)
+	if err != nil {
+		return nil, fmt.Errorf("create request: %w", err)
+	}
+	client := &http.Client{Timeout: wsImageDownloadTimeout}
+	resp, err := client.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("download: %w", err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("download HTTP %d", resp.StatusCode)
+	}
+	data, err := io.ReadAll(io.LimitReader(resp.Body, maxSize+1))
+	if err != nil {
+		return nil, fmt.Errorf("read body: %w", err)
+	}
+	if len(data) > maxSize {
+		return nil, fmt.Errorf("image too large (> %d MB)", maxSize>>20)
+	}
+	if aesKey == "" {
+		return data, nil
+	}
+	// WeCom per-image AES key: try standard base64 first, then WeCom 43-char format.
+	key, decErr := base64.StdEncoding.DecodeString(aesKey)
+	if decErr != nil || len(key) != 32 {
+		key, decErr = decodeWeComAESKey(aesKey)
+		if decErr != nil {
+			return nil, fmt.Errorf("decode image AES key: %w", decErr)
+		}
+	}
+	decrypted, err := decryptAESCBC(key, data)
+	if err != nil {
+		return nil, fmt.Errorf("decrypt image: %w", err)
+	}
+	return decrypted, nil
+}
+
+// storeWSImage downloads, optionally decrypts, and stores an inbound image.
+func (c *WeComAIBotWSChannel) storeWSImage(
+	ctx context.Context,
+	chatID, msgID, imageURL, aesKey string,
+) (string, error) {
+	store := c.GetMediaStore()
+	if store == nil {
+		return "", fmt.Errorf("no media store available")
+	}
+	data, err := c.downloadWSImage(ctx, imageURL, aesKey)
+	if err != nil {
+		return "", err
+	}
+	mediaDir := filepath.Join(os.TempDir(), "picoclaw_media")
+	if mkErr := os.MkdirAll(mediaDir, 0o700); mkErr != nil {
+		return "", fmt.Errorf("mkdir: %w", mkErr)
+	}
+	filename := msgID + ".jpg"
+	localPath := filepath.Join(mediaDir, utils.SanitizeFilename(filename))
+	if writeErr := os.WriteFile(localPath, data, 0o600); writeErr != nil {
+		return "", fmt.Errorf("write: %w", writeErr)
+	}
+	scope := channels.BuildMediaScope("wecom_aibot", chatID, msgID)
+	ref, err := store.Store(localPath, media.MediaMeta{
+		Filename: filename,
+		Source:   "wecom_aibot",
+	}, scope)
+	if err != nil {
+		os.Remove(localPath)
+		return "", fmt.Errorf("store: %w", err)
+	}
+	return ref, nil
+}