From 99fece268be5b2f4870a15566554226bd676d906 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= <arnaud.heritier@docker.com>
Date: Tue, 6 Oct 2026 15:26:33 +0200
Subject: [PATCH] fix(chatgpt): infer image input from OpenAI catalogue

---
 docs/configuration/models/index.md            |   8 +-
 docs/providers/chatgpt/index.md               |  27 +++++
 examples/README.md                            |   1 +
 examples/chatgpt-images.png                   | Bin 0 -> 306 bytes
 examples/chatgpt-images.yaml                  |  20 ++++
 .../provider/capability_override_test.go      |  16 +++
 pkg/model/provider/openai/chatgpt_test.go     |  93 ++++++++++++++++++
 pkg/modelinfo/catalog_caps.go                 |  27 +++++
 pkg/modelinfo/modelinfo.go                    |   9 +-
 pkg/modelinfo/resolve_caps_test.go            |  69 +++++++++++++
 pkg/runtime/transforms.go                     |   6 +-
 pkg/runtime/transforms_test.go                |  75 ++++++++++++++
 12 files changed, 344 insertions(+), 7 deletions(-)
 create mode 100644 examples/chatgpt-images.png
 create mode 100644 examples/chatgpt-images.yaml
 create mode 100644 pkg/modelinfo/catalog_caps.go

diff --git a/docs/configuration/models/index.md b/docs/configuration/models/index.md
index 2a82cdc0e..6a9ab65eb 100644
--- a/docs/configuration/models/index.md
+++ b/docs/configuration/models/index.md
@@ -90,7 +90,9 @@ For custom OpenAI-compatible providers, local models (Ollama, DMR), and any
 model the built-in catalogue does not describe, Docker Agent cannot
 auto-detect whether the endpoint accepts image, PDF, audio, or video
 attachments. When the model is absent from the catalogue, Docker Agent logs a
-diagnostic and falls back to text-only, silently dropping attachments.
+diagnostic and falls back to text-only, silently dropping attachments. For the
+`chatgpt` provider, a missing entry first falls back to the matching `openai`
+entry for **image input only**; PDF, audio, and video are not inferred.
 
 Declare `capabilities` to make the model's attachment support authoritative
 and skip the catalogue lookup entirely:
@@ -130,8 +132,8 @@ models:
 
 The flags must match what the endpoint actually accepts. Claiming a modality
 that the endpoint does not support leads to a provider-side API error. When
-`capabilities` is omitted the behaviour is unchanged (catalogue lookup then
-conservative text-only fallback).
+`capabilities` is omitted, input support comes from the catalogue, the scoped
+ChatGPT image fallback above, or conservative text-only fallback.
 
 ### Unsupported media is stripped before the call
 
diff --git a/docs/providers/chatgpt/index.md b/docs/providers/chatgpt/index.md
index ed5025005..7394783fc 100644
--- a/docs/providers/chatgpt/index.md
+++ b/docs/providers/chatgpt/index.md
@@ -101,6 +101,33 @@ The effort picker exposes Low/Medium/High/XHigh/Max on the GPT-5.6 family
   there. Client-side sampling parameters (`temperature`, `top_p`,
   `max_tokens`) are not supported by the backend and are dropped.
 
