Add per-user storage visibility and document delete flow

Author: owenisas

backend/lambda_handler.py

@@ -108,6 +108,7 @@ Resources:
           FRONTEND_EVENTS_TABLE: !Ref FrontendEventsTable
           WAITLIST_TABLE: !Ref WaitlistTable
           ASSETS_BUCKET: !Ref AssetsBucket
+          MAX_USER_STORAGE_BYTES: "5368709120"
       Policies:
         - DynamoDBCrudPolicy:
             TableName: !Ref SessionsTable
@@ -180,6 +181,12 @@ Resources:
             Path: /documents
             Method: GET
             RestApiId: !Ref FPAIApi
+        DeleteDocument:
+          Type: Api
+          Properties:
+            Path: /documents
+            Method: DELETE
+            RestApiId: !Ref FPAIApi
         DocumentDownload:
           Type: Api
           Properties:

backend/template.yaml

@@ -504,6 +504,168 @@ def test_get_presigned_url_success_includes_s3_refs():
     assert body["s3_key"].startswith("uploads/")
 
 
+def test_storage_quota_reservation_accepts_exact_boundary(monkeypatch):
+    monkeypatch.delenv("USERS_TABLE", raising=False)
+    monkeypatch.setenv("MAX_USER_STORAGE_BYTES", "100")
+    user_id = "quota-boundary-user"
+    lambda_handler.user_store.delete(user_id)
+
+    lambda_handler._ensure_user_storage_backfilled(user_id)
+    ok1, _ = lambda_handler._reserve_storage_quota_atomic(user_id, requested_bytes=60)
+    ok2, _ = lambda_handler._reserve_storage_quota_atomic(user_id, requested_bytes=40)
+    ok3, snapshot = lambda_handler._reserve_storage_quota_atomic(user_id, requested_bytes=1)
+
+    assert ok1 is True
+    assert ok2 is True
+    assert ok3 is False
+    assert snapshot["used_bytes"] + snapshot["reserved_bytes"] == snapshot["limit_bytes"]
+
+
+def test_get_presigned_url_returns_413_when_storage_limit_exceeded(monkeypatch):
+    event = {
+        "httpMethod": "POST",
+        "path": "/uploads/presigned",
+        "body": json.dumps(
+            {
+                "filename": "notes.md",
+                "file_type": "text/markdown",
+                "size_bytes": 1024,
+            }
+        ),
+    }
+    snapshot = {
+        "used_bytes": 90,
+        "reserved_bytes": 10,
+        "limit_bytes": 100,
+        "available_bytes": 0,
+    }
+
+    with patch.dict("os.environ", {"ASSETS_BUCKET": "bucket"}), patch(
+        "lambda_handler._reserve_storage_quota_atomic"
+    ) as mock_reserve:
+        mock_reserve.return_value = (False, snapshot)
+        response = handler(event, None)
+
+    assert response["statusCode"] == 413
+    body = json.loads(response["body"])
+    assert body["error"] == "Storage limit exceeded"
+    assert body["storage"] == snapshot
+
+
+def test_storage_backfill_runs_once_per_user(monkeypatch):
+    monkeypatch.delenv("USERS_TABLE", raising=False)
+    user_id = "storage-backfill-user"
+    lambda_handler.user_store.delete(user_id)
+
+    with patch("lambda_handler._sum_user_upload_bytes") as mock_sum:
+        mock_sum.return_value = 123
+        first = lambda_handler._ensure_user_storage_backfilled(user_id)
+        second = lambda_handler._ensure_user_storage_backfilled(user_id)
+
+    assert first["used_bytes"] == 123
+    assert second["used_bytes"] == 123
+    assert mock_sum.call_count == 1
+
+
+def test_cleanup_stale_upload_preps_releases_stale_reservation(monkeypatch):
+    stale_row = {
+        "upload_id": "resv_1",
+        "kind": "storage_reservation",
+        "owner_user_id": "user-1",
+        "status": "reserved",
+        "reserved_size_bytes": 500,
+        "s3_bucket": "bucket",
+        "s3_key": "uploads/user-1/u1/file.txt",
+        "created_at": "2020-01-01T00:00:00",
+    }
+    table = MagicMock()
+    table.scan.return_value = {"Items": [stale_row]}
+    s3_client = MagicMock()
+
+    monkeypatch.setattr(lambda_handler, "Key", None)
+    with patch("lambda_handler.get_upload_prep_table", return_value=table), patch(
+        "lambda_handler.get_aws_client", return_value=s3_client
+    ), patch("lambda_handler._release_reserved_storage", return_value=True) as mock_release:
+        result = lambda_handler._cleanup_stale_upload_preps({"source": "aws.events"})
+
+    assert result["success"] is True
+    assert result["deleted"] == 1
+    mock_release.assert_called_once_with("user-1", reservation_id="resv_1", reserved_size_bytes=500)
+    table.delete_item.assert_called_once_with(Key={"upload_id": "resv_1"})
+
+
+def test_cleanup_stale_upload_preps_decrements_used_for_stale_prepared_upload(monkeypatch):
+    stale_row = {
+        "upload_id": "up_1",
+        "kind": "prepared_upload",
+        "owner_user_id": "user-1",
+        "status": "ready",
+        "storage_committed": True,
+        "storage_committed_bytes": 321,
+        "created_at": "2020-01-01T00:00:00",
+    }
+    table = MagicMock()
+    table.scan.return_value = {"Items": [stale_row]}
+    s3_client = MagicMock()
+
+    monkeypatch.setattr(lambda_handler, "Key", None)
+    with patch("lambda_handler.get_upload_prep_table", return_value=table), patch(
+        "lambda_handler.get_aws_client", return_value=s3_client
+    ), patch("lambda_handler._decrement_used_storage", return_value=True) as mock_decrement:
+        result = lambda_handler._cleanup_stale_upload_preps({"source": "aws.events"})
+
+    assert result["success"] is True
+    mock_decrement.assert_called_once_with("user-1", used_size_bytes=321)
+    table.delete_item.assert_called_once_with(Key={"upload_id": "up_1"})
+
+
+def test_cleanup_stale_upload_preps_ignores_document_tombstones(monkeypatch):
+    stale_row = {
+        "upload_id": "docdel_1",
+        "kind": "document_tombstone",
+        "owner_user_id": "user-1",
+        "status": "deleted",
+        "created_at": "2020-01-01T00:00:00",
+    }
+    table = MagicMock()
+    table.scan.return_value = {"Items": [stale_row]}
+    s3_client = MagicMock()
+
+    monkeypatch.setattr(lambda_handler, "Key", None)
+    with patch("lambda_handler.get_upload_prep_table", return_value=table), patch(
+        "lambda_handler.get_aws_client", return_value=s3_client
+    ):
+        result = lambda_handler._cleanup_stale_upload_preps({"source": "aws.events"})
+
+    assert result["success"] is True
+    assert result["scanned"] == 1
+    assert result["deleted"] == 0
+    table.delete_item.assert_not_called()
+
+
+def test_status_includes_storage_usage():
+    event = {
+        "httpMethod": "GET",
+        "path": "/status",
+    }
+    snapshot = {
+        "used_bytes": 10,
+        "reserved_bytes": 5,
+        "limit_bytes": 100,
+        "available_bytes": 85,
+    }
+
+    with patch("lambda_handler._ensure_user_storage_backfilled", return_value=snapshot), patch(
+        "lambda_handler.user_store"
+    ) as mock_user_store:
+        mock_user_store.get.return_value = None
+        response = handler(event, None)
+
+    assert response["statusCode"] == 200
+    body = json.loads(response["body"])
+    assert body["storage"] == snapshot
+
+
 def test_documents_endpoint_lists_non_image_attachments():
     event = {
         "httpMethod": "GET",
@@ -670,6 +832,142 @@ def test_documents_endpoint_returns_partial_data_when_a_session_fails():
     assert body["documents"][0]["name"] == "good.txt"
 
 
+def test_documents_endpoint_excludes_tombstoned_documents():
+    event = {
+        "httpMethod": "GET",
+        "path": "/documents",
+    }
+
+    with patch("lambda_handler.session_store") as mock_store, patch(
+        "lambda_handler._is_document_tombstoned", return_value=True
+    ):
+        mock_store.list_by_user.return_value = {
+            "sessions": [{"id": "session-1"}],
+            "next_cursor": None,
+        }
+        mock_session = MagicMock()
+        mock_session.is_hidden = False
+        mock_session.title = "Roadmap notes"
+        mock_session.updated_at = "2026-02-27T00:00:00"
+        mock_session.messages = [
+            {
+                "role": "user",
+                "attachments": [
+                    {
+                        "name": "paper.pdf",
+                        "type": "application/pdf",
+                        "s3_key": "uploads/user-1/paper.pdf",
+                        "s3_bucket": "bucket",
+                        "size_bytes": 1024,
+                    }
+                ],
+            }
+        ]
+        mock_store.get.return_value = mock_session
+
+        response = handler(event, None)
+
+    assert response["statusCode"] == 200
+    body = json.loads(response["body"])
+    assert body["success"] is True
+    assert body["documents"] == []
+
+
+def test_delete_document_success_path_returns_storage_snapshot():
+    event = {
+        "queryStringParameters": {
+            "s3_key": "uploads/user-1/docs/file.pdf",
+            "s3_bucket": "assets-bucket",
+        }
+    }
+    snapshot = {
+        "used_bytes": 100,
+        "reserved_bytes": 0,
+        "limit_bytes": 1000,
+        "available_bytes": 900,
+    }
+    mock_table = MagicMock()
+    mock_s3 = MagicMock()
+    mock_s3.head_object.return_value = {"ContentLength": 321}
+
+    with patch("lambda_handler._get_authenticated_user") as mock_auth, patch(
+        "lambda_handler._upload_key_belongs_to_user", return_value=True
+    ), patch("lambda_handler._ensure_user_storage_backfilled"), patch(
+        "lambda_handler._get_document_tombstone", return_value={}
+    ), patch(
+        "lambda_handler.get_upload_prep_table", return_value=mock_table
+    ), patch(
+        "lambda_handler.get_aws_client", return_value=mock_s3
+    ), patch(
+        "lambda_handler._decrement_used_storage", return_value=True
+    ) as mock_decrement, patch(
+        "lambda_handler._load_user_storage_snapshot", return_value=snapshot
+    ):
+        mock_auth.return_value = lambda_handler.AuthenticatedUser(user_id="user-1")
+        response = lambda_handler.handle_delete_document(event)
+
+    assert response["statusCode"] == 200
+    body = json.loads(response["body"])
+    assert body["success"] is True
+    assert body["deleted"] is True
+    assert body["storage"] == snapshot
+    mock_s3.delete_object.assert_called_once_with(Bucket="assets-bucket", Key="uploads/user-1/docs/file.pdf")
+    mock_decrement.assert_called_once_with("user-1", used_size_bytes=321)
+    mock_table.put_item.assert_called_once()
+    mock_table.update_item.assert_called_once()
+
+
+def test_delete_document_idempotent_when_already_deleted():
+    event = {
+        "queryStringParameters": {
+            "s3_key": "uploads/user-1/docs/file.pdf",
+            "s3_bucket": "assets-bucket",
+        }
+    }
+    snapshot = {
+        "used_bytes": 10,
+        "reserved_bytes": 0,
+        "limit_bytes": 1000,
+        "available_bytes": 990,
+    }
+
+    with patch("lambda_handler._get_authenticated_user") as mock_auth, patch(
+        "lambda_handler._upload_key_belongs_to_user", return_value=True
+    ), patch(
+        "lambda_handler._get_document_tombstone", return_value={"status": "deleted"}
+    ), patch(
+        "lambda_handler._load_user_storage_snapshot", return_value=snapshot
+    ), patch("lambda_handler.get_upload_prep_table") as mock_get_table:
+        mock_auth.return_value = lambda_handler.AuthenticatedUser(user_id="user-1")
+        response = lambda_handler.handle_delete_document(event)
+
+    assert response["statusCode"] == 200
+    body = json.loads(response["body"])
+    assert body["success"] is True
+    assert body["deleted"] is False
+    assert body["storage"] == snapshot
+    mock_get_table.assert_not_called()
+
+
+def test_delete_document_forbidden_on_non_owned_key():
+    event = {
+        "queryStringParameters": {
+            "s3_key": "legacy/other-user/file.pdf",
+            "s3_bucket": "assets-bucket",
+        }
+    }
+
+    with patch("lambda_handler._get_authenticated_user") as mock_auth, patch(
+        "lambda_handler._upload_key_belongs_to_user", return_value=False
+    ):
+        mock_auth.return_value = lambda_handler.AuthenticatedUser(user_id="user-1")
+        response = lambda_handler.handle_delete_document(event)
+
+    assert response["statusCode"] == 403
+    body = json.loads(response["body"])
+    assert body["success"] is False
+
+
 def test_images_invalid_limit_post_body():
     """Invalid POST limit should return 400."""
     event = {
@@ -2002,6 +2300,27 @@ def test_documents_download_prefers_session_attachment_bucket():
     assert kwargs["Params"]["Bucket"] == "correct-bucket"
 
 
+def test_documents_download_returns_410_when_document_tombstoned():
+    event = {
+        "queryStringParameters": {
+            "s3_key": "uploads/user-1/example.pdf",
+            "s3_bucket": "assets-bucket",
+            "name": "example.pdf",
+        }
+    }
+
+    with patch("lambda_handler._get_authenticated_user") as mock_auth, patch(
+        "lambda_handler._upload_key_belongs_to_user", return_value=True
+    ), patch("lambda_handler._is_document_tombstoned", return_value=True):
+        mock_auth.return_value = lambda_handler.AuthenticatedUser(user_id="user-1")
+        response = lambda_handler.handle_get_document_download_url(event)
+
+    assert response["statusCode"] == 410
+    body = json.loads(response["body"])
+    assert body["success"] is False
+    assert "deleted" in body["error"].lower()
+
+
 def test_frontend_events_ingestion_accepts_valid_batch(monkeypatch):
     monkeypatch.setenv("FRONTEND_EVENTS_TABLE", "fpai-frontend-events-test")
     lambda_handler._frontend_events_table = None

backend/tests/test_lambda_handler.py

@@ -75,6 +75,7 @@ function ShellInner({ children }: { children: React.ReactNode }) {
           memory: memoryRes.memory,
           hasMemory: memoryRes.hasMemory,
           tokenUsage: statusRes.usage,
+          storageUsage: statusRes.storage,
         })
         setMemorySnapshot(memoryRes.memory)
         setOnboardingRequired(!hasCompletedOnboarding(memoryRes.memory))