diff --git a/pkg/model/provider/gemini/client.go b/pkg/model/provider/gemini/client.go index 498f6f966..60eba1423 100644 --- a/pkg/model/provider/gemini/client.go +++ b/pkg/model/provider/gemini/client.go @@ -425,7 +425,7 @@ func extractMimeType(dataURLPrefix string) string { return "image/jpeg" // Default fallback } -// buildConfig creates GenerateContentConfig from model config +// BuildConfig creates GenerateContentConfig from model config. func (c *Client) buildConfig() *genai.GenerateContentConfig { config := &genai.GenerateContentConfig{} if c.ModelConfig.MaxTokens != nil { @@ -453,7 +453,11 @@ func (c *Client) buildConfig() *genai.GenerateContentConfig { // Apply thinking configuration for Gemini models. // See https://ai.google.dev/gemini-api/docs/thinking if c.ModelOptions.NoThinking() { - // NoThinking requested (e.g. title generation). For Gemini 3+ models + if c.ModelOptions.GeneratingTitle() { + return config + } + + // NoThinking requested (e.g. MCP sampling). For Gemini 3+ models // that always think, use the lowest level and bump MaxOutputTokens so // internal reasoning doesn't consume the entire budget. Gemini 2.5 and // older can fully disable thinking with ThinkingBudget=0. diff --git a/pkg/model/provider/gemini/client_test.go b/pkg/model/provider/gemini/client_test.go index 9b344ddc0..e68adc7fe 100644 --- a/pkg/model/provider/gemini/client_test.go +++ b/pkg/model/provider/gemini/client_test.go @@ -10,10 +10,74 @@ import ( "github.com/docker/docker-agent/pkg/chat" "github.com/docker/docker-agent/pkg/config/latest" "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/tools" ) +func TestBuildConfig_NoThinking(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + model string + opts []options.Opt + wantThinking bool + wantMinTokens bool + }{ + { + name: "title generation omits thinking config", + model: "gemini-3-flash", + opts: []options.Opt{options.WithGeneratingTitle(), options.WithNoThinking()}, + wantThinking: false, + }, + { + name: "MCP sampling disables Gemini 3 thinking", + model: "gemini-3-flash", + opts: []options.Opt{options.WithNoThinking()}, + wantThinking: true, + wantMinTokens: true, + }, + { + name: "MCP sampling disables Gemini 2.5 thinking", + model: "gemini-2.5-flash", + opts: []options.Opt{options.WithNoThinking()}, + wantThinking: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + client := &Client{Config: base.Config{ + ModelConfig: latest.ModelConfig{ + Provider: "google", + Model: tt.model, + ThinkingBudget: &latest.ThinkingBudget{Effort: "high"}, + }, + ModelOptions: options.Apply(tt.opts...), + }} + + config := client.buildConfig() + if !tt.wantThinking { + assert.Nil(t, config.ThinkingConfig) + return + } + + require.NotNil(t, config.ThinkingConfig) + assert.False(t, config.ThinkingConfig.IncludeThoughts) + if tt.wantMinTokens { + assert.Equal(t, genai.ThinkingLevelLow, config.ThinkingConfig.ThinkingLevel) + assert.GreaterOrEqual(t, config.MaxOutputTokens, int32(200)) + return + } + require.NotNil(t, config.ThinkingConfig.ThinkingBudget) + assert.Zero(t, *config.ThinkingConfig.ThinkingBudget) + }) + } +} + func TestBuildConfig_Gemini25_ThinkingBudget(t *testing.T) { t.Parallel() diff --git a/pkg/model/provider/gemini/image_output_guard_test.go b/pkg/model/provider/gemini/image_output_guard_test.go index 346695d87..bebab0059 100644 --- a/pkg/model/provider/gemini/image_output_guard_test.go +++ b/pkg/model/provider/gemini/image_output_guard_test.go @@ -19,8 +19,8 @@ import ( ) // TestCheckImageOutputRequestCompatibility_GateConditions exhaustively -// covers when the guard does and does not apply: it must reject only on the -// gateway surface, only for a model with output_capabilities.image: true, +// covers when the guard does and does not apply: it must reject on each +// supported Google surface, only for a model with image output enabled, // and only when the request also carries custom function tools, a built-in // tool, or structured output. func TestCheckImageOutputRequestCompatibility_GateConditions(t *testing.T) { @@ -41,8 +41,8 @@ func TestCheckImageOutputRequestCompatibility_GateConditions(t *testing.T) { {name: "gateway declared true, no extras: allowed", apiSurface: apiSurfaceGateway, declared: declaredTrue}, {name: "gateway declared false: never rejects even with tools", apiSurface: apiSurfaceGateway, declared: declaredFalse, requestTools: 1}, {name: "gateway undeclared: never rejects even with tools", apiSurface: apiSurfaceGateway, declared: nil, requestTools: 1}, - {name: "direct Gemini API declared true: guard does not apply", apiSurface: apiSurfaceGeminiAPI, declared: declaredTrue, requestTools: 1}, - {name: "Vertex AI declared true: guard does not apply", apiSurface: apiSurfaceVertexAI, declared: declaredTrue, requestTools: 1}, + {name: "direct Gemini API declared true + tools: rejected", apiSurface: apiSurfaceGeminiAPI, declared: declaredTrue, requestTools: 1, wantReject: []imageOutputIncompatibility{imageOutputIncompatibleTools}}, + {name: "Vertex AI declared true + tools: rejected", apiSurface: apiSurfaceVertexAI, declared: declaredTrue, requestTools: 1, wantReject: []imageOutputIncompatibility{imageOutputIncompatibleTools}}, { name: "gateway declared true + custom function tools: rejected", apiSurface: apiSurfaceGateway, declared: declaredTrue, requestTools: 2, @@ -232,9 +232,8 @@ func TestCreateChatCompletionStream_ImageOutputGuard_RejectsBeforeDispatch(t *te // TestCreateChatCompletionStream_ImageOutputGuard_PreservesNormalBehavior // proves the guard is a no-op (request reaches the provider) for every route -// it must not touch: no extras on the declared route, tools/structured -// output when the declaration is false/missing, and tools/structured output -// on a direct (non-gateway) Gemini call even when declared true. +// it must not touch: no extras on the declared route and tools/structured +// output when the declaration is false or missing. func TestCreateChatCompletionStream_ImageOutputGuard_PreservesNormalBehavior(t *testing.T) { t.Parallel() @@ -336,31 +335,16 @@ func TestCreateChatCompletionStream_ImageOutputGuard_PreservesNormalBehavior(t * assert.Positive(t, counter.calls.Load()) }) - t.Run("direct (non-gateway) Gemini call with tools, declared true: guard does not apply", func(t *testing.T) { + t.Run("direct Gemini call with tools, resolved image output: rejected before dispatch", func(t *testing.T) { t.Parallel() var counter geminiCountingTransport - cfg := &latest.ModelConfig{ - Provider: "google", - Model: "gemini-2.5-flash-image", - BaseURL: server.URL, - OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}, - } - env := environment.NewMapEnvProvider(map[string]string{ - "GOOGLE_API_KEY": "test-key", - }) - client, err := NewClient(t.Context(), cfg, env, - options.WithHTTPTransportWrapper(func(base http.RoundTripper) http.RoundTripper { - counter.base = base - return &counter - }), - ) - require.NoError(t, err) - - 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"}}}) + cfg := &latest.ModelConfig{Provider: "google", Model: "gemini-2.5-flash-image", BaseURL: server.URL, OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}} + env := environment.NewMapEnvProvider(map[string]string{"GOOGLE_API_KEY": "test-key"}) + client, err := NewClient(t.Context(), cfg, env, options.WithHTTPTransportWrapper(func(base http.RoundTripper) http.RoundTripper { counter.base = base; return &counter })) require.NoError(t, err) - drain(t, stream) - assert.Positive(t, counter.calls.Load()) + 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"}}}) + require.Nil(t, stream) + require.Error(t, err) + assert.Zero(t, counter.calls.Load()) }) } diff --git a/pkg/model/provider/gemini/image_output_instruction.go b/pkg/model/provider/gemini/image_output_instruction.go index fa271eac4..1d086ada6 100644 --- a/pkg/model/provider/gemini/image_output_instruction.go +++ b/pkg/model/provider/gemini/image_output_instruction.go @@ -2,8 +2,8 @@ package gemini import "google.golang.org/genai" -// imageOutputMediaFileInstruction is appended as a system instruction on the -// explicit image-output gateway route (see wantsImageResponseModalities) so +// imageOutputMediaFileInstruction is appended as a system instruction on +// declared image-output chat requests (see wantsImageResponseModalities) so // generated images arrive with a machine-readable filename: the runtime // strips these exact marker lines from the reply and uses the paths to name // the materialized workspace files (pkg/runtime/generated_media_markers.go). @@ -21,7 +21,7 @@ Rules: // applyImageOutputMediaFileInstruction appends the marker-protocol // instruction to the request's system instruction, preserving any parts // already present. Callers gate it on wantsImageResponseModalities so only -// the explicit image-output gateway chat route ever carries it. +// declared image-output chat requests carry it. func applyImageOutputMediaFileInstruction(config *genai.GenerateContentConfig) { if config.SystemInstruction == nil { config.SystemInstruction = &genai.Content{} diff --git a/pkg/model/provider/gemini/image_output_instruction_test.go b/pkg/model/provider/gemini/image_output_instruction_test.go index 90919572b..6efd90591 100644 --- a/pkg/model/provider/gemini/image_output_instruction_test.go +++ b/pkg/model/provider/gemini/image_output_instruction_test.go @@ -45,46 +45,65 @@ func systemInstructionTextsInBody(t *testing.T, body []byte) []string { return out } -// TestCreateChatCompletionStream_MediaFileInstruction_PositiveRoute pins -// the sole route that must carry the media-file marker instruction — the -// gateway surface with an explicit output_capabilities.image declaration on -// an ordinary chat turn — and that it is sent exactly once. -func TestCreateChatCompletionStream_MediaFileInstruction_PositiveRoute(t *testing.T) { +// TestCreateChatCompletionStream_MediaFileInstruction_PositiveRoutes pins +// that ordinary image-output chat requests carry the media-file marker +// instruction exactly once on every supported Google surface. +func TestCreateChatCompletionStream_MediaFileInstruction_PositiveRoutes(t *testing.T) { t.Parallel() - server, captured := newBodyCapturingGeminiServer(t, writeGeminiSSEResponse) + 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() - cfg := &latest.ModelConfig{ - Provider: "google", - Model: "gemini-2.5-flash-image", - OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}, + server, captured := newBodyCapturingGeminiServer(t, writeGeminiSSEResponse) + opts := []options.Opt(nil) + 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) + texts := systemInstructionTextsInBody(t, bodies[0]) + require.Len(t, texts, 1, "the instruction must be sent exactly once") + assert.Equal(t, imageOutputMediaFileInstruction, texts[0]) + assert.Equal(t, 1, strings.Count(texts[0], "[media-file: "), "the instruction must show the marker format exactly once") + }) } - 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) - drainStream(t, stream) - - bodies := captured.all() - require.Len(t, bodies, 1) - texts := systemInstructionTextsInBody(t, bodies[0]) - require.Len(t, texts, 1, "the instruction must be sent exactly once") - assert.Equal(t, imageOutputMediaFileInstruction, texts[0]) - assert.Equal(t, 1, strings.Count(texts[0], "[media-file: "), "the instruction must show the marker format exactly once") } // TestCreateChatCompletionStream_MediaFileInstruction_AbsentOnOtherRoutes -// pins that every other route sends no marker instruction (and no system -// instruction at all, since nothing else sets one today): a direct -// (non-gateway) call even when declared image-capable, gateway calls -// without the explicit declaration, and gateway title-generation or -// compaction calls. +// pins that non-image-output and internal text-only requests send no marker +// instruction: gateway calls without image output enabled, plus gateway +// title-generation and compaction calls. func TestCreateChatCompletionStream_MediaFileInstruction_AbsentOnOtherRoutes(t *testing.T) { t.Parallel() @@ -95,16 +114,6 @@ func TestCreateChatCompletionStream_MediaFileInstruction_AbsentOnOtherRoutes(t * gateway bool opts []options.Opt }{ - { - name: "direct Gemini API, declared true: absent", - 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"}, - }, { name: "gateway, declared false: absent", cfg: func(string) *latest.ModelConfig { diff --git a/pkg/model/provider/gemini/image_response_modalities_test.go b/pkg/model/provider/gemini/image_response_modalities_test.go index 4380d2c7f..fa942db0b 100644 --- a/pkg/model/provider/gemini/image_response_modalities_test.go +++ b/pkg/model/provider/gemini/image_response_modalities_test.go @@ -28,8 +28,8 @@ import ( func TestWantsImageResponseModalities(t *testing.T) { t.Parallel() - declaredTrue := &latest.OutputCapabilitiesConfig{Image: latest.Bool(true)} - declaredFalse := &latest.OutputCapabilitiesConfig{Image: latest.Bool(false)} + declaredTrue := &latest.OutputCapabilitiesConfig{Image: new(true)} + declaredFalse := &latest.OutputCapabilitiesConfig{Image: new(false)} tests := []struct { name string @@ -174,7 +174,7 @@ func TestCreateChatCompletionStream_ImageResponseModalities_PositiveRoutes(t *te { name: "gateway", cfg: func(string) *latest.ModelConfig { - return &latest.ModelConfig{Provider: "google", Model: "gemini-2.5-flash-image", OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: latest.Bool(true)}} + 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, @@ -182,7 +182,7 @@ func TestCreateChatCompletionStream_ImageResponseModalities_PositiveRoutes(t *te { 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: latest.Bool(true)}} + 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"}, }, @@ -235,7 +235,7 @@ func TestCreateChatCompletionStream_ImageResponseModalities_AbsentOnOtherRoutes( cfg: func(string) *latest.ModelConfig { return &latest.ModelConfig{ Provider: "google", Model: "gemini-2.5-flash-image", - OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: latest.Bool(false)}, + OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(false)}, } }, env: map[string]string{environment.DockerDesktopTokenEnv: "test-dd-token"}, @@ -254,7 +254,7 @@ func TestCreateChatCompletionStream_ImageResponseModalities_AbsentOnOtherRoutes( cfg: func(string) *latest.ModelConfig { return &latest.ModelConfig{ Provider: "google", Model: "gemini-2.5-flash-image", - OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: latest.Bool(true)}, + OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}, } }, env: map[string]string{environment.DockerDesktopTokenEnv: "test-dd-token"}, @@ -266,7 +266,7 @@ func TestCreateChatCompletionStream_ImageResponseModalities_AbsentOnOtherRoutes( cfg: func(string) *latest.ModelConfig { return &latest.ModelConfig{ Provider: "google", Model: "gemini-2.5-flash-image", - OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: latest.Bool(true)}, + OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}, } }, env: map[string]string{environment.DockerDesktopTokenEnv: "test-dd-token"}, @@ -317,7 +317,7 @@ func TestRerank_NeverSetsResponseModalities(t *testing.T) { cfg := &latest.ModelConfig{ Provider: "google", Model: "gemini-2.5-flash-image", - OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: latest.Bool(true)}, + OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}, } env := environment.NewMapEnvProvider(map[string]string{ environment.DockerDesktopTokenEnv: "test-dd-token", @@ -347,7 +347,7 @@ func TestCreateChatCompletionStream_ImageResponseModalities_GuardRejectedRoutesN cfg := &latest.ModelConfig{ Provider: "google", Model: "gemini-2.5-flash-image", - OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: latest.Bool(true)}, + OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}, } env := environment.NewMapEnvProvider(map[string]string{ environment.DockerDesktopTokenEnv: "test-dd-token", diff --git a/pkg/runtime/image_output_guard_integration_test.go b/pkg/runtime/image_output_guard_integration_test.go index f251b3957..681331a86 100644 --- a/pkg/runtime/image_output_guard_integration_test.go +++ b/pkg/runtime/image_output_guard_integration_test.go @@ -20,7 +20,7 @@ import ( "github.com/docker/docker-agent/pkg/tools" ) -// TestRunStream_ImageOutputGuard_RejectsBeforeDispatch drives the gateway +// TestRunStream_ImageOutputGuard_RejectsBeforeDispatch drives the // image-output request guard (pkg/model/provider/gemini/image_output_guard.go) // through the real run loop: a real *gemini.Client, talking to an httptest // gateway, behind a real [agent.Agent] and [LocalRuntime], through RunStream. @@ -61,8 +61,7 @@ func TestRunStream_ImageOutputGuard_RejectsBeforeDispatch(t *testing.T) { // A custom function tool is enough to trip the guard on its own (no // ResponseModalities / rendering involved): declaring - // output_capabilities.image on the gateway route is incompatible with - // any custom tool. + // output_capabilities.image is incompatible with any custom tool. readFileTool := tools.Tool{ Name: "read_file", Description: "reads a file from disk", diff --git a/pkg/tui/page/chat/image_output_guard_integration_test.go b/pkg/tui/page/chat/image_output_guard_integration_test.go index cbd102cb6..55f93aefb 100644 --- a/pkg/tui/page/chat/image_output_guard_integration_test.go +++ b/pkg/tui/page/chat/image_output_guard_integration_test.go @@ -72,7 +72,7 @@ func (r *recordingMessages) RemoveSpinner() { } // TestImageOutputGuard_RuntimeToTUI_RejectsBeforeDispatchAndClearsSpinner -// drives the gateway image-output request guard +// drives the image-output request guard // (pkg/model/provider/gemini/image_output_guard.go) through its real // pre-dispatch path — a real *gemini.Client talking to an httptest gateway — // and then through the exact production event flow the run loop uses to @@ -124,8 +124,7 @@ func TestImageOutputGuard_RuntimeToTUI_RejectsBeforeDispatchAndClearsSpinner(t * // A custom function tool is enough to trip the guard on its own (no // ResponseModalities / rendering involved): declaring - // output_capabilities.image on the gateway route is incompatible with - // any custom tool. + // output_capabilities.image is incompatible with any custom tool. readFileTool := tools.Tool{ Name: "read_file", Description: "reads a file from disk",