fix: replace atexit with BackgroundTask for temp zip cleanup

Using atexit to clean up temporary zip files is unreliable because
atexit handlers only run when the process exits, not after each
download. This means temp files accumulate on disk, one per download,
until the server restarts.

Replace with Starlette's BackgroundTask which runs cleanup after
the response is fully sent, ensuring temp files are deleted promptly.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Test User 2026-05-14 20:17:09 +08:00
parent b23950d035
commit 958dc05b7f
No known key found for this signature in database

View File

@ -1,4 +1,3 @@
import atexit
import re import re
import shutil import shutil
import tempfile import tempfile
@ -6,6 +5,7 @@ from pathlib import Path
from fastapi import APIRouter, HTTPException from fastapi import APIRouter, HTTPException
from fastapi.responses import FileResponse from fastapi.responses import FileResponse
from starlette.background import BackgroundTask
from server.settings import WARE_HOUSE_DIR from server.settings import WARE_HOUSE_DIR
from utils.exceptions import ResourceNotFoundError, ValidationError from utils.exceptions import ResourceNotFoundError, ValidationError
@ -64,13 +64,12 @@ async def download_session(session_id: str):
if zip_path.exists(): if zip_path.exists():
zip_path.unlink() zip_path.unlink()
atexit.register(cleanup_zip)
return FileResponse( return FileResponse(
path=zip_path, path=zip_path,
filename=f"{dir_name}.zip", filename=f"{dir_name}.zip",
media_type="application/zip", media_type="application/zip",
headers={"Content-Disposition": f"attachment; filename={dir_name}.zip"}, headers={"Content-Disposition": f"attachment; filename={dir_name}.zip"},
background=BackgroundTask(cleanup_zip),
) )
except ValidationError as exc: except ValidationError as exc:
raise HTTPException(status_code=400, detail=str(exc)) raise HTTPException(status_code=400, detail=str(exc))