LangBot/src/langbot/pkg/agent/runner/orchestrator.py

"""Agent run orchestrator for coordinating runner execution."""

from __future__ import annotations

import time
import typing

from langbot_plugin.api.entities.builtin.provider import message as provider_message
from langbot_plugin.api.entities.builtin.pipeline import query as pipeline_query

from ...core import app
from .binding_resolver import AgentBindingResolver
from .context_builder import AgentRunContextBuilder, AgentRunContextPayload
from .descriptor import AgentRunnerDescriptor
from .host_models import AgentBinding, AgentEventEnvelope
from .invoker import AgentRunnerInvoker
from .query_bridge import QueryRunBridge
from .registry import AgentRunnerRegistry
from .resource_builder import AgentResourceBuilder
from .result_normalizer import AgentResultNormalizer
from .run_journal import AgentRunJournal
from .session_registry import AgentRunSessionRegistry, get_session_registry
from .state_scope import build_state_context
from ...provider.tools.loaders import skill as skill_loader


ACTIVATED_SKILL_NAMES_STATE_KEY = 'host.activated_skills'


class AgentRunOrchestrator:
    """Coordinate one AgentRunner execution.

    The orchestrator keeps the run state machine readable and delegates
    transport, Query bridging, and persistence side effects to narrower
    collaborators.
    """

    ap: app.Application
    registry: AgentRunnerRegistry
    context_builder: AgentRunContextBuilder
    resource_builder: AgentResourceBuilder
    result_normalizer: AgentResultNormalizer
    binding_resolver: AgentBindingResolver
    query_bridge: QueryRunBridge
    invoker: AgentRunnerInvoker
    journal: AgentRunJournal
    _session_registry: AgentRunSessionRegistry

    def __init__(
        self,
        ap: app.Application,
        registry: AgentRunnerRegistry,
    ):
        self.ap = ap
        self.registry = registry
        self.context_builder = AgentRunContextBuilder(ap)
        self.resource_builder = AgentResourceBuilder(ap)
        self.result_normalizer = AgentResultNormalizer(ap)
        self.binding_resolver = AgentBindingResolver()
        self.query_bridge = QueryRunBridge(self.binding_resolver)
        self.invoker = AgentRunnerInvoker(ap)
        self.journal = AgentRunJournal(ap)
        self._session_registry = get_session_registry()

    async def run(
        self,
        event: AgentEventEnvelope,
        binding: AgentBinding,
        bound_plugins: list[str] | None = None,
        adapter_context: dict[str, typing.Any] | None = None,
    ) -> typing.AsyncGenerator[provider_message.Message | provider_message.MessageChunk, None]:
        """Run an AgentRunner from an event-first envelope."""
        runner_id = binding.runner_id
        descriptor = await self.registry.get(runner_id, bound_plugins)

        resources = await self.resource_builder.build_resources_from_binding(
            event=event,
            binding=binding,
            descriptor=descriptor,
        )

        context = await self.context_builder.build_context_from_event(
            event=event,
            binding=binding,
            descriptor=descriptor,
            resources=resources,
        )

        session_query_id = None
        if adapter_context:
            query = adapter_context.get('_query')
            if query is not None:
                skill_loader.restore_activated_skills_from_state(
                    self.ap,
                    query,
                    context.get('state', {}),
                )
            session_query_id = adapter_context.get('query_id')
            if query is not None or session_query_id is not None:
                context['context']['available_apis']['prompt_get'] = True
            if 'params' in adapter_context:
                context['adapter']['extra']['params'] = adapter_context['params']

        state_context = build_state_context(event, binding, descriptor)
        run_id = context['run_id']
        available_apis = context.get('context', {}).get('available_apis')
        run_authorization = {
            'runner_id': descriptor.id,
            'binding_id': binding.binding_id,
            'plugin_identity': descriptor.get_plugin_id(),
            'resources': resources,
            'available_apis': available_apis,
            'conversation_id': event.conversation_id,
            'bot_id': event.bot_id,
            'workspace_id': event.workspace_id,
            'thread_id': event.thread_id,
            'state_policy': {
                'enable_state': binding.state_policy.enable_state,
                'state_scopes': list(binding.state_policy.state_scopes),
            },
            'state_context': state_context,
        }

        seen_sequences: set[int] = set()
        last_sequence = 0
        assistant_transcript_written = False
        terminal_status: str | None = None
        terminal_reason: str | None = None
        terminal_usage: dict[str, typing.Any] | None = None

        try:
            await self.journal.create_run(
                event=event,
                binding=binding,
                descriptor=descriptor,
                context=context,
                authorization=run_authorization,
            )
            await self._session_registry.register(
                run_id=run_id,
                runner_id=descriptor.id,
                query_id=session_query_id,
                plugin_identity=descriptor.get_plugin_id(),
                resources=resources,
                available_apis=context.get('context', {}).get('available_apis'),
                conversation_id=event.conversation_id,
                bot_id=event.bot_id,
                workspace_id=event.workspace_id,
                thread_id=event.thread_id,
                state_policy={
                    'enable_state': binding.state_policy.enable_state,
                    'state_scopes': list(binding.state_policy.state_scopes),
                },
                state_context=state_context,
            )

            event_log_id = await self.journal.write_event_log(
                event=event,
                binding=binding,
                run_id=run_id,
                runner_id=descriptor.id,
            )
            if event.event_type == 'message.received' and event.conversation_id:
                await self.journal.write_user_transcript(
                    event=event,
                    event_log_id=event_log_id,
                )

            async for result_dict in self.invoker.invoke(descriptor, context):
                result_dict = dict(result_dict)
                sequence = result_dict.get('sequence')
                if sequence is not None:
                    try:
                        sequence_int = int(sequence)
                    except (TypeError, ValueError):
                        self.ap.logger.warning(f'Runner {descriptor.id} returned invalid result sequence: {sequence}')
                        sequence_int = last_sequence + 1
                        result_dict['sequence'] = sequence_int
                    else:
                        if sequence_int in seen_sequences:
                            self.ap.logger.warning(
                                f'Runner {descriptor.id} returned duplicate result sequence '
                                f'{sequence_int} for run {run_id}; dropping duplicate'
                            )
                            continue
                        if sequence_int <= 0:
                            self.ap.logger.warning(
                                f'Runner {descriptor.id} returned non-positive result sequence '
                                f'{sequence_int} for run {run_id}'
                            )
                            sequence_int = last_sequence + 1
                            result_dict['sequence'] = sequence_int
                        elif last_sequence and sequence_int != last_sequence + 1:
                            self.ap.logger.warning(
                                f'Runner {descriptor.id} result sequence gap or out-of-order '
                                f'for run {run_id}: previous={last_sequence}, current={sequence_int}'
                            )
                        seen_sequences.add(sequence_int)
                        last_sequence = max(last_sequence, sequence_int)
                else:
                    sequence_int = last_sequence + 1
                    result_dict['sequence'] = sequence_int
                    seen_sequences.add(sequence_int)
                    last_sequence = sequence_int

                result_type = result_dict.get('type')
                if result_type and not self.result_normalizer.validate_payload(
                    result_type,
                    result_dict.get('data', {}),
                    descriptor,
                ):
                    continue

                await self.journal.append_run_result(
                    result_dict=result_dict,
                    run_id=run_id,
                    sequence=sequence_int,
                )

                if result_type == 'state.updated':
                    await self.journal.handle_state_updated_event(
                        result_dict,
                        event,
                        binding,
                        descriptor,
                        run_id=run_id,
                    )
                    await self.result_normalizer.normalize(result_dict, descriptor)
                    continue

                if result_type == 'run.completed':
                    terminal_status = 'completed'
                    terminal_reason = (
                        result_dict.get('data', {}).get('finish_reason')
                        if isinstance(result_dict.get('data'), dict)
                        else None
                    )
                    usage = result_dict.get('usage')
                    if isinstance(usage, dict):
                        terminal_usage = usage
                elif result_type == 'run.failed':
                    terminal_status = 'failed'
                    data = result_dict.get('data') if isinstance(result_dict.get('data'), dict) else {}
                    terminal_reason = data.get('error') or data.get('code')
                    usage = result_dict.get('usage')
                    if isinstance(usage, dict):
                        terminal_usage = usage

                has_completed_message = result_type == 'message.completed' or (
                    result_type == 'run.completed'
                    and isinstance(result_dict.get('data'), dict)
                    and bool(result_dict['data'].get('message'))
                )
                if has_completed_message and event.conversation_id and not assistant_transcript_written:
                    await self.journal.write_assistant_transcript(
                        result_dict=result_dict,
                        event=event,
                        run_id=run_id,
                        runner_id=descriptor.id,
                    )
                    assistant_transcript_written = True

                result = await self.result_normalizer.normalize(result_dict, descriptor)
                if result is not None:
                    yield result

                run_snapshot = await self.journal.get_run(run_id)
                if run_snapshot and run_snapshot.get('cancel_requested_at') is not None:
                    terminal_status = 'cancelled'
                    terminal_reason = run_snapshot.get('status_reason') or 'cancel_requested'
                    break
            await self.journal.finalize_run(
                run_id=run_id,
                status=terminal_status or 'completed',
                status_reason=terminal_reason,
                usage=terminal_usage,
            )
        except Exception as exc:
            failed_usage = terminal_usage
            await self.journal.finalize_run(
                run_id=run_id,
                status='timeout' if self._is_deadline_exhausted(context) else 'failed',
                status_reason=str(exc),
                usage=failed_usage,
            )
            raise
        finally:
            session = await self._session_registry.unregister(run_id)
            pending_steering = session.get('steering_queue', []) if session else []
            if pending_steering:
                try:
                    await self.journal.write_steering_dropped_audits(
                        pending_steering,
                        run_id,
                        descriptor.id,
                    )
                except Exception as exc:
                    self.ap.logger.warning(
                        f'Failed to write dropped steering audit for run {run_id}: {exc}',
                        exc_info=True,
                    )

    async def run_from_query(
        self,
        query: pipeline_query.Query,
    ) -> typing.AsyncGenerator[provider_message.Message | provider_message.MessageChunk, None]:
        """Run an AgentRunner from the current Pipeline Query entry point."""
        plan = self.query_bridge.build_plan(query)
        adapter_context = dict(plan.adapter_context)
        adapter_context['_query'] = query

        # Materialize inbound attachments into sandbox before running
        await self._materialize_inbound_attachments(query, plan.event)

        async for result in self.run(
            plan.event,
            plan.binding,
            bound_plugins=plan.bound_plugins,
            adapter_context=adapter_context,
        ):
            yield result

    async def _materialize_inbound_attachments(
        self,
        query: pipeline_query.Query,
        event: AgentEventEnvelope,
    ) -> None:
        """Persist inbound attachments into the sandbox and update event.input.attachments.

        No-op when the box service is unavailable or there are no attachments.
        On success, updates each attachment in event.input.attachments with the
        sandbox path so runners can tell the model where to find the files.
        """
        box_service = getattr(self.ap, 'box_service', None)
        if box_service is None or not getattr(box_service, 'available', False):
            return

        try:
            materialized = await box_service.materialize_inbound_attachments(query)
        except Exception as e:
            # Never break the chat turn over attachment IO
            self.ap.logger.warning(f'Inbound attachment materialization failed: {e}')
            return

        if not materialized:
            return

        # Build a lookup by name for matching
        materialized_by_name = {att.get('name'): att for att in materialized if att.get('name')}

        # Update event.input.attachments with sandbox paths
        if event.input and event.input.attachments:
            for attachment in event.input.attachments:
                name = attachment.name
                if name and name in materialized_by_name:
                    mat = materialized_by_name[name]
                    # Update the attachment with sandbox path
                    attachment.path = mat.get('path')
                    attachment.size = mat.get('size') or attachment.size
                    attachment.mime_type = attachment.mime_type or mat.get('mime_type')

        # Store materialized descriptors in query variables for downstream use
        query.variables['_sandbox_inbound_attachments'] = materialized

    def resolve_runner_id_for_telemetry(self, query: pipeline_query.Query) -> str | None:
        """Resolve runner ID for telemetry/logging without full execution."""
        return self.query_bridge.resolve_runner_id_for_telemetry(query)

    async def try_claim_steering_from_query(
        self,
        query: pipeline_query.Query,
    ) -> bool:
        """Claim a query as steering input for an active run when possible."""
        plan = self.query_bridge.build_plan(query)
        event = plan.event
        binding = plan.binding

        if event.event_type != 'message.received' or not event.conversation_id:
            return False

        descriptor = await self.registry.get(binding.runner_id, plan.bound_plugins)
        if not descriptor.supports_steering():
            return False

        target_run_id = await self._session_registry.find_steering_target(
            conversation_id=event.conversation_id,
            runner_id=descriptor.id,
            bot_id=event.bot_id,
            workspace_id=event.workspace_id,
            thread_id=event.thread_id,
        )
        if target_run_id is None:
            return False

        steering_item = self._build_steering_item(event, target_run_id, descriptor.id)
        if not await self._session_registry.enqueue_steering(target_run_id, steering_item):
            return False

        try:
            event_log_id = await self.journal.write_event_log(
                event=event,
                binding=binding,
                run_id=target_run_id,
                runner_id=descriptor.id,
                metadata={
                    'steering': {
                        'status': 'queued',
                        'trigger_behavior': 'absorbed_into_active_run',
                        'claimed_by_run_id': target_run_id,
                        'claimed_runner_id': descriptor.id,
                        'claimed_at': steering_item.get('claimed_at'),
                    },
                },
            )
            await self.journal.write_user_transcript(event, event_log_id)
        except Exception as exc:
            self.ap.logger.warning(
                f'Failed to persist steering event {event.event_id} for run {target_run_id}: {exc}',
                exc_info=True,
            )

        self.ap.logger.info(f'Claimed event {event.event_id} as steering input for run {target_run_id}')
        return True

    def _build_steering_item(
        self,
        event: AgentEventEnvelope,
        run_id: str,
        runner_id: str,
    ) -> dict[str, typing.Any]:
        """Build the run-scoped steering item returned by the Host pull API."""
        return {
            'claimed_run_id': run_id,
            'runner_id': runner_id,
            'claimed_at': int(time.time()),
            'event': {
                'event_id': event.event_id,
                'event_type': event.event_type,
                'event_time': event.event_time,
                'source': event.source,
                'source_event_type': event.source_event_type,
                'raw_ref': event.raw_ref.model_dump(mode='json') if event.raw_ref else None,
                'data': event.data,
            },
            'conversation': {
                'conversation_id': event.conversation_id,
                'thread_id': event.thread_id,
                'bot_id': event.bot_id,
                'workspace_id': event.workspace_id,
            },
            'actor': event.actor.model_dump(mode='json') if event.actor else None,
            'subject': event.subject.model_dump(mode='json') if event.subject else None,
            'input': {
                'text': event.input.text if event.input else None,
                'contents': [
                    c.model_dump(mode='json') if hasattr(c, 'model_dump') else c
                    for c in (event.input.contents if event.input else [])
                ],
                'attachments': [
                    a.model_dump(mode='json') if hasattr(a, 'model_dump') else a
                    for a in (event.input.attachments if event.input else [])
                ],
            },
        }

    async def _invoke_runner(
        self,
        descriptor: AgentRunnerDescriptor,
        context: AgentRunContextPayload,
    ) -> typing.AsyncGenerator[dict[str, typing.Any], None]:
        """Compatibility delegate for older tests and internal callers."""
        async for result in self.invoker.invoke(descriptor, context):
            yield result

    async def _next_with_deadline(
        self,
        gen: typing.AsyncGenerator[dict[str, typing.Any], None],
        descriptor: AgentRunnerDescriptor,
        context: AgentRunContextPayload,
    ) -> dict[str, typing.Any]:
        return await self.invoker._next_with_deadline(gen, descriptor, context)

    def _remaining_deadline_seconds(
        self,
        context: AgentRunContextPayload,
    ) -> float | None:
        return self.invoker._remaining_deadline_seconds(context)

    def _is_deadline_exhausted(self, context: AgentRunContextPayload) -> bool:
        return self.invoker._is_deadline_exhausted(context)

    async def _close_generator(
        self,
        gen: typing.AsyncGenerator[dict[str, typing.Any], None],
        descriptor: AgentRunnerDescriptor,
    ) -> None:
        await self.invoker._close_generator(gen, descriptor)

    async def _handle_state_updated_event(
        self,
        result_dict: dict[str, typing.Any],
        event: AgentEventEnvelope,
        binding: AgentBinding,
        descriptor: AgentRunnerDescriptor,
    ) -> None:
        await self.journal.handle_state_updated_event(result_dict, event, binding, descriptor)

    async def _write_event_log(
        self,
        event: AgentEventEnvelope,
        binding: AgentBinding,
        run_id: str,
        runner_id: str,
    ) -> str:
        return await self.journal.write_event_log(event, binding, run_id, runner_id)

    async def _write_user_transcript(
        self,
        event: AgentEventEnvelope,
        event_log_id: str,
    ) -> None:
        await self.journal.write_user_transcript(event, event_log_id)

    async def _write_assistant_transcript(
        self,
        result_dict: dict[str, typing.Any],
        event: AgentEventEnvelope,
        run_id: str,
        runner_id: str,
    ) -> None:
        await self.journal.write_assistant_transcript(
            result_dict=result_dict,
            event=event,
            run_id=run_id,
            runner_id=runner_id,
        )