Skip to content

Commit 9490c8a

Browse files
sallyomsallyom
committed
fix(image-gen): size bounded JSON caps for inline image payloads
Signed-off-by: sallyom <[email protected]>
1 parent bac9fd2 commit 9490c8a

18 files changed

Lines changed: 546 additions & 162 deletions

‎docs/.generated/plugin-sdk-api-baseline.sha256‎

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/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: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
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";
@@ -23,6 +25,8 @@ import { normalizeGoogleModelId, resolveGoogleGenerativeAiHttpRequestConfig } fr
2325
const DEFAULT_GOOGLE_IMAGE_MODEL = "gemini-3.1-flash-image-preview";
2426
const DEFAULT_IMAGE_TIMEOUT_MS = 180_000;
2527
const DEFAULT_OUTPUT_MIME = "image/png";
28+
const GOOGLE_MAX_IMAGE_RESULTS = 4;
29+
const MB = 1024 * 1024;
2630
const GOOGLE_SUPPORTED_SIZES = [
2731
"1024x1024",
2832
"1024x1536",
@@ -50,6 +54,16 @@ function normalizeGoogleImageModel(model: string | undefined): string {
5054
return normalizeGoogleModelId(trimmed || DEFAULT_GOOGLE_IMAGE_MODEL);
5155
}
5256

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+
5367
function mapSizeToImageConfig(
5468
size: string | undefined,
5569
): { aspectRatio?: string; imageSize?: "2K" | "4K" } | undefined {
@@ -150,14 +164,14 @@ export function buildGoogleImageGenerationProvider(): ImageGenerationProvider {
150164
}),
151165
capabilities: {
152166
generate: {
153-
maxCount: 4,
167+
maxCount: GOOGLE_MAX_IMAGE_RESULTS,
154168
supportsSize: true,
155169
supportsAspectRatio: true,
156170
supportsResolution: true,
157171
},
158172
edit: {
159173
enabled: true,
160-
maxCount: 4,
174+
maxCount: GOOGLE_MAX_IMAGE_RESULTS,
161175
maxInputImages: 5,
162176
supportsSize: true,
163177
supportsAspectRatio: true,
@@ -232,7 +246,12 @@ export function buildGoogleImageGenerationProvider(): ImageGenerationProvider {
232246
try {
233247
await assertOkOrThrowHttpError(res, "Google image generation failed");
234248

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

0 commit comments

Comments
 (0)