github/AutoGPT
1
1
Fork 0
mirror of https://github.com/Significant-Gravitas/AutoGPT.git synced 2026-04-30 03:00:41 -04:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
Nicholas Tindle
2026-04-14 13:31:53 -05:00
parent b06648de8c
commit 406ea4213d
6 changed files with 85 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
autogpt_platform/backend/backend/api/features/workspace/routes.py
View File

@@ -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,

70
autogpt_platform/backend/backend/api/features/workspace/routes_test.py
View File

@@ -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")

15
autogpt_platform/backend/backend/copilot/rate_limit.py
View File

@@ -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.

7
autogpt_platform/backend/backend/util/settings.py
View File

@@ -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,

19
autogpt_platform/backend/backend/util/workspace.py
View File

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

15
autogpt_platform/backend/backend/util/workspace_test.py
View File

@@ -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"):