diff --git a/pkg/httpclient/client.go b/pkg/httpclient/client.go index dfb932d6f..d4b123317 100644 --- a/pkg/httpclient/client.go +++ b/pkg/httpclient/client.go @@ -24,6 +24,10 @@ type HTTPOptions struct { Header http.Header Query url.Values + // dropSSEKeepaliveEvents enables keepalive-frame dropping in the SSE + // filter transport; see WithSSEKeepaliveFilter. + dropSSEKeepaliveEvents bool + // cagentID resolves the persistent install UUID stamped as // `X-Cagent-Id` on gateway-bound requests. It defaults to // [userid.Get]; tests inject their own source via @@ -89,7 +93,10 @@ func NewHTTPClient(ctx context.Context, opts ...Opt) *http.Client { var wrapped http.RoundTripper = &userAgentTransport{ httpOptions: httpOptions, - rt: &sseFilterTransport{base: rt}, + rt: &sseFilterTransport{ + base: rt, + dropKeepaliveEvents: httpOptions.dropSSEKeepaliveEvents, + }, } if httpOptions.refreshAuth != nil { // Outermost, so a replayed request goes through the whole chain again. @@ -235,6 +242,16 @@ func WithQuery(query url.Values) Opt { } } +// WithSSEKeepaliveFilter strips payload-free events named "keepalive". +// The Gemini gateway emits these transport frames, but the GenAI SDK rejects +// event-prefixed lines even when their only data is {}. Other names and +// keepalives with meaningful payloads are deliberately left unchanged. +func WithSSEKeepaliveFilter() Opt { + return func(o *HTTPOptions) { + o.dropSSEKeepaliveEvents = true + } +} + // newTransport returns an HTTP transport with automatic gzip compression disabled and Docker Desktop PAC support. func newTransport(_ context.Context) http.RoundTripper { rt := newAllowPrivateIPsTransport() diff --git a/pkg/httpclient/sse_filter.go b/pkg/httpclient/sse_filter.go index a4f7e1986..edc7b6908 100644 --- a/pkg/httpclient/sse_filter.go +++ b/pkg/httpclient/sse_filter.go @@ -10,7 +10,7 @@ import ( // sseFilterTransport wraps a base RoundTripper and, when the response is a // `text/event-stream`, replaces the body with one that strips SSE events -// containing no `data:` lines or named `keepalive`. +// containing no `data:` lines. // // Why this exists: some upstreams (notably OpenRouter) inject comment-only // keep-alive frames into their streams: @@ -28,12 +28,19 @@ import ( // // The filter normalises the byte stream so events with no `data:` lines // (comment-only events, or events bearing only `event:` / `id:` headers) -// never reach the SDK. Named `keepalive` events are also dropped, even when -// they carry `data: {}`, because Gemini's SDK rejects their `event:` header. -// Other data-bearing events pass through, and the filter is a no-op on -// non-SSE responses. +// never reach the SDK. Well-formed events pass through verbatim, and the +// filter is a no-op on non-SSE responses. +// +// dropKeepaliveEvents additionally drops whole `event: keepalive` frames +// whose data carries no payload (`data: {}` or empty). The Docker AI +// Gateway emits such frames during long generations; the genai SDK's SSE +// parser hard-fails on ANY `event:` line, so they must never reach it. The +// mode is opt-in (see WithSSEKeepaliveFilter) because other providers — +// Anthropic in particular — use `event:` headers as meaningful framing that +// must pass through untouched. type sseFilterTransport struct { - base http.RoundTripper + base http.RoundTripper + dropKeepaliveEvents bool } func (t *sseFilterTransport) RoundTrip(req *http.Request) (*http.Response, error) { @@ -44,32 +51,35 @@ func (t *sseFilterTransport) RoundTrip(req *http.Request) (*http.Response, error // Match the prefix so charset suffixes (e.g. "text/event-stream; // charset=utf-8") still trigger filtering. if strings.HasPrefix(strings.ToLower(res.Header.Get("Content-Type")), "text/event-stream") { - res.Body = newSSEFilterReader(res.Body) + res.Body = newSSEFilterReader(res.Body, t.dropKeepaliveEvents) } return res, err } // sseFilterReader buffers the lines of a single SSE event and only emits // them once it has seen the trailing blank line AND the event contained at -// least one `data:` line and was not a keepalive. A half-built event still -// pending at EOF is dropped silently — without the terminating blank line a -// downstream parser would not have dispatched it anyway. +// least one `data:` line. A half-built event still pending at EOF is +// dropped silently — without the terminating blank line a downstream parser +// would not have dispatched it anyway. type sseFilterReader struct { - src io.ReadCloser - scn *bufio.Scanner - out bytes.Buffer // bytes ready to hand back to the caller - pending bytes.Buffer // accumulated lines for the current event - hasData bool // saw at least one `data:` line in `pending` - keepalive bool // the last `event:` field names a keepalive + src io.ReadCloser + scn *bufio.Scanner + out bytes.Buffer // bytes ready to hand back to the caller + pending bytes.Buffer // accumulated lines for the current event + hasData bool // saw at least one `data:` line in `pending` + + dropKeepaliveEvents bool // see sseFilterTransport + isKeepalive bool // current event is named `keepalive` + hasMeaningfulData bool // saw a `data:` line whose payload isn't empty or `{}` } -func newSSEFilterReader(src io.ReadCloser) *sseFilterReader { +func newSSEFilterReader(src io.ReadCloser, dropKeepaliveEvents bool) *sseFilterReader { scn := bufio.NewScanner(src) // SSE events can be large (long completion tokens, image URLs, …). Match // the buffer size used by openai-go's own SSE decoder so we don't trip // `bufio.ErrTooLong` on payloads it would happily accept. scn.Buffer(make([]byte, 0, 64*1024), bufio.MaxScanTokenSize<<9) - return &sseFilterReader{src: src, scn: scn} + return &sseFilterReader{src: src, scn: scn, dropKeepaliveEvents: dropKeepaliveEvents} } func (r *sseFilterReader) Read(p []byte) (int, error) { @@ -88,28 +98,47 @@ func (r *sseFilterReader) Read(p []byte) (int, error) { func (r *sseFilterReader) consumeLine(line []byte) { switch { case len(line) == 0: - // Event boundary: emit data-bearing events except keepalives. - if r.hasData && !r.keepalive { + // Event boundary: emit the buffered event iff it had data and is + // not a payload-free keepalive frame in keepalive-dropping mode. + if r.hasData && (!r.isKeepalive || r.hasMeaningfulData) { r.out.Write(r.pending.Bytes()) r.out.WriteByte('\n') } r.pending.Reset() r.hasData = false - r.keepalive = false + r.isKeepalive = false + r.hasMeaningfulData = false case line[0] == ':': // SSE comment — drop entirely. default: r.pending.Write(line) r.pending.WriteByte('\n') - if bytes.HasPrefix(line, []byte("data:")) { + if value, ok := fieldValue(line, "data"); ok { r.hasData = true + if r.dropKeepaliveEvents { + if payload := bytes.TrimSpace(value); len(payload) > 0 && !bytes.Equal(payload, []byte("{}")) { + r.hasMeaningfulData = true + } + } + } else if r.dropKeepaliveEvents { + if field, value, _ := bytes.Cut(line, []byte(":")); bytes.Equal(field, []byte("event")) { + r.isKeepalive = bytes.Equal(bytes.TrimPrefix(value, []byte(" ")), []byte("keepalive")) + } } - if field, value, _ := bytes.Cut(line, []byte(":")); bytes.Equal(field, []byte("event")) { - r.keepalive = bytes.Equal(bytes.TrimPrefix(value, []byte(" ")), []byte("keepalive")) - } } } +// fieldValue returns the value of an SSE line whose field name is `name`, +// with the single optional leading space the SSE grammar allows already +// removed. +func fieldValue(line []byte, name string) ([]byte, bool) { + value, ok := bytes.CutPrefix(line, []byte(name+":")) + if !ok { + return nil, false + } + return bytes.TrimPrefix(value, []byte(" ")), true +} + func (r *sseFilterReader) Close() error { return r.src.Close() } diff --git a/pkg/httpclient/sse_filter_test.go b/pkg/httpclient/sse_filter_test.go index adc907422..5feb9ccd3 100644 --- a/pkg/httpclient/sse_filter_test.go +++ b/pkg/httpclient/sse_filter_test.go @@ -35,60 +35,6 @@ func TestSSEFilter_FiltersStream(t *testing.T) { "\n", want: "data: {\"id\":\"1\"}\n\n", }, - { - name: "drops keepalive events with data", - in: "event: keepalive\ndata: {}\n\ndata: ok\n\n", - want: "data: ok\n\n", - }, - { - name: "drops interleaved and trailing keepalives", - in: "data: first\n\n" + - "event: keepalive\ndata: {}\n\n" + - "event: keepalive\ndata: {}\n\n" + - "data: last\n\n" + - "event: keepalive\ndata: {}\n\n", - want: "data: first\n\ndata: last\n\n", - }, - { - name: "drops keepalive with data before event header", - in: "data: {}\nevent: keepalive\n\ndata: ok\n\n", - want: "data: ok\n\n", - }, - { - name: "drops keepalive without optional space and with CRLF", - in: "event:keepalive\r\ndata:{}\r\n\r\ndata: ok\r\n\r\n", - want: "data: ok\n\n", - }, - { - name: "only keepalives", - in: "event: keepalive\ndata: {}\n\nevent: keepalive\ndata: {}\n\n", - want: "", - }, - { - name: "uses last event header", - in: "event: chunk\nevent: keepalive\ndata: {}\n\n", - want: "", - }, - { - name: "later event header overrides keepalive", - in: "event: keepalive\nevent: chunk\ndata: {}\n\n", - want: "event: keepalive\nevent: chunk\ndata: {}\n\n", - }, - { - name: "empty event header overrides keepalive", - in: "event: keepalive\nevent\ndata: {}\n\n", - want: "event: keepalive\nevent\ndata: {}\n\n", - }, - { - name: "preserves unnamed empty JSON payload", - in: "data: {}\n\n", - want: "data: {}\n\n", - }, - { - name: "preserves other named events and errors", - in: "event: keepalive-extra\ndata: {}\n\nevent: error\ndata: {\"error\":\"failed\"}\n\n", - want: "event: keepalive-extra\ndata: {}\n\nevent: error\ndata: {\"error\":\"failed\"}\n\n", - }, { // Guard against the filter breaking ordinary streams that // don't contain comments. @@ -207,7 +153,7 @@ func TestSSEFilter_LargeEvent(t *testing.T) { t.Parallel() largeData := "data: " + strings.Repeat("x", 256*1024) + "\n\n" - r := newSSEFilterReader(io.NopCloser(strings.NewReader(largeData))) + r := newSSEFilterReader(io.NopCloser(strings.NewReader(largeData)), false) output, err := io.ReadAll(r) require.NoError(t, err) @@ -221,7 +167,7 @@ func TestSSEFilter_PartialReads(t *testing.T) { t.Parallel() input := "data: test1\n\ndata: test2\n\n" - r := newSSEFilterReader(io.NopCloser(strings.NewReader(input))) + r := newSSEFilterReader(io.NopCloser(strings.NewReader(input)), false) var output []byte buf := make([]byte, 5) @@ -246,7 +192,7 @@ func TestSSEFilter_IncompleteEventAtEOF(t *testing.T) { t.Parallel() input := "data: complete\n\ndata: incomplete" - r := newSSEFilterReader(io.NopCloser(strings.NewReader(input))) + r := newSSEFilterReader(io.NopCloser(strings.NewReader(input)), false) output, err := io.ReadAll(r) require.NoError(t, err) @@ -258,7 +204,7 @@ func TestSSEFilter_IncompleteEventAtEOF(t *testing.T) { func TestSSEFilter_EmptyInput(t *testing.T) { t.Parallel() - r := newSSEFilterReader(io.NopCloser(strings.NewReader(""))) + r := newSSEFilterReader(io.NopCloser(strings.NewReader("")), false) output, err := io.ReadAll(r) require.NoError(t, err) @@ -272,7 +218,7 @@ func TestSSEFilter_OnlyComments(t *testing.T) { t.Parallel() input := ": comment1\n\n: comment2\n\n" - r := newSSEFilterReader(io.NopCloser(strings.NewReader(input))) + r := newSSEFilterReader(io.NopCloser(strings.NewReader(input)), false) output, err := io.ReadAll(r) require.NoError(t, err) @@ -284,7 +230,7 @@ func TestSSEFilter_OnlyComments(t *testing.T) { func TestSSEFilter_ScannerError(t *testing.T) { t.Parallel() - r := newSSEFilterReader(io.NopCloser(&errorReader{err: io.ErrUnexpectedEOF})) + r := newSSEFilterReader(io.NopCloser(&errorReader{err: io.ErrUnexpectedEOF}), false) _, err := io.ReadAll(r) assert.ErrorIs(t, err, io.ErrUnexpectedEOF) @@ -310,7 +256,7 @@ func TestSSEFilter_CloseWithoutRead(t *testing.T) { onClose: func() { closed = true }, } - r := newSSEFilterReader(tracker) + r := newSSEFilterReader(tracker, false) require.NoError(t, r.Close()) assert.True(t, closed, "underlying reader should be closed") } @@ -398,3 +344,208 @@ func fetchThroughFilter(t *testing.T, url string) string { require.NoError(t, err) return string(body) } + +// Gemini-shaped data chunks used by the keepalive tests: a text delta and a +// media (inlineData) delta of the kind an image-output model streams. +const ( + geminiTextChunk = `data: {"candidates":[{"content":{"parts":[{"text":"hi"}],"role":"model"}}]}` + "\n\n" + geminiMediaChunk = `data: {"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"aGVsbG8="}}],"role":"model"},"finishReason":"STOP"}]}` + "\n\n" + keepaliveFrame = "event: keepalive\ndata: {}\n\n" +) + +// TestSSEFilter_KeepaliveMode covers the opt-in keepalive-dropping mode used +// by the Gemini gateway client: payload-free `event: keepalive` frames are +// removed while every other frame — including named events with meaningful +// data — passes through verbatim. +func TestSSEFilter_KeepaliveMode(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + in string + want string + }{ + { + name: "drops keepalive events with data", + in: "event: keepalive\ndata: {}\n\ndata: ok\n\n", + want: "data: ok\n\n", + }, + { + name: "drops interleaved and trailing keepalives", + in: "data: first\n\n" + + "event: keepalive\ndata: {}\n\n" + + "event: keepalive\ndata: {}\n\n" + + "data: last\n\n" + + "event: keepalive\ndata: {}\n\n", + want: "data: first\n\ndata: last\n\n", + }, + { + name: "drops keepalive with data before event header", + in: "data: {}\nevent: keepalive\n\ndata: ok\n\n", + want: "data: ok\n\n", + }, + { + name: "drops keepalive without optional space and with CRLF", + in: "event:keepalive\r\ndata:{}\r\n\r\ndata: ok\r\n\r\n", + want: "data: ok\n\n", + }, + { + name: "only keepalives", + in: "event: keepalive\ndata: {}\n\nevent: keepalive\ndata: {}\n\n", + want: "", + }, + { + name: "uses last event header", + in: "event: chunk\nevent: keepalive\ndata: {}\n\n", + want: "", + }, + { + name: "later event header overrides keepalive", + in: "event: keepalive\nevent: chunk\ndata: {}\n\n", + want: "event: keepalive\nevent: chunk\ndata: {}\n\n", + }, + { + name: "empty event header overrides keepalive", + in: "event: keepalive\nevent\ndata: {}\n\n", + want: "event: keepalive\nevent\ndata: {}\n\n", + }, + { + name: "preserves unnamed empty JSON payload", + in: "data: {}\n\n", + want: "data: {}\n\n", + }, + { + name: "preserves other named events and errors", + in: "event: keepalive-extra\ndata: {}\n\nevent: error\ndata: {\"error\":\"failed\"}\n\n", + want: "event: keepalive-extra\ndata: {}\n\nevent: error\ndata: {\"error\":\"failed\"}\n\n", + }, + { + // The gateway scenario: keepalive frames interleaved with + // real text and media chunks during a long image generation. + name: "drops keepalive frames interleaved with data and media chunks", + in: keepaliveFrame + geminiTextChunk + keepaliveFrame + keepaliveFrame + geminiMediaChunk, + want: geminiTextChunk + geminiMediaChunk, + }, + { + name: "drops keepalive without a space after data:", + in: "event: keepalive\ndata:{}\n\n" + geminiTextChunk, + want: geminiTextChunk, + }, + { + name: "drops keepalive with an empty data payload", + in: "event: keepalive\ndata:\n\n" + geminiTextChunk, + want: geminiTextChunk, + }, + { + // Anthropic-style framing: a named event with meaningful data + // must never be touched, even in keepalive mode. + name: "preserves named events with meaningful data", + in: "event: content_block_delta\ndata: {\"delta\":{\"text\":\"hi\"}}\n\n", + want: "event: content_block_delta\ndata: {\"delta\":{\"text\":\"hi\"}}\n\n", + }, + { + // Conservative: only payload-free keepalives are dropped. A + // keepalive-named event carrying real data is preserved. + name: "preserves keepalive-named event with meaningful data", + in: "event: keepalive\ndata: {\"note\":\"x\"}\n\n", + want: "event: keepalive\ndata: {\"note\":\"x\"}\n\n", + }, + { + // The base filter's behavior is unchanged by keepalive mode. + name: "still drops comment-only and no-data event frames", + in: ": ping\n\nevent: ping\nid: abc\n\n" + geminiTextChunk, + want: geminiTextChunk, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + r := newSSEFilterReader(io.NopCloser(strings.NewReader(tt.in)), true) + out, err := io.ReadAll(r) + require.NoError(t, err) + assert.Equal(t, tt.want, string(out)) + }) + } +} + +// TestSSEFilter_KeepaliveMode_OutputParsableByGenaiStyleParser feeds a +// keepalive-interleaved Gemini stream through the keepalive-mode filter and +// verifies the result against the constraint that made the fix necessary: +// genai's iterateResponseStream (google.golang.org/genai api_client.go) +// treats ANY non-blank line without a `data:` prefix as a fatal invalid +// chunk. Every data payload must survive, in order. +func TestSSEFilter_KeepaliveMode_OutputParsableByGenaiStyleParser(t *testing.T) { + t.Parallel() + + in := keepaliveFrame + geminiTextChunk + keepaliveFrame + geminiMediaChunk + keepaliveFrame + r := newSSEFilterReader(io.NopCloser(strings.NewReader(in)), true) + out, err := io.ReadAll(r) + require.NoError(t, err) + + var payloads []string + for line := range strings.Lines(string(out)) { + line = strings.TrimSuffix(line, "\n") + if line == "" { + continue + } + require.True(t, strings.HasPrefix(line, "data:"), "genai would reject this line as an invalid stream chunk: %q", line) + payloads = append(payloads, strings.TrimPrefix(line, "data: ")) + } + + require.Len(t, payloads, 2) + assert.Contains(t, payloads[0], `"text":"hi"`) + assert.Contains(t, payloads[1], `"inlineData"`) +} + +// TestSSEFilter_SharedPathKeepsKeepaliveFrames pins that the shared default +// filter (used by every other provider) does NOT gain keepalive dropping: +// an `event: keepalive` frame has a data line, so it passes through +// verbatim, exactly like Anthropic's meaningful named events. +func TestSSEFilter_SharedPathKeepsKeepaliveFrames(t *testing.T) { + t.Parallel() + + in := keepaliveFrame + + "event: content_block_delta\ndata: {\"delta\":{\"text\":\"hi\"}}\n\n" + + assert.Equal(t, in, fetchSSE(t, in)) +} + +// TestNewHTTPClient_SSEKeepaliveFilterOptIn verifies the option wiring end +// to end through NewHTTPClient: keepalive frames are dropped only when +// WithSSEKeepaliveFilter is passed, and the default client leaves them in. +func TestNewHTTPClient_SSEKeepaliveFilterOptIn(t *testing.T) { + t.Parallel() + + in := keepaliveFrame + geminiTextChunk + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, in) + })) + t.Cleanup(srv.Close) + + fetch := func(t *testing.T, client *http.Client) string { + t.Helper() + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, srv.URL, http.NoBody) + require.NoError(t, err) + res, err := client.Do(req) + require.NoError(t, err) + defer func() { _ = res.Body.Close() }() + body, err := io.ReadAll(res.Body) + require.NoError(t, err) + return string(body) + } + + t.Run("opted in drops keepalive frames", func(t *testing.T) { + t.Parallel() + client := NewHTTPClient(t.Context(), WithSSEKeepaliveFilter()) + assert.Equal(t, geminiTextChunk, fetch(t, client)) + }) + + t.Run("default keeps keepalive frames", func(t *testing.T) { + t.Parallel() + client := NewHTTPClient(t.Context()) + assert.Equal(t, in, fetch(t, client)) + }) +} diff --git a/pkg/model/provider/gemini/client.go b/pkg/model/provider/gemini/client.go index f009b67b0..59764553d 100644 --- a/pkg/model/provider/gemini/client.go +++ b/pkg/model/provider/gemini/client.go @@ -179,6 +179,12 @@ func NewClient(ctx context.Context, cfg *latest.ModelConfig, env environment.Pro } } + // The gateway keeps long generations alive with `event: keepalive` + // + `data: {}` frames, which genai's SSE parser rejects as fatal + // invalid chunks. Drop them here, on the gateway path only — direct + // Gemini/Vertex clients never receive them. + httpOptions = append(httpOptions, httpclient.WithSSEKeepaliveFilter()) + gatewayHTTPClient := httpclient.NewHTTPClient(ctx, httpOptions...) globalOptions.WrapTransport(ctx, gatewayHTTPClient) diff --git a/pkg/model/provider/gemini/gateway_sse_keepalive_test.go b/pkg/model/provider/gemini/gateway_sse_keepalive_test.go new file mode 100644 index 000000000..67548872f --- /dev/null +++ b/pkg/model/provider/gemini/gateway_sse_keepalive_test.go @@ -0,0 +1,110 @@ +package gemini + +import ( + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/config/latest" + "github.com/docker/docker-agent/pkg/environment" + "github.com/docker/docker-agent/pkg/model/provider/options" +) + +// writeGeminiSSEResponseWithKeepalives replays what the Docker AI Gateway +// sends during a long image generation: `event: keepalive` + `data: {}` +// frames interleaved with real text and inlineData chunks. +func writeGeminiSSEResponseWithKeepalives(w http.ResponseWriter) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, + "event: keepalive\ndata: {}\n\n"+ + `data: {"candidates":[{"content":{"parts":[{"text":"here it comes"}],"role":"model"}}]}`+"\n\n"+ + "event: keepalive\ndata: {}\n\n"+ + "event: keepalive\ndata: {}\n\n"+ + `data: {"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"aGVsbG8="}}],"role":"model"},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":1,"totalTokenCount":2}}`+"\n\n") +} + +// collectStream drains a stream, concatenating text deltas and collecting +// text deltas, and returns the first non-EOF error (nil on clean EOF). +func collectStream(t *testing.T, stream chat.MessageStream) (string, error) { + t.Helper() + defer stream.Close() + + var text strings.Builder + for { + resp, err := stream.Recv() + if errors.Is(err, io.EOF) { + return text.String(), nil + } + if err != nil { + return text.String(), err + } + for _, choice := range resp.Choices { + text.WriteString(choice.Delta.Content) + } + } +} + +// TestCreateChatCompletionStream_GatewaySurvivesKeepaliveFrames pins the +// keepalive fix end to end: a gateway stream interleaved with keepalive +// frames completes without error and delivers every text delta. +// Without the gateway-scoped filter, genai's SSE parser fails the whole +// stream with "invalid stream chunk: event: keepalive". +func TestCreateChatCompletionStream_GatewaySurvivesKeepaliveFrames(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + writeGeminiSSEResponseWithKeepalives(w) + })) + t.Cleanup(server.Close) + + cfg := &latest.ModelConfig{Provider: "google", Model: "gemini-3-pro-image-preview"} + env := environment.NewMapEnvProvider(map[string]string{ + environment.DockerDesktopTokenEnv: "test-dd-token", + }) + client, err := NewClient(t.Context(), cfg, env, options.WithGateway(server.URL)) + require.NoError(t, err) + + stream, err := client.CreateChatCompletionStream(t.Context(), []chat.Message{ + {Role: chat.MessageRoleUser, Content: "generate an image of a red panda"}, + }, nil) + require.NoError(t, err) + + text, err := collectStream(t, stream) + require.NoError(t, err, "keepalive frames must never reach the genai SSE parser on the gateway path") + assert.Equal(t, "here it comes", text) +} + +// TestCreateChatCompletionStream_DirectPathKeepaliveUnfiltered pins the +// scoping: the direct (non-gateway) Gemini API path does NOT get keepalive +// filtering, so the same frames still surface genai's invalid-chunk error. +// Direct Gemini never emits these frames — this test only guards against +// the filter accidentally widening beyond the gateway client. +func TestCreateChatCompletionStream_DirectPathKeepaliveUnfiltered(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + writeGeminiSSEResponseWithKeepalives(w) + })) + t.Cleanup(server.Close) + + cfg := &latest.ModelConfig{Provider: "google", Model: "gemini-3-pro-image-preview", BaseURL: server.URL} + env := environment.NewMapEnvProvider(map[string]string{"GOOGLE_API_KEY": "test-key"}) + client, err := NewClient(t.Context(), cfg, env) + require.NoError(t, err) + + stream, err := client.CreateChatCompletionStream(t.Context(), []chat.Message{ + {Role: chat.MessageRoleUser, Content: "hello"}, + }, nil) + require.NoError(t, err) + + _, err = collectStream(t, stream) + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid stream chunk") +} diff --git a/pkg/model/provider/gemini/image_response_modalities_test.go b/pkg/model/provider/gemini/image_response_modalities_test.go new file mode 100644 index 000000000..89b624ac9 --- /dev/null +++ b/pkg/model/provider/gemini/image_response_modalities_test.go @@ -0,0 +1,359 @@ +package gemini + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/config/latest" + "github.com/docker/docker-agent/pkg/environment" + "github.com/docker/docker-agent/pkg/model/provider/base" + "github.com/docker/docker-agent/pkg/model/provider/options" + "github.com/docker/docker-agent/pkg/modelsdev" + "github.com/docker/docker-agent/pkg/rag/types" + "github.com/docker/docker-agent/pkg/tools" +) + +func TestWantsImageResponseModalities(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + apiSurface string + enabled bool + opts []options.Opt + want bool + }{ + {name: "gateway, enabled: present", apiSurface: apiSurfaceGateway, enabled: true, want: true}, + {name: "direct Gemini API, enabled: present", apiSurface: apiSurfaceGeminiAPI, enabled: true, want: true}, + {name: "Vertex AI, enabled: present", apiSurface: apiSurfaceVertexAI, enabled: true, want: true}, + {name: "gateway, disabled: absent", apiSurface: apiSurfaceGateway}, + {name: "direct Gemini API, disabled: absent", apiSurface: apiSurfaceGeminiAPI}, + {name: "unknown surface, enabled: absent", apiSurface: "unknown", enabled: true}, + {name: "empty surface, enabled: absent", enabled: true}, + { + name: "gateway, enabled, generating title: absent", apiSurface: apiSurfaceGateway, enabled: true, + opts: []options.Opt{options.WithGeneratingTitle()}, + }, + { + name: "gateway, enabled, compacting: absent", apiSurface: apiSurfaceGateway, enabled: true, + opts: []options.Opt{options.WithCompacting()}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + client := &Client{ + Config: base.Config{ + ModelOptions: options.Apply(tt.opts...), + }, + apiSurface: tt.apiSurface, + } + + assert.Equal(t, tt.want, client.wantsImageResponseModalities(tt.enabled)) + }) + } +} + +// capturedRequests is a mutex-guarded log of raw request bodies received by +// a [newBodyCapturingGeminiServer], letting tests assert exactly what was +// (or, cheaply, was not — an empty log) serialized onto the wire. +type capturedRequests struct { + mu sync.Mutex + bodies [][]byte +} + +func (c *capturedRequests) add(b []byte) { + c.mu.Lock() + defer c.mu.Unlock() + c.bodies = append(c.bodies, b) +} + +func (c *capturedRequests) all() [][]byte { + c.mu.Lock() + defer c.mu.Unlock() + return append([][]byte(nil), c.bodies...) +} + +// newBodyCapturingGeminiServer starts an httptest server that records the +// raw request body of every call it receives before responding via +// respond, so tests can decode the exact generationConfig sent on the wire. +func newBodyCapturingGeminiServer(t *testing.T, respond func(w http.ResponseWriter)) (*httptest.Server, *capturedRequests) { + t.Helper() + captured := &capturedRequests{} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + captured.add(body) + respond(w) + })) + t.Cleanup(server.Close) + return server, captured +} + +// writeGeminiGenerateContentJSONResponse writes a minimal, non-streaming +// Gemini generateContent JSON response whose sole text part is text (used +// for Rerank, which does not use the SSE streaming endpoint). +func writeGeminiGenerateContentJSONResponse(w http.ResponseWriter, text string) { + w.Header().Set("Content-Type", "application/json") + payload, _ := json.Marshal(map[string]any{ + "candidates": []map[string]any{{ + "content": map[string]any{ + "role": "model", + "parts": []map[string]any{{"text": text}}, + }, + }}, + }) + _, _ = w.Write(payload) +} + +// responseModalitiesInBody decodes body's generationConfig.responseModalities +// (see genai's generateContentConfigToMldev, which nests the serialized +// GenerateContentConfig under a top-level "generationConfig" key for the +// Gemini Developer API), returning nil when either key is absent. +func responseModalitiesInBody(t *testing.T, body []byte) []string { + t.Helper() + + var req map[string]any + require.NoError(t, json.Unmarshal(body, &req)) + + genCfg, ok := req["generationConfig"].(map[string]any) + if !ok { + return nil + } + raw, ok := genCfg["responseModalities"].([]any) + if !ok { + return nil + } + out := make([]string, len(raw)) + for i, v := range raw { + out[i], _ = v.(string) + } + return out +} + +func drainStream(t *testing.T, stream chat.MessageStream) { + t.Helper() + defer stream.Close() + for { + if _, err := stream.Recv(); err != nil { + break + } + } +} + +// TestCreateChatCompletionStream_ImageResponseModalities_PositiveRoutes pins +// that supported Gemini surfaces request TEXT+IMAGE output in that order. +func TestCreateChatCompletionStream_ImageResponseModalities_PositiveRoutes(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + cfg func(serverURL string) *latest.ModelConfig + env map[string]string + gateway bool + }{ + { + name: "gateway", + cfg: func(string) *latest.ModelConfig { + return &latest.ModelConfig{Provider: "google", Model: "gemini-2.5-flash-image", OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}} + }, + env: map[string]string{environment.DockerDesktopTokenEnv: "test-dd-token"}, + gateway: true, + }, + { + name: "direct Gemini API", + cfg: func(serverURL string) *latest.ModelConfig { + return &latest.ModelConfig{Provider: "google", Model: "gemini-2.5-flash-image", BaseURL: serverURL, OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}} + }, + env: map[string]string{"GOOGLE_API_KEY": "test-key"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + server, captured := newBodyCapturingGeminiServer(t, writeGeminiSSEResponse) + var opts []options.Opt + if tt.gateway { + opts = append(opts, options.WithGateway(server.URL)) + } + client, err := NewClient(t.Context(), tt.cfg(server.URL), environment.NewMapEnvProvider(tt.env), opts...) + require.NoError(t, err) + + stream, err := client.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "generate an image of a red panda"}}, nil) + require.NoError(t, err) + drainStream(t, stream) + + bodies := captured.all() + require.Len(t, bodies, 1) + assert.Equal(t, []string{"TEXT", "IMAGE"}, responseModalitiesInBody(t, bodies[0])) + }) + } +} + +func TestCreateChatCompletionStream_ImageResponseModalities_GoogleSurfaces(t *testing.T) { + t.Parallel() + + store := modelsdev.NewDatabaseStore(&modelsdev.Database{Providers: map[string]modelsdev.Provider{ + "google": {Models: map[string]modelsdev.Model{ + "image-model": {Modalities: modelsdev.Modalities{Output: []string{"text", "image"}}}, + }}, + }}) + tests := []struct { + name string + model string + declared *latest.OutputCapabilitiesConfig + opts []options.Opt + want []string + }{ + {name: "catalogue enables omitted declaration", model: "image-model", want: []string{"TEXT", "IMAGE"}}, + {name: "catalogue enables empty block", model: "image-model", declared: &latest.OutputCapabilitiesConfig{}, want: []string{"TEXT", "IMAGE"}}, + {name: "explicit false beats catalogue", model: "image-model", declared: &latest.OutputCapabilitiesConfig{Image: new(false)}}, + {name: "unknown model stays off", model: "unknown-image-model"}, + {name: "explicit true enables unknown model", model: "unknown-image-model", declared: &latest.OutputCapabilitiesConfig{Image: new(true)}, want: []string{"TEXT", "IMAGE"}}, + {name: "catalogue enabled title omits modalities", model: "image-model", opts: []options.Opt{options.WithGeneratingTitle()}}, + {name: "catalogue enabled compaction omits modalities", model: "image-model", opts: []options.Opt{options.WithCompacting()}}, + } + for _, gateway := range []bool{false, true} { + for _, tt := range tests { + t.Run(fmt.Sprintf("gateway=%t/%s", gateway, tt.name), func(t *testing.T) { + t.Parallel() + server, captured := newBodyCapturingGeminiServer(t, writeGeminiSSEResponse) + cfg := &latest.ModelConfig{Provider: "google", Model: tt.model, BaseURL: server.URL, OutputCapabilities: tt.declared} + env := map[string]string{"GOOGLE_API_KEY": "test-key"} + opts := append([]options.Opt{options.WithModelsDevStore(store)}, tt.opts...) + if gateway { + cfg.BaseURL = "" + env = map[string]string{environment.DockerDesktopTokenEnv: "test-dd-token"} + opts = append(opts, options.WithGateway(server.URL)) + } + client, err := NewClient(t.Context(), cfg, environment.NewMapEnvProvider(env), opts...) + require.NoError(t, err) + stream, err := client.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "hello"}}, nil) + require.NoError(t, err) + drainStream(t, stream) + bodies := captured.all() + require.Len(t, bodies, 1) + assert.Equal(t, tt.want, responseModalitiesInBody(t, bodies[0])) + }) + } + } +} + +// TestRerank_NeverSetsResponseModalities pins that Rerank — which shares +// buildConfig with CreateChatCompletionStream but must retain its own +// structured-output-only config — never gains response modalities, even on +// a model explicitly declared image-output-capable. +func TestRerank_NeverSetsResponseModalities(t *testing.T) { + t.Parallel() + + server, captured := newBodyCapturingGeminiServer(t, func(w http.ResponseWriter) { + writeGeminiGenerateContentJSONResponse(w, `{"scores":[1]}`) + }) + + cfg := &latest.ModelConfig{ + Provider: "google", + Model: "gemini-2.5-flash-image", + OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}, + } + env := environment.NewMapEnvProvider(map[string]string{ + environment.DockerDesktopTokenEnv: "test-dd-token", + }) + client, err := NewClient(t.Context(), cfg, env, options.WithGateway(server.URL)) + require.NoError(t, err) + + scores, err := client.Rerank(t.Context(), "query", []types.Document{{Content: "doc1"}}, "") + require.NoError(t, err) + require.Len(t, scores, 1) + + bodies := captured.all() + require.Len(t, bodies, 1) + assert.Nil(t, responseModalitiesInBody(t, bodies[0]), "Rerank must never request response modalities") +} + +// TestCreateChatCompletionStream_ImageResponseModalities_GuardRejectedRoutesNeverDispatch +// preserves guard precedence: on every declared-image route, a request shape +// the guard rejects (custom function tools or structured output) must still +// make zero provider calls, so nothing — including response modalities — is +// ever serialized onto the wire. Server-side built-ins remain dispatchable. +func TestCreateChatCompletionStream_ImageResponseModalities_GuardRejectedRoutesNeverDispatch(t *testing.T) { + t.Parallel() + + newRejectedClient := func(t *testing.T, serverURL string, extraOpts ...options.Opt) *Client { + t.Helper() + cfg := &latest.ModelConfig{ + Provider: "google", + Model: "gemini-2.5-flash-image", + OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}, + } + env := environment.NewMapEnvProvider(map[string]string{ + environment.DockerDesktopTokenEnv: "test-dd-token", + }) + opts := append([]options.Opt{options.WithGateway(serverURL)}, extraOpts...) + client, err := NewClient(t.Context(), cfg, env, opts...) + require.NoError(t, err) + return client + } + + assertRejectedWithNoDispatch := func(t *testing.T, err error, stream chat.MessageStream, captured *capturedRequests) { + t.Helper() + require.Nil(t, stream) + var incompatible *ImageOutputRequestIncompatibleError + require.ErrorAs(t, err, &incompatible) + assert.Empty(t, captured.all(), "guard rejection must dispatch nothing, so no modalities are ever serialized") + } + + t.Run("custom function tools rejected", func(t *testing.T) { + t.Parallel() + server, captured := newBodyCapturingGeminiServer(t, writeGeminiSSEResponse) + client := newRejectedClient(t, server.URL) + + stream, err := client.CreateChatCompletionStream(t.Context(), []chat.Message{ + {Role: chat.MessageRoleUser, Content: "hello"}, + }, []tools.Tool{{Name: "read_file", Description: "reads a file", Parameters: map[string]any{"type": "object"}}}) + + assertRejectedWithNoDispatch(t, err, stream, captured) + }) + + t.Run("built-in tool dispatches", func(t *testing.T) { + t.Parallel() + server, captured := newBodyCapturingGeminiServer(t, writeGeminiSSEResponse) + client := newRejectedClient(t, server.URL) + client.ModelConfig.ProviderOpts = map[string]any{"google_search": true} + + stream, err := client.CreateChatCompletionStream(t.Context(), []chat.Message{ + {Role: chat.MessageRoleUser, Content: "hello"}, + }, nil) + + require.NoError(t, err) + for { + if _, err := stream.Recv(); err != nil { + break + } + } + require.Len(t, captured.all(), 1) + assert.Equal(t, []string{"TEXT", "IMAGE"}, responseModalitiesInBody(t, captured.all()[0])) + }) + + t.Run("structured output rejected", func(t *testing.T) { + t.Parallel() + server, captured := newBodyCapturingGeminiServer(t, writeGeminiSSEResponse) + client := newRejectedClient(t, server.URL, options.WithStructuredOutput(&latest.StructuredOutput{Schema: map[string]any{"type": "object"}})) + + stream, err := client.CreateChatCompletionStream(t.Context(), []chat.Message{ + {Role: chat.MessageRoleUser, Content: "hello"}, + }, nil) + + assertRejectedWithNoDispatch(t, err, stream, captured) + }) +} diff --git a/pkg/model/provider/gemini/streaming_test.go b/pkg/model/provider/gemini/streaming_test.go index 6c46f6c27..540ae90e2 100644 --- a/pkg/model/provider/gemini/streaming_test.go +++ b/pkg/model/provider/gemini/streaming_test.go @@ -62,6 +62,12 @@ func TestCreateChatCompletionStream_Keepalive(t *testing.T) { require.NoError(t, err) t.Cleanup(stream.Close) + if !gateway { + _, err := stream.Recv() + require.ErrorContains(t, err, "invalid stream chunk: event: keepalive") + return + } + for _, text := range []string{"hello", " world"} { response, err := stream.Recv() require.NoError(t, err)