fischer-agentkit/tests/integration/admin/test_quota_enforcement.py

"""Quota enforcement integration tests (U10).

Verifies quota enforcement end-to-end through the LLMGateway:

- Token limit exceeded → QuotaExceededError raised.
- Cost limit exceeded → QuotaExceededError raised.
- Model not in whitelist → QuotaExceededError raised.
- No quota set → request allowed.
- Multi-department: strictest-wins (one exceeds, other doesn't → rejected).
- Integration test with real LLMGateway + mock provider + InMemoryUsageStore.

These tests use a real :class:`LLMGateway` with a :class:`FakeProvider`
(mock LLM provider) and a real :class:`QuotaService` backed by a temp
SQLite auth DB. No external services (Redis, real LLM API) are required.
"""

from __future__ import annotations

import uuid
from pathlib import Path

import pytest

from agentkit.llm.gateway import LLMGateway, QuotaExceededError
from agentkit.llm.protocol import (
    LLMProvider,
    LLMRequest,
    LLMResponse,
    TokenUsage,
)
from agentkit.llm.providers.usage_store import InMemoryUsageStore
from agentkit.server.admin.quota_service import (
    get_quota_service,
    set_quota_service,
)
from agentkit.server.auth.models import init_auth_db


# ---------------------------------------------------------------------------
# Test doubles
# ---------------------------------------------------------------------------


class FakeProvider(LLMProvider):
    """A minimal LLMProvider that returns a fixed response.

    The response usage (prompt_tokens, completion_tokens) can be
    customized per-instance to simulate different token consumption.
    """

    def __init__(
        self,
        name: str = "fake",
        prompt_tokens: int = 100,
        completion_tokens: int = 50,
    ) -> None:
        self._name = name
        self._prompt_tokens = prompt_tokens
        self._completion_tokens = completion_tokens
        self.last_request: LLMRequest | None = None
        self.call_count = 0

    async def chat(self, request: LLMRequest) -> LLMResponse:
        self.last_request = request
        self.call_count += 1
        return LLMResponse(
            content=f"response from {self._name}",
            model=request.model,
            usage=TokenUsage(
                prompt_tokens=self._prompt_tokens,
                completion_tokens=self._completion_tokens,
            ),
        )


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture
def store() -> InMemoryUsageStore:
    return InMemoryUsageStore()


@pytest.fixture
def gateway(store: InMemoryUsageStore) -> LLMGateway:
    """A real LLMGateway with a FakeProvider registered as "openai"."""
    gw = LLMGateway(usage_store=store)
    gw.register_provider("openai", FakeProvider("openai"))
    return gw


@pytest.fixture
async def fresh_db(tmp_path: Path) -> Path:
    """A brand-new auth DB on a fresh path (no data)."""
    db_path = tmp_path / "quota_enforcement.db"
    await init_auth_db(db_path)
    return db_path


@pytest.fixture(autouse=True)
def _reset_quota_singleton():
    """Reset the QuotaService singleton before and after each test."""
    set_quota_service(None)
    yield
    set_quota_service(None)


def _random_dept_id() -> str:
    return str(uuid.uuid4())


# ---------------------------------------------------------------------------
# Quota enforcement tests
# ---------------------------------------------------------------------------


