akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(proxy): MemoryDecision contract + 3 bypass bugs + drop 500-char query cap

Browse source 
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.
This commit is contained in:
chopratejas 2026-05-19 10:04:21 -05:00
parent 10580439bb
commit 71d5a7b545
11 changed files with 1500 additions and 49 deletions
Download patch file Download diff file Expand all files Collapse all files

29
headroom/proxy/handlers/anthropic.py
View file

@ -27,6 +27,8 @@ from headroom.pipeline import PipelineStage, summarize_routing_markers
from headroom.proxy.auth_mode import classify_auth_mode, classify_client
from headroom.proxy.compression_decision import CompressionDecision
from headroom.proxy.helpers import extract_tags
from headroom.proxy.memory_decision import MemoryDecision
from headroom.proxy.memory_query import MemoryQuery
from headroom.proxy.outcome import RequestOutcome
logger = logging.getLogger("headroom.proxy")
@ -688,6 +690,23 @@ class AnthropicHandlerMixin:
),
)
# Canonical memory-injection gate. Reads `request.headers`
# so bypass detection sees the original inbound (the local
# `headers` dict was stripped of x-headroom-* above).
# Replaces the pre-PR-this raw `if self.memory_handler and
# memory_user_id:` conjunction that silently ignored
# `x-headroom-bypass: true` and mutated request bytes
# under the user's "don't touch my bytes" signal.
from headroom.proxy.helpers import get_memory_injection_mode
memory_decision = MemoryDecision.decide(
headers=request.headers,
memory_handler=self.memory_handler,
memory_user_id=memory_user_id,
mode_name=get_memory_injection_mode(),
)
memory_decision.apply_to_tags(tags)
# Check cache (non-streaming only)
cache_hit = False
if self.cache and not stream:
@ -1297,10 +1316,15 @@ class AnthropicHandlerMixin:
except Exception as e:
logger.debug(f"[{request_id}] Traffic learner: {e}")
# Memory: Inject context and tools
# Memory: Inject context and tools — gated on MemoryDecision.
# ``inject`` is False under bypass, missing handler, missing
# user_id, or HEADROOM_MEMORY_INJECTION_MODE in disabled/tool.
# Pre-PR-this the gate was a raw conjunction that silently
# ignored bypass; now bypass is honoured here on Anthropic
# /v1/messages just as on /v1/responses.
memory_context_injected = False
memory_tools_injected = False
if self.memory_handler and memory_user_id:
if memory_decision.inject:
# Search and inject memory context
if self.memory_handler.config.inject_context:
try:
@ -1310,6 +1334,7 @@ class AnthropicHandlerMixin:
memory_user_id,
optimized_messages,
request_context=memory_request_ctx,
query=MemoryQuery.from_messages(optimized_messages),
),
timeout=(
self.config.anthropic_pre_upstream_memory_context_timeout_seconds

42
headroom/proxy/handlers/gemini.py
View file

@ -5,6 +5,7 @@ Contains all Google Gemini API handlers including format conversion utilities.
from __future__ import annotations
import asyncio
import json
import logging
import os
@ -253,6 +254,22 @@ class GeminiHandlerMixin:
),
)
# Canonical memory-injection gate (parallels Anthropic + OpenAI).
# Pre-PR-this Gemini's memory site silently ignored
# `x-headroom-bypass: true`, mutating request bytes under the
# user's "don't touch my bytes" signal.
from headroom.proxy.helpers import get_memory_injection_mode
from headroom.proxy.memory_decision import MemoryDecision
from headroom.proxy.memory_query import MemoryQuery
memory_decision = MemoryDecision.decide(
headers=request.headers,
memory_handler=self.memory_handler,
memory_user_id=memory_user_id,
mode_name=get_memory_injection_mode(),
)
memory_decision.apply_to_tags(tags)
# Rate limiting (use Gemini API key)
if self.rate_limiter:
rate_key = headers.get("x-goog-api-key", "default")[:20]
@ -379,13 +396,28 @@ class GeminiHandlerMixin:
# the memory handler is in ``MemoryMode.TOOL`` its
# ``search_and_format_context`` returns ``None`` so nothing flows
# in here.
if self.memory_handler and memory_user_id:
if memory_decision.inject:
# Memory-handler is guaranteed present when inject=True.
# Add a timeout wrapping (matches Anthropic + Responses) so
# a slow memory backend can't stall Gemini requests — pre-
# PR-this Gemini was the only handler without one.
#
# The append uses provider="openai" because Gemini reuses
# OpenAI's user-message content shape after the proxy's
# gemini-contents → messages → gemini-contents round-trip.
# That's a real coupling, not a bug — `_append_to_latest_
# user_tail` only knows two surface shapes; openai matches
# the post-conversion structure exactly.
try:
if self.memory_handler.config.inject_context:
memory_context = await self.memory_handler.search_and_format_context(
memory_user_id,
optimized_messages,
request_context=memory_request_ctx,
memory_context = await asyncio.wait_for(
self.memory_handler.search_and_format_context(
memory_user_id,
optimized_messages,
request_context=memory_request_ctx,
query=MemoryQuery.from_messages(optimized_messages),
),
timeout=(self.config.anthropic_pre_upstream_memory_context_timeout_seconds),
)
if memory_context:
new_messages, bytes_appended = (

123
headroom/proxy/handlers/openai.py
View file

@ -1272,6 +1272,22 @@ class OpenAIHandlerMixin:
),
)
# Canonical memory-injection gate (parallels Anthropic). Pre-
# PR-this the inline conjunction at the memory site silently
# ignored `x-headroom-bypass: true`, mutating request bytes
# under the user's "don't touch my bytes" signal.
from headroom.proxy.helpers import get_memory_injection_mode
from headroom.proxy.memory_decision import MemoryDecision
from headroom.proxy.memory_query import MemoryQuery
memory_decision = MemoryDecision.decide(
headers=request.headers,
memory_handler=self.memory_handler,
memory_user_id=memory_user_id,
mode_name=get_memory_injection_mode(),
)
memory_decision.apply_to_tags(tags)
# Rate limiting
if self.rate_limiter:
rate_key = headers.get("authorization", "default")[:20]
@ -1617,13 +1633,21 @@ class OpenAIHandlerMixin:
# invariant I2. See REALIGNMENT/03-phase-A-lockdown.md PR-A3.
memory_context_injected = False
memory_tools_injected = False
if self.memory_handler and memory_user_id:
if memory_decision.inject:
# Memory-handler is guaranteed present when inject=True.
# Timeout-wrap (matches Anthropic /v1/messages and
# /v1/responses) — pre-PR-this site was the only chat
# path without one.
try:
if self.memory_handler.config.inject_context:
memory_context = await self.memory_handler.search_and_format_context(
memory_user_id,
optimized_messages,
request_context=memory_request_ctx,
memory_context = await asyncio.wait_for(
self.memory_handler.search_and_format_context(
memory_user_id,
optimized_messages,
request_context=memory_request_ctx,
query=MemoryQuery.from_messages(optimized_messages),
),
timeout=(self.config.anthropic_pre_upstream_memory_context_timeout_seconds),
)
if memory_context:
from headroom.proxy.helpers import (
@ -2423,8 +2447,26 @@ class OpenAIHandlerMixin:
transforms_applied: list[str] = []
optimization_latency = (time.time() - start_time) * 1000
# Memory: inject context and tools for Responses API requests
if self.memory_handler and memory_user_id and not _bypass:
# Memory: inject context and tools for Responses API requests.
# Gated on MemoryDecision — uniformly respects bypass across all
# five injection sites. The Responses path is the only one that
# injects BEFORE compression today (sites 1/2/3 inject after);
# bringing this into alignment is queued as a follow-up
# (FUTURE: move context injection to post-compression for
# uniform "memory text rides uncompressed across all
# handlers" semantics — separate PR with cache-stability tests).
from headroom.proxy.helpers import get_memory_injection_mode
from headroom.proxy.memory_decision import MemoryDecision
from headroom.proxy.memory_query import MemoryQuery
responses_memory_decision = MemoryDecision.decide(
headers=request.headers,
memory_handler=self.memory_handler,
memory_user_id=memory_user_id,
mode_name=get_memory_injection_mode(),
)
responses_memory_decision.apply_to_tags(tags)
if responses_memory_decision.inject:
try:
# Memory context now routes exclusively to the live-zone tail
# (latest non-frozen user item). Instructions are part of the
@ -2437,6 +2479,7 @@ class OpenAIHandlerMixin:
memory_user_id,
optimized_messages,
request_context=memory_request_ctx,
query=MemoryQuery.from_messages(optimized_messages),
),
timeout=RESPONSES_CONTEXT_SEARCH_TIMEOUT_SECONDS,
)
@ -3394,13 +3437,35 @@ class OpenAIHandlerMixin:
)
# --- Memory: inject context, tools, and instructions ---
# Gated on MemoryDecision — uniform bypass-respect across
# all five sites. WS sets memory_user_id only on the inject
# path (matches pre-PR behaviour); MemoryDecision is the
# canonical gate.
memory_user_id: str | None = None
memory_request_ctx = None
if self.memory_handler and body and not _ws_bypass:
memory_user_id = ws_headers.get(
if self.memory_handler and body:
_ws_memory_user_id_candidate = ws_headers.get(
"x-headroom-user-id",
os.environ.get("USER", os.environ.get("USERNAME", "default")),
)
else:
_ws_memory_user_id_candidate = None
from headroom.proxy.helpers import get_memory_injection_mode
from headroom.proxy.memory_decision import MemoryDecision
from headroom.proxy.memory_query import MemoryQuery
ws_memory_decision = MemoryDecision.decide(
headers=ws_headers,
memory_handler=self.memory_handler if body else None,
memory_user_id=_ws_memory_user_id_candidate,
mode_name=get_memory_injection_mode(),
)
# ws_tags was extracted at handler entry (L3028); applying
# the memory skip reason here so per-turn RequestOutcomes
# carry it for dashboard slicing.
ws_memory_decision.apply_to_tags(ws_tags)
if ws_memory_decision.inject:
memory_user_id = _ws_memory_user_id_candidate
try:
# Unwrap response.create envelope to access the response body
ws_response_body = body.get("response", body)
@ -3456,6 +3521,7 @@ class OpenAIHandlerMixin:
memory_user_id,
ws_msgs,
request_context=memory_request_ctx,
query=MemoryQuery.from_messages(ws_msgs),
),
timeout=RESPONSES_CONTEXT_SEARCH_TIMEOUT_SECONDS,
)
@ -3467,15 +3533,38 @@ class OpenAIHandlerMixin:
f"continuing without it"
)
if memory_context:
existing = ws_response_body.get("instructions") or ""
if existing:
ws_response_body["instructions"] = f"{existing}\n\n{memory_context}"
# Route memory into ws_response_body["input"]
# (the user-input field) rather than
# ws_response_body["instructions"] (the
# system/cache-hot-zone field). All other
# handlers inject at the user-message tail
# so the cache prefix bytes stay byte-
# stable across turns — invariant I2. The
# WS path was the lone outlier writing to
# instructions (system); fixed here for
# uniformity with sites 1/2/3/4.
ws_input_for_inject = ws_response_body.get("input", "")
if isinstance(ws_input_for_inject, str):
if ws_input_for_inject:
ws_response_body["input"] = (
ws_input_for_inject + "\n\n" + memory_context
)
else:
ws_response_body["input"] = memory_context
logger.info(
f"[{request_id}] WS Memory: Injected {len(memory_context)} chars "
f"into input tail (string-shaped input)"
)
else:
ws_response_body["instructions"] = memory_context
logger.info(
f"[{request_id}] WS Memory: Injected {len(memory_context)} chars "
f"of context into instructions"
)
# List-shaped WS input is owned by the
# Rust handler (per PR-C5 comment). The
# Python path leaves memory un-injected
# for list inputs rather than touching
# instructions.
logger.info(
f"[{request_id}] WS Memory: list-shaped input — "
f"injection deferred to Rust handler"
)
# Inject memory tools (Responses API format) — PR-A7 (P0-6).
# WS path uses a per-connection UUID; tracker scope is

