feat(wecom): implement byte-aware content splitting for WeCom AI Bot stream messages
This commit is contained in:
parent
a2a733038b
commit
8440fd1832
2 changed files with 155 additions and 13 deletions
|
|
@ -48,6 +48,11 @@ const (
|
||||||
|
|
||||||
// Keep req_id -> chat route for late fallback pushes after stream window closes.
|
// Keep req_id -> chat route for late fallback pushes after stream window closes.
|
||||||
wsLateReplyRouteTTL = 30 * time.Minute
|
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.
|
// wsImageHTTPClient is a shared HTTP client for downloading inbound images.
|
||||||
|
|
@ -865,8 +870,12 @@ func (c *WeComAIBotWSChannel) dispatchWSAgentTask(
|
||||||
for {
|
for {
|
||||||
select {
|
select {
|
||||||
case answer := <-task.answerCh:
|
case answer := <-task.answerCh:
|
||||||
// send final frame with finish=true; any media will come in subsequent frames (if at all)
|
// Split the answer into byte-bounded chunks and send as stream frames.
|
||||||
c.wsSendStreamFinish(reqID, streamID, answer)
|
// 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)
|
c.deleteReqState(reqID)
|
||||||
return
|
return
|
||||||
case <-ticker.C:
|
case <-ticker.C:
|
||||||
|
|
@ -994,22 +1003,29 @@ func (c *WeComAIBotWSChannel) wsSendWelcomeMsg(reqID, content string) {
|
||||||
}
|
}
|
||||||
|
|
||||||
// wsSendActivePush sends a proactive markdown message using aibot_send_msg.
|
// 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.
|
// 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 {
|
func (c *WeComAIBotWSChannel) wsSendActivePush(chatID string, chatType uint32, content string) error {
|
||||||
if chatID == "" {
|
if chatID == "" {
|
||||||
return fmt.Errorf("chatid is empty")
|
return fmt.Errorf("chatid is empty")
|
||||||
}
|
}
|
||||||
|
for _, chunk := range splitWSContent(content, wsStreamMaxContentBytes) {
|
||||||
reqID := wsGenerateID()
|
reqID := wsGenerateID()
|
||||||
return c.writeWSAndWait(wsCommand{
|
if err := c.writeWSAndWait(wsCommand{
|
||||||
Cmd: "aibot_send_msg",
|
Cmd: "aibot_send_msg",
|
||||||
Headers: wsHeaders{ReqID: reqID},
|
Headers: wsHeaders{ReqID: reqID},
|
||||||
Body: wsSendMsgBody{
|
Body: wsSendMsgBody{
|
||||||
ChatID: chatID,
|
ChatID: chatID,
|
||||||
ChatType: chatType,
|
ChatType: chatType,
|
||||||
MsgType: "markdown",
|
MsgType: "markdown",
|
||||||
Markdown: &wsMarkdownContent{Content: content},
|
Markdown: &wsMarkdownContent{Content: chunk},
|
||||||
},
|
},
|
||||||
}, wsSendMsgTimeout)
|
}, wsSendMsgTimeout); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// writeWSAndWait writes cmd to the active connection and validates the command response.
|
// writeWSAndWait writes cmd to the active connection and validates the command response.
|
||||||
|
|
@ -1280,3 +1296,52 @@ func wsLabelToDefaultExt(label string) string {
|
||||||
return ".bin"
|
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
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -6,6 +6,7 @@ import (
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httptest"
|
"net/http/httptest"
|
||||||
"os"
|
"os"
|
||||||
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
"github.com/sipeed/picoclaw/pkg/bus"
|
"github.com/sipeed/picoclaw/pkg/bus"
|
||||||
|
|
@ -216,3 +217,79 @@ func TestStoreWSMedia_ContentTypeExt(t *testing.T) {
|
||||||
t.Errorf("expected .mp4 extension from Content-Type, got %q", ext)
|
t.Errorf("expected .mp4 extension from Content-Type, got %q", ext)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TestSplitWSContent verifies byte-aware splitting of stream content.
|
||||||
|
func TestSplitWSContent(t *testing.T) {
|
||||||
|
t.Run("short content is not split", func(t *testing.T) {
|
||||||
|
chunks := splitWSContent("hello", 20480)
|
||||||
|
if len(chunks) != 1 || chunks[0] != "hello" {
|
||||||
|
t.Fatalf("unexpected chunks: %v", chunks)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("ASCII content split at byte boundary", func(t *testing.T) {
|
||||||
|
// Build a string just over the limit.
|
||||||
|
content := strings.Repeat("a", 20481)
|
||||||
|
chunks := splitWSContent(content, 20480)
|
||||||
|
if len(chunks) < 2 {
|
||||||
|
t.Fatalf("expected >= 2 chunks, got %d", len(chunks))
|
||||||
|
}
|
||||||
|
for i, c := range chunks {
|
||||||
|
if len(c) > 20480 {
|
||||||
|
t.Errorf("chunk %d has %d bytes, want <= 20480", i, len(c))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Reassembled content must equal the original (possibly without leading
|
||||||
|
// whitespace that splitWSContent trims between chunks).
|
||||||
|
joined := strings.Join(chunks, "")
|
||||||
|
if len(joined) < len(content)-len(chunks) {
|
||||||
|
t.Errorf("joined length %d too short (original %d)", len(joined), len(content))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("CJK content split within byte limit", func(t *testing.T) {
|
||||||
|
// Each CJK rune is 3 bytes in UTF-8.
|
||||||
|
// 7000 CJK chars = 21000 bytes, which exceeds 20480.
|
||||||
|
content := strings.Repeat("\u4e2d", 7000)
|
||||||
|
chunks := splitWSContent(content, 20480)
|
||||||
|
if len(chunks) < 2 {
|
||||||
|
t.Fatalf("expected >= 2 chunks for 21000-byte CJK content, got %d", len(chunks))
|
||||||
|
}
|
||||||
|
for i, c := range chunks {
|
||||||
|
if len(c) > 20480 {
|
||||||
|
t.Errorf("chunk %d has %d bytes, want <= 20480", i, len(c))
|
||||||
|
}
|
||||||
|
// Every chunk must be valid UTF-8.
|
||||||
|
if !strings.ContainsRune(c, '\u4e2d') && len(c) > 0 {
|
||||||
|
// quick plausibility check — content was pure CJK
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestSplitAtByteBoundary verifies the last-resort byte-boundary splitter.
|
||||||
|
func TestSplitAtByteBoundary(t *testing.T) {
|
||||||
|
t.Run("ASCII fits in one chunk", func(t *testing.T) {
|
||||||
|
parts := splitAtByteBoundary("hello world", 100)
|
||||||
|
if len(parts) != 1 {
|
||||||
|
t.Fatalf("expected 1 part, got %d", len(parts))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("splits at byte boundary, never mid-rune", func(t *testing.T) {
|
||||||
|
// 10 CJK characters = 30 bytes; split at 20 bytes.
|
||||||
|
s := strings.Repeat("\u6587", 10) // 10 × 3 bytes = 30 bytes
|
||||||
|
parts := splitAtByteBoundary(s, 20)
|
||||||
|
for i, p := range parts {
|
||||||
|
if len(p) > 20 {
|
||||||
|
t.Errorf("part %d has %d bytes, want <= 20", i, len(p))
|
||||||
|
}
|
||||||
|
// Must be valid UTF-8 (no torn multi-byte sequences).
|
||||||
|
for j, r := range p {
|
||||||
|
if r == '\uFFFD' {
|
||||||
|
t.Errorf("part %d has replacement rune at position %d: torn UTF-8", i, j)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue