修复代理与日志链路的可靠性问题

proxy/sse.go

@@ -552,11 +552,16 @@ func (g *geminiChunk) UnmarshalJSON(data []byte) error {
 }
 
 type geminiCandidate struct {
-	raw     map[string]json.RawMessage
-	index   int
-	role    string
-	text    strings.Builder
-	partRaw map[string]json.RawMessage
+	raw   map[string]json.RawMessage
+	index int
+	role  string
+	parts []*geminiPart
+}
+
+type geminiPart struct {
+	raw  map[string]json.RawMessage
+	text strings.Builder
+	kind string
 }
 
 func assembleGeminiSSE(payloads []string) ([]byte, bool) {
@@ -588,15 +593,31 @@ func assembleGeminiSSE(payloads []string) ([]byte, bool) {
 				assembled.role = candidate.Content.Role
 			}
 			if len(candidate.Content.Parts) > 0 {
-				for _, part := range candidate.Content.Parts {
-					assembled.text.WriteString(part.Text)
+				var contentRaw struct {
+					Parts []map[string]json.RawMessage `json:"parts"`
 				}
-				if len(assembled.partRaw) == 0 {
-					var contentRaw struct {
-						Parts []map[string]json.RawMessage `json:"parts"`
+				if rawContent, ok := candidate.Raw["content"]; ok {
+					_ = json.Unmarshal(rawContent, &contentRaw)
+				}
+				mergeByIndex := len(assembled.parts) == len(candidate.Content.Parts)
+				if mergeByIndex {
+					for i := range candidate.Content.Parts {
+						if i >= len(contentRaw.Parts) || assembled.parts[i].kind != geminiPartKind(contentRaw.Parts[i]) {
+							mergeByIndex = false
+							break
+						}
 					}
-					if rawContent, ok := candidate.Raw["content"]; ok && json.Unmarshal(rawContent, &contentRaw) == nil && len(contentRaw.Parts) > 0 {
-						assembled.partRaw = cloneRawMap(contentRaw.Parts[0])
+				}
+				for i, part := range candidate.Content.Parts {
+					target := i
+					if !mergeByIndex {
+						target = len(assembled.parts)
+						assembled.parts = append(assembled.parts, &geminiPart{})
+					}
+					assembled.parts[target].text.WriteString(part.Text)
+					if i < len(contentRaw.Parts) {
+						assembled.parts[target].kind = geminiPartKind(contentRaw.Parts[i])
+						assembled.parts[target].raw = mergeRawMap(assembled.parts[target].raw, contentRaw.Parts[i])
 					}
 				}
 			}
@@ -614,12 +635,15 @@ func assembleGeminiSSE(payloads []string) ([]byte, bool) {
 		candidate := candidates[index]
 		candidateRaw := cloneRawMap(candidate.raw)
 		candidateRaw["index"] = mustJSON(candidate.index)
-		contentRaw := map[string]any{"role": candidate.role, "parts": []any{map[string]any{"text": candidate.text.String()}}}
-		if len(candidate.partRaw) > 0 {
-			partRaw := cloneRawMap(candidate.partRaw)
-			partRaw["text"] = mustJSON(candidate.text.String())
-			contentRaw["parts"] = []any{rawMapToMap(partRaw)}
+		parts := make([]any, 0, len(candidate.parts))
+		for _, part := range candidate.parts {
+			partRaw := cloneRawMap(part.raw)
+			if part.text.Len() > 0 || len(partRaw) == 0 {
+				partRaw["text"] = mustJSON(part.text.String())
+			}
+			parts = append(parts, rawMapToMap(partRaw))
 		}
+		contentRaw := map[string]any{"role": candidate.role, "parts": parts}
 		candidateRaw["content"] = mustJSON(contentRaw)
 		assembled["candidates"] = append(assembled["candidates"].([]any), rawMapToMap(candidateRaw))
 	}
@@ -631,6 +655,28 @@ func assembleGeminiSSE(payloads []string) ([]byte, bool) {
 	return data, err == nil
 }
 
+func geminiPartKind(part map[string]json.RawMessage) string {
+	for _, key := range []string{"functionCall", "functionResponse", "inlineData", "fileData", "executableCode", "codeExecutionResult", "text"} {
+		if _, ok := part[key]; ok {
+			return key
+		}
+	}
+	return "unknown"
+}
+
+func mergeRawMap(dst, src map[string]json.RawMessage) map[string]json.RawMessage {
+	if dst == nil {
+		dst = make(map[string]json.RawMessage, len(src))
+	}
+	for key, value := range src {
+		if key == "text" {
+			continue
+		}
+		dst[key] = append(json.RawMessage(nil), value...)
+	}
+	return dst
+}
+
 func cloneRawMap(in map[string]json.RawMessage) map[string]json.RawMessage {
 	out := make(map[string]json.RawMessage, len(in))
 	for k, v := range in {