144
headroom/proxy/memory_decision.py Normal file
View file

@ -0,0 +1,144 @@
"""``MemoryDecision``: canonical "should we inject memory context?" gate.
Input-side analog of :class:`CompressionDecision`. Pre-this-PR, the
6 memory-injection sites across four handler files computed the gate
inline with subtle drift — most notably, 3 sites (Anthropic chat,
OpenAI chat, Gemini) never gated on ``x-headroom-bypass``, so
memory injection silently mutated requests when the user explicitly
asked for byte-faithful passthrough.
This is **decision-only**. It gates whether the request bytes get
mutated (memory context appended to user-tail). It does NOT gate
background memory STORAGE — traffic-learner runs on a separate path
and continues accumulating signal even under bypass. The user's
"don't touch my bytes" signal is for the INJECTION; the user's
working memory should still grow.
Precedence (highest first):
1. ``bypass_header`` — user's explicit "do not touch my bytes"
2. ``no_handler`` — no memory backend configured
3. ``no_user_id`` — per-request user_id missing
4. ``mode_disabled`` — operator HEADROOM_MEMORY_INJECTION_MODE=disabled
5. ``mode_tool`` — operator HEADROOM_MEMORY_INJECTION_MODE=tool
(auto-inject off; the agent calls memory tools explicitly)
6. otherwise → ``inject=True``
This module exposes one value type + one factory + one helper. Pure
function; same Rust-port shape as :class:`CompressionDecision`.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any
from headroom.proxy.helpers import _headroom_bypass_enabled
@dataclass(frozen=True)
class MemoryDecision:
"""Immutable, value-equal snapshot of the memory-injection decision.
Construction policy: use :meth:`decide`. Direct construction is
legal but unusual — tests use it; 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.
"""
inject: bool
# When ``inject`` is False, this is the canonical reason surfaced
# in logs and in RequestOutcome.tags["memory_skip_reason"] so the
# dashboard can slice memory-skip traffic by cause. One of:
# * "bypass_header" — user set x-headroom-bypass/x-headroom-mode
# * "no_handler" — no memory backend configured on the proxy
# * "no_user_id" — per-request user_id missing
# * "mode_disabled" — operator HEADROOM_MEMORY_INJECTION_MODE=disabled
# * "mode_tool" — operator HEADROOM_MEMORY_INJECTION_MODE=tool
# When ``inject`` is True, this is None.
skip_reason: str | None
# Observability: every constituent boolean exposed so debug tools
# answer "what did the decision see?" without re-running.
bypass_header_set: bool
memory_handler_present: bool
memory_user_id_present: bool
mode_name: str
@classmethod
def decide(
cls,
*,
headers: Any,
memory_handler: Any | None,
memory_user_id: str | None,
mode_name: str,
) -> MemoryDecision:
"""Compute the canonical memory-injection decision.
Parameters
----------
headers
Inbound request headers. Accepts any object with a
``.get(key)`` method (dict, starlette Headers, mapping).
Bypass detected via ``_headroom_bypass_enabled``.
memory_handler
The proxy's memory handler instance, or ``None`` if no
memory backend is configured. Presence only is checked —
no methods called.
memory_user_id
Per-request user_id from ``x-headroom-user-id`` header (or
env default). ``None`` or empty string treated as missing.
mode_name
One of ``"auto_tail"`` / ``"tool"`` / ``"disabled"``.
Comes from ``get_memory_injection_mode()`` which reads
``HEADROOM_MEMORY_INJECTION_MODE``.
"""
bypass = _headroom_bypass_enabled(headers)
has_handler = memory_handler is not None
has_user = bool(memory_user_id)
if bypass:
reason: str | None = "bypass_header"
inject = False
elif not has_handler:
reason = "no_handler"
inject = False
elif not has_user:
reason = "no_user_id"
inject = False
elif mode_name == "disabled":
reason = "mode_disabled"
inject = False
elif mode_name == "tool":
reason = "mode_tool"
inject = False
else:
reason = None
inject = True
return cls(
inject=inject,
skip_reason=reason,
bypass_header_set=bypass,
memory_handler_present=has_handler,
memory_user_id_present=has_user,
mode_name=mode_name,
)
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 ``inject=True`` —
absence vs presence of ``memory_skip_reason`` is the signal.
Mirror of :meth:`CompressionDecision.apply_to_tags`. Handlers
call this immediately after ``decide()`` so the resulting
``RequestOutcome.tags`` carries memory observability via the
same path the funnel already uses for ``client`` and
``passthrough_reason``.
"""
if self.skip_reason is not None:
tags["memory_skip_reason"] = self.skip_reason

