fix: preserve raw request bodies for custom content types

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]:

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):