From fb1aeab5c09f1e5f5b53142da111ad29c54b2164 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Thu, 30 Jul 2026 10:11:08 +0200 Subject: [PATCH] feat(#3996): reject incompatible image-output gateway requests Gemini gateway image-output requests fail with an opaque HTTP 400 when they also carry custom tools, built-in tools, or structured output. Resolve the model capability first, then reject those unsupported shapes locally before network dispatch with a bounded error naming only fixed safe categories. Keep ordinary text requests and compatible image-output requests unchanged. Request-shape diagnostics report the resolved capability and separately note whether configuration explicitly overrode catalogue metadata. Tests cover the gate, no-dispatch behavior, safe errors, and the runtime/TUI error seam. --- docs/providers/google/index.md | 28 ++ pkg/model/provider/gemini/client.go | 25 +- pkg/model/provider/gemini/diagnostics.go | 12 +- pkg/model/provider/gemini/diagnostics_test.go | 57 ++- .../provider/gemini/image_output_guard.go | 69 ++++ .../gemini/image_output_guard_test.go | 366 ++++++++++++++++++ .../image_output_guard_integration_test.go | 113 ++++++ pkg/tui/components/messages/messages.go | 17 + .../image_output_guard_integration_test.go | 202 ++++++++++ 9 files changed, 871 insertions(+), 18 deletions(-) create mode 100644 pkg/model/provider/gemini/image_output_guard.go create mode 100644 pkg/model/provider/gemini/image_output_guard_test.go create mode 100644 pkg/runtime/image_output_guard_integration_test.go create mode 100644 pkg/tui/page/chat/image_output_guard_integration_test.go diff --git a/docs/providers/google/index.md b/docs/providers/google/index.md index efbb278171..e03dbe3ef8 100644 --- a/docs/providers/google/index.md +++ b/docs/providers/google/index.md @@ -60,6 +60,34 @@ models: | `gemini-2.5-flash` | Fast inference, cost-effective | | `gemini-2.5-pro` | Strong reasoning, large context | +## Generated Images + +Some Gemini models (e.g. `gemini-2.5-flash-image`) are designed to generate +an image directly as part of their reply, not just describe one. Docker +Agent's Gemini request path doesn't yet ask for that image output — that +support is still being completed — so today a request like this gets a +text-only reply. See +[Generated Media](../../features/tui/index.md#generated-media) for the +current, verified state. + +```yaml +agents: + root: + model: google/gemini-2.5-flash-image +``` + +When the model is accessed through a Docker AI Gateway and explicitly +declared image-output-capable with +[`output_capabilities.image: true`](../../configuration/models/index.md#output-capabilities), +Docker Agent has verified that request combined with custom function tools, +a built-in tool (e.g. `google_search`), or structured output gets rejected +by the gateway with an opaque, empty-body HTTP 400. To avoid that, Docker +Agent rejects such a combination itself, before any request is sent, with a +clear error naming which feature is incompatible. Plain text requests to +that model (no tools, no structured output) are unaffected, as is every +other route: direct Gemini API/Vertex AI calls, and gateway calls to a model +without the declaration. + ## Thinking Budget Gemini supports two approaches depending on the model version: diff --git a/pkg/model/provider/gemini/client.go b/pkg/model/provider/gemini/client.go index 1a50f78000..29d50987af 100644 --- a/pkg/model/provider/gemini/client.go +++ b/pkg/model/provider/gemini/client.go @@ -729,6 +729,17 @@ func stringifyEnumValues(values []any) []string { return out } +// wantsImageResponseModalities reports whether this ordinary chat request +// should ask Gemini for TEXT+IMAGE output. +func (c *Client) wantsImageResponseModalities(imageOutputEnabled bool) bool { + switch c.apiSurface { + case apiSurfaceGateway, apiSurfaceGeminiAPI, apiSurfaceVertexAI: + default: + return false + } + return imageOutputEnabled && !c.ModelOptions.GeneratingTitle() && !c.ModelOptions.Compacting() +} + // CreateChatCompletionStream creates a streaming chat completion request func (c *Client) CreateChatCompletionStream( ctx context.Context, @@ -740,9 +751,15 @@ func (c *Client) CreateChatCompletionStream( } config := c.buildConfig() + imageOutputEnabled := c.ImageOutputEnabled(ctx) + + if c.wantsImageResponseModalities(imageOutputEnabled) { + config.ResponseModalities = []string{string(genai.ModalityText), string(genai.ModalityImage)} + } // Start with Google built-in tools (search, maps, code execution) from provider_opts - config.Tools = c.builtInTools() + builtInTools := c.builtInTools() + config.Tools = builtInTools // Add tools to config if provided if len(requestTools) > 0 { @@ -767,7 +784,11 @@ func (c *Client) CreateChatCompletionStream( } } - shape := newRequestShape(c, config, len(requestTools)) + if err := c.checkImageOutputRequestCompatibility(imageOutputEnabled, config, builtInTools, len(requestTools)); err != nil { + return nil, err + } + + shape := newRequestShape(c, config, len(requestTools), imageOutputEnabled) slog.DebugContext(ctx, "Gemini request shape", shape.LogAttrs()...) contents := convertMessagesToGemini(ctx, messages, c.ID(), c.ModelOptions.ModelsDevStore(), c.CapsOverride()) diff --git a/pkg/model/provider/gemini/diagnostics.go b/pkg/model/provider/gemini/diagnostics.go index 01fda9283c..ce8f2dd2aa 100644 --- a/pkg/model/provider/gemini/diagnostics.go +++ b/pkg/model/provider/gemini/diagnostics.go @@ -58,12 +58,8 @@ type RequestShape struct { // or "gateway". APISurface string - // OutputCapabilityKnown and OutputCapabilityEnabled report whether an - // authoritative source for the model's image/media *output* capability - // was consulted for this request. No such source exists yet — it is the - // subject of a later step — so both fields are always false today, - // deliberately reporting "unknown" rather than guessing from the model - // ID string. + // OutputCapabilityKnown records whether the capability was resolved from an + // explicit configuration override instead of the catalogue. OutputCapabilityKnown bool OutputCapabilityEnabled bool } @@ -71,7 +67,7 @@ type RequestShape struct { // newRequestShape captures a [RequestShape] from a fully-built // genai.GenerateContentConfig (i.e. after tools/ToolConfig have been // attached) and the client that built it. -func newRequestShape(c *Client, config *genai.GenerateContentConfig, functionToolCount int) RequestShape { +func newRequestShape(c *Client, config *genai.GenerateContentConfig, functionToolCount int, imageOutputEnabled bool) RequestShape { modalities := normalizeResponseModalities(config.ResponseModalities) kinds := builtInToolKinds(config.Tools) @@ -86,6 +82,8 @@ func newRequestShape(c *Client, config *genai.GenerateContentConfig, functionToo ThinkingConfigSet: config.ThinkingConfig != nil, NoThinkingRequested: c.ModelOptions.NoThinking(), APISurface: c.apiSurface, + OutputCapabilityKnown: c.ModelConfig.OutputCapabilities != nil && c.ModelConfig.OutputCapabilities.Image != nil, + OutputCapabilityEnabled: imageOutputEnabled, } if config.ToolConfig != nil { diff --git a/pkg/model/provider/gemini/diagnostics_test.go b/pkg/model/provider/gemini/diagnostics_test.go index b25bd56f93..1225d542f8 100644 --- a/pkg/model/provider/gemini/diagnostics_test.go +++ b/pkg/model/provider/gemini/diagnostics_test.go @@ -40,7 +40,7 @@ func TestNewRequestShape_Minimal(t *testing.T) { config := client.buildConfig() config.Tools = client.builtInTools() - shape := newRequestShape(client, config, 0) + shape := newRequestShape(client, config, 0, false) assert.False(t, shape.ResponseModalitiesSet) assert.Empty(t, shape.ResponseModalities) @@ -52,9 +52,8 @@ func TestNewRequestShape_Minimal(t *testing.T) { assert.False(t, shape.ThinkingConfigSet) assert.False(t, shape.NoThinkingRequested) assert.Equal(t, apiSurfaceGeminiAPI, shape.APISurface) - // No authoritative output-capability source exists yet: diagnostics must - // report "unknown", never guess from the model ID. - assert.False(t, shape.OutputCapabilityKnown) + // No output_capabilities declaration on this ModelConfig: diagnostics + // must report "unknown", never guess from the model ID. assert.False(t, shape.OutputCapabilityEnabled) } @@ -93,7 +92,7 @@ func TestNewRequestShape_ToolsAndBuiltIns(t *testing.T) { config.ToolConfig.IncludeServerSideToolInvocations = new(true) } - shape := newRequestShape(client, config, len(requestTools)) + shape := newRequestShape(client, config, len(requestTools), false) assert.Equal(t, 2, shape.BuiltInToolCount) assert.ElementsMatch(t, []string{"google_search", "google_maps"}, shape.BuiltInToolKinds) @@ -114,7 +113,7 @@ func TestNewRequestShape_ResponseModalitiesNormalized(t *testing.T) { config := client.buildConfig() config.ResponseModalities = []string{" text ", "IMAGE", "text", ""} - shape := newRequestShape(client, config, 0) + shape := newRequestShape(client, config, 0, false) assert.True(t, shape.ResponseModalitiesSet) assert.Equal(t, []string{"TEXT", "IMAGE"}, shape.ResponseModalities) @@ -135,7 +134,7 @@ func TestNewRequestShape_NoThinkingRequested(t *testing.T) { } config := client.buildConfig() - shape := newRequestShape(client, config, 0) + shape := newRequestShape(client, config, 0, false) assert.True(t, shape.ThinkingConfigSet) assert.True(t, shape.NoThinkingRequested) @@ -156,7 +155,7 @@ func TestNewRequestShape_StructuredOutputPresent(t *testing.T) { } config := client.buildConfig() - shape := newRequestShape(client, config, 0) + shape := newRequestShape(client, config, 0, false) require.True(t, shape.StructuredOutputPresent) @@ -168,6 +167,46 @@ func TestNewRequestShape_StructuredOutputPresent(t *testing.T) { } } +// TestNewRequestShape_OutputCapabilityUsesResolvedValue pins that diagnostics +// distinguish an explicit override while reporting the resolved capability. +func TestNewRequestShape_OutputCapabilityUsesResolvedValue(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + outputCapabilities *latest.OutputCapabilitiesConfig + resolvedEnabled bool + wantKnown bool + wantEnabled bool + }{ + {name: "catalogue enabled", outputCapabilities: nil, resolvedEnabled: true, wantKnown: false, wantEnabled: true}, + {name: "declared false", outputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(false)}, wantKnown: true, wantEnabled: false}, + {name: "declared true", outputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}, resolvedEnabled: true, wantKnown: true, wantEnabled: 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: "gemini-2.5-flash-image", + OutputCapabilities: tt.outputCapabilities, + }, + }, + apiSurface: apiSurfaceGeminiAPI, + } + config := client.buildConfig() + + shape := newRequestShape(client, config, 0, tt.resolvedEnabled) + assert.Equal(t, tt.wantKnown, shape.OutputCapabilityKnown) + assert.Equal(t, tt.wantEnabled, shape.OutputCapabilityEnabled) + }) + } +} + // TestRequestShape_LogAttrsNeverLeaksToolSchemas is the core safety // regression for this diagnostic: it builds a request with a function tool // carrying a marker description and parameter schema, then verifies that @@ -189,7 +228,7 @@ func TestRequestShape_LogAttrsNeverLeaksToolSchemas(t *testing.T) { require.NoError(t, err) config.Tools = allTools - shape := newRequestShape(client, config, len(requestTools)) + shape := newRequestShape(client, config, len(requestTools), false) for _, attr := range flattenAttrs(shape.LogAttrs()) { assert.NotContains(t, attr, marker, "RequestShape must never carry tool descriptions or schemas") diff --git a/pkg/model/provider/gemini/image_output_guard.go b/pkg/model/provider/gemini/image_output_guard.go new file mode 100644 index 0000000000..856b325423 --- /dev/null +++ b/pkg/model/provider/gemini/image_output_guard.go @@ -0,0 +1,69 @@ +package gemini + +import ( + "fmt" + "strings" + + "google.golang.org/genai" +) + +// imageOutputIncompatibility names a fixed, safe request-feature class +// rejected by the image-output request guard. Values are display-safe: +// never provider text, tool names, schema contents, or prompts. +type imageOutputIncompatibility string + +const ( + imageOutputIncompatibleTools imageOutputIncompatibility = "tools" + imageOutputIncompatibleBuiltInTools imageOutputIncompatibility = "built-in tools" + imageOutputIncompatibleStructuredOutput imageOutputIncompatibility = "structured output" +) + +// ImageOutputRequestIncompatibleError is returned before any provider +// dispatch when a request to an image-output-capable model +// (output_capabilities.image: true) combines custom function tools (with +// their required ToolConfig), a built-in tool, or structured output. The +// request shape is intentionally unsupported until live direct Gemini API and +// Vertex AI verification proves it is accepted there; gateway probing has +// already shown opaque, empty-body HTTP 400 responses. Rejecting locally keeps +// the verified minimal request byte-for-byte and gives the caller a specific, +// actionable error instead. +type ImageOutputRequestIncompatibleError struct { + // Incompatibilities is always non-empty. Its values are the fixed enum + // above — never provider text, tool names/schemas, or prompt content. + Incompatibilities []imageOutputIncompatibility +} + +func (e *ImageOutputRequestIncompatibleError) Error() string { + names := make([]string, len(e.Incompatibilities)) + for i, c := range e.Incompatibilities { + names[i] = string(c) + } + return fmt.Sprintf( + "this model is configured for image output (output_capabilities.image) and does not support %s in the same request; use a separate model or request for that combination", + strings.Join(names, ", "), + ) +} + +// checkImageOutputRequestCompatibility rejects, before any provider dispatch, +// an incompatible request when image output is enabled by configuration or the +// models.dev catalogue. +func (c *Client) checkImageOutputRequestCompatibility(imageOutputEnabled bool, config *genai.GenerateContentConfig, builtInTools []*genai.Tool, requestTools int) error { + if !imageOutputEnabled { + return nil + } + + var incompatibilities []imageOutputIncompatibility + if requestTools > 0 { + incompatibilities = append(incompatibilities, imageOutputIncompatibleTools) + } + if len(builtInTools) > 0 { + incompatibilities = append(incompatibilities, imageOutputIncompatibleBuiltInTools) + } + if config.ResponseMIMEType != "" || config.ResponseJsonSchema != nil { + incompatibilities = append(incompatibilities, imageOutputIncompatibleStructuredOutput) + } + if len(incompatibilities) == 0 { + return nil + } + return &ImageOutputRequestIncompatibleError{Incompatibilities: incompatibilities} +} diff --git a/pkg/model/provider/gemini/image_output_guard_test.go b/pkg/model/provider/gemini/image_output_guard_test.go new file mode 100644 index 0000000000..346695d877 --- /dev/null +++ b/pkg/model/provider/gemini/image_output_guard_test.go @@ -0,0 +1,366 @@ +package gemini + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/genai" + + "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/modelerrors" + "github.com/docker/docker-agent/pkg/tools" +) + +// 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, +// and only when the request also carries custom function tools, a built-in +// tool, or structured output. +func TestCheckImageOutputRequestCompatibility_GateConditions(t *testing.T) { + t.Parallel() + + declaredTrue := &latest.OutputCapabilitiesConfig{Image: new(true)} + declaredFalse := &latest.OutputCapabilitiesConfig{Image: new(false)} + + tests := []struct { + name string + apiSurface string + declared *latest.OutputCapabilitiesConfig + builtInTools []*genai.Tool + requestTools int + structured bool + wantReject []imageOutputIncompatibility + }{ + {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: "gateway declared true + custom function tools: rejected", + apiSurface: apiSurfaceGateway, declared: declaredTrue, requestTools: 2, + wantReject: []imageOutputIncompatibility{imageOutputIncompatibleTools}, + }, + { + name: "gateway declared true + built-in tool: rejected", + apiSurface: apiSurfaceGateway, declared: declaredTrue, builtInTools: []*genai.Tool{{GoogleSearch: &genai.GoogleSearch{}}}, + wantReject: []imageOutputIncompatibility{imageOutputIncompatibleBuiltInTools}, + }, + { + name: "gateway declared true + structured output: rejected", + apiSurface: apiSurfaceGateway, declared: declaredTrue, structured: true, + wantReject: []imageOutputIncompatibility{imageOutputIncompatibleStructuredOutput}, + }, + { + name: "gateway declared true + tools and structured output: both reported", + apiSurface: apiSurfaceGateway, declared: declaredTrue, requestTools: 1, structured: true, + wantReject: []imageOutputIncompatibility{imageOutputIncompatibleTools, imageOutputIncompatibleStructuredOutput}, + }, + } + + 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: "gemini-2.5-flash-image", + OutputCapabilities: tt.declared, + }, + }, + apiSurface: tt.apiSurface, + } + config := &genai.GenerateContentConfig{} + if tt.structured { + config.ResponseMIMEType = "application/json" + } + + err := client.checkImageOutputRequestCompatibility(tt.declared != nil && tt.declared.Image != nil && *tt.declared.Image, config, tt.builtInTools, tt.requestTools) + + if len(tt.wantReject) == 0 { + assert.NoError(t, err) + return + } + var incompatible *ImageOutputRequestIncompatibleError + require.ErrorAs(t, err, &incompatible, "expected an *ImageOutputRequestIncompatibleError, got %v", err) + assert.Equal(t, tt.wantReject, incompatible.Incompatibilities) + }) + } +} + +func TestImageOutputRequestIncompatibleError_MessageNamesCategoriesOnly(t *testing.T) { + t.Parallel() + + err := &ImageOutputRequestIncompatibleError{Incompatibilities: []imageOutputIncompatibility{ + imageOutputIncompatibleTools, imageOutputIncompatibleStructuredOutput, + }} + msg := err.Error() + assert.Contains(t, msg, "output_capabilities.image") + assert.Contains(t, msg, "tools") + assert.Contains(t, msg, "structured output") +} + +// TestImageOutputRequestIncompatibleError_RoutesThroughExistingErrorSeam +// drives the guard's error through the same modelerrors.FormatError call the +// runtime loop uses to build ErrorEvent.Error (pkg/runtime/loop_steps.go), +// which the TUI renders verbatim (pkg/tui/page/chat/runtime_events.go). No +// new plumbing is needed: the guard's error is a plain error, not an +// overflow/truncation-shaped one, so FormatError must pass it through +// unchanged, and ClassifyModelError must not mark it retryable (retrying +// this exact request would just reject again). +func TestImageOutputRequestIncompatibleError_RoutesThroughExistingErrorSeam(t *testing.T) { + t.Parallel() + + err := &ImageOutputRequestIncompatibleError{Incompatibilities: []imageOutputIncompatibility{imageOutputIncompatibleBuiltInTools}} + + visible := modelerrors.FormatError(err) + assert.Equal(t, err.Error(), visible, "a plain incompatibility error must pass through FormatError unchanged") + assert.Contains(t, visible, "output_capabilities.image") + assert.Contains(t, visible, "built-in tools") + + retryable, rateLimited, _ := modelerrors.ClassifyModelError(err) + assert.False(t, retryable, "a deterministic local rejection must not be retried") + assert.False(t, rateLimited) +} + +// TestCreateChatCompletionStream_ImageOutputGuard_RejectsBeforeDispatch drives +// the guard through the real CreateChatCompletionStream path against an +// httptest server, and asserts zero provider calls on rejection. +func TestCreateChatCompletionStream_ImageOutputGuard_RejectsBeforeDispatch(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + writeGeminiSSEResponse(w) + })) + defer server.Close() + + newClient := func(t *testing.T, counter *geminiCountingTransport) *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", + }) + client, err := NewClient(t.Context(), cfg, env, + options.WithGateway(server.URL), + options.WithHTTPTransportWrapper(func(base http.RoundTripper) http.RoundTripper { + counter.base = base + return counter + }), + ) + require.NoError(t, err) + return client + } + + t.Run("custom function tools rejected with zero provider calls", func(t *testing.T) { + t.Parallel() + var counter geminiCountingTransport + client := newClient(t, &counter) + + 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) + var incompatible *ImageOutputRequestIncompatibleError + require.ErrorAs(t, err, &incompatible) + assert.Equal(t, []imageOutputIncompatibility{imageOutputIncompatibleTools}, incompatible.Incompatibilities) + assert.Zero(t, counter.calls.Load(), "guard must reject before any provider dispatch") + }) + + t.Run("built-in tool rejected with zero provider calls", func(t *testing.T) { + t.Parallel() + var counter geminiCountingTransport + client := newClient(t, &counter) + client.ModelConfig.ProviderOpts = map[string]any{"google_search": true} + + stream, err := client.CreateChatCompletionStream(t.Context(), []chat.Message{ + {Role: chat.MessageRoleUser, Content: "hello"}, + }, nil) + + require.Nil(t, stream) + var incompatible *ImageOutputRequestIncompatibleError + require.ErrorAs(t, err, &incompatible) + assert.Equal(t, []imageOutputIncompatibility{imageOutputIncompatibleBuiltInTools}, incompatible.Incompatibilities) + assert.Zero(t, counter.calls.Load(), "guard must reject before any provider dispatch") + }) + + t.Run("structured output rejected with zero provider calls", func(t *testing.T) { + t.Parallel() + var counter geminiCountingTransport + 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), + options.WithStructuredOutput(&latest.StructuredOutput{Schema: map[string]any{"type": "object"}}), + 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"}, + }, nil) + + require.Nil(t, stream) + var incompatible *ImageOutputRequestIncompatibleError + require.ErrorAs(t, err, &incompatible) + assert.Equal(t, []imageOutputIncompatibility{imageOutputIncompatibleStructuredOutput}, incompatible.Incompatibilities) + assert.Zero(t, counter.calls.Load(), "guard must reject before any provider dispatch") + }) +} + +// 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. +func TestCreateChatCompletionStream_ImageOutputGuard_PreservesNormalBehavior(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + writeGeminiSSEResponse(w) + })) + defer server.Close() + + drain := func(t *testing.T, stream chat.MessageStream) { + t.Helper() + defer stream.Close() + for { + if _, err := stream.Recv(); err != nil { + break + } + } + } + + t.Run("gateway declared true, no tools/structured output: reaches provider", func(t *testing.T) { + t.Parallel() + var counter geminiCountingTransport + 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), + 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"}, + }, nil) + require.NoError(t, err) + drain(t, stream) + assert.Positive(t, counter.calls.Load()) + }) + + t.Run("gateway with tools, declaration false: reaches provider", func(t *testing.T) { + t.Parallel() + var counter geminiCountingTransport + cfg := &latest.ModelConfig{ + Provider: "google", + Model: "gemini-2.5-flash", + OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(false)}, + } + env := environment.NewMapEnvProvider(map[string]string{ + environment.DockerDesktopTokenEnv: "test-dd-token", + }) + client, err := NewClient(t.Context(), cfg, env, + options.WithGateway(server.URL), + 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"}}}) + require.NoError(t, err) + drain(t, stream) + assert.Positive(t, counter.calls.Load()) + }) + + t.Run("gateway with tools, declaration missing: reaches provider", func(t *testing.T) { + t.Parallel() + var counter geminiCountingTransport + cfg := &latest.ModelConfig{ + Provider: "google", + Model: "gemini-2.5-flash", + } + env := environment.NewMapEnvProvider(map[string]string{ + environment.DockerDesktopTokenEnv: "test-dd-token", + }) + client, err := NewClient(t.Context(), cfg, env, + options.WithGateway(server.URL), + 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"}}}) + require.NoError(t, err) + drain(t, stream) + 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.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"}}}) + require.NoError(t, err) + drain(t, stream) + assert.Positive(t, counter.calls.Load()) + }) +} diff --git a/pkg/runtime/image_output_guard_integration_test.go b/pkg/runtime/image_output_guard_integration_test.go new file mode 100644 index 0000000000..f251b39572 --- /dev/null +++ b/pkg/runtime/image_output_guard_integration_test.go @@ -0,0 +1,113 @@ +package runtime + +import ( + "fmt" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/agent" + "github.com/docker/docker-agent/pkg/config/latest" + "github.com/docker/docker-agent/pkg/environment" + "github.com/docker/docker-agent/pkg/model/provider/gemini" + "github.com/docker/docker-agent/pkg/model/provider/options" + "github.com/docker/docker-agent/pkg/session" + "github.com/docker/docker-agent/pkg/team" + "github.com/docker/docker-agent/pkg/tools" +) + +// TestRunStream_ImageOutputGuard_RejectsBeforeDispatch drives the gateway +// 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. +// +// It proves the guard's rejection survives the full fallback/loop machinery +// unchanged: zero HTTP requests ever reach the provider, the loop emits +// exactly one ErrorEvent whose text is the guard's safe, fixed message, a +// StreamStartedEvent precedes it and a StreamStoppedEvent closes the turn, +// and no assistant content (text or reasoning) is ever produced — i.e. no +// silent "success" alongside the error. The TUI-facing half of this seam +// (the same ErrorEvent reaching the message list and clearing the spinner) +// is covered by +// TestImageOutputGuard_RuntimeToTUI_RejectsBeforeDispatchAndClearsSpinner in +// pkg/tui/page/chat, which cannot import [team] (see +// e2e/dependencies_test.go's "TUI musn't know about teams"). +func TestRunStream_ImageOutputGuard_RejectsBeforeDispatch(t *testing.T) { + t.Parallel() + + var providerCalls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + providerCalls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + payload := `{"candidates":[{"content":{"parts":[{"text":"hi"}],"role":"model"},"finishReason":"STOP","index":0}]}` + _, _ = fmt.Fprintf(w, "data: %s\n\n", payload) + })) + defer server.Close() + + 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 := gemini.NewClient(t.Context(), cfg, env, options.WithGateway(server.URL)) + require.NoError(t, err) + + // 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. + readFileTool := tools.Tool{ + Name: "read_file", + Description: "reads a file from disk", + Parameters: map[string]any{"type": "object"}, + } + root := agent.New("root", "You are a test agent", agent.WithModel(client), agent.WithTools(readFileTool)) + tm := team.New(team.WithAgents(root)) + + rt, err := NewLocalRuntime(t.Context(), tm, WithSessionCompaction(false), WithModelStore(mockModelStore{})) + require.NoError(t, err) + + sess := session.New(session.WithUserMessage("draw a cat")) + sess.Title = "image output guard integration test" + + var events []Event + for ev := range rt.RunStream(t.Context(), sess) { + events = append(events, ev) + } + + assert.Zero(t, providerCalls.Load(), "the guard must reject before any request reaches the provider") + + var errEvent *ErrorEvent + var streamStarted *StreamStartedEvent + var streamStopped *StreamStoppedEvent + for _, ev := range events { + switch e := ev.(type) { + case *ErrorEvent: + require.Nil(t, errEvent, "expected exactly one ErrorEvent") + errEvent = e + case *StreamStartedEvent: + if streamStarted == nil { + streamStarted = e + } + case *StreamStoppedEvent: + streamStopped = e + case *AgentChoiceEvent: + t.Fatalf("guard rejection must not produce assistant content, got AgentChoiceEvent %q", e.Content) + case *AgentChoiceReasoningEvent: + t.Fatalf("guard rejection must not produce reasoning content, got AgentChoiceReasoningEvent %q", e.Content) + } + } + require.NotNil(t, streamStarted, "expected a StreamStartedEvent") + require.NotNil(t, errEvent, "expected an ErrorEvent for the rejected request") + require.NotNil(t, streamStopped, "expected a StreamStoppedEvent to close out the turn") + assert.Contains(t, errEvent.Error, "output_capabilities.image") + assert.Contains(t, errEvent.Error, "tools") + assert.Equal(t, ErrorCodeModelError, errEvent.Code) +} diff --git a/pkg/tui/components/messages/messages.go b/pkg/tui/components/messages/messages.go index 161f63603a..087828677a 100644 --- a/pkg/tui/components/messages/messages.go +++ b/pkg/tui/components/messages/messages.go @@ -105,6 +105,11 @@ type Model interface { AdjustBottomSlack(delta int) VisualGeneration() uint64 + // MessageTypeCount returns how many messages currently in the list have + // the given type. Read-only introspection for callers (e.g. tests) that + // need to observe real list state rather than trust a call was made. + MessageTypeCount(t types.MessageType) int + // IsScrollbarDragging returns true when the scrollbar thumb is being dragged. IsScrollbarDragging() bool @@ -1902,6 +1907,18 @@ func (m *model) RemoveSpinner() { m.removeSpinner() } +// MessageTypeCount returns how many messages currently in the list have the +// given type, by scanning the real message slice — never a call counter. +func (m *model) MessageTypeCount(t types.MessageType) int { + count := 0 + for _, msg := range m.messages { + if msg.Type == t { + count++ + } + } + return count +} + func (m *model) removeSpinner() { if len(m.messages) == 0 { return diff --git a/pkg/tui/page/chat/image_output_guard_integration_test.go b/pkg/tui/page/chat/image_output_guard_integration_test.go new file mode 100644 index 0000000000..cbd102cb6d --- /dev/null +++ b/pkg/tui/page/chat/image_output_guard_integration_test.go @@ -0,0 +1,202 @@ +package chat + +import ( + "fmt" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + tea "charm.land/bubbletea/v2" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/app" + "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/gemini" + "github.com/docker/docker-agent/pkg/model/provider/options" + "github.com/docker/docker-agent/pkg/modelerrors" + "github.com/docker/docker-agent/pkg/runtime" + "github.com/docker/docker-agent/pkg/session" + "github.com/docker/docker-agent/pkg/tools" + "github.com/docker/docker-agent/pkg/tui/animation" + "github.com/docker/docker-agent/pkg/tui/components/messages" + "github.com/docker/docker-agent/pkg/tui/service" + "github.com/docker/docker-agent/pkg/tui/types" +) + +// recordingMessages wraps the real [messages.Model], counting the calls this +// test cares about while forwarding every call to the embedded +// implementation so its real state (the spinner entry, the rendered error) +// mutates for real. Mirrors the recordingSidebar pattern in +// agent_switching_test.go. +type recordingMessages struct { + messages.Model + + assistantMessageCalls int + errorMessages []string + appendCalls int + appendReasoningCalls int + removeSpinnerCalls int +} + +func (r *recordingMessages) AddAssistantMessage(sender, label string) tea.Cmd { + r.assistantMessageCalls++ + return r.Model.AddAssistantMessage(sender, label) +} + +func (r *recordingMessages) AddErrorMessage(content string) tea.Cmd { + r.errorMessages = append(r.errorMessages, content) + return r.Model.AddErrorMessage(content) +} + +func (r *recordingMessages) AppendToLastMessage(agentName, content string) tea.Cmd { + r.appendCalls++ + return r.Model.AppendToLastMessage(agentName, content) +} + +func (r *recordingMessages) AppendReasoning(agentName, content string) tea.Cmd { + r.appendReasoningCalls++ + return r.Model.AppendReasoning(agentName, content) +} + +// RemoveSpinner counts calls and forwards to the real implementation, whose +// own removeSpinner is a no-op (skips invalidateView) when the last message +// isn't a spinner — so a VisualGeneration bump around the call proves a +// spinner actually existed and was removed, not just that the method fired. +func (r *recordingMessages) RemoveSpinner() { + r.removeSpinnerCalls++ + r.Model.RemoveSpinner() +} + +// TestImageOutputGuard_RuntimeToTUI_RejectsBeforeDispatchAndClearsSpinner +// drives the gateway 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 +// surface a fatal model error (pkg/runtime/loop_steps.go's +// handleStreamError: modelerrors.FormatError + ErrorWithCodeForSession) into +// a real chatPage.handleRuntimeEvent (pkg/tui/page/chat/runtime_events.go). +// +// This package cannot build a real [runtime.LocalRuntime] run itself: TUI +// code must not import pkg/team (see e2e/dependencies_test.go's "TUI musn't +// know about teams"), which a real run requires. That half — the guard's +// rejection surviving the full fallback/loop machinery unchanged, through +// RunStream — is covered by +// TestRunStream_ImageOutputGuard_RejectsBeforeDispatch in pkg/runtime. The +// two tests meet at the same seam: [modelerrors.FormatError] and the +// [runtime.ErrorEvent] it feeds, exercised here with the constructors +// production code actually calls, not a hand-rolled message. +// +// It proves, in one flow, against the real message-list state (never a call +// counter alone): zero HTTP requests reach the provider (the guard rejects +// before dispatch); StreamStartedEvent leaves exactly one real spinner +// message; handling the actual ErrorEvent through AddErrorMessage removes +// that spinner and adds the fixed safe error before StreamStoppedEvent ever +// runs; and the later StreamStoppedEvent's call to the exported RemoveSpinner +// finds nothing left to remove — a no-op, not the mechanism that cleared the +// spinner. No assistant text or reasoning is ever appended, i.e. no silent +// "success" alongside the error. +func TestImageOutputGuard_RuntimeToTUI_RejectsBeforeDispatchAndClearsSpinner(t *testing.T) { + t.Parallel() + + var providerCalls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + providerCalls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + payload := `{"candidates":[{"content":{"parts":[{"text":"hi"}],"role":"model"},"finishReason":"STOP","index":0}]}` + _, _ = fmt.Fprintf(w, "data: %s\n\n", payload) + })) + defer server.Close() + + 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 := gemini.NewClient(t.Context(), cfg, env, options.WithGateway(server.URL)) + require.NoError(t, err) + + // 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. + readFileTool := tools.Tool{ + Name: "read_file", + Description: "reads a file from disk", + Parameters: map[string]any{"type": "object"}, + } + stream, dispatchErr := client.CreateChatCompletionStream(t.Context(), []chat.Message{ + {Role: chat.MessageRoleUser, Content: "draw a cat"}, + }, []tools.Tool{readFileTool}) + require.Nil(t, stream) + require.Error(t, dispatchErr) + + var incompatible *gemini.ImageOutputRequestIncompatibleError + require.ErrorAs(t, dispatchErr, &incompatible, "expected the guard's rejection error") + require.Zero(t, providerCalls.Load(), "the guard must reject before any request reaches the provider") + + // Build the exact production event sequence: same constructors and the + // same modelerrors.FormatError call pkg/runtime/loop_steps.go's + // handleStreamError makes when a model call fails fatally. + const ( + agentName = "root" + sessionID = "sess-image-output-guard" + ) + visibleError := modelerrors.FormatError(dispatchErr) + + sessForPage := session.New() + p := New(animation.NewRuntime(), t.Context(), + app.New(t.Context(), queueTestRuntime{}, sessForPage), + service.NewSessionState(sessForPage)).(*chatPage) + + rec := &recordingMessages{Model: p.messages} + p.messages = rec + + handled, _ := p.handleRuntimeEvent(runtime.StreamStarted(sessionID, agentName)) + require.True(t, handled, "expected StreamStartedEvent to be a recognized runtime event") + assert.Equal(t, 1, rec.assistantMessageCalls, + "the stream-started spinner must have been requested exactly once") + require.Equal(t, 1, rec.MessageTypeCount(types.MessageTypeSpinner), + "a real spinner message must exist in the list right after StreamStartedEvent") + assert.Zero(t, rec.removeSpinnerCalls, "no spinner removal is expected before the stream stops") + + handled, _ = p.handleRuntimeEvent(runtime.ErrorWithCodeForSession(sessionID, runtime.ErrorCodeModelError, visibleError)) + require.True(t, handled, "expected ErrorEvent to be a recognized runtime event") + + // The spinner must already be gone here, before StreamStoppedEvent ever + // runs: this is the production AddErrorMessage -> internal removeSpinner + // path (pkg/tui/components/messages/messages.go), not the later + // StreamStoppedEvent -> exported RemoveSpinner cleanup. If + // AddErrorMessage's internal removal were disabled, this assertion would + // fail while the spinner count stayed at 1. + require.Zero(t, rec.MessageTypeCount(types.MessageTypeSpinner), + "the actual ErrorEvent must remove the real spinner via AddErrorMessage before the stream stops") + require.Equal(t, 1, rec.MessageTypeCount(types.MessageTypeError), + "the guard's error must be added as a real error message") + assert.Zero(t, rec.removeSpinnerCalls, + "the exported RemoveSpinner must not have been invoked yet; removal so far is AddErrorMessage's internal one") + + handled, _ = p.handleRuntimeEvent(runtime.StreamStopped(sessionID, agentName, "error")) + require.True(t, handled, "expected StreamStoppedEvent to be a recognized runtime event") + + assert.Equal(t, 1, rec.removeSpinnerCalls, + "the outermost stream-stop cleanup still calls the exported RemoveSpinner once") + assert.Zero(t, rec.MessageTypeCount(types.MessageTypeSpinner), + "the spinner count must stay at zero across StreamStoppedEvent: its RemoveSpinner call is a no-op here, "+ + "not the mechanism that removed the spinner") + + require.Len(t, rec.errorMessages, 1, "the guard's error must reach the message list exactly once") + assert.Equal(t, visibleError, rec.errorMessages[0], + "the TUI must show the exact same fixed error text modelerrors.FormatError produced") + assert.Contains(t, rec.errorMessages[0], "output_capabilities.image") + assert.Contains(t, rec.errorMessages[0], "tools") + assert.Zero(t, rec.appendCalls, "no assistant text may be appended after a rejected request") + assert.Zero(t, rec.appendReasoningCalls, "no reasoning may be appended after a rejected request") + assert.False(t, p.working, "the chat page must not be left in a working state after the stream stops") +}