96
headroom/proxy/memory_handler.py
View file

@ -622,37 +622,47 @@ class MemoryHandler:
user_id: str,
messages: list[dict[str, Any]],
request_context: RequestContext | None = None,
*,
query: Any | None = None,
budget: Any | None = None,
) -> str | None:
"""Search memories and format as context injection.
Args:
user_id: User identifier for memory scoping (the base user
id, derived from ``x-headroom-user-id`` upstream).
messages: Conversation messages (used to extract query).
messages: Conversation messages (used to extract query when
``query`` is not provided).
request_context: Optional request envelope (headers, system
prompt, base user id). When provided, memory retrieval
is scoped to the resolved workspace / project so memories
from unrelated projects can never bleed in (GH #462). When
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.
from unrelated projects can never bleed in (GH #462).
query: Optional :class:`MemoryQuery` — multi-source, full-
fidelity retrieval query. When provided, takes precedence
over the ``messages``-derived query. Constructed at the
handler from latest user msg + recent tool outputs +
recent assistant turns; preserves full input fidelity (no
500-char truncation).
budget: Optional :class:`MemoryInjectionBudget` — bounds the
returned formatted block by tokens / entries / min
similarity. When ``None``, defaults are taken from
``self.config`` so the existing top_k / min_similarity
contract is preserved.
Returns:
Formatted context string, or None if no relevant memories.
PR-B6: When ``self.config.mode == MemoryMode.TOOL``, this method
returns ``None`` unconditionally so the proxy never auto-injects.
The model must call ``memory_search`` explicitly to retrieve. This
is the single chokepoint that gates auto-injection across all
provider handlers (Anthropic /v1/messages, OpenAI /v1/chat/completions,
OpenAI /v1/responses).
The model must call ``memory_search`` explicitly to retrieve.
"""
from headroom.proxy.memory_injection import MemoryInjectionBudget
if not self.config.inject_context:
return None
# PR-B6: Tool mode disables auto-injection. The model calls
# ``memory_search`` to retrieve when it wants to. Log the skip
# decision so cache-affecting routing remains observable.
# ``memory_search`` to retrieve when it wants to.
if self.config.mode == MemoryMode.TOOL:
logger.info(
"event=memory_mode_skip mode=tool user_id=%s reason=tool_mode_no_auto_injection",
@ -666,18 +676,36 @@ class MemoryHandler:
backend, scope, effective_user_id = self._resolve_for_request(user_id, request_context)
# Extract query from last user message
query = self._extract_user_query(messages)
if not query:
logger.debug("Memory: No user query found for context search")
# Build the embedding query. When the handler provides a
# MemoryQuery, use its multi-source untruncated input; otherwise
# fall back to extracting from messages (kept for legacy callers
# / tests). Full fidelity in both paths.
if query is not None:
query_text = query.to_embedding_input()
else:
query_text = self._extract_user_query(messages)
if not query_text:
logger.debug("Memory: No query text for context search")
return None
# Compose the budget: explicit per-call wins; otherwise derive
# from self.config so existing top_k/min_similarity callers see
# no behaviour change.
effective_budget = (
budget
if budget is not None
else MemoryInjectionBudget(
max_entries=self.config.top_k,
min_similarity=self.config.min_similarity,
)
)
try:
# Search memories on the per-request resolved backend.
results = await backend.search_memories(
query=query,
query=query_text,
user_id=effective_user_id,
top_k=self.config.top_k,
top_k=effective_budget.max_entries,
include_related=True,
)
@ -689,17 +717,22 @@ class MemoryHandler:
)
return None
# Filter by minimum similarity
filtered_results = [r for r in results if r.score >= self.config.min_similarity]
# Filter by minimum similarity using the budget.
filtered_results = [r for r in results if r.score >= effective_budget.min_similarity]
if not filtered_results:
logger.debug(
f"Memory: {len(results)} memories found but none above threshold "
f"{self.config.min_similarity}"
f"{effective_budget.min_similarity}"
)
return None
# Format as context
# Cap entry count via the budget (defence-in-depth — backend
# already gets top_k=max_entries but this enforces it on
# post-filter results too).
filtered_results = filtered_results[: effective_budget.max_entries]
# Format as context.
memory_lines = []
for i, result in enumerate(filtered_results, 1):
memory_lines.append(f"{i}. {result.memory.content}")
@ -723,12 +756,19 @@ The following information was previously saved in this scope:
Use this context to provide personalized and contextually relevant responses."""
# Apply the token-budget cap on the formatted block. Pre-this-
# PR there was no cap — up to ~4000 tokens could be injected
# per request. The budget bounds the output without touching
# the input query (which stays full-fidelity per MemoryQuery).
context = effective_budget.apply_to_text(context)
logger.info(
"event=memory_inject user=%s scope=%s count=%d chars=%d",
"event=memory_inject user=%s scope=%s count=%d chars=%d budget_tokens=%d",
effective_user_id,
scope.display_name if scope else "<legacy>",
len(memory_lines),
len(context),
effective_budget.max_tokens,
)
return context
@ -796,7 +836,13 @@ Use this context to provide personalized and contextually relevant responses."""
raise ValueError(f"Unknown provider {provider!r}; expected 'anthropic' or 'openai'")
def _extract_user_query(self, messages: list[dict[str, Any]]) -> str:
"""Extract the user query from the last user message."""
"""Extract the user query from the last user message.
Returns the FULL message text — no truncation. The embedding
model handles its own context window. (Pre-this-PR this
method capped at 500 chars, silently throwing away signal —
none of Letta/Mem0/Cognee/Supermemory truncate.)
"""
for msg in reversed(messages):
if msg.get("role") != "user":
continue
@ -804,14 +850,14 @@ Use this context to provide personalized and contextually relevant responses."""
content = msg.get("content", "")
if isinstance(content, str):
return content[:500] # Limit query length
return content
if isinstance(content, list):
for block in content:
if isinstance(block, dict) and block.get("type") == "text":
text = str(block.get("text", ""))
if text:
return text[:500]
return text
return ""

93
headroom/proxy/memory_injection.py Normal file
View file

@ -0,0 +1,93 @@
"""``MemoryInjectionBudget``: uniform token/entry cap on retrieved memory.
Pre-this-PR Headroom had NO token cap on injected memory. Top-K=10
candidates × ~400 tokens each = up to ~4000 tokens injected per
request. None of Letta/Mem0/Cognee/Supermemory ship a token-uncapped
injection path on the hot wire.
This budget is applied at the formatting boundary in
``memory_handler.search_and_format_context`` (after the backend
returns candidates, before the formatted block is appended to the
request). One value type → consistent enforcement across all five
sites.
The budget bounds are configurable; the defaults are conservative
(1024 tokens / 10 entries / 0.3 similarity floor) so a missing config
can't accidentally restore the unbounded behaviour.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any
# Rough char-per-token heuristic for budget enforcement. Used only to
# bound the OUTPUT block (the formatted injection text) — INPUT
# fidelity is preserved by MemoryQuery (no truncation).
#
# 4 chars/token is the standard heuristic for English text. We use it
# for the cap; the actual token count at upstream is decided by the
# upstream provider's tokenizer.
_CHARS_PER_TOKEN_HEURISTIC = 4
@dataclass(frozen=True)
class MemoryInjectionBudget:
"""Frozen budget applied at the injection boundary.
Three independent dials:
* ``max_tokens`` — total bytes (heuristically converted) in the
formatted injection block. Default 1024 tokens (~4KB).
* ``max_entries`` — cap on the number of memory entries included.
Default 10 (matches backend top_k).
* ``min_similarity`` — floor on cosine similarity; entries below
are dropped. Default 0.3 (matches backend default).
Operators tune via constructor args. Defaults are hard-coded so a
misconfigured caller can't accidentally restore unbounded
behaviour (the pre-this-PR state).
"""
max_tokens: int = 1024
max_entries: int = 10
min_similarity: float = 0.3
def apply_to_text(self, text: str) -> str:
"""Bound a formatted injection block by ``max_tokens``.
Truncation prefers line boundaries (memory entries are
line-delimited so the dashboard renders intact bullet points
rather than mid-word cuts). Empty input → empty output.
This caps the OUTPUT block. The INPUT (the query going to the
embedder) is NOT truncated — that's MemoryQuery's contract.
"""
if not text:
return text
char_budget = self.max_tokens * _CHARS_PER_TOKEN_HEURISTIC
if len(text) <= char_budget:
return text
# Truncate at the last newline at or before the budget so the
# final included line is complete.
cut = text.rfind("\n", 0, char_budget)
if cut <= 0:
# No newline within budget — fall back to hard cut.
return text[:char_budget]
return text[: cut + 1]
def apply_to_entries(self, entries: list[dict[str, Any]]) -> list[dict[str, Any]]:
"""Cap a list of ranked memory candidates by entry count + min similarity.
Preserves order (assumes upstream has already ranked by
score). The budget does NOT re-rank — it only filters/caps.
Filtering precedence:
1. Drop entries with ``score < min_similarity``
2. Cap to ``max_entries``
Entries are dict-shaped (the backend returns dicts with at
minimum ``content`` and ``score`` keys).
"""
filtered = [e for e in entries if float(e.get("score", 0.0)) >= self.min_similarity]
return filtered[: self.max_entries]

