Merge pull request #8177 from feeeei/fix/gemini-negative-input-tokens

fix(gemini): prompt 被拦截只报缓存命中时 input_tokens 不再为负

backend/cmd/server/VERSION

@@ -1 +1 @@
-0.2.16
+0.2.17

backend/internal/pkg/antigravity/response_transformer.go

@@ -277,11 +277,12 @@ func (p *NonStreamingProcessor) buildResponse(geminiResp *GeminiResponse, respon
 	}
 
 	// 注意：Gemini 的 promptTokenCount 包含 cachedContentTokenCount，
-	// 但 Claude 的 input_tokens 不包含 cache_read_input_tokens，需要减去
+	// 但 Claude 的 input_tokens 不包含 cache_read_input_tokens，需要减去。
+	// prompt 被拦截时上游可能只报 cachedContentTokenCount，差值下限为 0。
 	usage := ClaudeUsage{}
 	if geminiResp.UsageMetadata != nil {
 		cached := geminiResp.UsageMetadata.CachedContentTokenCount
-		usage.InputTokens = geminiResp.UsageMetadata.PromptTokenCount - cached
+		usage.InputTokens = max(geminiResp.UsageMetadata.PromptTokenCount-cached, 0)
 		usage.OutputTokens = geminiResp.UsageMetadata.CandidatesTokenCount + geminiResp.UsageMetadata.ThoughtsTokenCount
 		usage.CacheReadInputTokens = cached
 		usage.ImageOutputTokens = geminiResp.UsageMetadata.ImageOutputTokens()

backend/internal/pkg/antigravity/stream_transformer.go

@@ -109,10 +109,11 @@ func (p *StreamingProcessor) ProcessLine(line string) []byte {
 
 	// 更新 usage
 	// 注意：Gemini 的 promptTokenCount 包含 cachedContentTokenCount，
-	// 但 Claude 的 input_tokens 不包含 cache_read_input_tokens，需要减去
+	// 但 Claude 的 input_tokens 不包含 cache_read_input_tokens，需要减去。
+	// prompt 被拦截时上游可能只报 cachedContentTokenCount，差值下限为 0。
 	if geminiResp.UsageMetadata != nil {
 		cached := geminiResp.UsageMetadata.CachedContentTokenCount
-		p.inputTokens = geminiResp.UsageMetadata.PromptTokenCount - cached
+		p.inputTokens = max(geminiResp.UsageMetadata.PromptTokenCount-cached, 0)
 		p.outputTokens = geminiResp.UsageMetadata.CandidatesTokenCount + geminiResp.UsageMetadata.ThoughtsTokenCount
 		p.cacheReadTokens = cached
 		p.imageOutputTokens = geminiResp.UsageMetadata.ImageOutputTokens()
@@ -190,7 +191,7 @@ func (p *StreamingProcessor) emitMessageStart(v1Resp *V1InternalResponse) []byte
 	usage := ClaudeUsage{}
 	if v1Resp.Response.UsageMetadata != nil {
 		cached := v1Resp.Response.UsageMetadata.CachedContentTokenCount
-		usage.InputTokens = v1Resp.Response.UsageMetadata.PromptTokenCount - cached
+		usage.InputTokens = max(v1Resp.Response.UsageMetadata.PromptTokenCount-cached, 0)
 		usage.OutputTokens = v1Resp.Response.UsageMetadata.CandidatesTokenCount + v1Resp.Response.UsageMetadata.ThoughtsTokenCount
 		usage.CacheReadInputTokens = cached
 		usage.ImageOutputTokens = v1Resp.Response.UsageMetadata.ImageOutputTokens()

backend/internal/pkg/antigravity/usage_contract_test.go

@@ -38,3 +38,29 @@ func TestGeminiUsageMapping_NoCacheCreationTokens(t *testing.T) {
 		require.Zero(t, usage.CacheCreationInputTokens)
 	})
 }
+
+// usageMetadata 缺 promptTokenCount、只有 cachedContentTokenCount 时，input_tokens 下限为 0，
+// 缓存命中仍计入 cache_read_input_tokens。
+func TestGeminiUsageMapping_MissingPromptTokenCountKeepsInputNonNegative(t *testing.T) {
+	const geminiBody = `{"candidates":[{"content":{"parts":[{"text":"hi"}],"role":"model"},"finishReason":"STOP"}],` +
+		`"usageMetadata":{"cachedContentTokenCount":30}}`
+
+	t.Run("non-stream", func(t *testing.T) {
+		_, usage, err := TransformGeminiToClaude([]byte(geminiBody), "gemini-3.1-pro-preview")
+		require.NoError(t, err)
+		require.NotNil(t, usage)
+		require.Zero(t, usage.InputTokens)
+		require.Equal(t, 30, usage.CacheReadInputTokens)
+	})
+
+	t.Run("stream", func(t *testing.T) {
+		p := NewStreamingProcessor("gemini-3.1-pro-preview")
+		out := p.ProcessLine(`data: {"response":` + geminiBody + `}`)
+		require.True(t, strings.Contains(string(out), `"message_start"`))
+		require.NotContains(t, string(out), `"input_tokens":-`, "message_start 的 input_tokens 不应为负")
+		_, usage := p.Finish()
+		require.NotNil(t, usage)
+		require.Zero(t, usage.InputTokens)
+		require.Equal(t, 30, usage.CacheReadInputTokens)
+	})
+}

backend/internal/service/gemini_chat_completions_compat_service.go

@@ -250,6 +250,8 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions(
 
 	var usage *ClaudeUsage
 	var firstTokenMs *int
+	clientDisconnect := false
+	streamFinished := false
 	if clientStream {
 		streamRes, err := s.handleChatCompletionsStreamingResponseFromGemini(c, resp, startTime, originalModel, account.Type == AccountTypeOAuth, includeUsage)
 		if err != nil {
@@ -257,6 +259,8 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions(
 		}
 		usage = streamRes.usage
 		firstTokenMs = streamRes.firstTokenMs
+		clientDisconnect = streamRes.clientDisconnect
+		streamFinished = streamRes.finished
 	} else if useUpstreamStream {
 		collected, usageObj, err := collectGeminiSSE(resp.Body, account.Type == AccountTypeOAuth)
 		if err != nil {
@@ -287,6 +291,10 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions(
 	if isImageGenerationModel(originalModel) {
 		imageCount = 1
 	}
+	// 客户端在收到结束分片前断开时只计已送达的图片：Chat 流不向客户端转发图片。
+	if clientDisconnect && !streamFinished {
+		imageCount = 0
+	}
 
 	return &ForwardResult{
 		RequestID:        requestID,
@@ -301,7 +309,7 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions(
 		ImageCount:       imageCount,
 		ImageSize:        imageSize,
 		ImageInputSize:   imageInputSize,
-		ClientDisconnect: false,
+		ClientDisconnect: clientDisconnect,
 	}, nil
 }
 
@@ -532,6 +540,11 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom
 	var usage ClaudeUsage
 	var firstTokenMs *int
 	firstChunk := true
+	finishReason := ""
+	clientDisconnect := false
+	streamResult := func() *geminiStreamResult {
+		return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs, clientDisconnect: clientDisconnect, finished: finishReason != ""}
+	}
 
 	writeChatChunk := func(chunk apicompat.ChatCompletionsChunk) bool {
 		sse, err := apicompat.ChatChunkToSSE(chunk)
@@ -539,6 +552,7 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom
 			return false
 		}
 		if _, err := io.WriteString(c.Writer, sse); err != nil {
+			clientDisconnect = true
 			return true
 		}
 		return false
@@ -571,10 +585,9 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom
 			Usage:      apicompat.AnthropicUsage{},
 		},
 	}) {
-		return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+		return streamResult(), nil
 	}
 
-	finishReason := ""
 	sawToolUse := false
 	nextBlockIndex := 0
 	openBlockIndex := -1
@@ -643,12 +656,12 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom
 								}
 								if openToolIndex >= 0 {
 									if closeOpenTool() {
-										return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+										return streamResult(), nil
 									}
 								}
 								if openBlockType != "thinking" {
 									if closeOpenBlock() {
-										return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+										return streamResult(), nil
 									}
 									idx := nextBlockIndex
 									nextBlockIndex++
@@ -659,21 +672,21 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom
 										Index:        &idx,
 										ContentBlock: &apicompat.AnthropicContentBlock{Type: "thinking"},
 									}) {
-										return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+										return streamResult(), nil
 									}
 								}
 								if emitAnthropicEvent(&apicompat.AnthropicStreamEvent{
 									Type:  "content_block_delta",
 									Delta: &apicompat.AnthropicDelta{Type: "thinking_delta", Thinking: delta},
 								}) {
-									return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+									return streamResult(), nil
 								}
 								continue
 							}
 							if text, ok := part["text"].(string); ok && text != "" {
 								if openToolIndex >= 0 {
 									if closeOpenTool() {
-										return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+										return streamResult(), nil
 									}
 								}
 								delta, newSeen := computeGeminiTextDelta(seenText, text)
@@ -683,7 +696,7 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom
 								}
 								if openBlockType != "text" {
 									if closeOpenBlock() {
-										return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+										return streamResult(), nil
 									}
 									idx := nextBlockIndex
 									nextBlockIndex++
@@ -697,7 +710,7 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom
 											Text: "",
 										},
 									}) {
-										return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+										return streamResult(), nil
 									}
 								}
 								if emitAnthropicEvent(&apicompat.AnthropicStreamEvent{
@@ -707,7 +720,7 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom
 										Text: delta,
 									},
 								}) {
-									return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+									return streamResult(), nil
 								}
 								continue
 							}
