feat(proxy): lazy memory embedder warmup + /debug/rss endpoint (memory-arch-v2)

docs/content/docs/configuration.mdx

@@ -341,6 +341,7 @@ headroom proxy --learn --min-evidence 3
 | `HEADROOM_SUBSCRIPTION_STATE_PATH` | Override subscription tracker state file. Always wins when set. | derived from `${HEADROOM_WORKSPACE_DIR}` |
 | `HEADROOM_PERIODIC_TOIN_STATS` | Controls periodic TOIN stats logging in long-lived proxy workers. Set to `0`, `false`, `off`, or `no` to disable the 5-minute stats loop without disabling TOIN learning or request-time feedback. | `true` |
 | `HEADROOM_MEMORY_INJECTION_MODE` | Memory-context routing mode: `live_zone_tail` (default) or `disabled`. The legacy `system_prompt` mode was retired by PR-A2; supplying it raises. | `live_zone_tail` |
+| `HEADROOM_SKIP_MEMORY_WARMUP` | Set to `1`, `true`, `yes`, or `on` to load the memory embedder lazily on the first memory request instead of during proxy startup. | `off` |
 | `HEADROOM_PROXY_PYTHON_FORWARDER_MODE` | Python forwarder serialization mode. `byte_faithful` (default) forwards original request bytes verbatim when no transform mutated the body and re-serializes canonically only when needed — keeps Anthropic prompt-cache hit-rate intact. `legacy_json_kwarg` is an explicit operator opt-in for emergency rollback to the historical `httpx ... json=body` behavior. NOT a fallback — only flip on explicit operator decision. | `byte_faithful` |
 | `HEADROOM_STRIP_INTERNAL_HEADERS` | Python proxy: whether to strip internal `x-headroom-*` request headers (e.g. `x-headroom-bypass`, `x-headroom-mode`, `x-headroom-user-id`, `x-headroom-stack`, `x-headroom-base-url`) before every upstream forwarder call (PR-A5, fixes P5-49). `enabled` (default) stops fingerprinting / leakage. `disabled` is an explicit operator opt-in for diagnostic shadow tracing — NOT a fallback. Inbound reads of these headers (bypass gating, memory user-id resolution) are unaffected because they read `request.headers` directly. | `enabled` |
 | `HEADROOM_PROXY_STRIP_INTERNAL_HEADERS` | Rust proxy: same policy as `HEADROOM_STRIP_INTERNAL_HEADERS` but for the Rust transparent proxy. Stripping happens inside `build_forward_request_headers` so both HTTP and WebSocket upstream calls are gated by one flag. `enabled` default; `disabled` operator opt-in for diagnostic shadow tracing. Response-side `X-Headroom-*` injection (e.g. `x-headroom-tokens-saved`) is unrelated and stays. | `enabled` |

headroom/proxy/server.py

@@ -27,12 +27,14 @@
 import asyncio
 import concurrent.futures
 import contextlib
+import gc
 import hmac
 import ipaddress
 import json
 import logging
 import math
 import os
+import platform
 import sys
 import threading
 import time
@@ -1886,17 +1888,23 @@ async def startup(self):
                     handle=self.memory_handler,
                     backend=memory_status.get("backend"),
                 )
-                # Force one embed call so the ONNX graph is compiled now,
-                # not lazily during the first request. Best-effort — any
-                # failure is swallowed inside warmup_embedder.
-                self.warmup.memory_embedder.mark_loading()
-                warmed = await self.memory_handler.warmup_embedder()
-                if warmed:
-                    self.warmup.memory_embedder.mark_loaded()
-                else:
-                    # Not an error — e.g. qdrant-neo4j has no embedder slot
-                    # we can reach, or the backend simply exposes no handle.
+                skip_embedder_warmup = os.environ.get(
+                    "HEADROOM_SKIP_MEMORY_WARMUP", ""
+                ).strip().lower() in {"1", "true", "yes", "on"}
+                if skip_embedder_warmup:
                     self.warmup.memory_embedder.mark_null()
