fix streaming: async generator keeps httpx client alive internally

Dockerfile

@@ -12,4 +12,4 @@ RUN pip install --no-cache-dir --upgrade -r requirements.txt
 COPY --chown=user . .
 
 EXPOSE 7860
-CMD uvicorn main:app --host 0.0.0.0 --port 7860
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

proxy/client.py

@@ -52,35 +52,29 @@ async def _poll_nvidia(client: httpx.AsyncClient, provider: Provider, request_id
     raise HTTPException(status_code=504, detail=f"NVIDIA async job {request_id} timed out after polling")
 
 
-async def _stream_chunks(response: httpx.Response) -> AsyncIterator[bytes]:
-    async for chunk in response.aiter_raw():
-        yield chunk
+async def _stream_response(provider: Provider, real_key: str, body: dict[str, Any]) -> AsyncIterator[bytes]:
+    """Async generator that manages its own httpx client for streaming."""
+    target_url = f"{provider.base_url}/v1/chat/completions"
+    headers = _build_headers(provider, real_key, {})
+    async with httpx.AsyncClient(timeout=_TIMEOUT, follow_redirects=True) as client:
+        async with client.stream("POST", target_url, headers=headers, json=body) as upstream:
+            if upstream.status_code >= 400:
+                content = await upstream.aread()
+                raise HTTPException(status_code=upstream.status_code, detail=content.decode("utf-8", errors="replace"))
+            async for chunk in upstream.aiter_raw():
+                yield chunk
 
 
 async def forward_request(provider: Provider, real_key: str, request: Request, body: dict[str, Any]) -> Any:
     target_url = f"{provider.base_url}/v1/chat/completions"
     headers = _build_headers(provider, real_key, dict(request.headers))
-    stream = bool(body.get("stream"))
+
+    if bool(body.get("stream")):
+        return StreamingResponse(_stream_response(provider, real_key, body), media_type="text/event-stream")
 
     async with httpx.AsyncClient(timeout=_TIMEOUT, follow_redirects=True) as client:
         try:
-            if stream:
-                req = client.build_request("POST", target_url, headers=headers, json=body)
-                upstream = await client.send(req, stream=True)
-
-                if upstream.status_code >= 400:
-                    content = await upstream.aread()
-                    logger.warning("Upstream error from %s: %d %s", provider.name, upstream.status_code, content[:500])
-                    raise HTTPException(status_code=upstream.status_code, detail=content.decode("utf-8", errors="replace"))
-
-                return StreamingResponse(
-                    _stream_chunks(upstream),
-                    status_code=upstream.status_code,
-                    media_type=upstream.headers.get("content-type", "text/event-stream"),
-                    headers={k: v for k, v in upstream.headers.items() if k.lower() in ("cache-control", "x-request-id")},
-                )
-            else:
-                upstream = await client.post(target_url, headers=headers, json=body)
+            upstream = await client.post(target_url, headers=headers, json=body)
         except httpx.TimeoutException as exc:
             logger.warning("Upstream timeout for %s: %s", provider.name, exc)
             raise HTTPException(status_code=504, detail="Upstream provider timed out")