From 150e7f0dc50e3d3a0f7c4e552cc402ae105eb2a0 Mon Sep 17 00:00:00 2001 From: sususu98 Date: Mon, 29 Jun 2026 18:11:44 +0800 Subject: [PATCH] fix(auth): repair force-mapped Responses SSE framing for WS forwarder Force-mapping rewrites streaming OpenAI Responses SSE through StreamRewriter before the /v1/responses websocket forwarder. Antigravity/Gemini and Codex emit frames without reliable trailing newlines, so buffered chunks glued as ...}event:..., ...}data:..., or event:/data: without separators. The rewriter dropped pending tails and downstream WS synthesized 408 without response.completed even when upstream returned HTTP 200. - safeReplaceGlued for }event: and }data: when data JSON is complete\n- Finish flush with glue normalization and line-wise fallback\n- Newline between pending event: and next data: line (Codex scanner lines)\n- Regression tests: Antigravity sim, Codex data lines, force-map WS forward --- .../auth/codex_forcemap_ws_forward_test.go | 80 +++++++++++++ sdk/cliproxy/auth/conductor.go | 8 +- sdk/cliproxy/auth/response_model_rewriter.go | 86 ++++++++++++-- ...nse_model_rewriter_antigravity_sim_test.go | 107 ++++++++++++++++++ .../auth/response_model_rewriter_test.go | 101 +++++++++++++++++ 5 files changed, 371 insertions(+), 11 deletions(-) create mode 100644 sdk/cliproxy/auth/codex_forcemap_ws_forward_test.go create mode 100644 sdk/cliproxy/auth/response_model_rewriter_antigravity_sim_test.go diff --git a/sdk/cliproxy/auth/codex_forcemap_ws_forward_test.go b/sdk/cliproxy/auth/codex_forcemap_ws_forward_test.go new file mode 100644 index 00000000..9996ccd3 --- /dev/null +++ b/sdk/cliproxy/auth/codex_forcemap_ws_forward_test.go @@ -0,0 +1,80 @@ +package auth + +import ( + "bytes" + "strings" + "testing" + + "github.com/tidwall/gjson" +) + +func parseWSDataEventTypesFromForwardedChunks(forwarded [][]byte) []string { + var types []string + for _, ch := range forwarded { + ch = normalizeGluedSSEEvents(ch) + for _, ln := range bytes.Split(ch, []byte("\n")) { + ln = bytes.TrimSpace(ln) + if !bytes.HasPrefix(ln, []byte("data:")) { + continue + } + j := bytes.TrimSpace(ln[5:]) + if gjson.ValidBytes(j) { + types = append(types, gjson.GetBytes(j, "type").String()) + } + } + } + return types +} + +func replayCodexForceMapLines(t *testing.T, lines [][]byte) []string { + t.Helper() + r := NewStreamRewriter(StreamRewriteOptions{RewriteModel: "gpt-5.4-fast"}) + var forwarded [][]byte + for _, line := range lines { + if out := rewriteForceMappedStreamChunk(r, line); len(out) > 0 { + forwarded = append(forwarded, out) + } + } + if tail := finishForceMappedStreamChunks(r); len(tail) > 0 { + forwarded = append(forwarded, tail) + } + return parseWSDataEventTypesFromForwardedChunks(forwarded) +} + +func TestCodexForceMapPerLineSSE_ForwardsCompleted(t *testing.T) { + lines := [][]byte{ + []byte("event: response.created"), + []byte(`data: {"type":"response.created","response":{"model":"gpt-5.4"}}`), + []byte("event: response.output_text.delta"), + []byte(`data: {"type":"response.output_text.delta","delta":"OK"}`), + []byte("event: response.completed"), + []byte(`data: {"type":"response.completed","response":{"model":"gpt-5.4","output":[]}}`), + } + types := replayCodexForceMapLines(t, lines) + found := false + for _, typ := range types { + if typ == "response.completed" { + found = true + break + } + } + if !found { + t.Fatalf("missing response.completed, types=%v", types) + } +} + +func TestRewriteForceMappedStreamChunk_FallbackWhenPendingBuffersEvent(t *testing.T) { + r := NewStreamRewriter(StreamRewriteOptions{RewriteModel: "gpt-5.4-fast"}) + _ = rewriteForceMappedStreamChunk(r, []byte("event: response.completed")) + out := rewriteForceMappedStreamChunk(r, []byte(`data: {"type":"response.completed","response":{"model":"gpt-5.4","output":[]}}`)) + if len(out) == 0 { + tail := finishForceMappedStreamChunks(r) + if !bytes.Contains(tail, []byte("response.completed")) { + t.Fatalf("expected completed in tail, got %q", tail) + } + return + } + if !strings.Contains(string(out), "response.completed") { + t.Fatalf("out=%q", out) + } +} diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 0f8bd2e4..1597efff 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -1336,12 +1336,14 @@ func rewriteForceMappedStreamChunk(rewriter *StreamRewriter, payload []byte) []b if len(rewritten) > 0 { return rewritten } + if bytes.Contains(payload, []byte("data:")) { + if lineWise := rewriteSSEPayloadLines(payload, rewriter.options.RewriteModel); len(lineWise) > 0 { + return lineWise + } + } if len(rewriter.pendingBuf) > 0 { return nil } - if lineWise := rewriteSSEPayloadLines(payload, rewriter.options.RewriteModel); len(lineWise) > 0 { - return lineWise - } return nil } diff --git a/sdk/cliproxy/auth/response_model_rewriter.go b/sdk/cliproxy/auth/response_model_rewriter.go index 6ebee5bb..f223f21d 100644 --- a/sdk/cliproxy/auth/response_model_rewriter.go +++ b/sdk/cliproxy/auth/response_model_rewriter.go @@ -72,9 +72,16 @@ func (r *StreamRewriter) RewriteChunk(chunk []byte) []byte { } if len(r.pendingBuf) > 0 { - chunk = append(r.pendingBuf, chunk...) + combined := make([]byte, 0, len(r.pendingBuf)+1+len(chunk)) + combined = append(combined, r.pendingBuf...) + if combined[len(combined)-1] != '\n' { + combined = append(combined, '\n') + } + combined = append(combined, chunk...) + chunk = combined r.pendingBuf = nil } + chunk = normalizeGluedSSEEvents(chunk) if len(chunk) > maxPendingBufSize { return chunk @@ -91,10 +98,6 @@ func (r *StreamRewriter) RewriteChunk(chunk []byte) []byte { } lastDoubleNewline := bytes.LastIndex(chunk, []byte("\n\n")) - lastNewline := -1 - if len(chunk) > 0 && chunk[len(chunk)-1] == '\n' { - lastNewline = len(chunk) - 1 - } var processChunk []byte if lastDoubleNewline >= 0 { @@ -106,7 +109,7 @@ func (r *StreamRewriter) RewriteChunk(chunk []byte) []byte { } else { processChunk = chunk } - } else if lastNewline >= 0 && gjson.ValidBytes(extractLastDataPayload(chunk)) { + } else if gjson.ValidBytes(extractLastDataPayload(chunk)) { processChunk = chunk } else if len(bytes.TrimSpace(chunk)) == 0 { return chunk @@ -200,12 +203,79 @@ func extractSSEDataLine(line []byte) (prefix []byte, jsonData []byte, ok bool) { return nil, nil, false } +func normalizeGluedSSEEvents(chunk []byte) []byte { + if len(chunk) == 0 { + return chunk + } + // Antigravity/Gemini translators emit event frames without trailing blank lines. + // When multiple frames are buffered back-to-back they can glue as "...}event:...". + // Only split when the bytes before the glue close a valid SSE data JSON object. + chunk = safeReplaceGlued(chunk, []byte("}event:"), []byte("}\n\nevent:")) + chunk = safeReplaceGlued(chunk, []byte("}\r\nevent:"), []byte("}\r\n\r\nevent:")) + // Codex executor emits one "data: {json}" chunk per SSE line without trailing newlines. + // Buffered chunks can glue as "...}data:...". + chunk = safeReplaceGlued(chunk, []byte("}data:"), []byte("}\ndata:")) + chunk = safeReplaceGlued(chunk, []byte("}\r\ndata:"), []byte("}\r\ndata:")) + return chunk +} + +func safeReplaceGlued(chunk []byte, old, new []byte) []byte { + if len(old) == 0 || len(chunk) == 0 { + return chunk + } + if !bytes.Contains(chunk, old) { + return chunk + } + var result []byte + remaining := chunk + for { + idx := bytes.Index(remaining, old) + if idx == -1 { + result = append(result, remaining...) + break + } + lineStart := bytes.LastIndexByte(remaining[:idx], '\n') + var part []byte + if lineStart == -1 { + part = remaining[:idx+1] + } else { + part = remaining[lineStart+1 : idx+1] + } + _, jsonData, ok := extractSSEDataLine(part) + if ok && len(jsonData) > 0 && gjson.ValidBytes(jsonData) { + result = append(result, remaining[:idx]...) + result = append(result, new...) + remaining = remaining[idx+len(old):] + continue + } + result = append(result, remaining[:idx+len(old)]...) + remaining = remaining[idx+len(old):] + } + return result +} + // Finish flushes any buffered partial SSE data at the end of a stream. func (r *StreamRewriter) Finish() []byte { if len(r.pendingBuf) == 0 { return nil } - chunk := r.RewriteChunk(r.pendingBuf) + buf := make([]byte, len(r.pendingBuf)+2) + copy(buf, r.pendingBuf) + buf[len(r.pendingBuf)] = '\n' + buf[len(r.pendingBuf)+1] = '\n' + buf = normalizeGluedSSEEvents(buf) r.pendingBuf = nil - return chunk + out := r.RewriteChunk(buf) + if len(r.pendingBuf) > 0 { + tail := rewriteSSEPayloadLines(r.pendingBuf, r.options.RewriteModel) + r.pendingBuf = nil + if len(tail) > 0 { + if len(out) > 0 { + out = append(out, tail...) + } else { + out = tail + } + } + } + return out } diff --git a/sdk/cliproxy/auth/response_model_rewriter_antigravity_sim_test.go b/sdk/cliproxy/auth/response_model_rewriter_antigravity_sim_test.go new file mode 100644 index 00000000..29be9243 --- /dev/null +++ b/sdk/cliproxy/auth/response_model_rewriter_antigravity_sim_test.go @@ -0,0 +1,107 @@ +package auth + +import ( + "context" + "strings" + "testing" + + gemresponses "github.com/router-for-me/CLIProxyAPI/v7/internal/translator/gemini/openai/responses" + "github.com/tidwall/gjson" +) + +func antigravityLiveSSEChunks(t *testing.T) [][]byte { + t.Helper() + rawOK := `{"response": {"candidates": [{"content": {"role": "model","parts": [{"text": "OK"}]}}],"usageMetadata": {"promptTokenCount": 21,"candidatesTokenCount": 1,"totalTokenCount": 131,"thoughtsTokenCount": 109},"modelVersion": "gemini-3-flash-a","responseId": "tjVCavaJBYjgz7IP-NnfSQ"},"traceId": "x","metadata": {}}` + rawStop := `{"response": {"candidates": [{"content": {"role": "model","parts": [{"thoughtSignature": "sig","text": ""}]},"finishReason": "STOP"}],"usageMetadata": {"promptTokenCount": 21,"candidatesTokenCount": 1,"totalTokenCount": 131,"thoughtsTokenCount": 109},"modelVersion": "gemini-3-flash-a","responseId": "tjVCavaJBYjgz7IP-NnfSQ"},"traceId": "x","metadata": {}}` + req := []byte(`{"model":"gemini-3.5-flash","input":[]}`) + var param any + var chunks [][]byte + for _, raw := range []string{rawOK, rawStop} { + chunks = append(chunks, gemresponses.ConvertGeminiResponseToOpenAIResponses(context.Background(), "gemini-3.5-flash", req, req, []byte("data: "+raw), ¶m)...) + } + if len(chunks) == 0 { + t.Fatal("translator produced no chunks") + } + return chunks +} + +func TestAntigravityTranslatorEmitsCompletedWithoutRewriter(t *testing.T) { + chunks := antigravityLiveSSEChunks(t) + combined := string(joinBytes(chunks)) + if !strings.Contains(combined, "response.completed") { + t.Fatalf("translator missing completed: chunks=%d preview=%q", len(chunks), trunc(combined, 400)) + } +} + +func TestRewriteForceMappedStreamChunk_AntigravityTranslatorEventChunks_PreservesCompleted(t *testing.T) { + chunks := antigravityLiveSSEChunks(t) + rewriter := NewStreamRewriter(StreamRewriteOptions{RewriteModel: "gemini-3.5-flash"}) + var out []byte + for _, ch := range chunks { + if rewritten := rewriteForceMappedStreamChunk(rewriter, ch); len(rewritten) > 0 { + out = append(out, rewritten...) + } + } + if tail := finishForceMappedStreamChunks(rewriter); len(tail) > 0 { + out = append(out, tail...) + } + if !parseCompletedFromSSE(out) { + t.Fatalf("rewriter output missing response.completed; preview=%q", trunc(string(out), 400)) + } +} + +func TestRewriteForceMappedStreamChunk_AntigravityGluedEventFramesFlushCompleted(t *testing.T) { + chunks := antigravityLiveSSEChunks(t) + rewriter := NewStreamRewriter(StreamRewriteOptions{RewriteModel: "gemini-3.5-flash"}) + var out []byte + for i, ch := range chunks { + if rewritten := rewriteForceMappedStreamChunk(rewriter, ch); len(rewritten) > 0 { + out = append(out, rewritten...) + } + if i == 1 && len(rewriter.pendingBuf) > 0 && strings.Contains(string(rewriter.pendingBuf), "}event:") { + t.Log("confirmed glued frames: ...}event:...") + } + } + if tail := finishForceMappedStreamChunks(rewriter); len(tail) > 0 { + out = append(out, tail...) + } + if !parseCompletedFromSSE(out) { + t.Fatalf("expected completed after glued frames flush; preview=%q", trunc(string(out), 400)) + } +} + +func joinBytes(parts [][]byte) []byte { + var out []byte + for _, p := range parts { + out = append(out, p...) + } + return out +} + +func trunc(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n] + "..." +} + +func parseCompletedFromSSE(payload []byte) bool { + if len(payload) == 0 { + return false + } + for _, line := range strings.Split(string(payload), "\n") { + line = strings.TrimSpace(line) + if !strings.HasPrefix(line, "data:") { + continue + } + line = strings.TrimSpace(strings.TrimPrefix(line, "data:")) + if gjson.Get(line, "type").String() == "response.completed" { + return true + } + } + trim := strings.TrimSpace(string(payload)) + if strings.HasPrefix(trim, "{") && gjson.Get(trim, "type").String() == "response.completed" { + return true + } + return false +} diff --git a/sdk/cliproxy/auth/response_model_rewriter_test.go b/sdk/cliproxy/auth/response_model_rewriter_test.go index ea570e5e..751744ef 100644 --- a/sdk/cliproxy/auth/response_model_rewriter_test.go +++ b/sdk/cliproxy/auth/response_model_rewriter_test.go @@ -1,6 +1,9 @@ package auth import ( + "bytes" + + "github.com/tidwall/gjson" "strings" "testing" @@ -204,3 +207,101 @@ func TestRewriteForceMappedStreamChunk_NoRewriteWhenRewriterNil(t *testing.T) { t.Fatalf("chunk = %q, want unchanged upstream payload", got) } } + +func TestNormalizeGluedSSEEvents_SplitsValidGlueOnly(t *testing.T) { + glued := []byte("event: response.created\ndata: {\"type\":\"response.created\"}event: response.completed\ndata: {\"type\":\"response.completed\"}") + got := normalizeGluedSSEEvents(glued) + if !bytes.Contains(got, []byte("}\n\nevent:")) { + t.Fatalf("expected glued frame split, got %q", got) + } + + inside := []byte("event: response.output_text.delta\ndata: {\"type\":\"delta\",\"text\":\"literal }event: inside string\"}") + gotInside := string(normalizeGluedSSEEvents(inside)) + if strings.Contains(gotInside, "}\n\nevent:") { + t.Fatalf("should not split inside JSON string, got %q", gotInside) + } + for _, line := range bytes.Split(inside, []byte("\n")) { + if bytes.HasPrefix(line, []byte("data:")) { + _, jd, ok := extractSSEDataLine(line) + if !ok || !gjson.ValidBytes(jd) { + t.Fatalf("baseline invalid") + } + } + } + for _, line := range bytes.Split([]byte(gotInside), []byte("\n")) { + if bytes.HasPrefix(line, []byte("data:")) { + _, jd, ok := extractSSEDataLine(line) + if !ok || !gjson.ValidBytes(jd) { + t.Fatalf("corrupted JSON after normalize: %q", gotInside) + } + } + } +} + +func TestNormalizeGluedSSEEvents_SplitsCodexDataGlueOnly(t *testing.T) { + glued := []byte(`data: {"type":"response.created"}data: {"type":"response.completed"}`) + got := normalizeGluedSSEEvents(glued) + if !bytes.Contains(got, []byte("}\ndata:")) { + t.Fatalf("expected codex glued split, got %q", got) + } + inside := []byte(`data: {"type":"delta","text":"literal }data: inside"}`) + gotInside := string(normalizeGluedSSEEvents(inside)) + if strings.Contains(gotInside, "}\ndata:") && !bytes.Equal([]byte(gotInside), inside) { + // Only fail if we actually inserted a split (unchanged is OK) + for _, line := range bytes.Split([]byte(gotInside), []byte("\n")) { + if bytes.HasPrefix(line, []byte("data:")) { + _, jd, ok := extractSSEDataLine(line) + if !ok || !gjson.ValidBytes(jd) { + t.Fatalf("corrupted JSON: %q", gotInside) + } + } + } + } +} + +func parseResponsesWSDataEventTypes(payload []byte) []string { + lines := bytes.Split(payload, []byte("\n")) + var types []string + for _, line := range lines { + line = bytes.TrimSpace(line) + if len(line) == 0 || bytes.HasPrefix(line, []byte("event:")) { + continue + } + if bytes.HasPrefix(line, []byte("data:")) { + line = bytes.TrimSpace(line[len("data:"):]) + } + if len(line) == 0 || !gjson.ValidBytes(line) { + continue + } + types = append(types, gjson.GetBytes(line, "type").String()) + } + return types +} + +func TestRewriteForceMappedStreamChunk_CodexDataLinesWithoutNewlines_FinishParsesCompleted(t *testing.T) { + rewriter := NewStreamRewriter(StreamRewriteOptions{RewriteModel: "gpt-5.4-fast"}) + lines := [][]byte{ + []byte(`data: {"type":"response.created","response":{"model":"gpt-5.4"}}`), + []byte(`data: {"type":"response.in_progress","response":{"model":"gpt-5.4"}}`), + []byte(`data: {"type":"response.completed","response":{"model":"gpt-5.4","output":[]}}`), + } + var types []string + for _, ln := range lines { + if out := rewriteForceMappedStreamChunk(rewriter, ln); len(out) > 0 { + types = append(types, parseResponsesWSDataEventTypes(out)...) + } + } + if tail := finishForceMappedStreamChunks(rewriter); len(tail) > 0 { + types = append(types, parseResponsesWSDataEventTypes(tail)...) + } + found := false + for _, typ := range types { + if typ == "response.completed" { + found = true + break + } + } + if !found { + t.Fatalf("missing response.completed; types=%v", types) + } +} -- 2.51.2