160
headroom/proxy/memory_query.py Normal file
View file

@ -0,0 +1,160 @@
"""``MemoryQuery``: multi-source, full-fidelity retrieval query.
Pre-this-PR, the retrieval query was "latest user message, truncated
to 500 chars" (memory_handler.py:807). The truncation was a real bug
— none of Letta / Mem0 / Cognee / Supermemory truncate the embedding
input. Tool outputs are often the strongest retrieval signal in
coding sessions, and they were ignored entirely.
This value type captures the query at full fidelity from three
sources:
* ``user_text`` — latest user message, untruncated
* ``recent_tool_outputs`` — last N tool results
* ``recent_assistant_turns`` — last K assistant turns for intent
The embedding model handles its own context window (MiniLM 512 tok;
BGE-small 8K tok). Long inputs that exceed the model window become
the model's problem to mean-pool or chunk — they don't get
truncated upstream.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any
# Section delimiters surfaced in the embedding input so the embedder
# sees structured context rather than a wall of run-on text. Kept short
# so they don't dominate the embedding signal.
_USER_DELIM = "### USER ###\n"
_ASSISTANT_DELIM = "\n### PRIOR_ASSISTANT ###\n"
_TOOL_DELIM = "\n### TOOL_OUTPUT ###\n"
@dataclass(frozen=True)
class MemoryQuery:
"""Frozen multi-source query for memory retrieval.
All fields preserve full input fidelity — no truncation, no
summarization. The caller assembles the sources; this type only
holds them. The retrieval backend decides how to embed
(mean-pool, chunk, model-side truncation, etc.) but cannot lose
information before it sees the data.
Tuples for the recent-* fields so the dataclass stays hashable
(frozen + value-equal).
"""
user_text: str
recent_tool_outputs: tuple[str, ...]
recent_assistant_turns: tuple[str, ...]
conversation_id: str | None
def to_embedding_input(self) -> str:
"""Concatenate sources into a delimited embedding input.
Order: prior assistant turns (oldest first) → tool outputs
(oldest first) → latest user text. User text last because the
embedder's positional weighting often emphasizes the tail of
the input.
"""
parts: list[str] = []
for asst in self.recent_assistant_turns:
if asst:
parts.append(_ASSISTANT_DELIM + asst)
for tool_out in self.recent_tool_outputs:
if tool_out:
parts.append(_TOOL_DELIM + tool_out)
if self.user_text:
parts.append(_USER_DELIM + self.user_text)
return "".join(parts)
@classmethod
def from_messages(
cls,
messages: list[dict[str, Any]] | None,
*,
lookback_assistant: int = 2,
lookback_tools: int = 3,
conversation_id: str | None = None,
) -> MemoryQuery:
"""Construct a MemoryQuery from a chat-style messages list.
Walks the message list once. Extracts:
* Latest ``role: user`` message → ``user_text``
* Up to ``lookback_assistant`` most recent assistant turns →
``recent_assistant_turns`` (chronological order)
* Up to ``lookback_tools`` most recent tool outputs →
``recent_tool_outputs`` (chronological order)
Handles both OpenAI shape (``role: tool``) and Anthropic shape
(``tool_result`` content block inside a ``role: user`` message).
"""
if not messages:
return cls(
user_text="",
recent_tool_outputs=(),
recent_assistant_turns=(),
conversation_id=conversation_id,
)
latest_user = ""
assistant_turns: list[str] = []
tool_outputs: list[str] = []
# Walk messages backward so we naturally find the LATEST entries
# first; preserve chronological order in the output by reversing
# the collected lists at the end.
for msg in reversed(messages):
role = msg.get("role")
content = msg.get("content", "")
if role == "user":
# Distinguish "real user text" from "Anthropic tool_result
# masquerading as a user message". Anthropic uses
# role=user with content=[{type: tool_result, ...}].
if isinstance(content, list):
for block in content:
if isinstance(block, dict) and block.get("type") == "tool_result":
tool_text = block.get("content", "")
if isinstance(tool_text, list):
# Nested content blocks; flatten text fields.
tool_text = "\n".join(
b.get("text", "") for b in tool_text if isinstance(b, dict)
)
if tool_text and len(tool_outputs) < lookback_tools:
tool_outputs.append(str(tool_text))
# Anthropic tool_result is NOT a real user turn —
# don't use it as the user_text source. Continue
# walking back for the actual user message.
elif isinstance(content, str):
if not latest_user:
latest_user = content
elif role == "assistant":
if isinstance(content, str) and content:
if len(assistant_turns) < lookback_assistant:
assistant_turns.append(content)
elif isinstance(content, list):
text_parts = [
b.get("text", "")
for b in content
if isinstance(b, dict) and b.get("type") == "text"
]
joined = "\n".join(p for p in text_parts if p)
if joined and len(assistant_turns) < lookback_assistant:
assistant_turns.append(joined)
elif role == "tool":
# OpenAI shape: tool messages carry the result.
if isinstance(content, str) and content and len(tool_outputs) < lookback_tools:
tool_outputs.append(content)
# Reverse to restore chronological order (we walked backward).
return cls(
user_text=latest_user,
recent_tool_outputs=tuple(reversed(tool_outputs)),
recent_assistant_turns=tuple(reversed(assistant_turns)),
conversation_id=conversation_id,
)

314
tests/test_memory_decision.py Normal file
View file

@ -0,0 +1,314 @@
"""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

