smart-org/apps/api/app/application/engine/protocol.py

"""The tag protocol: ``[[TAG attr="v" ...]]body[[/TAG]]`` and ``[[TAG attr="v"/]``.

One parser, one matrix. The rules that make an agent protocol safe all live
here so no call site can re-decide them:

1. **An allow-list per role and mode.** The manager triages, gives the final
   verdict, handles errors and chats, and each mode may emit a different set
   of tags; a planner plans, a developer implements. A tag outside the matrix
   is ignored and reported back to the model as an error line — the model
   learns the rule, and nothing dangerous executes by accident.
2. **The toolkit is the only way to touch files.** ``READ_FILE``,
   ``WRITE_FILE``, ``EDIT_FILE``, ``LIST_DIR``, ``SEARCH`` and ``RUN`` run
   inside the sandbox through ``SandboxProviderPort``; the platform still
   does git (``RepoWorkspacePort``), and ``RUN`` refuses git push/remote/config
   outright (see ``engine/tools.py``).
3. **Bounded turns.** A reply may create at most three tasks, change the team
   at most three times and record at most three knowledge items. A model that
   ignores the limit gets an error line and the extra tags are dropped, so the
   effect of one runaway reply stays finite.

Bodies may contain brackets (a source file does), so the parser never
re-scans a body it already consumed: an outer tag wins and its content is
literal.
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field
from enum import StrEnum

from app.domain.workforce.models import RoleKey
from app.domain.workforce.pipeline import Mode, PipelineKey


class TagName(StrEnum):
    """Every tag the protocol knows, whether or not a role may use it."""

    # manager / triage
    ASSIGN = "ASSIGN"
    REJECT = "REJECT"
    ASK_USER = "ASK_USER"
    KNOWLEDGE = "KNOWLEDGE"
    # manager / final
    COMPLETE = "COMPLETE"
    REVISE = "REVISE"
    # manager / error
    RESOLVE_ERROR = "RESOLVE_ERROR"
    CANCEL = "CANCEL"
    # manager / chat
    CREATE_TASK = "CREATE_TASK"
    PAUSE_TASK = "PAUSE_TASK"
    RESUME_TASK = "RESUME_TASK"
    CANCEL_TASK = "CANCEL_TASK"
    HIRE = "HIRE"
    UPDATE_EMPLOYEE = "UPDATE_EMPLOYEE"
    FIRE = "FIRE"
    # planner
    PLAN_SPEC = "PLAN_SPEC"
    SPLIT = "SPLIT"
    # developer
    WORK_REPORT = "WORK_REPORT"
    # reviewer
    REVIEW_PASS = "REVIEW_PASS"
    REVIEW_CHANGES = "REVIEW_CHANGES"
    # tester
    TEST_PASS = "TEST_PASS"
    TEST_FAIL = "TEST_FAIL"
    # toolkit (every working role)
    READ_FILE = "READ_FILE"
    WRITE_FILE = "WRITE_FILE"
    EDIT_FILE = "EDIT_FILE"
    LIST_DIR = "LIST_DIR"
    SEARCH = "SEARCH"
    RUN = "RUN"


#: The toolkit. Available to every working role, never to the manager: the
#: manager decides, it does not edit the repository.
TOOL_TAGS: frozenset[TagName] = frozenset(
    {
        TagName.READ_FILE,
        TagName.WRITE_FILE,
        TagName.EDIT_FILE,
        TagName.LIST_DIR,
        TagName.SEARCH,
        TagName.RUN,
    }
)

#: Team-changing tags (all manager/chat). Counted together, one limit per turn.
TEAM_CHANGE_TAGS: frozenset[TagName] = frozenset(
    {TagName.HIRE, TagName.UPDATE_EMPLOYEE, TagName.FIRE}
)

#: Tags that carry a body and therefore *must* be closed with ``[[/TAG]]``.
#: Every other tag is a single token (an order, a tool read, a decision), so
#: a forgotten trailing ``/`` is tolerated as self-closing instead of failing
#: the whole reply: ``[[RUN cmd="pytest"]]`` means what ``[[RUN cmd="pytest"/]]``
#: means.
BODY_TAGS: frozenset[TagName] = frozenset(
    {
        TagName.WRITE_FILE,
        TagName.EDIT_FILE,
        TagName.KNOWLEDGE,
        TagName.ASK_USER,
        TagName.CREATE_TASK,
        TagName.SPLIT,
        TagName.PLAN_SPEC,
        TagName.WORK_REPORT,
        TagName.REVIEW_PASS,
        TagName.REVIEW_CHANGES,
        TagName.TEST_PASS,
        TagName.TEST_FAIL,
    }
)

#: The per-role/mode matrix. Anything not listed is refused with an error line.
ALLOWED_TAGS: dict[tuple[RoleKey, Mode], frozenset[TagName]] = {
    (RoleKey.MANAGER, Mode.TRIAGE): frozenset(
        {TagName.ASSIGN, TagName.REJECT, TagName.ASK_USER, TagName.KNOWLEDGE}
    ),
    (RoleKey.MANAGER, Mode.FINAL): frozenset(
        {TagName.COMPLETE, TagName.REVISE, TagName.KNOWLEDGE}
    ),
    (RoleKey.MANAGER, Mode.ERROR): frozenset(
        {TagName.RESOLVE_ERROR, TagName.CANCEL, TagName.KNOWLEDGE}
    ),
    (RoleKey.MANAGER, Mode.CHAT): frozenset(
        {
            TagName.CREATE_TASK,
            TagName.PAUSE_TASK,
            TagName.RESUME_TASK,
            TagName.CANCEL_TASK,
            TagName.HIRE,
            TagName.UPDATE_EMPLOYEE,
            TagName.FIRE,
            TagName.KNOWLEDGE,
        }
    ),
    (RoleKey.PLANNER, Mode.WORK): frozenset({TagName.PLAN_SPEC, TagName.SPLIT, TagName.ASK_USER})
    | TOOL_TAGS,
    (RoleKey.DEVELOPER, Mode.WORK): frozenset({TagName.WORK_REPORT, TagName.ASK_USER}) | TOOL_TAGS,
    (RoleKey.REVIEWER, Mode.WORK): frozenset(
        {TagName.REVIEW_PASS, TagName.REVIEW_CHANGES}
    )
    | TOOL_TAGS,
    (RoleKey.TESTER, Mode.WORK): frozenset({TagName.TEST_PASS, TagName.TEST_FAIL}) | TOOL_TAGS,
}

#: Per-reply limits. Extras are dropped and reported, never executed.
MAX_CREATE_TASKS_PER_REPLY = 3
MAX_TEAM_CHANGES_PER_TURN = 3
MAX_KNOWLEDGE_PER_TURN = 3

#: Which pipeline names ``ASSIGN pipeline=`` may pick. Kept here so the
#: manager cannot invent a pipeline that does not exist.
ASSIGNABLE_PIPELINES: frozenset[str] = frozenset(p.value for p in PipelineKey)


class ProtocolError(ValueError):
    """A tag was parsed but its attributes are not usable."""


# ---------------------------------------------------------------------------
# Tag and reply
# ---------------------------------------------------------------------------


@dataclass(frozen=True, slots=True)
class Tag:
    """One parsed tag: its name, attributes and body (``""`` when self-closing)."""

    name: TagName
    attrs: dict[str, str] = field(default_factory=dict)
    body: str = ""
    #: The raw slice, so an error line can quote exactly what the model wrote.
    raw: str = ""

    def attr(self, key: str, default: str = "") -> str:
        return self.attrs.get(key, default)

    def require(self, key: str) -> str:
        value = self.attrs.get(key, "").strip()
        if not value:
            raise ProtocolError(f"[[{self.name}]] requires {key}=...")
        return value


@dataclass(frozen=True, slots=True)
class ParsedReply:
    """Everything one model reply contained, in order.

    ``errors`` are lines fed back to the model on its next turn: the protocol
    is learned from feedback, not from a crash. ``prose`` is the text outside
    any tag — for a manager chat reply that is the message the user reads.
    """

    tags: tuple[Tag, ...] = ()
    errors: tuple[str, ...] = ()
    prose: str = ""

    def names(self) -> tuple[TagName, ...]:
        return tuple(tag.name for tag in self.tags)

    def of(self, name: TagName | str) -> tuple[Tag, ...]:
        wanted = TagName(name)
        return tuple(tag for tag in self.tags if tag.name is wanted)

    def first(self, name: TagName | str) -> Tag | None:
        found = self.of(name)
        return found[0] if found else None

    def has(self, name: TagName | str) -> bool:
        return bool(self.of(name))

    def error_text(self) -> str:
        return "\n".join(self.errors)


# ---------------------------------------------------------------------------
# Typed views over tags
# ---------------------------------------------------------------------------


@dataclass(frozen=True, slots=True)
class AssignTag:
    """``ASSIGN employee="name" pipeline="plan_code_test"``."""

    employee: str
    pipeline: str = ""
    task_id: str = ""


@dataclass(frozen=True, slots=True)
class TaskDraft:
    """``CREATE_TASK``: the manager asks for a new task in chat."""

    title: str
    kind: str = "feature"
    priority: str = "normal"
    depends_on: tuple[str, ...] = ()
    pipeline: str = ""
    description: str = ""


@dataclass(frozen=True, slots=True)
class TeamChange:
    """``HIRE`` / ``UPDATE_EMPLOYEE`` / ``FIRE``."""

    action: TagName
    name: str = ""
    role: str = ""
    tier: str = ""
    task_id: str = ""
    employee_id: str = ""


@dataclass(frozen=True, slots=True)
class KnowledgeItem:
    """``KNOWLEDGE kind="rule|fact" key="..."`` with the fact as the body."""

    kind: str = "rule"
    key: str = ""
    content: str = ""


@dataclass(frozen=True, slots=True)
class EditBlock:
    """One ``<<<OLD … === … NEW>>>`` search/replace block inside an EDIT_FILE."""

    old: str
    new: str


_EDIT_BLOCK_RE = re.compile(
    r"<<<OLD\s*\n(?P<old>.*?)\n===\s*\n(?P<new>.*?)\n\s*NEW>>>",
    re.DOTALL,
)


def assign_of(tag: Tag) -> AssignTag:
    """Read an ``ASSIGN`` tag. ``employee=""`` means "pick for me"."""
    employee = tag.attr("employee", "").strip()
    if not employee and not tag.attr("task_id", "").strip():
        raise ProtocolError("[[ASSIGN]] needs employee=... or task_id=...")
    pipeline = tag.attr("pipeline", "").strip()
    if pipeline and pipeline not in ASSIGNABLE_PIPELINES:
        raise ProtocolError(f"unknown pipeline: {pipeline}")
    return AssignTag(
        employee=employee, pipeline=pipeline, task_id=tag.attr("task_id", "").strip()
    )


def task_draft_of(tag: Tag) -> TaskDraft:
    """Read a ``CREATE_TASK`` tag."""
    title = tag.attr("title", "").strip()
    if not title:
        raise ProtocolError("[[CREATE_TASK]] needs title=...")
    raw_deps = tag.attr("depends_on", "")
    depends = tuple(d.strip() for d in raw_deps.split(",") if d.strip())
    return TaskDraft(
        title=title,
        kind=tag.attr("kind", "feature").strip() or "feature",
        priority=tag.attr("priority", "normal").strip() or "normal",
        depends_on=depends,
        pipeline=tag.attr("pipeline", "").strip(),
        description=tag.body.strip()[:4000],
    )


def team_change_of(tag: Tag) -> TeamChange:
    """Read one of the team tags. Users choose role and tier, never a model."""
    return TeamChange(
        action=tag.name,
        name=tag.attr("name", "").strip(),
        role=tag.attr("role", "").strip(),
        tier=tag.attr("tier", "").strip(),
        task_id=tag.attr("task_id", "").strip(),
        employee_id=tag.attr("employee_id", "").strip(),
    )


def knowledge_of(tag: Tag) -> KnowledgeItem:
    """Read a ``KNOWLEDGE`` tag; the body is the content."""
    return KnowledgeItem(
        kind=tag.attr("kind", "rule").strip() or "rule",
        key=tag.attr("key", "").strip(),
        content=tag.body.strip(),
    )


def edit_blocks_of(tag: Tag) -> list[EditBlock]:
    """Read the search/replace blocks of an ``EDIT_FILE`` body.

    Blocks are ``<<<OLD`` … ``===`` … ``NEW>>>``. A malformed body raises
    rather than guessing: a half-applied edit is worse than a failed one.
    """
    blocks = [
        EditBlock(old=m.group("old"), new=m.group("new"))
        for m in _EDIT_BLOCK_RE.finditer(tag.body)
    ]
    if not blocks:
        raise ProtocolError("[[EDIT_FILE]] needs one or more <<<OLD ... === ... NEW>>> blocks")
    for block in blocks:
        if not block.old.strip():
            raise ProtocolError("[[EDIT_FILE]] OLD block is empty")
    return blocks


def task_drafts(reply: ParsedReply) -> list[TaskDraft]:
    return [task_draft_of(tag) for tag in reply.of(TagName.CREATE_TASK)]


def team_changes(reply: ParsedReply) -> list[TeamChange]:
    """Every team tag, in the order the model wrote them."""
    return [team_change_of(tag) for tag in reply.tags if tag.name in TEAM_CHANGE_TAGS]


def knowledge_items(reply: ParsedReply) -> list[KnowledgeItem]:
    return [knowledge_of(tag) for tag in reply.of(TagName.KNOWLEDGE)]


# ---------------------------------------------------------------------------
# Parser
# ---------------------------------------------------------------------------

#: ``[[ ... ]]`` whose content has no brackets inside: an opening token, a
#: closing token, or a self-closing one. Attribute values with a ``]`` are
#: refused by the attribute parser below instead of silently truncating.
_TOKEN_RE = re.compile(r"\[\[([^\[\]]*)\]\]")
_ATTR_RE = re.compile(r"([A-Za-z_][A-Za-z0-9_]*)\s*=\s*(?:\"([^\"]*)\"|'([^']*)')")


def _closing_re(name: str) -> re.Pattern[str]:
    return re.compile(r"\[\[/\s*" + re.escape(name) + r"\s*\]\]", re.IGNORECASE)


def _split_token(inner: str) -> tuple[str, str, bool]:
    """``'TAG a="1"'`` -> ``("TAG", ' a="1"', False)``; trailing ``/`` self-closes."""
    text = inner.strip()
    self_closing = text.endswith("/")
    if self_closing:
        text = text[:-1].strip()
    match = re.match(r"([A-Za-z_][A-Za-z0-9_]*)", text)
    if match is None:  # pragma: no cover - the token regex requires a name
        return "", "", self_closing
    name = match.group(1)
    return name.upper(), text[match.end() :], self_closing


def _parse_attrs(rest: str) -> tuple[dict[str, str], list[str]]:
    attrs: dict[str, str] = {}
    errors: list[str] = []
    consumed_end = 0
    for match in _ATTR_RE.finditer(rest):
        attrs[match.group(1)] = match.group(2) if match.group(2) is not None else match.group(3)
        consumed_end = match.end()
    leftover = rest[consumed_end:].strip()
    if leftover:
        # One line for the whole fragment: "employee سارا" reads better than two.
        errors.append(
            f"ERROR: malformed attribute fragment '{leftover}' (expected key=\"value\")"
        )
    return attrs, errors


def _consume(name: str, text: str, start: int, self_closing: bool) -> int:
    """Where the tag ends: its ``[[/NAME]]`` close, or just the token.

    Consuming the whole pair is what keeps a body literal: the scanner moves
    past the close, so brackets inside the body are never parsed as tags.
    """
    if self_closing:
        return start
    close = _closing_re(name).search(text, start)
    return close.end() if close else start


def parse_reply(text: str, *, role: RoleKey | str, mode: Mode | str) -> ParsedReply:
    """Parse one model reply under the role/mode matrix.

    Unknown tags, disallowed tags and unclosed tags are dropped (their body is
    consumed so it cannot re-appear as stray tags) and reported in
    ``ParsedReply.errors`` — one line per problem, in order, for the model to
    read on its next turn.
    """
    role_key = RoleKey(role)
    role_mode = Mode(mode)
    allowed = ALLOWED_TAGS.get((role_key, role_mode), frozenset())

    tags: list[Tag] = []
    errors: list[str] = []
    prose_parts: list[str] = []
    pos = 0

    while pos < len(text):
        match = _TOKEN_RE.search(text, pos)
        if match is None:
            prose_parts.append(text[pos:])
            break
        prose_parts.append(text[pos : match.start()])
        inner = match.group(1)
        name, rest, self_closing = _split_token(inner)

        if inner.lstrip().startswith("/"):
            # A close with no open: the model echoed a tag it never opened.
            errors.append(f"ERROR: stray closing tag [[{inner.strip()}]] ignored")
            pos = match.end()
            continue

        token_end = match.end()
        try:
            tag_name = TagName(name)
        except ValueError:
            # Unknown tag: consume the pair so its body stays literal text.
            end = _consume(name, text, token_end, self_closing)
            errors.append(f"ERROR: unknown tag [[{name}]] ignored")
            pos = end
            continue

        if tag_name not in allowed:
            end = _consume(name, text, token_end, self_closing)
            errors.append(
                f"ERROR: tag [[{tag_name}]] is not allowed for "
                f"role={role_key.value} mode={role_mode.value}"
            )
            pos = end
            continue

        attrs, attr_errors = _parse_attrs(rest)
        errors.extend(attr_errors)

        body = ""
        if not self_closing:
            if tag_name not in BODY_TAGS:
                # A bodyless tag written without its "/": read it as-is.
                tags.append(
                    Tag(
                        name=tag_name,
                        attrs=attrs,
                        body="",
                        raw=text[match.start() : token_end],
                    )
                )
                pos = token_end
                continue
            close = _closing_re(name).search(text, token_end)
            if close is None:
                # Unclosed: report it and ignore the tag, body included.
                errors.append(f"ERROR: unclosed tag [[{name}]] ignored")
                pos = token_end
                continue
            body = text[token_end : close.start()]
            token_end = close.end()

        tags.append(Tag(name=tag_name, attrs=attrs, body=body, raw=text[match.start() : token_end]))
        pos = token_end

    limited, limit_errors = _apply_limits(tags)
    errors.extend(limit_errors)
    return ParsedReply(
        tags=tuple(limited), errors=tuple(errors), prose="".join(prose_parts).strip()
    )


def _apply_limits(tags: list[Tag]) -> tuple[list[Tag], list[str]]:
    """Drop tags past their per-reply limit, keeping the first ones written.

    Team changes share one budget across all three tags: "max 3 changes per
    turn" means three hires, or two hires and a firing — not three of each.
    """
    errors: list[str] = []
    kept: list[Tag] = []
    created = 0
    team_used = 0
    knowledge_used = 0
    for tag in tags:
        if tag.name is TagName.CREATE_TASK:
            created += 1
            if created > MAX_CREATE_TASKS_PER_REPLY:
                errors.append(_limit_error(tag, MAX_CREATE_TASKS_PER_REPLY))
                continue
        elif tag.name in TEAM_CHANGE_TAGS:
            team_used += 1
            if team_used > MAX_TEAM_CHANGES_PER_TURN:
                errors.append(_limit_error(tag, MAX_TEAM_CHANGES_PER_TURN))
                continue
        elif tag.name is TagName.KNOWLEDGE:
            knowledge_used += 1
            if knowledge_used > MAX_KNOWLEDGE_PER_TURN:
                errors.append(_limit_error(tag, MAX_KNOWLEDGE_PER_TURN))
                continue
        kept.append(tag)
    return kept, errors


def _limit_error(tag: Tag, limit: int) -> str:
    return (
        f"ERROR: too many [[{tag.name}]] in one reply (max {limit}); the extra was ignored"
    )


def allowed_tags(*, role: RoleKey | str, mode: Mode | str) -> frozenset[TagName]:
    return ALLOWED_TAGS.get((RoleKey(role), Mode(mode)), frozenset())


__all__ = [
    "ALLOWED_TAGS",
    "ASSIGNABLE_PIPELINES",
    "BODY_TAGS",
    "MAX_CREATE_TASKS_PER_REPLY",
    "MAX_KNOWLEDGE_PER_TURN",
    "MAX_TEAM_CHANGES_PER_TURN",
    "TEAM_CHANGE_TAGS",
    "TOOL_TAGS",
    "AssignTag",
    "EditBlock",
    "KnowledgeItem",
    "ParsedReply",
    "ProtocolError",
    "Tag",
    "TagName",
    "TaskDraft",
    "TeamChange",
    "allowed_tags",
    "assign_of",
    "edit_blocks_of",
    "knowledge_items",
    "knowledge_of",
    "parse_reply",
    "task_drafts",
    "task_draft_of",
    "team_change_of",
    "team_changes",
]