@@ -718,11 +731,11 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom
 									name = "tool"
 								}
 								if closeOpenBlock() {
-									return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+									return streamResult(), nil
 								}
 								if openToolIndex >= 0 && openToolName != name {
 									if closeOpenTool() {
-										return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+										return streamResult(), nil
 									}
 								}
 								if openToolIndex < 0 {
@@ -741,7 +754,7 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom
 											Input: json.RawMessage(`{}`),
 										},
 									}) {
-										return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+										return streamResult(), nil
 									}
 								}
 
@@ -767,7 +780,7 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom
 											PartialJSON: delta,
 										},
 									}) {
-										return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+										return streamResult(), nil
 									}
 								}
 							}
@@ -781,15 +794,20 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom
 			break
 		}
 		if err != nil {
+			if geminiStreamAbortedWithBillableOutput(c, err, &usage) {
+				logger.LegacyPrintf("service.gemini_messages_compat", "[GeminiAPI] Client disconnected during chat completions stream, billing collected usage: %v", err)
+				clientDisconnect = true
+				return streamResult(), nil
+			}
 			return nil, fmt.Errorf("stream read error: %w", err)
 		}
 	}
 
 	if closeOpenBlock() {
-		return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+		return streamResult(), nil
 	}
 	if closeOpenTool() {
-		return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+		return streamResult(), nil
 	}
 
 	stopReason := mapGeminiFinishReasonToClaudeStopReason(finishReason)
@@ -810,30 +828,30 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom
 			CacheReadInputTokens: usage.CacheReadInputTokens,
 		},
 	}) {
-		return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+		return streamResult(), nil
 	}
 	if emitAnthropicEvent(&apicompat.AnthropicStreamEvent{Type: "message_stop"}) {
-		return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+		return streamResult(), nil
 	}
 
 	for _, resEvt := range apicompat.FinalizeAnthropicResponsesStream(anthState) {
 		chunks := apicompat.ResponsesEventToChatChunks(&resEvt, ccState)
 		for _, chunk := range chunks {
 			if disconnected := writeChatChunk(chunk); disconnected {
-				return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+				return streamResult(), nil
 			}
 		}
 	}
 	for _, chunk := range apicompat.FinalizeResponsesChatStream(ccState) {
 		if disconnected := writeChatChunk(chunk); disconnected {
-			return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+			return streamResult(), nil
 		}
 	}
 
 	_, _ = io.WriteString(c.Writer, "data: [DONE]\n\n")
 	flusher.Flush()
 
