okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

dawnkernel: a job_stop frame lets /jobs stop remote background jobs57a82be2026-10-11

"""Python plugin boundary; imported without starting a kernel."""

from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Callable, Sequence
import importlib
import inspect
import keyword
from dataclasses import dataclass
from typing import Protocol, cast

from albedo_protocol import (
    Execute as Execute,
    HostSuccess as HostSuccess,
    HostFailure as HostFailure,
    HostReply as HostReply,
    Reply as Reply,
    Interrupt as Interrupt,
    Shutdown as Shutdown,
    Invoke as Invoke,
    Introspect as Introspect,
    Release as Release,
    State as State,
    JobStop as JobStop,
    Incoming as Incoming,
    SavedCell as SavedCell,
)

RETAIN = 1024 * 1024  # output bytes kept per channel; a late pipe reads them
OUTPUT_PREVIEW = 64 * 1024

Host = Callable[[str, dict[str, object]], Awaitable[object]]
Send = Callable[[dict[str, object]], None]
# Shutdown callbacks run while the loop can still run code; returning an
# awaitable lets a plugin clean up asynchronously.
Cleanup = Callable[[], object]


def _ready(value):
    """`await` on a value that is already here: returns it without suspending."""
    yield from ()
    return value


class Text(str):
    """A result that is ready now. Awaiting it is harmless, so a synchronous
    call reads the same with or without `await`. `.content` and `.text` are the
    text itself, for code written against a response object."""

    def __await__(self):
        return _ready(self)

    @property
    def content(self) -> str:
        return str(self)

    @property
    def text(self) -> str:
        return str(self)


class ReadyList(list):
    """A list result that is ready now; like Text, it may be awaited."""

    def __await__(self):
        return _ready(self)


class Record(dict):
    """A result record: `item["id"]` and `item.id` both work, it prints and
    serializes as the plain dict it is, and it may be awaited like Text."""

    __slots__ = ()

    def __getattr__(self, name: str) -> object:
        try:
            return self[name]
        except KeyError:
            raise AttributeError(
                f"{type(self).__name__} has no field {name!r}; its fields are {', '.join(self) or 'none'}"
            ) from None

    def __await__(self):
        return _ready(self)


MAX_JOB_SECONDS = 86400  # a day: long builds fit, a deadline in ms does not


def check_timeout(timeout: float) -> None:
    """Raise unless a job timeout is within 0 < timeout <= MAX_JOB_SECONDS."""
    if not 0 < timeout <= MAX_JOB_SECONDS:
        raise ValueError(
            f"timeout is in seconds: 0 < timeout <= {MAX_JOB_SECONDS} required,"
            f" got {timeout:g}"
            + (" (milliseconds?)" if timeout > MAX_JOB_SECONDS else "")
        )


def excerpt(
    text: str, chars: int, lines: int | None, *, end: bool, clipped: bool = False
) -> Text:
    """A requested preview, never silently reduced to the retained window."""
    recovery = "use job.read() for complete text or job.save(path) for exact bytes"
    if lines is None:
        if not 0 <= chars <= OUTPUT_PREVIEW:
            raise ValueError(f"0 <= n <= {OUTPUT_PREVIEW} for previews; {recovery}")
        if clipped and chars > len(text):
            raise ValueError(
                f"the retained preview cannot supply {chars} characters; {recovery}"
            )
        return Text((text[-chars:] if end else text[:chars]) if chars else "")
    if lines < 0:
        raise ValueError("lines must be nonnegative")
    kept = text.splitlines(keepends=True)
    # A window cut through a line cannot count that fragment as a whole line.
    if clipped and kept:
        if end:
            kept = kept[1:]
        elif not kept[-1].endswith(("\n", "\r")):
            kept = kept[:-1]
    if clipped and lines > len(kept):
        raise ValueError(
            f"the retained preview cannot supply {lines} lines; {recovery}"
        )
    return Text("".join(kept[-lines:] if end else kept[:lines]) if lines else "")