130
tests/test_memory_injection_budget.py Normal file
View file

@ -0,0 +1,130 @@
"""Tests for :class:`headroom.proxy.memory_injection.MemoryInjectionBudget`.
Pre-PR Headroom had **no token cap on injected memory**: top_k=10
candidates × ~400 tokens each = up to ~4000 tokens injected per
request. None of Letta/Mem0/Cognee/Supermemory ship a token-uncapped
injection path on the hot wire.
``MemoryInjectionBudget`` is the single configurable cap applied to
every injection site so all 5 sites are uniformly bounded — set the
budget once, apply at every handler, dashboards see the same shape.
"""
from __future__ import annotations
from dataclasses import FrozenInstanceError
from headroom.proxy.memory_injection import MemoryInjectionBudget
# ── Value-type contract ───────────────────────────────────────────────
def test_budget_is_frozen() -> None:
b = MemoryInjectionBudget()
try:
b.max_tokens = 99 # type: ignore[misc]
except FrozenInstanceError:
pass
else:
raise AssertionError("MemoryInjectionBudget must be frozen")
def test_budget_defaults() -> None:
"""Default budget is conservative — 1024 tokens, 10 entries, 0.3
similarity floor. Operators can override via config; the default
is hard-set so a misconfiguration can't accidentally unbound
injection."""
b = MemoryInjectionBudget()
assert b.max_tokens == 1024
assert b.max_entries == 10
assert b.min_similarity == 0.3
def test_budget_value_equal() -> None:
a = MemoryInjectionBudget(max_tokens=512, max_entries=5, min_similarity=0.5)
b = MemoryInjectionBudget(max_tokens=512, max_entries=5, min_similarity=0.5)
assert a == b
# ── apply_to_text — bounding the formatted context block ─────────────
def test_apply_to_text_returns_input_when_under_budget() -> None:
"""Short context passes through unchanged — no spurious mutation."""
b = MemoryInjectionBudget(max_tokens=1024)
text = "## Relevant Memories\n1. small fact\n"
out = b.apply_to_text(text)
assert out == text
def test_apply_to_text_truncates_when_over_budget() -> None:
"""Large context is bounded — truncated at the budget. The
truncation here is on the OUTPUT (the formatted injection block),
NOT on the INPUT (which keeps full fidelity per MemoryQuery
contract)."""
# 4 tokens/char heuristic in our cap — make the input clearly
# over even the most generous budget.
b = MemoryInjectionBudget(max_tokens=128) # ~512 chars at 4 char/token
huge = "x" * 100000
out = b.apply_to_text(huge)
# Output should be substantially smaller than input.
assert len(out) < len(huge)
def test_apply_to_text_preserves_full_lines() -> None:
"""When truncating, prefer cutting at line boundaries so the
dashboard renders intact memory entries (no half-truncated bullet
point)."""
b = MemoryInjectionBudget(max_tokens=64) # very tight
text = "## Relevant Memories\n" + "".join(f"{i}. fact {i}\n" for i in range(100))
out = b.apply_to_text(text)
# No partial last line — every retained line ends in newline or is
# the final line.
if out and not out.endswith("\n"):
# The last char is the closing of the final line; it must not
# be in the middle of "fact " — easy heuristic: must not end
# mid-word with a hanging digit-then-period.
assert ". fact" not in out[-15:] or out.rstrip().endswith(("fact 0", "fact 1", "fact 2"))
def test_apply_to_text_handles_empty_input() -> None:
"""Empty input → empty output."""
assert MemoryInjectionBudget().apply_to_text("") == ""
# ── apply_to_entries — bounding the list before formatting ───────────
def test_apply_to_entries_caps_entry_count() -> None:
"""Even if the backend returns 100 candidates, the budget caps
entry count to ``max_entries``."""
b = MemoryInjectionBudget(max_entries=3)
entries = [{"content": f"entry {i}", "score": 0.9 - i * 0.01} for i in range(20)]
out = b.apply_to_entries(entries)
assert len(out) == 3
def test_apply_to_entries_preserves_order_of_input() -> None:
"""Budget doesn't re-rank — the backend's order is preserved. (The
backend should already have ranked by score; budget just caps.)"""
b = MemoryInjectionBudget(max_entries=2)
entries = [
{"content": "alpha", "score": 0.9},
{"content": "beta", "score": 0.8},
{"content": "gamma", "score": 0.7},
]
out = b.apply_to_entries(entries)
assert [e["content"] for e in out] == ["alpha", "beta"]
def test_apply_to_entries_filters_below_min_similarity() -> None:
"""Entries below ``min_similarity`` are dropped, regardless of
entry-count budget remaining."""
b = MemoryInjectionBudget(max_entries=10, min_similarity=0.5)
entries = [
{"content": "kept", "score": 0.9},
{"content": "dropped", "score": 0.3},
{"content": "kept2", "score": 0.55},
]
out = b.apply_to_entries(entries)
assert {e["content"] for e in out} == {"kept", "kept2"}

