Skip to content

Commit 2a91635

Browse files
committed
fix(#3996): drop gateway SSE keepalive frames before the genai parser
During long generations (e.g. Gemini image output) the Docker AI Gateway emits `event: keepalive` + `data: {}` frames. google.golang.org/genai treats any `event:` line as a fatal invalid stream chunk, so a single keepalive killed the whole stream. Add an opt-in httpclient.WithSSEKeepaliveFilter that makes the shared SSE filter transport also drop payload-free keepalive frames, and enable it only on the Gemini gateway HTTP client. The shared default path is unchanged: Anthropic-style named events with meaningful data still pass through verbatim, and only the gateway surface — never direct Gemini/Vertex — gets the new behavior.
1 parent cb2dfa0 commit 2a91635

5 files changed

Lines changed: 338 additions & 15 deletions

File tree

pkg/httpclient/client.go

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,10 @@ type HTTPOptions struct {
1919
Header http.Header
2020
Query url.Values
2121

22+
// dropSSEKeepaliveEvents enables keepalive-frame dropping in the SSE
23+
// filter transport; see WithSSEKeepaliveFilter.
24+
dropSSEKeepaliveEvents bool
25+
2226
// cagentID resolves the persistent install UUID stamped as
2327
// `X-Cagent-Id` on gateway-bound requests. It defaults to
2428
// [userid.Get]; tests inject their own source via
@@ -66,7 +70,10 @@ func NewHTTPClient(ctx context.Context, opts ...Opt) *http.Client {
6670

6771
var wrapped http.RoundTripper = &userAgentTransport{
6872
httpOptions: httpOptions,
69-
rt: &sseFilterTransport{base: rt},
73+
rt: &sseFilterTransport{
74+
base: rt,
75+
dropKeepaliveEvents: httpOptions.dropSSEKeepaliveEvents,
76+
},
7077
}
7178
if httpOptions.refreshAuth != nil {
7279
// Outermost, so a replayed request goes through the whole chain again.
@@ -201,6 +208,19 @@ func WithQuery(query url.Values) Opt {
201208
}
202209
}
203210

211+
// WithSSEKeepaliveFilter makes the SSE filter transport also drop whole
212+
// `event: keepalive` frames whose data carries no payload (`data: {}` or
213+
// empty). The Docker AI Gateway emits such frames during long generations
214+
// (e.g. Gemini image output), and google.golang.org/genai's SSE parser
215+
// treats any `event:` line as a fatal invalid chunk. Only enable this for
216+
// clients whose SDK cannot tolerate `event:` lines — providers like
217+
// Anthropic rely on `event:` headers as meaningful framing.
218+
func WithSSEKeepaliveFilter() Opt {
219+
return func(o *HTTPOptions) {
220+
o.dropSSEKeepaliveEvents = true
221+
}
222+
}
223+
204224
// newTransport returns an HTTP transport with automatic gzip compression disabled and Docker Desktop PAC support.
205225
func newTransport(_ context.Context) http.RoundTripper {
206226
rt := newAllowPrivateIPsTransport()

pkg/httpclient/sse_filter.go

Lines changed: 43 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,17 @@ import (
3030
// (comment-only events, or events bearing only `event:` / `id:` headers)
3131
// never reach the SDK. Well-formed events pass through verbatim, and the
3232
// filter is a no-op on non-SSE responses.
33+
//
34+
// dropKeepaliveEvents additionally drops whole `event: keepalive` frames
35+
// whose data carries no payload (`data: {}` or empty). The Docker AI
36+
// Gateway emits such frames during long generations; the genai SDK's SSE
37+
// parser hard-fails on ANY `event:` line, so they must never reach it. The
38+
// mode is opt-in (see WithSSEKeepaliveFilter) because other providers —
39+
// Anthropic in particular — use `event:` headers as meaningful framing that
40+
// must pass through untouched.
3341
type sseFilterTransport struct {
34-
base http.RoundTripper
42+
base http.RoundTripper
43+
dropKeepaliveEvents bool
3544
}
3645

3746
func (t *sseFilterTransport) RoundTrip(req *http.Request) (*http.Response, error) {
@@ -42,7 +51,7 @@ func (t *sseFilterTransport) RoundTrip(req *http.Request) (*http.Response, error
4251
// Match the prefix so charset suffixes (e.g. "text/event-stream;
4352
// charset=utf-8") still trigger filtering.
4453
if strings.HasPrefix(strings.ToLower(res.Header.Get("Content-Type")), "text/event-stream") {
45-
res.Body = newSSEFilterReader(res.Body)
54+
res.Body = newSSEFilterReader(res.Body, t.dropKeepaliveEvents)
4655
}
4756
return res, err
4857
}
@@ -58,15 +67,19 @@ type sseFilterReader struct {
5867
out bytes.Buffer // bytes ready to hand back to the caller
5968
pending bytes.Buffer // accumulated lines for the current event
6069
hasData bool // saw at least one `data:` line in `pending`
70+
71+
dropKeepaliveEvents bool // see sseFilterTransport
72+
isKeepalive bool // current event is named `keepalive`
73+
hasMeaningfulData bool // saw a `data:` line whose payload isn't empty or `{}`
6174
}
6275

63-
func newSSEFilterReader(src io.ReadCloser) *sseFilterReader {
76+
func newSSEFilterReader(src io.ReadCloser, dropKeepaliveEvents bool) *sseFilterReader {
6477
scn := bufio.NewScanner(src)
6578
// SSE events can be large (long completion tokens, image URLs, …). Match
6679
// the buffer size used by openai-go's own SSE decoder so we don't trip
6780
// `bufio.ErrTooLong` on payloads it would happily accept.
6881
scn.Buffer(make([]byte, 0, 64*1024), bufio.MaxScanTokenSize<<9)
69-
return &sseFilterReader{src: src, scn: scn}
82+
return &sseFilterReader{src: src, scn: scn, dropKeepaliveEvents: dropKeepaliveEvents}
7083
}
7184

7285
func (r *sseFilterReader) Read(p []byte) (int, error) {
@@ -85,24 +98,47 @@ func (r *sseFilterReader) Read(p []byte) (int, error) {
8598
func (r *sseFilterReader) consumeLine(line []byte) {
8699
switch {
87100
case len(line) == 0:
88-
// Event boundary: emit the buffered event iff it had data.
89-
if r.hasData {
101+
// Event boundary: emit the buffered event iff it had data and is
102+
// not a payload-free keepalive frame in keepalive-dropping mode.
103+
if r.hasData && (!r.isKeepalive || r.hasMeaningfulData) {
90104
r.out.Write(r.pending.Bytes())
91105
r.out.WriteByte('\n')
92106
}
93107
r.pending.Reset()
94108
r.hasData = false
109+
r.isKeepalive = false
110+
r.hasMeaningfulData = false
95111
case line[0] == ':':
96112
// SSE comment — drop entirely.
97113
default:
98114
r.pending.Write(line)
99115
r.pending.WriteByte('\n')
100-
if bytes.HasPrefix(line, []byte("data:")) {
116+
if value, ok := fieldValue(line, "data"); ok {
101117
r.hasData = true
118+
if r.dropKeepaliveEvents {
119+
if payload := bytes.TrimSpace(value); len(payload) > 0 && !bytes.Equal(payload, []byte("{}")) {
120+
r.hasMeaningfulData = true
121+
}
122+
}
123+
} else if r.dropKeepaliveEvents {
124+
if value, ok := fieldValue(line, "event"); ok && string(bytes.TrimSpace(value)) == "keepalive" {
125+
r.isKeepalive = true
126+
}
102127
}
103128
}
104129
}
105130

131+
// fieldValue returns the value of an SSE line whose field name is `name`,
132+
// with the single optional leading space the SSE grammar allows already
133+
// removed.
134+
func fieldValue(line []byte, name string) ([]byte, bool) {
135+
value, ok := bytes.CutPrefix(line, []byte(name+":"))
136+
if !ok {
137+
return nil, false
138+
}
139+
return bytes.TrimPrefix(value, []byte(" ")), true
140+
}
141+
106142
func (r *sseFilterReader) Close() error {
107143
return r.src.Close()
108144
}

pkg/httpclient/sse_filter_test.go

Lines changed: 158 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,7 @@ func TestSSEFilter_LargeEvent(t *testing.T) {
153153
t.Parallel()
154154

155155
largeData := "data: " + strings.Repeat("x", 256*1024) + "\n\n"
156-
r := newSSEFilterReader(io.NopCloser(strings.NewReader(largeData)))
156+
r := newSSEFilterReader(io.NopCloser(strings.NewReader(largeData)), false)
157157

158158
output, err := io.ReadAll(r)
159159
require.NoError(t, err)
@@ -167,7 +167,7 @@ func TestSSEFilter_PartialReads(t *testing.T) {
167167
t.Parallel()
168168

169169
input := "data: test1\n\ndata: test2\n\n"
170-
r := newSSEFilterReader(io.NopCloser(strings.NewReader(input)))
170+
r := newSSEFilterReader(io.NopCloser(strings.NewReader(input)), false)
171171

172172
var output []byte
173173
buf := make([]byte, 5)
@@ -192,7 +192,7 @@ func TestSSEFilter_IncompleteEventAtEOF(t *testing.T) {
192192
t.Parallel()
193193

194194
input := "data: complete\n\ndata: incomplete"
195-
r := newSSEFilterReader(io.NopCloser(strings.NewReader(input)))
195+
r := newSSEFilterReader(io.NopCloser(strings.NewReader(input)), false)
196196

197197
output, err := io.ReadAll(r)
198198
require.NoError(t, err)
@@ -204,7 +204,7 @@ func TestSSEFilter_IncompleteEventAtEOF(t *testing.T) {
204204
func TestSSEFilter_EmptyInput(t *testing.T) {
205205
t.Parallel()
206206

207-
r := newSSEFilterReader(io.NopCloser(strings.NewReader("")))
207+
r := newSSEFilterReader(io.NopCloser(strings.NewReader("")), false)
208208

209209
output, err := io.ReadAll(r)
210210
require.NoError(t, err)
@@ -218,7 +218,7 @@ func TestSSEFilter_OnlyComments(t *testing.T) {
218218
t.Parallel()
219219

220220
input := ": comment1\n\n: comment2\n\n"
221-
r := newSSEFilterReader(io.NopCloser(strings.NewReader(input)))
221+
r := newSSEFilterReader(io.NopCloser(strings.NewReader(input)), false)
222222

223223
output, err := io.ReadAll(r)
224224
require.NoError(t, err)
@@ -230,7 +230,7 @@ func TestSSEFilter_OnlyComments(t *testing.T) {
230230
func TestSSEFilter_ScannerError(t *testing.T) {
231231
t.Parallel()
232232

233-
r := newSSEFilterReader(io.NopCloser(&errorReader{err: io.ErrUnexpectedEOF}))
233+
r := newSSEFilterReader(io.NopCloser(&errorReader{err: io.ErrUnexpectedEOF}), false)
234234

235235
_, err := io.ReadAll(r)
236236
assert.ErrorIs(t, err, io.ErrUnexpectedEOF)
@@ -256,7 +256,7 @@ func TestSSEFilter_CloseWithoutRead(t *testing.T) {
256256
onClose: func() { closed = true },
257257
}
258258

259-
r := newSSEFilterReader(tracker)
259+
r := newSSEFilterReader(tracker, false)
260260
require.NoError(t, r.Close())
261261
assert.True(t, closed, "underlying reader should be closed")
262262
}
@@ -344,3 +344,154 @@ func fetchThroughFilter(t *testing.T, url string) string {
344344
require.NoError(t, err)
345345
return string(body)
346346
}
347+
348+
// Gemini-shaped data chunks used by the keepalive tests: a text delta and a
349+
// media (inlineData) delta of the kind an image-output model streams.
350+
const (
351+
geminiTextChunk = `data: {"candidates":[{"content":{"parts":[{"text":"hi"}],"role":"model"}}]}` + "\n\n"
352+
geminiMediaChunk = `data: {"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"aGVsbG8="}}],"role":"model"},"finishReason":"STOP"}]}` + "\n\n"
353+
keepaliveFrame = "event: keepalive\ndata: {}\n\n"
354+
)
355+
356+
// TestSSEFilter_KeepaliveMode covers the opt-in keepalive-dropping mode used
357+
// by the Gemini gateway client: payload-free `event: keepalive` frames are
358+
// removed while every other frame — including named events with meaningful
359+
// data — passes through verbatim.
360+
func TestSSEFilter_KeepaliveMode(t *testing.T) {
361+
t.Parallel()
362+
363+
tests := []struct {
364+
name string
365+
in string
366+
want string
367+
}{
368+
{
369+
// The gateway scenario: keepalive frames interleaved with
370+
// real text and media chunks during a long image generation.
371+
name: "drops keepalive frames interleaved with data and media chunks",
372+
in: keepaliveFrame + geminiTextChunk + keepaliveFrame + keepaliveFrame + geminiMediaChunk,
373+
want: geminiTextChunk + geminiMediaChunk,
374+
},
375+
{
376+
name: "drops keepalive without a space after data:",
377+
in: "event: keepalive\ndata:{}\n\n" + geminiTextChunk,
378+
want: geminiTextChunk,
379+
},
380+
{
381+
name: "drops keepalive with an empty data payload",
382+
in: "event: keepalive\ndata:\n\n" + geminiTextChunk,
383+
want: geminiTextChunk,
384+
},
385+
{
386+
// Anthropic-style framing: a named event with meaningful data
387+
// must never be touched, even in keepalive mode.
388+
name: "preserves named events with meaningful data",
389+
in: "event: content_block_delta\ndata: {\"delta\":{\"text\":\"hi\"}}\n\n",
390+
want: "event: content_block_delta\ndata: {\"delta\":{\"text\":\"hi\"}}\n\n",
391+
},
392+
{
393+
// Conservative: only payload-free keepalives are dropped. A
394+
// keepalive-named event carrying real data is preserved.
395+
name: "preserves keepalive-named event with meaningful data",
396+
in: "event: keepalive\ndata: {\"note\":\"x\"}\n\n",
397+
want: "event: keepalive\ndata: {\"note\":\"x\"}\n\n",
398+
},
399+
{
400+
// The base filter's behavior is unchanged by keepalive mode.
401+
name: "still drops comment-only and no-data event frames",
402+
in: ": ping\n\nevent: ping\nid: abc\n\n" + geminiTextChunk,
403+
want: geminiTextChunk,
404+
},
405+
}
406+
407+
for _, tt := range tests {
408+
t.Run(tt.name, func(t *testing.T) {
409+
t.Parallel()
410+
411+
r := newSSEFilterReader(io.NopCloser(strings.NewReader(tt.in)), true)
412+
out, err := io.ReadAll(r)
413+
require.NoError(t, err)
414+
assert.Equal(t, tt.want, string(out))
415+
})
416+
}
417+
}
418+
419+
// TestSSEFilter_KeepaliveMode_OutputParsableByGenaiStyleParser feeds a
420+
// keepalive-interleaved Gemini stream through the keepalive-mode filter and
421+
// verifies the result against the constraint that made the fix necessary:
422+
// genai's iterateResponseStream (google.golang.org/genai api_client.go)
423+
// treats ANY non-blank line without a `data:` prefix as a fatal invalid
424+
// chunk. Every data payload must survive, in order.
425+
func TestSSEFilter_KeepaliveMode_OutputParsableByGenaiStyleParser(t *testing.T) {
426+
t.Parallel()
427+
428+
in := keepaliveFrame + geminiTextChunk + keepaliveFrame + geminiMediaChunk + keepaliveFrame
429+
r := newSSEFilterReader(io.NopCloser(strings.NewReader(in)), true)
430+
out, err := io.ReadAll(r)
431+
require.NoError(t, err)
432+
433+
var payloads []string
434+
for line := range strings.Lines(string(out)) {
435+
line = strings.TrimSuffix(line, "\n")
436+
if line == "" {
437+
continue
438+
}
439+
require.True(t, strings.HasPrefix(line, "data:"), "genai would reject this line as an invalid stream chunk: %q", line)
440+
payloads = append(payloads, strings.TrimPrefix(line, "data: "))
441+
}
442+
443+
require.Len(t, payloads, 2)
444+
assert.Contains(t, payloads[0], `"text":"hi"`)
445+
assert.Contains(t, payloads[1], `"inlineData"`)
446+
}
447+
448+
// TestSSEFilter_SharedPathKeepsKeepaliveFrames pins that the shared default
449+
// filter (used by every other provider) does NOT gain keepalive dropping:
450+
// an `event: keepalive` frame has a data line, so it passes through
451+
// verbatim, exactly like Anthropic's meaningful named events.
452+
func TestSSEFilter_SharedPathKeepsKeepaliveFrames(t *testing.T) {
453+
t.Parallel()
454+
455+
in := keepaliveFrame +
456+
"event: content_block_delta\ndata: {\"delta\":{\"text\":\"hi\"}}\n\n"
457+
458+
assert.Equal(t, in, fetchSSE(t, in))
459+
}
460+
461+
// TestNewHTTPClient_SSEKeepaliveFilterOptIn verifies the option wiring end
462+
// to end through NewHTTPClient: keepalive frames are dropped only when
463+
// WithSSEKeepaliveFilter is passed, and the default client leaves them in.
464+
func TestNewHTTPClient_SSEKeepaliveFilterOptIn(t *testing.T) {
465+
t.Parallel()
466+
467+
in := keepaliveFrame + geminiTextChunk
468+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
469+
w.Header().Set("Content-Type", "text/event-stream")
470+
_, _ = io.WriteString(w, in)
471+
}))
472+
t.Cleanup(srv.Close)
473+
474+
fetch := func(t *testing.T, client *http.Client) string {
475+
t.Helper()
476+
req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, srv.URL, http.NoBody)
477+
require.NoError(t, err)
478+
res, err := client.Do(req)
479+
require.NoError(t, err)
480+
defer func() { _ = res.Body.Close() }()
481+
body, err := io.ReadAll(res.Body)
482+
require.NoError(t, err)
483+
return string(body)
484+
}
485+
486+
t.Run("opted in drops keepalive frames", func(t *testing.T) {
487+
t.Parallel()
488+
client := NewHTTPClient(t.Context(), WithSSEKeepaliveFilter())
489+
assert.Equal(t, geminiTextChunk, fetch(t, client))
490+
})
491+
492+
t.Run("default keeps keepalive frames", func(t *testing.T) {
493+
t.Parallel()
494+
client := NewHTTPClient(t.Context())
495+
assert.Equal(t, in, fetch(t, client))
496+
})
497+
}

pkg/model/provider/gemini/client.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,12 @@ func NewClient(ctx context.Context, cfg *latest.ModelConfig, env environment.Pro
179179
}
180180
}
181181

182+
// The gateway keeps long generations alive with `event: keepalive`
183+
// + `data: {}` frames, which genai's SSE parser rejects as fatal
184+
// invalid chunks. Drop them here, on the gateway path only — direct
185+
// Gemini/Vertex clients never receive them.
186+
httpOptions = append(httpOptions, httpclient.WithSSEKeepaliveFilter())
187+
182188
gatewayHTTPClient := httpclient.NewHTTPClient(ctx, httpOptions...)
183189
globalOptions.WrapTransport(ctx, gatewayHTTPClient)
184190

0 commit comments

Comments
 (0)