harness_zbt/backend/tests/test_agent_api.py

import re
from datetime import UTC, datetime

from fastapi.testclient import TestClient

from zbt.core.config import Settings
from zbt.core.passwords import PasswordService
from zbt.domains.agent.repository import InMemoryAgentRepository
from zbt.domains.agent.runtime import AgentReply
from zbt.domains.catalog.repository import InMemoryCatalogRepository
from zbt.domains.enrollment.repository import InMemoryEnrollmentRepository
from zbt.domains.identity.models import AdminUser
from zbt.domains.identity.repository import InMemoryIdentityRepository
from zbt.harness.kernel import AgentInvocation
from zbt.infrastructure.redis.agent_state import InMemoryAgentStateStore
from zbt.main import create_app


class FixedAgentRuntime:
    def reply(self, invocation: AgentInvocation) -> AgentReply:
        assert "65岁" in invocation.message
        assert invocation.persona == "customer"
        return AgentReply(
            text="可以优先考虑银龄守护医疗险，我会继续核对地区和职业。",
            cards=[],
            trace_id="trace_test_001",
        )


class FixedOperationRuntime:
    def reply(self, invocation: AgentInvocation) -> AgentReply:
        assert invocation.persona == "operation"
        assert invocation.admin_user is not None
        assert "订单" in invocation.message
        return AgentReply(
            text="当前共有24笔订单，其中18笔已出单。",
            cards=[
                {
                    "type": "metrics",
                    "version": "1.0",
                    "title": "经营总览",
                    "items": [
                        {"label": "订单总量", "value": 24, "unit": "笔"},
                    ],
                }
            ],
            trace_id="00000000-0000-0000-0000-000000000001",
            trace_url="https://smith.langchain.com/example-trace",
            invoked_tools=("get_operation_overview",),
        )


class StreamingAgentRuntime:
    def reply(self, invocation: AgentInvocation) -> AgentReply:
        assert invocation.h5_user is not None
        return AgentReply(
            text="已为你找到银龄守护医疗险。",
            cards=[
                {
                    "type": "product_recommendations",
                    "version": "1.0",
                    "title": "为你匹配的保障方案",
                    "items": [],
                }
            ],
            trace_id="00000000-0000-0000-0000-000000000002",
            actions=[],
            invoked_tools=("list_available_products",),
        )


def test_h5_user_can_create_customer_agent_thread() -> None:
    settings = Settings(
        app_env="test",
        jwt_access_secret="a" * 32,
        jwt_refresh_secret="b" * 32,
        field_encryption_key="c" * 32,
    )
    app = create_app(
        settings=settings,
        identity_repository=InMemoryIdentityRepository(),
        agent_repository=InMemoryAgentRepository(),
    )

    with TestClient(app) as client:
        token = client.post(
            "/api/v1/h5/auth/login",
            json={"mobile": "18800000001", "code": "147258"},
        ).json()["data"]["tokens"]["access_token"]
        response = client.post(
            "/api/v1/agent/threads",
            headers={"Authorization": f"Bearer {token}"},
            json={"title": "给父亲配置医疗险"},
        )

    assert response.status_code == 201
    assert response.json()["data"]["persona"] == "customer"
    assert response.json()["data"]["title"] == "给父亲配置医疗险"


