test_runtime_regressions.py 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546
  1. """执行、存储和人机交互的失败路径回归;不访问模型网络。"""
  2. import asyncio
  3. import json
  4. from types import SimpleNamespace
  5. import httpx
  6. import pytest
  7. from langchain_core.messages import AIMessage, AIMessageChunk
  8. import config
  9. from runtime import concurrency
  10. from api.interrupts import pack
  11. from db import billing_store, chat_store, db_conn, run_store, schedule_store
  12. @pytest.fixture(autouse=True)
  13. def isolation(monkeypatch, tmp_path):
  14. monkeypatch.setattr(config, "ROOT", tmp_path)
  15. (tmp_path / "skills").mkdir()
  16. concurrency.reset_for_tests()
  17. yield
  18. concurrency.reset_for_tests()
  19. class FakeAgent:
  20. def __init__(self, interrupt=None, wait=None):
  21. self.interrupt = interrupt
  22. self.wait = wait
  23. self.started = asyncio.Event()
  24. self.inputs = []
  25. async def aget_state(self, config):
  26. return SimpleNamespace(tasks=[])
  27. async def astream(self, graph_input, **kwargs):
  28. self.inputs.append((graph_input, kwargs))
  29. self.started.set()
  30. yield {"type": "messages", "data": (AIMessageChunk(content="部分回答", usage_metadata={"input_tokens": 12, "output_tokens": 3, "total_tokens": 15}), {})}
  31. if self.wait:
  32. await self.wait.wait()
  33. if self.interrupt:
  34. yield {"type": "updates", "data": {"__interrupt__": [self.interrupt]}}
  35. def test_sqlite_rollback_and_nested_savepoint():
  36. adapter = db_conn._SqliteAdapter(":memory:")
  37. adapter.execute("CREATE TABLE probe (id INTEGER)")
  38. with pytest.raises(RuntimeError):
  39. with adapter.transaction():
  40. adapter.execute("INSERT INTO probe VALUES (1)")
  41. adapter.execute_many("INSERT INTO probe VALUES (?)", [(2,), (3,)])
  42. raise RuntimeError()
  43. assert adapter.query("SELECT * FROM probe") == []
  44. with adapter.transaction():
  45. adapter.execute("INSERT INTO probe VALUES (4)")
  46. with pytest.raises(ValueError):
  47. with adapter.transaction():
  48. adapter.execute("INSERT INTO probe VALUES (5)")
  49. raise ValueError()
  50. adapter.execute("INSERT INTO probe VALUES (6)")
  51. assert adapter.query("SELECT id FROM probe") == [{"id": 4}, {"id": 6}]
  52. def test_billing_transaction_rolls_back_quota_failure(monkeypatch):
  53. def fail(*args):
  54. raise RuntimeError("quota database failure")
  55. monkeypatch.setattr(billing_store, "_accumulate_quota", fail)
  56. with pytest.raises(RuntimeError):
  57. billing_store.record_usage("u", "s", "chat", config.MODEL, 10, 2, 0, db_conn.day_key())
  58. assert billing_store.recent_usage("u") == []
  59. def test_failed_settlement_survives_new_meter_and_deduplicates(monkeypatch, pricing_ready):
  60. from billing.meter import UsageMeter, retry_pending
  61. real = billing_store.record_usage
  62. meter = UsageMeter("u", "s", "chat", config.MODEL)
  63. meter.add(10, 2)
  64. def fail(**kwargs):
  65. raise RuntimeError("temporary outage")
  66. monkeypatch.setattr(billing_store, "record_usage", fail)
  67. assert not meter.settle()["settled"]
  68. assert not meter.settled
  69. assert db_conn.scalar("SELECT COUNT(*) FROM billing_pending") == 1
  70. monkeypatch.setattr(billing_store, "record_usage", real)
  71. assert retry_pending() == 1
  72. assert meter.settle()["settled"] # old request's retry cannot double-charge
  73. assert len(billing_store.recent_usage("u")) == 1
  74. assert billing_store.get_quota("u", db_conn.day_key())["used_tokens"] == 12
  75. assert db_conn.scalar("SELECT COUNT(*) FROM billing_pending") == 0
  76. @pytest.mark.asyncio
  77. async def test_http_concurrency_returns_429(monkeypatch, pricing_ready):
  78. from main import app
  79. monkeypatch.setattr(config, "MAX_CONCURRENT_TASKS", 0)
  80. async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
  81. response = await c.post("/api/chat", data={"message": "hello"})
  82. assert response.status_code == 429
  83. assert response.headers["retry-after"]
  84. assert not concurrency._sessions
  85. @pytest.mark.asyncio
  86. async def test_cross_user_message_and_schedule_access(pricing_ready):
  87. from main import app
  88. session = chat_store.create_session("alice")
  89. message_id = chat_store.save_message(session["id"], "user", "secret")
  90. schedule = schedule_store.create_schedule("alice", "private", "0 9 * * *", "private")
  91. async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
  92. response = await c.delete(f"/api/messages/{message_id}", params={"user_name": "bob"})
  93. assert response.status_code == 404
  94. for method, suffix in [("GET", ""), ("DELETE", ""), ("POST", "/run"), ("GET", "/runs")]:
  95. response = await c.request(method, f"/api/schedules/{schedule['id']}{suffix}", params={"user_name": "bob"})
  96. assert response.status_code == 404
  97. assert chat_store.count_messages(session["id"]) == 1
  98. assert schedule_store.get_schedule(schedule["id"])
  99. @pytest.mark.asyncio
  100. async def test_same_session_concurrent_request_rejected(monkeypatch, pricing_ready):
  101. from main import app
  102. agent = FakeAgent(wait=asyncio.Event())
  103. monkeypatch.setattr("agents.agent_registry.get_agent", lambda: agent)
  104. session = chat_store.create_session("default")
  105. async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
  106. first = asyncio.create_task(c.post("/api/chat", data={"message": "one", "session_id": session["id"]}))
  107. await asyncio.wait_for(agent.started.wait(), 2)
  108. second = await c.post("/api/chat", data={"message": "two", "session_id": session["id"]})
  109. assert second.status_code == 409
  110. agent.wait.set()
  111. assert (await first).status_code == 200
  112. assert concurrency.stats()["global_active"] == 0
  113. assert len(run_store.list_runs(session["id"])) == 1
  114. assert "two" not in [m["content"] for m in chat_store.get_messages(session["id"])]
  115. @pytest.mark.asyncio
  116. @pytest.mark.parametrize("cancel", [False, True])
  117. async def test_stream_close_or_cancel_settles_and_releases(cancel, pricing_ready):
  118. from runtime.execution import admit, build_execution
  119. session = chat_store.create_session("default")
  120. agent = FakeAgent(wait=asyncio.Event())
  121. admission = await admit("default", session["id"])
  122. execution = await build_execution(admission, agent=agent, session_id=session["id"], thread_id=session["id"],
  123. user_name="default", message="test", config={"configurable": {"thread_id": session["id"]}})
  124. generator = execution.events()
  125. while True:
  126. frame = await anext(generator)
  127. if frame["type"] == "token":
  128. break
  129. if cancel:
  130. pending = asyncio.create_task(anext(generator))
  131. await asyncio.sleep(0)
  132. pending.cancel()
  133. with pytest.raises(asyncio.CancelledError):
  134. await pending
  135. await generator.aclose()
  136. assert not concurrency._sessions
  137. assert concurrency.stats()["global_active"] == 0
  138. assert billing_store.get_quota("default", db_conn.day_key())["used_tokens"] == 15
  139. assert run_store.list_runs(session["id"])[0]["status"] == "cancelled"
  140. assert "不完整" in chat_store.get_messages(session["id"])[-1]["content"]
  141. @pytest.mark.asyncio
  142. async def test_schedule_uses_shared_runtime_and_pauses(monkeypatch, pricing_ready):
  143. from scheduler.runner import run_schedule
  144. item = SimpleNamespace(id="approval-id", value={"action_requests": [{"name": "calculate", "args": {"expression": "1+1"}}]})
  145. agent = FakeAgent(interrupt=item)
  146. monkeypatch.setattr("agents.agent_registry.get_agent", lambda: agent)
  147. schedule = schedule_store.create_schedule("u", "plan", "0 9 * * *", "question", "plan")
  148. result = await run_schedule(schedule["id"])
  149. assert result["status"] == "waiting_input"
  150. assert schedule_store.latest_run(schedule["id"])["status"] == "waiting_input"
  151. assert run_store.list_runs(result["session_id"])[0]["status"] == "waiting_input"
  152. assert "request_plan_approval" in agent.inputs[0][0]["messages"][-1].content
  153. assert billing_store.get_quota("u", db_conn.day_key())["used_tokens"] == 15
  154. assert concurrency.stats()["global_active"] == 0
  155. @pytest.mark.asyncio
  156. async def test_resume_and_schedule_enforce_quota(pricing_ready):
  157. from main import app
  158. from scheduler.runner import run_schedule
  159. from billing.guard import QuotaExceededError
  160. session = chat_store.create_session("u")
  161. billing_store.set_quota_limit("u", db_conn.day_key(), 1)
  162. billing_store.record_usage("u", session["id"], "chat", config.MODEL, 1, 0, 0, db_conn.day_key())
  163. schedule = schedule_store.create_schedule("u", "x", "0 9 * * *", "x")
  164. with pytest.raises(QuotaExceededError):
  165. await run_schedule(schedule["id"])
  166. async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
  167. response = await c.post("/api/chat/resume", json={"session_id": session["id"], "user_name": "u", "approved": True})
  168. assert response.status_code == 429
  169. assert not concurrency._sessions
  170. def test_workspace_isolation_and_skills(monkeypatch, tmp_path):
  171. from agents.workspace import backend_for_runtime, session_root
  172. (tmp_path / "skills" / "guide.md").write_text("shared guide")
  173. def backend(user, session):
  174. return backend_for_runtime(SimpleNamespace(config={"configurable": {"workspace_user": user, "workspace_session": session}}))
  175. first, second, other_user = backend("alice", "one"), backend("alice", "two"), backend("bob", "one")
  176. first.write("/workspace/private.txt", "secret")
  177. assert "secret" in str(first.read("/workspace/private.txt"))
  178. assert second.read("/workspace/private.txt").error
  179. assert other_user.read("/workspace/private.txt").error
  180. assert "shared guide" in str(first.read("/skills/guide.md"))
  181. assert session_root("../escape", "../../other").is_relative_to(tmp_path / "data" / "workspaces")
  182. with pytest.raises(ValueError):
  183. backend_for_runtime(SimpleNamespace(config={}))
  184. def test_batch_requires_current_interrupt_and_shows_all():
  185. from api.interrupts import interrupt_frame, resume_value
  186. from api.models import ResumeRequest
  187. from api.sse_events import EventWriter
  188. from fastapi import HTTPException
  189. item = SimpleNamespace(id="a", value={"action_requests": [{"name": "one", "args": {}}, {"name": "two", "args": {}}]})
  190. frame = json.loads(interrupt_frame(EventWriter("s"), pack(item))[6:])
  191. assert len(frame["tool_approval"]["actions"]) == 2
  192. state = SimpleNamespace(tasks=[SimpleNamespace(interrupts=[item])])
  193. payload = ResumeRequest(session_id="s", approved=True, interrupt_token=frame["tool_approval"]["interrupt_token"])
  194. assert len(resume_value(payload, state)["a"]["decisions"]) == 2
  195. item.id = "newer"
  196. with pytest.raises(HTTPException) as error:
  197. resume_value(payload, state)
  198. assert error.value.status_code == 409
  199. def test_declared_side_effect_tools_require_approval():
  200. from tools.registry import register
  201. from agents.agent_factory import _build_interrupt_on
  202. @register(group="test", requires_approval=True)
  203. def external_write() -> str:
  204. """External change."""
  205. return "done"
  206. assert _build_interrupt_on()["external_write"] is True
  207. def test_unknown_model_is_not_silently_priced(monkeypatch):
  208. from billing import pricing, guard
  209. monkeypatch.setattr(config, "MODEL", "unpriced-test-model")
  210. pricing.seed_default_pricing()
  211. assert billing_store.get_pricing(config.MODEL) is None
  212. with pytest.raises(guard.BillingConfigError):
  213. guard.precheck("u")
  214. def test_startup_marks_stale_runs_interrupted():
  215. run_store.start("stale", "s", "u", "chat")
  216. run_store.recover_interrupted()
  217. assert run_store.list_runs("s")[0]["status"] == "interrupted"
  218. from langchain_core.language_models.chat_models import BaseChatModel
  219. from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
  220. from pydantic import Field
  221. class ScriptedModel(BaseChatModel):
  222. """真实图引擎,确定性模型;覆盖 middleware、工具和 checkpoint。"""
  223. responses: list[AIMessage] = Field(default_factory=list)
  224. calls: int = 0
  225. @property
  226. def _llm_type(self):
  227. return "offline-scripted"
  228. def bind_tools(self, tools, **kwargs):
  229. return self
  230. def _next(self):
  231. response = self.responses[self.calls]
  232. self.calls += 1
  233. return response
  234. def _generate(self, messages, stop=None, run_manager=None, **kwargs):
  235. return ChatResult(generations=[ChatGeneration(message=self._next())])
  236. def _stream(self, messages, stop=None, run_manager=None, **kwargs):
  237. message = self._next()
  238. yield ChatGenerationChunk(message=AIMessageChunk(content=message.content, tool_calls=message.tool_calls,
  239. usage_metadata={"input_tokens": 20, "output_tokens": 5, "total_tokens": 25}))
  240. @pytest.mark.asyncio
  241. @pytest.mark.parametrize("tool_name,args", [
  242. ("ask_user", {"questions": ["请选择区域"]}),
  243. ("request_plan_approval", {"plan": "读取文件后生成报告"}),
  244. ])
  245. async def test_real_graph_interrupt_refresh_resume(monkeypatch, pricing_ready, tool_name, args):
  246. from agents import agent_factory, agent_memory
  247. from main import app
  248. model = ScriptedModel(responses=[
  249. AIMessage(content="", tool_calls=[{"id": "call1", "name": tool_name, "args": args}]),
  250. AIMessage(content="已根据你的确认完成。"),
  251. ])
  252. monkeypatch.setattr("agents.model_factory.get_model", lambda: model)
  253. graph = agent_factory.build_agent()
  254. monkeypatch.setattr("agents.agent_registry.get_agent", lambda: graph)
  255. async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
  256. first = await c.post("/api/chat", data={"message": "test", "user_name": "u"})
  257. assert first.status_code == 200, first.text
  258. events = [json.loads(line[6:]) for line in first.text.splitlines() if line.startswith("data: ")]
  259. assert not [e for e in events if e["type"] == "error"], events
  260. interrupted = next(e for e in events if e["type"] == "interrupt")
  261. sid = first.headers["x-session-id"]
  262. token = interrupted.get("interrupt_token") or interrupted["tool_approval"]["interrupt_token"]
  263. refreshed = (await c.get("/api/chat/interrupt", params={"session_id": sid, "user_name": "u"})).json()["interrupts"][0]
  264. assert (refreshed.get("interrupt_token") or refreshed["tool_approval"]["interrupt_token"]) == token
  265. assert events[-1]["waiting_input"]
  266. response = await c.post("/api/chat/resume", json={"session_id": sid, "user_name": "u",
  267. "interrupt_token": token, "approved": True, "answers": ["一区"]})
  268. assert response.status_code == 200, response.text
  269. assert "已根据你的确认完成" in response.text
  270. assert run_store.list_runs(sid)[0]["status"] == "success"
  271. repeated = await c.post("/api/chat/resume", json={"session_id": sid, "user_name": "u", "interrupt_token": token, "approved": True})
  272. assert repeated.status_code == 409
  273. await agent_memory.delete_thread(sid)
  274. @pytest.mark.asyncio
  275. async def test_real_graph_workspace_write(monkeypatch, pricing_ready):
  276. from agents import agent_factory, agent_memory
  277. from agents.workspace import session_root
  278. from main import app
  279. model = ScriptedModel(responses=[
  280. AIMessage(content="", tool_calls=[{"id": "w1", "name": "write_file", "args": {"file_path": "/workspace/result.txt", "content": "隔离的产物"}}]),
  281. AIMessage(content="已写入"),
  282. ])
  283. monkeypatch.setattr("agents.model_factory.get_model", lambda: model)
  284. graph = agent_factory.build_agent()
  285. monkeypatch.setattr("agents.agent_registry.get_agent", lambda: graph)
  286. async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
  287. result = await c.post("/api/chat", data={"message": "write", "user_name": "u"})
  288. sid = result.headers["x-session-id"]
  289. assert (session_root("u", sid) / "workspace/result.txt").read_text() == "隔离的产物", result.text
  290. await agent_memory.delete_thread(sid)
  291. @pytest.mark.asyncio
  292. async def test_request_id_does_not_execute_twice(monkeypatch, pricing_ready):
  293. from main import app
  294. agent = FakeAgent()
  295. monkeypatch.setattr("agents.agent_registry.get_agent", lambda: agent)
  296. async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
  297. body = {"message": "once", "request_id": "client-generated-key"}
  298. first = await c.post("/api/chat", data=body)
  299. second = await c.post("/api/chat", data=body)
  300. assert first.status_code == 200
  301. assert second.status_code == 409
  302. assert second.json()["detail"]["run_id"] == first.headers["x-run-id"]
  303. assert len(agent.inputs) == 1
  304. assert len(chat_store.list_sessions("default")) == 1
  305. @pytest.mark.asyncio
  306. async def test_real_subagent_shares_only_parent_workspace(monkeypatch, pricing_ready):
  307. from agents import agent_factory, agent_memory
  308. from agents.workspace import session_root
  309. from main import app
  310. model = ScriptedModel(responses=[
  311. AIMessage(content="", tool_calls=[{"id": "t1", "name": "task", "args": {"subagent_type": "demo-searcher", "description": "write a report"}}]),
  312. AIMessage(content="", tool_calls=[{"id": "w2", "name": "write_file", "args": {"file_path": "/workspace/sub.txt", "content": "子任务产物"}}]),
  313. AIMessage(content="子任务完成"),
  314. AIMessage(content="汇总完成"),
  315. ])
  316. monkeypatch.setattr("agents.model_factory.get_model", lambda: model)
  317. graph = agent_factory.build_agent()
  318. monkeypatch.setattr("agents.agent_registry.get_agent", lambda: graph)
  319. async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
  320. result = await c.post("/api/chat", data={"message": "delegate", "user_name": "u"})
  321. sid = result.headers["x-session-id"]
  322. assert (session_root("u", sid) / "workspace/sub.txt").read_text() == "子任务产物", result.text
  323. await agent_memory.delete_thread(sid)
  324. @pytest.mark.asyncio
  325. async def test_schedule_timeout_is_failed_and_billed(monkeypatch, pricing_ready):
  326. from scheduler.runner import run_schedule
  327. agent = FakeAgent(wait=asyncio.Event())
  328. monkeypatch.setattr("agents.agent_registry.get_agent", lambda: agent)
  329. monkeypatch.setattr(config, "TASK_WALL_TIMEOUT", 0.05)
  330. schedule = schedule_store.create_schedule("u", "slow", "0 9 * * *", "wait")
  331. result = await run_schedule(schedule["id"])
  332. assert result["status"] == "failed"
  333. assert schedule_store.latest_run(schedule["id"])["status"] == "failed"
  334. assert billing_store.get_quota("u", db_conn.day_key())["used_tokens"] == 15
  335. assert concurrency.stats()["global_active"] == 0
  336. @pytest.mark.asyncio
  337. async def test_upload_limit_removes_partial_file(monkeypatch, pricing_ready):
  338. from fastapi import UploadFile, HTTPException
  339. from io import BytesIO
  340. from api.chat_routes import _save_upload
  341. from agents.workspace import session_root
  342. monkeypatch.setattr("api.chat_routes.MAX_UPLOAD_SIZE", 5)
  343. with pytest.raises(HTTPException) as error:
  344. await _save_upload(UploadFile(filename="data.txt", file=BytesIO(b"too long")), "u", "s")
  345. assert error.value.status_code == 413
  346. assert not list((session_root("u", "s") / "workspace/uploads").iterdir())
  347. def test_child_and_summary_usage_count_without_leaking_text_or_context():
  348. from api.sse_core import ChatStream, StreamRequest
  349. stream = ChatStream(StreamRequest(agent=None, session_id="s", thread_id="s", config={}, message="q"))
  350. main = AIMessageChunk(content="主回答", usage_metadata={"input_tokens": 100, "output_tokens": 10, "total_tokens": 110})
  351. child = AIMessageChunk(content="子任务内部材料", usage_metadata={"input_tokens": 10, "output_tokens": 2, "total_tokens": 12})
  352. summary = AIMessageChunk(content="内部摘要", usage_metadata={"input_tokens": 20, "output_tokens": 3, "total_tokens": 23})
  353. assert stream._handle_messages((main, {}), ())
  354. assert stream._handle_messages((child, {}), ("task:1",)) == []
  355. assert stream._handle_messages((summary, {"lc_source": "summarization"}), ()) == []
  356. assert stream.answer_parts == ["主回答"]
  357. assert stream.meter.total_tokens == 145
  358. assert stream.meter.context_tokens == 100
  359. assert not stream.meter.compacted
  360. @pytest.mark.asyncio
  361. async def test_scheduler_shutdown_waits_for_cleanup(monkeypatch):
  362. from scheduler import scheduler
  363. entered = asyncio.Event()
  364. cleaned = False
  365. calls = []
  366. async def work(*args, **kwargs):
  367. nonlocal cleaned
  368. entered.set()
  369. try:
  370. await asyncio.Event().wait()
  371. finally:
  372. await asyncio.sleep(0)
  373. cleaned = True
  374. class Scheduler:
  375. def pause(self):
  376. calls.append("pause")
  377. def shutdown(self, **kwargs):
  378. assert cleaned
  379. calls.append("shutdown")
  380. monkeypatch.setattr(scheduler, "run_schedule", work)
  381. monkeypatch.setattr(scheduler, "_scheduler", Scheduler())
  382. task = asyncio.create_task(scheduler._invoke("s"))
  383. await entered.wait()
  384. await scheduler.shutdown()
  385. assert task.done() and cleaned
  386. assert calls == ["pause", "shutdown"]
  387. assert not scheduler._active
  388. @pytest.mark.asyncio
  389. async def test_admission_revalidates_deleted_session(pricing_ready):
  390. from agents.execution import admit
  391. session = chat_store.create_session("u")
  392. chat_store.delete_session(session["id"])
  393. with pytest.raises(concurrency.SessionMissingError):
  394. await admit("u", session["id"])
  395. assert concurrency.stats()["global_active"] == 0
  396. assert not concurrency._sessions
  397. @pytest.mark.asyncio
  398. async def test_cancellation_during_persistence_does_not_duplicate_answer(monkeypatch, pricing_ready):
  399. import threading
  400. from runtime.execution import admit, build_execution
  401. session = chat_store.create_session("u")
  402. entered, release = threading.Event(), threading.Event()
  403. original = chat_store.save_message
  404. def slow_save(session_id, role, content, **kwargs):
  405. if role == "assistant":
  406. entered.set()
  407. assert release.wait(3)
  408. return original(session_id, role, content, **kwargs)
  409. monkeypatch.setattr(chat_store, "save_message", slow_save)
  410. admission = await admit("u", session["id"])
  411. execution = await build_execution(admission, agent=FakeAgent(), session_id=session["id"],
  412. thread_id=session["id"], user_name="u", message="question", config={"configurable": {"thread_id": session["id"]}})
  413. async def consume():
  414. async for _ in execution.events():
  415. pass
  416. task = asyncio.create_task(consume())
  417. assert await asyncio.to_thread(entered.wait, 3)
  418. task.cancel()
  419. await asyncio.sleep(0)
  420. release.set()
  421. with pytest.raises(asyncio.CancelledError):
  422. await task
  423. assert len([m for m in chat_store.get_messages(session["id"]) if m["msg_role"] == "assistant"]) == 1
  424. assert len(billing_store.recent_usage("u")) == 1
  425. assert concurrency.stats()["global_active"] == 0
  426. @pytest.mark.asyncio
  427. @pytest.mark.parametrize("delegated", [False, True])
  428. async def test_extension_approval_cannot_be_bypassed_by_main_or_subagent(monkeypatch, pricing_ready, delegated):
  429. import sys
  430. from extensions import registry as extensions
  431. from extensions import define_tool
  432. from tools import registry as tools
  433. from agents import agent_factory, agent_memory
  434. from main import app
  435. writes = []
  436. @define_tool(name="adjust_fan", group="通风")
  437. def write_fan(speed: int) -> str:
  438. """调整风机速度。"""
  439. writes.append(speed)
  440. return "完成"
  441. def setup(catalog):
  442. catalog.register_agent(name="ventilation", description="通风控制", prompt="使用工具调整风机。",
  443. tools=[write_fan], approval_tools=[write_fan])
  444. modules = extensions.discover_modules()
  445. monkeypatch.setitem(sys.modules, "extensions.ventilation.extension", SimpleNamespace(setup=setup))
  446. monkeypatch.setattr(extensions, "discover_modules", lambda: [*modules, "extensions.ventilation.extension"])
  447. tools.reset_for_tests()
  448. responses = []
  449. if delegated:
  450. responses.append(AIMessage(content="", tool_calls=[{"id": "delegate", "name": "task",
  451. "args": {"subagent_type": "ventilation", "description": "调整风机"}}]))
  452. responses.extend([
  453. AIMessage(content="", tool_calls=[{"id": "fan", "name": "adjust_fan", "args": {"speed": 10}}]),
  454. AIMessage(content="已完成"), AIMessage(content="已汇总"),
  455. ])
  456. model = ScriptedModel(responses=responses)
  457. monkeypatch.setattr("agents.model_factory.get_model", lambda: model)
  458. graph = agent_factory.build_agent()
  459. assert tools.get_spec("adjust_fan").requires_approval
  460. monkeypatch.setattr("agents.agent_registry.get_agent", lambda: graph)
  461. async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
  462. first = await client.post("/api/chat", data={"message": "调整风机", "user_name": "u"})
  463. events = [json.loads(line[6:]) for line in first.text.splitlines() if line.startswith("data: ")]
  464. assert not [e for e in events if e["type"] == "error"], events
  465. interrupted = next(e for e in events if e["type"] == "interrupt")
  466. assert writes == []
  467. sid = first.headers["x-session-id"]
  468. token = interrupted["tool_approval"]["interrupt_token"]
  469. response = await client.post("/api/chat/resume", json={"session_id": sid, "user_name": "u",
  470. "interrupt_token": token, "approved": True})
  471. assert response.status_code == 200, response.text
  472. assert writes == [10], response.text
  473. assert run_store.list_runs(sid)[0]["status"] == "success", response.text
  474. await agent_memory.delete_thread(sid)