diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 4d04f9d3a..e062040da 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -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 diff --git a/headroom/proxy/handlers/gemini.py b/headroom/proxy/handlers/gemini.py index e105e9bfd..0867eafae 100644 --- a/headroom/proxy/handlers/gemini.py +++ b/headroom/proxy/handlers/gemini.py @@ -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 = ( diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 04f27cc6a..10be7d828 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -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 diff --git a/headroom/proxy/memory_decision.py b/headroom/proxy/memory_decision.py new file mode 100644 index 000000000..02eaf4b3b --- /dev/null +++ b/headroom/proxy/memory_decision.py @@ -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 diff --git a/headroom/proxy/memory_handler.py b/headroom/proxy/memory_handler.py index 641de9caa..07f9e6f13 100644 --- a/headroom/proxy/memory_handler.py +++ b/headroom/proxy/memory_handler.py @@ -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 "", 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 "" diff --git a/headroom/proxy/memory_injection.py b/headroom/proxy/memory_injection.py new file mode 100644 index 000000000..459b83889 --- /dev/null +++ b/headroom/proxy/memory_injection.py @@ -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] diff --git a/headroom/proxy/memory_query.py b/headroom/proxy/memory_query.py new file mode 100644 index 000000000..f052f0a63 --- /dev/null +++ b/headroom/proxy/memory_query.py @@ -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, + ) diff --git a/tests/test_memory_decision.py b/tests/test_memory_decision.py new file mode 100644 index 000000000..5080f2422 --- /dev/null +++ b/tests/test_memory_decision.py @@ -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 diff --git a/tests/test_memory_injection_budget.py b/tests/test_memory_injection_budget.py new file mode 100644 index 000000000..87b0cd24f --- /dev/null +++ b/tests/test_memory_injection_budget.py @@ -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"} diff --git a/tests/test_memory_invariants.py b/tests/test_memory_invariants.py new file mode 100644 index 000000000..b7b737c65 --- /dev/null +++ b/tests/test_memory_invariants.py @@ -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." + ) diff --git a/tests/test_memory_query.py b/tests/test_memory_query.py new file mode 100644 index 000000000..02af32236 --- /dev/null +++ b/tests/test_memory_query.py @@ -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",)