X-Financial/server/tests/test_commercial_direct_operation.py

from __future__ import annotations

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

import pytest
from commercial_runtime_testkit import seed_meter
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 CommercialCostEvent, UsageMeterEvent
from app.models.commercial_runtime import CommercialRuntimeReservation
from app.services.commercial_direct_operation import (
    CommercialDirectOperationBridge,
    DirectOperationIdentity,
)


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


def _identity(now: datetime, suffix: str = "a") -> DirectOperationIdentity:
    return DirectOperationIdentity(
        tenant_id="tenant-a",
        operation_key=f"SENSITIVE-OPERATION-{suffix}",
        run_key=f"SENSITIVE-RUN-{suffix}",
        tool_type="llm",
        tool_name="chat.completions",
        provider="OpenAI",
        model_name="gpt-test",
        started_at=now,
    )


def test_unconfigured_direct_operation_is_compatible_and_writes_no_fact(
    factory: sessionmaker[Session],
) -> None:
    now = datetime.now(UTC)
    bridge = CommercialDirectOperationBridge(factory)

    permit = bridge.permit(_identity(now))

    assert permit.enforced is False
    assert permit.allowed is True
    assert permit.reason_code == "runtime_meter_not_configured"
    with factory() as db:
        assert db.query(CommercialRuntimeReservation).count() == 0
        assert db.query(UsageMeterEvent).count() == 0


def test_authoritative_tokens_settle_once_with_redacted_identity_and_cost(
    factory: sessionmaker[Session],
) -> None:
    now = datetime.now(UTC)
    with factory() as db:
        seed_meter(
            db,
            "tenant-a",
            now,
            basis="total_tokens",
            preflight_quantity=Decimal("20"),
            internal_cost={
                "enabled": True,
                "cost_category": "ai_inference",
                "unit": "token",
                "unit_cost": "0.002",
                "original_currency": "CNY",
                "reporting_currency": "CNY",
                "fx_rate": "1",
                "provider": "OpenAI",
                "model_name": "gpt-test",
            },
        )
        db.commit()

    identity = _identity(now, "tokens")
    bridge = CommercialDirectOperationBridge(factory)
    permit = bridge.permit(identity)
    result = bridge.complete(
        identity,
        outcome="succeeded",
        authoritative_quantities={
            "input_tokens": 3,
            "output_tokens": 2,
            "total_tokens": 5,
        },
        completed_at=now + timedelta(seconds=1),
        usage_source="openai_usage",
        usage_availability="available",
    )
    replay = bridge.complete(
        identity,
        outcome="succeeded",
        authoritative_quantities={
            "input_tokens": 3,
            "output_tokens": 2,
            "total_tokens": 5,
        },
        completed_at=now + timedelta(seconds=1),
        usage_source="openai_usage",
        usage_availability="available",
    )

    assert permit.allowed is True and permit.reservation_id
    assert permit.reserved_quantity == Decimal("20")
    assert result.status == "created"
    assert result.quantity == Decimal("5")
    assert replay.status == "replayed"
    with factory() as db:
        reservation = db.scalars(select(CommercialRuntimeReservation)).one()
        usage = db.scalars(select(UsageMeterEvent)).one()
        cost = db.scalars(select(CommercialCostEvent)).one()
        assert reservation.status == "committed"
        assert Decimal(reservation.actual_quantity or 0) == Decimal("5")
        assert Decimal(usage.quantity) == Decimal("5")
        assert Decimal(cost.cost_amount) == Decimal("0.0100")
        assert usage.billing_period_id == reservation.billing_period_id
        assert cost.billing_period_id == reservation.billing_period_id
        serialized = json.dumps(usage.metadata_json, ensure_ascii=False)
        assert "SENSITIVE-OPERATION" not in serialized
        assert "SENSITIVE-RUN" not in serialized
        assert db.query(UsageMeterEvent).count() == 1
        assert db.query(CommercialCostEvent).count() == 1