+## Image Input
+
+For image-capable models such as `gpt-6.1-sol`, you can attach an image:
+
+```bash
+docker agent run examples/chatgpt-images.yaml --exec \
+  --attach examples/chatgpt-images.png "Describe the shapes and colors."
+```
+
+The sample includes a red square centered on a blue background. Running it
+requires ChatGPT sign-in and access to the selected model. Image input for
+`gpt-6.1-sol` is inferred from the matching OpenAI catalogue entry; no
+`capabilities` override is needed.
+
+When models.dev has no direct `chatgpt/<model>` entry, Docker Agent uses the
+matching `openai/<model>` entry for **image input only**. This applies both to
+`--attach` images and image content returned by MCP tools (for example,
+screenshots). Tool images are sent in a follow-up user message alongside the
+text function-call output. PDF, audio, and video support, prices, context limits,
+tool support, output capabilities, and `max_tokens` are not inherited.
+A direct ChatGPT catalogue entry takes precedence; unknown models remain
+text-only.
+
+To intentionally override catalogue capabilities or declare support for an
+unknown model, see [Attachment Capability Overrides](../../configuration/models/index.md#attachment-capability-overrides).
+Image input support does not enable image generation or editing.
+
 ## Setting the Token Explicitly
 
 `CHATGPT_OAUTH_TOKEN` can also be set like any other credential (shell
diff --git a/examples/README.md b/examples/README.md
index 1f6196985..de8491267 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -210,6 +210,7 @@ remote MCP endpoints.
 | [`cloudflare-workers-ai.yaml`](cloudflare-workers-ai.yaml) | Cloudflare Workers AI edge-hosted open models. |
 | [`cloudflare-ai-gateway.yaml`](cloudflare-ai-gateway.yaml) | Cloudflare AI Gateway multi-provider router. |
 | [`grok.yaml`](grok.yaml) | xAI Grok model. |
+| [`chatgpt-images.yaml`](chatgpt-images.yaml) + [`chatgpt-images.png`](chatgpt-images.png) | ChatGPT image input with an attached sample image and automatic capability detection. |
 | [`github-copilot.yaml`](github-copilot.yaml) | GitHub Copilot models via OAuth device-flow. |
 | [`fallback_models.yaml`](fallback_models.yaml) | Automatic fallback to a secondary model when the primary fails. |
 | [`model_picker.yaml`](model_picker.yaml) | Lets the agent itself swap to a stronger model mid-conversation. |
diff --git a/examples/chatgpt-images.png b/examples/chatgpt-images.png
new file mode 100644
index 0000000000000000000000000000000000000000..21e47b4550b822d428d7a0a963355b66ab64c58c
GIT binary patch
literal 306
zcmeAS@N?(olHy`uVBq!ia0vp^2_VeD1SE5RJ;(=AKRsO>Ln`LHy<o`2V90aC;s5{3
zp7pM)8m2fexV0udEll>Y_zB9<?h=mwkAH_bFtUmSAQKOjz#@y7xHJ|t;2;z{k#&7x
iyHHQLdmf9gxIORvJ-*rAiAKOcVDNPHb6Mw<&;$TNf`VcI

literal 0
HcmV?d00001

diff --git a/examples/chatgpt-images.yaml b/examples/chatgpt-images.yaml
new file mode 100644
index 000000000..5a18ddeff
--- /dev/null
+++ b/examples/chatgpt-images.yaml
@@ -0,0 +1,20 @@
+# ChatGPT image input with automatic capability detection.
+# Requires ChatGPT sign-in and access to the selected model.
+# Image input for gpt-6.1-sol is inferred from the matching OpenAI catalogue entry.
+#
+# From the repository root:
+#   docker agent run examples/chatgpt-images.yaml --exec \
+#     --attach examples/chatgpt-images.png "Describe the shapes and colors."
+# The sample image is a red square centered on a blue background.
+# With a source build, use ./bin/docker-agent instead of docker agent.
+
+models:
+  vision:
+    provider: chatgpt
+    model: gpt-6.1-sol
+
+agents:
+  root:
+    model: vision
+    description: ChatGPT image-input example
+    instruction: Describe the attached image accurately and concisely.
diff --git a/pkg/model/provider/capability_override_test.go b/pkg/model/provider/capability_override_test.go
index 34fa73732..d96e07be6 100644
--- a/pkg/model/provider/capability_override_test.go
+++ b/pkg/model/provider/capability_override_test.go
@@ -74,3 +74,19 @@ func TestCapabilityOverride_SurvivesProviderConstruction(t *testing.T) {
 		assert.Equal(t, &modelinfo.CapsOverride{Image: true, PDF: true, Audio: true, Video: true}, got)
 	})
 }
+
+func TestCapabilityOverride_ChatGPT(t *testing.T) {
+	t.Parallel()
+
+	env := environment.NewMapEnvProvider(map[string]string{"CHATGPT_OAUTH_TOKEN": "test-token"})
+	for _, image := range []bool{false, true} {
+		cfg := &latest.ModelConfig{
+			Provider: "chatgpt", Model: "gpt-6.1-sol",
+			Capabilities: &latest.CapabilitiesConfig{Image: image},
+		}
+		p, err := fullTestRegistry().New(t.Context(), cfg, env)
+		require.NoError(t, err)
+		bc := p.BaseConfig()
+		assert.Equal(t, &modelinfo.CapsOverride{Image: image}, bc.CapsOverride())
+	}
+}
diff --git a/pkg/model/provider/openai/chatgpt_test.go b/pkg/model/provider/openai/chatgpt_test.go
index 14af2d5d6..1630848cc 100644
--- a/pkg/model/provider/openai/chatgpt_test.go
+++ b/pkg/model/provider/openai/chatgpt_test.go
@@ -19,6 +19,8 @@ import (
 	"github.com/docker/docker-agent/pkg/chatgpt"
 	"github.com/docker/docker-agent/pkg/config/latest"
 	"github.com/docker/docker-agent/pkg/environment"
+	"github.com/docker/docker-agent/pkg/model/provider/options"
+	"github.com/docker/docker-agent/pkg/modelsdev"
 	"github.com/docker/docker-agent/pkg/tools"
 )
 
@@ -260,3 +262,94 @@ func TestChatGPTFallsBackToStoredLogin(t *testing.T) {
 	assert.Equal(t, "Bearer "+token, got.header.Get("Authorization"))
 	assert.Equal(t, "acc_stored", got.header.Get("chatgpt-account-id"))
 }
+
+func TestChatGPTImageInputFromOpenAICatalog(t *testing.T) {
+	t.Parallel()
+
+	store := modelsdev.NewDatabaseStore(&modelsdev.Database{Providers: map[string]modelsdev.Provider{
+		"openai": {Models: map[string]modelsdev.Model{
+			"gpt-6.1-sol": {Modalities: modelsdev.Modalities{Input: []string{"text", "image", "pdf", "audio", "video"}}},
+		}},
+	}})
+	for _, tc := range []struct {
+		name      string
+		override  *latest.CapabilitiesConfig
+		wantImage bool
+	}{
+		{name: "no override", wantImage: true},
+		{name: "explicit image true", override: &latest.CapabilitiesConfig{Image: true}, wantImage: true},
+		{name: "explicit image false", override: &latest.CapabilitiesConfig{}},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			server, captured := startFakeCodexBackend(t)
+			cfg := &latest.ModelConfig{
+				Provider: "chatgpt", Model: "gpt-6.1-sol", BaseURL: server.URL,
+				TokenKey: chatgpt.TokenEnvVar, Capabilities: tc.override,
+			}
+			env := environment.NewMapEnvProvider(map[string]string{
+				chatgpt.TokenEnvVar: chatgptTestToken(t, "acc_images"),
+			})
+			client, err := NewClient(t.Context(), cfg, env, options.WithModelsDevStore(store))
+			require.NoError(t, err)
+
+			image := chat.MessagePart{Type: chat.MessagePartTypeDocument, Document: &chat.Document{
+				Name: "screenshot.png", MimeType: "image/png", Source: chat.DocumentSource{InlineData: []byte{1, 2, 3}},
+			}}
+			pdf := chat.MessagePart{Type: chat.MessagePartTypeDocument, Document: &chat.Document{
+				Name: "report.pdf", MimeType: "application/pdf", Source: chat.DocumentSource{InlineData: []byte("%PDF")},
+			}}
+			drainChatStream(t, client, []chat.Message{
+				{Role: chat.MessageRoleUser, MultiContent: []chat.MessagePart{
+					{Type: chat.MessagePartTypeText, Text: "describe this image"}, image, pdf,
+				}},
+				{Role: chat.MessageRoleAssistant, ToolCalls: []tools.ToolCall{{
+					ID: "call_image", Type: "function", Function: tools.FunctionCall{Name: "screenshot", Arguments: `{}`},
+				}}},
+				{Role: chat.MessageRoleTool, ToolCallID: "call_image", Content: "screenshot captured", MultiContent: []chat.MessagePart{image, pdf}},
+			})
+
+			input, ok := captured().body["input"].([]any)
+			require.True(t, ok)
+			wantLen := 3
+			if tc.wantImage {
+				wantLen++
+			}
+			require.Len(t, input, wantLen)
+			user := input[0].(map[string]any)
+			assert.Equal(t, "user", user["role"])
+			toolCall := input[1].(map[string]any)
+			assert.Equal(t, "function_call", toolCall["type"])
+			toolOutput := input[2].(map[string]any)
+			assert.Equal(t, "function_call_output", toolOutput["type"])
+			assert.Equal(t, "call_image", toolOutput["call_id"])
+			assert.Equal(t, "screenshot captured", toolOutput["output"])
+
+			var images int
+			for _, item := range input {
+				msg := item.(map[string]any)
+				content, _ := msg["content"].([]any)
+				for _, part := range content {
+					p := part.(map[string]any)
+					assert.NotEqual(t, "input_file", p["type"], "OpenAI PDF input must not be inherited")
+					if p["type"] == "input_image" {
+						assert.Equal(t, "user", msg["role"])
+						assert.Equal(t, "data:image/png;base64,AQID", p["image_url"])
+						images++
+					}
+				}
+			}
+			if tc.wantImage {
+				assert.Equal(t, 2, images, "attachment and tool image must both reach the backend")
+				followUp := input[3].(map[string]any)
+				assert.Equal(t, "user", followUp["role"])
+				content := followUp["content"].([]any)
+				require.Len(t, content, 2)
+				assert.Equal(t, "Attached content from tool result:", content[0].(map[string]any)["text"])
+				assert.Equal(t, "input_image", content[1].(map[string]any)["type"])
+			} else {
+				assert.Zero(t, images)
+			}
+		})
+	}
+}
diff --git a/pkg/modelinfo/catalog_caps.go b/pkg/modelinfo/catalog_caps.go
new file mode 100644
index 000000000..12de20e1f
--- /dev/null
+++ b/pkg/modelinfo/catalog_caps.go
@@ -0,0 +1,27 @@
+package modelinfo
+
+import (
+	"context"
+
+	"github.com/docker/docker-agent/pkg/modelsdev"
+)
+
+// ModelLookup is the catalogue lookup needed to resolve input capabilities.
+type ModelLookup interface {
+	GetModel(ctx context.Context, id modelsdev.ID) (*modelsdev.Model, error)
+}
+
+// AliasedCatalogCaps resolves image input support for ChatGPT from the matching
+// OpenAI entry after a direct catalogue miss. Other capabilities, limits, and
+// pricing are not inherited.
+func AliasedCatalogCaps(ctx context.Context, store ModelLookup, id modelsdev.ID) (ModelCapabilities, bool) {
+	if store == nil || !id.IsValid() || id.Provider != "chatgpt" || ctx.Err() != nil {
+		return ModelCapabilities{}, false
+	}
+	model, err := store.GetModel(ctx, modelsdev.NewID("openai", id.Model))
+	if err != nil || model == nil || ctx.Err() != nil {
+		return ModelCapabilities{}, false
+	}
+	caps := capsFromModalities(model.Modalities.Input)
+	return CapsWith(caps.SupportsImage(), false, false, false), true
+}
diff --git a/pkg/modelinfo/modelinfo.go b/pkg/modelinfo/modelinfo.go
index 111237624..4d7d56104 100644
--- a/pkg/modelinfo/modelinfo.go
+++ b/pkg/modelinfo/modelinfo.go
@@ -704,9 +704,9 @@ func warnCapsLookupMiss(ctx context.Context, id modelsdev.ID, cause error) {
 // model ID using the provided store.
 //
 // When the store is nil or the model is not found, LoadCaps returns a
-// conservative capability set that only allows text MIME types. A models.dev
-// miss is logged once per model via [warnCapsLookupMiss] so the degraded
-// behaviour is diagnosable rather than silent.
+// conservative capability set that only allows text MIME types. On a direct
+// catalogue miss, [AliasedCatalogCaps] may supply scoped input capabilities.
+// An unresolved models.dev miss is logged once per model via [warnCapsLookupMiss].
 //
 // The supplied ctx is wrapped with loadCapsTimeout so the lookup stays
 // cancellable with the caller and the underlying models.dev load is bounded.
@@ -722,6 +722,9 @@ func LoadCaps(ctx context.Context, store *modelsdev.Store, id modelsdev.ID) Mode
 
 	model, err := store.GetModel(ctx, id)
 	if err != nil {
+		if caps, ok := AliasedCatalogCaps(ctx, store, id); ok {
+			return caps
+		}
 		if ctx.Err() != nil {
 			slog.WarnContext(ctx, "modelinfo: models.dev lookup timed out, using conservative caps",
 				"model", id.String(), "timeout", loadCapsTimeout)
diff --git a/pkg/modelinfo/resolve_caps_test.go b/pkg/modelinfo/resolve_caps_test.go
index eb769dec6..bd2e26470 100644
--- a/pkg/modelinfo/resolve_caps_test.go
+++ b/pkg/modelinfo/resolve_caps_test.go
@@ -2,6 +2,7 @@ package modelinfo
 
 import (
 	"bytes"
+	"context"
 	"log/slog"
 	"strings"
 	"testing"
@@ -154,3 +155,71 @@ func TestResolveCapsFromModel(t *testing.T) {
 		})
 	}
 }
+
+func TestResolveCaps_ChatGPTCatalogFallback(t *testing.T) {
+	t.Parallel()
+
+	for _, tc := range []struct {
+		name     string
+		model    string
+		direct   []string
+		override *CapsOverride
+		want     ModelCapabilities
+	}{
+		{name: "image input only", model: "gpt-6.1-sol", want: CapsWith(true, false, false, false)},
+		{name: "unknown model", model: "unknown", want: CapsWith(false, false, false, false)},
+		{name: "direct entry wins", model: "gpt-6.1-sol", direct: []string{"text", "audio"}, want: CapsWith(false, false, true, false)},
+		{name: "explicit false wins", model: "gpt-6.1-sol", override: &CapsOverride{}, want: CapsWith(false, false, false, false)},
+		{name: "explicit true wins even for unknown", model: "unknown", override: &CapsOverride{Image: true, PDF: true}, want: CapsWith(true, true, false, false)},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			providers := map[string]modelsdev.Provider{
+				"openai": {Models: map[string]modelsdev.Model{
+					"gpt-6.1-sol": {Modalities: modelsdev.Modalities{Input: []string{"text", "image", "pdf", "audio", "video"}}},
+				}},
+			}
+			if tc.direct != nil {
+				providers["chatgpt"] = modelsdev.Provider{Models: map[string]modelsdev.Model{
+					tc.model: {Modalities: modelsdev.Modalities{Input: tc.direct}},
+				}}
+			}
+			store := modelsdev.NewDatabaseStore(&modelsdev.Database{Providers: providers})
+			got := ResolveCaps(t.Context(), store, modelsdev.NewID("chatgpt", tc.model), tc.override)
+			assert.Equal(t, tc.want, got)
+		})
+	}
+}
+
+func TestAliasedCatalogCaps_Misses(t *testing.T) {
+	t.Parallel()
+	store := modelsdev.NewDatabaseStore(&modelsdev.Database{Providers: map[string]modelsdev.Provider{
+		"openai": {Models: map[string]modelsdev.Model{
+			"vision":    {Modalities: modelsdev.Modalities{Input: []string{"image", "pdf"}}},
+			"text-only": {Modalities: modelsdev.Modalities{Input: []string{"text"}, Output: []string{"image"}}},
+		}},
+	}})
+	for _, id := range []modelsdev.ID{
+		modelsdev.NewID("chatgpt", "missing"),
+		modelsdev.NewID("other", "vision"),
+		modelsdev.NewID("", "vision"),
+		modelsdev.NewID("chatgpt", ""),
+	} {
+		caps, ok := AliasedCatalogCaps(t.Context(), store, id)
+		assert.False(t, ok, id.String())
+		assert.Equal(t, ModelCapabilities{}, caps)
+	}
+	caps, ok := AliasedCatalogCaps(t.Context(), store, modelsdev.NewID("chatgpt", "text-only"))
+	assert.True(t, ok)
+	assert.Equal(t, ModelCapabilities{}, caps, "output images must not imply image input")
+
+	caps, ok = AliasedCatalogCaps(t.Context(), nil, modelsdev.NewID("chatgpt", "vision"))
+	assert.False(t, ok)
+	assert.Equal(t, ModelCapabilities{}, caps)
+
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+	caps, ok = AliasedCatalogCaps(ctx, store, modelsdev.NewID("chatgpt", "vision"))
+	assert.False(t, ok)
+	assert.Equal(t, ModelCapabilities{}, caps)
+}
diff --git a/pkg/runtime/transforms.go b/pkg/runtime/transforms.go
index 76aae0e85..55a4e31d0 100644
--- a/pkg/runtime/transforms.go
+++ b/pkg/runtime/transforms.go
@@ -99,7 +99,11 @@ func (r *LocalRuntime) prepareMessagesForModel(
 	cfg := model.BaseConfig()
 	caps := modelinfo.ResolveCapsFromModel(catalogModel, cfg.CapsOverride())
 	if catalogModel == nil && cfg.CapsOverride() == nil {
-		caps = providerFallbackCaps(ctx, cfg.ModelConfig, modelID)
+		if aliased, ok := modelinfo.AliasedCatalogCaps(ctx, r.modelsStore, modelID); ok {
+			caps = aliased
+		} else {
+			caps = providerFallbackCaps(ctx, cfg.ModelConfig, modelID)
+		}
 	}
 	return r.applyBeforeLLMCallTransforms(ctx, sess, a, modelID.String(), &caps, msgs)
 }
