test_memory_search_concurrency.py

# Copyright (c) 2026 OceanBase.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import asyncio
import os
from contextlib import suppress
from pathlib import Path
from types import MethodType
from typing import Any, Literal
from uuid import uuid4

import pytest
from pydantic import SecretStr

from powercontext.builtin.artifacts.atomic_memory import AtomicMemoryContent
from powercontext.builtin.artifacts.memory import (
    EmbeddingProfile,
    MemoryRerankDecision,
    MemorySearchMode,
)
from powercontext.builtin.artifacts.memory.reranking import MemoryRerankText
from powercontext.builtin.inference import EmbeddingResult, InferenceUsage
from powercontext.builtin.persistence.oceanbase import OceanBaseConfig
from powercontext.builtin.persistence.sqlite import SQLiteConfig
from powercontext.builtin.runtime import (
    BuiltinConfig,
    open_builtin_runtime,
)
from powercontext.builtin.scope import ScopeDraft

DatabaseKind = Literal["sqlite", "oceanbase"]
TIMEOUT_SECONDS = 15
PROFILE = EmbeddingProfile(
    profile_id="concurrent-search-test-v1",
    model="test",
    dimension=3,
    distance="l2",
    normalization="unit",
)
EXPECTED_CHANNELS = {
    "fts": ("text",),
    "vector": ("vector",),
    "hybrid": ("text", "vector"),
}


class _KeywordEmbeddingModel:
    profile = PROFILE

    async def embed(self, texts: tuple[str, ...], /) -> EmbeddingResult:
        vectors = tuple((1.0, 0.0, 0.0) if "stable" in text.casefold() else (0.0, 1.0, 0.0) for text in texts)
        return EmbeddingResult(vectors=vectors)


class _PausingReranker:
    policy_id = "test.concurrent-memory-search.v1"
    supports_atomic_memory = True

    def __init__(self) -> None:
        self.paused = asyncio.Event()
        self.resume = asyncio.Event()

    async def rerank(
        self,
        _query: str,
        candidates: tuple[MemoryRerankText, ...],
        _limit: int,
        /,
    ) -> MemoryRerankDecision:
        self.paused.set()
        await self.resume.wait()
        return MemoryRerankDecision(
            selected_ranks=(1,),
            usage=InferenceUsage(requests=1),
        )


@pytest.mark.parametrize("database_kind", ["sqlite", "oceanbase"])
@pytest.mark.parametrize("mode", ["fts", "vector", "hybrid"])
def test_memory_search_stays_consistent_when_append_advances_head_before_index_query(
    database_kind: DatabaseKind,
    mode: MemorySearchMode,
    tmp_path: Path,
) -> None:
    async def scenario() -> None:
        database = _database_config(database_kind, mode, tmp_path)
        embedding_model = None if mode == "fts" else _KeywordEmbeddingModel()
        async with open_builtin_runtime(
            BuiltinConfig(database=database),
            embedding_model=embedding_model,
        ) as runtime:
            assert runtime.scopes is not None
            scope = await runtime.scopes.create(
                ScopeDraft(
                    title="Concurrent Memory search",
                    summary="Index snapshot consistency",
                    idempotency_key=f"concurrent-memory-search-{uuid4()}",
                )
            )
            assert runtime.atomic_memory is not None
            memory = runtime.atomic_memory.for_scope(scope.scope_id)
            initial = await memory.create((AtomicMemoryContent(kind="fact", text="Stable searchable fact."),))

            provider: Any = runtime._provider
            index = provider.atomic_memory.index
            original_search = index.search
            paused = asyncio.Event()
            resume = asyncio.Event()
            pending: asyncio.Task[Any] | None = None

            async def pause_first_search(
                _self: Any,
                connection: Any,
                scope_id: str,
                request: Any,
            ) -> Any:
                if not paused.is_set():
                    paused.set()
                    await resume.wait()
                return await original_search(connection, scope_id, request)

            index.search = MethodType(pause_first_search, index)
            try:
                pending = asyncio.create_task(
                    memory.search("stable searchable", mode="text" if mode == "fts" else mode)
                )
                await asyncio.wait_for(paused.wait(), timeout=TIMEOUT_SECONDS)
                new_head = await memory.create((AtomicMemoryContent(kind="fact", text="Unrelated appended fact."),))
                resume.set()
                result = await asyncio.wait_for(pending, timeout=TIMEOUT_SECONDS)
            finally:
                resume.set()
                index.search = original_search
                if pending is not None and not pending.done():
                    pending.cancel()
                    with suppress(asyncio.CancelledError):
                        await pending

            assert new_head.primary.ref.artifact_id != initial.primary.ref.artifact_id
            assert result.mode == ("text" if mode == "fts" else mode)
            assert tuple(hit.text for hit in result.hits) == ("Stable searchable fact.",)
            assert result.hits[0].hit.artifact_ref == initial.primary.ref
            assert result.hits[0].matched_by == EXPECTED_CHANNELS[mode]

    asyncio.run(scenario())


