chore: import upstream snapshot with attribution

tests/test_proxy_memory_integration.py

@@ -0,0 +1,698 @@
+"""Integration tests for proxy memory system with real API calls.
+
+These tests require:
+- ANTHROPIC_API_KEY environment variable set
+
+Run with:
+    ANTHROPIC_API_KEY=... uv run pytest tests/test_proxy_memory_integration.py -v
+
+Test categories:
+- TestMemoryHeaderValidation: User ID header validation
+- TestMemoryToolInjection: Memory tools are injected
+- TestMemorySaveAndSearch: End-to-end save/recall flow
+- TestMemoryUserIsolation: User memory isolation
+"""
+
+import os
+import tempfile
+import time
+from pathlib import Path
+
+import pytest
+
+# Set tokenizer parallelism before importing transformers
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+pytest.importorskip("fastapi")
+pytest.importorskip("httpx")
+
+from fastapi.testclient import TestClient
+
+from headroom.proxy.server import ProxyConfig, create_app
+
+
+@pytest.fixture
+def temp_memory_db():
+    """Create temporary memory database."""
+    with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
+        yield f.name
+    # Cleanup
+    Path(f.name).unlink(missing_ok=True)
+    # Also cleanup related files (HNSW index, etc.)
+    for suffix in ["-shm", "-wal", ".hnsw"]:
+        Path(f.name + suffix).unlink(missing_ok=True)
+
+
+@pytest.fixture
+def memory_client(temp_memory_db):
+    """Create test client with memory enabled."""
+    config = ProxyConfig(
+        optimize=False,  # Disable optimization for simpler tests
+        cache_enabled=False,
+        rate_limit_enabled=False,
+        cost_tracking_enabled=False,
+        memory_enabled=True,
+        memory_backend="local",
+        memory_db_path=temp_memory_db,
+        memory_inject_tools=True,
+        memory_inject_context=True,
+        memory_top_k=5,
+    )
+    app = create_app(config)
+    with TestClient(app) as client:
+        yield client
+
+
+@pytest.fixture
+def no_memory_client():
+    """Create test client with memory disabled."""
+    config = ProxyConfig(
+        optimize=False,
+        cache_enabled=False,
+        rate_limit_enabled=False,
+        cost_tracking_enabled=False,
+        memory_enabled=False,
+    )
+    app = create_app(config)
+    with TestClient(app) as client:
+        yield client
+
+
+@pytest.fixture
+def anthropic_api_key():
+    """Get Anthropic API key from environment."""
+    return os.environ.get("ANTHROPIC_API_KEY")
+
+
+class TestMemoryHeaderValidation:
+    """Test user ID header validation."""
+
+    def test_missing_user_id_uses_default(self, memory_client, anthropic_api_key):
+        """Request without x-headroom-user-id should use 'default' user for simple DevEx."""
+        if not anthropic_api_key:
+            pytest.skip("ANTHROPIC_API_KEY not set")
+
+        response = memory_client.post(
+            "/v1/messages",
+            headers={
+                "x-api-key": anthropic_api_key,
+                "anthropic-version": "2023-06-01",
+                # Note: NOT setting x-headroom-user-id - should default to "default"
+            },
+            json={
+                "model": "claude-sonnet-4-20250514",
+                "max_tokens": 100,
+                "messages": [{"role": "user", "content": "Hello"}],
+            },
+        )
+        # Should succeed, not return 400
+        assert response.status_code == 200
+
+    def test_with_user_id_succeeds(self, memory_client, anthropic_api_key):
+        """Request with x-headroom-user-id should succeed."""
+        if not anthropic_api_key:
+            pytest.skip("ANTHROPIC_API_KEY not set")
+
+        response = memory_client.post(
+            "/v1/messages",
+            headers={
+                "x-api-key": anthropic_api_key,
+                "anthropic-version": "2023-06-01",
+                "x-headroom-user-id": "test-user-123",
+            },
+            json={
+                "model": "claude-sonnet-4-20250514",
+                "max_tokens": 100,
+                "messages": [{"role": "user", "content": "Hello, just say hi back."}],
+            },
+        )
+        assert response.status_code == 200
+
+    def test_no_memory_client_doesnt_require_user_id(self, no_memory_client, anthropic_api_key):
+        """When memory is disabled, user ID header should not be required."""
+        if not anthropic_api_key:
+            pytest.skip("ANTHROPIC_API_KEY not set")
+
+        response = no_memory_client.post(
+            "/v1/messages",
+            headers={
+                "x-api-key": anthropic_api_key,
+                "anthropic-version": "2023-06-01",
+                # No x-headroom-user-id
+            },
+            json={
+                "model": "claude-sonnet-4-20250514",
+                "max_tokens": 100,
+                "messages": [{"role": "user", "content": "Hello, just say hi."}],
+            },
+        )
+        assert response.status_code == 200
+
+
+@pytest.mark.skipif(not os.environ.get("ANTHROPIC_API_KEY"), reason="ANTHROPIC_API_KEY not set")
+class TestMemoryToolInjection:
+    """Test memory tool injection."""
+
+    def test_memory_tools_are_available(self, memory_client, anthropic_api_key):
+        """Memory tools should be available to the LLM."""
+        response = memory_client.post(
+            "/v1/messages",
+            headers={
+                "x-api-key": anthropic_api_key,
+                "anthropic-version": "2023-06-01",
+                "x-headroom-user-id": "test-user-tool-check",
+            },
+            json={
+                "model": "claude-sonnet-4-20250514",
+                "max_tokens": 500,
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": "List the tools available to you. Just list the tool names.",
+                    }
+                ],
+            },
+        )
+        assert response.status_code == 200
+
+        # The response should mention memory tools
+        content = response.json().get("content", [])
+        text = ""
+        for block in content:
+            if isinstance(block, dict) and block.get("type") == "text":
+                text += block.get("text", "")
+
+        # At least one memory tool should be mentioned
+        assert any(tool in text.lower() for tool in ["memory_save", "memory_search", "memory"]), (
+            f"Memory tools not found in response: {text}"
+        )
+
+
+@pytest.mark.skipif(not os.environ.get("ANTHROPIC_API_KEY"), reason="ANTHROPIC_API_KEY not set")
+class TestMemorySaveAndSearch:
+    """Test memory save and search flow."""
+
+    def test_save_memory_via_explicit_instruction(self, memory_client, anthropic_api_key):
+        """LLM should be able to save memories when instructed."""
+        user_id = f"test-user-save-{int(time.time())}"
+
+        # Request that explicitly asks to save
+        response = memory_client.post(
+            "/v1/messages",
+            headers={
+                "x-api-key": anthropic_api_key,
+                "anthropic-version": "2023-06-01",
+                "x-headroom-user-id": user_id,
+            },
+            json={
+                "model": "claude-sonnet-4-20250514",
+                "max_tokens": 1000,
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": "Please save this to memory: My favorite programming language is Rust. "
+                        "Use the memory_save tool to save this information.",
+                    }
+                ],
+            },
+        )
+        assert response.status_code == 200
+
+        # Check if response indicates tool was used
+        resp_json = response.json()
+        content = resp_json.get("content", [])
+
+        # Response could be tool_use (if not handled) or text (if handled)
+        # Either way, it should complete successfully
+        assert content, "Response should have content"
+
+    def test_save_and_recall_memory(self, memory_client, anthropic_api_key):
+        """Save a memory and recall it in subsequent request."""
+        user_id = f"test-user-recall-{int(time.time())}"
+
+        # First request: save a memory with explicit instruction
+        save_response = memory_client.post(
+            "/v1/messages",
+            headers={
+                "x-api-key": anthropic_api_key,
+                "anthropic-version": "2023-06-01",
+                "x-headroom-user-id": user_id,
+            },
+            json={
+                "model": "claude-sonnet-4-20250514",
+                "max_tokens": 1000,
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": "Please remember this: My name is TestUser and I work at AcmeCorp. "
+                        "Save this information using the memory_save tool.",
+                    }
+                ],
+            },
+        )
+        assert save_response.status_code == 200
+
+        # Wait a moment for memory to be indexed
+        time.sleep(1)
+
+        # Second request: ask about saved info
+        # Memory context should be injected automatically
+        recall_response = memory_client.post(
+            "/v1/messages",
+            headers={
+                "x-api-key": anthropic_api_key,
+                "anthropic-version": "2023-06-01",
+                "x-headroom-user-id": user_id,
+            },
+            json={
+                "model": "claude-sonnet-4-20250514",
+                "max_tokens": 300,
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": "What is my name and where do I work? "
+                        "Answer based on what you know about me.",
+                    }
+                ],
+            },
+        )
+        assert recall_response.status_code == 200
+
+        # Check if response mentions the saved info
+        content = recall_response.json().get("content", [])
+        text = ""
+        for block in content:
+            if isinstance(block, dict) and block.get("type") == "text":
+                text += block.get("text", "")
+
+        # Should mention at least one of the saved facts
+        text_lower = text.lower()
+        assert "testuser" in text_lower or "acmecorp" in text_lower or "acme" in text_lower, (
+            f"Saved info not recalled: {text}"
+        )
+
+
+@pytest.mark.skipif(not os.environ.get("ANTHROPIC_API_KEY"), reason="ANTHROPIC_API_KEY not set")
+class TestMemoryUserIsolation:
+    """Test that memories are isolated per user."""
+
+    def test_different_users_have_isolated_memories(self, memory_client, anthropic_api_key):
+        """User A's memories should not appear for User B."""
+        timestamp = int(time.time())
+        user_a = f"user-a-isolation-{timestamp}"
+        user_b = f"user-b-isolation-{timestamp}"
+        secret_code = f"SECRETCODE{timestamp}"
+
+        # Save memory for user A
+        save_response = memory_client.post(
+            "/v1/messages",
+            headers={
+                "x-api-key": anthropic_api_key,
+                "anthropic-version": "2023-06-01",
+                "x-headroom-user-id": user_a,
+            },
+            json={
+                "model": "claude-sonnet-4-20250514",
+                "max_tokens": 1000,
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": f"Remember my secret code: {secret_code}. "
+                        "Save this using the memory_save tool.",
+                    }
+                ],
+            },
+        )
+        assert save_response.status_code == 200
+
+        # Wait for memory to be indexed
+        time.sleep(1)
+
+        # Query as user B - should NOT have access to user A's memory
+        response_b = memory_client.post(
+            "/v1/messages",
+            headers={
+                "x-api-key": anthropic_api_key,
+                "anthropic-version": "2023-06-01",
+                "x-headroom-user-id": user_b,
+            },
+            json={
+                "model": "claude-sonnet-4-20250514",
+                "max_tokens": 300,
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": "What is my secret code? Search your memory for it.",
+                    }
+                ],
+            },
+        )
+        assert response_b.status_code == 200
+
+        # User B should NOT see user A's secret code
+        content = response_b.json().get("content", [])
+        text = ""
+        for block in content:
+            if isinstance(block, dict) and block.get("type") == "text":
+                text += block.get("text", "")
+
+        assert secret_code not in text, f"User B should not see User A's secret: {text}"
+
+
+@pytest.mark.skipif(not os.environ.get("ANTHROPIC_API_KEY"), reason="ANTHROPIC_API_KEY not set")
+class TestMemoryStats:
+    """Test memory-related stats and health."""
+
+    def test_health_endpoint_works_with_memory(self, memory_client):
+        """Health endpoint should work when memory is enabled."""
+        response = memory_client.get("/health")
+        assert response.status_code == 200
+        data = response.json()
+        assert data.get("status") == "healthy"
+
+    def test_stats_endpoint_works_with_memory(self, memory_client):
+        """Stats endpoint should work when memory is enabled."""
+        response = memory_client.get("/stats")
+        assert response.status_code == 200
+        data = response.json()
+        assert "requests" in data
+
+
+@pytest.fixture
+def memory_client_global(temp_memory_db):
+    """Memory-enabled client with GLOBAL storage mode.
+
+    GLOBAL keeps every memory in a single SQLite file regardless of
+    project routing, so tests that pre-seed via direct backend access
+    are guaranteed to share the same DB as the proxy's runtime backend.
+    """
+    config = ProxyConfig(
+        optimize=False,
+        cache_enabled=False,
+        rate_limit_enabled=False,
+        cost_tracking_enabled=False,
+        memory_enabled=True,
+        memory_backend="local",
+        memory_db_path=temp_memory_db,
+        memory_inject_tools=True,
+        memory_inject_context=True,
+        memory_top_k=5,
+        memory_storage_mode="global",
+    )
+    app = create_app(config)
+    with TestClient(app) as client:
+        yield client
+
+
+# ---------------------------------------------------------------------------
+# Helpers shared by the live-API tests below. Each live test follows the same
+# three-step shape:
+#   1. seed a memory with known content via a fresh ONNX-backed LocalBackend
+#      pointed at the same db_path as the proxy backend (so the proxy reads
+#      our row, and we know its exact ID up front),
+#   2. install a recorder that captures every memory tool call the proxy
+#      dispatches downstream of the model's tool_use blocks,
+#   3. make a real Anthropic API request via TestClient and assert the
+#      recorded calls match the expected verb + memory_id contract.
+#
+# The helpers below factor out (1) and (2) so each test body reads as the
+# one-line intent it actually is.
+# ---------------------------------------------------------------------------
+
+
+def _seed_memory(*, db_path: str, user_id: str, content: str) -> str:
+    """Save ``content`` for ``user_id`` via a fresh ONNX-backed LocalBackend
+    pointed at ``db_path``. Returns the new memory's ID."""
+    import asyncio
+
+    from headroom.memory.backends.local import LocalBackend, LocalBackendConfig
+
+    async def _run() -> str:
+        backend = LocalBackend(
+            LocalBackendConfig(
+                db_path=db_path,
+                embedder_backend="onnx",
+                embedder_model="all-MiniLM-L6-v2",
+                vector_dimension=384,
+            )
+        )
+        mem = await backend.save_memory(content=content, user_id=user_id)
+        return mem.id
+
+    mem_id = asyncio.run(_run())
+    assert mem_id, "save_memory should return a usable id"
+    return mem_id
+
+
+def _install_tool_call_recorder(handler):
+    """Monkey-patch ``handler._execute_memory_tool`` to record every call.
+
+    Each recorded entry has ``tool_name``, ``input``, and ``result`` so the
+    test can assert both what the model called AND what the proxy returned
+    back to it (the dedup-hint path lives in the latter).
+
+    Returns ``(recorded_list, restore_callable)``. Call ``restore_callable()``
+    in a ``finally`` to put the original method back, no matter what the
+    request body does."""
+    recorded: list[dict] = []
+    original_execute = handler._execute_memory_tool
+
+    async def _capturing_execute(
+        tool_name, input_data, user_id_arg, provider, request_context=None
+    ):
+        result = await original_execute(
+            tool_name,
+            input_data,
+            user_id_arg,
+            provider,
+            request_context=request_context,
+        )
+        recorded.append({"tool_name": tool_name, "input": dict(input_data), "result": result})
+        return result
+
+    handler._execute_memory_tool = _capturing_execute  # type: ignore[assignment]
+
+    def _restore() -> None:
+        handler._execute_memory_tool = original_execute  # type: ignore[assignment]
+
+    return recorded, _restore
+
+
+@pytest.mark.skipif(not os.environ.get("ANTHROPIC_API_KEY"), reason="ANTHROPIC_API_KEY not set")
+class TestMemoryIdAutoTailAndUpdate:
+    """End-to-end live: model uses [memory_id] from auto-tail to call memory_update.
+
+    Validates that the IDs we added to the auto-injected memory block
+    (see ``MemoryHandler.search_and_format_context``) are extractable by
+    a real Claude model and can be passed directly to ``memory_update``
+    without an intervening ``memory_search`` round-trip.
+    """
+
+    def test_model_uses_memory_id_to_call_memory_update(
+        self,
+        memory_client_global,
+        anthropic_api_key,
+        temp_memory_db,
+    ):
+        memory_id = _seed_memory(
+            db_path=temp_memory_db,
+            user_id=(user_id := f"test-id-update-{int(time.time())}"),
+            content="The user's favorite color is blue.",
+        )
+        # Let the SQLite write + index settle before the proxy reads.
+        time.sleep(0.5)
+
+        proxy = memory_client_global.app.state.proxy
+        assert proxy.memory_handler is not None
+        recorded, restore = _install_tool_call_recorder(proxy.memory_handler)
+
+        try:
+            response = memory_client_global.post(
+                "/v1/messages",
+                headers={
+                    "x-api-key": anthropic_api_key,
+                    "anthropic-version": "2023-06-01",
+                    "x-headroom-user-id": user_id,
+                },
+                json={
+                    "model": "claude-sonnet-4-20250514",
+                    "max_tokens": 800,
+                    "messages": [
+                        {
+                            "role": "user",
+                            "content": (
+                                "Quick correction: my favorite color is actually "
+                                "green, not blue. Please call the memory_update "
+                                "tool to fix the relevant memory in your context. "
+                                "The relevant memories block lists each memory's "
+                                "ID in square brackets — use that ID for "
+                                "memory_id."
+                            ),
+                        }
+                    ],
+                },
+            )
+        finally:
+            restore()
+
+        assert response.status_code == 200, response.text
+
+        # The model should have called memory_update at least once.
+        update_calls = [c for c in recorded if c["tool_name"] == "memory_update"]
+        assert update_calls, f"Expected at least one memory_update call. Recorded: {recorded}"
+
+        # And it should reference the exact ID we seeded — i.e. the
+        # model used the [id] from the auto-tail block, not a guess.
+        assert any(c["input"].get("memory_id") == memory_id for c in update_calls), (
+            f"Expected memory_update(memory_id={memory_id!r}); got inputs: "
+            f"{[c['input'] for c in update_calls]}"
+        )
+
+    def test_model_uses_memory_id_to_call_memory_delete(
+        self,
+        memory_client_global,
+        anthropic_api_key,
+        temp_memory_db,
+    ):
+        """Same [id] handle, different destructive verb. Verifies the
+        auto-tail bracketed ID is usable for memory_delete just as it
+        is for memory_update — i.e. the handle is verb-agnostic."""
+
+        memory_id = _seed_memory(
+            db_path=temp_memory_db,
+            user_id=(user_id := f"test-id-delete-{int(time.time())}"),
+            content="The user used to work at AcmeCorp until 2024.",
+        )
+        time.sleep(0.5)
+
+        proxy = memory_client_global.app.state.proxy
+        assert proxy.memory_handler is not None
+        recorded, restore = _install_tool_call_recorder(proxy.memory_handler)
+
+        try:
+            response = memory_client_global.post(
+                "/v1/messages",
+                headers={
+                    "x-api-key": anthropic_api_key,
+                    "anthropic-version": "2023-06-01",
+                    "x-headroom-user-id": user_id,
+                },
+                json={
+                    "model": "claude-sonnet-4-20250514",
+                    "max_tokens": 800,
+                    "messages": [
+                        {
+                            "role": "user",
+                            "content": (
+                                "Please remove the memory about where I used to "
+                                "work (AcmeCorp). Call memory_delete directly — "
+                                "do NOT call memory_search or memory_list first. "
+                                "The memory's ID is shown in square brackets in "
+                                "the relevant memories block at the end of this "
+                                "message; pass that ID to memory_id."
+                            ),
+                        }
+                    ],
+                },
+            )
+        finally:
+            restore()
+
+        assert response.status_code == 200, response.text
+
+        delete_calls = [c for c in recorded if c["tool_name"] == "memory_delete"]
+        assert delete_calls, f"Expected at least one memory_delete call. Recorded: {recorded}"
+        assert any(c["input"].get("memory_id") == memory_id for c in delete_calls), (
+            f"Expected memory_delete(memory_id={memory_id!r}); got inputs: "
+            f"{[c['input'] for c in delete_calls]}"
+        )
+
+    def test_dedup_hint_surfaces_seeded_id_when_memory_save_runs_on_near_duplicate(
+        self,
+        memory_client_global,
+        anthropic_api_key,
+        temp_memory_db,
+    ):
+        """Live verification of the memory_save → dedup-hint mechanism.
+
+        Without this hint, ``memory_save`` on a near-duplicate would silently
+        accumulate parallel rows, polluting the cache prefix and confusing
+        the model on subsequent retrieval. The hint surfaces the existing
+        row's ID in the tool result so the model has a directly addressable
+        handle to consolidate via ``memory_update``.
+
+        We assert the MECHANISM end-to-end:
+
+          - Model fires ``memory_save`` on the prompted (similar) content.
+          - The proxy's ``_execute_save`` returns a ``note`` containing the
+            pre-seeded memory's exact ID.
+
+        We DO NOT assert that the model actually consolidates — the hint
+        text intentionally ends with "or ignore if these are distinct
+        facts", so the model is free to decline. Whether it consolidates
+        depends on its judgement about whether two phrasings are the same
+        fact, which is intentionally outside this contract."""
+
+        # Pre-seed a memory the new save will look semantically similar to.
+        # We use content close enough that cosine similarity comfortably
+        # clears DEDUP_HINT_THRESHOLD (0.75).
+        seeded_id = _seed_memory(
+            db_path=temp_memory_db,
+            user_id=(user_id := f"test-dedup-{int(time.time())}"),
+            content="The user prefers Python for data analysis work.",
+        )
+        time.sleep(0.5)
+
+        proxy = memory_client_global.app.state.proxy
+        assert proxy.memory_handler is not None
+        recorded, restore = _install_tool_call_recorder(proxy.memory_handler)
+
+        try:
+            response = memory_client_global.post(
+                "/v1/messages",
+                headers={
+                    "x-api-key": anthropic_api_key,
+                    "anthropic-version": "2023-06-01",
+                    "x-headroom-user-id": user_id,
+                },
+                json={
+                    "model": "claude-sonnet-4-20250514",
+                    "max_tokens": 1000,
+                    "messages": [
+                        {
+                            "role": "user",
+                            "content": (
+                                "Use memory_save DIRECTLY to store: "
+                                "'User prefers Python for data science.' "
+                                "Do NOT call memory_search or memory_list "
+                                "first — I want to exercise the save path."
+                            ),
+                        }
+                    ],
+                },
+            )
+        finally:
+            restore()
+
+        assert response.status_code == 200, response.text
+
+        # The model must have fired memory_save (we explicitly prompted
+        # that path).
+        save_calls = [c for c in recorded if c["tool_name"] == "memory_save"]
+        assert save_calls, f"Expected memory_save call. Recorded: {recorded}"
+
+        # The proxy's _execute_save must have returned a dedup hint
+        # surfacing the seeded memory's exact ID — that's the mechanism
+        # under test. The hint is a serialized JSON string with a "note"
+        # field; assert the seeded ID is present in it.
+        save_result = save_calls[0]["result"]
+        assert isinstance(save_result, str), (
+            f"Expected JSON-string tool result, got {type(save_result).__name__}: {save_result!r}"
+        )
+        assert "Similar memory exists" in save_result, (
+            "Expected dedup hint in memory_save result (similarity should clear "
+            f"DEDUP_HINT_THRESHOLD=0.75). Got: {save_result}"
+        )
+        assert seeded_id in save_result, (
+            f"Expected dedup hint to surface seeded memory_id={seeded_id!r}; got: {save_result}"
+        )