"""agent 护栏：取代以前固定的「最多 12 步」。 用户的原话是「12 太少了，撞上就判定运行失败」。真实运行里（719253eb、6310f291）agent 用满 12 步收了尾，结论缺数据，下游的结构化校验缺字段，整个运行失败，报错里看不出根源。 现在： - 步数只剩一个很高的兜底（默认 100，设置里可改）；以前画布写死的 12 按「没填」处理； - 同一个调用第二次不再执行，把上次结果交还；连续几步拿不到新信息就收尾； - 令牌 / 金额预算，可以设成不限； - 上下文接近窗口先压缩早期工具结果，再不够就收尾； - 临近上限提醒模型一次； - 下游校验失败时写明是上游哪个 agent 提前收了尾； - 升级前发起的运行（没有快照）一字不改走旧逻辑。 """ from __future__ import annotations import asyncio import pytest from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, ToolMessage from langchain_core.outputs import ChatGenerationChunk from sqlalchemy import select from app.db.base import SessionLocal from app.db.models import CustomTool, Run, RunEvent, Setting from app.engine import runner as runner_mod from app.engine.guards import node_limits, run_limits from app.engine.runner import run_manager from app.tools import custom as custom_tools FINAL = "结论：都查完了。" SETTLED = "收尾：按已经查到的部分。" TOOL = "shop_lookup" @pytest.fixture(autouse=True) async def engine_up(): await run_manager.setup() yield await run_manager.shutdown() @pytest.fixture(autouse=True) async def clean_run_settings(): async def drop(): async with SessionLocal() as session: row = await session.get(Setting, "run") if row: await session.delete(row) await session.commit() await drop() yield await drop() @pytest.fixture async def lookup(monkeypatch): """自定义工具：按参数返回一段文字，长度可调。真正的请求换成记账。""" calls: list[dict] = [] size = {"chars": 20} async def fake_http(row, args): calls.append(dict(args)) return f"门店 {args.get('q')}：" + "数" * size["chars"] monkeypatch.setattr(custom_tools, "_run_http", fake_http) async with SessionLocal() as session: session.add(CustomTool(name=TOOL, kind="http", description="查门店", parameters={"type": "object", "properties": {"q": {"type": "string"}}, "required": ["q"]}, config={"url": "https://example.com/shop"})) await session.commit() yield calls, size async with SessionLocal() as session: for row in (await session.execute(select(CustomTool).where(CustomTool.name == TOOL))).scalars(): await session.delete(row) await session.commit() def graph(**config): cfg = {"prompt": "查门店", "tools": [TOOL], "approval": "never", **config} return { "nodes": [ {"id": "in", "type": "input", "data": {"config": {}}}, {"id": "work", "type": "agent", "data": {"label": "查数", "config": cfg}}, {"id": "out", "type": "output", "data": {"config": {}}}, ], "edges": [{"source": "in", "target": "work"}, {"source": "work", "target": "out"}], } def script(monkeypatch, plan, *, usage=None): """plan(第几次调用, messages) → 查询参数（字符串）或 None（给结论）。 流式和非流式都换掉：mock 走流式时不报用量，预算的测试要自己带上（usage）。 seen 记下每次调用模型时看到的 messages。 """ from app.providers import mock_model seen: list[list] = [] def decide(self, messages): if not self.tools: return AIMessage(content=SETTLED) seen.append(list(messages)) q = plan(len(seen) - 1, messages) if q is None: return AIMessage(content=FINAL) return AIMessage(content="再查一个。", tool_calls=[ {"name": TOOL, "args": {"q": q}, "id": f"call_{len(seen)}"}]) async def astream(self, messages, stop=None, run_manager=None, **kwargs): reply = decide(self, messages) extra = {"usage_metadata": usage(messages, reply)} if usage else {} yield ChatGenerationChunk(message=AIMessageChunk( content=reply.content, tool_calls=reply.tool_calls, **extra)) monkeypatch.setattr(mock_model.MockChatModel, "_decide", decide) monkeypatch.setattr(mock_model.MockChatModel, "_astream", astream) return seen async def run(g, *, settings_run: dict | None = None): if settings_run is not None: async with SessionLocal() as session: session.add(Setting(key="run", value=settings_run)) await session.commit() started = await run_manager.start(graph=g, input_payload={"question": "有几家门店"}) for _ in range(400): await asyncio.sleep(0.05) async with SessionLocal() as session: row = await session.get(Run, started.id) if row.status in ("succeeded", "failed", "cancelled", "interrupted"): break else: pytest.fail("运行没有在 20 秒内结束") async with SessionLocal() as session: events = list((await session.execute( select(RunEvent).where(RunEvent.run_id == started.id).order_by(RunEvent.seq))).scalars()) return row, events def logs(events, code): return [e.data for e in events if e.type == "log" and e.data.get("code") == code] def node_out(row): return row.output # ---------------------------------------------------------------- 上限怎么定 def test_run_limits_from_settings(): assert run_limits({}, 100) == {"max_steps": 100, "budget_tokens": 2_000_000, "budget_usd": None} # 写 null / 0 / 空都是不限；存坏了的步数按缺省 assert run_limits({"agent_budget_tokens": None, "agent_budget_usd": 0}, 100)["budget_tokens"] is None assert run_limits({"agent_budget_tokens": ""}, 100)["budget_tokens"] is None assert run_limits({"agent_max_steps": "abc"}, 100)["max_steps"] == 100 # 环境变量的硬顶压过设置 assert run_limits({"agent_max_steps": 500}, 100)["max_steps"] == 100 assert run_limits({"agent_budget_usd": 1.5}, 100)["budget_usd"] == 1.5 def test_node_overrides_and_the_legacy_twelve(): snap = {"max_steps": 100, "budget_tokens": 2_000_000, "budget_usd": None} # 画布以前自动写进去的 12：当作没填 assert node_limits({"max_steps": 12}, snap, 100).max_steps == 100 assert node_limits({"max_steps": "12"}, snap, 100).max_steps == 100 assert node_limits({"max_steps": 20}, snap, 100).max_steps == 20 assert node_limits({"max_steps": 300}, snap, 100).max_steps == 100 got = node_limits({"budget_tokens": 5000, "budget_usd": 0.5}, snap, 100) assert (got.budget_tokens, got.budget_usd) == (5000, 0.5) assert node_limits({}, {**snap, "budget_tokens": None}, 100).budget_tokens is None async def test_settings_expose_the_defaults(): from httpx import ASGITransport, AsyncClient from app.main import app async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as c: got = (await c.get("/api/settings")).json()["run"] assert (got["agent_max_steps"], got["agent_budget_tokens"], got["agent_budget_usd"]) == (100, 2_000_000, None) # ---------------------------------------------------------------- 步数 async def test_twelve_steps_no_longer_cut_it_short(lookup, monkeypatch): """查 15 个门店再下结论：以前第 12 步就被掐断，现在跑完。节点上写着 12 也一样。""" calls, _ = lookup script(monkeypatch, lambda i, _: f"店{i}" if i < 15 else None) row, events = await run(graph(max_steps=12)) assert row.status == "succeeded", row.error assert len(calls) == 15 assert not logs(events, "step_limit_settled") assert "limited" not in str(row.output) async def test_an_explicit_step_limit_still_settles(lookup, monkeypatch): calls, _ = lookup script(monkeypatch, lambda i, _: f"店{i}") row, events = await run(graph(max_steps=4)) assert row.status == "succeeded", row.error assert len(calls) == 4 [settled] = logs(events, "step_limit_settled") assert settled["reason"] == "steps" and "4 步" in settled["message"] async def test_the_settings_default_applies(lookup, monkeypatch): calls, _ = lookup script(monkeypatch, lambda i, _: f"店{i}") row, events = await run(graph(), settings_run={"agent_max_steps": 6}) assert len(calls) == 6 assert logs(events, "step_limit_settled")[0]["reason"] == "steps" async def test_run_started_records_the_limits(lookup, monkeypatch): script(monkeypatch, lambda i, _: None) row, events = await run(graph(), settings_run={"agent_budget_tokens": None, "agent_max_steps": 30}) started = next(e for e in events if e.type == "run.started") assert started.data["agent_limits"] == {"max_steps": 30, "budget_tokens": None, "budget_usd": None} # ---------------------------------------------------------------- 重复与停滞 async def test_a_repeated_call_is_answered_from_the_first_result(lookup, monkeypatch): calls, _ = lookup seen = script(monkeypatch, lambda i, _: ["店A", "店A", None][i]) row, events = await run(graph()) assert row.status == "succeeded", row.error assert calls == [{"q": "店A"}], "同样的调用又执行了一次" [repeat] = logs(events, "tool_repeat") assert TOOL in repeat["message"] reused = [m for m in seen[2] if isinstance(m, ToolMessage)][-1].content assert "已经执行过" in reused and "门店 店A" in reused async def test_going_in_circles_settles(lookup, monkeypatch): """一直用同样的参数查：第一次真查，之后每步都拿不到新信息，连续 3 步就收尾。""" calls, _ = lookup script(monkeypatch, lambda i, _: "店A") row, events = await run(graph()) assert row.status == "succeeded", row.error assert calls == [{"q": "店A"}] assert len(logs(events, "tool_repeat")) == 3, "没进展的 3 步之后就该收尾，不能一路耗到步数兜底" [settled] = logs(events, "step_limit_settled") assert settled["reason"] == "stall" assert SETTLED in str(row.output) async def test_new_information_resets_the_stall_count(lookup, monkeypatch): calls, _ = lookup # A、A、A（两步没进展）、B（有进展，清零）、B、B（又两步）、结论 order = ["店A", "店A", "店A", "店B", "店B", "店B", None] script(monkeypatch, lambda i, _: order[i]) row, events = await run(graph()) assert row.status == "succeeded", row.error assert calls == [{"q": "店A"}, {"q": "店B"}] assert not logs(events, "step_limit_settled") # ---------------------------------------------------------------- 预算 def per_call(tokens): return lambda messages, reply: {"input_tokens": tokens, "output_tokens": 0, "total_tokens": tokens} async def test_token_budget_settles(lookup, monkeypatch): calls, _ = lookup script(monkeypatch, lambda i, _: f"店{i}", usage=per_call(1000)) row, events = await run(graph(budget_tokens=3500)) assert row.status == "succeeded", row.error # 第 4 次调用之后累计 4000，超过 3500：那一步的工具照常跑完，然后收尾 assert len(calls) == 4 [settled] = logs(events, "step_limit_settled") assert settled["reason"] == "budget_tokens" and "3,500" in settled["message"] async def test_budget_can_be_unlimited(lookup, monkeypatch): calls, _ = lookup # 用量记在输出上：输入一大就先撞上下文水位了，这里只测预算 script(monkeypatch, lambda i, _: f"店{i}" if i < 8 else None, usage=lambda m, r: {"input_tokens": 10, "output_tokens": 1_000_000, "total_tokens": 1_000_010}) row, events = await run(graph(), settings_run={"agent_budget_tokens": None}) assert row.status == "succeeded", row.error assert len(calls) == 8 and not logs(events, "step_limit_settled") async def test_usd_budget_settles_when_prices_are_known(lookup, monkeypatch): from app.providers import catalog monkeypatch.setattr(catalog, "estimate_cost", lambda model, i, o: i / 1000 * 0.01) calls, _ = lookup script(monkeypatch, lambda i, _: f"店{i}", usage=per_call(1000)) row, events = await run(graph(budget_usd=0.025)) assert len(calls) == 3 assert logs(events, "step_limit_settled")[0]["reason"] == "budget_usd" async def test_reminder_before_the_budget_runs_out(lookup, monkeypatch): seen = script(monkeypatch, lambda i, _: f"店{i}" if i < 5 else None, usage=per_call(1000)) row, _ = await run(graph(budget_tokens=5000)) notes = [m.content for call in seen for m in call if isinstance(m, HumanMessage) and "系统提醒" in m.content] assert notes and "token 预算已用 80%" in notes[0] assert len(set(notes)) == 1, "同一种提醒只该说一次" async def test_reminder_before_the_steps_run_out(lookup, monkeypatch): seen = script(monkeypatch, lambda i, _: f"店{i}") await run(graph(max_steps=8)) # 第 5 步（下标 4）之后还剩 3 步：第 6 次调用时看得到提醒 assert any("还剩 3 步" in m.content for m in seen[5] if isinstance(m, HumanMessage)) assert not any("系统提醒" in m.content for m in seen[4] if isinstance(m, HumanMessage)) # ---------------------------------------------------------------- 上下文 async def test_old_tool_results_are_compressed_near_the_window(lookup, monkeypatch): from app.providers import catalog calls, size = lookup size["chars"] = 3000 # 没报用量时按字数的一半估：每条结果约 1500 token，第 6 次调用前后到窗口的 60% monkeypatch.setattr(catalog, "find_context", lambda model: 12_000) seen = script(monkeypatch, lambda i, _: f"店{i}" if i < 7 else None) row, events = await run(graph()) assert row.status == "succeeded", row.error assert logs(events, "context_compressed") last = [m.content for m in seen[-1] if isinstance(m, ToolMessage)] assert any("早期的工具结果已压缩" in c for c in last[:-3]) assert all("早期的工具结果已压缩" not in c for c in last[-3:]), "最近 3 条不该压缩" async def test_too_close_to_the_window_settles(lookup, monkeypatch): from app.providers import catalog calls, size = lookup size["chars"] = 3000 monkeypatch.setattr(catalog, "find_context", lambda model: 1_000) script(monkeypatch, lambda i, _: f"店{i}") row, events = await run(graph()) assert row.status == "succeeded", row.error assert logs(events, "step_limit_settled")[0]["reason"] == "context" assert len(calls) <= 2 # ---------------------------------------------------------------- 下游校验与老运行 async def test_validation_failure_names_the_agent_that_was_cut_short(lookup, monkeypatch): script(monkeypatch, lambda i, _: f"店{i}") g = graph(max_steps=3) g["nodes"][2] = {"id": "check", "type": "validate", "data": {"label": "校验", "config": { "source": "{{ nodes.work.text }}", "repair_with_llm": False, "schema": {"type": "object", "required": ["stores"], "properties": {"stores": {"type": "array"}}}}}} g["edges"][1] = {"source": "work", "target": "check"} row, _ = await run(g) assert row.status == "failed" assert "上游「查数」因步数用完" in (row.error or ""), row.error async def test_runs_from_before_the_upgrade_keep_the_old_rules(lookup, monkeypatch): """没有快照的运行：默认 12 步，重复调用照样执行——它们可能正停在某个审批上， 多一次复用、少一次调用，重放时 checkpoint 就对错号了。""" async def no_snapshot(session): return None monkeypatch.setattr(runner_mod, "_agent_limits", no_snapshot) calls, _ = lookup script(monkeypatch, lambda i, _: "店A") row, events = await run(graph()) assert row.status == "succeeded", row.error assert len(calls) == 12 [settled] = logs(events, "step_limit_settled") assert "reason" not in settled and "已用完 12 步" in settled["message"] assert not logs(events, "tool_repeat") def test_review_signal_carries_the_reason(): """问数据页只在步数用满时给「放宽步数重跑」：复核信号要带上收尾原因。""" from app.engine.review import scan def log(code, reason=None): data = {"level": "warn", "code": code, "message": f"{code} {reason}"} if reason: data["reason"] = reason return {"type": "log", "node_id": "work", "data": data} got = {s.detail: s.as_dict() for s in scan([log("step_limit_settled", "stall"), log("step_limit_settled", "steps"), log("step_limit_settled")], {"a": "答案"})} assert got["step_limit_settled stall"]["reason"] == "stall" assert got["step_limit_settled steps"]["reason"] == "steps" assert "reason" not in got["step_limit_settled None"], "老运行的信号不该多出字段" async def test_settings_limits_reflect_the_real_config(): """库里存着旧版设置页写进去的 limits（见过 max_agent_steps: 25）：不生效，也不能盖住真实上限。""" from httpx import ASGITransport, AsyncClient from app.core.config import settings as app_settings from app.main import app async with SessionLocal() as session: row = await session.get(Setting, "limits") if row: row.value = {"max_agent_steps": 25} else: session.add(Setting(key="limits", value={"max_agent_steps": 25})) await session.commit() try: async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as c: assert (await c.get("/api/settings")).json()["limits"]["max_agent_steps"] == app_settings.max_agent_steps await c.put("/api/settings", json={"values": {"limits": {"max_agent_steps": 7}}}) async with SessionLocal() as session: assert (await session.get(Setting, "limits")).value == {"max_agent_steps": 25}, "limits 不该被写" finally: async with SessionLocal() as session: row = await session.get(Setting, "limits") if row: await session.delete(row) await session.commit()