From b681d7588779d67ec66f87b8f22710f95d01a7a0 Mon Sep 17 00:00:00 2001 From: zhangchi47 Date: Tue, 11 Aug 2026 11:18:27 +0800 Subject: [PATCH] fix(kb): clean up mental model when page creation fails Knowledge-page creation committed its mental model before inserting the page. Failures outside the existing duplicate-name path could leave an orphan. Add best-effort compensation for every page-insert failure. Shield cleanup from task cancellation, wait through repeated cancellation, and preserve the original page error when cleanup fails or cancels itself. Validate deterministic parent errors before model creation, while repeating the check inside the page transaction to observe current state. Keep the duplicate None result so HTTP callers continue to receive 409. Cover missing and invalid parents, ordinary failures, duplicate names, cleanup failures, and single or repeated task cancellation. --- .../hindsight_api/engine/memory_engine.py | 59 +++++- .../tests/test_knowledge_base.py | 176 ++++++++++++++++++ 2 files changed, 230 insertions(+), 5 deletions(-) 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):