mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 10:41:59 +00:00
ThearchyHelios
269 lines
11 KiB
Python
269 lines
11 KiB
Python
"""运行选项和来源：Run.options 是发起时定下、整次运行沿用的选项快照，Run.schedule_id 是发起它的定时任务。
|
|
|
|
- 旧库升级时补这两列和 schedule_id 的索引；存量运行的选项读出来是 {}，来源为空
|
|
- RunManager.start 收下两者并落库；选项非空时写进 run.started（进封存范围），为空时 run.started 和以前
|
|
逐字相同；恢复、继续运行沿用发起时的选项，节点从 RunContext.extra 读得到同一份
|
|
- 接口带出 schedule_id，运行列表能按定时任务筛；launch_run 原样往下交
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import uuid
|
|
|
|
import pytest
|
|
from httpx import ASGITransport, AsyncClient
|
|
from sqlalchemy import create_engine, select, text
|
|
|
|
from app.db.base import SessionLocal, engine as app_engine
|
|
from app.db.models import Run, RunEvent
|
|
from app.engine import compiler
|
|
from app.engine.runner import run_manager
|
|
from app.engine.schema import NodeType
|
|
from app.main import app
|
|
|
|
#: 加 options 之前 run.started 就有的键：选项为空时一个不多、一个不少
|
|
LEGACY_STARTED_KEYS = {"nodes", "resumed", "memory_scope", "collection", "approval_default", "tool_trust",
|
|
"agent_limits", "data_versions", "replay_protocol"}
|
|
OPTIONS = {"verify": {"dual_check": True, "max_queries": 6}, "note": "第 40 周例行核查"}
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
async def engine_up():
|
|
await run_manager.setup()
|
|
yield
|
|
await run_manager.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
async def client():
|
|
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as c:
|
|
yield c
|
|
|
|
|
|
def node(nid, ntype, **config):
|
|
return {"id": nid, "type": ntype, "position": {"x": 0, "y": 0}, "data": {"label": nid, "config": config}}
|
|
|
|
|
|
def chain(*nodes):
|
|
return {"nodes": list(nodes), "edges": [{"source": a["id"], "target": b["id"]} for a, b in zip(nodes, nodes[1:])]}
|
|
|
|
|
|
SIMPLE = chain(node("start", "input"), node("out", "output", fields=[{"name": "r", "value": "完成"}]))
|
|
GATED = chain(node("start", "input"),
|
|
node("before", "transform", mode="template", template="前"),
|
|
node("gate", "human", mode="approve", title="放行吗"),
|
|
node("after", "transform", mode="template", template="后"),
|
|
node("out", "output", fields=[{"name": "r", "value": "{{ nodes.after }}"}]))
|
|
|
|
|
|
async def _settle(run_id: str, statuses=("succeeded", "failed", "cancelled")) -> Run:
|
|
for _ in range(400):
|
|
async with SessionLocal() as session:
|
|
row = await session.get(Run, run_id)
|
|
if row.status in statuses and (row.status == "interrupted" or row.manifest_seq is not None):
|
|
await run_manager.wait_idle(run_id)
|
|
return row
|
|
await asyncio.sleep(0.03)
|
|
raise AssertionError(f"运行 {run_id} 停在 {row.status}")
|
|
|
|
|
|
async def _started(run_id: str) -> list[dict]:
|
|
async with SessionLocal() as session:
|
|
return [e.data for e in (await session.execute(
|
|
select(RunEvent).where(RunEvent.run_id == run_id, RunEvent.type == "run.started").order_by(RunEvent.seq)
|
|
)).scalars()]
|
|
|
|
|
|
def _schedule() -> str:
|
|
return uuid.uuid4().hex
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# 旧库升级
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
def test_old_databases_get_both_columns_and_the_index():
|
|
from app.db.base import _migrate
|
|
|
|
db = create_engine("sqlite://")
|
|
with db.begin() as conn:
|
|
conn.execute(text("CREATE TABLE runs (id VARCHAR(32) PRIMARY KEY, workflow_name VARCHAR(200), "
|
|
"status VARCHAR(20), last_seq INTEGER)"))
|
|
conn.execute(text("INSERT INTO runs (id, workflow_name, status, last_seq) VALUES ('r1', '周报', 'succeeded', 9)"))
|
|
_migrate(conn)
|
|
_migrate(conn) # 再执行一次不出错：列和索引都已经在了
|
|
cols = {row[1] for row in conn.execute(text("PRAGMA table_info(runs)"))}
|
|
assert {"options", "schedule_id"} <= cols
|
|
assert conn.execute(text("SELECT options, schedule_id FROM runs WHERE id = 'r1'")).one() == ("{}", None)
|
|
indexes = {row[1]: row for row in conn.execute(text("PRAGMA index_list(runs)"))}
|
|
assert "ix_runs_schedule_id" in indexes
|
|
assert [r[2] for r in conn.execute(text("PRAGMA index_info(ix_runs_schedule_id)"))] == ["schedule_id"]
|
|
|
|
|
|
def test_a_missing_table_is_left_alone():
|
|
"""还没有 runs 表的库（create_all 之前不会这样，这里只确认迁移不替它建索引、不报错）。"""
|
|
from app.db.base import _migrate
|
|
|
|
db = create_engine("sqlite://")
|
|
with db.begin() as conn:
|
|
_migrate(conn)
|
|
assert conn.execute(text("SELECT name FROM sqlite_master WHERE name = 'ix_runs_schedule_id'")).first() is None
|
|
|
|
|
|
async def test_new_databases_index_schedule_id():
|
|
async with app_engine.connect() as conn:
|
|
names = {row[1] for row in (await conn.execute(text("PRAGMA index_list(runs)"))).all()}
|
|
assert "ix_runs_schedule_id" in names
|
|
|
|
|
|
async def test_a_migrated_row_reads_back_empty_options():
|
|
"""迁移回填的是字符串 '{}'：经 ORM 读出来是空 dict，不是字符串，也不是 None。"""
|
|
async with SessionLocal() as session:
|
|
run = Run(workflow_name="迁移回填", status="succeeded", graph={}, input={})
|
|
session.add(run)
|
|
await session.commit()
|
|
run_id = run.id
|
|
async with app_engine.begin() as conn:
|
|
await conn.execute(text("UPDATE runs SET options = '{}' WHERE id = :id"), {"id": run_id})
|
|
async with SessionLocal() as session:
|
|
row = await session.get(Run, run_id)
|
|
assert row.options == {} and row.schedule_id is None
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# 发起：落库、run.started、沿用
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
async def test_options_and_schedule_are_stored_and_sealed():
|
|
sid = _schedule()
|
|
run = await run_manager.start(graph=SIMPLE, input_payload={}, options=OPTIONS, schedule_id=sid)
|
|
assert (run.options, run.schedule_id) == (OPTIONS, sid)
|
|
row = await _settle(run.id)
|
|
assert row.status == "succeeded", row.error
|
|
assert (row.options, row.schedule_id) == (OPTIONS, sid)
|
|
[started] = await _started(run.id)
|
|
assert started["options"] == OPTIONS
|
|
assert set(started) == LEGACY_STARTED_KEYS | {"options"}
|
|
|
|
|
|
@pytest.mark.parametrize("options", [None, {}])
|
|
async def test_without_options_run_started_is_unchanged(options):
|
|
run = await run_manager.start(graph=SIMPLE, input_payload={}, options=options)
|
|
row = await _settle(run.id)
|
|
assert row.options == {} and row.schedule_id is None
|
|
[started] = await _started(run.id)
|
|
assert set(started) == LEGACY_STARTED_KEYS, "选项为空时 run.started 多出了键"
|
|
|
|
|
|
async def test_the_options_belong_to_the_run_not_to_the_callers_dict():
|
|
mine = {"verify": {"dual_check": True}}
|
|
run = await run_manager.start(graph=SIMPLE, input_payload={}, options=mine)
|
|
mine["verify"]["dual_check"] = False
|
|
mine["late"] = 1
|
|
await _settle(run.id)
|
|
async with SessionLocal() as session:
|
|
assert (await session.get(Run, run.id)).options == {"verify": {"dual_check": True}}
|
|
|
|
|
|
@pytest.fixture
|
|
def seen_options(monkeypatch):
|
|
"""记下每次执行「数据整形」节点时它从运行上下文里读到的选项。"""
|
|
from app.engine.runner import RUN_OPTIONS_KEY
|
|
|
|
seen: list = []
|
|
original = compiler.RUNNERS[NodeType.TRANSFORM]
|
|
|
|
async def spy(state, ctx):
|
|
seen.append(ctx.run.extra.get(RUN_OPTIONS_KEY))
|
|
return await original(state, ctx)
|
|
|
|
monkeypatch.setitem(compiler.RUNNERS, NodeType.TRANSFORM, spy)
|
|
return seen
|
|
|
|
|
|
async def test_resume_keeps_the_options_the_run_started_with(seen_options):
|
|
run = await run_manager.start(graph=GATED, input_payload={}, options=OPTIONS)
|
|
assert (await _settle(run.id, ("interrupted",))).status == "interrupted"
|
|
await run_manager.resume(run.id, {"approved": True})
|
|
row = await _settle(run.id)
|
|
assert row.status == "succeeded", row.error
|
|
assert seen_options == [OPTIONS, OPTIONS], "恢复之后节点读到的选项变了"
|
|
first, second = await _started(run.id)
|
|
assert first["options"] == second["options"] == OPTIONS and second["resumed"] is True
|
|
|
|
|
|
async def test_continue_keeps_the_options_too(seen_options, monkeypatch):
|
|
original = compiler.RUNNERS[NodeType.TRANSFORM]
|
|
calls = {"n": 0}
|
|
|
|
async def flaky(state, ctx):
|
|
calls["n"] += 1
|
|
if ctx.node.id == "after" and calls["n"] <= 2:
|
|
raise RuntimeError("临时故障")
|
|
return await original(state, ctx)
|
|
|
|
monkeypatch.setitem(compiler.RUNNERS, NodeType.TRANSFORM, flaky)
|
|
graph = chain(node("start", "input"),
|
|
node("before", "transform", mode="template", template="前"),
|
|
node("after", "transform", mode="template", template="后"),
|
|
node("out", "output", fields=[{"name": "r", "value": "{{ nodes.after }}"}]))
|
|
run = await run_manager.start(graph=graph, input_payload={}, options=OPTIONS)
|
|
assert (await _settle(run.id)).status == "failed"
|
|
await run_manager.continue_failed(run.id)
|
|
row = await _settle(run.id)
|
|
assert row.status == "succeeded", row.error
|
|
assert [s.get("options") for s in await _started(run.id)] == [OPTIONS, OPTIONS]
|
|
|
|
|
|
async def test_nodes_read_an_empty_dict_when_there_are_no_options(seen_options):
|
|
run = await run_manager.start(graph=GATED, input_payload={})
|
|
await _settle(run.id, ("interrupted",))
|
|
assert seen_options == [{}]
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# 接口：带出来源、按定时任务筛、launch_run 照转
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
async def test_the_api_shows_and_filters_by_schedule(client):
|
|
sid, other = _schedule(), _schedule()
|
|
mine = [await run_manager.start(graph=SIMPLE, input_payload={}, schedule_id=sid) for _ in range(2)]
|
|
theirs = await run_manager.start(graph=SIMPLE, input_payload={}, schedule_id=other)
|
|
manual = await run_manager.start(graph=SIMPLE, input_payload={})
|
|
for run in (*mine, theirs, manual):
|
|
await _settle(run.id)
|
|
|
|
got = (await client.get(f"/api/runs/{mine[0].id}")).json()
|
|
assert got["schedule_id"] == sid
|
|
assert (await client.get(f"/api/runs/{manual.id}")).json()["schedule_id"] is None
|
|
|
|
listed = (await client.get("/api/runs", params={"schedule_id": sid})).json()
|
|
assert sorted(r["id"] for r in listed) == sorted(r.id for r in mine)
|
|
assert all(r["schedule_id"] == sid for r in listed)
|
|
# 不带这个参数时照旧全列
|
|
everything = {r["id"] for r in (await client.get("/api/runs", params={"limit": 200})).json()}
|
|
assert {theirs.id, manual.id} <= everything
|
|
|
|
|
|
async def test_launch_run_passes_both_through():
|
|
from app.engine.launch import launch_run
|
|
|
|
sid = _schedule()
|
|
async with SessionLocal() as session:
|
|
run = await launch_run(session, graph=SIMPLE, actor="定时任务", schedule_id=sid, options=OPTIONS)
|
|
row = await _settle(run.id)
|
|
assert (row.schedule_id, row.options, row.started_by) == (sid, OPTIONS, "定时任务")
|
|
[started] = await _started(run.id)
|
|
assert started["options"] == OPTIONS
|
|
|
|
|
|
async def test_an_http_launch_has_no_schedule_and_no_options(client):
|
|
res = await client.post("/api/runs", json={"graph": SIMPLE})
|
|
assert res.status_code == 201, res.text
|
|
assert res.json()["schedule_id"] is None
|
|
row = await _settle(res.json()["id"])
|
|
assert row.options == {}
|