Upload 34 files

DATASET_COLLECTION_GUIDANCE.md

@@ -1845,3 +1845,8 @@ This closes cross-replica **coordination authority** only. The proxy does not in
 Redis AOF/RDB/managed-service persistence guarantees, backup retention, or disaster
 recovery policy. Keep `CONTRIBUTION_REQUIRE_DURABLE` as a separate deployment
 requirement and verify the external Redis durability contract independently.
+
+
+## Run 25 / B44 provider response boundaries
+
+Storage-provider control traffic is a small metadata plane, not a bulk-download plane. Its response ceiling defaults to **4 MiB** and is hard-clamped to **16 MiB**. GitHub/GitLab/Bitbucket mutations that need no response body do not consume one; metadata reads are stream-bounded before JSON parsing. Hugging Face Hub control calls run with a scoped bounded lower-level HTTP client and restore the prior SDK client factory immediately afterward. Intentional large offline dataset/model downloads remain separately scoped (`SEC-P2-49`).

README.md

@@ -396,6 +396,7 @@ All values are in seconds.  Non-integer values silently fall back to the default
 | `SHARE_REQUIRE_DURABLE` | `false` | When true, Global Share writes fail closed unless the active store reports durable lifecycle storage. |
 | `SHARE_REQUIRE_SHARED` | `false` | When true, Global Share writes fail closed unless the active store is shared and authoritative. Use for horizontally scaled replicas. |
 | `MAX_BODY_BYTES` | `10485760` | Maximum accepted chat request body size. Enforced while streaming and hard-clamped to 16 MiB. |
+| `MAX_UPSTREAM_RESPONSE_BYTES` | `8388608` | Maximum decoded upstream response body accepted by the proxy. Enforced while streaming before whole-body buffering and hard-clamped to 32 MiB. Oversize or malformed declared lengths fail closed. |
 | `CHAT_RATE_LIMIT_PER_HOUR` | `30` | Chat requests per resolved client identity. Enforced by the selected local or Redis backend. |
 | `SHARE_RATE_LIMIT_PER_HOUR` | `10` | Global Share creates/updates per resolved client identity. Enforced by the selected backend. |
 | `FEEDBACK_RATE_LIMIT_PER_HOUR` | `30` | Feedback writes, including retractions, per resolved client identity. |
