headroom/tests/test_memory_decision.py
chopratejas 71d5a7b545 fix(proxy): MemoryDecision contract + 3 bypass bugs + drop 500-char query cap
Three bug classes fixed plus three architectural extension points,
together making the memory subsystem uniform across all five sites
and ready for future Mem0/Letta/Cognee backend integration.

## Bug fixes

* **3 sites silently ignored `x-headroom-bypass: true`** —
  ``anthropic.py:1303``, ``openai.py:1620`` (chat), ``gemini.py:382``
  injected memory under bypass, mutating request bytes when the user
  explicitly asked for byte-faithful passthrough. Now gated on
  ``MemoryDecision.decide(...)`` which honours bypass uniformly.

* **500-char query truncation** — ``memory_handler._extract_user_query``
  capped at 500 chars, silently throwing away signal. None of Letta /
  Mem0 / Cognee / Supermemory truncate. Removed; the embedding model
  handles its own window.

* **Gemini had no timeout** on ``search_and_format_context`` — the
  only chat handler without one. A slow backend could stall requests.
  Added ``asyncio.wait_for`` matching Anthropic + OpenAI Chat +
  Responses.

* **WS injected into ``body["instructions"]``** — the system /
  cache-hot-zone field, violating invariant I2 (all other handlers
  inject at user-message tail). Switched to ``ws_response_body["input"]``
  for string-shaped input; list-shaped input deferred to the Rust
  handler with a clear log.

## New value types (extension points)

* ``MemoryDecision`` — frozen dataclass + factory. Five-way skip
  reason enum (``bypass_header`` / ``no_handler`` / ``no_user_id`` /
  ``mode_disabled`` / ``mode_tool``). ``apply_to_tags()`` surfaces
  the skip reason in ``RequestOutcome.tags["memory_skip_reason"]``
  — dashboards can now slice memory-blind traffic by cause.

* ``MemoryQuery`` — multi-source retrieval query. ``from_messages()``
  walks the conversation and extracts latest user text + recent tool
  outputs + recent assistant turns at FULL fidelity (no truncation).
  Handles both OpenAI-shape ``role: tool`` and Anthropic-shape
  ``tool_result`` content blocks. ``to_embedding_input()`` produces
  a delimited concatenation the embedder sees as structured context.

* ``MemoryInjectionBudget`` — uniform token / entry / similarity
  bound on the formatted injection block. Pre-this-PR no cap (~4000
  tokens could land per request). Default 1024 tokens / 10 entries /
  0.3 similarity floor. ``apply_to_text()`` truncates at line
  boundaries so dashboard renders intact bullet points.

## Migration scope — all 5 sites uniform at the GATE level

| Site | Handler | Pre-PR gate | Post-PR gate |
|---|---|---|---|
| 1 | anthropic.py | `memory_handler and memory_user_id` | `memory_decision.inject` |
| 2 | gemini.py | `memory_handler and memory_user_id` | `memory_decision.inject` |
| 3 | openai.py chat | `memory_handler and memory_user_id` | `memory_decision.inject` |
| 4 | openai.py Responses | `memory_handler and memory_user_id and not _bypass` | `responses_memory_decision.inject` |
| 6 | openai.py WS | `memory_handler and body and not _ws_bypass` | `ws_memory_decision.inject` |

Site 5 (Responses bypass-elif log-only branch) is preserved verbatim.

## Deliberately deferred (separate PRs)

* **Memory injection order inversion** — sites 4 and 6 inject
  BEFORE compression; sites 1/2/3 inject AFTER. Moving 4 + 6 to
  post-compression needs its own focused cache-stability testing.
* **Importance scoring** — recency × source × access-count.
* **Per-memory atomize-and-split** — Mem0/Supermemory pattern.
* **AST-aware code chunking for tool outputs** — Supermemory's
  code-chunk approach.

The contracts shipped here (``MemoryQuery`` + ``MemoryInjectionBudget``)
are the extension points those will plug into.

## Test coverage

