multi-agent-mux/.agents/skills/lib_py/delegation/service.py

"""DelegationService — one façade over registering, publishing and supervising.

Phase 1 keeps every behaviour in the existing entry points; this only groups
them so a long-running caller (the phase 2a engine) has a single object to hold,
optionally with a persistent MQTT connection. CLI and shell callers always use
`persistent=False`, because a process that exits after one publish gains nothing
from connection reuse.
"""
from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional

from lib_py.delegation import daemon, mqtt_common, publish_event, registry

logger = logging.getLogger("delegate_job.service")


class DelegationService(object):
    def __init__(self, registry_dir=None, persistent=False):
        self.registry_dir = registry_dir or mqtt_common.default_registry_dir()
        self.persistent = persistent
        self._publisher = None

    # -- lifecycle ---------------------------------------------------------
    def __enter__(self):
        if self.persistent:
            from lib_py.delegation.publisher import PersistentPublisher
            self._publisher = PersistentPublisher()
            self._publisher.open()
        return self

    def __exit__(self, exc_type, exc, tb):
        self.close()
        return False

    def close(self):
        if self._publisher is not None:
            self._publisher.close()
            self._publisher = None

    # -- operations --------------------------------------------------------
    def submit_job(self, prompt, agent="claude-code", agent_session="herdr:claude", **kw):
        return registry.register_job(
            prompt=prompt, agent=agent, agent_session=agent_session,
            registry_dir=self.registry_dir, **kw
        )

    def publish_event(self, job_id, event, detail="", data=None, attempts=3,
                      expect_dispatch_token=None, attempt_id=None):
        """Publish one event. Returns the same exit status as the CLI (0 / 2 / 4)."""
        eff_dispatch_token = expect_dispatch_token
        eff_attempt_id = attempt_id
        if data and isinstance(data, dict):
            if eff_dispatch_token is None:
                eff_dispatch_token = data.get("dispatch_token")
            if eff_attempt_id is None:
                eff_attempt_id = data.get("attempt_id")

        try:
            rec = self.load_job(job_id)
            actual_tok = rec.get("dispatch_token")
            rec_attempt = rec.get("attempt_id", rec.get("iteration", 1))
        except FileNotFoundError:
            rec = None
            actual_tok = None
            rec_attempt = 1

        if rec is not None and actual_tok is not None:
            if eff_dispatch_token is None and eff_attempt_id is None:
                if event in ("completed", "error", "cancelled") and rec_attempt > 1:
                    logger.error("dispatch token required: redispatched job %s carries a dispatch token "
                                 "but none was provided", job_id)
                    return 4
            elif eff_dispatch_token is not None and actual_tok != eff_dispatch_token:
                return 4
            elif eff_attempt_id is not None and rec_attempt != eff_attempt_id:
                return 4
        elif eff_dispatch_token is not None and actual_tok != eff_dispatch_token:
            return 4

        prepared = publish_event.prepare_event(
            job_id=job_id, event=event, detail=detail,
            data=data, registry_dir=self.registry_dir,
            dispatch_token=eff_dispatch_token, attempt_id=eff_attempt_id,
        )
        # Local first, as in the CLI: the event reaches events.log before the
        # publish can block on the broker (see publish_event.record_local_event).
        publish_event.record_local_event(prepared, self.registry_dir)
        publisher = self._publisher
        if publisher is None:
            from lib_py.delegation.publisher import OneShotPublisher
            publisher = OneShotPublisher()
        ok, error = publisher.publish_with_retry(prepared, attempts=attempts)
        publish_event.record_outcome(
            prepared, self.registry_dir, ok, error,
            expected_dispatch_token=eff_dispatch_token,
            expected_attempt=eff_attempt_id,
            local_event_recorded=True,
        )
        return 0 if ok else 2

    def ensure_subscriber_daemon(self, job_id, workdir=None):
        return daemon.ensure_subscriber_daemon(job_id, workdir)

    def load_job(self, job_id) -> Dict[str, Any]:
        return mqtt_common.load_job(job_id, self.registry_dir)