diff --git a/python/src/typechat/_internal/model.py b/python/src/typechat/_internal/model.py index da52e306..471404f0 100644 --- a/python/src/typechat/_internal/model.py +++ b/python/src/typechat/_internal/model.py @@ -1,4 +1,5 @@ import asyncio +import json from types import TracebackType from typing_extensions import AsyncContextManager, Literal, Protocol, Self, TypedDict, cast, override @@ -35,6 +36,11 @@ async def complete(self, prompt: str | list[PromptSection]) -> Result[str]: 504, ] +class _ResponseTooLargeError(Exception): + """Raised when a model response body exceeds the configured maximum size.""" + def __init__(self, max_bytes: int): + super().__init__(f"REST API response exceeded the maximum allowed size of {max_bytes} bytes") + class HttpxLanguageModel(TypeChatLanguageModel, AsyncContextManager): url: str headers: dict[str, str] @@ -46,6 +52,11 @@ class HttpxLanguageModel(TypeChatLanguageModel, AsyncContextManager): # Specifies how long a request should wait in seconds # before timing out with a Failure. timeout_seconds = 10 + # Specifies the maximum size, in bytes, of a response body that will be read from the + # model endpoint (the default is 100 MB). A larger response is rejected without being fully + # buffered in memory, preventing a malicious or malfunctioning endpoint from exhausting memory. + # Set to 0 or a negative value to disable the limit. + max_response_bytes: int = 100 * 1024 * 1024 _async_client: httpx.AsyncClient def __init__(self, url: str, headers: dict[str, str], default_params: dict[str, str]): @@ -74,21 +85,25 @@ async def complete(self, prompt: str | list[PromptSection]) -> Success[str] | Fa retry_count = 0 while True: try: - response = await self._async_client.post( + async with self._async_client.stream( + "POST", self.url, headers=headers, json=body, - timeout=self.timeout_seconds - ) - if response.is_success: - json_result = cast( - dict[Literal["choices"], list[dict[Literal["message"], PromptSection]]], - response.json() - ) - return Success(json_result["choices"][0]["message"]["content"] or "") - - if response.status_code not in _TRANSIENT_ERROR_CODES or retry_count >= self.max_retry_attempts: - return Failure(f"REST API error {response.status_code}: {response.reason_phrase}") + timeout=self.timeout_seconds, + ) as response: + if response.is_success: + raw = await self._read_capped(response) + json_result = cast( + dict[Literal["choices"], list[dict[Literal["message"], PromptSection]]], + json.loads(raw) + ) + return Success(json_result["choices"][0]["message"]["content"] or "") + + if response.status_code not in _TRANSIENT_ERROR_CODES or retry_count >= self.max_retry_attempts: + return Failure(f"REST API error {response.status_code}: {response.reason_phrase}") + except _ResponseTooLargeError as e: + return Failure(str(e)) except Exception as e: if retry_count >= self.max_retry_attempts: return Failure(str(e) or f"{repr(e)} raised from within internal TypeChat language model.") @@ -96,6 +111,33 @@ async def complete(self, prompt: str | list[PromptSection]) -> Success[str] | Fa await asyncio.sleep(self.retry_pause_seconds) retry_count += 1 + async def _read_capped(self, response: httpx.Response) -> bytes: + """ + Reads a response body while enforcing `max_response_bytes`. The body is read incrementally + and the read is aborted as soon as the accumulated size exceeds the limit, so an oversized + response is never fully buffered in memory. A non-positive limit disables the check. + """ + max_bytes = self.max_response_bytes + if max_bytes <= 0: + return await response.aread() + + # Fail fast when the server advertises an oversized body up front. + content_length = response.headers.get("content-length") + if content_length is not None: + try: + advertised = int(content_length) + except ValueError: + advertised = None + if advertised is not None and advertised > max_bytes: + raise _ResponseTooLargeError(max_bytes) + + buffer = bytearray() + async for chunk in response.aiter_bytes(): + buffer.extend(chunk) + if len(buffer) > max_bytes: + raise _ResponseTooLargeError(max_bytes) + return bytes(buffer) + @override async def __aenter__(self) -> Self: return self diff --git a/python/tests/test_model.py b/python/tests/test_model.py new file mode 100644 index 00000000..5385b73c --- /dev/null +++ b/python/tests/test_model.py @@ -0,0 +1,83 @@ +""" +Tests for HttpxLanguageModel response-size limiting (DoS hardening). + +These use httpx.MockTransport to drive HttpxLanguageModel.complete without a real endpoint. +""" + +import asyncio +from collections.abc import Callable + +import httpx +import typechat +from typechat._internal.model import HttpxLanguageModel + + +class _MockHttpxLanguageModel(HttpxLanguageModel): + def use_mock_transport(self, handler: Callable[[httpx.Request], httpx.Response]) -> None: + self._async_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + + +def _make_model(handler: Callable[[httpx.Request], httpx.Response]) -> HttpxLanguageModel: + model = _MockHttpxLanguageModel( + url="https://example.invalid/v1/chat/completions", + headers={}, + default_params={"model": "gpt-test"}, + ) + # Route the model's requests through a mock transport instead of the network. + model.use_mock_transport(handler) + return model + + +def _completion_payload(content: str) -> dict[str, list[dict[str, dict[str, str]]]]: + return {"choices": [{"message": {"role": "assistant", "content": content}}]} + + +def test_reads_response_within_size_limit(): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=_completion_payload("Hello!")) + + model = _make_model(handler) + result = asyncio.run(model.complete("hi")) + assert isinstance(result, typechat.Success) + assert result.value == "Hello!" + + +def test_rejects_oversized_response_via_content_length(): + def handler(request: httpx.Request) -> httpx.Response: + # httpx sets Content-Length for an eager JSON body, exercising the fast-fail path. + return httpx.Response(200, json=_completion_payload("x" * 5000)) + + model = _make_model(handler) + model.max_response_bytes = 100 + result = asyncio.run(model.complete("hi")) + assert isinstance(result, typechat.Failure) + assert "maximum allowed size" in result.message + + +def test_rejects_oversized_streamed_response_without_content_length(): + async def body(): + # A chunked body (no Content-Length) forces the incremental accumulation check. + for _ in range(16): + yield b"x" * 256 + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, content=body()) + + model = _make_model(handler) + model.max_response_bytes = 512 + result = asyncio.run(model.complete("hi")) + assert isinstance(result, typechat.Failure) + assert "maximum allowed size" in result.message + + +def test_size_limit_disabled_allows_large_response(): + big_content = "x" * 5000 + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=_completion_payload(big_content)) + + model = _make_model(handler) + model.max_response_bytes = 0 + result = asyncio.run(model.complete("hi")) + assert isinstance(result, typechat.Success) + assert result.value == big_content diff --git a/typescript/src/model.ts b/typescript/src/model.ts index 39abf6cb..f20df5d8 100644 --- a/typescript/src/model.ts +++ b/typescript/src/model.ts @@ -68,6 +68,21 @@ export interface TypeChatLanguageModel { * Optional property that specifies the delay before retrying in milliseconds (the default is 1000ms). */ retryPauseMs?: number; + /** + * Optional property that specifies the timeout, in milliseconds, applied to each HTTP request + * made to the model endpoint (the default is 600000ms, i.e. 10 minutes). A request - including + * reading its response body - is aborted if it does not complete within this time, preventing a + * slow or unresponsive endpoint from causing `complete` to hang indefinitely and exhaust + * resources. Set to 0 or a negative value to disable the timeout. + */ + timeoutMs?: number; + /** + * Optional property that specifies the maximum size, in bytes, of a response body accepted from + * the model endpoint (the default is 104857600, i.e. 100 MB). A larger response is rejected + * without being fully buffered in memory, preventing a malicious or malfunctioning endpoint from + * exhausting memory. Set to 0 or a negative value to disable the limit. + */ + maxResponseBytes?: number; /** * Obtains a completion from the language model for the given prompt. * @param prompt A prompt string or an array of prompt sections. If a string is specified, @@ -205,6 +220,20 @@ export function createAzureOpenAILanguageModel(apiKey: string, endPoint: string, return createFetchLanguageModel(endPoint, headers, {}, resolveProxySettings(options)); } +/** + * Default per-request timeout applied when a model does not set {@link TypeChatLanguageModel.timeoutMs}. + * Ten minutes matches common LLM client defaults, allowing slow completions while still preventing a + * request from hanging indefinitely. + */ +const defaultTimeoutMs = 600000; + +/** + * Default maximum response body size (100 MB) applied when a model does not set + * {@link TypeChatLanguageModel.maxResponseBytes}. Generous enough for any legitimate completion while + * still preventing an endpoint from exhausting memory with an unbounded response. + */ +const defaultMaxResponseBytes = 100 * 1024 * 1024; + /** * Common OpenAI REST API endpoint encapsulation using the fetch API. */ @@ -221,6 +250,8 @@ function createFetchLanguageModel(url: string, headers: object, defaultParams: o let retryCount = 0; const retryMaxAttempts = model.retryMaxAttempts ?? 3; const retryPauseMs = model.retryPauseMs ?? 1000; + const timeoutMs = model.timeoutMs ?? defaultTimeoutMs; + const maxResponseBytes = model.maxResponseBytes ?? defaultMaxResponseBytes; const messages = typeof prompt === "string" ? [{ role: "user", content: prompt }] : prompt; while (true) { const options: RequestInit = { @@ -239,13 +270,29 @@ function createFetchLanguageModel(url: string, headers: object, defaultParams: o if (dispatcher) { options.dispatcher = dispatcher; } - const response = await fetch(url, options); + let response: Response; + try { + response = await fetchWithTimeout(url, options, timeoutMs); + } catch (e) { + if (retryCount >= retryMaxAttempts) { + return error(`REST API fetch error: ${getErrorMessage(e)}`); + } + await sleep(retryPauseMs); + retryCount++; + continue; + } if (response.ok) { - const json = await response.json() as { choices: { message: PromptSection }[] }; - if (typeof json.choices[0].message.content === "string") { - return success(json.choices[0].message.content ?? ""); + let json: { choices?: { message?: PromptSection }[] } | null; + try { + json = await readResponseJson(response, maxResponseBytes) as { choices?: { message?: PromptSection }[] } | null; + } catch (e) { + return error(`REST API response error: ${getErrorMessage(e)}`); + } + const content = json?.choices?.[0]?.message?.content; + if (typeof content === "string") { + return success(content); } else { - return error(`REST API unexpected response format: ${JSON.stringify(json.choices[0].message.content)}`); + return error(`REST API unexpected response format: ${JSON.stringify(json)}`); } } if (!isTransientHttpError(response.status) || retryCount >= retryMaxAttempts) { @@ -296,6 +343,8 @@ function createResponsesFetchLanguageModel(url: string, headers: object, default let retryCount = 0; const retryMaxAttempts = model.retryMaxAttempts ?? 3; const retryPauseMs = model.retryPauseMs ?? 1000; + const timeoutMs = model.timeoutMs ?? defaultTimeoutMs; + const maxResponseBytes = model.maxResponseBytes ?? defaultMaxResponseBytes; const input = typeof prompt === "string" ? prompt : (prompt as PromptSection[]); while (true) { const options: RequestInit = { @@ -313,15 +362,39 @@ function createResponsesFetchLanguageModel(url: string, headers: object, default if (dispatcher) { options.dispatcher = dispatcher; } - const response = await fetch(url, options); + let response: Response; + try { + response = await fetchWithTimeout(url, options, timeoutMs); + } catch (e) { + if (retryCount >= retryMaxAttempts) { + return error(`REST API fetch error: ${getErrorMessage(e)}`); + } + await sleep(retryPauseMs); + retryCount++; + continue; + } if (response.ok) { type ResponsesAPIOutputItem = { type: string; role?: string; content: { type: string; text: string }[]; }; - const json = await response.json() as { output: ResponsesAPIOutputItem[] }; - const message = json.output?.find(o => o.type === "message"); + let json: unknown; + try { + json = await readResponseJson(response, maxResponseBytes); + } catch (e) { + return error(`REST API response error: ${getErrorMessage(e)}`); + } + if ( + json === null || + typeof json !== "object" || + !("output" in json) || + !Array.isArray((json as { output?: unknown }).output) + ) { + return error(`REST API unexpected response format: ${JSON.stringify(json)}`); + } + const output = (json as { output: ResponsesAPIOutputItem[] }).output; + const message = output.find(o => o.type === "message"); const textContent = message?.content?.find(c => c.type === "output_text"); if (textContent?.text !== undefined) { return success(textContent.text); @@ -386,6 +459,59 @@ function isResponsesApiUrl(url: string): boolean { } } +/** + * Performs a `fetch` request that is aborted if it does not complete within `timeoutMs` + * milliseconds. Because the abort signal stays attached to the returned response, the timeout also + * bounds the time spent reading the response body. A non-positive `timeoutMs` disables the timeout. + * Enforcing a timeout prevents a slow or unresponsive endpoint from causing the request to hang + * indefinitely and exhaust resources. + */ +async function fetchWithTimeout(url: string, options: RequestInit, timeoutMs: number): Promise { + if (timeoutMs > 0) { + options.signal = AbortSignal.timeout(timeoutMs); + } + return fetch(url, options); +} + +/** + * Reads and JSON-parses a response body while enforcing a maximum size in bytes. When the runtime + * exposes a readable stream body (as Node's `fetch` does) and `maxBytes` is positive, the body is + * read incrementally and the read is aborted as soon as the accumulated size exceeds `maxBytes`, so + * an oversized response is never fully buffered in memory. When no stream body is available or the + * limit is disabled, this falls back to `response.json()`. + */ +async function readResponseJson(response: Response, maxBytes: number): Promise { + const body = response.body; + if (maxBytes > 0 && body && typeof body.getReader === "function") { + const reader = body.getReader(); + const decoder = new TextDecoder(); + let received = 0; + let text = ""; + while (true) { + const result = await reader.read(); + if (result.done) { + break; + } + received += result.value.byteLength; + if (received > maxBytes) { + await reader.cancel().catch(() => { /* ignore cancellation errors */ }); + throw new Error(`REST API response exceeded the maximum allowed size of ${maxBytes} bytes`); + } + text += decoder.decode(result.value, { stream: true }); + } + text += decoder.decode(); + return JSON.parse(text); + } + return response.json(); +} + +/** + * Extracts a human-readable message from a value thrown by `fetch` or response processing. + */ +function getErrorMessage(e: unknown): string { + return e instanceof Error ? e.message : String(e); +} + /** * Sleeps for the given number of milliseconds. */ diff --git a/typescript/tests/model.test.mjs b/typescript/tests/model.test.mjs index a5f3e296..d8417ac5 100644 --- a/typescript/tests/model.test.mjs +++ b/typescript/tests/model.test.mjs @@ -29,6 +29,15 @@ function makeChatCompletionsResponse(content) { }; } +function makeJsonResponse(body) { + return { + ok: true, + status: 200, + headers: { get: (_name) => null }, + json: () => Promise.resolve(body), + }; +} + function makeResponsesAPIResponse(text) { return { ok: true, @@ -145,6 +154,50 @@ describe("createOpenAILanguageModel (Chat Completions API)", () => { assert.ok(result.message.includes("401")); }); + // Regression tests: a 200 OK response whose body does not conform to the + // expected Chat Completions shape must return a recoverable error rather + // than throwing a TypeError (which would reject the promise and can lead to + // a Denial of Service in callers that expect a Result). + test("returns error when choices array is missing", async () => { + setupFetch([makeJsonResponse({ id: "chatcmpl-123" })]); + const model = createOpenAILanguageModel("sk-test", "gpt-4"); + const result = await model.complete("test"); + assert.equal(result.success, false); + assert.ok(result.message.includes("unexpected response format")); + }); + + test("returns error when choices array is empty", async () => { + setupFetch([makeJsonResponse({ choices: [] })]); + const model = createOpenAILanguageModel("sk-test", "gpt-4"); + const result = await model.complete("test"); + assert.equal(result.success, false); + assert.ok(result.message.includes("unexpected response format")); + }); + + test("returns error when message is missing from choice", async () => { + setupFetch([makeJsonResponse({ choices: [{}] })]); + const model = createOpenAILanguageModel("sk-test", "gpt-4"); + const result = await model.complete("test"); + assert.equal(result.success, false); + assert.ok(result.message.includes("unexpected response format")); + }); + + test("returns error when content is not a string", async () => { + setupFetch([makeJsonResponse({ choices: [{ message: { role: "assistant", content: null } }] })]); + const model = createOpenAILanguageModel("sk-test", "gpt-4"); + const result = await model.complete("test"); + assert.equal(result.success, false); + assert.ok(result.message.includes("unexpected response format")); + }); + + test("returns error when response body is JSON null", async () => { + setupFetch([makeJsonResponse(null)]); + const model = createOpenAILanguageModel("sk-test", "gpt-4"); + const result = await model.complete("test"); + assert.equal(result.success, false); + assert.ok(result.message.includes("unexpected response format")); + }); + test("auto-detects Responses API from a /responses endpoint URL", async () => { setupFetch([makeResponsesAPIResponse("Auto-detected!")]); const model = createOpenAILanguageModel("sk-test", "gpt-4", "https://api.openai.com/v1/responses"); @@ -325,3 +378,111 @@ describe("createLanguageModel environment variable routing", () => { }); }); +// --------------------------------------------------------------------------- +// Request timeout and response size limits (DoS hardening) +// --------------------------------------------------------------------------- + +// Mirrors the error fetch throws when an AbortSignal.timeout() fires. +function makeTimeoutError() { + return new DOMException("The operation timed out.", "TimeoutError"); +} + +// Builds a Response-like object whose body is a ReadableStream (as Node's fetch returns), so the +// size-limited streaming read path is exercised. json() throws to prove it is not used. +function makeStreamingResponse(payload, chunkSize = 16) { + const bytes = new TextEncoder().encode(payload); + return { + ok: true, + status: 200, + headers: { get: (_name) => null }, + json: () => { + throw new Error("json() should not be called when a stream body is present"); + }, + body: new ReadableStream({ + start(controller) { + for (let i = 0; i < bytes.length; i += chunkSize) { + controller.enqueue(bytes.subarray(i, i + chunkSize)); + } + controller.close(); + }, + }), + }; +} + +describe("request timeout and response size limits", () => { + after(teardownFetch); + + test("attaches an AbortSignal to each request by default", async () => { + setupFetch([makeChatCompletionsResponse("OK")]); + const model = createOpenAILanguageModel("sk-test", "gpt-4"); + await model.complete("test"); + const signal = capturedRequests[0].options.signal; + assert.ok(signal, "Expected an AbortSignal on the request options"); + assert.equal(typeof signal.aborted, "boolean", "Expected an AbortSignal-like object"); + }); + + test("omits the AbortSignal when timeoutMs is set to 0", async () => { + setupFetch([makeChatCompletionsResponse("OK")]); + const model = createOpenAILanguageModel("sk-test", "gpt-4"); + model.timeoutMs = 0; + await model.complete("test"); + assert.equal(capturedRequests[0].options.signal, undefined, "Expected no AbortSignal when the timeout is disabled"); + }); + + test("returns an error when every attempt times out", async () => { + let attempts = 0; + globalThis.fetch = async () => { + attempts++; + throw makeTimeoutError(); + }; + const model = createOpenAILanguageModel("sk-test", "gpt-4"); + model.retryMaxAttempts = 2; + model.retryPauseMs = 0; + const result = await model.complete("test"); + assert.equal(result.success, false); + assert.ok(result.message.includes("fetch error"), `Unexpected message: ${result.message}`); + assert.equal(attempts, 3, "Expected the initial attempt plus 2 retries"); + }); + + test("retries after a transient fetch failure and then succeeds", async () => { + let attempts = 0; + globalThis.fetch = async () => { + attempts++; + if (attempts === 1) { + throw makeTimeoutError(); + } + return makeChatCompletionsResponse("Recovered"); + }; + const model = createOpenAILanguageModel("sk-test", "gpt-4"); + model.retryMaxAttempts = 3; + model.retryPauseMs = 0; + const result = await model.complete("test"); + assert.equal(result.success, true); + assert.equal(result.data, "Recovered"); + assert.equal(attempts, 2, "Expected one failure followed by one success"); + }); + + test("rejects a response body that exceeds maxResponseBytes", async () => { + const payload = JSON.stringify({ + choices: [{ message: { role: "assistant", content: "x".repeat(500) } }], + }); + setupFetch([makeStreamingResponse(payload)]); + const model = createOpenAILanguageModel("sk-test", "gpt-4"); + model.maxResponseBytes = 50; + const result = await model.complete("test"); + assert.equal(result.success, false); + assert.ok(result.message.includes("maximum allowed size"), `Unexpected message: ${result.message}`); + }); + + test("reads a streaming response body within maxResponseBytes", async () => { + const payload = JSON.stringify({ + choices: [{ message: { role: "assistant", content: "Streamed!" } }], + }); + setupFetch([makeStreamingResponse(payload)]); + const model = createOpenAILanguageModel("sk-test", "gpt-4"); + const result = await model.complete("test"); + assert.equal(result.success, true); + assert.equal(result.data, "Streamed!"); + }); +}); +