diff --git a/docs/test-plans/assets/dingtalk-vision-acceptance.png b/docs/test-plans/assets/dingtalk-vision-acceptance.png new file mode 100644 index 00000000..2b736287 Binary files /dev/null and b/docs/test-plans/assets/dingtalk-vision-acceptance.png differ diff --git a/docs/test-plans/assets/generate_dingtalk_vision_acceptance.py b/docs/test-plans/assets/generate_dingtalk_vision_acceptance.py new file mode 100644 index 00000000..acbeb68e --- /dev/null +++ b/docs/test-plans/assets/generate_dingtalk_vision_acceptance.py @@ -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) diff --git a/docs/test-plans/dingtalk-image-multimodal-acceptance.md b/docs/test-plans/dingtalk-image-multimodal-acceptance.md new file mode 100644 index 00000000..73dd27a8 --- /dev/null +++ b/docs/test-plans/dingtalk-image-multimodal-acceptance.md @@ -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。 diff --git a/frontend/src/utils/weknoraCloudModels.ts b/frontend/src/utils/weknoraCloudModels.ts index 97fd2ae4..2f255491 100644 --- a/frontend/src/utils/weknoraCloudModels.ts +++ b/frontend/src/utils/weknoraCloudModels.ts @@ -57,7 +57,7 @@ export function buildWkcModelConfig( truncate_prompt_tokens: 0, } } - if (kind === 'vllm') { + if (kind === 'chat' || kind === 'vllm') { parameters.supports_vision = true } diff --git a/internal/agent/engine.go b/internal/agent/engine.go index a8710a47..52e82413 100644 --- a/internal/agent/engine.go +++ b/internal/agent/engine.go @@ -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 diff --git a/internal/agent/engine_test.go b/internal/agent/engine_test.go index 0b3da547..399ece31 100644 --- a/internal/agent/engine_test.go +++ b/internal/agent/engine_test.go @@ -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) diff --git a/internal/agent/finalize.go b/internal/agent/finalize.go index c80344b0..464217f5 100644 --- a/internal/agent/finalize.go +++ b/internal/agent/finalize.go @@ -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(), diff --git a/internal/agent/finalize_image_test.go b/internal/agent/finalize_image_test.go new file mode 100644 index 00000000..98ab47ee --- /dev/null +++ b/internal/agent/finalize_image_test.go @@ -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) + } +} diff --git a/internal/agent/think.go b/internal/agent/think.go index 3eae32cd..c88d2593 100644 --- a/internal/agent/think.go +++ b/internal/agent/think.go @@ -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) } diff --git a/internal/agent/think_stream_error_test.go b/internal/agent/think_stream_error_test.go new file mode 100644 index 00000000..084f2df1 --- /dev/null +++ b/internal/agent/think_stream_error_test.go @@ -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) + } +} diff --git a/internal/application/service/chat_pipeline/chat_completion_stream.go b/internal/application/service/chat_pipeline/chat_completion_stream.go index 9f519568..9710668d 100644 --- a/internal/application/service/chat_pipeline/chat_completion_stream.go +++ b/internal/application/service/chat_pipeline/chat_completion_stream.go @@ -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 + } } } } diff --git a/internal/application/service/chat_pipeline/query_understand.go b/internal/application/service/chat_pipeline/query_understand.go index 36c1271e..c791cc15 100644 --- a/internal/application/service/chat_pipeline/query_understand.go +++ b/internal/application/service/chat_pipeline/query_understand.go @@ -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. diff --git a/internal/application/service/chat_pipeline/query_understand_image_test.go b/internal/application/service/chat_pipeline/query_understand_image_test.go new file mode 100644 index 00000000..314555d1 --- /dev/null +++ b/internal/application/service/chat_pipeline/query_understand_image_test.go @@ -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) + } +} diff --git a/internal/application/service/session_agent_image_route.go b/internal/application/service/session_agent_image_route.go new file mode 100644 index 00000000..a26905d5 --- /dev/null +++ b/internal/application/service/session_agent_image_route.go @@ -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) +} diff --git a/internal/application/service/session_agent_image_test.go b/internal/application/service/session_agent_image_test.go new file mode 100644 index 00000000..2b960ca9 --- /dev/null +++ b/internal/application/service/session_agent_image_test.go @@ -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) + } +} diff --git a/internal/application/service/session_agent_qa.go b/internal/application/service/session_agent_qa.go index 9956966b..c57841e8 100644 --- a/internal/application/service/session_agent_qa.go +++ b/internal/application/service/session_agent_qa.go @@ -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( diff --git a/internal/application/service/session_agent_vlm_fallback_test.go b/internal/application/service/session_agent_vlm_fallback_test.go new file mode 100644 index 00000000..ea5ee845 --- /dev/null +++ b/internal/application/service/session_agent_vlm_fallback_test.go @@ -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) + } +} diff --git a/internal/application/service/session_knowledge_qa.go b/internal/application/service/session_knowledge_qa.go index 9e1ad706..77d56fcf 100644 --- a/internal/application/service/session_knowledge_qa.go +++ b/internal/application/service/session_knowledge_qa.go @@ -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() diff --git a/internal/application/service/session_knowledge_qa_test.go b/internal/application/service/session_knowledge_qa_test.go index 6800b1ca..c43c78d0 100644 --- a/internal/application/service/session_knowledge_qa_test.go +++ b/internal/application/service/session_knowledge_qa_test.go @@ -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 diff --git a/internal/im/adapter.go b/internal/im/adapter.go index 41ec2234..cef31f68 100644 --- a/internal/im/adapter.go +++ b/internal/im/adapter.go @@ -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 diff --git a/internal/im/dingtalk/adapter.go b/internal/im/dingtalk/adapter.go index f978dc52..28e6e34e 100644 --- a/internal/im/dingtalk/adapter.go +++ b/internal/im/dingtalk/adapter.go @@ -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) } } diff --git a/internal/im/dingtalk/adapter_security_edge_test.go b/internal/im/dingtalk/adapter_security_edge_test.go new file mode 100644 index 00000000..0a1351e4 --- /dev/null +++ b/internal/im/dingtalk/adapter_security_edge_test.go @@ -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) + } +} diff --git a/internal/im/dingtalk/adapter_test.go b/internal/im/dingtalk/adapter_test.go new file mode 100644 index 00000000..0ecda467 --- /dev/null +++ b/internal/im/dingtalk/adapter_test.go @@ -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) + } +} diff --git a/internal/im/dingtalk/finalize_fail_closed_test.go b/internal/im/dingtalk/finalize_fail_closed_test.go new file mode 100644 index 00000000..dea475c3 --- /dev/null +++ b/internal/im/dingtalk/finalize_fail_closed_test.go @@ -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") + } +} diff --git a/internal/im/dingtalk/longconn.go b/internal/im/dingtalk/longconn.go index 0931b912..c6e5203b 100644 --- a/internal/im/dingtalk/longconn.go +++ b/internal/im/dingtalk/longconn.go @@ -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) +} diff --git a/internal/im/dingtalk/longconn_test.go b/internal/im/dingtalk/longconn_test.go new file mode 100644 index 00000000..5579c0e6 --- /dev/null +++ b/internal/im/dingtalk/longconn_test.go @@ -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) + } +} diff --git a/internal/im/image_qa_test.go b/internal/im/image_qa_test.go new file mode 100644 index 00000000..6edbeef6 --- /dev/null +++ b/internal/im/image_qa_test.go @@ -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) + } +} diff --git a/internal/im/service.go b/internal/im/service.go index 31b17a07..b6112f36 100644 --- a/internal/im/service.go +++ b/internal/im/service.go @@ -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 } diff --git a/internal/im/stream_delivery_error_test.go b/internal/im/stream_delivery_error_test.go new file mode 100644 index 00000000..d67b0be5 --- /dev/null +++ b/internal/im/stream_delivery_error_test.go @@ -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") + } +} diff --git a/internal/logger/llm_logger.go b/internal/logger/llm_logger.go index 12ddd55b..b3417aea 100644 --- a/internal/logger/llm_logger.go +++ b/internal/logger/llm_logger.go @@ -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)) diff --git a/internal/logger/logger_test.go b/internal/logger/logger_test.go index 42da1f42..b1992d15 100644 --- a/internal/logger/logger_test.go +++ b/internal/logger/logger_test.go @@ -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{ diff --git a/internal/models/chat/anthropic.go b/internal/models/chat/anthropic.go index 8fbfc84b..c0a47902 100644 --- a/internal/models/chat/anthropic.go +++ b/internal/models/chat/anthropic.go @@ -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, diff --git a/internal/models/chat/anthropic_error_redaction_test.go b/internal/models/chat/anthropic_error_redaction_test.go new file mode 100644 index 00000000..b4dcabd9 --- /dev/null +++ b/internal/models/chat/anthropic_error_redaction_test.go @@ -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) + } +} diff --git a/internal/models/chat/anthropic_terminal_edge_test.go b/internal/models/chat/anthropic_terminal_edge_test.go new file mode 100644 index 00000000..88d467b5 --- /dev/null +++ b/internal/models/chat/anthropic_terminal_edge_test.go @@ -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") +} diff --git a/internal/models/chat/anthropic_test.go b/internal/models/chat/anthropic_test.go index f188177f..f4a9dbac 100644 --- a/internal/models/chat/anthropic_test.go +++ b/internal/models/chat/anthropic_test.go @@ -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}} diff --git a/internal/models/chat/chat.go b/internal/models/chat/chat.go index 47ed6941..d7568eb0 100644 --- a/internal/models/chat/chat.go +++ b/internal/models/chat/chat.go @@ -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, } } diff --git a/internal/models/chat/image_resolve.go b/internal/models/chat/image_resolve.go index 82b263df..0f2b2677 100644 --- a/internal/models/chat/image_resolve.go +++ b/internal/models/chat/image_resolve.go @@ -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 -} diff --git a/internal/models/chat/langfuse_image_test.go b/internal/models/chat/langfuse_image_test.go new file mode 100644 index 00000000..c3b54e85 --- /dev/null +++ b/internal/models/chat/langfuse_image_test.go @@ -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) + } +} diff --git a/internal/models/chat/langfuse_wrapper.go b/internal/models/chat/langfuse_wrapper.go index bc75fdc8..7cbbd1c1 100644 --- a/internal/models/chat/langfuse_wrapper.go +++ b/internal/models/chat/langfuse_wrapper.go @@ -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 diff --git a/internal/models/chat/llm_debug.go b/internal/models/chat/llm_debug.go index 698a6193..2aa76c85 100644 --- a/internal/models/chat/llm_debug.go +++ b/internal/models/chat/llm_debug.go @@ -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))) } } } diff --git a/internal/models/chat/llm_debug_image_test.go b/internal/models/chat/llm_debug_image_test.go new file mode 100644 index 00000000..01aa4f3e --- /dev/null +++ b/internal/models/chat/llm_debug_image_test.go @@ -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) + } +} diff --git a/internal/models/chat/llm_debug_wrapper.go b/internal/models/chat/llm_debug_wrapper.go index a676ef94..9a026b32 100644 --- a/internal/models/chat/llm_debug_wrapper.go +++ b/internal/models/chat/llm_debug_wrapper.go @@ -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() diff --git a/internal/models/chat/model_fail_closed_edge_test.go b/internal/models/chat/model_fail_closed_edge_test.go new file mode 100644 index 00000000..ce5916ab --- /dev/null +++ b/internal/models/chat/model_fail_closed_edge_test.go @@ -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") + } +} diff --git a/internal/models/chat/ollama.go b/internal/models/chat/ollama.go index 989a06fa..7e6ed9c1 100644 --- a/internal/models/chat/ollama.go +++ b/internal/models/chat/ollama.go @@ -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, } }() diff --git a/internal/models/chat/openai_stream.go b/internal/models/chat/openai_stream.go index cec9b1a5..87b619d3 100644 --- a/internal/models/chat/openai_stream.go +++ b/internal/models/chat/openai_stream.go @@ -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 的增量更新 diff --git a/internal/models/chat/plan_a_fail_closed_test.go b/internal/models/chat/plan_a_fail_closed_test.go new file mode 100644 index 00000000..21def1b7 --- /dev/null +++ b/internal/models/chat/plan_a_fail_closed_test.go @@ -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) + } +} diff --git a/internal/models/chat/remote_api.go b/internal/models/chat/remote_api.go index a936406a..f4843b09 100644 --- a/internal/models/chat/remote_api.go +++ b/internal/models/chat/remote_api.go @@ -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) diff --git a/internal/models/chat/remote_api_test.go b/internal/models/chat/remote_api_test.go index db755e52..e0ed56d3 100644 --- a/internal/models/chat/remote_api_test.go +++ b/internal/models/chat/remote_api_test.go @@ -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") + }) +} diff --git a/internal/models/chat/stream_raw_dump.go b/internal/models/chat/stream_raw_dump.go index fad3c20b..050a2ae3 100644 --- a/internal/models/chat/stream_raw_dump.go +++ b/internal/models/chat/stream_raw_dump.go @@ -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 { diff --git a/internal/models/chat/stream_raw_dump_image_test.go b/internal/models/chat/stream_raw_dump_image_test.go new file mode 100644 index 00000000..23231f72 --- /dev/null +++ b/internal/models/chat/stream_raw_dump_image_test.go @@ -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) + } +} diff --git a/internal/models/chat/weknora_cloud_image_test.go b/internal/models/chat/weknora_cloud_image_test.go new file mode 100644 index 00000000..b43eadc8 --- /dev/null +++ b/internal/models/chat/weknora_cloud_image_test.go @@ -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})) +} diff --git a/internal/types/chat_manage.go b/internal/types/chat_manage.go index 15d33314..85090577 100644 --- a/internal/types/chat_manage.go +++ b/internal/types/chat_manage.go @@ -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, diff --git a/internal/types/qa_request.go b/internal/types/qa_request.go index 7932403b..50eb841d 100644 --- a/internal/types/qa_request.go +++ b/internal/types/qa_request.go @@ -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) } diff --git a/internal/utils/log_sanitize.go b/internal/utils/log_sanitize.go index 9e04c48c..f577a15b 100644 --- a/internal/utils/log_sanitize.go +++ b/internal/utils/log_sanitize.go @@ -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] + "..." - }) - + 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 + } +} diff --git a/internal/utils/log_sanitize_edge_test.go b/internal/utils/log_sanitize_edge_test.go new file mode 100644 index 00000000..886b841d --- /dev/null +++ b/internal/utils/log_sanitize_edge_test.go @@ -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) + } + }) + } +} diff --git a/internal/utils/log_sanitize_test.go b/internal/utils/log_sanitize_test.go new file mode 100644 index 00000000..ba6921c9 --- /dev/null +++ b/internal/utils/log_sanitize_test.go @@ -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) + } + } + } +} diff --git a/internal/utils/security.go b/internal/utils/security.go index 0857444c..beb845ca 100644 --- a/internal/utils/security.go +++ b/internal/utils/security.go @@ -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, "" diff --git a/internal/utils/security_stage2_test.go b/internal/utils/security_stage2_test.go new file mode 100644 index 00000000..4740e9ad --- /dev/null +++ b/internal/utils/security_stage2_test.go @@ -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) + } + } +} diff --git a/internal/utils/security_test.go b/internal/utils/security_test.go index 22fba126..56f4b7ea 100644 --- a/internal/utils/security_test.go +++ b/internal/utils/security_test.go @@ -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, ""}, }