@pytest.mark.parametrize(
    ("basis", "tool_type", "tool_name"),
    [
        ("bytes", "storage", "attachment.upload"),
        ("pages", "ocr", "receipt.extract"),
        ("objects", "storage", "attachment.persist"),
        ("events", "connector", "financial.ingest"),
    ],
)
def test_authoritative_resource_quantities_settle_without_estimation(
    factory: sessionmaker[Session],
    basis: str,
    tool_type: str,
    tool_name: str,
) -> None:
    now = datetime.now(UTC)
    with factory() as db:
        seed_meter(
            db,
            "tenant-a",
            now,
            basis=basis,
            tool_type=tool_type,
            tool_name=tool_name,
            preflight_quantity=Decimal("100"),
        )
        db.commit()
    identity = DirectOperationIdentity(
        tenant_id="tenant-a",
        operation_key=f"resource-operation:{basis}",
        run_key=f"resource-run:{basis}",
        tool_type=tool_type,
        tool_name=tool_name,
        started_at=now,
    )
    bridge = CommercialDirectOperationBridge(factory)

    permit = bridge.permit(identity, requested_quantity=7)
    assert permit.allowed is True
    assert permit.reserved_quantity == Decimal("7")
    result = bridge.complete(
        identity,
        outcome="succeeded",
        authoritative_quantities={basis: 7},
        completed_at=now + timedelta(seconds=1),
        usage_source=f"authoritative_{basis}",
        usage_availability="available",
    )

    assert result.status == "created"
    assert result.quantity_basis == basis
    assert result.quantity == Decimal("7")
    with factory() as db:
        usage = db.scalars(select(UsageMeterEvent)).one()
        assert Decimal(usage.quantity) == Decimal("7")
        assert usage.metadata_json["usage_source"] == f"authoritative_{basis}"


@pytest.mark.parametrize(
    ("outcome", "quantities", "reason_code", "actual"),
    [
        ("outcome_unknown", {}, "authoritative_usage_unavailable", None),
        (
            "provider_rejected",
            {"total_tokens": 25},
            "actual_exceeds_preflight_reservation",
            Decimal("25"),
        ),
    ],
)
def test_uncertain_or_over_reservation_call_enters_durable_reconciliation(
    factory: sessionmaker[Session],
    outcome: str,
    quantities: dict[str, int],
    reason_code: str,
    actual: Decimal | None,
) -> None:
    now = datetime.now(UTC)
    with factory() as db:
        seed_meter(
            db,
            "tenant-a",
            now,
            basis="total_tokens",
            preflight_quantity=Decimal("20"),
        )
        db.commit()
    identity = _identity(now, outcome)
    bridge = CommercialDirectOperationBridge(factory)
    assert bridge.permit(identity).allowed is True

    result = bridge.complete(
        identity,
        outcome=outcome,  # type: ignore[arg-type]
        authoritative_quantities=quantities,
        completed_at=now + timedelta(seconds=1),
        usage_source="unavailable" if not quantities else "openai_usage",
        usage_availability="unavailable" if not quantities else "available",
    )

    assert result.status == "reconciliation_required"
    assert result.reason_code == reason_code
    with factory() as db:
        reservation = db.scalars(select(CommercialRuntimeReservation)).one()
        assert reservation.status == "reconciliation_required"
        assert (
            Decimal(reservation.actual_quantity)
            if reservation.actual_quantity is not None
            else None
        ) == actual
        assert db.query(UsageMeterEvent).count() == 0
        assert db.query(CommercialCostEvent).count() == 0


def test_not_sent_releases_reservation_and_never_meters(
    factory: sessionmaker[Session],
) -> None:
    now = datetime.now(UTC)
    with factory() as db:
        seed_meter(db, "tenant-a", now, basis="call")
        db.commit()
    identity = _identity(now, "not-sent")
    bridge = CommercialDirectOperationBridge(factory)
    assert bridge.permit(identity).allowed is True

    result = bridge.complete(
        identity,
        outcome="not_sent",
        authoritative_quantities={},
        completed_at=now + timedelta(milliseconds=10),
        usage_source="unavailable",
        usage_availability="unavailable",
    )

    assert result.status == "released"
    with factory() as db:
        reservation = db.scalars(select(CommercialRuntimeReservation)).one()
        assert reservation.status == "released"
        assert db.query(UsageMeterEvent).count() == 0


