iai-mcp-opencode/tests/test_guard.py

"""Tests for D-GUARD (BudgetLedger + RateLimitLedger + should_call_llm).

Covers:
- BudgetLedger daily/monthly caps + rollover
- RateLimitLedger cooldown window
- should_call_llm 7-step ladder ordering per CONTEXT.md D-GUARD
- Persistence across store reopen
"""
from __future__ import annotations

from datetime import datetime, timedelta, timezone
from uuid import uuid4

import pytest


# ------------------------------------------------------------- BudgetLedger


def test_budget_ledger_daily_cap_enforced(tmp_path):
    from iai_mcp.guard import BudgetLedger
    from iai_mcp.store import MemoryStore

    store = MemoryStore(path=tmp_path)
    bl = BudgetLedger(store, daily_usd_cap=0.10, monthly_usd_cap=3.00)

    ok, _ = bl.can_spend(0.05)
    assert ok is True

    bl.record_spend(0.08)
    ok, _ = bl.can_spend(0.03)
    # 0.08 + 0.03 = 0.11 > 0.10 -> NOT ok
    ok2, reason = bl.can_spend(0.03)
    assert ok2 is False
    assert "daily" in reason.lower()


def test_budget_ledger_daily_allows_under_cap(tmp_path):
    from iai_mcp.guard import BudgetLedger
    from iai_mcp.store import MemoryStore

    store = MemoryStore(path=tmp_path)
    bl = BudgetLedger(store, daily_usd_cap=0.10)
    bl.record_spend(0.05)
    ok, _ = bl.can_spend(0.04)
    assert ok is True


def test_budget_ledger_monthly_cap_enforced(tmp_path):
    """Daily small spends accumulate to monthly cap."""
    from iai_mcp.guard import BudgetLedger
    from iai_mcp.store import MemoryStore

    store = MemoryStore(path=tmp_path)
    bl = BudgetLedger(store, daily_usd_cap=10.0, monthly_usd_cap=0.20)
    bl.record_spend(0.15)
    ok, reason = bl.can_spend(0.10)
    # 0.15 + 0.10 = 0.25 > 0.20 -> NOT ok, but reason is monthly (daily cap 10.0 is fine)
    assert ok is False
    assert "monthly" in reason.lower()


def test_budget_ledger_daily_used(tmp_path):
    from iai_mcp.guard import BudgetLedger
    from iai_mcp.store import MemoryStore

    store = MemoryStore(path=tmp_path)
    bl = BudgetLedger(store)
    assert bl.daily_used() == 0.0
    bl.record_spend(0.01)
    bl.record_spend(0.02)
    assert abs(bl.daily_used() - 0.03) < 1e-5


def test_budget_ledger_monthly_used(tmp_path):
    from iai_mcp.guard import BudgetLedger
    from iai_mcp.store import MemoryStore

    store = MemoryStore(path=tmp_path)
    bl = BudgetLedger(store)
    bl.record_spend(0.05)
    bl.record_spend(0.03)
    assert abs(bl.monthly_used() - 0.08) < 1e-5


def test_budget_ledger_persists_across_reopen(tmp_path):
    """Ledger-backed by LanceDB -> survives store close/reopen (D-GUARD repudiation)."""
    from iai_mcp.guard import BudgetLedger
    from iai_mcp.store import MemoryStore

    store1 = MemoryStore(path=tmp_path)
    BudgetLedger(store1).record_spend(0.05)
    del store1

    store2 = MemoryStore(path=tmp_path)
    bl = BudgetLedger(store2)
    assert abs(bl.daily_used() - 0.05) < 1e-5


# ----------------------------------------------------------- RateLimitLedger


def test_ratelimit_ledger_no_history_not_in_cooldown(tmp_path):
    from iai_mcp.guard import RateLimitLedger
    from iai_mcp.store import MemoryStore

    store = MemoryStore(path=tmp_path)
    rl = RateLimitLedger(store)
    assert rl.in_cooldown() is False


def test_ratelimit_ledger_record_429_enters_cooldown(tmp_path):
    from iai_mcp.guard import RateLimitLedger
    from iai_mcp.store import MemoryStore

    store = MemoryStore(path=tmp_path)
    rl = RateLimitLedger(store)
    rl.record_429()
    assert rl.in_cooldown() is True


def test_ratelimit_ledger_persists_across_reopen(tmp_path):
    from iai_mcp.guard import RateLimitLedger
    from iai_mcp.store import MemoryStore

    store1 = MemoryStore(path=tmp_path)
    RateLimitLedger(store1).record_429()
    del store1

    store2 = MemoryStore(path=tmp_path)
    assert RateLimitLedger(store2).in_cooldown() is True


# -------------------------------------------------- should_call_llm ladder