@@ -870,3 +871,8 @@ the custom ZeroGPU Space that actually has the weights loaded.  This is why
 - [HuggingFace Inference API](https://huggingface.co/docs/api-inference/)
 - [HuggingFace Datasets — `huggingface_hub`](https://huggingface.co/docs/huggingface_hub/guides/repository)
 - [ZeroGPU documentation](https://huggingface.co/docs/hub/spaces-zerogpu)
+
+
+### Provider control-response boundary (Run 25 / B44)
+
+Record-storage provider control responses are bounded independently from intentional model/dataset downloads. The default ceiling is **4 MiB** and `AI_RECORD_STORAGE_CONTROL_RESPONSE_MAX_BYTES` is clamped to **16 MiB**. Provider mutations that do not require response content are streamed and closed without buffering a body; metadata JSON is byte-counted before parsing. Custom `api_base` is supported only for GitLab and must be HTTPS with a valid host and no userinfo, query, fragment, control characters, or traversal.

_utils/_share_contract.py

@@ -363,7 +363,7 @@ def render_share_viewer_shell(read_path: str = "/v1/share/read") -> str:
     conversation values with DOM ``textContent`` and never injects untrusted HTML.
     """
     read_path_json = json.dumps(str(read_path), ensure_ascii=False)
-    template = r'''<!doctype html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1"><meta name="referrer" content="no-referrer"><title>Shared AI conversation</title><style>:root{font-family:system-ui,sans-serif;color-scheme:light dark}body{margin:0;background:Canvas;color:CanvasText}.wrap{max-width:850px;margin:auto;padding:24px}.head{border-bottom:1px solid currentColor;padding-bottom:16px}.source{overflow-wrap:anywhere}.source a{color:inherit}.msg{margin:18px 0;padding:14px;border:1px solid currentColor;border-radius:12px}.msg.user{margin-left:10%}.msg.error{border-style:dashed}.role{font-weight:700;margin-bottom:8px}pre{white-space:pre-wrap;overflow-wrap:anywhere;font:inherit;margin:0}.error-note{border:1px dashed currentColor;padding:14px;border-radius:12px}</style></head><body><main id="app" class="wrap"><p>Loading shared conversation…</p></main><script>(()=>{'use strict';const app=document.getElementById('app');const fail=(m)=>{app.replaceChildren();const p=document.createElement('p');p.className='error-note';p.textContent=m;app.appendChild(p);};let raw=(location.hash||'').slice(1);if(raw.startsWith('share='))raw=raw.slice(6);try{raw=decodeURIComponent(raw)}catch(_e){}if(!/^(?:[0-9a-f]{32}|[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})$/i.test(raw)){fail('This Share link is invalid or incomplete.');return;}fetch(__READ_PATH__,{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({shareId:raw}),cache:'no-store',credentials:'omit',redirect:'error',referrerPolicy:'no-referrer'}).then(async r=>{let data=null;try{data=await r.json()}catch(_e){}if(!r.ok){throw new Error(r.status===410?'This Share has expired.':r.status===404?'This Share is unavailable.':'Could not load this Share.');}return data||{};}).then(data=>{app.replaceChildren();if(data.format==='html'&&data.snapshot&&data.snapshot.session&&Array.isArray(data.snapshot.records)){const snap=data.snapshot;const h=document.createElement('header');h.className='head';const h1=document.createElement('h1');h1.textContent=(snap.session.assistant_name||'AI Assistant')+' — Shared conversation';h.appendChild(h1);if(snap.session.page_title){const p=document.createElement('p');p.textContent=snap.session.page_title;h.appendChild(p);}if(snap.session.page_url){const p=document.createElement('p');p.className='source';p.append('Source: ');const a=document.createElement('a');a.href=snap.session.page_url;a.rel='noopener noreferrer';a.referrerPolicy='no-referrer';a.textContent=snap.session.page_url;p.appendChild(a);h.appendChild(p);}app.appendChild(h);for(const row of snap.records){const article=document.createElement('article');article.className='msg '+(row.role==='user'?'user':row.role==='error'?'error':'assistant');const role=document.createElement('div');role.className='role';role.textContent=row.role==='user'?'You':row.role==='error'?'Error':(snap.session.assistant_name||'AI Assistant');const pre=document.createElement('pre');pre.textContent=String(row.text||'');article.append(role,pre);app.appendChild(article);}return;}const pre=document.createElement('pre');pre.textContent=String(data.content||'');app.appendChild(pre);}).catch(e=>fail(e&&e.message?e.message:'Could not load this Share.'));})();</script></body></html>'''
+    template = r'''<!doctype html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1"><meta name="referrer" content="no-referrer"><title>Shared AI conversation</title><style>:root{font-family:system-ui,sans-serif;color-scheme:light dark}body{margin:0;background:Canvas;color:CanvasText}.wrap{max-width:850px;margin:auto;padding:24px}.head{border-bottom:1px solid currentColor;padding-bottom:16px}.source{overflow-wrap:anywhere}.source a{color:inherit}.msg{margin:18px 0;padding:14px;border:1px solid currentColor;border-radius:12px}.msg.user{margin-left:10%}.msg.error{border-style:dashed}.role{font-weight:700;margin-bottom:8px}pre{white-space:pre-wrap;overflow-wrap:anywhere;font:inherit;margin:0}.error-note{border:1px dashed currentColor;padding:14px;border-radius:12px}</style></head><body><main id="app" class="wrap"><p>Loading shared conversation…</p></main><script>(()=>{'use strict';const app=document.getElementById('app');const fail=(m)=>{app.replaceChildren();const p=document.createElement('p');p.className='error-note';p.textContent=m;app.appendChild(p);};let raw=(location.hash||'').slice(1);if(raw.startsWith('share='))raw=raw.slice(6);try{raw=decodeURIComponent(raw)}catch(_e){}if(!/^(?:[0-9a-f]{32}|[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})$/i.test(raw)){fail('This Share link is invalid or incomplete.');return;}const readJson=async(r)=>{const max=4*1024*1024;const h=r.headers&&r.headers.get?r.headers.get('content-length'):null;if(h!=null&&String(h).trim()!==''){if(!/^\d+$/.test(String(h).trim())||Number(h)>max)throw new Error('Share response is too large.');}if(!r.body||typeof r.body.getReader!=='function'||typeof TextDecoder!=='function')throw new Error('Bounded Share reader unavailable.');const rd=r.body.getReader(),dec=new TextDecoder(),parts=[];let n=0;try{for(;;){const x=await rd.read();if(x.done)break;const v=x.value||new Uint8Array(0);n+=Number(v.byteLength||v.length||0);if(n>max)throw new Error('Share response is too large.');parts.push(dec.decode(v,{stream:true}));}parts.push(dec.decode());}catch(e){try{await rd.cancel()}catch(_e){}throw e;}finally{try{rd.releaseLock()}catch(_e){}}return JSON.parse(parts.join(''));};fetch(__READ_PATH__,{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({shareId:raw}),cache:'no-store',credentials:'omit',redirect:'error',referrerPolicy:'no-referrer'}).then(async r=>{if(!r.ok){throw new Error(r.status===410?'This Share has expired.':r.status===404?'This Share is unavailable.':'Could not load this Share.');}return await readJson(r);}).then(data=>{app.replaceChildren();if(data.format==='html'&&data.snapshot&&data.snapshot.session&&Array.isArray(data.snapshot.records)){const snap=data.snapshot;const h=document.createElement('header');h.className='head';const h1=document.createElement('h1');h1.textContent=(snap.session.assistant_name||'AI Assistant')+' — Shared conversation';h.appendChild(h1);if(snap.session.page_title){const p=document.createElement('p');p.textContent=snap.session.page_title;h.appendChild(p);}if(snap.session.page_url){const p=document.createElement('p');p.className='source';p.append('Source: ');const a=document.createElement('a');a.href=snap.session.page_url;a.rel='noopener noreferrer';a.referrerPolicy='no-referrer';a.textContent=snap.session.page_url;p.appendChild(a);h.appendChild(p);}app.appendChild(h);for(const row of snap.records){const article=document.createElement('article');article.className='msg '+(row.role==='user'?'user':row.role==='error'?'error':'assistant');const role=document.createElement('div');role.className='role';role.textContent=row.role==='user'?'You':row.role==='error'?'Error':(snap.session.assistant_name||'AI Assistant');const pre=document.createElement('pre');pre.textContent=String(row.text||'');article.append(role,pre);app.appendChild(article);}return;}const pre=document.createElement('pre');pre.textContent=String(data.content||'');app.appendChild(pre);}).catch(e=>fail(e&&e.message?e.message:'Could not load this Share.'));})();</script></body></html>'''
     return template.replace("__READ_PATH__", read_path_json)
 
 def generate_edit_token() -> str:

_utils/_shared_logic.py

@@ -220,7 +220,7 @@ __all__ = [  # noqa: RUF022
 # ─────────────────────────────────────────────────────────────────────────────
 
 #: Proxy release version — bump on every breaking change.
-PROXY_VERSION: str = "7.0.0"
+PROXY_VERSION: str = "7.2.0"
 
 #: HuggingFace Inference Providers router base URL (no trailing slash).
 #: Only used for Path 3 (standard provider models) when ``BACKEND_URL`` is

_utils/_storage.py

@@ -15,10 +15,11 @@ import json
 import os
 import re
 import time
+import threading
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from typing import Any, Literal
-from urllib.parse import quote, urlencode
+from urllib.parse import quote, urlencode, urlsplit
 
 import httpx
 
@@ -34,6 +35,113 @@ _ID_RE = re.compile(r"^[a-z][a-z0-9_-]{0,47}$")
 _SEG_RE = re.compile(r"^[A-Za-z0-9._-]+$")
 _ALLOWED_PROVIDERS = {"huggingface", "github", "gitlab", "bitbucket"}
 _TRANSIENT_STATUS = {408, 409, 425, 429, 500, 502, 503, 504}
+_CONTROL_RESPONSE_DEFAULT = 4 * 1024 * 1024
+_CONTROL_RESPONSE_HARD_MAX = 16 * 1024 * 1024
+_HF_FACTORY_LOCK = threading.RLock()
+
+
+class _ProviderResponseTooLarge(RuntimeError):
+    pass
+
+
+def _control_response_limit() -> int:
+    raw = os.environ.get("AI_RECORD_STORAGE_CONTROL_RESPONSE_MAX_BYTES", "").strip()
+    try:
+        value = int(raw) if raw else _CONTROL_RESPONSE_DEFAULT
+    except ValueError:
+        value = _CONTROL_RESPONSE_DEFAULT
+    return max(1024, min(_CONTROL_RESPONSE_HARD_MAX, value))
+
+
+def _safe_gitlab_api_base(value: Any) -> str:
+    raw = str(value or "").strip().rstrip("/")
+    if not raw:
+        return "https://gitlab.com/api/v4"
+    if any(ord(ch) < 0x20 or ord(ch) == 0x7f for ch in raw):
+        raise StorageConfigError("TARGET_API_BASE")
+    try:
+        parsed = urlsplit(raw)
+    except Exception as exc:
+        raise StorageConfigError("TARGET_API_BASE") from exc
+    if parsed.scheme.lower() != "https" or not parsed.hostname:
+        raise StorageConfigError("TARGET_API_BASE")
+    if parsed.username is not None or parsed.password is not None or parsed.query or parsed.fragment:
+        raise StorageConfigError("TARGET_API_BASE")
+    if ".." in [seg for seg in parsed.path.split("/") if seg]:
+        raise StorageConfigError("TARGET_API_BASE")
+    if "\\" in parsed.path:
+        raise StorageConfigError("TARGET_API_BASE")
+    try:
+        _ = parsed.port
+    except ValueError as exc:
+        raise StorageConfigError("TARGET_API_BASE") from exc
+    return raw
+
+
+class _BoundedSyncStream(httpx.SyncByteStream):
+    def __init__(self, stream: httpx.SyncByteStream, limit: int) -> None:
+        self._stream = stream
+        self._limit = limit
+        self._seen = 0
+
+    def __iter__(self):
+        for chunk in self._stream:
+            self._seen += len(chunk)
+            if self._seen > self._limit:
+                raise _ProviderResponseTooLarge("provider control response exceeds configured limit")
+            yield chunk
+
+    def close(self) -> None:
+        self._stream.close()
+
+
+class _BoundedTransport(httpx.BaseTransport):
+    def __init__(self, transport: httpx.BaseTransport, limit: int) -> None:
+        self._transport = transport
+        self._limit = limit
+
+    def handle_request(self, request: httpx.Request) -> httpx.Response:
+        response = self._transport.handle_request(request)
+        length = response.headers.get("content-length")
+        if length and length.isdigit() and int(length) > self._limit:
+            response.close()
+            raise _ProviderResponseTooLarge("provider control response exceeds configured limit")
+        response.stream = _BoundedSyncStream(response.stream, self._limit)
+        return response
+
+    def close(self) -> None:
+        self._transport.close()
+
+
+def _with_bounded_hf_client(call):
+    """Run one huggingface_hub control-plane call with bounded HTTP bodies.
+
+    Minimal test doubles and old SDK shims may not expose the factory module;
+    production is pinned to a version that does. In that compatibility-only
+    case there is no underlying SDK HTTP client to wrap, so execute directly.
+    """
+    try:
+        from huggingface_hub.utils import _http as hf_http  # noqa: PLC0415
+    except (ImportError, ModuleNotFoundError):
+        return call()
+
+    with _HF_FACTORY_LOCK:
+        previous = hf_http._GLOBAL_CLIENT_FACTORY
+
+        def factory():
+            client = previous()
+            transport = getattr(client, "_transport", None)
+            if transport is None:
+                client.close()
+                raise RuntimeError("HF_CLIENT_TRANSPORT")
+            client._transport = _BoundedTransport(transport, _control_response_limit())
+            return client
+
+        hf_http.set_client_factory(factory)
+        try:
+            return call()
+        finally:
+            hf_http.set_client_factory(previous)
 
 
 class StorageConfigError(ValueError):
@@ -168,9 +276,13 @@ def _parse_target(raw: dict[str, Any], index: int) -> StorageTarget:
     paths = raw.get("paths") if isinstance(raw.get("paths"), dict) else {}
     token_env = _safe_token_env(raw.get("token_env"))
     token_type = _normalize_token_type(raw.get("token_type") or os.environ.get(token_env + "_TYPE"))
-    api_base = str(raw.get("api_base") or "").strip().rstrip("/")
-    if provider == "gitlab" and not api_base:
-        api_base = "https://gitlab.com/api/v4"
+    raw_api_base = raw.get("api_base")
+    if provider == "gitlab":
+        api_base = _safe_gitlab_api_base(raw_api_base)
+    else:
+        if str(raw_api_base or "").strip():
+            raise StorageConfigError("TARGET_API_BASE_UNSUPPORTED")
+        api_base = ""
     return StorageTarget(
         id=target_id,
         label=label,
@@ -353,10 +465,12 @@ class StorageCoordinator:
             auth_check = getattr(api, "auth_check", None)
             if auth_check and "write" in inspect.signature(auth_check).parameters:
                 await asyncio.to_thread(
-                    auth_check,
-                    repo_id=target.repo,
-                    repo_type="dataset",
-                    write=True,
+                    _with_bounded_hf_client,
+                    lambda: auth_check(
+                        repo_id=target.repo,
+                        repo_type="dataset",
+                        write=True,
+                    ),
                 )
                 return "verified"
         except Exception as exc:
@@ -589,12 +703,14 @@ class StorageCoordinator:
 
             api = HfApi(token=target.token)
             await asyncio.to_thread(
-                api.create_commit,
-                repo_id=target.repo,
-                repo_type="dataset",
-                revision=target.branch,
-                operations=[CommitOperationDelete(path_in_repo=path)],
-                commit_message=message,
+                _with_bounded_hf_client,
+                lambda: api.create_commit(
+                    repo_id=target.repo,
+                    repo_type="dataset",
+                    revision=target.branch,
+                    operations=[CommitOperationDelete(path_in_repo=path)],
+                    commit_message=message,
+                ),
             )
             return "removed-current-view"
         except Exception as exc:  # noqa: BLE001
@@ -611,23 +727,24 @@ class StorageCoordinator:
             "Accept": "application/vnd.github+json",
             "X-GitHub-Api-Version": "2022-11-28",
         }
-        current = await self._client().get(url, headers=headers, params={"ref": target.branch}, timeout=15.0)
-        if current.status_code == 404:
+        current_status, current_json = await self._request_bounded_json(
+            "GET", url, headers=headers, params={"ref": target.branch}, timeout=15.0)
+        if current_status == 404:
             return "already-absent"
-        if current.status_code != 200:
-            raise StorageWriteError("GITHUB_DELETE_LOOKUP", transient=current.status_code in _TRANSIENT_STATUS)
-        sha = str((current.json() or {}).get("sha") or "")
+        if current_status != 200:
+            raise StorageWriteError("GITHUB_DELETE_LOOKUP", transient=current_status in _TRANSIENT_STATUS)
+        sha = str((current_json or {}).get("sha") or "")
         if not sha:
             raise StorageWriteError("GITHUB_DELETE_SHA")
-        resp = await self._client().request(
+        status = await self._request_no_body(
             "DELETE", url, headers=headers,
             json={"message": message, "sha": sha, "branch": target.branch}, timeout=20.0,
         )
-        if resp.status_code in {200, 204}:
+        if status in {200, 204}:
             return "removed-current-view"
-        if resp.status_code == 404:
+        if status == 404:
             return "already-absent"
-        raise StorageWriteError("GITHUB_DELETE", transient=resp.status_code in _TRANSIENT_STATUS)
+        raise StorageWriteError("GITHUB_DELETE", transient=status in _TRANSIENT_STATUS)
 
     async def _delete_gitlab(self, target: StorageTarget, path: str, message: str) -> str:
         base = target.api_base or "https://gitlab.com/api/v4"
@@ -635,15 +752,15 @@ class StorageCoordinator:
         file_path = quote(path, safe="")
         url = f"{base}/projects/{project}/repository/files/{file_path}"
         headers = {"PRIVATE-TOKEN": target.token}
-        resp = await self._client().request(
+        status = await self._request_no_body(
             "DELETE", url, headers=headers,
             json={"branch": target.branch, "commit_message": message}, timeout=20.0,
         )
-        if resp.status_code in {200, 204}:
+        if status in {200, 204}:
             return "removed-current-view"
-        if resp.status_code == 404:
+        if status == 404:
             return "already-absent"
-        raise StorageWriteError("GITLAB_DELETE", transient=resp.status_code in _TRANSIENT_STATUS)
+        raise StorageWriteError("GITLAB_DELETE", transient=status in _TRANSIENT_STATUS)
 
     async def _delete_bitbucket(self, target: StorageTarget, path: str, message: str) -> str:
         workspace, repo = _repo_parts(target.repo)
@@ -652,20 +769,20 @@ class StorageCoordinator:
             f"/src/{quote(target.branch, safe='')}/{quote(path, safe='/')}"
         )
         headers = {"Authorization": f"Bearer {target.token}"}
-        current = await self._client().get(read_url, headers=headers, timeout=15.0)
-        if current.status_code == 404:
+        current_status = await self._request_no_body("GET", read_url, headers=headers, timeout=15.0)
+        if current_status == 404:
             return "already-absent"
-        if current.status_code != 200:
-            raise StorageWriteError("BITBUCKET_DELETE_LOOKUP", transient=current.status_code in _TRANSIENT_STATUS)
+        if current_status != 200:
+            raise StorageWriteError("BITBUCKET_DELETE_LOOKUP", transient=current_status in _TRANSIENT_STATUS)
         url = f"https://api.bitbucket.org/2.0/repositories/{quote(workspace)}/{quote(repo)}/src"
         form = urlencode([("branch", target.branch), ("message", message), ("files", "/" + path)])
         delete_headers = {**headers, "Content-Type": "application/x-www-form-urlencoded"}
-        resp = await self._client().post(
-            url, headers=delete_headers, content=form.encode("utf-8"), timeout=25.0
+        status = await self._request_no_body(
+            "POST", url, headers=delete_headers, content=form.encode("utf-8"), timeout=25.0
         )
-        if resp.status_code in {200, 201}:
+        if status in {200, 201}:
             return "removed-current-view"
-        raise StorageWriteError("BITBUCKET_DELETE", transient=resp.status_code in _TRANSIENT_STATUS)
+        raise StorageWriteError("BITBUCKET_DELETE", transient=status in _TRANSIENT_STATUS)
 
     async def _dispatch(self, target: StorageTarget, path: str, content: bytes, message: str) -> None:
         try:
@@ -692,12 +809,14 @@ class StorageCoordinator:
 
             api = HfApi(token=target.token)
             await asyncio.to_thread(
-                api.create_commit,
-                repo_id=target.repo,
-                repo_type="dataset",
-                revision=target.branch,
-                operations=[CommitOperationAdd(path_in_repo=path, path_or_fileobj=content)],
-                commit_message=message,
+                _with_bounded_hf_client,
+                lambda: api.create_commit(
+                    repo_id=target.repo,
+                    repo_type="dataset",
+                    revision=target.branch,
+                    operations=[CommitOperationAdd(path_in_repo=path, path_or_fileobj=content)],
+                    commit_message=message,
+                ),
             )
         except Exception as exc:  # noqa: BLE001
             # Keep the exception private.  HTTP permission/transient distinction is
@@ -716,6 +835,40 @@ class StorageCoordinator:
             raise StorageWriteError("NO_HTTP_CLIENT", transient=True)
         return self.client
 
+    async def _request_no_body(self, method: str, url: str, **kwargs: Any) -> int:
+        client = self._client()
+        request = client.build_request(method, url, **kwargs)
+        response = await client.send(request, stream=True)
+        try:
+            return response.status_code
+        finally:
+            await response.aclose()
+
+    async def _request_bounded_json(self, method: str, url: str, **kwargs: Any) -> tuple[int, Any]:
+        client = self._client()
+        request = client.build_request(method, url, **kwargs)
+        response = await client.send(request, stream=True)
+        try:
+            limit = _control_response_limit()
+            declared = response.headers.get("content-length")
+            if declared and declared.isdigit() and int(declared) > limit:
+                raise StorageWriteError("PROVIDER_RESPONSE_TOO_LARGE")
+            buf = bytearray()
+            async for chunk in response.aiter_bytes():
+                buf.extend(chunk)
+                if len(buf) > limit:
+                    raise StorageWriteError("PROVIDER_RESPONSE_TOO_LARGE")
+            if not buf:
+                payload = {}
+            else:
+                try:
+                    payload = json.loads(bytes(buf))
+                except Exception as exc:
+                    raise StorageWriteError("PROVIDER_RESPONSE_JSON") from exc
+            return response.status_code, payload
+        finally:
+            await response.aclose()
+
     async def _write_github(self, target: StorageTarget, path: str, content: bytes, message: str) -> None:
         owner, repo = _repo_parts(target.repo)
         url = f"https://api.github.com/repos/{quote(owner)}/{quote(repo)}/contents/{quote(path, safe='/')}"
@@ -729,21 +882,22 @@ class StorageCoordinator:
             "content": base64.b64encode(content).decode("ascii"),
             "branch": target.branch,
         }
-        resp = await self._client().put(url, headers=headers, json=payload, timeout=20.0)
-        if resp.status_code in {200, 201}:
+        status = await self._request_no_body("PUT", url, headers=headers, json=payload, timeout=20.0)
+        if status in {200, 201}:
             return
         # A retry may encounter an already-created idempotent path.  Confirm
         # content equality before treating the conflict as success.
-        if resp.status_code in {409, 422}:
-            get = await self._client().get(url, headers=headers, params={"ref": target.branch}, timeout=15.0)
-            if get.status_code == 200:
+        if status in {409, 422}:
+            get_status, get_json = await self._request_bounded_json(
+                "GET", url, headers=headers, params={"ref": target.branch}, timeout=15.0)
+            if get_status == 200:
                 try:
-                    existing = base64.b64decode((get.json().get("content") or "").replace("\n", ""))
+                    existing = base64.b64decode((get_json.get("content") or "").replace("\n", ""))
                     if existing == content:
                         return
                 except Exception:
                     pass
-        raise StorageWriteError("GITHUB_WRITE", transient=resp.status_code in _TRANSIENT_STATUS)
+        raise StorageWriteError("GITHUB_WRITE", transient=status in _TRANSIENT_STATUS)
 
     async def _write_gitlab(self, target: StorageTarget, path: str, content: bytes, message: str) -> None:
         base = target.api_base or "https://gitlab.com/api/v4"
@@ -752,18 +906,19 @@ class StorageCoordinator:
         url = f"{base}/projects/{project}/repository/files/{file_path}"
         headers = {"PRIVATE-TOKEN": target.token}
         payload = {"branch": target.branch, "commit_message": message, "content": content.decode("utf-8")}
-        resp = await self._client().post(url, headers=headers, json=payload, timeout=20.0)
-        if resp.status_code in {200, 201}:
+        status = await self._request_no_body("POST", url, headers=headers, json=payload, timeout=20.0)
+        if status in {200, 201}:
             return
-        if resp.status_code == 400:
-            get = await self._client().get(url, headers=headers, params={"ref": target.branch}, timeout=15.0)
-            if get.status_code == 200:
+        if status == 400:
+            get_status, get_json = await self._request_bounded_json(
+                "GET", url, headers=headers, params={"ref": target.branch}, timeout=15.0)
+            if get_status == 200:
                 try:
-                    if get.json().get("content_sha256") == hashlib.sha256(content).hexdigest():
+                    if get_json.get("content_sha256") == hashlib.sha256(content).hexdigest():
                         return
                 except Exception:
                     pass
-        raise StorageWriteError("GITLAB_WRITE", transient=resp.status_code in _TRANSIENT_STATUS)
+        raise StorageWriteError("GITLAB_WRITE", transient=status in _TRANSIENT_STATUS)
 
     async def _write_bitbucket(self, target: StorageTarget, path: str, content: bytes, message: str) -> None:
         workspace, repo = _repo_parts(target.repo)
@@ -771,7 +926,7 @@ class StorageCoordinator:
         headers = {"Authorization": f"Bearer {target.token}"}
         files = {"/" + path: (path.rsplit("/", 1)[-1], content, "application/x-ndjson")}
         data = {"branch": target.branch, "message": message}
-        resp = await self._client().post(url, headers=headers, data=data, files=files, timeout=25.0)
-        if resp.status_code in {200, 201}:
+        status = await self._request_no_body("POST", url, headers=headers, data=data, files=files, timeout=25.0)
+        if status in {200, 201}:
             return
-        raise StorageWriteError("BITBUCKET_WRITE", transient=resp.status_code in _TRANSIENT_STATUS)
+        raise StorageWriteError("BITBUCKET_WRITE", transient=status in _TRANSIENT_STATUS)

app.py

@@ -5,7 +5,7 @@
 # Authors: The scikit-plots developers
 # SPDX-License-Identifier: BSD-3-Clause
 #
-# scikit-plots/ai  ·  _hf_spaces_proxy/app.py  v7.0.0
+# scikit-plots/ai  ·  _hf_spaces_proxy/app.py  v7.2.0
 #
 # Server-authoritative chat proxy for sphinx-ai-assistant.
 #
@@ -93,6 +93,7 @@
 #                       prefix AI_RECORD_STORAGE_TOKEN_. When absent, legacy
 #                       TRAINING_DATASET_REPO + HF_* token settings are used.
 #   MAX_BODY_BYTES      Maximum accepted body size.  Default: 10485760.
+#   MAX_UPSTREAM_RESPONSE_BYTES Maximum decoded upstream response bytes. Default: 8388608.
 #   SHARE_MAX_BODY_BYTES Global Share request/canonical snapshot limit. Default: 512000.
 #   SHARE_MAX_ENTRIES   Maximum live in-memory Global Share entries. Default: 256.
 #   SHARE_MAX_TOTAL_BYTES Aggregate live Global Share byte budget. Default: 16777216.
@@ -459,6 +460,49 @@ def _origin_allowed(request: Request) -> bool:
     return False
 
 
+class _UpstreamResponseTooLarge(RuntimeError):
+    """Raised when decoded upstream response bytes exceed the hard application cap."""
+
+
+def _upstream_declared_length(response: httpx.Response) -> int | None:
+    raw = (response.headers.get("content-length") or "").strip()
+    if not raw:
+        return None
+    if not raw.isdigit():
+        raise _UpstreamResponseTooLarge("invalid upstream content-length")
+    value = int(raw)
+    if value < 0:
+        raise _UpstreamResponseTooLarge("invalid upstream content-length")
+    return value
+
+
+def _check_upstream_declared_length(
+    response: httpx.Response, max_bytes: int | None = None
+) -> None:
+    limit = MAX_UPSTREAM_RESPONSE_BYTES if max_bytes is None else int(max_bytes)
+    declared = _upstream_declared_length(response)
+    if declared is not None and declared > limit:
+        raise _UpstreamResponseTooLarge("declared upstream response too large")
+
+
+async def _read_upstream_limited(
+    response: httpx.Response, max_bytes: int | None = None
+) -> bytes:
+    """Read a streamed upstream response with a decoded-byte hard ceiling."""
+    limit = MAX_UPSTREAM_RESPONSE_BYTES if max_bytes is None else int(max_bytes)
+    _check_upstream_declared_length(response, limit)
+    out = bytearray()
+    total = 0
+    async for chunk in response.aiter_bytes():
+        if not chunk:
+            continue
+        total += len(chunk)
+        if total > limit:
+            raise _UpstreamResponseTooLarge("upstream response too large")
+        out.extend(chunk)
+    return bytes(out)
+
+
 async def _read_limited_body(request: Request, max_bytes: int, label: str = "Request") -> bytes:
     """Read an ASGI request incrementally and stop once *max_bytes* is crossed."""
     raw_cl = request.headers.get("content-length")
@@ -736,6 +780,14 @@ MAX_BODY_BYTES: int = max(16_384, min(
     16 * 1024 * 1024,
 ))
 
+#: Maximum decoded upstream response bytes accepted before the proxy aborts the
+#: provider stream. This is independent from request-body limits because a
+#: malicious/misconfigured upstream controls response size.
+MAX_UPSTREAM_RESPONSE_BYTES: int = max(64 * 1024, min(
+    _safe_int(os.environ.get("MAX_UPSTREAM_RESPONSE_BYTES"), 8 * 1024 * 1024),
+    32 * 1024 * 1024,
+))
+
 #: Maximum JSON request size accepted by Global Share.  Kept independent from
 #: chat bodies so a public share service cannot consume 10 MiB per request.
 SHARE_MAX_BODY_BYTES: int = max(16_384, min(
@@ -1845,9 +1897,21 @@ async def _forward(body: bytes, *, structured_body: bytes | None = None) -> Resp
         return json.dumps(clone, ensure_ascii=False, separators=(",", ":")).encode()
 
     async def _post_buffered(request_body: bytes) -> Response:
+        upstream: httpx.Response | None = None
         try:
-            upstream = await _http_client.post(
-                url, content=request_body, headers=headers, timeout=req_timeout
+            request = _http_client.build_request(
+                "POST", url, content=request_body, headers=headers, timeout=req_timeout
+            )
+            upstream = await _http_client.send(request, stream=True)
+            status = upstream.status_code
+            media_type = upstream.headers.get("content-type", "application/json")
+            content = await _read_upstream_limited(upstream)
+        except _UpstreamResponseTooLarge:
+            logger.warning("AI proxy [UPSTREAM_RESPONSE_LIMIT]: upstream response exceeded the application byte ceiling")
+            return _error_response(
+                502,
+                "upstream_response_too_large",
+                "The upstream response exceeded the proxy safety limit.",
             )
         except httpx.ReadTimeout:
             logger.warning(
@@ -1897,11 +1961,14 @@ async def _forward(body: bytes, *, structured_body: bytes | None = None) -> Resp
                 "upstream_request_error",
                 "Failed to reach the upstream service.",
             )
+        finally:
+            if upstream is not None and not upstream.is_closed:
+                await upstream.aclose()
 
         return Response(
-            content=upstream.content,
-            status_code=upstream.status_code,
-            media_type=upstream.headers.get("content-type", "application/json"),
+            content=content,
+            status_code=status,
+            media_type=media_type,
         )
 
     if not stream_requested:
@@ -1983,12 +2050,7 @@ async def _forward(body: bytes, *, structured_body: bytes | None = None) -> Resp
 
     if upstream.status_code != 200:  # noqa: PLR2004
         status = upstream.status_code
-        try:
-            await upstream.aread()
-        except httpx.HTTPError:
-            pass
-        finally:
-            await upstream.aclose()
+        await upstream.aclose()
         logger.warning(
             "AI proxy [UPSTREAM_STATUS]: upstream rejected streaming request status=%d",
             status,
@@ -1999,6 +2061,17 @@ async def _forward(body: bytes, *, structured_body: bytes | None = None) -> Resp
             f"The upstream service returned HTTP {status}.",
         )
 
+    try:
+        _check_upstream_declared_length(upstream)
+    except _UpstreamResponseTooLarge:
+        await upstream.aclose()
+        logger.warning("AI proxy [UPSTREAM_RESPONSE_LIMIT]: declared upstream response exceeded the application byte ceiling")
+        return _error_response(
+            502,
+            "upstream_response_too_large",
+            "The upstream response exceeded the proxy safety limit.",
+        )
+
     content_type = (upstream.headers.get("content-type") or "").lower()
     if "text/event-stream" not in content_type:
         # The custom ai-model Space currently returns an ordinary OpenAI JSON
@@ -2006,7 +2079,15 @@ async def _forward(body: bytes, *, structured_body: bytes | None = None) -> Resp
         # instead of falsely labelling raw JSON bytes as SSE (which the browser
         # would correctly ignore, producing an empty answer).
         try:
-            content = await upstream.aread()
+            content = await _read_upstream_limited(upstream)
+        except _UpstreamResponseTooLarge:
+            await upstream.aclose()
+            logger.warning("AI proxy [UPSTREAM_RESPONSE_LIMIT]: buffered JSON fallback exceeded the application byte ceiling")
+            return _error_response(
+                502,
+                "upstream_response_too_large",
+                "The upstream response exceeded the proxy safety limit.",
+            )
         except httpx.ReadTimeout:
             await upstream.aclose()
             logger.warning(
@@ -2078,10 +2159,20 @@ async def _forward(body: bytes, *, structured_body: bytes | None = None) -> Resp
         emitted = False
         saw_data_field = False
         probe_tail = b""
+        total_bytes = 0
         try:
             async for chunk in upstream.aiter_bytes():
                 if not chunk:
                     continue
+                total_bytes += len(chunk)
+                if total_bytes > MAX_UPSTREAM_RESPONSE_BYTES:
+                    logger.warning("AI proxy [UPSTREAM_RESPONSE_LIMIT]: SSE response exceeded the application byte ceiling")
+                    yield _sse_error_frame(
+                        502,
+                        "upstream_response_too_large",
+                        "The upstream response exceeded the proxy safety limit.",
+                    )
+                    return
                 emitted = True
                 probe = probe_tail + chunk.lower()
                 if b"data:" in probe:
@@ -2358,6 +2449,7 @@ async def root() -> JSONResponse:
                 "required_shared": RATE_LIMIT_REQUIRE_SHARED,
                 "scope": "single_redis_consistency_domain" if RATE_LIMIT_BACKEND == "redis" else "process_local",
             },
+            "limits": {"max_upstream_response_bytes": MAX_UPSTREAM_RESPONSE_BYTES},
             "cors": _cors_public_status(),
         },
         headers={"Cache-Control": "no-store"},
@@ -2464,6 +2556,7 @@ async def health() -> JSONResponse:
                 "ready": (RATE_LIMIT_BACKEND == "redis" and _SHARED_RATE_LIMITER_READY) or (RATE_LIMIT_BACKEND == "local" and not RATE_LIMIT_REQUIRE_SHARED),
                 "required_shared": RATE_LIMIT_REQUIRE_SHARED,
             },
+            "limits": {"max_upstream_response_bytes": MAX_UPSTREAM_RESPONSE_BYTES},
             "cors": _cors_public_status(),
         },
         headers={"Cache-Control": "no-store"},