From 93c7671c3ae80e2987677d25280de4a48f508b7d Mon Sep 17 00:00:00 2001 From: Krstan Vjestica Date: Wed, 19 Aug 2026 20:10:00 +0200 Subject: [PATCH] fix: nothing bounded the size or duration of a chat request `chat_max_concurrency` caps how many generations run at once and says nothing about how large or how long any one of them is. Four callers could hold every slot for the full Ollama timeout with a prompt the size of a book, and the service would look healthy the whole time. - The message is capped at 4000 characters (413 over it) and a blank one is rejected outright rather than sent to the model. - History is trimmed to the most recent 20 turns and 16000 characters rather than rejected: dropping the oldest turns costs a little context, while a 413 in the middle of a conversation ends it. History matters more than the message here -- every turn is re-sent to the model and paid for again on the next request. - The page context is truncated to 200 characters. It is a route, not a payload. - A generation gets a 180s deadline independent of the model's own timeout. One that will not stop still ends, because the slot it holds is one of only four. Defaults are settings, so an operator who wants a longer conversation can have one. Documented alongside the chat token, whose empty default leaves the endpoint open -- the chart now always sets it. --- app/config.py | 15 ++++++ app/main.py | 59 ++++++++++++++++++++-- docs/configuration.md | 11 +++++ tests/test_api.py | 110 ++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 192 insertions(+), 3 deletions(-) diff --git a/app/config.py b/app/config.py index c4a6992..0662cee 100644 --- a/app/config.py +++ b/app/config.py @@ -41,6 +41,21 @@ class Settings(BaseSettings): # exhausts GPU/CPU. Excess requests get 429. chat_max_concurrency: int = 4 + # Bounds on a single request (Codex M3). The concurrency cap limits how many + # generations run at once, but says nothing about how large or how long any + # one of them is -- four callers could hold every slot for the full Ollama + # timeout with a prompt the size of a book. + # + # A question is a question: 4000 characters is longer than anyone types. + chat_max_message_chars: int = 4000 + # Turns of prior conversation kept. Each one is re-sent to the model, so an + # unbounded history is an unbounded prompt, paid for on every request. + chat_max_history_turns: int = 20 + chat_max_history_chars: int = 16000 + # Hard ceiling on one streamed response, independent of the model's own + # timeout. A generation that will not stop still ends. + chat_deadline_seconds: int = 180 + model_config = {"env_prefix": "FORAIL_ASSISTANT_"} diff --git a/app/main.py b/app/main.py index abacf92..d2929d0 100644 --- a/app/main.py +++ b/app/main.py @@ -4,6 +4,7 @@ import hmac import json import logging +import time from fastapi import FastAPI, Header, HTTPException from fastapi.middleware.cors import CORSMiddleware @@ -72,6 +73,46 @@ class ChatRequest(BaseModel): history: list[dict] | None = None +def _bounded_request(req: "ChatRequest") -> tuple[str, list[dict]]: + """ + The message and history this request is allowed to spend, or 413. + + The concurrency cap limits how many generations run at once and says nothing + about how large any one of them is: four callers could hold every slot for + the full Ollama timeout with a prompt the size of a book. History matters + more than the message, because every turn is re-sent to the model and paid + for again on the next request. + """ + message = (req.message or "").strip() + if not message: + raise HTTPException(status_code=400, detail="message must not be empty") + if len(message) > settings.chat_max_message_chars: + raise HTTPException( + status_code=413, + detail=f"message must be at most {settings.chat_max_message_chars} characters", + ) + + history = req.history or [] + if not isinstance(history, list): + raise HTTPException(status_code=400, detail="history must be a list") + + # Trimmed rather than rejected: dropping the oldest turns degrades the answer + # a little, while a 413 in the middle of a conversation ends it. + history = history[-settings.chat_max_history_turns:] + budget = settings.chat_max_history_chars + kept: list[dict] = [] + for turn in reversed(history): + if not isinstance(turn, dict): + continue + cost = len(str(turn.get("content", ""))) + if cost > budget: + break + budget -= cost + kept.append(turn) + kept.reverse() + return message, kept + + class HealthResponse(BaseModel): status: str version: str @@ -129,18 +170,30 @@ async def chat(req: ChatRequest, authorization: str | None = Header(default=None if _chat_semaphore.locked(): raise HTTPException(status_code=429, detail="Assistant busy, retry shortly") + message, history = _bounded_request(req) + page_context = "" if req.context and req.context.get("page"): - page_context = req.context["page"] + page_context = str(req.context["page"])[:200] async def event_generator(): async with _chat_semaphore: + deadline = time.monotonic() + settings.chat_deadline_seconds try: async for token in stream_chat( - message=req.message, + message=message, page_context=page_context, - history=req.history, + history=history, ): + # A generation that will not stop still has to end: the + # slot it holds is one of only chat_max_concurrency. + if time.monotonic() > deadline: + logger.warning( + "Chat generation exceeded %ss deadline; cutting the stream", + settings.chat_deadline_seconds, + ) + yield {"data": json.dumps({"error": "response timed out", "done": True})} + return yield {"data": json.dumps({"token": token})} yield {"data": json.dumps({"done": True})} except Exception: diff --git a/docs/configuration.md b/docs/configuration.md index 9edf6bd..4dd39b0 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -23,6 +23,17 @@ All configuration is via environment variables with the `FORAIL_ASSISTANT_` pref | `FORAIL_ASSISTANT_CHROMA_PORT` | `8000` | ChromaDB port | | `FORAIL_ASSISTANT_CHROMA_COLLECTION` | `forail_docs` | Collection name for indexed documents | +### Request Limits + +| Variable | Default | Description | +|----------|---------|-------------| +| `FORAIL_ASSISTANT_CHAT_TOKEN` | `""` | Bearer token required on `/api/v1/chat`. **Empty means the endpoint is open** — set it whenever the service is reachable by anything you do not control | +| `FORAIL_ASSISTANT_CHAT_MAX_CONCURRENCY` | `4` | Concurrent generations; excess requests get 429 | +| `FORAIL_ASSISTANT_CHAT_MAX_MESSAGE_CHARS` | `4000` | Longest accepted question; over it returns 413 | +| `FORAIL_ASSISTANT_CHAT_MAX_HISTORY_TURNS` | `20` | Prior turns kept. Trimmed, not rejected — every turn is re-sent to the model on each request | +| `FORAIL_ASSISTANT_CHAT_MAX_HISTORY_CHARS` | `16000` | Total history size kept, oldest turns dropped first | +| `FORAIL_ASSISTANT_CHAT_DEADLINE_SECONDS` | `180` | Hard ceiling on one streamed response. A generation that will not stop still ends, so it cannot hold a concurrency slot indefinitely | + ### RAG Settings | Variable | Default | Description | diff --git a/tests/test_api.py b/tests/test_api.py index 2b60176..aae88a3 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -156,3 +156,113 @@ def test_openapi_schema(self, client): schema = resp.json() assert "/api/v1/health" in schema["paths"] assert "/api/v1/chat" in schema["paths"] + + +class TestChatRequestBounds: + """ + Codex M3: the concurrency cap limits how many generations run at once and + says nothing about how large or how long any one of them is. Four callers + could hold every slot for the full Ollama timeout with a prompt the size of + a book. + """ + + def _stream(self): + async def mock_stream(*args, **kwargs): + yield "ok" + return mock_stream + + def test_oversized_message_is_refused(self, client): + from app.config import settings + + resp = client.post( + "/api/v1/chat", + json={"message": "x" * (settings.chat_max_message_chars + 1)}, + ) + assert resp.status_code == 413 + + def test_message_at_the_limit_is_accepted(self, client): + from app.config import settings + + with patch("app.main.stream_chat", side_effect=self._stream()): + resp = client.post( + "/api/v1/chat", + json={"message": "x" * settings.chat_max_message_chars}, + ) + assert resp.status_code == 200 + + def test_blank_message_is_refused(self, client): + resp = client.post("/api/v1/chat", json={"message": " "}) + assert resp.status_code == 400 + + def test_history_is_trimmed_to_the_most_recent_turns(self, client): + # Trimmed rather than rejected: dropping the oldest turns costs a little + # context, while a 413 mid-conversation ends it. + from app.config import settings + + captured = {} + + async def mock_stream(*args, **kwargs): + captured.update(kwargs) + yield "ok" + + history = [{"role": "user", "content": f"turn {i}"} for i in range(200)] + with patch("app.main.stream_chat", side_effect=mock_stream): + resp = client.post("/api/v1/chat", json={"message": "hi", "history": history}) + + assert resp.status_code == 200 + assert len(captured["history"]) <= settings.chat_max_history_turns + # The turns kept are the recent ones, not the first ones. + assert captured["history"][-1]["content"] == "turn 199" + + def test_history_is_trimmed_by_total_size(self, client): + from app.config import settings + + captured = {} + + async def mock_stream(*args, **kwargs): + captured.update(kwargs) + yield "ok" + + history = [{"role": "user", "content": "x" * 5000} for _ in range(10)] + with patch("app.main.stream_chat", side_effect=mock_stream): + client.post("/api/v1/chat", json={"message": "hi", "history": history}) + + total = sum(len(t["content"]) for t in captured["history"]) + assert total <= settings.chat_max_history_chars + + def test_page_context_is_truncated(self, client): + captured = {} + + async def mock_stream(*args, **kwargs): + captured.update(kwargs) + yield "ok" + + with patch("app.main.stream_chat", side_effect=mock_stream): + client.post( + "/api/v1/chat", + json={"message": "hi", "context": {"page": "/x" * 5000}}, + ) + assert len(captured["page_context"]) <= 200 + + def test_a_generation_that_will_not_stop_is_cut(self, client): + # The slot it holds is one of only chat_max_concurrency, so an endless + # stream is a denial of service against the other three. A deadline + # already in the past is the same code path as one that runs out. + from app.config import settings + + emitted = 0 + + async def endless(*args, **kwargs): + nonlocal emitted + for _ in range(100): + emitted += 1 + yield "token" + + with patch("app.main.stream_chat", side_effect=endless), \ + patch.object(settings, "chat_deadline_seconds", -1): + resp = client.post("/api/v1/chat", json={"message": "hi"}) + + assert resp.status_code == 200 + assert "timed out" in resp.text + # Cut, not drained: the generator does not run to completion. + assert emitted < 100