Skip to content

Commit 6d26977

Browse files
committed
fix(image): bound image-generation provider response reads
Replace unbounded success-path `await response.json()` reads with the shared `readProviderJsonResponse` helper (cancels the stream on overflow) in three image-generation providers: - src/image-generation/openai-compatible-image-provider.ts (shared factory backing DeepInfra/LiteLLM/xAI) - extensions/openai/image-generation-provider.ts (OpenAI direct JSON) - extensions/microsoft-foundry/image-generation-provider.ts (MAI image JSON) A hostile or malfunctioning provider could otherwise stream an arbitrarily large untrusted body (image responses carry base64 payloads) and force the runtime to buffer the whole thing before parsing. The error path already used the bounded reader; this brings the success path in line. Pass an image-sized 64 MiB cap (`maxBytes`) instead of the generic 16 MiB provider-JSON default. Inline `b64_json` image responses with multiple results / high resolution can legitimately exceed 16 MiB, so the generic default would have failed valid responses closed and regressed existing users. 64 MiB matches the OpenAI Codex image read budget (`MAX_CODEX_IMAGE_SSE_BYTES` / `MAX_CODEX_IMAGE_BASE64_CHARS`) in the same module. Update the affected provider tests: the closed provider-http mocks now pass through the real `readProviderJsonResponse` (importActual) and serve real streamed `Response` bodies, plus focused coverage for a valid large inline payload that parses under the cap and an oversized body that is rejected at the cap.
1 parent ec737ee commit 6d26977

6 files changed

Lines changed: 316 additions & 39 deletions

File tree

extensions/microsoft-foundry/image-generation-provider.test.ts

Lines changed: 78 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -49,15 +49,22 @@ vi.mock("openclaw/plugin-sdk/provider-auth-runtime", () => ({
4949
resolveApiKeyForProvider: resolveApiKeyForProviderMock,
5050
}));
5151

52-
vi.mock("openclaw/plugin-sdk/provider-http", () => ({
53-
assertOkOrThrowHttpError: assertOkOrThrowHttpErrorMock,
54-
createProviderOperationDeadline: createProviderOperationDeadlineMock,
55-
postJsonRequest: postJsonRequestMock,
56-
postMultipartRequest: postMultipartRequestMock,
57-
resolveProviderHttpRequestConfig: resolveProviderHttpRequestConfigMock,
58-
resolveProviderOperationTimeoutMs: resolveProviderOperationTimeoutMsMock,
59-
sanitizeConfiguredModelProviderRequest: sanitizeConfiguredModelProviderRequestMock,
60-
}));
52+
// Transport helpers are mocked per-test, but readProviderJsonResponse is kept
53+
// REAL (via importActual) so the byte-bounded reader actually caps oversized
54+
// inline MAI image JSON instead of a stub.
55+
vi.mock("openclaw/plugin-sdk/provider-http", async (importActual) => {
56+
const actual = await importActual<typeof import("openclaw/plugin-sdk/provider-http")>();
57+
return {
58+
readProviderJsonResponse: actual.readProviderJsonResponse,
59+
assertOkOrThrowHttpError: assertOkOrThrowHttpErrorMock,
60+
createProviderOperationDeadline: createProviderOperationDeadlineMock,
61+
postJsonRequest: postJsonRequestMock,
62+
postMultipartRequest: postMultipartRequestMock,
63+
resolveProviderHttpRequestConfig: resolveProviderHttpRequestConfigMock,
64+
resolveProviderOperationTimeoutMs: resolveProviderOperationTimeoutMsMock,
65+
sanitizeConfiguredModelProviderRequest: sanitizeConfiguredModelProviderRequestMock,
66+
};
67+
});
6168

