From 406ea4213da94b5be94d5447095ede2e7c55a08e Mon Sep 17 00:00:00 2001 From: Nicholas Tindle Date: Tue, 14 Apr 2026 13:31:53 -0500 Subject: [PATCH] feat(backend): tier-based workspace file storage limits Replace the global `max_workspace_storage_mb` config with per-tier limits: - FREE: 250 MB, PRO: 1 GB, BUSINESS: 5 GB, ENTERPRISE: 15 GB Move quota enforcement into WorkspaceManager.write_file() so both REST uploads and CoPilot agent writes are covered (previously only REST had it). Co-Authored-By: Claude Opus 4.6 (1M context) --- .../backend/api/features/workspace/routes.py | 46 +++--------- .../api/features/workspace/routes_test.py | 70 +++++++------------ .../backend/backend/copilot/rate_limit.py | 15 ++++ .../backend/backend/util/settings.py | 7 -- .../backend/backend/util/workspace.py | 19 ++++- .../backend/backend/util/workspace_test.py | 15 ++++ 6 files changed, 85 insertions(+), 87 deletions(-) diff --git a/autogpt_platform/backend/backend/api/features/workspace/routes.py b/autogpt_platform/backend/backend/api/features/workspace/routes.py index 39bcc6c7c4..965a19ca21 100644 --- a/autogpt_platform/backend/backend/api/features/workspace/routes.py +++ b/autogpt_platform/backend/backend/api/features/workspace/routes.py @@ -14,6 +14,8 @@ from fastapi import Query, UploadFile from fastapi.responses import Response from pydantic import BaseModel, Field +from backend.api.features.store.exceptions import VirusDetectedError +from backend.copilot.rate_limit import get_workspace_storage_limit_bytes from backend.data.workspace import ( WorkspaceFile, count_workspace_files, @@ -24,7 +26,6 @@ from backend.data.workspace import ( 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 +241,26 @@ 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 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 +299,11 @@ 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 + limit_bytes = await get_workspace_storage_limit_bytes(user_id) return StorageUsageResponse( used_bytes=used_bytes, diff --git a/autogpt_platform/backend/backend/api/features/workspace/routes_test.py b/autogpt_platform/backend/backend/api/features/workspace/routes_test.py index 42726ba051..f0083753d4 100644 --- a/autogpt_platform/backend/backend/api/features/workspace/routes_test.py +++ b/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,23 +338,17 @@ 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, @@ -357,7 +356,6 @@ def test_upload_blocked_by_virus_scan(mocker): 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") diff --git a/autogpt_platform/backend/backend/copilot/rate_limit.py b/autogpt_platform/backend/backend/copilot/rate_limit.py index f72d36de23..d73301d946 100644 --- a/autogpt_platform/backend/backend/copilot/rate_limit.py +++ b/autogpt_platform/backend/backend/copilot/rate_limit.py @@ -58,6 +58,14 @@ TIER_MULTIPLIERS: dict[SubscriptionTier, int] = { 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.""" @@ -447,6 +455,13 @@ get_user_tier.cache_clear = _fetch_user_tier.cache_clear # type: ignore[attr-de 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. diff --git a/autogpt_platform/backend/backend/util/settings.py b/autogpt_platform/backend/backend/util/settings.py index 736219ea9b..128b17e1b9 100644 --- a/autogpt_platform/backend/backend/util/settings.py +++ b/autogpt_platform/backend/backend/util/settings.py @@ -400,13 +400,6 @@ class Config(UpdateTrackingModel["Config"], BaseSettings): description="Maximum file size in MB for workspace files (1-1024 MB)", ) - max_workspace_storage_mb: int = Field( - default=500, - ge=1, - le=10240, - description="Maximum total workspace storage per user in MB.", - ) - # AutoMod configuration automod_enabled: bool = Field( default=False, diff --git a/autogpt_platform/backend/backend/util/workspace.py b/autogpt_platform/backend/backend/util/workspace.py index 5ec4a5b336..84c32ef54f 100644 --- a/autogpt_platform/backend/backend/util/workspace.py +++ b/autogpt_platform/backend/backend/util/workspace.py @@ -12,8 +12,9 @@ from typing import Optional from prisma.errors import UniqueViolationError +from backend.copilot.rate_limit import get_workspace_storage_limit_bytes from backend.data.db_accessors import workspace_db -from backend.data.workspace import WorkspaceFile +from backend.data.workspace import WorkspaceFile, get_workspace_total_size from backend.util.settings import Config from backend.util.virus_scanner import scan_content_safe from backend.util.workspace_storage import compute_file_checksum, get_workspace_storage @@ -185,6 +186,22 @@ class WorkspaceManager: f"{Config().max_file_size_mb}MB limit" ) + # Enforce per-user workspace storage quota (tier-based). + storage_limit = await get_workspace_storage_limit_bytes(self.user_id) + current_usage = await get_workspace_total_size(self.workspace_id) + if current_usage + len(content) > storage_limit: + used_pct = (current_usage / storage_limit) * 100 + raise ValueError( + f"Storage limit exceeded: {current_usage:,} bytes used " + f"of {storage_limit:,} bytes ({used_pct:.1f}%)" + ) + if storage_limit and (current_usage + len(content)) / storage_limit >= 0.8: + logger.warning( + f"User {self.user_id} workspace storage at " + f"{(current_usage + len(content)) / storage_limit * 100:.1f}% " + f"({current_usage + len(content)} / {storage_limit} bytes)" + ) + # Scan here — callers must NOT duplicate this scan. # WorkspaceManager owns virus scanning for all persisted files. await scan_content_safe(content, filename=filename) diff --git a/autogpt_platform/backend/backend/util/workspace_test.py b/autogpt_platform/backend/backend/util/workspace_test.py index dcdbda648e..25f4849303 100644 --- a/autogpt_platform/backend/backend/util/workspace_test.py +++ b/autogpt_platform/backend/backend/util/workspace_test.py @@ -88,6 +88,11 @@ async def test_write_file_no_overwrite_unique_violation_raises_and_cleans_up( ), patch("backend.util.workspace.workspace_db", return_value=mock_db), patch("backend.util.workspace.scan_content_safe", new_callable=AsyncMock), + patch( + "backend.util.workspace.get_workspace_storage_limit_bytes", + return_value=250 * 1024 * 1024, + ), + patch("backend.util.workspace.get_workspace_total_size", return_value=0), ): with pytest.raises(ValueError, match="File already exists"): await manager.write_file( @@ -115,6 +120,11 @@ async def test_write_file_overwrite_conflict_then_retry_succeeds( ), patch("backend.util.workspace.workspace_db", return_value=mock_db), patch("backend.util.workspace.scan_content_safe", new_callable=AsyncMock), + patch( + "backend.util.workspace.get_workspace_storage_limit_bytes", + return_value=250 * 1024 * 1024, + ), + patch("backend.util.workspace.get_workspace_total_size", return_value=0), patch.object(manager, "delete_file", new_callable=AsyncMock) as mock_delete, ): result = await manager.write_file( @@ -148,6 +158,11 @@ async def test_write_file_overwrite_exhausted_retries_raises_and_cleans_up( ), patch("backend.util.workspace.workspace_db", return_value=mock_db), patch("backend.util.workspace.scan_content_safe", new_callable=AsyncMock), + patch( + "backend.util.workspace.get_workspace_storage_limit_bytes", + return_value=250 * 1024 * 1024, + ), + patch("backend.util.workspace.get_workspace_total_size", return_value=0), patch.object(manager, "delete_file", new_callable=AsyncMock), ): with pytest.raises(ValueError, match="Unable to overwrite.*concurrent write"):