* 20 new tests on ``MemoryDecision``
* 14 new tests on ``MemoryQuery`` (full-fidelity, multi-source)
* 10 new tests on ``MemoryInjectionBudget``
* 3 new AST contract tests (no raw gate; no system writes; every
  search call passes ``query=``)
* All existing memory + cache-stability tests still pass (222 passed)

## Rust portability

Every new value type ports cleanly to a frozen Rust struct. Pure
functions, no I/O, no global state. Same Python ↔ Rust parity-test
pattern that ``CompressionDecision`` already uses.

## Zero-regression contract

Existing chat/completion harnesses (Claude Code, Codex, Cursor,
Continue, Aider) see ZERO wire-byte changes when bypass is NOT set.
When bypass IS set, the 3 chat handlers now correctly skip memory
injection — that's the bug fix, not a regression.
2026-05-19 11:13:52 -05:00

314 lines
11 KiB
Python

"""Tests for :class:`headroom.proxy.memory_decision.MemoryDecision`.
The point of this file is the *contract* — every behavioural assertion
here is the canonical answer to "should this request have memory
context injected into it?" that today is computed inline across the
proxy's handlers with subtle drift.
Specifically locks (post-PR-this):
* Sites 1/2/3 (Anthropic ``/v1/messages``, Gemini
``:generateContent``, OpenAI ``/v1/chat/completions``) MUST gate on
bypass — pre-this-PR they didn't, so memory injection silently
mutated requests under ``x-headroom-bypass: true``.
* Site 4 (OpenAI ``/v1/responses``) already gated; locked here.
* Site 6 (OpenAI WS ``/v1/responses``) already gated; locked here.
* ``MemoryMode`` values (``auto_tail`` / ``tool``) and the env-driven
``HEADROOM_MEMORY_INJECTION_MODE`` (``disabled`` / ``auto_tail`` /
``tool``) all surface as explicit ``skip_reason`` values, not as
hidden conditional code.
The decision is **input-side only** — it gates whether mutation of the
request bytes happens. It does NOT gate background memory STORAGE
(traffic-learner runs on a separate path and is unaffected by this
decision).
"""
from __future__ import annotations
from dataclasses import FrozenInstanceError
from types import SimpleNamespace
from typing import Any
from headroom.proxy.memory_decision import MemoryDecision
def _memory_handler() -> Any:
"""Minimal stand-in for the proxy's ``self.memory_handler``."""
return SimpleNamespace(name="local")
# ── Value-type contract ───────────────────────────────────────────────
def test_decision_is_frozen() -> None:
"""Frozen dataclass — mutation would let a handler patch the
decision after handing it to the funnel."""
d = MemoryDecision.decide(
headers={}, memory_handler=_memory_handler(), memory_user_id="u1", mode_name="auto_tail"
)
try:
d.inject = False # type: ignore[misc]
except FrozenInstanceError:
pass
else:
raise AssertionError("MemoryDecision must be frozen")
def test_decision_is_value_equal() -> None:
"""Two decisions from identical inputs compare equal."""
a = MemoryDecision.decide(
headers={}, memory_handler=_memory_handler(), memory_user_id="u1", mode_name="auto_tail"
)
b = MemoryDecision.decide(
headers={}, memory_handler=_memory_handler(), memory_user_id="u1", mode_name="auto_tail"
)
assert a == b
# ── Precedence: bypass > no_handler > no_user_id > mode > INJECT ─────
def test_injects_when_every_gate_open() -> None:
"""Happy path: no bypass, handler wired, user_id set, mode=auto_tail."""
d = MemoryDecision.decide(
headers={}, memory_handler=_memory_handler(), memory_user_id="u1", mode_name="auto_tail"
)
assert d.inject is True
assert d.skip_reason is None
def test_bypass_header_wins_over_every_other_gate() -> None:
"""``x-headroom-bypass: true`` is the user's "do not touch my
bytes" signal — highest priority, even when memory is otherwise
fully wired. Memory injection mutates the request bytes; bypass
must skip it. (This was the 3-bug Gemini-class problem pre-PR
on Anthropic, OpenAI chat, Gemini.)"""
d = MemoryDecision.decide(
headers={"x-headroom-bypass": "true"},
memory_handler=_memory_handler(),
memory_user_id="u1",
mode_name="auto_tail",
)
assert d.inject is False
assert d.skip_reason == "bypass_header"
def test_passthrough_mode_header_also_triggers_bypass_skip() -> None:
"""``x-headroom-mode: passthrough`` is the alternate spelling of
the bypass signal — mirrors _headroom_bypass_enabled semantics."""
d = MemoryDecision.decide(
headers={"x-headroom-mode": "passthrough"},
memory_handler=_memory_handler(),
memory_user_id="u1",
mode_name="auto_tail",
)
assert d.inject is False
assert d.skip_reason == "bypass_header"
def test_no_handler_is_skip() -> None:
"""No memory backend configured → ``no_handler`` reason."""
d = MemoryDecision.decide(
headers={}, memory_handler=None, memory_user_id="u1", mode_name="auto_tail"
)
assert d.inject is False
assert d.skip_reason == "no_handler"
def test_no_user_id_is_skip() -> None:
"""Memory wired but per-request user_id missing → ``no_user_id``."""
d = MemoryDecision.decide(
headers={}, memory_handler=_memory_handler(), memory_user_id=None, mode_name="auto_tail"
)
assert d.inject is False
assert d.skip_reason == "no_user_id"
def test_empty_user_id_string_is_skip() -> None:
"""Empty string for user_id is treated as missing."""
d = MemoryDecision.decide(
headers={}, memory_handler=_memory_handler(), memory_user_id="", mode_name="auto_tail"
)
assert d.inject is False
assert d.skip_reason == "no_user_id"
def test_mode_disabled_is_skip() -> None:
"""Operator override via HEADROOM_MEMORY_INJECTION_MODE=disabled."""
d = MemoryDecision.decide(
headers={}, memory_handler=_memory_handler(), memory_user_id="u1", mode_name="disabled"
)
assert d.inject is False
assert d.skip_reason == "mode_disabled"
def test_mode_tool_is_skip() -> None:
"""TOOL mode: auto-injection disabled (the agent calls memory
tools explicitly instead). Distinct from disabled."""
d = MemoryDecision.decide(
headers={}, memory_handler=_memory_handler(), memory_user_id="u1", mode_name="tool"
)
assert d.inject is False
assert d.skip_reason == "mode_tool"
# ── Precedence ordering when multiple gates close ────────────────────
def test_bypass_beats_no_handler() -> None:
"""When both bypass AND no_handler would skip, surface bypass —
user's explicit signal is the more informative dashboard slice."""
d = MemoryDecision.decide(
headers={"x-headroom-bypass": "true"},
memory_handler=None,
memory_user_id=None,
mode_name="disabled",
)
assert d.skip_reason == "bypass_header"
def test_no_handler_beats_no_user_id() -> None:
"""no_handler is the more fundamental failure (no backend wired);
no_user_id only matters when a handler exists."""
d = MemoryDecision.decide(
headers={}, memory_handler=None, memory_user_id=None, mode_name="auto_tail"
)
assert d.skip_reason == "no_handler"
def test_no_user_id_beats_mode() -> None:
"""A request with no user_id can't be served regardless of mode."""
d = MemoryDecision.decide(
headers={}, memory_handler=_memory_handler(), memory_user_id=None, mode_name="disabled"
)
assert d.skip_reason == "no_user_id"
def test_mode_disabled_beats_mode_tool() -> None:
"""``disabled`` is operator-level kill; ``tool`` is mode pref.
Disabled is the more emphatic signal."""
# Construct two separate decisions on identical inputs but mode varying;
# check that each produces its own reason (no cross-contamination).
d_dis = MemoryDecision.decide(
headers={}, memory_handler=_memory_handler(), memory_user_id="u1", mode_name="disabled"
)
d_tool = MemoryDecision.decide(
headers={}, memory_handler=_memory_handler(), memory_user_id="u1", mode_name="tool"
)
assert d_dis.skip_reason == "mode_disabled"
assert d_tool.skip_reason == "mode_tool"
# ── Observability fields ─────────────────────────────────────────────
def test_observability_booleans_populated_when_injecting() -> None:
"""Even on the happy path, every constituent is exposed so debug
tooling can answer "what did the decision see?" without re-running."""
d = MemoryDecision.decide(
headers={}, memory_handler=_memory_handler(), memory_user_id="u1", mode_name="auto_tail"
)
assert d.bypass_header_set is False
assert d.memory_handler_present is True
assert d.memory_user_id_present is True
assert d.mode_name == "auto_tail"
def test_observability_booleans_populated_when_skipping() -> None:
"""Same on the skip path — every constituent must be visible."""
d = MemoryDecision.decide(
headers={"x-headroom-bypass": "true"},
memory_handler=None,
memory_user_id="u1",
mode_name="auto_tail",
)
assert d.bypass_header_set is True
assert d.memory_handler_present is False
assert d.memory_user_id_present is True
assert d.mode_name == "auto_tail"
# ── apply_to_tags — mirror CompressionDecision pattern ───────────────
def test_apply_to_tags_stamps_reason_when_skipping() -> None:
"""Skip decisions surface ``memory_skip_reason`` in tags so the
dashboard can slice memory-blind traffic by cause."""
d = MemoryDecision.decide(
headers={"x-headroom-bypass": "true"},
memory_handler=_memory_handler(),
memory_user_id="u1",
mode_name="auto_tail",
)
tags: dict[str, str] = {}
d.apply_to_tags(tags)
assert tags == {"memory_skip_reason": "bypass_header"}
def test_apply_to_tags_is_a_noop_when_injecting() -> None:
"""No tag when injecting — absence is the signal for "memory was
used". Avoids spurious ``memory_skip_reason=None`` strings."""
d = MemoryDecision.decide(
headers={}, memory_handler=_memory_handler(), memory_user_id="u1", mode_name="auto_tail"
)
tags: dict[str, str] = {"client": "codex"}
d.apply_to_tags(tags)
assert tags == {"client": "codex"}
assert "memory_skip_reason" not in tags
def test_apply_to_tags_preserves_pre_existing_entries() -> None:
"""Existing tags (client, passthrough_reason from CompressionDecision)
must survive unchanged."""
d = MemoryDecision.decide(
headers={}, memory_handler=None, memory_user_id="u1", mode_name="auto_tail"
)
tags: dict[str, str] = {"client": "claude-code", "passthrough_reason": "bypass_header"}
d.apply_to_tags(tags)
assert tags == {
"client": "claude-code",
"passthrough_reason": "bypass_header",
"memory_skip_reason": "no_handler",
}
def test_apply_to_tags_for_every_skip_reason() -> None:
"""Every skip reason name must round-trip through apply_to_tags."""
cases: dict[str, dict[str, Any]] = {
"bypass_header": {
"headers": {"x-headroom-bypass": "true"},
"memory_handler": _memory_handler(),
"memory_user_id": "u1",
"mode_name": "auto_tail",
},
"no_handler": {
"headers": {},
"memory_handler": None,
"memory_user_id": "u1",
"mode_name": "auto_tail",
},
"no_user_id": {
"headers": {},
"memory_handler": _memory_handler(),
"memory_user_id": None,
"mode_name": "auto_tail",
},
"mode_disabled": {
"headers": {},
"memory_handler": _memory_handler(),
"memory_user_id": "u1",
"mode_name": "disabled",
},
"mode_tool": {
"headers": {},
"memory_handler": _memory_handler(),
"memory_user_id": "u1",
"mode_name": "tool",
},
}
for expected, kwargs in cases.items():
d = MemoryDecision.decide(**kwargs)
tags: dict[str, str] = {}
d.apply_to_tags(tags)
assert tags.get("memory_skip_reason") == expected, expected