fix(#3996): preserve Gemini user thinking budgets outside title generation

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.

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()
 

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())
 	})
 }

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{}

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 {

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,15 +174,15 @@ 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,
 		},
 		{
 			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",

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",