Skip to content

fix(kb): compensate mental models when page creation fails #3370

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Sanderhoff-alt wants to merge 1 commit into from
+230 −5
Closed

fix(kb): compensate mental models when page creation fails #3370

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 54 additions & 5 deletions hindsight-api-slim/hindsight_api/engine/memory_engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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:
Expand All @@ -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.
Expand Down
176 changes: 176 additions & 0 deletions hindsight-api-slim/tests/test_knowledge_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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):
Expand Down