198
tests/test_memory_invariants.py Normal file
View file

@ -0,0 +1,198 @@
"""AST-walking contract tests locking memory-system invariants.
These tests introspect the four handler modules' source ASTs to assert
two structural properties that, if regressed, would silently re-
introduce the bug classes this PR fixed:
(a) **No raw memory-gate conjunction**: every memory injection block
must be gated by ``MemoryDecision``, not by an inline
``if self.memory_handler and memory_user_id`` conjunction. The
raw conjunction is what allowed sites 1/2/3 to silently ignore
``x-headroom-bypass: true``.
(b) **No memory writes to system/instructions**: memory injection
must target user-message-tail / body["input"] / messages, never
``body["instructions"]`` or system content. Pre-PR-this the WS
handler was the lone outlier writing to instructions; the AST
check ensures it doesn't sneak back.
The checks are static — no handler is invoked. They run in
milliseconds and catch future regressions at PR-review time.
"""
from __future__ import annotations
import ast
import re
from pathlib import Path
import pytest
HANDLER_FILES = [
Path("headroom/proxy/handlers/anthropic.py"),
Path("headroom/proxy/handlers/openai.py"),
Path("headroom/proxy/handlers/gemini.py"),
Path("headroom/proxy/handlers/batch.py"),
]
# ── Invariant A — no raw memory-gate conjunction ──────────────────────
def _file_contains_raw_memory_gate(file_path: Path) -> list[tuple[int, str]]:
"""Find ``if (self.)memory_handler and memory_user_id`` raw
conjunctions in the file. Returns list of (line, snippet).
The acceptable replacements are ``if memory_decision.inject:`` or
``if (responses|ws)_memory_decision.inject:``.
AST-walking the conditional itself is complex (BoolOp + Attribute),
so we use a regex to scan source lines. False positives are caught
by the test author at write time — this is a one-line invariant.
"""
text = file_path.read_text(encoding="utf-8")
# Match "if self.memory_handler and memory_user_id" but NOT inside
# the helper that defines the decision (which references both names
# for documentation purposes) — we only care about handler logic.
pattern = re.compile(r"^\s*if\s+self\.memory_handler\s+and\s+memory_user_id\b")
hits = []
for i, line in enumerate(text.splitlines(), start=1):
if pattern.match(line):
hits.append((i, line.rstrip()))
return hits
def test_no_raw_memory_handler_gate_in_handlers() -> None:
"""Pre-PR-this, sites 1/2/3 used ``if self.memory_handler and
memory_user_id:`` as the memory-injection gate — silently
ignoring bypass. After PR-this, every site routes through
``MemoryDecision.decide(...)`` and gates on
``memory_decision.inject``. This test ensures the raw conjunction
cannot return without explicit review."""
offenders = []
for f in HANDLER_FILES:
offenders.extend([(f, ln, src) for (ln, src) in _file_contains_raw_memory_gate(f)])
if offenders:
formatted = "\n".join(f" {f.name}:{ln} {src!r}" for f, ln, src in offenders)
pytest.fail(
f"{len(offenders)} handler site(s) use the pre-PR raw memory "
"gate `if self.memory_handler and memory_user_id`:\n"
f"{formatted}\n\n"
"Replace with `MemoryDecision.decide(...)` + "
"`if memory_decision.inject:`. See PR for the canonical pattern."
)
# ── Invariant B — memory never writes to system/instructions ─────────
_FORBIDDEN_SYSTEM_WRITES = (
# Direct mutation of cache-hot-zone system fields by the memory
# path. The patterns below are exact assignment forms — they match
# the pre-PR-this WS bug at openai.py:3517.
re.compile(r'ws_response_body\["instructions"\]\s*='),
re.compile(r'response_body\["instructions"\]\s*='),
re.compile(r'body\["instructions"\]\s*='),
re.compile(r'body\["system"\]\s*='),
)
def _line_is_memory_related(line: str) -> bool:
"""Heuristic: a line that contains ``memory_context`` or
``memory_inject`` is in the memory-injection code path."""
return "memory_context" in line or "memory_inject" in line
def _find_system_writes_in_memory_context(file_path: Path) -> list[tuple[int, str]]:
"""Find lines that both:
- Look like a write to body["instructions"] / body["system"]
- Live within ~10 lines of a ``memory_context`` reference
This is a windowed-context check — we don't want false positives
from unrelated instructions-writes (e.g. tool-result handling).
"""
text = file_path.read_text(encoding="utf-8").splitlines()
memory_line_indices = [i for i, line in enumerate(text) if _line_is_memory_related(line)]
hits = []
for i, line in enumerate(text):
if not any(p.search(line) for p in _FORBIDDEN_SYSTEM_WRITES):
continue
# Within 10 lines of any memory-related line?
if any(abs(i - mi) <= 10 for mi in memory_line_indices):
hits.append((i + 1, line.strip()))
return hits
def test_memory_never_writes_to_system_or_instructions() -> None:
"""Pre-PR-this, the WS handler wrote memory context to
``ws_response_body["instructions"]`` — the system / cache-hot-zone
field. That mutated the prefix cache bytes on every turn. All
other sites route to user-message tail / body["input"]. This
test asserts memory_context-related code paths never write to a
forbidden system-field assignment."""
offenders = []
for f in HANDLER_FILES:
offenders.extend([(f, ln, src) for (ln, src) in _find_system_writes_in_memory_context(f)])
if offenders:
formatted = "\n".join(f" {f.name}:{ln} {src!r}" for f, ln, src in offenders)
pytest.fail(
f"{len(offenders)} suspected memory→system write(s):\n"
f"{formatted}\n\n"
"Memory must append to user-message tail (e.g. body['input'] "
"for Responses, optimized_messages for chat). Never write to "
"body['instructions'] or body['system'] — they are the cache "
"hot zone (invariant I2)."
)
# ── Invariant C — every memory-search call passes a MemoryQuery ──────
def _find_search_and_format_context_calls_without_query(
file_path: Path,
) -> list[tuple[int, str]]:
"""Find ``search_and_format_context(...)`` invocations that DON'T
pass a ``query=`` kwarg.
Pre-PR-this no site passed a query — they all relied on the
handler's internal ``_extract_user_query(messages)`` with its
500-char truncation. The new contract: every handler builds a
full-fidelity ``MemoryQuery`` and passes it explicitly.
"""
text = file_path.read_text(encoding="utf-8")
tree = ast.parse(text, filename=str(file_path))
hits = []
for node in ast.walk(tree):
if not isinstance(node, ast.Call):
continue
if not isinstance(node.func, ast.Attribute):
continue
if node.func.attr != "search_and_format_context":
continue
kwarg_keys = {kw.arg for kw in node.keywords if kw.arg is not None}
if "query" not in kwarg_keys:
line_no = node.lineno
snippet = text.splitlines()[line_no - 1].strip()
hits.append((line_no, snippet))
return hits
def test_every_search_and_format_context_call_passes_query_kwarg() -> None:
"""Every handler that searches memory MUST pass a ``query=`` kwarg
(a :class:`MemoryQuery` instance), not rely on the handler's
internal ``_extract_user_query`` (which used to truncate to 500
chars). Locks the full-fidelity-query contract."""
offenders = []
for f in HANDLER_FILES:
offenders.extend(
[(f, ln, src) for (ln, src) in _find_search_and_format_context_calls_without_query(f)]
)
if offenders:
formatted = "\n".join(f" {f.name}:{ln} {src!r}" for f, ln, src in offenders)
pytest.fail(
f"{len(offenders)} search_and_format_context call(s) miss `query=`:\n"
f"{formatted}\n\n"
"Pass `query=MemoryQuery.from_messages(...)` — the multi-source, "
"untruncated query value type. See headroom/proxy/memory_query.py."
)

