From 1aec4a47e753e70a5e30d699eaf4823cbb55f51d Mon Sep 17 00:00:00 2001 From: zhangchi47 Date: Tue, 11 Aug 2026 11:18:27 +0800 Subject: [PATCH] refactor(kb): create page and mental model atomically Knowledge pages and their backing mental models share one lifecycle, but they were committed in separate transactions. A page insert failure could therefore leave an orphaned mental model. Extract mental-model embedding generation and insertion into typed helpers. Create the bank, mental model, and page through one connection and transaction, so any page failure rolls back all related writes. Preserve the duplicate-page contract by returning None only for uq_kp_folder_pagename. Add PostgreSQL and Oracle regression coverage for rollback after a real mental model insert, plus parent-validation and duplicate-name coverage. --- .../hindsight_api/engine/memory_engine.py | 218 +++++++++--------- .../tests/test_knowledge_base.py | 117 ++++++++++ .../tests/test_oracle_integration.py | 32 +++ 3 files changed, 262 insertions(+), 105 deletions(-) diff --git a/hindsight-api-slim/hindsight_api/engine/memory_engine.py b/hindsight-api-slim/hindsight_api/engine/memory_engine.py index afd392215a..1131fd1c57 100644 --- a/hindsight-api-slim/hindsight_api/engine/memory_engine.py +++ b/hindsight-api-slim/hindsight_api/engine/memory_engine.py @@ -51,7 +51,7 @@ from ..worker.stage import set_stage from .audit import AuditLogger, audit_context from .bank_stats_cache import BankStatsCache, DistributedBankStatsCache -from .db import DatabaseBackend, create_database_backend +from .db import DatabaseBackend, DatabaseConnection, ResultRow, create_database_backend from .db.ops_postgresql import pg_search_vector_expr from .db_budget import budgeted_operation from .llm_interface import ProviderRateLimitResetError @@ -12145,6 +12145,54 @@ async def get_mental_model_history( ) return result + async def _generate_mental_model_embedding(self, name: str, content: str) -> str | None: + embedding = await embedding_utils.generate_embeddings_batch(self.embeddings, [f"{name} {content}"]) + return str(embedding[0]) if embedding else None + + async def _insert_pinned_mental_model( + self, + conn: DatabaseConnection, + *, + mental_model_id: str, + bank_id: str, + name: str, + source_query: str, + content: str, + embedding: str | None, + tags: list[str] | None, + max_tokens: int | None, + trigger: dict[str, Any] | None, + ) -> ResultRow: + """Insert a pinned model using the caller's transaction.""" + # VectorChord needs mental_models.search_vector tokenized on write; every + # other backend either generates it or indexes the source columns. + sv_expr = pg_search_vector_expr( + get_config(), text_col="$3", context_col="$5", signals_col=None, native_inline=False + ) + sv_col = ", search_vector" if sv_expr else "" + sv_val = f", {sv_expr}" if sv_expr else "" + row = await conn.fetchrow( + f""" + INSERT INTO {fq_table("mental_models")} + (id, bank_id, subtype, name, description, source_query, content, embedding, tags, max_tokens, trigger{sv_col}) + VALUES ($1, $2, 'pinned', $3, ' ', $4, $5, $6, $7, COALESCE($8, 2048), COALESCE($9, '{{"refresh_after_consolidation": false}}'::jsonb){sv_val}) + RETURNING id, bank_id, name, source_query, content, tags, + last_refreshed_at, created_at, reflect_response, + max_tokens, trigger, structured_content + """, + mental_model_id, + bank_id, + name, + source_query, + content, + embedding, + tags or [], + max_tokens, + json.dumps(trigger) if trigger else None, + ) + assert row is not None + return row + async def create_mental_model( self, bank_id: str, @@ -12190,11 +12238,7 @@ async def create_mental_model( await self._validate_operation(self._operation_validator.validate_bank_write(ctx)) backend = await self._get_backend() - # Generate embedding for the content - embedding_text = f"{name} {content}" - embedding = await embedding_utils.generate_embeddings_batch(self.embeddings, [embedding_text]) - # Convert embedding to string for asyncpg vector type - embedding_str = str(embedding[0]) if embedding else None + embedding = await self._generate_mental_model_embedding(name, content) if not mental_model_id: mental_model_id = f"mm-{uuid.uuid4().hex}" @@ -12210,63 +12254,18 @@ async def create_mental_model( request_context, conn=conn, ) - # VectorChord needs mental_models.search_vector tokenized on write: - # its column is a plain bm25vector read by idx_mental_models_text_search - # (native's is GENERATED; pg_search/pg_textsearch/pgroonga index base - # columns), so every other backend leaves it out. Same tokenization the - # memory_units write path uses (pg_search_vector_expr / insert_facts_batch), - # over name + content — native_inline=False because mm's native column - # populates itself. - config = get_config() - if mental_model_id: - sv_expr = pg_search_vector_expr( - config, text_col="$3", context_col="$5", signals_col=None, native_inline=False - ) - sv_col = ", search_vector" if sv_expr else "" - sv_val = f", {sv_expr}" if sv_expr else "" - row = await conn.fetchrow( - f""" - INSERT INTO {fq_table("mental_models")} - (id, bank_id, subtype, name, description, source_query, content, embedding, tags, max_tokens, trigger{sv_col}) - VALUES ($1, $2, 'pinned', $3, ' ', $4, $5, $6, $7, COALESCE($8, 2048), COALESCE($9, '{{"refresh_after_consolidation": false}}'::jsonb){sv_val}) - RETURNING id, bank_id, name, source_query, content, tags, - last_refreshed_at, created_at, reflect_response, - max_tokens, trigger, structured_content - """, - mental_model_id, - bank_id, - name, - source_query, - content, - embedding_str, - tags or [], - max_tokens, - json.dumps(trigger) if trigger else None, - ) - else: - sv_expr = pg_search_vector_expr( - config, text_col="$2", context_col="$4", signals_col=None, native_inline=False - ) - sv_col = ", search_vector" if sv_expr else "" - sv_val = f", {sv_expr}" if sv_expr else "" - row = await conn.fetchrow( - f""" - INSERT INTO {fq_table("mental_models")} - (bank_id, subtype, name, description, source_query, content, embedding, tags, max_tokens, trigger{sv_col}) - VALUES ($1, 'pinned', $2, ' ', $3, $4, $5, $6, COALESCE($7, 2048), COALESCE($8, '{{"refresh_after_consolidation": false}}'::jsonb){sv_val}) - RETURNING id, bank_id, name, source_query, content, tags, - last_refreshed_at, created_at, reflect_response, - max_tokens, trigger, structured_content - """, - bank_id, - name, - source_query, - content, - embedding_str, - tags or [], - max_tokens, - json.dumps(trigger) if trigger else None, - ) + row = await self._insert_pinned_mental_model( + conn, + mental_model_id=mental_model_id, + bank_id=bank_id, + name=name, + source_query=source_query, + content=content, + embedding=embedding, + tags=tags, + max_tokens=max_tokens, + trigger=trigger, + ) # Best-effort default-template hook runs after the bank-create commits # (it opens its own connections and can create pinned models). @@ -13660,48 +13659,57 @@ async def create_knowledge_page( request_context=request_context, ) await self._validate_operation(self._operation_validator.validate_bank_write(ctx)) - # 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 - # without invoking the validator a second time. - with _authorize_nested_operations(): - mm = await self.create_mental_model( - bank_id=bank_id, - name=name, - source_query=source_query, - content=content, - mental_model_id=mental_model_id, - tags=tags, - max_tokens=max_tokens if max_tokens is not None else self.KNOWLEDGE_PAGE_DEFAULT_MAX_TOKENS, - 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: - async with conn.transaction(): - await self._kp_assert_folder_parent(conn, bank_id, parent_id) - row = await conn.fetchrow( - f""" - INSERT INTO {fq_table("knowledge_pages")} - (id, bank_id, parent_id, kind, name, mental_model_id, managed) - VALUES ($1, $2, $3, 'page', $4, $5, $6) - RETURNING {self._KP_COLUMNS} - """, - page_id, - bank_id, - parent_id, - name, - mm["id"], - 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) - return None + mental_model_id = mental_model_id or f"mm-{uuid.uuid4().hex}" + embedding = await self._generate_mental_model_embedding(name, content) + effective_max_tokens = max_tokens if max_tokens is not None else self.KNOWLEDGE_PAGE_DEFAULT_MAX_TOKENS + effective_trigger = trigger if trigger is not None else dict(self.KNOWLEDGE_PAGE_DEFAULT_TRIGGER) + backend = await self._get_backend() + page_id = f"kp-{uuid.uuid4().hex}" + try: + async with acquire_with_retry(backend) as conn: + # The page row and its backing model have one lifecycle, so they + # share a transaction instead of compensating after a partial commit. + async with conn.transaction(): + created = await self._ensure_bank_exists(bank_id, request_context, conn=conn) + await self._kp_assert_folder_parent(conn, bank_id, parent_id) + mm_row = await self._insert_pinned_mental_model( + conn, + mental_model_id=mental_model_id, + bank_id=bank_id, + name=name, + source_query=source_query, + content=content, + embedding=embedding, + tags=tags, + max_tokens=effective_max_tokens, + trigger=effective_trigger, + ) + row = await conn.fetchrow( + f""" + INSERT INTO {fq_table("knowledge_pages")} + (id, bank_id, parent_id, kind, name, mental_model_id, managed) + VALUES ($1, $2, $3, 'page', $4, $5, $6) + RETURNING {self._KP_COLUMNS} + """, + page_id, + bank_id, + parent_id, + name, + mental_model_id, + managed, + ) + except asyncpg.UniqueViolationError as exc: + if getattr(exc, "constraint_name", None) != "uq_kp_folder_pagename": + raise + # The transaction already rolled the MM back; preserve the existing + # API contract that a duplicate page is surfaced as HTTP 409. + return None + + # This hook opens its own connections and therefore must run after commit. + if created: + await self._apply_default_bank_template(bank_id, request_context) + logger.info(f"[MENTAL_MODELS] Created pinned mental model '{name}' for bank {bank_id}") + mm = self._row_to_mental_model(mm_row) 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 a2a24981f1..e835cc2185 100644 --- a/hindsight-api-slim/tests/test_knowledge_base.py +++ b/hindsight-api-slim/tests/test_knowledge_base.py @@ -8,9 +8,13 @@ import urllib.parse import uuid from datetime import datetime, timedelta, timezone +from typing import Any, NoReturn +import asyncpg +import pytest import pytest_asyncio +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 +369,119 @@ async def test_create_folder_bad_parent(self, api_client, kb_bank): ) assert resp.status_code == 400 + async def test_create_page_missing_parent_rolls_back_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_rolls_back_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_duplicate_page_rolls_back_mental_model(self, memory: MemoryEngine, request_context): + bank_id = f"test-kb-create-{uuid.uuid4().hex[:8]}" + parent = await memory.create_knowledge_folder(bank_id, "Root", request_context=request_context) + await memory.create_knowledge_page( + bank_id, + "Existing", + "What exists?", + "seed", + parent_id=parent["id"], + request_context=request_context, + ) + rolled_back_mm_id = f"mm-{uuid.uuid4().hex}" + + duplicate = await memory.create_knowledge_page( + bank_id, + "Existing", + "What is duplicated?", + "seed", + parent_id=parent["id"], + mental_model_id=rolled_back_mm_id, + request_context=request_context, + ) + + assert duplicate is None + assert await memory.get_mental_model(bank_id, rolled_back_mm_id, request_context=request_context) is None + await memory.delete_bank(bank_id, request_context=request_context) + + async def test_duplicate_mental_model_id_is_not_reported_as_duplicate_page( + self, memory: MemoryEngine, request_context + ): + bank_id = f"test-kb-create-{uuid.uuid4().hex[:8]}" + existing = await memory.create_mental_model( + bank_id, "Existing MM", "What exists?", "seed", request_context=request_context + ) + + with pytest.raises(asyncpg.UniqueViolationError): + await memory.create_knowledge_page( + bank_id, + "New page", + "What is new?", + "seed", + mental_model_id=existing["id"], + request_context=request_context, + ) + + await memory.delete_bank(bank_id, request_context=request_context) + + async def test_non_unique_failure_after_mental_model_insert_rolls_back( + self, memory: MemoryEngine, request_context, monkeypatch + ): + bank_id = f"test-kb-create-{uuid.uuid4().hex[:8]}" + mental_model_id = f"mm-{uuid.uuid4().hex}" + insert_mental_model = memory._insert_pinned_mental_model + + async def insert_then_fail(conn: DatabaseConnection, **kwargs: Any) -> NoReturn: + await insert_mental_model(conn, **kwargs) + raise RuntimeError("page write failed") + + monkeypatch.setattr(memory, "_insert_pinned_mental_model", insert_then_fail) + + with pytest.raises(RuntimeError, match="page write failed"): + await memory.create_knowledge_page( + bank_id, + "Rolled back", + "What is rolled back?", + "seed", + 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) + class TestExport: async def test_export_bundle_nested_index(self, api_client, kb_bank): diff --git a/hindsight-api-slim/tests/test_oracle_integration.py b/hindsight-api-slim/tests/test_oracle_integration.py index fdf28a9e4b..d593cc029e 100644 --- a/hindsight-api-slim/tests/test_oracle_integration.py +++ b/hindsight-api-slim/tests/test_oracle_integration.py @@ -12,11 +12,13 @@ import logging import uuid from datetime import datetime, timezone +from typing import Any, NoReturn import pytest import pytest_asyncio from hindsight_api import MemoryEngine, RequestContext +from hindsight_api.engine.db import DatabaseConnection from hindsight_api.engine.memory_engine import Budget pytestmark = pytest.mark.oracle @@ -876,6 +878,36 @@ async def test_mental_model_crud(self, oracle_memory: MemoryEngine, request_cont finally: await _safe_cleanup(oracle_memory, bank_id, request_context) + @pytest.mark.asyncio + async def test_knowledge_page_failure_rolls_back_mental_model( + self, oracle_memory: MemoryEngine, request_context: RequestContext, monkeypatch + ): + bank_id = _bank_id("kb-rollback") + mental_model_id = f"mm-{uuid.uuid4().hex}" + insert_mental_model = oracle_memory._insert_pinned_mental_model + + async def insert_then_fail(conn: DatabaseConnection, **kwargs: Any) -> NoReturn: + await insert_mental_model(conn, **kwargs) + raise RuntimeError("page write failed") + + monkeypatch.setattr(oracle_memory, "_insert_pinned_mental_model", insert_then_fail) + try: + with pytest.raises(RuntimeError, match="page write failed"): + await oracle_memory.create_knowledge_page( + bank_id, + "Rolled back", + "What is rolled back?", + "seed", + mental_model_id=mental_model_id, + request_context=request_context, + ) + + assert ( + await oracle_memory.get_mental_model(bank_id, mental_model_id, request_context=request_context) is None + ) + finally: + await _safe_cleanup(oracle_memory, bank_id, request_context) + @pytest.mark.asyncio async def test_mental_model_refresh(self, oracle_memory: MemoryEngine, request_context: RequestContext): bank_id = _bank_id("mmrefresh")