jiazhizhong/higress
1
0
Fork 0
mirror of https://github.com/alibaba/higress.git synced 2026-06-26 02:35:02 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(ai-cache): handle SSE first chunk with role only (fixes #3953) (#3962)

Browse Source 
Signed-off-by: CH3CHO <ch3cho@qq.com>
Co-authored-by: woody <yaodiwu618@gmail.com>
This commit is contained in:
Kent Dong
2026-06-23 17:38:50 +08:00
committed by GitHub
parent 7c94b5a822
commit 1694f48fd9
2 changed files with 205 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
plugins/wasm-go/extensions/ai-cache/util.go
View File

@@ -88,6 +88,13 @@ func processStreamLastChunk(ctx wrapper.HttpContext, c config.PluginConfig, chun
if err != nil {
return "", fmt.Errorf("[processStreamLastChunk] processSSEMessage failed, error: %v", err)
}
// 兜底：[DONE] 或其它尾部 chunk 无 content 时，processSSEMessage 返回空，
// 此时从 ctx 取已累积的缓存内容，避免缓存写空。
if value == "" {
if tempContentI := ctx.GetContext(CACHE_CONTENT_CONTEXT_KEY); tempContentI != nil {
value = tempContentI.(string)
}
}
return value, nil
}
tempContentI := ctx.GetContext(CACHE_CONTENT_CONTEXT_KEY)
@@ -99,6 +106,11 @@ func processStreamLastChunk(ctx wrapper.HttpContext, c config.PluginConfig, chun
func processSSEMessage(ctx wrapper.HttpContext, c config.PluginConfig, sseMessage string, log log.Log) (string, error) {
content := ""
// done 标记本次 sseMessage 是否遇到 [DONE]。当最后一段 content 与 [DONE]
// 处于同一 buffer 时，必须跳出循环后由循环外的 merge 逻辑统一合并到
// CACHE_CONTENT_CONTEXT_KEY，否则本次解析的最后一段 content 会被 [DONE]
// 早 return 丢弃（PR #3962 review）。
done := false
for _, chunk := range strings.Split(sseMessage, "\n\n") {
log.Debugf("single sse message: %s", chunk)
subMessages := strings.Split(chunk, "\n")
@@ -117,7 +129,9 @@ func processSSEMessage(ctx wrapper.HttpContext, c config.PluginConfig, sseMessag
bodyJson := message[5:]
if strings.TrimSpace(bodyJson) == "[DONE]" {
return content, nil
// 跳出循环，把已解析的局部 content 留到循环外统一合并。
done = true
break
}
// Extract values from JSON fields
@@ -130,22 +144,29 @@ func processSSEMessage(ctx wrapper.HttpContext, c config.PluginConfig, sseMessag
}
// Check if the ResponseBody field exists
if !responseBody.Exists() {
if ctx.GetContext(CACHE_CONTENT_CONTEXT_KEY) != nil {
log.Debugf("[processSSEMessage] unable to extract content from message; cache content is not nil: %s", message)
return content, nil
}
return content, fmt.Errorf("[processSSEMessage] unable to extract content from message; cache content is nil: %s", message)
} else {
if responseBody.Exists() {
content += responseBody.String()
}
}
tempContentI := ctx.GetContext(CACHE_CONTENT_CONTEXT_KEY)
// If there is no content in the cache, initialize and set the content
if tempContentI == nil {
ctx.SetContext(CACHE_CONTENT_CONTEXT_KEY, content)
} else {
ctx.SetContext(CACHE_CONTENT_CONTEXT_KEY, tempContentI.(string)+content)
// 本次 sseMessage 既没解析到 content 也没遇到 [DONE]：保持 ctx 不变，直接返回。
if content == "" && !done {
log.Debugf("[processSSEMessage] no content extracted; skipping cache update: %s", sseMessage)
return "", nil
}
return content, nil
if content != "" {
if v := ctx.GetContext(CACHE_CONTENT_CONTEXT_KEY); v == nil {
ctx.SetContext(CACHE_CONTENT_CONTEXT_KEY, content)
} else {
ctx.SetContext(CACHE_CONTENT_CONTEXT_KEY, v.(string)+content)
}
}
// handleStreamChunk 不使用返回值；processStreamLastChunk 把它作为 cacheResponse 的
// SET value，必须是完整累积值（避免最后一段 content 因 [DONE] 早 return 被丢）。
if v := ctx.GetContext(CACHE_CONTENT_CONTEXT_KEY); v != nil {
return v.(string), nil
}
return "", nil
}