diff --git a/pkg/runtime/transforms_test.go b/pkg/runtime/transforms_test.go
index beb7af5d3..dc3beeb05 100644
--- a/pkg/runtime/transforms_test.go
+++ b/pkg/runtime/transforms_test.go
@@ -1161,3 +1161,78 @@ func (p *queueRecordingProvider) CreateChatCompletionStream(ctx context.Context,
 	p.calls++
 	return p.recordingMsgProvider.CreateChatCompletionStream(ctx, msgs, tls)
 }
+
+func TestPrepareMessagesForModel_ChatGPTImageFallback(t *testing.T) {
+	t.Parallel()
+
+	for _, tc := range []struct {
+		name      string
+		model     string
+		direct    []string
+		override  *latest.CapabilitiesConfig
+		wantImage bool
+	}{
+		{name: "OpenAI image input", model: "gpt-6.1-sol", wantImage: true},
+		{name: "unknown stays text only", model: "unknown"},
+		{name: "direct entry wins", model: "gpt-6.1-sol", direct: []string{"text"}},
+		{name: "explicit false wins", model: "gpt-6.1-sol", override: &latest.CapabilitiesConfig{}},
+		{name: "explicit true wins", model: "unknown", override: &latest.CapabilitiesConfig{Image: true}, wantImage: true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			prov := &recordingMsgProvider{
+				mockProvider: mockProvider{id: "chatgpt/" + tc.model},
+				baseConfig: base.Config{ModelConfig: latest.ModelConfig{
+					Provider: "chatgpt", Model: tc.model, Capabilities: tc.override,
+				}},
+			}
+			a := agent.New("root", "instructions", agent.WithModel(prov))
+			providers := map[string]modelsdev.Provider{
+				"openai": {Models: map[string]modelsdev.Model{
+					"gpt-6.1-sol": {Modalities: modelsdev.Modalities{Input: []string{"text", "image", "pdf", "audio", "video"}}},
+				}},
+			}
+			if tc.direct != nil {
+				providers["chatgpt"] = modelsdev.Provider{Models: map[string]modelsdev.Model{
+					tc.model: {Modalities: modelsdev.Modalities{Input: tc.direct}},
+				}}
+			}
+			store := modelsdev.NewDatabaseStore(&modelsdev.Database{Providers: providers})
+			r, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(a)), WithModelStore(store))
+			require.NoError(t, err)
+
+			image := chat.MessagePart{Type: chat.MessagePartTypeDocument, Document: &chat.Document{
+				Name: "screenshot.png", MimeType: "image/png", Source: chat.DocumentSource{InlineData: []byte{1, 2, 3}},
+			}}
+			user := mixedMediaMsg()
+			user.MultiContent = append(user.MultiContent, image)
+			tool := chat.Message{Role: chat.MessageRoleTool, ToolCallID: "call_1", MultiContent: []chat.MessagePart{
+				{Type: chat.MessagePartTypeText, Text: "screenshot"},
+				image,
+				{Type: chat.MessagePartTypeImageURL, ImageURL: &chat.MessageImageURL{URL: "data:image/png;base64,AQID"}},
+			}}
+			got := r.prepareMessagesForModel(t.Context(), session.New(), a, prov, []chat.Message{user, tool})
+			require.Len(t, got, 2)
+			for _, msg := range got {
+				var images int
+				for _, part := range msg.MultiContent {
+					if part.Type == chat.MessagePartTypeImageURL {
+						images++
+					}
+					if part.Document != nil {
+						assert.NotContains(t, []string{"audio/wav", "video/mp4"}, part.Document.MimeType)
+						if part.Document.MimeType == "image/png" {
+							assert.Equal(t, image.Document, part.Document)
+							images++
+						}
+					}
+				}
+				if tc.wantImage {
+					assert.Equal(t, 2, images, "role %s", msg.Role)
+				} else {
+					assert.Zero(t, images, "role %s", msg.Role)
+				}
+			}
+		})
+	}
+}