6269
vi.mock("./runtime.js", () => ({
6370
prepareFoundryRuntimeAuth: prepareFoundryRuntimeAuthMock,
@@ -109,6 +116,36 @@ function releasedJson(payload: unknown) {
109116
};
110117
}
111118

119+
// The image-sized JSON cap (64 MiB) matches the source constant
120+
// MAX_IMAGE_GENERATION_JSON_BYTES.
121+
const IMAGE_JSON_CAP_BYTES = 64 * 1024 * 1024;
122+
123+
// Builds a JSON body larger than the image cap so the bounded reader cancels the
124+
// stream mid-flight instead of buffering the whole advertised payload.
125+
function releasedOversizedJson() {
126+
const ONE_MIB = 1024 * 1024;
127+
const TOTAL_CHUNKS = 65; // > 64 MiB cap.
128+
const chunk = new Uint8Array(ONE_MIB);
129+
let pulled = 0;
130+
const body = new ReadableStream<Uint8Array>({
131+
pull(controller) {
132+
if (pulled >= TOTAL_CHUNKS) {
133+
controller.close();
134+
return;
135+
}
136+
pulled += 1;
137+
controller.enqueue(chunk);
138+
},
139+
});
140+
return {
141+
response: new Response(body, {
142+
status: 200,
143+
headers: { "content-type": "application/json" },
144+
}),
145+
release: vi.fn(async () => {}),
146+
};
147+
}
148+
112149
function requirePostJsonRequest(): Record<string, unknown> {
113150
const request = postJsonRequestMock.mock.calls[0]?.[0];
114151
if (!request || typeof request !== "object") {
@@ -227,6 +264,38 @@ describe("microsoft foundry image generation provider", () => {
227264
expect(result.images[0]?.mimeType).toBe("image/png");
228265
});
229266

267+
it("parses a valid inline MAI image response within the image-sized JSON cap", async () => {
268+
// ~4 MiB of base64 image data: above bare metadata but well under the 64 MiB
269+
// image cap.
270+
const bigBase64 = Buffer.alloc(4 * 1024 * 1024, "a").toString("base64");
271+
postJsonRequestMock.mockResolvedValue(releasedJson({ data: [{ b64_json: bigBase64 }] }));
272+
const provider = buildMicrosoftFoundryImageGenerationProvider();
273+
274+
const result = await provider.generateImage({
275+
provider: PROVIDER_ID,
276+
model: "image-deployment",
277+
prompt: "large valid render",
278+
cfg: buildConfig(),
279+
});
280+
281+
expect(result.images).toHaveLength(1);
282+
expect(result.images[0]?.buffer.length).toBe(4 * 1024 * 1024);
283+
});
284+
285+
it("rejects MAI image responses that exceed the image-sized JSON cap", async () => {
286+
postJsonRequestMock.mockResolvedValue(releasedOversizedJson());
287+
const provider = buildMicrosoftFoundryImageGenerationProvider();
288+
289+
await expect(
290+
provider.generateImage({
291+
provider: PROVIDER_ID,
292+
model: "image-deployment",
293+
prompt: "oversized render",
294+
cfg: buildConfig(),
295+
}),
296+
).rejects.toThrow(`exceeds ${IMAGE_JSON_CAP_BYTES} bytes`);
297+
});
298+
230299
it("uses AZURE_OPENAI_ENDPOINT when env API-key auth has no configured base URL", async () => {
231300
vi.stubEnv("AZURE_OPENAI_ENDPOINT", "https://env.services.ai.azure.com");
232301
postJsonRequestMock.mockResolvedValue(

extensions/microsoft-foundry/image-generation-provider.ts

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import {
1818
createProviderOperationDeadline,
1919
postJsonRequest,
2020
postMultipartRequest,
21+
readProviderJsonResponse,
2122
resolveProviderHttpRequestConfig,
2223
resolveProviderOperationTimeoutMs,
2324
sanitizeConfiguredModelProviderRequest,
@@ -36,6 +37,10 @@ import {
3637
} from "./shared.js";
3738

3839
const DEFAULT_TIMEOUT_MS = 600_000;
40+
// MAI image responses return inline base64 image data, which can exceed the
41+
// generic 16 MiB provider JSON cap for multi-image / high-resolution outputs.
42+
// Match the OpenAI Codex image read budget so valid responses still parse.
43+
const MAX_IMAGE_GENERATION_JSON_BYTES = 64 * 1024 * 1024;
3944
const DEFAULT_IMAGE_SIZE = { width: 1024, height: 1024 };
4045
const MAI_MIN_IMAGE_SIDE_PX = 768;
4146
const MAI_MAX_IMAGE_PIXELS = 1_048_576;
@@ -368,7 +373,12 @@ export function buildMicrosoftFoundryImageGenerationProvider(): ImageGenerationP
368373
try {
369374
await assertOkOrThrowHttpError(response, `${label} failed`);
370375
return {
371-
images: parseMaiImageResponse(await response.json(), label),
376+
images: parseMaiImageResponse(
377+
await readProviderJsonResponse(response, label, {
378+
maxBytes: MAX_IMAGE_GENERATION_JSON_BYTES,
379+
}),
380+
label,
381+
),
372382
model,
373383
};
374384
} finally {

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

Lines changed: 98 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -60,13 +60,20 @@ vi.mock("openclaw/plugin-sdk/provider-auth-runtime", () => ({
6060
resolveApiKeyForProvider: resolveApiKeyForProviderMock,
6161
}));
6262

63-
vi.mock("openclaw/plugin-sdk/provider-http", () => ({
64-
assertOkOrThrowHttpError: assertOkOrThrowHttpErrorMock,
65-
postJsonRequest: postJsonRequestMock,
66-
postMultipartRequest: postMultipartRequestMock,
67-
resolveProviderHttpRequestConfig: resolveProviderHttpRequestConfigMock,
68-
sanitizeConfiguredModelProviderRequest: sanitizeConfiguredModelProviderRequestMock,
69-
}));
63+
// Transport helpers are mocked per-test, but readProviderJsonResponse is kept
64+
// REAL (via importActual) so the byte-bounded reader actually caps oversized
65+
// inline image JSON instead of a stub.
66+
vi.mock("openclaw/plugin-sdk/provider-http", async (importActual) => {
67+
const actual = await importActual<typeof import("openclaw/plugin-sdk/provider-http")>();
68+
return {
69+
readProviderJsonResponse: actual.readProviderJsonResponse,
70+
assertOkOrThrowHttpError: assertOkOrThrowHttpErrorMock,
71+
postJsonRequest: postJsonRequestMock,
72+
postMultipartRequest: postMultipartRequestMock,
73+
resolveProviderHttpRequestConfig: resolveProviderHttpRequestConfigMock,
74+
sanitizeConfiguredModelProviderRequest: sanitizeConfiguredModelProviderRequestMock,
75+
};
76+
});
7077

7178
vi.mock("openclaw/plugin-sdk/logging-core", () => ({
7279
createSubsystemLogger: vi.fn(() => ({
@@ -77,18 +84,57 @@ vi.mock("openclaw/plugin-sdk/logging-core", () => ({
7784
})),
7885
}));
7986

80-
function mockGeneratedPngResponse() {
81-
const response = {
82-
json: async () => ({
83-
data: [{ b64_json: Buffer.from("png-bytes").toString("base64") }],
87+
// Image success responses are now read through the byte-bounded reader, so the
88+
// mocked responses must expose a real readable body (not just a json() shim).
89+
function streamedJsonResponse(payload: unknown): Response {
90+
return new Response(
91+
new ReadableStream({
92+
start(controller) {
93+
controller.enqueue(new TextEncoder().encode(JSON.stringify(payload)));
94+
controller.close();
95+
},
8496
}),
97+
{ status: 200, headers: { "content-type": "application/json" } },
98+
);
99+
}
100+
101+
// The image-sized JSON cap (64 MiB) matches the source constant
102+
// MAX_IMAGE_GENERATION_JSON_BYTES.
103+
const IMAGE_JSON_CAP_BYTES = 64 * 1024 * 1024;
104+
105+
// Builds a JSON body larger than the image cap so the bounded reader cancels the
106+
// stream mid-flight instead of buffering the whole advertised payload.
107+
function makeOversizedJsonResponse(): Response {
108+
const ONE_MIB = 1024 * 1024;
109+
const TOTAL_CHUNKS = 65; // > 64 MiB cap.
110+
const chunk = new Uint8Array(ONE_MIB);
111+
let pulled = 0;
112+
const body = new ReadableStream<Uint8Array>({
113+
pull(controller) {
114+
if (pulled >= TOTAL_CHUNKS) {
115+
controller.close();
116+
return;
117+
}
118+
pulled += 1;
119+
controller.enqueue(chunk);
120+
},
121+
});
122+
return new Response(body, {
123+
status: 200,
124+
headers: { "content-type": "application/json" },
125+
});
126+
}
127+
128+
function mockGeneratedPngResponse() {
129+
const payload = {
130+
data: [{ b64_json: Buffer.from("png-bytes").toString("base64") }],
85131
};
86132
postJsonRequestMock.mockResolvedValue({
87-
response,
133+
response: streamedJsonResponse(payload),
88134
release: vi.fn(async () => {}),
89135
});
90136
postMultipartRequestMock.mockResolvedValue({
91-
response,
137+
response: streamedJsonResponse(payload),
92138
release: vi.fn(async () => {}),
93139
});
94140
}
@@ -579,7 +625,7 @@ describe("openai image generation provider", () => {
579625

580626
it("wraps malformed successful OpenAI image responses", async () => {
581627
postJsonRequestMock.mockResolvedValue({
582-
response: { json: async () => ({ data: { b64_json: "not-an-array" } }) },
628+
response: streamedJsonResponse({ data: { b64_json: "not-an-array" } }),
583629
release: vi.fn(async () => {}),
584630
});
585631

@@ -594,6 +640,44 @@ describe("openai image generation provider", () => {
594640
).rejects.toThrow("OpenAI image generation response malformed");
595641
});
596642

643+
it("parses a valid inline OpenAI image response within the image-sized JSON cap", async () => {
644+
// ~4 MiB of base64 image data: above bare metadata but well under the 64 MiB
645+
// image cap (would also exceed the generic 16 MiB default if combined).
646+
const bigBase64 = Buffer.alloc(4 * 1024 * 1024, "a").toString("base64");
647+
postJsonRequestMock.mockResolvedValue({
648+
response: streamedJsonResponse({ data: [{ b64_json: bigBase64 }] }),
649+
release: vi.fn(async () => {}),
650+
});
651+
652+
const provider = buildOpenAIImageGenerationProvider();
653+
const result = await provider.generateImage({
654+
provider: "openai",
655+
model: "gpt-image-2",
656+
prompt: "large valid image",
657+
cfg: {},
658+
});
659+
660+
expect(result.images).toHaveLength(1);
661+
expect(result.images[0]?.buffer.length).toBe(4 * 1024 * 1024);
662+
});
663+
664+
it("rejects OpenAI image responses that exceed the image-sized JSON cap", async () => {
665+
postJsonRequestMock.mockResolvedValue({
666+
response: makeOversizedJsonResponse(),
667+
release: vi.fn(async () => {}),
668+
});
669+
670+
const provider = buildOpenAIImageGenerationProvider();
671+
await expect(
672+
provider.generateImage({
673+
provider: "openai",
674+
model: "gpt-image-2",
675+
prompt: "oversized image",
676+
cfg: {},
677+
}),
678+
).rejects.toThrow(`exceeds ${IMAGE_JSON_CAP_BYTES} bytes`);
679+
});
680+
597681
it("forwards generation count and custom size overrides", async () => {
598682
mockGeneratedPngResponse();
599683

extensions/openai/image-generation-provider.ts

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import {
2424
assertOkOrThrowHttpError,
2525
postJsonRequest,
2626
postMultipartRequest,
27+
readProviderJsonResponse,
2728
resolveProviderHttpRequestConfig,
2829
sanitizeConfiguredModelProviderRequest,
2930
} from "openclaw/plugin-sdk/provider-http";
@@ -61,6 +62,10 @@ const OPENAI_MAX_IMAGE_RESULTS = 4;
6162
const MAX_CODEX_IMAGE_SSE_BYTES = 64 * 1024 * 1024;
6263
const MAX_CODEX_IMAGE_SSE_EVENTS = 512;
6364
const MAX_CODEX_IMAGE_BASE64_CHARS = 64 * 1024 * 1024;
65+
// Inline `b64_json` image responses (up to OPENAI_MAX_IMAGE_RESULTS results at
66+
// high resolution) can exceed the generic 16 MiB provider JSON cap, so the
67+
// success-path JSON read uses the same image budget as the Codex image paths.
68+
const MAX_IMAGE_GENERATION_JSON_BYTES = 64 * 1024 * 1024;
6469
const LOG_VALUE_MAX_CHARS = 256;
6570
const MOCK_OPENAI_PROVIDER_ID = "mock-openai";
6671
const OPENAI_OUTPUT_FORMATS = ["png", "jpeg", "webp"] as const;
@@ -1012,7 +1017,11 @@ export function buildOpenAIImageGenerationProvider(): ImageGenerationProvider {
10121017
isEdit ? "OpenAI image edit failed" : "OpenAI image generation failed",
10131018
);
10141019

1015-
const data = await response.json();
1020+
const data = await readProviderJsonResponse(
1021+
response,
1022+
isEdit ? "OpenAI image edit" : "OpenAI image generation",
1023+
{ maxBytes: MAX_IMAGE_GENERATION_JSON_BYTES },
1024+
);
10161025
const output = resolveOutputMime(req.outputFormat);
10171026
const images = parseOpenAiCompatibleImageResponse(data, {
10181027
defaultMimeType: output.mimeType,

0 commit comments

Comments
 (0)