velophex 0.2.0__py3-none-any.whl

velophex/worker/session.py

@@ -0,0 +1,519 @@
+"""One executor session: hello, optional prepare, start, telemetry, one terminal message.
+
+Threading, because it is the part a reviewer has to agree with: the protocol reader blocks on
+stdin, and a handler may block for hours, so the two cannot share a thread. The session reads
+``hello`` and ``start`` on the main thread, then hands the handler to a worker thread and
+keeps reading, so a ``stop`` that arrives mid-run is seen immediately. The frame writer is
+thread-safe, which is what lets log lines and the terminal message come from either thread.
+
+A run ends with exactly one of ``result``, ``canceled`` or ``fault``, after which the process
+exits. Anything the handler raises that is not a :class:`velophex.BusinessError` is an
+``Application`` fault; a package or handler that will not load is ``Application`` with
+``INVALID_PACKAGE``, matching the workflow executor so the two runtimes retry alike.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import logging
+import threading
+import traceback
+from dataclasses import replace
+from typing import Any, TextIO
+
+from velophex.contracts.executor import (
+    CURRENT_VERSION,
+    FAULT_ENVELOPE_VERSION,
+    MINIMUM_VERSION,
+    CanceledMessage,
+    ExecutionContextInfo,
+    ExecutorCapabilities,
+    ExecutorMessage,
+    ExecutorProtocolException,
+    FaultInfo,
+    FaultMessage,
+    FaultRetryScopes,
+    HelloAckMessage,
+    HelloMessage,
+    PreparedMessage,
+    PrepareMessage,
+    RequestStopMessage,
+    ResultMessage,
+    StartMessage,
+    StopMessage,
+    SubResultMessage,
+    select_version,
+)
+from velophex.contracts.runtimes import (
+    ExecutionPhases,
+    FaultCodes,
+    FaultKinds,
+    LogLevels,
+    RuntimeKind,
+    SignalReasons,
+)
+from velophex.entrypoints import EntryPointDefinition, EntryPointInputError, normalize_outputs
+from velophex.errors import BusinessError, VelophexError
+from velophex.runtime import _session as ambient
+from velophex.worker.framing import FrameReader
+from velophex.worker.loader import LoadedEntryPoint, PackageLoadError, load_entry_point
+from velophex.worker.outbox import ExecutorOutbox, WorkerLogHandler
+
+__all__ = [
+    "EXIT_OK",
+    "EXIT_PROTOCOL_VIOLATION",
+    "WORKER_NAME",
+    "WorkerSession",
+]
+
+#: The name the worker reports in ``hello-ack``.
+WORKER_NAME = "velophex.worker"
+
+EXIT_OK = 0
+EXIT_PROTOCOL_VIOLATION = 2
+
+#: How long to wait past the grace period before giving up on a cooperative stop. The parent
+#: kills the process tree at the grace deadline, so this only matters when the parent is slow.
+_STOP_SLACK_SECONDS = 5.0
+
+
+class WorkerSession:
+    """Serves one executor session over an authenticated frame pair."""
+
+    def __init__(
+        self,
+        reader: FrameReader,
+        outbox: ExecutorOutbox,
+        *,
+        runtime_version: str,
+        worker_version: str,
+        diagnostics: TextIO | None = None,
+    ) -> None:
+        self._reader = reader
+        self._outbox = outbox
+        self._runtime_version = runtime_version
+        self._worker_version = worker_version
+        self._diagnostics = diagnostics
+        self._session: ambient.RunSession | None = None
+        self._loop: asyncio.AbstractEventLoop | None = None
+        self._handler_task: asyncio.Task[Any] | None = None
+        self._parent_capabilities: frozenset[str] = frozenset()
+
+    def run(self) -> int:
+        """Serve the session and return the process exit code."""
+        try:
+            if not self._handshake():
+                return EXIT_PROTOCOL_VIOLATION
+
+            while True:
+                message = self._reader.read()
+
+                if message is None:
+                    return EXIT_OK
+
+                if isinstance(message, PrepareMessage):
+                    self._outbox.send(self._prepare(message))
+                elif isinstance(message, StartMessage):
+                    self._send_terminal(self._start(message))
+                    return EXIT_OK
+                elif isinstance(message, StopMessage):
+                    # Nothing runs yet: the session ends as canceled.
+                    self._outbox.send(CanceledMessage(reason=message.reason or SignalReasons.USER))
+                    return EXIT_OK
+                else:
+                    self._note(f"Ignoring unexpected '{message.type_name}' before start.")
+        except ExecutorProtocolException as exception:
+            self._note(f"Protocol violation ({exception.error}): {exception}")
+            return EXIT_PROTOCOL_VIOLATION
+
+    # ---------------------------------------------------------------- handshake
+
+    def _handshake(self) -> bool:
+        message = self._reader.read()
+
+        if not isinstance(message, HelloMessage):
+            self._note("The first message was not 'hello'.")
+            return False
+
+        version = select_version(message.minimum_version, message.maximum_version)
+        self._parent_capabilities = frozenset(message.capabilities)
+
+        if version is None:
+            self._outbox.send(
+                FaultMessage(
+                    FaultInfo(
+                        kind=FaultKinds.INFRASTRUCTURE,
+                        code=FaultCodes.EXECUTOR_PROTOCOL_VIOLATION,
+                        message=(
+                            f"The parent speaks executor protocol {message.minimum_version}-{message.maximum_version}; "
+                            f"this worker speaks {MINIMUM_VERSION}-{CURRENT_VERSION}."
+                        ),
+                    )
+                )
+            )
+            return False
+
+        self._outbox.send(
+            HelloAckMessage(
+                version=version,
+                runtime=RuntimeKind.PYTHON,
+                runtime_version=self._runtime_version,
+                executor_version=self._worker_version,
+                # Only should-stop: the Python environment is prepared by the Runner's uv
+                # adapter before this process starts (ADR-P008), there is no Python debugger
+                # over this protocol yet, and per-element progress has no meaning for a
+                # handler call. request-stop raises velophex.should_stop().
+                capabilities=[ExecutorCapabilities.SHOULD_STOP],
+            )
+        )
+
+        return True
+
+    # ---------------------------------------------------------------- prepare
+
+    def _prepare(self, message: PrepareMessage) -> ExecutorMessage:
+        """Answer a ``prepare`` a conforming parent should not have sent.
+
+        The worker does not offer the ``prepare`` capability: by the time it runs, ``uv sync``
+        has already built the environment it is running in. Answering ``prepared`` rather than
+        faulting keeps a parent that asks anyway from stalling.
+        """
+        self._note("Ignoring 'prepare': the Python environment is prepared before the worker starts.")
+        return PreparedMessage(environment_root=message.environment_root, fingerprint=None)
+
+    # ---------------------------------------------------------------- start
+
+    def _start(self, message: StartMessage) -> ExecutorMessage:
+        try:
+            loaded = load_entry_point(message.package_root, message.entry_point_id)
+        except PackageLoadError as error:
+            self._outbox.log(LogLevels.ERROR, str(error), ExecutionPhases.ENVIRONMENT_PREPARE)
+            return FaultMessage(
+                FaultInfo(
+                    kind=FaultKinds.APPLICATION,
+                    code=FaultCodes.INVALID_PACKAGE,
+                    message=str(error),
+                    retryable=False,
+                )
+            )
+
+        try:
+            arguments = loaded.definition.bind(message.inputs)
+        except EntryPointInputError as error:
+            return FaultMessage(
+                FaultInfo(
+                    kind=FaultKinds.APPLICATION,
+                    code=FaultCodes.INVALID_INPUT,
+                    message=str(error),
+                    detail=" ".join(error.pointers),
+                    retryable=False,
+                )
+            )
+
+        return self._invoke(loaded, arguments, message.context)
+
+    def _invoke(
+        self,
+        loaded: LoadedEntryPoint,
+        arguments: dict[str, Any],
+        context: ExecutionContextInfo,
+    ) -> ExecutorMessage:
+        session = ambient.RunSession.create(context, loaded.definition, loaded.project.name)
+        ambient.set_session(session)
+        self._session = session
+
+        if ExecutorCapabilities.SUB_INVOKE in self._parent_capabilities:
+            session.enable_sub_invoke(self._outbox.send)
+
+        handler_log = WorkerLogHandler(self._outbox)
+        root = logging.getLogger()
+        previous_level = root.level
+        root.addHandler(handler_log)
+
+        if previous_level > logging.INFO or previous_level == logging.NOTSET:
+            root.setLevel(logging.INFO)
+
+        outcome: list[ExecutorMessage] = []
+        finished = threading.Event()
+
+        def body() -> None:
+            try:
+                outcome.append(self._call(loaded.definition, arguments))
+            except BaseException as error:
+                outcome.append(self._fault_from(error, session))
+            finally:
+                finished.set()
+
+        worker = threading.Thread(target=body, name="velophex-handler", daemon=True)
+
+        self._outbox.log(LogLevels.INFORMATION, f"Started '{loaded.entry_point.id}'.", ExecutionPhases.EXECUTION)
+        worker.start()
+
+        try:
+            self._listen_in_background(finished)
+            finished.wait()
+        finally:
+            root.removeHandler(handler_log)
+            root.setLevel(previous_level)
+            ambient.clear_session()
+            self._session = None
+
+        if outcome:
+            terminal = outcome[0]
+        else:
+            terminal = FaultMessage(
+                FaultInfo(
+                    kind=FaultKinds.APPLICATION,
+                    code=FaultCodes.EXECUTOR_CRASHED,
+                    message="The handler thread ended without an outcome.",
+                )
+            )
+
+        # A handler that returned normally after a stop was asked for still ended because of
+        # the stop; the job is Canceled, not Succeeded.
+        if session.stop_requested and isinstance(terminal, ResultMessage):
+            return CanceledMessage(reason=session.stop_reason or SignalReasons.USER)
+
+        return terminal
+
+    def _call(self, definition: EntryPointDefinition, arguments: dict[str, Any]) -> ExecutorMessage:
+        if definition.is_async:
+            outputs = self._call_async(definition, arguments)
+        else:
+            outputs = definition.function(**arguments)
+
+            if inspect.isawaitable(outputs):
+                # A sync function that returns a coroutine (a wrapped handler) still has to
+                # be awaited, or the automation silently does nothing.
+                outputs = asyncio.run(outputs)
+
+        return ResultMessage(outputs=normalize_outputs(outputs))
+
+    def _call_async(self, definition: EntryPointDefinition, arguments: dict[str, Any]) -> Any:
+        async def main() -> Any:
+            self._loop = asyncio.get_running_loop()
+            self._handler_task = asyncio.ensure_future(definition.function(**arguments))
+
+            try:
+                return await self._handler_task
+            finally:
+                self._loop = None
+                self._handler_task = None
+
+        return asyncio.run(main())
+
+    # ---------------------------------------------------------------- control channel
+
+    def _listen_in_background(self, finished: threading.Event) -> None:
+        """Start reading control messages while the handler runs.
+
+        The reader blocks on stdin, so it needs a thread of its own; it is a daemon because
+        the process exits right after the terminal message, with the read still outstanding.
+        """
+        reader = threading.Thread(target=self._listen_loop, args=(finished,), name="velophex-control", daemon=True)
+        reader.start()
+
+    def _listen_loop(self, finished: threading.Event) -> None:
+        while not finished.is_set():
+            try:
+                message = self._reader.read()
+            except ExecutorProtocolException as exception:
+                self._note(f"Protocol violation ({exception.error}): {exception}")
+                message = None
+
+            if message is None:
+                # The parent closed stdin: it is gone, so the run is stopped and no step will
+                # ever be answered.
+                self._request_stop(SignalReasons.USER, grace_milliseconds=0)
+
+                if self._session is not None:
+                    self._session.fail_pending("The parent closed the executor channel before the step answered.")
+
+                return
+
+            if isinstance(message, SubResultMessage):
+                session = self._session
+
+                if session is not None and not session.deliver(message):
+                    # An unknown or repeated id is a protocol violation (EXECUTOR-PROTOCOL.md 4.1):
+                    # every waiting call fails rather than trusting the channel any further.
+                    self._note(f"Protocol violation: sub-result for unknown request {message.request_id!r}.")
+                    session.fail_pending(f"The parent answered an unknown step request {message.request_id!r}.")
+            elif isinstance(message, StopMessage):
+                self._outbox.log(LogLevels.INFORMATION, f"Stop requested ({message.reason}).")
+                self._request_stop(message.reason, message.grace_period_milliseconds)
+            elif isinstance(message, RequestStopMessage):
+                # The soft stop: only the flag velophex.should_stop() reads; nothing is canceled.
+                session = self._session
+
+                if session is not None and not session.should_stop:
+                    self._outbox.log(LogLevels.INFORMATION, "The run was asked to stop when its current work is done.")
+                    session.request_soft_stop()
+            else:
+                self._note(f"Ignoring unexpected '{message.type_name}' during a run.")
+
+    def _request_stop(self, reason: str, grace_milliseconds: int) -> None:
+        session = self._session
+
+        if session is None or session.stop_requested:
+            return
+
+        session.request_stop(reason or SignalReasons.USER)
+
+        grace = max(grace_milliseconds, 0) / 1000.0
+        threading.Timer(grace + _STOP_SLACK_SECONDS, self._grace_expired, args=(grace,)).start()
+
+    def _grace_expired(self, grace: float) -> None:
+        """Cancel what can be cancelled once the grace period has passed.
+
+        A synchronous handler cannot be interrupted from outside without corrupting whatever
+        it holds open, so the platform's answer is the parent's kill. An ``async`` handler is
+        different: cancelling its task unwinds through ``finally`` blocks, so the worker does
+        that and lets the handler decide how to end.
+        """
+        session = self._session
+
+        if session is None or not session.stop_requested:
+            return
+
+        loop, task = self._loop, self._handler_task
+
+        if loop is not None and task is not None and not task.done():
+            self._outbox.log(
+                LogLevels.WARNING,
+                f"The grace period of {grace:.0f}s expired; cancelling the handler task.",
+            )
+            loop.call_soon_threadsafe(task.cancel)
+        else:
+            self._outbox.log(
+                LogLevels.WARNING,
+                f"The grace period of {grace:.0f}s expired and the handler is still running; "
+                "the parent will end this process.",
+            )
+
+    # ---------------------------------------------------------------- faults
+
+    def _fault_from(self, error: BaseException, session: ambient.RunSession) -> ExecutorMessage:
+        if isinstance(error, BusinessError):
+            return FaultMessage(
+                FaultInfo(
+                    kind=FaultKinds.BUSINESS,
+                    code=error.code,
+                    message=error.message,
+                    detail=error.detail,
+                    retryable=False,
+                )
+            )
+
+        if isinstance(error, (asyncio.CancelledError, KeyboardInterrupt)) or (
+            isinstance(error, SystemExit) and session.stop_requested
+        ):
+            return CanceledMessage(reason=session.stop_reason or SignalReasons.USER)
+
+        self._note_exception(error)
+
+        # An error that knows its code (a Runtime API refusal such as ai_budget_exhausted, an
+        # agent guardrail) keeps it, with the envelope: it does not collapse to UNHANDLED_EXCEPTION.
+        fault = error.to_fault() if isinstance(error, VelophexError) else None
+
+        if fault is None:
+            fault = FaultInfo(
+                kind=FaultKinds.APPLICATION,
+                code=FaultCodes.UNHANDLED_EXCEPTION,
+                message=f"{type(error).__name__}: {error}",
+                # The type and message only: the traceback goes to stderr, which the parent
+                # captures and redacts.
+                detail=type(error).__module__ + "." + type(error).__qualname__,
+            )
+
+        if getattr(error, "effects_committed", False):
+            # The run changed something before it failed (an agent's tool call): running it
+            # again would repeat that, so it is never retried automatically.
+            fault = replace(
+                fault,
+                retryable=False,
+                version=FAULT_ENVELOPE_VERSION,
+                retry_scope=FaultRetryScopes.NONE,
+                effects_committed=True,
+            )
+
+        return FaultMessage(fault)
+
+    # ---------------------------------------------------------------- terminal delivery
+
+    def _send_terminal(self, message: ExecutorMessage) -> None:
+        """Send the one terminal message, or a fault that says why it could not be sent.
+
+        The handler's return value is only turned into JSON here, after the handler has
+        finished. Without this check a ``datetime`` in the outputs, a NaN, or outputs over the
+        frame limit raised out of the session with no terminal frame at all - which the parent
+        can only report as a retryable ``EXECUTOR_CRASHED``, so a run that will fail the same
+        way every time was retried. The replacement fault is ``Application`` and not
+        retryable, with a stable code, because re-running the same code on the same inputs
+        returns the same value.
+        """
+        try:
+            payload = self._outbox.encode(message)
+        except ExecutorProtocolException as error:
+            payload = self._outbox.encode(
+                _undeliverable(
+                    message,
+                    FaultCodes.RESULT_TOO_LARGE,
+                    f"The '{message.type_name}' message is larger than the "
+                    f"{self._outbox.maximum_payload_bytes}-byte frame limit ({error}).",
+                )
+            )
+        except (TypeError, ValueError) as error:
+            self._note_exception(error)
+            payload = self._outbox.encode(
+                _undeliverable(
+                    message,
+                    FaultCodes.RESULT_NOT_SERIALIZABLE,
+                    f"The '{message.type_name}' message cannot be written as JSON: {_bounded(str(error))}",
+                )
+            )
+
+        self._outbox.send_payload(payload)
+
+    # ---------------------------------------------------------------- diagnostics
+
+    def _note(self, message: str) -> None:
+        if self._diagnostics is not None:
+            print(message, file=self._diagnostics, flush=True)
+
+    def _note_exception(self, error: BaseException) -> None:
+        if self._diagnostics is not None:
+            traceback.print_exception(type(error), error, error.__traceback__, file=self._diagnostics)
+            self._diagnostics.flush()
+
+
+_MAXIMUM_REASON_CHARACTERS = 1024
+
+
+def _bounded(text: str) -> str:
+    return text if len(text) <= _MAXIMUM_REASON_CHARACTERS else text[:_MAXIMUM_REASON_CHARACTERS] + "..."
+
+
+def _undeliverable(message: ExecutorMessage, code: str, reason: str) -> FaultMessage:
+    """The fault sent in place of a terminal message that cannot be framed.
+
+    Built only from strings the worker controls (bounded), so it always encodes and always
+    fits. A handler's ``result`` becomes a plain output fault; any other terminal message
+    (a business fault whose detail is not a string, say) keeps its code in the detail so the
+    original outcome is not lost entirely.
+    """
+    detail = None
+
+    if isinstance(message, FaultMessage):
+        detail = _bounded(f"Original fault: {message.fault.kind} / {message.fault.code}")
+
+    return FaultMessage(
+        FaultInfo(
+            kind=FaultKinds.APPLICATION,
+            code=code,
+            message=_bounded(reason),
+            detail=detail,
+            retryable=False,
+        )
+    )