287 lines
12 KiB
Python
287 lines
12 KiB
Python
# ============================================================
|
||
# 黄金测试：四层上下文策略（plan.md §4.8 / 矩阵 64 行）
|
||
# 四层组装可检查 / 预算溢出可预测拒绝 / 摘要可溯源 / pin 生效
|
||
# ============================================================
|
||
from __future__ import annotations
|
||
|
||
import pytest
|
||
|
||
from server.agent_core.context_policy import (
|
||
ContextAssembly,
|
||
ContextBudget,
|
||
ContextBudgetExceeded,
|
||
assemble_context,
|
||
estimate_tokens,
|
||
get_policy,
|
||
pin_message,
|
||
scroll_summary,
|
||
unpin_message,
|
||
)
|
||
|
||
|
||
def _store_data() -> dict:
|
||
"""工作区快照：项目 + 会话 + 带 id 的消息。"""
|
||
return {
|
||
"systemPrompt": "你是 APS 排产智能体：提议权，写操作必须经门禁。",
|
||
"projects": [
|
||
{
|
||
"id": "p1", "name": "青岛Q3排产", "scopeLabel": "青岛工厂/总装车间",
|
||
"sharedContext": ["本项目用 V3 参数版本", "以 V20260701-005 为基准版本"],
|
||
"archived": False,
|
||
},
|
||
],
|
||
"sessions": [
|
||
{"id": "s1", "projectId": "p1", "title": "本周主排产", "status": "running"},
|
||
],
|
||
"messages": {
|
||
"s1": [
|
||
{"id": "m0", "role": "user", "text": "验收标准：插单必须满足交期优先"},
|
||
{"id": "m1", "role": "agent", "text": "好的，按交期优先试排"},
|
||
{"id": "m2", "role": "user", "text": "再看看产能是否足够"},
|
||
{"id": "m3", "role": "agent", "text": "产能利用率 92%，足够"},
|
||
],
|
||
},
|
||
}
|
||
|
||
|
||
def test_assemble_four_layers_inspectable():
|
||
sd = _store_data()
|
||
sd["ragResults"] = [{"assetId": "kb_1", "title": "换线标准SOP", "version": "2", "content": "换线前 30 分钟通知"}]
|
||
assembly = assemble_context(sd, "s1", "换线需要多久", ContextBudget())
|
||
assert isinstance(assembly, ContextAssembly)
|
||
# 四层齐全且可检查
|
||
assert list(assembly.layers.keys()) == ["fixed", "project", "session", "retrieval"]
|
||
assert "提议权" in assembly.layer("fixed")
|
||
assert "青岛Q3排产" in assembly.layer("project")
|
||
assert "V3 参数版本" in assembly.layer("project")
|
||
assert "验收标准" in assembly.layer("session")
|
||
assert "kb_1" in assembly.layer("retrieval")
|
||
# 整体文本按固定/项目/会话/检索拼接，可注入 prompt
|
||
text = assembly.text()
|
||
assert "[fixed]" in text and "[project]" in text
|
||
assert "[session]" in text and "[retrieval]" in text
|
||
# 用量与预算一一对应（可检查）
|
||
assert assembly.usage == {k: estimate_tokens(assembly.layers[k]) for k in assembly.layers}
|
||
assert assembly.budgets == {
|
||
"fixed": 1500, "project": 3000, "session": 6000, "retrieval": 4000,
|
||
}
|
||
|
||
|
||
def test_budget_overflow_predictable_rejection():
|
||
sd = _store_data()
|
||
sd["messages"]["s1"].append({"id": "m10", "role": "user", "text": "排产" * 120})
|
||
budget = ContextBudget(session=50)
|
||
|
||
def expect() -> ContextBudgetExceeded:
|
||
with pytest.raises(ContextBudgetExceeded) as exc_info:
|
||
assemble_context(sd, "s1", None, budget)
|
||
return exc_info.value
|
||
|
||
first = expect()
|
||
assert first.layer == "session"
|
||
assert first.required > first.budget == 50
|
||
# 可预测：同一输入重复触发得到完全一致的层/用量/限额
|
||
second = expect()
|
||
assert (second.layer, second.required, second.budget) == (first.layer, first.required, first.budget)
|
||
|
||
# 检索层独立限额：RAG 命中过长 -> 拒绝在 retrieval 层
|
||
sd2 = _store_data()
|
||
sd2["ragResults"] = [{"assetId": "kb_1", "title": "长文", "content": "换线" * 200}]
|
||
with pytest.raises(ContextBudgetExceeded) as exc_info:
|
||
assemble_context(sd2, "s1", None, ContextBudget(retrieval=10))
|
||
assert exc_info.value.layer == "retrieval"
|
||
# 提升预算后组装成功（预算控制可逆、可预测）
|
||
ok = assemble_context(sd2, "s1", None, ContextBudget(retrieval=10000))
|
||
assert "kb_1" in ok.layer("retrieval")
|
||
|
||
|
||
def test_context_budget_totals_and_token_estimate():
|
||
budget = ContextBudget(fixed=100, session=200)
|
||
assert budget.layer_limit("fixed") == 100
|
||
assert budget.layer_limit("session") == 200
|
||
assert budget.total() == 100 + 3000 + 200 + 4000
|
||
parsed = ContextBudget.from_dict({"fixed": "10", "project": 20, "bogus": 999})
|
||
assert (parsed.fixed, parsed.project) == (10, 20)
|
||
with pytest.raises(ValueError):
|
||
budget.layer_limit("bogus")
|
||
# 确定性 token 估算：CJK 1 token/字，其余 4 字符 1 token
|
||
assert estimate_tokens("你好ab") == 2 + 1
|
||
assert estimate_tokens("") == 0
|
||
|
||
|
||
def test_scroll_summary_traceable_to_source_messages():
|
||
sd = _store_data()
|
||
s1 = scroll_summary(
|
||
sd, "s1", "验收标准：插单必须满足交期优先；已决策采用方案B",
|
||
source_message_ids=["m0", "m1"],
|
||
decision_points=["采用方案B"],
|
||
open_items=["未决：产能上限"],
|
||
key_numbers=["利用率 92%"],
|
||
by="agent",
|
||
)
|
||
assert s1["sourceMessageIds"] == ["m0", "m1"]
|
||
assert s1["decisionPoints"] == ["采用方案B"]
|
||
|
||
# 滚动：新摘要压缩更早对话，supersedes 链指向旧摘要，旧摘要入历史
|
||
s2 = scroll_summary(sd, "s1", "插单已评估完毕，等待确认",
|
||
source_message_ids=["m2", "m3"])
|
||
assert s2["supersedes"] == s1["summaryId"]
|
||
assert s2["generation"] == 2
|
||
policy = get_policy(sd, "s1")
|
||
assert policy["summary"]["summaryId"] == s2["summaryId"]
|
||
assert len(policy["summaries"]) == 1
|
||
assert policy["summaries"][0]["sourceMessageIds"] == ["m0", "m1"]
|
||
|
||
# 组装后会话层注入摘要，且摘要可溯源到原消息 id
|
||
assembly = assemble_context(sd, "s1", None, ContextBudget())
|
||
session_text = assembly.layer("session")
|
||
assert "插单已评估完毕" in session_text
|
||
assert "来源消息：m2、m3" in session_text
|
||
assert assembly.summary is not None and assembly.summary["summaryId"] == s2["summaryId"]
|
||
# 空摘要拒绝（fail closed）
|
||
with pytest.raises(ValueError):
|
||
scroll_summary(sd, "s1", " ")
|
||
|
||
|
||
def test_pin_message_survives_window_and_unpin():
|
||
sd = _store_data()
|
||
sd["messages"]["s1"] = [
|
||
{"id": f"m{i}", "role": "user" if i % 2 == 0 else "agent", "text": f"消息{i}"}
|
||
for i in range(20)
|
||
]
|
||
pinned = pin_message(sd, "s1", "m0", note="验收标准不可淘汰")
|
||
assert pinned["messageId"] == "m0"
|
||
# 幂等：重复钉住不产生重复项
|
||
pin_message(sd, "s1", "m0", note="验收标准不可淘汰")
|
||
assert len(get_policy(sd, "s1")["pinned"]) == 1
|
||
|
||
assembly = assemble_context(sd, "s1", None, ContextBudget(), window=5)
|
||
session_text = assembly.layer("session")
|
||
# pin 生效：窗口外的最早消息仍注入
|
||
assert "消息0" in session_text
|
||
assert "（备注：验收标准不可淘汰）" in session_text
|
||
# 最近 N 轮原文在窗口内
|
||
assert "消息15" in session_text
|
||
# 窗口外未钉住的消息被淘汰
|
||
assert "消息10" not in session_text
|
||
assert [p["messageId"] for p in assembly.pinned] == ["m0"]
|
||
|
||
# 解除钉住：恢复窗口淘汰
|
||
assert unpin_message(sd, "s1", "m0") is True
|
||
after = assemble_context(sd, "s1", None, ContextBudget(), window=5)
|
||
assert "消息0" not in after.layer("session")
|
||
assert unpin_message(sd, "s1", "m0") is False
|
||
|
||
# 钉住不存在的消息：显式失败
|
||
with pytest.raises(ValueError):
|
||
pin_message(sd, "s1", "m_no_such")
|
||
|
||
|
||
def test_resolved_refs_render_into_retrieval_layer():
|
||
"""方向 A 内联：会话引用解析结果按需注入检索层（L-检索）。"""
|
||
sd = _store_data()
|
||
sd["resolvedRefs"] = [
|
||
{"kind": "session", "target_id": "sess_a", "snapshot_at": "2026-08-02 10:30"},
|
||
{"kind": "version", "target_id": "V20260716-003", "snapshot_at": "2026-08-02 10:30"},
|
||
]
|
||
assembly = assemble_context(sd, "s1", None, ContextBudget())
|
||
retrieval_text = assembly.layer("retrieval")
|
||
assert "[引用 session:sess_a @2026-08-02 10:30]" in retrieval_text
|
||
assert "[引用 version:V20260716-003 @2026-08-02 10:30]" in retrieval_text
|
||
# ============================================================
|
||
# 越权上下文测试（矩阵 64 行剩余项 · 方向 L）：
|
||
# 访问其他项目/会话的上下文 -> fail closed；预算超限拒绝可预测。
|
||
# ============================================================
|
||
|
||
def test_unauthorized_other_project_session_fails_closed():
|
||
"""越权上下文：会话归属项目不在可见项目列表 -> 显式拒绝（fail closed）。
|
||
|
||
可见快照由租户/成员过滤生成（ProjectStore.snapshot 只返回可见会话），
|
||
因此越权会话要么在快照中缺失（session not found），要么其项目不可访问
|
||
（project not accessible）——两条路径都绝不允许组装出任何一层上下文。
|
||
"""
|
||
sd = _store_data()
|
||
sd["projects"].append({
|
||
"id": "p2", "name": "别厂项目", "scopeLabel": "外部工厂/总装",
|
||
"sharedContext": ["别厂参数版本"], "archived": False,
|
||
})
|
||
sd["sessions"].append({"id": "s2", "projectId": "p2", "title": "别厂会话", "status": "running"})
|
||
sd["messages"]["s2"] = [
|
||
{"id": "x0", "role": "user", "text": "机密：别厂排产方案不可外泄"},
|
||
{"id": "x1", "role": "agent", "text": "别厂关键路径已锁定"},
|
||
]
|
||
|
||
# 路径 1：越权会话在可见快照中缺失（等同无权访问）——先于预算检查拒绝
|
||
visible = {
|
||
**sd,
|
||
"projects": sd["projects"][:1],
|
||
"sessions": sd["sessions"][:1],
|
||
"messages": {"s1": sd["messages"]["s1"]},
|
||
}
|
||
with pytest.raises(ValueError) as exc_info:
|
||
assemble_context(visible, "s2", None, ContextBudget(session=10 ** 9))
|
||
assert "session not found" in str(exc_info.value)
|
||
|
||
# 路径 2：会话在快照中但其项目不在可见项目列表——同样显式拒绝
|
||
half_visible = {
|
||
**sd,
|
||
"projects": sd["projects"][:1], # p2 对当前用户不可访问
|
||
}
|
||
with pytest.raises(ValueError) as exc_info2:
|
||
assemble_context(half_visible, "s2", None, ContextBudget())
|
||
assert "project not accessible" in str(exc_info2.value)
|
||
|
||
# 拒绝是可预测的：同一输入重复触发，异常信息一致
|
||
with pytest.raises(ValueError) as exc_info3:
|
||
assemble_context(half_visible, "s2", None, ContextBudget())
|
||
assert str(exc_info3.value) == str(exc_info2.value)
|
||
|
||
|
||
def test_authorized_session_never_leaks_other_session_context():
|
||
"""隔离：授权会话 s1 的组装绝不混入 s2 的消息/摘要（越权内容零泄漏）。"""
|
||
sd = _store_data()
|
||
sd["projects"].append({
|
||
"id": "p2", "name": "别厂项目", "scopeLabel": "外部工厂",
|
||
"sharedContext": [], "archived": False,
|
||
})
|
||
sd["sessions"].append({"id": "s2", "projectId": "p2", "title": "别厂会话", "status": "running"})
|
||
sd["messages"]["s2"] = [
|
||
{"id": "x0", "role": "user", "text": "机密：别厂排产方案不可外泄"},
|
||
{"id": "x1", "role": "agent", "text": "别厂关键路径已锁定"},
|
||
]
|
||
scroll_summary(sd, "s2", "别厂滚动摘要：方案不可外泄", source_message_ids=["x0", "x1"])
|
||
|
||
assembly = assemble_context(sd, "s1", None, ContextBudget())
|
||
blob = assembly.text()
|
||
assert "机密" not in blob
|
||
assert "别厂" not in blob
|
||
assert "x0" not in blob and "x1" not in blob
|
||
# 会话层只带 s1 自己的策略/消息
|
||
assert assembly.summary is None or "别厂" not in assembly.summary.get("text", "")
|
||
assert get_policy(sd, "s1") == {}
|
||
|
||
|
||
def test_budget_overflow_predictable_rejection_diagnostics():
|
||
"""预算超限拒绝可预测：异常携带 layer/required/budget，同一输入结果一致。
|
||
|
||
补齐「预算超限拒绝可预测」的验收：不仅拒绝，还给出可检查的诊断字段，
|
||
且重复触发得到完全相同的层/用量/限额（确定性）。
|
||
"""
|
||
sd = _store_data()
|
||
sd["messages"]["s1"].append({"id": "m20", "role": "user", "text": "排产" * 200})
|
||
budget = ContextBudget(session=50)
|
||
|
||
with pytest.raises(ContextBudgetExceeded) as exc_info:
|
||
assemble_context(sd, "s1", None, budget)
|
||
first = exc_info.value
|
||
assert (first.layer, first.budget) == ("session", 50)
|
||
assert first.required > 50
|
||
assert str(first) == f"context layer 'session' needs {first.required} tokens, budget 50"
|
||
|
||
with pytest.raises(ContextBudgetExceeded) as exc_info2:
|
||
assemble_context(sd, "s1", None, budget)
|
||
second = exc_info2.value
|
||
assert (second.layer, second.required, second.budget) == (
|
||
first.layer, first.required, first.budget,
|
||
)
|