class TestQuotaEnforcement:
    """Tests for quota enforcement in LLM calls."""

    async def test_token_limit_blocks_request(
        self, gateway: LLMGateway, store: InMemoryUsageStore, fresh_db: Path
    ):
        """When department exceeds token limit, LLM call raises QuotaExceededError."""
        dept_id = _random_dept_id()
        svc = get_quota_service()
        # Set a 100-token daily limit.
        await svc.set_quota(fresh_db, dept_id, "token_limit", 100, period="daily")

        # Pre-populate usage with 90 tokens (just under the limit).
        gateway._usage_tracker.record(
            agent_name="prev",
            model="openai/gpt-4o",
            usage=TokenUsage(prompt_tokens=60, completion_tokens=30),
            cost=0.0,
            latency_ms=10,
            user_id="u1",
            department_id=dept_id,
        )

        # The FakeProvider would use 100+50=150 tokens, but the quota
        # check happens BEFORE the provider call. Since current usage
        # (90) + nothing is checked — the gateway checks current_usage
        # >= limit, which is 90 < 100, so this would actually pass.
        #
        # To force a block, we pre-populate usage AT the limit (100
        # tokens). The check is `current_usage >= limit`, so 100 >= 100
        # → blocked.
        store._records.clear()
        gateway._usage_tracker.record(
            agent_name="prev",
            model="openai/gpt-4o",
            usage=TokenUsage(prompt_tokens=70, completion_tokens=30),
            cost=0.0,
            latency_ms=10,
            user_id="u1",
            department_id=dept_id,
        )

        with pytest.raises(QuotaExceededError) as exc_info:
            await gateway.chat(
                messages=[{"role": "user", "content": "hi"}],
                model="openai/gpt-4o",
                user_id="u1",
                department_ids=[dept_id],
                db_path=fresh_db,
            )
        err = exc_info.value
        assert err.department_id == dept_id
        assert err.quota_type == "token_limit"
        assert err.period == "daily"
        assert err.limit == 100
        assert err.current == 100  # 70 prompt + 30 completion

    async def test_cost_limit_blocks_request(
        self, gateway: LLMGateway, store: InMemoryUsageStore, fresh_db: Path
    ):
        """When department exceeds cost limit, LLM call raises QuotaExceededError."""
        dept_id = _random_dept_id()
        svc = get_quota_service()
        # cost_limit is in cents. Set 100 cents ($1.00) daily limit.
        await svc.set_quota(fresh_db, dept_id, "cost_limit", 100, period="daily")

        # Pre-populate usage with $1.50 cost = 150 cents, exceeding the
        # 100-cent limit.
        gateway._usage_tracker.record(
            agent_name="prev",
            model="openai/gpt-4o",
            usage=TokenUsage(prompt_tokens=100, completion_tokens=50),
            cost=1.50,  # $1.50 = 150 cents
            latency_ms=10,
            user_id="u1",
            department_id=dept_id,
        )

        with pytest.raises(QuotaExceededError) as exc_info:
            await gateway.chat(
                messages=[{"role": "user", "content": "hi"}],
                model="openai/gpt-4o",
                user_id="u1",
                department_ids=[dept_id],
                db_path=fresh_db,
            )
        err = exc_info.value
        assert err.quota_type == "cost_limit"
        assert err.period == "daily"
        assert err.limit == 100
        # current is in cents (150 cents = $1.50).
        assert err.current == 150.0

    async def test_model_whitelist_blocks_unlisted_model(self, gateway: LLMGateway, fresh_db: Path):
        """When model not in whitelist, LLM call is rejected."""
        dept_id = _random_dept_id()
        svc = get_quota_service()
        # Whitelist only allows "claude" — gateway is calling "gpt-4o".
        await svc.set_quota(fresh_db, dept_id, "model_whitelist", ["claude"], period="daily")

        with pytest.raises(QuotaExceededError) as exc_info:
            await gateway.chat(
                messages=[{"role": "user", "content": "hi"}],
                model="openai/gpt-4o",
                user_id="u1",
                department_ids=[dept_id],
                db_path=fresh_db,
            )
        err = exc_info.value
        assert err.quota_type == "model_whitelist"
        assert err.department_id == dept_id
        # For model_whitelist, current is the rejected model name.
        assert err.current == "openai/gpt-4o"

    async def test_no_quota_allows_all(self, gateway: LLMGateway, fresh_db: Path):
        """Without any quota set, all requests are allowed."""
        dept_id = _random_dept_id()
        # No quota set — request should succeed.
        response = await gateway.chat(
            messages=[{"role": "user", "content": "hi"}],
            model="openai/gpt-4o",
            user_id="u1",
            department_ids=[dept_id],
            db_path=fresh_db,
        )
        assert response.content == "response from openai"
        assert response.usage.total_tokens == 150  # 100 prompt + 50 completion

    async def test_multi_department_strictest_wins(
        self, gateway: LLMGateway, store: InMemoryUsageStore, fresh_db: Path
    ):
        """User in depts A+B: A has quota, B doesn't → A's quota applies.

        Strictest-wins: if ANY department fails ANY check, the request
        is rejected.
        """
        dept_a = _random_dept_id()
        dept_b = _random_dept_id()
        svc = get_quota_service()
        # Set a 1-token limit on dept A only; dept B has no quota.
        await svc.set_quota(fresh_db, dept_a, "token_limit", 1, period="daily")

        # Pre-populate usage for dept A so it exceeds the 1-token limit.
        gateway._usage_tracker.record(
            agent_name="prev",
            model="openai/gpt-4o",
            usage=TokenUsage(prompt_tokens=10, completion_tokens=5),
            cost=0.0,
            latency_ms=10,
            user_id="u1",
            department_id=dept_a,
        )

        with pytest.raises(QuotaExceededError) as exc_info:
            await gateway.chat(
                messages=[{"role": "user", "content": "hi"}],
                model="openai/gpt-4o",
                user_id="u1",
                department_ids=[dept_a, dept_b],
                db_path=fresh_db,
            )
        # The error should reference dept_a (the one that exceeded).
        assert exc_info.value.department_id == dept_a
        assert exc_info.value.quota_type == "token_limit"

    async def test_quota_check_with_real_gateway(
        self, gateway: LLMGateway, store: InMemoryUsageStore, fresh_db: Path
    ):
        """Integration test with real LLMGateway + mock provider.

        Verifies the full flow:
        1. Quota check happens before the provider call.
        2. On success, usage is recorded with the correct department_id.
        3. The usage record carries user_id + department_id.
        """
        dept_id = _random_dept_id()
        svc = get_quota_service()
        # Set a generous token limit (1M tokens) — should not block.
        await svc.set_quota(fresh_db, dept_id, "token_limit", 1_000_000, period="daily")

        # Make the LLM call.
        response = await gateway.chat(
            messages=[{"role": "user", "content": "hi"}],
            model="openai/gpt-4o",
            user_id="u1",
            department_ids=[dept_id],
            db_path=fresh_db,
        )
        assert response.content == "response from openai"

        # Verify usage was recorded with the correct attributes.
        summary = store.get_usage()
        assert len(summary.records) == 1
        rec = summary.records[0]
        assert rec.user_id == "u1"
        assert rec.department_id == dept_id
        assert rec.model == "gpt-4o"
        assert rec.total_tokens == 150  # 100 prompt + 50 completion

        # Verify the quota check counted this usage (next call should
        # still pass since the limit is 1M tokens).
        response2 = await gateway.chat(
            messages=[{"role": "user", "content": "hi again"}],
            model="openai/gpt-4o",
            user_id="u1",
            department_ids=[dept_id],
            db_path=fresh_db,
        )
        assert response2.content == "response from openai"

        # Now there should be 2 usage records.
        summary = store.get_usage()
        assert len(summary.records) == 2
        # All records should carry the department_id.
        assert all(r.department_id == dept_id for r in summary.records)

    async def test_quota_check_skipped_without_db_path(self, gateway: LLMGateway, fresh_db: Path):
        """When db_path is None, no quota check is performed."""
        dept_id = _random_dept_id()
        svc = get_quota_service()
        # Set a tiny quota that would normally block.
        await svc.set_quota(fresh_db, dept_id, "token_limit", 1, period="daily")

        # Call without db_path — should succeed (no quota check).
        response = await gateway.chat(
            messages=[{"role": "user", "content": "hi"}],
            model="openai/gpt-4o",
            user_id="u1",
            department_ids=[dept_id],
            db_path=None,
        )
        assert response.content == "response from openai"

    async def test_quota_check_skipped_without_department_ids(
        self, gateway: LLMGateway, fresh_db: Path
    ):
        """When department_ids is None, no quota check is performed."""
        response = await gateway.chat(
            messages=[{"role": "user", "content": "hi"}],
            model="openai/gpt-4o",
            user_id="u1",
            department_ids=None,
            db_path=fresh_db,
        )
        assert response.content == "response from openai"

    async def test_model_whitelist_allows_listed_model(self, gateway: LLMGateway, fresh_db: Path):
        """Model in whitelist → request allowed."""
        dept_id = _random_dept_id()
        svc = get_quota_service()
        # Whitelist uses the full resolved model identifier (provider/model).
        await svc.set_quota(
            fresh_db,
            dept_id,
            "model_whitelist",
            ["openai/gpt-4o"],
            period="daily",
        )
        response = await gateway.chat(
            messages=[{"role": "user", "content": "hi"}],
            model="openai/gpt-4o",
            user_id="u1",
            department_ids=[dept_id],
            db_path=fresh_db,
        )
        assert response.content == "response from openai"

    async def test_quota_check_uses_correct_period_window(
        self, gateway: LLMGateway, store: InMemoryUsageStore, fresh_db: Path
    ):
        """Quota check uses the daily window (since 00:00 UTC today).

        The quota check happens BEFORE the LLM call, using the current
        accumulated usage. So:
        - 1st call: usage=0, check 0 >= 150 → False → allowed. After
          the call, usage=150.
        - 2nd call: usage=150, check 150 >= 150 → True → blocked.
        """
        dept_id = _random_dept_id()
        svc = get_quota_service()
        # Set a 150-token daily limit (the FakeProvider uses 150 tokens
        # per call: 100 prompt + 50 completion).
        await svc.set_quota(fresh_db, dept_id, "token_limit", 150, period="daily")

        # First call: current usage is 0, under the 150 limit → allowed.
        response = await gateway.chat(
            messages=[{"role": "user", "content": "hi"}],
            model="openai/gpt-4o",
            user_id="u1",
            department_ids=[dept_id],
            db_path=fresh_db,
        )
        assert response.content == "response from openai"

        # Second call: current usage is now 150 (from the first call),
        # which is >= the 150-token limit → blocked.
        with pytest.raises(QuotaExceededError) as exc_info:
            await gateway.chat(
                messages=[{"role": "user", "content": "hi again"}],
                model="openai/gpt-4o",
                user_id="u1",
                department_ids=[dept_id],
                db_path=fresh_db,
            )
        assert exc_info.value.quota_type == "token_limit"
        assert exc_info.value.current == 150  # accumulated from first call
        assert exc_info.value.limit == 150