X-Financial/server/tests/test_commercial_runtime_reservations.py

from __future__ import annotations

import uuid
from datetime import UTC, datetime, timedelta
from decimal import Decimal

import pytest
from commercial_runtime_testkit import seed_meter, seed_run, seed_tool_call
from sqlalchemy import create_engine, select
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.pool import StaticPool

import app.models  # noqa: F401 - 注册完整 metadata
from app.db.base_class import Base
from app.models.commercial import UsageMeterEvent
from app.models.commercial_runtime import CommercialRuntimeReservation
from app.services.agent_runs import AgentRunService
from app.services.commercial_entitlements import CommercialEntitlementService
from app.services.commercial_runtime_bridge import CommercialRuntimeBridge
from app.services.commercial_runtime_reconciler import CommercialRuntimeReconciler


@pytest.fixture()
def db() -> Session:
    engine = create_engine(
        "sqlite+pysqlite:///:memory:",
        connect_args={"check_same_thread": False},
        poolclass=StaticPool,
    )
    Base.metadata.create_all(engine)
    factory = sessionmaker(bind=engine, expire_on_commit=False)
    with factory() as session:
        yield session
    Base.metadata.drop_all(engine)
    engine.dispose()


def test_reserved_success_settles_real_usage_and_replays_idempotently(
    db: Session,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    now = datetime.now(UTC)
    _, entitlement = seed_meter(db, "tenant-a", now, basis="call", hard_limit=Decimal("2"))
    run = seed_run(db, now, route_json={"tenant_id": "tenant-a"})
    tool_call_id = str(uuid.uuid4())
    bridge = CommercialRuntimeBridge(db)

    permit = bridge.reserve_tool(
        run.run_id,
        tool_call_id=tool_call_id,
        tool_type="llm",
        tool_name="chat.completions",
    )
    reservation = db.get(CommercialRuntimeReservation, permit.reservation_id)
    assert permit.gate.allowed is True
    assert reservation is not None and reservation.status == "reserved"
    assert db.query(UsageMeterEvent).count() == 0

    run_service = AgentRunService(db)
    monkeypatch.setattr(run_service, "_ensure_ready", lambda: None)
    run_service.record_tool_call(
        run_id=run.run_id,
        tool_call_id=tool_call_id,
        tool_type="llm",
        tool_name="chat.completions",
        status="succeeded",
    )
    replay = bridge.sync_tool_call(tool_call_id)
    db.refresh(reservation)
    quota = CommercialEntitlementService(db).get_account_for_tenant("tenant-a").quotas[0]

    assert reservation.status == "committed"
    assert Decimal(reservation.actual_quantity or 0) == Decimal("1")
    assert replay.status == "replayed"
    assert db.query(UsageMeterEvent).count() == 1
    assert quota.entitlement.id == entitlement.id
    assert quota.used_quantity == Decimal("1")
    assert quota.reserved_quantity == Decimal("0")


@pytest.mark.parametrize(
    ("tool_status", "business_occurred"),
    [("failed", True), ("blocked", False)],
)
def test_failed_or_blocked_tool_releases_reservation_without_usage(
    db: Session,
    monkeypatch: pytest.MonkeyPatch,
    tool_status: str,
    business_occurred: bool,
) -> None:
    now = datetime.now(UTC)
    seed_meter(db, "tenant-a", now, basis="call", hard_limit=Decimal("1"))
    run = seed_run(db, now, route_json={"tenant_id": "tenant-a"})
    tool_call_id = str(uuid.uuid4())
    permit = CommercialRuntimeBridge(db).reserve_tool(
        run.run_id,
        tool_call_id=tool_call_id,
        tool_type="llm",
        tool_name="chat.completions",
    )
    run_service = AgentRunService(db)
    monkeypatch.setattr(run_service, "_ensure_ready", lambda: None)

    run_service.record_tool_call(
        run_id=run.run_id,
        tool_call_id=tool_call_id,
        tool_type="llm",
        tool_name="chat.completions",
        status=tool_status,
    )
    result = CommercialRuntimeBridge(db).sync_tool_call(tool_call_id)
    reservation = db.get(CommercialRuntimeReservation, permit.reservation_id)

    assert reservation is not None and reservation.status == "released"
    assert result.reason_code == "tool_call_not_billable"
    assert result.business_call_occurred is business_occurred
    assert db.query(UsageMeterEvent).count() == 0


def test_active_reservation_holds_hard_quota_before_execution(db: Session) -> None:
    now = datetime.now(UTC)
    seed_meter(db, "tenant-a", now, basis="call", hard_limit=Decimal("1"))
    run = seed_run(db, now, route_json={"tenant_id": "tenant-a"})
    bridge = CommercialRuntimeBridge(db)

    first = bridge.reserve_tool(
        run.run_id,
        tool_call_id=str(uuid.uuid4()),
        tool_type="llm",
        tool_name="chat.completions",
    )
    second = bridge.reserve_tool(
        run.run_id,
        tool_call_id=str(uuid.uuid4()),
        tool_type="llm",
        tool_name="chat.completions",
    )
    quota = CommercialEntitlementService(db).get_account_for_tenant("tenant-a").quotas[0]

    assert first.reservation_id is not None
    assert second.reservation_id is None
    assert second.gate.allowed is False
    assert quota.reserved_quantity == Decimal("1")
    assert quota.hard_limit_remaining == Decimal("0")


def test_same_reservation_request_replays_without_consuming_its_own_quota(
    db: Session,
) -> None:
    now = datetime.now(UTC)
    seed_meter(db, "tenant-a", now, basis="call", hard_limit=Decimal("1"))
    run = seed_run(db, now, route_json={"tenant_id": "tenant-a"})
    bridge = CommercialRuntimeBridge(db)
    tool_call_id = str(uuid.uuid4())

    first = bridge.reserve_tool(
        run.run_id,
        tool_call_id=tool_call_id,
        tool_type="llm",
        tool_name="chat.completions",
    )
    replay = bridge.reserve_tool(
        run.run_id,
        tool_call_id=tool_call_id,
        tool_type="llm",
        tool_name="chat.completions",
    )

    assert replay.gate.allowed is True
    assert replay.gate.reason_code == "reservation_replayed"
    assert replay.reservation_id == first.reservation_id
    assert db.query(CommercialRuntimeReservation).count() == 1


def test_variable_quantity_requires_hard_max_and_refuses_actual_over_reservation(
    db: Session,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    now = datetime.now(UTC)
    seed_meter(db, "tenant-a", now, basis="duration_ms", hard_limit=Decimal("1000"))
    run = seed_run(db, now, route_json={"tenant_id": "tenant-a"})
    bridge = CommercialRuntimeBridge(db)
    rejected = bridge.reserve_tool(
        run.run_id,
        tool_call_id=str(uuid.uuid4()),
        tool_type="llm",
        tool_name="chat.completions",
        requested_quantity=Decimal("500"),
    )
    assert rejected.gate.reason_code == "hard_max_required"

    tool_call_id = str(uuid.uuid4())
    permit = bridge.reserve_tool(
        run.run_id,
        tool_call_id=tool_call_id,
        tool_type="llm",
        tool_name="chat.completions",
        requested_quantity=Decimal("500"),
        hard_max_confirmed=True,
    )
    run_service = AgentRunService(db)
    monkeypatch.setattr(run_service, "_ensure_ready", lambda: None)
    run_service.record_tool_call(
        run_id=run.run_id,
        tool_call_id=tool_call_id,
        tool_type="llm",
        tool_name="chat.completions",
        status="succeeded",
        duration_ms=501,
    )
    failed = bridge.sync_tool_call(tool_call_id)
    reservation = db.get(CommercialRuntimeReservation, permit.reservation_id)
    assert failed.status == "error"
    assert failed.requires_reconciliation is True
    assert reservation is not None and reservation.status == "reserved"
    assert db.query(UsageMeterEvent).count() == 0

    run_service.update_tool_call(tool_call_id, duration_ms=400, status="succeeded")
    db.refresh(reservation)
    assert reservation.status == "committed"
    assert Decimal(db.scalars(select(UsageMeterEvent)).one().quantity) == Decimal("400")


def test_expired_historical_meter_does_not_enable_runtime_enforcement(db: Session) -> None:
    now = datetime.now(UTC)
    subscription, _ = seed_meter(db, "tenant-a", now, basis="call")
    subscription.current_period_start = now - timedelta(days=31)
    subscription.current_period_end = now - timedelta(days=1)
    db.commit()
    run = seed_run(db, now, route_json={"tenant_id": "tenant-a"})

    permit = CommercialRuntimeBridge(db).reserve_tool(
        run.run_id,
        tool_call_id=str(uuid.uuid4()),
        tool_type="llm",
        tool_name="chat.completions",
    )

    assert permit.gate.enforced is False
    assert permit.gate.allowed is True
    assert permit.gate.reason_code == "runtime_meter_not_configured"
    assert permit.reservation_id is None


def test_legacy_run_without_tenant_skips_bridge_without_reconciliation(db: Session) -> None:
    now = datetime.now(UTC)
    _, tool_call = seed_tool_call(db, now, route_json={})

    result = CommercialRuntimeBridge(db).sync_tool_call(tool_call.id)

    assert result.status == "skipped"
    assert result.reason_code == "legacy_run_without_tenant"
    assert result.requires_reconciliation is False
    assert db.query(CommercialRuntimeReservation).count() == 0


def test_direct_call_on_suspended_contract_is_persisted_for_reconciliation(
    db: Session,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    now = datetime.now(UTC)
    subscription, _ = seed_meter(
        db,
        "tenant-a",
        now,
        basis="call",
    )
    subscription.status = "suspended"
    db.flush()
    run = seed_run(db, now, route_json={"tenant_id": "tenant-a"})
    run_service = AgentRunService(db)
    monkeypatch.setattr(run_service, "_ensure_ready", lambda: None)

    call = run_service.record_tool_call(
        run_id=run.run_id,
        tool_type="llm",
        tool_name="chat.completions",
        status="succeeded",
    )
    reservation = db.scalars(select(CommercialRuntimeReservation)).one()

    assert reservation.tool_call_id == call.id
    assert reservation.status == "reconciliation_required"
    assert db.query(UsageMeterEvent).count() == 0


def test_expired_reconciler_settles_releases_and_preserves_uncertain_holds(
    db: Session,
) -> None:
    now = datetime.now(UTC)
    seed_meter(db, "tenant-a", now, basis="call", hard_limit=Decimal("10"))
    bridge = CommercialRuntimeBridge(db)
    reservations: dict[str, CommercialRuntimeReservation] = {}
    for scenario, run_status in (
        ("success", "running"),
        ("failure", "running"),
        ("terminal_without_call", "failed"),
        ("uncertain", "running"),
    ):
        run = seed_run(
            db,
            now,
            route_json={"tenant_id": "tenant-a"},
            status=run_status,
        )
        tool_call_id = str(uuid.uuid4())
        permit = bridge.reserve_tool(
            run.run_id,
            tool_call_id=tool_call_id,
            tool_type="llm",
            tool_name="chat.completions",
        )
        reservation = db.get(CommercialRuntimeReservation, permit.reservation_id)
        assert reservation is not None
        reservation.created_at = now - timedelta(hours=1)
        reservation.expires_at = now - timedelta(minutes=30)
        reservations[scenario] = reservation
        if scenario == "success":
            seed_tool_call(
                db,
                now - timedelta(minutes=45),
                route_json={"tenant_id": "tenant-a"},
                run=run,
                status="succeeded",
            )[1].id = tool_call_id
        elif scenario == "failure":
            seed_tool_call(
                db,
                now - timedelta(minutes=45),
                route_json={"tenant_id": "tenant-a"},
                run=run,
                status="failed",
            )[1].id = tool_call_id
        db.commit()

    batch = CommercialRuntimeReconciler(db).reconcile_expired(as_of=now, limit=10)
    for row in reservations.values():
        db.refresh(row)

    assert batch.settled == 1
    assert batch.released == 2
    assert batch.deferred == 1
    assert batch.errors == 0
    assert reservations["success"].status == "committed"
    assert reservations["failure"].status == "released"
    assert reservations["terminal_without_call"].status == "expired"
    assert reservations["uncertain"].status == "reserved"
    assert db.query(UsageMeterEvent).count() == 1