-	return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
+	return streamResult(), nil
 }
 
 func (s *GeminiMessagesCompatService) writeGeminiChatCompletionsMappedError(

backend/internal/service/gemini_messages_compat_service.go

@@ -1015,6 +1015,8 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
 
 	var usage *ClaudeUsage
 	var firstTokenMs *int
+	clientDisconnect := false
+	streamFinished := false
 	if req.Stream {
 		streamRes, err := s.handleStreamingResponse(c, resp, startTime, originalModel)
 		if err != nil {
@@ -1022,6 +1024,8 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
 		}
 		usage = streamRes.usage
 		firstTokenMs = streamRes.firstTokenMs
+		clientDisconnect = streamRes.clientDisconnect
+		streamFinished = streamRes.finished
 	} else {
 		if useUpstreamStream {
 			collected, usageObj, err := collectGeminiSSE(resp.Body, true)
@@ -1049,6 +1053,10 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
 	imageInputSize := s.extractImageInputSize(body)
 	imageSize := normalizeOpenAIImageSizeTier(imageInputSize)
 	imageCount := resolveGeminiImageCount(c, originalModel, mappedModel)
+	// 客户端在收到结束分片前断开时只计已送达的图片：Claude 兼容流不向客户端转发图片。
+	if clientDisconnect && !streamFinished {
+		imageCount = 0
+	}
 
 	return &ForwardResult{
 		RequestID:                     requestID,
@@ -1062,6 +1070,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
 		Duration:                      time.Since(startTime),
 		FirstTokenMs:                  firstTokenMs,
 		ReasoningEffort:               extractGeminiReasoningEffortFromBody(geminiReq),
+		ClientDisconnect:              clientDisconnect,
 		ImageCount:                    imageCount,
 		ImageSize:                     imageSize,
 		ImageInputSize:                imageInputSize,
@@ -1530,6 +1539,8 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
 
 	var usage *ClaudeUsage
 	var firstTokenMs *int
+	clientDisconnect := false
+	streamFinished := false
 
 	if stream {
 		streamRes, err := s.handleNativeStreamingResponse(c, resp, startTime, isOAuth, account, requestID)
@@ -1538,6 +1549,8 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
 		}
 		usage = streamRes.usage
 		firstTokenMs = streamRes.firstTokenMs
+		clientDisconnect = streamRes.clientDisconnect
+		streamFinished = streamRes.finished
 	} else {
 		if useUpstreamStream {
 			var best geminiResponseSignal
@@ -1572,6 +1585,10 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
 	imageInputSize := s.extractImageInputSize(body)
 	imageSize := normalizeOpenAIImageSizeTier(imageInputSize)
 	imageCount := resolveGeminiImageCount(c, originalModel, mappedModel)
+	// 客户端在收到结束分片前断开时只计已送达的图片。
+	if clientDisconnect && !streamFinished {
+		imageCount = observedGeminiImageOutputs(c)
+	}
 
 	return &ForwardResult{
 		RequestID:                     requestID,
@@ -1585,6 +1602,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
 		ReasoningEffort:               extractGeminiReasoningEffortFromBody(body),
 		Duration:                      time.Since(startTime),
 		FirstTokenMs:                  firstTokenMs,
+		ClientDisconnect:              clientDisconnect,
 		ImageCount:                    imageCount,
 		ImageSize:                     imageSize,
 		ImageInputSize:                imageInputSize,
@@ -2068,8 +2086,10 @@ func mapGeminiStatusToClaudeErrorType(status string) string {
 }
 
 type geminiStreamResult struct {
-	usage        *ClaudeUsage
-	firstTokenMs *int
+	usage            *ClaudeUsage
+	firstTokenMs     *int
+	clientDisconnect bool // 客户端在流式传输中途断开，结果为断开前已收集的部分
+	finished         bool // 断开前已收到上游 finishReason
 }
 
 func (s *GeminiMessagesCompatService) handleNonStreamingResponse(c *gin.Context, resp *http.Response, originalModel string) (*ClaudeUsage, error) {
@@ -2151,6 +2171,15 @@ func (s *GeminiMessagesCompatService) handleStreamingResponse(c *gin.Context, re
 	for {
 		line, err := reader.ReadString('\n')
 		if err != nil && !errors.Is(err, io.EOF) {
+			if geminiStreamAbortedWithBillableOutput(c, err, &usage) {
+				logger.LegacyPrintf("service.gemini_messages_compat", "[GeminiAPI] Client disconnected during messages stream, billing collected usage: %v", err)
+				return &geminiStreamResult{
+					usage:            &usage,
+					firstTokenMs:     firstTokenMs,
+					clientDisconnect: true,
+					finished:         finishReason != "",
+				}, nil
+			}
 			return nil, fmt.Errorf("stream read error: %w", err)
 		}
 
@@ -2660,8 +2689,10 @@ func mergeCollectedTextParts(response map[string]any, textParts []string) map[st
 }
 
 type geminiNativeStreamResult struct {
-	usage        *ClaudeUsage
-	firstTokenMs *int
+	usage            *ClaudeUsage
+	firstTokenMs     *int
+	clientDisconnect bool // 客户端在流式传输中途断开，结果为断开前已收集的部分
+	finished         bool // 断开前已收到上游 finishReason
 }
 
 func isGeminiInsufficientScope(headers http.Header, body []byte) bool {
@@ -2838,10 +2869,13 @@ func (s *GeminiMessagesCompatService) handleNativeStreamingResponse(c *gin.Conte
 	var best geminiResponseSignal
 	sawDataEvent := false
 	fallback := &geminiSSEFallbackBody{}
+	clientDisconnect := false
+	finished := false
 
 	for {
 		line, err := reader.ReadString('\n')
-		if len(line) > 0 {
+		// 读错误时 line 是未以换行结尾的残行，不透传也不参与用量统计；EOF 时的末行完整有效。
+		if len(line) > 0 && (err == nil || errors.Is(err, io.EOF)) {
 			trimmed := strings.TrimRight(line, "\r\n")
 			if strings.HasPrefix(trimmed, "data:") {
 				payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:"))
@@ -2874,6 +2908,10 @@ func (s *GeminiMessagesCompatService) handleNativeStreamingResponse(c *gin.Conte
 					observer.ObserveGemini(rawBytes)
 					observeGeminiImageOutputs(c, rawBytes)
 
+					if gjson.GetBytes(rawBytes, "candidates.0.finishReason").String() != "" {
+						finished = true
+					}
+
 					if firstTokenMs == nil {
 						ms := int(time.Since(startTime).Milliseconds())
 						firstTokenMs = &ms
@@ -2901,13 +2939,18 @@ func (s *GeminiMessagesCompatService) handleNativeStreamingResponse(c *gin.Conte
 			break
 		}
 		if err != nil {
-			return nil, err
+			if !geminiStreamAbortedWithBillableOutput(c, err, usage) {
+				return nil, err
+			}
+			clientDisconnect = true
+			logger.LegacyPrintf("service.gemini_messages_compat", "[GeminiAPI] Client disconnected during native stream (account=%d), billing collected usage: %v", account.ID, err)
+			break
 		}
 	}
 
 	s.finalizeGeminiSSESignal(c, account, true, upstreamRequestID, best, sawDataEvent, fallback)
 
-	return &geminiNativeStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil
+	return &geminiNativeStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: clientDisconnect, finished: finished}, nil
 }
 
 // unwrapGeminiResponse 解包 Gemini OAuth 响应中的 response 字段
@@ -3036,9 +3079,10 @@ func extractGeminiUsage(data []byte) *ClaudeUsage {
 	}
 
 	// 注意：Gemini 的 promptTokenCount 包含 cachedContentTokenCount，
-	// 但 Claude 的 input_tokens 不包含 cache_read_input_tokens，需要减去
+	// 但 Claude 的 input_tokens 不包含 cache_read_input_tokens，需要减去。
+	// prompt 被拦截时上游可能只报 cachedContentTokenCount，差值下限为 0。
 	return &ClaudeUsage{
-		InputTokens:          prompt - cached,
+		InputTokens:          max(prompt-cached, 0),
 		OutputTokens:         cand + thoughts,
 		CacheReadInputTokens: cached,
 		ImageOutputTokens:    imageTokens,

backend/internal/service/gemini_messages_compat_service_test.go

@@ -947,6 +947,16 @@ func TestExtractGeminiUsage(t *testing.T) {
 				CacheReadInputTokens: 0,
 			},
 		},
+		{
+			name:    "prompt 被拦截只报 cachedContentTokenCount — input 下限为 0",
+			input:   `{"promptFeedback":{"blockReason":"PROHIBITED_CONTENT"},"usageMetadata":{"cachedContentTokenCount":12980}}`,
+			wantNil: false,
+			wantUsage: &ClaudeUsage{
+				InputTokens:          0,
+				OutputTokens:         0,
+				CacheReadInputTokens: 12980,
+			},
+		},
 		{
 			name:    "无 usageMetadata",
 			input:   `{"candidates":[]}`,

backend/internal/service/gemini_native_response_signal_test.go

@@ -180,6 +180,29 @@ func TestGeminiForwardNative_StreamPromptBlockedMarksContentPolicy(t *testing.T)
 	require.Empty(t, upstreamErrorEventsFromContext(t, c))
 }
 
+// Vertex 拦截 prompt 时 usageMetadata 可能只带 cachedContentTokenCount、缺 promptTokenCount：
+// input 下限为 0，缓存命中照上游所报计入 cache_read。
+func TestGeminiForwardNative_StreamPromptBlockedWithOnlyCachedTokensKeepsInputNonNegative(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	body := `data: {"promptFeedback":{"blockReason":"PROHIBITED_CONTENT"},"usageMetadata":{"cachedContentTokenCount":12980,"totalTokenCount":0,"trafficType":"ON_DEMAND"},"modelVersion":"gemini-3.5-flash-lite"}
+
+`
+	svc := newGeminiSignalService("text/event-stream", body)
+	c, rec := newGeminiNativeTestContext(t)
+
+	result, err := svc.ForwardNative(context.Background(), c, geminiSignalTestAccount(),
+		"gemini-3.5-flash-lite", "streamGenerateContent", true, geminiSignalTestRequest())
+	require.NoError(t, err)
+	require.Equal(t, body, rec.Body.String())
+	require.Equal(t, 0, result.Usage.InputTokens)
+	require.Equal(t, 12980, result.Usage.CacheReadInputTokens)
+	require.Equal(t, 0, result.Usage.OutputTokens)
+
+	streamErrs := GetOpsStreamErrors(c)
+	require.Len(t, streamErrs, 1)
+	require.Equal(t, "PROHIBITED_CONTENT", streamErrs[0].Code)
+}
+
 func TestGeminiForwardNative_EmptyStreamMarksUpstreamFailure(t *testing.T) {
 	gin.SetMode(gin.TestMode)
 	svc := newGeminiSignalService("text/event-stream", "")

backend/internal/service/gemini_stream_client_abort.go

@@ -0,0 +1,25 @@
+package service
+
+import (
+	"context"
+	"errors"
+	"io"
+
+	"github.com/gin-gonic/gin"
+)
+
+// geminiStreamAbortedWithBillableOutput 判断流式读取中断是否由客户端断开引起，且断开前
+// 已观测到可计费的输出（token 用量或已回吐的图片）。满足时调用方以已收集的结果收尾并照常计费。
+// usageMetadata 为累计值，按断开前收到的最后一份计费：读到结束分片后断开即完整用量；
+// 中途断开时中间分片可能只含输入计数或不含用量，此时按已有计数计费或不计费。
+func geminiStreamAbortedWithBillableOutput(c *gin.Context, err error, usage *ClaudeUsage) bool {
+	if err == nil || errors.Is(err, io.EOF) {
+		return false
+	}
+	clientGone := errors.Is(err, context.Canceled) ||
+		(c != nil && c.Request != nil && c.Request.Context().Err() != nil)
+	if !clientGone {
+		return false
+	}
+	return usage.hasObservedTokens() || observedGeminiImageOutputs(c) > 0
+}

backend/internal/service/gemini_stream_client_abort_test.go

@@ -0,0 +1,244 @@
+//go:build unit
+
+package service
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/require"
+)
+
+// geminiAbortingStreamBody 先吐出 data，读完后返回 abortErr；onAbort 用于模拟客户端断开时请求 context 被取消。
+type geminiAbortingStreamBody struct {
+	data     *bytes.Reader
+	abortErr error
+	onAbort  func()
+}
+
+func (b *geminiAbortingStreamBody) Read(p []byte) (int, error) {
+	n, err := b.data.Read(p)
+	if err == io.EOF {
+		if b.onAbort != nil {
+			b.onAbort()
+		}
+		return n, b.abortErr
+	}
+	return n, err
+}
+
+func (b *geminiAbortingStreamBody) Close() error { return nil }
+
+const (
+	geminiAbortTextChunk  = `data: {"candidates":[{"content":{"parts":[{"text":"hello"}],"role":"model"}}]}` + "\n\n"
+	geminiAbortFinalChunk = `data: {"candidates":[{"content":{"parts":[{"text":" world"}],"role":"model"},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5,"totalTokenCount":15}}` + "\n\n"
+	geminiAbortImageChunk = `data: {"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"aGVsbG8="}}],"role":"model"},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":1290,"totalTokenCount":1300}}` + "\n\n"
+	geminiAbortUsageChunk = `data: {"candidates":[{"content":{"parts":[{"text":"drawing"}],"role":"model"}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2,"totalTokenCount":12}}` + "\n\n"
+)
+
+// newGeminiAbortingStreamService 返回的上游流在吐完 sse 后中断；clientGone=true 时中断伴随请求 context 取消。
+func newGeminiAbortingStreamService(t *testing.T, path string, reqBody []byte, sse string, clientGone bool) (*GeminiMessagesCompatService, *gin.Context, *httptest.ResponseRecorder) {
+	t.Helper()
+	gin.SetMode(gin.TestMode)
+	rec := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(rec)
+	reqCtx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	c.Request = httptest.NewRequest(http.MethodPost, path, bytes.NewReader(reqBody)).WithContext(reqCtx)
+
+	body := &geminiAbortingStreamBody{data: bytes.NewReader([]byte(sse)), abortErr: io.ErrUnexpectedEOF}
+	if clientGone {
+		body.abortErr = context.Canceled
+		body.onAbort = cancel
+	}
+	svc, httpStub := newGeminiVertexAPIKeyTestService()
+	httpStub.response = &http.Response{
+		StatusCode: http.StatusOK,
+		Header:     http.Header{"Content-Type": []string{"text/event-stream"}},
+		Body:       body,
+	}
+	return svc, c, rec
+}
+
+func TestGeminiForwardNative_ClientDisconnectAfterFinalChunkBillsCollectedUsage(t *testing.T) {
+	reqBody := []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`)
+	svc, c, _ := newGeminiAbortingStreamService(t, "/v1beta/models/gemini-3-flash-preview:streamGenerateContent", reqBody, geminiAbortTextChunk+geminiAbortFinalChunk, true)
+
+	result, err := svc.ForwardNative(context.Background(), c, geminiVertexAPIKeyTestAccount(""), "gemini-3-flash-preview", "streamGenerateContent", true, reqBody)
+	require.NoError(t, err)
+	require.NotNil(t, result)
+	require.True(t, result.ClientDisconnect)
+	require.Equal(t, 10, result.Usage.InputTokens)
+	require.Equal(t, 5, result.Usage.OutputTokens)
+	require.Zero(t, result.ImageCount)
+}
+
+func TestGeminiForwardNative_ClientDisconnectWithoutUsageStaysUnbilled(t *testing.T) {
+	reqBody := []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`)
+	svc, c, _ := newGeminiAbortingStreamService(t, "/v1beta/models/gemini-3-flash-preview:streamGenerateContent", reqBody, geminiAbortTextChunk, true)
+
+	result, err := svc.ForwardNative(context.Background(), c, geminiVertexAPIKeyTestAccount(""), "gemini-3-flash-preview", "streamGenerateContent", true, reqBody)
+	require.ErrorIs(t, err, context.Canceled)
+	require.Nil(t, result)
+}
+
+func TestGeminiForwardNative_UpstreamReadErrorWithLiveClientStaysError(t *testing.T) {
+	reqBody := []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`)
+	svc, c, _ := newGeminiAbortingStreamService(t, "/v1beta/models/gemini-3-flash-preview:streamGenerateContent", reqBody, geminiAbortTextChunk+geminiAbortFinalChunk, false)
+
+	result, err := svc.ForwardNative(context.Background(), c, geminiVertexAPIKeyTestAccount(""), "gemini-3-flash-preview", "streamGenerateContent", true, reqBody)
+	require.ErrorIs(t, err, io.ErrUnexpectedEOF)
+	require.Nil(t, result)
+}
+
+func TestGeminiForwardNative_ClientDisconnectBillsOnlyDeliveredImages(t *testing.T) {
+	reqBody := []byte(`{"contents":[{"role":"user","parts":[{"text":"draw a cat"}]}]}`)
+
+	t.Run("image delivered before disconnect", func(t *testing.T) {
+		svc, c, _ := newGeminiAbortingStreamService(t, "/v1beta/models/gemini-3-pro-image:streamGenerateContent", reqBody, geminiAbortImageChunk, true)
+		result, err := svc.ForwardNative(context.Background(), c, geminiVertexAPIKeyTestAccount(""), "gemini-3-pro-image", "streamGenerateContent", true, reqBody)
+		require.NoError(t, err)
+		require.True(t, result.ClientDisconnect)
+		require.Equal(t, 1, result.ImageCount)
+	})
+
+	t.Run("disconnect before the image arrives", func(t *testing.T) {
+		svc, c, _ := newGeminiAbortingStreamService(t, "/v1beta/models/gemini-3-pro-image:streamGenerateContent", reqBody, geminiAbortUsageChunk, true)
+		result, err := svc.ForwardNative(context.Background(), c, geminiVertexAPIKeyTestAccount(""), "gemini-3-pro-image", "streamGenerateContent", true, reqBody)
+		require.NoError(t, err)
+		require.True(t, result.ClientDisconnect)
+		require.Zero(t, result.ImageCount)
+		require.Equal(t, 2, result.Usage.OutputTokens)
+	})
+}
+
+func TestGeminiCompatForward_ClientDisconnectAfterFinalChunkBillsCollectedUsage(t *testing.T) {
+	reqBody := []byte(`{"model":"gemini-3-flash-preview","max_tokens":16,"stream":true,"messages":[{"role":"user","content":"hi"}]}`)
+
+	t.Run("client gone", func(t *testing.T) {
+		svc, c, _ := newGeminiAbortingStreamService(t, "/v1/messages", reqBody, geminiAbortTextChunk+geminiAbortFinalChunk, true)
+		result, err := svc.Forward(context.Background(), c, geminiVertexAPIKeyTestAccount(""), reqBody)
+		require.NoError(t, err)
+		require.True(t, result.ClientDisconnect)
+		require.Equal(t, 10, result.Usage.InputTokens)
+		require.Equal(t, 5, result.Usage.OutputTokens)
+	})
+
+	t.Run("upstream read error with live client", func(t *testing.T) {
+		svc, c, _ := newGeminiAbortingStreamService(t, "/v1/messages", reqBody, geminiAbortTextChunk+geminiAbortFinalChunk, false)
+		result, err := svc.Forward(context.Background(), c, geminiVertexAPIKeyTestAccount(""), reqBody)
+		require.ErrorIs(t, err, io.ErrUnexpectedEOF)
+		require.Nil(t, result)
+	})
+}
+
+func TestGeminiForwardAsChatCompletions_ClientDisconnectAfterFinalChunkBillsCollectedUsage(t *testing.T) {
+	reqBody := []byte(`{"model":"gemini-3-flash-preview","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
+
+	t.Run("client gone", func(t *testing.T) {
+		svc, c, _ := newGeminiAbortingStreamService(t, "/v1/chat/completions", reqBody, geminiAbortTextChunk+geminiAbortFinalChunk, true)
+		result, err := svc.ForwardAsChatCompletions(context.Background(), c, geminiVertexAPIKeyTestAccount(""), reqBody)
+		require.NoError(t, err)
+		require.True(t, result.ClientDisconnect)
+		require.Equal(t, 10, result.Usage.InputTokens)
+		require.Equal(t, 5, result.Usage.OutputTokens)
+	})
+
+	t.Run("client gone before usage", func(t *testing.T) {
+		svc, c, _ := newGeminiAbortingStreamService(t, "/v1/chat/completions", reqBody, geminiAbortTextChunk, true)
+		result, err := svc.ForwardAsChatCompletions(context.Background(), c, geminiVertexAPIKeyTestAccount(""), reqBody)
+		require.ErrorIs(t, err, context.Canceled)
+		require.Nil(t, result)
+	})
+}
+
+func TestGeminiForwardNative_ClientDisconnectAfterFinishKeepsImageFallback(t *testing.T) {
+	// 结束分片已到达但图片不以 inlineData 回吐时，断开与正常完成同口径按模型名兜底。
+	finishWithoutInline := `data: {"candidates":[{"content":{"parts":[{"fileData":{"mimeType":"image/png","fileUri":"gs://bucket/cat.png"}}],"role":"model"},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":1290,"totalTokenCount":1300}}` + "\n\n"
+	reqBody := []byte(`{"contents":[{"role":"user","parts":[{"text":"draw a cat"}]}]}`)
+	svc, c, _ := newGeminiAbortingStreamService(t, "/v1beta/models/gemini-3-pro-image:streamGenerateContent", reqBody, finishWithoutInline, true)
+
+	result, err := svc.ForwardNative(context.Background(), c, geminiVertexAPIKeyTestAccount(""), "gemini-3-pro-image", "streamGenerateContent", true, reqBody)
+	require.NoError(t, err)
+	require.True(t, result.ClientDisconnect)
+	require.Equal(t, 1, result.ImageCount)
+}
+
+func TestGeminiForwardNative_ClientDisconnectIgnoresTruncatedTrailingLine(t *testing.T) {
+	complete := `data: {"candidates":[{"content":{"parts":[{"text":"hello"}],"role":"model"}}],"usageMetadata":{"promptTokenCount":1000,"candidatesTokenCount":300,"totalTokenCount":1300}}` + "\n\n"
+	truncated := `data: {"candidates":[{"content":{"parts":[{"text":" more"}],"role":"model"}}],"usageMetadata":{`
+	reqBody := []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`)
+	svc, c, _ := newGeminiAbortingStreamService(t, "/v1beta/models/gemini-3-flash-preview:streamGenerateContent", reqBody, complete+truncated, true)
+
+	result, err := svc.ForwardNative(context.Background(), c, geminiVertexAPIKeyTestAccount(""), "gemini-3-flash-preview", "streamGenerateContent", true, reqBody)
+	require.NoError(t, err)
+	require.True(t, result.ClientDisconnect)
+	require.Equal(t, 1000, result.Usage.InputTokens)
+	require.Equal(t, 300, result.Usage.OutputTokens)
+}
+
+func TestGeminiForwardAsChatCompletions_ClientDisconnectImageCount(t *testing.T) {
+	t.Run("disconnect before finish bills no image", func(t *testing.T) {
+		reqBody := []byte(`{"model":"gemini-3-pro-image","stream":true,"messages":[{"role":"user","content":"draw a cat"}]}`)
+		svc, c, _ := newGeminiAbortingStreamService(t, "/v1/chat/completions", reqBody, geminiAbortUsageChunk, true)
+		result, err := svc.ForwardAsChatCompletions(context.Background(), c, geminiVertexAPIKeyTestAccount(""), reqBody)
+		require.NoError(t, err)
+		require.True(t, result.ClientDisconnect)
+		require.Zero(t, result.ImageCount)
+		require.Equal(t, 2, result.Usage.OutputTokens)
+	})
+
+	t.Run("disconnect after finish keeps completed-request image count", func(t *testing.T) {
+		reqBody := []byte(`{"model":"gemini-3-pro-image","stream":true,"messages":[{"role":"user","content":"draw a cat"}]}`)
+		svc, c, _ := newGeminiAbortingStreamService(t, "/v1/chat/completions", reqBody, geminiAbortImageChunk, true)
+		result, err := svc.ForwardAsChatCompletions(context.Background(), c, geminiVertexAPIKeyTestAccount(""), reqBody)
+		require.NoError(t, err)
+		require.True(t, result.ClientDisconnect)
+		require.Equal(t, 1, result.ImageCount)
+	})
+}
+
+// geminiFailingResponseWriter 模拟客户端已断开：所有写入失败。
+type geminiFailingResponseWriter struct {
+	header http.Header
+}
+
+func (w *geminiFailingResponseWriter) Header() http.Header       { return w.header }
+func (w *geminiFailingResponseWriter) Write([]byte) (int, error) { return 0, io.ErrClosedPipe }
+func (w *geminiFailingResponseWriter) WriteHeader(int)           {}
+func (w *geminiFailingResponseWriter) Flush()                    {}
+
+func TestGeminiForwardAsChatCompletions_WriteFailureOnImageModelBillsNoImage(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	reqBody := []byte(`{"model":"gemini-3-pro-image","stream":true,"messages":[{"role":"user","content":"draw a cat"}]}`)
+	c, _ := gin.CreateTestContext(&geminiFailingResponseWriter{header: http.Header{}})
+	c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(reqBody))
+	svc, httpStub := newGeminiVertexAPIKeyTestService()
+	httpStub.response = &http.Response{
+		StatusCode: http.StatusOK,
+		Header:     http.Header{"Content-Type": []string{"text/event-stream"}},
+		Body:       io.NopCloser(bytes.NewReader([]byte(geminiAbortUsageChunk))),
+	}
+
+	result, err := svc.ForwardAsChatCompletions(context.Background(), c, geminiVertexAPIKeyTestAccount(""), reqBody)
+	require.NoError(t, err)
+	require.True(t, result.ClientDisconnect)
+	require.Zero(t, result.ImageCount)
+}
+
+func TestGeminiCompatForward_ClientDisconnectBeforeFinishBillsNoImage(t *testing.T) {
+	inlineWithoutFinish := `data: {"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"aGVsbG8="}}],"role":"model"}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":1290,"totalTokenCount":1300}}` + "\n\n"
+	reqBody := []byte(`{"model":"gemini-3-pro-image","max_tokens":16,"stream":true,"messages":[{"role":"user","content":"draw a cat"}]}`)
+	svc, c, _ := newGeminiAbortingStreamService(t, "/v1/messages", reqBody, inlineWithoutFinish, true)
+
+	result, err := svc.Forward(context.Background(), c, geminiVertexAPIKeyTestAccount(""), reqBody)
+	require.NoError(t, err)
+	require.True(t, result.ClientDisconnect)
+	require.Zero(t, result.ImageCount)
+	require.Equal(t, 1290, result.Usage.OutputTokens)
+}

backend/internal/service/image_storage.go

@@ -8,10 +8,16 @@ import (
 	"fmt"
 	"io"
 	"mime"
+	"net"
 	"net/http"
+	"net/url"
 	"strconv"
 	"strings"
+	"sync"
+	"syscall"
 	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
 )
 
 const defaultImageMaxDownloadBytes int64 = 32 << 20 // 32 MiB
@@ -53,8 +59,123 @@ func NewImageResultUploader(storage ImageStorage, prefix string, maxDownloadByte
 	}
 }
 
-func defaultImageDownloadHTTPClient() *http.Client {
-	return &http.Client{Timeout: 60 * time.Second}
+// defaultImageDownloadHTTPClient 是进程内共享的图片下载客户端，复用连接池。
+var defaultImageDownloadHTTPClient = sync.OnceValue(func() *http.Client {
+	return newImageDownloadHTTPClient(http.ProxyFromEnvironment)
+})
+
+// newImageDownloadHTTPClient 返回只允许下载公网目的地的客户端。
+// 图片 url 来自上游响应，是不可信输入：目的地为 localhost、回环、私网、链路本地、
+// 云元数据等地址时一律拒绝，重定向的每一跳都经同一 RoundTrip 校验。
+// 不经代理时在 socket 层校验真实连接 IP（防 DNS rebinding）；经环境代理时由代理解析目标，
+// 只能在发出前校验本地解析结果。
+func newImageDownloadHTTPClient(proxy func(*http.Request) (*url.URL, error)) *http.Client {
+	direct := newImageDownloadBaseTransport(imageDownloadDialControl)
+	proxied := newImageDownloadBaseTransport(nil)
+	proxied.Proxy = proxy
+	return &http.Client{
+		Timeout:   60 * time.Second,
+		Transport: &imageDownloadTransport{direct: direct, proxied: proxied, proxy: proxy, resolvesPrivate: isPrivateOrLoopbackHost},
+	}
+}
+
+// newImageDownloadBaseTransport 返回与 http.DefaultTransport 参数一致、但不带代理的 Transport。
+// control 非空时在每个 socket 连接前校验目标地址。
+func newImageDownloadBaseTransport(control func(network, address string, c syscall.RawConn) error) *http.Transport {
+	dialer := &net.Dialer{Timeout: 30 * time.Second, KeepAlive: 30 * time.Second, Control: control}
+	return &http.Transport{
+		DialContext:           dialer.DialContext,
+		ForceAttemptHTTP2:     true,
+		MaxIdleConns:          100,
+		IdleConnTimeout:       90 * time.Second,
+		TLSHandshakeTimeout:   10 * time.Second,
+		ExpectContinueTimeout: 1 * time.Second,
+	}
+}
+
+// imageDownloadDialControl 拒绝连接到非公网 IP。address 是 DNS 解析后的真实连接地址，
+// 多个解析结果时 Dialer 会跳过被拒地址继续尝试下一个。
+func imageDownloadDialControl(_, address string, _ syscall.RawConn) error {
+	host, _, err := net.SplitHostPort(address)
+	if err != nil {
+		return err
+	}
+	if ip := net.ParseIP(stripIPv6Zone(host)); ip == nil || isPrivateIP(ip) {
+		return &net.AddrError{Err: "address not allowed for image download", Addr: address}
+	}
+	return nil
+}
+
+// imageDownloadTransport 按请求是否走环境代理分派到对应 Transport，并在发出前拒绝非公网目的地。
+type imageDownloadTransport struct {
+	direct  http.RoundTripper
+	proxied http.RoundTripper
+	proxy   func(*http.Request) (*url.URL, error)
+	// resolvesPrivate 报告 host 的本地解析结果是否含非公网地址，仅用于经代理的请求。
+	resolvesPrivate func(ctx context.Context, host string) (bool, error)
+}
+
+func (t *imageDownloadTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	next, err := t.route(req)
+	if err != nil {
+		if req.Body != nil {
+			_ = req.Body.Close()
+		}
+		return nil, err
+	}
+	return next.RoundTrip(req)
+}
+
+func (t *imageDownloadTransport) route(req *http.Request) (http.RoundTripper, error) {
+	if req.URL == nil {
+		return nil, errors.New("image url is empty")
+	}
+	host := req.URL.Hostname()
+	if isBlockedImageDownloadHost(host) {
+		return nil, fmt.Errorf("image url host is not allowed: %s", host)
+	}
+	var proxyURL *url.URL
+	if t.proxy != nil {
+		var err error
+		if proxyURL, err = t.proxy(req); err != nil {
+			return nil, err
+		}
+	}
+	if proxyURL == nil {
+		return t.direct, nil
+	}
+	if t.resolvesPrivate == nil {
+		return nil, errors.New("image download proxy host check is not configured")
+	}
+	blocked, err := t.resolvesPrivate(req.Context(), host)
+	if err != nil {
+		return nil, fmt.Errorf("resolve image url host: %w", err)
+	}
+	if blocked {
+		return nil, fmt.Errorf("image url host is not allowed: %s", host)
+	}
+	return t.proxied, nil
+}
+
+// isBlockedImageDownloadHost 报告 host 是否为空、被拒绝的主机名，或非公网 IP 字面量。
+// 域名的解析结果由 socket 层或代理分支再校验。
+func isBlockedImageDownloadHost(host string) bool {
+	host = strings.TrimSuffix(strings.ToLower(strings.TrimSpace(host)), ".")
+	if urlvalidator.IsBlockedHost(host) || isBlockedHostname(host) {
+		return true
+	}
+	if ip := net.ParseIP(stripIPv6Zone(host)); ip != nil {
+		return isPrivateIP(ip)
+	}
+	return false
+}
+
+// stripIPv6Zone 去掉 IPv6 字面量的 zone（如 fe80::1%eth0），使其可被 net.ParseIP 解析。
+func stripIPv6Zone(host string) string {
+	if i := strings.LastIndexByte(host, '%'); i >= 0 {
+		return host[:i]
+	}
+	return host
 }
 
 // Rewrite 将 result（上游生图响应 JSON）里的每张图片转存到对象存储，
@@ -215,9 +336,12 @@ func (u *ImageResultUploader) download(ctx context.Context, rawURL string) ([]by
 	if int64(len(data)) > limit {
 		return nil, "", fmt.Errorf("downloaded image exceeds %d bytes", limit)
 	}
-	contentType := strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0])
+	contentType := strings.ToLower(strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0]))
 	if !strings.HasPrefix(contentType, "image/") {
-		contentType = detectImageContentType(data)
+		// 声明与嗅探都不是图片时拒绝，避免把任意响应体（内部接口的 JSON/HTML 等）转存给租户。
+		if contentType = detectedImageContentType(data); contentType == "" {
+			return nil, "", errors.New("download image: content is not an image")
+		}
 	}
 	return data, contentType, nil
 }

backend/internal/service/image_storage_test.go

@@ -1,12 +1,17 @@
 package service
 
 import (
+	"bytes"
 	"context"
 	"encoding/base64"
 	"encoding/json"
 	"errors"
+	"io"
+	"net"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -73,18 +78,24 @@ func TestImageResultUploaderRewritesB64JSON(t *testing.T) {
 }
 
 func TestImageResultUploaderRewritesURL(t *testing.T) {
-	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
-		w.Header().Set("Content-Type", "image/png")
-		_, _ = w.Write(pngBytes)
-	}))
-	defer upstream.Close()
+	var requested string
+	client := &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
+		requested = req.URL.String()
+		return &http.Response{
+			StatusCode: http.StatusOK,
+			Header:     http.Header{"Content-Type": []string{"image/png"}},
+			Body:       io.NopCloser(bytes.NewReader(pngBytes)),
+			Request:    req,
+		}, nil
+	})}
 
 	storage := &fakeImageStorage{}
-	uploader := NewImageResultUploader(storage, "images/", 0, nil)
+	uploader := NewImageResultUploader(storage, "images/", 0, client)
 
-	result := json.RawMessage(`{"created":1,"data":[{"url":"` + upstream.URL + `/pic.png"}]}`)
+	result := json.RawMessage(`{"created":1,"data":[{"url":"https://images.example.test/pic.png"}]}`)
 	out, err := uploader.Rewrite(context.Background(), "imgtask_xyz", result)
 	require.NoError(t, err)
+	require.Equal(t, "https://images.example.test/pic.png", requested)
 
 	require.Len(t, storage.saved, 1)
 	require.Equal(t, pngBytes, storage.saved[0].data)
@@ -244,3 +255,243 @@ func TestImageTaskServiceCompleteOffloadFailureMarksFailed(t *testing.T) {
 	require.Contains(t, string(got.Error), "object storage")
 	require.NotContains(t, string(got.Result), "b64_json", "failed offload must not persist base64 to Redis")
 }
+
+func TestImageResultUploaderDefaultClientRejectsPrivateURL(t *testing.T) {
+	var hits atomic.Int32
+	internal := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		hits.Add(1)
+		_, _ = w.Write([]byte("GIF89a-payload"))
+	}))
+	defer internal.Close()
+
+	storage := &fakeImageStorage{}
+	uploader := NewImageResultUploader(storage, "images/", 0, nil)
+	result := json.RawMessage(`{"data":[{"url":"` + internal.URL + `/a.gif"}]}`)
+
+	_, err := uploader.Rewrite(context.Background(), "imgtask_private", result)
+	require.ErrorContains(t, err, "image url host is not allowed")
+	require.Zero(t, hits.Load(), "internal service must not be contacted")
+	require.Empty(t, storage.saved)
+}
+
+func TestImageTaskServiceCompletePrivateURLMarksFailed(t *testing.T) {
+	var hits atomic.Int32
+	internal := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		hits.Add(1)
+		_, _ = w.Write(pngBytes)
+	}))
+	defer internal.Close()
+
+	store := &imageTaskMemoryStore{}
+	storage := &fakeImageStorage{}
+	svc := NewImageTaskServiceWithUploader(store, NewImageResultUploader(storage, "images/", 0, nil), time.Hour, time.Minute)
+	owner := ImageTaskOwner{UserID: 1, APIKeyID: 2}
+	created, err := svc.Create(context.Background(), owner)
+	require.NoError(t, err)
+
+	result := json.RawMessage(`{"data":[{"url":"` + internal.URL + `/a.gif"}]}`)
+	require.NoError(t, svc.Complete(context.Background(), created.ID, http.StatusOK, result))
+
+	got, err := svc.Get(context.Background(), owner, created.ID)
+	require.NoError(t, err)
+	require.Equal(t, ImageTaskStatusFailed, got.Status)
+	require.Empty(t, got.ImageURL)
+	require.NotContains(t, string(got.Result), internal.URL)
+	require.Zero(t, hits.Load())
+	require.Empty(t, storage.saved)
+}
+
+func newCountingImageDownloadTransport(proxyURL *url.URL) (*imageDownloadTransport, *atomic.Int32, *atomic.Int32) {
+	var directCalls, proxiedCalls atomic.Int32
+	respond := func(counter *atomic.Int32) http.RoundTripper {
+		return roundTripFunc(func(req *http.Request) (*http.Response, error) {
+			counter.Add(1)
+			return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(bytes.NewReader(pngBytes)), Request: req}, nil
+		})
+	}
+	return &imageDownloadTransport{
+		direct:  respond(&directCalls),
+		proxied: respond(&proxiedCalls),
+		proxy:   func(*http.Request) (*url.URL, error) { return proxyURL, nil },
+		resolvesPrivate: func(ctx context.Context, host string) (bool, error) {
+			if host == "intranet.example.test" {
+				return true, nil
+			}
+			if host == "unresolvable.example.test" {
+				return false, errors.New("no such host")
+			}
+			return isPrivateOrLoopbackHost(ctx, host)
+		},
+	}, &directCalls, &proxiedCalls
+}
+
+func TestImageDownloadTransportRejectsNonPublicHosts(t *testing.T) {
+	hosts := []string{
+		"http://127.0.0.1:18082/a.gif",
+		"http://localhost/a.png",
+		"http://LOCALHOST./a.png",
+		"http://foo.localhost/a.png",
+		"http://10.0.0.8/a.png",
+		"http://172.16.3.4/a.png",
+		"http://192.168.1.1/a.png",
+		"http://100.64.0.1/a.png",
+		"http://169.254.169.254/latest/meta-data/",
+		"http://metadata.google.internal/computeMetadata/v1/",
+		"http://0.0.0.0/a.png",
+		"http://[::1]/a.png",
+		"http://[::ffff:127.0.0.1]/a.png",
+		"http://[fd00::1]/a.png",
+		"http://[fe80::1%25eth0]/a.png",
+		"http:///a.png",
+	}
+	for _, proxied := range []bool{false, true} {
+		for _, rawURL := range hosts {
+			var proxyURL *url.URL
+			if proxied {
+				proxyURL = &url.URL{Scheme: "http", Host: "127.0.0.1:7890"}
+			}
+			transport, directCalls, proxiedCalls := newCountingImageDownloadTransport(proxyURL)
+			req, err := http.NewRequest(http.MethodGet, rawURL, nil)
+			require.NoError(t, err)
+
+			_, err = transport.RoundTrip(req)
+			require.Error(t, err, "proxied=%v url=%s", proxied, rawURL)
+			require.Zero(t, directCalls.Load(), rawURL)
+			require.Zero(t, proxiedCalls.Load(), rawURL)
+		}
+	}
+}
+
+func TestImageDownloadTransportRoutesPublicHosts(t *testing.T) {
+	transport, directCalls, proxiedCalls := newCountingImageDownloadTransport(nil)
+	req, err := http.NewRequest(http.MethodGet, "https://203.0.113.10/a.png", nil)
+	require.NoError(t, err)
+	resp, err := transport.RoundTrip(req)
+	require.NoError(t, err)
+	_ = resp.Body.Close()
+	require.EqualValues(t, 1, directCalls.Load())
+	require.Zero(t, proxiedCalls.Load())
+
+	transport, directCalls, proxiedCalls = newCountingImageDownloadTransport(&url.URL{Scheme: "http", Host: "127.0.0.1:7890"})
+	req, err = http.NewRequest(http.MethodGet, "https://203.0.113.10/a.png", nil)
+	require.NoError(t, err)
+	resp, err = transport.RoundTrip(req)
+	require.NoError(t, err)
+	_ = resp.Body.Close()
+	require.Zero(t, directCalls.Load())
+	require.EqualValues(t, 1, proxiedCalls.Load(), "a loopback env proxy must still be usable for public targets")
+}
+
+func TestImageDownloadClientRejectsRedirectToPrivateHost(t *testing.T) {
+	var calls atomic.Int32
+	transport := &imageDownloadTransport{
+		direct: roundTripFunc(func(req *http.Request) (*http.Response, error) {
+			calls.Add(1)
+			return &http.Response{
+				StatusCode: http.StatusFound,
+				Header:     http.Header{"Location": []string{"http://169.254.169.254/latest/meta-data/"}},
+				Body:       io.NopCloser(bytes.NewReader(nil)),
+				Request:    req,
+			}, nil
+		}),
+	}
+	uploader := NewImageResultUploader(&fakeImageStorage{}, "images/", 0, &http.Client{Transport: transport})
+
+	_, err := uploader.Rewrite(context.Background(), "imgtask_redirect", json.RawMessage(`{"data":[{"url":"https://images.example.test/a.png"}]}`))
+	require.ErrorContains(t, err, "image url host is not allowed: 169.254.169.254")
+	require.EqualValues(t, 1, calls.Load(), "the private redirect hop must not be sent")
+}
+
+func TestImageDownloadClientDialsOnlyPublicAddresses(t *testing.T) {
+	var hits atomic.Int32
+	internal := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		hits.Add(1)
+		_, _ = w.Write(pngBytes)
+	}))
+	defer internal.Close()
+
+	client := newImageDownloadHTTPClient(nil)
+	direct, ok := client.Transport.(*imageDownloadTransport).direct.(*http.Transport)
+	require.True(t, ok)
+	require.Nil(t, direct.Proxy, "the direct transport must not pick up an environment proxy")
+
+	// 绕过 RoundTrip 预检直接 dial：IP 字面量与经 DNS 解析到回环的域名都必须在 socket 层被拒。
+	conn, err := direct.DialContext(context.Background(), "tcp", internal.Listener.Addr().String())
+	if conn != nil {
+		_ = conn.Close()
+	}
+	require.ErrorContains(t, err, "address not allowed for image download")
+	_, port, err := net.SplitHostPort(internal.Listener.Addr().String())
+	require.NoError(t, err)
+	conn, err = direct.DialContext(context.Background(), "tcp", net.JoinHostPort("localhost", port))
+	if conn != nil {
+		_ = conn.Close()
+	}
+	require.ErrorContains(t, err, "address not allowed for image download")
+	require.Zero(t, hits.Load())
+}
+
+func newStaticImageDownloadClient(contentType string, body []byte) *http.Client {
+	return &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
+		header := http.Header{}
+		if contentType != "" {
+			header.Set("Content-Type", contentType)
+		}
+		return &http.Response{StatusCode: http.StatusOK, Header: header, Body: io.NopCloser(bytes.NewReader(body)), Request: req}, nil
+	})}
+}
+
+func TestImageResultUploaderDownloadContentTypes(t *testing.T) {
+	avif := []byte("\x00\x00\x00\x1cftypavif-not-sniffable")
+	tests := []struct {
+		name        string
+		contentType string
+		body        []byte
+		wantType    string
+		wantErr     string
+	}{
+		{name: "declared png", contentType: "image/png", body: pngBytes, wantType: "image/png"},
+		{name: "octet-stream sniffed png", contentType: "application/octet-stream", body: pngBytes, wantType: "image/png"},
+		{name: "missing header sniffed png", body: pngBytes, wantType: "image/png"},
+		{name: "declared avif not sniffable", contentType: "Image/AVIF; charset=binary", body: avif, wantType: "image/avif"},
+		{name: "json body", contentType: "application/json", body: []byte(`{"ok":true}`), wantErr: "content is not an image"},
+		{name: "html body", contentType: "text/html", body: []byte("<html>admin</html>"), wantErr: "content is not an image"},
+		{name: "empty body", body: nil, wantErr: "content is not an image"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			storage := &fakeImageStorage{}
+			uploader := NewImageResultUploader(storage, "images/", 0, newStaticImageDownloadClient(tt.contentType, tt.body))
+			_, err := uploader.Rewrite(context.Background(), "imgtask_ct", json.RawMessage(`{"data":[{"url":"https://images.example.test/a"}]}`))
+			if tt.wantErr != "" {
+				require.ErrorContains(t, err, tt.wantErr)
+				require.Empty(t, storage.saved)
+				return
+			}
+			require.NoError(t, err)
+			require.Len(t, storage.saved, 1)
+			require.Equal(t, tt.wantType, storage.saved[0].contentType)
+		})
+	}
+}
+
+func TestImageDownloadDialControl(t *testing.T) {
+	for _, address := range []string{"127.0.0.1:443", "10.1.2.3:80", "100.100.100.200:80", "169.254.169.254:80", "[::1]:443", "[fe80::1%eth0]:443", "[::ffff:192.168.0.1]:443", "bad"} {
+		require.Error(t, imageDownloadDialControl("tcp", address, nil), address)
+	}
+	for _, address := range []string{"203.0.113.10:443", "[2001:db8::1]:443"} {
+		require.NoError(t, imageDownloadDialControl("tcp", address, nil), address)
+	}
+}
+
+func TestImageDownloadTransportProxiedChecksResolvedHost(t *testing.T) {
+	for _, host := range []string{"intranet.example.test", "unresolvable.example.test"} {
+		transport, directCalls, proxiedCalls := newCountingImageDownloadTransport(&url.URL{Scheme: "http", Host: "127.0.0.1:7890"})
+		req, err := http.NewRequest(http.MethodGet, "https://"+host+"/a.png", nil)
+		require.NoError(t, err)
+		_, err = transport.RoundTrip(req)
+		require.Error(t, err, host)
+		require.Zero(t, directCalls.Load(), host)
+		require.Zero(t, proxiedCalls.Load(), "a proxied request must fail closed when the target is private or unresolvable: %s", host)
+	}
+}