tmux-amq/src/tamq/broker.py

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass

from .registry import RegistryError, validate_targets
from .routing import RoutedMessage, parse_address_line, parse_command_line
from .store import LimitBlock, Store
from .terminal import format_delivery


@dataclass(frozen=True)
class BrokerIdentity:
    endpoint_id: str
    source_repo: str


class InputBroker:
    """Convert attributed PTY lines into durable outbound messages."""

    def __init__(
        self,
        store: Store,
        identity: BrokerIdentity,
        *,
        limits: tuple[int, int, int] = (8, 1024, 32768),
        notify: Callable[[str], None] | None = None,
    ):
        self.store = store
        self.identity = identity
        self.notify = notify
        self.store.configure_window(identity.endpoint_id, identity.source_repo, limits)

    def _notice(self, text: str) -> None:
        if self.notify is not None:
            self.notify(f"From:tamq: {text}")

    def _blocked(self, blocked: LimitBlock) -> None:
        self._notice(
            "Messaging blocked! Running linecount "
            f"{blocked.count} limited by {blocked.limit} lines of {blocked.dimension} "
            "in this terminal. Use 'Cmd: reset-limits' to unblock."
        )

    def _route(self, line: str, provenance: str) -> RoutedMessage | None:
        routed = parse_address_line(line)
        if routed is None:
            return None
        try:
            validate_targets([routed.target_repo])
        except RegistryError:
            return None
        result = self.store.admit_message(
            self.identity.source_repo,
            routed.target_repo,
            routed.body,
            endpoint=self.identity.endpoint_id,
            source_repo=self.identity.source_repo,
            provenance=provenance,
        )
        if isinstance(result, LimitBlock):
            self._blocked(result)
            return None
        return routed

    def inspect_operator_line(self, line: str) -> RoutedMessage | None:
        """Observe one submitted operator line; it is still forwarded unchanged."""
        if line.startswith("From:"):
            return None
        self.store.count_line(
            self.identity.endpoint_id, self.identity.source_repo, "input"
        )
        command = parse_command_line(line)
        if command is not None:
            self._command(command)
            return None
        return self._route(line, "operator_input")

    def inspect_worker_line(self, line: str) -> RoutedMessage | None:
        """Observe worker output; Cmd is intentionally inert on this path."""
        if line.startswith("From:"):
            return None
        self.store.count_line(
            self.identity.endpoint_id, self.identity.source_repo, "output"
        )
        return self._route(line, "worker_output")

    def inspect_line(self, line: str) -> RoutedMessage | None:
        """Compatibility entry point for operator-input observers."""
        return self.inspect_operator_line(line)

    def _command(self, command: str) -> None:
        if command == "reset-limits":
            self.store.reset_limits(
                self.identity.endpoint_id, self.identity.source_repo
            )
            self._notice("Limits reset for this terminal.")
            return
        if "=" not in command:
            self._notice(f"Unknown command: {command}")
            return
        name, value = (part.strip() for part in command.split("=", 1))
        if name == "mode":
            if value not in {"inbox", "output", "pushy", "trigger"}:
                self._notice(f"Invalid mode: {value}")
                return
            if not self.store.set_endpoint_mode(self.identity.endpoint_id, value):
                self._notice("Cannot change mode before endpoint registration.")
                return
            self._notice(f"Mode set to {value}.")
            return
        if name in {"maxmsg", "maxin", "maxout"}:
            try:
                parsed = int(value)
                self.store.set_window_limit(
                    self.identity.endpoint_id,
                    self.identity.source_repo,
                    name,
                    parsed,
                )
            except ValueError as exc:
                self._notice(str(exc))
                return
            self._notice(f"{name} set to {parsed} for this terminal.")
            return
        self._notice(f"Unknown command: {name}")

    def deliver_pending(self, control, *, window_for_repo):
        """Place pending messages through the authoritative control client."""
        delivered = 0
        for row in self.store.list(state="pending"):
            if row["endpoint_id"] not in (None, self.identity.endpoint_id):
                continue
            target = window_for_repo(row["target_repo"])
            lease_id = self.store.claim(row["message_id"], self.identity.endpoint_id)
            if lease_id is None:
                continue
            try:
                control.inject(
                    target,
                    format_delivery(
                        row["sender_repo"], row["body"], row["provenance"]
                    ),
                )
            except Exception:
                continue
            self.store.release(row["message_id"], lease_id, "injected")
            delivered += 1
        return delivered