class OutputCapture(Protocol):
    data: bytearray  # the retained start of the output, as written
    seen: int  # bytes written, retained or not

    @property
    def retained(self) -> int: ...

    def read_bytes(self, offset: int, limit: int) -> bytes: ...

    def end_spill(self) -> None: ...

    def write(self, text: str) -> None: ...

    def write_bytes(self, data: bytes | bytearray | memoryview) -> None: ...

    def tail(self, limit: int = 65536) -> bytes: ...

    def read(self, offset: int = 0, limit: int = 4000) -> str: ...


@dataclass(frozen=True)
class PythonApi:
    loop: asyncio.AbstractEventLoop
    host: Host
    HostError: type[Exception]
    # An output channel for one piece of background work, retained under that id
    # so the session can read it back through `output`.
    capture: Callable[[str], OutputCapture]
    preview: int
    send: Send
    on_shutdown: Callable[[Cleanup], None]
    background_handle: Callable[[type[object]], None]
    version: int = 1
    # The module list this kernel was started with, so a plugin can learn the
    # session's tool selection (remote boots the same set on another machine).
    modules: Sequence[str] = ()
    # Registers one callback for a retained output channel, fired when the
    # model reads it, so background work can withdraw a completion wake the
    # read already satisfied. Absent when the host pre-dates the hook.
    watch_output: Callable[[str, Callable[[], None]], None] | None = None
    # Drops one retained output channel, so work a plugin runs for itself (a
    # search job, say) never crowds the model's own output out of retention.
    forget_output: Callable[[str], None] | None = None
    # Attaches PNG, JPEG, or WebP bytes to the running cell's result, so the
    # model sees the image; returns a short description, raises ValueError
    # past the per-cell limits. Absent when the host pre-dates images.
    attach_image: Callable[[bytes], str] | None = None
    # A host call that blocks until its answer, for synchronous APIs over
    # routes that answer at once. Absent when the host pre-dates it.
    host_now: Callable[[str, dict[str, object]], object] | None = None
    # Registers one handler to stop a background job (e.g. remote or fallback).
    # Handler returns {"ok": True} or {"ok": False, "error": "..."}, or None if not found.
    on_stop_job: (
        Callable[[Callable[[str], Awaitable[dict[str, object] | None]]], None] | None
    ) = None


class PythonPlugin(Protocol):
    def setup(
        self, api: PythonApi
    ) -> "dict[str, object] | Awaitable[dict[str, object]]": ...


async def load_plugins(
    names: Sequence[str], api: PythonApi, namespace: dict[str, object]
) -> None:
    """Load explicitly trusted modules before the workspace enters sys.path.

    setup(api) returns public REPL bindings, or an awaitable that resolves them,
    so a plugin may call api.host while the kernel boots. Plugins use api.host
    for host RPC, api.on_shutdown for cleanup, and api.background_handle for
    nonblocking jobs. This is composition, not isolation: installed plugins
    execute trusted code.
    """
    modules: list[str] = []
    for name in names:
        if not isinstance(name, str) or not all(
            part.isidentifier() and not keyword.iskeyword(part)
            for part in name.split(".")
        ):
            raise ValueError(f"invalid Python plugin module: {name!r}")
        module = name if "." in name else "albedo_plugins." + name
        if module in modules:
            raise ValueError(f"duplicate Python plugin module: {module}")
        modules.append(module)
    for module in modules:
        try:
            plugin = cast(PythonPlugin, cast(object, importlib.import_module(module)))
            exports = plugin.setup(api)
            if inspect.isawaitable(exports):
                exports = await cast(Awaitable[dict[str, object]], exports)
            if not isinstance(exports, dict) or not all(
                isinstance(name, str)
                and name.isidentifier()
                and not name.startswith("_")
                and not keyword.iskeyword(name)
                for name in exports
            ):
                raise ValueError(
                    "setup(api) must return a dict of public Python bindings"
                )
            collisions = namespace.keys() & exports.keys()
            if collisions:
                raise ValueError(
                    f"duplicate or reserved bindings: {', '.join(sorted(collisions))}"
                )
            namespace.update(exports)
        except Exception as error:
            raise RuntimeError(f"Python plugin {module}: {error}") from error