220
tests/test_memory_query.py Normal file
View file

@ -0,0 +1,220 @@
"""Tests for :class:`headroom.proxy.memory_query.MemoryQuery`.
``MemoryQuery`` is the multi-source query value type that replaces
the pre-PR pattern of "use the latest user message, truncated to 500
chars". The truncation was a real bug — none of Letta/Mem0/Cognee/
Supermemory truncate the embedding input.
The query is built from three sources, all preserved at full fidelity:
* ``user_text`` — latest user message, untruncated
* ``recent_tool_outputs`` — last N tool results (often the most
relevant signal in coding sessions)
* ``recent_assistant_turns`` — last K assistant turns for intent
Building the embedding input is a simple concatenation with delimiters
so the embedding model sees structured context, not a wall of text.
"""
from __future__ import annotations
from dataclasses import FrozenInstanceError
from headroom.proxy.memory_query import MemoryQuery
# ── Value-type contract ───────────────────────────────────────────────
def test_memory_query_is_frozen() -> None:
q = MemoryQuery(
user_text="hello",
recent_tool_outputs=(),
recent_assistant_turns=(),
conversation_id=None,
)
try:
q.user_text = "mutated" # type: ignore[misc]
except FrozenInstanceError:
pass
else:
raise AssertionError("MemoryQuery must be frozen")
def test_memory_query_value_equal() -> None:
a = MemoryQuery(
user_text="hi", recent_tool_outputs=(), recent_assistant_turns=(), conversation_id="c1"
)
b = MemoryQuery(
user_text="hi", recent_tool_outputs=(), recent_assistant_turns=(), conversation_id="c1"
)
assert a == b
# ── NO TRUNCATION — the entire point of this type ────────────────────
def test_full_user_message_is_preserved_no_500_char_cap() -> None:
"""Pre-PR: ``_extract_user_query`` capped at 500 chars. None of
the four memory systems we surveyed truncate. MemoryQuery must
preserve the full message — embedding models handle their own
window (MiniLM 512 tok; BGE-small 8K tok)."""
long_msg = "a" * 8000 # 8KB user message
q = MemoryQuery(
user_text=long_msg,
recent_tool_outputs=(),
recent_assistant_turns=(),
conversation_id=None,
)
embedding_input = q.to_embedding_input()
# Original content fully present — count actual occurrences of "a" run.
assert "a" * 8000 in embedding_input
def test_tool_outputs_preserved_at_full_fidelity() -> None:
"""Tool results — often the strongest retrieval signal in coding
sessions — must NOT be truncated."""
big_tool_output = "GREP RESULT\n" + "match line\n" * 1000 # large grep output
q = MemoryQuery(
user_text="how do I fix this?",
recent_tool_outputs=(big_tool_output,),
recent_assistant_turns=(),
conversation_id=None,
)
embedding_input = q.to_embedding_input()
assert "match line" * 1000 in embedding_input.replace("\n", "")
# ── Multi-source query construction ──────────────────────────────────
def test_embedding_input_includes_all_sources() -> None:
"""The query the embedder sees should include user msg + recent
tool outputs + recent assistant turns. Each source is delimited
so the embedder treats them as distinct context, not run-on text."""
q = MemoryQuery(
user_text="fix the auth bug",
recent_tool_outputs=("auth.py:42: KeyError",),
recent_assistant_turns=("I'll look at the auth flow",),
conversation_id=None,
)
txt = q.to_embedding_input()
assert "fix the auth bug" in txt
assert "auth.py:42: KeyError" in txt
assert "I'll look at the auth flow" in txt
def test_empty_sources_still_produce_valid_query() -> None:
"""A user-msg-only query (no tools, no prior assistant) is the
minimum viable case — common on first turn."""
q = MemoryQuery(
user_text="hello",
recent_tool_outputs=(),
recent_assistant_turns=(),
conversation_id=None,
)
txt = q.to_embedding_input()
assert txt
assert "hello" in txt
def test_empty_user_text_is_valid_when_only_tool_signal() -> None:
"""Edge case: agent-driven request with no new user text (e.g. a
tool-call follow-up). Query is the tool output."""
q = MemoryQuery(
user_text="",
recent_tool_outputs=("ls -la /home/user/projects/headroom",),
recent_assistant_turns=(),
conversation_id=None,
)
txt = q.to_embedding_input()
assert "ls -la /home/user/projects/headroom" in txt
# ── from_messages constructor ────────────────────────────────────────
def test_from_messages_extracts_latest_user_text() -> None:
"""Construct from a chat-style messages list — picks the most
recent ``role: user`` content."""
messages = [
{"role": "user", "content": "first turn"},
{"role": "assistant", "content": "ack"},
{"role": "user", "content": "second turn"},
]
q = MemoryQuery.from_messages(messages, lookback_assistant=0, lookback_tools=0)
assert q.user_text == "second turn"
def test_from_messages_extracts_recent_assistant_turns_in_order() -> None:
"""Recent assistant turns are pulled in chronological order
(oldest of the lookback window first, latest last)."""
messages = [
{"role": "user", "content": "u1"},
{"role": "assistant", "content": "a1"},
{"role": "user", "content": "u2"},
{"role": "assistant", "content": "a2"},
{"role": "user", "content": "u3"},
]
q = MemoryQuery.from_messages(messages, lookback_assistant=2, lookback_tools=0)
assert q.recent_assistant_turns == ("a1", "a2")
assert q.user_text == "u3"
def test_from_messages_caps_assistant_lookback() -> None:
"""``lookback_assistant=K`` keeps only the K most recent assistant
turns. With lookback=1 and three assistant turns, only the latest."""
messages = [
{"role": "assistant", "content": "a1"},
{"role": "assistant", "content": "a2"},
{"role": "assistant", "content": "a3"},
{"role": "user", "content": "u"},
]
q = MemoryQuery.from_messages(messages, lookback_assistant=1, lookback_tools=0)
assert q.recent_assistant_turns == ("a3",)
def test_from_messages_extracts_tool_outputs() -> None:
"""Tool results are pulled from ``role: tool`` messages (OpenAI
shape) — pre-PR these never participated in retrieval at all."""
messages = [
{"role": "user", "content": "list files"},
{"role": "assistant", "content": "I'll run ls"},
{"role": "tool", "content": "main.py\nREADME.md\n"},
{"role": "user", "content": "now read main.py"},
]
q = MemoryQuery.from_messages(messages, lookback_assistant=0, lookback_tools=2)
assert q.recent_tool_outputs == ("main.py\nREADME.md\n",)
def test_from_messages_handles_anthropic_tool_result_shape() -> None:
"""Anthropic shape: tool_result inside the user message as a
content block. The constructor should still extract it."""
messages = [
{"role": "user", "content": "go"},
{
"role": "user",
"content": [
{"type": "tool_result", "tool_use_id": "x", "content": "ANTHROPIC_TOOL_OUTPUT"}
],
},
{"role": "user", "content": "thanks"},
]
q = MemoryQuery.from_messages(messages, lookback_assistant=0, lookback_tools=2)
assert "ANTHROPIC_TOOL_OUTPUT" in q.recent_tool_outputs
def test_from_messages_empty_returns_empty_query() -> None:
"""No messages → empty query, no exception."""
q = MemoryQuery.from_messages([], lookback_assistant=2, lookback_tools=2)
assert q.user_text == ""
assert q.recent_assistant_turns == ()
assert q.recent_tool_outputs == ()
def test_from_messages_handles_assistant_only_messages() -> None:
"""Edge case: no user messages at all (rare; agent-driven). Should
still build a valid query."""
messages = [{"role": "assistant", "content": "assistant only"}]
q = MemoryQuery.from_messages(messages, lookback_assistant=2, lookback_tools=0)
assert q.user_text == ""
assert q.recent_assistant_turns == ("assistant only",)