@@ -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 \n data: 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 \n data: 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) {
204204func 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) {
230230func 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\n data: {}\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\n data:{}\n \n " + geminiTextChunk ,
378+ want : geminiTextChunk ,
379+ },
380+ {
381+ name : "drops keepalive with an empty data payload" ,
382+ in : "event: keepalive\n data:\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\n data: {\" delta\" :{\" text\" :\" hi\" }}\n \n " ,
390+ want : "event: content_block_delta\n data: {\" 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\n data: {\" note\" :\" x\" }\n \n " ,
397+ want : "event: keepalive\n data: {\" 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 \n event: ping\n id: 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\n data: {\" 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+ }
0 commit comments