#!/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 字段集合，record-run-evidence 与看板按此消费。 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()