LangBot/src/langbot/pkg/pipeline/pool.py

from __future__ import annotations

import asyncio
import dataclasses
import inspect
import typing
import uuid

import langbot_plugin.api.definition.abstract.platform.adapter as abstract_platform_adapter
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.platform.events as platform_events
import langbot_plugin.api.entities.builtin.platform.message as platform_message
import langbot_plugin.api.entities.builtin.provider.session as provider_session

from ..api.http.context import ExecutionContext
from . import plugin_diagnostics

QueryCacheKey = tuple[str, str]
LegacyQueryKey = tuple[str, int]
QueryCounterKey = tuple[str, str, int]
SingletonContextResolver = typing.Callable[
    [],
    ExecutionContext | typing.Awaitable[ExecutionContext],
]


class ExecutionContextRequiredError(ValueError):
    """Raised when runtime work is created without a trusted Workspace scope."""


class ExecutionContextMismatchError(ValueError):
    """Raised when entity fields conflict with their trusted execution scope."""


class QueryNotFoundError(LookupError):
    """Raised when a query does not exist inside the requested Workspace."""


class QueryPoolCapacityError(RuntimeError):
    """Raised when no queued query can be discarded to admit new work."""


def _validate_execution_context(execution_context: ExecutionContext) -> None:
    if not isinstance(execution_context, ExecutionContext):
        raise ExecutionContextRequiredError('A trusted ExecutionContext is required')
    if not isinstance(execution_context.instance_uuid, str) or not execution_context.instance_uuid.strip():
        raise ExecutionContextRequiredError('ExecutionContext.instance_uuid is required')
    if not isinstance(execution_context.workspace_uuid, str) or not execution_context.workspace_uuid.strip():
        raise ExecutionContextRequiredError('ExecutionContext.workspace_uuid is required')
    if (
        isinstance(execution_context.placement_generation, bool)
        or not isinstance(execution_context.placement_generation, int)
        or execution_context.placement_generation <= 0
    ):
        raise ExecutionContextRequiredError('ExecutionContext.placement_generation must be a positive integer')
    for field_name in ('bot_uuid', 'pipeline_uuid', 'query_uuid'):
        value = getattr(execution_context, field_name)
        if value is not None and (not isinstance(value, str) or not value.strip()):
            raise ExecutionContextRequiredError(f'ExecutionContext.{field_name} must be a non-empty string when set')


def bind_execution_context(
    execution_context: ExecutionContext,
    *,
    bot_uuid: str | None = None,
    pipeline_uuid: str | None = None,
    query_uuid: str | None = None,
) -> ExecutionContext:
    """Bind runtime entity identifiers without allowing scope substitution."""

    _validate_execution_context(execution_context)

    requested_fields = {
        'bot_uuid': bot_uuid,
        'pipeline_uuid': pipeline_uuid,
        'query_uuid': query_uuid,
    }
    updates: dict[str, str] = {}
    for field_name, requested_value in requested_fields.items():
        if requested_value is None:
            continue
        if not isinstance(requested_value, str) or not requested_value.strip():
            raise ExecutionContextRequiredError(f'{field_name} must be a non-empty string')
        current_value = getattr(execution_context, field_name)
        if current_value is not None and current_value != requested_value:
            raise ExecutionContextMismatchError(f'ExecutionContext.{field_name} does not match the runtime entity')
        if current_value is None:
            updates[field_name] = requested_value

    if not updates:
        return execution_context
    return dataclasses.replace(execution_context, **updates)


def get_query_execution_context(query: pipeline_query.Query) -> ExecutionContext:
    """Return and validate the trusted context attached to a Query."""

    attached_context = getattr(query, '_execution_context', None)
    bot_uuid = getattr(query, 'bot_uuid', None)
    pipeline_uuid = getattr(query, 'pipeline_uuid', None)
    query_uuid = getattr(query, 'query_uuid', None)

    if isinstance(attached_context, ExecutionContext):
        return bind_execution_context(
            attached_context,
            bot_uuid=bot_uuid,
            pipeline_uuid=pipeline_uuid,
            query_uuid=query_uuid,
        )

    raise ExecutionContextRequiredError('Query is missing its trusted ExecutionContext')


