X-Financial/server/tests/test_runtime_chat_commercial.py

from __future__ import annotations

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.agent_run import AgentRun
from app.models.commercial import UsageMeterEvent
from app.models.commercial_runtime import CommercialRuntimeReservation
from app.services.commercial_direct_operation import CommercialDirectOperationBridge
from app.services.runtime_chat_attempts import (
    RuntimeChatAttemptCompletedEvent,
    RuntimeChatAttemptPermitEvent,
    RuntimeChatAuthoritativeUsage,
    RuntimeChatOperationContext,
)
from app.services.runtime_chat_commercial import (
    CommercialDirectReconciliationRequired,
    CommercialRuntimeChatAttemptObserver,
    trusted_runtime_chat_operation_context,
)


@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(*, tenant_id: str = "tenant-a", operation_id: str = "operation-a"):
    return RuntimeChatOperationContext(
        tenant_id=tenant_id,
        operation_id=operation_id,
        run_id="run-a",
        invocation_seq=3,
        attempt_scope="user-agent-response",
    ).build_attempt_identity(
        slot="main",
        provider="Ollama",
        model="qwen-test",
        attempt=1,
    )


def test_runtime_chat_observer_settles_authoritative_provider_tokens(
    factory: sessionmaker[Session],
) -> None:
    now = datetime.now(UTC)
    with factory() as db:
        seed_meter(
            db,
            "tenant-a",
            now,
            basis="total_tokens",
            preflight_quantity=Decimal("50"),
        )
        db.commit()
    observer = CommercialRuntimeChatAttemptObserver(
        CommercialDirectOperationBridge(factory)
    )
    identity = _identity()

    permit = observer.on_permit(
        RuntimeChatAttemptPermitEvent(identity=identity, started_at=now)
    )
    observer.on_completed(
        RuntimeChatAttemptCompletedEvent(
            identity=identity,
            started_at=now,
            completed_at=now + timedelta(milliseconds=250),
            outcome="succeeded",
            response_id=None,
            response_model="qwen-test",
            provider_status_code=200,
            usage=RuntimeChatAuthoritativeUsage(
                source="ollama_response",
                availability="available",
                prompt_eval_count=12,
                eval_count=8,
            ),
        )
    )

    assert permit.allowed is True
    with factory() as db:
        reservation = db.scalars(select(CommercialRuntimeReservation)).one()
        usage = db.scalars(select(UsageMeterEvent)).one()
        assert reservation.status == "committed"
        assert Decimal(reservation.actual_quantity or 0) == Decimal("20")
        assert Decimal(usage.quantity) == Decimal("20")
        assert usage.metadata_json["usage_source"] == "ollama_response"


def test_runtime_chat_observer_persists_reconciliation_when_usage_is_missing(
    factory: sessionmaker[Session],
) -> None:
    now = datetime.now(UTC)
    with factory() as db:
        seed_meter(
            db,
            "tenant-a",
            now,
            basis="total_tokens",
            preflight_quantity=Decimal("50"),
        )
        db.commit()
    observer = CommercialRuntimeChatAttemptObserver(
        CommercialDirectOperationBridge(factory)
    )
    identity = _identity(operation_id="operation-missing-usage")
    assert observer.on_permit(
        RuntimeChatAttemptPermitEvent(identity=identity, started_at=now)
    ).allowed

    with pytest.raises(
        CommercialDirectReconciliationRequired,
        match="authoritative_usage_unavailable",
    ):
        observer.on_completed(
            RuntimeChatAttemptCompletedEvent(
                identity=identity,
                started_at=now,
                completed_at=now + timedelta(seconds=1),
                outcome="succeeded",
                response_id="response-without-usage",
                response_model="qwen-test",
                provider_status_code=200,
                usage=RuntimeChatAuthoritativeUsage(),
            )
        )

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


def test_operation_context_uses_only_persisted_agent_run_tenant(
    factory: sessionmaker[Session],
) -> None:
    now = datetime.now(UTC)
    with factory() as db:
        db.add_all(
            [
                AgentRun(
                    run_id="trusted-run",
                    agent="user_agent",
                    source="chat",
                    user_id="user-a",
                    route_json={"tenant_id": "tenant-trusted"},
                    permission_level="write",
                    status="running",
                    started_at=now,
                ),
                AgentRun(
                    run_id="unscoped-run",
                    agent="user_agent",
                    source="chat",
                    user_id="user-a",
                    route_json={"context_json": {"tenant_id": "tenant-spoofed"}},
                    permission_level="write",
                    status="running",
                    started_at=now,
                ),
            ]
        )
        db.commit()

        context = trusted_runtime_chat_operation_context(
            db,
            run_id="trusted-run",
            attempt_scope="user-agent-response",
            invocation_seq=2,
        )

        assert context is not None
        assert context.tenant_id == "tenant-trusted"
        assert context.operation_id == "agent-run:trusted-run"
        assert context.invocation_seq == 2
        assert (
            trusted_runtime_chat_operation_context(
                db,
                run_id="unscoped-run",
                attempt_scope="user-agent-response",
            )
            is None
        )
        assert (
            trusted_runtime_chat_operation_context(
                db,
                run_id="missing-run",
                attempt_scope="user-agent-response",
            )
            is None
        )