mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 16:14:10 +00:00
fix(api): address review on the event-loop streams
- Per-user SSE cap uses per-connection leases in a sorted set instead of a shared INCR/DECR counter with a TTL. A stream that outlived the TTL could let the counter expire, then decrement another stream's slot or drive the count negative. Leases refresh while a stream sends frames and age out when a stream dies without cleanup. - Shielded cleanup is bounded per step: stream close and on_close in ClosingStreamingResponse, unsubscribe and close in AsyncTopic, lease release, and the replay-budget check, so a dead Redis connection can't hold a request or a graceful shutdown. - Artifact downloads send an ASCII filename plus an RFC 5987 filename* for non-ASCII names, build headers before opening the file, and close disk handles off the event loop. - Type hints and docstrings on the new helpers.
This commit is contained in:
1 parent
54b540ea42
commit
e889392eba
17 files changed
+911
-255
No files matched your search
@@ -11,9 +11,12 @@ from __future__ import annotations
|
||||
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import unicodedata
|
||||
from functools import partial
|
||||
from typing import AsyncIterator, BinaryIO, Mapping, Optional, Tuple
|
||||
from urllib.parse import quote
|
||||
|
||||
import anyio
|
||||
from starlette.requests import Request
|
||||
@@ -47,6 +50,25 @@ def _sanitize_header_filename(filename: Optional[str], fallback: str) -> str:
|
||||
return cleaned or fallback
|
||||
|
||||
|
||||
def _ascii_fold(value: str) -> str:
|
||||
return unicodedata.normalize("NFKD", value).encode("ascii", "ignore").decode("ascii")
|
||||
|
||||
|
||||
def _content_disposition(filename: Optional[str], artifact_id: str) -> str:
|
||||
"""Build an ``attachment`` Content-Disposition value that any filename can travel in.
|
||||
|
||||
Header values are Latin-1, so a non-ASCII name gets an ASCII ``filename``
|
||||
fallback plus an RFC 5987 ``filename*`` carrying the real name.
|
||||
"""
|
||||
name = _sanitize_header_filename(filename, f"artifact-{artifact_id}")
|
||||
if name.isascii():
|
||||
return f'attachment; filename="{name}"'
|
||||
stem, ext = os.path.splitext(name)
|
||||
ascii_stem = _ascii_fold(stem).strip()
|
||||
fallback = f"{ascii_stem or f'artifact-{artifact_id}'}{_ascii_fold(ext)}"
|
||||
return f"attachment; filename=\"{fallback}\"; filename*=UTF-8''{quote(name, safe='')}"
|
||||
|
||||
|
||||
def _load_version(
|
||||
artifact_id: str, decoded_token: Optional[dict], query: Mapping[str, str]
|
||||
) -> Tuple[Optional[dict], Optional[Tuple[str, int]]]:
|
||||
@@ -90,9 +112,14 @@ async def _read_chunks(file_obj: BinaryIO) -> AsyncIterator[bytes]:
|
||||
|
||||
|
||||
async def _close_file(file_obj: BinaryIO) -> None:
|
||||
"""Close the handle; closing anything but an in-memory buffer runs in a worker thread."""
|
||||
close = getattr(file_obj, "close", None)
|
||||
if callable(close):
|
||||
if not callable(close):
|
||||
return
|
||||
if isinstance(file_obj, io.BytesIO):
|
||||
close()
|
||||
else:
|
||||
await anyio.to_thread.run_sync(close)
|
||||
|
||||
|
||||
async def download_artifact(request: Request) -> Response:
|
||||
@@ -118,7 +145,6 @@ async def download_artifact(request: Request) -> Response:
|
||||
if not storage_path:
|
||||
return json_error("No file for this version", 404)
|
||||
|
||||
filename = _sanitize_header_filename(version_row.get("filename"), f"artifact-{artifact_id}")
|
||||
mime_type = version_row.get("mime_type") or "application/octet-stream"
|
||||
storage = await anyio.to_thread.run_sync(StorageCreator.get_storage)
|
||||
|
||||
@@ -157,6 +183,9 @@ async def download_artifact(request: Request) -> Response:
|
||||
)
|
||||
return RedirectResponse(url, status_code=302)
|
||||
|
||||
# Build the headers before opening the file, so nothing between the
|
||||
# open and the response can fail and leave the handle unclosed.
|
||||
headers = {"Content-Disposition": _content_disposition(version_row.get("filename"), artifact_id)}
|
||||
# Stream the bytes in chunks instead of buffering the whole object in
|
||||
# memory (artifacts can be many MB); the handle closes when the
|
||||
# response ends, including on a client disconnect.
|
||||
@@ -164,7 +193,7 @@ async def download_artifact(request: Request) -> Response:
|
||||
return ClosingStreamingResponse(
|
||||
_read_chunks(file_obj),
|
||||
media_type=mime_type,
|
||||
headers={"Content-Disposition": f'attachment; filename="{filename}"'},
|
||||
headers=headers,
|
||||
on_close=partial(_close_file, file_obj),
|
||||
)
|
||||
except FileNotFoundError:
|
||||
|
||||
Reference in new issue
Block a user