Weknora/Enterprise-Knowledge-Base
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(dingtalk): enable secure multimodal image QA

Browse Source
This commit is contained in:
姚章浩
2026-07-16 00:17:16 +08:00
parent 4022040e98
commit c9c4f0aaa5
59 changed files with 3279 additions and 429 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/test-plans/assets/dingtalk-vision-acceptance.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 70 KiB

docs/test-plans/assets/generate_dingtalk_vision_acceptance.py
+34
View File
@@ -0,0 +1,34 @@
from pathlib import Path
from PIL import Image, ImageDraw, ImageFont
OUT = Path(__file__).with_name("dingtalk-vision-acceptance.png")
FONT_PATH = "/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc"
image = Image.new("RGB", (1280, 720), "#F7FAFC")
draw = ImageDraw.Draw(image)
font_title = ImageFont.truetype(FONT_PATH, 46)
font_code = ImageFont.truetype(FONT_PATH, 58)
font_label = ImageFont.truetype(FONT_PATH, 34)
font_question = ImageFont.truetype(FONT_PATH, 38)
# Header and unique OCR target.
draw.rounded_rectangle((50, 40, 1230, 190), radius=24, fill="#E6F4FF", outline="#1677FF", width=4)
draw.text((85, 65), "WeKnora 钉钉多模态验收图", font=font_title, fill="#102A43")
draw.text((85, 120), "唯一验收码：WK-VISION-7264", font=font_code, fill="#C53030")
# Three visually distinct objects.
draw.ellipse((120, 280, 340, 500), fill="#1677FF", outline="#0B4F9C", width=8)
draw.text((160, 520), "蓝色圆形", font=font_label, fill="#0B4F9C")
draw.polygon([(520, 500), (650, 275), (780, 500)], fill="#FFD43B", outline="#9C7600", width=8)
draw.text((548, 520), "黄色三角形", font=font_label, fill="#725500")
draw.rounded_rectangle((940, 295, 1140, 495), radius=12, fill="#2FB344", outline="#176B2A", width=8)
draw.text((958, 520), "绿色方形", font=font_label, fill="#176B2A")
question = "问题：蓝色圆形右侧紧邻的图形是什么颜色和形状？"
draw.rounded_rectangle((80, 610, 1200, 690), radius=18, fill="#FFF4E6", outline="#E67700", width=3)
draw.text((110, 625), question, font=font_question, fill="#7C3E00")
image.save(OUT, format="PNG", optimize=True)
print(OUT)
docs/test-plans/dingtalk-image-multimodal-acceptance.md
+170
View File
@@ -0,0 +1,170 @@
# 阶段2：钉钉图片传输与多模态识别验收计划
## 1. 验收目标
验证图片从钉钉用户消息进入 WeKnora，并由视觉模型真实读取图片内容后返回结果；不得把“图片已上传”“图片作为知识文件入库”或“删除图片后按纯文字回答”判定为多模态通过。
## 2. 当前自动化验证基线
测试分支：`feature/dingtalk-image-only`
生产基线：`4022040e`（不包含阶段3～5提交）
阶段2提交：待本轮代码审查通过后创建。
已执行的目标测试：
```bash
GOTOOLCHAIN=local GOPROXY=off GOMAXPROCS=4 \
/home/zheromiku/.local/toolchains/go1.26.0/bin/go test -p 1 \
./internal/im/dingtalk \
-run 'Test(ParseCallbackMessagePicture|AdapterDownloadFileUsesMessageFilesAPI|ValidateDingTalkDownloadURLRejectsUnsafeTargets|ParseDingTalkStreamPayloadPicture|ParseDingTalkStreamPayloadPictureJSONString)$' \
-count=1
GOTOOLCHAIN=local GOPROXY=off GOMAXPROCS=4 \
/home/zheromiku/.local/toolchains/go1.26.0/bin/go test -p 1 \
./internal/im -run 'Test(PrepareIMImageForQA|IMQAErrorReply|FailClosedIMImageAnswer)' -count=1
GOTOOLCHAIN=local GOPROXY=off GOMAXPROCS=4 \
/home/zheromiku/.local/toolchains/go1.26.0/bin/go test -p 1 \
./internal/models/chat \
-run '^TestRemoteAPIChat_MultimodalUnsupportedDoesNotStripImages$' \
-count=1
```
结果：钉钉图片解析、两段式下载、Token请求禁重定向、DNS解析后固定IP拨号、保留/Teredo/受限6to4地址拒绝、签名URL错误脱敏、所有data URI及图片调试日志全量脱敏、JPEG/PNG格式/大小/维度校验、Worker延迟下载、`QARequest.ImageURLs`传递，以及视觉失败/流式部分输出时禁止删图或伪成功，全部通过。
已执行相关包回归：
```bash
GOTOOLCHAIN=local GOPROXY=off GOMAXPROCS=4 \
/home/zheromiku/.local/toolchains/go1.26.0/bin/go test -p 1 \
./internal/agent ./internal/im/dingtalk ./internal/im ./internal/models/chat ./internal/utils ./internal/logger \
./internal/application/service/chat_pipeline ./internal/application/service -count=1
```
结果：八个相关包全部通过，八包 `go vet` 同步通过。
## 3. 真实验收前置条件
- 采用已授权的 App-only 可回滚部署：先备份当前 App 镜像标签，仅重建 `app`，不重建 PostgreSQL、Redis、前端、DocReader 或 MinerU。
- 部署后先检查 `WeKnora-app` 健康状态；失败立即恢复备份镜像并仅重建 `app`。
- 钉钉应用已开通机器人消息文件下载权限。
- WeKnora Agent 已开启“图片上传/多模态”。
- Agent 当前问答模型支持视觉输入，或配置了明确的独立 VLM；VLM不可用或未返回图片描述时必须失败关闭。
- 不允许在视觉能力不足时静默删除图片并继续生成答案。
## 4. 测试数据
标准测试图：
```text
docs/test-plans/assets/dingtalk-vision-acceptance.png
```
图片包含：
- 唯一验收码：`WK-VISION-7264`
- 蓝色圆形
- 黄色三角形
- 绿色方形
- 明确问题：蓝色圆形右侧紧邻的图形是什么颜色和形状？
标准答案必须同时包含：
```text
WK-VISION-7264
黄色三角形
```
## 5. 人工验收用例
### DT-IMG-01：PNG图片真实识别
操作：在钉钉中把标准测试图发送给测试机器人。
通过标准：
- 机器人成功回复；
- 回复包含 `WK-VISION-7264`；
- 回复包含“黄色三角形”；
- 回复不是“请用文字描述”“当前不支持图片”或无关答案。
### DT-IMG-02：普通截图OCR
操作：发送一张包含中文、英文和数字的实际设备截图。
通过标准：
- 能提取截图里的关键文字；
- 能结合图片回答具体问题；
- 不编造截图中不存在的错误码、型号或参数。
### DT-IMG-03：纯文字回归
操作：向同一机器人发送普通文字问题。
通过标准：正常进入原知识库问答，图片功能不影响文字消息。
### DT-IMG-04：伪图片拒绝
操作：将非图片内容伪装为 `.png` 后发送。
通过标准：拒绝处理，提示图片格式校验失败；伪图片不得进入模型。
### DT-IMG-05：超限图片拒绝
操作：发送超过10MB的JPG/PNG。
通过标准：明确提示10MB限制，不进入模型请求。
### DT-IMG-06：模型不支持视觉
操作：测试环境临时选择不支持图片的问答模型后发送标准图。
通过标准：明确提示视觉模型未配置或不支持；禁止删除图片后继续给出看似正常的答案。
### DT-IMG-07：钉钉权限缺失
操作：在没有消息文件下载权限的测试应用中发送图片，或使用接口模拟对应403响应。
通过标准：日志记录脱敏后的权限错误，用户收到可理解提示；不得记录 access token、downloadCode 或临时下载URL。
## 6. 日志观察点
发送标准图后，检查测试环境日志：
```bash
sudo docker logs WeKnora-app --since 5m 2>&1 | \
grep -E '\[DingTalk\]|Failed to prepare image|message file|multimodal|vision|image'
```
日志不得出现：
- access token明文；
- `downloadCode`明文；
- 完整临时下载URL；
- 完整Base64图片。
## 7. 数据检查
阶段2只验收图片传输和模型识别。消息审计字段的完整性、Token统计、知识库命中片段和渠道统计属于阶段3，本阶段不提前判定通过。
本阶段允许确认：
- 用户消息以默认视觉问题进入原有会话历史；
- 图片只作为瞬态 `QARequest.ImageURLs` 传给模型；
- 数据库不保存完整Base64、临时下载码和临时签名URL。
## 8. 通过门槛
必须同时满足：
1. 图片自动化测试全部通过；
2. 相关IM包回归通过；
3. `DT-IMG-01`和`DT-IMG-02`真实钉钉测试通过；
4. `DT-IMG-03`文字回归通过；
5. `DT-IMG-04`～`DT-IMG-07`至少完成目标验证或提供自动化等价证据；
6. 用户确认验收。
在以上条件完成前，阶段2状态保持“进行中”，不得启动阶段3。
frontend/src/utils/weknoraCloudModels.ts
+1 -1
View File
@@ -57,7 +57,7 @@ export function buildWkcModelConfig(
truncate_prompt_tokens: 0,
}
}
if (kind === 'vllm') {
if (kind === 'chat' || kind === 'vllm') {
parameters.supports_vision = true
}
internal/agent/engine.go
+4 -3
View File
@@ -269,7 +269,7 @@ func (e *AgentEngine) Execute(
"tools": toolListStr,
})
_, err := e.executeLoop(ctx, state, query, messages, tools, sessionID, messageID)
_, err := e.executeLoop(ctx, state, query, messages, tools, sessionID, messageID, imgs)
if err != nil {
logger.Errorf(ctx, "[Agent] Execution failed: %v", err)
e.eventBus.Emit(ctx, event.Event{
@@ -349,6 +349,7 @@ func (e *AgentEngine) executeLoop(
tools []chat.Tool,
sessionID string,
messageID string,
imageURLs []string,
) (*types.AgentState, error) {
startTime := time.Now()
common.PipelineInfo(ctx, "Agent", "loop_start", map[string]interface{}{
@@ -386,7 +387,7 @@ loop:
if totalTC := countTotalToolCalls(state.RoundSteps); totalTC > 0 {
logger.Infof(ctx, "[Agent] Synthesizing final answer from %d existing tool results",
totalTC)
_ = e.streamFinalAnswerToEventBus(ctx, query, state, sessionID)
_ = e.streamFinalAnswerToEventBus(ctx, query, state, sessionID, imageURLs)
state.IsComplete = true
}
return state, ctx.Err()
@@ -419,7 +420,7 @@ loop:
// complete answer." message, which then leaks to the UI as the final
// answer for a conversation the user deliberately stopped.
if !state.IsComplete && ctx.Err() == nil {
e.handleMaxIterations(ctx, query, state, sessionID)
e.handleMaxIterations(ctx, query, state, sessionID, imageURLs)
}
return state, nil
internal/agent/engine_test.go
+5 -5
View File
@@ -120,7 +120,7 @@ func TestExecuteLoop_EmptyContentWithStop_ShouldNotCompleteWithEmpty(t *testing.
state := &types.AgentState{}
ctx := context.Background()
_, err := engine.executeLoop(ctx, state, "test query", emptyMessages(), emptyTools(), "sess-1", "msg-1")
_, err := engine.executeLoop(ctx, state, "test query", emptyMessages(), emptyTools(), "sess-1", "msg-1", nil)
assert.NoError(t, err)
assert.True(t, state.IsComplete)
@@ -146,7 +146,7 @@ func TestExecuteLoop_NonEmptyContentWithStop_ShouldComplete(t *testing.T) {
state := &types.AgentState{}
ctx := context.Background()
_, err := engine.executeLoop(ctx, state, "test query", emptyMessages(), emptyTools(), "sess-1", "msg-1")
_, err := engine.executeLoop(ctx, state, "test query", emptyMessages(), emptyTools(), "sess-1", "msg-1", nil)
assert.NoError(t, err)
assert.True(t, state.IsComplete)
@@ -173,7 +173,7 @@ func TestExecuteLoop_EmptyThenNonEmpty_ShouldRetryAndComplete(t *testing.T) {
state := &types.AgentState{}
ctx := context.Background()
_, err := engine.executeLoop(ctx, state, "test query", emptyMessages(), emptyTools(), "sess-1", "msg-1")
_, err := engine.executeLoop(ctx, state, "test query", emptyMessages(), emptyTools(), "sess-1", "msg-1", nil)
assert.NoError(t, err)
assert.True(t, state.IsComplete)
@@ -328,7 +328,7 @@ func TestExecuteLoop_NaturalStop_DoesNotDuplicateAnswer(t *testing.T) {
state := &types.AgentState{}
_, err := engine.executeLoop(context.Background(), state, "test query",
emptyMessages(), emptyTools(), "sess-1", "msg-1")
emptyMessages(), emptyTools(), "sess-1", "msg-1", nil)
require.NoError(t, err)
assert.True(t, state.IsComplete)
@@ -358,7 +358,7 @@ func TestStreamFinalAnswerToEventBus_EmitsDoneWhenProviderEndsWithEmptyChunk(t *
})
state := &types.AgentState{}
err := engine.streamFinalAnswerToEventBus(context.Background(), "test query", state, "sess-1")
err := engine.streamFinalAnswerToEventBus(context.Background(), "test query", state, "sess-1", nil)
require.NoError(t, err)
require.Len(t, finalAnswerEvents, 2)
internal/agent/finalize.go
+17 -3
View File
@@ -19,6 +19,7 @@ func (e *AgentEngine) streamFinalAnswerToEventBus(
query string,
state *types.AgentState,
sessionID string,
imageURLs []string,
) error {
totalToolCalls := countTotalToolCalls(state.RoundSteps)
logger.Infof(ctx, "[Agent][FinalAnswer] Synthesizing from %d steps, %d tool calls",
@@ -36,7 +37,7 @@ func (e *AgentEngine) streamFinalAnswerToEventBus(
messages := []chat.Message{
{Role: "system", Content: systemPrompt},
{Role: "user", Content: userTurn},
{Role: "user", Content: userTurn, Images: append([]string(nil), imageURLs...)},
}
// Add all tool call results as context
@@ -138,10 +139,23 @@ Now generate the final answer:`, query)
return nil
}
func collectAgentImageURLs(messages []chat.Message) []string {
var images []string
for _, message := range messages {
images = append(images, message.Images...)
for _, part := range message.MultiContent {
if part.Type == "image_url" && part.ImageURL != nil && part.ImageURL.URL != "" {
images = append(images, part.ImageURL.URL)
}
}
}
return images
}
// handleMaxIterations generates a final answer when the agent loop exhausted all iterations
// without the LLM producing a natural stop. It marks state.IsComplete = true.
func (e *AgentEngine) handleMaxIterations(
ctx context.Context, query string, state *types.AgentState, sessionID string,
ctx context.Context, query string, state *types.AgentState, sessionID string, imageURLs []string,
) {
logger.Info(ctx, "Reached max iterations, generating final answer")
common.PipelineWarn(ctx, "Agent", "max_iterations_reached", map[string]interface{}{
@@ -150,7 +164,7 @@ func (e *AgentEngine) handleMaxIterations(
})
// Stream final answer generation through EventBus
if err := e.streamFinalAnswerToEventBus(ctx, query, state, sessionID); err != nil {
if err := e.streamFinalAnswerToEventBus(ctx, query, state, sessionID, imageURLs); err != nil {
logger.Errorf(ctx, "Failed to synthesize final answer: %v", err)
common.PipelineError(ctx, "Agent", "final_answer_failed", map[string]interface{}{
"error": err.Error(),
internal/agent/finalize_image_test.go
+18
View File
@@ -0,0 +1,18 @@
package agent
import (
"testing"
"github.com/Tencent/WeKnora/internal/models/chat"
)
func TestCollectAgentImageURLsForFinalSynthesis(t *testing.T) {
messages := []chat.Message{
{Role: "system", Content: "system"},
{Role: "user", Content: "question", Images: []string{"data:image/png;base64,YWJj"}},
}
got := collectAgentImageURLs(messages)
if len(got) != 1 || got[0] != "data:image/png;base64,YWJj" {
t.Fatalf("collectAgentImageURLs() = %#v", got)
}
}
internal/agent/think.go
+10 -6
View File
@@ -2,6 +2,7 @@ package agent
import (
"context"
"errors"
"fmt"
"strings"
"time"
@@ -60,7 +61,7 @@ func (e *AgentEngine) streamLLMToEventBus(
// as if they were part of the LLM answer.
if chunk.ResponseType == types.ResponseTypeError {
result.StreamError = chunk.Content
continue
break
}
if chunk.Content != "" {
@@ -89,6 +90,9 @@ func (e *AgentEngine) streamLLMToEventBus(
if emitFunc != nil {
emitFunc(&chunk, result.Content)
}
if chunk.Done && chunk.ResponseType == types.ResponseTypeAnswer {
break
}
}
// Stream diagnostic summary: helps identify non-streaming patterns
@@ -101,10 +105,10 @@ func (e *AgentEngine) streamLLMToEventBus(
chunkCount, len(result.Content), len(result.ToolCalls),
streamDuration.Milliseconds(), responseTypeCounts)
// If the stream produced an error and no usable content/tool calls,
// surface it as a Go error so the caller can retry or degrade gracefully.
if result.StreamError != "" && result.Content == "" && len(result.ToolCalls) == 0 {
return result, fmt.Errorf("LLM stream error: %s", result.StreamError)
// A provider error invalidates the entire stream, even if partial content or
// tool calls arrived first. The caller must replace/discard that partial data.
if result.StreamError != "" {
return result, errors.New("LLM stream failed")
}
return result, nil
@@ -406,7 +410,7 @@ func (e *AgentEngine) callLLMWithRetry(
"steps": len(state.RoundSteps),
"tool_calls": totalTC,
})
if synthErr := e.streamFinalAnswerToEventBus(ctx, query, state, sessionID); synthErr != nil {
if synthErr := e.streamFinalAnswerToEventBus(ctx, query, state, sessionID, collectAgentImageURLs(messages)); synthErr != nil {
logger.Errorf(ctx, "[Agent] Final answer synthesis also failed: %v", synthErr)
return nil, fmt.Errorf("LLM call failed: %w (synthesis also failed: %v)", err, synthErr)
}
internal/agent/think_stream_error_test.go
+26
View File
@@ -0,0 +1,26 @@
package agent
import (
"context"
"testing"
"github.com/Tencent/WeKnora/internal/models/chat"
"github.com/Tencent/WeKnora/internal/types"
)
func TestStreamLLMToEventBusRejectsPartialContentAfterStreamError(t *testing.T) {
model := &mockChat{responses: []mockResponse{{chunks: []types.StreamResponse{
{ResponseType: types.ResponseTypeAnswer, Content: "unverified partial"},
{ResponseType: types.ResponseTypeError, Content: "provider failed", Done: true},
{ResponseType: types.ResponseTypeAnswer, Content: "must be ignored", Done: true},
}}}}
engine := &AgentEngine{config: &types.AgentConfig{}, chatModel: model}
result, err := engine.streamLLMToEventBus(context.Background(), []chat.Message{{Role: "user", Content: "inspect image"}}, nil, nil)
if err == nil {
t.Fatal("expected stream error after partial content")
}
if result == nil || result.Content != "unverified partial" {
t.Fatalf("result = %#v; caller needs partial content identified for replacement", result)
}
}
internal/application/service/chat_pipeline/chat_completion_stream.go
+12 -4
View File
@@ -79,8 +79,10 @@ func (p *PluginChatCompletionStream) OnEvent(ctx context.Context,
pipelineInfo(ctx, "Stream", "model_call", map[string]interface{}{
"chat_model": chatManage.ChatModelID,
})
responseChan, err := chatModel.ChatStream(ctx, chatMessages, opt)
streamCtx, streamCancel := context.WithCancel(ctx)
responseChan, err := chatModel.ChatStream(streamCtx, chatMessages, opt)
if err != nil {
streamCancel()
pipelineError(ctx, "Stream", "model_call", map[string]interface{}{
"chat_model": chatManage.ChatModelID,
"error": err.Error(),
@@ -88,6 +90,7 @@ func (p *PluginChatCompletionStream) OnEvent(ctx context.Context,
return ErrModelCall.WithError(err)
}
if responseChan == nil {
streamCancel()
pipelineError(ctx, "Stream", "model_call", map[string]interface{}{
"chat_model": chatManage.ChatModelID,
"error": "nil_channel",
@@ -105,6 +108,7 @@ func (p *PluginChatCompletionStream) OnEvent(ctx context.Context,
// The goroutine monitors ctx.Done() to avoid leaking when the context is cancelled
// and the upstream channel is not closed promptly.
go func() {
defer streamCancel()
thinkingID := fmt.Sprintf("%s-thinking", uuid.New().String()[:8])
answerID := fmt.Sprintf("%s-answer", uuid.New().String()[:8])
thinkingOpen := false
@@ -145,19 +149,20 @@ func (p *PluginChatCompletionStream) OnEvent(ctx context.Context,
if response.ResponseType == types.ResponseTypeError {
pipelineError(ctx, "Stream", "stream_error", map[string]interface{}{
"session_id": chatManage.SessionID,
"error": response.Content,
"error": "chat completion stream failed",
})
eventBus.Emit(ctx, types.Event{
ID: fmt.Sprintf("%s-error", uuid.New().String()[:8]),
Type: types.EventType(event.EventError),
SessionID: chatManage.SessionID,
Data: event.ErrorData{
Error: response.Content,
Error: "chat completion stream failed",
Stage: "chat_completion_stream",
SessionID: chatManage.SessionID,
},
})
continue
closeThinking()
return
}
if response.ResponseType == types.ResponseTypeThinking {
@@ -190,6 +195,9 @@ func (p *PluginChatCompletionStream) OnEvent(ctx context.Context,
Done: response.Done,
},
})
if response.Done {
return
}
}
}
}
internal/application/service/chat_pipeline/query_understand.go
+39 -15
View File
@@ -91,7 +91,14 @@ func (p *PluginQueryUnderstand) OnEvent(ctx context.Context,
}
// --- Select the appropriate model ---
rewriteModel, useImages := p.selectModel(ctx, chatManage, hasImages)
rewriteModel, useImages, selectErr := p.selectModel(ctx, chatManage, hasImages)
if selectErr != nil {
pipelineError(ctx, "QueryUnderstand", "get_visual_model", map[string]interface{}{
"session_id": chatManage.SessionID,
"error": selectErr.Error(),
})
return &PluginError{Err: selectErr, Description: "Visual model is unavailable", ErrorType: "image_understanding_failed"}
}
if rewriteModel == nil {
pipelineError(ctx, "QueryUnderstand", "get_model", map[string]interface{}{
"session_id": chatManage.SessionID,
@@ -127,11 +134,18 @@ func (p *PluginQueryUnderstand) OnEvent(ctx context.Context,
"session_id": chatManage.SessionID,
"error": err.Error(),
})
if hasImages {
return &PluginError{Err: err, Description: "Visual model request failed", ErrorType: "image_understanding_failed"}
}
return next()
}
// --- Parse structured output ---
p.parseOutput(chatManage, response.Content)
if hasImages && !chatManage.ChatModelSupportsVision && strings.TrimSpace(chatManage.ImageDescription) == "" {
err := fmt.Errorf("visual model returned no image description")
return &PluginError{Err: err, Description: "Visual model returned no usable image description", ErrorType: "image_understanding_failed"}
}
// Persist image description asynchronously — this DB write does not affect
// the current pipeline result, so it can run in the background.
@@ -222,33 +236,43 @@ func (p *PluginQueryUnderstand) loadHistory(ctx context.Context, chatManage *typ
}
// selectModel picks the model for query understanding. When images are present
// it prefers a vision-capable model. Returns (model, useImages).
func (p *PluginQueryUnderstand) selectModel(ctx context.Context, chatManage *types.ChatManage, hasImages bool) (chat.Chat, bool) {
// it requires a vision-capable model and never falls back to text-only processing.
func (p *PluginQueryUnderstand) selectModel(ctx context.Context, chatManage *types.ChatManage, hasImages bool) (chat.Chat, bool, error) {
if hasImages {
var lastErr error
if chatManage.ChatModelSupportsVision {
m, err := p.modelService.GetChatModel(ctx, chatManage.ChatModelID)
if err == nil {
return m, true
if err == nil && chat.SupportsImageInput(m) {
return m, true, nil
}
pipelineWarn(ctx, "QueryUnderstand", "vision_model_fallback", map[string]interface{}{
if err == nil {
err = fmt.Errorf("configured Chat transport does not support image input")
}
lastErr = err
pipelineWarn(ctx, "QueryUnderstand", "vision_model_unavailable", map[string]interface{}{
"session_id": chatManage.SessionID,
"error": err.Error(),
})
}
if chatManage.VLMModelID != "" {
m, err := p.modelService.GetChatModel(ctx, chatManage.VLMModelID)
if err == nil {
return m, true
if err == nil && chat.SupportsImageInput(m) {
return m, true, nil
}
pipelineWarn(ctx, "QueryUnderstand", "vlm_model_fallback", map[string]interface{}{
if err == nil {
err = fmt.Errorf("configured VLM transport does not support image input")
}
lastErr = err
pipelineWarn(ctx, "QueryUnderstand", "vlm_model_unavailable", map[string]interface{}{
"session_id": chatManage.SessionID,
"vlm_model_id": chatManage.VLMModelID,
"error": err.Error(),
})
}
pipelineWarn(ctx, "QueryUnderstand", "no_vision_model", map[string]interface{}{
"session_id": chatManage.SessionID,
})
if lastErr != nil {
return nil, false, fmt.Errorf("no usable visual model: %w", lastErr)
}
return nil, false, fmt.Errorf("no visual model is configured")
}
textModelID := chatManage.ChatModelID
@@ -266,7 +290,7 @@ func (p *PluginQueryUnderstand) selectModel(ctx context.Context, chatManage *typ
"error": err.Error(),
})
if fallback, fbErr := p.modelService.GetChatModel(ctx, chatManage.ChatModelID); fbErr == nil {
return fallback, false
return fallback, false, nil
}
}
pipelineError(ctx, "QueryUnderstand", "get_model", map[string]interface{}{
@@ -274,9 +298,9 @@ func (p *PluginQueryUnderstand) selectModel(ctx context.Context, chatManage *typ
"chat_model_id": textModelID,
"error": err.Error(),
})
return nil, false
return nil, false, nil
}
return m, false
return m, false, nil
}
// buildPrompts constructs system and user prompts with placeholder replacement.
internal/application/service/chat_pipeline/query_understand_image_test.go
+54
View File
@@ -0,0 +1,54 @@
package chatpipeline
import (
"context"
"testing"
modelchat "github.com/Tencent/WeKnora/internal/models/chat"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
)
type visionModelStub struct{}
func (visionModelStub) Chat(context.Context, []modelchat.Message, *modelchat.ChatOptions) (*types.ChatResponse, error) {
return nil, nil
}
func (visionModelStub) ChatStream(context.Context, []modelchat.Message, *modelchat.ChatOptions) (<-chan types.StreamResponse, error) {
return nil, nil
}
func (visionModelStub) GetModelName() string { return "test-vlm" }
func (visionModelStub) GetModelID() string { return "test-vlm" }
func (visionModelStub) SupportsImageInput() bool { return true }
type modelServiceStub struct {
interfaces.ModelService
model modelchat.Chat
}
func (s modelServiceStub) GetChatModel(context.Context, string) (modelchat.Chat, error) {
return s.model, nil
}
func TestSelectModelFailsClosedWhenImageHasNoVisualModel(t *testing.T) {
plugin := &PluginQueryUnderstand{}
manage := &types.ChatManage{PipelineRequest: types.PipelineRequest{ChatModelSupportsVision: false}}
model, useImages, err := plugin.selectModel(context.Background(), manage, true)
if err == nil || model != nil || useImages {
t.Fatalf("selectModel = model %v useImages %v err %v; want fail closed", model, useImages, err)
}
}
func TestSelectModelUsesVLMWhenChatTransportIsNotVisual(t *testing.T) {
vlm := visionModelStub{}
plugin := &PluginQueryUnderstand{modelService: modelServiceStub{model: vlm}}
manage := &types.ChatManage{PipelineRequest: types.PipelineRequest{
ChatModelSupportsVision: false,
VLMModelID: "document-parser-vlm",
}}
model, useImages, err := plugin.selectModel(context.Background(), manage, true)
if err != nil || model == nil || !useImages {
t.Fatalf("selectModel = model %v useImages %v err %v; want visual VLM fallback", model, useImages, err)
}
}
internal/application/service/session_agent_image_route.go
+63
View File
@@ -0,0 +1,63 @@
package service
import (
"context"
"encoding/base64"
"errors"
"fmt"
"strings"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
)
const maxAgentFallbackImageBytes = 10 << 20
// prepareAgentImageInput routes image input according to the resolved primary
// transport capability. A configured VLM is required only when the actual
// primary Chat transport cannot carry images directly.
func prepareAgentImageInput(
ctx context.Context,
req *types.QARequest,
primarySupportsVision bool,
modelService interfaces.ModelService,
) (string, []string, error) {
if req == nil {
return "", nil, errors.New("QA request is required")
}
if len(req.ImageURLs) == 0 || primarySupportsVision || strings.TrimSpace(req.ImageDescription) != "" {
return routeAgentImageInput(req, primarySupportsVision)
}
if req.CustomAgent == nil || strings.TrimSpace(req.CustomAgent.Config.VLMModelID) == "" {
return "", nil, errors.New("agent model does not support visual input and no VLM is configured")
}
if modelService == nil {
return "", nil, errors.New("agent VLM service is unavailable")
}
images := make([][]byte, 0, len(req.ImageURLs))
for _, raw := range req.ImageURLs {
comma := strings.IndexByte(raw, ',')
if comma <= 0 || !strings.HasPrefix(strings.ToLower(raw[:comma]), "data:image/") ||
!strings.HasSuffix(strings.ToLower(raw[:comma]), ";base64") {
return "", nil, errors.New("agent image payload is invalid")
}
decoded, err := base64.StdEncoding.Strict().DecodeString(raw[comma+1:])
if err != nil || len(decoded) == 0 || len(decoded) > maxAgentFallbackImageBytes {
return "", nil, errors.New("agent image payload is invalid")
}
images = append(images, decoded)
}
model, err := modelService.GetVLMModel(ctx, req.CustomAgent.Config.VLMModelID)
if err != nil || model == nil {
return "", nil, errors.New("agent VLM is unavailable")
}
prompt := fmt.Sprintf("Analyze the uploaded image(s) for the user's question. Return a concise factual description only.\nQuestion: %s", req.Query)
description, err := model.Predict(ctx, images, prompt)
if err != nil || strings.TrimSpace(description) == "" {
return "", nil, errors.New("agent VLM analysis failed")
}
req.ImageDescription = strings.TrimSpace(description)
return routeAgentImageInput(req, false)
}
internal/application/service/session_agent_image_test.go
+26
View File
@@ -0,0 +1,26 @@
package service
import (
"strings"
"testing"
"github.com/Tencent/WeKnora/internal/types"
)
func TestRouteAgentImageInputFailsClosedWithoutVision(t *testing.T) {
req := &types.QARequest{Query: "what is this", ImageURLs: []string{"data:image/png;base64,abc"}}
if _, images, err := routeAgentImageInput(req, false); err == nil || len(images) != 0 {
t.Fatalf("routeAgentImageInput should reject unsupported visual input: images=%v err=%v", images, err)
}
query, images, err := routeAgentImageInput(req, true)
if err != nil || len(images) != 1 || query != req.Query {
t.Fatalf("vision route = query %q images %d err %v", query, len(images), err)
}
req.ImageDescription = "a blue robot"
query, images, err = routeAgentImageInput(req, false)
if err != nil || len(images) != 0 || !strings.Contains(query, req.ImageDescription) {
t.Fatalf("description fallback = query %q images %d err %v", query, len(images), err)
}
}
internal/application/service/session_agent_qa.go
+26 -12
View File
@@ -6,6 +6,7 @@ import (
"errors"
"fmt"
"os"
"strings"
"github.com/Tencent/WeKnora/internal/agent/tools"
"github.com/Tencent/WeKnora/internal/event"
@@ -156,21 +157,18 @@ func (s *sessionService) AgentQA(
return err
}
// Route image data based on agent model's vision capability
var agentModelSupportsVision bool
if effectiveModelID != "" {
if modelInfo, err := s.modelService.GetModelByID(ctx, effectiveModelID); err == nil && modelInfo != nil {
agentModelSupportsVision = modelInfo.Parameters.SupportsVision
}
}
// The resolved concrete transport is authoritative. Stored model metadata can
// be stale (notably older WeKnora Cloud records), while unknown transports
// still fail closed through SupportsImageInput.
agentModelSupportsVision := modelSupportsDirectImages(summaryModel)
agentQuery := req.Query
var agentImageURLs []string
if agentModelSupportsVision && len(req.ImageURLs) > 0 {
agentImageURLs = req.ImageURLs
agentQuery, agentImageURLs, err := prepareAgentImageInput(ctx, req, agentModelSupportsVision, s.modelService)
if err != nil {
return err
}
if len(agentImageURLs) > 0 {
logger.Infof(ctx, "Agent model supports vision, passing %d image(s) directly", len(agentImageURLs))
} else if req.ImageDescription != "" {
agentQuery = req.Query + "\n\n[用户上传图片内容]\n" + req.ImageDescription
logger.Infof(ctx, "Agent model does not support vision, appending image description (%d chars)", len(req.ImageDescription))
}
if req.QuotedContext != "" {
@@ -208,6 +206,22 @@ func (s *sessionService) AgentQA(
return nil
}
func routeAgentImageInput(req *types.QARequest, modelSupportsVision bool) (string, []string, error) {
if req == nil {
return "", nil, errors.New("QA request is required")
}
if len(req.ImageURLs) == 0 {
return req.Query, nil, nil
}
if modelSupportsVision {
return req.Query, req.ImageURLs, nil
}
if strings.TrimSpace(req.ImageDescription) != "" {
return req.Query + "\n\n[用户上传图片内容]\n" + req.ImageDescription, nil, nil
}
return "", nil, errors.New("agent model does not support visual input and no image description is available")
}
// buildAgentConfig creates a runtime AgentConfig from the QARequest's custom agent configuration,
// tenant info, and resolved knowledge bases / search targets.
func (s *sessionService) buildAgentConfig(
internal/application/service/session_agent_vlm_fallback_test.go
+75
View File
@@ -0,0 +1,75 @@
package service
import (
"context"
"errors"
"testing"
"github.com/Tencent/WeKnora/internal/models/vlm"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
)
type agentImageVLMStub struct {
calls int
images [][]byte
result string
err error
}
func (s *agentImageVLMStub) Predict(_ context.Context, images [][]byte, _ string) (string, error) {
s.calls++
s.images = images
return s.result, s.err
}
func (*agentImageVLMStub) GetModelName() string { return "vlm-stub" }
func (*agentImageVLMStub) GetModelID() string { return "vlm-1" }
type agentImageModelServiceStub struct {
interfaces.ModelService
model vlm.VLM
err error
calls int
}
func (s *agentImageModelServiceStub) GetVLMModel(context.Context, string) (vlm.VLM, error) {
s.calls++
return s.model, s.err
}
func TestPrepareAgentImageInputUsesVLMOnlyForNonVisualPrimary(t *testing.T) {
visual := &agentImageVLMStub{result: "机械臂报警E01"}
models := &agentImageModelServiceStub{model: visual}
req := &types.QARequest{
Query: "什么故障？",
ImageURLs: []string{"data:image/png;base64,YWJj"},
CustomAgent: &types.CustomAgent{Config: types.CustomAgentConfig{VLMModelID: "vlm-1"}},
}
query, images, err := prepareAgentImageInput(context.Background(), req, false, models)
if err != nil {
t.Fatalf("prepareAgentImageInput: %v", err)
}
if len(images) != 0 || query == req.Query || req.ImageDescription != visual.result {
t.Fatalf("query=%q images=%v description=%q", query, images, req.ImageDescription)
}
if models.calls != 1 || visual.calls != 1 || len(visual.images) != 1 {
t.Fatalf("model calls=%d VLM calls=%d images=%d", models.calls, visual.calls, len(visual.images))
}
}
func TestPrepareAgentImageInputSkipsBrokenVLMForVisualPrimary(t *testing.T) {
models := &agentImageModelServiceStub{err: errors.New("stale VLM")}
req := &types.QARequest{
Query: "describe",
ImageURLs: []string{"data:image/png;base64,YWJj"},
CustomAgent: &types.CustomAgent{Config: types.CustomAgentConfig{VLMModelID: "broken"}},
}
primary := &captureChatModel{supportsImage: true}
query, images, err := prepareAgentImageInput(context.Background(), req, modelSupportsDirectImages(primary), models)
if err != nil {
t.Fatalf("visual primary should bypass VLM: %v", err)
}
if query != req.Query || len(images) != 1 || models.calls != 0 {
t.Fatalf("query=%q images=%v model_calls=%d", query, images, models.calls)
}
}
internal/application/service/session_knowledge_qa.go
+20 -5
View File
@@ -16,6 +16,12 @@ import (
"github.com/Tencent/WeKnora/internal/types"
)
// modelSupportsDirectImages asks the resolved concrete Chat transport. Unknown
// transports fail closed because chat.SupportsImageInput returns false.
func modelSupportsDirectImages(model chat.Chat) bool {
return chat.SupportsImageInput(model)
}
// KnowledgeQA performs knowledge base question answering with LLM summarization
// Events are emitted through eventBus (references, answer chunks, completion)
// customAgent is optional - if provided, uses custom agent configuration for multiTurnEnabled and historyTurns
@@ -68,12 +74,17 @@ func (s *sessionService) KnowledgeQA(
logger.Infof(ctx, "Fallback strategy not set, using default: %v", fallbackStrategy)
}
// Resolve chat model vision capability and VLM model ID for image routing
// Resolve the concrete Chat transport for image routing. Runtime transport
// capability is authoritative; stored metadata may be stale for existing
// WeKnora Cloud records.
var chatModelSupportsVision bool
var vlmModelID string
if chatModelID != "" {
if chatModelInfo, err := s.modelService.GetModelByID(ctx, chatModelID); err == nil && chatModelInfo != nil {
chatModelSupportsVision = chatModelInfo.Parameters.SupportsVision
if chatModelID != "" && len(req.ImageURLs) > 0 {
chatModel, modelErr := s.modelService.GetChatModel(ctx, chatModelID)
if modelErr != nil {
logger.Warnf(ctx, "Failed to initialize Chat model %s for image routing: %v", chatModelID, modelErr)
} else {
chatModelSupportsVision = modelSupportsDirectImages(chatModel)
}
}
if req.CustomAgent != nil {
@@ -136,7 +147,8 @@ func (s *sessionService) KnowledgeQA(
VLMModelID: vlmModelID,
ChatModelSupportsVision: chatModelSupportsVision,
Attachments: req.Attachments,
Language: types.LanguageNameFromContext(ctx),
Language: types.LanguageNameFromContext(ctx),
},
PipelineState: types.PipelineState{
RewriteQuery: req.Query,
@@ -158,6 +170,7 @@ func (s *sessionService) KnowledgeQA(
hasKB := len(knowledgeBaseIDs) > 0 || len(knowledgeIDs) > 0
needsRAG := hasKB || req.WebSearchEnabled
hasHistory := chatManage.MaxRounds > 0
needsVisualDescription := len(req.ImageURLs) > 0 && !chatModelSupportsVision
var pipeline []types.EventType
if !needsRAG {
@@ -178,6 +191,8 @@ func (s *sessionService) KnowledgeQA(
pipeline = types.NewPipelineBuilder().
AddIf(hasHistory, types.LOAD_HISTORY).
AddIf(chatManage.EnableMemory, types.MEMORY_RETRIEVAL).
AddIf(needsVisualDescription, types.QUERY_UNDERSTAND).
AddIf(needsVisualDescription, types.INTO_CHAT_MESSAGE).
Add(types.CHAT_COMPLETION_STREAM).
AddIf(chatManage.EnableMemory, types.MEMORY_STORAGE).
Build()
internal/application/service/session_knowledge_qa_test.go
+17 -7
View File
@@ -16,15 +16,19 @@ import (
)
type captureChatModel struct {
lastMessages []chat.Message
lastMessages []chat.Message
supportsImage bool
chatResponse *types.ChatResponse
chatErr error
}
func (m *captureChatModel) Chat(
context.Context,
[]chat.Message,
*chat.ChatOptions,
_ context.Context,
messages []chat.Message,
_ *chat.ChatOptions,
) (*types.ChatResponse, error) {
return nil, nil
m.lastMessages = append([]chat.Message(nil), messages...)
return m.chatResponse, m.chatErr
}
func (m *captureChatModel) ChatStream(
@@ -44,8 +48,14 @@ func (m *captureChatModel) ChatStream(
return ch, nil
}
func (m *captureChatModel) GetModelName() string { return "capture" }
func (m *captureChatModel) GetModelID() string { return "capture" }
func (m *captureChatModel) GetModelName() string { return "capture" }
func (m *captureChatModel) GetModelID() string { return "capture" }
func (m *captureChatModel) SupportsImageInput() bool { return m.supportsImage }
func TestModelSupportsDirectImagesUsesResolvedTransportAuthority(t *testing.T) {
assert.False(t, modelSupportsDirectImages(&captureChatModel{supportsImage: false}))
assert.True(t, modelSupportsDirectImages(&captureChatModel{supportsImage: true}))
}
type stubModelService struct {
chatModel chat.Chat
internal/im/adapter.go
+2
View File
@@ -64,6 +64,8 @@ type IncomingMessage struct {
FileName string
// FileSize is the file size in bytes (for file messages, optional).
FileSize int64
// ImageURLs carries transient model-ready image inputs (usually data URIs). Never log these values.
ImageURLs []string
// ThreadID is the platform-specific thread identifier.
// - Slack: thread_ts (top-level message uses its own timestamp)
// - Mattermost: root_id, or post_id if top-level
internal/im/dingtalk/adapter.go
+346 -37
View File
@@ -9,7 +9,12 @@ import (
"encoding/json"
"fmt"
"io"
"mime"
"net"
"net/http"
"net/netip"
"net/url"
"path/filepath"
"strconv"
"strings"
"sync"
@@ -20,6 +25,7 @@ import (
"github.com/Tencent/WeKnora/internal/im"
"github.com/Tencent/WeKnora/internal/logger"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
// httpClient is a shared HTTP client with a reasonable timeout for DingTalk API calls.
@@ -33,15 +39,18 @@ const dingtalkConvTypeGroup = "2"
// Compile-time checks.
var (
_ im.Adapter = (*Adapter)(nil)
_ im.StreamSender = (*Adapter)(nil)
_ im.Adapter = (*Adapter)(nil)
_ im.StreamSender = (*Adapter)(nil)
_ im.FileDownloader = (*Adapter)(nil)
)
// Adapter implements im.Adapter for DingTalk.
type Adapter struct {
clientID string
clientSecret string
cardTemplateID string // optional: enables AI card streaming when set
clientID string
clientSecret string
cardTemplateID string // optional: enables AI card streaming when set
messageFileDownloadURL string
allowInsecureDownload bool // tests only; production requires HTTPS
// accessToken cache
tokenMu sync.RWMutex
@@ -114,19 +123,20 @@ func (a *Adapter) VerifyCallback(c *gin.Context) error {
// DingTalk callback message structure.
type callbackMessage struct {
ConversationID string `json:"conversationId"`
ConversationType string `json:"conversationType"`
MsgID string `json:"msgId"`
Msgtype string `json:"msgtype"`
Text *textContent `json:"text"`
SenderNick string `json:"senderNick"`
SenderStaffId string `json:"senderStaffId"`
SenderID string `json:"senderId"`
SessionWebhook string `json:"sessionWebhook"`
RobotCode string `json:"robotCode"`
AtUsers []atUser `json:"atUsers"`
IsInAtList bool `json:"isInAtList"`
ChatbotCorpId string `json:"chatbotCorpId"`
ConversationID string `json:"conversationId"`
ConversationType string `json:"conversationType"`
MsgID string `json:"msgId"`
Msgtype string `json:"msgtype"`
Text *textContent `json:"text"`
Content json.RawMessage `json:"content"`
SenderNick string `json:"senderNick"`
SenderStaffId string `json:"senderStaffId"`
SenderID string `json:"senderId"`
SessionWebhook string `json:"sessionWebhook"`
RobotCode string `json:"robotCode"`
AtUsers []atUser `json:"atUsers"`
IsInAtList bool `json:"isInAtList"`
ChatbotCorpId string `json:"chatbotCorpId"`
}
type textContent struct {
@@ -170,6 +180,20 @@ func parseCallbackMessage(msg *callbackMessage) *im.IncomingMessage {
if msg.Text != nil {
content = strings.TrimSpace(msg.Text.Content)
}
messageType := im.MessageTypeText
fileKey := ""
fileName := ""
if strings.EqualFold(msg.Msgtype, "picture") {
messageType = im.MessageTypeImage
fileName = msg.MsgID + ".png"
content = ""
var picture struct {
DownloadCode string `json:"downloadCode"`
}
if json.Unmarshal(msg.Content, &picture) == nil {
fileKey = strings.TrimSpace(picture.DownloadCode)
}
}
return &im.IncomingMessage{
Platform: im.PlatformDingtalk,
@@ -178,14 +202,278 @@ func parseCallbackMessage(msg *callbackMessage) *im.IncomingMessage {
ChatID: chatID,
ChatType: chatType,
MessageID: msg.MsgID,
MessageType: im.MessageTypeText,
MessageType: messageType,
Content: content,
FileKey: fileKey,
FileName: fileName,
Extra: map[string]string{
"session_webhook": msg.SessionWebhook,
},
}
}
const (
dingtalkMessageFileDownloadURL = "https://api.dingtalk.com/v1.0/robot/messageFiles/download"
maxDingTalkAPIResponseBytes = 1 << 20
)
func (a *Adapter) newDingTalkAPIClient(base *http.Client) *http.Client {
client := *base
client.CheckRedirect = func(_ *http.Request, _ []*http.Request) error {
return http.ErrUseLastResponse
}
if !a.allowInsecureDownload {
transport := http.DefaultTransport.(*http.Transport).Clone()
transport.Proxy = nil
dialer := &net.Dialer{Timeout: 30 * time.Second, KeepAlive: 30 * time.Second}
transport.DialContext = newDingTalkSafeDialContext(net.DefaultResolver.LookupIPAddr, dialer.DialContext)
client.Transport = transport
}
return &client
}
func readBoundedDingTalkAPIResponse(resp *http.Response) ([]byte, error) {
if resp.ContentLength > maxDingTalkAPIResponseBytes {
return nil, fmt.Errorf("dingtalk API response exceeds size limit")
}
payload, err := io.ReadAll(io.LimitReader(resp.Body, maxDingTalkAPIResponseBytes+1))
if err != nil {
return nil, fmt.Errorf("read dingtalk API response: %w", err)
}
if len(payload) > maxDingTalkAPIResponseBytes {
return nil, fmt.Errorf("dingtalk API response exceeds size limit")
}
return payload, nil
}
func (a *Adapter) DownloadFile(ctx context.Context, msg *im.IncomingMessage) (io.ReadCloser, string, error) {
if msg == nil || strings.TrimSpace(msg.FileKey) == "" {
return nil, "", fmt.Errorf("download code is required")
}
token, err := a.getAccessToken(ctx)
if err != nil {
return nil, "", fmt.Errorf("get access token: %w", err)
}
body, err := json.Marshal(map[string]string{
"downloadCode": msg.FileKey,
"robotCode": a.clientID,
})
if err != nil {
return nil, "", fmt.Errorf("marshal message file request: %w", err)
}
apiURL := a.messageFileDownloadURL
if apiURL == "" {
apiURL = dingtalkMessageFileDownloadURL
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(body))
if err != nil {
return nil, "", fmt.Errorf("create message file request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("x-acs-dingtalk-access-token", token)
apiClient := a.newDingTalkAPIClient(httpClient)
resp, err := apiClient.Do(req)
if err != nil {
return nil, "", fmt.Errorf("request message file URL: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
return nil, "", fmt.Errorf("dingtalk message file API returned %d", resp.StatusCode)
}
payload, err := readBoundedDingTalkAPIResponse(resp)
if err != nil {
return nil, "", err
}
var result struct {
DownloadURL string `json:"downloadUrl"`
}
if err := json.Unmarshal(payload, &result); err != nil {
return nil, "", fmt.Errorf("decode message file response: %w", err)
}
if !a.allowInsecureDownload {
result.DownloadURL, err = normalizeDingTalkInitialDownloadURL(result.DownloadURL)
if err != nil {
return nil, "", err
}
}
if err := validateDingTalkDownloadURL(result.DownloadURL, a.allowInsecureDownload); err != nil {
return nil, "", err
}
downloadReq, err := http.NewRequestWithContext(ctx, http.MethodGet, result.DownloadURL, nil)
if err != nil {
return nil, "", fmt.Errorf("create image download request: %w", err)
}
downloadClient := *httpClient
if !a.allowInsecureDownload {
transport := http.DefaultTransport.(*http.Transport).Clone()
transport.Proxy = nil
dialer := &net.Dialer{Timeout: 30 * time.Second, KeepAlive: 30 * time.Second}
transport.DialContext = newDingTalkSafeDialContext(net.DefaultResolver.LookupIPAddr, dialer.DialContext)
downloadClient.Transport = transport
}
downloadClient.CheckRedirect = func(req *http.Request, via []*http.Request) error {
return validateDingTalkRedirect(req, via, a.allowInsecureDownload)
}
downloadResp, err := downloadClient.Do(downloadReq)
if err != nil {
if ctx.Err() != nil {
return nil, "", ctx.Err()
}
return nil, "", fmt.Errorf("download dingtalk image request failed")
}
if downloadResp.StatusCode < http.StatusOK || downloadResp.StatusCode >= http.StatusMultipleChoices {
downloadResp.Body.Close()
return nil, "", fmt.Errorf("dingtalk image download returned %d", downloadResp.StatusCode)
}
contentType := strings.ToLower(strings.TrimSpace(strings.Split(downloadResp.Header.Get("Content-Type"), ";")[0]))
if msg.MessageType == im.MessageTypeImage && contentType != "" && contentType != "application/octet-stream" && !strings.HasPrefix(contentType, "image/") {
downloadResp.Body.Close()
return nil, "", fmt.Errorf("dingtalk picture response is not an image")
}
return downloadResp.Body, dingtalkDownloadedFileName(downloadResp, msg, contentType), nil
}
var dingtalkBlockedPublicRanges = []netip.Prefix{
netip.MustParsePrefix("100.64.0.0/10"),
netip.MustParsePrefix("192.0.0.0/24"),
netip.MustParsePrefix("192.0.2.0/24"),
netip.MustParsePrefix("198.18.0.0/15"),
netip.MustParsePrefix("198.51.100.0/24"),
netip.MustParsePrefix("203.0.113.0/24"),
netip.MustParsePrefix("2001:db8::/32"),
}
func isSafeDingTalkDownloadIP(ip net.IP) bool {
if ip == nil || !secutils.IsPublicIP(ip) {
return false
}
addr, ok := netip.AddrFromSlice(ip)
if !ok {
return false
}
addr = addr.Unmap()
if !addr.IsGlobalUnicast() || addr.IsPrivate() || addr.IsLoopback() || addr.IsLinkLocalUnicast() || addr.IsUnspecified() || addr.IsMulticast() {
return false
}
for _, prefix := range dingtalkBlockedPublicRanges {
if prefix.Contains(addr) {
return false
}
}
return true
}
func newDingTalkSafeDialContext(
lookup func(context.Context, string) ([]net.IPAddr, error),
dial func(context.Context, string, string) (net.Conn, error),
) func(context.Context, string, string) (net.Conn, error) {
return func(ctx context.Context, network, addr string) (net.Conn, error) {
host, port, err := net.SplitHostPort(addr)
if err != nil {
return nil, fmt.Errorf("invalid dingtalk media address")
}
var ips []net.IPAddr
if literal := net.ParseIP(host); literal != nil {
ips = []net.IPAddr{{IP: literal}}
} else {
ips, err = lookup(ctx, host)
if err != nil {
return nil, fmt.Errorf("resolve dingtalk media host failed")
}
}
if len(ips) == 0 {
return nil, fmt.Errorf("dingtalk media host resolved without addresses")
}
for _, resolved := range ips {
if !isSafeDingTalkDownloadIP(resolved.IP) {
return nil, fmt.Errorf("dingtalk media host resolved to a blocked address")
}
}
var lastErr error
for _, resolved := range ips {
conn, dialErr := dial(ctx, network, net.JoinHostPort(resolved.IP.String(), port))
if dialErr == nil {
return conn, nil
}
lastErr = dialErr
}
return nil, fmt.Errorf("connect to validated dingtalk media address: %w", lastErr)
}
}
func validateDingTalkRedirect(req *http.Request, via []*http.Request, allowInsecure bool) error {
if len(via) >= 10 {
return fmt.Errorf("stopped after 10 dingtalk image redirects")
}
// Go may synthesize Referer from the credential-bearing source URL and copy
// headers across redirects. Media redirects need no inherited headers.
req.Header = make(http.Header)
return validateDingTalkDownloadURL(req.URL.String(), allowInsecure)
}
// normalizeDingTalkInitialDownloadURL upgrades the one-time media URL returned
// directly by DingTalk's authenticated message-file API. Some DingTalk media
// backends still return an http scheme even though the same signed resource is
// available over TLS. Only this initial trusted API response is upgraded;
// redirects remain HTTPS-only and are validated independently.
func normalizeDingTalkInitialDownloadURL(rawURL string) (string, error) {
u, err := url.Parse(strings.TrimSpace(rawURL))
if err != nil || u.Hostname() == "" || u.User != nil {
return "", fmt.Errorf("invalid dingtalk download URL")
}
switch u.Scheme {
case "http":
u.Scheme = "https"
case "https":
default:
return "", fmt.Errorf("invalid dingtalk download URL scheme")
}
return u.String(), nil
}
func validateDingTalkDownloadURL(rawURL string, allowInsecure bool) error {
u, err := url.Parse(strings.TrimSpace(rawURL))
if err != nil || u.Hostname() == "" || u.User != nil {
return fmt.Errorf("invalid dingtalk download URL")
}
if allowInsecure {
if u.Scheme != "http" && u.Scheme != "https" {
return fmt.Errorf("invalid dingtalk download URL scheme")
}
return nil
}
if u.Scheme != "https" {
return fmt.Errorf("dingtalk download URL must use HTTPS")
}
if ip := net.ParseIP(u.Hostname()); ip != nil && (ip.IsLoopback() || ip.IsPrivate() || ip.IsLinkLocalUnicast() || ip.IsUnspecified()) {
return fmt.Errorf("dingtalk download URL points to a private address")
}
if err := secutils.ValidateURLForSSRF(u.String()); err != nil {
return fmt.Errorf("dingtalk download URL failed SSRF validation: %w", err)
}
return nil
}
func dingtalkDownloadedFileName(resp *http.Response, msg *im.IncomingMessage, contentType string) string {
if _, params, err := mime.ParseMediaType(resp.Header.Get("Content-Disposition")); err == nil {
if name := filepath.Base(strings.TrimSpace(params["filename"])); name != "" && name != "." && name != ".." {
return name
}
}
if msg != nil && strings.TrimSpace(msg.FileName) != "" {
return filepath.Base(msg.FileName)
}
ext := ".bin"
if extensions, _ := mime.ExtensionsByType(contentType); len(extensions) > 0 {
ext = extensions[0]
}
if msg != nil && msg.MessageID != "" {
return msg.MessageID + ext
}
return "dingtalk-file" + ext
}
// ── Send reply ──
func (a *Adapter) SendReply(ctx context.Context, incoming *im.IncomingMessage, reply *im.ReplyMessage) error {
@@ -203,6 +491,9 @@ func (a *Adapter) SendReply(ctx context.Context, incoming *im.IncomingMessage, r
}
func (a *Adapter) replyViaSessionWebhook(ctx context.Context, webhookURL, content string) error {
if err := validateDingTalkDownloadURL(webhookURL, a.allowInsecureDownload); err != nil {
return fmt.Errorf("invalid DingTalk session webhook")
}
body := map[string]interface{}{
"msgtype": "markdown",
"markdown": map[string]string{
@@ -218,19 +509,22 @@ func (a *Adapter) replyViaSessionWebhook(ctx context.Context, webhookURL, conten
req, err := http.NewRequestWithContext(ctx, http.MethodPost, webhookURL, bytes.NewReader(jsonBody))
if err != nil {
return fmt.Errorf("create request: %w", err)
return fmt.Errorf("create DingTalk session webhook request")
}
req.Header.Set("Content-Type", "application/json")
resp, err := httpClient.Do(req)
client := a.newDingTalkAPIClient(httpClient)
resp, err := client.Do(req)
if err != nil {
return fmt.Errorf("send reply: %w", err)
return fmt.Errorf("send DingTalk session webhook failed")
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
respBody, _ := io.ReadAll(resp.Body)
return fmt.Errorf("dingtalk sessionWebhook returned %d: %s", resp.StatusCode, string(respBody))
return fmt.Errorf("dingtalk session webhook returned %d", resp.StatusCode)
}
if _, err := readBoundedDingTalkAPIResponse(resp); err != nil {
return err
}
return nil
}
@@ -273,15 +567,18 @@ func (a *Adapter) replyViaOpenAPI(ctx context.Context, incoming *im.IncomingMess
req.Header.Set("Content-Type", "application/json")
req.Header.Set("x-acs-dingtalk-access-token", token)
resp, err := httpClient.Do(req)
apiClient := a.newDingTalkAPIClient(httpClient)
resp, err := apiClient.Do(req)
if err != nil {
return fmt.Errorf("send reply: %w", err)
return fmt.Errorf("send reply failed")
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
respBody, _ := io.ReadAll(resp.Body)
return fmt.Errorf("dingtalk OpenAPI returned %d: %s", resp.StatusCode, string(respBody))
return fmt.Errorf("dingtalk OpenAPI returned %d", resp.StatusCode)
}
if _, err := readBoundedDingTalkAPIResponse(resp); err != nil {
return err
}
return nil
}
@@ -317,22 +614,26 @@ func (a *Adapter) getAccessToken(ctx context.Context) (string, error) {
}
req.Header.Set("Content-Type", "application/json")
resp, err := httpClient.Do(req)
apiClient := a.newDingTalkAPIClient(httpClient)
resp, err := apiClient.Do(req)
if err != nil {
return "", err
return "", fmt.Errorf("dingtalk access token request failed")
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
respBody, _ := io.ReadAll(resp.Body)
return "", fmt.Errorf("dingtalk accessToken returned %d: %s", resp.StatusCode, string(respBody))
return "", fmt.Errorf("dingtalk access token returned %d", resp.StatusCode)
}
payload, err := readBoundedDingTalkAPIResponse(resp)
if err != nil {
return "", err
}
var result struct {
AccessToken string `json:"accessToken"`
ExpireIn int64 `json:"expireIn"`
}
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
if err := json.Unmarshal(payload, &result); err != nil {
return "", fmt.Errorf("decode token response: %w", err)
}
if result.AccessToken == "" {
@@ -366,15 +667,19 @@ func (a *Adapter) dingtalkAPI(ctx context.Context, method, path string, body int
req.Header.Set("Content-Type", "application/json")
req.Header.Set("x-acs-dingtalk-access-token", token)
resp, err := httpClient.Do(req)
apiClient := a.newDingTalkAPIClient(httpClient)
resp, err := apiClient.Do(req)
if err != nil {
return nil, fmt.Errorf("do request: %w", err)
return nil, fmt.Errorf("dingtalk API request failed")
}
defer resp.Body.Close()
respBody, _ := io.ReadAll(resp.Body)
respBody, readErr := readBoundedDingTalkAPIResponse(resp)
if readErr != nil {
return nil, readErr
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("dingtalk API %s returned %d: %s", path, resp.StatusCode, string(respBody))
return nil, fmt.Errorf("dingtalk API %s returned %d", path, resp.StatusCode)
}
return respBody, nil
}
@@ -573,6 +878,7 @@ func (a *Adapter) FinalizeStream(ctx context.Context, incoming *im.IncomingMessa
if outTrackID != "" {
if err := a.streamingUpdateCard(ctx, outTrackID, finalContent, false); err != nil {
logger.Warnf(ctx, "[DingTalk] Failed to finalize card stream: %v", err)
return fmt.Errorf("finalize DingTalk card stream: %w", err)
}
}
return nil
@@ -601,14 +907,17 @@ func (a *Adapter) EndStream(ctx context.Context, incoming *im.IncomingMessage, s
if outTrackID != "" {
if err := a.streamingUpdateCard(ctx, outTrackID, fullContent, true); err != nil {
logger.Warnf(ctx, "[DingTalk] Failed to finalize card stream: %v", err)
return fmt.Errorf("end DingTalk card stream: %w", err)
}
} else if sessionWebhook != "" {
if err := a.replyViaSessionWebhook(ctx, sessionWebhook, fullContent); err != nil {
logger.Warnf(ctx, "[DingTalk] Failed to end stream: %v", err)
return fmt.Errorf("end DingTalk session webhook stream: %w", err)
}
} else {
if err := a.replyViaOpenAPI(ctx, incoming, fullContent); err != nil {
logger.Warnf(ctx, "[DingTalk] Failed to end stream via OpenAPI: %v", err)
return fmt.Errorf("end DingTalk OpenAPI stream: %w", err)
}
}
internal/im/dingtalk/adapter_security_edge_test.go
+45
View File
@@ -0,0 +1,45 @@
package dingtalk
import (
"context"
"net/http"
"net/http/httptest"
"strings"
"testing"
)
func TestDingTalkRedirectClearsRefererAndInheritedHeaders(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, "https://cdn-two.example/image.png", nil)
req.Header.Set("Referer", "https://cdn-one.example/image.png?signature=TOPSECRET")
req.Header.Set("Authorization", "Bearer TOPSECRET")
if err := validateDingTalkRedirect(req, nil, true); err != nil {
t.Fatalf("validateDingTalkRedirect: %v", err)
}
if len(req.Header) != 0 {
t.Fatalf("redirect headers = %#v, want empty", req.Header)
}
}
func TestReplyViaSessionWebhookRejectsRedirectWithoutLeakingURL(t *testing.T) {
var reached bool
target := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { reached = true }))
defer target.Close()
source := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Location", target.URL)
w.WriteHeader(http.StatusTemporaryRedirect)
}))
defer source.Close()
a := &Adapter{allowInsecureDownload: true}
signedURL := source.URL + "/hook?access_token=TOPSECRET"
err := a.replyViaSessionWebhook(context.Background(), signedURL, "reply")
if err == nil {
t.Fatal("expected webhook redirect rejection")
}
if reached {
t.Fatal("session webhook followed redirect")
}
if strings.Contains(err.Error(), "TOPSECRET") || strings.Contains(err.Error(), signedURL) {
t.Fatalf("credential URL leaked in error: %v", err)
}
}
internal/im/dingtalk/adapter_test.go
+276
View File
@@ -0,0 +1,276 @@
package dingtalk
import (
"context"
"errors"
"io"
"net"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/Tencent/WeKnora/internal/im"
)
func TestParseCallbackMessagePicture(t *testing.T) {
msg := &callbackMessage{
ConversationID: "cid-group",
ConversationType: dingtalkConvTypeGroup,
MsgID: "msg-picture-1",
Msgtype: "picture",
Content: []byte(`{"downloadCode":"download-code-1"}`),
SenderStaffId: "staff-1",
SenderNick: "Alice",
SessionWebhook: "https://example.invalid/session",
}
got := parseCallbackMessage(msg)
if got.MessageType != im.MessageTypeImage {
t.Fatalf("MessageType = %q, want image", got.MessageType)
}
if got.FileKey != "download-code-1" {
t.Fatalf("FileKey = %q, want download code", got.FileKey)
}
if got.FileName != "msg-picture-1.png" {
t.Fatalf("FileName = %q", got.FileName)
}
if got.Content != "" {
t.Fatalf("Content = %q, want empty", got.Content)
}
}
func TestAdapterDownloadFileUsesMessageFilesAPI(t *testing.T) {
var sawToken bool
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/v1.0/robot/messageFiles/download":
sawToken = r.Header.Get("x-acs-dingtalk-access-token") == "cached-access-token"
body, _ := io.ReadAll(r.Body)
if !strings.Contains(string(body), `"downloadCode":"download-code-1"`) || !strings.Contains(string(body), `"robotCode":"robot-code-1"`) {
t.Errorf("unexpected request body: %s", body)
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"downloadUrl":"` + serverURLForRequest(r) + `/media/picture"}`))
case "/media/picture":
w.Header().Set("Content-Type", "image/png")
w.Header().Set("Content-Disposition", `attachment; filename="question.png"`)
_, _ = w.Write([]byte("PNGDATA"))
default:
http.NotFound(w, r)
}
}))
defer server.Close()
adapter := NewAdapter("robot-code-1", "secret", "")
adapter.token = "cached-access-token"
adapter.tokenExpAt = time.Now().Add(time.Hour)
adapter.messageFileDownloadURL = server.URL + "/v1.0/robot/messageFiles/download"
adapter.allowInsecureDownload = true
reader, fileName, err := adapter.DownloadFile(context.Background(), &im.IncomingMessage{
MessageType: im.MessageTypeImage,
MessageID: "msg-picture-1",
FileKey: "download-code-1",
})
if err != nil {
t.Fatalf("DownloadFile: %v", err)
}
defer reader.Close()
body, _ := io.ReadAll(reader)
if string(body) != "PNGDATA" {
t.Fatalf("body = %q", body)
}
if fileName != "question.png" {
t.Fatalf("fileName = %q", fileName)
}
if !sawToken {
t.Fatal("missing DingTalk access token header")
}
}
func serverURLForRequest(r *http.Request) string {
return "http://" + r.Host
}
func TestNormalizeDingTalkInitialDownloadURLUpgradesHTTPToHTTPS(t *testing.T) {
rawURL := "http://media.example.com/path/image.png?signature=abc%2F123&name=a+b.png"
got, err := normalizeDingTalkInitialDownloadURL(rawURL)
if err != nil {
t.Fatalf("normalizeDingTalkInitialDownloadURL: %v", err)
}
want := "https://media.example.com/path/image.png?signature=abc%2F123&name=a+b.png"
if got != want {
t.Fatalf("normalized URL = %q, want %q", got, want)
}
}
func TestValidateDingTalkDownloadURLRejectsUnsafeTargets(t *testing.T) {
for _, rawURL := range []string{
"http://example.com/image.png",
"https://user:pass@example.com/image.png",
"https://127.0.0.1/image.png",
"https://10.0.0.8/image.png",
} {
t.Run(rawURL, func(t *testing.T) {
if err := validateDingTalkDownloadURL(rawURL, false); err == nil {
t.Fatalf("unsafe URL accepted: %s", rawURL)
}
})
}
}
func TestParseCallbackMessagePictureWithoutDownloadCodeStaysImage(t *testing.T) {
for _, content := range [][]byte{[]byte(`{}`), []byte(`not-json`)} {
msg := parseCallbackMessage(&callbackMessage{MsgID: "bad-picture", Msgtype: "picture", Content: content})
if msg.MessageType != im.MessageTypeImage {
t.Fatalf("MessageType = %q, want image", msg.MessageType)
}
if msg.FileKey != "" {
t.Fatalf("FileKey = %q, want empty", msg.FileKey)
}
}
}
func TestAdapterDownloadFileDoesNotFollowCredentialedAPIRedirect(t *testing.T) {
var redirectedToken string
target := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
redirectedToken = r.Header.Get("x-acs-dingtalk-access-token")
}))
defer target.Close()
api := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Location", target.URL)
w.WriteHeader(http.StatusTemporaryRedirect)
}))
defer api.Close()
adapter := NewAdapter("robot-code-1", "secret", "")
adapter.token = "cached-access-token"
adapter.tokenExpAt = time.Now().Add(time.Hour)
adapter.messageFileDownloadURL = api.URL
adapter.allowInsecureDownload = true
_, _, err := adapter.DownloadFile(context.Background(), &im.IncomingMessage{MessageType: im.MessageTypeImage, FileKey: "code"})
if err == nil {
t.Fatal("expected redirect rejection")
}
if redirectedToken != "" {
t.Fatal("access token leaked across redirect")
}
}
func TestAdapterDownloadFileAcceptsGenericMediaContentType(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/api" {
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"downloadUrl":"` + serverURLForRequest(r) + `/media"}`))
return
}
w.Header().Set("Content-Type", "application/octet-stream")
_, _ = w.Write([]byte("PNGDATA"))
}))
defer server.Close()
adapter := NewAdapter("robot", "secret", "")
adapter.token = "token"
adapter.tokenExpAt = time.Now().Add(time.Hour)
adapter.messageFileDownloadURL = server.URL + "/api"
adapter.allowInsecureDownload = true
body, _, err := adapter.DownloadFile(context.Background(), &im.IncomingMessage{MessageType: im.MessageTypeImage, FileKey: "code"})
if err != nil {
t.Fatalf("DownloadFile: %v", err)
}
body.Close()
}
func TestDingTalkAPIClientDisablesProxyAndRedirects(t *testing.T) {
a := &Adapter{}
client := a.newDingTalkAPIClient(httpClient)
transport, ok := client.Transport.(*http.Transport)
if !ok {
t.Fatalf("transport type = %T, want *http.Transport", client.Transport)
}
if transport.Proxy != nil {
t.Fatal("DingTalk API client must ignore environment proxies")
}
req := httptest.NewRequest(http.MethodPost, "https://api.dingtalk.com/redirect", nil)
if err := client.CheckRedirect(req, nil); !errors.Is(err, http.ErrUseLastResponse) {
t.Fatalf("redirect error = %v, want http.ErrUseLastResponse", err)
}
}
func TestDingTalkRedirectLimitFailsClosed(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, "https://example.com/image.png", nil)
via := make([]*http.Request, 10)
if err := validateDingTalkRedirect(req, via, false); err == nil {
t.Fatal("validateDingTalkRedirect should reject more than 10 requests")
}
}
func TestDingTalkSafeDialUsesValidatedConcreteIP(t *testing.T) {
stop := errors.New("stop after address capture")
var dialed string
dial := newDingTalkSafeDialContext(
func(context.Context, string) ([]net.IPAddr, error) {
return []net.IPAddr{{IP: net.ParseIP("1.1.1.1")}}, nil
},
func(_ context.Context, _ string, addr string) (net.Conn, error) {
dialed = addr
return nil, stop
},
)
_, err := dial(context.Background(), "tcp", "cdn.example.test:443")
if !errors.Is(err, stop) {
t.Fatalf("dial error = %v, want sentinel", err)
}
if dialed != "1.1.1.1:443" {
t.Fatalf("dialed address = %q, want validated concrete IP", dialed)
}
}
func TestDingTalkSafeDialRejectsEveryRestrictedResolvedIP(t *testing.T) {
for _, rawIP := range []string{
"0.0.0.1",
"240.0.0.1",
"2001:0000::1",
"2002:7f00:0001::1",
} {
t.Run(rawIP, func(t *testing.T) {
dialCalled := false
dial := newDingTalkSafeDialContext(
func(context.Context, string) ([]net.IPAddr, error) {
return []net.IPAddr{{IP: net.ParseIP(rawIP)}}, nil
},
func(context.Context, string, string) (net.Conn, error) {
dialCalled = true
return nil, errors.New("must not dial")
},
)
if _, err := dial(context.Background(), "tcp", "cdn.example.test:443"); err == nil {
t.Fatal("restricted resolved IP was accepted")
}
if dialCalled {
t.Fatal("restricted resolved IP reached network dial")
}
})
}
}
func TestAdapterDownloadFileRedactsSignedURLFromErrors(t *testing.T) {
api := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"downloadUrl":"http://127.0.0.1:1/media?X-Amz-Signature=TOPSECRET"}`))
}))
defer api.Close()
adapter := NewAdapter("robot", "secret", "")
adapter.token = "token"
adapter.tokenExpAt = time.Now().Add(time.Hour)
adapter.messageFileDownloadURL = api.URL
adapter.allowInsecureDownload = true
_, _, err := adapter.DownloadFile(context.Background(), &im.IncomingMessage{MessageType: im.MessageTypeImage, FileKey: "code"})
if err == nil {
t.Fatal("expected download error")
}
if strings.Contains(err.Error(), "TOPSECRET") || strings.Contains(err.Error(), "X-Amz-Signature") {
t.Fatalf("signed URL leaked in error: %v", err)
}
}
internal/im/dingtalk/finalize_fail_closed_test.go
+43
View File
@@ -0,0 +1,43 @@
package dingtalk
import (
"context"
"errors"
"net/http"
"testing"
"time"
"github.com/Tencent/WeKnora/internal/im"
)
type failingCardRoundTripper struct{}
func (failingCardRoundTripper) RoundTrip(*http.Request) (*http.Response, error) {
return nil, errors.New("simulated delivery failure")
}
func TestFinalizeStreamPropagatesCardDeliveryFailure(t *testing.T) {
oldClient := httpClient
httpClient = &http.Client{Transport: failingCardRoundTripper{}, Timeout: time.Second}
t.Cleanup(func() { httpClient = oldClient })
adapter := &Adapter{
clientID: "robot",
token: "cached-token",
tokenExpAt: time.Now().Add(time.Hour),
}
streamID := "dt:fail-closed-test"
streamsMu.Lock()
dStreams[streamID] = &streamState{outTrackID: "track-id", createdAt: time.Now()}
streamsMu.Unlock()
t.Cleanup(func() {
streamsMu.Lock()
delete(dStreams, streamID)
streamsMu.Unlock()
})
err := adapter.FinalizeStream(context.Background(), &im.IncomingMessage{}, streamID, "image recognition failed")
if err == nil {
t.Fatal("expected final card delivery failure to propagate")
}
}
internal/im/dingtalk/longconn.go
+37 -2
View File
@@ -2,6 +2,7 @@ package dingtalk
import (
"context"
"encoding/json"
"strings"
"github.com/open-dingtalk/dingtalk-stream-sdk-go/chatbot"
@@ -71,7 +72,11 @@ func (c *LongConnClient) onChatBotMessage(ctx context.Context, data *chatbot.Bot
userID = data.SenderId
}
content := strings.TrimSpace(data.Text.Content)
messageType, content, fileKey := parseDingTalkStreamPayload(data.Msgtype, data.Content, data.Text.Content)
fileName := ""
if messageType == im.MessageTypeImage {
fileName = data.MsgId + ".png"
}
incoming := &im.IncomingMessage{
Platform: im.PlatformDingtalk,
@@ -80,8 +85,10 @@ func (c *LongConnClient) onChatBotMessage(ctx context.Context, data *chatbot.Bot
ChatID: chatID,
ChatType: chatType,
MessageID: data.MsgId,
MessageType: im.MessageTypeText,
MessageType: messageType,
Content: content,
FileKey: fileKey,
FileName: fileName,
Extra: map[string]string{
"session_webhook": data.SessionWebhook,
},
@@ -93,3 +100,31 @@ func (c *LongConnClient) onChatBotMessage(ctx context.Context, data *chatbot.Bot
return nil, nil
}
func parseDingTalkStreamPayload(msgType string, rawContent interface{}, text string) (im.MessageType, string, string) {
if !strings.EqualFold(strings.TrimSpace(msgType), "picture") {
return im.MessageTypeText, strings.TrimSpace(text), ""
}
var payload []byte
switch value := rawContent.(type) {
case string:
payload = []byte(value)
case []byte:
payload = value
case json.RawMessage:
payload = value
default:
var err error
payload, err = json.Marshal(rawContent)
if err != nil {
return im.MessageTypeImage, "", ""
}
}
var picture struct {
DownloadCode string `json:"downloadCode"`
}
if json.Unmarshal(payload, &picture) != nil {
return im.MessageTypeImage, "", ""
}
return im.MessageTypeImage, "", strings.TrimSpace(picture.DownloadCode)
}
internal/im/dingtalk/longconn_test.go
+37
View File
@@ -0,0 +1,37 @@
package dingtalk
import (
"testing"
"github.com/Tencent/WeKnora/internal/im"
)
func TestParseDingTalkStreamPayloadPicture(t *testing.T) {
messageType, content, fileKey := parseDingTalkStreamPayload("picture", map[string]interface{}{
"downloadCode": "stream-download-code",
}, "ignored text")
if messageType != im.MessageTypeImage {
t.Fatalf("MessageType = %q, want image", messageType)
}
if content != "" {
t.Fatalf("content = %q, want empty", content)
}
if fileKey != "stream-download-code" {
t.Fatalf("fileKey = %q", fileKey)
}
}
func TestParseDingTalkStreamPayloadText(t *testing.T) {
messageType, content, fileKey := parseDingTalkStreamPayload("text", nil, " hello ")
if messageType != im.MessageTypeText || content != "hello" || fileKey != "" {
t.Fatalf("got type=%q content=%q fileKey=%q", messageType, content, fileKey)
}
}
func TestParseDingTalkStreamPayloadPictureJSONString(t *testing.T) {
messageType, _, fileKey := parseDingTalkStreamPayload("picture", `{"downloadCode":"string-code"}`, "")
if messageType != im.MessageTypeImage || fileKey != "string-code" {
t.Fatalf("got type=%q fileKey=%q", messageType, fileKey)
}
}
internal/im/image_qa_test.go
+132
View File
@@ -0,0 +1,132 @@
package im
import (
"bytes"
"context"
"errors"
"image"
"image/color"
"image/png"
"io"
"strings"
"testing"
"github.com/Tencent/WeKnora/internal/types"
)
type staticImageDownloader struct {
body string
fileName string
}
func (d staticImageDownloader) DownloadFile(context.Context, *IncomingMessage) (io.ReadCloser, string, error) {
return io.NopCloser(strings.NewReader(d.body)), d.fileName, nil
}
func testPNG(t *testing.T, width, height int) string {
t.Helper()
img := image.NewNRGBA(image.Rect(0, 0, width, height))
img.SetNRGBA(0, 0, color.NRGBA{R: 20, G: 120, B: 240, A: 255})
var buf bytes.Buffer
if err := png.Encode(&buf, img); err != nil {
t.Fatalf("encode png: %v", err)
}
return buf.String()
}
func TestPrepareIMImageForQAAndBuildRequest(t *testing.T) {
png := testPNG(t, 32, 32)
msg := &IncomingMessage{
Platform: PlatformDingtalk,
MessageType: MessageTypeImage,
MessageID: "picture-message-1",
FileKey: "temporary-download-code",
}
if err := prepareIMImageForQA(context.Background(), msg, staticImageDownloader{body: png, fileName: "question.png"}); err != nil {
t.Fatalf("prepareIMImageForQA: %v", err)
}
if len(msg.ImageURLs) != 1 || !strings.HasPrefix(msg.ImageURLs[0], "data:image/png;base64,") {
t.Fatalf("ImageURLs = %#v", msg.ImageURLs)
}
if msg.Content == "" {
t.Fatal("image-only message should receive a default visual question")
}
req := buildIMQARequest(&types.Session{ID: "session-1"}, msg.Content, "assistant-1", "user-1", nil, nil, nil, msg.ImageURLs)
if len(req.ImageURLs) != 1 || req.ImageURLs[0] != msg.ImageURLs[0] {
t.Fatalf("request ImageURLs = %#v", req.ImageURLs)
}
}
func TestPrepareIMImageForQARejectsInvalidOrOversizedPayload(t *testing.T) {
tests := []struct {
name string
body string
want string
}{
{name: "not an image", body: strings.Repeat("plain text", 80), want: "unsupported image content type"},
{name: "oversized", body: strings.Repeat("x", maxIMImageBytes+1), want: "exceeds 10 MiB limit"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
msg := &IncomingMessage{Platform: PlatformDingtalk, MessageType: MessageTypeImage, MessageID: "image-1"}
err := prepareIMImageForQA(context.Background(), msg, staticImageDownloader{body: tt.body, fileName: "payload.bin"})
if err == nil || !strings.Contains(err.Error(), tt.want) {
t.Fatalf("error = %v, want containing %q", err, tt.want)
}
if len(msg.ImageURLs) != 0 {
t.Fatal("rejected image must not reach the model")
}
})
}
t.Run("dimensions", func(t *testing.T) {
msg := &IncomingMessage{MessageType: MessageTypeImage}
err := prepareIMImageForQA(context.Background(), msg, staticImageDownloader{body: testPNG(t, maxIMImageDimension+1, 1), fileName: "wide.png"})
if err == nil || !strings.Contains(err.Error(), "dimensions") {
t.Fatalf("error = %v, want dimensions rejection", err)
}
})
t.Run("animated formats rejected", func(t *testing.T) {
msg := &IncomingMessage{MessageType: MessageTypeImage}
err := prepareIMImageForQA(context.Background(), msg, staticImageDownloader{body: "GIF89a" + strings.Repeat("x", 600), fileName: "animated.gif"})
if err == nil || !strings.Contains(err.Error(), "only JPEG and PNG") {
t.Fatalf("error = %v, want JPEG/PNG rejection", err)
}
})
}
func TestIMQAErrorReplyDistinguishesImageFailures(t *testing.T) {
imageReply := imQAErrorReply([]string{"data:image/png;base64,abc"})
if !strings.Contains(imageReply, "图片识别失败") || !strings.Contains(imageReply, "视觉模型") {
t.Fatalf("image error reply = %q", imageReply)
}
if got := imQAErrorReply(nil); got != "抱歉，处理您的问题时出现了异常，请稍后再试。" {
t.Fatalf("text error reply = %q", got)
}
}
func TestImageStreamDoesNotPublishPartialContent(t *testing.T) {
if allowIncrementalIMFlush(&IncomingMessage{ImageURLs: []string{"data:image/png;base64,YWJj"}}) {
t.Fatal("image streams must buffer until successful completion")
}
if !allowIncrementalIMFlush(&IncomingMessage{Content: "hello"}) {
t.Fatal("text streams must preserve incremental updates")
}
}
func TestFailClosedIMImageAnswerDiscardsPartialOutput(t *testing.T) {
partial := "I think the image might show"
err := errors.New("multimodal request rejected")
got, overridden := failClosedIMImageAnswer(partial, err, []string{"data:image/png;base64,abc"})
if !overridden || strings.Contains(got, partial) || !strings.Contains(got, "图片识别失败") {
t.Fatalf("image override = %q overridden=%v", got, overridden)
}
got, overridden = failClosedIMImageAnswer(partial, err, nil)
if overridden || got != partial {
t.Fatalf("text answer changed = %q overridden=%v", got, overridden)
}
}
internal/im/service.go
+165 -29
View File
@@ -3,10 +3,15 @@ package im
import (
"bytes"
"context"
"encoding/base64"
"errors"
"fmt"
"image"
_ "image/jpeg"
_ "image/png"
"io"
"mime/multipart"
"net/http"
"net/textproto"
"os"
"regexp"
@@ -483,12 +488,17 @@ func buildIMQARequest(
customAgent *types.CustomAgent,
kbIDs []string,
quote *QuotedMessage,
imageURLs ...[]string,
) *types.QARequest {
// WebSearchEnabled: the web handler passes this per-request from the
// frontend toggle; for IM channels the user has no per-message toggle,
// so we derive it from the agent config (the single source of truth).
webSearchEnabled := customAgent != nil && customAgent.Config.WebSearchEnabled
quotedContext := formatQuotedContext(quote)
var images []string
if len(imageURLs) > 0 {
images = append([]string(nil), imageURLs[0]...)
}
return &types.QARequest{
Session: session,
Query: query,
@@ -498,9 +508,78 @@ func buildIMQARequest(
UserMessageID: userMessageID,
WebSearchEnabled: webSearchEnabled,
QuotedContext: quotedContext,
ImageURLs: images,
}
}
const (
maxIMImageBytes = 10 << 20
maxIMImageDimension = 16_384
maxIMImagePixels = 40_000_000
)
func prepareIMImageForQA(ctx context.Context, msg *IncomingMessage, downloader FileDownloader) error {
if msg == nil || downloader == nil {
return fmt.Errorf("image message and downloader are required")
}
reader, fileName, err := downloader.DownloadFile(ctx, msg)
if err != nil {
return fmt.Errorf("download image: %w", err)
}
defer reader.Close()
data, err := io.ReadAll(io.LimitReader(reader, maxIMImageBytes+1))
if err != nil {
return fmt.Errorf("read image: %w", err)
}
if len(data) == 0 {
return fmt.Errorf("downloaded image is empty")
}
if len(data) > maxIMImageBytes {
return fmt.Errorf("image exceeds %d MiB limit", maxIMImageBytes>>20)
}
contentType := strings.ToLower(http.DetectContentType(data))
switch contentType {
case "image/jpeg", "image/png":
case "image/gif", "image/webp":
return fmt.Errorf("only JPEG and PNG images are supported")
default:
return fmt.Errorf("unsupported image content type %q", contentType)
}
config, _, err := image.DecodeConfig(bytes.NewReader(data))
if err != nil {
return fmt.Errorf("decode image metadata: %w", err)
}
if config.Width <= 0 || config.Height <= 0 || config.Width > maxIMImageDimension || config.Height > maxIMImageDimension || int64(config.Width)*int64(config.Height) > maxIMImagePixels {
return fmt.Errorf("image dimensions %dx%d exceed safety limits", config.Width, config.Height)
}
msg.FileName = fileName
msg.ImageURLs = []string{"data:" + contentType + ";base64," + base64.StdEncoding.EncodeToString(data)}
if strings.TrimSpace(msg.Content) == "" {
msg.Content = "请识别这张图片的内容，并回答图片中的问题。"
}
return nil
}
func imQAErrorReply(imageURLs []string) string {
if len(imageURLs) > 0 {
return "图片识别失败：当前模型可能不支持视觉输入，请检查智能体的图片上传开关和视觉模型配置。"
}
return "抱歉，处理您的问题时出现了异常，请稍后再试。"
}
func allowIncrementalIMFlush(msg *IncomingMessage) bool {
return msg == nil || len(msg.ImageURLs) == 0
}
func failClosedIMImageAnswer(answer string, err error, imageURLs []string) (string, bool) {
if err != nil && len(imageURLs) > 0 {
return imQAErrorReply(imageURLs), true
}
return answer, false
}
func buildIMLastRequestState(agentID string, customAgent *types.CustomAgent, kbIDs []string) *types.SessionLastRequestState {
state := &types.SessionLastRequestState{
AgentID: agentID,
@@ -1232,21 +1311,17 @@ func (s *Service) HandleMessage(ctx context.Context, msg *IncomingMessage, chann
logger.Infof(ctx, "[IM] HandleMessage: channel=%s platform=%s user=%s chat=%s msgtype=%s content_len=%d",
channelID, msg.Platform, msg.UserID, msg.ChatID, msg.MessageType, len(msg.Content))
logger.Debugf(ctx, "[IM] HandleMessage detail: msgid=%s filekey=%s filename=%s",
msg.MessageID, msg.FileKey, msg.FileName)
logger.Debugf(ctx, "[IM] HandleMessage detail: msgid=%s has_filekey=%t filename=%s",
msg.MessageID, strings.TrimSpace(msg.FileKey) != "", msg.FileName)
// ── File/Image message shortcut ──
// If the message is a file or image and the channel has a knowledge_base_id configured,
// handle it separately without entering the QA pipeline.
if (msg.MessageType == MessageTypeFile || msg.MessageType == MessageTypeImage) && channel.KnowledgeBaseID != "" {
// Files keep the legacy knowledge-base ingestion path. Images are routed
// below either to multimodal QA or, when vision is disabled, to this path.
if msg.MessageType == MessageTypeFile && channel.KnowledgeBaseID != "" {
return s.handleFileMessage(ctx, msg, adapter, channel)
}
// ── Non-text message without text content ──
// If the message is an image/file/video but has no text content, the QA pipeline
// cannot do anything useful (no vision support in IM yet). Sending an empty query
// to KB retrieval would return irrelevant results and cause hallucination.
if msg.Content == "" && (msg.MessageType == MessageTypeImage || msg.MessageType == MessageTypeFile) {
// Non-text files without a configured knowledge base still cannot enter QA.
if msg.Content == "" && msg.MessageType == MessageTypeFile {
logger.Infof(ctx, "[IM] Skipping QA for non-text message without content: type=%s", msg.MessageType)
if err := adapter.SendReply(ctx, msg, &ReplyMessage{
Content: "当前渠道未配置文件知识库，无法处理图片/文件消息。请在渠道设置中配置文件知识库后再发送，或直接用文字描述您的问题。",
@@ -1282,6 +1357,24 @@ func (s *Service) HandleMessage(ctx context.Context, msg *IncomingMessage, chann
}
}
if msg.MessageType == MessageTypeImage {
if customAgent == nil || !customAgent.Config.ImageUploadEnabled {
return adapter.SendReply(ctx, msg, &ReplyMessage{
Content: "当前机器人未启用图片识别，请在智能体设置中开启图片上传并配置视觉模型。",
IsFinal: true,
})
}
if strings.TrimSpace(msg.FileKey) == "" {
return adapter.SendReply(ctx, msg, &ReplyMessage{Content: "图片消息缺少有效下载凭证，请重新发送图片。", IsFinal: true})
}
if _, ok := adapter.(FileDownloader); !ok {
return adapter.SendReply(ctx, msg, &ReplyMessage{Content: "当前钉钉渠道暂不支持图片下载。", IsFinal: true})
}
if strings.TrimSpace(msg.Content) == "" {
msg.Content = "请识别这张图片的内容，并回答图片中的问题。"
}
}
// ── Slash-command dispatch ──
// Commands are handled before the QA pipeline so they respond instantly.
if cmd, args, ok := s.cmdRegistry.Parse(msg.Content); ok {
@@ -1415,6 +1508,21 @@ func (s *Service) executeQARequest(req *qaRequest) {
// runQA after the assistant message is created (that's when we have the
// sessionID + messageID needed to poll StreamManager).
// Keep large image bytes out of the bounded queue. Download, validate and
// encode only after this worker owns a concurrency slot.
if req.msg.MessageType == MessageTypeImage && len(req.msg.ImageURLs) == 0 {
downloader, ok := req.adapter.(FileDownloader)
if !ok {
_ = req.adapter.SendReply(ctx, req.msg, &ReplyMessage{Content: "当前钉钉渠道暂不支持图片下载。", IsFinal: true})
return
}
if err := prepareIMImageForQA(ctx, req.msg, downloader); err != nil {
logger.Warnf(ctx, "[IM] Failed to prepare image for QA: %v", err)
_ = req.adapter.SendReply(ctx, req.msg, &ReplyMessage{Content: "图片下载或格式校验失败，请发送10MB以内的 JPG 或 PNG 图片。", IsFinal: true})
return
}
}
// kbIDs is left empty so the QA pipeline resolves them from the agent config.
var kbIDs []string
@@ -1432,10 +1540,10 @@ func (s *Service) executeQARequest(req *qaRequest) {
}
// Non-streaming fallback: collect full answer then send.
answer, err := s.runQA(ctx, req.session, req.msg.Content, req.agent, kbIDs, req.userKey, req.msg.Quote)
answer, err := s.runQA(ctx, req.session, req.msg.Content, req.agent, kbIDs, req.userKey, req.msg.Quote, req.msg.ImageURLs)
if err != nil {
logger.Errorf(ctx, "[IM] QA failed: %v, sending fallback reply", err)
answer = "抱歉，处理您的问题时出现了异常，请稍后再试。"
answer = imQAErrorReply(req.msg.ImageURLs)
}
reply := &ReplyMessage{
@@ -2132,7 +2240,7 @@ func (s *Service) handleMessageStream(ctx context.Context, msg *IncomingMessage,
// Run QA async
go func() {
var err error
req := buildIMQARequest(session, msg.Content, assistantMsg.ID, userMsg.ID, customAgent, kbIDs, msg.Quote)
req := buildIMQARequest(session, msg.Content, assistantMsg.ID, userMsg.ID, customAgent, kbIDs, msg.Quote, msg.ImageURLs)
if req.QuotedContext != "" {
logger.Debugf(qaCtx, "[IM] QuotedContext set: length=%d", len(req.QuotedContext))
}
@@ -2158,6 +2266,9 @@ func (s *Service) handleMessageStream(ctx context.Context, msg *IncomingMessage,
defer ticker.Stop()
flush := func() {
if !allowIncrementalIMFlush(msg) {
return
}
bufMu.Lock()
parts := getStreamParts()
agentRunning := useAgent && !agentDone
@@ -2211,10 +2322,13 @@ loop:
bufMu.Unlock()
finalDisplay := cleanIMContent(ctx, FormatIMFinalFromParts(parts), tenant, s.defaultFileSvc)
if noVisibleContent || finalDisplay == "" {
if replacement, override := failClosedIMImageAnswer(answer, finalErr, msg.ImageURLs); override {
answer = replacement
finalDisplay = replacement
} else if noVisibleContent || finalDisplay == "" {
fallback := "抱歉，我暂时无法回答这个问题。"
if finalErr != nil {
fallback = "抱歉，处理您的问题时出现了异常，请稍后再试。"
fallback = imQAErrorReply(msg.ImageURLs)
}
finalDisplay = fallback
if answer == "" {
@@ -2222,14 +2336,7 @@ loop:
}
}
if err := streamer.FinalizeStream(ctx, msg, streamID, finalDisplay); err != nil {
logger.Warnf(ctx, "[IM] FinalizeStream failed: %v", err)
}
// End the stream
if err := streamer.EndStream(ctx, msg, streamID); err != nil {
logger.Warnf(ctx, "[IM] EndStream failed: %v", err)
}
deliveryErr := finalizeAndEndIMStream(ctx, streamer, msg, streamID, finalDisplay)
if answer == "" {
answer = "抱歉，我暂时无法回答这个问题。"
@@ -2242,28 +2349,49 @@ loop:
}
logger.Infof(ctx, "[IM] Stream reply sent: platform=%s user=%s answer_len=%d", msg.Platform, msg.UserID, len(answer))
return nil
return deliveryErr
}
func finalizeAndEndIMStream(ctx context.Context, streamer StreamSender, msg *IncomingMessage, streamID, finalDisplay string) error {
var deliveryErr error
if err := streamer.FinalizeStream(ctx, msg, streamID, finalDisplay); err != nil {
logger.Warnf(ctx, "[IM] FinalizeStream failed: %v", err)
deliveryErr = fmt.Errorf("finalize IM stream: %w", err)
}
// End the stream even when finalization failed so adapter resources are released.
if err := streamer.EndStream(ctx, msg, streamID); err != nil {
logger.Warnf(ctx, "[IM] EndStream failed: %v", err)
if deliveryErr == nil {
deliveryErr = fmt.Errorf("end IM stream: %w", err)
}
}
return deliveryErr
}
// fallbackNonStream is used when streaming initialization fails.
func (s *Service) fallbackNonStream(ctx context.Context, msg *IncomingMessage, session *types.Session, customAgent *types.CustomAgent, kbIDs []string, adapter Adapter, userKey string, tenant *types.Tenant) error {
answer, err := s.runQA(ctx, session, msg.Content, customAgent, kbIDs, userKey, msg.Quote)
answer, err := s.runQA(ctx, session, msg.Content, customAgent, kbIDs, userKey, msg.Quote, msg.ImageURLs)
if err != nil {
logger.Errorf(ctx, "[IM] QA fallback failed: %v", err)
answer = "抱歉，处理您的问题时出现了异常，请稍后再试。"
answer = imQAErrorReply(msg.ImageURLs)
}
return adapter.SendReply(ctx, msg, &ReplyMessage{Content: formatIMOutboundAnswer(ctx, answer, tenant, s.defaultFileSvc), IsFinal: true})
}
// runQA executes the WeKnora QA pipeline and returns the full answer text.
func (s *Service) runQA(ctx context.Context, session *types.Session, query string, customAgent *types.CustomAgent, kbIDs []string, userKey string, quote *QuotedMessage) (string, error) {
func (s *Service) runQA(ctx context.Context, session *types.Session, query string, customAgent *types.CustomAgent, kbIDs []string, userKey string, quote *QuotedMessage, imageURLs ...[]string) (string, error) {
// Cancellable context (no hard deadline): each agent round has its own
// LLMCallTimeout. The context can still be cancelled by /stop.
ctx, cancel := context.WithCancel(ctx)
defer cancel()
eventBus := event.NewEventBus()
var images []string
if len(imageURLs) > 0 {
images = imageURLs[0]
}
// Thread-safe answer collection
var answerMu sync.Mutex
@@ -2365,7 +2493,7 @@ func (s *Service) runQA(ctx context.Context, session *types.Session, query strin
// Run QA async
go func() {
var err error
req := buildIMQARequest(session, query, assistantMsg.ID, userMsg.ID, customAgent, kbIDs, quote)
req := buildIMQARequest(session, query, assistantMsg.ID, userMsg.ID, customAgent, kbIDs, quote, images)
if req.QuotedContext != "" {
logger.Debugf(ctx, "[IM] QuotedContext set: length=%d", len(req.QuotedContext))
}
@@ -2406,6 +2534,14 @@ func (s *Service) runQA(ctx context.Context, session *types.Session, query strin
qaError := qaErr
answerMu.Unlock()
if replacement, override := failClosedIMImageAnswer(answer, qaError, images); override {
assistantMsg.Content = replacement
assistantMsg.IsCompleted = true
if updateErr := s.messageService.UpdateMessage(ctx, assistantMsg); updateErr != nil {
logger.Warnf(ctx, "[IM] Failed to persist image QA error: %v", updateErr)
}
return "", qaError
}
if answer == "" && qaError != nil {
return "", qaError
}
internal/im/stream_delivery_error_test.go
+42
View File
@@ -0,0 +1,42 @@
package im
import (
"context"
"errors"
"strings"
"testing"
)
type failingFinalStreamSender struct {
finalizeErr error
endErr error
ended bool
}
func (*failingFinalStreamSender) StartStream(context.Context, *IncomingMessage) (string, error) {
return "stream", nil
}
func (*failingFinalStreamSender) UpdateStreamContent(context.Context, *IncomingMessage, string, string) error {
return nil
}
func (s *failingFinalStreamSender) FinalizeStream(context.Context, *IncomingMessage, string, string) error {
return s.finalizeErr
}
func (s *failingFinalStreamSender) EndStream(context.Context, *IncomingMessage, string) error {
s.ended = true
return s.endErr
}
func TestFinalizeAndEndIMStreamPropagatesDeliveryFailure(t *testing.T) {
sender := &failingFinalStreamSender{
finalizeErr: errors.New("simulated finalize failure"),
endErr: errors.New("simulated end failure"),
}
err := finalizeAndEndIMStream(context.Background(), sender, &IncomingMessage{}, "stream", "final")
if err == nil || !strings.Contains(err.Error(), "finalize IM stream") {
t.Fatalf("error = %v; want finalize delivery error", err)
}
if !sender.ended {
t.Fatal("EndStream must run even after FinalizeStream fails")
}
}
internal/logger/llm_logger.go
+18 -11
View File
@@ -8,6 +8,8 @@ import (
"strings"
"sync"
"time"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
var llmDebug struct {
@@ -36,10 +38,14 @@ func configureLLMDebugLog() {
dir = "llm_debug"
}
if err := os.MkdirAll(dir, 0o755); err != nil {
if err := os.MkdirAll(dir, 0o700); err != nil {
fmt.Fprintf(os.Stderr, "llm_debug: failed to create dir %s: %v\n", dir, err)
return
}
if err := os.Chmod(dir, 0o700); err != nil {
fmt.Fprintf(os.Stderr, "llm_debug: failed to secure dir %s: %v\n", dir, err)
return
}
llmDebug.dir = dir
llmDebug.enabled = true
@@ -101,12 +107,16 @@ func LLMDebugLog(ctx context.Context, record *LLMCallRecord) {
defer llmDebug.mu.Unlock()
path := filepath.Join(llmDebug.dir, filename)
f, err := os.OpenFile(path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644)
f, err := os.OpenFile(path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o600)
if err != nil {
fmt.Fprintf(os.Stderr, "llm_debug: open %s: %v\n", path, err)
return
}
defer f.Close()
if err := f.Chmod(0o600); err != nil {
fmt.Fprintf(os.Stderr, "llm_debug: secure %s: %v\n", path, err)
return
}
_, _ = f.WriteString(text)
}
@@ -132,16 +142,17 @@ func formatRecord(r *LLMCallRecord) string {
}
for _, s := range r.Sections {
content := secutils.RedactDataURLsForLog(s.Content)
b.WriteString(fmt.Sprintf("\n---------- %s ----------\n", s.Title))
b.WriteString(s.Content)
if !strings.HasSuffix(s.Content, "\n") {
b.WriteString(content)
if !strings.HasSuffix(content, "\n") {
b.WriteString("\n")
}
}
if r.Error != "" {
b.WriteString("\n---------- Error ----------\n")
b.WriteString(r.Error)
b.WriteString(secutils.RedactDataURLsForLog(r.Error))
b.WriteString("\n")
}
@@ -163,7 +174,7 @@ func extractRequestID(ctx context.Context) string {
// LLMCallRecord holds all information for one model API call.
type LLMCallRecord struct {
CallType string // "Chat", "Chat Stream", "Embedding", "Rerank", "VLM"
CallType string // "Chat", "Chat Stream", "Embedding", "Rerank", "VLM"
Model string
Duration time.Duration
Sections []RecordSection
@@ -213,11 +224,7 @@ func FormatMessages(messages []LLMMessage) string {
b.WriteString("\n")
}
for _, img := range m.Images {
if len([]rune(img)) > 80 {
b.WriteString(fmt.Sprintf("[image: %s (%d bytes)]\n", TruncateRunes(img, 80), len(img)))
} else {
b.WriteString(fmt.Sprintf("[image: %s]\n", img))
}
b.WriteString(fmt.Sprintf("[image: redacted (%d bytes)]\n", len(img)))
}
for _, tc := range m.ToolCalls {
b.WriteString(fmt.Sprintf(" -> tool_call: id=%s, func=%s, args=%s\n", tc.ID, tc.FuncName, tc.Arguments))
internal/logger/logger_test.go
+26
View File
@@ -17,6 +17,32 @@ func newEntry(level logrus.Level, msg string, data logrus.Fields) *logrus.Entry
return e
}
func TestLLMCallRecordFormatRedactsDataURIInSectionsAndErrors(t *testing.T) {
secret := "TOPSECRETBASE64"
record := &LLMCallRecord{
Sections: []RecordSection{{Title: "Response", Content: "data:image/png;base64," + secret}},
Error: "provider echoed data:text/plain," + secret,
}
got := formatRecord(record)
if strings.Contains(got, secret) || strings.Contains(strings.ToLower(got), "data:") {
t.Fatalf("record leaked data URI: %q", got)
}
}
func TestFormatMessagesRedactsImageValues(t *testing.T) {
secret := "TOPSECRETBASE64"
got := FormatMessages([]LLMMessage{{
Role: "user",
Images: []string{"data:image/png;base64," + secret},
}})
if strings.Contains(got, secret) || strings.Contains(strings.ToLower(got), "data:image") {
t.Fatalf("FormatMessages leaked image value: %q", got)
}
if !strings.Contains(got, "redacted") {
t.Fatalf("FormatMessages missing redaction marker: %q", got)
}
}
func TestFormat_DefaultModeUnchanged(t *testing.T) {
f := &CustomFormatter{} // no template, no color
entry := newEntry(logrus.InfoLevel, "hello", logrus.Fields{
internal/models/chat/anthropic.go
+271 -63
View File
@@ -3,7 +3,9 @@ package chat
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
@@ -15,19 +17,37 @@ import (
secutils "github.com/Tencent/WeKnora/internal/utils"
)
const anthropicVersion = "2023-06-01"
const (
anthropicVersion = "2023-06-01"
anthropicResponseBodyLimit = 1 << 20
)
type AnthropicChat struct {
modelName string
modelID string
baseURL string
apiKey string
customHeaders map[string]string
modelName string
modelID string
baseURL string
apiKey string
supportsVision bool
customHeaders map[string]string
}
func (c *AnthropicChat) SupportsImageInput() bool { return c.supportsVision }
type anthropicMessage struct {
Role string `json:"role"`
Content string `json:"content"`
Content any `json:"content"`
}
type anthropicContentBlock struct {
Type string `json:"type"`
Text string `json:"text,omitempty"`
Source *anthropicImageSource `json:"source,omitempty"`
}
type anthropicImageSource struct {
Type string `json:"type"`
MediaType string `json:"media_type"`
Data string `json:"data"`
}
type anthropicRequest struct {
@@ -61,6 +81,7 @@ type anthropicResponse struct {
type anthropicStreamEvent struct {
Type string `json:"type"`
Index *int `json:"index,omitempty"`
Message *struct {
Usage struct {
InputTokens int `json:"input_tokens"`
@@ -69,6 +90,9 @@ type anthropicStreamEvent struct {
CacheReadInputTokens int `json:"cache_read_input_tokens"`
} `json:"usage"`
} `json:"message,omitempty"`
ContentBlock *struct {
Type string `json:"type"`
} `json:"content_block,omitempty"`
Delta *struct {
Type string `json:"type"`
Text string `json:"text"`
@@ -86,6 +110,117 @@ type anthropicStreamEvent struct {
} `json:"error,omitempty"`
}
type anthropicStreamProtocolState struct {
messageStarted bool
finishReason string
messageStopped bool
activeBlocks map[int]string
completedBlocks map[int]struct{}
nextBlockIndex int
}
func validAnthropicStopReason(reason string) bool {
switch reason {
case "end_turn", "max_tokens", "stop_sequence", "tool_use", "pause_turn", "refusal", "model_context_window_exceeded":
return true
default:
return false
}
}
func anthropicDeltaMatchesBlock(blockType, deltaType string) bool {
switch deltaType {
case "text_delta", "citations_delta":
return blockType == "text"
case "thinking_delta", "signature_delta":
return blockType == "thinking"
case "input_json_delta":
return blockType == "tool_use" || blockType == "server_tool_use" || blockType == "mcp_tool_use"
default:
return false
}
}
func (s *anthropicStreamProtocolState) observe(event *anthropicStreamEvent) error {
if event == nil {
return errors.New("anthropic stream event is missing")
}
if event.Type == "error" || event.Error != nil {
return errors.New("anthropic stream failed")
}
if s.messageStopped {
if event.Type == "ping" {
return nil
}
return errors.New("anthropic event followed message_stop")
}
if event.Delta != nil && event.Delta.StopReason != "" && event.Type != "message_delta" {
return errors.New("anthropic stop reason appeared on an invalid event")
}
switch event.Type {
case "message_start":
if s.messageStarted || s.finishReason != "" {
return errors.New("anthropic message_start is out of order")
}
s.messageStarted = true
s.activeBlocks = make(map[int]string)
s.completedBlocks = make(map[int]struct{})
case "message_delta":
if !s.messageStarted || s.finishReason != "" || len(s.activeBlocks) != 0 || event.Delta == nil || !validAnthropicStopReason(event.Delta.StopReason) {
return errors.New("anthropic message_delta has an invalid stop reason")
}
s.finishReason = event.Delta.StopReason
case "message_stop":
if !s.messageStarted || s.finishReason == "" || len(s.activeBlocks) != 0 {
return errors.New("anthropic message_stop preceded valid completion")
}
s.messageStopped = true
case "content_block_start":
if !s.messageStarted || s.finishReason != "" || event.Index == nil || *event.Index < 0 || event.ContentBlock == nil || strings.TrimSpace(event.ContentBlock.Type) == "" {
return errors.New("anthropic content_block_start is invalid")
}
index := *event.Index
if len(s.activeBlocks) != 0 || index != s.nextBlockIndex {
return errors.New("anthropic content block index is out of order")
}
if _, active := s.activeBlocks[index]; active {
return errors.New("anthropic content block started twice")
}
if _, completed := s.completedBlocks[index]; completed {
return errors.New("anthropic completed content block restarted")
}
s.activeBlocks[index] = event.ContentBlock.Type
case "content_block_delta":
if !s.messageStarted || s.finishReason != "" || event.Index == nil || event.Delta == nil {
return errors.New("anthropic content_block_delta is invalid")
}
blockType, active := s.activeBlocks[*event.Index]
if !active || !anthropicDeltaMatchesBlock(blockType, event.Delta.Type) {
return errors.New("anthropic content block delta is out of order")
}
case "content_block_stop":
if !s.messageStarted || s.finishReason != "" || event.Index == nil {
return errors.New("anthropic content_block_stop is invalid")
}
index := *event.Index
if _, active := s.activeBlocks[index]; !active {
return errors.New("anthropic content block stop is out of order")
}
delete(s.activeBlocks, index)
s.completedBlocks[index] = struct{}{}
s.nextBlockIndex++
case "ping":
default:
return errors.New("anthropic stream event type is invalid")
}
return nil
}
func (s *anthropicStreamProtocolState) complete() bool {
return s.messageStarted && s.messageStopped && len(s.activeBlocks) == 0 && validAnthropicStopReason(s.finishReason)
}
func NewAnthropicChat(config *ChatConfig) (*AnthropicChat, error) {
if config.BaseURL != "" {
if err := secutils.ValidateURLForSSRF(config.BaseURL); err != nil {
@@ -102,16 +237,31 @@ func NewAnthropicChat(config *ChatConfig) (*AnthropicChat, error) {
}
return &AnthropicChat{
modelName: config.ModelName,
modelID: config.ModelID,
baseURL: baseURL,
apiKey: config.APIKey,
customHeaders: config.CustomHeaders,
modelName: config.ModelName,
modelID: config.ModelID,
baseURL: baseURL,
apiKey: config.APIKey,
supportsVision: config.SupportsVision,
customHeaders: config.CustomHeaders,
}, nil
}
func readAnthropicResponseBody(body io.ReadCloser) ([]byte, error) {
data, err := io.ReadAll(io.LimitReader(body, anthropicResponseBodyLimit+1))
if err != nil {
return nil, errors.New("read Anthropic response failed")
}
if len(data) > anthropicResponseBodyLimit {
return nil, errors.New("Anthropic response body too large")
}
return data, nil
}
func (c *AnthropicChat) Chat(ctx context.Context, messages []Message, opts *ChatOptions) (*types.ChatResponse, error) {
reqBody := c.buildRequest(messages, opts)
reqBody, err := c.buildRequest(messages, opts)
if err != nil {
return nil, err
}
jsonData, err := json.Marshal(reqBody)
if err != nil {
return nil, fmt.Errorf("marshal request: %w", err)
@@ -136,13 +286,13 @@ func (c *AnthropicChat) Chat(ctx context.Context, messages []Message, opts *Chat
resp, err := rawHTTPClient.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("send request: %w", err)
return nil, errors.New("send Anthropic request failed")
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
body, err := readAnthropicResponseBody(resp.Body)
if err != nil {
return nil, fmt.Errorf("read response: %w", err)
return nil, err
}
if strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") {
@@ -151,7 +301,7 @@ func (c *AnthropicChat) Chat(ctx context.Context, messages []Message, opts *Chat
return nil, err
}
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, chatResp.Content)
return nil, remoteAPIStatusError(resp.StatusCode, body)
}
logUsage(ctx, c.modelName, &chatResp.Usage)
return chatResp, nil
@@ -159,13 +309,10 @@ func (c *AnthropicChat) Chat(ctx context.Context, messages []Message, opts *Chat
var chatResp anthropicResponse
if err := json.Unmarshal(body, &chatResp); err != nil {
return nil, fmt.Errorf("decode response: %w", err)
return nil, errors.New("decode Anthropic response failed")
}
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
if chatResp.Error != nil && chatResp.Error.Message != "" {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, chatResp.Error.Message)
}
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
return nil, remoteAPIStatusError(resp.StatusCode, body)
}
result := c.parseResponse(&chatResp)
@@ -174,7 +321,10 @@ func (c *AnthropicChat) Chat(ctx context.Context, messages []Message, opts *Chat
}
func (c *AnthropicChat) ChatStream(ctx context.Context, messages []Message, opts *ChatOptions) (<-chan types.StreamResponse, error) {
reqBody := c.buildRequest(messages, opts)
reqBody, err := c.buildRequest(messages, opts)
if err != nil {
return nil, err
}
reqBody.Stream = true
jsonData, err := json.Marshal(reqBody)
if err != nil {
@@ -198,12 +348,15 @@ func (c *AnthropicChat) ChatStream(ctx context.Context, messages []Message, opts
resp, err := rawHTTPClient.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("send request: %w", err)
return nil, errors.New("send Anthropic stream request failed")
}
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
body, _ := io.ReadAll(resp.Body)
body, readErr := readAnthropicResponseBody(resp.Body)
resp.Body.Close()
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
if readErr != nil {
return nil, readErr
}
return nil, remoteAPIStatusError(resp.StatusCode, body)
}
streamChan := make(chan types.StreamResponse)
@@ -248,7 +401,7 @@ func isAnthropicVersionedBaseURL(baseURL string) bool {
return strings.HasSuffix(path, "/v1") || strings.HasSuffix(path, "/v1beta")
}
func (c *AnthropicChat) buildRequest(messages []Message, opts *ChatOptions) anthropicRequest {
func (c *AnthropicChat) buildRequest(messages []Message, opts *ChatOptions) (anthropicRequest, error) {
req := anthropicRequest{
Model: c.modelName,
MaxTokens: 1024,
@@ -276,22 +429,75 @@ func (c *AnthropicChat) buildRequest(messages []Message, opts *ChatOptions) anth
if content == "" {
content = textFromMultiContent(msg.MultiContent)
}
if content == "" {
images := append([]string(nil), msg.Images...)
for _, part := range msg.MultiContent {
if part.Type == "image_url" && part.ImageURL != nil && strings.TrimSpace(part.ImageURL.URL) != "" {
images = append(images, part.ImageURL.URL)
}
}
if len(images) == 0 && content == "" {
continue
}
switch msg.Role {
case "system":
if len(images) > 0 {
return anthropicRequest{}, fmt.Errorf("anthropic system messages cannot contain images")
}
systemParts = append(systemParts, content)
case "assistant":
if len(images) > 0 {
return anthropicRequest{}, fmt.Errorf("anthropic assistant messages cannot contain images")
}
req.Messages = append(req.Messages, anthropicMessage{Role: "assistant", Content: content})
case "user":
req.Messages = append(req.Messages, anthropicMessage{Role: "user", Content: content})
default:
req.Messages = append(req.Messages, anthropicMessage{Role: "user", Content: content})
if len(images) == 0 {
req.Messages = append(req.Messages, anthropicMessage{Role: "user", Content: content})
continue
}
blocks := make([]anthropicContentBlock, 0, len(images)+1)
if content != "" {
blocks = append(blocks, anthropicContentBlock{Type: "text", Text: content})
}
for _, imageURL := range images {
block, err := anthropicImageBlock(imageURL)
if err != nil {
return anthropicRequest{}, err
}
blocks = append(blocks, block)
}
req.Messages = append(req.Messages, anthropicMessage{Role: "user", Content: blocks})
}
}
req.System = strings.Join(systemParts, "\n\n")
return req
return req, nil
}
func anthropicImageBlock(raw string) (anthropicContentBlock, error) {
metadata, payload, ok := strings.Cut(strings.TrimSpace(raw), ",")
if !ok {
return anthropicContentBlock{}, fmt.Errorf("anthropic image input must be a JPEG or PNG data URI")
}
metadata = strings.ToLower(metadata)
var mediaType string
switch metadata {
case "data:image/jpeg;base64", "data:image/jpg;base64":
mediaType = "image/jpeg"
case "data:image/png;base64":
mediaType = "image/png"
default:
return anthropicContentBlock{}, fmt.Errorf("anthropic image input must be a JPEG or PNG data URI")
}
if _, err := base64.StdEncoding.DecodeString(payload); err != nil {
return anthropicContentBlock{}, fmt.Errorf("anthropic image data is invalid")
}
return anthropicContentBlock{
Type: "image",
Source: &anthropicImageSource{
Type: "base64",
MediaType: mediaType,
Data: payload,
},
}, nil
}
func textFromMultiContent(parts []MessageContentPart) string {
@@ -330,7 +536,7 @@ func (c *AnthropicChat) parseResponse(resp *anthropicResponse) *types.ChatRespon
func parseAnthropicSSE(reader io.Reader) (*types.ChatResponse, error) {
sseReader := NewSSEReader(reader)
var contentParts []string
var finishReason string
var protocol anthropicStreamProtocolState
var inputTokens int
var outputTokens int
@@ -353,8 +559,8 @@ func parseAnthropicSSE(reader io.Reader) (*types.ChatResponse, error) {
if err := json.Unmarshal(event.Data, &streamEvent); err != nil {
return nil, fmt.Errorf("decode SSE response: %w", err)
}
if streamEvent.Error != nil && streamEvent.Error.Message != "" {
return nil, fmt.Errorf("API stream error: %s", streamEvent.Error.Message)
if err := protocol.observe(&streamEvent); err != nil {
return nil, errors.New("anthropic stream protocol failed")
}
if streamEvent.Message != nil {
inputTokens = max(inputTokens, streamEvent.Message.Usage.InputTokens)
@@ -364,19 +570,19 @@ func parseAnthropicSSE(reader io.Reader) (*types.ChatResponse, error) {
if streamEvent.Delta.Type == "text_delta" && streamEvent.Delta.Text != "" {
contentParts = append(contentParts, streamEvent.Delta.Text)
}
if streamEvent.Delta.StopReason != "" {
finishReason = streamEvent.Delta.StopReason
}
}
if streamEvent.Usage != nil {
inputTokens = max(inputTokens, streamEvent.Usage.InputTokens)
outputTokens = max(outputTokens, streamEvent.Usage.OutputTokens)
}
}
if !protocol.complete() {
return nil, errors.New("anthropic stream ended before completion")
}
return &types.ChatResponse{
Content: strings.Join(contentParts, ""),
FinishReason: finishReason,
FinishReason: protocol.finishReason,
Usage: types.TokenUsage{
PromptTokens: inputTokens,
CompletionTokens: outputTokens,
@@ -391,38 +597,43 @@ func processAnthropicStream(ctx context.Context, model string, resp *http.Respon
sseReader := NewSSEReader(resp.Body)
var usage *types.TokenUsage
var finishReason string
var protocol anthropicStreamProtocolState
emitTerminal := func() {
if !protocol.complete() {
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeError,
Content: "anthropic stream ended before completion",
Done: true,
}
return
}
logUsage(ctx, model, usage)
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Content: "",
Done: true,
Usage: usage,
FinishReason: protocol.finishReason,
}
}
for {
event, err := sseReader.ReadEvent()
if err != nil {
if err == io.EOF {
logUsage(ctx, model, usage)
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Content: "",
Done: true,
Usage: usage,
FinishReason: finishReason,
}
emitTerminal()
} else {
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeError,
Content: err.Error(),
Content: "anthropic stream read failed",
Done: true,
}
}
return
}
if event.Done {
logUsage(ctx, model, usage)
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Content: "",
Done: true,
Usage: usage,
FinishReason: finishReason,
}
emitTerminal()
return
}
if len(event.Data) == 0 {
@@ -433,15 +644,15 @@ func processAnthropicStream(ctx context.Context, model string, resp *http.Respon
if err := json.Unmarshal(event.Data, &streamEvent); err != nil {
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeError,
Content: fmt.Sprintf("decode SSE response: %v", err),
Content: "anthropic stream response was invalid",
Done: true,
}
return
}
if streamEvent.Error != nil && streamEvent.Error.Message != "" {
if err := protocol.observe(&streamEvent); err != nil {
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeError,
Content: streamEvent.Error.Message,
Content: "anthropic stream protocol failed",
Done: true,
}
return
@@ -450,9 +661,6 @@ func processAnthropicStream(ctx context.Context, model string, resp *http.Respon
usage = mergeAnthropicUsage(usage, streamEvent.Message.Usage.InputTokens, streamEvent.Message.Usage.OutputTokens)
}
if streamEvent.Delta != nil {
if streamEvent.Delta.StopReason != "" {
finishReason = streamEvent.Delta.StopReason
}
if streamEvent.Delta.Type == "text_delta" && streamEvent.Delta.Text != "" {
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
internal/models/chat/anthropic_error_redaction_test.go
+38
View File
@@ -0,0 +1,38 @@
package chat
import (
"context"
"io"
"net/http"
"strings"
"testing"
"github.com/Tencent/WeKnora/internal/types"
)
func TestAnthropicProviderErrorDoesNotExposeImagePayload(t *testing.T) {
secret := "TOPSECRETIMAGEPAYLOAD"
payload := "data: {\"type\":\"error\",\"error\":{\"message\":\"data:image/png;base64," + secret + "\"}}\n\n"
_, err := parseAnthropicSSE(strings.NewReader(payload))
if err == nil {
t.Fatal("expected Anthropic SSE error")
}
if strings.Contains(err.Error(), secret) || strings.Contains(strings.ToLower(err.Error()), "data:image") {
t.Fatalf("non-stream error leaked provider payload: %v", err)
}
resp := &http.Response{Body: io.NopCloser(strings.NewReader(payload))}
stream := make(chan types.StreamResponse)
go processAnthropicStream(context.Background(), "vision-test", resp, stream)
var terminal types.StreamResponse
for chunk := range stream {
terminal = chunk
}
if terminal.ResponseType != types.ResponseTypeError {
t.Fatalf("terminal response = %#v; want error", terminal)
}
if strings.Contains(terminal.Content, secret) || strings.Contains(strings.ToLower(terminal.Content), "data:image") {
t.Fatalf("stream error leaked provider payload: %q", terminal.Content)
}
}
internal/models/chat/anthropic_terminal_edge_test.go
+223
View File
@@ -0,0 +1,223 @@
package chat
import (
"context"
"io"
"net/http"
"strings"
"testing"
"github.com/Tencent/WeKnora/internal/types"
)
const anthropicTestMessageStart = "data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":1,\"output_tokens\":0}}}\n\n"
const anthropicTestBlockStart0 = "data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n"
const anthropicTestBlockDelta0 = "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"x\"}}\n\n"
const anthropicTestBlockStop0 = "data: {\"type\":\"content_block_stop\",\"index\":0}\n\n"
func invalidAnthropicBlockPayloads() map[string]string {
return map[string]string{
"stop before start": anthropicTestMessageStart + anthropicTestBlockStop0,
"delta before start": anthropicTestMessageStart + anthropicTestBlockDelta0,
"delta after stop": anthropicTestMessageStart + anthropicTestBlockStart0 + anthropicTestBlockStop0 + anthropicTestBlockDelta0,
"duplicate start": anthropicTestMessageStart + anthropicTestBlockStart0 + anthropicTestBlockStart0,
"duplicate stop": anthropicTestMessageStart + anthropicTestBlockStart0 + anthropicTestBlockStop0 + anthropicTestBlockStop0,
"mismatched index": anthropicTestMessageStart + anthropicTestBlockStart0 + "data: {\"type\":\"content_block_delta\",\"index\":1,\"delta\":{\"type\":\"text_delta\",\"text\":\"x\"}}\n\n",
"index gap": anthropicTestMessageStart + "data: {\"type\":\"content_block_start\",\"index\":1,\"content_block\":{\"type\":\"text\"}}\n\n",
"concurrent block start": anthropicTestMessageStart + anthropicTestBlockStart0 + "data: {\"type\":\"content_block_start\",\"index\":1,\"content_block\":{\"type\":\"text\"}}\n\n",
"mismatched delta type": anthropicTestMessageStart + anthropicTestBlockStart0 + "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\"}}\n\n",
}
}
func TestAnthropicContentBlockLifecycleRejectedByBothParsers(t *testing.T) {
for name, payload := range invalidAnthropicBlockPayloads() {
t.Run(name+"/aggregate", func(t *testing.T) {
if _, err := parseAnthropicSSE(strings.NewReader(payload)); err == nil {
t.Fatal("aggregate parser accepted invalid content-block lifecycle")
}
})
t.Run(name+"/stream", func(t *testing.T) {
resp := &http.Response{Body: io.NopCloser(strings.NewReader(payload))}
stream := make(chan types.StreamResponse, 8)
go processAnthropicStream(context.Background(), "test", resp, stream)
sawError := false
for item := range stream {
if item.Done && item.ResponseType == types.ResponseTypeAnswer {
t.Fatalf("stream accepted invalid content-block lifecycle: %#v", item)
}
if item.Done && item.ResponseType == types.ResponseTypeError {
sawError = true
}
}
if !sawError {
t.Fatal("stream did not emit an error terminal")
}
})
}
}
func TestAnthropicContentBlockLifecycleAcceptsSequentialBlocks(t *testing.T) {
payload := anthropicTestMessageStart + anthropicTestBlockStart0 + anthropicTestBlockDelta0 + anthropicTestBlockStop0 +
"data: {\"type\":\"content_block_start\",\"index\":1,\"content_block\":{\"type\":\"thinking\",\"thinking\":\"\"}}\n\n" +
"data: {\"type\":\"content_block_delta\",\"index\":1,\"delta\":{\"type\":\"thinking_delta\"}}\n\n" +
"data: {\"type\":\"content_block_stop\",\"index\":1}\n\n" +
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"}}\n\n" +
"data: {\"type\":\"message_stop\"}\n\n"
response, err := parseAnthropicSSE(strings.NewReader(payload))
if err != nil || response.FinishReason != "end_turn" || response.Content != "x" {
t.Fatalf("valid sequential blocks rejected: response=%#v err=%v", response, err)
}
resp := &http.Response{Body: io.NopCloser(strings.NewReader(payload))}
stream := make(chan types.StreamResponse, 8)
go processAnthropicStream(context.Background(), "test", resp, stream)
successTerminals := 0
for item := range stream {
if item.Done && item.ResponseType == types.ResponseTypeError {
t.Fatalf("valid stream rejected: %#v", item)
}
if item.Done && item.ResponseType == types.ResponseTypeAnswer {
successTerminals++
}
}
if successTerminals != 1 {
t.Fatalf("success terminals=%d, want 1", successTerminals)
}
}
func TestAnthropicProtocolAcceptsOfficialStopReasons(t *testing.T) {
reasons := []string{"end_turn", "max_tokens", "stop_sequence", "tool_use", "pause_turn", "refusal", "model_context_window_exceeded"}
for _, reason := range reasons {
t.Run(reason, func(t *testing.T) {
var state anthropicStreamProtocolState
start := anthropicStreamEvent{Type: "message_start"}
if err := state.observe(&start); err != nil {
t.Fatalf("message_start: %v", err)
}
delta := anthropicStreamEvent{Type: "message_delta"}
delta.Delta = &struct {
Type string `json:"type"`
Text string `json:"text"`
StopReason string `json:"stop_reason"`
}{StopReason: reason}
if err := state.observe(&delta); err != nil {
t.Fatalf("message_delta: %v", err)
}
stop := anthropicStreamEvent{Type: "message_stop"}
if err := state.observe(&stop); err != nil || !state.complete() {
t.Fatalf("message_stop err=%v complete=%v", err, state.complete())
}
})
}
}
func TestParseAnthropicSSERejectsPrematureEOF(t *testing.T) {
payload := anthropicTestMessageStart + anthropicTestBlockStart0 + "event: content_block_delta\n" +
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"partial\"}}\n\n"
if _, err := parseAnthropicSSE(strings.NewReader(payload)); err == nil {
t.Fatal("expected premature Anthropic EOF to fail closed")
}
}
func TestParseAnthropicSSERejectsMalformedPacketAfterFinishReason(t *testing.T) {
payload := anthropicTestMessageStart + "event: message_delta\n" +
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"}}\n\n" +
"event: broken\n" +
"data: {not-json}\n\n"
if _, err := parseAnthropicSSE(strings.NewReader(payload)); err == nil {
t.Fatal("expected trailing malformed packet to override finish reason")
}
}
func TestParseAnthropicSSERejectsInvalidOrOutOfOrderCompletion(t *testing.T) {
cases := map[string]string{
"invalid stop reason": anthropicTestMessageStart + "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"made_up\"}}\n\n" +
"data: {\"type\":\"message_stop\"}\n\n",
"message stop first": "data: {\"type\":\"message_stop\"}\n\n" +
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"}}\n\n",
"stop reason wrong event": anthropicTestMessageStart + "data: {\"type\":\"content_block_delta\",\"delta\":{\"stop_reason\":\"end_turn\"}}\n\n" +
"data: {\"type\":\"message_stop\"}\n\n",
"empty provider error after completion": anthropicTestMessageStart + "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"}}\n\n" +
"data: {\"type\":\"message_stop\"}\n\n" +
"data: {\"type\":\"error\",\"error\":{\"type\":\"api_error\",\"message\":\"\"}}\n\n",
}
for name, payload := range cases {
t.Run(name, func(t *testing.T) {
if _, err := parseAnthropicSSE(strings.NewReader(payload)); err == nil {
t.Fatal("expected invalid Anthropic completion protocol to fail")
}
})
}
}
func TestProcessAnthropicStreamRejectsPrematureDone(t *testing.T) {
payload := anthropicTestMessageStart + anthropicTestBlockStart0 + anthropicTestBlockDelta0 +
"data: [DONE]\n\n"
resp := &http.Response{Body: io.NopCloser(strings.NewReader(payload))}
stream := make(chan types.StreamResponse, 4)
go processAnthropicStream(context.Background(), "test", resp, stream)
for item := range stream {
if item.Done && item.ResponseType == types.ResponseTypeAnswer {
t.Fatalf("premature [DONE] produced success: %#v", item)
}
}
}
func TestProcessAnthropicStreamRejectsEmptyProviderErrorAfterCompletion(t *testing.T) {
payload := anthropicTestMessageStart + "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"}}\n\n" +
"data: {\"type\":\"message_stop\"}\n\n" +
"data: {\"type\":\"error\",\"error\":{\"type\":\"api_error\",\"message\":\"\"}}\n\n"
resp := &http.Response{Body: io.NopCloser(strings.NewReader(payload))}
stream := make(chan types.StreamResponse, 4)
go processAnthropicStream(context.Background(), "test", resp, stream)
for item := range stream {
if item.Done && item.ResponseType == types.ResponseTypeAnswer {
t.Fatalf("empty-message provider error was ignored: %#v", item)
}
if item.Done && item.ResponseType == types.ResponseTypeError {
return
}
}
t.Fatal("expected provider error terminal")
}
func TestProcessAnthropicStreamRejectsPrematureEOF(t *testing.T) {
payload := anthropicTestMessageStart + anthropicTestBlockStart0 + "event: content_block_delta\n" +
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"partial\"}}\n\n"
resp := &http.Response{Body: io.NopCloser(strings.NewReader(payload))}
stream := make(chan types.StreamResponse, 4)
go processAnthropicStream(context.Background(), "test", resp, stream)
successTerminals, errorTerminals := 0, 0
for item := range stream {
if item.Done && item.ResponseType == types.ResponseTypeAnswer {
successTerminals++
}
if item.Done && item.ResponseType == types.ResponseTypeError {
errorTerminals++
}
}
if successTerminals != 0 || errorTerminals != 1 {
t.Fatalf("terminal counts success=%d error=%d; want 0/1", successTerminals, errorTerminals)
}
}
func TestProcessAnthropicStreamRejectsMalformedPacketAfterFinishReason(t *testing.T) {
payload := anthropicTestMessageStart + "event: message_delta\n" +
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"}}\n\n" +
"event: broken\n" +
"data: {not-json}\n\n"
resp := &http.Response{Body: io.NopCloser(strings.NewReader(payload))}
stream := make(chan types.StreamResponse, 4)
go processAnthropicStream(context.Background(), "test", resp, stream)
for item := range stream {
if item.Done && item.ResponseType == types.ResponseTypeAnswer {
t.Fatalf("unexpected success terminal before trailing malformed packet: %#v", item)
}
if item.ResponseType == types.ResponseTypeError {
return
}
}
t.Fatal("expected terminal stream error")
}
internal/models/chat/anthropic_test.go
+41 -2
View File
@@ -13,6 +13,33 @@ import (
"github.com/stretchr/testify/require"
)
func TestAnthropicBuildRequestPreservesImageDataURI(t *testing.T) {
chat := &AnthropicChat{modelName: "claude-test"}
req, err := chat.buildRequest([]Message{{
Role: "user",
Content: "describe this image",
Images: []string{"data:image/png;base64,YWJj"},
}}, nil)
require.NoError(t, err)
require.Len(t, req.Messages, 1)
blocks, ok := req.Messages[0].Content.([]anthropicContentBlock)
require.True(t, ok, "content type = %T", req.Messages[0].Content)
require.Len(t, blocks, 2)
assert.Equal(t, "text", blocks[0].Type)
assert.Equal(t, "describe this image", blocks[0].Text)
assert.Equal(t, "image", blocks[1].Type)
require.NotNil(t, blocks[1].Source)
assert.Equal(t, "base64", blocks[1].Source.Type)
assert.Equal(t, "image/png", blocks[1].Source.MediaType)
assert.Equal(t, "YWJj", blocks[1].Source.Data)
}
func TestAnthropicBuildRequestRejectsMalformedImage(t *testing.T) {
chat := &AnthropicChat{modelName: "claude-test"}
_, err := chat.buildRequest([]Message{{Role: "user", Content: "describe", Images: []string{"not-an-image"}}}, nil)
require.Error(t, err)
}
func TestAnthropicChat(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
@@ -147,8 +174,14 @@ func TestAnthropicChat_SSEResponse(t *testing.T) {
_, _ = w.Write([]byte(`event: message_start
data: {"type":"message_start","message":{"usage":{"input_tokens":0,"output_tokens":0}}}
event: content_block_start
data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
event: content_block_delta
data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"pong"}}
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"pong"}}
event: content_block_stop
data: {"type":"content_block_stop","index":0}
event: message_delta
data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"input_tokens":114,"output_tokens":5}}
@@ -191,8 +224,14 @@ func TestAnthropicChat_ChatStream(t *testing.T) {
_, _ = w.Write([]byte(`event: message_start
data: {"type":"message_start","message":{"usage":{"input_tokens":0,"output_tokens":0}}}
event: content_block_start
data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
event: content_block_delta
data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"pong"}}
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"pong"}}
event: content_block_stop
data: {"type":"content_block_stop","index":0}
event: message_delta
data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"input_tokens":114,"output_tokens":5}}
internal/models/chat/chat.go
+36 -17
View File
@@ -98,14 +98,27 @@ type Chat interface {
GetModelID() string
}
type imageInputCapable interface {
SupportsImageInput() bool
}
// SupportsImageInput reports whether the concrete Chat transport can preserve
// image input. Unknown implementations fail closed instead of trusting model
// metadata alone.
func SupportsImageInput(model Chat) bool {
capable, ok := model.(imageInputCapable)
return ok && capable.SupportsImageInput()
}
type ChatConfig struct {
Source types.ModelSource
BaseURL string
ModelName string
APIKey string
ModelID string
Provider string
ExtraConfig map[string]string
Source types.ModelSource
BaseURL string
ModelName string
APIKey string
ModelID string
Provider string
SupportsVision bool
ExtraConfig map[string]string
// CustomHeaders 允许在调用远程 OpenAI 兼容 API 时附加自定义 HTTP 请求头（类似 OpenAI Python SDK 的 extra_headers）。
CustomHeaders map[string]string
AppID string
@@ -120,17 +133,23 @@ func ConfigFromModel(m *types.Model, appID, appSecret string) *ChatConfig {
if m == nil {
return nil
}
supportsVision := m.Parameters.SupportsVision
if provider.ProviderName(m.Parameters.Provider) == provider.ProviderWeKnoraCloud &&
m.Type == types.ModelTypeKnowledgeQA && strings.EqualFold(strings.TrimSpace(m.Name), "chat") {
supportsVision = true
}
return &ChatConfig{
ModelID: m.ID,
APIKey: m.Parameters.APIKey,
BaseURL: m.Parameters.BaseURL,
ModelName: m.Name,
Source: m.Source,
Provider: m.Parameters.Provider,
ExtraConfig: m.Parameters.ExtraConfig,
CustomHeaders: m.Parameters.CustomHeaders,
AppID: appID,
AppSecret: appSecret,
ModelID: m.ID,
APIKey: m.Parameters.APIKey,
BaseURL: m.Parameters.BaseURL,
ModelName: m.Name,
Source: m.Source,
Provider: m.Parameters.Provider,
SupportsVision: supportsVision,
ExtraConfig: m.Parameters.ExtraConfig,
CustomHeaders: m.Parameters.CustomHeaders,
AppID: appID,
AppSecret: appSecret,
}
}
internal/models/chat/image_resolve.go
-21
View File
@@ -74,24 +74,3 @@ func readLocalStorageBytes(storagePath string) []byte {
}
return data
}
// isMultimodalNotSupportedError checks if an error indicates the model does not
// support multimodal/image input.
func isMultimodalNotSupportedError(err error) bool {
if err == nil {
return false
}
msg := strings.ToLower(err.Error())
return (strings.Contains(msg, "multimodal") || strings.Contains(msg, "image") || strings.Contains(msg, "vision")) &&
(strings.Contains(msg, "not support") || strings.Contains(msg, "unsupported") || strings.Contains(msg, "400"))
}
// stripImagesFromMessages returns a copy of messages with all image data removed.
func stripImagesFromMessages(messages []Message) []Message {
cleaned := make([]Message, len(messages))
for i, msg := range messages {
cleaned[i] = msg
cleaned[i].Images = nil
}
return cleaned
}
internal/models/chat/langfuse_image_test.go
+26
View File
@@ -0,0 +1,26 @@
package chat
import (
"encoding/json"
"strings"
"testing"
)
func TestBuildLangfuseMessagesRedactsMultiContentImage(t *testing.T) {
secret := "TOPSECRETBASE64"
msgs := buildLangfuseMessages([]Message{{
Role: "user",
MultiContent: []MessageContentPart{
{Type: "image_url", ImageURL: &ImageURL{URL: "data:image/png;base64," + secret}},
{Type: "text", Text: "describe"},
},
}})
encoded, err := json.Marshal(msgs)
if err != nil {
t.Fatal(err)
}
serialized := string(encoded)
if strings.Contains(serialized, secret) || strings.Contains(strings.ToLower(serialized), "data:image") {
t.Fatalf("Langfuse input leaked image payload: %s", serialized)
}
}
internal/models/chat/langfuse_wrapper.go
+45 -11
View File
@@ -2,10 +2,12 @@ package chat
import (
"context"
"errors"
"time"
"github.com/Tencent/WeKnora/internal/tracing/langfuse"
"github.com/Tencent/WeKnora/internal/types"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
// langfuseChat wraps a Chat implementation and emits a Langfuse generation
@@ -16,8 +18,9 @@ type langfuseChat struct {
inner Chat
}
func (l *langfuseChat) GetModelName() string { return l.inner.GetModelName() }
func (l *langfuseChat) GetModelID() string { return l.inner.GetModelID() }
func (l *langfuseChat) GetModelName() string { return l.inner.GetModelName() }
func (l *langfuseChat) GetModelID() string { return l.inner.GetModelID() }
func (l *langfuseChat) SupportsImageInput() bool { return SupportsImageInput(l.inner) }
func (l *langfuseChat) Chat(ctx context.Context, messages []Message, opts *ChatOptions) (*types.ChatResponse, error) {
mgr := langfuse.GetManager()
@@ -47,7 +50,7 @@ func (l *langfuseChat) Chat(ctx context.Context, messages []Message, opts *ChatO
resp.Content, resp.ReasoningContent, resp.FinishReason, resp.ToolCalls,
)
}
gen.Finish(output, usage, err)
gen.Finish(output, usage, sanitizeLangfuseError(err))
return resp, err
}
@@ -71,7 +74,7 @@ func (l *langfuseChat) ChatStream(ctx context.Context, messages []Message, opts
ch, err := l.inner.ChatStream(genCtx, messages, opts)
if err != nil {
gen.Finish(nil, nil, err)
gen.Finish(nil, nil, sanitizeLangfuseError(err))
return ch, err
}
if ch == nil {
@@ -131,10 +134,21 @@ func buildLangfuseMessages(messages []Message) []map[string]interface{} {
"role": m.Role,
}
if m.Content != "" {
entry["content"] = m.Content
entry["content"] = secutils.RedactDataURLsForLog(m.Content)
}
if len(m.MultiContent) > 0 {
entry["content"] = m.MultiContent
parts := make([]MessageContentPart, len(m.MultiContent))
copy(parts, m.MultiContent)
for i := range parts {
parts[i].Text = secutils.RedactDataURLsForLog(parts[i].Text)
if parts[i].ImageURL == nil {
continue
}
imageURL := *parts[i].ImageURL
imageURL.URL = "[image-url-redacted]"
parts[i].ImageURL = &imageURL
}
entry["content"] = parts
}
if m.Name != "" {
entry["name"] = m.Name
@@ -143,10 +157,15 @@ func buildLangfuseMessages(messages []Message) []map[string]interface{} {
entry["tool_call_id"] = m.ToolCallID
}
if len(m.ToolCalls) > 0 {
entry["tool_calls"] = m.ToolCalls
calls := make([]ToolCall, len(m.ToolCalls))
copy(calls, m.ToolCalls)
for i := range calls {
calls[i].Function.Arguments = secutils.RedactDataURLsForLog(calls[i].Function.Arguments)
}
entry["tool_calls"] = calls
}
if m.ReasoningContent != "" {
entry["reasoning_content"] = m.ReasoningContent
entry["reasoning_content"] = secutils.RedactDataURLsForLog(m.ReasoningContent)
}
out = append(out, entry)
}
@@ -157,17 +176,32 @@ func buildLangfuseGenerationOutput(
content, reasoningContent, finishReason string,
toolCalls []types.LLMToolCall,
) map[string]interface{} {
var safeToolCalls []types.LLMToolCall
if toolCalls != nil {
safeToolCalls = make([]types.LLMToolCall, len(toolCalls))
copy(safeToolCalls, toolCalls)
}
for i := range safeToolCalls {
safeToolCalls[i].Function.Arguments = secutils.RedactDataURLsForLog(safeToolCalls[i].Function.Arguments)
}
output := map[string]interface{}{
"content": content,
"tool_calls": toolCalls,
"content": secutils.RedactDataURLsForLog(content),
"tool_calls": safeToolCalls,
"finish_reason": finishReason,
}
if reasoningContent != "" {
output["reasoning_content"] = reasoningContent
output["reasoning_content"] = secutils.RedactDataURLsForLog(reasoningContent)
}
return output
}
func sanitizeLangfuseError(err error) error {
if err == nil {
return nil
}
return errors.New(secutils.RedactDataURLsForLog(err.Error()))
}
func buildLangfuseModelParams(opts *ChatOptions) map[string]interface{} {
if opts == nil {
return nil
internal/models/chat/llm_debug.go
+1 -1
View File
@@ -29,7 +29,7 @@ func buildLLMMessages(messages []Message) []logger.LLMMessage {
parts = append(parts, mc.Text)
case "image_url":
if mc.ImageURL != nil {
parts = append(parts, fmt.Sprintf("[image_url: %s]", truncateForDebug(mc.ImageURL.URL, 120)))
parts = append(parts, fmt.Sprintf("[image_url: redacted (%d bytes)]", len(mc.ImageURL.URL)))
}
}
}
internal/models/chat/llm_debug_image_test.go
+27
View File
@@ -0,0 +1,27 @@
package chat
import (
"strings"
"testing"
)
func TestBuildLLMMessagesRedactsMultiContentImageURL(t *testing.T) {
secret := "TOPSECRETBASE64"
messages := buildLLMMessages([]Message{{
Role: "user",
MultiContent: []MessageContentPart{{
Type: "image_url",
ImageURL: &ImageURL{URL: "data:image/png;base64," + secret},
}},
}})
if len(messages) != 1 {
t.Fatalf("messages length = %d", len(messages))
}
got := messages[0].Content
if strings.Contains(got, secret) || strings.Contains(strings.ToLower(got), "data:image") {
t.Fatalf("buildLLMMessages leaked image URL: %q", got)
}
if !strings.Contains(got, "redacted") {
t.Fatalf("buildLLMMessages missing redaction marker: %q", got)
}
}
internal/models/chat/llm_debug_wrapper.go
+3 -2
View File
@@ -16,8 +16,9 @@ type debugChat struct {
inner Chat
}
func (d *debugChat) GetModelName() string { return d.inner.GetModelName() }
func (d *debugChat) GetModelID() string { return d.inner.GetModelID() }
func (d *debugChat) GetModelName() string { return d.inner.GetModelName() }
func (d *debugChat) GetModelID() string { return d.inner.GetModelID() }
func (d *debugChat) SupportsImageInput() bool { return SupportsImageInput(d.inner) }
func (d *debugChat) Chat(ctx context.Context, messages []Message, opts *ChatOptions) (*types.ChatResponse, error) {
callStart := time.Now()
internal/models/chat/model_fail_closed_edge_test.go
+29
View File
@@ -0,0 +1,29 @@
package chat
import (
"io"
"strings"
"testing"
)
func TestOllamaBuildChatRequestRejectsAnyUnresolvableImage(t *testing.T) {
c := &OllamaChat{supportsVision: true}
_, err := c.buildChatRequest([]Message{{
Role: "user",
Content: "describe",
Images: []string{"data:image/png;base64,%%%INVALID%%%"},
}}, nil, false)
if err == nil {
t.Fatal("expected malformed image to fail closed")
}
if strings.Contains(err.Error(), "%%%INVALID%%%") {
t.Fatalf("image value leaked in error: %v", err)
}
}
func TestReadAnthropicResponseBodyRejectsOversize(t *testing.T) {
_, err := readAnthropicResponseBody(io.NopCloser(strings.NewReader(strings.Repeat("x", anthropicResponseBodyLimit+1))))
if err == nil {
t.Fatal("expected oversized Anthropic response rejection")
}
}
internal/models/chat/ollama.go
+77 -41
View File
@@ -3,7 +3,7 @@ package chat
import (
"context"
"encoding/json"
"fmt"
"errors"
"io"
"net/http"
"strconv"
@@ -13,27 +13,33 @@ import (
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/models/utils/ollama"
"github.com/Tencent/WeKnora/internal/types"
secutils "github.com/Tencent/WeKnora/internal/utils"
ollamaapi "github.com/ollama/ollama/api"
)
// OllamaChat 实现了基于 Ollama 的聊天
type OllamaChat struct {
modelName string
modelID string
ollamaService *ollama.OllamaService
modelName string
modelID string
supportsVision bool
ollamaService *ollama.OllamaService
}
func (c *OllamaChat) SupportsImageInput() bool { return c.supportsVision }
// NewOllamaChat 创建 Ollama 聊天实例
func NewOllamaChat(config *ChatConfig, ollamaService *ollama.OllamaService) (*OllamaChat, error) {
return &OllamaChat{
modelName: config.ModelName,
modelID: config.ModelID,
ollamaService: ollamaService,
modelName: config.ModelName,
modelID: config.ModelID,
supportsVision: config.SupportsVision,
ollamaService: ollamaService,
}, nil
}
// convertMessages 转换消息格式为Ollama API格式
func (c *OllamaChat) convertMessages(messages []Message) []ollamaapi.Message {
// convertMessages converts messages and fails closed if any declared image
// cannot be resolved. Silently dropping one image would change the user's input.
func (c *OllamaChat) convertMessages(messages []Message) ([]ollamaapi.Message, error) {
ollamaMessages := make([]ollamaapi.Message, 0, len(messages))
for _, msg := range messages {
msgOllama := ollamaapi.Message{
@@ -46,47 +52,69 @@ func (c *OllamaChat) convertMessages(messages []Message) []ollamaapi.Message {
}
if len(msg.Images) > 0 && msg.Role == "user" {
for _, imgURL := range msg.Images {
if imgData := resolveImageForOllama(imgURL); imgData != nil {
msgOllama.Images = append(msgOllama.Images, imgData)
imgData, err := resolveImageForOllama(imgURL)
if err != nil {
return nil, err
}
msgOllama.Images = append(msgOllama.Images, imgData)
}
}
ollamaMessages = append(ollamaMessages, msgOllama)
}
return ollamaMessages
return ollamaMessages, nil
}
// resolveImageForOllama resolves an image URL into raw bytes for Ollama.
// Handles local serving paths (/files/...), data URIs, and remote HTTP URLs.
func resolveImageForOllama(imageURL string) ollamaapi.ImageData {
func resolveImageForOllama(imageURL string) (ollamaapi.ImageData, error) {
if data := resolveImageURLForOllama(imageURL); data != nil {
return data
return data, nil
}
if strings.HasPrefix(imageURL, "http://") || strings.HasPrefix(imageURL, "https://") {
client := &http.Client{Timeout: 30 * time.Second}
if err := secutils.ValidateURLForSSRF(imageURL); err != nil {
return nil, errors.New("ollama image URL is not allowed")
}
client := &http.Client{
Timeout: 30 * time.Second,
Transport: &http.Transport{
Proxy: nil,
DialContext: secutils.SSRFSafeDialContext,
},
CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse },
}
resp, err := client.Get(imageURL)
if err != nil {
return nil
return nil, errors.New("ollama image download failed")
}
defer resp.Body.Close()
data, err := io.ReadAll(io.LimitReader(resp.Body, 20*1024*1024))
if err != nil {
return nil
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
return nil, errors.New("ollama image download failed")
}
return data
const maxOllamaImageBytes = 20 * 1024 * 1024
data, err := io.ReadAll(io.LimitReader(resp.Body, maxOllamaImageBytes+1))
if err != nil {
return nil, errors.New("ollama image download failed")
}
if len(data) == 0 || len(data) > maxOllamaImageBytes {
return nil, errors.New("ollama image payload is invalid")
}
return data, nil
}
return nil
return nil, errors.New("ollama image payload is invalid")
}
// buildChatRequest 构建聊天请求参数
func (c *OllamaChat) buildChatRequest(messages []Message, opts *ChatOptions, isStream bool) *ollamaapi.ChatRequest {
func (c *OllamaChat) buildChatRequest(messages []Message, opts *ChatOptions, isStream bool) (*ollamaapi.ChatRequest, error) {
// 设置流式标志
streamFlag := isStream
// 构建请求参数
convertedMessages, err := c.convertMessages(messages)
if err != nil {
return nil, err
}
chatReq := &ollamaapi.ChatRequest{
Model: c.modelName,
Messages: c.convertMessages(messages),
Messages: convertedMessages,
Stream: &streamFlag,
Options: make(map[string]interface{}),
}
@@ -113,7 +141,7 @@ func (c *OllamaChat) buildChatRequest(messages []Message, opts *ChatOptions, isS
}
}
return chatReq
return chatReq, nil
}
// Chat 进行非流式聊天
@@ -124,7 +152,10 @@ func (c *OllamaChat) Chat(ctx context.Context, messages []Message, opts *ChatOpt
}
// 构建请求参数
chatReq := c.buildChatRequest(messages, opts, false)
chatReq, err := c.buildChatRequest(messages, opts, false)
if err != nil {
return nil, err
}
// 记录请求日志
logger.GetLogger(ctx).Infof("发送聊天请求到模型 %s", c.modelName)
@@ -134,7 +165,7 @@ func (c *OllamaChat) Chat(ctx context.Context, messages []Message, opts *ChatOpt
var promptTokens, completionTokens int
// 使用 Ollama 客户端发送请求
err := c.ollamaService.Chat(ctx, chatReq, func(resp ollamaapi.ChatResponse) error {
err = c.ollamaService.Chat(ctx, chatReq, func(resp ollamaapi.ChatResponse) error {
responseContent = resp.Message.Content
// 当 Content 为空但 Thinking 有内容时（如推理模型未正确配置 thinking 参数），使用 Thinking 作为兜底
if responseContent == "" && resp.Message.Thinking != "" {
@@ -151,7 +182,7 @@ func (c *OllamaChat) Chat(ctx context.Context, messages []Message, opts *ChatOpt
return nil
})
if err != nil {
return nil, fmt.Errorf("聊天请求失败: %w", err)
return nil, errors.New("ollama chat failed")
}
usage := types.TokenUsage{
@@ -180,7 +211,10 @@ func (c *OllamaChat) ChatStream(
}
// 构建请求参数
chatReq := c.buildChatRequest(messages, opts, true)
chatReq, err := c.buildChatRequest(messages, opts, true)
if err != nil {
return nil, err
}
// 记录请求日志
logger.GetLogger(ctx).Infof("发送流式聊天请求到模型 %s", c.modelName)
@@ -193,6 +227,8 @@ func (c *OllamaChat) ChatStream(
defer close(streamChan)
var thinking thinkingEmitter
var terminalUsage *types.TokenUsage
providerDone := false
err := c.ollamaService.Chat(ctx, chatReq, func(resp ollamaapi.ChatResponse) error {
// 发送思考内容（支持 Qwen3、DeepSeek 等推理模型）
if resp.Message.Thinking != "" {
@@ -247,32 +283,32 @@ func (c *OllamaChat) ChatStream(
}
if resp.Done {
var usage *types.TokenUsage
providerDone = true
if resp.PromptEvalCount > 0 || resp.EvalCount > 0 {
usage = &types.TokenUsage{
terminalUsage = &types.TokenUsage{
PromptTokens: resp.PromptEvalCount,
CompletionTokens: resp.EvalCount,
TotalTokens: resp.PromptEvalCount + resp.EvalCount,
}
}
logUsage(ctx, c.modelName, usage)
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Done: true,
Usage: usage,
}
}
return nil
})
if err != nil {
logger.GetLogger(ctx).Errorf("流式聊天请求失败: %v", err)
// 发送错误响应
if err != nil || !providerDone {
logger.GetLogger(ctx).Error("流式聊天请求失败")
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeError,
Content: err.Error(),
Content: "ollama stream failed",
Done: true,
}
return
}
logUsage(ctx, c.modelName, terminalUsage)
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Done: true,
Usage: terminalUsage,
}
}()
internal/models/chat/openai_stream.go
+36 -38
View File
@@ -122,7 +122,7 @@ func (c *RemoteAPIChat) processStream(
for {
response, err := stream.Recv()
if err != nil {
if err == io.EOF {
if err == io.EOF && state.lastFinishReason != "" {
logUsage(ctx, c.modelName, state.usage)
toolCalls := state.buildOrderedToolCalls()
streamChan <- types.StreamResponse{
@@ -134,9 +134,10 @@ func (c *RemoteAPIChat) processStream(
FinishReason: state.lastFinishReason,
}
} else {
logger.Warnf(ctx, "Chat completion stream ended before finish reason")
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeError,
Content: err.Error(),
Content: "chat completion stream failed",
Done: true,
}
}
@@ -178,23 +179,26 @@ func (c *RemoteAPIChat) processRawHTTPStream(
for {
event, err := reader.ReadEvent()
if err != nil {
if err == io.EOF {
if err == io.EOF && state.lastFinishReason != "" {
logUsage(ctx, c.modelName, state.usage)
toolCalls := state.buildOrderedToolCalls()
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Content: "",
Done: true,
ToolCalls: toolCalls,
ToolCalls: state.buildOrderedToolCalls(),
Usage: state.usage,
FinishReason: state.lastFinishReason,
}
return
}
if err != io.EOF {
logger.Warnf(ctx, "Chat stream read failed")
} else {
logger.Errorf(ctx, "Stream read error: %v", err)
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeError,
Content: err.Error(),
Done: true,
}
logger.Warnf(ctx, "Chat stream ended before finish reason")
}
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeError,
Content: "chat stream ended before completion",
Done: true,
}
return
}
@@ -204,6 +208,14 @@ func (c *RemoteAPIChat) processRawHTTPStream(
}
if event.Done {
if state.lastFinishReason == "" {
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeError,
Content: "chat stream ended before completion",
Done: true,
}
return
}
logUsage(ctx, c.modelName, state.usage)
toolCalls := state.buildOrderedToolCalls()
streamChan <- types.StreamResponse{
@@ -212,6 +224,7 @@ func (c *RemoteAPIChat) processRawHTTPStream(
Done: true,
ToolCalls: toolCalls,
Usage: state.usage,
FinishReason: state.lastFinishReason,
}
return
}
@@ -241,8 +254,13 @@ func (c *RemoteAPIChat) processRawHTTPStream(
}
if err := json.Unmarshal(event.Data, &streamResp); err != nil {
logger.Errorf(ctx, "Failed to parse stream response: %v", err)
continue
logger.Warnf(ctx, "Failed to parse chat stream response")
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeError,
Content: "chat stream response was invalid",
Done: true,
}
return
}
if streamResp.Usage != nil {
@@ -453,38 +471,18 @@ func (c *RemoteAPIChat) processStreamDelta(
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Content: delta.Content,
Done: isDone,
Done: false,
ToolCalls: state.buildOrderedToolCalls(),
FinishReason: string(choice.FinishReason),
}
}
if isDone && len(state.toolCallMap) > 0 {
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Content: "",
Done: true,
ToolCalls: state.buildOrderedToolCalls(),
FinishReason: string(choice.FinishReason),
}
}
// Ensure thinking done is sent when stream finishes without any answer content
// (e.g., model only produced reasoning then hit finish_reason with empty content).
// Ensure thinking done is sent when the provider supplied a finish reason.
// The outer stream reader emits the sole terminal answer after it confirms
// the provider ended successfully.
if isDone {
state.finish(streamChan)
}
// Catch-all: isDone but none of the above branches sent a response with
// FinishReason (empty content, no tool calls, no thinking). This prevents
// the finish_reason from being lost in the streaming pipeline.
if isDone && delta.Content == "" && len(state.toolCallMap) == 0 {
streamChan <- types.StreamResponse{
ResponseType: types.ResponseTypeAnswer,
Done: true,
FinishReason: string(choice.FinishReason),
}
}
}
// processToolCallsDelta 处理 tool calls 的增量更新
internal/models/chat/plan_a_fail_closed_test.go
+129
View File
@@ -0,0 +1,129 @@
package chat
import (
"context"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/Tencent/WeKnora/internal/types"
)
func TestRemoteAPIProviderErrorDoesNotExposeImagePayload(t *testing.T) {
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
secret := "TOPSECRETIMAGEPAYLOAD"
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusBadRequest)
_, _ = w.Write([]byte(`{"error":{"message":"data:image/png;base64,` + secret + `"}}`))
}))
defer server.Close()
model, err := NewRemoteAPIChat(&ChatConfig{
Source: types.ModelSourceRemote,
BaseURL: server.URL + "/v1",
ModelName: "vision-test",
APIKey: "test-key",
Provider: "openai",
SupportsVision: true,
})
if err != nil {
t.Fatal(err)
}
_, err = model.Chat(context.Background(), []Message{{
Role: "user", Content: "inspect", Images: []string{"data:image/png;base64," + secret},
}}, nil)
if err == nil {
t.Fatal("expected provider error")
}
if strings.Contains(err.Error(), secret) || strings.Contains(strings.ToLower(err.Error()), "data:image") {
t.Fatalf("provider error leaked image payload: %v", err)
}
}
func TestProcessRawHTTPStreamPrematureEOFFailsClosed(t *testing.T) {
body := "data: {\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n\n"
resp := &http.Response{Body: io.NopCloser(strings.NewReader(body))}
stream := make(chan types.StreamResponse)
model := &RemoteAPIChat{modelName: "cloud-vision"}
go model.processRawHTTPStream(context.Background(), resp, stream, nil)
var got []types.StreamResponse
for chunk := range stream {
got = append(got, chunk)
}
if len(got) == 0 || got[len(got)-1].ResponseType != types.ResponseTypeError {
t.Fatalf("stream = %#v; want terminal error on premature EOF", got)
}
if got[len(got)-1].Content != "chat stream ended before completion" {
t.Fatalf("terminal error = %q", got[len(got)-1].Content)
}
}
func TestProcessRawHTTPStreamFinishReasonThenEOFCompletes(t *testing.T) {
body := "data: {\"choices\":[{\"delta\":{\"content\":\"answer\"},\"finish_reason\":\"stop\"}]}\n\n"
resp := &http.Response{Body: io.NopCloser(strings.NewReader(body))}
stream := make(chan types.StreamResponse)
model := &RemoteAPIChat{modelName: "cloud-vision"}
go model.processRawHTTPStream(context.Background(), resp, stream, nil)
var got []types.StreamResponse
for chunk := range stream {
got = append(got, chunk)
}
if len(got) == 0 {
t.Fatal("expected stream chunks")
}
last := got[len(got)-1]
if last.ResponseType == types.ResponseTypeError || !last.Done || last.FinishReason != "stop" {
t.Fatalf("stream = %#v; want successful stop completion", got)
}
}
func TestProcessRawHTTPStreamFinishReasonThenMalformedPacketFailsWithoutSuccessTerminal(t *testing.T) {
body := "data: {\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n\n" +
"data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n" +
"data: {not-json}\n\n"
resp := &http.Response{Body: io.NopCloser(strings.NewReader(body))}
stream := make(chan types.StreamResponse, 8)
model := &RemoteAPIChat{modelName: "stream-test"}
go model.processRawHTTPStream(context.Background(), resp, stream, nil)
successTerminals, errorTerminals := 0, 0
for item := range stream {
if item.Done && item.ResponseType == types.ResponseTypeAnswer {
successTerminals++
}
if item.Done && item.ResponseType == types.ResponseTypeError {
errorTerminals++
}
}
if successTerminals != 0 || errorTerminals != 1 {
t.Fatalf("terminal counts success=%d error=%d; want 0/1", successTerminals, errorTerminals)
}
}
func TestProcessRawHTTPStreamEmitsExactlyOneSuccessTerminal(t *testing.T) {
body := "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"}}]}\n\n" +
"data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n" +
"data: [DONE]\n\n"
resp := &http.Response{Body: io.NopCloser(strings.NewReader(body))}
stream := make(chan types.StreamResponse, 8)
model := &RemoteAPIChat{modelName: "stream-test"}
go model.processRawHTTPStream(context.Background(), resp, stream, nil)
successTerminals := 0
for item := range stream {
if item.Done && item.ResponseType == types.ResponseTypeAnswer {
successTerminals++
}
if item.ResponseType == types.ResponseTypeError {
t.Fatalf("unexpected stream error: %#v", item)
}
}
if successTerminals != 1 {
t.Fatalf("success terminals = %d, want 1", successTerminals)
}
}
internal/models/chat/remote_api.go
+31 -35
View File
@@ -21,14 +21,15 @@ import (
// providerAdapter（见 provider.go），thinking 编码委托给 ThinkingStrategy
// （见 thinking.go）。
type RemoteAPIChat struct {
modelName string
client *openai.Client
modelID string
baseURL string
apiKey string
provider provider.ProviderName
appID string
appSecret string
modelName string
client *openai.Client
modelID string
baseURL string
apiKey string
provider provider.ProviderName
supportsVision bool
appID string
appSecret string
// customHeaders 为用户在模型配置中指定的自定义 HTTP 请求头（类似 OpenAI Python SDK 的 extra_headers）。
customHeaders map[string]string
@@ -38,6 +39,8 @@ type RemoteAPIChat struct {
thinkingOverride ThinkingStrategy
}
func (c *RemoteAPIChat) SupportsImageInput() bool { return c.supportsVision }
// NewRemoteAPIChat 创建远程 API 聊天实例
func NewRemoteAPIChat(chatConfig *ChatConfig) (*RemoteAPIChat, error) {
if chatConfig.BaseURL != "" {
@@ -103,6 +106,7 @@ func NewRemoteAPIChat(chatConfig *ChatConfig) (*RemoteAPIChat, error) {
baseURL: chatConfig.BaseURL,
apiKey: apiKey,
provider: providerName,
supportsVision: chatConfig.SupportsVision,
appID: chatConfig.AppID,
appSecret: chatConfig.AppSecret,
customHeaders: chatConfig.CustomHeaders,
@@ -120,7 +124,11 @@ func (c *RemoteAPIChat) authCreds() authCreds {
// transform and parameter shaping (but not thinking, which may wrap the body).
func (c *RemoteAPIChat) shapedRequest(messages []Message, opts *ChatOptions, isStream bool) openai.ChatCompletionRequest {
req := c.BuildChatCompletionRequest(messages, opts, isStream)
req.Messages = c.adapter.TransformMessages(req.Messages)
// WeKnora Cloud text models retain the legacy plain-text downgrade, while
// vision-capable Chat models must keep the OpenAI-compatible image parts.
if c.provider != provider.ProviderWeKnoraCloud || !c.supportsVision {
req.Messages = c.adapter.TransformMessages(req.Messages)
}
c.adapter.ShapeRequest(&req, opts, isStream)
return req
}
@@ -180,15 +188,7 @@ func (c *RemoteAPIChat) Chat(ctx context.Context, messages []Message, opts *Chat
c.logRequest(timeoutCtx, req, false)
resp, err := c.client.CreateChatCompletion(timeoutCtx, req)
if err != nil {
if isMultimodalNotSupportedError(err) {
logger.Warnf(timeoutCtx, "[LLM Request] Model %s does not support multimodal, retrying without images", c.modelName)
cleaned := stripImagesFromMessages(messages)
req = c.shapedRequest(cleaned, opts, false)
resp, err = c.client.CreateChatCompletion(timeoutCtx, req)
}
if err != nil {
return nil, fmt.Errorf("create chat completion: %w", err)
}
return nil, fmt.Errorf("create chat completion failed")
}
result, err := c.parseCompletionResponse(&resp)
@@ -231,13 +231,13 @@ func (c *RemoteAPIChat) chatWithRawHTTP(ctx context.Context, endpoint string, cu
resp, err := rawHTTPClient.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("send request: %w", err)
return nil, fmt.Errorf("send chat request failed")
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
return nil, remoteAPIStatusError(resp.StatusCode, body)
}
body, err := io.ReadAll(resp.Body)
@@ -259,6 +259,10 @@ func (c *RemoteAPIChat) chatWithRawHTTP(ctx context.Context, endpoint string, cu
return result, nil
}
func remoteAPIStatusError(status int, _ []byte) error {
return fmt.Errorf("API request failed with status %d", status)
}
// ChatStream 进行流式聊天
func (c *RemoteAPIChat) ChatStream(ctx context.Context, messages []Message, opts *ChatOptions) (<-chan types.StreamResponse, error) {
// 仅在调用方未设置 deadline 时附加兜底超时；流式调用默认超时更长，
@@ -287,17 +291,9 @@ func (c *RemoteAPIChat) ChatStream(ctx context.Context, messages []Message, opts
stream, err := c.client.CreateChatCompletionStream(timeoutCtx, req)
if err != nil {
if isMultimodalNotSupportedError(err) {
logger.Warnf(timeoutCtx, "[LLM Stream] Model %s does not support multimodal, retrying without images", c.modelName)
cleaned := stripImagesFromMessages(messages)
req = c.shapedRequest(cleaned, opts, true)
stream, err = c.client.CreateChatCompletionStream(timeoutCtx, req)
}
if err != nil {
cancel()
close(streamChan)
return nil, fmt.Errorf("create chat completion stream: %w", err)
}
cancel()
close(streamChan)
return nil, fmt.Errorf("create chat completion stream failed")
}
go func() {
@@ -363,13 +359,13 @@ func (c *RemoteAPIChat) chatStreamWithRawHTTP(ctx context.Context, endpoint stri
resp, err := rawHTTPClient.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("send request: %w", err)
return nil, fmt.Errorf("send chat request failed")
}
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
resp.Body.Close()
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
return nil, remoteAPIStatusError(resp.StatusCode, body)
}
streamChan := make(chan types.StreamResponse)
internal/models/chat/remote_api_test.go
+56
View File
@@ -3,7 +3,11 @@ package chat
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"os"
"strings"
"sync/atomic"
"testing"
"time"
@@ -13,6 +17,17 @@ import (
"github.com/stretchr/testify/require"
)
func TestRemoteAPIStatusErrorDoesNotExposeResponseBody(t *testing.T) {
secret := "data:image/png;base64,TOPSECRET"
err := remoteAPIStatusError(http.StatusBadRequest, []byte(secret))
if strings.Contains(err.Error(), secret) || strings.Contains(err.Error(), "TOPSECRET") {
t.Fatalf("remote API error leaked response body: %v", err)
}
if !strings.Contains(err.Error(), "400") {
t.Fatalf("remote API error = %v, want status code", err)
}
}
func newTestRemoteChat(t *testing.T) *RemoteAPIChat {
t.Helper()
@@ -533,3 +548,44 @@ func TestTokenUsage_CachedTokensJSONOmitempty(t *testing.T) {
assert.Contains(t, string(b), `"cached_tokens":7`)
})
}
func newMultimodalUnsupportedTestChat(t *testing.T, calls *atomic.Int32) *RemoteAPIChat {
t.Helper()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
calls.Add(1)
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusBadRequest)
_, _ = w.Write([]byte(`{"error":{"message":"image input is not supported","type":"invalid_request_error","code":"unsupported"}}`))
}))
t.Cleanup(server.Close)
chat := newTestRemoteChat(t)
cfg := openai.DefaultConfig("test-key")
cfg.BaseURL = server.URL
chat.client = openai.NewClientWithConfig(cfg)
return chat
}
func TestRemoteAPIChat_MultimodalUnsupportedDoesNotStripImages(t *testing.T) {
messages := []Message{{
Role: "user",
Content: "read this image",
Images: []string{"data:image/png;base64,iVBORw0KGgo="},
}}
t.Run("non-streaming", func(t *testing.T) {
var calls atomic.Int32
chat := newMultimodalUnsupportedTestChat(t, &calls)
_, err := chat.Chat(context.Background(), messages, &ChatOptions{})
require.Error(t, err)
assert.EqualValues(t, 1, calls.Load(), "image input must not be removed and retried as text")
})
t.Run("streaming", func(t *testing.T) {
var calls atomic.Int32
chat := newMultimodalUnsupportedTestChat(t, &calls)
_, err := chat.ChatStream(context.Background(), messages, &ChatOptions{})
require.Error(t, err)
assert.EqualValues(t, 1, calls.Load(), "image input must not be removed and retried as text")
})
}
internal/models/chat/stream_raw_dump.go
+23 -4
View File
@@ -8,6 +8,8 @@ import (
"strings"
"sync"
"time"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
// streamRawDumpDir returns the directory for per-stream raw packet dumps.
@@ -43,7 +45,10 @@ func newStreamPacketDumper(modelName string, request any) *streamPacketDumper {
if dir == "" {
return nil
}
if err := os.MkdirAll(dir, 0o755); err != nil {
if err := os.MkdirAll(dir, 0o700); err != nil {
return nil
}
if err := os.Chmod(dir, 0o700); err != nil {
return nil
}
@@ -70,10 +75,14 @@ func newStreamPacketDumper(modelName string, request any) *streamPacketDumper {
return nil
}
f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o644)
f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600)
if err != nil {
return nil
}
if err := f.Chmod(0o600); err != nil {
_ = f.Close()
return nil
}
d := &streamPacketDumper{file: f, path: path, model: modelName}
_ = d.writeRequest(request)
@@ -81,11 +90,19 @@ func newStreamPacketDumper(modelName string, request any) *streamPacketDumper {
}
func (d *streamPacketDumper) writeRequest(request any) error {
raw, err := json.Marshal(request)
if err != nil {
return err
}
var redacted any
if err := json.Unmarshal([]byte(secutils.CompactImageDataURLForLog(string(raw))), &redacted); err != nil {
return err
}
line, err := json.Marshal(map[string]any{
"type": "request",
"model": d.model,
"timestamp": time.Now().UTC().Format(time.RFC3339Nano),
"data": request,
"data": redacted,
})
if err != nil {
return err
@@ -112,6 +129,7 @@ func (d *streamPacketDumper) WritePacketRaw(raw []byte) {
if len(raw) == 0 {
return
}
raw = []byte(secutils.CompactImageDataURLForLog(string(raw)))
d.mu.Lock()
defer d.mu.Unlock()
@@ -139,6 +157,7 @@ func (d *streamPacketDumper) WriteError(message string) {
if d == nil || d.file == nil || strings.TrimSpace(message) == "" {
return
}
message = secutils.CompactImageDataURLForLog(message)
line, err := json.Marshal(map[string]any{
"type": "error",
"model": d.model,
@@ -163,7 +182,7 @@ func (d *streamPacketDumper) WriteHTTPError(statusCode int, body []byte) {
"status_code": statusCode,
}
if len(body) > 0 {
trimmed := bytesTrimSpace(body)
trimmed := []byte(secutils.CompactImageDataURLForLog(string(bytesTrimSpace(body))))
if json.Valid(trimmed) {
var parsed any
if json.Unmarshal(trimmed, &parsed) == nil {
internal/models/chat/stream_raw_dump_image_test.go
+51
View File
@@ -0,0 +1,51 @@
package chat
import (
"os"
"strings"
"testing"
)
func TestStreamPacketDumperRedactsImageRequestAndUsesPrivatePermissions(t *testing.T) {
dir := t.TempDir()
if err := os.Chmod(dir, 0o755); err != nil {
t.Fatal(err)
}
t.Setenv("WEKNORA_LLM_STREAM_RAW_DUMP_DIR", dir)
secret := strings.Repeat("TOPSECRET", 30)
dumper := newStreamPacketDumper("vision-model", map[string]any{
"messages": []any{map[string]any{
"role": "user",
"image_url": map[string]any{"url": "data:image/png;base64," + secret},
}},
})
if dumper == nil {
t.Fatal("expected dumper")
}
path := dumper.Path()
dumper.Close()
data, err := os.ReadFile(path)
if err != nil {
t.Fatalf("read dump: %v", err)
}
if strings.Contains(string(data), secret) || strings.Contains(string(data), "data:image/") {
t.Fatalf("image data leaked to dump: %s", data)
}
if !strings.Contains(string(data), "[image-data-redacted]") {
t.Fatalf("missing redaction marker: %s", data)
}
info, err := os.Stat(path)
if err != nil {
t.Fatalf("stat dump: %v", err)
}
if got := info.Mode().Perm(); got != 0o600 {
t.Fatalf("dump permissions = %o, want 600", got)
}
dirInfo, err := os.Stat(dir)
if err != nil {
t.Fatalf("stat dump directory: %v", err)
}
if got := dirInfo.Mode().Perm(); got != 0o700 {
t.Fatalf("dump directory permissions = %o, want 700", got)
}
}
internal/models/chat/weknora_cloud_image_test.go
+80
View File
@@ -0,0 +1,80 @@
package chat
import (
"testing"
"github.com/Tencent/WeKnora/internal/models/provider"
"github.com/Tencent/WeKnora/internal/types"
"github.com/stretchr/testify/require"
)
func TestRemoteAPIChat_WeKnoraCloudTransformsImagesByCapability(t *testing.T) {
messages := []Message{{
Role: "user",
Content: "请识别图片",
Images: []string{"data:image/png;base64,iVBORw0KGgo="},
}}
build := func(t *testing.T, supportsVision bool) *RemoteAPIChat {
t.Helper()
model, err := NewRemoteAPIChat(&ChatConfig{
Source: types.ModelSourceRemote,
ModelName: "cloud-chat",
ModelID: "cloud-chat",
Provider: string(provider.ProviderWeKnoraCloud),
AppID: "test-app",
AppSecret: "test-secret",
SupportsVision: supportsVision,
})
require.NoError(t, err)
return model
}
t.Run("vision chat preserves image parts", func(t *testing.T) {
req := build(t, true).shapedRequest(messages, &ChatOptions{}, false)
require.Len(t, req.Messages, 1)
require.Len(t, req.Messages[0].MultiContent, 2)
require.NotNil(t, req.Messages[0].MultiContent[0].ImageURL)
require.Equal(t, messages[0].Images[0], req.Messages[0].MultiContent[0].ImageURL.URL)
})
t.Run("text chat keeps legacy downgrade", func(t *testing.T) {
req := build(t, false).shapedRequest(messages, &ChatOptions{}, false)
require.Len(t, req.Messages, 1)
require.Empty(t, req.Messages[0].MultiContent)
require.Equal(t, messages[0].Content, req.Messages[0].Content)
})
}
func TestConfigFromModelRecognizesExistingWeKnoraCloudChatAsVision(t *testing.T) {
cloudChat := &types.Model{
Name: "chat",
Type: types.ModelTypeKnowledgeQA,
Source: types.ModelSourceRemote,
Parameters: types.ModelParameters{
Provider: string(provider.ProviderWeKnoraCloud),
},
}
genericChat := &types.Model{
Name: "chat",
Type: types.ModelTypeKnowledgeQA,
Source: types.ModelSourceRemote,
}
require.True(t, ConfigFromModel(cloudChat, "", "").SupportsVision)
require.False(t, ConfigFromModel(genericChat, "", "").SupportsVision)
}
func TestSupportsImageInputUsesRuntimeTransportCapability(t *testing.T) {
visionRemote := &RemoteAPIChat{supportsVision: true}
textRemote := &RemoteAPIChat{supportsVision: false}
anthropicText := &AnthropicChat{}
anthropicVision := &AnthropicChat{supportsVision: true}
require.True(t, SupportsImageInput(visionRemote))
require.False(t, SupportsImageInput(textRemote))
require.False(t, SupportsImageInput(anthropicText))
require.True(t, SupportsImageInput(anthropicVision))
require.True(t, SupportsImageInput(&debugChat{inner: visionRemote}))
require.True(t, SupportsImageInput(&langfuseChat{inner: visionRemote}))
}
internal/types/chat_manage.go
+10 -9
View File
@@ -221,15 +221,16 @@ func (c *ChatManage) Clone() *ChatManage {
Images: append([]string(nil), c.Images...),
VLMModelID: c.VLMModelID,
ChatModelSupportsVision: c.ChatModelSupportsVision,
Attachments: append(MessageAttachments(nil), c.Attachments...),
TenantID: c.TenantID,
WebSearchEnabled: c.WebSearchEnabled,
WebSearchProviderID: c.WebSearchProviderID,
WebSearchMaxResults: c.WebSearchMaxResults,
WebFetchEnabled: c.WebFetchEnabled,
WebFetchTopN: c.WebFetchTopN,
Language: c.Language,
IntentPromptOverrides: maps.Clone(c.IntentPromptOverrides),
Attachments: append(MessageAttachments(nil), c.Attachments...),
TenantID: c.TenantID,
WebSearchEnabled: c.WebSearchEnabled,
WebSearchProviderID: c.WebSearchProviderID,
WebSearchMaxResults: c.WebSearchMaxResults,
WebFetchEnabled: c.WebFetchEnabled,
WebFetchTopN: c.WebFetchTopN,
Language: c.Language,
IntentPromptOverrides: maps.Clone(c.IntentPromptOverrides),
},
PipelineState: PipelineState{
RewriteQuery: c.RewriteQuery,
internal/types/qa_request.go
+18 -17
View File
@@ -4,21 +4,22 @@ package types
// replacing the previous 14-parameter method signatures.
// EventBus is passed separately to avoid circular dependency with the event package.
type QARequest struct {
Session *Session // The conversation session
Query string // User query text
AssistantMessageID string // Pre-created assistant message ID
SummaryModelID string // Optional model override; empty = use agent/KB default
CustomAgent *CustomAgent // Optional custom agent for config override
KnowledgeBaseIDs []string // Knowledge base IDs to search (from request + @mentions)
KnowledgeIDs []string // Specific knowledge (file) IDs to search
TagScopes []TagScope // Tag-constrained KB scopes from @mentions
MCPServiceIDs []string // Per-request MCP service IDs from @mentions
SkillNames []string // Per-request preloaded skill names from @mentions
ImageURLs []string // Image URLs for multimodal input
ImageDescription string // VLM-generated image description (fallback for non-vision models)
UserMessageID string // Created user message ID
WebSearchEnabled bool // Whether web search is enabled for this request
EnableMemory bool // Whether memory feature is enabled
QuotedContext string // Quoted message content from IM quote-reply (appended at LLM prompt stage, not used for retrieval)
Attachments MessageAttachments // File attachments (processed and ready for prompt injection)
Session *Session // The conversation session
Query string // User query text
AssistantMessageID string // Pre-created assistant message ID
SummaryModelID string // Optional model override; empty = use agent/KB default
CustomAgent *CustomAgent // Optional custom agent for config override
KnowledgeBaseIDs []string // Knowledge base IDs to search (from request + @mentions)
KnowledgeIDs []string // Specific knowledge (file) IDs to search
TagScopes []TagScope // Tag-constrained KB scopes from @mentions
MCPServiceIDs []string // Per-request MCP service IDs from @mentions
SkillNames []string // Per-request preloaded skill names from @mentions
ImageURLs []string // Image URLs for multimodal input
ImageDescription string // VLM-generated image description (fallback for non-vision models)
UserMessageID string // Created user message ID
WebSearchEnabled bool // Whether web search is enabled for this request
EnableMemory bool // Whether memory feature is enabled
QuotedContext string // Quoted message content from IM quote-reply (appended at LLM prompt stage, not used for retrieval)
Attachments MessageAttachments // File attachments (processed and ready for prompt injection)
}
internal/utils/log_sanitize.go
+121 -14
View File
@@ -1,29 +1,136 @@
package utils
import (
"regexp"
"strconv"
"strings"
)
var imageDataURLPatternForLog = regexp.MustCompile(`data:image\/[a-zA-Z0-9.+-]+;base64,[A-Za-z0-9+/=]+`)
const (
defaultMaxLogChars = 12000
defaultMaxDataURLPreview = 96
defaultMaxLogChars = 12000
dataURLRedaction = "[image-data-redacted]"
)
// CompactImageDataURLForLog shortens large image data URLs for log output.
// CompactImageDataURLForLog removes every data URI before log or dump output.
// It deliberately handles non-canonical forms because debug paths may receive
// JSON-escaped or percent-encoded requests in addition to ordinary data: URLs.
func CompactImageDataURLForLog(raw string) string {
masked := imageDataURLPatternForLog.ReplaceAllStringFunc(raw, func(match string) string {
if len(match) <= defaultMaxDataURLPreview {
return match
}
hidden := len(match) - defaultMaxDataURLPreview
return match[:defaultMaxDataURLPreview] + "...<omitted " + strconv.Itoa(hidden) + " chars>"
})
masked := RedactDataURLsForLog(raw)
if len(masked) <= defaultMaxLogChars {
return masked
}
return masked[:defaultMaxLogChars] + "... (truncated, total " + strconv.Itoa(len(masked)) + " chars)"
}
// RedactDataURLsForLog removes data URIs without truncating surrounding text.
func RedactDataURLsForLog(raw string) string {
return redactDataURLs(raw)
}
func redactDataURLs(raw string) string {
var out strings.Builder
out.Grow(len(raw))
for i := 0; i < len(raw); {
prefixLen := dataURLPrefixLen(raw[i:])
if prefixLen == 0 {
out.WriteByte(raw[i])
i++
continue
}
out.WriteString(dataURLRedaction)
i += prefixLen
for i < len(raw) && !isDataURLTerminator(raw, i) {
i++
}
}
return out.String()
}
// dataURLPrefixLen recognizes "data:" when each ASCII character is either
// literal, percent encoded, or represented as one-or-more JSON escape layers.
func dataURLPrefixLen(value string) int {
pos := 0
for i := 0; i < len("data:"); i++ {
consumed := encodedASCIILen(value[pos:], "data:"[i])
if consumed == 0 {
return 0
}
pos += consumed
}
return pos
}
func encodedASCIILen(value string, want byte) int {
if len(value) > 0 && strings.EqualFold(value[:1], string(want)) {
return 1
}
if len(value) >= 3 && value[0] == '%' {
pos := 0
for len(value) >= pos+3 && strings.EqualFold(value[pos:pos+3], "%25") {
pos += 3
}
if pos == 0 {
if decoded, ok := decodeHexByte(value[1], value[2]); ok && decoded == want {
return 3
}
} else if len(value) >= pos+2 {
if decoded, ok := decodeHexByte(value[pos], value[pos+1]); ok && decoded == want {
return pos + 2
}
}
}
backslashes := 0
for backslashes < len(value) && value[backslashes] == '\\' {
backslashes++
}
if backslashes > 0 && len(value) >= backslashes+5 && (value[backslashes] == 'u' || value[backslashes] == 'U') {
hex := value[backslashes+1 : backslashes+5]
if strings.EqualFold(hex[:2], "00") {
if decoded, ok := decodeHexByte(hex[2], hex[3]); ok && decoded == want {
return backslashes + 5
}
}
}
return 0
}
func decodeHexByte(high, low byte) (byte, bool) {
hex := func(ch byte) (byte, bool) {
switch {
case ch >= '0' && ch <= '9':
return ch - '0', true
case ch >= 'a' && ch <= 'f':
return ch - 'a' + 10, true
case ch >= 'A' && ch <= 'F':
return ch - 'A' + 10, true
default:
return 0, false
}
}
hi, ok := hex(high)
if !ok {
return 0, false
}
lo, ok := hex(low)
if !ok {
return 0, false
}
return hi<<4 | lo, true
}
// Data characters may legally include punctuation such as parentheses and
// backslashes. Stop only at whitespace or an unescaped serialized-string quote.
func isDataURLTerminator(raw string, pos int) bool {
switch raw[pos] {
case ' ', '\t', '\r', '\n':
return true
case '"':
backslashes := 0
for i := pos - 1; i >= 0 && raw[i] == '\\'; i-- {
backslashes++
}
return backslashes%2 == 0
default:
return false
}
}
internal/utils/log_sanitize_edge_test.go
+29
View File
@@ -0,0 +1,29 @@
package utils
import (
"strings"
"testing"
)
func TestCompactImageDataURLForLogRedactsEscapedAndLegalPayloadForms(t *testing.T) {
tests := []struct {
name string
raw string
secret string
}{
{"encoded leading d", `\u0064ata:image/png;base64,ESCAPEDLEADING`, "ESCAPEDLEADING"},
{"fully percent encoded prefix", `%64%61%74%61%3Aimage/png;base64,PERCENTPREFIX`, "PERCENTPREFIX"},
{"escaped payload sequence", `data:image/png;base64,BEFORE\u002bAFTER`, "AFTER"},
{"parenthesis is legal payload", `data:text/plain,BEFORE(AFTER)SUFFIX next`, "SUFFIX"},
{"apostrophe is legal payload", `data:image/png;name=foo'bar;base64,APOSTROPHESECRET next`, "APOSTROPHESECRET"},
{"double percent encoded prefix", `%2564%2561%2574%2561%253Aimage/png;base64,DOUBLEPERCENT`, "DOUBLEPERCENT"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := CompactImageDataURLForLog(tt.raw)
if strings.Contains(got, tt.secret) || !strings.Contains(got, dataURLRedaction) {
t.Fatalf("CompactImageDataURLForLog(%q) = %q", tt.raw, got)
}
})
}
}
internal/utils/log_sanitize_test.go
+29
View File
@@ -0,0 +1,29 @@
package utils
import (
"strings"
"testing"
)
func TestCompactImageDataURLForLogFullyRedactsImages(t *testing.T) {
cases := []string{
"data:image/png;base64,YWJj",
"DATA:image/png;BASE64," + strings.Repeat("QUJD", 100),
`{"image":"data:image\\/png;base64,TOPSECRET"}`,
"data:image/png,%89PNG%0D%0A",
"data:text/plain,NONIMAGESECRET",
"data%3Aimage%2Fpng%3Bbase64%2CENCODEDSECRET",
`data\u003aimage/png;base64,ESCAPEDSECRET`,
}
for _, raw := range cases {
got := CompactImageDataURLForLog(raw)
if !strings.Contains(got, "[image-data-redacted]") {
t.Fatalf("CompactImageDataURLForLog(%q) = %q", raw, got)
}
for _, secret := range []string{"YWJj", "QUJD", "TOPSECRET", "%89PNG", "NONIMAGESECRET", "ENCODEDSECRET", "ESCAPEDSECRET"} {
if strings.Contains(got, secret) {
t.Fatalf("sanitized output retained %q: %q", secret, got)
}
}
}
}
internal/utils/security.go
+37 -13
View File
@@ -226,6 +226,14 @@ var restrictedHostSuffixes = []string{
var restrictedIPv4Ranges = []*net.IPNet{
// 100.64.0.0/10 - Carrier-grade NAT (RFC 6598)
mustParseCIDR("100.64.0.0/10"),
// 192.31.196.0/24 - AS112-v4 special-purpose service
mustParseCIDR("192.31.196.0/24"),
// 192.52.193.0/24 - AMT special-purpose service
mustParseCIDR("192.52.193.0/24"),
// 192.88.99.0/24 - Deprecated 6to4 relay anycast
mustParseCIDR("192.88.99.0/24"),
// 192.175.48.0/24 - Direct Delegation AS112 service
mustParseCIDR("192.175.48.0/24"),
// 198.18.0.0/15 - Network device benchmark testing (RFC 2544)
mustParseCIDR("198.18.0.0/15"),
// 198.51.100.0/24 - TEST-NET-2 for documentation (RFC 5737)
@@ -250,6 +258,22 @@ var restrictedIPv4Ranges = []*net.IPNet{
mustParseCIDR("172.20.0.0/16"),
}
// restrictedIPv6Ranges contains special-purpose address space that must not be
// used as an outbound media-fetch destination. These ranges are not ordinary
// public unicast destinations and some can translate or encapsulate IPv4.
var restrictedIPv6Ranges = []*net.IPNet{
mustParseCIDR("64:ff9b::/96"),
mustParseCIDR("64:ff9b:1::/48"),
mustParseCIDR("100::/64"),
mustParseCIDR("100:0:0:1::/64"),
mustParseCIDR("2001::/23"),
mustParseCIDR("2001:db8::/32"),
mustParseCIDR("2002::/16"),
mustParseCIDR("2620:4f:8000::/48"),
mustParseCIDR("3fff::/20"),
mustParseCIDR("5f00::/16"),
}
// mustParseCIDR parses a CIDR string and panics on error
func mustParseCIDR(s string) *net.IPNet {
_, ipNet, err := net.ParseCIDR(s)
@@ -289,6 +313,19 @@ func isRestrictedIP(ip net.IP) (bool, string) {
// Check IPv6-specific restrictions
if ip.To4() == nil && len(ip) == 16 {
// Preserve explicit diagnostics for tunneling prefixes before the broader
// IANA special-purpose range table.
if ip[0] == 0x20 && ip[1] == 0x01 && ip[2] == 0x00 && ip[3] == 0x00 {
return true, "Teredo tunneling address"
}
if ip[0] == 0x20 && ip[1] == 0x02 {
return true, "6to4 special-purpose address"
}
for _, cidr := range restrictedIPv6Ranges {
if cidr.Contains(ip) {
return true, fmt.Sprintf("restricted range %s", cidr.String())
}
}
// Site-local (deprecated but still blocked): fec0::/10
if ip[0] == 0xfe && (ip[1]&0xc0) == 0xc0 {
return true, "site-local IPv6 address"
@@ -304,19 +341,6 @@ func isRestrictedIP(ip net.IP) (bool, string) {
return true, fmt.Sprintf("IPv4-mapped %s", reason)
}
}
// Teredo tunneling addresses: 2001:0000::/32
// Embed arbitrary IPv4 in the payload; can reach internal hosts via relay.
if ip[0] == 0x20 && ip[1] == 0x01 && ip[2] == 0x00 && ip[3] == 0x00 {
return true, "Teredo tunneling address"
}
// 6to4 addresses: 2002::/16
// Bits 16-47 carry an IPv4 address; block when embedded IPv4 is restricted.
if ip[0] == 0x20 && ip[1] == 0x02 {
embeddedIP := net.IP(ip[2:6])
if restricted, reason := isRestrictedIP(embeddedIP); restricted {
return true, fmt.Sprintf("6to4 embedded %s", reason)
}
}
}
return false, ""
internal/utils/security_stage2_test.go
+25
View File
@@ -0,0 +1,25 @@
package utils
import "testing"
func TestIsPublicIPRejectsSpecialPurposeIPv6(t *testing.T) {
for _, raw := range []string{
"192.31.196.1",
"192.52.193.1",
"192.88.99.1",
"192.175.48.1",
"100::1",
"100:0:0:1::1",
"2001:1::1",
"2001:2::1",
"2001:10::1",
"2001:20::1",
"64:ff9b:1::1",
"2002:0808:0808::1",
"2620:4f:8000::1",
} {
if IsPublicIP(parseIPForTest(t, raw)) {
t.Errorf("IsPublicIP(%s) = true; want false for special-purpose IPv6", raw)
}
}
}
internal/utils/security_test.go
+1 -1
View File
@@ -266,7 +266,7 @@ func TestIsRestrictedIP_IPv6(t *testing.T) {
{"site-local", "fec0::1", true, "site-local"},
{"Teredo", "2001:0000:4136:e378:8000:63bf:3fff:fdd2", true, "Teredo"},
{"6to4 private", "2002:c0a8:0101::1", true, "6to4"},
{"6to4 public", "2002:0808:0808::1", false, ""},
{"6to4 public endpoint still special-purpose", "2002:0808:0808::1", true, "6to4"},
{"public IPv6", "2001:4860:4860::8888", false, ""},
}