diff --git a/hindsight-api-slim/hindsight_api/engine/memory_engine.py b/hindsight-api-slim/hindsight_api/engine/memory_engine.py index afd392215..1d974c847 100644 --- a/hindsight-api-slim/hindsight_api/engine/memory_engine.py +++ b/hindsight-api-slim/hindsight_api/engine/memory_engine.py @@ -13624,6 +13624,37 @@ async def create_knowledge_folder( ) return self._row_to_knowledge_node(row) + async def _compensate_knowledge_page_mental_model( + self, + bank_id: str, + mental_model_id: str, + request_context: "RequestContext", + ) -> None: + """Best-effort removal of a committed MM after page creation fails.""" + cleanup_task = asyncio.create_task( + self.delete_mental_model(bank_id, mental_model_id, request_context=request_context) + ) + deferred_cancellation: asyncio.CancelledError | None = None + while not cleanup_task.done(): + try: + await asyncio.shield(cleanup_task) + except asyncio.CancelledError as exc: + if cleanup_task.cancelled(): + break + # Shield alone would leave cleanup running in the background. + # Delay repeated cancellation until the compensating write ends. + deferred_cancellation = exc + except Exception: + break + try: + cleanup_task.result() + except asyncio.CancelledError: + logger.error("Cleanup was cancelled for mental model %s after page creation failed", mental_model_id) + except Exception: + logger.exception("Failed to clean up mental model %s after page creation failed", mental_model_id) + if deferred_cancellation is not None: + raise deferred_cancellation + async def create_knowledge_page( self, bank_id: str, @@ -13660,6 +13691,13 @@ async def create_knowledge_page( request_context=request_context, ) await self._validate_operation(self._operation_validator.validate_bank_write(ctx)) + backend = await self._get_backend() + # Reject deterministic parent errors before creating the independently + # committed mental model. The page transaction repeats this validation so + # its INSERT still observes the latest parent state. + if parent_id is not None: + async with acquire_with_retry(backend) as conn: + await self._kp_assert_folder_parent(conn, bank_id, parent_id) # The mental model carries the content (and is created by the existing # path, including lazy bank creation); the node only refs it. The write is # already authorized above, so the nested mental-model create/delete run @@ -13676,7 +13714,6 @@ async def create_knowledge_page( trigger=trigger if trigger is not None else dict(self.KNOWLEDGE_PAGE_DEFAULT_TRIGGER), request_context=request_context, ) - backend = await self._get_backend() page_id = f"kp-{uuid.uuid4().hex}" try: async with acquire_with_retry(backend) as conn: @@ -13697,11 +13734,23 @@ async def create_knowledge_page( managed, ) except asyncpg.UniqueViolationError: - # Duplicate page name in this folder (uq_kp_folder_pagename). Roll back - # by deleting the orphan mental model we just created, then signal the - # caller that the page already exists. - await self.delete_mental_model(bank_id, mm["id"], request_context=request_context) + # Preserve the duplicate-page 409 contract even if best-effort + # cleanup of the independently committed model fails. + await self._compensate_knowledge_page_mental_model(bank_id, mm["id"], request_context) return None + except asyncio.CancelledError: + # Cancellation is delivered outside Exception on Python 3.11. + try: + await self._compensate_knowledge_page_mental_model(bank_id, mm["id"], request_context) + except asyncio.CancelledError: + pass + raise + except Exception: + # create_mental_model commits independently, so every later failure + # triggers best-effort compensation. Preserve the page error even if + # cleanup also fails so the API reports the actual failed operation. + await self._compensate_knowledge_page_mental_model(bank_id, mm["id"], request_context) + raise node = self._row_to_knowledge_node(row) # Surface the mental-model metadata so the caller can render markdown or # schedule a content refresh without a second fetch. diff --git a/hindsight-api-slim/tests/test_knowledge_base.py b/hindsight-api-slim/tests/test_knowledge_base.py index a2a24981f..6a5e48ccd 100644 --- a/hindsight-api-slim/tests/test_knowledge_base.py +++ b/hindsight-api-slim/tests/test_knowledge_base.py @@ -5,12 +5,17 @@ without consolidation. """ +import asyncio import urllib.parse import uuid from datetime import datetime, timedelta, timezone +from unittest.mock import AsyncMock +import pytest import pytest_asyncio +from hindsight_api import RequestContext +from hindsight_api.engine.db import DatabaseConnection from hindsight_api.engine.memory_engine import MemoryEngine, _may_need_refresh from hindsight_api.extensions import ( BankReadContext, @@ -365,6 +370,177 @@ async def test_create_folder_bad_parent(self, api_client, kb_bank): ) assert resp.status_code == 400 + async def test_create_page_missing_parent_does_not_leak_mental_model(self, memory: MemoryEngine, request_context): + bank_id = f"test-kb-create-{uuid.uuid4().hex[:8]}" + await memory.create_knowledge_folder(bank_id, "Root", request_context=request_context) + before = await memory.list_mental_models(bank_id, request_context=request_context) + + with pytest.raises(ValueError, match="not found"): + await memory.create_knowledge_page( + bank_id, + "Orphan", + "What is orphaned?", + "seed", + parent_id="missing-parent", + request_context=request_context, + ) + + after = await memory.list_mental_models(bank_id, request_context=request_context) + assert {mm["id"] for mm in after} == {mm["id"] for mm in before} + await memory.delete_bank(bank_id, request_context=request_context) + + async def test_create_page_under_page_does_not_leak_mental_model(self, memory: MemoryEngine, request_context): + bank_id = f"test-kb-create-{uuid.uuid4().hex[:8]}" + parent = await memory.create_knowledge_page( + bank_id, "Parent page", "What is the parent?", "seed", request_context=request_context + ) + before = await memory.list_mental_models(bank_id, request_context=request_context) + + with pytest.raises(ValueError, match="is not a folder"): + await memory.create_knowledge_page( + bank_id, + "Orphan", + "What is orphaned?", + "seed", + parent_id=parent["id"], + request_context=request_context, + ) + + after = await memory.list_mental_models(bank_id, request_context=request_context) + assert {mm["id"] for mm in after} == {mm["id"] for mm in before} + await memory.delete_bank(bank_id, request_context=request_context) + + async def test_create_page_post_mental_model_failure_is_compensated( + self, memory: MemoryEngine, request_context, monkeypatch + ): + bank_id = f"test-kb-create-{uuid.uuid4().hex[:8]}" + parent = await memory.create_knowledge_folder(bank_id, "Root", request_context=request_context) + mental_model_id = f"mm-{uuid.uuid4().hex}" + assert_parent = AsyncMock(side_effect=[None, RuntimeError("page write failed")]) + monkeypatch.setattr(memory, "_kp_assert_folder_parent", assert_parent) + + with pytest.raises(RuntimeError, match="page write failed"): + await memory.create_knowledge_page( + bank_id, + "Orphan", + "What is orphaned?", + "seed", + parent_id=parent["id"], + mental_model_id=mental_model_id, + request_context=request_context, + ) + + assert await memory.get_mental_model(bank_id, mental_model_id, request_context=request_context) is None + await memory.delete_bank(bank_id, request_context=request_context) + + async def test_create_page_cancellation_after_mental_model_creation_is_compensated( + self, memory: MemoryEngine, request_context, monkeypatch + ): + bank_id = f"test-kb-create-{uuid.uuid4().hex[:8]}" + parent = await memory.create_knowledge_folder(bank_id, "Root", request_context=request_context) + mental_model_id = f"mm-{uuid.uuid4().hex}" + entered_page_phase = asyncio.Event() + never_complete = asyncio.Event() + cleanup_started = asyncio.Event() + allow_cleanup = asyncio.Event() + validation_count = 0 + delete_mental_model = memory.delete_mental_model + + async def controlled_assert_parent(_conn: DatabaseConnection, _bank_id: str, _parent_id: str | None) -> None: + nonlocal validation_count + validation_count += 1 + if validation_count == 1: + return + entered_page_phase.set() + await never_complete.wait() + + async def delayed_delete( + delete_bank_id: str, + delete_mental_model_id: str, + *, + request_context: RequestContext, + ) -> bool: + cleanup_started.set() + await allow_cleanup.wait() + return await delete_mental_model( + delete_bank_id, + delete_mental_model_id, + request_context=request_context, + ) + + monkeypatch.setattr(memory, "_kp_assert_folder_parent", controlled_assert_parent) + monkeypatch.setattr(memory, "delete_mental_model", delayed_delete) + + create_task = asyncio.create_task( + memory.create_knowledge_page( + bank_id, + "Cancelled", + "What was cancelled?", + "seed", + parent_id=parent["id"], + mental_model_id=mental_model_id, + request_context=request_context, + ) + ) + await entered_page_phase.wait() + create_task.cancel() + await cleanup_started.wait() + create_task.cancel() + allow_cleanup.set() + with pytest.raises(asyncio.CancelledError): + await create_task + + assert await memory.get_mental_model(bank_id, mental_model_id, request_context=request_context) is None + await memory.delete_bank(bank_id, request_context=request_context) + + @pytest.mark.parametrize("cleanup_error", [RuntimeError("cleanup failed"), asyncio.CancelledError()]) + async def test_cleanup_failure_preserves_page_creation_error( + self, memory: MemoryEngine, request_context, monkeypatch, cleanup_error: BaseException + ): + bank_id = f"test-kb-create-{uuid.uuid4().hex[:8]}" + parent = await memory.create_knowledge_folder(bank_id, "Root", request_context=request_context) + mental_model_id = f"mm-{uuid.uuid4().hex}" + assert_parent = AsyncMock(side_effect=[None, RuntimeError("page write failed")]) + delete_mental_model = AsyncMock(side_effect=cleanup_error) + monkeypatch.setattr(memory, "_kp_assert_folder_parent", assert_parent) + monkeypatch.setattr(memory, "delete_mental_model", delete_mental_model) + + with pytest.raises(RuntimeError, match="page write failed"): + await memory.create_knowledge_page( + bank_id, + "Orphan", + "What is orphaned?", + "seed", + parent_id=parent["id"], + mental_model_id=mental_model_id, + request_context=request_context, + ) + + assert await memory.get_mental_model(bank_id, mental_model_id, request_context=request_context) is not None + await memory.delete_bank(bank_id, request_context=request_context) + + async def test_duplicate_page_cleanup_failure_preserves_conflict_result( + self, memory: MemoryEngine, request_context, monkeypatch + ): + bank_id = f"test-kb-create-{uuid.uuid4().hex[:8]}" + await memory.create_knowledge_page(bank_id, "Existing", "What exists?", "seed", request_context=request_context) + mental_model_id = f"mm-{uuid.uuid4().hex}" + delete_mental_model = AsyncMock(side_effect=RuntimeError("cleanup failed")) + monkeypatch.setattr(memory, "delete_mental_model", delete_mental_model) + + duplicate = await memory.create_knowledge_page( + bank_id, + "Existing", + "What is duplicated?", + "seed", + mental_model_id=mental_model_id, + request_context=request_context, + ) + + assert duplicate is None + assert await memory.get_mental_model(bank_id, mental_model_id, request_context=request_context) is not None + await memory.delete_bank(bank_id, request_context=request_context) + class TestExport: async def test_export_bundle_nested_index(self, api_client, kb_bank):