lisone/ds2api
1
0
Fork 0
mirror of https://github.com/CJackHwang/ds2api.git synced 2026-05-05 17:05:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

Preserve SSE frame delimiters when injecting Gemini usage

Browse Source
This commit is contained in:
CJACK.
2026-04-07 10:59:27 +08:00
parent 668b9c26bd
commit 86ecbc89bd
2 changed files with 23 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
internal/translatorcliproxy/stream_writer.go
View File

@@ -162,6 +162,13 @@ func injectStreamUsageMetadata(chunk []byte, target sdktranslator.Format, usage
if target != sdktranslator.FormatGemini {
return chunk
}
suffix := ""
switch {
case bytes.HasSuffix(chunk, []byte("\n\n")):
suffix = "\n\n"
case bytes.HasSuffix(chunk, []byte("\n")):
suffix = "\n"
}
text := strings.TrimSpace(string(chunk))
if text == "" {
return chunk
@@ -194,7 +201,10 @@ func injectStreamUsageMetadata(chunk []byte, target sdktranslator.Format, usage
return chunk
}
if hasDataPrefix {
return []byte("data: " + string(b))
return []byte("data: " + string(b) + suffix)
}
if suffix != "" {
return append(b, []byte(suffix)...)
}
return b
}

12
internal/translatorcliproxy/stream_writer_test.go
View File

@@ -63,3 +63,15 @@ func TestOpenAIStreamTranslatorWriterPreservesKeepAliveComment(t *testing.T) {
t.Fatalf("expected keep-alive comment passthrough, got %q", body)
}
}
func TestInjectStreamUsageMetadataPreservesSSEFrameTerminator(t *testing.T) {
chunk := []byte("data: {\"candidates\":[{\"index\":0}],\"model\":\"gemini-2.5-pro\"}\n\n")
usage := openAIUsage{PromptTokens: 11, CompletionTokens: 29, TotalTokens: 40}
got := injectStreamUsageMetadata(chunk, sdktranslator.FormatGemini, usage)
if !strings.HasSuffix(string(got), "\n\n") {
t.Fatalf("expected injected chunk to preserve \\n\\n frame terminator, got %q", string(got))
}
if !strings.Contains(string(got), `"usageMetadata"`) {
t.Fatalf("expected usageMetadata injected, got %q", string(got))
}
}