diff --git a/recognization/code/app.py b/recognization/code/app.py index 5c16dbd..bcbd1ff 100644 --- a/recognization/code/app.py +++ b/recognization/code/app.py @@ -270,7 +270,8 @@ def create_app( async def lifespan(_app: FastAPI): yield for session in list(sessions.values()): - session.close() + if session is not None: + session.close() engine.shutdown() app = FastAPI( @@ -1111,6 +1112,11 @@ def main() -> None: LOGGER.info("Loading model and catalog (device=%s, imgsz=%s)...", engine_settings.device or "auto", engine_settings.imgsz) engine = build_default_engine(engine_settings) + if (engine_settings.identity_backend or "").lower() == "gallery": + if engine.identity_gallery() is None: + raise RuntimeError( + "gallery identity backend requires local DINOv2 weights; refusing to start" + ) warm_ms = engine.warm_up() LOGGER.info( "Ready: %d SKUs, catalog %s, warmup %.0f ms, classify workers %d", diff --git a/recognization/code/test_api.py b/recognization/code/test_api.py index bc21027..f7cf60a 100644 --- a/recognization/code/test_api.py +++ b/recognization/code/test_api.py @@ -196,6 +196,31 @@ class RestContractTests(unittest.TestCase): self.assertEqual(sent[0]["status"], 413) self.assertIn(b"PAYLOAD_TOO_LARGE", sent[1]["body"]) + def test_overflow_after_response_start_sends_no_second_start(self): + """A body that overflows after headers began must not start again.""" + import anyio + + from transport import BodyLimit + + async def app(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await receive() + await send({"type": "http.response.body", "body": b"ok"}) + + sent = [] + + async def receive(): + return {"type": "http.request", "body": b"x" * 40, "more_body": False} + + async def send(message): + sent.append(message) + + anyio.run(BodyLimit(app, 8), {"type": "http"}, receive, send) + starts = [message for message in sent if message.get("type") == "http.response.start"] + self.assertEqual(len(starts), 1) + self.assertEqual(starts[0]["status"], 200) + self.assertFalse(any(message.get("status") == 413 for message in sent)) + def test_classify_and_inspect_roundtrip(self): recognized = self.client.post( "/v1/recognize?classify=false", content=jpeg_bytes(), @@ -360,6 +385,16 @@ class WebSocketProtocolTests(unittest.TestCase): finally: engine.shutdown() + def test_lifespan_shutdown_skips_a_reserved_slot(self): + """A reservation is None until setup finishes. Shutdown must not call it.""" + client, engine = make_app() + client.app.state.sessions["reserved"] = None + try: + with client: + pass + finally: + engine.shutdown() + class LegacyRouteTests(unittest.TestCase): def test_legacy_sessions_still_respond(self): diff --git a/recognization/code/transport.py b/recognization/code/transport.py index 66e94a1..0322306 100644 --- a/recognization/code/transport.py +++ b/recognization/code/transport.py @@ -104,7 +104,10 @@ class BodyLimit: A handler that checks Content-Length, or reads request.body(), runs after the server has already buffered the stream. This counts bytes at receive - and returns 413 without handing the rest to the app. WebSocket frames are + and returns 413 without handing the rest to the app. Content-Length is + not trusted: the stream is what counts, and reading stops at the cap. + If the app already started the response, the connection is aborted + instead of sending a second http.response.start. WebSocket frames are not HTTP bodies and are left alone. """ @@ -116,10 +119,13 @@ class BodyLimit: if scope.get("type") != "http": await self.app(scope, receive, send) return + class TooLarge(Exception): pass seen = 0 + started = False + async def limited() -> Dict[str, Any]: nonlocal seen message = await receive() @@ -128,12 +134,24 @@ class BodyLimit: body = message.get("body") or b"" seen += len(body) if seen > self.max_bytes: + # The app must not see this chunk or any later one. raise TooLarge return message + async def tracked(message: Dict[str, Any]) -> None: + nonlocal started + if message.get("type") == "http.response.start": + started = True + await send(message) + try: - await self.app(scope, limited, send) + await self.app(scope, limited, tracked) except TooLarge: + if started: + # The status line is already on the wire. Returning leaves the + # response unfinished so the server closes the connection + # instead of accepting a second http.response.start. + return await send( { "type": "http.response.start",