fix: guard reserved sessions and bounded response state

This commit is contained in:
kiyo9w 2026-09-28 13:41:35 +09:00
parent 0e85c399bb
commit 745a29d979
3 changed files with 62 additions and 3 deletions

View file

@ -270,7 +270,8 @@ def create_app(
async def lifespan(_app: FastAPI): async def lifespan(_app: FastAPI):
yield yield
for session in list(sessions.values()): for session in list(sessions.values()):
session.close() if session is not None:
session.close()
engine.shutdown() engine.shutdown()
app = FastAPI( app = FastAPI(
@ -1111,6 +1112,11 @@ def main() -> None:
LOGGER.info("Loading model and catalog (device=%s, imgsz=%s)...", LOGGER.info("Loading model and catalog (device=%s, imgsz=%s)...",
engine_settings.device or "auto", engine_settings.imgsz) engine_settings.device or "auto", engine_settings.imgsz)
engine = build_default_engine(engine_settings) 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() warm_ms = engine.warm_up()
LOGGER.info( LOGGER.info(
"Ready: %d SKUs, catalog %s, warmup %.0f ms, classify workers %d", "Ready: %d SKUs, catalog %s, warmup %.0f ms, classify workers %d",

View file

@ -196,6 +196,31 @@ class RestContractTests(unittest.TestCase):
self.assertEqual(sent[0]["status"], 413) self.assertEqual(sent[0]["status"], 413)
self.assertIn(b"PAYLOAD_TOO_LARGE", sent[1]["body"]) 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): def test_classify_and_inspect_roundtrip(self):
recognized = self.client.post( recognized = self.client.post(
"/v1/recognize?classify=false", content=jpeg_bytes(), "/v1/recognize?classify=false", content=jpeg_bytes(),
@ -360,6 +385,16 @@ class WebSocketProtocolTests(unittest.TestCase):
finally: finally:
engine.shutdown() 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): class LegacyRouteTests(unittest.TestCase):
def test_legacy_sessions_still_respond(self): def test_legacy_sessions_still_respond(self):

View file

@ -104,7 +104,10 @@ class BodyLimit:
A handler that checks Content-Length, or reads request.body(), runs after A handler that checks Content-Length, or reads request.body(), runs after
the server has already buffered the stream. This counts bytes at receive 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. not HTTP bodies and are left alone.
""" """
@ -116,10 +119,13 @@ class BodyLimit:
if scope.get("type") != "http": if scope.get("type") != "http":
await self.app(scope, receive, send) await self.app(scope, receive, send)
return return
class TooLarge(Exception): class TooLarge(Exception):
pass pass
seen = 0 seen = 0
started = False
async def limited() -> Dict[str, Any]: async def limited() -> Dict[str, Any]:
nonlocal seen nonlocal seen
message = await receive() message = await receive()
@ -128,12 +134,24 @@ class BodyLimit:
body = message.get("body") or b"" body = message.get("body") or b""
seen += len(body) seen += len(body)
if seen > self.max_bytes: if seen > self.max_bytes:
# The app must not see this chunk or any later one.
raise TooLarge raise TooLarge
return message 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: try:
await self.app(scope, limited, send) await self.app(scope, limited, tracked)
except TooLarge: 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( await send(
{ {
"type": "http.response.start", "type": "http.response.start",