From f68404b9f8316fdff81afe87fad3a74406f90946 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Fri, 31 Jul 2026 22:55:15 +0200 Subject: [PATCH 1/2] feat(#3996): name generated media via [media-file:] response markers Explicitly image-output-capable Gemini gateway requests now carry a system instruction asking the model to emit one exact "[media-file: relative/path]" line per generated image, echoing an explicitly requested name or choosing a meaningful kebab-case one, and to produce a single image unless variations are requested. handleStream filters those marker lines out of both the live event stream and the persisted assistant text with a chunk-split-safe, strictly bounded line parser, then pairs the extracted paths positionally onto the accumulated media blobs via MediaDelta.RequestedPath. Naming precedence stays marker -> provider display name -> generated-N, and marker paths remain untrusted input: they flow through the existing workspacemedia classification, MIME extension correction, escape elicitation/redirect, and never-overwrite writer unchanged. Extra blobs still materialize under their fallback names; extra markers are stripped but ignored. --- pkg/chat/media.go | 9 +- pkg/model/provider/gemini/client.go | 1 + .../gemini/image_output_instruction.go | 31 ++ .../gemini/image_output_instruction_test.go | 179 ++++++++++++ pkg/runtime/generated_media_markers.go | 177 ++++++++++++ pkg/runtime/generated_media_markers_test.go | 270 ++++++++++++++++++ pkg/runtime/streaming.go | 52 +++- pkg/runtime/streaming_test.go | 182 ++++++++++++ 8 files changed, 885 insertions(+), 16 deletions(-) create mode 100644 pkg/model/provider/gemini/image_output_instruction.go create mode 100644 pkg/model/provider/gemini/image_output_instruction_test.go create mode 100644 pkg/runtime/generated_media_markers.go create mode 100644 pkg/runtime/generated_media_markers_test.go diff --git a/pkg/chat/media.go b/pkg/chat/media.go index f09859061..bf4ea6e02 100644 --- a/pkg/chat/media.go +++ b/pkg/chat/media.go @@ -27,10 +27,11 @@ type MediaDelta struct { // (e.g. echoed from an "as sunshine.jpg" instruction), when one exists. // It is untrusted model input: the runtime routes it through // workspacemedia.ClassifyRequestedPath, and a path escaping the workspace - // requires an explicit user confirmation before it is honored. Response - // marker extraction (the "[media-file: ...]" protocol) will populate it; - // until that lands, providers leave it empty and materialization falls - // back to Name. + // requires an explicit user confirmation before it is honored. The + // runtime's response marker filter (the "[media-file: ...]" protocol, + // pkg/runtime/generated_media_markers.go) populates it by pairing marker + // paths with blobs in response order; blobs no marker names keep it empty + // and materialization falls back to Name, then a generic name. RequestedPath string `json:"requested_path,omitempty"` // Size is the byte length of Data, cached because Data itself is diff --git a/pkg/model/provider/gemini/client.go b/pkg/model/provider/gemini/client.go index 4acefa9a3..498f6f966 100644 --- a/pkg/model/provider/gemini/client.go +++ b/pkg/model/provider/gemini/client.go @@ -761,6 +761,7 @@ func (c *Client) CreateChatCompletionStream( if c.wantsImageResponseModalities(imageOutputEnabled) { config.ResponseModalities = []string{string(genai.ModalityText), string(genai.ModalityImage)} + applyImageOutputMediaFileInstruction(config) } // Start with Google built-in tools (search, maps, code execution) from provider_opts diff --git a/pkg/model/provider/gemini/image_output_instruction.go b/pkg/model/provider/gemini/image_output_instruction.go new file mode 100644 index 000000000..fa271eac4 --- /dev/null +++ b/pkg/model/provider/gemini/image_output_instruction.go @@ -0,0 +1,31 @@ +package gemini + +import "google.golang.org/genai" + +// imageOutputMediaFileInstruction is appended as a system instruction on the +// explicit image-output gateway route (see wantsImageResponseModalities) so +// generated images arrive with a machine-readable filename: the runtime +// strips these exact marker lines from the reply and uses the paths to name +// the materialized workspace files (pkg/runtime/generated_media_markers.go). +// The single-image steering keeps one request yielding one predictably named +// file; every blob the model actually returns is still persisted. +const imageOutputMediaFileInstruction = `When you generate images, name each one with a marker line, placed alone on its own line, in this exact format: +[media-file: relative/path.ext] +Rules: +- Emit exactly one marker line per generated image, in the same order as the images. +- If the user asked for a specific file name or path, echo it in the marker exactly as requested. +- Otherwise choose a short, meaningful, kebab-case file name. +- Generate a single image unless the user explicitly asks for multiple images or variations. +- Never emit a marker line for an image you did not generate.` + +// applyImageOutputMediaFileInstruction appends the marker-protocol +// instruction to the request's system instruction, preserving any parts +// already present. Callers gate it on wantsImageResponseModalities so only +// the explicit image-output gateway chat route ever carries it. +func applyImageOutputMediaFileInstruction(config *genai.GenerateContentConfig) { + if config.SystemInstruction == nil { + config.SystemInstruction = &genai.Content{} + } + config.SystemInstruction.Parts = append(config.SystemInstruction.Parts, + genai.NewPartFromText(imageOutputMediaFileInstruction)) +} diff --git a/pkg/model/provider/gemini/image_output_instruction_test.go b/pkg/model/provider/gemini/image_output_instruction_test.go new file mode 100644 index 000000000..90919572b --- /dev/null +++ b/pkg/model/provider/gemini/image_output_instruction_test.go @@ -0,0 +1,179 @@ +package gemini + +import ( + "encoding/json" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "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/options" +) + +// systemInstructionTextsInBody decodes body's systemInstruction part texts +// (genai serializes GenerateContentConfig.SystemInstruction under the +// top-level "systemInstruction" key for the Gemini Developer API), returning +// nil when the key is absent. +func systemInstructionTextsInBody(t *testing.T, body []byte) []string { + t.Helper() + + var req map[string]any + require.NoError(t, json.Unmarshal(body, &req)) + + si, ok := req["systemInstruction"].(map[string]any) + if !ok { + return nil + } + parts, ok := si["parts"].([]any) + if !ok { + return nil + } + var out []string + for _, p := range parts { + partMap, ok := p.(map[string]any) + if !ok { + continue + } + if text, ok := partMap["text"].(string); ok { + out = append(out, text) + } + } + return out +} + +// TestCreateChatCompletionStream_MediaFileInstruction_PositiveRoute pins +// the sole route that must carry the media-file marker instruction — the +// gateway surface with an explicit output_capabilities.image declaration on +// an ordinary chat turn — and that it is sent exactly once. +func TestCreateChatCompletionStream_MediaFileInstruction_PositiveRoute(t *testing.T) { + t.Parallel() + + server, captured := newBodyCapturingGeminiServer(t, writeGeminiSSEResponse) + + 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)) + require.NoError(t, err) + + stream, err := client.CreateChatCompletionStream(t.Context(), []chat.Message{ + {Role: chat.MessageRoleUser, Content: "generate an image of a red panda"}, + }, nil) + require.NoError(t, err) + drainStream(t, stream) + + bodies := captured.all() + require.Len(t, bodies, 1) + texts := systemInstructionTextsInBody(t, bodies[0]) + require.Len(t, texts, 1, "the instruction must be sent exactly once") + assert.Equal(t, imageOutputMediaFileInstruction, texts[0]) + assert.Equal(t, 1, strings.Count(texts[0], "[media-file: "), "the instruction must show the marker format exactly once") +} + +// TestCreateChatCompletionStream_MediaFileInstruction_AbsentOnOtherRoutes +// pins that every other route sends no marker instruction (and no system +// instruction at all, since nothing else sets one today): a direct +// (non-gateway) call even when declared image-capable, gateway calls +// without the explicit declaration, and gateway title-generation or +// compaction calls. +func TestCreateChatCompletionStream_MediaFileInstruction_AbsentOnOtherRoutes(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + cfg func(serverURL string) *latest.ModelConfig + env map[string]string + gateway bool + opts []options.Opt + }{ + { + name: "direct Gemini API, declared true: absent", + cfg: func(serverURL string) *latest.ModelConfig { + return &latest.ModelConfig{ + Provider: "google", Model: "gemini-2.5-flash-image", BaseURL: serverURL, + OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}, + } + }, + env: map[string]string{"GOOGLE_API_KEY": "test-key"}, + }, + { + name: "gateway, declared false: absent", + cfg: func(string) *latest.ModelConfig { + return &latest.ModelConfig{ + Provider: "google", Model: "gemini-2.5-flash-image", + OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(false)}, + } + }, + env: map[string]string{environment.DockerDesktopTokenEnv: "test-dd-token"}, + gateway: true, + }, + { + name: "gateway, declaration missing: absent", + cfg: func(string) *latest.ModelConfig { + return &latest.ModelConfig{Provider: "google", Model: "gemini-2.5-flash-image"} + }, + env: map[string]string{environment.DockerDesktopTokenEnv: "test-dd-token"}, + gateway: true, + }, + { + name: "gateway, declared true, generating title: absent", + cfg: func(string) *latest.ModelConfig { + return &latest.ModelConfig{ + Provider: "google", Model: "gemini-2.5-flash-image", + OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}, + } + }, + env: map[string]string{environment.DockerDesktopTokenEnv: "test-dd-token"}, + gateway: true, + opts: []options.Opt{options.WithGeneratingTitle()}, + }, + { + name: "gateway, declared true, compacting: absent", + cfg: func(string) *latest.ModelConfig { + return &latest.ModelConfig{ + Provider: "google", Model: "gemini-2.5-flash-image", + OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}, + } + }, + env: map[string]string{environment.DockerDesktopTokenEnv: "test-dd-token"}, + gateway: true, + opts: []options.Opt{options.WithCompacting()}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + server, captured := newBodyCapturingGeminiServer(t, writeGeminiSSEResponse) + + cfg := tt.cfg(server.URL) + env := environment.NewMapEnvProvider(tt.env) + opts := tt.opts + if tt.gateway { + opts = append(opts, options.WithGateway(server.URL)) + } + client, err := NewClient(t.Context(), cfg, env, opts...) + require.NoError(t, err) + + stream, err := client.CreateChatCompletionStream(t.Context(), []chat.Message{ + {Role: chat.MessageRoleUser, Content: "hello"}, + }, nil) + require.NoError(t, err) + drainStream(t, stream) + + bodies := captured.all() + require.Len(t, bodies, 1) + assert.Nil(t, systemInstructionTextsInBody(t, bodies[0]), "the marker instruction must be absent on this route") + }) + } +} diff --git a/pkg/runtime/generated_media_markers.go b/pkg/runtime/generated_media_markers.go new file mode 100644 index 000000000..718e74408 --- /dev/null +++ b/pkg/runtime/generated_media_markers.go @@ -0,0 +1,177 @@ +package runtime + +import ( + "bytes" + "strings" + "unicode/utf8" + + "github.com/docker/docker-agent/pkg/chat" +) + +// Generated-media filename markers: an explicitly image-output-capable model +// is asked, via a provider-scoped system instruction (see +// pkg/model/provider/gemini/image_output_instruction.go), to emit one exact +// +// [media-file: relative/path] +// +// line per generated image, in response order. handleStream filters those +// lines out of the visible and persisted assistant text and pairs the +// extracted paths positionally with the accumulated media blobs. A marker +// path is untrusted model input: it flows into [chat.MediaDelta.RequestedPath] +// and through the existing workspacemedia classification, escape-confirmation, +// and never-overwrite policy — the parser itself never touches the filesystem. + +const ( + mediaFileMarkerPrefix = "[media-file: " + mediaFileMarkerSuffix = "]" + + // maxMediaFileMarkerLineBytes bounds a marker line (line terminator + // excluded); longer lines are ordinary text. This also caps how much + // text the filter may withhold while a line could still become a marker. + maxMediaFileMarkerLineBytes = 512 +) + +// parseMediaFileMarkerLine reports whether line — one complete line without +// its trailing "\n" but possibly with a trailing "\r" — is exactly a +// media-file marker, returning the raw requested path. The grammar is +// deliberately strict (exact lowercase prefix at column zero, closing bracket +// at end of line, no path-edge whitespace, no control characters, valid +// UTF-8, bounded length) so ordinary prose is virtually never swallowed. +func parseMediaFileMarkerLine(line string) (string, bool) { + line = strings.TrimSuffix(line, "\r") + if len(line) > maxMediaFileMarkerLineBytes { + return "", false + } + rest, ok := strings.CutPrefix(line, mediaFileMarkerPrefix) + if !ok { + return "", false + } + path, ok := strings.CutSuffix(rest, mediaFileMarkerSuffix) + if !ok || path == "" { + return "", false + } + if strings.TrimSpace(path) != path { + return "", false + } + if !utf8.ValidString(path) { + return "", false + } + for _, r := range path { + if r < 0x20 || r == 0x7f { + return "", false + } + } + return path, true +} + +// mediaFileMarkerFilter incrementally strips exact media-file marker lines +// from streamed assistant text, robust to chunks split at any byte boundary. +// Push returns the chunk's visible text — it withholds only bytes that could +// still become a marker line — and Finish flushes the final unterminated +// line, honoring an end-of-stream marker. Extracted paths accumulate in +// paths in response order. +type mediaFileMarkerFilter struct { + // candidate buffers the current line while it can still become a marker. + candidate []byte + // passthrough is set once the current line has diverged from the marker + // grammar; its remaining bytes then stream through until the newline. + passthrough bool + paths []string +} + +func (f *mediaFileMarkerFilter) Push(chunk string) string { + if chunk == "" { + return "" + } + var out strings.Builder + for i := range len(chunk) { + b := chunk[i] + if f.passthrough { + out.WriteByte(b) + if b == '\n' { + f.passthrough = false + } + continue + } + f.candidate = append(f.candidate, b) + if b == '\n' { + f.flushLine(&out) + continue + } + if !canBecomeMediaFileMarker(f.candidate) { + out.Write(f.candidate) + f.candidate = f.candidate[:0] + f.passthrough = true + } + } + return out.String() +} + +// flushLine consumes the buffered complete line (ending in "\n"): a valid +// marker is recorded, anything else is emitted verbatim. +func (f *mediaFileMarkerFilter) flushLine(out *strings.Builder) { + line := strings.TrimSuffix(string(f.candidate), "\n") + if path, ok := parseMediaFileMarkerLine(line); ok { + f.paths = append(f.paths, path) + } else { + out.Write(f.candidate) + } + f.candidate = f.candidate[:0] +} + +// Finish resolves the final unterminated line: an exact end-of-stream marker +// is recorded, anything else is returned as visible text. +func (f *mediaFileMarkerFilter) Finish() string { + if len(f.candidate) == 0 { + return "" + } + line := string(f.candidate) + f.candidate = nil + if path, ok := parseMediaFileMarkerLine(line); ok { + f.paths = append(f.paths, path) + return "" + } + return line +} + +// canBecomeMediaFileMarker reports whether the partial line (no newline yet) +// could still grow into a valid marker line; false diverges the filter to +// passthrough so the withheld bytes are released immediately. +func canBecomeMediaFileMarker(line []byte) bool { + // +1 tolerates a trailing '\r' still awaiting its '\n' at the bound. + if len(line) > maxMediaFileMarkerLineBytes+1 { + return false + } + if len(line) <= len(mediaFileMarkerPrefix) { + return string(line) == mediaFileMarkerPrefix[:len(line)] + } + if !bytes.HasPrefix(line, []byte(mediaFileMarkerPrefix)) { + return false + } + for i := len(mediaFileMarkerPrefix); i < len(line); i++ { + b := line[i] + if b == '\r' { + // A carriage return can only precede the terminating newline. + return i == len(line)-1 + } + if b < 0x20 || b == 0x7f { + return false + } + } + return true +} + +// applyMediaFileRequestedPaths pairs extracted marker paths with the turn's +// media blobs by position: marker i names blob i, overriding any +// provider-supplied requested path for that slot. Blobs beyond the last +// marker keep their existing naming fallback (provider display name, then +// the generic generated-N); extra markers are already stripped from the text +// and are simply ignored. +func applyMediaFileRequestedPaths(media []chat.MediaDelta, paths []string) { + for i := range media { + if i >= len(paths) { + return + } + media[i].RequestedPath = paths[i] + } +} diff --git a/pkg/runtime/generated_media_markers_test.go b/pkg/runtime/generated_media_markers_test.go new file mode 100644 index 000000000..a021f521a --- /dev/null +++ b/pkg/runtime/generated_media_markers_test.go @@ -0,0 +1,270 @@ +package runtime + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/agent" + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/session" +) + +// TestParseMediaFileMarkerLine pins the exact marker grammar: whole line, +// lowercase prefix at column zero, closing bracket at end, non-empty path +// with no edge whitespace, no control characters, valid UTF-8, bounded. +func TestParseMediaFileMarkerLine(t *testing.T) { + t.Parallel() + + atBoundPath := strings.Repeat("a", maxMediaFileMarkerLineBytes-len(mediaFileMarkerPrefix)-len(mediaFileMarkerSuffix)) + + tests := []struct { + name string + line string + wantPath string + wantOK bool + }{ + {name: "simple", line: "[media-file: cat.png]", wantPath: "cat.png", wantOK: true}, + {name: "nested path with internal space", line: "[media-file: images/my cat.png]", wantPath: "images/my cat.png", wantOK: true}, + {name: "trailing CR from a CRLF line", line: "[media-file: cat.png]\r", wantPath: "cat.png", wantOK: true}, + {name: "bracket inside the path", line: "[media-file: a]b.png]", wantPath: "a]b.png", wantOK: true}, + {name: "traversal is a valid untrusted path", line: "[media-file: ../escape.png]", wantPath: "../escape.png", wantOK: true}, + {name: "absolute is a valid untrusted path", line: "[media-file: /tmp/x.png]", wantPath: "/tmp/x.png", wantOK: true}, + {name: "exactly at the byte bound", line: mediaFileMarkerPrefix + atBoundPath + mediaFileMarkerSuffix, wantPath: atBoundPath, wantOK: true}, + + {name: "leading indentation", line: " [media-file: cat.png]"}, + {name: "uppercase prefix", line: "[MEDIA-FILE: cat.png]"}, + {name: "missing space after colon", line: "[media-file:cat.png]"}, + {name: "trailing text after bracket", line: "[media-file: cat.png] extra"}, + {name: "backtick-quoted", line: "`[media-file: cat.png]`"}, + {name: "empty path", line: "[media-file:]"}, + {name: "whitespace-only path", line: "[media-file: ]"}, + {name: "trailing space in path", line: "[media-file: cat.png ]"}, + {name: "leading space in path", line: "[media-file: cat.png]"}, + {name: "missing closing bracket", line: "[media-file: cat.png"}, + {name: "control character in path", line: "[media-file: a\tb.png]"}, + {name: "invalid UTF-8 in path", line: "[media-file: \xff.png]"}, + {name: "over the byte bound", line: mediaFileMarkerPrefix + atBoundPath + "a" + mediaFileMarkerSuffix}, + {name: "bare prefix only", line: "[media-file: "}, + {name: "empty line", line: ""}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + path, ok := parseMediaFileMarkerLine(tt.line) + assert.Equal(t, tt.wantOK, ok) + assert.Equal(t, tt.wantPath, path) + }) + } +} + +// runMarkerFilter feeds input to a fresh filter in chunks of chunkSize bytes +// and returns the concatenated visible text (including the Finish tail) and +// the extracted paths. +func runMarkerFilter(input string, chunkSize int) (string, []string) { + var f mediaFileMarkerFilter + var out strings.Builder + for start := 0; start < len(input); start += chunkSize { + end := min(start+chunkSize, len(input)) + out.WriteString(f.Push(input[start:end])) + } + out.WriteString(f.Finish()) + return out.String(), f.paths +} + +// TestMediaFileMarkerFilter pins stripping and extraction on whole inputs, +// then proves chunk-split invariance: every chunk size from a single byte +// upward must yield byte-identical visible text and identical paths. +func TestMediaFileMarkerFilter(t *testing.T) { + t.Parallel() + + overBoundLine := mediaFileMarkerPrefix + strings.Repeat("a", maxMediaFileMarkerLineBytes) + mediaFileMarkerSuffix + + tests := []struct { + name string + input string + wantText string + wantPaths []string + }{ + {name: "marker alone with newline", input: "[media-file: cat.png]\n", wantText: "", wantPaths: []string{"cat.png"}}, + {name: "marker terminated by end of stream", input: "[media-file: cat.png]", wantText: "", wantPaths: []string{"cat.png"}}, + {name: "CRLF marker", input: "[media-file: cat.png]\r\n", wantText: "", wantPaths: []string{"cat.png"}}, + {name: "marker between prose lines", input: "before\n[media-file: a.png]\nafter", wantText: "before\nafter", wantPaths: []string{"a.png"}}, + {name: "multiple markers keep response order", input: "one\n[media-file: first.png]\n[media-file: second.png]\ntwo\n", wantText: "one\ntwo\n", wantPaths: []string{"first.png", "second.png"}}, + {name: "indented lookalike stays visible", input: " [media-file: cat.png]\n", wantText: " [media-file: cat.png]\n"}, + {name: "suffixed lookalike stays visible", input: "[media-file: cat.png] done\n", wantText: "[media-file: cat.png] done\n"}, + {name: "uppercase lookalike stays visible", input: "[MEDIA-FILE: cat.png]\n", wantText: "[MEDIA-FILE: cat.png]\n"}, + {name: "empty-path lookalike stays visible", input: "[media-file:]\n", wantText: "[media-file:]\n"}, + {name: "control character diverges to text", input: "[media-file: a\tb.png]\n", wantText: "[media-file: a\tb.png]\n"}, + {name: "over-bound line streams as text", input: overBoundLine + "\n", wantText: overBoundLine + "\n"}, + {name: "unterminated non-marker tail is flushed", input: "trailing [media", wantText: "trailing [media"}, + {name: "unterminated prefix-only tail is flushed", input: "[media-file: half", wantText: "[media-file: half"}, + {name: "blank lines survive", input: "a\n\n[media-file: x.png]\n\nb\n", wantText: "a\n\n\nb\n", wantPaths: []string{"x.png"}}, + {name: "CRLF prose preserved byte-for-byte", input: "line one\r\nline two\r\n", wantText: "line one\r\nline two\r\n"}, + {name: "plain prose untouched", input: "no markers here, just text.", wantText: "no markers here, just text."}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + for chunkSize := 1; chunkSize <= max(len(tt.input), 1); chunkSize++ { + text, paths := runMarkerFilter(tt.input, chunkSize) + require.Equal(t, tt.wantText, text, "visible text must be split-invariant (chunk size %d)", chunkSize) + require.Equal(t, tt.wantPaths, paths, "extracted paths must be split-invariant (chunk size %d)", chunkSize) + } + }) + } +} + +// TestApplyMediaFileRequestedPaths pins positional pairing: marker i names +// blob i (overriding any pre-set requested path), unpaired blobs keep their +// existing value, extra markers are ignored. +func TestApplyMediaFileRequestedPaths(t *testing.T) { + t.Parallel() + + t.Run("fewer markers than blobs", func(t *testing.T) { + t.Parallel() + media := []chat.MediaDelta{{Name: "a"}, {Name: "b", RequestedPath: "pre-set.png"}} + applyMediaFileRequestedPaths(media, []string{"named.png"}) + assert.Equal(t, "named.png", media[0].RequestedPath) + assert.Equal(t, "pre-set.png", media[1].RequestedPath, "an unpaired blob keeps its existing requested path") + }) + + t.Run("more markers than blobs", func(t *testing.T) { + t.Parallel() + media := []chat.MediaDelta{{Name: "a"}} + applyMediaFileRequestedPaths(media, []string{"one.png", "two.png"}) + assert.Equal(t, "one.png", media[0].RequestedPath) + }) + + t.Run("marker overrides a pre-set path", func(t *testing.T) { + t.Parallel() + media := []chat.MediaDelta{{RequestedPath: "provider.png"}} + applyMediaFileRequestedPaths(media, []string{"marker.png"}) + assert.Equal(t, "marker.png", media[0].RequestedPath) + }) + + t.Run("no markers is a no-op", func(t *testing.T) { + t.Parallel() + media := []chat.MediaDelta{{RequestedPath: "keep.png"}} + applyMediaFileRequestedPaths(media, nil) + assert.Equal(t, "keep.png", media[0].RequestedPath) + }) +} + +// markerTurnMedia runs handleStream over a marker-bearing stream and returns +// the aggregated media, ready for materialization. It also proves the +// persisted text is marker-free. +func markerTurnMedia(t *testing.T, stream *mockStream, wantContent string) []chat.MediaDelta { + t.Helper() + + a := agent.New("root", "test", agent.WithModel(&mockProvider{id: "test/mock-model", stream: stream})) + sess := session.New(session.WithUserMessage("go")) + evCh := make(chan Event, 64) + res, err := handleStream( + t.Context(), nil, stream, a, nil, sess, nil, + defaultTelemetry{}, NewChannelSink(evCh), defaultStreamIdleTimeout, + ) + require.NoError(t, err) + assert.Equal(t, wantContent, res.Content, "the persisted text must carry no marker line") + return res.Media +} + +// TestMarkerNamedMediaMaterializesEndToEnd is the full naming-protocol +// integration: a streamed "[media-file: sunshine.jpg]" marker names the PNG +// blob, the writer corrects the extension to the actual MIME type, the +// correction notice is visible, and the persisted part references the final +// workspace path. +func TestMarkerNamedMediaMaterializesEndToEnd(t *testing.T) { + r, _, _ := newMediaTestRuntime(t) + sess, root := workspaceSession(t, "sess-marker-e2e") + + stream := newStreamBuilder(). + AddContent("Here you go!\n[media-file: sunshine.jpg]\n"). + AddMedia([]byte{0x89, 0x50, 0x4e, 0x47}, "image/png", "provider-name.png"). + AddStopWithUsage(1, 1). + Build() + media := markerTurnMedia(t, stream, "Here you go!\n") + + sink := &collectingSink{} + parts := r.materializeGeneratedMedia(t.Context(), sess, media, "root", sink) + + require.Len(t, parts, 1) + assert.Equal(t, "sunshine.png", parts[0].Document.Source.ArtifactPath, "the marker name must win over the provider name, with the MIME-derived extension") + _, err := os.Stat(filepath.Join(root, "sunshine.png")) + require.NoError(t, err) + + warnings := sink.warnings() + require.Len(t, warnings, 1, "the extension correction must be user-visible") + assert.Contains(t, warnings[0].Message, "sunshine.png") + assert.Contains(t, warnings[0].Message, `".jpg"`) +} + +// TestMarkerEscapedPathRedirectsThroughEscapePolicy proves a traversing +// marker path reaches the existing escape policy unchanged: with no user to +// ask (non-interactive), the bytes are redirected into the workspace under +// the sanitized basename instead of being written outside or discarded. +func TestMarkerEscapedPathRedirectsThroughEscapePolicy(t *testing.T) { + r, _ := newEscapeTestRuntime(t) + r.nonInteractive = true + sess, root := workspaceSession(t, "sess-marker-escape") + + stream := newStreamBuilder(). + AddContent("[media-file: ../outside.png]\n"). + AddMedia([]byte{0xAA}, "image/png", ""). + AddStopWithUsage(1, 1). + Build() + media := markerTurnMedia(t, stream, "") + + sink := &collectingSink{} + parts := r.materializeGeneratedMedia(t.Context(), sess, media, "root", sink) + + require.Len(t, parts, 1) + assert.Equal(t, "outside.png", parts[0].Document.Source.ArtifactPath) + assert.Equal(t, chat.ArtifactRootWorkspace, parts[0].Document.Source.ArtifactRoot) + _, err := os.Stat(filepath.Join(root, "outside.png")) + require.NoError(t, err) + _, err = os.Stat(filepath.Join(filepath.Dir(root), "outside.png")) + assert.True(t, os.IsNotExist(err), "nothing may land outside the workspace without confirmation") + require.Len(t, sink.warnings(), 1, "the redirect must be explained to the user") +} + +// TestUnmarkedBlobFallsBackToProviderThenGenericName proves the precedence +// tail end to end: in a two-blob turn with one marker, the second blob still +// materializes under its provider display name, and with neither marker nor +// provider name the generic generated-N name is used. +func TestUnmarkedBlobFallsBackToProviderThenGenericName(t *testing.T) { + r, _, _ := newMediaTestRuntime(t) + sess, root := workspaceSession(t, "sess-marker-fallback") + + stream := newStreamBuilder(). + AddContent("[media-file: named.png]\n"). + AddMultiMedia( + chat.MediaDelta{Data: []byte{0x01}, MimeType: "image/png", Name: "first", Size: 1}, + chat.MediaDelta{Data: []byte{0x02}, MimeType: "image/png", Name: "provider-pick", Size: 1}, + chat.MediaDelta{Data: []byte{0x03}, MimeType: "image/png", Size: 1}, + ). + AddStopWithUsage(1, 1). + Build() + media := markerTurnMedia(t, stream, "") + + sink := &collectingSink{} + parts := r.materializeGeneratedMedia(t.Context(), sess, media, "root", sink) + + require.Len(t, parts, 3, "every returned blob must materialize, marker or not") + assert.Empty(t, sink.warnings()) + assert.Equal(t, "named.png", parts[0].Document.Source.ArtifactPath) + assert.Equal(t, "provider-pick.png", parts[1].Document.Source.ArtifactPath) + assert.Equal(t, "generated-3.png", parts[2].Document.Source.ArtifactPath) + for _, name := range []string{"named.png", "provider-pick.png", "generated-3.png"} { + _, err := os.Stat(filepath.Join(root, name)) + require.NoError(t, err) + } +} diff --git a/pkg/runtime/streaming.go b/pkg/runtime/streaming.go index 1ae41a5fb..42b1bfdcb 100644 --- a/pkg/runtime/streaming.go +++ b/pkg/runtime/streaming.go @@ -129,6 +129,42 @@ func handleStream(ctx context.Context, cancelStream context.CancelCauseFunc, str toolDefMap[t.Name] = t } + // markerFilter strips [media-file: ...] naming markers from the assistant + // text BEFORE it is emitted or accumulated, so markers never flash in the + // TUI and never reach the persisted message. Provider-neutral: the strict + // line grammar is a no-op on streams that never emit markers. + var markerFilter mediaFileMarkerFilter + + // appendContent accumulates and emits assistant text that survived the + // marker filter, keeping the live event text and fullContent identical + // while gating raw XML out of the event stream. + appendContent := func(content string) { + if content == "" { + return + } + if !xmlToolCallGate { + tagIdx := strings.Index(content, "") + if tagIdx < 0 { + events.Emit(AgentChoice(a.Name(), sess.ID, content)) + } else { + xmlToolCallGate = true + if tagIdx > 0 { + events.Emit(AgentChoice(a.Name(), sess.ID, content[:tagIdx])) + } + } + } + fullContent.WriteString(content) + } + + // finishAssistantText flushes the marker filter's withheld tail and pairs + // the extracted requested paths onto the accumulated media. Called on + // every successful completion path (terminal finish reason or bare EOF), + // before any XML tool-call fallback parsing. + finishAssistantText := func() { + appendContent(markerFilter.Finish()) + applyMediaFileRequestedPaths(media, markerFilter.paths) + } + // applyXMLFallback extracts blocks from accumulated content when // no structured tool calls were received. Called from both the early-return // and EOF paths. @@ -291,6 +327,7 @@ mainLoop: } if choice.FinishReason == chat.FinishReasonStop || choice.FinishReason == chat.FinishReasonLength || choice.FinishReason == chat.FinishReasonRefusal { + finishAssistantText() recordUsage() finishReason := choice.FinishReason if finishReason == chat.FinishReasonRefusal { @@ -340,18 +377,7 @@ mainLoop: } if choice.Delta.Content != "" { - if !xmlToolCallGate { - tagIdx := strings.Index(choice.Delta.Content, "") - if tagIdx < 0 { - events.Emit(AgentChoice(a.Name(), sess.ID, choice.Delta.Content)) - } else { - xmlToolCallGate = true - if tagIdx > 0 { - events.Emit(AgentChoice(a.Name(), sess.ID, choice.Delta.Content[:tagIdx])) - } - } - } - fullContent.WriteString(choice.Delta.Content) + appendContent(markerFilter.Push(choice.Delta.Content)) } case <-ctx.Done(): @@ -375,6 +401,8 @@ mainLoop: } } + finishAssistantText() + recordUsage() applyXMLFallback() diff --git a/pkg/runtime/streaming_test.go b/pkg/runtime/streaming_test.go index 321b89f65..d8890b325 100644 --- a/pkg/runtime/streaming_test.go +++ b/pkg/runtime/streaming_test.go @@ -3,6 +3,7 @@ package runtime import ( "context" "io" + "strings" "sync" "testing" "time" @@ -495,3 +496,184 @@ func TestHandleStream_ContextCancellation(t *testing.T) { require.ErrorIs(t, err, context.Canceled, "error must be context.Canceled") assert.True(t, res.Stopped) } + +// agentChoiceText drains every buffered event and concatenates the +// AgentChoice content, i.e. exactly what a live consumer (TUI/API) rendered. +func agentChoiceText(ch chan Event) string { + var b strings.Builder + for { + select { + case e := <-ch: + if c, ok := e.(*AgentChoiceEvent); ok { + b.WriteString(c.Content) + } + default: + return b.String() + } + } +} + +// runMarkerStream runs handleStream over stream and returns the result plus +// the concatenated live AgentChoice text. +func runMarkerStream(t *testing.T, stream *mockStream) (streamResult, string) { + t.Helper() + + a := agent.New("root", "test", agent.WithModel(&mockProvider{id: "test/mock-model", stream: stream})) + sess := session.New(session.WithUserMessage("go")) + evCh := make(chan Event, 64) + res, err := handleStream( + t.Context(), nil, stream, a, nil, sess, nil, + defaultTelemetry{}, NewChannelSink(evCh), defaultStreamIdleTimeout, + ) + require.NoError(t, err) + return res, agentChoiceText(evCh) +} + +// TestHandleStream_MediaFileMarkerStrippedAndPaired is the core streaming +// contract of the naming protocol: a marker line split across chunks never +// reaches the live event stream or the aggregated content, and its path is +// paired onto the blob in [chat.MediaDelta.RequestedPath]. +func TestHandleStream_MediaFileMarkerStrippedAndPaired(t *testing.T) { + t.Parallel() + + imgBytes := []byte{0x89, 0x50, 0x4e, 0x47} + stream := newStreamBuilder(). + AddContent("Here you go!\n[media-fi"). + AddContent("le: red-panda.png]\n"). + AddMedia(imgBytes, "image/png", "provider-name.png"). + AddStopWithUsage(1, 1). + Build() + + res, live := runMarkerStream(t, stream) + + assert.Equal(t, "Here you go!\n", res.Content, "the marker line must be stripped from the persisted text") + assert.Equal(t, res.Content, live, "live event text and aggregated content must be identical") + require.Len(t, res.Media, 1) + assert.Equal(t, "red-panda.png", res.Media[0].RequestedPath) + assert.Equal(t, "provider-name.png", res.Media[0].Name, "the provider display name must survive for fallback") +} + +// TestHandleStream_MarkerAtEOFWithoutNewline: a marker terminated by the end +// of the stream (bare EOF, no trailing newline, media arrived first) is +// still stripped and paired. +func TestHandleStream_MarkerAtEOFWithoutNewline(t *testing.T) { + t.Parallel() + + stream := newStreamBuilder(). + AddMedia([]byte{0x01}, "image/png", ""). + AddContent("[media-file: cat.png]"). + Build() + + res, live := runMarkerStream(t, stream) + + assert.Empty(t, res.Content) + assert.Empty(t, live) + require.Len(t, res.Media, 1) + assert.Equal(t, "cat.png", res.Media[0].RequestedPath) + assert.True(t, res.Stopped) +} + +// TestHandleStream_MarkerBlobCountMismatch pins the pairing rules when the +// model misbehaves: markers pair positionally, extra blobs keep their +// fallback naming, and extra markers are stripped but ignored. +func TestHandleStream_MarkerBlobCountMismatch(t *testing.T) { + t.Parallel() + + t.Run("fewer markers than blobs", func(t *testing.T) { + t.Parallel() + + stream := newStreamBuilder(). + AddContent("[media-file: only.png]\n"). + AddMultiMedia( + chat.MediaDelta{Data: []byte{0x01}, MimeType: "image/png", Name: "a", Size: 1}, + chat.MediaDelta{Data: []byte{0x02}, MimeType: "image/png", Name: "b", Size: 1}, + ). + AddStopWithUsage(1, 1). + Build() + + res, live := runMarkerStream(t, stream) + + assert.Empty(t, res.Content) + assert.Empty(t, live) + require.Len(t, res.Media, 2, "every blob must survive, marker or not") + assert.Equal(t, "only.png", res.Media[0].RequestedPath) + assert.Empty(t, res.Media[1].RequestedPath, "the unpaired blob falls back to its provider name") + assert.Equal(t, []byte{0x01}, res.Media[0].Data, "blob order must be preserved") + }) + + t.Run("more markers than blobs", func(t *testing.T) { + t.Parallel() + + stream := newStreamBuilder(). + AddContent("[media-file: one.png]\n[media-file: two.png]\n"). + AddMedia([]byte{0x01}, "image/png", ""). + AddStopWithUsage(1, 1). + Build() + + res, live := runMarkerStream(t, stream) + + assert.Empty(t, res.Content, "every valid marker line is stripped, even unpaired ones") + assert.Empty(t, live) + require.Len(t, res.Media, 1) + assert.Equal(t, "one.png", res.Media[0].RequestedPath) + }) +} + +// TestHandleStream_MultipleMarkersPairInOrder: marker i names blob i, in +// response order, across separate chunks. +func TestHandleStream_MultipleMarkersPairInOrder(t *testing.T) { + t.Parallel() + + stream := newStreamBuilder(). + AddContent("Two variations:\n[media-file: variant-one.png]\n"). + AddMedia([]byte{0x01}, "image/png", ""). + AddContent("[media-file: variant-two.png]\n"). + AddMedia([]byte{0x02}, "image/png", ""). + AddStopWithUsage(1, 1). + Build() + + res, live := runMarkerStream(t, stream) + + assert.Equal(t, "Two variations:\n", res.Content) + assert.Equal(t, res.Content, live) + require.Len(t, res.Media, 2) + assert.Equal(t, "variant-one.png", res.Media[0].RequestedPath) + assert.Equal(t, "variant-two.png", res.Media[1].RequestedPath) +} + +// TestHandleStream_MalformedMarkerStaysVisible: near-miss lines are ordinary +// prose — visible live, persisted, and never consuming a pairing slot. +func TestHandleStream_MalformedMarkerStaysVisible(t *testing.T) { + t.Parallel() + + stream := newStreamBuilder(). + AddContent(" [media-file: indented.png]\n[media-file: real.png]\n"). + AddMedia([]byte{0x01}, "image/png", ""). + AddStopWithUsage(1, 1). + Build() + + res, live := runMarkerStream(t, stream) + + assert.Equal(t, " [media-file: indented.png]\n", res.Content) + assert.Equal(t, res.Content, live) + require.Len(t, res.Media, 1) + assert.Equal(t, "real.png", res.Media[0].RequestedPath, "the malformed line must not consume the pairing slot") +} + +// TestHandleStream_TextWithoutMarkersUnchanged guards against the filter +// perturbing ordinary streamed text, including bracketed prose. +func TestHandleStream_TextWithoutMarkersUnchanged(t *testing.T) { + t.Parallel() + + stream := newStreamBuilder(). + AddContent("see [media docs] and "). + AddContent("[media-file spec] for details\n"). + AddStopWithUsage(1, 1). + Build() + + res, live := runMarkerStream(t, stream) + + assert.Equal(t, "see [media docs] and [media-file spec] for details\n", res.Content) + assert.Equal(t, res.Content, live) + assert.Empty(t, res.Media) +} From 25c9d9a8124f3dae48b487223567cb2a0f89fb9d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Sat, 1 Aug 2026 00:04:36 +0200 Subject: [PATCH 2/2] feat(#3996): fall back to an explicit user-prompt filename for one unnamed image MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A live image-output model can ignore the [media-file:] marker instruction entirely, leaving a prompt like "Generate an image as sunshine.jpg" to land as generated-1.png. Add a deterministic fallback: when a turn returns exactly ONE media blob that marker pairing left unnamed, parse a single unambiguous explicit output filename from the triggering user message. The cue grammar is strict, deliberately not NLP: save (it) as / save to / write to / output to / name it / call it / filename:=, plus bare "as" only inside a narrow imperative output context (generation verb + optionally-articled media noun, e.g. "Generate an image as sunshine.jpg"), and a companion of-phrase form (generation verb + media noun + "of " + "as ") so "Generate an image of a red panda as assets/red-panda.jpg" extracts the intended name. The of-phrase subject cannot cross quotes, clause punctuation, or CR/LF, and any subject containing "as", "with", or "in" is refused because those prepositions introduce open-ended attribute phrases whose trailing "as" compares — a false reject only costs the generic generated-N name, while a false capture could engage the escape-confirmation policy for a merely referenced file. Bare "called" and unanchored "as" are NOT cues, so comparative references to existing files never extract a name. Candidates are keyed by the filename capture's position: one occurrence matched by both grammars counts once; distinct occurrences stay ambiguous and extract nothing. Names may be quoted, backticked, or unquoted with a known image extension; zero or multiple candidates extract nothing. Precedence stays marker -> user-prompt filename -> provider display name -> generic generated-N, and the extracted path only fills MediaDelta.RequestedPath, so the existing untrusted-path pipeline (MIME/extension correction, collision suffixing, workspace containment and escape confirmation) applies unchanged. --- pkg/chat/media.go | 12 +- .../generated_media_prompt_filename.go | 170 +++++++ .../generated_media_prompt_filename_test.go | 443 ++++++++++++++++++ pkg/runtime/loop.go | 15 +- pkg/runtime/media_escape.go | 6 +- 5 files changed, 634 insertions(+), 12 deletions(-) create mode 100644 pkg/runtime/generated_media_prompt_filename.go create mode 100644 pkg/runtime/generated_media_prompt_filename_test.go diff --git a/pkg/chat/media.go b/pkg/chat/media.go index bf4ea6e02..d541c59ef 100644 --- a/pkg/chat/media.go +++ b/pkg/chat/media.go @@ -23,15 +23,17 @@ type MediaDelta struct { // runtime accumulator synthesizes one when needed. Name string `json:"name,omitempty"` - // RequestedPath is the prompt-directed target path the model asked for - // (e.g. echoed from an "as sunshine.jpg" instruction), when one exists. - // It is untrusted model input: the runtime routes it through + // RequestedPath is the prompt-directed target path for this blob, when + // one exists. It is untrusted input: the runtime routes it through // workspacemedia.ClassifyRequestedPath, and a path escaping the workspace // requires an explicit user confirmation before it is honored. The // runtime's response marker filter (the "[media-file: ...]" protocol, // pkg/runtime/generated_media_markers.go) populates it by pairing marker - // paths with blobs in response order; blobs no marker names keep it empty - // and materialization falls back to Name, then a generic name. + // paths with blobs in response order; a single otherwise-unnamed blob may + // instead get it from deterministic explicit-filename extraction on the + // triggering user message (pkg/runtime/generated_media_prompt_filename.go). + // Blobs neither source names keep it empty and materialization falls back + // to Name, then a generic name. RequestedPath string `json:"requested_path,omitempty"` // Size is the byte length of Data, cached because Data itself is diff --git a/pkg/runtime/generated_media_prompt_filename.go b/pkg/runtime/generated_media_prompt_filename.go new file mode 100644 index 000000000..57dfa08af --- /dev/null +++ b/pkg/runtime/generated_media_prompt_filename.go @@ -0,0 +1,170 @@ +package runtime + +import ( + "path" + "regexp" + "strings" + "unicode/utf8" + + "github.com/docker/docker-agent/pkg/chat" +) + +// Deterministic user-prompt filename fallback: when a turn returns exactly +// one media blob and no [media-file:] marker named it (a model may ignore +// the marker instruction entirely), the triggering user message — never +// model text or history — is scanned for a single unambiguous explicit +// output filename ("Generate an image as sunshine.jpg", "Generate an image +// of a red panda as assets/red-panda.jpg", "save it as `pics/cat.png`", +// "filename: x.webp", ...). The naming precedence is +// therefore: marker → +// user-prompt explicit filename → provider display name → generic +// generated-N. An extracted path is untrusted exactly like a marker path: +// it only fills [chat.MediaDelta.RequestedPath] and flows through the same +// workspacemedia classification, MIME/extension correction, collision +// suffixing, and escape-confirmation pipeline. This is a strict grammar, +// deliberately not generic NLP: zero or multiple candidates mean no +// extraction. + +// maxExplicitOutputFilenameBytes bounds an extracted filename; anything +// longer is not treated as a candidate. +const maxExplicitOutputFilenameBytes = 256 + +// explicitOutputFilenameRE matches an explicit output-naming cue followed by +// a quoted, backticked, or unquoted filename. Bare "to" is deliberately not +// a cue so input-file mentions ("add a border to photo.jpg") never match, +// and bare "called" is not a cue because it usually references an existing +// input file ("similar to the one called old-render.png"). Bare "as" only +// counts inside a narrow imperative output context — a generation verb +// directly followed by an optionally-articled media noun ("Generate an +// image as sunshine.jpg"; explicitOutputFilenameOfPhraseRE below adds the +// "of " variant of the same context) — because RE2 has no +// lookbehind to exclude the comparative form ("in the same style as +// sunshine.jpg", "the same background as assets/bg.png") any other way. +// The unquoted form +// additionally anchors on a known image extension so trailing punctuation +// is excluded. Quoted/backticked contents may include spaces and are +// validated separately by isExplicitImageFilename. +var explicitOutputFilenameRE = regexp.MustCompile( + `(?i)\b(?:(?:save\s+it\s+as|save\s+as|save\s+to|write\s+to|output\s+to|name\s+it|call\s+it)\s+|filename\s*[:=]\s*|` + + `(?:re)?(?:generate|create|make|draw|render|produce)\s+(?:an?\s+|the\s+)?` + + `(?:image|picture|photo|banner|logo|icon|graphic|drawing|illustration|thumbnail|sticker|avatar|gif)\s+as\s+)` + + `(?:"([^"]+)"|'([^']+)'|` + "`([^`]+)`" + `|([^\s"'` + "`" + `]+\.(?:png|jpe?g|webp|gif))\b)`) + +// explicitOutputFilenameOfPhraseRE extends the imperative form above to a +// media noun carrying an "of " description before the bare "as" +// cue ("Generate an image of a red panda coding at a terminal as +// assets/red-panda-terminal.jpg"). The subject cannot cross quotes, +// clause punctuation, or CR/LF — so "Generate an image of a cat. Save it +// as x.png" and the newline-separated equivalent stay a single +// save-it-as candidate — and it is captured so extraction can reject +// comparative phrasings inside it ("of a cat in the same style as +// old.png", "of a beach with the exact palette as ref.png"; see +// comparativeCueRE): RE2 has no lookbehind to express that exclusion in +// the pattern itself. The subject MAY swallow a conjunction ahead of an +// explicit cue ("of a wolf and save it as logo.png"); extraction +// deduplicates that overlap with explicitOutputFilenameRE by capture +// position. +var explicitOutputFilenameOfPhraseRE = regexp.MustCompile( + `(?i)\b(?:re)?(?:generate|create|make|draw|render|produce)\s+(?:an?\s+|the\s+)?` + + `(?:image|picture|photo|banner|logo|icon|graphic|drawing|illustration|thumbnail|sticker|avatar|gif)\s+` + + `(of\s+[^"'` + "`" + `.,;:!?\r\n]+?)\s+as\s+` + + `(?:"([^"]+)"|'([^']+)'|` + "`([^`]+)`" + `|([^\s"'` + "`" + `]+\.(?:png|jpe?g|webp|gif))\b)`) + +// comparativeCueRE flags an of-phrase subject whose trailing "as" compares +// against an existing file instead of naming the output ("of a cat in the +// same style as old.png", "of something like ref.png"). Beyond explicit +// comparative words, it conservatively rejects any subject containing +// "as", "with", or "in": those introduce attribute phrases whose trailing +// "as" compares ("with the exact palette as ref.png", "with identical +// colors as ref.png", "as tall as tree.png") and comparative adjectives +// are open-ended. The asymmetry justifies over-rejecting: a false reject +// only costs the generic generated-N name, while a false capture can +// engage the escape-confirmation policy for a merely referenced file. +var comparativeCueRE = regexp.MustCompile( + `(?i)\b(?:same|style|similar|like|such|as|with|in|exact|identical|matching|equivalent)\b`) + +// extractExplicitOutputFilename returns the single unambiguous explicit +// output filename in the triggering user prompt, if there is exactly one. +// Candidates that fail validation (unknown extension, control characters, +// invalid UTF-8, over-long, extension-only, comparative of-phrase) are +// ignored rather than counted. The two grammars overlap: an of-phrase +// subject may swallow a conjunction ahead of an explicit cue ("make a +// logo of a wolf and save it as logo.png"), so candidates are keyed by +// the filename capture's position — the same occurrence matched by both +// grammars counts once, while distinct occurrences (even of the same +// name) stay ambiguous and yield no extraction. +func extractExplicitOutputFilename(prompt string) (string, bool) { + candidates := make(map[int]string) + for _, m := range explicitOutputFilenameRE.FindAllStringSubmatchIndex(prompt, -1) { + if offset, name := firstMatchedGroup(prompt, m, 1); isExplicitImageFilename(name) { + candidates[offset] = name + } + } + for _, m := range explicitOutputFilenameOfPhraseRE.FindAllStringSubmatchIndex(prompt, -1) { + if comparativeCueRE.MatchString(prompt[m[2]:m[3]]) { + continue + } + if offset, name := firstMatchedGroup(prompt, m, 2); isExplicitImageFilename(name) { + candidates[offset] = name + } + } + if len(candidates) != 1 { + return "", false + } + for _, name := range candidates { + return name, true + } + return "", false +} + +// firstMatchedGroup returns the start offset and text of the one capture +// group the quoting alternation filled, scanning submatch index pairs from +// firstGroup on; every group requires at least one character, so a filled +// group has a non-negative start. +func firstMatchedGroup(prompt string, m []int, firstGroup int) (int, string) { + for g := firstGroup; 2*g+1 < len(m); g++ { + if start, end := m[2*g], m[2*g+1]; start >= 0 { + return start, prompt[start:end] + } + } + return -1, "" +} + +// isExplicitImageFilename applies the same field constraints as the marker +// grammar (valid UTF-8, no control characters, no edge whitespace, bounded) +// plus a known image extension and a non-empty stem. Traversing or absolute +// paths are valid candidates on purpose: the untrusted-path pipeline owns +// containment and escape confirmation. +func isExplicitImageFilename(name string) bool { + if name == "" || len(name) > maxExplicitOutputFilenameBytes || !utf8.ValidString(name) { + return false + } + if strings.TrimSpace(name) != name { + return false + } + for _, r := range name { + if r < 0x20 || r == 0x7f { + return false + } + } + switch strings.ToLower(path.Ext(name)) { + case ".png", ".jpg", ".jpeg", ".webp", ".gif": + default: + return false + } + return len(path.Base(name)) > len(path.Ext(name)) +} + +// applyUserPromptRequestedPath fills RequestedPath from the triggering user +// prompt for a turn that returned exactly ONE media blob that marker pairing +// left unnamed. Markers keep precedence (a non-empty RequestedPath is never +// overwritten) and multi-blob turns are skipped entirely — a single prompt +// filename cannot unambiguously name one blob among several. +func applyUserPromptRequestedPath(media []chat.MediaDelta, prompt string) { + if len(media) != 1 || media[0].RequestedPath != "" { + return + } + if name, ok := extractExplicitOutputFilename(prompt); ok { + media[0].RequestedPath = name + } +} diff --git a/pkg/runtime/generated_media_prompt_filename_test.go b/pkg/runtime/generated_media_prompt_filename_test.go new file mode 100644 index 000000000..700fbba65 --- /dev/null +++ b/pkg/runtime/generated_media_prompt_filename_test.go @@ -0,0 +1,443 @@ +package runtime + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/session" +) + +// TestExtractExplicitOutputFilename pins the strict explicit-naming grammar: +// a known cue phrase followed by one quoted/backticked/unquoted filename +// with a known image extension, exactly one candidate, bounded and +// control-character/UTF-8 validated. Bare "as" only counts right after a +// generation verb plus media noun; comparative references to existing files +// ("same style as X", "the one called X"), ordinary input-file mentions, +// and ambiguous prompts must extract nothing. +func TestExtractExplicitOutputFilename(t *testing.T) { + t.Parallel() + + atBoundName := strings.Repeat("a", maxExplicitOutputFilenameBytes-len(".png")) + ".png" + + tests := []struct { + name string + prompt string + wantName string + wantOK bool + }{ + {name: "live repro prompt", prompt: "Generate an image as sunshine.jpg", wantName: "sunshine.jpg", wantOK: true}, + {name: "of-phrase live repro prompt", prompt: "Generate an image of a red panda coding at a terminal as assets/red-panda-terminal.jpg", wantName: "assets/red-panda-terminal.jpg", wantOK: true}, + {name: "of-phrase with quoted name", prompt: "draw a picture of my dog as `pets/rex holiday.png`", wantName: "pets/rex holiday.png", wantOK: true}, + {name: "of-phrase across a sentence boundary is not a second candidate", prompt: "Generate an image of a cat. Save it as x.png", wantName: "x.png", wantOK: true}, + {name: "of-phrase across a newline is not a second candidate", prompt: "Generate an image of a cat\nSave it as x.png", wantName: "x.png", wantOK: true}, + {name: "of-phrase across a CRLF is not a second candidate", prompt: "Generate an image of a cat\r\nSave it as x.png", wantName: "x.png", wantOK: true}, + {name: "of-phrase conjunction into a save-it-as cue dedupes to one candidate", prompt: "make a logo of a wolf and save it as logo.png", wantName: "logo.png", wantOK: true}, + {name: "image-noun conjunction into a save-it-as cue dedupes to one candidate", prompt: "Generate an image of a cat and save it as cat.png", wantName: "cat.png", wantOK: true}, + {name: "save as", prompt: "please save as cat.png", wantName: "cat.png", wantOK: true}, + {name: "save it as", prompt: "make a logo and save it as logo.png", wantName: "logo.png", wantOK: true}, + {name: "name it", prompt: "name it banner.webp", wantName: "banner.webp", wantOK: true}, + {name: "call it", prompt: "call it pic.jpeg", wantName: "pic.jpeg", wantOK: true}, + {name: "verb anchor with 'the'", prompt: "create the logo as logo.png", wantName: "logo.png", wantOK: true}, + {name: "verb anchor without article", prompt: "regenerate image as fixed.png", wantName: "fixed.png", wantOK: true}, + {name: "verb anchor with quoted name", prompt: `draw a picture as "my sun.png"`, wantName: "my sun.png", wantOK: true}, + {name: "save to subdir", prompt: "save to images/out.png", wantName: "images/out.png", wantOK: true}, + {name: "write to", prompt: "write to pics/x.png", wantName: "pics/x.png", wantOK: true}, + {name: "output to", prompt: "output to result.png", wantName: "result.png", wantOK: true}, + {name: "filename colon", prompt: "filename: sunset.png", wantName: "sunset.png", wantOK: true}, + {name: "filename equals", prompt: "filename=sunset.png", wantName: "sunset.png", wantOK: true}, + {name: "filename spaced equals", prompt: "filename = sunset.png", wantName: "sunset.png", wantOK: true}, + {name: "double-quoted with space", prompt: `save it as "my picture.png"`, wantName: "my picture.png", wantOK: true}, + {name: "single-quoted", prompt: "name it 'logo.webp'", wantName: "logo.webp", wantOK: true}, + {name: "backticked subdir", prompt: "save it as `pics/cat.png`", wantName: "pics/cat.png", wantOK: true}, + {name: "trailing sentence punctuation", prompt: "save it as sunshine.jpg.", wantName: "sunshine.jpg", wantOK: true}, + {name: "trailing comma", prompt: "make an image as sunshine.jpg, please", wantName: "sunshine.jpg", wantOK: true}, + {name: "uppercase cue and extension", prompt: "Save As PHOTO.JPG", wantName: "PHOTO.JPG", wantOK: true}, + {name: "traversal is a valid untrusted candidate", prompt: "save to ../evil.png", wantName: "../evil.png", wantOK: true}, + {name: "absolute is a valid untrusted candidate", prompt: "write to /tmp/out.png", wantName: "/tmp/out.png", wantOK: true}, + {name: "exactly at the byte bound", prompt: "save as " + atBoundName, wantName: atBoundName, wantOK: true}, + + {name: "input-file mention is not a cue", prompt: "add a border to photo.jpg"}, + {name: "comparative 'same style as' is a reference, not output intent", prompt: "Generate an image in the same style as sunshine.jpg"}, + {name: "comparative inside an of-phrase is a reference, not output intent", prompt: "Generate an image of a cat in the same style as old-render.png"}, + {name: "'like' inside an of-phrase is a reference, not output intent", prompt: "Generate an image of a panda just like my avatar as panda.png"}, + {name: "of-phrase without an 'as' cue", prompt: "Generate an image of assets/red-panda-terminal.jpg"}, + {name: "comparative 'same background as' is a reference, not output intent", prompt: "Make a banner with the same background as assets/bg.png"}, + {name: "comparative 'exact palette as' is a reference, not output intent", prompt: "Generate an image with the exact palette as sunshine.jpg"}, + {name: "comparative 'identical colors as' is a reference, not output intent", prompt: "Generate an image with identical colors as assets/ref.png"}, + {name: "'exact palette' inside an of-phrase is a reference, not output intent", prompt: "Generate an image of a beach with the exact palette as sunshine.jpg"}, + {name: "'identical colors' inside an of-phrase is a reference, not output intent", prompt: "Generate an image of a beach with identical colors as assets/ref.png"}, + {name: "comparative 'as tall as' inside an of-phrase is a reference, not output intent", prompt: "Generate an image of a tower as tall as tree.png"}, + {name: "attribute prepositions in an of-phrase subject are conservatively rejected", prompt: "Generate an image of a cat in a hat as cat-hat.png"}, + {name: "'the one called' references an existing file", prompt: "Generate an image similar to the one called old-render.png"}, + {name: "'called' without imperative output context", prompt: "an image called loop.gif"}, + {name: "bare 'as' without a generation-verb anchor", prompt: "as sunshine.jpg, please"}, + {name: "comparative 'as wide as' after a media noun", prompt: "make the image as wide as banner.png"}, + {name: "plural noun breaks the verb anchor", prompt: "generate images such as sunset.png"}, + {name: "plain filename mention without cue", prompt: "here is photo.jpg"}, + {name: "no filename at all", prompt: "generate an image of a sunset"}, + {name: "prose 'as' without a filename", prompt: "make it as good as new"}, + {name: "cue with non-image extension", prompt: `save as "notes.txt"`}, + {name: "cue with bare word", prompt: "call it Bob"}, + {name: "extension only", prompt: "save it as .png"}, + {name: "quoted extension only", prompt: `call it ".png"`}, + {name: "two candidates are ambiguous", prompt: "save it as a.png and call it b.png"}, + {name: "duplicate candidates are still ambiguous", prompt: "save as x.png, yes save as x.png"}, + {name: "distinct of-phrase and cue candidates are ambiguous", prompt: "Generate an image of a dog as dog.png and save it as cat.png"}, + {name: "same name at distinct occurrences is still ambiguous", prompt: "make an image of a sun as sun.png and save it as sun.png"}, + {name: "control character in name", prompt: "save as bad\x01name.png"}, + {name: "invalid UTF-8 in name", prompt: "save as \xff.png"}, + {name: "over the byte bound", prompt: "save as " + strings.Repeat("a", maxExplicitOutputFilenameBytes) + ".png"}, + {name: "empty prompt", prompt: ""}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + name, ok := extractExplicitOutputFilename(tt.prompt) + assert.Equal(t, tt.wantOK, ok) + assert.Equal(t, tt.wantName, name) + }) + } +} + +// TestApplyUserPromptRequestedPath pins the fallback's guard conditions: +// exactly one blob, still unnamed after marker pairing, unambiguous prompt. +func TestApplyUserPromptRequestedPath(t *testing.T) { + t.Parallel() + + t.Run("names the single unnamed blob", func(t *testing.T) { + t.Parallel() + media := []chat.MediaDelta{{Name: "provider"}} + applyUserPromptRequestedPath(media, "Generate an image as sunshine.jpg") + assert.Equal(t, "sunshine.jpg", media[0].RequestedPath) + }) + + t.Run("a marker-named blob keeps its path", func(t *testing.T) { + t.Parallel() + media := []chat.MediaDelta{{RequestedPath: "marker.png"}} + applyUserPromptRequestedPath(media, "Generate an image as sunshine.jpg") + assert.Equal(t, "marker.png", media[0].RequestedPath, "marker precedence must win over the prompt") + }) + + t.Run("multi-blob turns get no fallback", func(t *testing.T) { + t.Parallel() + media := []chat.MediaDelta{{}, {}} + applyUserPromptRequestedPath(media, "Generate an image as sunshine.jpg") + assert.Empty(t, media[0].RequestedPath) + assert.Empty(t, media[1].RequestedPath) + }) + + t.Run("partially marker-named multi-blob turns get no fallback", func(t *testing.T) { + t.Parallel() + media := []chat.MediaDelta{{RequestedPath: "named.png"}, {}} + applyUserPromptRequestedPath(media, "Generate an image as sunshine.jpg") + assert.Empty(t, media[1].RequestedPath) + }) + + t.Run("ambiguous prompt leaves the blob unnamed", func(t *testing.T) { + t.Parallel() + media := []chat.MediaDelta{{}} + applyUserPromptRequestedPath(media, "save as a.png or save as b.png") + assert.Empty(t, media[0].RequestedPath) + }) + + t.Run("comparative reference leaves the blob unnamed", func(t *testing.T) { + t.Parallel() + media := []chat.MediaDelta{{}} + applyUserPromptRequestedPath(media, "Generate an image in the same style as sunshine.jpg") + assert.Empty(t, media[0].RequestedPath) + }) +} + +// promptWorkspaceSession is workspaceSession plus the triggering user +// message, so materialization sees the prompt the fallback must parse. +func promptWorkspaceSession(t *testing.T, id, prompt string) (*session.Session, string) { + t.Helper() + sess, root := workspaceSession(t, id) + sess.AddMessage(session.UserMessage(prompt)) + return sess, root +} + +// TestUserPromptFilenameNamesSingleUnmarkedBlob is the live-repro +// integration: the exact prompt "Generate an image as sunshine.jpg", a model +// that ignores the marker instruction, and one PNG blob. The prompt filename +// must name the file, with the writer still correcting the extension to the +// actual MIME type and surfacing the correction notice. +func TestUserPromptFilenameNamesSingleUnmarkedBlob(t *testing.T) { + r, _, _ := newMediaTestRuntime(t) + sess, root := promptWorkspaceSession(t, "sess-prompt-name", "Generate an image as sunshine.jpg") + + stream := newStreamBuilder(). + AddContent("Here you go!\n"). + AddMedia([]byte{0x89, 0x50, 0x4e, 0x47}, "image/png", "provider-name.png"). + AddStopWithUsage(1, 1). + Build() + media := markerTurnMedia(t, stream, "Here you go!\n") + + sink := &collectingSink{} + parts := r.materializeGeneratedMedia(t.Context(), sess, media, "root", sink) + + require.Len(t, parts, 1) + assert.Equal(t, "sunshine.png", parts[0].Document.Source.ArtifactPath, "the prompt filename must win over the provider name, with the MIME-derived extension") + _, err := os.Stat(filepath.Join(root, "sunshine.png")) + require.NoError(t, err) + + warnings := sink.warnings() + require.Len(t, warnings, 1, "the extension correction must be user-visible") + assert.Contains(t, warnings[0].Message, "sunshine.png") + assert.Contains(t, warnings[0].Message, `".jpg"`) +} + +// TestUserPromptOfPhraseFilenameMaterializesEndToEnd is the live repro of +// the "of as " prompt form: a model that ignores the marker +// instruction returns one generic unnamed PNG blob, and the prompt-directed +// subdirectory path must still become the persisted final path — with the +// writer's MIME/extension correction — and the part's Document.Name the +// final basename the UI labels the image with. +func TestUserPromptOfPhraseFilenameMaterializesEndToEnd(t *testing.T) { + r, _, _ := newMediaTestRuntime(t) + sess, root := promptWorkspaceSession(t, "sess-prompt-of-phrase", + "Generate an image of a red panda coding at a terminal as assets/red-panda-terminal.jpg") + + stream := newStreamBuilder(). + AddContent("Here is your red panda coding at a terminal:\n"). + AddMedia([]byte{0x89, 0x50, 0x4e, 0x47}, "image/png", ""). + AddStopWithUsage(1, 1). + Build() + media := markerTurnMedia(t, stream, "Here is your red panda coding at a terminal:\n") + + sink := &collectingSink{} + parts := r.materializeGeneratedMedia(t.Context(), sess, media, "root", sink) + + require.Len(t, parts, 1) + assert.Equal(t, "assets/red-panda-terminal.png", parts[0].Document.Source.ArtifactPath, + "the prompt-directed path must win over the generic generated-N name, with the MIME-derived extension") + assert.Equal(t, "red-panda-terminal.png", parts[0].Document.Name, + "the persisted display name must be the final basename") + _, err := os.Stat(filepath.Join(root, "assets", "red-panda-terminal.png")) + require.NoError(t, err) + + warnings := sink.warnings() + require.Len(t, warnings, 1, "the extension correction must be user-visible") + assert.Contains(t, warnings[0].Message, "assets/red-panda-terminal.png") + assert.Contains(t, warnings[0].Message, `".jpg"`) +} + +// TestUserPromptOfPhraseFilenameCollisionSuffixPersists: when the +// prompt-directed target already exists, the collision-suffixed path the +// writer actually used is what the part persists — ArtifactPath and +// Document.Name both name the suffixed file, never the requested one. +func TestUserPromptOfPhraseFilenameCollisionSuffixPersists(t *testing.T) { + r, _, _ := newMediaTestRuntime(t) + sess, root := promptWorkspaceSession(t, "sess-prompt-of-phrase-collision", + "Generate an image of a red panda coding at a terminal as assets/red-panda-terminal.png") + require.NoError(t, os.MkdirAll(filepath.Join(root, "assets"), 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(root, "assets", "red-panda-terminal.png"), []byte{0x01}, 0o644)) + + stream := newStreamBuilder(). + AddMedia([]byte{0x89, 0x50, 0x4e, 0x47}, "image/png", ""). + AddStopWithUsage(1, 1). + Build() + media := markerTurnMedia(t, stream, "") + + sink := &collectingSink{} + parts := r.materializeGeneratedMedia(t.Context(), sess, media, "root", sink) + + require.Len(t, parts, 1) + assert.Equal(t, "assets/red-panda-terminal-1.png", parts[0].Document.Source.ArtifactPath) + assert.Equal(t, "red-panda-terminal-1.png", parts[0].Document.Name) + _, err := os.Stat(filepath.Join(root, "assets", "red-panda-terminal-1.png")) + require.NoError(t, err) + assert.Empty(t, sink.warnings(), "a collision suffix alone is not a user-visible correction") +} + +// TestOverlappingGrammarPromptsMaterializeEndToEnd pins the end-to-end +// behavior of prompts where the of-phrase grammar overlaps the explicit-cue +// grammar: one occurrence captured by both grammars (conjunction swallowed +// into the subject, or nothing once CR/LF stops the subject) must still +// name the file, while genuinely distinct candidates stay ambiguous and +// keep the generic generated-N name. +func TestOverlappingGrammarPromptsMaterializeEndToEnd(t *testing.T) { + tests := map[string]struct { + prompt string + wantPath string + }{ + "of-phrase conjunction": {"make a logo of a wolf and save it as logo.png", "logo.png"}, + "image-noun conjunction": {"Generate an image of a cat and save it as cat.png", "cat.png"}, + "newline-separated cue": {"Generate an image of a cat\nSave it as x.png", "x.png"}, + "distinct candidates": {"Generate an image of a dog as dog.png and save it as cat.png", "generated-1.png"}, + } + + for name, tt := range tests { + t.Run(name, func(t *testing.T) { + r, _, _ := newMediaTestRuntime(t) + sess, root := promptWorkspaceSession(t, "sess-prompt-overlap-"+name, tt.prompt) + + stream := newStreamBuilder(). + AddMedia([]byte{0x89, 0x50, 0x4e, 0x47}, "image/png", ""). + AddStopWithUsage(1, 1). + Build() + media := markerTurnMedia(t, stream, "") + + sink := &collectingSink{} + parts := r.materializeGeneratedMedia(t.Context(), sess, media, "root", sink) + + require.Len(t, parts, 1) + assert.Equal(t, tt.wantPath, parts[0].Document.Source.ArtifactPath) + assert.Equal(t, tt.wantPath, parts[0].Document.Name) + _, err := os.Stat(filepath.Join(root, tt.wantPath)) + require.NoError(t, err) + assert.Empty(t, sink.warnings()) + }) + } +} + +// TestMarkerOverridesUserPromptFilename proves precedence: a compliant model +// emitting a marker wins over the explicit filename in the prompt. +func TestMarkerOverridesUserPromptFilename(t *testing.T) { + r, _, _ := newMediaTestRuntime(t) + sess, root := promptWorkspaceSession(t, "sess-prompt-marker", "Generate an image as sunshine.jpg") + + stream := newStreamBuilder(). + AddContent("[media-file: marker-pick.png]\n"). + AddMedia([]byte{0x89, 0x50, 0x4e, 0x47}, "image/png", ""). + AddStopWithUsage(1, 1). + Build() + media := markerTurnMedia(t, stream, "") + + sink := &collectingSink{} + parts := r.materializeGeneratedMedia(t.Context(), sess, media, "root", sink) + + require.Len(t, parts, 1) + assert.Equal(t, "marker-pick.png", parts[0].Document.Source.ArtifactPath) + _, err := os.Stat(filepath.Join(root, "sunshine.png")) + assert.True(t, os.IsNotExist(err), "the prompt filename must not be used when a marker named the blob") +} + +// TestUserPromptFilenameSkipsMultiBlobTurns: one prompt filename cannot +// unambiguously name one of several blobs, so both fall back to the +// provider-then-generic naming unchanged. +func TestUserPromptFilenameSkipsMultiBlobTurns(t *testing.T) { + r, _, _ := newMediaTestRuntime(t) + sess, root := promptWorkspaceSession(t, "sess-prompt-multi", "Generate an image as sunshine.jpg") + + stream := newStreamBuilder(). + AddMultiMedia( + chat.MediaDelta{Data: []byte{0x01}, MimeType: "image/png", Name: "provider-pick", Size: 1}, + chat.MediaDelta{Data: []byte{0x02}, MimeType: "image/png", Size: 1}, + ). + AddStopWithUsage(1, 1). + Build() + media := markerTurnMedia(t, stream, "") + + sink := &collectingSink{} + parts := r.materializeGeneratedMedia(t.Context(), sess, media, "root", sink) + + require.Len(t, parts, 2) + assert.Equal(t, "provider-pick.png", parts[0].Document.Source.ArtifactPath) + assert.Equal(t, "generated-2.png", parts[1].Document.Source.ArtifactPath) + _, err := os.Stat(filepath.Join(root, "sunshine.png")) + assert.True(t, os.IsNotExist(err)) +} + +// TestUserPromptWithoutExplicitFilenameKeepsProviderFallback: a prompt with +// no explicit filename leaves the existing provider-name fallback untouched. +func TestUserPromptWithoutExplicitFilenameKeepsProviderFallback(t *testing.T) { + r, _, _ := newMediaTestRuntime(t) + sess, _ := promptWorkspaceSession(t, "sess-prompt-none", "Draw me a sunny landscape please") + + stream := newStreamBuilder(). + AddMedia([]byte{0x01}, "image/png", "provider-pick"). + AddStopWithUsage(1, 1). + Build() + media := markerTurnMedia(t, stream, "") + + sink := &collectingSink{} + parts := r.materializeGeneratedMedia(t.Context(), sess, media, "root", sink) + + require.Len(t, parts, 1) + assert.Equal(t, "provider-pick.png", parts[0].Document.Source.ArtifactPath) + assert.Empty(t, sink.warnings()) +} + +// TestUserPromptEscapingFilenameRedirectsThroughEscapePolicy proves an +// extracted traversing filename reaches the existing escape policy +// unchanged: with no user to ask (non-interactive), the bytes are redirected +// into the workspace under the sanitized basename. +func TestUserPromptEscapingFilenameRedirectsThroughEscapePolicy(t *testing.T) { + r, _ := newEscapeTestRuntime(t) + r.nonInteractive = true + sess, root := workspaceSession(t, "sess-prompt-escape") + sess.AddMessage(session.UserMessage("save to ../outside.png")) + + stream := newStreamBuilder(). + AddMedia([]byte{0xAA}, "image/png", ""). + AddStopWithUsage(1, 1). + Build() + media := markerTurnMedia(t, stream, "") + + sink := &collectingSink{} + parts := r.materializeGeneratedMedia(t.Context(), sess, media, "root", sink) + + require.Len(t, parts, 1) + assert.Equal(t, "outside.png", parts[0].Document.Source.ArtifactPath) + assert.Equal(t, chat.ArtifactRootWorkspace, parts[0].Document.Source.ArtifactRoot) + _, err := os.Stat(filepath.Join(root, "outside.png")) + require.NoError(t, err) + _, err = os.Stat(filepath.Join(filepath.Dir(root), "outside.png")) + assert.True(t, os.IsNotExist(err), "nothing may land outside the workspace without confirmation") + require.Len(t, sink.warnings(), 1, "the redirect must be explained to the user") +} + +// TestComparativeReferencePromptsNeverNameGeneratedMedia is the +// false-capture regression for prompts that mention a filename only as a +// comparison or reference to an existing file: no extraction, so the blob +// keeps the generic name and the escape policy is never engaged (no +// confirmation, no redirect warning, no file under the referenced name). +func TestComparativeReferencePromptsNeverNameGeneratedMedia(t *testing.T) { + prompts := map[string]string{ + "same style": "Generate an image in the same style as sunshine.jpg", + "same background": "Make a banner with the same background as assets/bg.png", + "the one called": "Generate an image similar to the one called old-render.png", + "same style traversing": "Generate an image in the same style as ../sunshine.jpg", + "exact palette": "Generate an image with the exact palette as sunshine.jpg", + "identical colors": "Generate an image with identical colors as assets/ref.png", + "exact palette of-phrase": "Generate an image of a beach with the exact palette as sunshine.jpg", + "identical colors of-phrase": "Generate an image of a beach with identical colors as assets/ref.png", + "exact palette traversing": "Generate an image of a beach with the exact palette as ../sunshine.jpg", + } + + for name, prompt := range prompts { + t.Run(name, func(t *testing.T) { + r, _ := newEscapeTestRuntime(t) + // Non-interactive: a falsely extracted escaping path would surface + // as a redirect warning instead of deadlocking on a confirmation. + r.nonInteractive = true + sess, root := promptWorkspaceSession(t, "sess-prompt-comparative-"+name, prompt) + + stream := newStreamBuilder(). + AddMedia([]byte{0x89, 0x50, 0x4e, 0x47}, "image/png", ""). + AddStopWithUsage(1, 1). + Build() + media := markerTurnMedia(t, stream, "") + + sink := &collectingSink{} + parts := r.materializeGeneratedMedia(t.Context(), sess, media, "root", sink) + + require.Len(t, parts, 1) + assert.Equal(t, "generated-1.png", parts[0].Document.Source.ArtifactPath, "a comparative filename mention must never name the output") + assert.Empty(t, sink.warnings(), "the escape policy must never be engaged for a comparative mention") + entries, err := os.ReadDir(root) + require.NoError(t, err) + require.Len(t, entries, 1, "only the generically named file may exist in the workspace") + assert.Equal(t, "generated-1.png", entries[0].Name()) + }) + } +} diff --git a/pkg/runtime/loop.go b/pkg/runtime/loop.go index 082b2c499..4d8dc2f91 100644 --- a/pkg/runtime/loop.go +++ b/pkg/runtime/loop.go @@ -1315,10 +1315,12 @@ func sanitizeToolCallName(name string) string { // generated-media manifest ([session.GeneratedMediaManifest]), the trust // anchor a resolver must consult before reading a workspace path back. // -// The requested filename is the prompt-directed path when the provider -// surfaced one ([chat.MediaDelta.RequestedPath], populated by marker -// extraction once that lands), otherwise the sanitized provider display -// name, otherwise a generic "generated-N"; the writer owns MIME/extension +// The requested filename is the prompt-directed path when one exists +// ([chat.MediaDelta.RequestedPath], populated by media-file marker pairing +// in handleStream, or — for a turn whose single blob no marker named — by +// deterministic explicit-filename extraction from the triggering user +// message, see [applyUserPromptRequestedPath]), otherwise the sanitized +// provider display name, otherwise a generic "generated-N"; the writer owns MIME/extension // correction and collision suffixing, and the part persists the exact final // path it returns. A corrected extension additionally surfaces a bounded // user-visible notice naming the final path. A prompt-directed path that @@ -1359,6 +1361,11 @@ func sanitizeToolCallName(name string) string { // successfully in the same reply, and must not lose the (already generated) // accompanying text either. func (r *LocalRuntime) materializeGeneratedMedia(ctx context.Context, sess *session.Session, media []chat.MediaDelta, agentName string, events EventSink) []chat.MessagePart { + // Runs after marker pairing (handleStream already filled RequestedPath + // for marker-named blobs) and before any write, so marker precedence and + // the untrusted-path pipeline below apply unchanged. + applyUserPromptRequestedPath(media, sess.GetLastUserMessageContent()) + root, rootErr := session.ResolveWorkingDir(ctx, sess, r.sessionLookup()) if rootErr != nil { slog.DebugContext(ctx, "No workspace root for generated media; dropping every media item, keeping the rest of the turn", diff --git a/pkg/runtime/media_escape.go b/pkg/runtime/media_escape.go index 6c75985d6..ebde8b16d 100644 --- a/pkg/runtime/media_escape.go +++ b/pkg/runtime/media_escape.go @@ -20,9 +20,9 @@ import ( // generated media blob through the naming and workspace-escape policy. type generatedMediaItem struct { workspaceRoot string - // requestedPath is the prompt-directed target (untrusted model input, - // see chat.MediaDelta.RequestedPath); empty when the model named - // nothing explicitly. + // requestedPath is the prompt-directed target (untrusted marker or + // user-prompt input, see chat.MediaDelta.RequestedPath); empty when + // nothing named the blob explicitly. requestedPath string // providerName is the sanitized provider display name; genericName the // deterministic "generated-N" fallback.