From 628608655124b14bca3982efe1f59dbc36eac6c8 Mon Sep 17 00:00:00 2001
From: iomgaa <iomgaaycz@gmail.com>
Date: Mon, 20 Jul 2026 07:10:42 -0400
Subject: [PATCH] feat: add response cache middleware with poisoning-safe keys

---
 src/polygateway/backends/memory/cache.py |  30 ++++
 src/polygateway/backends/redis_cache.py  |  44 +++++
 src/polygateway/middleware/cache.py      | 155 +++++++++++++++++
 tests/unit/test_cache.py                 | 205 +++++++++++++++++++++++
 4 files changed, 434 insertions(+)
 create mode 100644 src/polygateway/backends/memory/cache.py
 create mode 100644 src/polygateway/backends/redis_cache.py
 create mode 100644 src/polygateway/middleware/cache.py
 create mode 100644 tests/unit/test_cache.py

diff --git a/src/polygateway/backends/memory/cache.py b/src/polygateway/backends/memory/cache.py
new file mode 100644
index 0000000..2e4984c
--- /dev/null
+++ b/src/polygateway/backends/memory/cache.py
@@ -0,0 +1,30 @@
+"""进程内缓存后端: dict + 过期时刻;测试与无 Redis 场景使用。"""
+
+from __future__ import annotations
+
+import time
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+class InMemoryCache:
+    """笨 KV(CacheBackend 契约);容量不设限——库内默认 TTL 必填防无界增长。"""
+
+    def __init__(self, now: Callable[[], float] = time.monotonic) -> None:
+        self._now = now
+        self._store: dict[str, tuple[float, str]] = {}
+
+    async def get(self, key: str) -> str | None:
+        item = self._store.get(key)
+        if item is None:
+            return None
+        expires_at, value = item
+        if self._now() >= expires_at:
+            del self._store[key]
+            return None
+        return value
+
+    async def set(self, key: str, value: str, ttl_s: int) -> None:
+        self._store[key] = (self._now() + ttl_s, value)
diff --git a/src/polygateway/backends/redis_cache.py b/src/polygateway/backends/redis_cache.py
new file mode 100644
index 0000000..a755b34
--- /dev/null
+++ b/src/polygateway/backends/redis_cache.py
@@ -0,0 +1,44 @@
+"""Redis 缓存后端: 笨 KV(CacheBackend 契约);redis 是 optional extra。
+
+降级不在这里做——本类忠实上抛异常,静默降级是 CacheMW 的职责(算法一份)。
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+try:
+    import redis.asyncio as aioredis
+except ImportError as _exc:  # pragma: no cover - 依赖缺失路径
+    aioredis = None
+    _IMPORT_ERROR = _exc
+else:
+    _IMPORT_ERROR = None
+
+
+class RedisCache:
+    """基于 redis.asyncio 的 KV;value 为 JSON 字符串,TTL 由调用方传入。"""
+
+    def __init__(self, client: Any) -> None:
+        if aioredis is None:
+            raise ImportError(
+                "Redis 缓存后端需要 redis 包: pip install 'polygateway[redis]'"
+            ) from _IMPORT_ERROR
+        self._client = client
+
+    @classmethod
+    def from_url(cls, url: str) -> RedisCache:
+        if aioredis is None:
+            raise ImportError(
+                "Redis 缓存后端需要 redis 包: pip install 'polygateway[redis]'"
+            ) from _IMPORT_ERROR
+        return cls(aioredis.from_url(url, decode_responses=True))
+
+    async def get(self, key: str) -> str | None:
+        return await self._client.get(key)
+
+    async def set(self, key: str, value: str, ttl_s: int) -> None:
+        await self._client.set(key, value, ex=ttl_s)
+
+    async def aclose(self) -> None:
+        await self._client.aclose()
diff --git a/src/polygateway/middleware/cache.py b/src/polygateway/middleware/cache.py
new file mode 100644
index 0000000..c9f8255
--- /dev/null
+++ b/src/polygateway/middleware/cache.py
@@ -0,0 +1,155 @@
+"""CacheMW: 响应缓存中间件;key 公式在此(算法一份),后端是笨 KV。
+
+ARCH §7.5: key = sha256(canonical_json({model, messages_digest, namespace,
+salt}));多模态 part 先摘要再 hash(防 Video-Tree 整段 base64 进 hash 的
+开销);namespace 必填防跨项目/租户毒化;只缓存阶梯通过的成功响应
+(StructuredMW 在内层,能返回即已通过)。读写失败静默降级(铁律)。
+"""
+
+from __future__ import annotations
+
+import asyncio
+import dataclasses
+import hashlib
+import json
+import uuid
+from typing import TYPE_CHECKING, Any
+
+from loguru import logger
+
+from polygateway.types import ChatRequest, LLMResponse
+
+if TYPE_CHECKING:
+    from polygateway.ports import CacheBackend, CallNext, StructuredOutputStrategy
+
+_KEY_PREFIX = "pgw:cache:"
+_RESPONSE_FIELDS = {f.name for f in dataclasses.fields(LLMResponse)}
+
+
+def digest_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """多模态 content part 先各自 sha256 摘要再参与序列化;文本原文参与。
+
+    与遥测落库共用同一函数(ARCH §7.8),保证缓存 key 与遥测口径一致。
+    """
+    digested = []
+    for msg in messages:
+        content = msg.get("content")
+        if isinstance(content, list):
+            parts = [_digest_part(part) for part in content]
+            digested.append({**msg, "content": parts})
+        else:
+            digested.append(msg)
+    return digested
+
+
+def _digest_part(part: Any) -> Any:
+    if isinstance(part, dict) and part.get("type") == "image_url":
+        url = str(part.get("image_url", {}).get("url", ""))
+        return {"type": "image_url", "sha256": hashlib.sha256(url.encode()).hexdigest()}
+    return part
+
+
+def build_cache_key(
+    model_fingerprint: str, messages: list[dict[str, Any]], namespace: str, salt: str | None
+) -> str:
+    """缓存 key 公式;salt 仅非 None 时参与(VT 旧键语义: 不传 salt 键形不变)。"""
+    key_obj: dict[str, Any] = {
+        "model": model_fingerprint,
+        "messages": digest_messages(messages),
+        "namespace": namespace,
+    }
+    if salt is not None:
+        key_obj["salt"] = salt
+    payload = json.dumps(key_obj, sort_keys=True, ensure_ascii=False)
+    return _KEY_PREFIX + hashlib.sha256(payload.encode("utf-8")).hexdigest()
+
+
+class CacheMW:
+    """洋葱第二层(遥测内、结构化外)。
+
+    model_fingerprint 由装配层从源列表计算(多源 scope = 排序去重的 model
+    名合集);源集合变化 → key 变化 → 一次性冷启动,换取跨源集合零毒化。
+    """
+
+    def __init__(
+        self,
+        *,
+        backend: CacheBackend,
+        model_fingerprint: str,
+        default_namespace: str,
+        ttl_s: int,
+        strategy: StructuredOutputStrategy | None = None,
+    ) -> None:
+        if not default_namespace.strip():
+            raise ValueError("缓存 namespace 不能为空(防跨项目/租户毒化)")
+        if ttl_s <= 0:
+            raise ValueError("缓存 TTL 必须 > 0(禁止永不过期)")
+        self._backend = backend
+        self._fingerprint = model_fingerprint
+        self._namespace = default_namespace
+        self._ttl_s = ttl_s
+        self._strategy = strategy
+
+    async def __call__(self, request: ChatRequest, call_next: CallNext) -> LLMResponse:
+        namespace = request.cache_namespace or self._namespace
+        key = build_cache_key(self._fingerprint, request.messages, namespace, request.cache_salt)
+        cached = await self._safe_get(key)
+        if cached is not None:
+            hit = self._rehydrate(cached, request)
+            if hit is not None:
+                return hit
+        response = await call_next(request)
+        await self._safe_set(key, self._serialize(response))
+        return response
+
+    # —— 命中路径 ——
+
+    def _rehydrate(self, raw: str, request: ChatRequest) -> LLMResponse | None:
+        """反序列化 + 按本次调用的 structured 档零网络重建;失败按未命中。"""
+        try:
+            data = json.loads(raw)
+            fields = {k: v for k, v in data.items() if k in _RESPONSE_FIELDS}
+            structured_data = self._rebuild_structured(fields.get("content", ""), request)
+            fields.update(
+                cache_hit=True, latency_ms=0, ttft_ms=None, max_inter_token_ms=None,
+                call_id=str(uuid.uuid4()), structured_data=structured_data,
+            )
+            return LLMResponse(**fields)
+        except Exception as exc:
+            logger.warning("缓存命中重建失败,按未命中回源: {}", exc)
+            return None
+
+    def _rebuild_structured(self, content: str, request: ChatRequest) -> Any | None:
+        """对缓存 content 重跑阶梯②③(schema 变更后旧缓存自动重校验,设计 §2.1)。"""
+        if request.structured is None:
+            return None
+        if self._strategy is None:
+            raise ValueError("structured 调用需要装配 strategy 才能命中重建")
+        parsed = self._strategy.parse(content)
+        if request.structured == "json":
+            return parsed
+        return request.structured.model_validate(parsed)
+
+    # —— 写路径与降级 ——
+
+    def _serialize(self, response: LLMResponse) -> str:
+        data = dataclasses.asdict(response)
+        data.pop("structured_data", None)  # pydantic 实例不可 JSON 往返(设计 §2.1)
+        return json.dumps(data, ensure_ascii=False)
+
+    async def _safe_get(self, key: str) -> str | None:
+        try:
+            return await self._backend.get(key)
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            logger.warning("缓存读取失败,降级为未命中: {}", exc)
+            return None
+
+    async def _safe_set(self, key: str, value: str) -> None:
+        try:
+            await self._backend.set(key, value, self._ttl_s)
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            logger.warning("缓存写入失败,跳过缓存: {}", exc)
diff --git a/tests/unit/test_cache.py b/tests/unit/test_cache.py
new file mode 100644
index 0000000..a8b2c80
--- /dev/null
+++ b/tests/unit/test_cache.py
@@ -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)