refactor(proxy): MemoryRanker + ImageCompressionDecision + branch-aware version-sync

Three independent contract-pattern follow-ons bundled into one PR.
Same frozen-dataclass + factory + apply_to_tags + Rust-portable
shape that PR #473 / #477 / #483 established.

## (1) MemoryRanker + 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.

* ``MemoryRanker`` Protocol — pluggable re-ranker; future PRs add
  source-weight + access-count rankers behind the same interface.
* ``RecencyBoostRanker`` — first concrete impl. Final score is
  ``cosine × exp(-age_days / decay_days)``. Default decay 30 days
  (half-life ~21 days; 60-day-old factor 0.135, 90-day-old 0.050).
* ``MemoryCandidate`` — backend-agnostic frozen value type that
  flows through the ranker. ``MemoryCandidate.from_backend_result``
  adapter converts the existing ``MemoryResult`` shape (with nested
  ``memory.created_at``) into the ranker's flatter form.
* Wired into ``memory_handler.search_and_format_context`` as an
  optional ``ranker=`` kwarg — backwards-compat: ``None`` (default)
  preserves the pure-cosine path identically.

Defensive:
* ``created_at=None`` → factor 1.0 (recency-neutral, back-compat with
  legacy rows / migrating backends)
* Negative age (clock skew) → clamped to factor 1.0 (a future-dated
  row can't outrank a real fresh memory)
* Sort is stable on ties — same input → same output every turn, so
  consecutive turns inject memories in the same order (prefix-cache
  friendly)

Performance: O(N) over candidates where N=top_k≈10. One ``math.exp``
per candidate. Sub-microsecond. Zero new I/O.

## (2) ImageCompressionDecision

Mirror of :class:`CompressionDecision` for image compression. Two
sites today (``openai.py:1203``, ``anthropic.py:868``) gate inline;
both already respect bypass (no Gemini-class drift bug like text
compression had), but consolidating into a value type:

* Locks bypass-respect via AST contract test — future sites can't
  drift on it
* Surfaces ``image_skip_reason`` in ``RequestOutcome.tags`` for
  dashboard slicing (same observability surface as
  ``passthrough_reason`` and ``memory_skip_reason``)
* Same Rust-port shape as the other decision types

Precedence: ``bypass_header`` > ``image_optimize_disabled`` >
``no_messages`` > ``should_compress=True``.

Anthropic's extra ``is_cache_mode`` check stays inline because it's
Anthropic-specific (openai/gemini don't have it). Documented in a
code comment.

## (3) Branch-aware sync-plugin-versions hook

Pre-this-fix the pre-commit ``sync-plugin-versions`` hook ran on
every commit and bumped manifests to the predicted-next-release
version. Every PR ended up carrying the prediction as collateral
("Why are we bumping ``.claude-plugin/marketplace.json`` — we
should not, right??" - user, on PR #483).

Fix: the hook is now a NO-OP unless EITHER:
* We're on the ``main`` branch, OR
* ``HEADROOM_SYNC_VERSIONS=1`` is set explicitly (release workflow)

On feature branches the hook prints a single line explaining the
skip and exits cleanly. The release workflow opts in via the env
var; behaviour on main / at release time is unchanged.

## Test coverage

* 16 new tests on ``MemoryRanker`` / ``RecencyBoostRanker``
  (frozen, equal cosine wins by recency, decay configurable, NULL
  timestamp neutral, no-mutation contract, Rust-port shape)
* 17 new tests on ``ImageCompressionDecision`` (frozen, all 3
  skip reasons, precedence, observability fields, apply_to_tags)
* 1 new AST invariant test (extends
  ``test_handler_outcome_tag_invariant.py``) — locks "no raw
  ``if self.config.image_optimize and messages and not _bypass:``
  conjunction in any handler"

All existing memory + cache-stability + handler tests pass (203 ✓).
``make ci-precheck`` clean.

## Rust portability

All three new value types port cleanly to frozen Rust structs +
pure functions. Same migration pattern as ``CompressionDecision``
(already locked in for the SmartCrusher Rust port).

## Zero-regression contract

* Default ``ranker=None`` → memory_handler behaves identically to
  pre-this-PR (pure cosine; no perf change)
* Image decision migration is identity at the bypass/optimize/messages
  gate — no behaviour change, just contract consolidation
* Hook fix is no-op on feature branches (less churn) and unchanged
  on main (release flow preserved)
This commit is contained in:
chopratejas 2026-05-19 12:13:09 -05:00
parent 10580439bb
commit a7b197c6ec
9 changed files with 982 additions and 24 deletions

View file

@ -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()

View file

@ -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

View file

@ -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

View file

@ -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}")

View file

@ -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)

View file

@ -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 "<unknown>"
print(
f"sync-plugin-versions: skipping on branch '{branch}' (set HEADROOM_SYNC_VERSIONS=1 to force)"
)
return
version = compute_repo_semver(root)
subprocess.run(
[

View file

@ -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."
)

View file

@ -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

243
tests/test_memory_ranker.py Normal file
View file

@ -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