feat(backend): tier-based workspace file storage limits

autogpt_platform/backend/backend/api/features/workspace/routes.py

@@ -2,6 +2,7 @@
 Workspace API routes for managing user file storage.
 """
 
+import asyncio
 import logging
 import os
 import re
@@ -14,6 +15,8 @@
 from fastapi.responses import Response
 from pydantic import BaseModel, Field
 
+from backend.api.features.store.exceptions import VirusDetectedError, VirusScanError
+from backend.copilot.rate_limit import get_workspace_storage_limit_bytes
 from backend.data.workspace import (
     WorkspaceFile,
     count_workspace_files,
@@ -24,7 +27,6 @@
     soft_delete_workspace_file,
 )
 from backend.util.settings import Config
-from backend.util.virus_scanner import scan_content_safe
 from backend.util.workspace import WorkspaceManager
 from backend.util.workspace_storage import get_workspace_storage
 
@@ -240,50 +242,28 @@ async def upload_file(
     # Get or create workspace
     workspace = await get_or_create_workspace(user_id)
 
-    # Pre-write storage cap check (soft check — final enforcement is post-write)
-    storage_limit_bytes = config.max_workspace_storage_mb * 1024 * 1024
-    current_usage = await get_workspace_total_size(workspace.id)
-    if storage_limit_bytes and current_usage + len(content) > storage_limit_bytes:
-        used_percent = (current_usage / storage_limit_bytes) * 100
-        raise fastapi.HTTPException(
-            status_code=413,
-            detail={
-                "message": "Storage limit exceeded",
-                "used_bytes": current_usage,
-                "limit_bytes": storage_limit_bytes,
-                "used_percent": round(used_percent, 1),
-            },
-        )
-
-    # Warn at 80% usage
-    if (
-        storage_limit_bytes
-        and (usage_ratio := (current_usage + len(content)) / storage_limit_bytes) >= 0.8
-    ):
-        logger.warning(
-            f"User {user_id} workspace storage at {usage_ratio * 100:.1f}% "
-            f"({current_usage + len(content)} / {storage_limit_bytes} bytes)"
-        )
-
-    # Virus scan
-    await scan_content_safe(content, filename=filename)
-
-    # Write file via WorkspaceManager
+    # Write file via WorkspaceManager (handles virus scan, per-file size,
+    # and per-user tier-based storage quota internally).
     manager = WorkspaceManager(user_id, workspace.id, session_id)
     try:
         workspace_file = await manager.write_file(
             content, filename, overwrite=overwrite, metadata={"origin": "user-upload"}
         )
+    except VirusDetectedError as e:
+        raise fastapi.HTTPException(status_code=400, detail=str(e)) from e
+    except VirusScanError as e:
+        raise fastapi.HTTPException(status_code=500, detail=str(e)) from e
     except ValueError as e:
-        # write_file raises ValueError for both path-conflict and size-limit
-        # cases; map each to its correct HTTP status.
+        # write_file raises ValueError for path-conflict, size-limit, and
+        # storage-quota cases; map each to its correct HTTP status.
         message = str(e)
-        if message.startswith("File too large"):
+        if message.startswith(("File too large", "Storage limit exceeded")):
             raise fastapi.HTTPException(status_code=413, detail=message) from e
         raise fastapi.HTTPException(status_code=409, detail=message) from e
 
     # Post-write storage check — eliminates TOCTOU race on the quota.
     # If a concurrent upload pushed us over the limit, undo this write.
+    storage_limit_bytes = await get_workspace_storage_limit_bytes(user_id)
     new_total = await get_workspace_total_size(workspace.id)
     if storage_limit_bytes and new_total > storage_limit_bytes:
         try:
@@ -322,12 +302,13 @@ async def get_storage_usage(
     """
     Get storage usage information for the user's workspace.
     """
-    config = Config()
     workspace = await get_or_create_workspace(user_id)
 
-    used_bytes = await get_workspace_total_size(workspace.id)
-    file_count = await count_workspace_files(workspace.id)
-    limit_bytes = config.max_workspace_storage_mb * 1024 * 1024
+    used_bytes, file_count, limit_bytes = await asyncio.gather(
+        get_workspace_total_size(workspace.id),
+        count_workspace_files(workspace.id),
+        get_workspace_storage_limit_bytes(user_id),
+    )
 
     return StorageUsageResponse(
         used_bytes=used_bytes,

autogpt_platform/backend/backend/api/features/workspace/routes_test.py

@@ -151,15 +151,16 @@ def test_list_files_null_metadata_coerced_to_empty_dict(
 # -- upload_file metadata tests --
 
 
+@patch("backend.api.features.workspace.routes.get_workspace_storage_limit_bytes")
 @patch("backend.api.features.workspace.routes.get_or_create_workspace")
 @patch("backend.api.features.workspace.routes.get_workspace_total_size")
-@patch("backend.api.features.workspace.routes.scan_content_safe")
 @patch("backend.api.features.workspace.routes.WorkspaceManager")
 def test_upload_passes_user_upload_origin_metadata(
-    mock_manager_cls, mock_scan, mock_total_size, mock_get_workspace
+    mock_manager_cls, mock_total_size, mock_get_workspace, mock_storage_limit
 ):
     mock_get_workspace.return_value = _make_workspace()
     mock_total_size.return_value = 100
+    mock_storage_limit.return_value = 250 * 1024 * 1024
     written = _make_file(id="new-file", name="doc.pdf")
     mock_instance = AsyncMock()
     mock_instance.write_file.return_value = written
@@ -178,10 +179,9 @@ def test_upload_passes_user_upload_origin_metadata(
 
 @patch("backend.api.features.workspace.routes.get_or_create_workspace")
 @patch("backend.api.features.workspace.routes.get_workspace_total_size")
-@patch("backend.api.features.workspace.routes.scan_content_safe")
 @patch("backend.api.features.workspace.routes.WorkspaceManager")
 def test_upload_returns_409_on_file_conflict(
-    mock_manager_cls, mock_scan, mock_total_size, mock_get_workspace
+    mock_manager_cls, mock_total_size, mock_get_workspace
 ):
     mock_get_workspace.return_value = _make_workspace()
     mock_total_size.return_value = 100
@@ -234,8 +234,8 @@ def test_upload_happy_path(mocker):
         return_value=0,
     )
     mocker.patch(
-        "backend.api.features.workspace.routes.scan_content_safe",
-        return_value=None,
+        "backend.api.features.workspace.routes.get_workspace_storage_limit_bytes",
+        return_value=250 * 1024 * 1024,
     )
     mock_manager = mocker.MagicMock()
     mock_manager.write_file = mocker.AsyncMock(return_value=_MOCK_FILE)
@@ -256,20 +256,24 @@ def test_upload_exceeds_max_file_size(mocker):
     """Files larger than max_file_size_mb should be rejected with 413."""
     cfg = mocker.patch("backend.api.features.workspace.routes.Config")
     cfg.return_value.max_file_size_mb = 0  # 0 MB → any content is too big
-    cfg.return_value.max_workspace_storage_mb = 500
 
     response = _upload(content=b"x" * 1024)
     assert response.status_code == 413
 
 
 def test_upload_storage_quota_exceeded(mocker):
+    """WorkspaceManager.write_file raises ValueError when quota exceeded → 413."""
     mocker.patch(
         "backend.api.features.workspace.routes.get_or_create_workspace",
         return_value=_make_workspace(),
     )
+    mock_manager = mocker.MagicMock()
+    mock_manager.write_file = mocker.AsyncMock(
+        side_effect=ValueError("Storage limit exceeded: 500 MB used of 250 MB (200.0%)")
+    )
     mocker.patch(
-        "backend.api.features.workspace.routes.get_workspace_total_size",
-        return_value=500 * 1024 * 1024,
+        "backend.api.features.workspace.routes.WorkspaceManager",
+        return_value=mock_manager,
     )
 
     response = _upload()
@@ -283,13 +287,14 @@ def test_upload_post_write_quota_race(mocker):
         "backend.api.features.workspace.routes.get_or_create_workspace",
         return_value=_make_workspace(),
     )
+    # Post-write total exceeds the tier-based limit (250 MB for FREE).
     mocker.patch(
         "backend.api.features.workspace.routes.get_workspace_total_size",
-        side_effect=[0, 600 * 1024 * 1024],
+        return_value=600 * 1024 * 1024,
     )
     mocker.patch(
-        "backend.api.features.workspace.routes.scan_content_safe",
-        return_value=None,
+        "backend.api.features.workspace.routes.get_workspace_storage_limit_bytes",
+        return_value=250 * 1024 * 1024,
     )
     mock_manager = mocker.MagicMock()
     mock_manager.write_file = mocker.AsyncMock(return_value=_MOCK_FILE)
@@ -318,8 +323,8 @@ def test_upload_any_extension(mocker):
         return_value=0,
     )
     mocker.patch(
-        "backend.api.features.workspace.routes.scan_content_safe",
-        return_value=None,
+        "backend.api.features.workspace.routes.get_workspace_storage_limit_bytes",
+        return_value=250 * 1024 * 1024,
     )
     mock_manager = mocker.MagicMock()
     mock_manager.write_file = mocker.AsyncMock(return_value=_MOCK_FILE)
@@ -333,31 +338,24 @@ def test_upload_any_extension(mocker):
 
 
 def test_upload_blocked_by_virus_scan(mocker):
-    """Files flagged by ClamAV should be rejected and never written to storage."""
+    """Files flagged by ClamAV should be rejected via WorkspaceManager."""
     from backend.api.features.store.exceptions import VirusDetectedError
 
     mocker.patch(
         "backend.api.features.workspace.routes.get_or_create_workspace",
         return_value=_make_workspace(),
     )
-    mocker.patch(
-        "backend.api.features.workspace.routes.get_workspace_total_size",
-        return_value=0,
-    )
-    mocker.patch(
-        "backend.api.features.workspace.routes.scan_content_safe",
+    mock_manager = mocker.MagicMock()
+    mock_manager.write_file = mocker.AsyncMock(
         side_effect=VirusDetectedError("Eicar-Test-Signature"),
     )
-    mock_manager = mocker.MagicMock()
-    mock_manager.write_file = mocker.AsyncMock(return_value=_MOCK_FILE)
     mocker.patch(
         "backend.api.features.workspace.routes.WorkspaceManager",
         return_value=mock_manager,
     )
 
     response = _upload(filename="evil.exe", content=b"X5O!P%@AP...")
     assert response.status_code == 400
-    mock_manager.write_file.assert_not_called()
 
 
 def test_upload_file_without_extension(mocker):
@@ -371,8 +369,8 @@ def test_upload_file_without_extension(mocker):
         return_value=0,
     )
     mocker.patch(
-        "backend.api.features.workspace.routes.scan_content_safe",
-        return_value=None,
+        "backend.api.features.workspace.routes.get_workspace_storage_limit_bytes",
+        return_value=250 * 1024 * 1024,
     )
     mock_manager = mocker.MagicMock()
     mock_manager.write_file = mocker.AsyncMock(return_value=_MOCK_FILE)
@@ -402,8 +400,8 @@ def test_upload_strips_path_components(mocker):
         return_value=0,
     )
     mocker.patch(
-        "backend.api.features.workspace.routes.scan_content_safe",
-        return_value=None,
+        "backend.api.features.workspace.routes.get_workspace_storage_limit_bytes",
+        return_value=250 * 1024 * 1024,
     )
     mock_manager = mocker.MagicMock()
     mock_manager.write_file = mocker.AsyncMock(return_value=_MOCK_FILE)
@@ -488,14 +486,6 @@ def test_upload_write_file_too_large_returns_413(mocker):
         "backend.api.features.workspace.routes.get_or_create_workspace",
         return_value=_make_workspace(),
     )
-    mocker.patch(
-        "backend.api.features.workspace.routes.get_workspace_total_size",
-        return_value=0,
-    )
-    mocker.patch(
-        "backend.api.features.workspace.routes.scan_content_safe",
-        return_value=None,
-    )
     mock_manager = mocker.MagicMock()
     mock_manager.write_file = mocker.AsyncMock(
         side_effect=ValueError("File too large: 900 bytes exceeds 1MB limit")
@@ -516,14 +506,6 @@ def test_upload_write_file_conflict_returns_409(mocker):
         "backend.api.features.workspace.routes.get_or_create_workspace",
         return_value=_make_workspace(),
     )
-    mocker.patch(
-        "backend.api.features.workspace.routes.get_workspace_total_size",
-        return_value=0,
-    )
-    mocker.patch(
-        "backend.api.features.workspace.routes.scan_content_safe",
-        return_value=None,
-    )
     mock_manager = mocker.MagicMock()
     mock_manager.write_file = mocker.AsyncMock(
         side_effect=ValueError("File already exists at path: /sessions/x/a.txt")
@@ -600,3 +582,51 @@ def test_list_files_offset_is_echoed_back(mock_manager_cls, mock_get_workspace):
     mock_instance.list_files.assert_called_once_with(
         limit=11, offset=50, include_all_sessions=True
     )
+
+
+def test_upload_virus_scan_infrastructure_error_returns_500(mocker):
+    """VirusScanError (ClamAV outage) should return 500, not 409."""
+    from backend.api.features.store.exceptions import VirusScanError
+
+    mocker.patch(
+        "backend.api.features.workspace.routes.get_or_create_workspace",
+        return_value=_make_workspace(),
+    )
+    mock_manager = mocker.MagicMock()
+    mock_manager.write_file = mocker.AsyncMock(
+        side_effect=VirusScanError("ClamAV connection refused"),
+    )
+    mocker.patch(
+        "backend.api.features.workspace.routes.WorkspaceManager",
+        return_value=mock_manager,
+    )
+
+    response = _upload()
+    assert response.status_code == 500
+
+
+def test_get_storage_usage_returns_tier_based_limit(mocker):
+    """get_storage_usage should return the user's tier-based limit, not a static config."""
+    mocker.patch(
+        "backend.api.features.workspace.routes.get_or_create_workspace",
+        return_value=_make_workspace(),
+    )
+    mocker.patch(
+        "backend.api.features.workspace.routes.get_workspace_total_size",
+        return_value=100 * 1024 * 1024,  # 100 MB used
+    )
+    mocker.patch(
+        "backend.api.features.workspace.routes.count_workspace_files",
+        return_value=5,
+    )
+    mocker.patch(
+        "backend.api.features.workspace.routes.get_workspace_storage_limit_bytes",
+        return_value=1024 * 1024 * 1024,  # 1 GB (PRO tier)
+    )
+
+    response = client.get("/storage/usage")
+    assert response.status_code == 200
+    data = response.json()
+    assert data["limit_bytes"] == 1024 * 1024 * 1024
+    assert data["used_bytes"] == 100 * 1024 * 1024
+    assert data["file_count"] == 5