def test_h5_user_can_send_message_and_receive_agent_reply() -> None:
    settings = Settings(
        app_env="test",
        jwt_access_secret="a" * 32,
        jwt_refresh_secret="b" * 32,
        field_encryption_key="c" * 32,
    )
    app = create_app(
        settings=settings,
        identity_repository=InMemoryIdentityRepository(),
        agent_repository=InMemoryAgentRepository(),
        agent_runtime=FixedAgentRuntime(),
    )

    with TestClient(app) as client:
        token = client.post(
            "/api/v1/h5/auth/login",
            json={"mobile": "18800000001", "code": "147258"},
        ).json()["data"]["tokens"]["access_token"]
        headers = {"Authorization": f"Bearer {token}"}
        thread_id = client.post(
            "/api/v1/agent/threads",
            headers=headers,
            json={"title": "给父亲配置医疗险"},
        ).json()["data"]["id"]
        response = client.post(
            f"/api/v1/agent/threads/{thread_id}/messages",
            headers=headers,
            json={
                "content": {
                    "type": "text",
                    "text": "想给65岁的父亲买医疗险",
                },
                "client_message_id": "client-message-001",
            },
        )

    assert response.status_code == 202
    assert response.json()["data"]["status"] == "COMPLETED"
    assert "银龄守护医疗险" in response.json()["data"]["assistant_message"]["text"]
    assert re.fullmatch(
        r"ZBT-\d{8}-\d{4}",
        response.json()["data"]["service_no"],
    )
    assert (
        response.json()["data"]["assistant_message"]["service_no"]
        == response.json()["data"]["service_no"]
    )
    assert response.json()["data"]["trace_id"] == "trace_test_001"
    run_id = response.json()["data"]["run_id"]

    with TestClient(app) as client:
        token = client.post(
            "/api/v1/h5/auth/login",
            json={"mobile": "18800000001", "code": "147258"},
        ).json()["data"]["tokens"]["access_token"]
        headers = {"Authorization": f"Bearer {token}"}
        messages = client.get(
            f"/api/v1/agent/threads/{thread_id}/messages",
            headers=headers,
        )
        run = client.get(f"/api/v1/agent/runs/{run_id}", headers=headers)

    assert [item["role"] for item in messages.json()["data"]["items"]] == [
        "USER",
        "ASSISTANT",
    ]
    assert run.json()["data"]["status"] == "COMPLETED"
    assert run.json()["data"]["trace_id"] == "trace_test_001"


def test_admin_can_use_operation_persona_with_the_same_agent_kernel_contract() -> None:
    settings = Settings(
        app_env="test",
        jwt_access_secret="a" * 32,
        jwt_refresh_secret="b" * 32,
        field_encryption_key="c" * 32,
    )
    identities = InMemoryIdentityRepository()
    identities.save_admin_user(
        AdminUser(
            id="01ADMIN0000000000000000001",
            username="admin",
            password_hash=PasswordService().hash("zaq1XSW@"),
            display_name="系统管理员",
            status="ACTIVE",
            roles=("SUPER_ADMIN",),
            permissions=("*",),
            data_scope="ALL",
        )
    )
    app = create_app(
        settings=settings,
        identity_repository=identities,
        catalog_repository=InMemoryCatalogRepository(),
        agent_repository=InMemoryAgentRepository(),
        agent_runtime=FixedOperationRuntime(),
        enrollment_repository=InMemoryEnrollmentRepository(),
        clock=lambda: datetime(2026, 7, 26, 8, tzinfo=UTC),
    )

    with TestClient(app) as client:
        token = client.post(
            "/api/v1/admin/auth/login",
            json={"username": "admin", "password": "zaq1XSW@"},
        ).json()["data"]["tokens"]["access_token"]
        headers = {"Authorization": f"Bearer {token}"}
        thread = client.post(
            "/api/v1/admin/agent/threads",
            headers=headers,
            json={"title": "订单经营分析"},
        )
        response = client.post(
            f"/api/v1/admin/agent/threads/{thread.json()['data']['id']}/messages",
            headers=headers,
            json={
                "content": {"type": "text", "text": "分析一下当前订单情况"},
                "client_message_id": "operation-message-001",
            },
        )
        runs = client.get("/api/v1/admin/agent/runs", headers=headers)

    assert thread.status_code == 201
    assert thread.json()["data"]["persona"] == "operation"
    assert response.status_code == 202
    data = response.json()["data"]
    assert data["invoked_tools"] == ["get_operation_overview"]
    assert data["trace_url"] == "https://smith.langchain.com/example-trace"
    assert data["assistant_message"]["cards"][0]["type"] == "metrics"
    assert runs.status_code == 200
    assert runs.json()["data"]["total"] == 1
    assert runs.json()["data"]["items"][0]["id"] == data["run_id"]


