Skip to content

Commit 527f8f0

Browse files
hugenshencursoragentsallyom
authored
fix(image-gen): bound image generation provider JSON response reads (#96495)
* fix(image-gen): bound image generation provider JSON response reads Route success JSON reads through readProviderJsonResponse (16 MiB cap) in openrouter, google, fal, minimax, openai, and vydra image generation providers to prevent OOM from oversized or hostile endpoint responses. Mirrors the response-limit campaign already applied to other provider paths. AI-assisted. Co-authored-by: Cursor <[email protected]> * fix(image-gen): size bounded JSON caps for inline image payloads Signed-off-by: sallyom <[email protected]> --------- Signed-off-by: sallyom <[email protected]> Co-authored-by: Cursor <[email protected]> Co-authored-by: sallyom <[email protected]>
1 parent c05d0d5 commit 527f8f0

22 files changed

Lines changed: 568 additions & 166 deletions
Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,2 @@
1-
abdff20b710c6b0fecb5af25603d7cfad7ade80600ca374ebe38f69d78933b50 plugin-sdk-api-baseline.json
2-
630367961e4d14463020f588564c23308159ae2de6e4301418b2b0c471797e70 plugin-sdk-api-baseline.jsonl
1+
e9fb501204b6c4c0e08c09174311f85ae8a129bf35e18edea0fa217ee4203ad8 plugin-sdk-api-baseline.json
2+
50db19cf60c5465b22d342ccdabe465ae08d0c475e6403b86360f775bfc3513a plugin-sdk-api-baseline.jsonl

extensions/deepinfra/image-generation-provider.test.ts

Lines changed: 34 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -31,15 +31,21 @@ vi.mock("openclaw/plugin-sdk/provider-auth-runtime", () => ({
3131
resolveApiKeyForProvider: resolveApiKeyForProviderMock,
3232
}));
3333

34-
vi.mock("openclaw/plugin-sdk/provider-http", () => ({
35-
assertOkOrThrowHttpError: assertOkOrThrowHttpErrorMock,
36-
createProviderOperationDeadline: createProviderOperationDeadlineMock,
37-
postJsonRequest: postJsonRequestMock,
38-
postMultipartRequest: postMultipartRequestMock,
39-
resolveProviderHttpRequestConfig: resolveProviderHttpRequestConfigMock,
40-
resolveProviderOperationTimeoutMs: resolveProviderOperationTimeoutMsMock,
41-
sanitizeConfiguredModelProviderRequest: vi.fn((request) => request),
42-
}));
34+
vi.mock("openclaw/plugin-sdk/provider-http", async () => {
35+
const actual = await vi.importActual<typeof import("openclaw/plugin-sdk/provider-http")>(
36+
"openclaw/plugin-sdk/provider-http",
37+
);
38+
return {
39+
assertOkOrThrowHttpError: assertOkOrThrowHttpErrorMock,
40+
createProviderOperationDeadline: createProviderOperationDeadlineMock,
41+
postJsonRequest: postJsonRequestMock,
42+
postMultipartRequest: postMultipartRequestMock,
43+
readProviderJsonResponse: actual.readProviderJsonResponse,
44+
resolveProviderHttpRequestConfig: resolveProviderHttpRequestConfigMock,
45+
resolveProviderOperationTimeoutMs: resolveProviderOperationTimeoutMsMock,
46+
sanitizeConfiguredModelProviderRequest: vi.fn((request) => request),
47+
};
48+
});
4349

4450
afterAll(() => {
4551
vi.doUnmock("openclaw/plugin-sdk/provider-auth-runtime");
@@ -63,6 +69,13 @@ function requireFirstMockObjectArg(mock: ReturnType<typeof vi.fn>, label: string
6369
return value;
6470
}
6571

72+
function jsonResponse(payload: unknown): Response {
73+
return new Response(JSON.stringify(payload), {
74+
status: 200,
75+
headers: { "Content-Type": "application/json" },
76+
});
77+
}
78+
6679
describe("deepinfra image generation provider", () => {
6780
afterEach(() => {
6881
assertOkOrThrowHttpErrorMock.mockClear();
@@ -86,11 +99,9 @@ describe("deepinfra image generation provider", () => {
8699
const release = vi.fn(async () => {});
87100
const jpegBytes = Buffer.from([0xff, 0xd8, 0xff, 0x00]);
88101
postJsonRequestMock.mockResolvedValue({
89-
response: {
90-
json: async () => ({
91-
data: [{ b64_json: jpegBytes.toString("base64"), revised_prompt: "red square" }],
92-
}),
93-
},
102+
response: jsonResponse({
103+
data: [{ b64_json: jpegBytes.toString("base64"), revised_prompt: "red square" }],
104+
}),
94105
release,
95106
});
96107

@@ -168,17 +179,15 @@ describe("deepinfra image generation provider", () => {
168179

169180
it("sends image edits as multipart OpenAI-compatible requests", async () => {
170181
postMultipartRequestMock.mockResolvedValue({
171-
response: {
172-
json: async () => ({
173-
data: [
174-
{
175-
b64_json: Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]).toString(
176-
"base64",
177-
),
178-
},
179-
],
180-
}),
181-
},
182+
response: jsonResponse({
183+
data: [
184+
{
185+
b64_json: Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]).toString(
186+
"base64",
187+
),
188+
},
189+
],
190+
}),
182191
release: vi.fn(async () => {}),
183192
});
184193

extensions/fal/image-generation-provider.ts

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import { isProviderApiKeyConfigured } from "openclaw/plugin-sdk/provider-auth";
1212
import {
1313
assertOkOrThrowHttpError,
1414
assertOkOrThrowProviderError,
15+
readProviderJsonResponse,
1516
} from "openclaw/plugin-sdk/provider-http";
1617
import { readResponseWithLimit } from "openclaw/plugin-sdk/response-limit-runtime";
1718
import {
@@ -645,7 +646,9 @@ export function buildFalImageGenerationProvider(): ImageGenerationProvider {
645646
try {
646647
await assertOkOrThrowHttpError(response, "fal image generation failed");
647648

648-
const payload = parseFalImageGenerationResponse(await response.json());
649+
const payload = parseFalImageGenerationResponse(
650+
await readProviderJsonResponse(response, "fal.image-generation"),
651+
);
649652
const images: GeneratedImageAsset[] = [];
650653
let imageIndex = 0;
651654
for (const entry of payload.images) {

extensions/google/image-generation-provider.test.ts

Lines changed: 88 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,13 @@ import { testing as geminiWebSearchTesting } from "./src/gemini-web-search-provi
88

99
let ssrfMock: { mockRestore: () => void } | undefined;
1010

11+
function jsonResponse(payload: unknown): Response {
12+
return new Response(JSON.stringify(payload), {
13+
status: 200,
14+
headers: { "Content-Type": "application/json" },
15+
});
16+
}
17+
1118
function mockGoogleApiKeyAuth() {
1219
vi.spyOn(providerAuthRuntime, "resolveApiKeyForProvider").mockResolvedValue({
1320
apiKey: "google-test-key",
@@ -24,9 +31,8 @@ function installGoogleFetchMock(params?: {
2431
const mimeType = params?.mimeType ?? "image/png";
2532
const data = params?.data ?? "png-data";
2633
const inlineDataKey = params?.inlineDataKey ?? "inlineData";
27-
const fetchMock = vi.fn().mockResolvedValue({
28-
ok: true,
29-
json: async () => ({
34+
const fetchMock = vi.fn().mockResolvedValue(
35+
jsonResponse({
3036
candidates: [
3137
{
3238
content: {
@@ -42,7 +48,7 @@ function installGoogleFetchMock(params?: {
4248
},
4349
],
4450
}),
45-
});
51+
);
4652
vi.stubGlobal("fetch", fetchMock);
4753
return fetchMock;
4854
}
@@ -100,9 +106,8 @@ describe("Google image-generation provider", () => {
100106
source: "env",
101107
mode: "api-key",
102108
});
103-
const fetchMock = vi.fn().mockResolvedValue({
104-
ok: true,
105-
json: async () => ({
109+
const fetchMock = vi.fn().mockResolvedValue(
110+
jsonResponse({
106111
candidates: [
107112
{
108113
content: {
@@ -119,7 +124,7 @@ describe("Google image-generation provider", () => {
119124
},
120125
],
121126
}),
122-
});
127+
);
123128
vi.stubGlobal("fetch", fetchMock);
124129

125130
const provider = buildGoogleImageGenerationProvider();
@@ -208,10 +213,7 @@ describe("Google image-generation provider", () => {
208213
mockGoogleApiKeyAuth();
209214
vi.stubGlobal(
210215
"fetch",
211-
vi.fn().mockResolvedValue({
212-
ok: true,
213-
json: async () => ({ candidates: { content: { parts: [] } } }),
214-
}),
216+
vi.fn().mockResolvedValue(jsonResponse({ candidates: { content: { parts: [] } } })),
215217
);
216218

217219
const provider = buildGoogleImageGenerationProvider();
@@ -229,9 +231,8 @@ describe("Google image-generation provider", () => {
229231
mockGoogleApiKeyAuth();
230232
vi.stubGlobal(
231233
"fetch",
232-
vi.fn().mockResolvedValue({
233-
ok: true,
234-
json: async () => ({
234+
vi.fn().mockResolvedValue(
235+
jsonResponse({
235236
candidates: [
236237
{
237238
content: {
@@ -240,7 +241,7 @@ describe("Google image-generation provider", () => {
240241
},
241242
],
242243
}),
243-
}),
244+
),
244245
);
245246

246247
const provider = buildGoogleImageGenerationProvider();
@@ -260,9 +261,8 @@ describe("Google image-generation provider", () => {
260261
source: "profile",
261262
mode: "token",
262263
});
263-
const fetchMock = vi.fn().mockResolvedValue({
264-
ok: true,
265-
json: async () => ({
264+
const fetchMock = vi.fn().mockResolvedValue(
265+
jsonResponse({
266266
candidates: [
267267
{
268268
content: {
@@ -278,7 +278,7 @@ describe("Google image-generation provider", () => {
278278
},
279279
],
280280
}),
281-
});
281+
);
282282
vi.stubGlobal("fetch", fetchMock);
283283

284284
const provider = buildGoogleImageGenerationProvider();
@@ -305,6 +305,74 @@ describe("Google image-generation provider", () => {
305305
});
306306
});
307307

308+
it("accepts valid multi-image inline JSON responses above the generic provider JSON cap", async () => {
309+
mockGoogleApiKeyAuth();
310+
const imageBytes = Buffer.alloc(6 * 1024 * 1024, 1);
311+
const imagePayload = imageBytes.toString("base64");
312+
vi.stubGlobal(
313+
"fetch",
314+
vi.fn().mockResolvedValue(
315+
jsonResponse({
316+
candidates: [
317+
{
318+
content: {
319+
parts: Array.from({ length: 3 }, () => ({
320+
inlineData: {
321+
mimeType: "image/png",
322+
data: imagePayload,
323+
},
324+
})),
325+
},
326+
},
327+
],
328+
}),
329+
),
330+
);
331+
332+
const provider = buildGoogleImageGenerationProvider();
333+
const result = await provider.generateImage({
334+
provider: "google",
335+
model: "gemini-3.1-flash-image-preview",
336+
prompt: "draw a cat",
337+
cfg: {},
338+
});
339+
340+
expect(result.images).toHaveLength(3);
341+
expect(result.images.map((image) => image.buffer.byteLength)).toEqual([
342+
imageBytes.byteLength,
343+
imageBytes.byteLength,
344+
imageBytes.byteLength,
345+
]);
346+
});
347+
348+
it("still rejects oversized Google image JSON responses", async () => {
349+
mockGoogleApiKeyAuth();
350+
vi.stubGlobal(
351+
"fetch",
352+
vi.fn().mockResolvedValue(
353+
jsonResponse({
354+
candidates: [
355+
{
356+
content: {
357+
parts: [{ text: "x".repeat(35 * 1024 * 1024) }],
358+
},
359+
},
360+
],
361+
}),
362+
),
363+
);
364+
365+
const provider = buildGoogleImageGenerationProvider();
366+
await expect(
367+
provider.generateImage({
368+
provider: "google",
369+
model: "gemini-3.1-flash-image-preview",
370+
prompt: "draw a cat",
371+
cfg: {},
372+
}),
373+
).rejects.toThrow("google.image-generation: JSON response exceeds");
374+
});
375+
308376
it("sends reference images and explicit resolution for edit flows", async () => {
309377
mockGoogleApiKeyAuth();
310378
const fetchMock = installGoogleFetchMock();

extensions/google/image-generation-provider.ts

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,18 @@
11
// Google provider module implements model/runtime integration.
22
import {
33
generatedImageAssetFromBase64,
4+
resolveInlineImageJsonResponseMaxBytes,
45
type GeneratedImageAsset,
56
type ImageGenerationProvider,
67
} from "openclaw/plugin-sdk/image-generation";
8+
import { MAX_IMAGE_BYTES } from "openclaw/plugin-sdk/media-runtime";
79
import { parseStrictPositiveInteger } from "openclaw/plugin-sdk/number-runtime";
810
import { isProviderApiKeyConfigured } from "openclaw/plugin-sdk/provider-auth";
911
import { resolveApiKeyForProvider } from "openclaw/plugin-sdk/provider-auth-runtime";
1012
import {
1113
assertOkOrThrowHttpError,
1214
postJsonRequest,
15+
readProviderJsonResponse,
1316
sanitizeConfiguredModelProviderRequest,
1417
} from "openclaw/plugin-sdk/provider-http";
1518
import {
@@ -22,6 +25,8 @@ import { normalizeGoogleModelId, resolveGoogleGenerativeAiHttpRequestConfig } fr
2225
const DEFAULT_GOOGLE_IMAGE_MODEL = "gemini-3.1-flash-image-preview";
2326
const DEFAULT_IMAGE_TIMEOUT_MS = 180_000;
2427
const DEFAULT_OUTPUT_MIME = "image/png";
28+
const GOOGLE_MAX_IMAGE_RESULTS = 4;
29+
const MB = 1024 * 1024;
2530
const GOOGLE_SUPPORTED_SIZES = [
2631
"1024x1024",
2732
"1024x1536",
@@ -49,6 +54,16 @@ function normalizeGoogleImageModel(model: string | undefined): string {
4954
return normalizeGoogleModelId(trimmed || DEFAULT_GOOGLE_IMAGE_MODEL);
5055
}
5156

57+
function resolveGeneratedImageMaxBytes(req: {
58+
cfg: { agents?: { defaults?: { mediaMaxMb?: number } } };
59+
}): number {
60+
const configured = req.cfg.agents?.defaults?.mediaMaxMb;
61+
if (typeof configured === "number" && Number.isFinite(configured) && configured > 0) {
62+
return Math.floor(configured * MB);
63+
}
64+
return MAX_IMAGE_BYTES;
65+
}
66+
5267
function mapSizeToImageConfig(
5368
size: string | undefined,
5469
): { aspectRatio?: string; imageSize?: "2K" | "4K" } | undefined {
@@ -149,14 +164,14 @@ export function buildGoogleImageGenerationProvider(): ImageGenerationProvider {
149164
}),
150165
capabilities: {
151166
generate: {
152-
maxCount: 4,
167+
maxCount: GOOGLE_MAX_IMAGE_RESULTS,
153168
supportsSize: true,
154169
supportsAspectRatio: true,
155170
supportsResolution: true,
156171
},
157172
edit: {
158173
enabled: true,
159-
maxCount: 4,
174+
maxCount: GOOGLE_MAX_IMAGE_RESULTS,
160175
maxInputImages: 5,
161176
supportsSize: true,
162177
supportsAspectRatio: true,
@@ -231,7 +246,12 @@ export function buildGoogleImageGenerationProvider(): ImageGenerationProvider {
231246
try {
232247
await assertOkOrThrowHttpError(res, "Google image generation failed");
233248

234-
const payload = await res.json();
249+
const payload = await readProviderJsonResponse(res, "google.image-generation", {
250+
maxBytes: resolveInlineImageJsonResponseMaxBytes(
251+
GOOGLE_MAX_IMAGE_RESULTS,
252+
resolveGeneratedImageMaxBytes(req),
253+
),
254+
});
235255
let imageIndex = 0;
236256
const images: GeneratedImageAsset[] = [];
237257
for (const part of googleResponseParts(payload)) {

0 commit comments

Comments
 (0)