506 lines
20 KiB
Python
506 lines
20 KiB
Python
#!/usr/bin/env python3
|
||
"""muse_role 离线单测：provider-neutral 角色执行与回执合同。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import unittest
|
||
from decimal import Decimal
|
||
from unittest import mock
|
||
|
||
import muse_llm
|
||
|
||
from muse_role import (
|
||
CONTENT_POLICY_ALIAS,
|
||
CONTENT_POLICY_VERSION,
|
||
FIXED_OPUS_MODEL_ID,
|
||
FIXED_OPUS_POLICY_ALIAS,
|
||
FIXED_OPUS_POLICY_VERSION,
|
||
RoleExecutionProfile,
|
||
RoleRuntimeError,
|
||
build_dispatch_system_prompt,
|
||
canonical_json,
|
||
compose_role_system_prompt,
|
||
contains_path_traversal,
|
||
run_role,
|
||
sha256_json,
|
||
sha256_text,
|
||
)
|
||
|
||
WRITER_SCHEMA = {
|
||
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
||
"type": "object",
|
||
"additionalProperties": False,
|
||
"required": ["candidateBody"],
|
||
"properties": {"candidateBody": {"type": "string", "minLength": 1}},
|
||
}
|
||
|
||
# 既有回执合同的 camelCase 字段集合，记录运行证据 与看板按此消费。
|
||
RECEIPT_KEYS = {
|
||
"adapterRole",
|
||
"invocationId",
|
||
"executionProfileSha256",
|
||
"requestedModelId",
|
||
"actualModelId",
|
||
"modelMatch",
|
||
"effort",
|
||
"maxBudgetUsdPerCall",
|
||
"totalCostUsd",
|
||
"usage",
|
||
"modelUsage",
|
||
"stopReason",
|
||
"terminalReason",
|
||
"isError",
|
||
"apiErrorStatus",
|
||
"exitCode",
|
||
"durationMs",
|
||
"inputSha256",
|
||
"structuredOutputSha256",
|
||
"jsonSchemaSha256",
|
||
}
|
||
|
||
|
||
def _profile(schema=None, prompt="你是测试角色。", **overrides):
|
||
schema = schema or WRITER_SCHEMA
|
||
fields = dict(
|
||
profile_version="role-writer-v4",
|
||
adapter_role="writer",
|
||
model_alias=FIXED_OPUS_POLICY_ALIAS,
|
||
model_policy_version=FIXED_OPUS_POLICY_VERSION,
|
||
resolved_model_id=FIXED_OPUS_MODEL_ID,
|
||
max_budget_usd_per_call=Decimal("1.00"),
|
||
timeout_seconds=60.0,
|
||
max_context_chars=100_000,
|
||
json_schema_id="writer-draft-v2",
|
||
json_schema=schema,
|
||
json_schema_sha256=sha256_json(schema),
|
||
system_prompt_id="writer-system-prompt-v2",
|
||
system_prompt=prompt,
|
||
system_prompt_sha256=sha256_text(prompt),
|
||
)
|
||
fields.update(overrides)
|
||
return RoleExecutionProfile(**fields)
|
||
|
||
|
||
def _fake_chat(payload, *, actual_model=muse_llm.FIXED_OPUS_CANONICAL_MODEL):
|
||
"""构造受控的 governed_chat 假实现。payload 是返回内容或 None（模拟全链耗尽）。"""
|
||
|
||
calls = []
|
||
|
||
def chat(prompt, model=None, system=None, temperature=0.2, top_p=None, timeout=900,
|
||
run_id=None, caller=None, persist_call=None, **_kwargs):
|
||
calls.append({"prompt": prompt, "model": model, "system": system,
|
||
"timeout": timeout, "run_id": run_id, "caller": caller})
|
||
if payload is None:
|
||
return None, None, None
|
||
return payload, {"prompt_tokens": 10, "completion_tokens": 20}, actual_model
|
||
|
||
chat.calls = calls
|
||
return chat
|
||
|
||
|
||
class _FakeStreamResponse:
|
||
def __init__(self, events, *, status_code=200, text=""):
|
||
self.status_code = status_code
|
||
self.text = text
|
||
self._events = events
|
||
self.iter_lines_kwargs = None
|
||
self.closed = False
|
||
|
||
def raise_for_status(self):
|
||
if self.status_code >= 400:
|
||
raise muse_llm.requests.HTTPError(f"HTTP {self.status_code}")
|
||
|
||
def iter_lines(self, **kwargs):
|
||
self.iter_lines_kwargs = kwargs
|
||
for event in self._events:
|
||
if event == "[DONE]":
|
||
yield b"data: [DONE]"
|
||
else:
|
||
yield b"data: " + json.dumps(
|
||
event, ensure_ascii=False, separators=(",", ":")
|
||
).encode("utf-8")
|
||
yield b""
|
||
|
||
def close(self):
|
||
self.closed = True
|
||
|
||
|
||
class FixedOpusTransportTests(unittest.TestCase):
|
||
def test_fixed_opus_streams_utf8_without_splitting_c1_bytes(self):
|
||
response = _FakeStreamResponse([
|
||
{
|
||
"type": "message_start",
|
||
"message": {
|
||
"model": muse_llm.FIXED_OPUS_CANONICAL_MODEL,
|
||
"usage": {"input_tokens": 10, "output_tokens": 0},
|
||
},
|
||
},
|
||
{
|
||
"type": "content_block_start",
|
||
"index": 0,
|
||
"content_block": {"type": "text", "text": ""},
|
||
},
|
||
{
|
||
"type": "content_block_delta",
|
||
"index": 0,
|
||
"delta": {
|
||
"type": "text_delta",
|
||
"text": '{"candidateBody":"内文"}',
|
||
},
|
||
},
|
||
{"type": "content_block_stop", "index": 0},
|
||
{
|
||
"type": "message_delta",
|
||
"delta": {"stop_reason": "end_turn"},
|
||
"usage": {"output_tokens": 5},
|
||
},
|
||
{"type": "message_stop"},
|
||
"[DONE]",
|
||
])
|
||
session = mock.Mock()
|
||
session.post.return_value = response
|
||
events = []
|
||
env = {
|
||
"MUSE_ROLE_OPUS_BASE_URL": "https://role.example",
|
||
"MUSE_ROLE_OPUS_AUTH_TOKEN": "secret",
|
||
"MUSE_ROLE_OPUS_MODEL": FIXED_OPUS_MODEL_ID,
|
||
}
|
||
with (
|
||
mock.patch.dict(os.environ, env, clear=False),
|
||
mock.patch.object(muse_llm.requests, "Session", return_value=session),
|
||
):
|
||
content, usage, actual = muse_llm.chat_fixed_opus(
|
||
"{}",
|
||
model=FIXED_OPUS_POLICY_ALIAS,
|
||
system="system",
|
||
max_tokens=123,
|
||
retries=0,
|
||
run_id="probe-1",
|
||
caller="writer",
|
||
persist_call=events.append,
|
||
)
|
||
self.assertFalse(session.trust_env)
|
||
request = session.post.call_args
|
||
self.assertEqual(request.args[0], "https://role.example/v1/messages")
|
||
self.assertEqual(request.kwargs["json"]["model"], FIXED_OPUS_MODEL_ID)
|
||
self.assertEqual(request.kwargs["json"]["max_tokens"], 123)
|
||
self.assertIs(request.kwargs["json"]["stream"], True)
|
||
self.assertIs(request.kwargs["stream"], True)
|
||
self.assertEqual(
|
||
response.iter_lines_kwargs,
|
||
{"decode_unicode": False, "delimiter": b"\n"},
|
||
)
|
||
self.assertTrue(response.closed)
|
||
self.assertEqual(content, '{"candidateBody":"内文"}')
|
||
self.assertEqual(usage["output_tokens"], 5)
|
||
self.assertEqual(actual, muse_llm.FIXED_OPUS_CANONICAL_MODEL)
|
||
self.assertEqual(events[0]["requested_model_id"], FIXED_OPUS_POLICY_ALIAS)
|
||
self.assertEqual(events[0]["actual_model_id"], muse_llm.FIXED_OPUS_CANONICAL_MODEL)
|
||
self.assertTrue(events[0]["model_match"])
|
||
|
||
def test_fixed_opus_rejects_stream_that_ends_without_message_stop(self):
|
||
response = _FakeStreamResponse([
|
||
{
|
||
"type": "message_start",
|
||
"message": {
|
||
"model": muse_llm.FIXED_OPUS_CANONICAL_MODEL,
|
||
"usage": {"input_tokens": 10, "output_tokens": 0},
|
||
},
|
||
},
|
||
{
|
||
"type": "content_block_start",
|
||
"index": 0,
|
||
"content_block": {"type": "text", "text": ""},
|
||
},
|
||
])
|
||
session = mock.Mock()
|
||
session.post.return_value = response
|
||
env = {
|
||
"MUSE_ROLE_OPUS_BASE_URL": "https://role.example",
|
||
"MUSE_ROLE_OPUS_AUTH_TOKEN": "secret",
|
||
}
|
||
with (
|
||
mock.patch.dict(os.environ, env, clear=False),
|
||
mock.patch.object(muse_llm.requests, "Session", return_value=session),
|
||
):
|
||
with self.assertRaisesRegex(RuntimeError, "message_stop"):
|
||
muse_llm.chat_fixed_opus(
|
||
"{}",
|
||
model=FIXED_OPUS_POLICY_ALIAS,
|
||
retries=0,
|
||
)
|
||
self.assertTrue(response.closed)
|
||
|
||
def test_fixed_opus_retry_respects_single_total_deadline(self):
|
||
response = mock.Mock(status_code=429, text="rate limited")
|
||
session = mock.Mock()
|
||
session.post.return_value = response
|
||
env = {
|
||
"MUSE_ROLE_OPUS_BASE_URL": "https://role.example",
|
||
"MUSE_ROLE_OPUS_AUTH_TOKEN": "secret",
|
||
}
|
||
with (
|
||
mock.patch.dict(os.environ, env, clear=False),
|
||
mock.patch.object(muse_llm.requests, "Session", return_value=session),
|
||
mock.patch.object(muse_llm, "_monotonic", side_effect=[0.0, 0.0, 0.9]),
|
||
mock.patch.object(muse_llm.time, "sleep") as sleep,
|
||
):
|
||
with self.assertRaisesRegex(RuntimeError, "deadline"):
|
||
muse_llm.chat_fixed_opus(
|
||
"{}",
|
||
model=FIXED_OPUS_POLICY_ALIAS,
|
||
retries=2,
|
||
timeout=1,
|
||
)
|
||
sleep.assert_not_called()
|
||
self.assertEqual(session.post.call_count, 1)
|
||
|
||
def test_fixed_opus_rejects_env_model_drift_before_http(self):
|
||
env = {
|
||
"MUSE_ROLE_OPUS_BASE_URL": "https://role.example",
|
||
"MUSE_ROLE_OPUS_AUTH_TOKEN": "secret",
|
||
"MUSE_ROLE_OPUS_MODEL": "other-model",
|
||
}
|
||
with (
|
||
mock.patch.dict(os.environ, env, clear=False),
|
||
mock.patch.object(muse_llm.requests, "Session") as session,
|
||
):
|
||
with self.assertRaisesRegex(RuntimeError, "冻结完整模型 ID"):
|
||
muse_llm.chat_fixed_opus(
|
||
"{}",
|
||
model=FIXED_OPUS_POLICY_ALIAS,
|
||
resolved_model_id=FIXED_OPUS_MODEL_ID,
|
||
retries=0,
|
||
)
|
||
session.assert_not_called()
|
||
|
||
def test_fixed_opus_requires_explicit_role_credentials(self):
|
||
with mock.patch.dict(
|
||
os.environ,
|
||
{"MUSE_ROLE_OPUS_BASE_URL": "", "MUSE_ROLE_OPUS_AUTH_TOKEN": ""},
|
||
clear=False,
|
||
):
|
||
with self.assertRaisesRegex(RuntimeError, "MUSE_ROLE_OPUS"):
|
||
muse_llm.chat_fixed_opus("{}", model=FIXED_OPUS_POLICY_ALIAS, retries=0)
|
||
|
||
def test_fixed_opus_cost_and_budget_output_cap_are_conservative(self):
|
||
cost = muse_llm.cost_fixed_opus({"input_tokens": 10, "output_tokens": 5})
|
||
self.assertGreater(cost, 0)
|
||
cap = muse_llm.fixed_opus_max_tokens(Decimal("1.000000"), 1000)
|
||
expected = int(
|
||
(
|
||
1.0
|
||
- 1000 * muse_llm.FIXED_OPUS_CACHE_WRITE_USD_PER_MTOK / 1_000_000
|
||
)
|
||
* 1_000_000
|
||
/ muse_llm.FIXED_OPUS_OUTPUT_USD_PER_MTOK
|
||
)
|
||
self.assertEqual(cap, expected)
|
||
self.assertLessEqual(cap, muse_llm.FIXED_OPUS_MAX_OUTPUT_TOKENS)
|
||
|
||
|
||
class RunRoleContractTests(unittest.TestCase):
|
||
def test_role_prompt_composition_is_full_and_idempotent(self):
|
||
role = "---\nname: writer\n---\n你是写手。\n"
|
||
task = "只返回正文草稿。"
|
||
composed = compose_role_system_prompt(role, task)
|
||
self.assertTrue(composed.startswith(role.rstrip()))
|
||
self.assertIn(task, composed)
|
||
self.assertEqual(compose_role_system_prompt(role, composed), composed)
|
||
|
||
def test_success_returns_structured_output_and_camelcase_receipt(self):
|
||
chat = _fake_chat('{"candidateBody": "正文内容"}')
|
||
profile = _profile()
|
||
result = run_role(profile, {"task": "write"}, governed_chat=chat)
|
||
self.assertEqual(result.structured_output, {"candidateBody": "正文内容"})
|
||
receipt = result.receipt.as_dict()
|
||
self.assertEqual(set(receipt), RECEIPT_KEYS)
|
||
self.assertEqual(receipt["adapterRole"], "writer")
|
||
self.assertEqual(receipt["requestedModelId"], FIXED_OPUS_POLICY_ALIAS)
|
||
self.assertEqual(receipt["actualModelId"], muse_llm.FIXED_OPUS_CANONICAL_MODEL)
|
||
self.assertEqual(receipt["terminalReason"], "completed")
|
||
self.assertFalse(receipt["isError"])
|
||
self.assertIsNone(receipt["exitCode"])
|
||
self.assertEqual(receipt["jsonSchemaSha256"], profile.json_schema_sha256)
|
||
self.assertEqual(receipt["inputSha256"], sha256_text(canonical_json({"task": "write"})))
|
||
self.assertEqual(
|
||
receipt["structuredOutputSha256"],
|
||
sha256_json({"candidateBody": "正文内容"}),
|
||
)
|
||
# 系统提示词必须原样注入治理调用（派发合同：已装配的身份与角色合同进 system）。
|
||
self.assertEqual(chat.calls[0]["model"], profile.model_alias)
|
||
self.assertEqual(chat.calls[0]["system"], build_dispatch_system_prompt(profile))
|
||
self.assertTrue(chat.calls[0]["system"].startswith("你是测试角色。"))
|
||
self.assertIn("candidateBody", chat.calls[0]["system"])
|
||
self.assertEqual(chat.calls[0]["timeout"], profile.timeout_seconds)
|
||
|
||
def test_schema_violation_fails_closed_with_receipt(self):
|
||
chat = _fake_chat('{"unexpected": true}')
|
||
with self.assertRaises(RoleRuntimeError) as ctx:
|
||
run_role(_profile(), {"task": "write"}, governed_chat=chat)
|
||
self.assertEqual(ctx.exception.primary_code, "WRITER_SCHEMA_INVALID")
|
||
# 单一原因时 causes 不重复主码。
|
||
self.assertIsNotNone(ctx.exception.receipt)
|
||
self.assertTrue(ctx.exception.receipt.as_dict()["isError"])
|
||
|
||
def test_minlength_enforced(self):
|
||
chat = _fake_chat('{"candidateBody": ""}')
|
||
with self.assertRaises(RoleRuntimeError) as ctx:
|
||
run_role(_profile(), {"task": "write"}, governed_chat=chat)
|
||
self.assertEqual(ctx.exception.primary_code, "WRITER_SCHEMA_INVALID")
|
||
|
||
def test_model_chain_exhausted_is_model_unavailable(self):
|
||
chat = _fake_chat(None)
|
||
with self.assertRaises(RoleRuntimeError) as ctx:
|
||
run_role(_profile(), {"task": "write"}, governed_chat=chat)
|
||
self.assertEqual(ctx.exception.primary_code, "WRITER_MODEL_UNAVAILABLE")
|
||
|
||
def test_model_outside_frozen_policy_fails_closed(self):
|
||
chat = _fake_chat('{"candidateBody": "正文内容"}', actual_model="unknown-model")
|
||
with self.assertRaises(RoleRuntimeError) as ctx:
|
||
run_role(_profile(), {"task": "write"}, governed_chat=chat)
|
||
self.assertEqual(ctx.exception.primary_code, "WRITER_MODEL_POLICY_MISMATCH")
|
||
self.assertFalse(ctx.exception.receipt.model_match)
|
||
|
||
def test_budget_exceeded_after_call_fails_closed(self):
|
||
# 冻结预算 0.000001，任何真实用量都会超出。
|
||
chat = _fake_chat('{"candidateBody": "正文内容"}')
|
||
profile = _profile(max_budget_usd_per_call=Decimal("0.000001"))
|
||
with self.assertRaises(RoleRuntimeError) as ctx:
|
||
run_role(profile, {"task": "write"}, governed_chat=chat)
|
||
self.assertEqual(ctx.exception.primary_code, "WRITER_BUDGET_EXCEEDED")
|
||
self.assertIsNotNone(ctx.exception.receipt)
|
||
|
||
def test_context_cap_checked_before_model_call(self):
|
||
chat = _fake_chat('{"candidateBody": "x"}')
|
||
profile = _profile(max_context_chars=10)
|
||
with self.assertRaises(RoleRuntimeError) as ctx:
|
||
run_role(profile, {"task": "write" * 100}, governed_chat=chat)
|
||
self.assertEqual(ctx.exception.primary_code, "WRITER_BUDGET_EXCEEDED")
|
||
self.assertEqual(chat.calls, [])
|
||
|
||
def test_business_validator_transforms_and_rejects(self):
|
||
chat = _fake_chat('{"candidateBody": "正文内容"}')
|
||
|
||
def validator(value):
|
||
return {**value, "validated": True}
|
||
|
||
# additionalProperties False 的 schema 下，validator 追加字段仍按输出合同校验后通过：
|
||
# validator 返回即最终 structured_output（与 run_claude 语义一致）。
|
||
result = run_role(
|
||
_profile(), {"task": "write"}, governed_chat=chat,
|
||
business_validator=validator,
|
||
)
|
||
self.assertEqual(result.structured_output["validated"], True)
|
||
|
||
def reject(_value):
|
||
raise ValueError("业务合同不满足")
|
||
|
||
with self.assertRaises(RoleRuntimeError) as ctx:
|
||
run_role(_profile(), {"task": "write"}, governed_chat=chat,
|
||
business_validator=reject)
|
||
self.assertEqual(ctx.exception.primary_code, "WRITER_SCHEMA_INVALID")
|
||
|
||
def test_full_draft202012_keywords_fail_closed(self):
|
||
schema = {
|
||
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
||
"type": "object",
|
||
"additionalProperties": False,
|
||
"required": ["schemaVersion", "verdict", "scores"],
|
||
"properties": {
|
||
"schemaVersion": {"const": "v1"},
|
||
"verdict": {"type": "string", "enum": ["pass", "fail"]},
|
||
"scores": {
|
||
"type": "array",
|
||
"minItems": 2,
|
||
"uniqueItems": True,
|
||
"items": {
|
||
"type": "number",
|
||
"minimum": 0,
|
||
"maximum": 10,
|
||
"multipleOf": 0.5,
|
||
},
|
||
},
|
||
},
|
||
}
|
||
invalid = (
|
||
{"schemaVersion": "wrong", "verdict": "pass", "scores": [1, 2]},
|
||
{"schemaVersion": "v1", "verdict": "INVALID", "scores": [1, 2]},
|
||
{"schemaVersion": "v1", "verdict": "pass", "scores": []},
|
||
{"schemaVersion": "v1", "verdict": "pass", "scores": [1, 1]},
|
||
{"schemaVersion": "v1", "verdict": "pass", "scores": [1, 10.25]},
|
||
{"schemaVersion": "v1", "verdict": "pass", "scores": [1, 10.5]},
|
||
)
|
||
for payload in invalid:
|
||
with self.subTest(payload=payload), self.assertRaises(RoleRuntimeError) as caught:
|
||
run_role(
|
||
_profile(schema=schema),
|
||
{"task": "judge"},
|
||
governed_chat=_fake_chat(canonical_json(payload)),
|
||
)
|
||
self.assertEqual(caught.exception.primary_code, "WRITER_SCHEMA_INVALID")
|
||
|
||
def test_invalid_json_schema_is_rejected_at_profile_construction(self):
|
||
schema = {"type": 123}
|
||
with self.assertRaisesRegex(ValueError, "Draft 2020-12"):
|
||
_profile(schema=schema)
|
||
|
||
def test_writer_cannot_silently_switch_to_content_model_chain(self):
|
||
with self.assertRaisesRegex(ValueError, "固定 Opus"):
|
||
_profile(
|
||
model_alias=CONTENT_POLICY_ALIAS,
|
||
model_policy_version=CONTENT_POLICY_VERSION,
|
||
resolved_model_id=None,
|
||
)
|
||
|
||
def test_detector_may_use_explicit_content_policy(self):
|
||
profile = _profile(
|
||
adapter_role="semantic_detector",
|
||
model_alias=CONTENT_POLICY_ALIAS,
|
||
model_policy_version=CONTENT_POLICY_VERSION,
|
||
resolved_model_id=None,
|
||
)
|
||
self.assertEqual(profile.model_alias, CONTENT_POLICY_ALIAS)
|
||
|
||
def test_profile_hash_drift_rejected_at_construction(self):
|
||
with self.assertRaises(ValueError):
|
||
_profile(system_prompt_sha256="sha256:" + "0" * 64)
|
||
with self.assertRaises(ValueError):
|
||
_profile(json_schema_sha256="sha256:" + "0" * 64)
|
||
|
||
def test_unsupported_role_rejected(self):
|
||
with self.assertRaises(ValueError):
|
||
_profile(adapter_role="attacker")
|
||
|
||
def test_profile_identity_stable_and_hashable(self):
|
||
profile = _profile()
|
||
self.assertTrue(profile.execution_profile_sha256.startswith("sha256:"))
|
||
self.assertEqual(
|
||
profile.execution_profile_sha256, _profile().execution_profile_sha256
|
||
)
|
||
|
||
def test_error_prefix_follows_role(self):
|
||
chat = _fake_chat(None)
|
||
profile = _profile(adapter_role="semantic_detector")
|
||
with self.assertRaises(RoleRuntimeError) as ctx:
|
||
run_role(profile, {"task": "check"}, governed_chat=chat)
|
||
self.assertEqual(
|
||
ctx.exception.primary_code, "SEMANTIC_DETECTOR_MODEL_UNAVAILABLE"
|
||
)
|
||
|
||
|
||
class PathTraversalTests(unittest.TestCase):
|
||
def test_ellipsis_is_not_traversal(self):
|
||
self.assertFalse(contains_path_traversal("他说：等等...然后离开了"))
|
||
self.assertFalse(contains_path_traversal("..."))
|
||
|
||
def test_component_traversal_detected(self):
|
||
for value in ("a/../b", "../x", "a/..", "..", "a\\..\\b"):
|
||
self.assertTrue(contains_path_traversal(value), value)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|