headroom/tests/test_proxy_openai_responses_bypass.py
chopratejas 7694f050fe fix: per-project memory storage so projects can no longer bleed memories (GH #462)
Memory retrieval was partitioned only by `x-headroom-user-id`. Claude
Code never sets that header, so every project a user worked on landed
in one global `default` bucket; the proxy then injected semantically
similar memories from that mixed bucket into every `/v1/messages`
request, regardless of which repo the session was actually about. The
injected `## Relevant Memories` block reads like a prompt-injection
payload and Claude has been seen to refuse to act on it, defeating the
feature.

This change makes leakage structurally impossible by giving each
resolved workspace its own SQLite database file. The wrong DB is
simply not open during a request.

- `headroom/memory/storage_router.py` (new) — `MemoryStorageMode`
  (project/user/global), `ProjectResolver` (x-headroom-project-id →
  x-headroom-cwd → --memory-project-root CLI override → env-block
  parse: `Primary working directory:` / `Working directory:` / `cwd:`,
  no regex), and `BackendRouter` with an LRU of open `LocalBackend`s
  keyed by db_path.
- `proxy/memory_handler.py` — `MemoryConfig.storage_mode` defaults to
  `PROJECT`. Provider handlers build a `RequestContext` once and pass
  it through; `search_and_format_context`, `handle_memory_tool_calls`,
  and the `_execute_*` methods route save/search/update/delete on the
  per-project backend. Qdrant-neo4j gets a composite
  `user::project_key` partition so external Mem0-style deployments
  also isolate per project without a parallel collection.
- Fix C — injected block carries provenance:
  `## Relevant Memories (workspace: <basename>, scope: project)`.
  CCR proactive-expansion block gets a matching workspace tag.
- `memory/factory.py` — process-wide embedder cache so opening N
  project DBs doesn't load the embedder N times. OpenAI key
  validation runs ahead of the cache.
- CLI — `--memory-storage={project,user,global}` (default `project`),
  `--memory-project-root` override, rewritten `--memory` help text,
  banner reports storage mode.
- Migration UX — if the legacy single-file DB has content while
  project mode is active, an INFO log points users at
  `--memory-storage=global`. Bridge currently only syncs the legacy
  DB; a WARN fires when bridge + project mode are combined.

Backward-compatible: legacy `~/.headroom/memory.db` untouched and
reachable via `--memory-storage=global`. `request_context` is
keyword-only on entry points so existing tests/mocks keep working.

Tests: 24 new (resolver tiers, LRU eviction, two-cwd isolation,
user-mode partition, legacy fallback, provenance headers); full
suite 5260 passing, ci-precheck green.
2026-05-13 15:27:41 -07:00

114 lines
3.3 KiB
Python

from __future__ import annotations
from types import SimpleNamespace
from typing import Any
import httpx
import pytest
pytest.importorskip("fastapi")
from fastapi.testclient import TestClient # noqa: E402
from headroom.proxy.loopback_guard import require_loopback # noqa: E402
from headroom.proxy.server import ProxyConfig, create_app # noqa: E402
class _MemoryHandler:
def __init__(self) -> None:
self.search_calls = 0
self.tool_calls = 0
self.config = SimpleNamespace(inject_context=True, inject_tools=True)
async def search_and_format_context(
self,
user_id: str,
messages: list[dict[str, Any]],
**_kwargs: Any,
) -> str:
self.search_calls += 1
return "memory context that must not be injected"
def compute_memory_tool_definitions(self, provider: str) -> list[dict[str, Any]]:
self.tool_calls += 1
return [
{
"type": "function",
"function": {
"name": "memory_search",
"description": "search memory",
"parameters": {"type": "object", "properties": {}},
},
}
]
def has_memory_tool_calls(self, response: dict[str, Any], provider: str) -> bool:
return False
def test_responses_bypass_skips_memory_and_compression_mutation() -> None:
app = create_app(
ProxyConfig(
optimize=True,
cache_enabled=False,
rate_limit_enabled=False,
cost_tracking_enabled=False,
log_requests=False,
)
)
app.dependency_overrides[require_loopback] = lambda: None
original_input = [
{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "summarize"}],
},
{
"type": "function_call_output",
"call_id": "call_1",
"output": "large tool output " * 200,
},
]
captured: dict[str, Any] = {}
with TestClient(app) as client:
proxy = client.app.state.proxy
memory_handler = _MemoryHandler()
proxy.memory_handler = memory_handler
async def _fake_retry(
method: str,
url: str,
headers: dict[str, str],
body: dict[str, Any],
stream: bool = False,
**kwargs: Any,
) -> httpx.Response:
captured["body"] = body
return httpx.Response(
200,
json={
"id": "resp_1",
"output": [],
"usage": {"input_tokens": 10, "output_tokens": 1},
},
)
proxy._retry_request = _fake_retry
response = client.post(
"/v1/responses",
headers={
"authorization": "Bearer test-key",
"x-headroom-bypass": "true",
"x-headroom-user-id": "user-1",
},
json={"model": "gpt-4o-mini", "input": original_input},
)
assert response.status_code == 200
assert captured["body"]["input"] == original_input
assert "tools" not in captured["body"]
assert memory_handler.search_calls == 0
assert memory_handler.tool_calls == 0