Improve error handling and add response size cap with timeout

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,28 +85,59 @@ 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.")
 
             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

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

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<Response> {
+    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<unknown> {
+    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.
  */