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)