diff --git a/src/llm/providers/mistral.bounded-stream.test.ts b/src/llm/providers/mistral.bounded-stream.test.ts new file mode 100644 index 000000000000..019ae45b8ea5 --- /dev/null +++ b/src/llm/providers/mistral.bounded-stream.test.ts @@ -0,0 +1,184 @@ +// Mistral provider tests cover bounded-stream-read helper (`createBoundedMistralFetcher`). +import http from "node:http"; +import type { AddressInfo } from "node:net"; +import { describe, expect, it } from "vitest"; +import { createBoundedMistralFetcher } from "./mistral.js"; + +const MAX = 16 * 1024 * 1024; +const TOTAL = 18 * 1024 * 1024; + +async function readAllChunks(body: ReadableStream | null): Promise<{ total: number }> { + if (!body) { + return { total: 0 }; + } + const reader = body.getReader(); + let total = 0; + while (true) { + const { done, value } = await reader.read(); + if (done) { + break; + } + if (value) { + total += value.byteLength; + } + } + return { total }; +} + +describe("Mistral bounded-stream-read real wire proof (loopback http.createServer)", () => { + it("caps an oversized body streamed chunked over real wire", async () => { + const fetcher = createBoundedMistralFetcher(MAX); + const CHUNK = 1024 * 1024; + const server = http.createServer((req, res) => { + res.writeHead(200, { "content-type": "application/octet-stream" }); + let sent = 0; + const tick = setInterval(() => { + if (sent < 18) { + res.write(Buffer.alloc(CHUNK)); + sent++; + } else { + clearInterval(tick); + res.end(); + } + }, 1); + }); + await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(0, "127.0.0.1", () => { + resolve(); + }); + }); + const port = (server.address() as AddressInfo).port; + + let captured: Error | undefined; + let totalGot = 0; + try { + const response = await fetcher(`http://127.0.0.1:${port}/`); + // Wire framing merges TCP packets, so the reported size at throw time + // is between MAX (cap) and TOTAL (cap + last merged packet). Both + // bounds prove (a) cap fired (got > MAX) and (b) we did not buffer + // beyond the server's full 18 MiB (got < TOTAL). + try { + const result = await readAllChunks(response.body); + totalGot = result.total; + } catch (err) { + captured = err as Error; + } + expect(captured).toBeInstanceOf(Error); + const match = (captured as Error).message.match( + /mistral: stream body exceeds \d+ bytes \(got (\d+)\)/, + ); + expect(match).not.toBeNull(); + const got = Number(match![1]); + expect(got).toBeGreaterThan(MAX); + expect(got).toBeLessThan(TOTAL); + // Print to vitest stdout for PR-body real behavior proof capture. + console.log( + `[mistral bounded-stream proof] oversized path: cap=${MAX} reported=${got} server_total=${TOTAL}`, + ); + } finally { + await new Promise((resolve) => { + server.close(() => { + resolve(); + }); + }); + if (totalGot > 0) { + // Use the value to satisfy strict unused rules without affecting asserts. + expect(totalGot).toBeGreaterThan(0); + } + } + }); + + it("returns a Response with exact bytes for normal-size responses on real wire", async () => { + const fetcher = createBoundedMistralFetcher(MAX); + const bodyText = 'data: {"choices":[{"delta":{"content":"hello"}}]}\n\ndata: [DONE]\n\n'; + const server = http.createServer((req, res) => { + res.writeHead(200, { "content-type": "text/event-stream" }); + res.end(bodyText); + }); + await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(0, "127.0.0.1", () => { + resolve(); + }); + }); + const port = (server.address() as AddressInfo).port; + + try { + const response = await fetcher(`http://127.0.0.1:${port}/`); + expect(response.status).toBe(200); + const { total } = await readAllChunks(response.body); + expect(total).toBe(Buffer.byteLength(bodyText, "utf8")); + console.log( + `[mistral bounded-stream proof] normal path: cap=${MAX} returned=${total} body=${JSON.stringify(bodyText)}`, + ); + } finally { + await new Promise((resolve) => { + server.close(() => { + resolve(); + }); + }); + } + }); +}); + +// Drive the bounded fetcher directly against a synthetic ReadableStream that +// exceeds the cap. Bypasses any HTTP layer; proves the cap fires against an +// unbounded chunked source, mirroring what the Mistral SDK's internal SSE +// parser (`EventStream`) would see when a streaming body exceeds 16 MiB. +describe("Mistral bounded-stream-read direct (synthetic ReadableStream)", () => { + it("caps an oversized synthetic ReadableStream at 16 MiB", async () => { + const fetcher = createBoundedMistralFetcher(MAX); + const CHUNK = 1024 * 1024; + let sent = 0; + const synthetic = new ReadableStream({ + pull(controller) { + if (sent < 18) { + controller.enqueue(new Uint8Array(CHUNK)); + sent++; + } else { + controller.close(); + } + }, + }); + // Build the same shape `fetcher` expects from a real fetch(): a + // `Response` whose `body` is a ReadableStream. + const syntheticResponse = new Response(synthetic, { + status: 200, + headers: { "content-type": "application/octet-stream" }, + }); + + let captured: Error | undefined; + try { + // Replace the fetcher's internal `fetch` call by exercising the + // post-fetchResponse code path directly: build a `Wrapped` + // that re-enters `fetcher` as if a real fetch returned our + // synthetic Response, by patching the global fetch. + const originalFetch = globalThis.fetch; + globalThis.fetch = (() => Promise.resolve(syntheticResponse)) as typeof globalThis.fetch; + try { + const wrapped = await fetcher("http://unused.invalid/"); + try { + await readAllChunks(wrapped.body); + } catch (err) { + captured = err as Error; + } + } finally { + globalThis.fetch = originalFetch; + } + expect(captured).toBeInstanceOf(Error); + const match = (captured as Error).message.match( + /mistral: stream body exceeds \d+ bytes \(got (\d+)\)/, + ); + expect(match).not.toBeNull(); + const got = Number(match![1]); + // Synthetic stream chunks are exactly 1 MiB aligned, so cap+1 reads + // give exactly cap + 1 MiB = 16 MiB + 1 MiB = 17 825 792 bytes. + expect(got).toBe(16777216 + CHUNK); + } finally { + // Best-effort cleanup if the test threw mid-flight. + // No intervals to clear for this test; the synthetic stream closes + // automatically when `sent >= 18`. + } + }); +}); diff --git a/src/llm/providers/mistral.test.ts b/src/llm/providers/mistral.test.ts index 8af0ec919ae2..4e50e08483d3 100644 --- a/src/llm/providers/mistral.test.ts +++ b/src/llm/providers/mistral.test.ts @@ -7,16 +7,26 @@ const mistralMockState = vi.hoisted(() => ({ payloads: [] as unknown[], })); -vi.mock("@mistralai/mistralai", () => ({ - Mistral: class MockMistral { - chat = { - stream: vi.fn(async (payload: unknown) => { - mistralMockState.payloads.push(payload); - throw new Error("stop before network"); - }), - }; - }, -})); +vi.mock("@mistralai/mistralai", async () => { + // Preserve real exports for everything except `Mistral`, so the new + // imports of `HTTPClient` and `Fetcher` introduced by the bounded-stream + // helper (`createBoundedMistralHttpClient`) resolve correctly. Only + // `Mistral` itself is overridden so the test can capture payloads without + // any actual HTTP traffic. + const actual = + await vi.importActual("@mistralai/mistralai"); + return { + ...actual, + Mistral: class MockMistral { + chat = { + stream: vi.fn(async (payload: unknown) => { + mistralMockState.payloads.push(payload); + throw new Error("stop before network"); + }), + }; + }, + }; +}); import { streamMistral, streamSimpleMistral } from "./mistral.js"; diff --git a/src/llm/providers/mistral.ts b/src/llm/providers/mistral.ts index c87536188757..6ff13cb61527 100644 --- a/src/llm/providers/mistral.ts +++ b/src/llm/providers/mistral.ts @@ -1,5 +1,5 @@ // Mistral provider adapts Mistral streams and tool calls to the runtime. -import { Mistral } from "@mistralai/mistralai"; +import { HTTPClient, Mistral, type Fetcher } from "@mistralai/mistralai"; import type { ChatCompletionStreamRequest, ChatCompletionStreamRequestMessage, @@ -7,6 +7,7 @@ import type { ContentChunk, FunctionTool, } from "@mistralai/mistralai/models/components"; +import { createSseByteGuard } from "../../agents/streaming-byte-guard.js"; import { stripSystemPromptCacheBoundary } from "../../agents/system-prompt-cache-boundary.js"; import { getEnvApiKey } from "../env-api-keys.js"; import { calculateCost, clampThinkingLevel } from "../model-utils.js"; @@ -34,6 +35,63 @@ import { transformMessages } from "./transform-messages.js"; const MISTRAL_TOOL_CALL_ID_LENGTH = 9; const MAX_MISTRAL_ERROR_BODY_CHARS = 4000; +// 16 MiB cap on Mistral streaming success bodies, matching the +// `PROVIDER_TEXT_RESPONSE_MAX_BYTES` / `PROVIDER_JSON_RESPONSE_MAX_BYTES` +// 16 MiB cap used elsewhere. A hostile or malfunctioning Mistral-compatible +// endpoint cannot exhaust memory by streaming an unbounded SSE body; +// `createSseByteGuard` cancels the upstream reader and throws once the +// accumulated byte count exceeds this cap. +const MISTRAL_STREAM_BODY_MAX_BYTES = 16 * 1024 * 1024; + +/** + * Builds a `Fetcher` that wraps the default `fetch` with a 16 MiB byte cap + * on streamed response bodies. The wrapped `Response.body` exposes a + * `ReadableStream` whose chunks flow through `createSseByteGuard`, so the + * SDK's internal SSE parser (`EventStream` in + * `@mistralai/mistralai/lib/event-streams.ts`) reads exactly as it would on + * an unbounded body — but bounded. + * + * Bodyless responses (no `body` or no `getReader`) are returned unchanged so + * the SDK's error-path `res.arrayBuffer()` call still works. + */ +export function createBoundedMistralFetcher( + maxBytes: number = MISTRAL_STREAM_BODY_MAX_BYTES, +): Fetcher { + return async (input, init) => { + const response = init == null ? await fetch(input) : await fetch(input, init); + if (!response.body || typeof response.body.getReader !== "function") { + return response; + } + const reader = response.body.getReader(); + const guard = createSseByteGuard(reader, { + maxBytes, + onOverflow: ({ size, maxBytes: cap }) => + new Error(`mistral: stream body exceeds ${cap} bytes (got ${size})`), + }); + // Re-shape the response body so the SDK's `responseBody.getReader()` + // call inside `EventStream` resolves to a stream whose `read()` is + // routed through `guard.read()`. Cancellation is also forwarded. + const guardedStream = new ReadableStream({ + async pull(controller) { + const { done, value } = await guard.read(); + if (done) { + controller.close(); + return; + } + controller.enqueue(value); + }, + async cancel(reason) { + await guard.cancel(reason); + }, + }); + return new Response(guardedStream, { + status: response.status, + statusText: response.statusText, + headers: response.headers, + }); + }; +} + /** * Provider-specific options for the Mistral API. */ @@ -73,6 +131,14 @@ export const streamMistral: StreamFunction<"mistral-conversations", MistralOptio const mistral = new Mistral({ apiKey, serverURL: model.baseUrl, + // Bound the streamed Mistral response body at 16 MiB so a hostile or + // malfunctioning endpoint cannot exhaust memory. The fetcher is + // injected via the SDK's `HTTPClient` (see + // `@mistralai/mistralai/lib/sdks.ts` `ClientSDK` constructor: when + // `httpClient` is passed, `ClientSDK.#httpClient` is set from it and + // every `chat.stream` / `complete` call routes through + // `HTTPClient.request` → `this.fetcher(req)`). + httpClient: new HTTPClient({ fetcher: createBoundedMistralFetcher() }), }); const normalizeMistralToolCallId = createMistralToolCallIdNormalizer();