def test_memory_search_keeps_the_completed_revision_snapshot_when_head_advances_during_reranking(
    tmp_path: Path,
) -> None:
    async def scenario() -> None:
        reranker = _PausingReranker()
        database = SQLiteConfig(url=f"sqlite+aiosqlite:///{tmp_path / 'rerank-snapshot.db'}")
        async with open_builtin_runtime(BuiltinConfig(database=database), memory_reranker=reranker) as runtime:
            assert runtime.scopes is not None
            scope = await runtime.scopes.create(
                ScopeDraft(
                    title="Rerank snapshot", summary="Stable revision snapshot", idempotency_key="rerank-snapshot"
                )
            )
            assert runtime.atomic_memory is not None
            memory = runtime.atomic_memory.for_scope(scope.scope_id)
            initial = await memory.create((AtomicMemoryContent(kind="fact", text="Stable searchable fact."),))
            pending = asyncio.create_task(memory.search("stable searchable", mode="text", limit=1))
            try:
                await asyncio.wait_for(reranker.paused.wait(), timeout=TIMEOUT_SECONDS)
                new_head = await memory.create((AtomicMemoryContent(kind="fact", text="Unrelated appended fact."),))
                reranker.resume.set()
                result = await asyncio.wait_for(pending, timeout=TIMEOUT_SECONDS)
            finally:
                reranker.resume.set()
                if not pending.done():
                    pending.cancel()
                    with suppress(asyncio.CancelledError):
                        await pending

            assert new_head.primary.ref.artifact_id != initial.primary.ref.artifact_id
            assert tuple(hit.text for hit in result.hits) == ("Stable searchable fact.",)
            assert result.hits[0].hit.artifact_ref == initial.primary.ref
            assert result.rerank is not None

    asyncio.run(scenario())


def test_memory_search_keeps_exact_identity_when_an_unrelated_memory_is_added_before_query() -> None:
    async def scenario() -> None:
        async with open_builtin_runtime(BuiltinConfig(database=SQLiteConfig())) as runtime:
            assert runtime.scopes is not None
            scope = await runtime.scopes.create(
                ScopeDraft(
                    title="Perpetually stale search",
                    summary="Revision conflict retry limit",
                    idempotency_key="perpetually-stale-search",
                )
            )
            scope_id = scope.scope_id
            assert runtime.atomic_memory is not None
            memory = runtime.atomic_memory.for_scope(scope_id)
            initial = await memory.create((AtomicMemoryContent(kind="fact", text="Stable searchable fact."),))

            provider: Any = runtime._provider
            service = provider.atomic_memory.index
            original_search = service.search
            update_number = 0

            async def advance_head_before_search(_self: Any, *args: Any, **kwargs: Any) -> Any:
                nonlocal update_number
                update_number += 1
                await memory.create((AtomicMemoryContent(kind="fact", text=f"Concurrent update {update_number}."),))
                return await original_search(*args, **kwargs)

            service.search = MethodType(advance_head_before_search, service)
            try:
                result = await asyncio.wait_for(
                    memory.search("stable searchable", mode="text"),
                    timeout=TIMEOUT_SECONDS,
                )
                assert tuple(hit.hit.artifact_ref for hit in result.hits) == (initial.primary.ref,)
                assert update_number == 1
                assert len((await memory.list()).items) == 2
            finally:
                service.search = original_search

    asyncio.run(scenario())


def _database_config(
    database_kind: DatabaseKind,
    mode: MemorySearchMode,
    tmp_path: Path,
) -> SQLiteConfig | OceanBaseConfig:
    if database_kind == "oceanbase":
        url = os.environ.get("POWERCONTEXT_TEST_OCEANBASE_URL")
        if url is None:
            pytest.skip("set POWERCONTEXT_TEST_OCEANBASE_URL to a dedicated OceanBase MySQL-mode test database")
        return OceanBaseConfig(url=SecretStr(url))

    return SQLiteConfig(url=f"sqlite+aiosqlite:///{tmp_path / f'memory-search-{mode}.db'}")