test_agent_api.py 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534
  1. import re
  2. from datetime import UTC, datetime
  3. from fastapi.testclient import TestClient
  4. from zbt.core.config import Settings
  5. from zbt.core.errors import AppError
  6. from zbt.core.passwords import PasswordService
  7. from zbt.domains.agent.repository import InMemoryAgentRepository
  8. from zbt.domains.agent.runtime import AgentReply
  9. from zbt.domains.catalog.repository import InMemoryCatalogRepository
  10. from zbt.domains.enrollment.repository import InMemoryEnrollmentRepository
  11. from zbt.domains.identity.models import AdminUser
  12. from zbt.domains.identity.repository import InMemoryIdentityRepository
  13. from zbt.harness.kernel import AgentInvocation
  14. from zbt.infrastructure.redis.agent_state import InMemoryAgentStateStore
  15. from zbt.main import create_app
  16. class FixedAgentRuntime:
  17. def reply(self, invocation: AgentInvocation) -> AgentReply:
  18. assert "65岁" in invocation.message
  19. assert invocation.persona == "customer"
  20. return AgentReply(
  21. text="可以优先考虑银龄守护医疗险,我会继续核对地区和职业。",
  22. cards=[],
  23. trace_id="trace_test_001",
  24. )
  25. class FixedOperationRuntime:
  26. def reply(self, invocation: AgentInvocation) -> AgentReply:
  27. assert invocation.persona == "operation"
  28. assert invocation.admin_user is not None
  29. assert "订单" in invocation.message
  30. return AgentReply(
  31. text="当前共有24笔订单,其中18笔已出单。",
  32. cards=[
  33. {
  34. "type": "metrics",
  35. "version": "1.0",
  36. "title": "经营总览",
  37. "items": [
  38. {"label": "订单总量", "value": 24, "unit": "笔"},
  39. ],
  40. }
  41. ],
  42. trace_id="00000000-0000-0000-0000-000000000001",
  43. trace_url="https://smith.langchain.com/example-trace",
  44. invoked_tools=("get_operation_overview",),
  45. )
  46. class StreamingAgentRuntime:
  47. def reply(self, invocation: AgentInvocation) -> AgentReply:
  48. assert invocation.h5_user is not None
  49. return AgentReply(
  50. text="已为你找到银龄守护医疗险。",
  51. cards=[
  52. {
  53. "type": "product_recommendations",
  54. "version": "1.0",
  55. "title": "为你匹配的保障方案",
  56. "items": [],
  57. }
  58. ],
  59. trace_id="00000000-0000-0000-0000-000000000002",
  60. actions=[],
  61. invoked_tools=("list_available_products",),
  62. )
  63. class FailingAgentRuntime:
  64. def reply(self, invocation: AgentInvocation) -> AgentReply:
  65. raise AppError(
  66. "KNOWLEDGE_RERANK_FAILED",
  67. "知识重排失败",
  68. 503,
  69. retryable=True,
  70. )
  71. def test_h5_user_can_create_customer_agent_thread() -> None:
  72. settings = Settings(
  73. app_env="test",
  74. jwt_access_secret="a" * 32,
  75. jwt_refresh_secret="b" * 32,
  76. field_encryption_key="c" * 32,
  77. )
  78. app = create_app(
  79. settings=settings,
  80. identity_repository=InMemoryIdentityRepository(),
  81. agent_repository=InMemoryAgentRepository(),
  82. )
  83. with TestClient(app) as client:
  84. token = client.post(
  85. "/api/v1/h5/auth/login",
  86. json={"mobile": "18800000001", "code": "147258"},
  87. ).json()["data"]["tokens"]["access_token"]
  88. response = client.post(
  89. "/api/v1/agent/threads",
  90. headers={"Authorization": f"Bearer {token}"},
  91. json={"title": "给父亲配置医疗险"},
  92. )
  93. assert response.status_code == 201
  94. assert response.json()["data"]["persona"] == "customer"
  95. assert response.json()["data"]["title"] == "给父亲配置医疗险"
  96. def test_h5_user_can_send_message_and_receive_agent_reply() -> None:
  97. settings = Settings(
  98. app_env="test",
  99. jwt_access_secret="a" * 32,
  100. jwt_refresh_secret="b" * 32,
  101. field_encryption_key="c" * 32,
  102. )
  103. app = create_app(
  104. settings=settings,
  105. identity_repository=InMemoryIdentityRepository(),
  106. agent_repository=InMemoryAgentRepository(),
  107. agent_runtime=FixedAgentRuntime(),
  108. )
  109. with TestClient(app) as client:
  110. token = client.post(
  111. "/api/v1/h5/auth/login",
  112. json={"mobile": "18800000001", "code": "147258"},
  113. ).json()["data"]["tokens"]["access_token"]
  114. headers = {"Authorization": f"Bearer {token}"}
  115. thread_id = client.post(
  116. "/api/v1/agent/threads",
  117. headers=headers,
  118. json={"title": "给父亲配置医疗险"},
  119. ).json()["data"]["id"]
  120. response = client.post(
  121. f"/api/v1/agent/threads/{thread_id}/messages",
  122. headers=headers,
  123. json={
  124. "content": {
  125. "type": "text",
  126. "text": "想给65岁的父亲买医疗险",
  127. },
  128. "client_message_id": "client-message-001",
  129. },
  130. )
  131. assert response.status_code == 202
  132. assert response.json()["data"]["status"] == "COMPLETED"
  133. assert "银龄守护医疗险" in response.json()["data"]["assistant_message"]["text"]
  134. assert re.fullmatch(
  135. r"ZBT-\d{8}-\d{4}",
  136. response.json()["data"]["service_no"],
  137. )
  138. assert (
  139. response.json()["data"]["assistant_message"]["service_no"]
  140. == response.json()["data"]["service_no"]
  141. )
  142. assert response.json()["data"]["trace_id"] == "trace_test_001"
  143. run_id = response.json()["data"]["run_id"]
  144. with TestClient(app) as client:
  145. token = client.post(
  146. "/api/v1/h5/auth/login",
  147. json={"mobile": "18800000001", "code": "147258"},
  148. ).json()["data"]["tokens"]["access_token"]
  149. headers = {"Authorization": f"Bearer {token}"}
  150. messages = client.get(
  151. f"/api/v1/agent/threads/{thread_id}/messages",
  152. headers=headers,
  153. )
  154. run = client.get(f"/api/v1/agent/runs/{run_id}", headers=headers)
  155. assert [item["role"] for item in messages.json()["data"]["items"]] == [
  156. "USER",
  157. "ASSISTANT",
  158. ]
  159. assert run.json()["data"]["status"] == "COMPLETED"
  160. assert run.json()["data"]["trace_id"] == "trace_test_001"
  161. def test_h5_user_can_list_and_restore_own_conversation_history() -> None:
  162. settings = Settings(
  163. app_env="test",
  164. jwt_access_secret="a" * 32,
  165. jwt_refresh_secret="b" * 32,
  166. field_encryption_key="c" * 32,
  167. )
  168. repository = InMemoryAgentRepository()
  169. app = create_app(
  170. settings=settings,
  171. identity_repository=InMemoryIdentityRepository(),
  172. agent_repository=repository,
  173. agent_runtime=FixedAgentRuntime(),
  174. )
  175. with TestClient(app) as client:
  176. token = client.post(
  177. "/api/v1/h5/auth/login",
  178. json={"mobile": "18800000001", "code": "147258"},
  179. ).json()["data"]["tokens"]["access_token"]
  180. headers = {"Authorization": f"Bearer {token}"}
  181. thread_id = client.post(
  182. "/api/v1/agent/threads",
  183. headers=headers,
  184. json={"title": "父母医疗险咨询"},
  185. ).json()["data"]["id"]
  186. client.post(
  187. f"/api/v1/agent/threads/{thread_id}/messages",
  188. headers=headers,
  189. json={
  190. "content": {"type": "text", "text": "想给65岁的父亲买医疗险"},
  191. "client_message_id": "history-message-001",
  192. },
  193. )
  194. threads = client.get("/api/v1/agent/threads", headers=headers)
  195. messages = client.get(
  196. f"/api/v1/agent/threads/{thread_id}/messages",
  197. headers=headers,
  198. )
  199. assert threads.status_code == 200
  200. assert threads.json()["data"]["total"] == 1
  201. assert threads.json()["data"]["items"][0]["title"] == "父母医疗险咨询"
  202. assert threads.json()["data"]["items"][0]["message_count"] == 2
  203. assert threads.json()["data"]["items"][0]["last_message"]["role"] == "ASSISTANT"
  204. assert len(messages.json()["data"]["items"]) == 2
  205. def test_admin_can_use_operation_persona_with_the_same_agent_kernel_contract() -> None:
  206. settings = Settings(
  207. app_env="test",
  208. jwt_access_secret="a" * 32,
  209. jwt_refresh_secret="b" * 32,
  210. field_encryption_key="c" * 32,
  211. )
  212. identities = InMemoryIdentityRepository()
  213. identities.save_admin_user(
  214. AdminUser(
  215. id="01ADMIN0000000000000000001",
  216. username="admin",
  217. password_hash=PasswordService().hash("zaq1XSW@"),
  218. display_name="系统管理员",
  219. status="ACTIVE",
  220. roles=("SUPER_ADMIN",),
  221. permissions=("*",),
  222. data_scope="ALL",
  223. )
  224. )
  225. app = create_app(
  226. settings=settings,
  227. identity_repository=identities,
  228. catalog_repository=InMemoryCatalogRepository(),
  229. agent_repository=InMemoryAgentRepository(),
  230. agent_runtime=FixedOperationRuntime(),
  231. enrollment_repository=InMemoryEnrollmentRepository(),
  232. clock=lambda: datetime(2026, 7, 26, 8, tzinfo=UTC),
  233. )
  234. with TestClient(app) as client:
  235. token = client.post(
  236. "/api/v1/admin/auth/login",
  237. json={"username": "admin", "password": "zaq1XSW@"},
  238. ).json()["data"]["tokens"]["access_token"]
  239. headers = {"Authorization": f"Bearer {token}"}
  240. thread = client.post(
  241. "/api/v1/admin/agent/threads",
  242. headers=headers,
  243. json={"title": "订单经营分析"},
  244. )
  245. response = client.post(
  246. f"/api/v1/admin/agent/threads/{thread.json()['data']['id']}/messages",
  247. headers=headers,
  248. json={
  249. "content": {"type": "text", "text": "分析一下当前订单情况"},
  250. "client_message_id": "operation-message-001",
  251. },
  252. )
  253. runs = client.get("/api/v1/admin/agent/runs", headers=headers)
  254. assert thread.status_code == 201
  255. assert thread.json()["data"]["persona"] == "operation"
  256. assert response.status_code == 202
  257. data = response.json()["data"]
  258. assert data["invoked_tools"] == ["get_operation_overview"]
  259. assert data["trace_url"] == "https://smith.langchain.com/example-trace"
  260. assert data["assistant_message"]["cards"][0]["type"] == "metrics"
  261. assert runs.status_code == 200
  262. assert runs.json()["data"]["total"] == 1
  263. assert runs.json()["data"]["items"][0]["id"] == data["run_id"]
  264. def test_admin_can_list_only_own_operation_conversations() -> None:
  265. settings = Settings(
  266. app_env="test",
  267. jwt_access_secret="a" * 32,
  268. jwt_refresh_secret="b" * 32,
  269. field_encryption_key="c" * 32,
  270. )
  271. identities = InMemoryIdentityRepository()
  272. for user_id, username in (
  273. ("01ADMIN0000000000000000011", "admin-alpha"),
  274. ("01ADMIN0000000000000000012", "admin-beta"),
  275. ):
  276. identities.save_admin_user(
  277. AdminUser(
  278. id=user_id,
  279. username=username,
  280. password_hash=PasswordService().hash("zaq1XSW@"),
  281. display_name=username,
  282. status="ACTIVE",
  283. roles=("OPERATOR",),
  284. permissions=("dashboard:read",),
  285. data_scope="ALL",
  286. )
  287. )
  288. app = create_app(
  289. settings=settings,
  290. identity_repository=identities,
  291. catalog_repository=InMemoryCatalogRepository(),
  292. agent_repository=InMemoryAgentRepository(),
  293. enrollment_repository=InMemoryEnrollmentRepository(),
  294. )
  295. with TestClient(app) as client:
  296. alpha_token = client.post(
  297. "/api/v1/admin/auth/login",
  298. json={"username": "admin-alpha", "password": "zaq1XSW@"},
  299. ).json()["data"]["tokens"]["access_token"]
  300. beta_token = client.post(
  301. "/api/v1/admin/auth/login",
  302. json={"username": "admin-beta", "password": "zaq1XSW@"},
  303. ).json()["data"]["tokens"]["access_token"]
  304. alpha_headers = {"Authorization": f"Bearer {alpha_token}"}
  305. beta_headers = {"Authorization": f"Bearer {beta_token}"}
  306. alpha_thread = client.post(
  307. "/api/v1/admin/agent/threads",
  308. headers=alpha_headers,
  309. json={"title": "Alpha订单趋势"},
  310. ).json()["data"]
  311. client.post(
  312. "/api/v1/admin/agent/threads",
  313. headers=alpha_headers,
  314. json={"title": "评测-operation-analytics-trend"},
  315. )
  316. client.post(
  317. "/api/v1/admin/agent/threads",
  318. headers=alpha_headers,
  319. json={"title": "__EVAL__:operation-overview"},
  320. )
  321. beta_thread = client.post(
  322. "/api/v1/admin/agent/threads",
  323. headers=beta_headers,
  324. json={"title": "Beta保单趋势"},
  325. ).json()["data"]
  326. alpha_threads = client.get(
  327. "/api/v1/admin/agent/threads",
  328. headers=alpha_headers,
  329. )
  330. beta_threads = client.get(
  331. "/api/v1/admin/agent/threads",
  332. headers=beta_headers,
  333. )
  334. cross_account_messages = client.get(
  335. f"/api/v1/admin/agent/threads/{alpha_thread['id']}/messages",
  336. headers=beta_headers,
  337. )
  338. assert alpha_threads.status_code == 200
  339. assert [item["id"] for item in alpha_threads.json()["data"]["items"]] == [
  340. alpha_thread["id"]
  341. ]
  342. assert beta_threads.status_code == 200
  343. assert [item["id"] for item in beta_threads.json()["data"]["items"]] == [
  344. beta_thread["id"]
  345. ]
  346. assert cross_account_messages.status_code == 404
  347. assert cross_account_messages.json()["error"]["code"] == "AGENT_THREAD_NOT_FOUND"
  348. def test_h5_agent_stream_emits_and_replays_runtime_events() -> None:
  349. settings = Settings(
  350. app_env="test",
  351. jwt_access_secret="a" * 32,
  352. jwt_refresh_secret="b" * 32,
  353. field_encryption_key="c" * 32,
  354. )
  355. state_store = InMemoryAgentStateStore()
  356. app = create_app(
  357. settings=settings,
  358. identity_repository=InMemoryIdentityRepository(),
  359. catalog_repository=InMemoryCatalogRepository(),
  360. agent_repository=InMemoryAgentRepository(),
  361. agent_runtime=StreamingAgentRuntime(),
  362. enrollment_repository=InMemoryEnrollmentRepository(),
  363. agent_state_store=state_store,
  364. )
  365. with TestClient(app) as client:
  366. token = client.post(
  367. "/api/v1/h5/auth/login",
  368. json={"mobile": "18800000001", "code": "147258"},
  369. ).json()["data"]["tokens"]["access_token"]
  370. headers = {"Authorization": f"Bearer {token}"}
  371. thread_id = client.post(
  372. "/api/v1/agent/threads",
  373. headers=headers,
  374. json={"title": "SSE测试"},
  375. ).json()["data"]["id"]
  376. response = client.post(
  377. f"/api/v1/agent/threads/{thread_id}/messages/stream",
  378. headers=headers,
  379. json={
  380. "content": {"type": "text", "text": "推荐父母医疗险"},
  381. "client_message_id": "stream-message-001",
  382. },
  383. )
  384. assert response.status_code == 200
  385. assert response.headers["content-type"].startswith("text/event-stream")
  386. assert response.text.index("event: run.started") < response.text.index(
  387. "event: tool.completed"
  388. )
  389. assert response.text.index("event: tool.completed") < response.text.index("event: ui.ready")
  390. assert response.text.index("event: ui.ready") < response.text.index("event: run.completed")
  391. run_id = next(
  392. line.removeprefix("id: ").split(":")[0]
  393. for line in response.text.splitlines()
  394. if line.startswith("id: ")
  395. )
  396. replay = client.get(
  397. f"/api/v1/agent/runs/{run_id}/events",
  398. headers=headers,
  399. )
  400. assert replay.status_code == 200
  401. assert replay.text.count("event: ") == 4
  402. assert "list_available_products" in replay.text
  403. def test_failed_agent_run_does_not_leave_a_dangling_user_message() -> None:
  404. settings = Settings(
  405. app_env="test",
  406. jwt_access_secret="a" * 32,
  407. jwt_refresh_secret="b" * 32,
  408. field_encryption_key="c" * 32,
  409. )
  410. app = create_app(
  411. settings=settings,
  412. identity_repository=InMemoryIdentityRepository(),
  413. catalog_repository=InMemoryCatalogRepository(),
  414. agent_repository=InMemoryAgentRepository(),
  415. agent_runtime=FailingAgentRuntime(),
  416. enrollment_repository=InMemoryEnrollmentRepository(),
  417. agent_state_store=InMemoryAgentStateStore(),
  418. )
  419. with TestClient(app) as client:
  420. token = client.post(
  421. "/api/v1/h5/auth/login",
  422. json={"mobile": "18800000001", "code": "147258"},
  423. ).json()["data"]["tokens"]["access_token"]
  424. headers = {"Authorization": f"Bearer {token}"}
  425. thread_id = client.post(
  426. "/api/v1/agent/threads",
  427. headers=headers,
  428. json={"title": "失败消息回滚"},
  429. ).json()["data"]["id"]
  430. response = client.post(
  431. f"/api/v1/agent/threads/{thread_id}/messages",
  432. headers=headers,
  433. json={
  434. "content": {"type": "text", "text": "医疗险等待期是多少天?"},
  435. "client_message_id": "failed-message-001",
  436. },
  437. )
  438. messages = client.get(
  439. f"/api/v1/agent/threads/{thread_id}/messages",
  440. headers=headers,
  441. )
  442. assert response.status_code == 503
  443. assert response.json()["error"]["code"] == "KNOWLEDGE_RERANK_FAILED"
  444. assert messages.json()["data"]["items"] == []
  445. def test_agent_rate_limit_is_enforced_before_model_invocation() -> None:
  446. settings = Settings(
  447. app_env="test",
  448. jwt_access_secret="a" * 32,
  449. jwt_refresh_secret="b" * 32,
  450. field_encryption_key="c" * 32,
  451. )
  452. app = create_app(
  453. settings=settings,
  454. identity_repository=InMemoryIdentityRepository(),
  455. catalog_repository=InMemoryCatalogRepository(),
  456. agent_repository=InMemoryAgentRepository(),
  457. agent_runtime=StreamingAgentRuntime(),
  458. enrollment_repository=InMemoryEnrollmentRepository(),
  459. agent_state_store=InMemoryAgentStateStore(rate_limit=1),
  460. )
  461. with TestClient(app) as client:
  462. token = client.post(
  463. "/api/v1/h5/auth/login",
  464. json={"mobile": "18800000001", "code": "147258"},
  465. ).json()["data"]["tokens"]["access_token"]
  466. headers = {"Authorization": f"Bearer {token}"}
  467. thread_id = client.post(
  468. "/api/v1/agent/threads",
  469. headers=headers,
  470. json={"title": "限流测试"},
  471. ).json()["data"]["id"]
  472. payload = {
  473. "content": {"type": "text", "text": "推荐医疗险"},
  474. "client_message_id": "rate-message",
  475. }
  476. first = client.post(
  477. f"/api/v1/agent/threads/{thread_id}/messages",
  478. headers=headers,
  479. json=payload,
  480. )
  481. second = client.post(
  482. f"/api/v1/agent/threads/{thread_id}/messages",
  483. headers=headers,
  484. json=payload,
  485. )
  486. assert first.status_code == 202
  487. assert second.status_code == 429
  488. assert second.json()["error"]["code"] == "AGENT_RATE_LIMITED"