diff --git a/tests/test_perception.py b/tests/test_perception.py index 9d3293af..7e21a13d 100644 --- a/tests/test_perception.py +++ b/tests/test_perception.py @@ -89,3 +89,25 @@ def test_describe_cached_does_not_cache_failures() -> None: assert perception.describe_cached(**kw) == "" # backend down → "" (uncached) assert perception.describe_cached(**kw) == "recovered" # retried, succeeds assert prov.calls == 2 + + +def test_describe_peek_returns_none_when_absent() -> None: + assert perception.describe_peek(alias="omni", content_hash="missing") is None + + +def test_describe_peek_returns_cached_without_recompute() -> None: + prov = _StubProvider(content="desc") + kw: dict[str, Any] = { + "provider": prov, + "client": object(), + "model": "m", + "alias": "omni", + "content_hash": "h", + "parts": _parts(), + } + perception.describe_cached(**kw) # populate the memo + assert prov.calls == 1 + # Peek serves the memoized text and never re-invokes the backend — this is + # what lets the wire resolver skip the PDF rasterize on a cross-send hit. + assert perception.describe_peek(alias="omni", content_hash="h") == "desc" + assert prov.calls == 1 diff --git a/tests/test_session_attachments.py b/tests/test_session_attachments.py index d92c72cb..ee2f4307 100644 --- a/tests/test_session_attachments.py +++ b/tests/test_session_attachments.py @@ -650,6 +650,73 @@ class TestResolveAttachmentsCapsThreading: assert out["aT"]["document"]["media_type"] == "application/pdf" +class TestResolveAttachmentsPerSendCache: + """The per-send wire-part memo collapses the re-fetch + re-rasterize that the + resolver would otherwise repeat on every agentic round-trip within one send.""" + + def _att(self): + return { + "attachment_id": "aT", + "filename": "r.pdf", + "mime_type": "application/pdf", + "kind": "pdf", + "content": b"%PDF", + } + + def test_cache_collapses_repeat_resolves(self, tmp_db, mock_openai_client, monkeypatch): + s = _make_session(mock_openai_client) + fetches = {"n": 0} + rasters = {"n": 0} + + def _fetch(ids): + fetches["n"] += 1 + return [self._att()] if ids else [] + + monkeypatch.setattr("turnstone.core.session.get_attachments", _fetch) + monkeypatch.setattr( + "turnstone.core.pdf.rasterize_pdf", + lambda data: rasters.__setitem__("n", rasters["n"] + 1) or [b"pg"], + ) + caps = ModelCapabilities(supports_pdf=False, supports_vision=True) + s._wire_part_cache = {} # simulate being inside send() + first = s._resolve_attachments(["aT"], caps) + second = s._resolve_attachments(["aT"], caps) + assert first == second + assert isinstance(first["aT"], list) + # Fetched + rasterized once despite two resolver passes. + assert fetches["n"] == 1 + assert rasters["n"] == 1 + + def test_no_cache_outside_send_rematerializes(self, tmp_db, mock_openai_client, monkeypatch): + s = _make_session(mock_openai_client) + fetches = {"n": 0} + + def _fetch(ids): + fetches["n"] += 1 + return [self._att()] + + monkeypatch.setattr("turnstone.core.session.get_attachments", _fetch) + monkeypatch.setattr("turnstone.core.pdf.rasterize_pdf", lambda data: [b"pg"]) + caps = ModelCapabilities(supports_pdf=False, supports_vision=True) + assert s._wire_part_cache is None # default outside a send → no caching + s._resolve_attachments(["aT"], caps) + s._resolve_attachments(["aT"], caps) + assert fetches["n"] == 2 + + def test_cache_keyed_by_caps(self, tmp_db, mock_openai_client, monkeypatch): + s = _make_session(mock_openai_client) + monkeypatch.setattr("turnstone.core.session.get_attachments", lambda ids: [self._att()]) + monkeypatch.setattr("turnstone.core.pdf.rasterize_pdf", lambda data: [b"pg"]) + s._wire_part_cache = {} + native = s._resolve_attachments(["aT"], ModelCapabilities(supports_pdf=True)) + rasterized = s._resolve_attachments( + ["aT"], ModelCapabilities(supports_pdf=False, supports_vision=True) + ) + # Different caps → different materialization, not a stale same-id hit. + assert native["aT"]["type"] == "document" + assert isinstance(rasterized["aT"], list) + + class TestByReferenceMediaBudget: """bug-2: by-reference pdf/audio are charged a bounded budget — not zero (over-context), not the full multi-MB source blob (over-trim).""" diff --git a/turnstone/core/perception.py b/turnstone/core/perception.py index 13301dbe..a70dc645 100644 --- a/turnstone/core/perception.py +++ b/turnstone/core/perception.py @@ -140,3 +140,15 @@ def describe_cached( _cache.pop(next(iter(_cache)), None) _cache[key] = text return text + + +def describe_peek(*, alias: str, content_hash: str) -> str | None: + """Return the memoized description for ``(alias, content_hash)`` without + computing, or ``None`` if absent. + + Lets the wire resolver skip the expensive parts build (a PDF rasterize) when + the description is already memoized from an earlier send — :func:`describe_cached` + ignores ``parts`` on a hit, so building them first would be pure waste. + """ + with _cache_lock: + return _cache.get(f"{alias}:{content_hash}") diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 2ae24a78..e49c5be4 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -1057,6 +1057,14 @@ class ChatSession: # many times within a turn, so the injected set is touched at most once # per memory per turn. Cleared alongside the search cache. self._touched_memory_keys: set[tuple[str, str, str]] = set() + # Per-send memo for the wire attachment resolver (set in send(), None + # outside a send). _resolve_attachments re-runs on every agentic + # round-trip, so this caches the materialized part by + # (attachment_id, caps-signature) to avoid re-fetching + re-rasterizing + + # re-base64'ing the same blob once per round-trip. + self._wire_part_cache: ( + dict[tuple[str, tuple[bool, bool, bool]], dict[str, Any] | list[dict[str, Any]]] | None + ) = None self._ws_id = ws_id or uuid.uuid4().hex self._title_generated = False self._read_files: set[str] = set() @@ -2955,11 +2963,30 @@ class ChatSession: # right caps; default to the primary only when called without one. if caps is None: caps = self._get_capabilities() + # Per-send memo (see send()): the wire resolver is re-invoked on every + # round-trip and per fallback model, so without this a PDF in history is + # re-rasterized / a blob re-base64'd once per round-trip. Key on + # (id, caps-signature): the same stored blob materializes differently per + # capability set, and a fallback to a different-caps model can resolve + # within one send. ``cache`` is None outside a send → original behavior. + cache = self._wire_part_cache + caps_sig = (caps.supports_pdf, caps.supports_vision, caps.supports_audio_input) out: dict[str, Any] = {} - for att in get_attachments(ids): - part = self._wire_content_part(att, caps) - if part is not None: - out[str(att["attachment_id"])] = part + missing: list[str] = [] + for att_id in ids: + hit = cache.get((att_id, caps_sig)) if cache is not None else None + if hit is not None: + out[att_id] = hit + else: + missing.append(att_id) + if missing: + for att in get_attachments(missing): + part = self._wire_content_part(att, caps) + if part is not None: + aid = str(att["attachment_id"]) + out[aid] = part + if cache is not None: + cache[(aid, caps_sig)] = part return out def _wire_content_part( @@ -3149,19 +3176,26 @@ class ChatSession: return None if kind == "audio" and not caps.supports_audio_input: return None - parts = self._perception_parts(att, kind) - if not parts: - return None - from turnstone.core.perception import describe_cached + from turnstone.core.perception import describe_cached, describe_peek - text = describe_cached( - provider=provider, - client=client, - model=model, - alias=alias, - content_hash=str(att.get("attachment_id")), - parts=parts, - ) + # Peek the (alias, content_hash) memo BEFORE building parts: for a PDF, + # _perception_parts rasterizes every page, but describe_cached returns a + # memoized description without touching parts on a hit — so on a cross-send + # hit the rasterize would be pure waste. + content_hash = str(att.get("attachment_id")) + text = describe_peek(alias=alias, content_hash=content_hash) + if text is None: + parts = self._perception_parts(att, kind) + if not parts: + return None + text = describe_cached( + provider=provider, + client=client, + model=model, + alias=alias, + content_hash=content_hash, + parts=parts, + ) if not text: return None name = str(att.get("filename") or kind) @@ -4084,6 +4118,10 @@ class ChatSession: # reference so subprocesses from old generations are still killed. self._cancel_event = threading.Event() self._cancelled_partial_msg = None + # Fresh per-send attachment wire-part memo (see __init__): bounds the + # heavy rasterized-page parts to one send and picks up any mid-session + # capability / config change. + self._wire_part_cache = {} # Metacognitive nudge: check for correction/completion signals # before _append_user_turn so any fired nudge (plus any nudges