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

"""Workspace-scoped message aggregation and debounce support."""

from __future__ import annotations

import asyncio
import time
import typing
from dataclasses import dataclass, field

import langbot_plugin.api.definition.abstract.platform.adapter as abstract_platform_adapter
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 ..core.task_boundary import create_detached_task, run_in_workspace_uow
from .pool import ExecutionContextMismatchError
from ..workspace.errors import WorkspaceError, WorkspaceInvariantError

if typing.TYPE_CHECKING:
    from ..core import app

MAX_BUFFER_MESSAGES = 10

AggregationKey = tuple[
    str,
    str,
    int,
    str,
    str | None,
    str,
    int | str,
]


@dataclass
class PendingMessage:
    """A pending message carrying its trusted execution scope."""

    execution_context: ExecutionContext
    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
    routed_by_rule: bool = False
    timestamp: float = field(default_factory=time.time)


@dataclass
class SessionBuffer:
    """Pending messages for one scoped aggregation key."""

    aggregation_key: AggregationKey
    execution_context: ExecutionContext
    messages: list[PendingMessage] = field(default_factory=list)
    timer_task: asyncio.Task | None = None
    last_message_time: float = field(default_factory=time.time)


class MessageAggregator:
    """Debounce consecutive messages without crossing Workspace boundaries."""

    ap: app.Application
    buffers: dict[AggregationKey, SessionBuffer]
    lock: asyncio.Lock

    def __init__(self, ap: app.Application):
        self.ap = ap
        self.buffers = {}
        self._buffer_counts_by_scope: dict[
            tuple[str, str, int],
            int,
        ] = {}
        self.lock = asyncio.Lock()
        concurrency = self.ap.instance_config.data.get('concurrency', {})
        self.max_buffers = max(int(concurrency.get('pending_queries', 1000)), 1)
        self.max_buffers_per_workspace = max(
            int(concurrency.get('pending_queries_per_workspace', 100)),
            1,
        )

    def _get_aggregation_key(
        self,
        execution_context: ExecutionContext,
        bot_uuid: str,
        launcher_type: provider_session.LauncherTypes,
        launcher_id: int | str,
        pipeline_uuid: str | None,
    ) -> AggregationKey:
        """Build a key that cannot alias another Workspace, bot, or pipeline."""

        return (
            execution_context.instance_uuid,
            execution_context.workspace_uuid,
            execution_context.placement_generation,
            bot_uuid,
            pipeline_uuid,
            launcher_type.value,
            launcher_id,
        )

    async def _get_aggregation_config(
        self,
        execution_context: ExecutionContext,
        pipeline_uuid: str | None,
    ) -> tuple[bool, float]:
        """Return aggregation enablement and a clamped debounce delay."""

        default_enabled = False
        default_delay = 1.5

        if pipeline_uuid is None:
            return default_enabled, default_delay

        pipeline = await self.ap.pipeline_mgr.get_pipeline_by_uuid(
            execution_context,
            pipeline_uuid,
        )
        if pipeline is None:
            return default_enabled, default_delay

        config = pipeline.pipeline_entity.config or {}
        trigger_config = config.get('trigger', {})
        aggregation_config = trigger_config.get('message-aggregation', {})
        enabled = aggregation_config.get('enabled', default_enabled)

        delay_raw = aggregation_config.get('delay', default_delay)
        try:
            delay = float(delay_raw)
        except (TypeError, ValueError):
            delay = default_delay

        return enabled, max(1.0, min(10.0, delay))

    async def add_message(
        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,
        execution_context: ExecutionContext | None = None,
    ) -> None:
        """Buffer or directly enqueue a message in its trusted Workspace."""

        execution_context = await self.ap.query_pool.resolve_execution_context(
            execution_context,
            bot_uuid=bot_uuid,
            pipeline_uuid=pipeline_uuid,
        )
        enabled, delay = await self._get_aggregation_config(execution_context, pipeline_uuid)

        if not enabled:
            await self.ap.query_pool.add_query(
                bot_uuid=bot_uuid,
                launcher_type=launcher_type,
                launcher_id=launcher_id,
                sender_id=sender_id,
                message_event=message_event,
                message_chain=message_chain,
                adapter=adapter,
                pipeline_uuid=pipeline_uuid,
                routed_by_rule=routed_by_rule,
                execution_context=execution_context,
            )
            return

        aggregation_key = self._get_aggregation_key(
            execution_context,
            bot_uuid,
            launcher_type,
            launcher_id,
            pipeline_uuid,
        )
        pending_msg = PendingMessage(
            execution_context=execution_context,
            bot_uuid=bot_uuid,
            launcher_type=launcher_type,
            launcher_id=launcher_id,
            sender_id=sender_id,
            message_event=message_event,
            message_chain=message_chain,
            adapter=adapter,
            pipeline_uuid=pipeline_uuid,
            routed_by_rule=routed_by_rule,
        )

        force_flush = False
        bypass_aggregation = False
        async with self.lock:
            buffer = self.buffers.get(aggregation_key)
            if buffer is None:
                scope_key = (
                    execution_context.instance_uuid,
                    execution_context.workspace_uuid,
                    execution_context.placement_generation,
                )
                workspace_buffer_count = self._buffer_counts_by_scope.get(
                    scope_key,
                    0,
                )
                if len(self.buffers) >= self.max_buffers or workspace_buffer_count >= self.max_buffers_per_workspace:
                    bypass_aggregation = True
                else:
                    buffer = SessionBuffer(
                        aggregation_key=aggregation_key,
                        execution_context=execution_context,
                        messages=[pending_msg],
                    )
                    self.buffers[aggregation_key] = buffer
                    self._buffer_counts_by_scope[scope_key] = workspace_buffer_count + 1
            else:
                if buffer.execution_context != execution_context:
                    raise ExecutionContextMismatchError('Aggregation buffer ExecutionContext changed for the same key')
                if buffer.timer_task and not buffer.timer_task.done():
                    buffer.timer_task.cancel()
                buffer.messages.append(pending_msg)

            if not bypass_aggregation:
                buffer.last_message_time = time.time()
                if len(buffer.messages) >= MAX_BUFFER_MESSAGES:
                    force_flush = True
                else:
                    buffer.timer_task = create_detached_task(
                        self._delayed_flush(aggregation_key, delay, execution_context),
                        after_commit_manager=getattr(self.ap, 'persistence_mgr', None),
                        workspace_uuid=execution_context.workspace_uuid,
                    )

        if bypass_aggregation:
            await self.ap.query_pool.add_query(
                bot_uuid=bot_uuid,
                launcher_type=launcher_type,
                launcher_id=launcher_id,
                sender_id=sender_id,
                message_event=message_event,
                message_chain=message_chain,
                adapter=adapter,
                pipeline_uuid=pipeline_uuid,
                routed_by_rule=routed_by_rule,
                execution_context=execution_context,
            )
            return

        if force_flush:
            await self._flush_buffer(aggregation_key, execution_context)

    async def _delayed_flush(
        self,
        aggregation_key: AggregationKey,
        delay: float,
        execution_context: ExecutionContext,
    ) -> None:
        """Flush after the debounce delay using the captured context."""

        try:
            await asyncio.sleep(delay)
            await run_in_workspace_uow(
                self.ap,
                execution_context.workspace_uuid,
                lambda: self._flush_buffer(aggregation_key, execution_context),
            )
        except asyncio.CancelledError:
            pass
        except WorkspaceError as exc:
            self.ap.logger.info(
                f'Dropped an aggregated message because its Workspace execution binding is stale: {exc}'
            )

    async def _flush_buffer(
        self,
        aggregation_key: AggregationKey,
        execution_context: ExecutionContext,
    ) -> None:
        """Flush one buffer only when the captured scope still matches."""

        async with self.lock:
            buffer = self.buffers.get(aggregation_key)
            if buffer is None:
                return
            if buffer.execution_context != execution_context:
                raise ExecutionContextMismatchError('Timer ExecutionContext does not match the aggregation buffer')
            self.buffers.pop(aggregation_key)
            scope_key = aggregation_key[:3]
            scope_count = self._buffer_counts_by_scope.get(scope_key, 0)
            if scope_count <= 1:
                self._buffer_counts_by_scope.pop(scope_key, None)
            else:
                self._buffer_counts_by_scope[scope_key] = scope_count - 1

        if not buffer.messages:
            return

        message = buffer.messages[0] if len(buffer.messages) == 1 else self._merge_messages(buffer.messages)
        binding = await self.ap.workspace_service.get_execution_binding(
            execution_context.workspace_uuid,
            expected_generation=execution_context.placement_generation,
        )
        if binding.instance_uuid != execution_context.instance_uuid:
            raise WorkspaceInvariantError('Aggregation buffer instance does not match the active Workspace binding')
        await self.ap.query_pool.add_query(
            bot_uuid=message.bot_uuid,
            launcher_type=message.launcher_type,
            launcher_id=message.launcher_id,
            sender_id=message.sender_id,
            message_event=message.message_event,
            message_chain=message.message_chain,
            adapter=message.adapter,
            pipeline_uuid=message.pipeline_uuid,
            routed_by_rule=message.routed_by_rule,
            execution_context=message.execution_context,
        )

    def _merge_messages(self, messages: list[PendingMessage]) -> PendingMessage:
        """Merge message chains after proving all messages share one scope."""

        if not messages:
            raise ValueError('At least one pending message is required')
        if len(messages) == 1:
            return messages[0]

        base_msg = messages[0]
        base_key = self._get_aggregation_key(
            base_msg.execution_context,
            base_msg.bot_uuid,
            base_msg.launcher_type,
            base_msg.launcher_id,
            base_msg.pipeline_uuid,
        )
        for message in messages[1:]:
            message_key = self._get_aggregation_key(
                message.execution_context,
                message.bot_uuid,
                message.launcher_type,
                message.launcher_id,
                message.pipeline_uuid,
            )
            if message_key != base_key or message.execution_context != base_msg.execution_context:
                raise ExecutionContextMismatchError('Cannot merge pending messages from different execution scopes')

        merged_chain = platform_message.MessageChain([])
        for index, message in enumerate(messages):
            if index > 0:
                merged_chain.append(platform_message.Plain(text='\n'))
            for component in message.message_chain:
                merged_chain.append(component)

        return PendingMessage(
            execution_context=base_msg.execution_context,
            bot_uuid=base_msg.bot_uuid,
            launcher_type=base_msg.launcher_type,
            launcher_id=base_msg.launcher_id,
            sender_id=base_msg.sender_id,
            message_event=base_msg.message_event,
            message_chain=merged_chain,
            adapter=base_msg.adapter,
            pipeline_uuid=base_msg.pipeline_uuid,
            routed_by_rule=any(message.routed_by_rule for message in messages),
        )

    async def flush_all(self) -> None:
        """Flush all pending buffers without dropping their captured scopes."""

        async with self.lock:
            pending_buffers = [(key, buffer.execution_context) for key, buffer in self.buffers.items()]
            for buffer in self.buffers.values():
                if buffer.timer_task and not buffer.timer_task.done():
                    buffer.timer_task.cancel()

        for aggregation_key, execution_context in pending_buffers:
            try:
                await self._flush_buffer(aggregation_key, execution_context)
            except WorkspaceError as exc:
                self.ap.logger.info(
                    'Dropped an aggregated message during shutdown because its '
                    f'Workspace execution binding is stale: {exc}'
                )