Initial curated RAGcore source from c16667f0c9397e2bfc8dd0255b114fba23c8d05e

tests/security/rate_limits/test_concurrency_budget.py

@@ -0,0 +1,78 @@
+"""ConcurrencyBudget: overload is bounded, not silently queued forever."""
+
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from ragcore.application.limits import ConcurrencyBudget, ConcurrencyExceededError
+
+
+async def test_overload_is_bounded_not_queued_without_limit() -> None:
+    budget = ConcurrencyBudget(name="model-calls", max_concurrent=1, max_queued=1)
+    release = asyncio.Event()
+
+    async def occupy() -> None:
+        async with budget.acquire():
+            await release.wait()
+
+    first = asyncio.create_task(occupy())
+    await asyncio.sleep(0)
+
+    second = asyncio.create_task(occupy())
+    await asyncio.sleep(0)
+    assert budget.waiting == 1
+
+    with pytest.raises(ConcurrencyExceededError):
+        async with budget.acquire():
+            pass  # pragma: no cover -- must never be entered
+
+    release.set()
+    await first
+    await second
+
+
+async def test_a_call_the_semaphore_can_serve_immediately_is_never_queued() -> None:
+    budget = ConcurrencyBudget(name="model-calls", max_concurrent=1, max_queued=0)
+    async with budget.acquire():
+        pass  # the only slot is free, so this must succeed even with max_queued=0
+    async with budget.acquire():
+        pass
+
+
+async def test_cancellation_while_waiting_releases_the_queue_slot() -> None:
+    budget = ConcurrencyBudget(name="model-calls", max_concurrent=1, max_queued=1)
+    release = asyncio.Event()
+
+    async def occupy() -> None:
+        async with budget.acquire():
+            await release.wait()
+
+    holder = asyncio.create_task(occupy())
+    await asyncio.sleep(0)
+
+    async def waiter() -> None:
+        async with budget.acquire():
+            pass  # pragma: no cover -- cancelled before ever acquiring
+
+    waiting_task = asyncio.create_task(waiter())
+    await asyncio.sleep(0)
+    assert budget.waiting == 1
+
+    waiting_task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await waiting_task
+    assert budget.waiting == 0
+
+    release.set()
+    await holder
+
+
+def test_invalid_construction_arguments_are_rejected() -> None:
+    with pytest.raises(ValueError, match="name"):
+        ConcurrencyBudget(name="", max_concurrent=1, max_queued=0)
+    with pytest.raises(ValueError, match="max_concurrent"):
+        ConcurrencyBudget(name="x", max_concurrent=0, max_queued=0)
+    with pytest.raises(ValueError, match="max_queued"):
+        ConcurrencyBudget(name="x", max_concurrent=1, max_queued=-1)

tests/security/rate_limits/test_rate_limit_middleware.py

