diff --git a/tests/test_proxy_memory_integration.py b/tests/test_proxy_memory_integration.py index df9d61d5c..60554b1a3 100644 --- a/tests/test_proxy_memory_integration.py +++ b/tests/test_proxy_memory_integration.py @@ -404,6 +404,80 @@ def memory_client_global(temp_memory_db): 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. @@ -420,57 +494,17 @@ class TestMemoryIdAutoTailAndUpdate: anthropic_api_key, temp_memory_db, ): - import asyncio - - from headroom.memory.backends.local import LocalBackend, LocalBackendConfig - - user_id = f"test-id-update-{int(time.time())}" - - # Pre-seed a known memory directly via a fresh backend so we - # control its content and learn its ID up front. Same db_path - # AND same embedder (ONNX) as the proxy backend → both read - # the same SQLite file and produce comparable vectors. - async def _seed() -> str: - backend = LocalBackend( - LocalBackendConfig( - db_path=temp_memory_db, - embedder_backend="onnx", - embedder_model="all-MiniLM-L6-v2", - vector_dimension=384, - ) - ) - mem = await backend.save_memory( - content="The user's favorite color is blue.", - user_id=user_id, - ) - return mem.id - - memory_id = asyncio.run(_seed()) - assert memory_id, "save_memory should return a usable id" - + 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) - # Capture tool calls the proxy executes, so we can assert the - # model picked the right tool with the right memory_id. proxy = memory_client_global.app.state.proxy assert proxy.memory_handler is not None - recorded: list[dict] = [] - original_execute = proxy.memory_handler._execute_memory_tool - - async def _capturing_execute( - tool_name, input_data, user_id_arg, provider, request_context=None - ): - recorded.append({"tool_name": tool_name, "input": dict(input_data)}) - return await original_execute( - tool_name, - input_data, - user_id_arg, - provider, - request_context=request_context, - ) - - proxy.memory_handler._execute_memory_tool = _capturing_execute # type: ignore[assignment] + recorded, restore = _install_tool_call_recorder(proxy.memory_handler) try: response = memory_client_global.post( @@ -499,7 +533,7 @@ class TestMemoryIdAutoTailAndUpdate: }, ) finally: - proxy.memory_handler._execute_memory_tool = original_execute # type: ignore[assignment] + restore() assert response.status_code == 200, response.text @@ -513,3 +547,152 @@ class TestMemoryIdAutoTailAndUpdate: 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}" + )