nisse/nisse-dev/src/nisse_dev/control.py

"""Control loop and state machine (spec A4, B1 subset).

Gating is structural, not advisory: the agent's only actions are ``SubmitPlan``
(posts the plan for approval and ends the agent's turn) and ``AskHuman``
(pause, ask, resume on answer). There is deliberately no proceed/start action —
advancement past a gate happens only here, when an operator approval or command
arrives over transport. Every position change is checkpointed through the state
store before the loop continues, so a restart lands mid-gate, not mid-air.
"""

import threading
import uuid
from collections.abc import Iterable, Mapping
from dataclasses import dataclass
from typing import Any, Protocol

from nisse_dev.canonical import GateOpened, PlanSubmitted, UserMessage
from nisse_dev.gates import Approved, Gate, Ignored, PlanRevision, Rejected
from nisse_dev.git import GitError
from nisse_dev.gitea import GiteaError
from nisse_dev.llm import RoleMap
from nisse_dev.reconcile import Reconciler, reconcile_on_boot
from nisse_dev.state import StateStore
from nisse_dev.transport import InboundMessage, Transport

IDLE = "IDLE"
PLANNING_ARCH = "PLANNING_ARCH"
AWAIT_ARCH_SIGNOFF = "AWAIT_ARCH_SIGNOFF"
MILESTONE_PLANNING = "MILESTONE_PLANNING"
AWAIT_MILESTONE_SIGNOFF = "AWAIT_MILESTONE_SIGNOFF"
EXECUTING = "EXECUTING"
PR_OPEN = "PR_OPEN"
AWAIT_INPUT = "AWAIT_INPUT"

TRANSITIONS = {
    (IDLE, PLANNING_ARCH),
    (PLANNING_ARCH, AWAIT_ARCH_SIGNOFF),
    (AWAIT_ARCH_SIGNOFF, PLANNING_ARCH),  # rejected: replan with feedback
    (AWAIT_ARCH_SIGNOFF, MILESTONE_PLANNING),
    (MILESTONE_PLANNING, AWAIT_MILESTONE_SIGNOFF),
    (AWAIT_MILESTONE_SIGNOFF, MILESTONE_PLANNING),
    (AWAIT_MILESTONE_SIGNOFF, EXECUTING),
    (EXECUTING, PR_OPEN),
    # PR acceptance is the milestone's final gate (B1): a merged PR is the
    # durable record, so completion returns to planning or to idle rather
    # than passing through a separate FINISHED state.
    (PR_OPEN, MILESTONE_PLANNING),
    (PR_OPEN, IDLE),
}


class TransitionError(RuntimeError):
    pass


class PlanningError(RuntimeError):
    """A planning turn produced nothing usable. Part of the PlanningAgent
    contract: the orchestrator gates to AWAIT_INPUT rather than inventing a
    plan or retrying forever."""


@dataclass(frozen=True)
class SubmitPlan:
    content: str


@dataclass(frozen=True)
class AskHuman:
    question: str


class MilestoneExecutor(Protocol):
    """What the orchestrator needs from an executor loop; the concrete one
    lives in nisse_dev.executor (imported there, not here, to keep the
    dependency pointing one way)."""

    def run(self, milestone: int) -> Any: ...


class MilestonePublisher(Protocol):
    """What the orchestrator needs to publish a milestone; the concrete one
    lives in nisse_dev.publish."""

    def reconcilers(self) -> Mapping[str, Reconciler]: ...

    def publish(
        self, *, task_id: str, milestone: int, project: str, plan: str
    ) -> Any: ...

    def push_followup(
        self, *, task_id: str, milestone: int, message: str
    ) -> str | None: ...

    def comment(self, index: int, body: str) -> None: ...


class PlanBoard(Protocol):
    """The Gitea side of plan gating; the concrete one is in
    nisse_dev.approvals."""

    def open_project_issue(self, project: str, task_id: str) -> int: ...

    def post_plan(self, issue: int, gate: str, rev: int, plan: str) -> Any: ...

    def echo(self, issue: int, text: str) -> None: ...


class PlanningAgent(Protocol):
    def plan(
        self,
        gate: str,
        project: str,
        feedback: str | None = None,
        answer: str | None = None,
    ) -> SubmitPlan | AskHuman: ...


