LangBot/plugins/KnowledgeEngine/RAGFlowConnector/components/shared_state.py

"""Installation- and knowledge-base-scoped state using the SDK-bound Host proxy.

Bundled in each independent archive. No filesystem/global tenant registry.

Locks and fences are keyed by ``(installation binding, knowledge-base identity)``
so an ambiguous failure for one knowledge base never blocks the other knowledge
bases that share the same installation; the platform storage key is the same for
a whole installation, hence the knowledge-base id is always folded into our keys.
Fences are persisted in installation-bound Host storage (key includes the
knowledge-base id) with an in-process cache in front, so a restarted worker still
recognises them. An unreadable fence record refuses the operation instead of
being read as "no fence", and such a read is retried rather than cached. Only
dispatched mutations whose outcome is unknown fence; a deterministic validation
error raised before any remote call and ordinary cancellation before a call is
dispatched do not. A fence that cannot be persisted aborts the mutation before
it is dispatched, and a fence is released only after the provider answer was
confirmed and its Host record written.
"""
import asyncio
import hashlib
import inspect
import json
import logging
from functools import wraps
from contextlib import asynccontextmanager
import httpx
from weakref import WeakValueDictionary

logger = logging.getLogger(__name__)

HTTP_TOTAL_TIMEOUT = 150.0
FENCE_KEY_PREFIX = "ke.fence.v1."
# Fallback fence scope when the wrapped method carries no knowledge-base identity.
_INSTALLATION_FENCE_KEY = FENCE_KEY_PREFIX + "installation"
_IDENTITY_ARGUMENTS = ('kb_id', 'collection_id', 'knowledge_base_id')


class AmbiguousMutationError(RuntimeError):
    """A dispatched mutation's outcome is unknown and requires reconciliation.

    Raise this from a call site when a remote mutation may have taken effect but
    its result (or its durable record) is unknown. Only these failures fence a
    knowledge base; deterministic errors raised before any remote call, and
    ordinary cancellation before a call is dispatched, do not fence.
    """


class FenceStateUnavailableError(RuntimeError):
    """Persisted fence state could not be read, so the operation is refused.

    An unreadable fence record cannot be told apart from a knowledge base that
    is fenced after an ambiguous mutation, so an unknown fence state fails
    closed. The read is not cached: the next attempt re-reads the durable state,
    and nothing has been dispatched while the state stays unknown.
    """


class FenceLifecycleError(RuntimeError):
    """A mutation was aborted before dispatch because its fence is not in place.

    Both failures below mean "this call dispatched nothing": the pre-dispatch
    protection could not be persisted, or an earlier unresolved mutation is
    still fenced. Call sites must report them instead of folding them into an
    ordinary provider failure.
    """


class FencePersistError(FenceLifecycleError):
    """The pre-dispatch fence record could not be persisted.

    Raised by ``fence`` when the Host write of the fence record fails: the fence
    is not in place, so ``dispatched`` aborts before its request instead of
    dispatching a mutation whose protection does not exist.
    """


class FencedKnowledgeBaseError(FenceLifecycleError):
    """A mutation was refused: an earlier mutation's outcome is unresolved.

    Raised when a serialized operation, or a dispatch inside one, is attempted
    while the knowledge base (or the whole installation) is still fenced. A
    later successful operation therefore never clears an earlier unresolved
    operation's fence.
    """


def is_ambiguous_http_failure(exc):
    """Whether a failed provider call had an unknown outcome.

    Only a rejection the protocol makes unambiguous releases the caller's
    protection: a 4xx response means the provider received and refused the
    request, and a connection that was never established means nothing was
    dispatched. A 5xx or gateway response is NOT proof that the request was not
    applied — the provider or an intermediary may have committed the change and
    then failed — so it stays ambiguous together with timeouts, lost connections
    and unparsable responses.
    """
    if isinstance(exc, httpx.HTTPStatusError):
        response = getattr(exc, "response", None)
        return response is None or response.status_code >= 500
    if isinstance(exc, (httpx.ConnectError, httpx.ConnectTimeout)):
        return False
    return True


async def _settle_commit(task):
    """Wait for a dispatched commit and report its outcome.

    Repeated caller cancellation never abandons the wait, so a Host commit that
    was already dispatched settles before its lock is released. Returns
    ``(caller_cancelled, failure)``; ``failure`` is the commit's exception, or
    ``asyncio.CancelledError`` when the commit itself was cancelled.
    """
    cancelled = False
    while not task.done():
        try:
            await asyncio.shield(task)
        except asyncio.CancelledError:
            cancelled = True
        except Exception:
            # The commit raised; its outcome is read from the state below.
            pass
    if task.cancelled():
        return cancelled, asyncio.CancelledError()
    return cancelled, task.exception()


