bot-bottle/bot_bottle/orchestrator/lifecycle.py

"""The orchestration lifecycle: forge events -> bottle transitions.

`Orchestrator.handle(event)` is the single entry point the webhook layer
calls. `on_done_signal(...)` is called by the sidecar relay when an agent
signals completion. All collaborators (forge, store, runner) are
injected and duck-typed; `now` and `label_for` are injectable for tests.

Transitions:
  IssueAssigned (targeted, new)  -> start bottle, record = running
  signal_done (running)          -> freeze bottle, record = frozen
  CommentCreated (frozen)        -> resume bottle, record = running
  PullRequestClosed (tracked)    -> destroy bottle, record removed
"""

from __future__ import annotations

from collections.abc import Callable
from datetime import datetime

from .model import (
    STATUS_DESTROYED,
    STATUS_FROZEN,
    STATUS_RUNNING,
    CommentCreated,
    ForgeEvent,
    IssueAssigned,
    PullRequestClosed,
    RunRecord,
)
from .runner import BottleRunner
from .store import StateStore
from .targeting import Membership, Target, resolve_target


def _iso_now() -> str:
    return datetime.now().astimezone().isoformat(timespec="seconds")


def _default_label(agent: str, event: IssueAssigned) -> str:
    # Embed the issue identity so slugs are unique per issue and never
    # get renamed on collision.
    return f"{agent}-{event.owner}-{event.repo}-{event.issue_number}"


class Orchestrator:
    def __init__(
        self,
        *,
        forge: Membership,
        store: StateStore,
        runner: BottleRunner,
        org: str,
        gitea_api: str = "",
        forge_env_base: dict[str, str] | None = None,
        now: Callable[[], str] = _iso_now,
        label_for: Callable[[str, IssueAssigned], str] = _default_label,
    ) -> None:
        self._forge = forge
        self._store = store
        self._runner = runner
        self._org = org
        self._gitea_api = gitea_api
        self._forge_env_base = forge_env_base or {}
        self._now = now
        self._label_for = label_for

    # --- entry points ------------------------------------------------------

    def handle(self, event: ForgeEvent) -> None:
        if isinstance(event, IssueAssigned):
            self._on_issue_assigned(event)
        elif isinstance(event, CommentCreated):
            self._on_comment(event)
        else:
            self._on_pr_closed(event)

    def on_done_signal(  # pylint: disable=unused-argument
        self, owner: str, repo: str, issue_number: int, status: str, summary: str
    ) -> None:
        """Sidecar relay: an agent signalled completion. Freeze the bottle.
        `status`/`summary` are recorded by provenance (via the op log), not
        acted on here."""
        record = self._store.get(owner, repo, issue_number)
        if record is None or record.status != STATUS_RUNNING:
            return
        self._runner.freeze(record.slug)
        record.status = STATUS_FROZEN
        record.last_checkin_at = self._now()
        self._store.upsert(record)

    def link_pr(self, owner: str, repo: str, issue_number: int, pr_number: int) -> None:
        """Record the PR a tracked issue produced, so PR comments and the
        PR-close event route back to this record."""
        record = self._store.get(owner, repo, issue_number)
        if record is not None:
            record.pr_number = pr_number
            self._store.upsert(record)

    # --- handlers ----------------------------------------------------------

    def _on_issue_assigned(self, event: IssueAssigned) -> None:
        target = resolve_target(event, self._forge, self._org)
        if target is None:
            return
        # Idempotent: a webhook redelivery must not launch a second bottle.
        if self._store.get(event.owner, event.repo, event.issue_number) is not None:
            return
        self._launch(event, target)

    def _launch(self, event: IssueAssigned, target: Target) -> None:
        label = self._label_for(target.agent_name, event)
        bottles = [target.bottle_override] if target.bottle_override else []
        result = self._runner.start(
            agent=target.agent_name,
            bottles=bottles,
            label=label,
            prompt=event.body,
            forge_env=self._forge_env(event.owner, event.repo, event.issue_number),
        )
        self._store.upsert(
            RunRecord(
                owner=event.owner,
                repo=event.repo,
                issue_number=event.issue_number,
                slug=result.slug,
                agent_name=target.agent_name,
                bottle_names=bottles,
                status=STATUS_RUNNING,
                last_checkin_at=self._now(),
            )
        )

    def _on_comment(self, event: CommentCreated) -> None:
        record = self._route_comment(event)
        if record is None or record.status != STATUS_FROZEN:
            return
        # Echo-loop guard: ignore the agent's own comments.
        if record.agent_git_user and event.author == record.agent_git_user:
            return
        self._runner.resume(record.slug, event.body)
        record.status = STATUS_RUNNING
        record.last_checkin_at = self._now()
        self._store.upsert(record)

    def _route_comment(self, event: CommentCreated) -> RunRecord | None:
        # A comment on the issue routes by issue number; a comment on a PR
        # routes by the recorded pr_number.
        direct = self._store.get(event.owner, event.repo, event.issue_number)
        if direct is not None:
            return direct
        if event.is_pull:
            return self._find_by_pr(event.owner, event.repo, event.issue_number)
        return None

    def _on_pr_closed(self, event: PullRequestClosed) -> None:
        record = self._find_by_pr(event.owner, event.repo, event.pr_number)
        if record is None:
            return
        self._runner.destroy(record.slug)
        record.status = STATUS_DESTROYED
        self._store.delete(record.owner, record.repo, record.issue_number)

    def _find_by_pr(self, owner: str, repo: str, pr_number: int) -> RunRecord | None:
        for record in self._store.all():
            if (
                record.owner == owner
                and record.repo == repo
                and record.pr_number == pr_number
            ):
                return record
        return None

    def _forge_env(self, owner: str, repo: str, issue_number: int) -> dict[str, str]:
        env = dict(self._forge_env_base)
        if self._gitea_api:
            env["FORGE_GITEA_API"] = self._gitea_api
        env["FORGE_OWNER"] = owner
        env["FORGE_REPO"] = repo
        env["FORGE_ISSUE_NUMBER"] = str(issue_number)
        return env