@@ -0,0 +1,125 @@
+"""RateLimitMiddleware: operation-scoped limits that survive an alternate endpoint."""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime
+
+from fastapi import FastAPI, Request
+from fastapi.testclient import TestClient
+
+from ragcore.api.middleware import RateLimitMiddleware
+from ragcore.application.audit import AuditEventService
+from ragcore.application.limits import RateLimitRegistry, RateLimitRule
+from ragcore.application.limits.rate_limiter import RateLimiter
+
+
+class _NullAuditRepository:
+    async def append(self, event: object) -> None:  # pragma: no cover -- unused here
+        raise AssertionError("no admin overrides are exercised in these tests")
+
+    async def list_events(self, *args: object, **kwargs: object) -> tuple[object, ...]:
+        return ()
+
+
+def _fixed_subject(request: Request) -> str:
+    return request.headers.get("x-test-subject", "anonymous")
+
+
+def _app(*, rules: dict[str, RateLimitRule]) -> tuple[FastAPI, RateLimiter]:
+    audit = AuditEventService(_NullAuditRepository(), clock=lambda: datetime.now(UTC))
+    registry = RateLimitRegistry(default_rules=rules, audit=audit)
+    limiter = RateLimiter(registry=registry)
+
+    app = FastAPI()
+
+    @app.get("/v1/search", name="search")
+    async def search_route() -> dict[str, str]:
+        return {"endpoint": "search"}
+
+    @app.get("/v1/search-alias", name="search_alias")
+    async def search_alias_route() -> dict[str, str]:
+        return {"endpoint": "search_alias"}
+
+    @app.get("/v1/context", name="context")
+    async def context_route() -> dict[str, str]:
+        return {"endpoint": "context"}
+
+    app.add_middleware(
+        RateLimitMiddleware,
+        limiter=limiter,
+        route_operations={
+            "search": "search_operation",
+            "search_alias": "search_operation",
+            "context": "context_operation",
+        },
+        subject_resolver=_fixed_subject,
+    )
+    return app, limiter
+
+
+def test_limit_cannot_be_bypassed_by_an_alternate_endpoint_for_the_same_operation() -> None:
+    app, _ = _app(rules={"search_operation": RateLimitRule(max_requests=2, window_seconds=60)})
+    client = TestClient(app)
+    headers = {"x-test-subject": "alice"}
+
+    assert client.get("/v1/search", headers=headers).status_code == 200
+    assert client.get("/v1/search-alias", headers=headers).status_code == 200
+    # A third call to a *different* route mapped to the *same* operation
+    # still hits the shared budget -- the limit cannot be dodged by calling
+    # whichever endpoint happens to still have quota left.
+    response = client.get("/v1/search", headers=headers)
+    assert response.status_code == 429
+
+
+def test_a_429_carries_a_retry_after_header_and_a_stable_json_body() -> None:
+    app, _ = _app(rules={"search_operation": RateLimitRule(max_requests=1, window_seconds=60)})
+    client = TestClient(app)
+    headers = {"x-test-subject": "alice"}
+
+    client.get("/v1/search", headers=headers)
+    response = client.get("/v1/search", headers=headers)
+
+    assert response.status_code == 429
+    assert "retry-after" in response.headers
+    assert int(response.headers["retry-after"]) >= 1
+    body = response.json()
+    assert body["code"] == "RATE_LIMITED"
+    assert body["status"] == 429
+    assert body["retry_after_seconds"] >= 1
+
+
+def test_unrelated_operations_are_not_affected_by_an_exhausted_budget() -> None:
+    app, _ = _app(
+        rules={
+            "search_operation": RateLimitRule(max_requests=1, window_seconds=60),
+            "context_operation": RateLimitRule(max_requests=1, window_seconds=60),
+        }
+    )
+    client = TestClient(app)
+    headers = {"x-test-subject": "alice"}
+
+    client.get("/v1/search", headers=headers)
+    assert client.get("/v1/search", headers=headers).status_code == 429
+    # The context operation has its own, independent budget.
+    assert client.get("/v1/context", headers=headers).status_code == 200
+
+
+def test_different_subjects_have_independent_budgets_through_the_middleware() -> None:
+    app, _ = _app(rules={"search_operation": RateLimitRule(max_requests=1, window_seconds=60)})
+    client = TestClient(app)
+
+    assert client.get("/v1/search", headers={"x-test-subject": "alice"}).status_code == 200
+    assert client.get("/v1/search", headers={"x-test-subject": "alice"}).status_code == 429
+    assert client.get("/v1/search", headers={"x-test-subject": "bob"}).status_code == 200
+
+
+def test_a_route_with_no_configured_operation_is_never_rate_limited() -> None:
+    app, _ = _app(rules={"search_operation": RateLimitRule(max_requests=1, window_seconds=60)})
+
+    @app.get("/unmapped")
+    async def unmapped_route() -> dict[str, str]:
+        return {"ok": "true"}
+
+    client = TestClient(app)
+    for _ in range(5):
+        assert client.get("/unmapped").status_code == 200

tests/security/rate_limits/test_rate_limiter.py