def _kb_identity(method, instance, args, kwargs):
    """Extract the knowledge-base identity from a serialized method's arguments.

    Reads a ``kb_id`` / ``collection_id`` / ``knowledge_base_id`` argument first,
    then an argument exposing ``knowledge_base_id`` or ``get_collection_id()``
    (ingestion passes the whole context). The knowledge-base id is preferred: the
    collection id is a vector-store identity that need not equal it, and every
    other entry point of the connector is addressed by the knowledge-base id, so
    locking and fencing on the collection id would split one knowledge base into
    two independently protected scopes. Returns ``None`` when the wrapped method
    carries no knowledge-base identity, in which case the caller falls back to an
    installation-wide lock and fence.
    """
    try:
        bound = inspect.signature(method).bind_partial(instance, *args, **kwargs)
    except (TypeError, ValueError):
        return None
    for name in _IDENTITY_ARGUMENTS:
        value = bound.arguments.get(name)
        if isinstance(value, str) and value:
            return value
    for value in bound.arguments.values():
        identity = getattr(value, 'knowledge_base_id', None)
        if isinstance(identity, str) and identity:
            return identity
        getter = getattr(value, 'get_collection_id', None)
        if callable(getter):
            try:
                identity = getter()
            except Exception:
                continue
            if isinstance(identity, str) and identity:
                return identity
    return None


def _kb_fence_key(kb_identity):
    if kb_identity is None:
        return _INSTALLATION_FENCE_KEY
    return FENCE_KEY_PREFIX + hashlib.sha256(kb_identity.encode()).hexdigest()


