From fc76f6a54d9dc7f9e55183dca2f075518013ed0d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Sun, 16 Aug 2026 10:25:34 +0200 Subject: [PATCH] feat(#3996): stream model-generated media deltas MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extend the shared chat.MessageDelta streaming representation with typed MediaDelta entries (bytes + MIME + name + size) so generated binary media has a provider-agnostic path instead of a Gemini-only side channel. The Gemini adapter captures inline image blobs from response candidates the same way it already captures text/reasoning/tool calls/usage — retaining every blob in a chunk, not just the last — and the run() content gate forwards image-only chunks that previously would have been dropped. The runtime's stream accumulator (pkg/runtime/streaming.go) collects the deltas into streamResult.Media, accumulating before any terminal check so a final chunk that carries both media and a finish reason is never lost. Persistence of the accumulated media ships in the follow-up materialization commits. --- pkg/chat/chat.go | 7 + pkg/chat/media.go | 29 ++++ pkg/model/provider/gemini/adapter.go | 35 +++- pkg/model/provider/gemini/adapter_test.go | 145 +++++++++++++++++ pkg/runtime/streaming.go | 29 +++- pkg/runtime/streaming_test.go | 185 ++++++++++++++++++++++ 6 files changed, 421 insertions(+), 9 deletions(-) create mode 100644 pkg/chat/media.go diff --git a/pkg/chat/chat.go b/pkg/chat/chat.go index 99727ff2f6..ebedf2c6f6 100644 --- a/pkg/chat/chat.go +++ b/pkg/chat/chat.go @@ -153,6 +153,13 @@ type MessageDelta struct { ThoughtSignature []byte `json:"thought_signature,omitempty"` FunctionCall *tools.FunctionCall `json:"function_call,omitempty"` ToolCalls []tools.ToolCall `json:"tool_calls,omitempty"` + // Media carries provider-generated binary content (e.g. inline image + // blobs) streamed alongside or instead of text. A single chunk can carry + // more than one blob (e.g. Gemini returning multiple inline parts across + // candidates in one chunk), so this is a slice rather than a single + // pointer — a scalar field silently dropped every blob but the last. See + // [MediaDelta]. + Media []MediaDelta `json:"media,omitempty"` } // MessageStreamChoice represents a choice in a streaming response diff --git a/pkg/chat/media.go b/pkg/chat/media.go new file mode 100644 index 0000000000..0ea0c48f99 --- /dev/null +++ b/pkg/chat/media.go @@ -0,0 +1,29 @@ +package chat + +// MediaDelta carries binary media generated by a provider while a response +// is streaming (e.g. an inline image blob from Gemini). It is the shared, +// provider-agnostic representation for this content: adapters for any +// provider populate it the same way, and the runtime accumulator is the only +// consumer, so no provider gets a side channel of its own. +// +// Data holds the live bytes only for the duration of the stream. Once the +// turn completes, the runtime materializes them into a session artifact and +// the persisted assistant message keeps only a reference — see +// [DocumentSource.ArtifactPath]. Data is never written to session storage. +type MediaDelta struct { + // Data is the raw generated bytes. Present only on the streaming delta; + // never persisted. + Data []byte `json:"data,omitempty"` + + // MimeType is the media's MIME type as reported by the provider (e.g. + // "image/png"). + MimeType string `json:"mime_type,omitempty"` + + // Name is the provider-supplied display name, if any. May be empty; the + // runtime accumulator synthesizes one when needed. + Name string `json:"name,omitempty"` + + // Size is the byte length of Data, cached because Data itself is + // dropped once the artifact is materialized. + Size int64 `json:"size,omitempty"` +} diff --git a/pkg/model/provider/gemini/adapter.go b/pkg/model/provider/gemini/adapter.go index 77552a813e..0140e924ca 100644 --- a/pkg/model/provider/gemini/adapter.go +++ b/pkg/model/provider/gemini/adapter.go @@ -112,18 +112,22 @@ func (g *StreamAdapter) run() { } if resp != nil { - // Check for text content without using Text() to avoid warnings + // Check for text content and generated inline media without using + // Text() to avoid warnings hasText := false + hasMedia := false for _, candidate := range resp.Candidates { if candidate.Content != nil { for _, part := range candidate.Content.Parts { if part.Text != "" { hasText = true - break + } + if part.InlineData != nil && len(part.InlineData.Data) > 0 { + hasMedia = true } } } - if hasText { + if hasText && hasMedia { break } } @@ -134,9 +138,9 @@ func (g *StreamAdapter) run() { // calls. Forward such chunks so downstream can capture token usage. hasUsage := resp.UsageMetadata != nil - // Send response if it has content, function calls, or usage metadata - if hasText || hasFuncs || hasUsage { - hasContent = hasContent || hasText + // Send response if it has content, generated media, function calls, or usage metadata + if hasText || hasMedia || hasFuncs || hasUsage { + hasContent = hasContent || hasText || hasMedia hasToolCalls = hasToolCalls || hasFuncs lastResponse = resp // Store for final message if !g.send(result{resp: resp}) { @@ -232,6 +236,7 @@ func (g *StreamAdapter) Recv() (chat.MessageStreamResponse, error) { var reasoningTextSb strings.Builder var textContentSb strings.Builder var thoughtSignature []byte + var media []chat.MediaDelta for _, candidate := range res.resp.Candidates { if candidate.Content != nil { for _, part := range candidate.Content.Parts { @@ -246,6 +251,21 @@ func (g *StreamAdapter) Recv() (chat.MessageStreamResponse, error) { textContentSb.WriteString(part.Text) } } + + // Inline generated media (e.g. an image from an + // image-output model). Gemini can return more than one + // inline blob per chunk — multiple parts in one candidate, + // or multiple candidates — so every blob is appended + // rather than overwriting a single field, which used to + // silently drop all but the last one. + if part.InlineData != nil && len(part.InlineData.Data) > 0 { + media = append(media, chat.MediaDelta{ + Data: part.InlineData.Data, + MimeType: part.InlineData.MIMEType, + Name: part.InlineData.DisplayName, + Size: int64(len(part.InlineData.Data)), + }) + } } } } @@ -260,6 +280,9 @@ func (g *StreamAdapter) Recv() (chat.MessageStreamResponse, error) { if len(thoughtSignature) > 0 { resp.Choices[0].Delta.ThoughtSignature = thoughtSignature } + if len(media) > 0 { + resp.Choices[0].Delta.Media = media + } // Handle function calls if funcs := res.resp.FunctionCalls(); len(funcs) > 0 { diff --git a/pkg/model/provider/gemini/adapter_test.go b/pkg/model/provider/gemini/adapter_test.go index fccdd8a61c..ad7844b07c 100644 --- a/pkg/model/provider/gemini/adapter_test.go +++ b/pkg/model/provider/gemini/adapter_test.go @@ -4,6 +4,7 @@ import ( "io" "testing" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "google.golang.org/genai" @@ -183,3 +184,147 @@ func TestStreamAdapter_FunctionCalls(t *testing.T) { require.Empty(t, finalResp.Choices[0].Delta.ToolCalls) }) } + +func TestStreamAdapter_GeneratedImage(t *testing.T) { + t.Parallel() + + t.Run("inline image-only chunk is forwarded as Delta.Media", func(t *testing.T) { + imgBytes := []byte{0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a} + mockResp := &genai.GenerateContentResponse{ + Candidates: []*genai.Candidate{ + { + Content: &genai.Content{ + Parts: []*genai.Part{ + { + InlineData: &genai.Blob{ + Data: imgBytes, + MIMEType: "image/png", + DisplayName: "cat.png", + }, + }, + }, + }, + }, + }, + } + + iter := func(fn func(*genai.GenerateContentResponse, error) bool) { + fn(mockResp, nil) + } + adapter := NewStreamAdapter(iter, "test-model", true) + + // A chunk carrying only inline media (no text, no function calls, no + // usage) must still be forwarded rather than dropped by the run() + // content gate. + resp, err := adapter.Recv() + require.NoError(t, err) + require.Len(t, resp.Choices[0].Delta.Media, 1, "image-only chunk must be forwarded") + assert.Equal(t, imgBytes, resp.Choices[0].Delta.Media[0].Data) + assert.Equal(t, "image/png", resp.Choices[0].Delta.Media[0].MimeType) + assert.Equal(t, "cat.png", resp.Choices[0].Delta.Media[0].Name) + assert.Equal(t, int64(len(imgBytes)), resp.Choices[0].Delta.Media[0].Size) + assert.Empty(t, resp.Choices[0].Delta.Content, "no text was in this chunk") + + // The synthesized done event must still report a normal stop — + // generated media does not change finish-reason handling. + done, err := adapter.Recv() + require.NoError(t, err) + assert.Equal(t, chat.FinishReasonStop, done.Choices[0].FinishReason) + }) + + t.Run("text and inline image in the same chunk both surface", func(t *testing.T) { + imgBytes := []byte{0x01, 0x02, 0x03} + mockResp := &genai.GenerateContentResponse{ + Candidates: []*genai.Candidate{ + { + Content: &genai.Content{ + Parts: []*genai.Part{ + {Text: "here you go"}, + {InlineData: &genai.Blob{Data: imgBytes, MIMEType: "image/jpeg"}}, + }, + }, + }, + }, + } + + iter := func(fn func(*genai.GenerateContentResponse, error) bool) { + fn(mockResp, nil) + } + adapter := NewStreamAdapter(iter, "test-model", true) + + resp, err := adapter.Recv() + require.NoError(t, err) + assert.Equal(t, "here you go", resp.Choices[0].Delta.Content, "text handling must remain intact alongside media") + require.Len(t, resp.Choices[0].Delta.Media, 1) + assert.Equal(t, imgBytes, resp.Choices[0].Delta.Media[0].Data) + assert.Equal(t, "image/jpeg", resp.Choices[0].Delta.Media[0].MimeType) + // Provider omitted a display name; the adapter must not invent one — + // that is the runtime accumulator's job. + assert.Empty(t, resp.Choices[0].Delta.Media[0].Name) + }) + + t.Run("empty inline data is not surfaced as media", func(t *testing.T) { + // A Blob with zero bytes must not produce a spurious empty Media + // delta (and, since it carries no text either, must not be + // forwarded as a chunk at all). + mockResp := &genai.GenerateContentResponse{ + Candidates: []*genai.Candidate{ + { + Content: &genai.Content{ + Parts: []*genai.Part{ + {InlineData: &genai.Blob{Data: nil, MIMEType: "image/png"}}, + }, + }, + }, + }, + } + + iter := func(fn func(*genai.GenerateContentResponse, error) bool) { + fn(mockResp, nil) + } + adapter := NewStreamAdapter(iter, "test-model", true) + + // Nothing at all was produced (no text, no media, no tool calls, no + // usage), so the stream ends directly with EOF — same as any other + // content-free chunk, no synthesized done event. + _, err := adapter.Recv() + require.ErrorIs(t, err, io.EOF) + }) + + t.Run("multiple inline blobs across parts and candidates in one chunk are all retained", func(t *testing.T) { + // Gemini can pack more than one generated blob into a single chunk: + // multiple parts within a candidate, and/or multiple candidates. + // Every blob must survive, not just the last one seen. + mockResp := &genai.GenerateContentResponse{ + Candidates: []*genai.Candidate{ + { + Content: &genai.Content{ + Parts: []*genai.Part{ + {InlineData: &genai.Blob{Data: []byte{0x01}, MIMEType: "image/png", DisplayName: "first.png"}}, + {InlineData: &genai.Blob{Data: []byte{0x02}, MIMEType: "image/jpeg", DisplayName: "second.jpg"}}, + }, + }, + }, + { + Content: &genai.Content{ + Parts: []*genai.Part{ + {InlineData: &genai.Blob{Data: []byte{0x03}, MIMEType: "image/webp", DisplayName: "third.webp"}}, + }, + }, + }, + }, + } + + iter := func(fn func(*genai.GenerateContentResponse, error) bool) { + fn(mockResp, nil) + } + adapter := NewStreamAdapter(iter, "test-model", true) + + resp, err := adapter.Recv() + require.NoError(t, err) + require.Len(t, resp.Choices[0].Delta.Media, 3, "every blob in the chunk must be retained") + assert.Equal(t, "first.png", resp.Choices[0].Delta.Media[0].Name) + assert.Equal(t, "second.jpg", resp.Choices[0].Delta.Media[1].Name) + assert.Equal(t, "third.webp", resp.Choices[0].Delta.Media[2].Name) + }) +} diff --git a/pkg/runtime/streaming.go b/pkg/runtime/streaming.go index abd43bb022..1ae41a5fbb 100644 --- a/pkg/runtime/streaming.go +++ b/pkg/runtime/streaming.go @@ -47,9 +47,13 @@ type streamResult struct { ReasoningContent string ThinkingSignature string ThoughtSignature []byte - Stopped bool - FinishReason chat.FinishReason - Usage *chat.Usage + // Media accumulates every [chat.MediaDelta] streamed during the turn + // (e.g. generated images). Populated regardless of provider — see + // chat.MessageDelta.Media. + Media []chat.MediaDelta + Stopped bool + FinishReason chat.FinishReason + Usage *chat.Usage } // handleStream reads a chat.MessageStream to completion, emitting streaming @@ -110,6 +114,7 @@ func handleStream(ctx context.Context, cancelStream context.CancelCauseFunc, str var thinkingSignature string var thoughtSignature []byte var toolCalls []tools.ToolCall + var media []chat.MediaDelta var messageUsage *chat.Usage var providerFinishReason chat.FinishReason @@ -205,6 +210,20 @@ mainLoop: thoughtSignature = choice.Delta.ThoughtSignature } + // Provider-generated binary media (e.g. inline image blobs) is + // accumulated up front, before the terminal finish-reason check + // below. A provider (e.g. Gemini) can pack generated media and a + // terminal finish_reason ("stop"/"length"/"refusal") into the SAME + // chunk; accumulating after that check would return before this + // chunk's media was ever added, silently dropping it. Not emitted + // as a streaming event: no consumer needs a live partial-media + // event yet, and the runtime materializes the final bytes into a + // session artifact once the turn completes (see + // recordAssistantMessage). + if len(choice.Delta.Media) > 0 { + media = append(media, choice.Delta.Media...) + } + // Accumulate tool call deltas from this chunk *before* evaluating the // finish reason below. Some OpenAI-compatible providers (e.g. LiteLLM // in front of Gemini) pack a complete tool call and a terminal @@ -296,6 +315,7 @@ mainLoop: ReasoningContent: fullReasoningContent.String(), ThinkingSignature: thinkingSignature, ThoughtSignature: thoughtSignature, + Media: media, Stopped: len(toolCalls) == 0, // stop only when there are no tool calls to execute FinishReason: finishReason, Usage: messageUsage, @@ -362,6 +382,8 @@ mainLoop: // Invariant: a bare-EOF turn (no per-choice finish_reason) is terminal // whenever there are no tool calls — the outer loop has nothing to continue // on. Turns with tool calls keep Stopped=false so the loop executes them. + // Media is irrelevant to this decision: with no tool calls pending, the turn + // is over either way (media or not) — stopping is what ends it correctly. // NOTE(krissetto): this can likely be removed once compaction works properly with all providers (aka dmr) stoppedNoToolCalls := len(toolCalls) == 0 @@ -396,6 +418,7 @@ mainLoop: ReasoningContent: fullReasoningContent.String(), ThinkingSignature: thinkingSignature, ThoughtSignature: thoughtSignature, + Media: media, Stopped: stoppedNoToolCalls, FinishReason: finishReason, Usage: messageUsage, diff --git a/pkg/runtime/streaming_test.go b/pkg/runtime/streaming_test.go index 129aa66e72..321b89f659 100644 --- a/pkg/runtime/streaming_test.go +++ b/pkg/runtime/streaming_test.go @@ -16,6 +16,57 @@ import ( "github.com/docker/docker-agent/pkg/tools" ) +// AddMedia appends a chunk carrying a generated-media delta (e.g. an inline +// image blob), the way the Gemini adapter surfaces InlineData parts. name may +// be empty to exercise the provider-omits-a-name path. +func (b *streamBuilder) AddMedia(data []byte, mimeType, name string) *streamBuilder { + b.responses = append(b.responses, chat.MessageStreamResponse{ + Choices: []chat.MessageStreamChoice{{ + Index: 0, + Delta: chat.MessageDelta{Media: []chat.MediaDelta{{ + Data: data, + MimeType: mimeType, + Name: name, + Size: int64(len(data)), + }}}, + }}, + }) + return b +} + +// AddMultiMedia appends a SINGLE chunk carrying multiple generated-media +// blobs at once, the way Gemini can pack several inline parts (across parts +// or candidates) into one chunk. +func (b *streamBuilder) AddMultiMedia(blobs ...chat.MediaDelta) *streamBuilder { + b.responses = append(b.responses, chat.MessageStreamResponse{ + Choices: []chat.MessageStreamChoice{{ + Index: 0, + Delta: chat.MessageDelta{Media: blobs}, + }}, + }) + return b +} + +// AddMediaWithStop appends a SINGLE terminal chunk carrying both a +// generated-media blob and a terminal finish_reason, the way a provider can +// pack the final image and "stop" into one chunk. +func (b *streamBuilder) AddMediaWithStop(data []byte, mimeType, name string, finishReason chat.FinishReason) *streamBuilder { + b.responses = append(b.responses, chat.MessageStreamResponse{ + Choices: []chat.MessageStreamChoice{{ + Index: 0, + FinishReason: finishReason, + Delta: chat.MessageDelta{Media: []chat.MediaDelta{{ + Data: data, + MimeType: mimeType, + Name: name, + Size: int64(len(data)), + }}}, + }}, + Usage: &chat.Usage{InputTokens: 1, OutputTokens: 1}, + }) + return b +} + // AddToolCallWithStop appends a single chunk that carries BOTH a complete tool // call AND a terminal finish_reason ("stop"), the way LiteLLM/Gemini emit a // function call atomically. The OpenAI-native streaming protocol never does @@ -163,6 +214,140 @@ func TestHandleStream_ToolCallThenSeparateStop(t *testing.T) { assert.False(t, res.Stopped) } +// TestHandleStream_MediaAccumulatesAlongsideText verifies that a +// generated-media delta streamed alongside text is accumulated into +// streamResult.Media without disturbing the existing text/finish-reason +// handling. +func TestHandleStream_MediaAccumulatesAlongsideText(t *testing.T) { + t.Parallel() + + imgBytes := []byte{0x89, 0x50, 0x4e, 0x47} + stream := newStreamBuilder(). + AddContent("here is your image"). + AddMedia(imgBytes, "image/png", "cat.png"). + AddStopWithUsage(1, 1). + Build() + + 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, "here is your image", res.Content, "text must survive alongside media") + require.Len(t, res.Media, 1) + assert.Equal(t, imgBytes, res.Media[0].Data) + assert.Equal(t, "image/png", res.Media[0].MimeType) + assert.Equal(t, "cat.png", res.Media[0].Name) + assert.Equal(t, chat.FinishReasonStop, res.FinishReason) + assert.True(t, res.Stopped) +} + +// TestHandleStream_MediaOnlyTurnNotTreatedAsEmpty is a regression test: a +// turn that streams ONLY a generated image (no text, no tool calls) and +// ends with a bare EOF must not be misclassified as the "no output" stall +// case — it is a normal completion and must report Stopped=true (turn +// ends) without going through the no-output warning path. +func TestHandleStream_MediaOnlyTurnNotTreatedAsEmpty(t *testing.T) { + t.Parallel() + + imgBytes := []byte{0x89, 0x50, 0x4e, 0x47} + stream := newStreamBuilder(). + AddMedia(imgBytes, "image/png", ""). + Build() // no terminal chunk: bare EOF, no finish reason + + 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.Empty(t, res.Content) + require.Len(t, res.Media, 1, "the generated image must be accumulated") + assert.True(t, res.Stopped, "a media-only turn is a normal completion, not a stall") +} + +// TestHandleStream_MultipleMediaBlobsInOneChunk verifies that every inline +// blob a provider packs into a SINGLE chunk is retained, not just the last +// one — a provider (Gemini in particular) can return more than one +// generated image across parts/candidates in the same streaming chunk. +func TestHandleStream_MultipleMediaBlobsInOneChunk(t *testing.T) { + t.Parallel() + + blob1 := chat.MediaDelta{Data: []byte{0x01}, MimeType: "image/png", Name: "one.png", Size: 1} + blob2 := chat.MediaDelta{Data: []byte{0x02}, MimeType: "image/jpeg", Name: "two.jpg", Size: 1} + blob3 := chat.MediaDelta{Data: []byte{0x03}, MimeType: "image/webp", Name: "three.webp", Size: 1} + + stream := newStreamBuilder(). + AddMultiMedia(blob1, blob2, blob3). + AddStopWithUsage(1, 1). + Build() + + 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) + + require.Len(t, res.Media, 3, "every blob in the chunk must be retained, not just the last one") + assert.Equal(t, blob1, res.Media[0]) + assert.Equal(t, blob2, res.Media[1]) + assert.Equal(t, blob3, res.Media[2]) +} + +// TestHandleStream_MediaInTerminalChunkIsAccumulated verifies that a +// generated-media blob packed into the SAME chunk as a terminal finish +// reason ("stop", "length", or "refusal") is accumulated before the early +// return, matching the same-chunk tool-call fix above. Accumulating after +// the terminal-finish-reason check would return before this chunk's media +// was ever added, silently dropping it. +func TestHandleStream_MediaInTerminalChunkIsAccumulated(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + finishReason chat.FinishReason + }{ + {"stop", chat.FinishReasonStop}, + {"length", chat.FinishReasonLength}, + {"refusal", chat.FinishReasonRefusal}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + imgBytes := []byte{0x89, 0x50, 0x4e, 0x47} + stream := newStreamBuilder(). + AddMediaWithStop(imgBytes, "image/png", "cat.png", tc.finishReason). + Build() + + 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) + + require.Len(t, res.Media, 1, "media sharing a chunk with the terminal finish reason must not be dropped") + assert.Equal(t, imgBytes, res.Media[0].Data) + assert.Equal(t, tc.finishReason, res.FinishReason) + }) + } +} + // TestHandleStream_WhitespaceOnlyContentStops is a regression test for an // infinite-loop risk surfaced while reviewing #3145. A turn that streams only // whitespace content and ends with a bare EOF (no finish reason) must report