From 716d10c04d6b53e3ceef2b6802c0d9937f0c1662 Mon Sep 17 00:00:00 2001 From: zhangchi47 Date: Mon, 10 Aug 2026 20:26:20 +0800 Subject: [PATCH] fix(kb): serialize knowledge tree writes Knowledge-tree create, delete, and move operations previously read and updated the hierarchy without a common lock. Concurrent opposite moves could both pass cycle detection against the same snapshot and commit a parent loop. Lock the bank row with FOR NO KEY UPDATE before structural reads and writes. This serializes tree writers for one bank without conflicting with the FOR KEY SHARE locks taken by unrelated foreign-key inserts. The second opposite move now observes the first committed parent link and is rejected by the existing cycle guard. Add a deterministic concurrency test that holds the first bank lock, verifies the second move cannot finish early, and checks that only the first move succeeds. --- .../hindsight_api/engine/memory_engine.py | 19 +++++ .../tests/test_knowledge_base.py | 71 +++++++++++++++++++ 2 files changed, 90 insertions(+) diff --git a/hindsight-api-slim/hindsight_api/engine/memory_engine.py b/hindsight-api-slim/hindsight_api/engine/memory_engine.py index afd392215..c12c4eb46 100644 --- a/hindsight-api-slim/hindsight_api/engine/memory_engine.py +++ b/hindsight-api-slim/hindsight_api/engine/memory_engine.py @@ -13580,6 +13580,19 @@ async def _kp_assert_folder_parent(self, conn, bank_id: str, parent_id: str | No if row["kind"] != "folder": raise ValueError(f"Parent '{parent_id}' is not a folder") + async def _kp_lock_bank(self, conn, bank_id: str) -> None: + """Serialize structural knowledge-tree writers on the bank row. + + FOR NO KEY UPDATE conflicts with other tree writers but not with the + FOR KEY SHARE locks taken by inserts into tables that reference banks. + Oracle rewrites it to FOR UPDATE, which does not block indexed-FK child + inserts there. + """ + await conn.fetchrow( + f"SELECT bank_id FROM {fq_table('banks')} WHERE bank_id = $1 FOR NO KEY UPDATE", + bank_id, + ) + async def create_knowledge_folder( self, bank_id: str, @@ -13609,6 +13622,7 @@ async def create_knowledge_folder( async with acquire_with_retry(backend) as conn: async with conn.transaction(): await self._ensure_bank_exists(bank_id, request_context, conn=conn) + await self._kp_lock_bank(conn, bank_id) await self._kp_assert_folder_parent(conn, bank_id, parent_id) row = await conn.fetchrow( f""" @@ -13681,6 +13695,7 @@ async def create_knowledge_page( try: async with acquire_with_retry(backend) as conn: async with conn.transaction(): + await self._kp_lock_bank(conn, bank_id) await self._kp_assert_folder_parent(conn, bank_id, parent_id) row = await conn.fetchrow( f""" @@ -14010,6 +14025,9 @@ async def move_knowledge_node( backend = await self._get_backend() async with acquire_with_retry(backend) as conn: async with conn.transaction(): + # Structural tree writers lock this bank row before reading or + # changing the hierarchy, serializing moves with creates/deletes. + await self._kp_lock_bank(conn, bank_id) await self._kp_assert_folder_parent(conn, bank_id, new_parent_id) # Cycle guard: walk up from the new parent; if we reach node_id, # the move would create a loop. Done in Python so the check stays @@ -14060,6 +14078,7 @@ async def delete_knowledge_node(self, bank_id: str, node_id: str, *, request_con backend = await self._get_backend() async with acquire_with_retry(backend) as conn: async with conn.transaction(): + await self._kp_lock_bank(conn, bank_id) all_rows = await conn.fetch( f"SELECT id, parent_id, mental_model_id FROM {fq_table('knowledge_pages')} WHERE bank_id = $1", bank_id, diff --git a/hindsight-api-slim/tests/test_knowledge_base.py b/hindsight-api-slim/tests/test_knowledge_base.py index a2a24981f..ec18bdb69 100644 --- a/hindsight-api-slim/tests/test_knowledge_base.py +++ b/hindsight-api-slim/tests/test_knowledge_base.py @@ -5,12 +5,16 @@ without consolidation. """ +import asyncio import urllib.parse import uuid +from contextlib import asynccontextmanager from datetime import datetime, timedelta, timezone +import pytest import pytest_asyncio +import hindsight_api.engine.memory_engine as memory_engine_module from hindsight_api.engine.memory_engine import MemoryEngine, _may_need_refresh from hindsight_api.extensions import ( BankReadContext, @@ -434,6 +438,73 @@ async def test_move_cycle_rejected(self, api_client, kb_bank): ) assert resp.status_code == 400 + async def test_concurrent_opposite_moves_are_serialized(self, memory: MemoryEngine, request_context, monkeypatch): + """The second opposite move must observe the first committed parent link. + + Pause the first request only after its real ``FOR NO KEY UPDATE`` query has + acquired the bank lock. The second request then reaches the same query + but cannot finish it until the first commits, making this a deterministic + regression test rather than a timing-dependent concurrent test. + """ + bank_id = f"test-kb-move-race-{uuid.uuid4().hex[:8]}" + folder_a = await memory.create_knowledge_folder(bank_id, "A", request_context=request_context) + folder_b = await memory.create_knowledge_folder(bank_id, "B", request_context=request_context) + first_lock_acquired = asyncio.Event() + second_lock_attempted = asyncio.Event() + release_first_lock = asyncio.Event() + original_acquire = memory_engine_module.acquire_with_retry + lock_query = "FOR NO KEY UPDATE" + lock_query_count = 0 + + class _PausingConnection: + def __init__(self, conn): + self._conn = conn + + async def fetchrow(self, query, *args, **kwargs): + nonlocal lock_query_count + if lock_query in query: + lock_query_count += 1 + if lock_query_count == 1: + row = await self._conn.fetchrow(query, *args, **kwargs) + first_lock_acquired.set() + await release_first_lock.wait() + return row + second_lock_attempted.set() + return await self._conn.fetchrow(query, *args, **kwargs) + + async def fetch(self, query, *args, **kwargs): + return await self._conn.fetch(query, *args, **kwargs) + + def __getattr__(self, name): + return getattr(self._conn, name) + + @asynccontextmanager + async def pausing_acquire(backend, *args, **kwargs): + async with original_acquire(backend, *args, **kwargs) as conn: + yield _PausingConnection(conn) + + monkeypatch.setattr(memory_engine_module, "acquire_with_retry", pausing_acquire) + try: + first = asyncio.create_task( + memory.move_knowledge_node(bank_id, folder_a["id"], folder_b["id"], request_context=request_context) + ) + await asyncio.wait_for(first_lock_acquired.wait(), timeout=5) + + second = asyncio.create_task( + memory.move_knowledge_node(bank_id, folder_b["id"], folder_a["id"], request_context=request_context) + ) + await asyncio.wait_for(second_lock_attempted.wait(), timeout=5) + assert not second.done(), "second move must not finish before the first commits" + + release_first_lock.set() + first_result = await first + assert first_result["parent_id"] == folder_b["id"] + with pytest.raises(ValueError, match="own subtree"): + await second + finally: + release_first_lock.set() + await memory.delete_bank(bank_id, request_context=request_context) + async def test_delete_folder_cascades(self, api_client, kb_bank, memory, request_context): bank_id, ids = kb_bank # deleting Runbooks removes Sub + Orders (and Orders' mental model)