Bound the size and duration of a chat request

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_"}
 
 

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:

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 |

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