Fix token quota rollback for streaming

middleware/quota.py

@@ -268,6 +268,7 @@ async def _gen(self) -> AsyncIterator[bytes]:
                         if chunk_tokens:
                             total, _ = await self.store.adjust(self.user, chunk_tokens)
                             if total > self.limit:
+                                await self.store.adjust(self.user, -chunk_tokens)
                                 self.quota_exceeded = True
                                 logger.warning(
                                     "User %s quota breached mid-stream", self.user
@@ -354,10 +355,6 @@ async def finalize() -> None:
             if getattr(streamer, "quota_exceeded", False):
                 usage["detail"] = "token quota exceeded mid-stream"
                 retry_after = max(0, int(self.window - (time.time() - oldest)))
-                response.status_code = 429
-                response.body_iterator = async_iter([])
-                response.headers["Retry-After"] = str(retry_after)
-                response.headers["content-type"] = "application/json"
                 usage["ts"] = time.time()
                 await request.app.state.usage.record(**usage)
                 return

tests/test_token_quota_middleware.py

@@ -6,7 +6,9 @@
 
 pytest.importorskip("tiktoken")
 
-from middleware.quota import TokenQuotaMiddleware
+from starlette.responses import StreamingResponse
+
+from middleware.quota import InMemoryMeterStore, TokenQuotaMiddleware
 
 
 @pytest.mark.asyncio
@@ -44,3 +46,30 @@ async def echo(request: Request):
         resp = await client.post("/echo", json={"msg": "hello"}, headers=headers)
     assert resp.status_code == 429
     assert "retry_after" in resp.json()
+
+
+@pytest.mark.asyncio
+async def test_midstream_over_limit_rolls_back(monkeypatch):
+    monkeypatch.setenv("MAX_TOKENS_PER_MIN", "5")
+    store = InMemoryMeterStore()
+    app = FastAPI()
+    app.add_middleware(TokenQuotaMiddleware, store=store)
+
+    @app.get("/stream")
+    async def stream():
+        async def gen():
+            yield b"hi"
+            yield b"aaaaaaa"
+
+        return StreamingResponse(gen(), media_type="text/plain")
+
+    headers = {"X-Attach-User": "carol"}
+    transport = ASGITransport(app=app)
+    async with AsyncClient(transport=transport, base_url="http://test") as client:
+        resp = await client.get("/stream", headers=headers)
+        text = resp.text
+
+    assert resp.status_code == 200
+    assert text == "hi"
+    total = await store.peek_total("carol")
+    assert total == 2