0.37.2: stream folder zip-downloads to fix 504 on large folders

A big folder hit a 504 Gateway Timeout: the zip was built into a temp
file *before* any response was sent, so for large folders the backend
stayed silent past nginx's proxy_read_timeout.

Now the zip is streamed as it's built, end to end:
- file_service.open_archive() returns (filename, byte iterator); _iter_zip
  walks the dir and yields zip bytes incrementally via a small drain
  buffer, writing each file in 1 MiB chunks (bounded memory, valid CRCs).
  Same hardening as before — only real regular files; FIFOs/sockets/
  devices/symlinks skipped without open(); per-file read errors skipped.
- /api/files/download and /agent/files/download return a StreamingResponse
  (no temp file). The agent proxy streams the agent response straight
  through (agent_service.stream_download), pulling the first chunk eagerly
  so an offline/bad-token agent still yields a clean status before 200.
- Files page: streamed downloads have no Content-Length, so the progress
  bar shows the running downloaded byte count ("Downloading … 12.3 MB")
  instead of a percentage, after the initial "Preparing …".

Verified end to end via TestClient (200, application/zip, valid zip,
2 MiB file intact, FIFO skipped, no hang).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
menzelj
2026-06-21 20:28:49 +00:00
co-authored by Claude Opus 4.8
parent 0dc430bb2a
commit 5ac9f15de4
8 changed files with 149 additions and 66 deletions
+6 -6
View File
@@ -32,8 +32,7 @@ from fastapi import (
WebSocket, WebSocket,
WebSocketDisconnect, WebSocketDisconnect,
) )
from fastapi.responses import FileResponse, JSONResponse from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
from starlette.background import BackgroundTask
from pydantic import BaseModel from pydantic import BaseModel
from config import settings from config import settings
@@ -668,10 +667,11 @@ def files_read(path: str = Query(...)) -> dict:
@app.get("/agent/files/download", dependencies=[Depends(verify_token)]) @app.get("/agent/files/download", dependencies=[Depends(verify_token)])
def files_download(path: str = Query(...)): def files_download(path: str = Query(...)):
if _file_guard(file_service.is_dir, path): if _file_guard(file_service.is_dir, path):
tmp, filename = _file_guard(file_service.archive_dir, path) filename, chunks = _file_guard(file_service.open_archive, path)
return FileResponse( # Stream the zip as it's built (no temp file, starts immediately).
tmp, filename=filename, media_type="application/zip", return StreamingResponse(
background=BackgroundTask(os.unlink, tmp), chunks, media_type="application/zip",
headers={"Content-Disposition": f'attachment; filename="{filename}"'},
) )
real, filename = _file_guard(file_service.resolve_download, path) real, filename = _file_guard(file_service.resolve_download, path)
return FileResponse(real, filename=filename, media_type="application/octet-stream") return FileResponse(real, filename=filename, media_type="application/octet-stream")
+22 -11
View File
@@ -6,7 +6,7 @@ import os
import tempfile import tempfile
from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, Request, UploadFile from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, Request, UploadFile
from fastapi.responses import FileResponse from fastapi.responses import FileResponse, StreamingResponse
from pydantic import BaseModel from pydantic import BaseModel
from sqlmodel import Session, select from sqlmodel import Session, select
from starlette.background import BackgroundTask from starlette.background import BackgroundTask
@@ -791,19 +791,30 @@ async def agent_files_download(
_user: User = Depends(get_current_user), _user: User = Depends(get_current_user),
): ):
agent = _get_or_404(session, agent_id) agent = _get_or_404(session, agent_id)
tmp = tempfile.NamedTemporaryFile(delete=False) # Stream the agent's response straight through (works for single files and
tmp.close() # for on-the-fly folder zips), so nothing is staged to disk and the
# download starts immediately. Pull the first chunk eagerly so a failed
# agent (offline / bad token / 404) still surfaces a clean HTTP status
# before we commit to a 200 streaming response.
chunks = agent_service.stream_download(
session, agent, "/agent/files/download", params={"path": path}
)
try: try:
await agent_service.download_to_file( first = await chunks.__anext__()
session, agent, "/agent/files/download", tmp.name, params={"path": path} except StopAsyncIteration:
) first = b""
except AgentError as exc: except AgentError as exc:
if os.path.exists(tmp.name):
os.unlink(tmp.name)
_raise(exc) _raise(exc)
return FileResponse(
tmp.name, media_type="application/octet-stream", filename=os.path.basename(path), async def body():
background=BackgroundTask(os.unlink, tmp.name), yield first
async for chunk in chunks:
yield chunk
return StreamingResponse(
body(),
media_type="application/octet-stream",
headers={"Content-Disposition": f'attachment; filename="{os.path.basename(path)}"'},
) )
+13 -6
View File
@@ -19,8 +19,7 @@ from fastapi import (
Request, Request,
UploadFile, UploadFile,
) )
from fastapi.responses import FileResponse from fastapi.responses import FileResponse, StreamingResponse
from starlette.background import BackgroundTask
from pydantic import BaseModel from pydantic import BaseModel
from sqlmodel import Session from sqlmodel import Session
@@ -32,6 +31,12 @@ from services import audit_service, device_service, file_service
router = APIRouter(prefix="/api/files", tags=["files"]) router = APIRouter(prefix="/api/files", tags=["files"])
def _attachment(filename: str) -> str:
"""A safe ``Content-Disposition`` value for an arbitrary filename."""
safe = filename.replace("\\", "_").replace('"', "_")
return f'attachment; filename="{safe}"'
def _ip(request: Request) -> str: def _ip(request: Request) -> str:
return request.client.host if request.client else "unknown" return request.client.host if request.client else "unknown"
@@ -71,10 +76,12 @@ def download(
_user: User = Depends(get_current_user), _user: User = Depends(get_current_user),
): ):
if _guard(file_service.is_dir, path): if _guard(file_service.is_dir, path):
tmp, filename = _guard(file_service.archive_dir, path) filename, chunks = _guard(file_service.open_archive, path)
return FileResponse( # Stream the zip as it's built so the response starts immediately
tmp, filename=filename, media_type="application/zip", # (large folders no longer hit the proxy's read timeout).
background=BackgroundTask(os.unlink, tmp), return StreamingResponse(
chunks, media_type="application/zip",
headers={"Content-Disposition": _attachment(filename)},
) )
real, filename = _guard(file_service.resolve_download, path) real, filename = _guard(file_service.resolve_download, path)
return FileResponse(real, filename=filename, media_type="application/octet-stream") return FileResponse(real, filename=filename, media_type="application/octet-stream")
+29
View File
@@ -137,6 +137,35 @@ async def download_to_file(
raise AgentError(502, "agent_unreachable", str(exc)) from exc raise AgentError(502, "agent_unreachable", str(exc)) from exc
async def stream_download(
session: Session,
agent: Agent,
path: str,
*,
params: Optional[dict] = None,
):
"""Stream a GET from the agent straight through, yielding chunks.
Unlike :func:`download_to_file` this never buffers to disk, so a large
response (e.g. a folder zip the agent builds on the fly) starts flowing to
the browser immediately instead of being staged first.
"""
url = agent.url.rstrip("/") + path
headers = {"Authorization": f"Bearer {agent.token}"}
try:
async with httpx.AsyncClient(follow_redirects=True) as client:
async with client.stream("GET", url, headers=headers, params=params, timeout=None) as resp:
if resp.status_code >= 400:
text = (await resp.aread()).decode("utf-8", "replace")
_handle_status(session, agent, resp.status_code, text)
_handle_status(session, agent, resp.status_code)
async for chunk in resp.aiter_bytes(1024 * 256):
yield chunk
except httpx.HTTPError as exc:
_mark(session, agent, "offline")
raise AgentError(502, "agent_unreachable", str(exc)) from exc
async def upload_file( async def upload_file(
session: Session, session: Session,
agent: Agent, agent: Agent,
+69 -39
View File
@@ -10,8 +10,8 @@ from __future__ import annotations
import os import os
import shutil import shutil
import tempfile
import zipfile import zipfile
from collections.abc import Iterator
from services.device_service import BrowseError, _is_allowed, _real_root from services.device_service import BrowseError, _is_allowed, _real_root
@@ -182,54 +182,84 @@ def is_dir(path: str) -> bool:
return os.path.isdir(_safe_real(path)) return os.path.isdir(_safe_real(path))
def archive_dir(path: str) -> tuple[str, str]: class _ZipBuffer:
"""Zip a directory (recursively) into a temp file. """A writable sink that hands out and clears whatever was written to it.
Returns ``(tmp_zip_path, download_filename)``. The caller is responsible Lets us drive ``zipfile`` while draining its output incrementally so the
for deleting the temp file once it has been streamed to the client. archive can be streamed to the client instead of buffered to disk.
"""
def __init__(self) -> None:
self._buf = bytearray()
def write(self, data: bytes) -> int:
self._buf += data
return len(data)
def flush(self) -> None: # pragma: no cover - zipfile calls this
pass
def take(self) -> bytes:
data = bytes(self._buf)
self._buf.clear()
return data
def open_archive(path: str) -> tuple[str, "Iterator[bytes]"]:
"""Validate a directory and return ``(download_filename, byte_iterator)``.
The iterator zips the directory recursively **on the fly**, yielding bytes
as they are produced so the response starts immediately (no waiting for the
whole archive to build → no gateway timeout) and memory stays bounded.
Only regular files and real subdirectories are archived. Symlinks are Only regular files and real subdirectories are archived. Symlinks are
skipped (no sandbox escape / loops); special files (FIFOs, sockets, skipped (no sandbox escape / loops); special files (FIFOs, sockets,
devices) are skipped too — opening a FIFO would block forever and a devices) are skipped too — opening a FIFO would block forever and a socket
socket can't be read at all. Files that can't be read (permissions, or can't be read at all. Files that can't be read (permissions, or that vanish
that vanish mid-walk) are skipped individually rather than aborting the mid-walk) are skipped individually rather than aborting the whole archive.
whole archive.
""" """
real = _safe_real(path) real = _safe_real(path)
if not os.path.isdir(real): if not os.path.isdir(real):
raise BrowseError(f"Not a directory: {path}") raise BrowseError(f"Not a directory: {path}")
name = os.path.basename(path.rstrip("/")) or "root" name = os.path.basename(path.rstrip("/")) or "root"
return f"{name}.zip", _iter_zip(real, name)
fd, tmp = tempfile.mkstemp(suffix=".zip")
os.close(fd) def _iter_zip(real: str, name: str):
try: sink = _ZipBuffer()
with zipfile.ZipFile(tmp, "w", zipfile.ZIP_DEFLATED) as zf: with zipfile.ZipFile(sink, "w", zipfile.ZIP_DEFLATED) as zf:
for root, dirs, files in os.walk(real): for root, dirs, files in os.walk(real):
# Don't follow symlinked directories (avoids loops / escapes). # Don't follow symlinked directories (avoids loops / escapes).
dirs[:] = [d for d in dirs if not os.path.islink(os.path.join(root, d))] dirs[:] = [d for d in dirs if not os.path.islink(os.path.join(root, d))]
rel_root = os.path.relpath(root, real) rel_root = os.path.relpath(root, real)
if not files and not dirs and rel_root != ".": if not files and not dirs and rel_root != ".":
# Preserve otherwise-empty directories. # Preserve otherwise-empty directories.
zf.writestr(os.path.join(name, rel_root) + "/", "") zf.writestr(os.path.join(name, rel_root) + "/", "")
for f in files: if chunk := sink.take():
full = os.path.join(root, f) yield chunk
# os.path.isfile follows symlinks; combined with the islink for f in files:
# check it admits only real regular files (skips FIFOs, full = os.path.join(root, f)
# sockets, devices and symlinks without ever open()-ing them). # os.path.isfile follows symlinks; combined with the islink
if os.path.islink(full) or not os.path.isfile(full): # check it admits only real regular files (skips FIFOs, sockets,
continue # devices and symlinks without ever open()-ing them).
arc = (os.path.join(name, rel_root, f) if rel_root != "." if os.path.islink(full) or not os.path.isfile(full):
else os.path.join(name, f)) continue
try: arc = (os.path.join(name, rel_root, f) if rel_root != "."
zf.write(full, arc) else os.path.join(name, f))
except OSError: try:
# Unreadable or vanished mid-walk — skip just this file. info = zipfile.ZipInfo.from_file(full, arc)
continue info.compress_type = zipfile.ZIP_DEFLATED
except OSError: with open(full, "rb") as src, zf.open(info, "w") as dest:
if os.path.exists(tmp): while buf := src.read(1024 * 1024):
os.unlink(tmp) dest.write(buf)
raise if chunk := sink.take():
return tmp, f"{name}.zip" yield chunk
except OSError:
# Unreadable or vanished mid-walk — skip just this file.
continue
if chunk := sink.take():
yield chunk
yield sink.take()
def upload_target( def upload_target(
+1 -1
View File
@@ -1,3 +1,3 @@
"""Single source of truth for the StackPilot release version.""" """Single source of truth for the StackPilot release version."""
APP_VERSION = "0.37.1" APP_VERSION = "0.37.2"
+1 -1
View File
@@ -1,7 +1,7 @@
{ {
"name": "stackpilot-frontend", "name": "stackpilot-frontend",
"private": true, "private": true,
"version": "0.37.1", "version": "0.37.2",
"type": "module", "type": "module",
"scripts": { "scripts": {
"dev": "vite", "dev": "vite",
+8 -2
View File
@@ -88,6 +88,7 @@ export function Files() {
pct: number; pct: number;
kind: "upload" | "download"; kind: "upload" | "download";
indeterminate?: boolean; indeterminate?: boolean;
loaded?: number;
} | null>(null); } | null>(null);
const fileInput = useRef<HTMLInputElement>(null); const fileInput = useRef<HTMLInputElement>(null);
const folderInput = useRef<HTMLInputElement>(null); const folderInput = useRef<HTMLInputElement>(null);
@@ -222,6 +223,7 @@ export function Files() {
pct: total ? (loaded / total) * 100 : 0, pct: total ? (loaded / total) * 100 : 0,
kind: "download", kind: "download",
indeterminate: !total, indeterminate: !total,
loaded,
}), }),
) )
.catch((err) => toast.error(apiErrorMessage(err))) .catch((err) => toast.error(apiErrorMessage(err)))
@@ -329,14 +331,18 @@ export function Files() {
)} )}
<span className="truncate"> <span className="truncate">
{progress.kind === "download" {progress.kind === "download"
? progress.indeterminate ? progress.indeterminate && !progress.loaded
? `Preparing ${progress.label}` ? `Preparing ${progress.label}`
: `Downloading ${progress.label}` : `Downloading ${progress.label}`
: `Uploading ${progress.label}`} : `Uploading ${progress.label}`}
</span> </span>
</span> </span>
<span className="tabular-nums"> <span className="tabular-nums">
{progress.indeterminate ? "" : `${Math.round(progress.pct)}%`} {progress.indeterminate
? progress.loaded
? formatBytes(progress.loaded)
: ""
: `${Math.round(progress.pct)}%`}
</span> </span>
</div> </div>
<div className="h-2 w-full overflow-hidden rounded-full bg-slate-200 dark:bg-slate-700"> <div className="h-2 w-full overflow-hidden rounded-full bg-slate-200 dark:bg-slate-700">