fix(examples): validate generated image responses - #1478
karolpiotrowicz merged 10 commits into
Conversation
karolpiotrowicz
left a comment
There was a problem hiding this comment.
This does what #1477 asks, and I checked the panic is real rather than theoretical. On main both call sites still evaluate response.GeneratedImages[0].Image.ImageBytes unguarded, and both shapes the issue names crash: a safety-filtered result nil-derefs, an empty result indexes out of range. Routed through the new helper, the same two responses come back as errors carrying the filter reason. Setting IncludeRAIReason: true is the right call and not scope creep — without it the server never populates RAIFilteredReason, so the reason branch in the helper would be near-dead code.
Reproduction — the two shapes, before and after
// examples/internal/reprocheck/repro_test.go, in a scratch copy of each tree.
// Run: go test -v ./examples/internal/reprocheck/...
package reprocheck
import (
"testing"
"google.golang.org/genai"
)
func TestSafetyFiltered(t *testing.T) {
response := &genai.GenerateImagesResponse{
GeneratedImages: []*genai.GeneratedImage{{RAIFilteredReason: "filtered"}},
}
// Verbatim from both call sites before this PR.
_ = genai.NewPartFromBytes(response.GeneratedImages[0].Image.ImageBytes, "image/png")
}
func TestNoImages(t *testing.T) {
response := &genai.GenerateImagesResponse{}
_ = genai.NewPartFromBytes(response.GeneratedImages[0].Image.ImageBytes, "image/png")
}Before the fix:
--- FAIL: TestSafetyFiltered
panic: runtime error: invalid memory address or nil pointer dereference
--- FAIL: TestNoImages
panic: runtime error: index out of range [0] with length 0
With the PR applied, the same two responses through imagegen.ImageBytes:
image generation returned no image: filtered
image generation returned no images
The one thing worth discussing: the error shape in the web example
examples/web/agents/image_generator.go#L60-L63 returns generateImageResult{}, err, while the three sibling failure paths in the same function return generateImageResult{Status: "fail"}, nil. So the same outcome — no image was produced — reaches the model two different ways depending on whether the block arrived as a transport error or as a filtered HTTP 200. One consequence worth knowing: when a handler returns a non-nil error the framework discards the result value, so the generateImageResult{} literal on that line can never be observed.
I would not "fix" this by making the new line match its neighbours. Changing it to return generateImageResult{Status: "fail"}, nil throws away err, and with it the RAIFilteredReason this PR exists to surface — that would leave the panic fixed and the rest of #1477 unmet. If anything the three older branches are the odd ones out: Status: "fail" appears in exactly three places in the whole examples/ tree, all inside this one function, and every other example tool handler returns the error, including the sibling image example at examples/vertexai/imagegenerator/main.go#L101-L125, which is internally consistent. Converting all four paths in the web example to return errors is the coherent direction, and it is more than this PR signed up for. Leaving it as-is is defensible. I would not block on it either way, but a one-line comment saying the divergence is deliberate would save the next reader the same detour.
One related gap, and I could not settle it here because it needs a live API: if Vertex reports a safety block as a transport error rather than a filtered 200, examples/web/agents/image_generator.go#L54-L58 swallows it and the reason you just asked the server for reaches nobody. That path predates this PR. Worth knowing which shape Vertex actually produces before assuming the flag pays off at this call site.
Worth a separate issue, not this PR
saveImage in the same file takes a model-controlled filename straight into a write: examples/vertexai/imagegenerator/main.go#L167-L168 does filepath.Join(outputDir, localFilename) then os.WriteFile(..., 0o644). filepath.Join cleans the path but does not confine it, so a filename of ../../something.png writes outside output/. The forced .png suffix bounds the extension, not the directory. This is identical on main and is not yours to fix here — flagging it because it is the larger untrusted-input hole in a file you are already hardening, and an example is exactly where a reader copies this pattern.
Smaller notes, on the lines themselves
The remaining four are test-completeness and documentation, all left as line comments. The short version: the helper's error contract is undocumented, and three of its behaviors — the reason-selection rule, "first" in "first usable image", and what comes back alongside an error — can each be changed without the table test noticing. Breaking the panic fix itself does fail the suite, so these are specific gaps rather than a weak test.
|
Thanks for the detailed review and for verifying both panic cases. I addressed all of the follow-up comments:
Validation completed successfully:
|
karolpiotrowicz
left a comment
There was a problem hiding this comment.
Correcting something I got wrong in my last review
I flagged a path traversal in saveImage and said a model-supplied filename of ../../something.png would write outside output/. That is wrong, and I would rather say so than let it sit.
Half of it holds: filepath.Join cleans a path without confining it. What I missed is that a name needing confinement never reaches that line. saveImage loads the artifact first and returns on error at main.go#L143-L147, and validateFileName rejects any name containing a path separator. It runs from LoadRequest.Validate at service.go#L157, which inMemoryService.Load calls before anything else, and which the GCS service calls at six sites.
Driven against the same artifact.InMemoryService() the example configures:
Load("../../something.png") rejected: invalid name: filename cannot contain path separators
Load("../evil.png") rejected: invalid name: filename cannot contain path separators
Load("/etc/cron.d/evil.png") rejected: invalid name: filename cannot contain path separators
Load("..\\..\\windows.png") rejected: invalid name: filename cannot contain path separators
Load("output/../../escape.png") rejected: invalid name: filename cannot contain path separators
filepath.Join("output", "../../something.png") = "../something.png" <- escapes, but is unreachable
name ".." -> "output/...png" confined name "~" -> "output/~.png" confined
name "." -> "output/..png" confined name "-rf" -> "output/-rf.png" confined
So it is not reachable as the example ships. One boundary worth naming, since it is the case the original note was really aimed at: this drives the artifact layer's own gate rather than calling saveImage, and it covers the one service the example configures. Someone who copies this example and supplies an artifact service that skips Validate gets no protection from saveImage itself, which still has none of its own.
The new commit does what I asked
The doc comment now states all three error outcomes, the nil-bytes-on-error rule and the aliasing note. The "first usable" case now holds two usable entries with distinct payloads and asserts the earlier one. The error branch now asserts nil bytes. And the comment at image_generator.go#L60-L64 records that the divergent return is deliberate, which is the right call — matching the three sibling paths would mean discarding err, and with it the filtering reason this change exists to surface.
I checked each of those additions can actually fail rather than just being present. Mutating the helper to return the last usable image, to let a later filtering reason win, or to hand back bytes alongside an error each turns the suite red, and each is caught by a different one of the three new assertions. Beyond that, the helper does not panic on any response shape I could construct: 1887 combinations of twelve distinguishable entry kinds at slice lengths zero to three, including nil entries, a nil Image, an empty byte slice, a GCS URI in place of bytes, and PNG bytes labelled image/jpeg.
Where the validation stops
Absent data is fully covered — nil response, empty slice, nil entry, nil Image, empty bytes, and a safety-filtered entry all come back as errors. Three things pass through, none of them reachable with the configs as committed, all of them one edited line away for a reader who copies this. Two are here, and the third — an entry carrying a GCS URI instead of inline bytes — is on its line:
- A MIME type that disagrees with the bytes.
imagegen.go#L44-L45reads onlyImageBytesand dropsImage.MIMEType, and both call sites then assert"image/png"—image_generator.go#L66andmain.go#L122. SetOutputMIMEType: "image/jpeg"and you store JPEG bytes labelled and named as PNG. The hardcode predates this PR, so this is not a regression — but the helper was the natural place to stop discarding the field, and the response is still in scope at both call sites if you would rather fix it there. - A truncated payload. One byte satisfies the usability test. That matches the contract the doc comment states, so it is only worth knowing that "usable" means non-empty and nothing more.
What no test in the repository holds in place
The helper is well pinned. The wiring is not. Deleting IncludeRAIReason: true from either config, or putting response.GeneratedImages[0].Image.ImageBytes back at either call site, leaves go test ./examples/... green.
I am not asking you to fix that. There is no seam to test through — client.Models.GenerateImages is called on a client constructed in-line, so pinning it means either a live Vertex call, which CONTRIBUTING.md rules out, or injecting a client into both examples, which would make them teach dependency injection instead of image generation. examples/internal/imagegen/imagegen_test.go is the only test file anywhere under examples/, so this change took example coverage from zero to one. It is worth stating plainly rather than letting a green board imply more than it covers.
One question I could not answer without a live call, and you may already know it: does IncludeRAIReason cause the server to return a GeneratedImage entry carrying the reason for a fully blocked request, or does it come back with an empty GeneratedImages? In the second case imagegen.go#L35-L36 emits the generic "returned no images" and the reason is lost. The field's own documentation reads the way you assumed, so I expect you are right.
A pre-existing bug in a file you are already touching
Not yours to fix here, and I am not asking you to. main.go#L149-L151 returns saveImageResult{}, err on the "artifact has no inline data" path, where err is the error already checked and found nil at line 144. So an artifact with no inline data is reported to the model as an ordinary success with an empty status, and nothing is written to disk. That is the same "empty data reported as success" shape as #1477, twenty lines below the fix. Worth its own issue, and I am happy to file it if you would rather not.
The remaining notes are small and on the lines they concern.
- Return Image.MIMEType from imagegen.ImageBytes and use it at both call sites, falling back to image/png when the response omits it. - Trim RAI filtering reasons so a whitespace-only reason no longer latches. - Pin the first-filtering-reason guard in the table test by putting the entry with a non-nil empty Image first. - Document that the helper assumes inline image bytes (a GCS-URI-only entry is treated as having no usable image data).
|
Thanks for the thorough review and for the follow-up. Line comments (4 threads): addressed inline —
MIME type: the helper now returns The Wiring / truncated-payload notes: I left these as documented limitations rather than changes — the "usable" contract is "non-empty bytes" (as the doc comment states), and there's no test seam through the inline-built client to pin Validation (golangci-lint v2.3.1, CI-pinned):
|
There was a problem hiding this comment.
Everything I asked for last time is in, and the change is correct on every response shape I could construct against it. What holds it up is the test file rather than the code, and one of the two gaps is my fault.
Splitting the saveImage nil-error bug out into #1509 was the right call and keeps this PR focused. I have not looked at that one yet.
The row I asked you to swap traded one gap for another
I asked you to swap the two entries in imagegen_test.go#L48-L60 so the first-reason-wins guard would actually be exercised, and said I had tried it. That part worked. Narrowing the guard to filteredReason == "" && reason != "" && generatedImage.Image == nil now turns the table red, where before the swap it passed.
What I missed is what the swap cost. Before it, the first entry was {RAIFilteredReason: "blocked"} with no Image at all, which is the shape #1477 describes. After it, every entry whose reason has to reach the error carries a non-nil Image, and the only nil-Image reason left in the table is one that must be ignored. So this rewrite of imagegen.go#L53-L55 passes all seven rows:
if generatedImage.Image != nil {
if reason := strings.TrimSpace(generatedImage.RAIFilteredReason); filteredReason == "" && reason != "" {
filteredReason = reason
}
}and with it applied the response the issue actually names comes back with the wrong message:
response := &genai.GenerateImagesResponse{
GeneratedImages: []*genai.GeneratedImage{{RAIFilteredReason: "blocked"}},
}
as written image generation returned no image: blocked
with the rewrite image generation returned no usable image data
The code you shipped gives the first answer. Nothing in the table would notice if it stopped.
The MIME type you now return is not held in place
imagegen_test.go#L79-L91 is the only row that reaches a successful return, and its MIMEType is "image/png" — the same literal both call sites fall back to at main.go#L121-L123 and image_generator.go#L65-L67. Replace the return at imagegen.go#L51 with a hardcoded "image/png" and go test ./examples/... stays green, so the whole point of that part of the commit could be undone without a signal.
Credit where it is due: returning an empty string instead is caught, and so is returning a different constant. It is specifically the value that matches the fallback that slips through, which is also the value the pre-change code hardcoded.
Both close with three lines in one file
I ran each of these rather than suggesting them cold.
Add a row that gives a nil-Image entry a real reason and requires it to win:
{
name: "whitespace reason yields to a later real one",
response: &genai.GenerateImagesResponse{
GeneratedImages: []*genai.GeneratedImage{
{RAIFilteredReason: " "},
{RAIFilteredReason: "blocked"},
},
},
wantErr: "image generation returned no image: blocked",
},That one row rejects the rewrite above, and it also covers the case the current whitespace row at imagegen_test.go#L61-L69 cannot, since that row has a single entry and so never shows a whitespace reason stepping aside for a real one.
Then change the success row's MIMEType and wantMIME from "image/png" to "image/jpeg", and add one more success row whose response omits the MIME type entirely:
{
name: "mime type omitted by the response",
response: &genai.GenerateImagesResponse{
GeneratedImages: []*genai.GeneratedImage{
{Image: &genai.Image{ImageBytes: []byte("bytes")}},
},
},
want: []byte("bytes"),
wantMIME: "",
},The second row is worth the extra four lines: with only the "image/jpeg" change, a hardcoded "image/jpeg" slips through in exactly the way "image/png" does today. It also pins the empty-MIME case that both call-site fallbacks exist for, which nothing currently exercises.
Two things worth knowing rather than changing
The MIME string now travels further than it used to. Returning the response's MIMEType is the right call and I asked for it. Three places it reaches that the old constant did not, none of which needs a change here:
- Onto a GCS object's HTTP
Content-Type, atgcsartifact/service.go#L213-L215. Neither example configures that service, so it is out of reach as these ship, and nothing inartifact/validates the string for any caller. - Into the next model request, at
load_artifacts_tool.go#L232-L238, which forwards the stored part with its MIME. Both examples register that tool. - Not into the local filename.
main.go#L157-L161still forces.png, so a JPEG would be stored asimage/jpegand written tooutput/<name>.png.
None of that is a regression. Those lines predate the PR, and the old hardcoded "image/png" was not better — it mislabelled in one fixed direction instead. The change made the artifact's own label truthful, which is one untruth removed rather than one added.
Nothing in the repository holds generateImage in place at either call site. Deleting IncludeRAIReason: true from either config, or putting response.GeneratedImages[0].Image.ImageBytes back at either site, leaves go test ./examples/... green — I tried each of those four edits separately. I am not asking you to fix it, though the reason is narrower than the one I gave last round. generateImage builds its genai.Client in-line and then calls the network, so there is nothing to inject without restructuring the example, and CONTRIBUTING.md rules out a live call. saveImage is different, since it only reaches the artifact service, which is exactly the seam #1509 uses. Worth saying plainly so a green board is not read as covering more than it does.
The remaining notes are small and sit on the lines they concern.
Add tests for filtered responses, GCS URI-only responses, and MIME type propagation. Clarify how filtering reasons are normalized.
|
Thanks for the review. I added the missing test cases for filtered responses with a nil I ran the package tests, the full race test suite, build, lint, tidy, and formatting checks. All passed. |
karolpiotrowicz
left a comment
There was a problem hiding this comment.
Both of the things I asked for last round are in, and I have nothing blocking left. From my side this is ready to go.
The nil-Image reason now reaches the error. The row at imagegen_test.go#L70-L79 is the one that was missing: both entries carry no Image at all, and the expected message is image generation returned no image: blocked. Moving the reason capture inside the generatedImage.Image != nil branch at imagegen.go#L51-L56 now turns the table red, where before it passed every row. That is the shape #1477 actually describes, so it is the right thing to have pinned. Padding the winning reason to " blocked " rather than plain "blocked" was a better answer than the one I suggested, since it also catches storing the untrimmed value — which was a separate note of mine, so that one is closed too.
The MIME type is held from both directions now. imagegen_test.go#L104 asserting image/jpeg rules out a hardcoded "image/png", and the new row at imagegen_test.go#L111-L120 asserting wantMIME: "" rules out a hardcoded "image/jpeg". Either row on its own would leave a gap. Together they close it, and "image/png" no longer appears as a MIME anywhere in the file — its only occurrence left is inside the GCS URI. The GCS-URI row and the trimming clause in the doc comment close the remaining two notes from last round.
I also ran the change against four response shapes nobody had tried, and it is correct on all four:
- an entry whose
ImageBytesis present but zero-length - a
GeneratedImagesslice that is empty rather than nil - an entry carrying both usable bytes and a filtering reason
- a reason made of a tab and a newline
Two of those are worth a row and two are not. The two that are, I have left on the lines they concern, and neither is a condition of merging.
One thing that is still true and still not yours to fix: neither generateImage call site is exercised by anything, so the mimeType == "" fallback at main.go#L121-L123 and image_generator.go#L65-L67 could be deleted without a signal. I raised this last round and I am not asking again — generateImage builds its client in-line and then calls the network, so there is nothing to inject without restructuring the example. It is worth recording as follow-up for the repository rather than for this PR, and a separate PR is the right home for it along with the two smaller rows below — happy to review that as its own change. Worth noting the direction of travel too: examples/ had no test file at all before this change and now has one.
One administrative note: my previous review is still sitting on this PR as "changes requested", and the findings behind it are the two closed above. It no longer describes the code and needs clearing separately — a comment does not lift it.
|
Thanks for taking another look and for confirming the changes. I’ll leave the two non-blocking test cases for a follow-up. It looks like the previous changes-requested review is still active; could you clear it when convenient? |
Verify that non-nil empty image bytes do not count as usable data and cover tab and newline handling in filtering reasons.
|
I ended up adding both cases here since they fit cleanly into the existing table. |
Superseded: both findings this review gated on are closed at 4a1bc0d, each verified by re-applying the mutation that previously slipped past the tests.
karolpiotrowicz
left a comment
There was a problem hiding this comment.
Nothing here blocks. Both points I raised last round are closed, and I checked each one by re-applying the change that slipped past the tests before.
The whitespace row now uses " \t\n " (imagegen_test.go:65) rather than spaces alone, so narrowing strings.TrimSpace to trim only spaces — or only spaces and tabs, or only spaces and newlines — now turns the suite red. With the old literal all three of those slipped through. The new row at imagegen_test.go:89-100 does the same job for the usable-image check: weakening len(...) > 0 to != nil at imagegen.go:51 is caught only because that row exists. Delete either one and the corresponding regression goes back to landing silently.
The commit takes nothing away, which was the other thing worth confirming. It removes exactly one line, replacing the spaces-only literal in place, and the all-spaces case is still asserted by the row at imagegen_test.go:70-79, whose expected error is only reachable if the first entry's " " is ignored. The assertion block itself is unchanged.
A few optional things for a later PR, none of them reasons to hold this one up:
- An entry carrying usable bytes and a filtering reason is the one response shape the table never builds. The two are always on separate entries today. A row asserting that the bytes still win would pin the rule that a reason never suppresses a usable image, which is currently true by construction rather than by test. I tried this row and it passes as-is and fails if the predicate is narrowed.
- A non-nil but zero-length
GeneratedImagesslice. The "empty response" case leaves the field nil, so thelen(...) == 0guard and a== nilguard are indistinguishable. Only the error string differs between them, so this one is close to cosmetic. - Two clauses in the doc comment. It doesn't mention that inline bytes present but zero-length count as unusable, nor that a nil response is accepted rather than a programming error. Both behaviours are real and tested — the prose just doesn't cover them.
One last thing, to know rather than to fix: no test reaches generateImage in either example. I reverted a call site to the old unchecked expression, dropped the now-unused import, and the build and the whole examples/ suite stayed green. That isn't this PR's problem to solve — generateImage builds its client inline and calls out to the API, so there's no seam to test against without reshaping the example.
|
Thanks for the thorough follow-up and approval. I’ll keep these optional cases in mind for a separate follow-up. |
* fix(examples): validate generated image responses * test(examples): strengthen image response validation coverage * fix(examples): address review feedback on image response validation - Return Image.MIMEType from imagegen.ImageBytes and use it at both call sites, falling back to image/png when the response omits it. - Trim RAI filtering reasons so a whitespace-only reason no longer latches. - Pin the first-filtering-reason guard in the table test by putting the entry with a non-nil empty Image first. - Document that the helper assumes inline image bytes (a GCS-URI-only entry is treated as having no usable image data). * test(examples): cover image response edge cases Add tests for filtered responses, GCS URI-only responses, and MIME type propagation. Clarify how filtering reasons are normalized. * test(examples): cover additional image response cases Verify that non-nil empty image bytes do not count as usable data and cover tab and newline handling in filtering reasons. --------- Co-authored-by: Karol Piotrowicz <karol.piotrowicz@gmail.com>
* fix(examples): validate generated image responses * test(examples): strengthen image response validation coverage * fix(examples): address review feedback on image response validation - Return Image.MIMEType from imagegen.ImageBytes and use it at both call sites, falling back to image/png when the response omits it. - Trim RAI filtering reasons so a whitespace-only reason no longer latches. - Pin the first-filtering-reason guard in the table test by putting the entry with a non-nil empty Image first. - Document that the helper assumes inline image bytes (a GCS-URI-only entry is treated as having no usable image data). * test(examples): cover image response edge cases Add tests for filtered responses, GCS URI-only responses, and MIME type propagation. Clarify how filtering reasons are normalized. * test(examples): cover additional image response cases Verify that non-nil empty image bytes do not count as usable data and cover tab and newline handling in filtering reasons. --------- Co-authored-by: Karol Piotrowicz <karol.piotrowicz@gmail.com>
Link to Issue or Description of Change
Problem:
The image generation examples accessed the first generated image without
validating that the response contained usable image data. Empty, nil, or
safety-filtered responses could therefore cause an index-out-of-range or
nil-pointer panic.
Solution:
Testing Plan
Unit Tests:
The tests cover nil responses, empty responses, nil generated-image entries,
safety-filtered results, empty image data, and later usable images.
Commands run successfully:
go build -mod=readonly workgo test -race -mod=readonly -count=1 -shuffle=on workgo vet ./examples/...golangci-lint runin both modulesgo mod tidy -diffin both modulesManual End-to-End (E2E) Tests:
Not run. A live test requires Vertex AI credentials and a request that produces
a safety-filtered image response. The response handling is covered with
constructed SDK responses in unit tests.
Checklist
CONTRIBUTING.mddocument.