From eb5be2ab637e1bc7f7869e0c7dd4578ebca3e640 Mon Sep 17 00:00:00 2001 From: diwu-sf Date: Wed, 19 Feb 2025 07:51:52 -0800 Subject: [PATCH] Fix download workspace zip file event loop hanging (#6722) --- openhands/runtime/action_execution_server.py | 20 ++++++++----------- .../action_execution_client.py | 11 +++++----- openhands/server/routes/files.py | 20 +++++++------------ 3 files changed, 20 insertions(+), 31 deletions(-) diff --git a/openhands/runtime/action_execution_server.py b/openhands/runtime/action_execution_server.py index d9ee74c4293a..1a2d4f1d4a7b 100644 --- a/openhands/runtime/action_execution_server.py +++ b/openhands/runtime/action_execution_server.py @@ -8,7 +8,6 @@ import argparse import asyncio import base64 -import io import mimetypes import os import shutil @@ -21,12 +20,13 @@ from fastapi import Depends, FastAPI, HTTPException, Request, UploadFile from fastapi.exceptions import RequestValidationError -from fastapi.responses import JSONResponse, StreamingResponse +from fastapi.responses import FileResponse, JSONResponse from fastapi.security import APIKeyHeader from openhands_aci.editor.editor import OHEditor from openhands_aci.editor.exceptions import ToolError from openhands_aci.editor.results import ToolResult from pydantic import BaseModel +from starlette.background import BackgroundTask from starlette.exceptions import HTTPException as StarletteHTTPException from uvicorn import run @@ -631,7 +631,7 @@ async def upload_file( raise HTTPException(status_code=500, detail=str(e)) @app.get('/download_files') - async def download_file(path: str): + def download_file(path: str): logger.debug('Downloading files') try: if not os.path.isabs(path): @@ -642,7 +642,7 @@ async def download_file(path: str): if not os.path.exists(path): raise HTTPException(status_code=404, detail='File not found') - with tempfile.TemporaryFile() as temp_zip: + with tempfile.NamedTemporaryFile(suffix='.zip', delete=False) as temp_zip: with ZipFile(temp_zip, 'w') as zipf: for root, _, files in os.walk(path): for file in files: @@ -650,15 +650,11 @@ async def download_file(path: str): zipf.write( file_path, arcname=os.path.relpath(file_path, path) ) - temp_zip.seek(0) # Rewind the file to the beginning after writing - content = temp_zip.read() - # Good for small to medium-sized files. For very large files, streaming directly from the - # file chunks may be more memory-efficient. - zip_stream = io.BytesIO(content) - return StreamingResponse( - content=zip_stream, + return FileResponse( + path=temp_zip.name, media_type='application/zip', - headers={'Content-Disposition': f'attachment; filename={path}.zip'}, + filename=f'{os.path.basename(path)}.zip', + background=BackgroundTask(lambda: os.unlink(temp_zip.name)), ) except Exception as e: diff --git a/openhands/runtime/impl/action_execution/action_execution_client.py b/openhands/runtime/impl/action_execution/action_execution_client.py index 258dcf3a85f7..683c16578dad 100644 --- a/openhands/runtime/impl/action_execution/action_execution_client.py +++ b/openhands/runtime/impl/action_execution/action_execution_client.py @@ -1,4 +1,5 @@ import os +import shutil import tempfile import threading from abc import abstractmethod @@ -143,12 +144,10 @@ def copy_from(self, path: str) -> Path: stream=True, timeout=30, ) as response: - with tempfile.NamedTemporaryFile(delete=False) as temp_file: - total_length = 0 - for chunk in response.iter_content(chunk_size=8192): - if chunk: # filter out keep-alive new chunks - total_length += len(chunk) - temp_file.write(chunk) + with tempfile.NamedTemporaryFile( + suffix='.zip', delete=False + ) as temp_file: + shutil.copyfileobj(response.raw, temp_file, length=16 * 1024) return Path(temp_file.name) except requests.Timeout: raise TimeoutError('Copy operation timed out') diff --git a/openhands/server/routes/files.py b/openhands/server/routes/files.py index 1581d3abaddb..eaa805174f9b 100644 --- a/openhands/server/routes/files.py +++ b/openhands/server/routes/files.py @@ -3,7 +3,6 @@ from fastapi import ( APIRouter, - BackgroundTasks, HTTPException, Request, UploadFile, @@ -12,6 +11,7 @@ from fastapi.responses import FileResponse, JSONResponse from pathspec import PathSpec from pathspec.patterns import GitWildMatchPattern +from starlette.background import BackgroundTask from openhands.core.exceptions import AgentRuntimeUnavailableError from openhands.core.logger import openhands_logger as logger @@ -309,31 +309,25 @@ async def save_file(request: Request): @app.get('/zip-directory') -async def zip_current_workspace( - request: Request, conversation_id: str, background_tasks: BackgroundTasks -): +def zip_current_workspace(request: Request, conversation_id: str): try: logger.debug('Zipping workspace') runtime: Runtime = request.state.conversation.runtime path = runtime.config.workspace_mount_path_in_sandbox try: - zip_file = await call_sync_from_async(runtime.copy_from, path) + zip_file_path = runtime.copy_from(path) except AgentRuntimeUnavailableError as e: logger.error(f'Error zipping workspace: {e}') return JSONResponse( status_code=500, content={'error': f'Error zipping workspace: {e}'}, ) - response = FileResponse( - path=zip_file, + return FileResponse( + path=zip_file_path, filename='workspace.zip', - media_type='application/x-zip-compressed', + media_type='application/zip', + background=BackgroundTask(lambda: os.unlink(zip_file_path)), ) - - # This will execute after the response is sent (So the file is not deleted before being sent) - background_tasks.add_task(zip_file.unlink) - - return response except Exception as e: logger.error(f'Error zipping workspace: {e}') raise HTTPException(