From 19e867693ff5d9012d8a1225e3d537d77fb33b19 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 Google requests Request TEXT+IMAGE on supported Google surfaces when an explicit override or models.dev enables image output. Before ordinary dispatch, reject custom function tools and structured output while allowing server-side built-ins. Title generation and compaction omit image response modalities and bypass the guard; omission does not explicitly force a TEXT modality. Use catalogue tool_call metadata to distinguish unsupported tools from an image-output combination conflict. Keep unknown support conservative and make the diagnostic accurate for either override or catalogue resolution. Remove dead ToolConfig when no custom tools remain and keep the TUI inline error visible without a misleading footer toast. Cover provider, runtime, and TUI paths, explicit false, unknown support and no-dispatch failures. --- docs/providers/google/index.md | 28 ++ pkg/model/provider/base/base.go | 5 + pkg/model/provider/base/base_test.go | 16 + pkg/model/provider/gemini/client.go | 25 +- pkg/model/provider/gemini/diagnostics.go | 16 +- pkg/model/provider/gemini/diagnostics_test.go | 55 ++- .../provider/gemini/image_output_guard.go | 79 ++++ .../gemini/image_output_guard_test.go | 432 ++++++++++++++++++ pkg/modelinfo/modelinfo.go | 36 ++ pkg/modelinfo/output_capabilities_test.go | 32 ++ .../image_output_guard_integration_test.go | 112 +++++ pkg/tui/components/messages/messages.go | 17 + .../image_output_guard_integration_test.go | 201 ++++++++ 13 files changed, 1034 insertions(+), 20 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 1bdc06e26..ff5958e19 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/base/base.go b/pkg/model/provider/base/base.go index 12edfe3ad..1520276b5 100644 --- a/pkg/model/provider/base/base.go +++ b/pkg/model/provider/base/base.go @@ -79,6 +79,11 @@ func (c *Config) CapsOverride() *modelinfo.CapsOverride { return &modelinfo.CapsOverride{Image: caps.Image, PDF: caps.PDF, Audio: caps.Audio, Video: caps.Video} } +// ToolCallSupport resolves the model's tool-call capability from models.dev. +func (c *Config) ToolCallSupport(ctx context.Context) modelinfo.ToolCallSupport { + return modelinfo.ResolveToolCallSupport(ctx, c.ModelOptions.ModelsDevStore(), c.ID()) +} + // ImageOutputEnabled resolves the model's image-output capability from its // explicit tri-state configuration and, when unset, the models.dev catalogue. func (c *Config) ImageOutputEnabled(ctx context.Context) bool { diff --git a/pkg/model/provider/base/base_test.go b/pkg/model/provider/base/base_test.go index b481d3514..591e6c17a 100644 --- a/pkg/model/provider/base/base_test.go +++ b/pkg/model/provider/base/base_test.go @@ -59,6 +59,22 @@ func TestConfigCapsOverride(t *testing.T) { }) } +func TestConfigToolCallSupport(t *testing.T) { + t.Parallel() + + store := modelsdev.NewDatabaseStore(&modelsdev.Database{Providers: map[string]modelsdev.Provider{ + "google": {Models: map[string]modelsdev.Model{ + "tool-model": {ToolCall: true}, + }}, + }}) + cfg := Config{ + ModelConfig: latest.ModelConfig{Provider: "google", Model: "tool-model"}, + ModelOptions: options.Apply(options.WithModelsDevStore(store)), + } + + assert.Equal(t, modelinfo.ToolCallSupported, cfg.ToolCallSupport(t.Context())) +} + func TestConfigImageOutputEnabled(t *testing.T) { t.Parallel() diff --git a/pkg/model/provider/gemini/client.go b/pkg/model/provider/gemini/client.go index 1a50f7800..f009b67b0 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(ctx, imageOutputEnabled, config, 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 411dade66..5029799ad 100644 --- a/pkg/model/provider/gemini/diagnostics.go +++ b/pkg/model/provider/gemini/diagnostics.go @@ -6,8 +6,8 @@ import ( "google.golang.org/genai" ) -// apiSurface* classifies which backend/transport a Client talks to. Used -// only for diagnostics; never derived from or containing request content. +// These transport categories gate response modalities and label diagnostics; +// they never contain request content. const ( apiSurfaceGeminiAPI = "gemini_api" apiSurfaceVertexAI = "vertex_ai" @@ -55,12 +55,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 } @@ -68,7 +64,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) @@ -83,6 +79,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 b25bd56f9..8c0564ab7 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,6 @@ 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) assert.False(t, shape.OutputCapabilityEnabled) } @@ -93,7 +90,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 +111,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 +132,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 +153,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 +165,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 +226,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 000000000..9801a5dcc --- /dev/null +++ b/pkg/model/provider/gemini/image_output_guard.go @@ -0,0 +1,79 @@ +package gemini + +import ( + "context" + "fmt" + "strings" + + "google.golang.org/genai" + + "github.com/docker/docker-agent/pkg/modelinfo" +) + +// 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" + 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) or structured output. Gemini server-side built-in +// tools remain allowed. Title generation and compaction are always text-only +// and bypass this guard. +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 + ToolCallSupport modelinfo.ToolCallSupport +} + +func (e *ImageOutputRequestIncompatibleError) Error() string { + if len(e.Incompatibilities) == 1 && e.Incompatibilities[0] == imageOutputIncompatibleTools { + switch e.ToolCallSupport { + case modelinfo.ToolCallUnsupported: + return "this model does not support tool calls; use a tool-capable model or remove tools from the request" + case modelinfo.ToolCallSupported: + return "this model supports tool calls, but not while image output is enabled; use a separate model or request for that combination" + } + } + + names := make([]string, len(e.Incompatibilities)) + for i, c := range e.Incompatibilities { + names[i] = string(c) + } + return fmt.Sprintf( + "image output is enabled for this model (by output_capabilities.image or models.dev) and is incompatible with %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(ctx context.Context, imageOutputEnabled bool, config *genai.GenerateContentConfig, requestTools int) error { + if !imageOutputEnabled || c.ModelOptions.GeneratingTitle() || c.ModelOptions.Compacting() { + return nil + } + + var incompatibilities []imageOutputIncompatibility + if requestTools > 0 { + incompatibilities = append(incompatibilities, imageOutputIncompatibleTools) + } + if config.ResponseMIMEType != "" || config.ResponseJsonSchema != nil { + incompatibilities = append(incompatibilities, imageOutputIncompatibleStructuredOutput) + } + if len(incompatibilities) == 0 { + return nil + } + incompatible := &ImageOutputRequestIncompatibleError{Incompatibilities: incompatibilities} + if len(incompatibilities) == 1 && incompatibilities[0] == imageOutputIncompatibleTools { + incompatible.ToolCallSupport = c.ToolCallSupport(ctx) + } + return incompatible +} 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 000000000..1b8e5285c --- /dev/null +++ b/pkg/model/provider/gemini/image_output_guard_test.go @@ -0,0 +1,432 @@ +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/modelinfo" + "github.com/docker/docker-agent/pkg/modelsdev" + "github.com/docker/docker-agent/pkg/tools" +) + +// TestCheckImageOutputRequestCompatibility_GateConditions exhaustively +// 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 an ordinary outgoing request also carries custom function +// tools or structured output. Server-side built-ins stay allowed. +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 + generatingTitle bool + compacting bool + wantReject []imageOutputIncompatibility + }{ + {name: "gateway declared true, no extras: allowed", apiSurface: apiSurfaceGateway, declared: declaredTrue}, + {name: "title generation bypasses custom tools", apiSurface: apiSurfaceGateway, declared: declaredTrue, requestTools: 1, generatingTitle: true}, + {name: "compaction bypasses structured output", apiSurface: apiSurfaceGateway, declared: declaredTrue, structured: true, compacting: true}, + {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 + 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, + wantReject: []imageOutputIncompatibility{imageOutputIncompatibleTools}, + }, + { + name: "gateway declared true + built-in tool: allowed", + apiSurface: apiSurfaceGateway, declared: declaredTrue, builtInTools: []*genai.Tool{{GoogleSearch: &genai.GoogleSearch{}}}, + }, + { + 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, + } + if tt.generatingTitle { + client.ModelOptions = options.Apply(options.WithGeneratingTitle()) + } + if tt.compacting { + client.ModelOptions = options.Apply(options.WithCompacting()) + } + config := &genai.GenerateContentConfig{} + if tt.structured { + config.ResponseMIMEType = "application/json" + } + _ = tt.builtInTools + + err := client.checkImageOutputRequestCompatibility(t.Context(), tt.declared != nil && tt.declared.Image != nil && *tt.declared.Image, config, 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_ToolMessages(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + incompatibilities []imageOutputIncompatibility + toolCallSupport modelinfo.ToolCallSupport + want string + }{ + { + name: "model does not support tools", + incompatibilities: []imageOutputIncompatibility{imageOutputIncompatibleTools}, + toolCallSupport: modelinfo.ToolCallUnsupported, + want: "this model does not support tool calls; use a tool-capable model or remove tools from the request", + }, + { + name: "tool-capable image model", + incompatibilities: []imageOutputIncompatibility{imageOutputIncompatibleTools}, + toolCallSupport: modelinfo.ToolCallSupported, + want: "this model supports tool calls, but not while image output is enabled; use a separate model or request for that combination", + }, + { + name: "unknown support keeps conservative message", + incompatibilities: []imageOutputIncompatibility{imageOutputIncompatibleTools}, + toolCallSupport: modelinfo.ToolCallSupportUnknown, + want: "image output is enabled for this model (by output_capabilities.image or models.dev) and is incompatible with tools in the same request; use a separate model or request for that combination", + }, + { + name: "structured output ignores tool support", + incompatibilities: []imageOutputIncompatibility{imageOutputIncompatibleStructuredOutput}, + toolCallSupport: modelinfo.ToolCallUnsupported, + want: "image output is enabled for this model (by output_capabilities.image or models.dev) and is incompatible with structured output in the same request; use a separate model or request for that combination", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + err := (&ImageOutputRequestIncompatibleError{ + Incompatibilities: tt.incompatibilities, + ToolCallSupport: tt.toolCallSupport, + }).Error() + assert.Equal(t, tt.want, err) + }) + } +} + +func TestCheckImageOutputRequestCompatibility_ToolCallSupport(t *testing.T) { + t.Parallel() + + store := modelsdev.NewDatabaseStore(&modelsdev.Database{Providers: map[string]modelsdev.Provider{ + "google": {Models: map[string]modelsdev.Model{ + "tool-model": {ToolCall: true}, + "no-tool-model": {ToolCall: false}, + }}, + }}) + + tests := []struct { + name string + model string + want modelinfo.ToolCallSupport + }{ + {name: "catalogue reports no tool calls", model: "no-tool-model", want: modelinfo.ToolCallUnsupported}, + {name: "catalogue reports tool calls", model: "tool-model", want: modelinfo.ToolCallSupported}, + {name: "missing catalogue model stays unknown", model: "missing", want: modelinfo.ToolCallSupportUnknown}, + } + + 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}, + ModelOptions: options.Apply(options.WithModelsDevStore(store)), + }} + err := client.checkImageOutputRequestCompatibility(t.Context(), true, &genai.GenerateContentConfig{}, 1) + + var incompatible *ImageOutputRequestIncompatibleError + require.ErrorAs(t, err, &incompatible) + assert.Equal(t, tt.want, incompatible.ToolCallSupport) + }) + } +} + +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, "models.dev") + assert.NotContains(t, msg, "is configured for image output") + 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{imageOutputIncompatibleTools}} + + 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, "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("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 and tools/structured +// output when the declaration is false or missing. +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 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"}}}) + require.Nil(t, stream) + require.Error(t, err) + assert.Zero(t, counter.calls.Load()) + }) +} diff --git a/pkg/modelinfo/modelinfo.go b/pkg/modelinfo/modelinfo.go index 2c3695ced..17ae2630a 100644 --- a/pkg/modelinfo/modelinfo.go +++ b/pkg/modelinfo/modelinfo.go @@ -713,6 +713,42 @@ func LoadCaps(ctx context.Context, store *modelsdev.Store, id modelsdev.ID) Mode return capsFromModalities(model.Modalities.Input) } +// ToolCallSupport is the models.dev catalogue's tri-state tool-call capability. +type ToolCallSupport uint8 + +const ( + ToolCallSupportUnknown ToolCallSupport = iota + ToolCallUnsupported + ToolCallSupported +) + +// ResolveToolCallSupport reports whether models.dev says a model supports tool +// calls. Missing catalogue data remains unknown rather than being treated as a +// negative capability claim. +func ResolveToolCallSupport(ctx context.Context, store *modelsdev.Store, id modelsdev.ID) ToolCallSupport { + if store == nil { + return ToolCallSupportUnknown + } + + ctx, cancel := context.WithTimeout(ctx, loadCapsTimeout) + defer cancel() + + model, err := store.GetModel(ctx, id) + if err != nil { + if ctx.Err() != nil { + slog.WarnContext(ctx, "modelinfo: models.dev tool-call lookup timed out, leaving support unknown", + "model", id.String(), "timeout", loadCapsTimeout) + } else { + warnCapsLookupMiss(ctx, id, err) + } + return ToolCallSupportUnknown + } + if model.ToolCall { + return ToolCallSupported + } + return ToolCallUnsupported +} + // ResolveOutputImage applies an explicit image-output override when present; // otherwise it derives support from the models.dev output modalities. Missing // catalogue data conservatively disables image output. diff --git a/pkg/modelinfo/output_capabilities_test.go b/pkg/modelinfo/output_capabilities_test.go index 3090a1529..b77e0fd13 100644 --- a/pkg/modelinfo/output_capabilities_test.go +++ b/pkg/modelinfo/output_capabilities_test.go @@ -9,6 +9,38 @@ import ( "github.com/docker/docker-agent/pkg/modelsdev" ) +func TestResolveToolCallSupport(t *testing.T) { + t.Parallel() + + store := modelsdev.NewDatabaseStore(&modelsdev.Database{Providers: map[string]modelsdev.Provider{ + "google": {Models: map[string]modelsdev.Model{ + "tool-model": {ToolCall: true}, + "no-tool-model": {ToolCall: false}, + }}, + }}) + + tests := []struct { + name string + store *modelsdev.Store + model string + want ToolCallSupport + }{ + {name: "catalogue true", store: store, model: "tool-model", want: ToolCallSupported}, + {name: "catalogue false", store: store, model: "no-tool-model", want: ToolCallUnsupported}, + {name: "missing model", store: store, model: "missing", want: ToolCallSupportUnknown}, + {name: "nil store", model: "tool-model", want: ToolCallSupportUnknown}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + got := ResolveToolCallSupport(t.Context(), tt.store, modelsdev.NewID("google", tt.model)) + assert.Equal(t, tt.want, got) + }) + } +} + func TestResolveOutputImage(t *testing.T) { t.Parallel() 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 000000000..681331a86 --- /dev/null +++ b/pkg/runtime/image_output_guard_integration_test.go @@ -0,0 +1,112 @@ +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 +// 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 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 ea2198b50..ab700fd25 100644 --- a/pkg/tui/components/messages/messages.go +++ b/pkg/tui/components/messages/messages.go @@ -109,6 +109,11 @@ type Model interface { // VisualGeneration increments only when Update changes rendered output. 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 @@ -2220,6 +2225,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 000000000..55f93aefb --- /dev/null +++ b/pkg/tui/page/chat/image_output_guard_integration_test.go @@ -0,0 +1,201 @@ +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 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 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") +}