autogpt_platform/backend/backend/copilot/rate_limit.py

@@ -58,6 +58,14 @@ class SubscriptionTier(str, Enum):
 
 DEFAULT_TIER = SubscriptionTier.FREE
 
+# Per-tier workspace storage caps in MB.
+TIER_WORKSPACE_STORAGE_MB: dict[SubscriptionTier, int] = {
+    SubscriptionTier.FREE: 250,  # 250 MB
+    SubscriptionTier.PRO: 1024,  # 1 GB
+    SubscriptionTier.BUSINESS: 5 * 1024,  # 5 GB
+    SubscriptionTier.ENTERPRISE: 15 * 1024,  # 15 GB
+}
+
 
 class UsageWindow(BaseModel):
     """Usage within a single time window."""
@@ -456,6 +464,13 @@ async def get_user_tier(user_id: str) -> SubscriptionTier:
 get_user_tier.cache_delete = _fetch_user_tier.cache_delete  # type: ignore[attr-defined]
 
 
+async def get_workspace_storage_limit_bytes(user_id: str) -> int:
+    """Return the workspace storage cap in bytes for the user's subscription tier."""
+    tier = await get_user_tier(user_id)
+    mb = TIER_WORKSPACE_STORAGE_MB.get(tier, TIER_WORKSPACE_STORAGE_MB[DEFAULT_TIER])
+    return mb * 1024 * 1024
+
+
 async def set_user_tier(user_id: str, tier: SubscriptionTier) -> None:
     """Persist the user's rate-limit tier to the database.