class SerialState:
    def __init__(self):
        self._locks = WeakValueDictionary()
        # binding -> set of fence keys; cache in front of persisted fences.
        self._fenced = {}
        self._fence_loaded = set()
        # binding -> single-flight lock and mutation counter for the fence read.
        self._fence_load_locks: dict = {}
        self._fence_generation: dict = {}

    def _bump_fence_generation(self, binding):
        self._fence_generation[binding] = self._fence_generation.get(binding, 0) + 1

    @staticmethod
    def binding(plugin):
        try:
            binding = plugin.get_installation_binding()
        except Exception:
            # Revocation has no invocation context; fall back to dedicated scope.
            binding = None
        return binding if binding is not None else ('dedicated',)

    @staticmethod
    def _scope(binding, kb_identity):
        return (binding, kb_identity)

    def _fenced_keys(self, binding):
        return self._fenced.setdefault(binding, set())

    def is_fenced(self, binding, kb_identity):
        fenced = self._fenced_keys(binding)
        return (
            _kb_fence_key(kb_identity) in fenced
            or _INSTALLATION_FENCE_KEY in fenced
        )

    async def _load_fences(self, plugin, binding):
        """Read persisted fences once per binding, or refuse the operation.

        A failed read means the pending ambiguity of the installation is
        unknown, so it must not be mistaken for "no fences": the operation is
        refused with ``FenceStateUnavailableError`` and nothing is dispatched.
        The binding is not marked as loaded, so the next attempt re-reads and
        the knowledge base becomes usable again once storage answers. The read
        is single-flight per installation, and a snapshot read before a sibling
        operation changed the fence state is discarded instead of installed.
        """
        if binding in self._fence_loaded:
            return
        lock = self._fence_load_locks.get(binding)
        if lock is None:
            lock = self._fence_load_locks[binding] = asyncio.Lock()
        async with lock:
            if binding in self._fence_loaded:
                return
            for _ in range(8):
                generation = self._fence_generation.get(binding, 0)
                snapshot = await self._read_fences(plugin, binding)
                if generation == self._fence_generation.get(binding, 0):
                    self._fenced_keys(binding).update(snapshot)
                    self._fence_loaded.add(binding)
                    return
            raise FenceStateUnavailableError(
                f"Fence state for {binding!r} kept changing while it was being read"
            )

    async def _read_fences(self, plugin, binding) -> set:
        """Read the persisted fence keys for one installation.

        An unreadable marker counts as fenced and a failed read refuses the
        operation with ``FenceStateUnavailableError``; the read is never cached.
        """
        try:
            keys = await plugin.get_plugin_storage_keys()
            snapshot = set()
            for key in keys:
                if not key.startswith(FENCE_KEY_PREFIX):
                    continue
                try:
                    marker = json.loads(await plugin.get_plugin_storage(key))
                except Exception:
                    # Unreadable marker: fail closed on this knowledge base.
                    marker = {'reason': 'unreadable fence record'}
                if marker is not None:
                    snapshot.add(key)
            return snapshot
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            logger.warning("Could not read persisted fences for %r", binding, exc_info=True)
            raise FenceStateUnavailableError(
                f"Could not read persisted fences for {binding!r}; pending ambiguous "
                f"mutations are unknown: {exc!r}"
            ) from exc

    async def run(self, plugin, kb_identity, operation, *, allow_fenced=False):
        binding = self.binding(plugin)
        key = self._scope(binding, kb_identity)
        lock = self._locks.get(key)
        if lock is None:
            lock = self._locks[key] = asyncio.Lock()
        async with lock:
            await self._load_fences(plugin, binding)
            if not allow_fenced and self.is_fenced(binding, kb_identity):
                raise FencedKnowledgeBaseError(
                    f"Knowledge base {kb_identity!r} fenced after an ambiguous mutation; "
                    "reconcile with the provider before retrying"
                )
            try:
                return await operation()
            except AmbiguousMutationError:
                await self.fence(plugin, kb_identity, 'ambiguous remote mutation')
                raise

    async def dispatched(self, plugin, kb_identity, operation, reason):
        """Dispatch one remote mutation under a pre-persisted fence.

        ``operation`` must cover the whole mutation: the provider request, the
        response validation and the durable result/mapping write. The fence is
        persisted before the request goes out — a failed persist aborts here,
        before any remote call — and it is released only once ``operation`` has
        returned, so an ambiguous outcome or a caller cancellation keeps the
        knowledge base fenced for reconciliation. A knowledge base that is
        already fenced is refused before anything is written or dispatched, so a
        later successful operation never clears an earlier unresolved
        operation's fence.
        """
        binding = self.binding(plugin)
        if self.is_fenced(binding, kb_identity):
            # The enclosing serialized call loaded and validated the fence state,
            # so this is a real unresolved mutation rather than a stale view.
            raise FencedKnowledgeBaseError(
                f"Knowledge base {kb_identity!r} fenced by an unresolved mutation; "
                "reconcile with the provider before dispatching another mutation"
            )
        await self.fence(plugin, kb_identity, reason)
        try:
            result = await operation()
        except asyncio.CancelledError:
            # Dispatched, outcome unknown: keep the fence and let the caller see it.
            raise
        except Exception as exc:
            if not is_ambiguous_http_failure(exc):
                await self.clear_fence(plugin, kb_identity)
            raise
        await self.clear_fence(plugin, kb_identity)
        return result

    async def fence(self, plugin, kb_identity, reason):
        """Persist one fence, then mark that knowledge base in this process.

        The durable record is written first and a failed write is fatal
        (``FencePersistError``), so a caller that fences before dispatching a
        mutation aborts before its request instead of dispatching under a fence
        that does not exist. Only a persisted fence sets the in-process marker,
        so a failed write and a caller cancellation both leave no marker behind.
        The residual race — the Host write landed but its acknowledgement was
        lost — can make a later read load a fence for an operation that was
        never dispatched; that is the conservative direction and is accepted.
        """
        binding = self.binding(plugin)
        fence_key = _kb_fence_key(kb_identity)
        payload = json.dumps({'kb_id': kb_identity, 'reason': reason}).encode()
        cancelled, failure = await _settle_commit(
            asyncio.create_task(plugin.set_plugin_storage(fence_key, payload))
        )
        if cancelled:
            raise asyncio.CancelledError
        if failure is not None:
            raise FencePersistError(
                f"Could not persist fence for {kb_identity!r}: {failure!r}"
            ) from failure
        self._fenced_keys(binding).add(fence_key)
        self._bump_fence_generation(binding)

    async def clear_fence(self, plugin, kb_identity):
        """Drop the fence for one knowledge base, in process and in storage.

        Used when a knowledge base is deleted: its pending ambiguity is moot, so
        the operator path out of a fence is deletion. A lingering persisted
        marker is inert because the configuration tombstone is already written.
        """
        binding = self.binding(plugin)
        fence_key = _kb_fence_key(kb_identity)
        self._fenced_keys(binding).discard(fence_key)
        self._bump_fence_generation(binding)
        cancelled, failure = await _settle_commit(
            asyncio.create_task(plugin.set_plugin_storage(fence_key, b'null'))
        )
        if failure is not None:
            logger.warning("Could not clear fence for %r: %r", kb_identity, failure)
        if cancelled:
            raise asyncio.CancelledError

    async def write(self, plugin, kb_identity, key, value):
        # A caller cancellation must not abandon a dispatched Host commit while
        # its lock is released; settle it and fence only a real ambiguity.
        task = asyncio.create_task(plugin.set_plugin_storage(key, value))
        cancelled, failure = await _settle_commit(task)
        if failure is not None:
            try:
                await self.fence(plugin, kb_identity, f'Host write outcome unknown: {failure!r}')
            except FencePersistError:
                # The write failure is what the caller must see; a fence that
                # could not be persisted is reported alongside it, not instead.
                logger.error(
                    "Could not persist fence for %r after a failed Host write",
                    kb_identity, exc_info=True,
                )
        if cancelled:
            raise asyncio.CancelledError
        if failure is not None:
            raise failure
        return task.result()

    def release_binding(self, binding):
        """Drop process-local state keyed by one installation binding."""
        for key in [k for k in self._locks if k[0] == binding]:
            self._locks.pop(key, None)
        self._fenced.pop(binding, None)
        self._fence_loaded.discard(binding)
        self._fence_load_locks.pop(binding, None)
        self._fence_generation.pop(binding, None)