def test_h5_agent_stream_emits_and_replays_runtime_events() -> None:
    settings = Settings(
        app_env="test",
        jwt_access_secret="a" * 32,
        jwt_refresh_secret="b" * 32,
        field_encryption_key="c" * 32,
    )
    state_store = InMemoryAgentStateStore()
    app = create_app(
        settings=settings,
        identity_repository=InMemoryIdentityRepository(),
        catalog_repository=InMemoryCatalogRepository(),
        agent_repository=InMemoryAgentRepository(),
        agent_runtime=StreamingAgentRuntime(),
        enrollment_repository=InMemoryEnrollmentRepository(),
        agent_state_store=state_store,
    )

    with TestClient(app) as client:
        token = client.post(
            "/api/v1/h5/auth/login",
            json={"mobile": "18800000001", "code": "147258"},
        ).json()["data"]["tokens"]["access_token"]
        headers = {"Authorization": f"Bearer {token}"}
        thread_id = client.post(
            "/api/v1/agent/threads",
            headers=headers,
            json={"title": "SSE测试"},
        ).json()["data"]["id"]
        response = client.post(
            f"/api/v1/agent/threads/{thread_id}/messages/stream",
            headers=headers,
            json={
                "content": {"type": "text", "text": "推荐父母医疗险"},
                "client_message_id": "stream-message-001",
            },
        )

        assert response.status_code == 200
        assert response.headers["content-type"].startswith("text/event-stream")
        assert response.text.index("event: run.started") < response.text.index(
            "event: tool.completed"
        )
        assert response.text.index("event: tool.completed") < response.text.index("event: ui.ready")
        assert response.text.index("event: ui.ready") < response.text.index("event: run.completed")
        run_id = next(
            line.removeprefix("id: ").split(":")[0]
            for line in response.text.splitlines()
            if line.startswith("id: ")
        )
        replay = client.get(
            f"/api/v1/agent/runs/{run_id}/events",
            headers=headers,
        )

    assert replay.status_code == 200
    assert replay.text.count("event: ") == 4
    assert "list_available_products" in replay.text


def test_agent_rate_limit_is_enforced_before_model_invocation() -> None:
    settings = Settings(
        app_env="test",
        jwt_access_secret="a" * 32,
        jwt_refresh_secret="b" * 32,
        field_encryption_key="c" * 32,
    )
    app = create_app(
        settings=settings,
        identity_repository=InMemoryIdentityRepository(),
        catalog_repository=InMemoryCatalogRepository(),
        agent_repository=InMemoryAgentRepository(),
        agent_runtime=StreamingAgentRuntime(),
        enrollment_repository=InMemoryEnrollmentRepository(),
        agent_state_store=InMemoryAgentStateStore(rate_limit=1),
    )

    with TestClient(app) as client:
        token = client.post(
            "/api/v1/h5/auth/login",
            json={"mobile": "18800000001", "code": "147258"},
        ).json()["data"]["tokens"]["access_token"]
        headers = {"Authorization": f"Bearer {token}"}
        thread_id = client.post(
            "/api/v1/agent/threads",
            headers=headers,
            json={"title": "限流测试"},
        ).json()["data"]["id"]
        payload = {
            "content": {"type": "text", "text": "推荐医疗险"},
            "client_message_id": "rate-message",
        }
        first = client.post(
            f"/api/v1/agent/threads/{thread_id}/messages",
            headers=headers,
            json=payload,
        )
        second = client.post(
            f"/api/v1/agent/threads/{thread_id}/messages",
            headers=headers,
            json=payload,
        )

    assert first.status_code == 202
    assert second.status_code == 429
    assert second.json()["error"]["code"] == "AGENT_RATE_LIMITED"