feat: add response cache middleware with poisoning-safe keys
This commit is contained in:
@@ -0,0 +1,205 @@
|
||||
"""CacheMW 与缓存 key 公式测试(ARCH §7.5: 防毒化 key、命中重建、静默降级)。"""
|
||||
|
||||
import dataclasses
|
||||
import hashlib
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from polygateway.backends.memory.cache import InMemoryCache
|
||||
from polygateway.errors import ResultInvalidError, TransientError
|
||||
from polygateway.middleware.cache import CacheMW, build_cache_key, digest_messages
|
||||
from polygateway.types import ChatRequest, LLMResponse
|
||||
|
||||
_MSGS = [{"role": "user", "content": "hi"}]
|
||||
|
||||
|
||||
def _resp(content="cached", **overrides):
|
||||
base = dict(
|
||||
content=content, thinking="", model="m", provider="p",
|
||||
prompt_tokens=1, completion_tokens=2, latency_ms=30,
|
||||
ttft_ms=5.0, max_inter_token_ms=2.0, cache_hit=False, call_id="orig",
|
||||
source_name="s1", usage_source="measured",
|
||||
)
|
||||
base.update(overrides)
|
||||
return LLMResponse(**base)
|
||||
|
||||
|
||||
class TestKeyFormula:
|
||||
def test_same_request_same_key(self):
|
||||
k1 = build_cache_key("m", _MSGS, "proj", None)
|
||||
k2 = build_cache_key("m", _MSGS, "proj", None)
|
||||
assert k1 == k2 and k1.startswith("pgw:cache:")
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("a", "b"),
|
||||
[
|
||||
(("m1", _MSGS, "proj", None), ("m2", _MSGS, "proj", None)),
|
||||
(("m", _MSGS, "proj", None), ("m", _MSGS, "tenant2", None)),
|
||||
(("m", _MSGS, "proj", None), ("m", _MSGS, "proj", "epoch2")),
|
||||
(("m", _MSGS, "proj", "s1"), ("m", _MSGS, "proj", "s2")),
|
||||
(("m", _MSGS, "proj", None), ("m", [{"role": "user", "content": "yo"}], "proj", None)),
|
||||
],
|
||||
)
|
||||
def test_any_dimension_change_changes_key(self, a, b):
|
||||
assert build_cache_key(*a) != build_cache_key(*b)
|
||||
|
||||
def test_multimodal_part_digested_not_inlined(self):
|
||||
big_b64 = "data:image/png;base64," + "A" * 1_000_000
|
||||
messages = [{"role": "user", "content": [
|
||||
{"type": "image_url", "image_url": {"url": big_b64}},
|
||||
{"type": "text", "text": "describe"},
|
||||
]}]
|
||||
digested = digest_messages(messages)
|
||||
payload = json.dumps(digested, ensure_ascii=False)
|
||||
assert len(payload) < 500 # 大图不进 canonical_json
|
||||
expected = hashlib.sha256(big_b64.encode()).hexdigest()
|
||||
assert expected in payload # 但字节变化仍改变 key
|
||||
# 图像字节变化 → key 变
|
||||
messages2 = [{"role": "user", "content": [
|
||||
{"type": "image_url", "image_url": {"url": big_b64[:-1] + "B"}},
|
||||
{"type": "text", "text": "describe"},
|
||||
]}]
|
||||
assert build_cache_key("m", messages, "p", None) != build_cache_key("m", messages2, "p", None)
|
||||
|
||||
|
||||
class _Terminal:
|
||||
def __init__(self, response):
|
||||
self.response = response
|
||||
self.calls = 0
|
||||
|
||||
async def __call__(self, request):
|
||||
self.calls += 1
|
||||
if isinstance(self.response, Exception):
|
||||
raise self.response
|
||||
return self.response
|
||||
|
||||
|
||||
def _mw(backend, **kwargs):
|
||||
defaults = dict(backend=backend, model_fingerprint="m", default_namespace="proj", ttl_s=3600)
|
||||
defaults.update(kwargs)
|
||||
return CacheMW(**defaults)
|
||||
|
||||
|
||||
class TestCacheFlow:
|
||||
async def test_miss_then_hit_with_fresh_call_id(self):
|
||||
backend = InMemoryCache()
|
||||
mw = _mw(backend)
|
||||
terminal = _Terminal(_resp())
|
||||
first = await mw(ChatRequest(messages=_MSGS), terminal)
|
||||
assert first.cache_hit is False and terminal.calls == 1
|
||||
second = await mw(ChatRequest(messages=_MSGS), terminal)
|
||||
assert second.cache_hit is True and second.latency_ms == 0
|
||||
assert second.content == "cached"
|
||||
assert second.call_id != first.call_id # 命中生成独立 cache_call_id
|
||||
assert terminal.calls == 1 # 未再触达内层
|
||||
|
||||
async def test_per_call_namespace_overrides_default(self):
|
||||
backend = InMemoryCache()
|
||||
mw = _mw(backend)
|
||||
terminal = _Terminal(_resp())
|
||||
await mw(ChatRequest(messages=_MSGS, cache_namespace="tenant-a"), terminal)
|
||||
# 另一租户不得命中
|
||||
await mw(ChatRequest(messages=_MSGS, cache_namespace="tenant-b"), terminal)
|
||||
assert terminal.calls == 2
|
||||
|
||||
async def test_failure_not_cached(self):
|
||||
backend = InMemoryCache()
|
||||
mw = _mw(backend)
|
||||
failing = _Terminal(TransientError("boom"))
|
||||
with pytest.raises(TransientError):
|
||||
await mw(ChatRequest(messages=_MSGS), failing)
|
||||
ok = _Terminal(_resp())
|
||||
await mw(ChatRequest(messages=_MSGS), ok)
|
||||
assert ok.calls == 1 # 失败未被固化,正常回源
|
||||
|
||||
async def test_ttl_expiry(self):
|
||||
t = {"now": 0.0}
|
||||
backend = InMemoryCache(now=lambda: t["now"])
|
||||
mw = _mw(backend, ttl_s=100)
|
||||
terminal = _Terminal(_resp())
|
||||
await mw(ChatRequest(messages=_MSGS), terminal)
|
||||
t["now"] = 101.0
|
||||
await mw(ChatRequest(messages=_MSGS), terminal)
|
||||
assert terminal.calls == 2
|
||||
|
||||
|
||||
class _BrokenBackend:
|
||||
async def get(self, key):
|
||||
raise ConnectionError("redis down")
|
||||
|
||||
async def set(self, key, value, ttl_s):
|
||||
raise ConnectionError("redis down")
|
||||
|
||||
|
||||
class TestDegradation:
|
||||
async def test_backend_failure_degrades_silently(self):
|
||||
mw = _mw(_BrokenBackend())
|
||||
terminal = _Terminal(_resp())
|
||||
resp = await mw(ChatRequest(messages=_MSGS), terminal)
|
||||
assert resp.content == "cached" and terminal.calls == 1 # 读写全降级,调用照常
|
||||
|
||||
async def test_corrupt_cache_value_treated_as_miss(self):
|
||||
backend = InMemoryCache()
|
||||
mw = _mw(backend)
|
||||
key = build_cache_key("m", _MSGS, "proj", None)
|
||||
await backend.set(key, "{not json", 3600)
|
||||
terminal = _Terminal(_resp())
|
||||
resp = await mw(ChatRequest(messages=_MSGS), terminal)
|
||||
assert terminal.calls == 1 and resp.cache_hit is False
|
||||
|
||||
|
||||
class _FakeStrategy:
|
||||
"""fake StructuredOutputStrategy(T2 冻结的 Protocol,不依赖 T11)。"""
|
||||
|
||||
def request_overlay(self, schema):
|
||||
return {}
|
||||
|
||||
def parse(self, text):
|
||||
data = json.loads(text) # 简化: 直接 json.loads
|
||||
if not isinstance(data, dict):
|
||||
raise ResultInvalidError("非对象", raw_text=text)
|
||||
return data
|
||||
|
||||
|
||||
class _StrictModel:
|
||||
"""鸭子型 pydantic 模型: model_validate 要求含 answer 键。"""
|
||||
|
||||
@classmethod
|
||||
def model_validate(cls, data):
|
||||
if "answer" not in data:
|
||||
raise ValueError("missing answer")
|
||||
return {"validated": data["answer"]}
|
||||
|
||||
|
||||
class TestStructuredRehydration:
|
||||
async def test_hit_rebuilds_structured_data(self):
|
||||
backend = InMemoryCache()
|
||||
mw = _mw(backend, strategy=_FakeStrategy())
|
||||
terminal = _Terminal(_resp(content='{"answer": 42}'))
|
||||
req = ChatRequest(messages=_MSGS, structured=_StrictModel)
|
||||
await mw(req, terminal)
|
||||
hit = await mw(req, terminal)
|
||||
assert hit.cache_hit is True
|
||||
assert hit.structured_data == {"validated": 42}
|
||||
assert terminal.calls == 1
|
||||
|
||||
async def test_schema_change_revalidation_failure_falls_back_to_source(self):
|
||||
backend = InMemoryCache()
|
||||
mw = _mw(backend, strategy=_FakeStrategy())
|
||||
terminal = _Terminal(_resp(content='{"other": 1}')) # 缓存内容不含 answer
|
||||
await mw(ChatRequest(messages=_MSGS), terminal) # 无 structured 写入
|
||||
# 换 schema 读: 重校验失败 → 按未命中回源
|
||||
again = await mw(ChatRequest(messages=_MSGS, structured=_StrictModel), terminal)
|
||||
assert terminal.calls == 2 and again.cache_hit is False
|
||||
|
||||
async def test_structured_data_not_serialized_into_cache(self):
|
||||
backend = InMemoryCache()
|
||||
mw = _mw(backend, strategy=_FakeStrategy())
|
||||
terminal = _Terminal(
|
||||
dataclasses.replace(_resp(content='{"answer": 1}'), structured_data={"x": object()})
|
||||
)
|
||||
await mw(ChatRequest(messages=_MSGS), terminal) # 不可 JSON 的 structured_data 不阻塞写缓存
|
||||
key = build_cache_key("m", _MSGS, "proj", None)
|
||||
raw = await backend.get(key)
|
||||
assert raw is not None and "structured_data" not in json.loads(raw)
|
||||
Reference in New Issue
Block a user