test_core_units.py 38 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004
  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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api.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 api import concurrency
  397. concurrency.reset_for_tests()
  398. def teardown_method(self):
  399. from api import concurrency
  400. concurrency.reset_for_tests()
  401. def test_acquire_and_release(self):
  402. from api 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 api 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 api 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 api 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 api 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.setattr(pricing.config, "get_bool", lambda *a, **k: False)
  450. with pytest.raises(pricing.UnpricedModelError):
  451. pricing.compute_cost("unknown-model", 100, 100)
  452. def test_unpriced_model_allowed_when_configured(self, pricing_ready, monkeypatch):
  453. from billing import pricing
  454. monkeypatch.setattr(pricing.config, "get_bool", lambda *a, **k: True)
  455. cost, _ = pricing.compute_cost("unknown-model", 100, 100)
  456. assert cost == 0.0
  457. def test_quota_blocks_when_exhausted(self, pricing_ready, unique_user):
  458. from billing import guard
  459. from db import billing_store, db_conn
  460. day = db_conn.day_key()
  461. billing_store.set_quota_limit(unique_user, day, 100)
  462. billing_store.record_usage(unique_user, "s", "chat", "m", 100, 0, 0.0, day)
  463. with pytest.raises(guard.QuotaExceededError):
  464. guard.precheck(unique_user, day)
  465. def test_quota_zero_means_unlimited(self, pricing_ready, unique_user):
  466. from billing import guard
  467. from db import billing_store, db_conn
  468. day = db_conn.day_key()
  469. billing_store.set_quota_limit(unique_user, day, 0)
  470. billing_store.record_usage(unique_user, "s", "chat", "m", 10_000_000, 0, 0.0, day)
  471. guard.precheck(unique_user, day) # 不抛即通过
  472. def test_meter_settle_idempotent(self, pricing_ready, unique_user):
  473. from billing import guard
  474. from billing.meter import UsageMeter
  475. meter = UsageMeter(unique_user, "s1", "chat", "deepseek-flash")
  476. meter.add(1000, 100)
  477. first = meter.settle()
  478. second = meter.settle()
  479. assert first["settled"] is True
  480. assert second.get("duplicated") is True
  481. assert guard.snapshot(unique_user)["requests"] == 1
  482. def test_meter_no_usage_skips(self, pricing_ready, unique_user):
  483. from billing import guard
  484. from billing.meter import UsageMeter
  485. meter = UsageMeter(unique_user, "s1", "chat", "deepseek-flash")
  486. result = meter.settle()
  487. assert result["settled"] is False and result["reason"] == "no_usage"
  488. assert guard.snapshot(unique_user)["requests"] == 0
  489. def test_meter_source_downgrades_to_estimate(self, pricing_ready, unique_user):
  490. from billing.meter import UsageMeter
  491. meter = UsageMeter(unique_user, "s1", "chat", "deepseek-flash")
  492. meter.add(100, 10, "provider")
  493. meter.add(50, 5, "estimate")
  494. assert meter.usage_source == "estimate"
  495. def test_snapshot_has_total_tokens(self, pricing_ready, unique_user):
  496. """字段名与 billing_store.usage_for_day 对齐,消费方不用记两套名字。"""
  497. from billing import guard
  498. snapshot = guard.snapshot(unique_user)
  499. assert "total_tokens" in snapshot and "used_tokens" in snapshot
  500. def test_usage_day_bound_to_request_start(self, pricing_ready, unique_user):
  501. """跨零点的长请求不能被算到第二天。"""
  502. from billing import guard
  503. from billing.meter import UsageMeter
  504. from db import db_conn
  505. yesterday = "20200101"
  506. meter = UsageMeter(unique_user, "s1", "chat", "deepseek-flash", usage_day=yesterday)
  507. meter.add(100, 10)
  508. meter.settle()
  509. assert guard.snapshot(unique_user, yesterday)["requests"] == 1
  510. assert guard.snapshot(unique_user, db_conn.day_key())["requests"] == 0
  511. # ============================================================
  512. # 存储
  513. # ============================================================
  514. class TestStores:
  515. def test_message_ordering_stable_within_same_second(self):
  516. """同秒写入的多条消息顺序必须稳定——排序错乱会让补种的上下文角色错位。"""
  517. from db import chat_store
  518. sid = chat_store.create_session("u", "full")["id"]
  519. for i in range(6):
  520. chat_store.save_message(sid, "user", f"Q{i}")
  521. chat_store.save_message(sid, "assistant", f"A{i}")
  522. rows = chat_store.get_messages(sid)
  523. assert [r["msg_role"] for r in rows] == ["user", "assistant"] * 6
  524. assert [r["content"] for r in rows] == [
  525. c for i in range(6) for c in (f"Q{i}", f"A{i}")
  526. ]
  527. def test_get_messages_limit_returns_most_recent_in_order(self):
  528. from db import chat_store
  529. sid = chat_store.create_session("u", "full")["id"]
  530. for i in range(5):
  531. chat_store.save_message(sid, "user", f"M{i}")
  532. recent = chat_store.get_messages(sid, limit=3)
  533. assert [r["content"] for r in recent] == ["M2", "M3", "M4"]
  534. def test_message_json_fields_roundtrip(self):
  535. from db import chat_store
  536. sid = chat_store.create_session("u", "full")["id"]
  537. chat_store.save_message(
  538. sid, "assistant", "答",
  539. tool_calls=[{"name": "calculate", "args": {"expression": "1+1"}}],
  540. citations=[{"title": "T", "url": "u"}],
  541. )
  542. row = chat_store.get_messages(sid)[0]
  543. assert row["tool_calls"][0]["name"] == "calculate"
  544. assert row["citations"][0]["title"] == "T"
  545. def test_session_ownership_check(self):
  546. from db import chat_store
  547. sid = chat_store.create_session("alice", "full")["id"]
  548. assert chat_store.get_session_for_user(sid, "alice") is not None
  549. assert chat_store.get_session_for_user(sid, "bob") is None
  550. def test_delete_session_cascades(self):
  551. from db import chat_store
  552. sid = chat_store.create_session("u", "full")["id"]
  553. chat_store.save_message(sid, "user", "x")
  554. chat_store.save_context_usage(sid, 1, 10, 5)
  555. chat_store.delete_session(sid)
  556. assert chat_store.get_session(sid) is None
  557. assert chat_store.get_messages(sid) == []
  558. assert chat_store.get_context_usage(sid)["history"] == []
  559. def test_preference_upsert(self):
  560. from db import chat_store
  561. chat_store.save_preference("u", "单位", "米")
  562. chat_store.save_preference("u", "单位", "千米")
  563. prefs = {p["pref_key"]: p["pref_value"] for p in chat_store.list_preferences("u")}
  564. assert prefs == {"单位": "千米"}
  565. def test_preferences_as_text_empty(self):
  566. from db import chat_store
  567. assert chat_store.preferences_as_text("nobody") == ""
  568. def test_search_sessions_by_title(self):
  569. from db import chat_store
  570. chat_store.create_session("u", "full", "通风计算讨论")
  571. chat_store.create_session("u", "full", "其他话题")
  572. found = chat_store.search_sessions("u", "通风")
  573. assert len(found) == 1 and "通风" in found[0]["title"]
  574. def test_pinned_sessions_first(self):
  575. from db import chat_store
  576. a = chat_store.create_session("u", "full", "A")["id"]
  577. chat_store.create_session("u", "full", "B")
  578. chat_store.update_session_pin(a, True)
  579. assert chat_store.list_sessions("u")[0]["id"] == a
  580. def test_context_usage_aggregates(self):
  581. from db import chat_store
  582. sid = chat_store.create_session("u", "full")["id"]
  583. chat_store.save_context_usage(sid, 1, 100, 10, usage_source="provider")
  584. chat_store.save_context_usage(sid, 2, 200, 20, compacted=True, compact_count=1,
  585. usage_source="estimate")
  586. usage = chat_store.get_context_usage(sid)
  587. assert usage["compact_count"] == 1
  588. assert usage["latest"]["total_tokens"] == 220
  589. assert usage["latest"]["usage_source"] == "estimate"
  590. def test_schedule_crud_and_runs(self):
  591. from db import schedule_store
  592. schedule = schedule_store.create_schedule("u", "每日", "0 9 * * *", "汇报")
  593. sid = schedule["id"]
  594. assert schedule["enabled"] is True
  595. run_id = schedule_store.start_run(sid, "thread-1")
  596. schedule_store.finish_run(run_id, "success", output="ok")
  597. runs = schedule_store.list_runs(sid)
  598. assert len(runs) == 1 and runs[0]["status"] == "success"
  599. updated = schedule_store.update_schedule(sid, enabled=False, title="改名")
  600. assert updated["enabled"] is False and updated["title"] == "改名"
  601. assert schedule_store.delete_schedule(sid) is True
  602. assert schedule_store.get_schedule(sid) is None
  603. assert schedule_store.list_runs(sid) == []
  604. def test_schedule_update_ignores_unknown_fields(self):
  605. from db import schedule_store
  606. sid = schedule_store.create_schedule("u", "T", "0 9 * * *", "P")["id"]
  607. schedule_store.update_schedule(sid, title="新", 恶意字段="x")
  608. assert schedule_store.get_schedule(sid)["title"] == "新"
  609. def test_init_schema_idempotent(self):
  610. from db import db_conn
  611. db_conn.init_schema()
  612. db_conn.init_schema(force=True)
  613. assert db_conn.health()["status"] == "ok"
  614. # ============================================================
  615. # 路由面(把「删掉 /chat/only」这个决定变成可执行的守卫)
  616. # ============================================================
  617. class TestRouting:
  618. """断言**实际对外服务**的路由面。
  619. 这类断言的价值不在于「测试代码」,而在于把架构决定固化下来:
  620. 没有它,将来有人「为了兼容 1.5」把 /chat/only 加回来时,只有 code review
  621. 能拦住;有了它,CI 会拦。
  622. 刻意断言 ``main.app.openapi()["paths"]`` 而不是 ``api.router.routes``:
  623. 后者在本版 FastAPI 里是 ``_IncludedRouter`` 包装(不含 ``.path``),
  624. 而且真正要守住的是**服务出去的**那套路径。顺带还能验证路由定义合法——
  625. schema 生成失败会在这里报出来。
  626. """
  627. def _paths(self) -> set[str]:
  628. from main import app
  629. return set(app.openapi()["paths"])
  630. def test_chat_entry_exists(self):
  631. assert "/api/chat" in self._paths()
  632. def test_chat_only_is_gone(self):
  633. """2.0 刻意删除 /chat/only:它与 /chat 行为完全相同,唯一的 kind 标签
  634. 无人消费,且 1.5 为此付过计费 bug 与重复验证的代价。见 AGENTS.md §六。"""
  635. assert "/api/chat/only" not in self._paths()
  636. def test_interrupt_endpoints_kept(self):
  637. """/chat/resume 与 /chat/interrupt 承载的是不同能力,不是同一能力的
  638. 别名,所以保留——别把它们和 /chat/only 一起删。"""
  639. paths = self._paths()
  640. assert "/api/chat/resume" in paths
  641. assert "/api/chat/interrupt" in paths
  642. def test_detail_routes_not_lost_when_removing(self):
  643. """删入口时最容易连带弄丢同族路由,这里把对话族钉死。"""
  644. chat_paths = {p for p in self._paths() if p.startswith("/api/chat")}
  645. assert chat_paths == {
  646. "/api/chat",
  647. "/api/chat/resume",
  648. "/api/chat/interrupt",
  649. "/api/chat/history/{session_id}",
  650. }, chat_paths
  651. def test_chat_only_route_not_defined_in_source(self):
  652. """源码里不得再出现该路由的装饰器。
  653. 只在装饰器形态上断言,不搜 ``chat_only`` 这个词——它在注释与文档里
  654. 正当地出现(用来解释为什么没有它),搜词会把说明文字也判成违规。
  655. """
  656. import inspect
  657. from api import chat_routes
  658. src = inspect.getsource(chat_routes)
  659. assert '@router.post("/chat/only")' not in src
  660. assert '"chat_only"' not in src
  661. def test_kind_values_are_meaningful(self):
  662. """kind 只剩三个有意义取值:chat / resume / scheduled。"""
  663. import inspect
  664. from api import chat_routes
  665. from scheduler import runner
  666. chat_src = inspect.getsource(chat_routes)
  667. assert 'kind="chat"' in chat_src
  668. assert 'kind="resume"' in chat_src
  669. assert 'kind="scheduled"' in inspect.getsource(runner)
  670. # ============================================================
  671. # 模型工厂
  672. # ============================================================
  673. class TestModelFactory:
  674. def test_context_window_known(self):
  675. from agents import model_factory
  676. assert model_factory.context_window("deepseek-flash") == 128_000
  677. def test_context_limit_override_wins(self, monkeypatch):
  678. from agents import model_factory
  679. monkeypatch.setattr(model_factory.config, "CONTEXT_LIMIT", 5000)
  680. assert model_factory.context_window("deepseek-flash") == 5000
  681. def test_unknown_model_uses_fallback(self):
  682. from agents import model_factory
  683. assert model_factory.context_window("some-unknown-model-xyz") == \
  684. model_factory.FALLBACK_CONTEXT_WINDOW
  685. def test_thinking_levels_normalized(self):
  686. from agents import model_factory
  687. assert model_factory.normalize_thinking("HIGH") == "high"
  688. assert model_factory.normalize_thinking("bogus") == "off"
  689. assert model_factory.normalize_thinking(None) in model_factory.VALID_THINKING_LEVELS
  690. def test_off_sends_no_extra_body(self):
  691. """off 时什么都不发:不同 provider 关闭思考的字段名不一致,发错会被拒。"""
  692. from agents import model_factory
  693. assert model_factory._thinking_payload("off") is None
  694. assert model_factory._thinking_payload("high") is not None
  695. def test_build_model_declares_profile(self):
  696. """必须显式声明上下文窗口,否则摘要中间件首次触发时会报错。"""
  697. from agents import model_factory
  698. model = model_factory.build_model()
  699. assert model.profile and model.profile.get("max_input_tokens")
  700. # ============================================================
  701. # 中断恢复翻译
  702. # ============================================================
  703. class TestResumeTranslation:
  704. def _payload(self, **kw):
  705. from api.models import ResumeRequest
  706. base = {"session_id": "s", "status": "answered"}
  707. base.update(kw)
  708. return ResumeRequest(**base)
  709. def test_approve(self):
  710. from api.chat_routes import _to_official_resume
  711. assert _to_official_resume(self._payload(approved=True), 1) == \
  712. {"decisions": [{"type": "approve"}]}
  713. def test_reject(self):
  714. from api.chat_routes import _to_official_resume
  715. decisions = _to_official_resume(self._payload(approved=False), 1)["decisions"]
  716. assert decisions[0]["type"] == "reject" and decisions[0]["message"]
  717. def test_cancelled_becomes_reject(self):
  718. from api.chat_routes import _to_official_resume
  719. decisions = _to_official_resume(self._payload(status="cancelled"), 1)["decisions"]
  720. assert decisions[0]["type"] == "reject"
  721. def test_answers_become_respond(self):
  722. from api.chat_routes import _to_official_resume
  723. decisions = _to_official_resume(self._payload(answers=["甲", "乙"]), 1)["decisions"]
  724. assert decisions[0]["type"] == "respond" and "甲" in decisions[0]["message"]
  725. def test_multiple_pending_actions_get_matching_decisions(self):
  726. """官方要求 decisions 与 action_requests 数量严格一致。"""
  727. from api.chat_routes import _to_official_resume
  728. decisions = _to_official_resume(self._payload(approved=True, interrupt_token="reviewed-batch"), 3)["decisions"]
  729. assert len(decisions) == 3
  730. assert all(d["type"] == "approve" for d in decisions)
  731. def test_unknown_payload_defaults_to_safe_reject(self):
  732. """没有明确答复时宁可拒绝执行。"""
  733. from api.chat_routes import _to_official_resume
  734. decisions = _to_official_resume(self._payload(), 1)["decisions"]
  735. assert decisions[0]["type"] == "reject"