fix(aistudio): strip usage metadata from non-final stream chunks

This commit is contained in:
hkfires
2025-10-25 16:53:49 +08:00
parent 8aaed4cf09
commit ea6065f1b1

View File

@@ -150,14 +150,16 @@ func (e *AistudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth
case wsrelay.MessageTypeStreamChunk: case wsrelay.MessageTypeStreamChunk:
if len(event.Payload) > 0 { if len(event.Payload) > 0 {
appendAPIResponseChunk(ctx, e.cfg, bytes.Clone(event.Payload)) appendAPIResponseChunk(ctx, e.cfg, bytes.Clone(event.Payload))
if detail, ok := parseGeminiStreamUsage(event.Payload); ok { filtered := filterAistudioUsageMetadata(event.Payload)
if detail, ok := parseGeminiStreamUsage(filtered); ok {
reporter.publish(ctx, detail) reporter.publish(ctx, detail)
} }
} lines := sdktranslator.TranslateStream(ctx, body.toFormat, opts.SourceFormat, req.Model, bytes.Clone(opts.OriginalRequest), translatedReq, bytes.Clone(filtered), &param)
lines := sdktranslator.TranslateStream(ctx, body.toFormat, opts.SourceFormat, req.Model, bytes.Clone(opts.OriginalRequest), translatedReq, bytes.Clone(event.Payload), &param)
for i := range lines { for i := range lines {
out <- cliproxyexecutor.StreamChunk{Payload: []byte(lines[i])} out <- cliproxyexecutor.StreamChunk{Payload: []byte(lines[i])}
} }
break
}
case wsrelay.MessageTypeStreamEnd: case wsrelay.MessageTypeStreamEnd:
return return
case wsrelay.MessageTypeHTTPResp: case wsrelay.MessageTypeHTTPResp:
@@ -281,3 +283,62 @@ func (e *AistudioExecutor) buildEndpoint(model, action, alt string) string {
} }
return base return base
} }
// filterAistudioUsageMetadata removes usageMetadata from intermediate SSE events so that
// only the terminal chunk retains token statistics.
func filterAistudioUsageMetadata(payload []byte) []byte {
if len(payload) == 0 {
return payload
}
lines := bytes.Split(payload, []byte("\n"))
modified := false
for idx, line := range lines {
trimmed := bytes.TrimSpace(line)
if len(trimmed) == 0 || !bytes.HasPrefix(trimmed, []byte("data:")) {
continue
}
dataIdx := bytes.Index(line, []byte("data:"))
if dataIdx < 0 {
continue
}
rawJSON := bytes.TrimSpace(line[dataIdx+5:])
cleaned, changed := stripUsageMetadataFromJSON(rawJSON)
if !changed {
continue
}
var rebuilt []byte
rebuilt = append(rebuilt, line[:dataIdx]...)
rebuilt = append(rebuilt, []byte("data:")...)
if len(cleaned) > 0 {
rebuilt = append(rebuilt, ' ')
rebuilt = append(rebuilt, cleaned...)
}
lines[idx] = rebuilt
modified = true
}
if !modified {
return payload
}
return bytes.Join(lines, []byte("\n"))
}
// stripUsageMetadataFromJSON drops usageMetadata when no finishReason is present.
func stripUsageMetadataFromJSON(rawJSON []byte) ([]byte, bool) {
jsonBytes := bytes.TrimSpace(rawJSON)
if len(jsonBytes) == 0 || !gjson.ValidBytes(jsonBytes) {
return rawJSON, false
}
finishReason := gjson.GetBytes(jsonBytes, "candidates.0.finishReason")
if finishReason.Exists() && finishReason.String() != "" {
return rawJSON, false
}
if !gjson.GetBytes(jsonBytes, "usageMetadata").Exists() {
return rawJSON, false
}
cleaned, err := sjson.DeleteBytes(jsonBytes, "usageMetadata")
if err != nil {
return rawJSON, false
}
return cleaned, true
}