foldit-plugin-sdk 0.1.5

"""Unified plugin interface for the foldit plugin SDK.

Plugins implement :class:`PluginInterface` to expose ops + queries to the
foldit orchestrator via the plugin protocol (see ``proto/plugin.proto``).
Each plugin runs in its own worker process; the worker host dispatches lifecycle endpoints
(Init, UpdateAssembly, Drop), op dispatch (Invoke, StartStream,
PollStream, UpdateStream, CancelStream), and query dispatch (Query) to
methods on the plugin instance.

Ops vs queries
--------------

Two categories of plugin-exposed work, with different semantics:

- **Ops** mutate state. Take entity locks. Return assembly bytes
  (orchestrator copies locked-entity slices into canonical state).
  Examples: predict, design, wiggle, mutate. Implemented in
  :meth:`invoke` (single-shot) or :meth:`start_stream` (long-running).
- **Queries** read state. No locks. Return query-defined opaque bytes.
  Examples: score breakdown, rama colors, sequence_design candidates.
  Implemented in :meth:`query` (single-shot only).

Lifecycle a plugin sees:

1. Worker process spawns; imports the plugin module; instantiates the
   ``Plugin`` class with config.
2. Worker calls :meth:`init` with the canonical Assembly bytes; plugin
   returns a SessionId.
3. Worker calls :meth:`register` to retrieve the plugin's op + query
   catalog (``PluginRegistration``); orchestrator caches it.
4. Per user action, worker calls :meth:`invoke` (op, single-shot) or
   :meth:`start_stream` / :meth:`poll_stream` / :meth:`update_stream` /
   :meth:`cancel_stream` (op, long-running) or :meth:`query`
   (read).
5. On orchestrator-driven Assembly changes, worker calls
   :meth:`update_assembly`. Plugin treats incoming Assembly as latest
   authority.
6. On session teardown, worker calls :meth:`drop`.

Bound-type surface
-------------------

The dispatch context and poll outcome are native types compiled into the
SDK's pyo3 extension, not Python dataclasses:

- :meth:`invoke` / :meth:`query` / :meth:`start_stream` receive a
  :class:`~foldit_plugin_sdk.DispatchContext` (read-only; ``focused_entity_id``,
  ``selection``, ``designable``) plus the flattened native-Python params dict.
- :meth:`poll_stream` returns a :class:`~foldit_plugin_sdk.PollOutcome`, which
  the plugin builds via its static factories
  (``PollOutcome.pending`` / ``.checkpoint`` / ``.cancelled`` / ``.final_`` /
  ``.error``), optionally attaching a
  :class:`~foldit_plugin_sdk.ScoreReport`.

Streaming model
---------------

Streams are *polling-based*, not push. ``start_stream`` kicks off the
work (typically in a background thread) under the host-assigned
``request_id``. The plugin tracks per-request_id state in an internal
dict. ``poll_stream`` returns the latest snapshot from that dict.

Plugins MUST coalesce: between polls (and between successive
``UpdateStream`` calls) only the latest snapshot / latest params
survives. No frame queue, no client-side cursor. This is critical for
interactive-drag patterns where high-frequency mouse events outpace the
plugin's apply rate.

See ``poll_stream`` for the return-shape contract.
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any

# plugin_pb2 is generated by `just generate-proto`; lives in the sibling
# `proto/` subpackage.
from .proto import plugin_pb2  # type: ignore[import-not-found]

if TYPE_CHECKING:
    # The dispatch context and poll outcome are compiled into the SDK's
    # pyo3 extension and only resolve at runtime, post-maturin. Importing
    # them under TYPE_CHECKING keeps the annotations honest without
    # requiring the built extension at module-import time.
    from foldit_plugin_sdk import DispatchContext, PollOutcome


def make_param_value(value: Any):
    """Construct a :class:`plugin_pb2.ParamValue` from a native Python value.

    Supported: int, float, bool, str, and 3-tuples / lists for Vec3.
    Plugins building :class:`plugin_pb2.ParamSpec` defaults use this.
    """
    pv = plugin_pb2.ParamValue()
    if isinstance(value, bool):
        pv.bool_value = value
    elif isinstance(value, int):
        pv.int_value = value
    elif isinstance(value, float):
        pv.float_value = value
    elif isinstance(value, str):
        pv.string_value = value
    elif isinstance(value, (tuple, list)) and len(value) == 3:
        pv.vec3_value.x = float(value[0])
        pv.vec3_value.y = float(value[1])
        pv.vec3_value.z = float(value[2])
    else:
        raise TypeError(f"Unsupported ParamValue type: {type(value).__name__}")
    return pv


class PluginInterface(ABC):
    """Base class all foldit plugins implement.

    Required overrides:

    - :meth:`__init__` receives plugin-private config dict.
    - :meth:`register` returns the plugin's op catalog as a
      :class:`plugin_pb2.PluginRegistration`.
    - :meth:`init` starts a session given an Assembly.
    - :meth:`update_assembly` replaces the working Assembly.
    - :meth:`drop` tears down a session.

    Optional overrides (raise NotImplementedError by default; plugins
    declare which ops they support via :meth:`register`):

    - :meth:`invoke` single-shot ops.
    - :meth:`start_stream` / :meth:`poll_stream` / :meth:`update_stream`
      / :meth:`cancel_stream` streaming ops.
    """

    @abstractmethod
    def __init__(self, config: dict[str, Any]) -> None:
        """Initialize the plugin with its host-process config.

        Per ``PLUGIN_PROTOCOL.md``: no init params on the wire; config is
        plugin-private. The host process passes whatever the plugin's
        spawn descriptor carries.
        """
        ...

    # Lifecycle: required

    @abstractmethod
    def register(self) -> "plugin_pb2.PluginRegistration":
        """Return this plugin's :class:`plugin_pb2.PluginRegistration`.

        Called by the worker host after :meth:`init` succeeds; the
        orchestrator caches the resulting catalog and routes ops by
        ``PluginOp.id``.
        """
        ...

    @abstractmethod
    def init(self, assembly_bytes: bytes) -> int:
        """Start a session. Returns a SessionId (uint64), the plugin's choice;
        commonly ``1`` for single-session plugins.

        ``assembly_bytes`` is the canonical Assembly (assembly wire format) as
        of the moment of session start. Plugins decode this into whatever
        internal working state they need (loaded model + cached
        embeddings for ML, ``RISession`` for Rosetta, etc.).
        """
        ...

    # Payload-kind tags matching `FolditPluginAssemblyPayloadKind` on
    # the Rust ABI side. Used as the `payload_kind` arg to
    # :meth:`update_assembly`.
    PAYLOAD_KIND_FULL: int = 0
    PAYLOAD_KIND_DELTA: int = 1

    @abstractmethod
    def update_assembly(
        self,
        session: int,
        payload_kind: int,
        bytes: bytes,
        from_gen: int,
        to_gen: int,
    ) -> None:
        """Push an Assembly update.

        ``payload_kind`` is :attr:`PAYLOAD_KIND_FULL` (fresh assembly
        snapshot) or :attr:`PAYLOAD_KIND_DELTA` (delta edit list;
        decode via ``molex.deserialize_delta``). ``from_gen`` / ``to_gen``
        are the host's broadcast generation counters; a plugin whose
        local gen does not match ``from_gen`` must arm a ``STALE_GEN``
        error to return on its next dispatch so the host re-syncs.

        Plugins MUST treat the incoming Assembly (full, or the state
        after applying the delta) as latest authority and discard any
        internal state inconsistent with it.
        """
        ...

    @abstractmethod
    def drop(self, session: int) -> None:
        """Tear down a session. Free any resources keyed to this session.
        Idempotent: calling twice with the same SessionId is a no-op."""
        ...

    # Op dispatch: single-shot mutation

    def invoke(
        self,
        session: int,
        op: str,
        context: "DispatchContext",
        params: dict[str, Any],
    ) -> bytes:
        """Single-shot mutating op (kind=INVOKE).

        Ops mutate the plugin's working assembly. ``op`` matches a
        registered :class:`plugin_pb2.PluginOp` id with
        ``kind=OP_KIND_INVOKE``. ``context`` is the bound
        :class:`~foldit_plugin_sdk.DispatchContext` carrying the focus +
        selection captured by the orchestrator at trigger time. ``params``
        is the flattened native-Python dict of typed params.

        Returns the plugin's working assembly bytes post-op (assembly
        wire format). The orchestrator copies locked-entity slices from the
        returned assembly into canonical state.

        Ops with no state mutation should be queries, not invokes; see
        :meth:`query`.
        """
        raise NotImplementedError(f"Plugin does not implement Invoke (op={op!r})")

    # Query dispatch: read state without mutation

    def query(
        self,
        session: int,
        query: str,
        context: "DispatchContext",
        params: dict[str, Any],
    ) -> bytes:
        """Single-shot read query.

        Queries READ state; no entity mutation, no entity locks.
        ``query`` matches a registered :class:`plugin_pb2.PluginQuery` id.
        ``context`` and ``params`` mirror :meth:`invoke`.

        Returns query-defined opaque bytes. The plugin and the consuming
        Tier 2 panel agree on the encoding (e.g. UTF-8 JSON for
        sequence-design candidates, packed float arrays for rama colors).

        Concurrent queries on the same entity are safe; queries racing
        against ops read what's there at execution time.
        """
        raise NotImplementedError(f"Plugin does not implement Query (query={query!r})")

    # Generic dispatch: streaming

    def start_stream(
        self,
        session: int,
        op: str,
        context: "DispatchContext",
        params: dict[str, Any],
        request_id: int,
    ) -> None:
        """Begin a long-running op (kind=STREAM) under the host-assigned
        ``request_id`` (uint64). The plugin keys its per-stream state on
        that id; it does not choose its own.

        The plugin starts the actual work (typically in a background
        thread) and tracks per-``request_id`` state in an internal dict.
        Subsequent :meth:`poll_stream` calls read that state and return
        the latest snapshot.

        Plugins MUST coalesce: only the latest snapshot survives between
        polls. See module docstring for the streaming model.
        """
        raise NotImplementedError(f"Plugin does not implement StartStream (op={op!r})")

    def poll_stream(self, request_id: int) -> "PollOutcome":
        """Return the latest snapshot for a running stream as a bound
        :class:`~foldit_plugin_sdk.PollOutcome`, built via its static
        factories:

        - ``PollOutcome.pending(...)`` op still running. Caller should
          continue polling.
        - ``PollOutcome.checkpoint(...)`` accepted intermediate the host
          commits while the stream keeps running. Polling continues.
        - ``PollOutcome.cancelled(...)`` op stopped at host request,
          returning a usable working pose. No further polls.
        - ``PollOutcome.final_(...)`` op finished successfully. The
          orchestrator promotes the assembly into canonical state. No
          further polls.
        - ``PollOutcome.error(...)`` op failed. No further polls.
        """
        raise NotImplementedError("Plugin does not implement PollStream")

    def update_stream(self, request_id: int, params: dict[str, Any]) -> None:
        """Push new params to a running stream.

        Used for pull-target updates, rama-drag ticks, etc. ``params``
        carries only the fields the plugin's metadata declared as
        updatable; values for non-updatable params are silently ignored.
        """
        raise NotImplementedError("Plugin does not implement UpdateStream")

    def cancel_stream(self, request_id: int) -> None:
        """Stop a running stream. Idempotent: cancelling an already-finished
        stream is a no-op."""
        raise NotImplementedError("Plugin does not implement CancelStream")


def find_plugin_class(module):
    """Find a class implementing :class:`PluginInterface` in ``module``.

    Used by the worker host after dynamically importing the plugin module.
    The plugin is expected to expose exactly one ``PluginInterface``
    subclass (conventionally named ``Plugin``).

    Raises :class:`ValueError` if no subclass found.
    """
    for attr_name in dir(module):
        attr = getattr(module, attr_name)
        if (
            isinstance(attr, type)
            and issubclass(attr, PluginInterface)
            and attr is not PluginInterface
        ):
            return attr
    raise ValueError(
        f"No class implementing PluginInterface found in {module.__name__}"
    )