mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 10:11:58 +00:00
ThearchyHelios
7739340607
- 问数据和个人应用的图放行绑定只读查询等合规工具的调用工具节点，只读和遮罩在这条路径上生效；被治理成写入的工具在只读运行里直接拒；知识库范围在工具层固定（图级默认值、参数名纠错、Agent 都改不了） - 个人应用的刷新和个人蓝图的试运行只有拥有者能读：运行、事件、证据、工件、文件、运行与审批列表、对比、可验证报告、质询、假设分析、提取工作流一并收紧；个人蓝图的试运行不论谁发起都按个人应用约束 - 共享蓝图只给拥有者、对应应用的管理者、管理员，看不到的一律 404；审核副本归审核人；管理员看别人的个人应用只有元数据、可以归档 - 记录表跨版本汇总（只读挂载旧版本、按当前列 UNION ALL、各版本各自固定高水位，缺的版本照实写进印章）；停用前发起的运行可以继续写旧版本；更正沿用原提交人、更正人记审计 - 约定收口：待办种类 review、事件 file.extracted、表单错误 file、模板定义带示例与预置；前端类型一次补齐、去掉临时的局部类型；审核待办卡不再给打不开的链接；读附件事件有中文说法；检查脚本补断言
666 lines
36 KiB
Python
666 lines
36 KiB
Python
"""开登录后使用者的问数据（D-c）：能发起什么样的运行、查询结果在工具边界遮、运行的 WebSocket 按运行判、数据源名单。
|
||
|
||
- 图：只许输入、模型、Agent、报告、成果、知识检索、数据整形、条件分支、口径卡、合并查询、调用工具；Agent 和调用工具
|
||
节点只许数据源工具和只读的内置工具（app/auth/chat_policy.py）。沙箱代码、自定义工具、MCP、网页、HTTP、文件、平台的
|
||
写入能力一律 403；
|
||
- 调用工具节点查数据源：和 Agent 同一条路，结果在工具边界遮、可写源上的写语句被守卫拒掉；能力治理里改成写入的工具
|
||
在只读运行里不执行；
|
||
- 使用者发起的问数据运行一定带 MASK_RESULTS_OPTION：真的执行一次（演示模型调数据源查询），事件、快照、实时流里都没有
|
||
遮罩列的原值；
|
||
- WebSocket：没登录的握手前关掉；使用者只能连自己能读、查询结果在工具边界遮过的运行（握手时判，1008 拒）；
|
||
- GET /api/datasources/names：开了问数据的使用者能读，只有名字、说明、种类，只列启用着的。
|
||
|
||
数据源是现造的 SQLite 文件，姓名、手机号都是虚构的。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import sqlite3
|
||
import uuid
|
||
|
||
import pytest
|
||
from langchain_core.messages import AIMessage, ToolMessage
|
||
from sqlalchemy import delete, select
|
||
|
||
from app.auth.chat_policy import chat_graph_problem
|
||
from app.data import introspect
|
||
from app.db.base import SessionLocal
|
||
from app.db.models import Artifact, DataSource, Run, RunEvent
|
||
from app.engine.runner import run_manager
|
||
from app.tools.datasource import MASK_RESULTS_OPTION
|
||
from tests.fakes import auth_env
|
||
|
||
PHONE = "13900000001"
|
||
MASKED = "已遮罩"
|
||
|
||
|
||
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:])]}
|
||
|
||
|
||
def chat_graph(*tools: str, middle: dict | None = None) -> dict:
|
||
return chain(node("start", "input", fields=[{"name": "question"}]),
|
||
middle or node("ask", "agent", prompt="{{ input.question }}", tools=list(tools), max_steps=4),
|
||
node("out", "output", fields=[{"name": "回答", "value": "{{ nodes.ask.text }}"}]))
|
||
|
||
|
||
# ---------------------------------------------------------------- 图的白名单（纯函数）
|
||
|
||
|
||
@pytest.mark.parametrize("graph", [
|
||
chat_graph("db_query__crm", "db_schema__crm", "calculator", "current_time", "json_query", "kb_search",
|
||
"memory_recall"),
|
||
chain(node("start", "input"), node("kb", "retrieve", query="x"), node("t", "transform", expression="1"),
|
||
node("out", "output")),
|
||
# 调用工具节点：绑定的工具过同一份白名单就行（D1）
|
||
chat_graph(middle=node("ask", "tool", tool="db_query__crm", args={"sql": "SELECT city FROM customers"})),
|
||
chat_graph(middle=node("ask", "tool", tool="calculator", args={"expression": "1 + 1"})),
|
||
chat_graph(middle=node("ask", "tool", tool="kb_search", args={"query": "退款规则"})),
|
||
{"nodes": "坏的"}, # 不成形的图交给 launch_run 照常报校验问题
|
||
])
|
||
def test_graphs_a_user_may_run_in_chat(graph):
|
||
assert chat_graph_problem(graph) is None
|
||
|
||
|
||
@pytest.mark.parametrize("graph, said", [
|
||
(chat_graph(middle=node("ask", "code", code="print(1)")), "沙箱代码"),
|
||
(chat_graph(middle=node("ask", "tool", tool="http_request")), "http_request"),
|
||
(chat_graph(middle=node("ask", "tool", tool="records_append")), "records_append"),
|
||
(chat_graph(middle=node("ask", "tool", tool="memory_write")), "memory_write"),
|
||
(chat_graph(middle=node("ask", "tool", tool="mcp:fs/read_file")), "mcp:fs/read_file"),
|
||
(chat_graph(middle=node("ask", "tool", tool="{{ vars.tool }}")), "模板"),
|
||
(chat_graph(middle=node("ask", "supervisor", agents=[])), "多 Agent 协作"),
|
||
(chat_graph(middle=node("ask", "subgraph", workflow_id="w")), "子工作流"),
|
||
(chat_graph(middle=node("ask", "human", mode="approve")), "人工审批"),
|
||
(chat_graph(middle=node("ask", "memory", action="write")), "长期记忆"),
|
||
(chat_graph("http_request"), "http_request"),
|
||
(chat_graph("web_fetch"), "web_fetch"),
|
||
(chat_graph("python_exec"), "python_exec"),
|
||
(chat_graph("file_read"), "file_read"),
|
||
(chat_graph("memory_write"), "memory_write"),
|
||
(chat_graph("mcp:fs/read_file"), "mcp:fs/read_file"),
|
||
(chat_graph("my_custom_http"), "my_custom_http"),
|
||
(chat_graph("{{ vars.tool }}"), "模板"),
|
||
])
|
||
def test_graphs_a_user_may_not_run_in_chat(graph, said):
|
||
problem = chat_graph_problem(graph)
|
||
assert problem and problem.startswith("没有权限") and said in problem, problem
|
||
|
||
|
||
@pytest.mark.parametrize("graph", [
|
||
chat_graph(middle=node("ask", "agent", prompt="x", tools=["memory_recall"], memory_scope="builders")),
|
||
chat_graph(middle=node("kb", "retrieve", query="x", collection="hr")),
|
||
{**chat_graph("kb_search"), "defaults": {"collection": "hr"}},
|
||
# 调用工具节点的参数直接交给工具：在参数里挑知识库同样不许
|
||
chat_graph(middle=node("kb", "tool", tool="kb_search", args={"query": "x", "collection": "hr"})),
|
||
{**chat_graph("memory_recall"), "defaults": {"memory_scope": " "}},
|
||
# 节点上没写参数、运行时取图级 defaults 里的 args（和运行时同一个取值规则，engine/launch.node_setting）
|
||
{**chat_graph(middle=node("kb", "tool", tool="kb_search")), "defaults": {"args": {"query": "x", "collection": "hr"}}},
|
||
])
|
||
def test_users_may_not_pick_memory_or_knowledge_scope(graph):
|
||
"""记忆、知识库的内容不经遮罩：节点上、图级 defaults 里都不许指定范围，用管理员设好的默认值。"""
|
||
problem = chat_graph_problem(graph)
|
||
assert problem and "记忆范围和知识库" in problem
|
||
|
||
|
||
def test_tools_from_graph_defaults_count_too():
|
||
"""Agent 上没写工具时运行时用图级默认值里的：白名单同样按它查（engine/launch.tool_bindings）。"""
|
||
graph = chain(node("start", "input"), node("ask", "agent", prompt="x"), node("out", "output"))
|
||
graph["defaults"] = {"tools": ["web_search"]}
|
||
problem = chat_graph_problem(graph)
|
||
assert problem and "web_search" in problem
|
||
|
||
|
||
# ---------------------------------------------------------------- 接口
|
||
|
||
|
||
async def _masked_source(tmp_path) -> str:
|
||
path = tmp_path / f"{uuid.uuid4().hex[:8]}.db"
|
||
db = sqlite3.connect(path)
|
||
db.executescript("CREATE TABLE customers (id INTEGER PRIMARY KEY, name TEXT, phone TEXT, city TEXT);"
|
||
f"INSERT INTO customers VALUES (1, '张三', '{PHONE}', '杭州'), (2, '李四', '13900000002', '宁波');")
|
||
db.commit()
|
||
db.close()
|
||
name = f"crm_{uuid.uuid4().hex[:6]}"
|
||
async with SessionLocal() as session:
|
||
row = DataSource(name=name, kind="sqlite", database=str(path), readonly=True, enabled=True,
|
||
description="客户", options={"mask_columns": "phone"})
|
||
session.add(row)
|
||
await session.flush()
|
||
row.schema_cache = await introspect.introspect(row)
|
||
session.add(DataSource(name=f"off_{uuid.uuid4().hex[:6]}", kind="sqlite", database=str(path), enabled=False))
|
||
await session.commit()
|
||
return name
|
||
|
||
|
||
@pytest.fixture
|
||
async def world(tmp_path):
|
||
await auth_env.reset()
|
||
await run_manager.setup()
|
||
admin = await auth_env.turn_on_local()
|
||
await auth_env.make_user("uma", "user")
|
||
await auth_env.make_user("bob", "builder")
|
||
source = await _masked_source(tmp_path)
|
||
uma = await auth_env.login("uma")
|
||
bob = await auth_env.login("bob")
|
||
try:
|
||
yield {"admin": admin, "uma": uma, "bob": bob, "source": source}
|
||
finally:
|
||
await run_manager.shutdown()
|
||
async with SessionLocal() as s:
|
||
await s.execute(delete(Run).where(Run.started_by.in_(["uma", "bob"])))
|
||
await s.execute(delete(DataSource))
|
||
await s.commit()
|
||
await auth_env.reset()
|
||
|
||
|
||
async def _allow_chat(world) -> None:
|
||
r = await world["admin"].put("/api/auth/settings", json={"user_can_chat": True})
|
||
assert r.status_code == 200, r.text
|
||
|
||
|
||
async def _finish(run_id: str) -> Run:
|
||
for _ in range(400):
|
||
async with SessionLocal() as session:
|
||
row = await session.get(Run, run_id)
|
||
if row.status in ("succeeded", "failed", "interrupted", "cancelled"):
|
||
await run_manager.wait_idle(run_id)
|
||
return row
|
||
await asyncio.sleep(0.05)
|
||
raise AssertionError("运行没有结束")
|
||
|
||
|
||
def _script(monkeypatch, tool: str) -> None:
|
||
"""演示模型：第一轮查遮罩列，看到结果后作答。"""
|
||
from app.providers import mock_model
|
||
|
||
def _decide(self, messages):
|
||
if any(isinstance(m, ToolMessage) for m in messages):
|
||
return AIMessage(content="杭州有一位客户。")
|
||
return AIMessage(content="", tool_calls=[{"name": tool, "args": {"sql": "SELECT name, phone, city FROM customers"},
|
||
"id": "call_1"}])
|
||
|
||
monkeypatch.setattr(mock_model.MockChatModel, "_decide", _decide)
|
||
|
||
|
||
async def test_a_user_can_only_launch_whitelisted_chat_graphs(world):
|
||
await _allow_chat(world)
|
||
uma = world["uma"]
|
||
r = await uma.post("/api/runs", json={"origin": "chat", "graph": chat_graph("http_request"), "input": {}})
|
||
assert r.status_code == 403 and "http_request" in r.json()["detail"]
|
||
r = await uma.post("/api/runs", json={"origin": "chat", "graph": chat_graph(middle=node("ask", "code"))})
|
||
assert r.status_code == 403 and "沙箱代码" in r.json()["detail"]
|
||
r = await uma.post("/api/runs", json={"origin": "chat", "run_class": "formal", "workflow_id": "x"})
|
||
assert r.status_code == 403, "正式运行只经应用入口发起"
|
||
# 搭建者照旧不受这份白名单约束（发起前检查另说：这里只看不是 403）
|
||
r = await world["bob"].post("/api/runs", json={"origin": "chat", "graph": chat_graph(middle=node("ask", "code"))})
|
||
assert r.status_code != 403
|
||
|
||
|
||
async def test_a_users_chat_run_is_masked_end_to_end(world, monkeypatch):
|
||
"""真的执行一次：使用者发起、Agent 查了遮罩列。运行选项带着遮罩开关，事件、查询快照、实时流里都没有原值。"""
|
||
await _allow_chat(world)
|
||
tool = f"db_query__{world['source']}"
|
||
_script(monkeypatch, tool)
|
||
r = await world["uma"].post("/api/runs", json={"origin": "chat", "graph": chat_graph(tool),
|
||
"input": {"question": "杭州有几位客户"}})
|
||
assert r.status_code == 201, r.text
|
||
run = await _finish(r.json()["id"])
|
||
assert run.status == "succeeded", run.error
|
||
assert run.options.get(MASK_RESULTS_OPTION) is True and run.started_by == "uma" and run.origin_kind == "chat"
|
||
async with SessionLocal() as s:
|
||
events = list((await s.execute(select(RunEvent).where(RunEvent.run_id == run.id))).scalars())
|
||
snaps = list((await s.execute(select(Artifact).where(Artifact.run_id == run.id,
|
||
Artifact.kind == "query_snapshot"))).scalars())
|
||
dumped = json.dumps([e.data for e in events], ensure_ascii=False)
|
||
assert "tool.end" in {e.type for e in events}
|
||
assert PHONE not in dumped and "张三" in dumped, "遮罩列换掉，别的列照旧"
|
||
assert snaps, "查询照常落快照"
|
||
from app.core.artifact_store import load
|
||
|
||
for snap in snaps:
|
||
assert PHONE not in json.dumps(load(snap.id), ensure_ascii=False)
|
||
|
||
# 使用者能读自己这次运行的记录、事件，也能连实时流；实时流里同样没有原值
|
||
assert (await world["uma"].get(f"/api/runs/{run.id}/events")).status_code == 200
|
||
frames = await _ws(f"/api/runs/{run.id}/stream", _cookie(world["uma"]))
|
||
assert frames[0]["type"] == "websocket.accept"
|
||
texts = [f.get("text") or "" for f in frames if f["type"] == "websocket.send"]
|
||
assert texts and '"stream.end"' in texts[-1] and PHONE not in "".join(texts)
|
||
|
||
|
||
async def test_a_users_chat_run_cannot_write_a_writable_source(world, monkeypatch, tmp_path):
|
||
"""可写源 + Agent 自己写了 approval=never：写语句照样被守卫拒掉、不执行（只读不交给审批），表里的行一条没少。
|
||
搭建者发起的同一张图照旧能写——只读只给 user 角色的问数据运行。"""
|
||
from app.providers import mock_model
|
||
from app.tools.datasource import READ_ONLY_OPTION
|
||
|
||
path = tmp_path / "ledger.db"
|
||
db = sqlite3.connect(path)
|
||
db.executescript("CREATE TABLE ledger (id INTEGER PRIMARY KEY, amount INTEGER);"
|
||
"INSERT INTO ledger VALUES (1, 10), (2, 20);")
|
||
db.commit()
|
||
db.close()
|
||
name = f"ledger_{uuid.uuid4().hex[:6]}"
|
||
async with SessionLocal() as session:
|
||
row = DataSource(name=name, kind="sqlite", database=str(path), readonly=False, enabled=True)
|
||
session.add(row)
|
||
await session.flush()
|
||
row.schema_cache = await introspect.introspect(row)
|
||
await session.commit()
|
||
replies: list[str] = []
|
||
|
||
def _decide(self, messages):
|
||
if seen := [m for m in messages if isinstance(m, ToolMessage)]:
|
||
replies.append(str(seen[-1].content))
|
||
return AIMessage(content="好了。")
|
||
return AIMessage(content="", tool_calls=[{"name": f"db_query__{name}", "args": {"sql": "DELETE FROM ledger"},
|
||
"id": "call_1"}])
|
||
|
||
monkeypatch.setattr(mock_model.MockChatModel, "_decide", _decide)
|
||
graph = chat_graph(middle=node("ask", "agent", prompt="{{ input.question }}", tools=[f"db_query__{name}"],
|
||
max_steps=4, approval="never"))
|
||
await _allow_chat(world)
|
||
r = await world["uma"].post("/api/runs", json={"origin": "chat", "graph": graph, "input": {"question": "清空"}})
|
||
assert r.status_code == 201, r.text
|
||
run = await _finish(r.json()["id"])
|
||
assert run.status == "succeeded", run.error
|
||
assert run.options.get(READ_ONLY_OPTION) is True
|
||
assert replies and replies[-1].startswith("SQL 被拒绝"), replies
|
||
assert sqlite3.connect(path).execute("SELECT COUNT(*) FROM ledger").fetchone()[0] == 2
|
||
|
||
replies.clear()
|
||
r = await world["bob"].post("/api/runs", json={"origin": "chat", "graph": graph, "input": {"question": "清空"}})
|
||
assert r.status_code == 201, r.text
|
||
run = await _finish(r.json()["id"])
|
||
assert run.status == "succeeded", run.error
|
||
assert not (run.options or {}).get(READ_ONLY_OPTION)
|
||
assert sqlite3.connect(path).execute("SELECT COUNT(*) FROM ledger").fetchone()[0] == 0, "搭建者照旧能写"
|
||
|
||
|
||
async def _writable_masked_source(tmp_path) -> tuple[str, object]:
|
||
"""可写的客户库（手机号是遮罩列）：调用工具节点查它、想删它。"""
|
||
path = tmp_path / f"crm_rw_{uuid.uuid4().hex[:6]}.db"
|
||
db = sqlite3.connect(path)
|
||
db.executescript("CREATE TABLE customers (id INTEGER PRIMARY KEY, name TEXT, phone TEXT, city TEXT);"
|
||
f"INSERT INTO customers VALUES (1, '张三', '{PHONE}', '杭州'), (2, '李四', '13900000002', '宁波');")
|
||
db.commit()
|
||
db.close()
|
||
name = f"crm_rw_{uuid.uuid4().hex[:6]}"
|
||
async with SessionLocal() as session:
|
||
row = DataSource(name=name, kind="sqlite", database=str(path), readonly=False, enabled=True,
|
||
description="客户（可写）", options={"mask_columns": "phone"})
|
||
session.add(row)
|
||
await session.flush()
|
||
row.schema_cache = await introspect.introspect(row)
|
||
await session.commit()
|
||
return name, path
|
||
|
||
|
||
def _tool_chat(tool: str, sql: str) -> dict:
|
||
"""问数据当场搭的图：取数用调用工具节点（D1），不经 Agent、不调模型。approval=never 是发起人自己写的：只读不交给
|
||
审批，写成 never 也照样拒写。"""
|
||
return chain(node("start", "input", fields=[{"name": "question"}]),
|
||
node("q", "tool", tool=tool, args={"sql": sql}, fail_fast=False, approval="never"),
|
||
node("out", "output", fields=[{"name": "结果", "value": "{{ nodes.q }}"}]))
|
||
|
||
|
||
async def test_a_users_tool_node_is_masked_and_read_only(world, tmp_path):
|
||
"""调用工具节点这条路和 Agent 一样（D1）：使用者问数据的运行里，结果在工具边界就遮（事件、查询快照、运行产出里都没有
|
||
原值），可写源上的写语句被守卫拒掉、一行没少（节点上写了免审批也一样）；搭建者发起的同一张写图照旧能写。"""
|
||
from app.core.artifact_store import load
|
||
from app.tools.datasource import FORCED_READONLY, READ_ONLY_OPTION
|
||
|
||
await _allow_chat(world)
|
||
name, path = await _writable_masked_source(tmp_path)
|
||
tool = f"db_query__{name}"
|
||
|
||
r = await world["uma"].post("/api/runs", json={"origin": "chat", "input": {"question": "客户"},
|
||
"graph": _tool_chat(tool, "SELECT name, phone, city FROM customers")})
|
||
assert r.status_code == 201, r.text
|
||
run = await _finish(r.json()["id"])
|
||
assert run.status == "succeeded", run.error
|
||
assert run.options[MASK_RESULTS_OPTION] is True and run.options[READ_ONLY_OPTION] is True
|
||
async with SessionLocal() as s:
|
||
events = list((await s.execute(select(RunEvent).where(RunEvent.run_id == run.id))).scalars())
|
||
snaps = list((await s.execute(select(Artifact).where(Artifact.run_id == run.id,
|
||
Artifact.kind == "query_snapshot"))).scalars())
|
||
dumped = json.dumps([e.data for e in events], ensure_ascii=False) + json.dumps(run.output, ensure_ascii=False)
|
||
assert "tool.end" in {e.type for e in events}
|
||
assert PHONE not in dumped and MASKED in dumped and "张三" in dumped, "遮罩列换掉，别的列照旧"
|
||
assert snaps and all(PHONE not in json.dumps(load(snap.id), ensure_ascii=False) for snap in snaps)
|
||
|
||
delete_all = _tool_chat(tool, "DELETE FROM customers")
|
||
r = await world["uma"].post("/api/runs", json={"origin": "chat", "input": {"question": "清空"}, "graph": delete_all})
|
||
assert r.status_code == 201, r.text
|
||
run = await _finish(r.json()["id"])
|
||
async with SessionLocal() as s:
|
||
events = list((await s.execute(select(RunEvent).where(RunEvent.run_id == run.id))).scalars())
|
||
assert FORCED_READONLY[:12] in json.dumps([e.data for e in events], ensure_ascii=False), "写语句被守卫拒掉"
|
||
assert sqlite3.connect(path).execute("SELECT COUNT(*) FROM customers").fetchone()[0] == 2
|
||
|
||
r = await world["bob"].post("/api/runs", json={"origin": "chat", "input": {"question": "清空"}, "graph": delete_all})
|
||
assert r.status_code == 201, r.text
|
||
await _finish(r.json()["id"])
|
||
assert sqlite3.connect(path).execute("SELECT COUNT(*) FROM customers").fetchone()[0] == 0, "搭建者照旧能写"
|
||
|
||
|
||
async def test_a_tool_governed_as_write_is_not_run_in_a_read_only_run(world, tmp_path):
|
||
"""能力治理把数据源查询改成「写入」：使用者问数据里的调用工具节点不进动作流程（探索运行在那里会演练、交回假回执），
|
||
直接拒、说清是只读运行；表里的行一条没动。"""
|
||
from app.db.models import ActionRecord, Capability
|
||
|
||
await _allow_chat(world)
|
||
name, path = await _writable_masked_source(tmp_path)
|
||
tool = f"db_query__{name}"
|
||
async with SessionLocal() as s:
|
||
s.add(Capability(name=tool, system=name, title="客户库", effect_system="dynamic", effect="write"))
|
||
await s.commit()
|
||
try:
|
||
r = await world["uma"].post("/api/runs", json={"origin": "chat", "input": {"question": "清空"},
|
||
"graph": _tool_chat(tool, "DELETE FROM customers")})
|
||
assert r.status_code == 201, r.text
|
||
run = await _finish(r.json()["id"])
|
||
assert run.status == "failed" and "只能读数据" in (run.error or "") and tool in run.error, run.error
|
||
async with SessionLocal() as s:
|
||
assert not list((await s.execute(select(ActionRecord).where(ActionRecord.run_id == run.id))).scalars())
|
||
assert sqlite3.connect(path).execute("SELECT COUNT(*) FROM customers").fetchone()[0] == 2
|
||
finally:
|
||
async with SessionLocal() as s:
|
||
await s.execute(delete(Capability).where(Capability.name == tool))
|
||
await s.commit()
|
||
|
||
|
||
async def test_kb_search_keeps_to_the_runs_knowledge_base_when_the_scope_is_fixed(monkeypatch):
|
||
"""工具层的那一道（D1）：范围固定的运行里 kb_search 不认 collection 参数——直接写的、参数名写错被纠成 collection 的都拒，
|
||
不悄悄改查默认库；没写的照常查运行定下的那一个。范围不固定（搭建者的运行）照旧按参数查。"""
|
||
from app.memory import kb
|
||
from app.tools.registry import ToolArgsError, ToolContext, Versions, call_tool
|
||
|
||
searched: list[str] = []
|
||
|
||
async def _search(session, *, collection, query, limit, **_):
|
||
searched.append(collection)
|
||
return []
|
||
|
||
monkeypatch.setattr(kb, "search", _search)
|
||
fixed = ToolContext(collection="default", fixed_scope=True, data_versions=Versions.CURRENT)
|
||
for args in ({"query": "x", "collection": "hr"}, {"query": "x", "limit": 5, "coll": "hr"}):
|
||
with pytest.raises(ToolArgsError) as refused:
|
||
await call_tool("kb_search", args, fixed)
|
||
assert "只能检索管理员设好的知识库" in str(refused.value) and "hr" in str(refused.value)
|
||
assert searched == []
|
||
await call_tool("kb_search", {"query": "x"}, fixed)
|
||
await call_tool("kb_search", {"query": "x", "collection": "default"}, fixed)
|
||
await call_tool("kb_search", {"query": "x", "collection": "hr"},
|
||
ToolContext(collection="default", data_versions=Versions.CURRENT))
|
||
assert searched == ["default", "default", "hr"]
|
||
|
||
|
||
def _kb_search_spy(monkeypatch) -> list[str]:
|
||
from app.memory import kb
|
||
|
||
searched: list[str] = []
|
||
|
||
async def _search(session, *, collection, query, limit, **_):
|
||
searched.append(collection)
|
||
return [{"content": f"{collection} 库里的一段", "score": 1.0}]
|
||
|
||
monkeypatch.setattr(kb, "search", _search)
|
||
return searched
|
||
|
||
|
||
async def test_a_users_run_cannot_reach_another_knowledge_base_through_the_tool_layer(world, monkeypatch):
|
||
"""白名单认不出的两种写法，运行时在工具层拒（D1）：调用工具节点的参数名写错（coll 被纠成 collection）、Agent 在参数里
|
||
挑知识库。使用者的问数据只查得到运行定下的知识库；搭建者发起的同一张图照旧按参数查。"""
|
||
from app.providers import mock_model
|
||
|
||
await _allow_chat(world)
|
||
searched = _kb_search_spy(monkeypatch)
|
||
typo = chain(node("start", "input", fields=[{"name": "question"}]),
|
||
node("kb", "tool", tool="kb_search", args={"query": "x", "limit": 5, "coll": "hr"}),
|
||
node("out", "output", fields=[{"name": "结果", "value": "{{ nodes.kb }}"}]))
|
||
assert chat_graph_problem(typo) is None, "白名单只看字面写的参数名，这一道在工具层"
|
||
r = await world["uma"].post("/api/runs", json={"origin": "chat", "graph": typo, "input": {"question": "x"}})
|
||
assert r.status_code == 201, r.text
|
||
run = await _finish(r.json()["id"])
|
||
assert run.status == "failed" and "只能检索管理员设好的知识库" in (run.error or ""), run.error
|
||
assert "hr" not in searched
|
||
|
||
replies: list[str] = []
|
||
|
||
def _decide(self, messages):
|
||
if seen := [m for m in messages if isinstance(m, ToolMessage)]:
|
||
replies.append(str(seen[-1].content))
|
||
return AIMessage(content="好了。")
|
||
return AIMessage(content="", tool_calls=[{"name": "kb_search", "args": {"query": "薪酬", "collection": "hr"},
|
||
"id": "call_1"}])
|
||
|
||
monkeypatch.setattr(mock_model.MockChatModel, "_decide", _decide)
|
||
r = await world["uma"].post("/api/runs", json={"origin": "chat", "graph": chat_graph("kb_search"),
|
||
"input": {"question": "薪酬"}})
|
||
assert r.status_code == 201, r.text
|
||
run = await _finish(r.json()["id"])
|
||
assert run.status == "succeeded", run.error
|
||
assert replies and "只能检索管理员设好的知识库" in replies[-1], replies
|
||
assert "hr" not in searched
|
||
|
||
replies.clear()
|
||
r = await world["bob"].post("/api/runs", json={"origin": "chat", "graph": chat_graph("kb_search"),
|
||
"input": {"question": "薪酬"}})
|
||
assert r.status_code == 201, r.text
|
||
assert (await _finish(r.json()["id"])).status == "succeeded"
|
||
assert searched == ["hr"] and "hr 库里的一段" in replies[-1], "搭建者的运行照旧按参数查"
|
||
|
||
|
||
async def test_a_users_chat_run_keeps_to_its_own_memory_scope(world):
|
||
"""请求体里带记忆范围、知识库的拒；不带的，记忆是这位使用者自己的范围，知识库是设置里的默认值。"""
|
||
await _allow_chat(world)
|
||
uma = world["uma"]
|
||
graph = chat_graph("memory_recall")
|
||
for extra in ({"memory_scope": "default"}, {"collection": "hr"}):
|
||
r = await uma.post("/api/runs", json={"origin": "chat", "graph": graph, "input": {}, **extra})
|
||
assert r.status_code == 403 and "记忆范围和知识库" in r.json()["detail"]
|
||
r = await uma.post("/api/runs", json={"origin": "chat", "graph": graph, "input": {"question": "x"}})
|
||
assert r.status_code == 201, r.text
|
||
body = r.json()
|
||
assert body["memory_scope"] == "chat:uma" and body["collection"] == "default"
|
||
await _finish(body["id"])
|
||
|
||
|
||
# ---------------------------------------------------------------- 问数据的生成流
|
||
|
||
|
||
class _Scripted:
|
||
"""第一轮按节点手册的标准写法搭（tool 节点调 fetch_tool，缺省是查库），收到自查请求时改成 Agent 查库。记下 system 和
|
||
每一轮的请求。"""
|
||
|
||
def __init__(self, source: str, fetch_tool: str | None = None) -> None:
|
||
self.source = source
|
||
self.fetch_tool = fetch_tool or f"db_query__{source}"
|
||
self.systems: list[str] = []
|
||
self.humans: list[str] = []
|
||
|
||
def _lines(self, repair: bool) -> list[str]:
|
||
if repair:
|
||
ops = [{"op": "remove_node", "id": "fetch"},
|
||
{"op": "add_node", "node": {"id": "ask", "type": "agent", "label": "查库",
|
||
"config": {"prompt": "{{ input.question }}",
|
||
"tools": [f"db_query__{self.source}"]}}},
|
||
{"op": "add_edge", "edge": {"source": "start", "target": "ask"}},
|
||
{"op": "add_edge", "edge": {"source": "ask", "target": "out"}},
|
||
{"op": "done", "explanation": "改成 Agent 查库"}]
|
||
else:
|
||
ops = [{"op": "plan", "summary": "查客户"},
|
||
{"op": "add_node", "node": {"id": "start", "type": "input", "label": "输入",
|
||
"config": {"fields": [{"name": "question"}]}}},
|
||
{"op": "add_node", "node": {"id": "fetch", "type": "tool", "label": "查库",
|
||
"config": {"tool": self.fetch_tool,
|
||
"args": {"sql": "SELECT city FROM customers"}}}},
|
||
{"op": "add_node", "node": {"id": "out", "type": "output", "label": "成果",
|
||
"config": {"fields": [{"name": "回答", "value": "{{ last_message }}"}]}}},
|
||
{"op": "add_edge", "edge": {"source": "start", "target": "fetch"}},
|
||
{"op": "add_edge", "edge": {"source": "fetch", "target": "out"}},
|
||
{"op": "done", "explanation": "查好了", "run": True}]
|
||
return [json.dumps(o, ensure_ascii=False) for o in ops]
|
||
|
||
async def astream(self, messages):
|
||
from langchain_core.messages import AIMessageChunk
|
||
|
||
self.systems.append(next(text for role, text in messages if role == "system"))
|
||
human = next(text for role, text in messages if role == "human")
|
||
self.humans.append(human)
|
||
for line in self._lines("上线前自查" in human):
|
||
yield AIMessageChunk(content=line + "\n")
|
||
|
||
|
||
async def _generate(client, monkeypatch, model: _Scripted) -> list[dict]:
|
||
import app.api.copilot as copilot
|
||
|
||
async def _model(*_a, **_k):
|
||
return model, "mock-fast"
|
||
|
||
monkeypatch.setattr(copilot, "get_chat_model", _model)
|
||
events: list[dict] = []
|
||
async with client.stream("POST", "/api/copilot/generate-stream", json={"instruction": "杭州有几位客户"}) as r:
|
||
assert r.status_code == 200
|
||
async for line in r.aiter_lines():
|
||
if line.startswith("data: "):
|
||
events.append(json.loads(line[6:]))
|
||
return events
|
||
|
||
|
||
async def test_generation_for_a_user_keeps_to_the_chat_policy(world, monkeypatch):
|
||
"""使用者问数据时，提示词先说清只能用数据源查询和只读的内置工具；tool 节点查库照标准写法搭，一次就过、自动运行（D1）。
|
||
tool 节点绑了不许的工具（http_request）按 chat_policy 交回去改，改好才交付。搭建者照旧按节点手册来。"""
|
||
await _allow_chat(world)
|
||
model = _Scripted(world["source"])
|
||
events = await _generate(world["uma"], monkeypatch, model)
|
||
assert "这位用户是「使用者」" in model.systems[0] and "- http_request" not in model.systems[0]
|
||
assert "tool 节点和 agent 都只能用" in model.systems[0]
|
||
assert len(model.humans) == 1, "tool 节点查库不算问题"
|
||
final = next(e for e in events if e.get("op") == "final")
|
||
assert final["autorun"] is True and "tool" in {n["type"] for n in final["graph"]["nodes"]}
|
||
assert chat_graph_problem(final["graph"]) is None
|
||
|
||
model = _Scripted(world["source"], fetch_tool="http_request")
|
||
events = await _generate(world["uma"], monkeypatch, model)
|
||
assert len(model.humans) == 2 and "http_request" in model.humans[1], "绑了不许的工具按 chat_policy 交回去改"
|
||
check = [e for e in events if e.get("op") == "check"][-1]
|
||
final = next(e for e in events if e.get("op") == "final")
|
||
assert check["status"] == "passed" and final["autorun"] is True
|
||
assert {n["type"] for n in final["graph"]["nodes"]} == {"input", "agent", "output"}
|
||
assert chat_graph_problem(final["graph"]) is None
|
||
|
||
model = _Scripted(world["source"])
|
||
events = await _generate(world["bob"], monkeypatch, model)
|
||
assert "这位用户是「使用者」" not in model.systems[0] and len(model.humans) == 1
|
||
final = next(e for e in events if e.get("op") == "final")
|
||
assert "tool" in {n["type"] for n in final["graph"]["nodes"]}
|
||
|
||
|
||
async def test_generation_for_a_user_reports_what_it_could_not_fix(world, monkeypatch):
|
||
"""改不好的照实交付：自查失败、final 里列出 chat.not_allowed，不自动运行。"""
|
||
await _allow_chat(world)
|
||
model = _Scripted(world["source"], fetch_tool="http_request")
|
||
model._lines = lambda repair: _Scripted(world["source"], fetch_tool="http_request")._lines(False) if not repair \
|
||
else [json.dumps({"op": "done", "explanation": "没改"}, ensure_ascii=False)]
|
||
events = await _generate(world["uma"], monkeypatch, model)
|
||
check = [e for e in events if e.get("op") == "check"][-1]
|
||
final = next(e for e in events if e.get("op") == "final")
|
||
assert check["status"] == "failed" and any(i["code"] == "chat.not_allowed" for i in check["issues"])
|
||
assert final["autorun"] is False and any(i["code"] == "chat.not_allowed" for i in final["issues"])
|
||
|
||
|
||
# ---------------------------------------------------------------- WebSocket
|
||
|
||
|
||
def _cookie(client) -> str:
|
||
return f"agentlab_session={client.cookies.get('agentlab_session')}"
|
||
|
||
|
||
async def _ws(path: str, cookie: str | None) -> list[dict]:
|
||
"""把一次 WebSocket 握手原样交给整个应用（中间件 + 路由级依赖 + 路由），记下发出去的帧。"""
|
||
from app.main import app
|
||
|
||
inbox: asyncio.Queue = asyncio.Queue()
|
||
await inbox.put({"type": "websocket.connect"})
|
||
sent: list[dict] = []
|
||
|
||
async def receive():
|
||
return await inbox.get()
|
||
|
||
async def send(message):
|
||
sent.append(message)
|
||
if message["type"] == "websocket.close":
|
||
await inbox.put({"type": "websocket.disconnect", "code": 1000})
|
||
|
||
headers = [(b"host", b"test")] + ([(b"cookie", cookie.encode())] if cookie else [])
|
||
scope = {"type": "websocket", "asgi": {"version": "3.0"}, "scheme": "ws", "http_version": "1.1", "path": path,
|
||
"raw_path": path.encode(), "query_string": b"", "root_path": "", "headers": headers,
|
||
"client": ("127.0.0.1", 50000), "server": ("test", 80), "subprotocols": [], "state": {}}
|
||
await asyncio.wait_for(app(scope, receive, send), 10)
|
||
return sent
|
||
|
||
|
||
async def _settled_run(**fields) -> str:
|
||
async with SessionLocal() as s:
|
||
run = Run(workflow_name="问数据", graph={}, input={}, status="succeeded", run_class="exploratory", **fields)
|
||
s.add(run)
|
||
await s.flush()
|
||
s.add(RunEvent(run_id=run.id, seq=1, type="run.started", ts=1.0, data={}))
|
||
await s.commit()
|
||
return run.id
|
||
|
||
|
||
def _refused(frames: list[dict]) -> bool:
|
||
return [f["type"] for f in frames] == ["websocket.close"] and frames[0].get("code") == 1008
|
||
|
||
|
||
async def test_websocket_is_judged_per_run_at_the_handshake(world):
|
||
mine = await _settled_run(started_by="uma", origin_kind="chat", options={MASK_RESULTS_OPTION: True})
|
||
others = await _settled_run(started_by="bob", origin_kind="chat", options={MASK_RESULTS_OPTION: True})
|
||
raw = await _settled_run(started_by="uma", origin_kind="chat") # 关登录时自报同名留下的，没在边界遮过
|
||
uma, bob = _cookie(world["uma"]), _cookie(world["bob"])
|
||
|
||
assert _refused(await _ws(f"/api/runs/{mine}/stream", None)), "没登录的握手前关掉"
|
||
assert _refused(await _ws(f"/api/runs/{mine}/stream", uma)), "没开问数据时使用者连不了运行"
|
||
await _allow_chat(world)
|
||
frames = await _ws(f"/api/runs/{mine}/stream", uma)
|
||
assert frames[0]["type"] == "websocket.accept" and '"stream.end"' in frames[-2]["text"]
|
||
assert _refused(await _ws(f"/api/runs/{others}/stream", uma)), "别人的运行"
|
||
assert _refused(await _ws(f"/api/runs/{raw}/stream", uma)), "查询结果没在工具边界遮过的运行"
|
||
assert _refused(await _ws("/api/runs/nope/stream", uma))
|
||
assert (await _ws(f"/api/runs/{raw}/stream", bob))[0]["type"] == "websocket.accept", "搭建者照旧"
|
||
|
||
|
||
# ---------------------------------------------------------------- 数据源名单
|
||
|
||
|
||
async def test_source_names_for_the_chat_page(world):
|
||
uma, bob = world["uma"], world["bob"]
|
||
assert (await uma.get("/api/datasources/names")).status_code == 403, "没开问数据时不给使用者"
|
||
assert (await uma.get("/api/datasources")).status_code == 403
|
||
await _allow_chat(world)
|
||
body = (await uma.get("/api/datasources/names")).json()
|
||
# 同一个测试进程里别的测试留下的数据源也会在名单里：只核对这里建的两个（一个启用、一个停用）
|
||
names = [i["name"] for i in body["items"]]
|
||
assert world["source"] in names and not [n for n in names if n.startswith("off_")], "只列启用着的"
|
||
item = next(i for i in body["items"] if i["name"] == world["source"])
|
||
assert set(item) == {"id", "name", "description", "kind"}, "连接信息不在名单里"
|
||
assert item["description"] == "客户" and item["kind"] == "sqlite"
|
||
assert (await uma.get("/api/datasources")).status_code == 403, "完整的数据源列表照旧只给搭建者"
|
||
assert (await bob.get("/api/datasources/names")).status_code == 200
|
||
|
||
|
||
async def test_source_names_with_login_off():
|
||
await auth_env.reset()
|
||
async with auth_env.client() as c:
|
||
r = await c.get("/api/datasources/names")
|
||
assert r.status_code == 200 and "items" in r.json()
|