From bc7cac78ae0b2573eff7bcfcdaed20ffc8efa620 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 8 Feb 2026 08:43:24 +0000 Subject: [PATCH 1/2] Add user-specific rate limits with admin API management Introduces per-user submission rate limiting (hourly and daily caps) configurable via the admin API. Users without an override are unrestricted. - New migration: leaderboard.user_rate_limits table - DB methods: get/set/delete user rate limits + submission rate checking - Admin endpoints: GET/PUT/DELETE /admin/rate-limits/{user_id} - Both submission endpoints check user rate limits (return 429 if exceeded) - Tests for DB methods and all API endpoints https://claude.ai/code/session_01LiSWKfcKe4riGZwMYdubiJ --- src/kernelbot/api/main.py | 91 ++++++++ src/libkernelbot/leaderboard_db.py | 212 ++++++++++++++++++ .../20260208_01_RkLmN-add-user-rate-limits.py | 27 +++ tests/conftest.py | 2 +- tests/test_admin_api.py | 132 +++++++++++ tests/test_leaderboard_db.py | 118 ++++++++++ 6 files changed, 581 insertions(+), 1 deletion(-) create mode 100644 src/migrations/20260208_01_RkLmN-add-user-rate-limits.py diff --git a/src/kernelbot/api/main.py b/src/kernelbot/api/main.py index 2ae2bf97..bced59b0 100644 --- a/src/kernelbot/api/main.py +++ b/src/kernelbot/api/main.py @@ -388,6 +388,13 @@ async def run_submission( # noqa: C901 StreamingResponse: A streaming response containing the status and results of the submission. """ await simple_rate_limit() + + # Check user-specific rate limits + with db_context as db: + rate_check = db.check_user_submission_rate(user_info["user_id"]) + if not rate_check["allowed"]: + raise HTTPException(status_code=429, detail=f"Rate limit exceeded. {rate_check['retry_after']}") + submission_request, submission_mode_enum = await to_submit_info( user_info, submission_mode, file, leaderboard_name, gpu_type, db_context ) @@ -449,6 +456,11 @@ async def run_submission_async( await simple_rate_limit() logger.info(f"Received submission request for {leaderboard_name} {gpu_type} {submission_mode}") + # Check user-specific rate limits + with db_context as db: + rate_check = db.check_user_submission_rate(user_info["user_id"]) + if not rate_check["allowed"]: + raise HTTPException(status_code=429, detail=f"Rate limit exceeded. {rate_check['retry_after']}") # throw error if submission request is invalid try: @@ -643,6 +655,85 @@ async def admin_update_problems( } +@app.get("/admin/rate-limits") +async def get_all_rate_limits( + _: Annotated[None, Depends(require_admin)], + db_context=Depends(get_db), +) -> dict: + """Get all user rate limit overrides.""" + with db_context as db: + rate_limits = db.get_all_user_rate_limits() + return {"status": "ok", "rate_limits": rate_limits} + + +@app.get("/admin/rate-limits/{user_id}") +async def get_user_rate_limit( + user_id: str, + _: Annotated[None, Depends(require_admin)], + db_context=Depends(get_db), +) -> dict: + """Get rate limit for a specific user.""" + with db_context as db: + rate_limit = db.get_user_rate_limit(user_id) + if rate_limit is None: + raise HTTPException(status_code=404, detail="No rate limit override found for this user") + return {"status": "ok", "rate_limit": rate_limit} + + +@app.put("/admin/rate-limits/{user_id}") +async def set_user_rate_limit( + user_id: str, + payload: dict, + _: Annotated[None, Depends(require_admin)], + db_context=Depends(get_db), +) -> dict: + """Set or update rate limit for a user. + + Payload fields: + max_submissions_per_hour (int, optional): Max submissions per hour + max_submissions_per_day (int, optional): Max submissions per day + note (str, optional): Admin note about why the limit was set + """ + max_per_hour = payload.get("max_submissions_per_hour") + max_per_day = payload.get("max_submissions_per_day") + note = payload.get("note") + + if max_per_hour is None and max_per_day is None: + raise HTTPException( + status_code=400, + detail="At least one of max_submissions_per_hour or max_submissions_per_day is required", + ) + + if max_per_hour is not None and (not isinstance(max_per_hour, int) or max_per_hour < 0): + raise HTTPException(status_code=400, detail="max_submissions_per_hour must be a non-negative integer") + + if max_per_day is not None and (not isinstance(max_per_day, int) or max_per_day < 0): + raise HTTPException(status_code=400, detail="max_submissions_per_day must be a non-negative integer") + + with db_context as db: + rate_limit = db.set_user_rate_limit( + user_id=user_id, + max_submissions_per_hour=max_per_hour, + max_submissions_per_day=max_per_day, + note=note, + ) + return {"status": "ok", "rate_limit": rate_limit} + + +@app.delete("/admin/rate-limits/{user_id}") +async def delete_user_rate_limit( + user_id: str, + _: Annotated[None, Depends(require_admin)], + db_context=Depends(get_db), +) -> dict: + """Delete a user's rate limit override.""" + with db_context as db: + deleted = db.delete_user_rate_limit(user_id) + if not deleted: + raise HTTPException(status_code=404, detail="No rate limit override found for this user") + return {"status": "ok", "user_id": user_id} + + @app.get("/leaderboards") async def get_leaderboards(db_context=Depends(get_db)): """An endpoint that returns all leaderboards. diff --git a/src/libkernelbot/leaderboard_db.py b/src/libkernelbot/leaderboard_db.py index 334ad633..1cd08174 100644 --- a/src/libkernelbot/leaderboard_db.py +++ b/src/libkernelbot/leaderboard_db.py @@ -1172,6 +1172,218 @@ def cleanup_temp_users(self): logger.exception("Could not cleanup temp users", exc_info=e) raise KernelBotError("Database error while cleaning up temp users") from e + def get_user_rate_limit(self, user_id: str) -> Optional[dict]: + """ + Get the rate limit override for a specific user. + + Returns: + Optional[dict]: Rate limit info or None if no override exists. + """ + try: + self.cursor.execute( + """ + SELECT user_id, max_submissions_per_hour, max_submissions_per_day, + note, created_at, updated_at + FROM leaderboard.user_rate_limits + WHERE user_id = %s + """, + (user_id,), + ) + row = self.cursor.fetchone() + if row is None: + return None + return { + "user_id": row[0], + "max_submissions_per_hour": row[1], + "max_submissions_per_day": row[2], + "note": row[3], + "created_at": row[4], + "updated_at": row[5], + } + except psycopg2.Error as e: + self.connection.rollback() + logger.exception("Error fetching rate limit for user %s", user_id, exc_info=e) + raise KernelBotError("Error fetching user rate limit") from e + + def get_all_user_rate_limits(self) -> list[dict]: + """ + Get all user rate limit overrides. + + Returns: + list[dict]: All user rate limit entries. + """ + try: + self.cursor.execute( + """ + SELECT rl.user_id, rl.max_submissions_per_hour, rl.max_submissions_per_day, + rl.note, rl.created_at, rl.updated_at, ui.user_name + FROM leaderboard.user_rate_limits rl + LEFT JOIN leaderboard.user_info ui ON rl.user_id = ui.id + ORDER BY rl.updated_at DESC + """ + ) + return [ + { + "user_id": row[0], + "max_submissions_per_hour": row[1], + "max_submissions_per_day": row[2], + "note": row[3], + "created_at": row[4], + "updated_at": row[5], + "user_name": row[6], + } + for row in self.cursor.fetchall() + ] + except psycopg2.Error as e: + self.connection.rollback() + logger.exception("Error fetching all user rate limits", exc_info=e) + raise KernelBotError("Error fetching user rate limits") from e + + def set_user_rate_limit( + self, + user_id: str, + max_submissions_per_hour: Optional[int] = None, + max_submissions_per_day: Optional[int] = None, + note: Optional[str] = None, + ) -> dict: + """ + Set or update rate limit for a user (upsert). + + Returns: + dict: The created/updated rate limit entry. + """ + try: + self.cursor.execute( + """ + INSERT INTO leaderboard.user_rate_limits + (user_id, max_submissions_per_hour, max_submissions_per_day, note) + VALUES (%s, %s, %s, %s) + ON CONFLICT (user_id) DO UPDATE SET + max_submissions_per_hour = EXCLUDED.max_submissions_per_hour, + max_submissions_per_day = EXCLUDED.max_submissions_per_day, + note = EXCLUDED.note, + updated_at = NOW() + RETURNING user_id, max_submissions_per_hour, max_submissions_per_day, + note, created_at, updated_at + """, + (user_id, max_submissions_per_hour, max_submissions_per_day, note), + ) + row = self.cursor.fetchone() + self.connection.commit() + return { + "user_id": row[0], + "max_submissions_per_hour": row[1], + "max_submissions_per_day": row[2], + "note": row[3], + "created_at": row[4], + "updated_at": row[5], + } + except psycopg2.Error as e: + self.connection.rollback() + logger.exception("Error setting rate limit for user %s", user_id, exc_info=e) + raise KernelBotError("Error setting user rate limit") from e + + def delete_user_rate_limit(self, user_id: str) -> bool: + """ + Delete a user's rate limit override. + + Returns: + bool: True if a row was deleted, False if no override existed. + """ + try: + self.cursor.execute( + """ + DELETE FROM leaderboard.user_rate_limits + WHERE user_id = %s + """, + (user_id,), + ) + deleted = self.cursor.rowcount > 0 + self.connection.commit() + return deleted + except psycopg2.Error as e: + self.connection.rollback() + logger.exception("Error deleting rate limit for user %s", user_id, exc_info=e) + raise KernelBotError("Error deleting user rate limit") from e + + def check_user_submission_rate(self, user_id: str) -> dict: + """ + Check a user's current submission counts against their rate limits. + + Returns: + dict with keys: + - allowed: bool, whether the user can submit + - hourly_count: int, submissions in the last hour + - daily_count: int, submissions in the last day + - hourly_limit: int or None + - daily_limit: int or None + - retry_after: str or None, human-readable wait time if blocked + """ + try: + # Get user's rate limits (None means no override) + rate_limit = self.get_user_rate_limit(user_id) + if rate_limit is None: + return { + "allowed": True, + "hourly_count": 0, + "daily_count": 0, + "hourly_limit": None, + "daily_limit": None, + "retry_after": None, + } + + hourly_limit = rate_limit["max_submissions_per_hour"] + daily_limit = rate_limit["max_submissions_per_day"] + + # If both limits are None, user is unrestricted + if hourly_limit is None and daily_limit is None: + return { + "allowed": True, + "hourly_count": 0, + "daily_count": 0, + "hourly_limit": None, + "daily_limit": None, + "retry_after": None, + } + + # Count submissions in the last hour and day + self.cursor.execute( + """ + SELECT + COUNT(*) FILTER (WHERE submission_time > NOW() - INTERVAL '1 hour') AS hourly_count, + COUNT(*) FILTER (WHERE submission_time > NOW() - INTERVAL '1 day') AS daily_count + FROM leaderboard.submission + WHERE user_id = %s + """, + (user_id,), + ) + row = self.cursor.fetchone() + hourly_count = row[0] + daily_count = row[1] + + # Check limits + hourly_exceeded = hourly_limit is not None and hourly_count >= hourly_limit + daily_exceeded = daily_limit is not None and daily_count >= daily_limit + + retry_after = None + if hourly_exceeded: + retry_after = "Try again in up to 1 hour" + if daily_exceeded: + retry_after = "Try again in up to 24 hours" + + return { + "allowed": not (hourly_exceeded or daily_exceeded), + "hourly_count": hourly_count, + "daily_count": daily_count, + "hourly_limit": hourly_limit, + "daily_limit": daily_limit, + "retry_after": retry_after, + } + except psycopg2.Error as e: + self.connection.rollback() + logger.exception("Error checking rate limit for user %s", user_id, exc_info=e) + raise KernelBotError("Error checking user rate limit") from e + def validate_cli_id(self, cli_id: str) -> Optional[dict[str, str]]: """ Validates a CLI ID and returns the associated user ID if valid. diff --git a/src/migrations/20260208_01_RkLmN-add-user-rate-limits.py b/src/migrations/20260208_01_RkLmN-add-user-rate-limits.py new file mode 100644 index 00000000..8c334dec --- /dev/null +++ b/src/migrations/20260208_01_RkLmN-add-user-rate-limits.py @@ -0,0 +1,27 @@ +""" +add-user-rate-limits +""" + +from yoyo import step + +__depends__ = {'20260108_01_gzSm3-add-submission-status'} + +steps = [ + step( + # forward + """ + CREATE TABLE leaderboard.user_rate_limits ( + user_id TEXT PRIMARY KEY REFERENCES leaderboard.user_info(id), + max_submissions_per_hour INTEGER, + max_submissions_per_day INTEGER, + note TEXT, + created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(), + updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW() + ); + """, + # backward + """ + DROP TABLE leaderboard.user_rate_limits; + """ + ) +] diff --git a/tests/conftest.py b/tests/conftest.py index 1a049250..72a8c569 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -50,7 +50,7 @@ def _nuke_contents(db): db.cursor.execute( "TRUNCATE leaderboard.code_files, leaderboard.submission, leaderboard.runs, " "leaderboard.leaderboard, leaderboard.user_info, leaderboard.templates, " - "leaderboard.gpu_type RESTART IDENTITY CASCADE" + "leaderboard.gpu_type, leaderboard.user_rate_limits RESTART IDENTITY CASCADE" ) db.connection.commit() diff --git a/tests/test_admin_api.py b/tests/test_admin_api.py index ecd1b33c..dbfab612 100644 --- a/tests/test_admin_api.py +++ b/tests/test_admin_api.py @@ -405,3 +405,135 @@ def test_update_problems_with_errors(self, test_client, mock_backend): assert data["status"] == "ok" assert len(data["errors"]) == 1 assert data["errors"][0]["name"] == "bad-problem" + + +class TestAdminRateLimits: + """Test admin rate limit endpoints.""" + + def _setup_db_mock(self, mock_backend): + mock_backend.db.__enter__ = MagicMock(return_value=mock_backend.db) + mock_backend.db.__exit__ = MagicMock(return_value=None) + + def test_get_all_rate_limits(self, test_client, mock_backend): + """GET /admin/rate-limits returns all rate limits.""" + self._setup_db_mock(mock_backend) + mock_backend.db.get_all_user_rate_limits = MagicMock(return_value=[ + {"user_id": "123", "max_submissions_per_hour": 10, "max_submissions_per_day": 50, + "note": None, "created_at": None, "updated_at": None, "user_name": "testuser"}, + ]) + + response = test_client.get( + "/admin/rate-limits", + headers={"Authorization": "Bearer test_token"} + ) + assert response.status_code == 200 + data = response.json() + assert data["status"] == "ok" + assert len(data["rate_limits"]) == 1 + assert data["rate_limits"][0]["user_id"] == "123" + + def test_get_all_rate_limits_requires_auth(self, test_client): + """GET /admin/rate-limits requires authentication.""" + response = test_client.get("/admin/rate-limits") + assert response.status_code == 401 + + def test_get_user_rate_limit(self, test_client, mock_backend): + """GET /admin/rate-limits/{user_id} returns user's rate limit.""" + self._setup_db_mock(mock_backend) + mock_backend.db.get_user_rate_limit = MagicMock(return_value={ + "user_id": "123", "max_submissions_per_hour": 10, "max_submissions_per_day": 50, + "note": "heavy user", "created_at": None, "updated_at": None, + }) + + response = test_client.get( + "/admin/rate-limits/123", + headers={"Authorization": "Bearer test_token"} + ) + assert response.status_code == 200 + data = response.json() + assert data["status"] == "ok" + assert data["rate_limit"]["user_id"] == "123" + assert data["rate_limit"]["note"] == "heavy user" + + def test_get_user_rate_limit_not_found(self, test_client, mock_backend): + """GET /admin/rate-limits/{user_id} returns 404 for missing.""" + self._setup_db_mock(mock_backend) + mock_backend.db.get_user_rate_limit = MagicMock(return_value=None) + + response = test_client.get( + "/admin/rate-limits/999", + headers={"Authorization": "Bearer test_token"} + ) + assert response.status_code == 404 + + def test_set_user_rate_limit(self, test_client, mock_backend): + """PUT /admin/rate-limits/{user_id} sets rate limit.""" + self._setup_db_mock(mock_backend) + mock_backend.db.set_user_rate_limit = MagicMock(return_value={ + "user_id": "123", "max_submissions_per_hour": 5, "max_submissions_per_day": 20, + "note": "restricted", "created_at": None, "updated_at": None, + }) + + response = test_client.put( + "/admin/rate-limits/123", + headers={"Authorization": "Bearer test_token"}, + json={"max_submissions_per_hour": 5, "max_submissions_per_day": 20, "note": "restricted"}, + ) + assert response.status_code == 200 + data = response.json() + assert data["status"] == "ok" + assert data["rate_limit"]["max_submissions_per_hour"] == 5 + mock_backend.db.set_user_rate_limit.assert_called_once_with( + user_id="123", + max_submissions_per_hour=5, + max_submissions_per_day=20, + note="restricted", + ) + + def test_set_user_rate_limit_requires_at_least_one_limit(self, test_client, mock_backend): + """PUT /admin/rate-limits/{user_id} requires at least one limit field.""" + self._setup_db_mock(mock_backend) + + response = test_client.put( + "/admin/rate-limits/123", + headers={"Authorization": "Bearer test_token"}, + json={"note": "just a note"}, + ) + assert response.status_code == 400 + assert "At least one of" in response.json()["detail"] + + def test_set_user_rate_limit_validates_negative(self, test_client, mock_backend): + """PUT /admin/rate-limits/{user_id} rejects negative values.""" + self._setup_db_mock(mock_backend) + + response = test_client.put( + "/admin/rate-limits/123", + headers={"Authorization": "Bearer test_token"}, + json={"max_submissions_per_hour": -1}, + ) + assert response.status_code == 400 + assert "non-negative integer" in response.json()["detail"] + + def test_delete_user_rate_limit(self, test_client, mock_backend): + """DELETE /admin/rate-limits/{user_id} deletes rate limit.""" + self._setup_db_mock(mock_backend) + mock_backend.db.delete_user_rate_limit = MagicMock(return_value=True) + + response = test_client.delete( + "/admin/rate-limits/123", + headers={"Authorization": "Bearer test_token"} + ) + assert response.status_code == 200 + assert response.json()["status"] == "ok" + assert response.json()["user_id"] == "123" + + def test_delete_user_rate_limit_not_found(self, test_client, mock_backend): + """DELETE /admin/rate-limits/{user_id} returns 404 when not found.""" + self._setup_db_mock(mock_backend) + mock_backend.db.delete_user_rate_limit = MagicMock(return_value=False) + + response = test_client.delete( + "/admin/rate-limits/999", + headers={"Authorization": "Bearer test_token"} + ) + assert response.status_code == 404 diff --git a/tests/test_leaderboard_db.py b/tests/test_leaderboard_db.py index 1b349816..74c56f2b 100644 --- a/tests/test_leaderboard_db.py +++ b/tests/test_leaderboard_db.py @@ -775,3 +775,121 @@ def test_get_user_submissions_with_multiple_runs(database, submit_leaderboard): assert 1.5 in scores assert 2.0 in scores + +def test_user_rate_limit_crud(database, submit_leaderboard): + """Test basic CRUD operations for user rate limits.""" + with database as db: + # Create a user first + db.cursor.execute( + "INSERT INTO leaderboard.user_info (id, user_name) VALUES (%s, %s)", + ("rate-user-1", "rate_user"), + ) + db.connection.commit() + + # No rate limit initially + assert db.get_user_rate_limit("rate-user-1") is None + assert db.get_all_user_rate_limits() == [] + + # Set rate limit + result = db.set_user_rate_limit("rate-user-1", max_submissions_per_hour=10, max_submissions_per_day=50) + assert result["user_id"] == "rate-user-1" + assert result["max_submissions_per_hour"] == 10 + assert result["max_submissions_per_day"] == 50 + assert result["note"] is None + + # Read it back + rl = db.get_user_rate_limit("rate-user-1") + assert rl["user_id"] == "rate-user-1" + assert rl["max_submissions_per_hour"] == 10 + assert rl["max_submissions_per_day"] == 50 + + # List all + all_limits = db.get_all_user_rate_limits() + assert len(all_limits) == 1 + assert all_limits[0]["user_name"] == "rate_user" + + # Update (upsert) + updated = db.set_user_rate_limit("rate-user-1", max_submissions_per_hour=5, note="reduced limit") + assert updated["max_submissions_per_hour"] == 5 + assert updated["max_submissions_per_day"] is None + assert updated["note"] == "reduced limit" + + # Delete + assert db.delete_user_rate_limit("rate-user-1") is True + assert db.get_user_rate_limit("rate-user-1") is None + + # Delete non-existent returns False + assert db.delete_user_rate_limit("rate-user-1") is False + + +def test_user_rate_limit_no_override(database, submit_leaderboard): + """Users without a rate limit override are always allowed.""" + with database as db: + result = db.check_user_submission_rate("nonexistent-user") + assert result["allowed"] is True + + +def test_user_rate_limit_enforcement(database, submit_leaderboard): + """Test that rate limit enforcement blocks users who exceed limits.""" + with database as db: + # Create a user + db.cursor.execute( + "INSERT INTO leaderboard.user_info (id, user_name) VALUES (%s, %s)", + ("limited-user", "limited"), + ) + db.connection.commit() + + # Set a tight limit: 2 per hour + db.set_user_rate_limit("limited-user", max_submissions_per_hour=2, max_submissions_per_day=100) + + # No submissions yet, should be allowed + check = db.check_user_submission_rate("limited-user") + assert check["allowed"] is True + assert check["hourly_count"] == 0 + assert check["hourly_limit"] == 2 + + # Create 2 submissions + for i in range(2): + db.create_submission( + "submit-leaderboard", + f"file_{i}.py", + "limited-user", + f"code {i}", + datetime.datetime.now(tz=datetime.timezone.utc), + user_name="limited", + ) + + # Now should be blocked + check = db.check_user_submission_rate("limited-user") + assert check["allowed"] is False + assert check["hourly_count"] == 2 + assert check["retry_after"] is not None + + +def test_user_rate_limit_daily(database, submit_leaderboard): + """Test daily rate limit enforcement.""" + with database as db: + db.cursor.execute( + "INSERT INTO leaderboard.user_info (id, user_name) VALUES (%s, %s)", + ("daily-user", "daily"), + ) + db.connection.commit() + + # Set a daily limit only + db.set_user_rate_limit("daily-user", max_submissions_per_day=1) + + # Create 1 submission + db.create_submission( + "submit-leaderboard", + "file.py", + "daily-user", + "code", + datetime.datetime.now(tz=datetime.timezone.utc), + user_name="daily", + ) + + check = db.check_user_submission_rate("daily-user") + assert check["allowed"] is False + assert check["daily_count"] == 1 + assert check["daily_limit"] == 1 + From cf49a7ac14774d3dc7ae278539491ccb61a398be Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 8 Feb 2026 08:44:16 +0000 Subject: [PATCH 2/2] Add *.egg-info/ to .gitignore https://claude.ai/code/session_01LiSWKfcKe4riGZwMYdubiJ --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 5c184087..885eaf40 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,4 @@ reference-kernels/ yoyo.ini .venv .claude/ +*.egg-info/