test_core_units.py 39 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006
  1. """核心单元测试:跟踪器、计量器、来源、提示词、工具、并发、计费、存储。
  2. 这些用例全部离线可跑(不起服务、不连模型)。端到端的 HTTP 验证见
  3. ``scripts/http_e2e.py``(需要一个真实可用的模型密钥)。
  4. """
  5. from __future__ import annotations
  6. import json
  7. import pytest
  8. # ============================================================
  9. # 子智能体跟踪
  10. # ============================================================
  11. class TestSubagentTracker:
  12. def test_namespace_prefix_parsing(self):
  13. from runtime.subagent_tracker import parse_namespace
  14. assert parse_namespace(("tools:abc123",)) == "abc123"
  15. assert parse_namespace(["tools:x"]) == "x"
  16. assert parse_namespace("tools:y") == "y"
  17. def test_empty_namespace_is_main_agent(self):
  18. from runtime.subagent_tracker import parse_namespace
  19. assert parse_namespace(()) is None
  20. assert parse_namespace(None) is None
  21. assert parse_namespace([]) is None
  22. def test_control_segments_not_mistaken_for_ids(self):
  23. """带冒号/斜杠的控制段不能被当成 tool_call id。"""
  24. from runtime.subagent_tracker import parse_namespace
  25. assert parse_namespace(("model_request:1",)) is None
  26. assert parse_namespace(("a/b",)) is None
  27. def test_lifecycle_started_to_done(self):
  28. from runtime.subagent_tracker import SubagentTracker
  29. tracker = SubagentTracker()
  30. calls = [{"name": "task", "id": "c1", "args": {"subagent_type": "demo-searcher", "description": "查资料"}}]
  31. started = tracker.note_task_calls(calls)
  32. assert len(started) == 1
  33. assert started[0].subagent_type == "demo-searcher"
  34. assert started[0].status == "started"
  35. run = tracker.note_namespace(("tools:c1",))
  36. assert run is not None and run.status == "running" and run.tool_calls == 1
  37. done = tracker.note_result("c1", ok=True)
  38. assert done is not None and done.status == "done"
  39. assert done.duration_ms is not None
  40. def test_non_task_calls_ignored(self):
  41. from runtime.subagent_tracker import SubagentTracker
  42. tracker = SubagentTracker()
  43. assert tracker.note_task_calls([{"name": "calculate", "id": "c1", "args": {}}]) == []
  44. assert len(tracker) == 0
  45. def test_call_without_id_skipped(self):
  46. """没有 id 就无法关联后续子图事件,跳过而不是编一个假 id。"""
  47. from runtime.subagent_tracker import SubagentTracker
  48. tracker = SubagentTracker()
  49. assert tracker.note_task_calls([{"name": "task", "args": {"subagent_type": "x"}}]) == []
  50. assert len(tracker) == 0
  51. def test_unknown_namespace_does_not_invent_run(self):
  52. """表里没有的 id 不能凭空造一条 subagent_type=unknown 的记录。"""
  53. from runtime.subagent_tracker import SubagentTracker
  54. tracker = SubagentTracker()
  55. assert tracker.note_namespace(("tools:never-seen",)) is None
  56. assert len(tracker) == 0
  57. def test_finish_all_flushes_active(self):
  58. """流异常退出时不能让子智能体永远挂在 started,前端卡片会一直转。"""
  59. from runtime.subagent_tracker import SubagentTracker
  60. tracker = SubagentTracker()
  61. tracker.note_task_calls([
  62. {"name": "task", "id": "c1", "args": {"subagent_type": "a"}},
  63. {"name": "task", "id": "c2", "args": {"subagent_type": "b"}},
  64. ])
  65. tracker.note_result("c1", ok=True)
  66. flushed = tracker.finish_all()
  67. assert len(flushed) == 1 and flushed[0].call_id == "c2"
  68. assert not tracker.active()
  69. def test_duplicate_task_call_not_counted_twice(self):
  70. from runtime.subagent_tracker import SubagentTracker
  71. tracker = SubagentTracker()
  72. call = {"name": "task", "id": "c1", "args": {"subagent_type": "a"}}
  73. tracker.note_task_calls([call])
  74. tracker.note_task_calls([call])
  75. assert len(tracker) == 1
  76. # ============================================================
  77. # 上下文计量
  78. # ============================================================
  79. class _Msg:
  80. """最小的消息替身。"""
  81. def __init__(self, content="", usage_metadata=None, response_metadata=None,
  82. additional_kwargs=None, name=None) -> None:
  83. self.content = content
  84. self.usage_metadata = usage_metadata
  85. self.response_metadata = response_metadata or {}
  86. self.additional_kwargs = additional_kwargs or {}
  87. self.name = name
  88. class TestContextMeter:
  89. def test_extract_usage_from_langchain_metadata(self):
  90. from runtime.context_meter import extract_usage
  91. usage = extract_usage(_Msg(usage_metadata={
  92. "input_tokens": 100, "output_tokens": 20, "total_tokens": 120,
  93. "input_token_details": {"cache_read": 30},
  94. }))
  95. assert usage == {"input_tokens": 100, "output_tokens": 20, "cache_read_tokens": 30}
  96. def test_extract_usage_from_provider_raw(self):
  97. from runtime.context_meter import extract_usage
  98. usage = extract_usage(_Msg(response_metadata={
  99. "token_usage": {"prompt_tokens": 50, "completion_tokens": 5,
  100. "prompt_tokens_details": {"cached_tokens": 8}}
  101. }))
  102. assert usage == {"input_tokens": 50, "output_tokens": 5, "cache_read_tokens": 8}
  103. def test_missing_usage_returns_none_not_zero(self):
  104. """「用量为 0」与「没有用量信息」必须能区分,否则估算兜底会失效。"""
  105. from runtime.context_meter import extract_usage
  106. assert extract_usage(_Msg()) is None
  107. assert extract_usage(_Msg(usage_metadata={})) is None
  108. def test_accumulates_across_model_calls(self):
  109. from runtime.context_meter import ContextMeter
  110. meter = ContextMeter(window=1000)
  111. meter.note_message(_Msg(usage_metadata={"input_tokens": 100, "output_tokens": 10}))
  112. meter.note_message(_Msg(usage_metadata={"input_tokens": 150, "output_tokens": 20}))
  113. assert meter.input_tokens == 250
  114. assert meter.output_tokens == 30
  115. assert meter.model_calls == 2
  116. def test_context_tokens_is_last_input(self):
  117. """上下文占用取最近一次调用的 input,而不是累加值。"""
  118. from runtime.context_meter import ContextMeter
  119. meter = ContextMeter(window=1000)
  120. meter.note_message(_Msg(usage_metadata={"input_tokens": 100, "output_tokens": 10}))
  121. meter.note_message(_Msg(usage_metadata={"input_tokens": 180, "output_tokens": 10}))
  122. assert meter.context_tokens == 180
  123. assert meter.input_tokens == 280
  124. def test_source_is_provider_when_usage_present(self):
  125. from runtime.context_meter import ContextMeter
  126. meter = ContextMeter(window=1000)
  127. meter.note_message(_Msg(usage_metadata={"input_tokens": 5, "output_tokens": 5}))
  128. assert meter.usage_source == "provider"
  129. def test_estimate_marks_source_estimate(self):
  130. from runtime.context_meter import ContextMeter
  131. meter = ContextMeter(window=1000)
  132. meter.note_estimated(300)
  133. assert meter.usage_source == "estimate"
  134. assert meter.input_tokens == 300
  135. def test_estimate_ignored_when_real_usage_seen(self):
  136. """已有真实用量时估算值只会污染口径,丢弃。"""
  137. from runtime.context_meter import ContextMeter
  138. meter = ContextMeter(window=1000)
  139. meter.note_message(_Msg(usage_metadata={"input_tokens": 100, "output_tokens": 10}))
  140. meter.note_estimated(9999)
  141. assert meter.input_tokens == 100
  142. assert meter.usage_source == "provider"
  143. def test_compaction_detected_by_lc_source(self):
  144. from runtime.context_meter import ContextMeter
  145. meter = ContextMeter(window=1000)
  146. meter.note_message(_Msg(content="摘要", response_metadata={"lc_source": "summarization"}))
  147. assert meter.compacted and meter.compact_count == 1
  148. def test_compaction_detected_by_context_drop(self):
  149. """占用断崖式下跌也视为压缩。"""
  150. from runtime.context_meter import ContextMeter
  151. meter = ContextMeter(window=10000)
  152. meter.note_message(_Msg(usage_metadata={"input_tokens": 9000, "output_tokens": 10}))
  153. assert not meter.compacted
  154. meter.note_message(_Msg(usage_metadata={"input_tokens": 2000, "output_tokens": 10}))
  155. assert meter.compacted and meter.compact_count == 1
  156. def test_compaction_counted_once(self):
  157. """两个信号同时命中不能重复计数。"""
  158. from runtime.context_meter import ContextMeter
  159. meter = ContextMeter(window=10000)
  160. meter.note_message(_Msg(usage_metadata={"input_tokens": 9000, "output_tokens": 1}))
  161. meter.note_message(_Msg(usage_metadata={"input_tokens": 100, "output_tokens": 1}))
  162. meter.note_message(_Msg(content="x", response_metadata={"lc_source": "summarization"}))
  163. assert meter.compact_count == 1
  164. def test_archive_text_captured_on_compaction(self):
  165. from runtime.context_meter import ContextMeter
  166. meter = ContextMeter(window=1000)
  167. meter.note_text("被压缩掉的早期对话")
  168. meter.note_message(_Msg(content="摘要", response_metadata={"lc_source": "summarization"}))
  169. snapshot = meter.snapshot()
  170. assert snapshot.compacted and snapshot.archive_text == "被压缩掉的早期对话"
  171. def test_fill_ratio(self):
  172. from runtime.context_meter import ContextMeter
  173. meter = ContextMeter(window=1000)
  174. meter.note_message(_Msg(usage_metadata={"input_tokens": 850, "output_tokens": 0}))
  175. assert meter.fill_ratio() == pytest.approx(0.85)
  176. def test_estimate_tokens_positive_for_cjk(self):
  177. from runtime.context_meter import estimate_tokens
  178. assert estimate_tokens("矿井通风") > 0
  179. assert estimate_tokens("") == 0
  180. # ============================================================
  181. # 来源标注
  182. # ============================================================
  183. class TestSourceTracker:
  184. def test_extract_from_sourced_payload(self):
  185. from runtime.source_tracker import SourceTracker
  186. from tools.registry import source, sourced
  187. tracker = SourceTracker()
  188. payload = json.dumps(sourced(data=[1], sources=[source("T", "https://a.invalid/x", "摘要", 0.5)]))
  189. added = tracker.note_tool_result("search_demo_knowledge", payload)
  190. assert len(added) == 1 and added[0]["title"] == "T"
  191. def test_extract_from_dict_payload(self):
  192. from runtime.source_tracker import SourceTracker
  193. from tools.registry import source, sourced
  194. tracker = SourceTracker()
  195. tracker.note_tool_result("t", sourced(data=[], sources=[source("T", "https://a.invalid")]))
  196. assert tracker.count == 1
  197. def test_dedup_by_url(self):
  198. from runtime.source_tracker import SourceTracker
  199. from tools.registry import source, sourced
  200. tracker = SourceTracker()
  201. payload = sourced(data=[], sources=[source("T", "https://same.invalid"), source("T2", "https://same.invalid")])
  202. tracker.note_tool_result("t", payload)
  203. assert tracker.count == 1
  204. def test_records_without_title_or_url_dropped(self):
  205. """既无标题也无链接的来源无法核实,不展示空壳。"""
  206. from runtime.source_tracker import SourceTracker
  207. from tools.registry import sourced
  208. tracker = SourceTracker()
  209. tracker.note_tool_result("t", sourced(data=[], sources=[{"snippet": "只有摘要"}]))
  210. assert tracker.count == 0
  211. def test_plain_text_yields_nothing(self):
  212. from runtime.source_tracker import SourceTracker
  213. tracker = SourceTracker()
  214. tracker.note_tool_result("t", "这是一段普通文本,没有 sources")
  215. assert tracker.count == 0
  216. def test_classification_web_vs_kb(self):
  217. from runtime.source_tracker import SourceTracker
  218. from tools.registry import source, sourced
  219. tracker = SourceTracker()
  220. tracker.note_tool_result("t", sourced(data=[], sources=[
  221. source("网页", "https://news.example.com/a"),
  222. source("知识库", "https://kb.example.com/b"),
  223. ]))
  224. types = {i["title"]: i["type"] for i in tracker.build_citations()}
  225. assert types["网页"] == "web"
  226. assert types["知识库"] == "knowledge_base"
  227. def test_nested_payload_lookup(self):
  228. from runtime.source_tracker import SourceTracker
  229. from tools.registry import source
  230. tracker = SourceTracker()
  231. tracker.note_tool_result("t", {"data": {"sources": [source("深层", "https://d.invalid")]}})
  232. assert tracker.count == 1 and tracker.build_citations()[0]["title"] == "深层"
  233. def test_citations_sorted_by_score(self):
  234. from runtime.source_tracker import SourceTracker
  235. from tools.registry import source, sourced
  236. tracker = SourceTracker()
  237. tracker.note_tool_result("t", sourced(data=[], sources=[
  238. source("低", "https://a.invalid", score=0.1),
  239. source("高", "https://b.invalid", score=0.9),
  240. ]))
  241. assert [i["title"] for i in tracker.build_citations()] == ["高", "低"]
  242. def test_label_attached(self):
  243. from runtime.source_tracker import SourceTracker
  244. from tools.registry import source, sourced
  245. tracker = SourceTracker()
  246. tracker.note_tool_result("search_demo_knowledge",
  247. sourced(data=[], sources=[source("T", "https://a.invalid")]))
  248. item = tracker.build_citations()[0]
  249. assert item["tool"] == "search_demo_knowledge"
  250. assert item["tool_cn"] == "检索演示知识库"
  251. # ============================================================
  252. # 提示词与消息前缀
  253. # ============================================================
  254. class TestPrompts:
  255. def test_skill_prefix_extracted(self):
  256. from agents.prompts import extract_skill_prefix
  257. text, skill = extract_skill_prefix("#skill:demo-skill 帮我写一份简报")
  258. assert skill == "demo-skill" and text == "帮我写一份简报"
  259. def test_skill_prefix_only_at_start(self):
  260. """正文中间的 #skill: 是普通文本,不能当指令。"""
  261. from agents.prompts import extract_skill_prefix
  262. text, skill = extract_skill_prefix("请看这个 #skill:demo-skill 的用法")
  263. assert skill is None and text.startswith("请看这个")
  264. def test_skill_prefix_without_message(self):
  265. from agents.prompts import extract_skill_prefix
  266. text, skill = extract_skill_prefix("#skill:demo-skill")
  267. assert skill == "demo-skill" and text == ""
  268. def test_path_traversal_rejected(self):
  269. """技能名带路径分隔符不能当技能解析——否则能读任意文件。"""
  270. from agents.prompts import extract_skill_prefix
  271. for bad in ("#skill:../../etc/passwd x", "#skill:.hidden x", "#skill:a/b x"):
  272. _, skill = extract_skill_prefix(bad)
  273. assert skill is None, bad
  274. def test_session_ref_extracted(self):
  275. from agents.prompts import extract_session_ref
  276. text, ref = extract_session_ref("#session:abc123 接着上次说")
  277. assert ref == "abc123" and text == "接着上次说"
  278. def test_both_prefixes_stacked(self):
  279. from agents.prompts import extract_session_ref, extract_skill_prefix
  280. text, skill = extract_skill_prefix("#skill:demo-skill #session:abc 继续")
  281. text, ref = extract_session_ref(text)
  282. assert skill == "demo-skill" and ref == "abc" and text == "继续"
  283. def test_mode_instruction_differs(self):
  284. from agents.prompts import build_mode_instruction
  285. plan = build_mode_instruction("plan")
  286. full = build_mode_instruction("full")
  287. assert "plan" in plan and "计划" in plan
  288. assert "full" in full and plan != full
  289. def test_preference_instruction_empty_when_no_prefs(self):
  290. from agents.prompts import build_preference_instruction
  291. assert build_preference_instruction("") == ""
  292. def test_time_block_format(self):
  293. from agents.prompts import build_time_block
  294. block = build_time_block()
  295. assert "## 系统时间" in block and "当前时间" in block
  296. def test_system_prompt_includes_subagents(self):
  297. from agents.prompts import build_system_prompt
  298. prompt = build_system_prompt([
  299. {"name": "demo-searcher", "description": "检索专员"}
  300. ])
  301. assert "demo-searcher" in prompt
  302. assert "检索专员" in prompt
  303. def test_system_prompt_without_subagents(self):
  304. from agents.prompts import build_system_prompt
  305. prompt = build_system_prompt([])
  306. assert "子智能体" in prompt # 通用说明里提到,但不该有具体名单
  307. assert "可委派的子智能体" not in prompt
  308. # ============================================================
  309. # 工具
  310. # ============================================================
  311. class TestTools:
  312. def test_builtin_tools_registered(self):
  313. from tools import registry
  314. names = set(registry.get_tool_names())
  315. assert {"get_current_time", "search_demo_knowledge", "calculate",
  316. "request_plan_approval"} <= names
  317. def test_tools_are_basetool_instances(self):
  318. """子智能体的 tools 必须是 BaseTool(ToolNode 要读 .name)。"""
  319. from tools import registry
  320. for tool in registry.get_all_tools():
  321. assert hasattr(tool, "name") and hasattr(tool, "description")
  322. def test_get_tools_by_name_skips_unknown(self):
  323. from tools import registry
  324. tools = registry.get_tools_by_name(["calculate", "does-not-exist"])
  325. assert [t.name for t in tools] == ["calculate"]
  326. def test_duplicate_registration_raises(self):
  327. from tools.registry import register, reset_for_tests
  328. reset_for_tests()
  329. @register(group="X")
  330. def dup_tool() -> str:
  331. """d"""
  332. return ""
  333. with pytest.raises(ValueError, match="工具名冲突"):
  334. @register(group="Y")
  335. def dup_tool() -> str: # noqa: F811
  336. """d"""
  337. return ""
  338. reset_for_tests()
  339. from tools import registry as r
  340. r.load_builtin_tools()
  341. def test_calculate_basic(self):
  342. from tools.demo_tools import calculate
  343. assert calculate("(128000 * 0.85)")["result"] == 108800
  344. def test_calculate_integer_result_is_int(self):
  345. from tools.demo_tools import calculate
  346. assert isinstance(calculate("1+1")["result"], int)
  347. def test_calculate_rejects_code_execution(self):
  348. """表达式求值绝不能变成代码执行入口。"""
  349. from tools.demo_tools import calculate
  350. for evil in [
  351. "__import__('os').system('echo x')",
  352. "open('/etc/passwd').read()",
  353. "(lambda: 1)()",
  354. "1 if True else 2",
  355. "[x for x in range(3)]",
  356. ]:
  357. result = calculate(evil)
  358. assert "error" in result, evil
  359. def test_calculate_rejects_huge_power(self):
  360. from tools.demo_tools import calculate
  361. assert "error" in calculate("9 ** 9 ** 9")
  362. def test_calculate_zero_division(self):
  363. from tools.demo_tools import calculate
  364. assert "零" in calculate("1/0")["error"]
  365. def test_calculate_empty(self):
  366. from tools.demo_tools import calculate
  367. assert "error" in calculate("")
  368. def test_calculate_returns_error_not_raises(self):
  369. """工具失败要返回可读原因,抛异常会打断整轮工具调用。"""
  370. from tools.demo_tools import calculate
  371. result = calculate("1 +")
  372. assert isinstance(result, dict) and "error" in result
  373. def test_search_returns_sources_contract(self):
  374. from tools.demo_tools import search_demo_knowledge
  375. result = search_demo_knowledge("需风量")
  376. assert result["data"] and result["sources"]
  377. assert result["sources"][0]["url"]
  378. def test_search_empty_query_returns_error(self):
  379. from tools.demo_tools import search_demo_knowledge
  380. result = search_demo_knowledge(" ")
  381. assert result["error"] and result["sources"] == []
  382. def test_search_no_match(self):
  383. from tools.demo_tools import search_demo_knowledge
  384. result = search_demo_knowledge("不存在的关键词zzz")
  385. assert result["data"] == [] and result["sources"] == []
  386. def test_registry_describe_shape(self):
  387. from tools import registry
  388. described = registry.describe()
  389. assert described["count"] >= 4
  390. assert all(isinstance(v, list) for v in described["groups"].values())
  391. # ============================================================
  392. # 并发闸
  393. # ============================================================
  394. class TestConcurrency:
  395. def setup_method(self):
  396. from runtime import concurrency
  397. concurrency.reset_for_tests()
  398. def teardown_method(self):
  399. from runtime import concurrency
  400. concurrency.reset_for_tests()
  401. def test_acquire_and_release(self):
  402. from runtime import concurrency
  403. slot = concurrency.acquire_task_slot("u1")
  404. assert concurrency.stats()["global_active"] == 1
  405. assert slot.release() is True
  406. assert concurrency.stats()["global_active"] == 0
  407. def test_release_is_idempotent(self):
  408. """SSE 有多条退出路径,重复 release 不能把计数减成负数。"""
  409. from runtime import concurrency
  410. slot = concurrency.acquire_task_slot("u1")
  411. slot.release()
  412. slot.release()
  413. assert concurrency.stats()["global_active"] == 0
  414. def test_per_user_limit(self, monkeypatch):
  415. from runtime import concurrency
  416. monkeypatch.setattr(concurrency.config, "MAX_USER_TASKS", 1)
  417. concurrency.acquire_task_slot("u1")
  418. with pytest.raises(concurrency.TooManyTasksError) as exc:
  419. concurrency.acquire_task_slot("u1")
  420. assert exc.value.scope == "user"
  421. def test_global_limit(self, monkeypatch):
  422. from runtime import concurrency
  423. monkeypatch.setattr(concurrency.config, "MAX_CONCURRENT_TASKS", 1)
  424. monkeypatch.setattr(concurrency.config, "MAX_USER_TASKS", 10)
  425. concurrency.acquire_task_slot("u1")
  426. with pytest.raises(concurrency.TooManyTasksError) as exc:
  427. concurrency.acquire_task_slot("u2")
  428. assert exc.value.scope == "global"
  429. def test_different_users_independent(self, monkeypatch):
  430. from runtime import concurrency
  431. monkeypatch.setattr(concurrency.config, "MAX_USER_TASKS", 1)
  432. concurrency.acquire_task_slot("u1")
  433. concurrency.acquire_task_slot("u2")
  434. assert concurrency.stats()["global_active"] == 2
  435. # ============================================================
  436. # 计费
  437. # ============================================================
  438. class TestBilling:
  439. def test_cost_computation(self, pricing_ready):
  440. from billing import pricing
  441. from db import billing_store
  442. billing_store.set_pricing("m1", 0.001, 0.002)
  443. cost, currency = pricing.compute_cost("m1", 1000, 500)
  444. assert cost == pytest.approx(0.001 + 0.001)
  445. assert currency == "CNY"
  446. def test_unpriced_model_fails_closed(self, pricing_ready, monkeypatch):
  447. """静默按 0 计费等于白送额度,所以默认必须报错。"""
  448. from billing import pricing
  449. monkeypatch.setenv("BILLING_ENABLED", "1")
  450. monkeypatch.setenv("BILLING_ALLOW_UNPRICED", "0")
  451. with pytest.raises(pricing.UnpricedModelError):
  452. pricing.compute_cost("unknown-model", 100, 100)
  453. def test_unpriced_model_allowed_when_configured(self, pricing_ready, monkeypatch):
  454. from billing import pricing
  455. monkeypatch.setenv("BILLING_ENABLED", "1")
  456. monkeypatch.setenv("BILLING_ALLOW_UNPRICED", "1")
  457. cost, _ = pricing.compute_cost("unknown-model", 100, 100)
  458. assert cost == 0.0
  459. def test_quota_blocks_when_exhausted(self, pricing_ready, unique_user):
  460. from billing import guard
  461. from db import billing_store, db_conn
  462. day = db_conn.day_key()
  463. billing_store.set_quota_limit(unique_user, day, 100)
  464. billing_store.record_usage(unique_user, "s", "chat", "m", 100, 0, 0.0, day)
  465. with pytest.raises(guard.QuotaExceededError):
  466. guard.precheck(unique_user, day)
  467. def test_quota_zero_means_unlimited(self, pricing_ready, unique_user):
  468. from billing import guard
  469. from db import billing_store, db_conn
  470. day = db_conn.day_key()
  471. billing_store.set_quota_limit(unique_user, day, 0)
  472. billing_store.record_usage(unique_user, "s", "chat", "m", 10_000_000, 0, 0.0, day)
  473. guard.precheck(unique_user, day) # 不抛即通过
  474. def test_meter_settle_idempotent(self, pricing_ready, unique_user):
  475. from billing import guard
  476. from billing.meter import UsageMeter
  477. meter = UsageMeter(unique_user, "s1", "chat", "deepseek-flash")
  478. meter.add(1000, 100)
  479. first = meter.settle()
  480. second = meter.settle()
  481. assert first["settled"] is True
  482. assert second.get("duplicated") is True
  483. assert guard.snapshot(unique_user)["requests"] == 1
  484. def test_meter_no_usage_skips(self, pricing_ready, unique_user):
  485. from billing import guard
  486. from billing.meter import UsageMeter
  487. meter = UsageMeter(unique_user, "s1", "chat", "deepseek-flash")
  488. result = meter.settle()
  489. assert result["settled"] is False and result["reason"] == "no_usage"
  490. assert guard.snapshot(unique_user)["requests"] == 0
  491. def test_meter_source_downgrades_to_estimate(self, pricing_ready, unique_user):
  492. from billing.meter import UsageMeter
  493. meter = UsageMeter(unique_user, "s1", "chat", "deepseek-flash")
  494. meter.add(100, 10, "provider")
  495. meter.add(50, 5, "estimate")
  496. assert meter.usage_source == "estimate"
  497. def test_snapshot_has_total_tokens(self, pricing_ready, unique_user):
  498. """字段名与 billing_store.usage_for_day 对齐,消费方不用记两套名字。"""
  499. from billing import guard
  500. snapshot = guard.snapshot(unique_user)
  501. assert "total_tokens" in snapshot and "used_tokens" in snapshot
  502. def test_usage_day_bound_to_request_start(self, pricing_ready, unique_user):
  503. """跨零点的长请求不能被算到第二天。"""
  504. from billing import guard
  505. from billing.meter import UsageMeter
  506. from db import db_conn
  507. yesterday = "20200101"
  508. meter = UsageMeter(unique_user, "s1", "chat", "deepseek-flash", usage_day=yesterday)
  509. meter.add(100, 10)
  510. meter.settle()
  511. assert guard.snapshot(unique_user, yesterday)["requests"] == 1
  512. assert guard.snapshot(unique_user, db_conn.day_key())["requests"] == 0
  513. # ============================================================
  514. # 存储
  515. # ============================================================
  516. class TestStores:
  517. def test_message_ordering_stable_within_same_second(self):
  518. """同秒写入的多条消息顺序必须稳定——排序错乱会让补种的上下文角色错位。"""
  519. from db import chat_store
  520. sid = chat_store.create_session("u", "full")["id"]
  521. for i in range(6):
  522. chat_store.save_message(sid, "user", f"Q{i}")
  523. chat_store.save_message(sid, "assistant", f"A{i}")
  524. rows = chat_store.get_messages(sid)
  525. assert [r["msg_role"] for r in rows] == ["user", "assistant"] * 6
  526. assert [r["content"] for r in rows] == [
  527. c for i in range(6) for c in (f"Q{i}", f"A{i}")
  528. ]
  529. def test_get_messages_limit_returns_most_recent_in_order(self):
  530. from db import chat_store
  531. sid = chat_store.create_session("u", "full")["id"]
  532. for i in range(5):
  533. chat_store.save_message(sid, "user", f"M{i}")
  534. recent = chat_store.get_messages(sid, limit=3)
  535. assert [r["content"] for r in recent] == ["M2", "M3", "M4"]
  536. def test_message_json_fields_roundtrip(self):
  537. from db import chat_store
  538. sid = chat_store.create_session("u", "full")["id"]
  539. chat_store.save_message(
  540. sid, "assistant", "答",
  541. tool_calls=[{"name": "calculate", "args": {"expression": "1+1"}}],
  542. citations=[{"title": "T", "url": "u"}],
  543. )
  544. row = chat_store.get_messages(sid)[0]
  545. assert row["tool_calls"][0]["name"] == "calculate"
  546. assert row["citations"][0]["title"] == "T"
  547. def test_session_ownership_check(self):
  548. from db import chat_store
  549. sid = chat_store.create_session("alice", "full")["id"]
  550. assert chat_store.get_session_for_user(sid, "alice") is not None
  551. assert chat_store.get_session_for_user(sid, "bob") is None
  552. def test_delete_session_cascades(self):
  553. from db import chat_store
  554. sid = chat_store.create_session("u", "full")["id"]
  555. chat_store.save_message(sid, "user", "x")
  556. chat_store.save_context_usage(sid, 1, 10, 5)
  557. chat_store.delete_session(sid)
  558. assert chat_store.get_session(sid) is None
  559. assert chat_store.get_messages(sid) == []
  560. assert chat_store.get_context_usage(sid)["history"] == []
  561. def test_preference_upsert(self):
  562. from db import chat_store
  563. chat_store.save_preference("u", "单位", "米")
  564. chat_store.save_preference("u", "单位", "千米")
  565. prefs = {p["pref_key"]: p["pref_value"] for p in chat_store.list_preferences("u")}
  566. assert prefs == {"单位": "千米"}
  567. def test_preferences_as_text_empty(self):
  568. from db import chat_store
  569. assert chat_store.preferences_as_text("nobody") == ""
  570. def test_search_sessions_by_title(self):
  571. from db import chat_store
  572. chat_store.create_session("u", "full", "通风计算讨论")
  573. chat_store.create_session("u", "full", "其他话题")
  574. found = chat_store.search_sessions("u", "通风")
  575. assert len(found) == 1 and "通风" in found[0]["title"]
  576. def test_pinned_sessions_first(self):
  577. from db import chat_store
  578. a = chat_store.create_session("u", "full", "A")["id"]
  579. chat_store.create_session("u", "full", "B")
  580. chat_store.update_session_pin(a, True)
  581. assert chat_store.list_sessions("u")[0]["id"] == a
  582. def test_context_usage_aggregates(self):
  583. from db import chat_store
  584. sid = chat_store.create_session("u", "full")["id"]
  585. chat_store.save_context_usage(sid, 1, 100, 10, usage_source="provider")
  586. chat_store.save_context_usage(sid, 2, 200, 20, compacted=True, compact_count=1,
  587. usage_source="estimate")
  588. usage = chat_store.get_context_usage(sid)
  589. assert usage["compact_count"] == 1
  590. assert usage["latest"]["total_tokens"] == 220
  591. assert usage["latest"]["usage_source"] == "estimate"
  592. def test_schedule_crud_and_runs(self):
  593. from db import schedule_store
  594. schedule = schedule_store.create_schedule("u", "每日", "0 9 * * *", "汇报")
  595. sid = schedule["id"]
  596. assert schedule["enabled"] is True
  597. run_id = schedule_store.start_run(sid, "thread-1")
  598. schedule_store.finish_run(run_id, "success", output="ok")
  599. runs = schedule_store.list_runs(sid)
  600. assert len(runs) == 1 and runs[0]["status"] == "success"
  601. updated = schedule_store.update_schedule(sid, enabled=False, title="改名")
  602. assert updated["enabled"] is False and updated["title"] == "改名"
  603. assert schedule_store.delete_schedule(sid) is True
  604. assert schedule_store.get_schedule(sid) is None
  605. assert schedule_store.list_runs(sid) == []
  606. def test_schedule_update_ignores_unknown_fields(self):
  607. from db import schedule_store
  608. sid = schedule_store.create_schedule("u", "T", "0 9 * * *", "P")["id"]
  609. schedule_store.update_schedule(sid, title="新", 恶意字段="x")
  610. assert schedule_store.get_schedule(sid)["title"] == "新"
  611. def test_init_schema_idempotent(self):
  612. from db import db_conn
  613. db_conn.init_schema()
  614. db_conn.init_schema(force=True)
  615. assert db_conn.health()["status"] == "ok"
  616. # ============================================================
  617. # 路由面(把「删掉 /chat/only」这个决定变成可执行的守卫)
  618. # ============================================================
  619. class TestRouting:
  620. """断言**实际对外服务**的路由面。
  621. 这类断言的价值不在于「测试代码」,而在于把架构决定固化下来:
  622. 没有它,将来有人「为了兼容 1.5」把 /chat/only 加回来时,只有 code review
  623. 能拦住;有了它,CI 会拦。
  624. 刻意断言 ``main.app.openapi()["paths"]`` 而不是 ``api.router.routes``:
  625. 后者在本版 FastAPI 里是 ``_IncludedRouter`` 包装(不含 ``.path``),
  626. 而且真正要守住的是**服务出去的**那套路径。顺带还能验证路由定义合法——
  627. schema 生成失败会在这里报出来。
  628. """
  629. def _paths(self) -> set[str]:
  630. from main import app
  631. return set(app.openapi()["paths"])
  632. def test_chat_entry_exists(self):
  633. assert "/api/chat" in self._paths()
  634. def test_chat_only_is_gone(self):
  635. """2.0 刻意删除 /chat/only:它与 /chat 行为完全相同,唯一的 kind 标签
  636. 无人消费,且 1.5 为此付过计费 bug 与重复验证的代价。见 AGENTS.md §六。"""
  637. assert "/api/chat/only" not in self._paths()
  638. def test_interrupt_endpoints_kept(self):
  639. """/chat/resume 与 /chat/interrupt 承载的是不同能力,不是同一能力的
  640. 别名,所以保留——别把它们和 /chat/only 一起删。"""
  641. paths = self._paths()
  642. assert "/api/chat/resume" in paths
  643. assert "/api/chat/interrupt" in paths
  644. def test_detail_routes_not_lost_when_removing(self):
  645. """删入口时最容易连带弄丢同族路由,这里把对话族钉死。"""
  646. chat_paths = {p for p in self._paths() if p.startswith("/api/chat")}
  647. assert chat_paths == {
  648. "/api/chat",
  649. "/api/chat/resume",
  650. "/api/chat/interrupt",
  651. "/api/chat/history/{session_id}",
  652. }, chat_paths
  653. def test_chat_only_route_not_defined_in_source(self):
  654. """源码里不得再出现该路由的装饰器。
  655. 只在装饰器形态上断言,不搜 ``chat_only`` 这个词——它在注释与文档里
  656. 正当地出现(用来解释为什么没有它),搜词会把说明文字也判成违规。
  657. """
  658. import inspect
  659. from api import chat_routes
  660. src = inspect.getsource(chat_routes)
  661. assert '@router.post("/chat/only")' not in src
  662. assert '"chat_only"' not in src
  663. def test_kind_values_are_meaningful(self):
  664. """kind 只剩三个有意义取值:chat / resume / scheduled。"""
  665. import inspect
  666. from api import chat_routes
  667. from scheduler import runner
  668. chat_src = inspect.getsource(chat_routes)
  669. assert 'kind="chat"' in chat_src
  670. assert 'kind="resume"' in chat_src
  671. assert 'kind="scheduled"' in inspect.getsource(runner)
  672. # ============================================================
  673. # 模型工厂
  674. # ============================================================
  675. class TestModelFactory:
  676. def test_context_window_known(self):
  677. from agents import model_factory
  678. assert model_factory.context_window("deepseek-flash") == 128_000
  679. def test_context_limit_override_wins(self, monkeypatch):
  680. from agents import model_factory
  681. monkeypatch.setattr(model_factory.config, "CONTEXT_LIMIT", 5000)
  682. assert model_factory.context_window("deepseek-flash") == 5000
  683. def test_unknown_model_uses_fallback(self):
  684. from agents import model_factory
  685. assert model_factory.context_window("some-unknown-model-xyz") == \
  686. model_factory.FALLBACK_CONTEXT_WINDOW
  687. def test_thinking_levels_normalized(self):
  688. from agents import model_factory
  689. assert model_factory.normalize_thinking("HIGH") == "high"
  690. assert model_factory.normalize_thinking("bogus") == "off"
  691. assert model_factory.normalize_thinking(None) in model_factory.VALID_THINKING_LEVELS
  692. def test_off_sends_no_extra_body(self):
  693. """off 时什么都不发:不同 provider 关闭思考的字段名不一致,发错会被拒。"""
  694. from agents import model_factory
  695. assert model_factory._thinking_payload("off") is None
  696. assert model_factory._thinking_payload("high") is not None
  697. def test_build_model_declares_profile(self):
  698. """必须显式声明上下文窗口,否则摘要中间件首次触发时会报错。"""
  699. from agents import model_factory
  700. model = model_factory.build_model()
  701. assert model.profile and model.profile.get("max_input_tokens")
  702. # ============================================================
  703. # 中断恢复翻译
  704. # ============================================================
  705. class TestResumeTranslation:
  706. def _payload(self, **kw):
  707. from api.models import ResumeRequest
  708. base = {"session_id": "s", "status": "answered"}
  709. base.update(kw)
  710. return ResumeRequest(**base)
  711. def test_approve(self):
  712. from api.chat_routes import _to_official_resume
  713. assert _to_official_resume(self._payload(approved=True), 1) == \
  714. {"decisions": [{"type": "approve"}]}
  715. def test_reject(self):
  716. from api.chat_routes import _to_official_resume
  717. decisions = _to_official_resume(self._payload(approved=False), 1)["decisions"]
  718. assert decisions[0]["type"] == "reject" and decisions[0]["message"]
  719. def test_cancelled_becomes_reject(self):
  720. from api.chat_routes import _to_official_resume
  721. decisions = _to_official_resume(self._payload(status="cancelled"), 1)["decisions"]
  722. assert decisions[0]["type"] == "reject"
  723. def test_answers_become_respond(self):
  724. from api.chat_routes import _to_official_resume
  725. decisions = _to_official_resume(self._payload(answers=["甲", "乙"]), 1)["decisions"]
  726. assert decisions[0]["type"] == "respond" and "甲" in decisions[0]["message"]
  727. def test_multiple_pending_actions_get_matching_decisions(self):
  728. """官方要求 decisions 与 action_requests 数量严格一致。"""
  729. from api.chat_routes import _to_official_resume
  730. decisions = _to_official_resume(self._payload(approved=True, interrupt_token="reviewed-batch"), 3)["decisions"]
  731. assert len(decisions) == 3
  732. assert all(d["type"] == "approve" for d in decisions)
  733. def test_unknown_payload_defaults_to_safe_reject(self):
  734. """没有明确答复时宁可拒绝执行。"""
  735. from api.chat_routes import _to_official_resume
  736. decisions = _to_official_resume(self._payload(), 1)["decisions"]
  737. assert decisions[0]["type"] == "reject"