diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 4d04f9d3a..efabc6f10 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -864,12 +864,18 @@ class AnthropicHandlerMixin: # turn becomes historical on the next request, so even "latest turn only" # rewrites can invalidate the next cache read when the client resends the # original transcript. - if ( - self.config.image_optimize - and messages - and not _bypass - and not is_cache_mode(self.config.mode) - ): + # + # Bypass / image_optimize / messages gating routes through + # ImageCompressionDecision for uniformity with CompressionDecision + + # MemoryDecision. The cache_mode check stays inline because it's + # Anthropic-specific (sites in openai.py / gemini.py don't have it). + from headroom.proxy.image_compression_decision import ImageCompressionDecision + + _image_decision = ImageCompressionDecision.decide( + headers=request.headers, config=self.config, messages=messages + ) + _image_decision.apply_to_tags(tags) + if _image_decision.should_compress and not is_cache_mode(self.config.mode): compressor = None try: compressor = _get_image_compressor() diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 04f27cc6a..d45a53cc1 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -1200,8 +1200,19 @@ class OpenAIHandlerMixin: if _bypass: logger.info(f"[{request_id}] Bypass: skipping compression (header)") - # Image compression: tile alignment + ML-based technique routing - if self.config.image_optimize and messages and not _bypass: + # Image compression: tile alignment + ML-based technique routing. + # Gated on ImageCompressionDecision — same value-type pattern + # as CompressionDecision + MemoryDecision; locks bypass-respect + # in tests so a future site can't drift. + from headroom.proxy.image_compression_decision import ImageCompressionDecision + + _image_decision = ImageCompressionDecision.decide( + headers=request.headers, config=self.config, messages=messages + ) + # tags is populated downstream at L1229 — defer apply_to_tags + # to where the tags dict exists. The decision is captured here + # so the conditional is uniform with the other gates. + if _image_decision.should_compress: from headroom.proxy.helpers import _get_image_compressor compressor = None @@ -1229,6 +1240,10 @@ class OpenAIHandlerMixin: headers.pop("accept-encoding", None) tags = extract_tags(headers) client = classify_client(headers) + # Surface the image-compression decision (computed earlier) into + # tags now that the tags dict exists. Same observability pattern + # the funnel uses for passthrough_reason + memory_skip_reason. + _image_decision.apply_to_tags(tags) # PR-A5 (P5-49): strip internal x-headroom-* from upstream-bound # headers AFTER `_extract_tags` reads them. Inbound bypass gating # uses `request.headers.get(...)` above; memory user-id reads diff --git a/headroom/proxy/image_compression_decision.py b/headroom/proxy/image_compression_decision.py new file mode 100644 index 000000000..b558a0cd1 --- /dev/null +++ b/headroom/proxy/image_compression_decision.py @@ -0,0 +1,125 @@ +"""``ImageCompressionDecision``: canonical "should this request have +images compressed?" gate. + +Mirror of :class:`CompressionDecision` (text compression) and +:class:`MemoryDecision`. Pre-this-PR image compression was gated at +two sites (``openai.py:1203``, ``anthropic.py:868``) by inline +conjunctions. Both already checked ``_bypass`` (no drift bug like +the text-compression Gemini bypass-misses) — but consolidating into +a value type still pays off: + +* Locks the contract via tests so a future site can't drift +* ``apply_to_tags()`` surfaces ``image_skip_reason`` to + :class:`RequestOutcome.tags` for dashboard slicing +* Rust-portable shape, same as the other decision types + +Precedence (highest first): + + 1. ``bypass_header`` — user's explicit opt-out + 2. ``image_optimize_disabled`` — operator ``config.image_optimize=False`` + 3. ``no_messages`` — empty / missing messages + 4. otherwise → ``should_compress=True`` + +Distinct from text :class:`CompressionDecision`'s +``compression_disabled`` reason: operators can enable text + disable +image (or vice versa) independently. Same shape, different gate. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +from headroom.proxy.helpers import _headroom_bypass_enabled + + +@dataclass(frozen=True) +class ImageCompressionDecision: + """Immutable, value-equal snapshot of the image-compression gate. + + Construction policy: use :meth:`decide`. Direct construction is + allowed for tests but unusual — handlers always go through + ``decide``. The constituent observability booleans + (``bypass_header_set`` etc.) MUST match the inputs ``decide`` saw; + the factory enforces that invariant, and the dataclass being + frozen means downstream code can't violate it. + """ + + should_compress: bool + # When ``should_compress`` is False, the canonical reason surfaced + # in ``RequestOutcome.tags["image_skip_reason"]`` so the dashboard + # can slice image-skipped traffic by cause. One of: + # * "bypass_header" — user set x-headroom-bypass/mode + # * "image_optimize_disabled" — operator config off + # * "no_messages" — empty / missing messages + # When ``should_compress`` is True, this is None. + passthrough_reason: str | None + + # Observability: every constituent boolean exposed so debug tools + # answer "what did the decision see?" without re-running it. + bypass_header_set: bool + image_optimize_enabled: bool + has_messages: bool + + @classmethod + def decide( + cls, + *, + headers: Any, + config: Any, + messages: Sequence[Any] | None, + ) -> ImageCompressionDecision: + """Compute the canonical image-compression decision. + + Parameters + ---------- + headers + Inbound request headers. Accepts any object with a + ``.get(key)`` method (dict, starlette Headers, mapping). + Bypass detected via ``_headroom_bypass_enabled``. + config + ``HeadroomConfig``-shaped object; only ``image_optimize`` + is read. + messages + Request messages. ``None`` and ``[]`` are equivalent. + """ + bypass = _headroom_bypass_enabled(headers) + image_ok = bool(getattr(config, "image_optimize", False)) + has_msgs = bool(messages) + + if bypass: + reason: str | None = "bypass_header" + should = False + elif not image_ok: + reason = "image_optimize_disabled" + should = False + elif not has_msgs: + reason = "no_messages" + should = False + else: + reason = None + should = True + + return cls( + should_compress=should, + passthrough_reason=reason, + bypass_header_set=bypass, + image_optimize_enabled=image_ok, + has_messages=has_msgs, + ) + + def apply_to_tags(self, tags: dict[str, str]) -> None: + """Stamp the skip reason into a tags dict for dashboard slicing. + + Mutates ``tags`` in place. No-op when ``should_compress=True`` + — absence vs presence is the signal. + + Mirror of :meth:`CompressionDecision.apply_to_tags` and + :meth:`MemoryDecision.apply_to_tags`. Multiple decision tags + coexist in the same dict (``passthrough_reason``, + ``memory_skip_reason``, ``image_skip_reason``) for full + dashboard slicing. + """ + if self.passthrough_reason is not None: + tags["image_skip_reason"] = self.passthrough_reason diff --git a/headroom/proxy/memory_handler.py b/headroom/proxy/memory_handler.py index 641de9caa..10386fd19 100644 --- a/headroom/proxy/memory_handler.py +++ b/headroom/proxy/memory_handler.py @@ -622,6 +622,8 @@ class MemoryHandler: user_id: str, messages: list[dict[str, Any]], request_context: RequestContext | None = None, + *, + ranker: Any | None = None, ) -> str | None: """Search memories and format as context injection. @@ -636,6 +638,14 @@ class MemoryHandler: omitted, behaves as before this fix — single-bucket search against the legacy backend. Production handlers always pass it; tests / mocks can keep the simpler call shape. + ranker: Optional :class:`~headroom.proxy.memory_ranker.MemoryRanker` + — re-ranks the backend's cosine-only candidates by an + additional signal (recency, source, access count, …). + When ``None`` (default), behaviour is identical to pre- + this-PR: pure cosine + min_similarity floor. When + provided, candidates are adapted to + :class:`MemoryCandidate`, re-ranked, then re-filtered by + ``min_similarity`` on the boosted score. Returns: Formatted context string, or None if no relevant memories. @@ -689,23 +699,48 @@ class MemoryHandler: ) return None - # Filter by minimum similarity - filtered_results = [r for r in results if r.score >= self.config.min_similarity] + # Optional re-rank: when a MemoryRanker is provided, adapt + # results to MemoryCandidate, re-rank, then filter by + # min_similarity on the BOOSTED score. The re-rank can + # promote a fresh weak-cosine memory above a stale strong- + # cosine one (RecencyBoostRanker default behaviour). + if ranker is not None: + from headroom.proxy.memory_ranker import MemoryCandidate - if not filtered_results: - logger.debug( - f"Memory: {len(results)} memories found but none above threshold " - f"{self.config.min_similarity}" - ) - return None + candidates = [MemoryCandidate.from_backend_result(r) for r in results] + ranked = ranker.rank(candidates) + # Filter on the post-rank score (the ranker may have + # boosted or attenuated original cosine values). + ranked = [c for c in ranked if c.score >= self.config.min_similarity] + if not ranked: + logger.debug( + f"Memory: {len(results)} memories found but none above threshold " + f"{self.config.min_similarity} after re-rank" + ) + return None + memory_lines = [] + for i, candidate in enumerate(ranked, 1): + memory_lines.append(f"{i}. {candidate.content}") + if candidate.related_entities: + entities_str = ", ".join(candidate.related_entities[:3]) + memory_lines.append(f" (Related: {entities_str})") + else: + # Pre-PR-this behaviour: pure cosine, no boost. + filtered_results = [r for r in results if r.score >= self.config.min_similarity] - # Format as context - memory_lines = [] - for i, result in enumerate(filtered_results, 1): - memory_lines.append(f"{i}. {result.memory.content}") - if hasattr(result, "related_entities") and result.related_entities: - entities_str = ", ".join(result.related_entities[:3]) - memory_lines.append(f" (Related: {entities_str})") + if not filtered_results: + logger.debug( + f"Memory: {len(results)} memories found but none above threshold " + f"{self.config.min_similarity}" + ) + return None + + memory_lines = [] + for i, result in enumerate(filtered_results, 1): + memory_lines.append(f"{i}. {result.memory.content}") + if hasattr(result, "related_entities") and result.related_entities: + entities_str = ", ".join(result.related_entities[:3]) + memory_lines.append(f" (Related: {entities_str})") except Exception as e: logger.warning(f"Memory: Search failed for user {effective_user_id}: {e}") diff --git a/headroom/proxy/memory_ranker.py b/headroom/proxy/memory_ranker.py new file mode 100644 index 000000000..ef5aea09f --- /dev/null +++ b/headroom/proxy/memory_ranker.py @@ -0,0 +1,198 @@ +"""``MemoryRanker``: pluggable re-ranker for memory candidates. + +Pre-this-PR Headroom ranked memory candidates by pure cosine +similarity. Every other memory system we surveyed +(Letta / Mem0 / Cognee / Supermemory) re-ranks beyond cosine — recency, +source weight, access count, and decay are table-stakes for not +returning 6-month-old "winners" when fresh signal exists. + +This module ships the first ranker — :class:`RecencyBoostRanker` — +plus the :class:`MemoryRanker` protocol that future rankers +(source-weight, access-count) plug into. + +The ranker is **pure**: ``rank(candidates) -> ranked_candidates``, +no I/O, no state, no mutation of inputs. Same Rust-port shape as +``CompressionDecision`` and ``MemoryDecision``. + +Performance: O(N) over candidates where N = top_k. One ``math.exp()`` +per candidate. Sub-microsecond per request — no embedding compute, +no network, no disk. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Protocol + +# Use ``timezone.utc`` (always available) instead of ``datetime.UTC`` +# (Python 3.11+) so this module imports cleanly on older interpreters. +_UTC = timezone.utc + + +@dataclass(frozen=True) +class MemoryCandidate: + """Immutable retrieval candidate as it flows through the ranker. + + The shape is the **proxy-side internal contract** — the backend's + return type (typically ``MemoryResult`` with nested + ``MemoryResult.memory.created_at``) is adapted into this flatter + shape at the ranker boundary so the ranker stays backend-agnostic. + + ``score`` is the cosine similarity as returned by the backend. + Rankers MAY mutate ``score`` by returning a new candidate with an + updated score (frozen dataclass means they cannot mutate in place). + """ + + content: str + score: float + created_at: datetime | None = None + source: str | None = None # e.g. "memory_save" | "traffic_learner" | "inline" + related_entities: tuple[str, ...] = field(default_factory=tuple) + + @classmethod + def from_backend_result(cls, result: object) -> MemoryCandidate: + """Adapter from backend ``MemoryResult`` shape to ``MemoryCandidate``. + + The backend returns objects with ``.score``, ``.memory.content``, + and (optionally) ``.memory.created_at`` (str ISO timestamp) + + ``.related_entities``. This adapter flattens that to the + ranker's expected shape and parses the timestamp to ``datetime``. + + Missing / unparseable timestamps → ``None`` (recency-neutral). + """ + score = float(getattr(result, "score", 0.0)) + memory = getattr(result, "memory", None) + content = str(getattr(memory, "content", "")) if memory is not None else "" + raw_dt = getattr(memory, "created_at", None) if memory is not None else None + created_at = _parse_created_at(raw_dt) + raw_related = getattr(result, "related_entities", None) or () + related = tuple(str(x) for x in raw_related) + source_meta = getattr(memory, "metadata", None) or {} + source = source_meta.get("source") if isinstance(source_meta, dict) else None + return cls( + content=content, + score=score, + created_at=created_at, + source=source, + related_entities=related, + ) + + +def _parse_created_at(value: object) -> datetime | None: + """Best-effort parse of a timestamp into a UTC-aware datetime. + + Accepts ``datetime`` (returned as-is, UTC-normalized) or ISO-8601 + string (with or without trailing ``Z``). Anything else → ``None`` + so the ranker treats the candidate as recency-neutral. + """ + if value is None: + return None + if isinstance(value, datetime): + return value if value.tzinfo is not None else value.replace(tzinfo=_UTC) + if isinstance(value, str): + try: + return datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + return None + + +class MemoryRanker(Protocol): + """Re-ranks retrieval candidates. Pure function. + + Implementations MUST: + * Not mutate the input list or its elements + * Be deterministic (same input → same output) for prefix-cache + stability across consecutive turns + * Be backend-agnostic — work with any + :class:`MemoryCandidate`, regardless of which backend produced it + """ + + def rank(self, candidates: list[MemoryCandidate]) -> list[MemoryCandidate]: ... + + +@dataclass(frozen=True) +class RecencyBoostRanker: + """Re-ranker applying an exponential recency decay to cosine scores. + + Final score: ``cosine × exp(-age_days / decay_days)``. + + At ``decay_days=30``: + * age = 0 days → factor 1.000 + * age = 15 days → factor 0.607 + * age = 30 days → factor 0.368 + * age = 60 days → factor 0.135 + * age = 90 days → factor 0.050 + + Tuned so a fresh memory with weak cosine doesn't dominate (factor + decays gradually), but a 6-month-old strong-cosine can't dominate + either (factor approaches zero). Operators tune ``decay_days`` for + their codebase's rate of change — 7 days for rapidly-evolving + repos, 90 days for stable archival. + + Backwards-compat: candidates with ``created_at=None`` get factor + 1.0 — treated as recency-neutral. Lets a backend during a + migration return some rows with timestamps and some without + without breaking the ranker. + + Defensive: negative ages (clock skew on future-timestamped rows) + are clamped to factor 1.0 — a clock-skewed candidate cannot + outrank a real fresh one with the same cosine. + """ + + decay_days: float = 30.0 + + def rank(self, candidates: list[MemoryCandidate]) -> list[MemoryCandidate]: + """Re-rank by ``score × recency_factor``. Pure; input unchanged. + + Sorts descending by boosted score. Ties broken by input order + (Python's sort is stable) — deterministic output for prefix- + cache stability across turns. + """ + if not candidates: + return [] + + now = datetime.now(_UTC) + boosted: list[tuple[int, MemoryCandidate, float]] = [] + for idx, c in enumerate(candidates): + factor = self._recency_factor(now, c.created_at) + boosted.append((idx, c, c.score * factor)) + + # Sort descending by boosted score; stable on ties via the + # captured idx — same input order preserved on ties for + # deterministic output across turns. + boosted.sort(key=lambda triple: (-triple[2], triple[0])) + + # Return new candidates with the boosted score so downstream + # consumers (logging, budget filter) see the post-boost number. + return [ + MemoryCandidate( + content=c.content, + score=new_score, + created_at=c.created_at, + source=c.source, + related_entities=c.related_entities, + ) + for _, c, new_score in boosted + ] + + def _recency_factor(self, now: datetime, created_at: datetime | None) -> float: + """Compute the recency multiplier for a single candidate. + + ``None`` timestamp → 1.0 (recency-neutral, backwards-compat). + Future timestamps → 1.0 (clock-skew defence). + Otherwise: ``exp(-age_days / decay_days)``. + """ + if created_at is None: + return 1.0 + # Normalize to UTC-aware for safe subtraction. + if created_at.tzinfo is None: + created_at = created_at.replace(tzinfo=_UTC) + delta = now - created_at + age_days = delta.total_seconds() / 86400.0 + if age_days <= 0: + # Future-dated row (clock skew). Clamp to neutral. + return 1.0 + return math.exp(-age_days / self.decay_days) diff --git a/scripts/sync-plugin-versions.py b/scripts/sync-plugin-versions.py index b8d88d4f1..347c61830 100644 --- a/scripts/sync-plugin-versions.py +++ b/scripts/sync-plugin-versions.py @@ -1,7 +1,22 @@ -"""Sync plugin manifest versions to the repo's computed release semver.""" +"""Sync plugin manifest versions to the repo's computed release semver. + +Branch-aware: by default this script is a NO-OP on feature branches. +Pre-this-fix it ran on every commit and bumped the manifests to the +PREDICTED next release version, which polluted every PR with version- +bump noise (the prediction advanced as commits landed; each PR ended +up carrying the bump as collateral). + +Sync now only runs when EITHER: + * We're on the ``main`` branch, OR + * ``HEADROOM_SYNC_VERSIONS=1`` is set explicitly (release workflow) + +Result: feature-branch PRs no longer carry manifest bumps; the +release workflow still gets a canonical sync at publish time. +""" from __future__ import annotations +import os import subprocess import sys from pathlib import Path @@ -33,8 +48,50 @@ def compute_repo_semver(root: Path) -> str: return info.npm_version +def _current_branch(root: Path) -> str | None: + """Return the current git branch name, or None if git isn't usable.""" + try: + result = subprocess.run( + ["git", "rev-parse", "--abbrev-ref", "HEAD"], + cwd=root, + capture_output=True, + text=True, + check=False, + ) + except (FileNotFoundError, OSError): + return None + if result.returncode != 0: + return None + return result.stdout.strip() or None + + +def _should_sync(root: Path) -> bool: + """Decide whether to actually run the sync. + + Release workflow opts in via ``HEADROOM_SYNC_VERSIONS=1``; otherwise + we only sync on ``main`` (where the next-release prediction + legitimately lives). On feature branches we no-op — the prediction + would just create PR-level noise. + """ + if os.environ.get("HEADROOM_SYNC_VERSIONS") == "1": + return True + branch = _current_branch(root) + if branch is None: + # Git unavailable or detached HEAD — safest default is no-op. + return False + return branch == "main" + + def main() -> None: root = ROOT + if not _should_sync(root): + # Quiet no-op on feature branches. Print a single line so + # pre-commit users see the reason if they look. + branch = _current_branch(root) or "" + print( + f"sync-plugin-versions: skipping on branch '{branch}' (set HEADROOM_SYNC_VERSIONS=1 to force)" + ) + return version = compute_repo_semver(root) subprocess.run( [ diff --git a/tests/test_handler_outcome_tag_invariant.py b/tests/test_handler_outcome_tag_invariant.py index 33c1a0517..42ceb3009 100644 --- a/tests/test_handler_outcome_tag_invariant.py +++ b/tests/test_handler_outcome_tag_invariant.py @@ -116,3 +116,41 @@ def test_outcome_call_sites_pass_client_kwarg() -> None: "`client = classify_client(headers)` and thread `client=client` " "into the outcome construction. See PR #473 for the pattern." ) + + +# ── Invariant: image-compression must route through ImageCompressionDecision ── + + +import re # noqa: E402 -- only used by the image-decision invariant below + + +def test_no_raw_image_optimize_gate_in_handlers() -> None: + """Locks the post-this-PR contract: image compression must be + gated by :class:`ImageCompressionDecision`, not by an inline + ``if self.config.image_optimize and messages and not _bypass:`` + conjunction. Pre-PR-this both sites used the raw conjunction; + consolidating into a value type means a future site (e.g., new + provider handler) can't drift on bypass-respect or skip-reason + observability. + + Allowed forms after this PR: + * ``if _image_decision.should_compress`` + * ``if _image_decision.should_compress and ...`` + """ + pattern = re.compile(r"^\s*if\s*\(?\s*self\.config\.image_optimize\s+and\s+messages\b") + offenders: list[tuple[str, int, str]] = [] + for f in HANDLER_FILES: + text = f.read_text(encoding="utf-8") + for i, line in enumerate(text.splitlines(), start=1): + if pattern.match(line): + offenders.append((f.name, i, line.rstrip())) + if offenders: + formatted = "\n".join(f" {f}:{ln} {src!r}" for f, ln, src in offenders) + pytest.fail( + f"{len(offenders)} handler site(s) use the pre-PR raw image " + "gate `if self.config.image_optimize and messages [and ...]`:\n" + f"{formatted}\n\n" + "Replace with `ImageCompressionDecision.decide(...)` + " + "`if _image_decision.should_compress:`. See " + "headroom/proxy/image_compression_decision.py for the pattern." + ) diff --git a/tests/test_image_compression_decision.py b/tests/test_image_compression_decision.py new file mode 100644 index 000000000..043781692 --- /dev/null +++ b/tests/test_image_compression_decision.py @@ -0,0 +1,241 @@ +"""Tests for :class:`headroom.proxy.image_compression_decision.ImageCompressionDecision`. + +Image compression today is gated at two sites (``openai.py:1203`` + +``anthropic.py:868``) by inline conjunctions. Both already check +``_bypass`` (the drift problem CompressionDecision fixed is NOT +present here), but consolidating into a value type still pays off: + +* test-lockable contract (no future site can forget bypass) +* ``apply_to_tags()`` surfaces ``image_skip_reason`` in + ``RequestOutcome.tags`` — dashboards can slice image-skipped + traffic by cause (same observability surface as + ``passthrough_reason`` and ``memory_skip_reason``) +* Rust-portable shape mirrors :class:`CompressionDecision` exactly + +Precedence (highest first): + 1. ``bypass_header`` — user opt-out + 2. ``image_optimize_disabled`` — operator config off + 3. ``no_messages`` — nothing to inspect +""" + +from __future__ import annotations + +from dataclasses import FrozenInstanceError +from types import SimpleNamespace +from typing import Any + +from headroom.proxy.image_compression_decision import ImageCompressionDecision + +# ── Helpers ─────────────────────────────────────────────────────────── + + +def _config(*, image_optimize: bool = True) -> Any: + """Minimal stand-in for ``HeadroomConfig`` — only the field the + decision reads.""" + return SimpleNamespace(image_optimize=image_optimize) + + +def _msgs(n: int = 1) -> list[dict[str, str]]: + return [{"role": "user", "content": f"hi-{i}"} for i in range(n)] + + +# ── Value-type contract ─────────────────────────────────────────────── + + +def test_decision_is_frozen() -> None: + d = ImageCompressionDecision.decide(headers={}, config=_config(), messages=_msgs()) + try: + d.should_compress = False # type: ignore[misc] + except FrozenInstanceError: + pass + else: + raise AssertionError("ImageCompressionDecision must be frozen") + + +def test_decision_is_value_equal() -> None: + a = ImageCompressionDecision.decide(headers={}, config=_config(), messages=_msgs()) + b = ImageCompressionDecision.decide(headers={}, config=_config(), messages=_msgs()) + assert a == b + + +# ── Precedence ──────────────────────────────────────────────────────── + + +def test_compresses_when_every_gate_open() -> None: + d = ImageCompressionDecision.decide( + headers={}, config=_config(image_optimize=True), messages=_msgs() + ) + assert d.should_compress is True + assert d.passthrough_reason is None + + +def test_bypass_header_wins() -> None: + """Bypass is the user's explicit "don't touch my bytes" signal — + image compression mutates bytes (tile-aligns / re-encodes), so + bypass must skip it. Mirror of CompressionDecision.""" + d = ImageCompressionDecision.decide( + headers={"x-headroom-bypass": "true"}, + config=_config(image_optimize=True), + messages=_msgs(), + ) + assert d.should_compress is False + assert d.passthrough_reason == "bypass_header" + + +def test_passthrough_mode_header_also_triggers_bypass_skip() -> None: + """``x-headroom-mode: passthrough`` alt spelling — mirrors + ``_headroom_bypass_enabled`` semantics across all proxy gates.""" + d = ImageCompressionDecision.decide( + headers={"x-headroom-mode": "passthrough"}, + config=_config(image_optimize=True), + messages=_msgs(), + ) + assert d.should_compress is False + assert d.passthrough_reason == "bypass_header" + + +def test_image_optimize_disabled_is_skip() -> None: + """Operator config — ``config.image_optimize = False``. Distinct + reason from ``compression_disabled`` (text compression's gate); + operators can enable text + disable image independently.""" + d = ImageCompressionDecision.decide( + headers={}, config=_config(image_optimize=False), messages=_msgs() + ) + assert d.should_compress is False + assert d.passthrough_reason == "image_optimize_disabled" + + +def test_no_messages_is_skip() -> None: + """Empty or missing messages — nothing to look at. Same shape as + CompressionDecision's no_messages reason.""" + d = ImageCompressionDecision.decide(headers={}, config=_config(), messages=[]) + assert d.should_compress is False + assert d.passthrough_reason == "no_messages" + + +def test_messages_none_is_skip() -> None: + """``messages=None`` is treated identically to empty list.""" + d = ImageCompressionDecision.decide(headers={}, config=_config(), messages=None) + assert d.should_compress is False + assert d.passthrough_reason == "no_messages" + + +# ── Precedence ordering ────────────────────────────────────────────── + + +def test_bypass_beats_image_optimize_disabled() -> None: + """User signal beats operator signal — bypass is the more + informative dashboard slice.""" + d = ImageCompressionDecision.decide( + headers={"x-headroom-bypass": "true"}, + config=_config(image_optimize=False), + messages=_msgs(), + ) + assert d.passthrough_reason == "bypass_header" + + +def test_bypass_beats_no_messages() -> None: + """Bypass+no-messages surfaces bypass — user opted out, the + empty body is incidental.""" + d = ImageCompressionDecision.decide( + headers={"x-headroom-bypass": "true"}, + config=_config(), + messages=[], + ) + assert d.passthrough_reason == "bypass_header" + + +def test_image_optimize_disabled_beats_no_messages() -> None: + """When config is off AND messages empty, surface the operator + decision — more meaningful for dashboards.""" + d = ImageCompressionDecision.decide( + headers={}, config=_config(image_optimize=False), messages=[] + ) + assert d.passthrough_reason == "image_optimize_disabled" + + +# ── Observability fields ───────────────────────────────────────────── + + +def test_observability_booleans_populated_when_compressing() -> None: + d = ImageCompressionDecision.decide( + headers={}, config=_config(image_optimize=True), messages=_msgs(2) + ) + assert d.bypass_header_set is False + assert d.image_optimize_enabled is True + assert d.has_messages is True + + +def test_observability_booleans_populated_when_passthrough() -> None: + d = ImageCompressionDecision.decide( + headers={"x-headroom-bypass": "true"}, + config=_config(image_optimize=True), + messages=_msgs(), + ) + assert d.bypass_header_set is True + assert d.image_optimize_enabled is True + assert d.has_messages is True + + +# ── apply_to_tags ──────────────────────────────────────────────────── + + +def test_apply_to_tags_stamps_reason_when_passthrough() -> None: + d = ImageCompressionDecision.decide( + headers={"x-headroom-bypass": "true"}, config=_config(), messages=_msgs() + ) + tags: dict[str, str] = {} + d.apply_to_tags(tags) + assert tags == {"image_skip_reason": "bypass_header"} + + +def test_apply_to_tags_is_a_noop_when_compressing() -> None: + d = ImageCompressionDecision.decide(headers={}, config=_config(), messages=_msgs()) + tags: dict[str, str] = {"client": "codex"} + d.apply_to_tags(tags) + assert tags == {"client": "codex"} + + +def test_apply_to_tags_preserves_pre_existing_entries() -> None: + """Image skip reason coexists with other slicing tags (client, + passthrough_reason, memory_skip_reason) — they all live in the + same RequestOutcome.tags dict.""" + d = ImageCompressionDecision.decide( + headers={}, config=_config(image_optimize=False), messages=_msgs() + ) + tags: dict[str, str] = { + "client": "claude-code", + "passthrough_reason": "compression_disabled", + "memory_skip_reason": "no_user_id", + } + d.apply_to_tags(tags) + assert tags["client"] == "claude-code" + assert tags["passthrough_reason"] == "compression_disabled" + assert tags["memory_skip_reason"] == "no_user_id" + assert tags["image_skip_reason"] == "image_optimize_disabled" + + +def test_apply_to_tags_for_every_skip_reason() -> None: + """Every reason name round-trips cleanly into the tag dict.""" + cases: dict[str, dict[str, Any]] = { + "bypass_header": { + "headers": {"x-headroom-bypass": "true"}, + "config": _config(), + "messages": _msgs(), + }, + "image_optimize_disabled": { + "headers": {}, + "config": _config(image_optimize=False), + "messages": _msgs(), + }, + "no_messages": { + "headers": {}, + "config": _config(), + "messages": [], + }, + } + for expected_reason, kwargs in cases.items(): + d = ImageCompressionDecision.decide(**kwargs) + tags: dict[str, str] = {} + d.apply_to_tags(tags) + assert tags.get("image_skip_reason") == expected_reason, expected_reason diff --git a/tests/test_memory_ranker.py b/tests/test_memory_ranker.py new file mode 100644 index 000000000..eacfcabb2 --- /dev/null +++ b/tests/test_memory_ranker.py @@ -0,0 +1,243 @@ +"""Tests for :class:`headroom.proxy.memory_ranker.MemoryRanker` + +:class:`RecencyBoostRanker`. + +Pre-this-PR Headroom ranked memory candidates by pure cosine +similarity. Every other memory system we surveyed (Letta, Mem0, +Cognee, Supermemory) re-ranks beyond cosine — recency / source / +access-count / decay are table-stakes. The pure-cosine baseline +returns 6-month-old memories with 0.9 similarity ahead of fresh +memories with 0.5 — wrong for most use cases. + +``RecencyBoostRanker`` is the first ranker we ship: a pure-function +``score = cosine × exp(-age_days / decay_days)`` re-ranker. Default +``decay_days=30`` (half-life ~21 days). Other rankers (source-weight, +access-count) plug into the same :class:`MemoryRanker` protocol in +follow-on PRs. + +Performance: O(N) over candidates where N = top_k = ~10. One ``exp()`` +per candidate. Sub-microsecond per request — no embedding compute, no +I/O. The ranker is pure and Rust-portable. +""" + +from __future__ import annotations + +from dataclasses import FrozenInstanceError +from datetime import datetime, timedelta, timezone + +from headroom.proxy.memory_ranker import ( + MemoryCandidate, + RecencyBoostRanker, +) + +_UTC = timezone.utc + +# ── Helpers ─────────────────────────────────────────────────────────── + + +def _candidate(content: str, score: float, age_days: float = 0.0) -> MemoryCandidate: + """Build a MemoryCandidate at the given cosine score and age.""" + created = datetime.now(_UTC) - timedelta(days=age_days) + return MemoryCandidate(content=content, score=score, created_at=created) + + +def _candidate_no_timestamp(content: str, score: float) -> MemoryCandidate: + """Build a MemoryCandidate without a created_at (back-compat shape).""" + return MemoryCandidate(content=content, score=score, created_at=None) + + +# ── MemoryCandidate value-type contract ────────────────────────────── + + +def test_candidate_is_frozen() -> None: + """Frozen so a ranker can't mutate a candidate's score and lie about + which candidates it returned.""" + c = _candidate("x", 0.9, age_days=0) + try: + c.score = 0.1 # type: ignore[misc] + except FrozenInstanceError: + pass + else: + raise AssertionError("MemoryCandidate must be frozen") + + +# ── RecencyBoostRanker contract ────────────────────────────────────── + + +def test_ranker_is_frozen() -> None: + """The ranker config is itself immutable — operators set + ``decay_days`` at construction; runtime cannot edit it.""" + r = RecencyBoostRanker() + try: + r.decay_days = 99 # type: ignore[misc] + except FrozenInstanceError: + pass + else: + raise AssertionError("RecencyBoostRanker must be frozen") + + +def test_ranker_default_decay_is_thirty_days() -> None: + """30-day decay is the conservative default. At 30 days, factor is + ~0.37 (e^{-1}); at 90 days, ~0.05. Tuned so a fresh memory with + weak cosine doesn't dominate, but a 6-month-old strong-cosine + can't dominate either.""" + assert RecencyBoostRanker().decay_days == 30.0 + + +def test_ranker_default_decay_is_configurable() -> None: + """Operators can tune decay; e.g., 7 days for an aggressive + recency bias on rapidly-evolving codebases.""" + r = RecencyBoostRanker(decay_days=7.0) + assert r.decay_days == 7.0 + + +def test_ranker_returns_list_preserving_shape() -> None: + """Output is a list of candidates (re-ranked). Length matches + input length — the ranker does NOT filter, only re-orders. The + budget filters; the ranker ranks.""" + candidates = [_candidate("a", 0.9), _candidate("b", 0.5)] + out = RecencyBoostRanker().rank(candidates) + assert len(out) == 2 + assert {c.content for c in out} == {"a", "b"} + + +# ── Recency boost behaviour ────────────────────────────────────────── + + +def test_equal_cosine_younger_wins() -> None: + """Two candidates with identical cosine score — the younger one + wins because its recency factor is closer to 1.0.""" + fresh = _candidate("fresh", 0.5, age_days=0) + old = _candidate("old", 0.5, age_days=60) + out = RecencyBoostRanker().rank([old, fresh]) + assert out[0].content == "fresh" + assert out[1].content == "old" + + +def test_old_strong_cosine_can_still_beat_young_weak_cosine() -> None: + """The boost is multiplicative, not absolute — a 60-day-old memory + with 0.9 cosine (0.9 × 0.135 ≈ 0.12) still loses to a 0-day-old + memory with 0.5 cosine (0.5 × 1.0 = 0.5). But a 5-day-old memory + with 0.9 (0.9 × 0.847 ≈ 0.76) beats a 0-day-old with 0.5.""" + very_old_strong = _candidate("old_strong", 0.9, age_days=60) + fresh_weak = _candidate("fresh_weak", 0.5, age_days=0) + out = RecencyBoostRanker().rank([very_old_strong, fresh_weak]) + # fresh_weak should win because 60-day decay flattens the strong cosine + assert out[0].content == "fresh_weak" + + # Versus: slightly-old strong beats fresh weak + slightly_old_strong = _candidate("slightly_old_strong", 0.9, age_days=5) + fresh_weak2 = _candidate("fresh_weak2", 0.5, age_days=0) + out2 = RecencyBoostRanker().rank([fresh_weak2, slightly_old_strong]) + assert out2[0].content == "slightly_old_strong" + + +def test_decay_rate_changes_winner() -> None: + """An aggressive decay_days=7 makes a 30-day-old memory much + weaker than a default decay_days=30. Locks the configurability + contract.""" + old_strong = _candidate("old_strong", 0.9, age_days=30) + fresh_weak = _candidate("fresh_weak", 0.6, age_days=0) + + # decay_days=30: old × e^{-1} ≈ 0.331; fresh = 0.6 → fresh wins + r_default = RecencyBoostRanker(decay_days=30.0) + out_default = r_default.rank([old_strong, fresh_weak]) + assert out_default[0].content == "fresh_weak" + + # decay_days=120 (loose): old × e^{-0.25} ≈ 0.701; fresh = 0.6 → old wins + r_loose = RecencyBoostRanker(decay_days=120.0) + out_loose = r_loose.rank([old_strong, fresh_weak]) + assert out_loose[0].content == "old_strong" + + +def test_zero_age_memory_keeps_full_cosine() -> None: + """At age=0 days, the recency factor is e^0 = 1.0 — the boosted + score equals the original cosine. Fresh memories see no penalty.""" + fresh = _candidate("fresh", 0.7, age_days=0) + out = RecencyBoostRanker().rank([fresh]) + # Compare with tolerance — datetime.now() drift between + # _candidate() and rank() is microseconds, so factor ~ 1.0. + assert out[0].score == 0.7 or abs(out[0].score - 0.7) < 1e-3 + + +def test_candidate_without_timestamp_keeps_pure_cosine() -> None: + """Backwards-compat: pre-this-PR candidates may not have a + ``created_at`` (older rows / older backends). NULL timestamp + means "treat as recency-neutral" — factor 1.0. Pure cosine.""" + no_ts = _candidate_no_timestamp("legacy", 0.8) + out = RecencyBoostRanker().rank([no_ts]) + assert out[0].score == 0.8 + + +def test_mixed_with_and_without_timestamps() -> None: + """A backend that returns SOME candidates with timestamps and + SOME without (e.g., during a migration) must still produce a + sensible ranking. NULL-timestamp candidates get factor 1.0, + timestamped ones get their decay.""" + fresh_ts = _candidate("fresh_ts", 0.6, age_days=0) + old_ts = _candidate("old_ts", 0.6, age_days=60) + no_ts_neutral = _candidate_no_timestamp("no_ts", 0.6) + out = RecencyBoostRanker().rank([old_ts, no_ts_neutral, fresh_ts]) + # fresh_ts (~0.6) and no_ts (=0.6) tied at top — old_ts decayed. + assert out[-1].content == "old_ts" + assert {out[0].content, out[1].content} == {"fresh_ts", "no_ts"} + + +# ── Stability + edge cases ─────────────────────────────────────────── + + +def test_empty_input_returns_empty_output() -> None: + """No candidates → no candidates. Boundary case.""" + assert RecencyBoostRanker().rank([]) == [] + + +def test_ranking_is_stable_for_identical_candidates() -> None: + """Two candidates with identical content + score + age → stable + order (no spurious reshuffling). Important for prefix-cache + stability: a deterministic ranker means consecutive turns inject + the same memory in the same order, preserving byte-stable + output.""" + a = _candidate("same", 0.5, age_days=10) + b = _candidate("same", 0.5, age_days=10) + out = RecencyBoostRanker().rank([a, b]) + assert len(out) == 2 + + +def test_negative_age_treated_as_zero() -> None: + """Defensive: a candidate with a future ``created_at`` (clock + skew) shouldn't crash or give a > 1.0 factor. ``exp(-age/decay)`` + with negative age gives > 1; we clamp to 1.0 so a clock-skewed + candidate can't outrank a real fresh one with score=1.0 + artifically.""" + future = _candidate("future", 0.5, age_days=-10) # 10 days in future + fresh = _candidate("fresh", 0.5, age_days=0) + out = RecencyBoostRanker().rank([future, fresh]) + # Both should have factor 1.0 (clamped) — score equal → stable order + assert {out[0].content, out[1].content} == {"future", "fresh"} + assert out[0].score == 0.5 or abs(out[0].score - 0.5) < 1e-3 + + +# ── Rust-port shape ───────────────────────────────────────────────── + + +def test_ranker_is_pure_no_side_effects() -> None: + """Calling rank() twice with the same input gives the same + output. No state on the ranker; no I/O. Rust-portable.""" + candidates = [_candidate("a", 0.7, age_days=5), _candidate("b", 0.5, age_days=20)] + r = RecencyBoostRanker() + out1 = r.rank(candidates) + out2 = r.rank(candidates) + assert [c.content for c in out1] == [c.content for c in out2] + # Inputs preserved — ranker did not mutate + assert candidates[0].content == "a" + assert candidates[1].content == "b" + + +def test_ranker_does_not_mutate_input_list() -> None: + """Defence-in-depth: the input list and its elements must be + unchanged after ranking. Frozen candidates make element mutation + impossible; the list order itself must also be preserved.""" + a = _candidate("a", 0.5, age_days=20) + b = _candidate("b", 0.5, age_days=5) + candidates = [a, b] + RecencyBoostRanker().rank(candidates) + assert candidates == [a, b] # original list order preserved