+                    logger.info(
+                        "Memory embedder warmup: SKIPPED; first memory request will load it lazily"
+                    )
+                else:
+                    # Compile the ONNX graph before the first real request.
+                    # Best-effort: warmup_embedder handles failures internally.
+                    self.warmup.memory_embedder.mark_loading()
+                    warmed = await self.memory_handler.warmup_embedder()
+                    if warmed:
+                        self.warmup.memory_embedder.mark_loaded()
+                    else:
+                        self.warmup.memory_embedder.mark_null()
             else:
                 if self.warmup.memory_backend.status != "error":
                     self.warmup.memory_backend.mark_null()
@@ -2359,6 +2367,39 @@ def _register_memory_components(proxy: HeadroomProxy, tracker: MemoryTracker) ->
     # registered when the memory system is initialized with specific backends.
 
 
+def _peak_rss_mb() -> float | None:
+    """Return the process RSS high-water mark when the platform exposes it."""
+    try:
+        import importlib
+
+        resource_module: Any = importlib.import_module("resource")
+    except ImportError:  # Windows
+        return None
+
+    peak = float(resource_module.getrusage(resource_module.RUSAGE_SELF).ru_maxrss)
+    divisor = 1024 * 1024 if platform.system() == "Darwin" else 1024
+    return round(peak / divisor, 1)
+
+
+def _gc_snapshot() -> tuple[list[dict[str, int]], list[dict[str, int | str]]]:
+    """Inspect tracked objects without triggering a stop-the-world collection."""
+    from collections import Counter
+
+    stats: list[dict[str, int]] = [
+        {
+            "generation": generation,
+            "collections": int(values["collections"]),
+            "collected": int(values["collected"]),
+        }
+        for generation, values in enumerate(gc.get_stats())
+    ]
+    counts = Counter(type(obj).__name__ for obj in gc.get_objects())
+    top_types: list[dict[str, int | str]] = [
+        {"type": name, "count": count} for name, count in counts.most_common(10)
+    ]
+    return stats, top_types
+
+
 def _request_is_loopback(request: Request) -> bool:
     """Return True iff the caller is on loopback by *both* peer IP and Host header.
 
@@ -4652,6 +4693,43 @@ async def debug_memory():
         report = tracker.get_report()
         return report.to_dict()
 
+    @app.get("/debug/rss", dependencies=[Depends(_require_loopback)])
+    async def debug_rss():
+        """Return a loopback-only process and Python heap diagnostic snapshot."""
+        from ..memory.tracker import MemoryTracker
+        from ..models.ml_models import MLModelRegistry
+        from ..telemetry.toin import get_toin
+
+        tracker = MemoryTracker.get()
+        _register_memory_components(proxy, tracker)
+        report = tracker.get_report()
+        gc_stats, top_types = await asyncio.to_thread(_gc_snapshot)
+
+        request_logger = report.components.get("request_logger")
+        try:
+            toin_patterns = get_toin().get_stats().get("patterns_tracked")
+        except Exception:  # pragma: no cover - diagnostic endpoint is best-effort
+            toin_patterns = None
+
+        with proxy._compression_caches_lock:
+            compression_cache_sessions = len(proxy._compression_caches)
+
+        return {
+            "pid": os.getpid(),
+            "rss_mb": report.process.to_dict()["rss_mb"],
+            "peak_rss_mb": _peak_rss_mb(),
+            "python_version": sys.version,
+            "gc_stats": gc_stats,
+            "top_types": top_types,
+            "ml_models": MLModelRegistry.get_memory_stats(),
+            # Per-user vector indexes are not globally enumerable on current backends.
+            "hnsw_elements": None,
+            "compression_cache_sessions": compression_cache_sessions,
+            "toin_patterns": toin_patterns,
+            "request_log_count": request_logger.entry_count if request_logger else None,
+            "memory_embedder_warmed": proxy.warmup.memory_embedder.status == "loaded",
+        }
+
     @app.post(
         "/cache/clear",
         dependencies=[Depends(_require_loopback), Depends(_require_same_origin)],

tests/test_memory_arch_v2.py

@@ -0,0 +1,117 @@
+"""Regression tests for lazy memory warmup and RSS diagnostics."""
+
+from __future__ import annotations
+
+import builtins
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+
+pytest.importorskip("fastapi")
+pytest.importorskip("httpx")
+
+from fastapi.testclient import TestClient
+
+from headroom.proxy.server import (
+    HeadroomProxy,
+    ProxyConfig,
+    _gc_snapshot,
+    _peak_rss_mb,
+    create_app,
+)
+
+
+def _proxy_config() -> ProxyConfig:
+    return ProxyConfig(
+        optimize=False,
+        cache_enabled=False,
+        rate_limit_enabled=False,
+        cost_tracking_enabled=False,
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("value", ["1", "true", "YES", "on"])
+async def test_skip_memory_warmup_truthy_values(monkeypatch, value: str) -> None:
+    monkeypatch.setenv("HEADROOM_SKIP_MEMORY_WARMUP", value)
+    proxy = HeadroomProxy(_proxy_config())
+    handler = SimpleNamespace(
+        ensure_initialized=AsyncMock(),
+        warmup_embedder=AsyncMock(return_value=True),
+        health_status=lambda: {"initialized": True, "backend": "local"},
+    )
+    proxy.memory_handler = handler
+
+    await proxy.startup()
+
+    handler.warmup_embedder.assert_not_awaited()
+    assert proxy.warmup.memory_embedder.status == "null"
+
+
+@pytest.mark.asyncio
+async def test_memory_warmup_remains_eager_by_default(monkeypatch) -> None:
+    monkeypatch.delenv("HEADROOM_SKIP_MEMORY_WARMUP", raising=False)
+    proxy = HeadroomProxy(_proxy_config())
+    handler = SimpleNamespace(
+        ensure_initialized=AsyncMock(),
+        warmup_embedder=AsyncMock(return_value=True),
+        health_status=lambda: {"initialized": True, "backend": "local"},
+    )
+    proxy.memory_handler = handler
+
+    await proxy.startup()
+
+    handler.warmup_embedder.assert_awaited_once_with()
+    assert proxy.warmup.memory_embedder.status == "loaded"
+
+
+def test_peak_rss_is_optional_without_resource(monkeypatch) -> None:
+    real_import = builtins.__import__
+
+    def import_without_resource(name, *args, **kwargs):
+        if name == "resource":
+            raise ImportError
+        return real_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", import_without_resource)
+    assert _peak_rss_mb() is None
+
+
+def test_gc_snapshot_does_not_force_collection(monkeypatch) -> None:
+    monkeypatch.setattr(
+        "headroom.proxy.server.gc.collect", lambda: pytest.fail("gc.collect called")
+    )
+    stats, top_types = _gc_snapshot()
+    assert stats
+    assert all({"generation", "collections", "collected"} <= row.keys() for row in stats)
+    assert len(top_types) <= 10
+
+
+def test_debug_rss_schema_and_loopback_guard() -> None:
+    app = create_app(_proxy_config())
+    with TestClient(app, client=("127.0.0.1", 12345)) as client:
+        response = client.get("/debug/rss", headers={"host": "127.0.0.1"})
+
+    assert response.status_code == 200
+    data = response.json()
+    assert {
+        "pid",
+        "rss_mb",
+        "peak_rss_mb",
+        "python_version",
+        "gc_stats",
+        "top_types",
+        "ml_models",
+        "hnsw_elements",
+        "compression_cache_sessions",
+        "toin_patterns",
+        "request_log_count",
+        "memory_embedder_warmed",
+    } <= data.keys()
+    assert data["pid"] > 0
+    assert isinstance(data["memory_embedder_warmed"], bool)
+
+    with TestClient(app, client=("10.0.0.1", 12345)) as client:
+        denied = client.get("/debug/rss", headers={"host": "127.0.0.1"})
+    assert denied.status_code == 404