feat(upload): implement WeCom media upload via WebSocket API
This commit is contained in:
parent
9a25fad20a
commit
e46ef82aac
6 changed files with 550 additions and 42 deletions
|
|
@ -193,6 +193,7 @@ func registerSharedTools(
|
||||||
Channel: channel,
|
Channel: channel,
|
||||||
ChatID: chatID,
|
ChatID: chatID,
|
||||||
Content: content,
|
Content: content,
|
||||||
|
Final: true,
|
||||||
})
|
})
|
||||||
})
|
})
|
||||||
agent.Tools.Register(messageTool)
|
agent.Tools.Register(messageTool)
|
||||||
|
|
@ -315,6 +316,7 @@ func (al *AgentLoop) Run(ctx context.Context) error {
|
||||||
Channel: msg.Channel,
|
Channel: msg.Channel,
|
||||||
ChatID: msg.ChatID,
|
ChatID: msg.ChatID,
|
||||||
Content: response,
|
Content: response,
|
||||||
|
Final: true,
|
||||||
})
|
})
|
||||||
logger.InfoCF("agent", "Published outbound response",
|
logger.InfoCF("agent", "Published outbound response",
|
||||||
map[string]any{
|
map[string]any{
|
||||||
|
|
|
||||||
|
|
@ -64,7 +64,7 @@ func (al *AgentLoop) ensureMCPInitialized(ctx context.Context) error {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
if al.cfg.Tools.MCP.Servers == nil || len(al.cfg.Tools.MCP.Servers) == 0 {
|
if len(al.cfg.Tools.MCP.Servers) == 0 {
|
||||||
logger.WarnCF("agent", "MCP is enabled but no servers are configured, skipping MCP initialization", nil)
|
logger.WarnCF("agent", "MCP is enabled but no servers are configured, skipping MCP initialization", nil)
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -34,6 +34,7 @@ type OutboundMessage struct {
|
||||||
ChatID string `json:"chat_id"`
|
ChatID string `json:"chat_id"`
|
||||||
Content string `json:"content"`
|
Content string `json:"content"`
|
||||||
ReplyToMessageID string `json:"reply_to_message_id,omitempty"`
|
ReplyToMessageID string `json:"reply_to_message_id,omitempty"`
|
||||||
|
Final bool `json:"final,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// MediaPart describes a single media attachment to send.
|
// MediaPart describes a single media attachment to send.
|
||||||
|
|
|
||||||
|
|
@ -38,9 +38,7 @@ const (
|
||||||
wsInitialReconnect = time.Second
|
wsInitialReconnect = time.Second
|
||||||
|
|
||||||
// WeCom requires finish=true within 6 minutes of the first stream frame.
|
// 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.
|
// wsStreamMaxDuration is a safety margin below the 6-minute hard limit.
|
||||||
wsStreamTickInterval = 30 * time.Second
|
|
||||||
wsStreamMaxDuration = 5*time.Minute + 30*time.Second
|
wsStreamMaxDuration = 5*time.Minute + 30*time.Second
|
||||||
|
|
||||||
// wsImageDownloadTimeout caps the time we spend downloading an inbound image.
|
// wsImageDownloadTimeout caps the time we spend downloading an inbound image.
|
||||||
|
|
@ -286,6 +284,11 @@ func (c *WeComAIBotWSChannel) Send(ctx context.Context, msg bus.OutboundMessage)
|
||||||
// and there will be no matching entry in reqStates; fall through to proactive push.
|
// and there will be no matching entry in reqStates; fall through to proactive push.
|
||||||
task, route, ok := c.getReqState(msg.ChatID)
|
task, route, ok := c.getReqState(msg.ChatID)
|
||||||
if !ok {
|
if !ok {
|
||||||
|
if !msg.Final {
|
||||||
|
// Intermediate message (tool feedback, etc.) with no active WS session.
|
||||||
|
// Nothing to stream to; silently drop it.
|
||||||
|
return nil
|
||||||
|
}
|
||||||
// No req_id record found — this is a cron/scheduler-originated message.
|
// No req_id record found — this is a cron/scheduler-originated message.
|
||||||
// Send it as a proactive markdown push using the chat ID directly.
|
// 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)",
|
logger.InfoCF("wecom_aibot", "Send: no req_id state, delivering via proactive push (cron/scheduler)",
|
||||||
|
|
@ -299,6 +302,12 @@ func (c *WeComAIBotWSChannel) Send(ctx context.Context, msg bus.OutboundMessage)
|
||||||
}
|
}
|
||||||
|
|
||||||
if task == nil {
|
if task == nil {
|
||||||
|
if !msg.Final {
|
||||||
|
// Intermediate message: stream for this req_id has already closed (deadline
|
||||||
|
// or new-message cancel). Nothing to do.
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
if time.Now().Before(route.ReadyAt) {
|
if time.Now().Before(route.ReadyAt) {
|
||||||
// Keep using aibot_respond_msg within stream window; do not proactively
|
// Keep using aibot_respond_msg within stream window; do not proactively
|
||||||
// push unless wsStreamMaxDuration has elapsed.
|
// push unless wsStreamMaxDuration has elapsed.
|
||||||
|
|
@ -318,6 +327,19 @@ func (c *WeComAIBotWSChannel) Send(ctx context.Context, msg bus.OutboundMessage)
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if !msg.Final {
|
||||||
|
// Intermediate message (tool feedback, context notices, etc.) with an active
|
||||||
|
// stream. Send directly as a non-finishing stream chunk so the user sees
|
||||||
|
// live progress without touching answerCh (which the goroutine reserves for
|
||||||
|
// the single final answer).
|
||||||
|
for _, chunk := range splitWSContent(msg.Content, wsStreamMaxContentBytes) {
|
||||||
|
c.wsSendStreamChunk(task.ReqID, task.StreamID, false, chunk)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Final answer: deliver to the waiting goroutine via answerCh, which will
|
||||||
|
// send it with finish=true and clean up the stream state.
|
||||||
// Non-blocking fast path: when answerCh has space, deliver without racing
|
// 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
|
// against task.ctx.Done() (which fires when the task is canceled by a new
|
||||||
// incoming message, but the response must still be sent).
|
// incoming message, but the response must still be sent).
|
||||||
|
|
@ -852,37 +874,34 @@ func (c *WeComAIBotWSChannel) dispatchWSAgentTask(
|
||||||
c.HandleMessage(taskCtx, peer, reqID, userID, reqID,
|
c.HandleMessage(taskCtx, peer, reqID, userID, reqID,
|
||||||
content, mediaRefs, metadata, sender)
|
content, mediaRefs, metadata, sender)
|
||||||
|
|
||||||
// Wait for the agent reply. While waiting, send periodic finish=false
|
// Wait for the agent reply.
|
||||||
// hints so the user knows processing is still in progress.
|
|
||||||
// WeCom requires finish=true within 6 minutes of the first stream frame;
|
// WeCom requires finish=true within 6 minutes of the first stream frame;
|
||||||
// wsStreamMaxDuration enforces that limit with a safety margin.
|
// 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)
|
deadlineTimer := time.NewTimer(wsStreamMaxDuration)
|
||||||
defer deadlineTimer.Stop()
|
defer deadlineTimer.Stop()
|
||||||
tickCount := 0
|
|
||||||
for {
|
for {
|
||||||
select {
|
select {
|
||||||
case answer := <-task.answerCh:
|
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.
|
// Split the answer into byte-bounded chunks.
|
||||||
|
// Each chunk must fit within wsStreamMaxContentBytes (20480 bytes).
|
||||||
|
// WeCom's stream update semantics REPLACE content (not append), so all
|
||||||
|
// chunks beyond the first need a fresh stream.id — otherwise only the
|
||||||
|
// last chunk would be visible to the user.
|
||||||
|
//
|
||||||
|
// Delivery plan:
|
||||||
|
// chunk[0] — close the already-opened stream (finish=true, same streamID)
|
||||||
|
// chunk[1+] — create + immediately close a new stream per chunk
|
||||||
chunks := splitWSContent(answer, wsStreamMaxContentBytes)
|
chunks := splitWSContent(answer, wsStreamMaxContentBytes)
|
||||||
for i, chunk := range chunks {
|
for i, chunk := range chunks {
|
||||||
c.wsSendStreamChunk(reqID, streamID, i == len(chunks)-1, chunk)
|
sid := streamID
|
||||||
|
if i > 0 {
|
||||||
|
sid = wsGenerateID()
|
||||||
|
}
|
||||||
|
c.wsSendStreamChunk(reqID, sid, true, chunk)
|
||||||
}
|
}
|
||||||
c.deleteReqState(reqID)
|
c.deleteReqState(reqID)
|
||||||
return
|
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:
|
case <-deadlineTimer.C:
|
||||||
logger.WarnCF("wecom_aibot",
|
logger.WarnCF("wecom_aibot",
|
||||||
"Stream response deadline reached, closing stream; late reply will be pushed",
|
"Stream response deadline reached, closing stream; late reply will be pushed",
|
||||||
|
|
@ -1028,27 +1047,11 @@ func (c *WeComAIBotWSChannel) wsSendActivePush(chatID string, chatType uint32, c
|
||||||
}
|
}
|
||||||
|
|
||||||
// writeWSAndWait writes cmd to the active connection and validates the command response.
|
// writeWSAndWait writes cmd to the active connection and validates the command response.
|
||||||
|
// It is a thin wrapper around callWSCommand that discards the response body.
|
||||||
func (c *WeComAIBotWSChannel) writeWSAndWait(cmd wsCommand, timeout time.Duration) error {
|
func (c *WeComAIBotWSChannel) writeWSAndWait(cmd wsCommand, timeout time.Duration) error {
|
||||||
if cmd.Headers.ReqID == "" {
|
_, err := c.callWSCommand(cmd, timeout)
|
||||||
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
|
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.
|
// 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
|
// It also expires each task's stream window (ReadyAt = now) so that when the agent
|
||||||
|
|
|
||||||
234
pkg/channels/wecom/aibot_ws_upload.go
Normal file
234
pkg/channels/wecom/aibot_ws_upload.go
Normal file
|
|
@ -0,0 +1,234 @@
|
||||||
|
package wecom
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"crypto/md5"
|
||||||
|
"encoding/base64"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Upload temporary media constants.
|
||||||
|
// Ref: https://developer.work.weixin.qq.com/document/path/101463#上传临时素材
|
||||||
|
const (
|
||||||
|
// wsUploadChunkSize is the maximum raw (pre-base64) bytes per chunk.
|
||||||
|
// The API limit is 512 KB before base64 encoding.
|
||||||
|
wsUploadChunkSize = 512 << 10 // 512 KB
|
||||||
|
|
||||||
|
// wsUploadMaxChunks is the maximum number of chunks per upload session.
|
||||||
|
wsUploadMaxChunks = 100
|
||||||
|
|
||||||
|
wsUploadInitTimeout = 30 * time.Second
|
||||||
|
wsUploadChunkTimeout = 60 * time.Second // larger: base64 payload may be big
|
||||||
|
wsUploadFinishTimeout = 30 * time.Second
|
||||||
|
)
|
||||||
|
|
||||||
|
// ---- Request / response body types ----
|
||||||
|
|
||||||
|
// wsUploadInitBody is the body for aibot_upload_media_init.
|
||||||
|
type wsUploadInitBody struct {
|
||||||
|
Type string `json:"type"`
|
||||||
|
Filename string `json:"filename"`
|
||||||
|
TotalSize int64 `json:"total_size"`
|
||||||
|
TotalChunks int `json:"total_chunks"`
|
||||||
|
MD5 string `json:"md5,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// wsUploadInitResponse is the body received in the aibot_upload_media_init response.
|
||||||
|
type wsUploadInitResponse struct {
|
||||||
|
UploadID string `json:"upload_id"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// wsUploadChunkBody is the body for aibot_upload_media_chunk.
|
||||||
|
type wsUploadChunkBody struct {
|
||||||
|
UploadID string `json:"upload_id"`
|
||||||
|
ChunkIndex int `json:"chunk_index"`
|
||||||
|
Base64Data string `json:"base64_data"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// wsUploadFinishBody is the body for aibot_upload_media_finish.
|
||||||
|
type wsUploadFinishBody struct {
|
||||||
|
UploadID string `json:"upload_id"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// wsUploadFinishResponse is the body received in the aibot_upload_media_finish response.
|
||||||
|
type wsUploadFinishResponse struct {
|
||||||
|
Type string `json:"type"`
|
||||||
|
MediaID string `json:"media_id"`
|
||||||
|
CreatedAt string `json:"created_at"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- Public API ----
|
||||||
|
|
||||||
|
// UploadWSMedia uploads a local file as a temporary WeCom media asset via the
|
||||||
|
// WebSocket long-connection API and returns the resulting media_id (valid 3 days).
|
||||||
|
//
|
||||||
|
// mediaType must be one of:
|
||||||
|
// - "image" — PNG / JPG(JPEG) / GIF, ≤ 2 MB
|
||||||
|
// - "voice" — AMR, ≤ 2 MB
|
||||||
|
// - "video" — MP4, ≤ 10 MB
|
||||||
|
// - "file" — any format, ≤ 20 MB
|
||||||
|
//
|
||||||
|
// filename is the logical file name sent to WeCom. When empty it is derived
|
||||||
|
// from the base name of filePath.
|
||||||
|
//
|
||||||
|
// The upload is split into at most wsUploadMaxChunks chunks of wsUploadChunkSize
|
||||||
|
// bytes each. The full operation must complete within wsUploadSessionTTL (30 min).
|
||||||
|
// Each individual command call is also independently time-bounded.
|
||||||
|
func (c *WeComAIBotWSChannel) UploadWSMedia(
|
||||||
|
ctx context.Context,
|
||||||
|
filePath, filename, mediaType string,
|
||||||
|
) (string, error) {
|
||||||
|
// ---- Validate media type ----
|
||||||
|
switch mediaType {
|
||||||
|
case "image", "voice", "video", "file":
|
||||||
|
default:
|
||||||
|
return "", fmt.Errorf("unsupported media type %q: must be image, voice, video, or file", mediaType)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- Read file ----
|
||||||
|
data, err := os.ReadFile(filePath)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("read file: %w", err)
|
||||||
|
}
|
||||||
|
totalSize := int64(len(data))
|
||||||
|
if totalSize == 0 {
|
||||||
|
return "", fmt.Errorf("file is empty: %s", filePath)
|
||||||
|
}
|
||||||
|
|
||||||
|
if filename == "" {
|
||||||
|
filename = filepath.Base(filePath)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- Compute chunk count ----
|
||||||
|
totalChunks := int((totalSize + wsUploadChunkSize - 1) / wsUploadChunkSize)
|
||||||
|
if totalChunks > wsUploadMaxChunks {
|
||||||
|
return "", fmt.Errorf(
|
||||||
|
"file too large: requires %d chunks but limit is %d (max ~%d MB)",
|
||||||
|
totalChunks, wsUploadMaxChunks,
|
||||||
|
int64(wsUploadMaxChunks)*wsUploadChunkSize>>20,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- Compute MD5 ----
|
||||||
|
rawMD5 := md5.Sum(data) //nolint:gosec // MD5 is required by the WeCom API spec
|
||||||
|
fileMD5 := fmt.Sprintf("%x", rawMD5)
|
||||||
|
|
||||||
|
// ---- Step 1: Init ----
|
||||||
|
initEnv, err := c.callWSCommand(wsCommand{
|
||||||
|
Cmd: "aibot_upload_media_init",
|
||||||
|
Headers: wsHeaders{ReqID: wsGenerateID()},
|
||||||
|
Body: wsUploadInitBody{
|
||||||
|
Type: mediaType,
|
||||||
|
Filename: filename,
|
||||||
|
TotalSize: totalSize,
|
||||||
|
TotalChunks: totalChunks,
|
||||||
|
MD5: fileMD5,
|
||||||
|
},
|
||||||
|
}, wsUploadInitTimeout)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("upload init: %w", err)
|
||||||
|
}
|
||||||
|
var initResp wsUploadInitResponse
|
||||||
|
if err = json.Unmarshal(initEnv.Body, &initResp); err != nil {
|
||||||
|
return "", fmt.Errorf("parse upload init response: %w", err)
|
||||||
|
}
|
||||||
|
if initResp.UploadID == "" {
|
||||||
|
return "", fmt.Errorf("upload init returned empty upload_id")
|
||||||
|
}
|
||||||
|
uploadID := initResp.UploadID
|
||||||
|
|
||||||
|
logger.InfoCF("wecom_aibot", "Media upload initialized", map[string]any{
|
||||||
|
"upload_id": uploadID,
|
||||||
|
"type": mediaType,
|
||||||
|
"filename": filename,
|
||||||
|
"total_size": totalSize,
|
||||||
|
"total_chunks": totalChunks,
|
||||||
|
})
|
||||||
|
|
||||||
|
// ---- Step 2: Upload chunks ----
|
||||||
|
for i := 0; i < totalChunks; i++ {
|
||||||
|
if ctxErr := ctx.Err(); ctxErr != nil {
|
||||||
|
return "", fmt.Errorf("upload aborted before chunk %d: %w", i, ctxErr)
|
||||||
|
}
|
||||||
|
|
||||||
|
start := int64(i) * wsUploadChunkSize
|
||||||
|
end := start + wsUploadChunkSize
|
||||||
|
if end > totalSize {
|
||||||
|
end = totalSize
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err = c.callWSCommand(wsCommand{
|
||||||
|
Cmd: "aibot_upload_media_chunk",
|
||||||
|
Headers: wsHeaders{ReqID: wsGenerateID()},
|
||||||
|
Body: wsUploadChunkBody{
|
||||||
|
UploadID: uploadID,
|
||||||
|
ChunkIndex: i,
|
||||||
|
Base64Data: base64.StdEncoding.EncodeToString(data[start:end]),
|
||||||
|
},
|
||||||
|
}, wsUploadChunkTimeout)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("upload chunk %d/%d: %w", i, totalChunks-1, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.DebugCF("wecom_aibot", "Media chunk uploaded", map[string]any{
|
||||||
|
"upload_id": uploadID,
|
||||||
|
"chunk": i,
|
||||||
|
"total": totalChunks,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- Step 3: Finish ----
|
||||||
|
finishEnv, err := c.callWSCommand(wsCommand{
|
||||||
|
Cmd: "aibot_upload_media_finish",
|
||||||
|
Headers: wsHeaders{ReqID: wsGenerateID()},
|
||||||
|
Body: wsUploadFinishBody{UploadID: uploadID},
|
||||||
|
}, wsUploadFinishTimeout)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("upload finish: %w", err)
|
||||||
|
}
|
||||||
|
var finishResp wsUploadFinishResponse
|
||||||
|
if err := json.Unmarshal(finishEnv.Body, &finishResp); err != nil {
|
||||||
|
return "", fmt.Errorf("parse upload finish response: %w", err)
|
||||||
|
}
|
||||||
|
if finishResp.MediaID == "" {
|
||||||
|
return "", fmt.Errorf("upload finish returned empty media_id")
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.InfoCF("wecom_aibot", "Media upload complete", map[string]any{
|
||||||
|
"upload_id": uploadID,
|
||||||
|
"media_id": finishResp.MediaID,
|
||||||
|
"type": finishResp.Type,
|
||||||
|
})
|
||||||
|
return finishResp.MediaID, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- Internal helper ----
|
||||||
|
|
||||||
|
// callWSCommand sends a WebSocket command and returns the raw server envelope.
|
||||||
|
// It validates the errcode field and returns an error when non-zero.
|
||||||
|
// Use callWSCommand (over writeWSAndWait) when the response body must be inspected.
|
||||||
|
func (c *WeComAIBotWSChannel) callWSCommand(cmd wsCommand, timeout time.Duration) (wsEnvelope, error) {
|
||||||
|
if cmd.Headers.ReqID == "" {
|
||||||
|
return wsEnvelope{}, fmt.Errorf("req_id is empty")
|
||||||
|
}
|
||||||
|
c.connMu.Lock()
|
||||||
|
conn := c.conn
|
||||||
|
c.connMu.Unlock()
|
||||||
|
if conn == nil {
|
||||||
|
return wsEnvelope{}, fmt.Errorf("websocket not connected")
|
||||||
|
}
|
||||||
|
env, err := c.sendAndWait(conn, cmd.Headers.ReqID, cmd, timeout)
|
||||||
|
if err != nil {
|
||||||
|
return wsEnvelope{}, err
|
||||||
|
}
|
||||||
|
if env.ErrCode != 0 {
|
||||||
|
return wsEnvelope{}, fmt.Errorf("%s rejected (errcode=%d): %s", cmd.Cmd, env.ErrCode, env.ErrMsg)
|
||||||
|
}
|
||||||
|
return env, nil
|
||||||
|
}
|
||||||
268
pkg/channels/wecom/aibot_ws_upload_test.go
Normal file
268
pkg/channels/wecom/aibot_ws_upload_test.go
Normal file
|
|
@ -0,0 +1,268 @@
|
||||||
|
package wecom
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/base64"
|
||||||
|
"encoding/json"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// UploadWSMedia — pre-flight validation tests (no real WS required)
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
// TestUploadWSMedia_InvalidMediaType verifies that unsupported media types are
|
||||||
|
// rejected before any network activity occurs.
|
||||||
|
func TestUploadWSMedia_InvalidMediaType(t *testing.T) {
|
||||||
|
ch := newTestWSChannel(t)
|
||||||
|
|
||||||
|
for _, bad := range []string{"", "pdf", "audio", "IMAGE", "gif"} {
|
||||||
|
_, err := ch.UploadWSMedia(context.Background(), "/dev/null", "test.pdf", bad)
|
||||||
|
if err == nil {
|
||||||
|
t.Errorf("expected error for media type %q, got nil", bad)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestUploadWSMedia_FileNotFound verifies that a non-existent path returns an error.
|
||||||
|
func TestUploadWSMedia_FileNotFound(t *testing.T) {
|
||||||
|
ch := newTestWSChannel(t)
|
||||||
|
_, err := ch.UploadWSMedia(context.Background(), "/nonexistent/path/to/file.bin", "", "file")
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for missing file, got nil")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestUploadWSMedia_EmptyFile verifies that an empty file is rejected before
|
||||||
|
// the upload starts.
|
||||||
|
func TestUploadWSMedia_EmptyFile(t *testing.T) {
|
||||||
|
ch := newTestWSChannel(t)
|
||||||
|
|
||||||
|
tmp, err := os.CreateTemp(t.TempDir(), "empty-*.bin")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
tmp.Close()
|
||||||
|
|
||||||
|
_, err = ch.UploadWSMedia(context.Background(), tmp.Name(), "", "file")
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for empty file, got nil")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestUploadWSMedia_FileTooLarge verifies that a file requiring more than
|
||||||
|
// wsUploadMaxChunks chunks is rejected before the upload starts.
|
||||||
|
func TestUploadWSMedia_FileTooLarge(t *testing.T) {
|
||||||
|
ch := newTestWSChannel(t)
|
||||||
|
|
||||||
|
// Synthesize a file that would need 101 chunks (one byte over the limit).
|
||||||
|
size := int64(wsUploadMaxChunks)*int64(wsUploadChunkSize) + 1
|
||||||
|
|
||||||
|
tmp, err := os.CreateTemp(t.TempDir(), "large-*.bin")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
truncErr := tmp.Truncate(size)
|
||||||
|
tmp.Close()
|
||||||
|
if truncErr != nil {
|
||||||
|
t.Fatal(truncErr)
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err = ch.UploadWSMedia(context.Background(), tmp.Name(), "", "file")
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for oversized file, got nil")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestUploadWSMedia_NotConnected verifies that calling UploadWSMedia when no
|
||||||
|
// WebSocket connection is active returns a clear error.
|
||||||
|
func TestUploadWSMedia_NotConnected(t *testing.T) {
|
||||||
|
ch := newTestWSChannel(t)
|
||||||
|
|
||||||
|
tmp := filepath.Join(t.TempDir(), "test.txt")
|
||||||
|
if err := os.WriteFile(tmp, []byte("hello world"), 0o600); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// conn is nil by default in a test channel — callWSCommand must return an error.
|
||||||
|
_, err := ch.UploadWSMedia(context.Background(), tmp, "", "file")
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error when conn is nil, got nil")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestUploadWSMedia_ContextCanceledPreFlight verifies that a pre-canceled
|
||||||
|
// context is detected during chunk iteration (even before any WS call, when
|
||||||
|
// the connection is nil the WS error takes precedence for a single-chunk file,
|
||||||
|
// so we test this path with a file that triggers ctx.Err() first).
|
||||||
|
func TestUploadWSMedia_ContextCanceledPreFlight(t *testing.T) {
|
||||||
|
ch := newTestWSChannel(t)
|
||||||
|
|
||||||
|
tmp := filepath.Join(t.TempDir(), "test.txt")
|
||||||
|
if err := os.WriteFile(tmp, []byte("hello"), 0o600); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
cancel() // pre-cancel before the call
|
||||||
|
|
||||||
|
// With conn == nil, callWSCommand will return "websocket not connected".
|
||||||
|
// The pre-canceled context is also an error; either one satisfies the test.
|
||||||
|
_, err := ch.UploadWSMedia(ctx, tmp, "", "file")
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for canceled context or nil conn, got nil")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// callWSCommand — unit tests without a live connection
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
// TestCallWSCommand_EmptyReqID verifies that callWSCommand rejects a command
|
||||||
|
// whose req_id is empty before touching the connection.
|
||||||
|
func TestCallWSCommand_EmptyReqID(t *testing.T) {
|
||||||
|
ch := newTestWSChannel(t)
|
||||||
|
cmd := wsCommand{
|
||||||
|
Cmd: "aibot_upload_media_init",
|
||||||
|
Body: map[string]string{"type": "file"},
|
||||||
|
// Headers.ReqID intentionally left empty
|
||||||
|
}
|
||||||
|
_, err := ch.callWSCommand(cmd, 5*time.Second)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for empty req_id, got nil")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestCallWSCommand_NotConnected verifies that callWSCommand returns a clear
|
||||||
|
// error when the WebSocket connection is nil.
|
||||||
|
func TestCallWSCommand_NotConnected(t *testing.T) {
|
||||||
|
ch := newTestWSChannel(t)
|
||||||
|
cmd := wsCommand{
|
||||||
|
Cmd: "aibot_upload_media_init",
|
||||||
|
Headers: wsHeaders{ReqID: "test-req-id"},
|
||||||
|
Body: map[string]string{"type": "file"},
|
||||||
|
}
|
||||||
|
_, err := ch.callWSCommand(cmd, 5*time.Second)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error when conn is nil, got nil")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Upload body / response type round-trip tests
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
// TestWsUploadChunkCount verifies the integer chunk-count formula used inside
|
||||||
|
// UploadWSMedia without needing the math package.
|
||||||
|
func TestWsUploadChunkCount(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
size int64
|
||||||
|
wantChunks int
|
||||||
|
}{
|
||||||
|
{"one byte", 1, 1},
|
||||||
|
{"exactly one chunk", int64(wsUploadChunkSize), 1},
|
||||||
|
{"one byte over", int64(wsUploadChunkSize) + 1, 2},
|
||||||
|
{"max allowed", int64(wsUploadMaxChunks) * int64(wsUploadChunkSize), wsUploadMaxChunks},
|
||||||
|
}
|
||||||
|
for _, tc := range tests {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
got := int((tc.size + int64(wsUploadChunkSize) - 1) / int64(wsUploadChunkSize))
|
||||||
|
if got != tc.wantChunks {
|
||||||
|
t.Errorf("size=%d: got %d chunks, want %d", tc.size, got, tc.wantChunks)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestWsUploadInitBodyJSON checks that wsUploadInitBody serializes and
|
||||||
|
// deserializes correctly, including the omitempty MD5 field.
|
||||||
|
func TestWsUploadInitBodyJSON(t *testing.T) {
|
||||||
|
b := wsUploadInitBody{
|
||||||
|
Type: "image",
|
||||||
|
Filename: "photo.jpg",
|
||||||
|
TotalSize: 1024,
|
||||||
|
TotalChunks: 1,
|
||||||
|
MD5: "deadbeefdeadbeef",
|
||||||
|
}
|
||||||
|
data, err := json.Marshal(b)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("marshal: %v", err)
|
||||||
|
}
|
||||||
|
var got wsUploadInitBody
|
||||||
|
if err := json.Unmarshal(data, &got); err != nil {
|
||||||
|
t.Fatalf("unmarshal: %v", err)
|
||||||
|
}
|
||||||
|
if got != b {
|
||||||
|
t.Errorf("round-trip mismatch: got %+v, want %+v", got, b)
|
||||||
|
}
|
||||||
|
|
||||||
|
// MD5 must be omitted when empty.
|
||||||
|
b.MD5 = ""
|
||||||
|
data, _ = json.Marshal(b)
|
||||||
|
if contains(string(data), `"md5"`) {
|
||||||
|
t.Errorf("expected md5 key to be absent when empty, got: %s", data)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestWsUploadChunkBodyBase64 checks that wsUploadChunkBody correctly
|
||||||
|
// round-trips arbitrary binary data through base64.
|
||||||
|
func TestWsUploadChunkBodyBase64(t *testing.T) {
|
||||||
|
raw := []byte("binary\x00data\xff\xfe")
|
||||||
|
body := wsUploadChunkBody{
|
||||||
|
UploadID: "uid-test-1",
|
||||||
|
ChunkIndex: 0,
|
||||||
|
Base64Data: base64.StdEncoding.EncodeToString(raw),
|
||||||
|
}
|
||||||
|
encoded, err := json.Marshal(body)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("marshal: %v", err)
|
||||||
|
}
|
||||||
|
var got wsUploadChunkBody
|
||||||
|
if err = json.Unmarshal(encoded, &got); err != nil {
|
||||||
|
t.Fatalf("unmarshal: %v", err)
|
||||||
|
}
|
||||||
|
decoded, err := base64.StdEncoding.DecodeString(got.Base64Data)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("decode base64: %v", err)
|
||||||
|
}
|
||||||
|
if string(decoded) != string(raw) {
|
||||||
|
t.Errorf("base64 round-trip mismatch: got %v, want %v", decoded, raw)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestWsUploadFinishResponseJSON checks that the finish response unmarshals
|
||||||
|
// the fields WeCom returns.
|
||||||
|
func TestWsUploadFinishResponseJSON(t *testing.T) {
|
||||||
|
raw := `{"type":"file","media_id":"MEDIAID_ABCDE","created_at":"1680000000"}`
|
||||||
|
var resp wsUploadFinishResponse
|
||||||
|
if err := json.Unmarshal([]byte(raw), &resp); err != nil {
|
||||||
|
t.Fatalf("unmarshal: %v", err)
|
||||||
|
}
|
||||||
|
if resp.MediaID != "MEDIAID_ABCDE" {
|
||||||
|
t.Errorf("media_id: got %q, want %q", resp.MediaID, "MEDIAID_ABCDE")
|
||||||
|
}
|
||||||
|
if resp.Type != "file" {
|
||||||
|
t.Errorf("type: got %q, want %q", resp.Type, "file")
|
||||||
|
}
|
||||||
|
if resp.CreatedAt != "1680000000" {
|
||||||
|
t.Errorf("created_at: got %q, want %q", resp.CreatedAt, "1680000000")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// contains is a helper to avoid importing "strings" for a single check.
|
||||||
|
func contains(s, sub string) bool {
|
||||||
|
return len(s) >= len(sub) && (s == sub || len(s) > 0 && containsHelper(s, sub))
|
||||||
|
}
|
||||||
|
|
||||||
|
func containsHelper(s, sub string) bool {
|
||||||
|
for i := 0; i <= len(s)-len(sub); i++ {
|
||||||
|
if s[i:i+len(sub)] == sub {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
Loading…
Add table
Reference in a new issue