def test_released_operation_can_be_preflighted_again_after_business_rollback(
    factory: sessionmaker[Session],
) -> None:
    now = datetime.now(UTC)
    with factory() as db:
        seed_meter(db, "tenant-a", now, basis="call")
        db.commit()
    identity = _identity(now, "released-retry")
    bridge = CommercialDirectOperationBridge(factory)
    assert bridge.permit(identity).allowed is True
    bridge.complete(
        identity,
        outcome="not_sent",
        authoritative_quantities={},
        completed_at=now + timedelta(milliseconds=10),
        usage_source="business_transaction_rolled_back",
        usage_availability="unavailable",
    )

    retry = bridge.permit(identity)
    result = bridge.complete(
        identity,
        outcome="succeeded",
        authoritative_quantities={"call": 1},
        completed_at=now + timedelta(milliseconds=20),
        usage_source="accepted_retry",
        usage_availability="available",
    )

    assert retry.allowed is True
    assert result.status == "created"
    with factory() as db:
        reservation = db.scalars(select(CommercialRuntimeReservation)).one()
        assert reservation.status == "committed"
        assert db.query(UsageMeterEvent).count() == 1


def test_variable_meter_without_explicit_hard_preflight_is_denied(
    factory: sessionmaker[Session],
) -> None:
    now = datetime.now(UTC)
    with factory() as db:
        seed_meter(db, "tenant-a", now, basis="total_tokens")
        db.commit()

    permit = CommercialDirectOperationBridge(factory).permit(_identity(now, "no-max"))

    assert permit.enforced is True
    assert permit.allowed is False
    assert "preflight_quantity" in permit.reason
    with factory() as db:
        assert db.query(CommercialRuntimeReservation).count() == 0


def test_known_resource_quantity_above_configured_maximum_is_denied_before_call(
    factory: sessionmaker[Session],
) -> None:
    now = datetime.now(UTC)
    with factory() as db:
        seed_meter(
            db,
            "tenant-a",
            now,
            basis="pages",
            tool_type="ocr",
            tool_name="receipt.extract",
            preflight_quantity=Decimal("10"),
        )
        db.commit()
    identity = DirectOperationIdentity(
        tenant_id="tenant-a",
        operation_key="resource-operation:too-many-pages",
        run_key="resource-run:too-many-pages",
        tool_type="ocr",
        tool_name="receipt.extract",
        started_at=now,
    )

    permit = CommercialDirectOperationBridge(factory).permit(
        identity,
        requested_quantity=11,
    )

    assert permit.allowed is False
    assert "硬上限" in permit.reason
    with factory() as db:
        assert db.query(CommercialRuntimeReservation).count() == 0


def test_call_without_observed_permit_creates_reconciliation_backlog(
    factory: sessionmaker[Session],
) -> None:
    now = datetime.now(UTC)
    with factory() as db:
        seed_meter(db, "tenant-a", now, basis="call")
        db.commit()
    identity = _identity(now, "missing-permit")

    result = CommercialDirectOperationBridge(factory).complete(
        identity,
        outcome="provider_rejected",
        authoritative_quantities={},
        completed_at=now + timedelta(milliseconds=50),
        usage_source="unavailable",
        usage_availability="unavailable",
    )

    assert result.status == "reconciliation_required"
    assert result.reason_code == "missing_pre_execution_reservation"
    assert result.quantity == Decimal("1")
    with factory() as db:
        reservation = db.scalars(select(CommercialRuntimeReservation)).one()
        assert reservation.status == "reconciliation_required"
        assert reservation.resolution_code == "missing_pre_execution_reservation"