def serialized(method=None, *, allow_fenced=False):
    def decorate(func):
        @wraps(func)
        async def call(self, *args, **kwargs):
            return await self._state.run(
                self.plugin,
                _kb_identity(func, self, args, kwargs),
                lambda: func(self, *args, **kwargs),
                allow_fenced=allow_fenced,
            )
        return call
    return decorate if method is None else decorate(method)


class ConfigStore:
    """No credential cache: every delete resolves durable installation storage.

    A JSON null tombstone makes missing KBs explicit without ambiguous delete RPCs.
    The one mutation lock also orders ingest/delete/create inside this worker.
    Configuration and document mappings are keyed by the Host knowledge-base id;
    a record written under another knowledge base's id is never adopted, and only
    an exact mapping for the requested knowledge base is honoured.
    """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._state = SerialState()
        self._http_slots = asyncio.Semaphore(4)

    async def initialize(self) -> None:
        await super().initialize()
        # The runtime dispatches installation revocation to the plugin instance,
        # which holds no back-reference to its components; hand it this
        # component's state so the revocation hook can drop per-binding locks and
        # fences. Without the hook (langbot-plugin 0.7.4) this attribute is simply
        # never read.
        plugin = getattr(self, 'plugin', None)
        if plugin is not None:
            plugin.knowledge_engine_serial_state = self._state

    @asynccontextmanager
    async def http_client(self):
        # Include queueing and the whole response body in the deadline, not just
        # individual socket reads. Clients/credentials never survive a call.
        async with asyncio.timeout(HTTP_TOTAL_TIMEOUT):
            async with self._http_slots:
                async with httpx.AsyncClient() as client:
                    yield client

    @staticmethod
    def _key(kb_id):
        return "ke.config.v1." + hashlib.sha256(kb_id.encode()).hexdigest()

    async def _save_config(self, kb_id, config):
        value = json.dumps(config, ensure_ascii=False).encode()
        if len(value) > 65536:
            raise ValueError("Knowledge-base configuration exceeds 64 KiB")
        await self._state.write(self.plugin, kb_id, self._key(kb_id), value)

    async def _load_config(self, kb_id):
        key = self._key(kb_id)
        # Only an authoritative keys response establishes absence; RPC errors propagate.
        if key not in await self.plugin.get_plugin_storage_keys():
            return None
        value = json.loads(await self.plugin.get_plugin_storage(key))
        if value is not None and not isinstance(value, dict):
            raise ValueError("Invalid persisted knowledge-base configuration")
        return value

    @staticmethod
    def _document_key(kb_id, host_document_id):
        return "ke.document.v1." + hashlib.sha256(
            json.dumps([kb_id, host_document_id]).encode()
        ).hexdigest()

    async def _load_document(self, kb_id, host_document_id):
        """Resolve this knowledge base's mapping, never another knowledge base's.

        A record written under another knowledge base's id is never adopted:
        attributing historic collection-keyed data requires explicit
        reconciliation. An absent exact mapping therefore refuses the operation,
        and this path performs no writes.
        """
        key = self._document_key(kb_id, host_document_id)
        keys = await self.plugin.get_plugin_storage_keys()
        if key in keys:
            value = json.loads(await self.plugin.get_plugin_storage(key))
        else:
            # Current Host passes the durable upstream ID on delete; older Host
            # revisions pass its file UUID. Resolve either only for a record this
            # knowledge base itself wrote, without trusting an arbitrary
            # unrecorded identifier.
            matches = []
            for candidate in keys:
                if candidate.startswith('ke.document.v1.'):
                    item = json.loads(await self.plugin.get_plugin_storage(candidate))
                    if item.get('kb_id') == kb_id and item.get('upstream_id') == host_document_id:
                        matches.append(item)
            if len(matches) > 1:
                raise ValueError('Ambiguous upstream document ID')
            return matches[0] if matches else None
        if not isinstance(value, dict) or not isinstance(value.get('upstream_id'), str):
            raise ValueError('Invalid upstream document mapping; reconcile installation')
        return value

    async def _save_document(self, kb_id, host_document_id, value):
        value = {**value, 'kb_id': kb_id, 'host_document_id': host_document_id}
        await self._state.write(self.plugin, kb_id, self._document_key(kb_id, host_document_id),
                                json.dumps(value).encode())