class StubAgent:
    """Echoes canned plans so the approval loop is exercisable with no LLM."""

    def plan(
        self,
        gate: str,
        project: str,
        feedback: str | None = None,
        answer: str | None = None,
    ) -> SubmitPlan | AskHuman:
        notes = ""
        if feedback:
            notes += f" (revised after feedback: {feedback})"
        if answer:
            notes += f" (operator answered: {answer})"
        return SubmitPlan(f"[stub] {gate} plan for {project}{notes}")


def _fresh_data() -> dict[str, Any]:
    return {
        "revs": {},
        "gate": None,
        # What an operator reply in AWAIT_INPUT should resume. Every path into
        # AWAIT_INPUT must set this, or the reply has nowhere to go.
        "pending": None,
        "milestone": 1,
        "task_id": None,
        "roles": None,  # A11 mapping, pinned at task start
        # What each pinned alias resolved to, discovered from LiteLLM. An
        # alias is free to point somewhere else tomorrow, so the transcript
        # has to say what actually answered (A11, A16).
        "models": None,
        "pr": None,  # the open PR for the current milestone, once published
        # Pointer to the Gitea plan record (Part H): the state store holds
        # only this, so gate state can be re-derived from Gitea after a
        # state-volume loss.
        "board": None,
    }


class Orchestrator:
    def __init__(
        self,
        store: StateStore,
        transport: Transport,
        agent: PlanningAgent,
        operators: Iterable[str],
        roles: RoleMap | None = None,
        executor: "MilestoneExecutor | None" = None,
        publisher: "MilestonePublisher | None" = None,
        gitea_operators: Iterable[str] = (),
        board: "PlanBoard | None" = None,
        delegate: Any = None,
        poller_factory: Any = None,
        catalog: Any = None,
    ) -> None:
        self._store = store
        self._transport = transport
        self._agent = agent
        self._roles = roles
        self._catalog = catalog
        self._executor = executor
        self._publisher = publisher
        self._operators = set(operators)
        # Per-channel allowlists (A4): a Matrix MXID and a Gitea login
        # are different names for the same person, and neither implies
        # the other.
        self._gitea_operators = set(gitea_operators)
        self._board = board
        self._delegate = delegate
        # Builds a poller bound to a comment handler; the orchestrator
        # owns *when* to watch, which is a state-machine concern, while
        # the factory owns how to build one.
        self._poller_factory = poller_factory
        self._pollers: dict[str, Any] = {}
        # Transport and the comment poller deliver on separate threads;
        # state transitions must not interleave.
        self._lock = threading.RLock()
        self._data = _fresh_data()
        self._project: str | None = None
        transport.on_message(self._handle)

    def start(self, reconcilers: Mapping[str, Reconciler] | None = None) -> None:
        # A3 step 5 runs before the generic sweep: an interrupted delegation is
        # resumed or discarded by inspecting the branch and the log, and only a
        # genuinely unclean cut is left for the sweep below to gate on.
        if self._delegate is not None:
            self._delegate.recover(self._transport.send_message)
        # The publisher owns the external-truth query for its own intents, so
        # a configured publisher always contributes its reconciler (A3).
        registered: dict[str, Reconciler] = dict(reconcilers or {})
        if self._publisher is not None:
            registered = {**self._publisher.reconcilers(), **registered}
        report = reconcile_on_boot(self._store, registered)
        state = self._store.get_state()
        if state is None:
            self._transition(IDLE, cause="boot: fresh state")
            self._transport.send_message("nisse-dev ready (idle)")
            return
        self._data = {**_fresh_data(), **state.data}
        self._project = state.project
        self._store.oplog.task_id = self._data.get("task_id")
        gate = self._data.get("gate")
        if report.action == "await_input":
            self._await_input(
                cause="boot: unresolved intents",
                prompt=(
                    "restart found unresolved work: "
                    + "; ".join(report.notes)
                    + " — how should I proceed?"
                ),
                pending={"kind": "boot"},
            )
        elif state.state == IDLE or (
            state.state in (AWAIT_ARCH_SIGNOFF, AWAIT_MILESTONE_SIGNOFF) and gate
        ):
            # A gate is a safe place to have been interrupted: nothing was in
            # flight, the plan and revision are durable. Announce and resume.
            if gate:
                self._transport.send_message(
                    f"restarted while awaiting sign-off on {gate['name']} "
                    f"rev-{gate['rev']}; the gate is still open"
                )
        else:
            # Interrupted mid-turn (planning, executing, unanswered question):
            # never blind-resume (A3.6).
            self._await_input(
                cause=f"boot: restarted in {state.state}",
                prompt=f"restarted mid-{state.state}; continue or re-plan?",
                pending={"kind": "boot"},
            )

    def _handle(self, msg: InboundMessage) -> None:
        with self._lock:
            self._handle_locked(msg)

    def _handle_locked(self, msg: InboundMessage) -> None:
        if msg.sender not in self._operators:
            return
        # Every operator message enters the canonical transcript, so planning
        # turns see feedback and answers as conversation rather than as
        # out-of-band parameters (A10).
        self._store.append_event(
            UserMessage(text=msg.text.strip(), source=f"transport:{msg.sender}")
        )
        state = self._store.get_state()
        current = state.state if state else IDLE
        text = msg.text.strip()
        if text == "status":
            self._send_status(current)
        elif current == IDLE:
            if text.startswith("use "):
                self._begin_project(text[4:].strip())
            else:
                self._transport.send_message("no active project; say 'use <project>'")
        elif current in (AWAIT_ARCH_SIGNOFF, AWAIT_MILESTONE_SIGNOFF):
            self._handle_gate_reply(msg, text)
        elif current == PR_OPEN:
            self._handle_pr_open(text)
        elif current == AWAIT_INPUT:
            self._handle_answer(text)
        else:
            self._transport.send_message(f"busy ({current}); 'status' for details")

    def _send_status(self, current: str) -> None:
        gate = self._data.get("gate")
        gate_note = (
            f", awaiting sign-off on {gate['name']} rev-{gate['rev']}" if gate else ""
        )
        self._transport.send_message(
            f"state={current} project={self._project}{gate_note}"
        )

    def _begin_project(self, project: str) -> None:
        task_id = f"{project}-{uuid.uuid4().hex[:8]}"
        # Pin the role mapping now (A11): a mid-task config change must not
        # split one task across two planners.
        pinned = self._roles.pinned() if self._roles is not None else None
        resolved = (
            self._catalog.pinned(pinned)
            if self._catalog is not None and pinned is not None
            else None
        )
        self._data = {
            **_fresh_data(),
            "task_id": task_id,
            "roles": pinned,
            "models": resolved,
        }
        self._project = project
        self._store.oplog.task_id = task_id
        if self._board is not None:
            try:
                self._data["board"] = {
                    "issue": self._board.open_project_issue(project, task_id),
                    "last_comment": 0,
                }
            except GiteaError as exc:
                # Gitea being down must not block planning; Matrix remains a
                # complete approval channel on its own.
                self._transport.send_message(
                    f"note: could not open the Gitea plan issue ({exc}); "
                    "approvals will be Matrix-only for this task"
                )
        self._transition(PLANNING_ARCH, cause=f"operator command: use {project}")
        self._run_planning_turn("arch")

    def _run_planning_turn(
        self, gate_name: str, feedback: str | None = None, answer: str | None = None
    ) -> None:
        assert self._project is not None
        try:
            action = self._agent.plan(
                gate_name, self._project, feedback=feedback, answer=answer
            )
        except PlanningError as exc:
            self._await_input(
                cause=f"planning failed: {exc}",
                prompt=f"planning for {gate_name} failed ({exc}); how should I proceed?",
                pending={"kind": "planning_failed", "gate": gate_name},
            )
            return
        if isinstance(action, AskHuman):
            self._await_input(
                cause=f"agent question during {gate_name} planning",
                prompt=action.question,
                pending={"kind": "question", "gate": gate_name},
            )
            return
        rev = self._data["revs"].get(gate_name, 0) + 1
        self._data["revs"][gate_name] = rev
        self._data["gate"] = {
            "name": gate_name,
            "rev": rev,
            "plan": action.content,
            "approved": False,
        }
        await_state = (
            AWAIT_ARCH_SIGNOFF if gate_name == "arch" else AWAIT_MILESTONE_SIGNOFF
        )
        roles = self._data.get("roles") or {}
        self._store.append_event(
            PlanSubmitted(plan=action.content, model=roles.get("planner", "stub"))
        )
        self._store.append_event(
            GateOpened(gate=gate_name, reason=f"plan submitted (rev-{rev})")
        )
        self._transition(await_state, cause=f"plan submitted ({gate_name} rev-{rev})")
        self._transport.send_approval_request(
            f"{action.content}\n\nreply 'approve' to approve rev-{rev}; "
            "any other reply is rejection feedback"
        )
        self._post_plan_to_board(gate_name, rev, action.content)
        self._store.oplog.emit(
            "gate.open",
            cause=f"plan submitted (rev-{rev})",
            gate=gate_name,
            rev=f"rev-{rev}",
            channel="transport",
        )

    def _handle_gate_reply(self, msg: InboundMessage, text: str) -> None:
        self._apply_gate_reply(sender=msg.sender, text=text, channel="transport")

    def _apply_gate_reply(self, *, sender: str, text: str, channel: str) -> None:
        """Shared by every approval source (A4). The allowlist is per channel:
        a Matrix MXID and a Gitea login are different names, and holding one
        says nothing about holding the other."""
        gate_data = self._data["gate"]
        name, rev = gate_data["name"], gate_data["rev"]
        milestone = self._data["milestone"]
        if gate_data["approved"]:
            if text == f"start milestone {milestone}":
                self._transition(
                    EXECUTING, cause=f"operator command: start milestone {milestone}"
                )
                self._execute_milestone(milestone)
            else:
                self._transport.send_message(
                    f"plan approved; say 'start milestone {milestone}' to begin"
                )
            return
        revision = PlanRevision(gate=name, rev=rev, content=gate_data["plan"])
        allowlist = self._operators if channel == "transport" else self._gitea_operators
        result = Gate(revision, allowlist).evaluate(sender, text, channel=channel)
        if isinstance(result, Ignored):
            if result.notice:
                self._transport.send_message(result.notice)
            return
        if isinstance(result, Rejected):
            self._store.oplog.emit(
                "gate.reject",
                cause=f"rejected by {result.approver}",
                gate=name,
                rev=f"rev-{rev}",
                channel=result.channel,
            )
            self._echo_outcome(
                channel,
                f"{name} rev-{rev} rejected by {result.approver}: {result.feedback}",
            )
            self._data["gate"] = None
            back = PLANNING_ARCH if name == "arch" else MILESTONE_PLANNING
            self._transition(back, cause=f"plan rejected ({name} rev-{rev})")
            self._run_planning_turn(name, feedback=result.feedback)
            return
        assert isinstance(result, Approved)
        self._echo_outcome(
            channel, f"{name} rev-{rev} approved by {result.approver} via {channel}"
        )
        self._store.oplog.emit(
            "gate.approve",
            cause=f"approved by {result.approver}",
            gate=name,
            rev=f"rev-{rev}",
            channel=result.channel,
        )
        if name == "arch":
            self._data["gate"] = None
            self._transition(
                MILESTONE_PLANNING, cause=f"arch plan approved (rev-{rev})"
            )
            self._transport.send_message(
                f"architecture approved; planning milestone {milestone}"
            )
            self._run_planning_turn(f"milestone-{milestone}")
        else:
            gate_data["approved"] = True
            self._transition(
                AWAIT_MILESTONE_SIGNOFF,
                cause=f"milestone plan approved (rev-{rev})",
            )
            self._transport.send_message(
                f"milestone plan approved; say 'start milestone {milestone}' to begin"
            )

    def _execute_milestone(self, milestone: int) -> None:
        if self._executor is None:
            self._transport.send_message(f"[stub] executing milestone {milestone}")
            return
        self._transport.send_message(f"executing milestone {milestone}")
        outcome = self._executor.run(milestone)
        if outcome.status == "completed":
            self._transport.send_message(f"milestone {milestone} done: {outcome.note}")
            self._publish_milestone(milestone)
            return
        self._await_input(
            cause=f"milestone {milestone} {outcome.status}",
            prompt=f"milestone {milestone} stopped ({outcome.status}): {outcome.note}",
            pending={"kind": "milestone_stopped", "milestone": milestone},
        )

    def _echo_outcome(self, deciding_channel: str, text: str) -> None:
        """A gate satisfied on one channel is echoed to the other, so neither
        record silently disagrees about what was approved (A4)."""
        board_state = self._data.get("board")
        if deciding_channel != "gitea" and self._board is not None and board_state:
            self._board.echo(board_state["issue"], text)
        elif deciding_channel == "gitea":
            self._transport.send_message(text)

    def _post_plan_to_board(self, gate_name: str, rev: int, plan: str) -> None:
        board_state = self._data.get("board")
        if self._board is None or not board_state:
            return
        try:
            post = self._board.post_plan(board_state["issue"], gate_name, rev, plan)
        except GiteaError:
            # Matrix already carries the plan; a Gitea outage degrades the
            # durability of the record, not the ability to approve.
            return
        board_state["comment_id"] = post.comment_id
        # Approvals can arrive on the issue as well as over transport (A4),
        # so watch it for as long as the gate stands.
        self._watch("plan", board_state["issue"], self.handle_plan_comment)

    def handle_plan_comment(self, comment: Any) -> None:
        """A comment on the project issue may satisfy an open plan gate (A4).

        Approvals bind to a revision here exactly as over transport, and the
        first valid approval on either channel wins.
        """
        with self._lock:
            self._handle_plan_comment_locked(comment)

    def _handle_plan_comment_locked(self, comment: Any) -> None:
        if comment.author not in self._gitea_operators:
            return
        state = self._store.get_state()
        if state is None or state.state not in (
            AWAIT_ARCH_SIGNOFF,
            AWAIT_MILESTONE_SIGNOFF,
        ):
            return
        self._store.append_event(
            UserMessage(text=comment.body, source=f"gitea:comment:{comment.id}")
        )
        self._apply_gate_reply(
            sender=comment.author, text=comment.body.strip(), channel="gitea"
        )

    def handle_review_comment(self, comment: Any) -> None:
        """React to a PR review comment (B2). Called from the poller thread.

        Comment authors are checked against the Gitea operator allowlist:
        a PR comment is an instruction channel into an agent that writes
        code, so it is operator-scoped like any other (A4, and the same
        reasoning C4 applies to tickets).
        """
        with self._lock:
            self._handle_review_comment_locked(comment)

    def _handle_review_comment_locked(self, comment: Any) -> None:
        if comment.author not in self._gitea_operators:
            self._store.oplog.emit(
                "gitea.comment_ignored",
                cause=f"author {comment.author!r} not an operator",
                comment_id=comment.id,
            )
            return
        state = self._store.get_state()
        if state is None or state.state != PR_OPEN:
            self._store.oplog.emit(
                "gitea.comment_ignored",
                cause=f"not in PR_OPEN (state={state.state if state else None})",
                comment_id=comment.id,
            )
            return
        if self._executor is None or self._publisher is None:
            return

        pr = self._data.get("pr") or {}
        milestone = int(pr.get("milestone", self._data.get("milestone", 1)))
        # The comment enters the transcript, so the executor sees review
        # feedback as conversation rather than as an out-of-band instruction.
        self._store.append_event(
            UserMessage(text=comment.body, source=f"gitea:comment:{comment.id}")
        )
        outcome = self._executor.run(milestone)
        if outcome.status != "completed":
            self._await_input(
                cause=f"review response {outcome.status}",
                prompt=(
                    f"could not address the review comment ({outcome.status}): "
                    f"{outcome.note}"
                ),
                pending={"kind": "milestone_stopped", "milestone": milestone},
            )
            return
        try:
            sha = self._publisher.push_followup(
                task_id=str(self._data.get("task_id")),
                milestone=milestone,
                message=f"Address review comment {comment.id}",
            )
        except (GiteaError, GitError) as exc:
            self._transport.send_message(f"could not push the fix: {exc}")
            return
        note = (
            f"pushed {sha[:8]} addressing the review comment"
            if sha
            else "no code change was needed for that comment"
        )
        self._transport.send_message(note)
        if pr.get("number"):
            self._publisher.comment(int(pr["number"]), note)

    def _publish_milestone(self, milestone: int) -> None:
        if self._publisher is None:
            return
        gate = self._data.get("gate") or {}
        try:
            result = self._publisher.publish(
                task_id=str(self._data.get("task_id")),
                milestone=milestone,
                project=str(self._project),
                plan=str(gate.get("plan", "")),
            )
        except (GiteaError, GitError) as exc:
            # The intent log already records what was in flight, so boot
            # reconciliation can settle it against Gitea. Gate rather than
            # retry blindly into a possible duplicate.
            self._await_input(
                cause=f"publishing milestone {milestone} failed: {exc}",
                prompt=f"could not publish milestone {milestone}: {exc}",
                pending={"kind": "milestone_stopped", "milestone": milestone},
            )
            return
        if result.pull_request is None:
            self._transport.send_message(result.note)
            return
        self._data["pr"] = {
            "number": result.pull_request.number,
            "url": result.pull_request.url,
            "milestone": milestone,
        }
        self._unwatch("plan")  # the gate is behind us; the PR is the thread now
        self._watch("review", result.pull_request.number, self.handle_review_comment)
        self._transition(PR_OPEN, cause=f"milestone {milestone} published")
        self._transport.send_message(
            f"{result.note} — review it, then say 'next milestone' or "
            "'project complete'"
        )

    def _handle_pr_open(self, text: str) -> None:
        if text == "next milestone":
            self._unwatch("review")
            self._data["milestone"] = int(self._data.get("milestone", 1)) + 1
            self._data["gate"] = None
            self._data["pr"] = None
            self._transition(MILESTONE_PLANNING, cause="operator: next milestone")
            self._run_planning_turn(f"milestone-{self._data['milestone']}")
        elif text == "project complete":
            self.stop_watching()
            self._data = {**_fresh_data()}
            self._project = None
            self._store.oplog.task_id = None
            self._transition(IDLE, cause="operator: project complete")
            self._transport.send_message(
                "project closed; say 'use <project>' to start another"
            )
        else:
            pr = self._data.get("pr") or {}
            self._transport.send_message(
                f"PR #{pr.get('number')} is open; say 'next milestone' or "
                "'project complete'"
            )

    def _watch(self, kind: str, index: int, handler: Any) -> None:
        """Start watching one Gitea thread, replacing any previous watch of
        the same kind. Priming first, so attaching does not replay comments
        already addressed."""
        if self._poller_factory is None:
            return
        self._unwatch(kind)
        poller = self._poller_factory(handler)
        poller.prime(index)
        poller.start(index)
        self._pollers[kind] = poller

    def _unwatch(self, kind: str) -> None:
        poller = self._pollers.pop(kind, None)
        if poller is not None:
            poller.stop()

    def stop_watching(self) -> None:
        for kind in list(self._pollers):
            self._unwatch(kind)

    def _await_input(self, *, cause: str, prompt: str, pending: dict[str, Any]) -> None:
        self._data["pending"] = pending
        self._transition(AWAIT_INPUT, cause=cause)
        self._transport.ask_question(prompt)

    def _handle_answer(self, text: str) -> None:
        pending = self._data.get("pending")
        if pending is None:
            # Reached only if state was forced externally; say what works
            # rather than dead-ending the operator.
            self._transport.send_message(
                "waiting for direction, but I have no pending item to resume; "
                "try 'status', or 'use <project>' to start over"
            )
            return
        self._data["pending"] = None
        match pending.get("kind"):
            case "question" | "planning_failed":
                gate_name = pending["gate"]
                back = PLANNING_ARCH if gate_name == "arch" else MILESTONE_PLANNING
                self._transition(back, cause="operator answered; resuming planning")
                self._run_planning_turn(gate_name, answer=text)
            case "milestone_stopped":
                milestone = pending["milestone"]
                self._transition(
                    AWAIT_MILESTONE_SIGNOFF,
                    cause="operator answered; plan still approved",
                )
                self._transport.send_message(
                    f"noted. the milestone {milestone} plan is still approved — "
                    f"say 'start milestone {milestone}' to run it again"
                )
            case _:
                self._transport.send_message(
                    "noted. say 'status' for where things stand"
                )

    def _transition(self, new_state: str, *, cause: str) -> None:
        state = self._store.get_state()
        current = state.state if state else None
        allowed = (
            current is None
            or current == new_state
            or new_state == AWAIT_INPUT
            or current == AWAIT_INPUT
            or (current, new_state) in TRANSITIONS
        )
        if not allowed:
            raise TransitionError(f"illegal transition {current} -> {new_state}")
        self._store.transition(
            new_state, cause=cause, project=self._project, data=self._data
        )