Spaces:
Running
Running
import os | |
from typing import List, Optional | |
from fastapi import FastAPI, UploadFile, File, HTTPException, APIRouter, Form, Depends, Request | |
from fastapi.responses import Response, JSONResponse | |
from pydantic import BaseModel | |
from utils.logger import logger | |
from utils.auth_utils import get_current_user_id, get_user_id_from_stream_auth, get_optional_user_id | |
from sandbox.sandbox import get_or_start_sandbox | |
from services.supabase import DBConnection | |
from agent.api import get_or_create_project_sandbox | |
# Initialize shared resources | |
router = APIRouter(tags=["sandbox"]) | |
db = None | |
def initialize(_db: DBConnection): | |
"""Initialize the sandbox API with resources from the main API.""" | |
global db | |
db = _db | |
logger.info("Initialized sandbox API with database connection") | |
class FileInfo(BaseModel): | |
"""Model for file information""" | |
name: str | |
path: str | |
is_dir: bool | |
size: int | |
mod_time: str | |
permissions: Optional[str] = None | |
async def verify_sandbox_access(client, sandbox_id: str, user_id: Optional[str] = None): | |
""" | |
Verify that a user has access to a specific sandbox based on account membership. | |
Args: | |
client: The Supabase client | |
sandbox_id: The sandbox ID to check access for | |
user_id: The user ID to check permissions for. Can be None for public resource access. | |
Returns: | |
dict: Project data containing sandbox information | |
Raises: | |
HTTPException: If the user doesn't have access to the sandbox or sandbox doesn't exist | |
""" | |
# Find the project that owns this sandbox | |
project_result = await client.table('projects').select('*').filter('sandbox->>id', 'eq', sandbox_id).execute() | |
if not project_result.data or len(project_result.data) == 0: | |
raise HTTPException(status_code=404, detail="Sandbox not found") | |
project_data = project_result.data[0] | |
if project_data.get('is_public'): | |
return project_data | |
# For private projects, we must have a user_id | |
if not user_id: | |
raise HTTPException(status_code=401, detail="Authentication required for this resource") | |
account_id = project_data.get('account_id') | |
# Verify account membership | |
if account_id: | |
account_user_result = await client.schema('basejump').from_('account_user').select('account_role').eq('user_id', user_id).eq('account_id', account_id).execute() | |
if account_user_result.data and len(account_user_result.data) > 0: | |
return project_data | |
raise HTTPException(status_code=403, detail="Not authorized to access this sandbox") | |
async def get_sandbox_by_id_safely(client, sandbox_id: str): | |
""" | |
Safely retrieve a sandbox object by its ID, using the project that owns it. | |
Args: | |
client: The Supabase client | |
sandbox_id: The sandbox ID to retrieve | |
Returns: | |
Sandbox: The sandbox object | |
Raises: | |
HTTPException: If the sandbox doesn't exist or can't be retrieved | |
""" | |
# Find the project that owns this sandbox | |
project_result = await client.table('projects').select('project_id').filter('sandbox->>id', 'eq', sandbox_id).execute() | |
if not project_result.data or len(project_result.data) == 0: | |
logger.error(f"No project found for sandbox ID: {sandbox_id}") | |
raise HTTPException(status_code=404, detail="Sandbox not found - no project owns this sandbox ID") | |
project_id = project_result.data[0]['project_id'] | |
logger.debug(f"Found project {project_id} for sandbox {sandbox_id}") | |
try: | |
# Get the sandbox | |
sandbox, retrieved_sandbox_id, sandbox_pass = await get_or_create_project_sandbox(client, project_id) | |
# Verify we got the right sandbox | |
if retrieved_sandbox_id != sandbox_id: | |
logger.warning(f"Retrieved sandbox ID {retrieved_sandbox_id} doesn't match requested ID {sandbox_id} for project {project_id}") | |
# Fall back to the direct method if IDs don't match (shouldn't happen but just in case) | |
sandbox = await get_or_start_sandbox(sandbox_id) | |
return sandbox | |
except Exception as e: | |
logger.error(f"Error retrieving sandbox {sandbox_id}: {str(e)}") | |
raise HTTPException(status_code=500, detail=f"Failed to retrieve sandbox: {str(e)}") | |
async def create_file( | |
sandbox_id: str, | |
path: str = Form(...), | |
file: UploadFile = File(...), | |
request: Request = None, | |
user_id: Optional[str] = Depends(get_optional_user_id) | |
): | |
"""Create a file in the sandbox using direct file upload""" | |
logger.info(f"Received file upload request for sandbox {sandbox_id}, path: {path}, user_id: {user_id}") | |
client = await db.client | |
# Verify the user has access to this sandbox | |
await verify_sandbox_access(client, sandbox_id, user_id) | |
try: | |
# Get sandbox using the safer method | |
sandbox = await get_sandbox_by_id_safely(client, sandbox_id) | |
# Read file content directly from the uploaded file | |
content = await file.read() | |
# Create file using raw binary content | |
sandbox.fs.upload_file(path, content) | |
logger.info(f"File created at {path} in sandbox {sandbox_id}") | |
return {"status": "success", "created": True, "path": path} | |
except Exception as e: | |
logger.error(f"Error creating file in sandbox {sandbox_id}: {str(e)}") | |
raise HTTPException(status_code=500, detail=str(e)) | |
# For backward compatibility, keep the JSON version too | |
async def create_file_json( | |
sandbox_id: str, | |
file_request: dict, | |
request: Request = None, | |
user_id: Optional[str] = Depends(get_optional_user_id) | |
): | |
"""Create a file in the sandbox using JSON (legacy support)""" | |
logger.info(f"Received JSON file creation request for sandbox {sandbox_id}, user_id: {user_id}") | |
client = await db.client | |
# Verify the user has access to this sandbox | |
await verify_sandbox_access(client, sandbox_id, user_id) | |
try: | |
# Get sandbox using the safer method | |
sandbox = await get_sandbox_by_id_safely(client, sandbox_id) | |
# Get file path and content | |
path = file_request.get("path") | |
content = file_request.get("content", "") | |
if not path: | |
logger.error(f"Missing file path in request for sandbox {sandbox_id}") | |
raise HTTPException(status_code=400, detail="File path is required") | |
# Convert string content to bytes | |
if isinstance(content, str): | |
content = content.encode('utf-8') | |
# Create file | |
sandbox.fs.upload_file(path, content) | |
logger.info(f"File created at {path} in sandbox {sandbox_id}") | |
return {"status": "success", "created": True, "path": path} | |
except Exception as e: | |
logger.error(f"Error creating file in sandbox {sandbox_id}: {str(e)}") | |
raise HTTPException(status_code=500, detail=str(e)) | |
async def list_files( | |
sandbox_id: str, | |
path: str, | |
request: Request = None, | |
user_id: Optional[str] = Depends(get_optional_user_id) | |
): | |
"""List files and directories at the specified path""" | |
logger.info(f"Received list files request for sandbox {sandbox_id}, path: {path}, user_id: {user_id}") | |
client = await db.client | |
# Verify the user has access to this sandbox | |
await verify_sandbox_access(client, sandbox_id, user_id) | |
try: | |
# Get sandbox using the safer method | |
sandbox = await get_sandbox_by_id_safely(client, sandbox_id) | |
# List files | |
files = sandbox.fs.list_files(path) | |
result = [] | |
for file in files: | |
# Convert file information to our model | |
# Ensure forward slashes are used for paths, regardless of OS | |
full_path = f"{path.rstrip('/')}/{file.name}" if path != '/' else f"/{file.name}" | |
file_info = FileInfo( | |
name=file.name, | |
path=full_path, # Use the constructed path | |
is_dir=file.is_dir, | |
size=file.size, | |
mod_time=str(file.mod_time), | |
permissions=getattr(file, 'permissions', None) | |
) | |
result.append(file_info) | |
logger.info(f"Successfully listed {len(result)} files in sandbox {sandbox_id}") | |
return {"files": [file.dict() for file in result]} | |
except Exception as e: | |
logger.error(f"Error listing files in sandbox {sandbox_id}: {str(e)}") | |
raise HTTPException(status_code=500, detail=str(e)) | |
async def read_file( | |
sandbox_id: str, | |
path: str, | |
request: Request = None, | |
user_id: Optional[str] = Depends(get_optional_user_id) | |
): | |
"""Read a file from the sandbox""" | |
logger.info(f"Received file read request for sandbox {sandbox_id}, path: {path}, user_id: {user_id}") | |
client = await db.client | |
# Verify the user has access to this sandbox | |
await verify_sandbox_access(client, sandbox_id, user_id) | |
try: | |
# Get sandbox using the safer method | |
sandbox = await get_sandbox_by_id_safely(client, sandbox_id) | |
# Read file | |
content = sandbox.fs.download_file(path) | |
# Return a Response object with the content directly | |
filename = os.path.basename(path) | |
logger.info(f"Successfully read file {filename} from sandbox {sandbox_id}") | |
return Response( | |
content=content, | |
media_type="application/octet-stream", | |
headers={"Content-Disposition": f"attachment; filename={filename}"} | |
) | |
except Exception as e: | |
logger.error(f"Error reading file in sandbox {sandbox_id}: {str(e)}") | |
raise HTTPException(status_code=500, detail=str(e)) | |
async def ensure_project_sandbox_active( | |
project_id: str, | |
request: Request = None, | |
user_id: Optional[str] = Depends(get_optional_user_id) | |
): | |
""" | |
Ensure that a project's sandbox is active and running. | |
Checks the sandbox status and starts it if it's not running. | |
""" | |
logger.info(f"Received ensure sandbox active request for project {project_id}, user_id: {user_id}") | |
client = await db.client | |
# Find the project and sandbox information | |
project_result = await client.table('projects').select('*').eq('project_id', project_id).execute() | |
if not project_result.data or len(project_result.data) == 0: | |
logger.error(f"Project not found: {project_id}") | |
raise HTTPException(status_code=404, detail="Project not found") | |
project_data = project_result.data[0] | |
# For public projects, no authentication is needed | |
if not project_data.get('is_public'): | |
# For private projects, we must have a user_id | |
if not user_id: | |
logger.error(f"Authentication required for private project {project_id}") | |
raise HTTPException(status_code=401, detail="Authentication required for this resource") | |
account_id = project_data.get('account_id') | |
# Verify account membership | |
if account_id: | |
account_user_result = await client.schema('basejump').from_('account_user').select('account_role').eq('user_id', user_id).eq('account_id', account_id).execute() | |
if not (account_user_result.data and len(account_user_result.data) > 0): | |
logger.error(f"User {user_id} not authorized to access project {project_id}") | |
raise HTTPException(status_code=403, detail="Not authorized to access this project") | |
try: | |
# Get or create the sandbox | |
logger.info(f"Ensuring sandbox is active for project {project_id}") | |
sandbox, sandbox_id, sandbox_pass = await get_or_create_project_sandbox(client, project_id) | |
logger.info(f"Successfully ensured sandbox {sandbox_id} is active for project {project_id}") | |
return { | |
"status": "success", | |
"sandbox_id": sandbox_id, | |
"message": "Sandbox is active" | |
} | |
except Exception as e: | |
logger.error(f"Error ensuring sandbox is active for project {project_id}: {str(e)}") | |
raise HTTPException(status_code=500, detail=str(e)) | |