def test_should_call_llm_tier_0_fallback_llm_disabled(tmp_path):
    """Step 1: llm_enabled=False -> (False, 'sleep.llm_enabled=false')."""
    from iai_mcp.guard import BudgetLedger, RateLimitLedger, should_call_llm
    from iai_mcp.store import MemoryStore

    store = MemoryStore(path=tmp_path)
    bl = BudgetLedger(store)
    rl = RateLimitLedger(store)
    ok, reason = should_call_llm(bl, rl, llm_enabled=False, has_api_key=True)
    assert ok is False
    assert "llm_enabled" in reason


def test_should_call_llm_no_api_key(tmp_path):
    """Step 2: no api key -> (False, 'no api key')."""
    from iai_mcp.guard import BudgetLedger, RateLimitLedger, should_call_llm
    from iai_mcp.store import MemoryStore

    store = MemoryStore(path=tmp_path)
    bl = BudgetLedger(store)
    rl = RateLimitLedger(store)
    ok, reason = should_call_llm(bl, rl, llm_enabled=True, has_api_key=False)
    assert ok is False
    assert "api key" in reason.lower()


def test_should_call_llm_daily_cap_hit(tmp_path):
    """Step 3: daily cap exhausted -> (False, ... daily cap ...)."""
    from iai_mcp.guard import BudgetLedger, RateLimitLedger, should_call_llm
    from iai_mcp.store import MemoryStore

    store = MemoryStore(path=tmp_path)
    bl = BudgetLedger(store, daily_usd_cap=0.01, monthly_usd_cap=3.0)
    bl.record_spend(0.009)
    rl = RateLimitLedger(store)
    ok, reason = should_call_llm(
        bl, rl, llm_enabled=True, has_api_key=True, estimated_usd=0.005
    )
    assert ok is False
    assert "daily" in reason.lower()


def test_should_call_llm_monthly_cap_hit(tmp_path):
    """Step 4: daily ok, monthly cap exhausted."""
    from iai_mcp.guard import BudgetLedger, RateLimitLedger, should_call_llm
    from iai_mcp.store import MemoryStore

    store = MemoryStore(path=tmp_path)
    bl = BudgetLedger(store, daily_usd_cap=10.0, monthly_usd_cap=0.02)
    bl.record_spend(0.015)
    rl = RateLimitLedger(store)
    ok, reason = should_call_llm(
        bl, rl, llm_enabled=True, has_api_key=True, estimated_usd=0.01
    )
    assert ok is False
    assert "monthly" in reason.lower()


def test_should_call_llm_in_cooldown(tmp_path):
    """Step 5: budget ok, but rate limiter in cooldown."""
    from iai_mcp.guard import BudgetLedger, RateLimitLedger, should_call_llm
    from iai_mcp.store import MemoryStore

    store = MemoryStore(path=tmp_path)
    bl = BudgetLedger(store)
    rl = RateLimitLedger(store)
    rl.record_429()
    ok, reason = should_call_llm(bl, rl, llm_enabled=True, has_api_key=True)
    assert ok is False
    assert "cooldown" in reason.lower()


def test_should_call_llm_all_green(tmp_path):
    """All 7 steps pass -> (True, 'ok')."""
    from iai_mcp.guard import BudgetLedger, RateLimitLedger, should_call_llm
    from iai_mcp.store import MemoryStore

    store = MemoryStore(path=tmp_path)
    bl = BudgetLedger(store)
    rl = RateLimitLedger(store)
    ok, reason = should_call_llm(bl, rl, llm_enabled=True, has_api_key=True)
    assert ok is True
    assert reason == "ok"


def test_should_call_llm_ordering_llm_enabled_first(tmp_path):
    """Ladder ordering: llm_enabled takes precedence over budget+cooldown+apikey."""
    from iai_mcp.guard import BudgetLedger, RateLimitLedger, should_call_llm
    from iai_mcp.store import MemoryStore

    store = MemoryStore(path=tmp_path)
    bl = BudgetLedger(store, daily_usd_cap=0.01)
    bl.record_spend(0.02)  # over cap
    rl = RateLimitLedger(store)
    rl.record_429()        # in cooldown

    # llm_enabled=False short-circuits BEFORE cap + cooldown checks
    ok, reason = should_call_llm(bl, rl, llm_enabled=False, has_api_key=False)
    assert ok is False
    assert "llm_enabled" in reason


def test_should_call_llm_ordering_cap_before_cooldown(tmp_path):
    """With llm_enabled+api_key, budget cap check precedes cooldown."""
    from iai_mcp.guard import BudgetLedger, RateLimitLedger, should_call_llm
    from iai_mcp.store import MemoryStore

    store = MemoryStore(path=tmp_path)
    bl = BudgetLedger(store, daily_usd_cap=0.01)
    bl.record_spend(0.02)  # over cap
    rl = RateLimitLedger(store)
    rl.record_429()        # also in cooldown

    ok, reason = should_call_llm(
        bl, rl, llm_enabled=True, has_api_key=True, estimated_usd=0.001
    )
    assert ok is False
    # "daily" message means cap was checked before cooldown
    assert "daily" in reason.lower()