Upload main.py

main.py

@@ -67,10 +67,10 @@ app.add_middleware(
 )
 
 # to ensure we don’t leak hop-by-hop headers
-def _filter_resp_headers(h):
+def _filter_resp_headers(headers: dict) -> dict:
     # pass through useful headers but strip hop-by-hop
     allowed = {"content-disposition"}
-    return {k: v for k, v in h.items() if k.lower() in allowed}
+    return {k: v for k, v in headers.items() if k.lower() in allowed}
 
 def _extract_user_token(req: Request) -> str | None:
     """
@@ -120,23 +120,44 @@ async def _forward_stream(path: str, files=None, data=None, user_token: str | No
         headers["X-API-Key"] = user_token
 
     timeout = httpx.Timeout(600.0)
-    async with httpx.AsyncClient(timeout=timeout) as client:
-        async with client.stream(method, url, headers=headers, files=files, data=data) as resp:
-            if resp.status_code >= 400:
-                text = await resp.aread()
-                return Response(
-                    content=text,
-                    status_code=resp.status_code,
-                    media_type=resp.headers.get("content-type","text/plain")
-                )
-            media_type = resp.headers.get("content-type", "application/octet-stream")
-            safe_headers = _filter_resp_headers(resp.headers)  # Content-Disposition only
-            return StreamingResponse(
-                resp.aiter_raw(),
-                status_code=resp.status_code,
-                media_type=media_type,
-                headers=safe_headers,
-            )
+
+    client = httpx.AsyncClient(timeout=timeout, follow_redirects=True)
+
+    # Don't open the context yet; the iterator must own the context lifetime.
+    stream_ctx = client.stream(method, url, headers=headers, files=files, data=data)
+
+    # Mutable holders we can fill once the stream opens
+    status_holder = {"code": 200}
+    media_type_holder = {"ct": "application/octet-stream"}
+    headers_holder = {}
+
+    async def body_iter():
+        try:
+            async with stream_ctx as resp:
+                status_holder["code"] = resp.status_code
+                media_type_holder["ct"] = resp.headers.get("content-type", "application/octet-stream")
+                headers_holder.update(_filter_resp_headers(resp.headers))
+
+                # If upstream already returned an error, buffer it (short text/json)
+                if resp.status_code >= 400:
+                    # Buffer entire payload and yield once
+                    chunk = await resp.aread()
+                    yield chunk
+                    return
+
+                async for chunk in resp.aiter_raw():
+                    yield chunk
+        finally:
+            await client.aclose()
+
+    response = StreamingResponse(
+        body_iter(),
+        status_code=status_holder["code"],
+        media_type=media_type_holder["ct"],
+        headers=headers_holder,
+    )
+    return response
+
 
 # for demo
 async def _forward_demo_stream(path: str, *, files: dict | None, data: dict | None, method: str = "POST"):
@@ -149,26 +170,41 @@ async def _forward_demo_stream(path: str, *, files: dict | None, data: dict | No
         "Authorization": f"Bearer {HF_TOKEN}",
         # App-level demo token (checked by the private API):
         "X-DATFID-Token": DATFID_DEMO_TOKEN,
+        "Accept": "*/*",
     }
 
     timeout = httpx.Timeout(120.0)
-    async with httpx.AsyncClient(timeout=timeout) as client:
-        async with client.stream(method, url, headers=headers, files=files, data=data) as resp:
-            if resp.status_code >= 400:
-                text = await resp.aread()
-                return Response(
-                    content=text,
-                    status_code=resp.status_code,
-                    media_type=resp.headers.get("content-type","text/plain")
-                )
-            media_type = resp.headers.get("content-type", "application/octet-stream")
-            safe_headers = _filter_resp_headers(resp.headers)  # Content-Disposition only
-            return StreamingResponse(
-                resp.aiter_raw(),
-                status_code=resp.status_code,
-                media_type=media_type,
-                headers=safe_headers,
-            )
+    client = httpx.AsyncClient(timeout=timeout, follow_redirects=True)
+    stream_ctx = client.stream(method, url, headers=headers, files=files, data=data)
+
+    status_holder = {"code": 200}
+    media_type_holder = {"ct": "application/octet-stream"}
+    headers_holder = {}
+
+    async def body_iter():
+        try:
+            async with stream_ctx as resp:
+                status_holder["code"] = resp.status_code
+                media_type_holder["ct"] = resp.headers.get("content-type", "application/octet-stream")
+                headers_holder.update(_filter_resp_headers(resp.headers))
+
+                if resp.status_code >= 400:
+                    chunk = await resp.aread()
+                    yield chunk
+                    return
+
+                async for chunk in resp.aiter_raw():
+                    yield chunk
+        finally:
+            await client.aclose()
+
+    response = StreamingResponse(
+        body_iter(),
+        status_code=status_holder["code"],
+        media_type=media_type_holder["ct"],
+        headers=headers_holder,
+    )
+    return response
 
 @app.get("/")
 async def root(req: Request):