@@ -0,0 +1,203 @@
+"""RateLimiter / RateLimitRegistry: fixed windows, retry hints, audited overrides."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable
+from datetime import UTC, datetime
+from uuid import UUID, uuid4
+
+import pytest
+
+from ragcore.application.audit import AuditEventService
+from ragcore.application.limits import (
+    RateLimiter,
+    RateLimitExceededError,
+    RateLimitRegistry,
+    RateLimitRule,
+)
+from ragcore.domain.audit import AuditEvent, AuditEventType
+
+WORKSPACE_ID = UUID("00000000-0000-0000-0000-000000000001")
+
+
+class _InMemoryAuditRepository:
+    def __init__(self) -> None:
+        self.events: list[AuditEvent] = []
+
+    async def append(self, event: AuditEvent) -> None:
+        self.events.append(event)
+
+    async def list_events(
+        self,
+        workspace_id: UUID,
+        *,
+        event_types: Iterable[AuditEventType] = (),
+        actor_id: UUID | None = None,
+        application_id: UUID | None = None,
+        before: datetime | None = None,
+        limit: int = 100,
+    ) -> tuple[AuditEvent, ...]:
+        return tuple(self.events)
+
+
+def _clock(value: dict[str, float]) -> Callable[[], float]:
+    return lambda: value["now"]
+
+
+def _registry(
+    *, default_rules: dict[str, RateLimitRule] | None = None
+) -> tuple[RateLimitRegistry, _InMemoryAuditRepository]:
+    repository = _InMemoryAuditRepository()
+    audit = AuditEventService(repository, clock=lambda: datetime.now(UTC))
+    return RateLimitRegistry(default_rules=default_rules or {}, audit=audit), repository
+
+
+def test_no_rule_configured_means_unlimited() -> None:
+    registry, _ = _registry()
+    limiter = RateLimiter(registry=registry)
+    for _ in range(1000):
+        limiter.check(subject="alice", operation="unconfigured")
+
+
+def test_requests_within_the_window_are_allowed() -> None:
+    registry, _ = _registry(
+        default_rules={"search": RateLimitRule(max_requests=3, window_seconds=60)}
+    )
+    limiter = RateLimiter(registry=registry)
+    for _ in range(3):
+        limiter.check(subject="alice", operation="search")
+
+
+def test_exceeding_the_window_raises_with_a_positive_retry_after() -> None:
+    registry, _ = _registry(
+        default_rules={"search": RateLimitRule(max_requests=2, window_seconds=60)}
+    )
+    limiter = RateLimiter(registry=registry)
+    limiter.check(subject="alice", operation="search")
+    limiter.check(subject="alice", operation="search")
+    with pytest.raises(RateLimitExceededError) as excinfo:
+        limiter.check(subject="alice", operation="search")
+    assert excinfo.value.retry_after_seconds > 0
+
+
+def test_the_window_resets_after_it_elapses() -> None:
+    time_value = {"now": 0.0}
+    registry, _ = _registry(
+        default_rules={"search": RateLimitRule(max_requests=1, window_seconds=10)}
+    )
+    limiter = RateLimiter(registry=registry, clock=_clock(time_value))
+
+    limiter.check(subject="alice", operation="search")
+    with pytest.raises(RateLimitExceededError):
+        limiter.check(subject="alice", operation="search")
+
+    time_value["now"] = 10.1
+    limiter.check(subject="alice", operation="search")  # a fresh window, must not raise
+
+
+def test_different_subjects_have_independent_budgets() -> None:
+    registry, _ = _registry(
+        default_rules={"search": RateLimitRule(max_requests=1, window_seconds=60)}
+    )
+    limiter = RateLimiter(registry=registry)
+
+    limiter.check(subject="alice", operation="search")
+    limiter.check(subject="bob", operation="search")  # must not raise -- separate subject
+
+
+def test_different_operations_have_independent_budgets() -> None:
+    registry, _ = _registry(
+        default_rules={
+            "search": RateLimitRule(max_requests=1, window_seconds=60),
+            "upload": RateLimitRule(max_requests=1, window_seconds=60),
+        }
+    )
+    limiter = RateLimiter(registry=registry)
+
+    limiter.check(subject="alice", operation="search")
+    limiter.check(subject="alice", operation="upload")  # must not raise -- separate operation
+
+
+async def test_admin_override_replaces_the_default_rule_and_is_audited() -> None:
+    registry, repository = _registry(
+        default_rules={"search": RateLimitRule(max_requests=1, window_seconds=60)}
+    )
+    limiter = RateLimiter(registry=registry)
+    limiter.check(subject="alice", operation="search")
+    with pytest.raises(RateLimitExceededError):
+        limiter.check(subject="alice", operation="search")
+
+    await registry.set_override(
+        subject="alice",
+        operation="search",
+        rule=RateLimitRule(max_requests=5, window_seconds=60),
+        actor_id=uuid4(),
+        workspace_id=WORKSPACE_ID,
+        trace_id="trace-1",
+    )
+
+    limiter.check(subject="alice", operation="search")  # the raised limit now allows this
+
+    assert len(repository.events) == 1
+    event = repository.events[0]
+    assert event.event_type is AuditEventType.ADMIN_ACTION
+    assert event.resource_type == "rate_limit_override"
+    assert event.details["subject"] == "alice"
+    assert event.details["max_requests"] == 5
+
+
+async def test_override_only_affects_the_named_subject() -> None:
+    registry, _ = _registry(
+        default_rules={"search": RateLimitRule(max_requests=1, window_seconds=60)}
+    )
+    limiter = RateLimiter(registry=registry)
+
+    await registry.set_override(
+        subject="alice",
+        operation="search",
+        rule=RateLimitRule(max_requests=5, window_seconds=60),
+        actor_id=uuid4(),
+        workspace_id=WORKSPACE_ID,
+        trace_id="trace-1",
+    )
+
+    limiter.check(subject="bob", operation="search")
+    with pytest.raises(RateLimitExceededError):
+        limiter.check(subject="bob", operation="search")  # bob still has the default limit
+
+
+async def test_clear_override_restores_the_default_rule_and_is_audited() -> None:
+    registry, repository = _registry(
+        default_rules={"search": RateLimitRule(max_requests=1, window_seconds=60)}
+    )
+    await registry.set_override(
+        subject="alice",
+        operation="search",
+        rule=RateLimitRule(max_requests=5, window_seconds=60),
+        actor_id=uuid4(),
+        workspace_id=WORKSPACE_ID,
+        trace_id="trace-1",
+    )
+
+    await registry.clear_override(
+        subject="alice",
+        operation="search",
+        actor_id=uuid4(),
+        workspace_id=WORKSPACE_ID,
+        trace_id="trace-2",
+    )
+
+    limiter = RateLimiter(registry=registry)
+    limiter.check(subject="alice", operation="search")
+    with pytest.raises(RateLimitExceededError):
+        limiter.check(subject="alice", operation="search")  # back to the default limit of 1
+
+    assert len(repository.events) == 2
+    assert repository.events[-1].details["cleared"] is True
+
+
+def test_invalid_rule_construction_is_rejected() -> None:
+    with pytest.raises(ValueError, match="max_requests"):
+        RateLimitRule(max_requests=0, window_seconds=60)
+    with pytest.raises(ValueError, match="window_seconds"):
+        RateLimitRule(max_requests=1, window_seconds=0)