From e80c41623b8cc71ea75f9f612c500fbf1cd5c8e2 Mon Sep 17 00:00:00 2001 From: agent-bot Date: Sun, 5 Jul 2026 01:53:58 +0000 Subject: [PATCH] Implement issue #27 --- README.md | 4 +- app/main.py | 8 + app/request_limits.py | 41 ++++- tests/contract/test_http_contract.py | 197 +++++++++++++++++++--- tests/contract/test_mcp_http_transport.py | 86 +++++++++- 5 files changed, 295 insertions(+), 41 deletions(-) diff --git a/README.md b/README.md index 5cba1fb..85e6fab 100644 --- a/README.md +++ b/README.md @@ -207,13 +207,13 @@ Runtime limits are configured with environment variables: | `MEMORIES_RATE_LIMIT_WRITES_PER_MINUTE` | `30` | Applies to REST create, update, and delete requests. | | `MEMORIES_RATE_LIMIT_BATCH_PER_MINUTE` | `10` | Applies to `POST /memories/batch`. | | `MEMORIES_RATE_LIMIT_MCP_PER_MINUTE` | `240` | Applies to MCP streamable HTTP traffic under `/mcp`. | -| `MEMORIES_REQUEST_BODY_MAX_BYTES` | `1048576` | Applies to `POST` and `PATCH` REST/MCP HTTP requests with `Content-Length`. | +| `MEMORIES_REQUEST_BODY_MAX_BYTES` | `1048576` | Applies to `POST` and `PATCH` REST/MCP HTTP request bodies. | Rate limiting uses a fixed 60-second in-memory window per process. Clients are identified by a trimmed `X-Client-Id` header, capped at 128 characters, when present; otherwise the client IP is used. A limited request returns `429` with `Retry-After`, `X-RateLimit-Limit`, `X-RateLimit-Remaining`, and `X-RateLimit-Reset` headers. `/health` and `/ready` are exempt from normal rate limiting. -Request body-size enforcement uses the `Content-Length` header and returns `413` with `X-Request-Body-Limit`. This covers normal local agent clients but does not fully cover oversized chunked or missing-length bodies. +Request body-size enforcement returns `413` with `X-Request-Body-Limit` when the declared or streamed request body exceeds the configured limit. For behavioral load checks, use the standalone stress harness at [scripts/rate_limit_stress.py](scripts/rate_limit_stress.py). It exercises mixed agent usage, fast bursts, concurrent floods, and two-agent isolation across REST and MCP surfaces, then writes an HTML report. See [scripts/README.md](scripts/README.md) for setup, tuning options, and result interpretation. diff --git a/app/main.py b/app/main.py index ad2b4f8..878abb1 100644 --- a/app/main.py +++ b/app/main.py @@ -18,6 +18,7 @@ FixedWindowRateLimiter, reject_request_body_if_too_large, reject_request_if_rate_limited, + reject_request_stream_if_too_large, ) from app.request_logging import current_duration_ms, log_http_request from app.schemas import ( @@ -159,6 +160,13 @@ def complete_response(response): if body_limit_response is not None: return complete_response(body_limit_response) + body_limit_response = await reject_request_stream_if_too_large( + request, + safety_config.request_body_max_bytes, + ) + if body_limit_response is not None: + return complete_response(body_limit_response) + rate_limit_response = reject_request_if_rate_limited( request, application.state.rate_limiter, diff --git a/app/request_limits.py b/app/request_limits.py index eebb38a..ee68d58 100644 --- a/app/request_limits.py +++ b/app/request_limits.py @@ -38,10 +38,13 @@ class _RateLimitWindow: def request_body_limit_applies(request: Request) -> bool: - if request.method.upper() not in BODY_LIMITED_METHODS: + return _body_limit_applies(request.method, request.url.path) + + +def _body_limit_applies(method: str, path: str) -> bool: + if method.upper() not in BODY_LIMITED_METHODS: return False - path = request.url.path return path == "/memories" or path.startswith("/memories/") or path.startswith("/mcp") @@ -68,12 +71,36 @@ def reject_request_body_if_too_large(request: Request, max_body_bytes: int) -> J and request_body_limit_applies(request) and content_length > max_body_bytes ): - return JSONResponse( - status_code=413, - content={"detail": "Request body too large"}, - headers={REQUEST_BODY_LIMIT_HEADER: str(max_body_bytes)}, - ) + return request_body_too_large_response(max_body_bytes) + + return None + + +def request_body_too_large_response(max_body_bytes: int) -> JSONResponse: + return JSONResponse( + status_code=413, + content={"detail": "Request body too large"}, + headers={REQUEST_BODY_LIMIT_HEADER: str(max_body_bytes)}, + ) + + +async def reject_request_stream_if_too_large( + request: Request, + max_body_bytes: int, +) -> JSONResponse | None: + if not request_body_limit_applies(request): + return None + + chunks: list[bytes] = [] + total_bytes = 0 + async for chunk in request.stream(): + total_bytes += len(chunk) + if total_bytes > max_body_bytes: + return request_body_too_large_response(max_body_bytes) + if chunk: + chunks.append(chunk) + request._body = b"".join(chunks) return None diff --git a/tests/contract/test_http_contract.py b/tests/contract/test_http_contract.py index a611633..239505a 100644 --- a/tests/contract/test_http_contract.py +++ b/tests/contract/test_http_contract.py @@ -1,3 +1,4 @@ +import asyncio import json import logging import sqlite3 @@ -40,6 +41,76 @@ def request_log_payloads(caplog): return payloads +def asgi_request( + application, + method: str, + path: str, + chunks: list[bytes], + headers: dict[str, str] | None = None, +): + return asyncio.run(_asgi_request(application, method, path, chunks, headers or {})) + + +async def _asgi_request( + application, + method: str, + path: str, + chunks: list[bytes], + headers: dict[str, str], +): + response_messages = [] + request_messages = [ + { + "type": "http.request", + "body": chunk, + "more_body": index < len(chunks) - 1, + } + for index, chunk in enumerate(chunks) + ] + request_index = 0 + + async def receive(): + nonlocal request_index + if request_index < len(request_messages): + message = request_messages[request_index] + request_index += 1 + return message + return {"type": "http.disconnect"} + + async def send(message): + response_messages.append(message) + + await application( + { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": method, + "scheme": "http", + "path": path, + "raw_path": path.encode(), + "query_string": b"", + "headers": [(name.lower().encode(), value.encode()) for name, value in headers.items()], + "client": ("testclient", 50000), + "server": ("testserver", 80), + "root_path": "", + }, + receive, + send, + ) + + response_start = next( + message for message in response_messages if message["type"] == "http.response.start" + ) + response_body = b"".join( + message.get("body", b"") + for message in response_messages + if message["type"] == "http.response.body" + ) + response_headers = {name.decode(): value.decode() for name, value in response_start["headers"]} + return response_start["status"], response_headers, response_body + + def test_health_check_returns_ping_without_initializing_database( client: TestClient, data_file: Path ): @@ -189,18 +260,24 @@ def test_post_memory_rejects_oversized_body_before_json_parsing( ): monkeypatch.setenv("MEMORIES_REQUEST_BODY_MAX_BYTES", "10") caplog.set_level(logging.INFO, logger="app.main") - test_client = TestClient(app_main.create_app()) + test_app = app_main.create_app() - response = test_client.post( + status_code, headers, body = asgi_request( + test_app, + "POST", "/memories", - content="x" * 11, - headers={"Content-Type": "application/json", "X-Request-Id": "oversized-request"}, + [b"x" * 11], + { + "Content-Type": "application/json", + "Content-Length": "11", + "X-Request-Id": "oversized-request", + }, ) - assert response.status_code == 413 - assert response.json() == {"detail": "Request body too large"} - assert response.headers["X-Request-Body-Limit"] == "10" - assert response.headers["X-Request-Id"] == "oversized-request" + assert status_code == 413 + assert json.loads(body) == {"detail": "Request body too large"} + assert headers["x-request-body-limit"] == "10" + assert headers["x-request-id"] == "oversized-request" request_logs = request_log_payloads(caplog) assert len(request_logs) == 1 record, payload = request_logs[0] @@ -212,6 +289,56 @@ def test_post_memory_rejects_oversized_body_before_json_parsing( assert read_database(data_file) == [] +def test_post_memory_rejects_oversized_body_without_content_length( + monkeypatch, data_file: Path, caplog +): + monkeypatch.setenv("MEMORIES_REQUEST_BODY_MAX_BYTES", "10") + caplog.set_level(logging.INFO, logger="app.main") + test_app = app_main.create_app() + + status_code, headers, body = asgi_request( + test_app, + "POST", + "/memories", + [b"x" * 6, b"x" * 5], + {"Content-Type": "application/json", "X-Request-Id": "streamed-request"}, + ) + + assert status_code == 413 + assert json.loads(body) == {"detail": "Request body too large"} + assert headers["x-request-body-limit"] == "10" + assert headers["x-request-id"] == "streamed-request" + request_logs = request_log_payloads(caplog) + assert len(request_logs) == 1 + record, payload = request_logs[0] + assert record.levelno == logging.INFO + assert payload["method"] == "POST" + assert payload["path"] == "/memories" + assert payload["status"] == 413 + assert payload["request_id"] == "streamed-request" + assert read_database(data_file) == [] + + +def test_post_memory_replays_body_without_content_length_under_limit(monkeypatch, data_file: Path): + monkeypatch.setenv("MEMORIES_REQUEST_BODY_MAX_BYTES", "100") + test_app = app_main.create_app() + + status_code, _headers, body = asgi_request( + test_app, + "POST", + "/memories", + [b'{"content": "Learning FastAPI testing"}'], + {"Content-Type": "application/json"}, + ) + + assert status_code == 422 + assert any( + error["loc"] == ["body", "tags"] and error["type"] == "missing" + for error in json.loads(body)["detail"] + ) + assert read_database(data_file) == [] + + def test_post_memory_rate_limit_uses_client_id_identity(monkeypatch, data_file: Path): monkeypatch.setenv("MEMORIES_RATE_LIMIT_WRITES_PER_MINUTE", "1") test_client = TestClient(app_main.create_app()) @@ -298,35 +425,57 @@ def test_rate_limiting_can_be_disabled(monkeypatch, data_file: Path): def test_body_size_limit_runs_before_rate_limit(monkeypatch, data_file: Path): monkeypatch.setenv("MEMORIES_REQUEST_BODY_MAX_BYTES", "10") monkeypatch.setenv("MEMORIES_RATE_LIMIT_WRITES_PER_MINUTE", "1") - test_client = TestClient(app_main.create_app()) + test_app = app_main.create_app() - oversized_response = test_client.post( + oversized_status_code, _oversized_headers, oversized_body = asgi_request( + test_app, + "POST", "/memories", - content="x" * 11, - headers={"Content-Type": "application/json"}, + [b"x" * 11], + {"Content-Type": "application/json", "Content-Length": "11"}, + ) + first_status_code, _first_headers, _first_body = asgi_request( + test_app, + "POST", + "/memories", + [b"{}"], + {"Content-Type": "application/json", "Content-Length": "2"}, + ) + second_status_code, second_headers, second_body = asgi_request( + test_app, + "POST", + "/memories", + [b"{}"], + {"Content-Type": "application/json", "Content-Length": "2"}, ) - first_counted_response = test_client.post("/memories", json={}) - second_counted_response = test_client.post("/memories", json={}) - assert oversized_response.status_code == 413 - assert first_counted_response.status_code == 422 - assert_rate_limited(second_counted_response, 1) + assert oversized_status_code == 413 + assert json.loads(oversized_body) == {"detail": "Request body too large"} + assert first_status_code == 422 + assert second_status_code == 429 + assert json.loads(second_body) == {"detail": "Rate limit exceeded"} + assert second_headers["retry-after"].isdigit() + assert second_headers["x-ratelimit-limit"] == "1" + assert second_headers["x-ratelimit-remaining"] == "0" + assert second_headers["x-ratelimit-reset"].isdigit() assert read_database(data_file) == [] def test_patch_memory_rejects_oversized_body_before_json_parsing(monkeypatch, data_file: Path): monkeypatch.setenv("MEMORIES_REQUEST_BODY_MAX_BYTES", "10") - test_client = TestClient(app_main.create_app()) + test_app = app_main.create_app() - response = test_client.patch( + status_code, headers, body = asgi_request( + test_app, + "PATCH", "/memories/1", - content="x" * 11, - headers={"Content-Type": "application/json"}, + [b"x" * 11], + {"Content-Type": "application/json", "Content-Length": "11"}, ) - assert response.status_code == 413 - assert response.json() == {"detail": "Request body too large"} - assert response.headers["X-Request-Body-Limit"] == "10" + assert status_code == 413 + assert json.loads(body) == {"detail": "Request body too large"} + assert headers["x-request-body-limit"] == "10" assert read_database(data_file) == [] diff --git a/tests/contract/test_mcp_http_transport.py b/tests/contract/test_mcp_http_transport.py index 00e8faa..8edcd44 100644 --- a/tests/contract/test_mcp_http_transport.py +++ b/tests/contract/test_mcp_http_transport.py @@ -1,3 +1,4 @@ +import asyncio import importlib import json import logging @@ -31,6 +32,74 @@ def request_log_payloads(caplog): return payloads +def asgi_post( + application, + path: str, + chunks: list[bytes], + headers: dict[str, str] | None = None, +): + return asyncio.run(_asgi_post(application, path, chunks, headers or {})) + + +async def _asgi_post( + application, + path: str, + chunks: list[bytes], + headers: dict[str, str], +): + response_messages = [] + request_messages = [ + { + "type": "http.request", + "body": chunk, + "more_body": index < len(chunks) - 1, + } + for index, chunk in enumerate(chunks) + ] + request_index = 0 + + async def receive(): + nonlocal request_index + if request_index < len(request_messages): + message = request_messages[request_index] + request_index += 1 + return message + return {"type": "http.disconnect"} + + async def send(message): + response_messages.append(message) + + await application( + { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": path, + "raw_path": path.encode(), + "query_string": b"", + "headers": [(name.lower().encode(), value.encode()) for name, value in headers.items()], + "client": ("testclient", 50000), + "server": ("testserver", 80), + "root_path": "", + }, + receive, + send, + ) + + response_start = next( + message for message in response_messages if message["type"] == "http.response.start" + ) + response_body = b"".join( + message.get("body", b"") + for message in response_messages + if message["type"] == "http.response.body" + ) + response_headers = {name.decode(): value.decode() for name, value in response_start["headers"]} + return response_start["status"], response_headers, response_body + + def build_client_with_fresh_mcp(): importlib.reload(mcp_server_module) reloaded_main = importlib.reload(main_module) @@ -104,18 +173,19 @@ def test_mcp_http_allows_configured_browser_origin(monkeypatch, tmp_path: Path): def test_mcp_http_rejects_oversized_body_before_transport_parsing(monkeypatch): monkeypatch.setenv("MEMORIES_REQUEST_BODY_MAX_BYTES", "10") - client = TestClient(main_module.create_app()) + test_app = main_module.create_app() - response = client.post( + status_code, headers, body = asgi_post( + test_app, "/mcp", - content="x" * 11, - headers={"Content-Type": "application/json"}, + [b"x" * 11], + {"Content-Type": "application/json", "Content-Length": "11"}, ) - assert response.status_code == 413 - assert response.json() == {"detail": "Request body too large"} - assert response.headers["X-Request-Body-Limit"] == "10" - assert_uuid(response.headers["X-Request-Id"]) + assert status_code == 413 + assert json.loads(body) == {"detail": "Request body too large"} + assert headers["x-request-body-limit"] == "10" + assert_uuid(headers["x-request-id"]) def test_mcp_http_rate_limit_returns_stable_429(monkeypatch):