diff --git a/src/litserve/server.py b/src/litserve/server.py index 9a6c253d..f770be3e 100644 --- a/src/litserve/server.py +++ b/src/litserve/server.py @@ -32,7 +32,7 @@ from collections.abc import Callable, Iterable, Mapping, Sequence from contextlib import asynccontextmanager from queue import Queue -from typing import TYPE_CHECKING, Literal, Optional, Union +from typing import TYPE_CHECKING, Any, Literal, Optional, Union import uvicorn import uvicorn.server @@ -294,13 +294,15 @@ def __init__(self, lit_api: LitAPI, server: "LitServer"): self.lit_api = lit_api self.server = server - async def _prepare_request(self, request, request_type) -> dict: + async def _prepare_request(self, request, request_type) -> Any: """Common request preparation logic.""" if request_type == Request: - content_type = request.headers.get("Content-Type", "") + content_type = request.headers.get("Content-Type", "").split(";", 1)[0].strip().lower() if content_type == "application/x-www-form-urlencoded" or content_type.startswith("multipart/form-data"): return await request.form() - return await request.json() + if content_type == "application/json" or content_type.endswith("+json"): + return await request.json() + return await request.body() return request async def _submit_request(self, payload: dict) -> tuple[str, asyncio.Event]: diff --git a/tests/unit/test_request_handlers.py b/tests/unit/test_request_handlers.py index c879d7e7..995ce5c6 100644 --- a/tests/unit/test_request_handlers.py +++ b/tests/unit/test_request_handlers.py @@ -47,19 +47,29 @@ def _get_request_queue(self, api_path): class MockRequest: """Mock FastAPI Request object for testing.""" - def __init__(self, json_data=None, form_data=None, content_type="application/json"): + def __init__(self, json_data=None, form_data=None, body=b"", content_type="application/json"): self._json_data = json_data or {} self._form_data = form_data or {} + self._body = body self.headers = {"Content-Type": content_type} + self.json_called = False + self.form_called = False + self.body_called = False async def json(self): + self.json_called = True if self._json_data is None: raise json.JSONDecodeError("Invalid JSON", "", 0) return self._json_data async def form(self): + self.form_called = True return self._form_data + async def body(self): + self.body_called = True + return self._body + class TestRequestHandler(BaseRequestHandler): def __init__(self, lit_api, server): @@ -81,6 +91,44 @@ async def test_request_handler(mock_lit_api): assert response_queue_id == 0 +@pytest.mark.asyncio +async def test_request_handler_preserves_raw_body(mock_lit_api): + mock_server = MockServer(mock_lit_api) + handler = TestRequestHandler(mock_lit_api, mock_server) + body = b"\x00protobuf-payload\xff" + mock_request = MockRequest(body=body, content_type="application/x-recordio-protobuf") + + await handler.handle_request(mock_request, Request) + + assert mock_server.request_queue.get()[3] == body + assert mock_request.body_called + assert not mock_request.json_called + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "content_type", ["application/json", "application/vnd.api+json", "application/json; charset=utf-8"] +) +async def test_request_handler_preserves_json_decoding(mock_lit_api, content_type): + handler = TestRequestHandler(mock_lit_api, MockServer(mock_lit_api)) + request = MockRequest(json_data={"input": 1}, content_type=content_type) + + assert await handler._prepare_request(request, Request) == {"input": 1} + assert request.json_called + assert not request.body_called + + +@pytest.mark.asyncio +@pytest.mark.parametrize("content_type", ["application/x-www-form-urlencoded", "multipart/form-data; boundary=test"]) +async def test_request_handler_preserves_form_decoding(mock_lit_api, content_type): + handler = TestRequestHandler(mock_lit_api, MockServer(mock_lit_api)) + request = MockRequest(form_data={"input": "1"}, content_type=content_type) + + assert await handler._prepare_request(request, Request) == {"input": "1"} + assert request.form_called + assert not request.body_called + + @pytest.mark.asyncio @patch("litserve.server.asyncio.Event") async def test_request_handler_streaming(mock_event, mock_lit_api):