Fix download workspace zip file event loop hanging (#6722)

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,23 +642,19 @@ 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:
                             file_path = os.path.join(root, file)
                             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:

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')

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(