class QueryPool:
    """Workspace-scoped queue of requests waiting for pipeline scheduling."""

    query_id_counter: int
    pool_lock: asyncio.Lock
    queries: list[pipeline_query.Query]
    cached_queries: dict[QueryCacheKey, pipeline_query.Query]
    legacy_query_index: dict[LegacyQueryKey, str]
    query_count_by_scope: dict[QueryCounterKey, int]
    condition: asyncio.Condition

    def __init__(
        self,
        singleton_context_resolver: SingletonContextResolver | None = None,
        *,
        max_queries: int = 1000,
        max_queries_per_workspace: int = 100,
    ):
        if max_queries < 1:
            raise ValueError('max_queries must be positive')
        if max_queries_per_workspace < 1:
            raise ValueError('max_queries_per_workspace must be positive')
        if max_queries_per_workspace > max_queries:
            raise ValueError('max_queries_per_workspace cannot exceed max_queries')
        self.query_id_counter = 0
        self.pool_lock = asyncio.Lock()
        self.queries = []
        self.cached_queries = {}
        self.active_query_count_by_workspace: dict[str, int] = {}
        self.legacy_query_index = {}
        self.query_count_by_scope = {}
        self.dropped_query_count_by_scope: dict[QueryCounterKey, int] = {}
        self.condition = asyncio.Condition(self.pool_lock)
        self._singleton_context_resolver = singleton_context_resolver
        self.max_queries = max_queries
        self.max_queries_per_workspace = max_queries_per_workspace

    def _discard_queued_query_locked(
        self,
        *,
        workspace_uuid: str | None = None,
    ) -> pipeline_query.Query | None:
        """Discard the oldest queued query from one scope and all indexes."""

        for index, query in enumerate(self.queries):
            execution_context = get_query_execution_context(query)
            if workspace_uuid is not None and execution_context.workspace_uuid != workspace_uuid:
                continue
            self.queries.pop(index)
            query_uuid = execution_context.query_uuid
            if query_uuid is not None:
                self.cached_queries.pop((execution_context.workspace_uuid, query_uuid), None)
            self.legacy_query_index.pop((execution_context.workspace_uuid, query.query_id), None)
            query_workspace_uuid = execution_context.workspace_uuid
            remaining = self.active_query_count_by_workspace.get(query_workspace_uuid, 0) - 1
            if remaining > 0:
                self.active_query_count_by_workspace[query_workspace_uuid] = remaining
            else:
                self.active_query_count_by_workspace.pop(query_workspace_uuid, None)
            plugin_diagnostics.discard_query_state(query)
            counter_key = (
                execution_context.instance_uuid,
                execution_context.workspace_uuid,
                execution_context.placement_generation,
            )
            self.dropped_query_count_by_scope[counter_key] = self.dropped_query_count_by_scope.get(counter_key, 0) + 1
            return query
        return None

    def _admit_query_locked(self, workspace_uuid: str) -> None:
        workspace_query_count = self.active_query_count_by_workspace.get(workspace_uuid, 0)
        if workspace_query_count >= self.max_queries_per_workspace:
            if self._discard_queued_query_locked(workspace_uuid=workspace_uuid) is None:
                raise QueryPoolCapacityError(f'Workspace query capacity reached ({self.max_queries_per_workspace})')

        if len(self.cached_queries) >= self.max_queries:
            if self._discard_queued_query_locked() is None:
                raise QueryPoolCapacityError(f'Global query capacity reached ({self.max_queries})')

    def mark_query_running_locked(self, query: pipeline_query.Query) -> None:
        """Remove a scheduled query from the overload-discardable queue."""

        if not self.pool_lock.locked():
            raise RuntimeError('Query pool lock is required to schedule a query')
        for index, queued_query in enumerate(self.queries):
            if queued_query is query:
                self.queries.pop(index)
                return
        raise QueryNotFoundError('Scheduled query is no longer queued')

    def _make_scope_counter_room_locked(self, counter_key: QueryCounterKey) -> None:
        """Retain recent counters without pinning every historical Workspace."""

        if counter_key in self.query_count_by_scope:
            return
        while len(self.query_count_by_scope) >= self.max_queries:
            stale_key = next(
                (
                    existing_key
                    for existing_key in self.query_count_by_scope
                    if self.active_query_count_by_workspace.get(existing_key[1], 0) <= 0
                ),
                None,
            )
            if stale_key is None:
                raise QueryPoolCapacityError('Query counter capacity reached while every Workspace is active')
            self.query_count_by_scope.pop(stale_key, None)
            self.dropped_query_count_by_scope.pop(stale_key, None)

    async def resolve_execution_context(
        self,
        execution_context: ExecutionContext | None,
        *,
        bot_uuid: str,
        pipeline_uuid: str | None,
        query_uuid: str | None = None,
    ) -> ExecutionContext:
        """Resolve an explicit scope or the opt-in OSS singleton scope."""

        if execution_context is None:
            if self._singleton_context_resolver is None:
                raise ExecutionContextRequiredError('ExecutionContext is required; no singleton resolver is configured')
            resolved_context = self._singleton_context_resolver()
            if inspect.isawaitable(resolved_context):
                resolved_context = await resolved_context
            execution_context = resolved_context

        return bind_execution_context(
            execution_context,
            bot_uuid=bot_uuid,
            pipeline_uuid=pipeline_uuid,
            query_uuid=query_uuid,
        )

    async def add_query(
        self,
        bot_uuid: str,
        launcher_type: provider_session.LauncherTypes,
        launcher_id: int | str,
        sender_id: int | str,
        message_event: platform_events.MessageEvent,
        message_chain: platform_message.MessageChain,
        adapter: abstract_platform_adapter.AbstractMessagePlatformAdapter,
        pipeline_uuid: str | None = None,
        routed_by_rule: bool = False,
        variables: dict[str, typing.Any] | None = None,
        execution_context: ExecutionContext | None = None,
    ) -> pipeline_query.Query:
        """Create a query and cache it under an opaque, Workspace-scoped key."""

        query_uuid = str(uuid.uuid4())
        execution_context = await self.resolve_execution_context(
            execution_context,
            bot_uuid=bot_uuid,
            pipeline_uuid=pipeline_uuid,
            query_uuid=query_uuid,
        )

        async with self.condition:
            self._admit_query_locked(execution_context.workspace_uuid)
            query_id = self.query_id_counter
            initial_variables: dict[str, typing.Any] = {'_routed_by_rule': routed_by_rule}
            if variables:
                initial_variables.update(variables)
            query = pipeline_query.Query(
                instance_uuid=execution_context.instance_uuid,
                workspace_uuid=execution_context.workspace_uuid,
                placement_generation=execution_context.placement_generation,
                bot_uuid=bot_uuid,
                query_id=query_id,
                query_uuid=query_uuid,
                launcher_type=launcher_type,
                launcher_id=launcher_id,
                sender_id=sender_id,
                message_event=message_event,
                message_chain=message_chain,
                variables=initial_variables,
                resp_messages=[],
                resp_message_chain=[],
                adapter=adapter,
                pipeline_uuid=pipeline_uuid,
            )

            # langbot-plugin 0.4.13 ignores these forward-compatible fields.
            # Attach them explicitly until the Workspace-aware SDK is released.
            object.__setattr__(query, 'instance_uuid', execution_context.instance_uuid)
            object.__setattr__(query, 'workspace_uuid', execution_context.workspace_uuid)
            object.__setattr__(
                query,
                'placement_generation',
                execution_context.placement_generation,
            )
            object.__setattr__(query, 'query_uuid', query_uuid)
            object.__setattr__(query, '_execution_context', execution_context)

            self.queries.append(query)
            self.cached_queries[(execution_context.workspace_uuid, query_uuid)] = query
            self.active_query_count_by_workspace[execution_context.workspace_uuid] = (
                self.active_query_count_by_workspace.get(execution_context.workspace_uuid, 0) + 1
            )
            self.legacy_query_index[(execution_context.workspace_uuid, query_id)] = query_uuid
            self.query_id_counter += 1
            counter_key = (
                execution_context.instance_uuid,
                execution_context.workspace_uuid,
                execution_context.placement_generation,
            )
            # A Workspace has only one active placement. Drop obsolete
            # generation counters so deployment churn cannot grow these maps.
            for existing_key in tuple(self.query_count_by_scope):
                if existing_key[:2] == counter_key[:2] and existing_key != counter_key:
                    self.query_count_by_scope.pop(existing_key, None)
                    self.dropped_query_count_by_scope.pop(existing_key, None)
            self._make_scope_counter_room_locked(counter_key)
            self.query_count_by_scope[counter_key] = self.query_count_by_scope.get(counter_key, 0) + 1
            self.condition.notify_all()
            return query

    def get_query_count(self, execution_context: ExecutionContext) -> int:
        """Return the lifetime query count for one active placement scope."""

        _validate_execution_context(execution_context)
        return self.query_count_by_scope.get(
            (
                execution_context.instance_uuid,
                execution_context.workspace_uuid,
                execution_context.placement_generation,
            ),
            0,
        )

    def get_dropped_query_count(self, execution_context: ExecutionContext) -> int:
        """Return overload drops for one active placement scope."""

        _validate_execution_context(execution_context)
        return self.dropped_query_count_by_scope.get(
            (
                execution_context.instance_uuid,
                execution_context.workspace_uuid,
                execution_context.placement_generation,
            ),
            0,
        )

    async def get_query(
        self,
        workspace_uuid: str,
        query_uuid: str,
    ) -> pipeline_query.Query | None:
        """Return a query only from the explicitly selected Workspace."""

        async with self.pool_lock:
            return self.cached_queries.get((workspace_uuid, query_uuid))

    async def require_query(
        self,
        workspace_uuid: str,
        query_uuid: str,
    ) -> pipeline_query.Query:
        """Return a scoped query or raise without checking other Workspaces."""

        query = await self.get_query(workspace_uuid, query_uuid)
        if query is None:
            raise QueryNotFoundError(f'Query {query_uuid!r} was not found in Workspace {workspace_uuid!r}')
        return query

    async def get_query_by_legacy_id(
        self,
        workspace_uuid: str,
        query_id: int,
    ) -> pipeline_query.Query | None:
        """Resolve a legacy integer ID within one explicit Workspace."""

        async with self.pool_lock:
            query_uuid = self.legacy_query_index.get((workspace_uuid, query_id))
            if query_uuid is None:
                return None
            return self.cached_queries.get((workspace_uuid, query_uuid))

    async def remove_query(self, query: pipeline_query.Query) -> bool:
        """Remove a query and both of its Workspace-scoped indexes."""

        execution_context = get_query_execution_context(query)
        query_uuid = execution_context.query_uuid
        if query_uuid is None:
            raise ExecutionContextRequiredError('Query.query_uuid is required for removal')

        async with self.pool_lock:
            cache_key = (execution_context.workspace_uuid, query_uuid)
            cached_query = self.cached_queries.get(cache_key)
            if cached_query is not query:
                return False
            del self.cached_queries[cache_key]
            remaining = self.active_query_count_by_workspace.get(execution_context.workspace_uuid, 0) - 1
            if remaining > 0:
                self.active_query_count_by_workspace[execution_context.workspace_uuid] = remaining
            else:
                self.active_query_count_by_workspace.pop(execution_context.workspace_uuid, None)
            self.legacy_query_index.pop(
                (execution_context.workspace_uuid, query.query_id),
                None,
            )
            for index, queued_query in enumerate(self.queries):
                if queued_query is query:
                    self.queries.pop(index)
                    break
            plugin_diagnostics.discard_query_state(query)
            return True

    async def __aenter__(self):
        await self.pool_lock.acquire()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        self.pool_lock.release()