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

View file

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

View file

@ -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",