feat(engine): workforce engine core — tag protocol, pipelines, optimistic state
Phase 1 of TODO-agent-2: pure logic plus the tables it writes to. - domain/workforce/pipeline.py: Stage/PipelineKey/Outcome, the four pipelines, TRANSITIONS, next_stage, STAGE_ROLE (role+mode), resolve_pipeline with the no-reviewer fallback, task-kind defaults, stage -> coarse tasks.status, and the TaskPipelineState record (with the optimistic version). - application/engine/protocol.py: the [[TAG attr="v"]]body[[/TAG]] parser and its self-closing form, the per-role/mode allow-list, literal bodies (a bracket inside a body is not a tag), typed readers (ASSIGN/CREATE_TASK/ team changes/KNOWLEDGE/EDIT_FILE blocks) and the per-reply limits (3 tasks, 3 team changes, 3 knowledge items). - application/engine/state.py: advance() with the optimistic compare-and-set and park_dependents() so a failed dependency never deadlocks its children. - db/models_engine.py + migration a2b3c4d5e6f2: task_pipeline_state, employee_profiles, engine_turns, gate_results as side tables (tasks and employees stay locked). Linear head; alembic check reports no drift. - repositories/engine.py + 58 new tests (test_engine_protocol, test_engine_pipeline) covering every pipeline's happy path, illegal transitions, exactly-one-winner concurrent advance, dependency parking and all parser edge cases. Suite: 594 passed, ruff clean, mypy clean.
This commit is contained in:
mjalmoudy94
@@ -0,0 +1 @@
|
||||
"""The AI-workforce engine: the tag protocol and the worker-facing helpers."""
|
||||
@@ -0,0 +1,576 @@
|
||||
"""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",
|
||||
]
|
||||
@@ -0,0 +1,114 @@
|
||||
"""Moving a task between stages — the only place the stage column changes.
|
||||
|
||||
Two rules live here rather than in the scheduler:
|
||||
|
||||
1. **Optimistic moves.** ``advance`` refuses to guess: it takes the stage the
|
||||
caller believes the task is at, and the repository's compare-and-set either
|
||||
applies the move or raises :class:`~app.core.errors.ConflictError` because
|
||||
another worker got there first. A lost race is a normal outcome, not a
|
||||
crash — the caller re-reads and continues.
|
||||
2. **Dependents never get stuck.** When a task fails or is cancelled, every
|
||||
task that declared it in ``depends_on`` is parked in
|
||||
``pending_user_decision`` instead of waiting for a dependency that will
|
||||
never finish (agybot bug M1: a dead task deadlocked its children).
|
||||
|
||||
The coarse ``tasks.status`` is written by the caller in the same transaction
|
||||
(see ``engine/scheduler.py`` in a later phase): the stage is the engine's
|
||||
truth, the status column is the summary existing APIs keep reading.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from app.core.errors import ConflictError, NotFoundError
|
||||
from app.domain.workforce.pipeline import (
|
||||
SIDE_STAGES,
|
||||
TERMINAL_STAGES,
|
||||
Stage,
|
||||
TaskPipelineState,
|
||||
can_transition_stage,
|
||||
)
|
||||
from app.infrastructure.repositories.engine import (
|
||||
InvalidTransitionError,
|
||||
PipelineStateRepository,
|
||||
StaleStageError,
|
||||
UnknownTaskError,
|
||||
)
|
||||
|
||||
|
||||
async def advance(
|
||||
states: PipelineStateRepository,
|
||||
*,
|
||||
task_id: str,
|
||||
expected_from: Stage,
|
||||
to: Stage,
|
||||
) -> TaskPipelineState:
|
||||
"""Move one task from ``expected_from`` to ``to``, or fail cleanly.
|
||||
|
||||
Raises ``ConflictError`` when the transition is illegal or when another
|
||||
worker moved the task first, and ``NotFoundError`` when the task has no
|
||||
engine state row at all.
|
||||
"""
|
||||
if not can_transition_stage(expected_from, to):
|
||||
raise ConflictError(
|
||||
f"گذار نابجا از {expected_from.value} به {to.value}."
|
||||
)
|
||||
try:
|
||||
return await states.advance(task_id=task_id, expected_from=expected_from, to=to)
|
||||
except UnknownTaskError as exc:
|
||||
raise NotFoundError("این کار وضعیت موتور ندارد.") from exc
|
||||
except InvalidTransitionError as exc:
|
||||
raise ConflictError(str(exc)) from exc
|
||||
except StaleStageError as exc:
|
||||
raise ConflictError(
|
||||
"این کار هماکنون توسط موتور دیگری در حرکت است؛ دوباره تلاش کنید."
|
||||
) from exc
|
||||
|
||||
|
||||
async def park_dependents(
|
||||
states: PipelineStateRepository,
|
||||
*,
|
||||
task_id: str,
|
||||
) -> list[TaskPipelineState]:
|
||||
"""Park every dependent of a failed or cancelled task with a human.
|
||||
|
||||
A dependency that failed leaves its dependents unable to ever finish, so
|
||||
they move to ``pending_user_decision`` instead of waiting for work that
|
||||
will never happen (agybot bug M1: a dead task deadlocked its children).
|
||||
Dependents that already sit in another side state are left alone — a task
|
||||
a human already paused must not be re-parked.
|
||||
|
||||
The guard is deliberate: a task that merely *exists* is not a reason to
|
||||
park anything. Only a dependency that is actually ``failed`` or
|
||||
``cancelled`` strands its children, so the function refuses to act on a
|
||||
healthy task however often a tick calls it.
|
||||
"""
|
||||
state = await states.get(task_id)
|
||||
if state is None or state.stage not in (Stage.FAILED, Stage.CANCELLED):
|
||||
return []
|
||||
parked: list[TaskPipelineState] = []
|
||||
for dependent in await states.list_dependents(task_id):
|
||||
# A dependent that already waits on a human (paused, another decision)
|
||||
# is left alone: it is already parked, just elsewhere.
|
||||
if dependent.stage in SIDE_STAGES:
|
||||
continue
|
||||
if not can_transition_stage(dependent.stage, Stage.PENDING_USER_DECISION):
|
||||
continue
|
||||
try:
|
||||
parked.append(
|
||||
await advance(
|
||||
states,
|
||||
task_id=dependent.task_id,
|
||||
expected_from=dependent.stage,
|
||||
to=Stage.PENDING_USER_DECISION,
|
||||
)
|
||||
)
|
||||
except ConflictError:
|
||||
# A worker is already moving it; the next tick will see the truth.
|
||||
continue
|
||||
return parked
|
||||
|
||||
|
||||
async def is_terminal_state(states: PipelineStateRepository, task_id: str) -> bool:
|
||||
"""True when the task's stage can no longer move (done/cancelled/failed)."""
|
||||
state = await states.get(task_id)
|
||||
return state is None or state.stage in TERMINAL_STAGES
|
||||
@@ -0,0 +1,123 @@
|
||||
"""Domain records for the AI-workforce engine.
|
||||
|
||||
Pure data: the rows the engine writes about itself (one turn, one gate, one
|
||||
employee's model binding) plus the error vocabulary the runner classifies
|
||||
with. Kept separate from ``pipeline.py`` because that module is the state
|
||||
machine and this one is the ledger the state machine leaves behind.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from enum import StrEnum
|
||||
|
||||
from app.domain.base import new_id, utcnow
|
||||
from app.domain.workforce.models import RoleKey
|
||||
from app.domain.workforce.pipeline import Mode, Stage
|
||||
|
||||
|
||||
class ErrorKind(StrEnum):
|
||||
"""How a failed model call should be treated.
|
||||
|
||||
``rate_limit``, ``unavailable``, ``network`` and ``timeout`` are
|
||||
transient: the platform retries later instead of burning a cycle. The
|
||||
others are the caller's problem — a bad key never fixes itself.
|
||||
"""
|
||||
|
||||
RATE_LIMIT = "rate_limit"
|
||||
UNAVAILABLE = "unavailable"
|
||||
NETWORK = "network"
|
||||
TIMEOUT = "timeout"
|
||||
AUTH = "auth"
|
||||
MODEL = "model"
|
||||
PROTOCOL = "protocol"
|
||||
BUDGET = "budget"
|
||||
|
||||
|
||||
TRANSIENT_ERROR_KINDS: frozenset[ErrorKind] = frozenset(
|
||||
{
|
||||
ErrorKind.RATE_LIMIT,
|
||||
ErrorKind.UNAVAILABLE,
|
||||
ErrorKind.NETWORK,
|
||||
ErrorKind.TIMEOUT,
|
||||
}
|
||||
)
|
||||
|
||||
#: Retries a transient failure gets before the task goes to error review.
|
||||
MAX_TRANSIENT_RETRIES = 2
|
||||
|
||||
|
||||
def is_transient(kind: ErrorKind | str) -> bool:
|
||||
return ErrorKind(kind) in TRANSIENT_ERROR_KINDS
|
||||
|
||||
|
||||
class GateName(StrEnum):
|
||||
"""The gates that decide whether work may move forward."""
|
||||
|
||||
TESTS = "tests"
|
||||
DIFF_GUARD = "diff_guard"
|
||||
NO_PROGRESS = "no_progress"
|
||||
GREEN = "green"
|
||||
LINT = "lint"
|
||||
|
||||
|
||||
@dataclass(slots=True, kw_only=True)
|
||||
class EmployeeProfile:
|
||||
"""The engine's per-employee record: which model it runs on, and its roles.
|
||||
|
||||
A 1:1 side table for the same reason as ``TaskPipelineState`` — the
|
||||
``employees`` table is locked. Model identity lives here and never
|
||||
crosses an API boundary: users pick a role and a neutral tier.
|
||||
"""
|
||||
|
||||
employee_id: str
|
||||
ai_model_id: str | None = None
|
||||
effort: str = "medium"
|
||||
#: CSV of role keys the employee may be hired for (its own role plus extras).
|
||||
roles: list[str] = field(default_factory=list)
|
||||
hired_by: str = "user" # user | manager
|
||||
busy_task_id: str | None = None
|
||||
updated_at: datetime = field(default_factory=utcnow)
|
||||
|
||||
def role_list(self) -> list[str]:
|
||||
return [r.strip() for r in self.roles if r.strip()]
|
||||
|
||||
|
||||
@dataclass(slots=True, kw_only=True)
|
||||
class EngineTurn:
|
||||
"""One employee turn: who ran, what came back, what it cost.
|
||||
|
||||
``reply_excerpt`` is bounded (4,000 chars) on purpose — the full reply
|
||||
can be enormous, and this table exists to answer "what happened", not to
|
||||
archive the conversation.
|
||||
"""
|
||||
|
||||
task_id: str
|
||||
run_id: str
|
||||
employee_id: str | None = None
|
||||
role: RoleKey = RoleKey.DEVELOPER
|
||||
mode: Mode = Mode.WORK
|
||||
stage: Stage = Stage.PENDING_IMPLEMENTATION
|
||||
prompt_chars: int = 0
|
||||
reply_excerpt: str = ""
|
||||
#: tag name -> count, so the UI can show "3 files edited" without a join.
|
||||
tags: dict = field(default_factory=dict)
|
||||
tool_calls: int = 0
|
||||
ok: bool = True
|
||||
error_kind: str = ""
|
||||
id: str = field(default_factory=new_id)
|
||||
started_at: datetime = field(default_factory=utcnow)
|
||||
finished_at: datetime | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True, kw_only=True)
|
||||
class GateResult:
|
||||
"""One gate verdict, with the detail a human needs to trust it."""
|
||||
|
||||
task_id: str
|
||||
gate: GateName
|
||||
ok: bool
|
||||
detail: str = ""
|
||||
id: str = field(default_factory=new_id)
|
||||
created_at: datetime = field(default_factory=utcnow)
|
||||
@@ -0,0 +1,494 @@
|
||||
"""Dynamic task pipelines: stages, transitions and the manager's choices.
|
||||
|
||||
Pure logic: no I/O, no framework, no database. Everything here is a decision
|
||||
the engine makes *before* it touches a row, which is what makes the state
|
||||
machine testable without a session.
|
||||
|
||||
The vocabulary comes from agybot (concepts only, never its code): one task
|
||||
walks an ordered list of stages, every stage needs exactly one role in one
|
||||
mode, and the manager decides which pipeline a task gets. The differences
|
||||
that matter here are that every move is optimistic (``version`` in
|
||||
``TaskPipelineState``) and that a failed or cancelled dependency parks its
|
||||
dependents instead of leaving them stuck forever (agybot bug M1).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from enum import StrEnum
|
||||
from typing import NamedTuple
|
||||
|
||||
from app.domain.base import utcnow
|
||||
from app.domain.workforce.models import RoleKey, TaskStatus
|
||||
|
||||
|
||||
class Stage(StrEnum):
|
||||
"""Where a task is in its pipeline.
|
||||
|
||||
The first eight values are pipeline stages (they appear in
|
||||
:data:`PIPELINES`); the rest are side states a task can be pushed into
|
||||
from anywhere, which is why they are not part of any pipeline list.
|
||||
"""
|
||||
|
||||
PENDING_MANAGER_REVIEW = "pending_manager_review"
|
||||
PENDING_ANALYSIS = "pending_analysis"
|
||||
PENDING_IMPLEMENTATION = "pending_implementation"
|
||||
PENDING_REVIEW = "pending_review"
|
||||
PENDING_TEST = "pending_test"
|
||||
PENDING_FINAL_REVIEW = "pending_final_review"
|
||||
MERGE_REQUEST = "merge_request"
|
||||
COMPLETED = "completed"
|
||||
PENDING_USER_DECISION = "pending_user_decision"
|
||||
ERROR_REVIEW = "error_review"
|
||||
PAUSED = "paused"
|
||||
CANCELLED = "cancelled"
|
||||
FAILED = "failed"
|
||||
|
||||
|
||||
#: Side states and terminals: a task in one of these is not walking a pipeline.
|
||||
SIDE_STAGES: frozenset[Stage] = frozenset(
|
||||
{
|
||||
Stage.PENDING_USER_DECISION,
|
||||
Stage.ERROR_REVIEW,
|
||||
Stage.PAUSED,
|
||||
Stage.CANCELLED,
|
||||
Stage.FAILED,
|
||||
}
|
||||
)
|
||||
TERMINAL_STAGES: frozenset[Stage] = frozenset(
|
||||
{Stage.COMPLETED, Stage.CANCELLED, Stage.FAILED}
|
||||
)
|
||||
|
||||
|
||||
class PipelineKey(StrEnum):
|
||||
"""The four pipelines. Names match the plan (agybot's, minus its extras)."""
|
||||
|
||||
PLAN_CODE_TEST = "plan_code_test"
|
||||
PLAN_CODE_REVIEW_TEST = "plan_code_review_test"
|
||||
CODE_TEST = "code_test"
|
||||
ANALYSIS_ONLY = "analysis_only"
|
||||
|
||||
|
||||
class TaskKind(StrEnum):
|
||||
FEATURE = "feature"
|
||||
BUGFIX = "bugfix"
|
||||
REFACTOR = "refactor"
|
||||
DOCS = "docs"
|
||||
ANALYSIS = "analysis"
|
||||
|
||||
|
||||
class Risk(StrEnum):
|
||||
LOW = "low"
|
||||
MEDIUM = "medium"
|
||||
HIGH = "high"
|
||||
|
||||
|
||||
class Mode(StrEnum):
|
||||
"""The hat a role wears in one turn.
|
||||
|
||||
Only the manager has more than one: the same employee triages, gives the
|
||||
final verdict and handles errors, and each mode allows a different tag set
|
||||
(see ``engine/protocol.py``). ``work`` is what every other role uses.
|
||||
"""
|
||||
|
||||
TRIAGE = "triage"
|
||||
FINAL = "final"
|
||||
ERROR = "error"
|
||||
CHAT = "chat"
|
||||
WORK = "work"
|
||||
|
||||
|
||||
class Outcome(StrEnum):
|
||||
"""The verdict of a finished turn, as the pipeline understands it."""
|
||||
|
||||
#: The turn produced its tag: move to the next pipeline stage.
|
||||
OK = "ok"
|
||||
#: Reviewer/tester/manager sent the work back: cycle again.
|
||||
REVISE = "revise"
|
||||
#: The manager rejected the task outright.
|
||||
REJECT = "reject"
|
||||
#: The turn needs a human answer before it can continue.
|
||||
ASK_USER = "ask_user"
|
||||
#: Something failed hard enough for the manager's error turn.
|
||||
ERROR = "error"
|
||||
#: The error turn decided to retry with a hint.
|
||||
RESOLVE = "resolve"
|
||||
|
||||
|
||||
#: Ordered stages per pipeline. ``completed`` is part of every list so that
|
||||
#: "advance past the last stage" is just an index lookup, not a special case.
|
||||
PIPELINES: dict[PipelineKey, tuple[Stage, ...]] = {
|
||||
PipelineKey.PLAN_CODE_TEST: (
|
||||
Stage.PENDING_MANAGER_REVIEW,
|
||||
Stage.PENDING_ANALYSIS,
|
||||
Stage.PENDING_IMPLEMENTATION,
|
||||
Stage.PENDING_TEST,
|
||||
Stage.PENDING_FINAL_REVIEW,
|
||||
Stage.MERGE_REQUEST,
|
||||
Stage.COMPLETED,
|
||||
),
|
||||
PipelineKey.PLAN_CODE_REVIEW_TEST: (
|
||||
Stage.PENDING_MANAGER_REVIEW,
|
||||
Stage.PENDING_ANALYSIS,
|
||||
Stage.PENDING_IMPLEMENTATION,
|
||||
Stage.PENDING_REVIEW,
|
||||
Stage.PENDING_TEST,
|
||||
Stage.PENDING_FINAL_REVIEW,
|
||||
Stage.MERGE_REQUEST,
|
||||
Stage.COMPLETED,
|
||||
),
|
||||
PipelineKey.CODE_TEST: (
|
||||
Stage.PENDING_MANAGER_REVIEW,
|
||||
Stage.PENDING_IMPLEMENTATION,
|
||||
Stage.PENDING_TEST,
|
||||
Stage.PENDING_FINAL_REVIEW,
|
||||
Stage.MERGE_REQUEST,
|
||||
Stage.COMPLETED,
|
||||
),
|
||||
#: No worktree, no merge: the planner answers and the task is done.
|
||||
PipelineKey.ANALYSIS_ONLY: (
|
||||
Stage.PENDING_MANAGER_REVIEW,
|
||||
Stage.PENDING_ANALYSIS,
|
||||
Stage.COMPLETED,
|
||||
),
|
||||
}
|
||||
|
||||
#: The pipeline used when nobody expressed a preference.
|
||||
APP_DEFAULT_PIPELINE: PipelineKey = PipelineKey.PLAN_CODE_TEST
|
||||
|
||||
#: Default pipeline and risk per task kind. The manager's choice wins over both.
|
||||
TASK_KIND_DEFAULTS: dict[TaskKind, tuple[PipelineKey, Risk]] = {
|
||||
TaskKind.FEATURE: (PipelineKey.PLAN_CODE_TEST, Risk.MEDIUM),
|
||||
TaskKind.BUGFIX: (PipelineKey.CODE_TEST, Risk.LOW),
|
||||
TaskKind.REFACTOR: (PipelineKey.PLAN_CODE_REVIEW_TEST, Risk.HIGH),
|
||||
TaskKind.DOCS: (PipelineKey.CODE_TEST, Risk.LOW),
|
||||
TaskKind.ANALYSIS: (PipelineKey.ANALYSIS_ONLY, Risk.LOW),
|
||||
}
|
||||
|
||||
#: Legal stage moves. Anything not listed here is a bug in the caller, and
|
||||
#: ``engine/state.py`` refuses it before it reaches the database.
|
||||
TRANSITIONS: dict[Stage, frozenset[Stage]] = {
|
||||
Stage.PENDING_MANAGER_REVIEW: frozenset(
|
||||
{
|
||||
Stage.PENDING_ANALYSIS,
|
||||
Stage.PENDING_IMPLEMENTATION,
|
||||
Stage.PENDING_USER_DECISION,
|
||||
Stage.ERROR_REVIEW,
|
||||
Stage.PAUSED,
|
||||
Stage.CANCELLED,
|
||||
}
|
||||
),
|
||||
Stage.PENDING_ANALYSIS: frozenset(
|
||||
{
|
||||
Stage.PENDING_IMPLEMENTATION,
|
||||
# analysis_only ends here; the planner's answer *is* the deliverable.
|
||||
Stage.COMPLETED,
|
||||
Stage.PENDING_USER_DECISION,
|
||||
Stage.ERROR_REVIEW,
|
||||
Stage.PAUSED,
|
||||
Stage.CANCELLED,
|
||||
}
|
||||
),
|
||||
Stage.PENDING_IMPLEMENTATION: frozenset(
|
||||
{
|
||||
Stage.PENDING_REVIEW,
|
||||
Stage.PENDING_TEST,
|
||||
Stage.PENDING_FINAL_REVIEW,
|
||||
Stage.PENDING_USER_DECISION,
|
||||
Stage.ERROR_REVIEW,
|
||||
Stage.PAUSED,
|
||||
Stage.CANCELLED,
|
||||
}
|
||||
),
|
||||
Stage.PENDING_REVIEW: frozenset(
|
||||
{
|
||||
Stage.PENDING_TEST,
|
||||
Stage.PENDING_IMPLEMENTATION,
|
||||
Stage.PENDING_USER_DECISION,
|
||||
Stage.ERROR_REVIEW,
|
||||
Stage.PAUSED,
|
||||
Stage.CANCELLED,
|
||||
}
|
||||
),
|
||||
Stage.PENDING_TEST: frozenset(
|
||||
{
|
||||
Stage.PENDING_FINAL_REVIEW,
|
||||
Stage.PENDING_IMPLEMENTATION,
|
||||
Stage.PENDING_USER_DECISION,
|
||||
Stage.ERROR_REVIEW,
|
||||
Stage.PAUSED,
|
||||
Stage.CANCELLED,
|
||||
}
|
||||
),
|
||||
Stage.PENDING_FINAL_REVIEW: frozenset(
|
||||
{
|
||||
Stage.MERGE_REQUEST,
|
||||
Stage.PENDING_IMPLEMENTATION,
|
||||
Stage.PENDING_USER_DECISION,
|
||||
Stage.ERROR_REVIEW,
|
||||
Stage.PAUSED,
|
||||
Stage.CANCELLED,
|
||||
}
|
||||
),
|
||||
#: merge_request is executed by the platform, never by an AI turn: git
|
||||
#: lives behind RepoWorkspacePort (PLAN-v2 D6).
|
||||
Stage.MERGE_REQUEST: frozenset(
|
||||
{
|
||||
Stage.COMPLETED,
|
||||
Stage.PENDING_IMPLEMENTATION,
|
||||
Stage.ERROR_REVIEW,
|
||||
Stage.PAUSED,
|
||||
Stage.CANCELLED,
|
||||
}
|
||||
),
|
||||
Stage.PENDING_USER_DECISION: frozenset(
|
||||
{
|
||||
# The user answered, so the task goes back to triage: a manager
|
||||
# decides the pipeline again with the new information.
|
||||
Stage.PENDING_MANAGER_REVIEW,
|
||||
Stage.ERROR_REVIEW,
|
||||
Stage.PAUSED,
|
||||
Stage.CANCELLED,
|
||||
}
|
||||
),
|
||||
Stage.ERROR_REVIEW: frozenset(
|
||||
{
|
||||
Stage.PENDING_ANALYSIS,
|
||||
Stage.PENDING_IMPLEMENTATION,
|
||||
Stage.PENDING_USER_DECISION,
|
||||
Stage.FAILED,
|
||||
Stage.PAUSED,
|
||||
Stage.CANCELLED,
|
||||
}
|
||||
),
|
||||
#: A paused task resumes through triage as well: the manager re-reads the
|
||||
# task before spending more budget on it.
|
||||
Stage.PAUSED: frozenset(
|
||||
{Stage.PENDING_MANAGER_REVIEW, Stage.ERROR_REVIEW, Stage.CANCELLED}
|
||||
),
|
||||
Stage.COMPLETED: frozenset(),
|
||||
Stage.CANCELLED: frozenset(),
|
||||
Stage.FAILED: frozenset(),
|
||||
}
|
||||
|
||||
|
||||
class StageActor(NamedTuple):
|
||||
"""Who acts in a stage: a role and the mode it must use."""
|
||||
|
||||
role: RoleKey
|
||||
mode: Mode
|
||||
|
||||
|
||||
#: stage -> (role, mode). ``None`` means "no AI turn acts here": the stage is
|
||||
#: either terminal, platform-executed (merge_request) or waiting on a human.
|
||||
STAGE_ROLE: dict[Stage, StageActor | None] = {
|
||||
Stage.PENDING_MANAGER_REVIEW: StageActor(RoleKey.MANAGER, Mode.TRIAGE),
|
||||
Stage.PENDING_ANALYSIS: StageActor(RoleKey.PLANNER, Mode.WORK),
|
||||
Stage.PENDING_IMPLEMENTATION: StageActor(RoleKey.DEVELOPER, Mode.WORK),
|
||||
Stage.PENDING_REVIEW: StageActor(RoleKey.REVIEWER, Mode.WORK),
|
||||
Stage.PENDING_TEST: StageActor(RoleKey.TESTER, Mode.WORK),
|
||||
Stage.PENDING_FINAL_REVIEW: StageActor(RoleKey.MANAGER, Mode.FINAL),
|
||||
Stage.MERGE_REQUEST: None,
|
||||
Stage.COMPLETED: None,
|
||||
Stage.PENDING_USER_DECISION: None,
|
||||
Stage.ERROR_REVIEW: StageActor(RoleKey.MANAGER, Mode.ERROR),
|
||||
Stage.PAUSED: None,
|
||||
Stage.CANCELLED: None,
|
||||
Stage.FAILED: None,
|
||||
}
|
||||
|
||||
#: The coarse ``tasks.status`` each stage maps onto. Existing APIs and the old
|
||||
#: worker read that column, so the engine keeps it truthful at every move
|
||||
#: instead of introducing a second source of truth.
|
||||
STAGE_STATUS: dict[Stage, TaskStatus] = {
|
||||
Stage.PENDING_MANAGER_REVIEW: TaskStatus.RUNNING,
|
||||
Stage.PENDING_ANALYSIS: TaskStatus.RUNNING,
|
||||
Stage.PENDING_IMPLEMENTATION: TaskStatus.RUNNING,
|
||||
Stage.PENDING_REVIEW: TaskStatus.RUNNING,
|
||||
Stage.PENDING_TEST: TaskStatus.RUNNING,
|
||||
Stage.PENDING_FINAL_REVIEW: TaskStatus.RUNNING,
|
||||
Stage.MERGE_REQUEST: TaskStatus.RUNNING,
|
||||
Stage.COMPLETED: TaskStatus.DONE,
|
||||
Stage.PENDING_USER_DECISION: TaskStatus.WAITING_APPROVAL,
|
||||
Stage.ERROR_REVIEW: TaskStatus.RUNNING,
|
||||
Stage.PAUSED: TaskStatus.WAITING_APPROVAL,
|
||||
Stage.CANCELLED: TaskStatus.CANCELLED,
|
||||
Stage.FAILED: TaskStatus.FAILED,
|
||||
}
|
||||
|
||||
#: Fixed targets for outcomes that are not "the next stage": work that is sent
|
||||
#: back always returns to the developer, because that is the role that can
|
||||
#: actually change the diff.
|
||||
_OUTCOME_TARGET: dict[Outcome, Stage] = {
|
||||
Outcome.REVISE: Stage.PENDING_IMPLEMENTATION,
|
||||
Outcome.RESOLVE: Stage.PENDING_IMPLEMENTATION,
|
||||
Outcome.ASK_USER: Stage.PENDING_USER_DECISION,
|
||||
Outcome.ERROR: Stage.ERROR_REVIEW,
|
||||
Outcome.REJECT: Stage.CANCELLED,
|
||||
}
|
||||
|
||||
|
||||
def is_terminal(stage: Stage) -> bool:
|
||||
return stage in TERMINAL_STAGES
|
||||
|
||||
|
||||
def is_side_stage(stage: Stage) -> bool:
|
||||
"""True when the stage is not part of any pipeline's ordered list."""
|
||||
return stage in SIDE_STAGES
|
||||
|
||||
|
||||
def can_transition_stage(current: Stage, target: Stage) -> bool:
|
||||
return target in TRANSITIONS.get(current, frozenset())
|
||||
|
||||
|
||||
def pipeline_stages(pipeline: PipelineKey | str) -> tuple[Stage, ...]:
|
||||
"""The ordered stages of a pipeline (including the terminal)."""
|
||||
return PIPELINES[PipelineKey(pipeline)]
|
||||
|
||||
|
||||
def requires_worktree(pipeline: PipelineKey | str) -> bool:
|
||||
"""``analysis_only`` never creates a worktree and never merges."""
|
||||
return PipelineKey(pipeline) is not PipelineKey.ANALYSIS_ONLY
|
||||
|
||||
|
||||
def next_stage(pipeline: PipelineKey | str, current: Stage | str, outcome: Outcome | str) -> Stage:
|
||||
"""The stage a task moves to after a turn ends.
|
||||
|
||||
``Outcome.OK`` walks the pipeline list one step forward (past the last
|
||||
stage there is nothing left, which only happens for a malformed call);
|
||||
every other outcome has a fixed target. The result is *not* checked
|
||||
against :data:`TRANSITIONS` here — callers go through
|
||||
``engine/state.py``, which refuses illegal moves.
|
||||
"""
|
||||
stage = Stage(current)
|
||||
verdict = Outcome(outcome)
|
||||
if verdict is Outcome.OK:
|
||||
stages = PIPELINES[PipelineKey(pipeline)]
|
||||
index = stages.index(stage)
|
||||
return stages[index + 1] if index + 1 < len(stages) else stage
|
||||
try:
|
||||
return _OUTCOME_TARGET[verdict]
|
||||
except KeyError as exc: # pragma: no cover - the enum is closed
|
||||
raise ValueError(f"نتیجهٔ ناشناخته برای مرحله: {verdict}") from exc
|
||||
|
||||
|
||||
def resolve_pipeline(
|
||||
*,
|
||||
manager_choice: PipelineKey | str | None = None,
|
||||
task_kind: TaskKind | str | None = None,
|
||||
app_default: PipelineKey | str | None = None,
|
||||
has_active_reviewer: bool = True,
|
||||
) -> PipelineKey:
|
||||
"""Decide which pipeline a task gets.
|
||||
|
||||
Priority: the manager's explicit choice, then the task kind's default,
|
||||
then the app profile's default, then ``plan_code_test``.
|
||||
|
||||
A review pipeline without an active reviewer falls back to
|
||||
``plan_code_test`` — including an explicit manager choice, because a
|
||||
review stage nobody can act on is a dead end. The manager who really
|
||||
wants a reviewer should hire one (``HIRE role="reviewer"``), not start a
|
||||
task that will wait forever.
|
||||
"""
|
||||
chosen: PipelineKey | None = None
|
||||
if manager_choice:
|
||||
chosen = PipelineKey(manager_choice)
|
||||
elif task_kind:
|
||||
chosen = TASK_KIND_DEFAULTS[TaskKind(task_kind)][0]
|
||||
elif app_default:
|
||||
chosen = PipelineKey(app_default)
|
||||
if chosen is None:
|
||||
chosen = APP_DEFAULT_PIPELINE
|
||||
if chosen is PipelineKey.PLAN_CODE_REVIEW_TEST and not has_active_reviewer:
|
||||
return PipelineKey.PLAN_CODE_TEST
|
||||
return chosen
|
||||
|
||||
|
||||
def risk_for(task_kind: TaskKind | str) -> Risk:
|
||||
"""The risk a task kind carries when the manager did not set one."""
|
||||
return TASK_KIND_DEFAULTS[TaskKind(task_kind)][1]
|
||||
|
||||
|
||||
def coarse_status(stage: Stage | str) -> TaskStatus:
|
||||
return STAGE_STATUS[Stage(stage)]
|
||||
|
||||
|
||||
@dataclass(slots=True, kw_only=True)
|
||||
class TaskPipelineState:
|
||||
"""The engine's per-task record: pipeline, position and loop guards.
|
||||
|
||||
Lives in a 1:1 side table (``task_pipeline_state``) because ``tasks`` is
|
||||
a locked table: the engine must not add columns to it.
|
||||
|
||||
``version`` is the optimistic lock. ``advance`` only writes when the
|
||||
row still has the version the caller read, so two workers moving the
|
||||
same task cannot both win.
|
||||
"""
|
||||
|
||||
task_id: str
|
||||
organization_id: str
|
||||
project_id: str | None = None
|
||||
pipeline: PipelineKey = APP_DEFAULT_PIPELINE
|
||||
kind: TaskKind = TaskKind.FEATURE
|
||||
risk: Risk = Risk.MEDIUM
|
||||
stage: Stage = Stage.PENDING_MANAGER_REVIEW
|
||||
cycle: int = 0
|
||||
max_cycles: int = 6
|
||||
branch: str = ""
|
||||
#: Hash of the last diff, for the no-progress gate.
|
||||
last_diff_hash: str = ""
|
||||
same_diff_count: int = 0
|
||||
retry_count: int = 0
|
||||
waiver_reason: str = ""
|
||||
parent_task_id: str | None = None
|
||||
depends_on: list[str] = field(default_factory=list)
|
||||
version: int = 0
|
||||
updated_at: datetime = field(default_factory=utcnow)
|
||||
|
||||
@property
|
||||
def status(self) -> TaskStatus:
|
||||
return coarse_status(self.stage)
|
||||
|
||||
@property
|
||||
def needs_turn(self) -> bool:
|
||||
"""True when the stage is waiting for an AI turn."""
|
||||
actor = STAGE_ROLE.get(self.stage)
|
||||
return actor is not None and not is_terminal(self.stage)
|
||||
|
||||
@property
|
||||
def is_blocked(self) -> bool:
|
||||
"""True while the task waits on a human (decision, pause or budget)."""
|
||||
return self.stage in SIDE_STAGES and not is_terminal(self.stage)
|
||||
|
||||
def actor(self) -> StageActor | None:
|
||||
return STAGE_ROLE.get(self.stage)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"APP_DEFAULT_PIPELINE",
|
||||
"PIPELINES",
|
||||
"SIDE_STAGES",
|
||||
"STAGE_ROLE",
|
||||
"STAGE_STATUS",
|
||||
"TASK_KIND_DEFAULTS",
|
||||
"TERMINAL_STAGES",
|
||||
"TRANSITIONS",
|
||||
"Mode",
|
||||
"Outcome",
|
||||
"PipelineKey",
|
||||
"Risk",
|
||||
"Stage",
|
||||
"StageActor",
|
||||
"TaskKind",
|
||||
"TaskPipelineState",
|
||||
"can_transition_stage",
|
||||
"coarse_status",
|
||||
"is_side_stage",
|
||||
"is_terminal",
|
||||
"next_stage",
|
||||
"pipeline_stages",
|
||||
"requires_worktree",
|
||||
"resolve_pipeline",
|
||||
"risk_for",
|
||||
]
|
||||
@@ -5,6 +5,170 @@ Use ``from app.core.database import Base`` and the helpers in
|
||||
``app.infrastructure.db.models`` here: that module imports this one (import cycle).
|
||||
Foreign keys to existing tables are written as plain strings, e.g.
|
||||
``ForeignKey("organizations.id", ondelete="CASCADE")``. See docs/plan/PLAN-v2.md §6.
|
||||
|
||||
Four tables, deliberately split so ``tasks`` and ``employees`` stay untouched
|
||||
(they are locked):
|
||||
|
||||
* ``task_pipeline_state`` — the engine's own view of a task: pipeline, stage,
|
||||
cycle counters, branch and the optimistic ``version``. One row per task.
|
||||
* ``employee_profiles`` — which granted model an employee runs on, and how
|
||||
busy it is. Model identity never leaves the server.
|
||||
* ``engine_turns`` — one row per agent turn: the audit trail of who thought
|
||||
what, with a bounded reply excerpt.
|
||||
* ``gate_results`` — one row per gate verdict (tests, diff guard, no-progress,
|
||||
green), with the detail that justifies it.
|
||||
|
||||
Both side tables cascade from their parent, so deleting a task or an employee
|
||||
never leaves an orphaned engine row behind.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import (
|
||||
JSON,
|
||||
Boolean,
|
||||
DateTime,
|
||||
ForeignKey,
|
||||
Index,
|
||||
Integer,
|
||||
String,
|
||||
Text,
|
||||
)
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from app.core.database import Base
|
||||
from app.domain.base import utcnow
|
||||
from app.infrastructure.db.columns import created, pk
|
||||
|
||||
#: The engine's default limits, mirroring the settings block in ``config.py``.
|
||||
#: They are column defaults (not settings) so that a database written by an
|
||||
#: older deployment still carries a sane value.
|
||||
DEFAULT_MAX_CYCLES = 6
|
||||
|
||||
|
||||
class TaskPipelineStateRow(Base):
|
||||
__tablename__ = "task_pipeline_state"
|
||||
|
||||
task_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("tasks.id", ondelete="CASCADE"), primary_key=True
|
||||
)
|
||||
#: Denormalised from the task so the tenant filter stays a single-table
|
||||
#: WHERE: the scheduler's hot query never joins back to ``tasks``.
|
||||
organization_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("organizations.id", ondelete="CASCADE"), nullable=False, index=True
|
||||
)
|
||||
project_id: Mapped[str | None] = mapped_column(
|
||||
ForeignKey("projects.id", ondelete="CASCADE"), nullable=True
|
||||
)
|
||||
pipeline: Mapped[str] = mapped_column(
|
||||
String(32), default="plan_code_test", nullable=False
|
||||
)
|
||||
kind: Mapped[str] = mapped_column(String(16), default="feature", nullable=False)
|
||||
risk: Mapped[str] = mapped_column(String(8), default="medium", nullable=False)
|
||||
#: Indexed: every scheduler tick asks "which tasks are at a turn stage".
|
||||
stage: Mapped[str] = mapped_column(
|
||||
String(32), default="pending_manager_review", nullable=False, index=True
|
||||
)
|
||||
cycle: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
|
||||
max_cycles: Mapped[int] = mapped_column(
|
||||
Integer, default=DEFAULT_MAX_CYCLES, nullable=False
|
||||
)
|
||||
branch: Mapped[str] = mapped_column(String(120), default="", nullable=False)
|
||||
last_diff_hash: Mapped[str] = mapped_column(String(64), default="", nullable=False)
|
||||
same_diff_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
|
||||
retry_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
|
||||
#: A waived test pass must always be explainable: the reason is stored,
|
||||
#: never inferred (PLAN-v2 §3, green gate).
|
||||
waiver_reason: Mapped[str] = mapped_column(Text, default="", nullable=False)
|
||||
parent_task_id: Mapped[str | None] = mapped_column(
|
||||
ForeignKey("tasks.id", ondelete="SET NULL"), nullable=True
|
||||
)
|
||||
depends_on: Mapped[list] = mapped_column(JSON, default=list, nullable=False)
|
||||
#: Optimistic lock. Every stage move bumps it; a stale write is refused.
|
||||
version: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), default=utcnow, onupdate=utcnow, nullable=False
|
||||
)
|
||||
|
||||
|
||||
class EmployeeProfileRow(Base):
|
||||
__tablename__ = "employee_profiles"
|
||||
__table_args__ = (Index("ix_employee_profile_busy", "busy_task_id"),)
|
||||
|
||||
employee_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("employees.id", ondelete="CASCADE"), primary_key=True
|
||||
)
|
||||
#: A granted model id (``org_model_grants.ai_model_id``). Nullable until the
|
||||
#: owner grants something: an employee exists but is inactive, which is a
|
||||
#: state the UI shows instead of hiding the employee.
|
||||
ai_model_id: Mapped[str | None] = mapped_column(String(32), nullable=True)
|
||||
effort: Mapped[str] = mapped_column(String(16), default="medium", nullable=False)
|
||||
roles: Mapped[str] = mapped_column(String(255), default="", nullable=False)
|
||||
hired_by: Mapped[str] = mapped_column(String(8), default="user", nullable=False)
|
||||
busy_task_id: Mapped[str | None] = mapped_column(
|
||||
ForeignKey("tasks.id", ondelete="SET NULL"), nullable=True
|
||||
)
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), default=utcnow, onupdate=utcnow, nullable=False
|
||||
)
|
||||
|
||||
|
||||
class EngineTurnRow(Base):
|
||||
__tablename__ = "engine_turns"
|
||||
__table_args__ = (
|
||||
Index("ix_engine_turn_task_stage", "task_id", "stage"),
|
||||
Index("ix_engine_turn_run", "run_id"),
|
||||
)
|
||||
|
||||
id: Mapped[str] = pk()
|
||||
task_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("tasks.id", ondelete="CASCADE"), nullable=False
|
||||
)
|
||||
run_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("runs.id", ondelete="CASCADE"), nullable=False
|
||||
)
|
||||
employee_id: Mapped[str | None] = mapped_column(
|
||||
ForeignKey("employees.id", ondelete="SET NULL"), nullable=True
|
||||
)
|
||||
role: Mapped[str] = mapped_column(String(24), default="developer", nullable=False)
|
||||
mode: Mapped[str] = mapped_column(String(16), default="work", nullable=False)
|
||||
stage: Mapped[str] = mapped_column(String(32), default="", nullable=False)
|
||||
prompt_chars: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
|
||||
#: Bounded on write: a reply is unbounded, a row is not.
|
||||
reply_excerpt: Mapped[str] = mapped_column(Text, default="", nullable=False)
|
||||
tags: Mapped[dict] = mapped_column(JSON, default=dict, nullable=False)
|
||||
tool_calls: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
|
||||
ok: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
|
||||
error_kind: Mapped[str] = mapped_column(String(16), default="", nullable=False)
|
||||
started_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), default=utcnow, nullable=False
|
||||
)
|
||||
finished_at: Mapped[datetime | None] = mapped_column(
|
||||
DateTime(timezone=True), nullable=True
|
||||
)
|
||||
created_at: Mapped[datetime] = created()
|
||||
|
||||
|
||||
class GateResultRow(Base):
|
||||
__tablename__ = "gate_results"
|
||||
__table_args__ = (Index("ix_gate_result_task_gate", "task_id", "gate"),)
|
||||
|
||||
id: Mapped[str] = pk()
|
||||
task_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("tasks.id", ondelete="CASCADE"), nullable=False
|
||||
)
|
||||
gate: Mapped[str] = mapped_column(String(16), default="tests", nullable=False)
|
||||
ok: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
|
||||
detail: Mapped[str] = mapped_column(Text, default="", nullable=False)
|
||||
created_at: Mapped[datetime] = created()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_MAX_CYCLES",
|
||||
"EmployeeProfileRow",
|
||||
"EngineTurnRow",
|
||||
"GateResultRow",
|
||||
"TaskPipelineStateRow",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,475 @@
|
||||
"""Repositories for the AI-workforce engine's own tables.
|
||||
|
||||
Four repositories, one per engine table, all following the house rule: the
|
||||
ORM row never leaves this module — the domain dataclasses in
|
||||
``app.domain.workforce.pipeline`` / ``engine`` are what the application layer
|
||||
holds.
|
||||
|
||||
The interesting one is :class:`PipelineStateRepository`. Its ``advance`` is an
|
||||
optimistic compare-and-set: the ``UPDATE`` carries the caller's expected stage
|
||||
and version in its ``WHERE`` clause, so two workers moving the same task
|
||||
cannot both win — the loser sees zero updated rows and gets a
|
||||
:class:`StaleStageError` instead of silently overwriting the winner.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy import func, select, update
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.domain.base import utcnow
|
||||
from app.domain.workforce.engine import (
|
||||
EmployeeProfile,
|
||||
EngineTurn,
|
||||
GateName,
|
||||
GateResult,
|
||||
)
|
||||
from app.domain.workforce.models import RoleKey
|
||||
from app.domain.workforce.pipeline import (
|
||||
TRANSITIONS,
|
||||
Mode,
|
||||
PipelineKey,
|
||||
Risk,
|
||||
Stage,
|
||||
TaskKind,
|
||||
TaskPipelineState,
|
||||
)
|
||||
from app.infrastructure.db.models_engine import (
|
||||
EmployeeProfileRow,
|
||||
EngineTurnRow,
|
||||
GateResultRow,
|
||||
TaskPipelineStateRow,
|
||||
)
|
||||
|
||||
|
||||
class StaleStageError(Exception):
|
||||
"""Another worker moved the task first; the caller must re-read it."""
|
||||
|
||||
|
||||
class UnknownTaskError(LookupError):
|
||||
"""The task has no pipeline state row (it was never registered)."""
|
||||
|
||||
|
||||
class InvalidTransitionError(ValueError):
|
||||
"""The requested move is not in the stage transition table."""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# task_pipeline_state
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
#: Columns ``patch`` may touch. An allow-list, so a typo can never become SQL.
|
||||
_PATCHABLE = frozenset(
|
||||
{
|
||||
"pipeline",
|
||||
"kind",
|
||||
"risk",
|
||||
"stage",
|
||||
"cycle",
|
||||
"max_cycles",
|
||||
"branch",
|
||||
"last_diff_hash",
|
||||
"same_diff_count",
|
||||
"retry_count",
|
||||
"waiver_reason",
|
||||
"parent_task_id",
|
||||
"depends_on",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _state(row: TaskPipelineStateRow) -> TaskPipelineState:
|
||||
return TaskPipelineState(
|
||||
task_id=row.task_id,
|
||||
organization_id=row.organization_id,
|
||||
project_id=row.project_id,
|
||||
pipeline=PipelineKey(row.pipeline),
|
||||
kind=TaskKind(row.kind),
|
||||
risk=Risk(row.risk),
|
||||
stage=Stage(row.stage),
|
||||
cycle=row.cycle,
|
||||
max_cycles=row.max_cycles,
|
||||
branch=row.branch,
|
||||
last_diff_hash=row.last_diff_hash,
|
||||
same_diff_count=row.same_diff_count,
|
||||
retry_count=row.retry_count,
|
||||
waiver_reason=row.waiver_reason,
|
||||
parent_task_id=row.parent_task_id,
|
||||
depends_on=list(row.depends_on or []),
|
||||
version=row.version,
|
||||
updated_at=row.updated_at,
|
||||
)
|
||||
|
||||
|
||||
def _legal_targets(stage: Stage) -> frozenset[Stage]:
|
||||
"""The stage moves the domain allows from here."""
|
||||
return TRANSITIONS.get(stage, frozenset())
|
||||
|
||||
|
||||
class PipelineStateRepository:
|
||||
"""1:1 with a task: pipeline, stage, cycle guards and the version lock."""
|
||||
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get(self, task_id: str) -> TaskPipelineState | None:
|
||||
row = await self.session.get(TaskPipelineStateRow, task_id)
|
||||
return _state(row) if row else None
|
||||
|
||||
async def get_or_raise(self, task_id: str) -> TaskPipelineState:
|
||||
state = await self.get(task_id)
|
||||
if state is None:
|
||||
raise UnknownTaskError(task_id)
|
||||
return state
|
||||
|
||||
async def list_for_org(
|
||||
self, organization_id: str, *, limit: int = 500
|
||||
) -> list[TaskPipelineState]:
|
||||
stmt = (
|
||||
select(TaskPipelineStateRow)
|
||||
.where(TaskPipelineStateRow.organization_id == organization_id)
|
||||
.order_by(TaskPipelineStateRow.updated_at.desc())
|
||||
.limit(min(limit, 1000))
|
||||
)
|
||||
return [_state(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def list_at_stages(
|
||||
self, stages: tuple[Stage, ...], *, limit: int = 50
|
||||
) -> list[TaskPipelineState]:
|
||||
"""The scheduler's hot query: tasks currently at a turn stage."""
|
||||
if not stages:
|
||||
return []
|
||||
stmt = (
|
||||
select(TaskPipelineStateRow)
|
||||
.where(TaskPipelineStateRow.stage.in_([str(s) for s in stages]))
|
||||
.order_by(TaskPipelineStateRow.updated_at.asc())
|
||||
.limit(min(limit, 500))
|
||||
)
|
||||
return [_state(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def list_dependents(self, task_id: str) -> list[TaskPipelineState]:
|
||||
"""Tasks that declare ``task_id`` in ``depends_on`` and are not closed.
|
||||
|
||||
JSON columns have no portable containment operator (SQLite differs from
|
||||
PostgreSQL), so the match happens in Python. The set is small — it is
|
||||
the tasks waiting on one specific task — so this is not a scan of the
|
||||
table.
|
||||
"""
|
||||
stmt = select(TaskPipelineStateRow).where(
|
||||
TaskPipelineStateRow.stage.notin_(
|
||||
[str(Stage.COMPLETED), str(Stage.CANCELLED), str(Stage.FAILED)]
|
||||
)
|
||||
)
|
||||
dependents: list[TaskPipelineState] = []
|
||||
for row in (await self.session.execute(stmt)).scalars():
|
||||
if task_id in list(row.depends_on or []):
|
||||
dependents.append(_state(row))
|
||||
return dependents
|
||||
|
||||
async def create(self, state: TaskPipelineState) -> TaskPipelineState:
|
||||
row = TaskPipelineStateRow(
|
||||
task_id=state.task_id,
|
||||
organization_id=state.organization_id,
|
||||
project_id=state.project_id,
|
||||
pipeline=str(state.pipeline),
|
||||
kind=str(state.kind),
|
||||
risk=str(state.risk),
|
||||
stage=str(state.stage),
|
||||
cycle=state.cycle,
|
||||
max_cycles=state.max_cycles,
|
||||
branch=state.branch,
|
||||
last_diff_hash=state.last_diff_hash,
|
||||
same_diff_count=state.same_diff_count,
|
||||
retry_count=state.retry_count,
|
||||
waiver_reason=state.waiver_reason,
|
||||
parent_task_id=state.parent_task_id,
|
||||
depends_on=list(state.depends_on),
|
||||
version=state.version,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _state(row)
|
||||
|
||||
async def advance(
|
||||
self, *, task_id: str, expected_from: Stage, to: Stage
|
||||
) -> TaskPipelineState:
|
||||
"""Move a task one stage forward, but only if nobody moved it yet.
|
||||
|
||||
The row is read for its current version, then written with both the
|
||||
expected stage and that version in the ``WHERE``. If another worker
|
||||
got there first, zero rows match and :class:`StaleStageError` is
|
||||
raised — the caller re-reads and retries (or gives up cleanly).
|
||||
"""
|
||||
row = await self.session.get(TaskPipelineStateRow, task_id)
|
||||
if row is None:
|
||||
raise UnknownTaskError(task_id)
|
||||
if Stage(row.stage) is not Stage(expected_from):
|
||||
raise StaleStageError(
|
||||
f"task {task_id} is at {row.stage}, not {expected_from.value}"
|
||||
)
|
||||
if Stage(to) not in _legal_targets(Stage(row.stage)):
|
||||
raise InvalidTransitionError(f"{row.stage} -> {to.value} is not allowed")
|
||||
expected_version = row.version
|
||||
# synchronize_session="fetch" keeps the identity-mapped object truthful
|
||||
# (SQLite has no RETURNING for this), so the check below reads what the
|
||||
# database actually did rather than what we hoped it would do.
|
||||
await self.session.execute(
|
||||
update(TaskPipelineStateRow)
|
||||
.where(
|
||||
TaskPipelineStateRow.task_id == task_id,
|
||||
TaskPipelineStateRow.stage == str(expected_from),
|
||||
TaskPipelineStateRow.version == expected_version,
|
||||
)
|
||||
.values(stage=str(to), version=expected_version + 1, updated_at=utcnow())
|
||||
.execution_options(synchronize_session="fetch")
|
||||
)
|
||||
await self.session.flush()
|
||||
# A version that did not move means the WHERE matched nothing: another
|
||||
# worker wrote this row first and this move never happened.
|
||||
if row.version == expected_version or Stage(row.stage) is not Stage(to):
|
||||
raise StaleStageError(f"task {task_id} was moved by another worker")
|
||||
return _state(row)
|
||||
|
||||
async def patch(
|
||||
self, *, task_id: str, expected_version: int, **fields: object
|
||||
) -> TaskPipelineState:
|
||||
"""Change counters without moving the stage, under the version lock."""
|
||||
unknown = set(fields) - _PATCHABLE
|
||||
if unknown:
|
||||
raise ValueError(f"non-patchable columns: {sorted(unknown)}")
|
||||
row = await self.session.get(TaskPipelineStateRow, task_id)
|
||||
if row is None:
|
||||
raise UnknownTaskError(task_id)
|
||||
if not fields:
|
||||
return _state(row)
|
||||
before = row.version
|
||||
await self.session.execute(
|
||||
update(TaskPipelineStateRow)
|
||||
.where(
|
||||
TaskPipelineStateRow.task_id == task_id,
|
||||
TaskPipelineStateRow.version == expected_version,
|
||||
)
|
||||
.values(**fields, version=expected_version + 1, updated_at=utcnow())
|
||||
.execution_options(synchronize_session="fetch")
|
||||
)
|
||||
await self.session.flush()
|
||||
if row.version == before:
|
||||
raise StaleStageError(f"task {task_id} was moved by another worker")
|
||||
return _state(row)
|
||||
|
||||
async def count_at(self, stage: Stage) -> int:
|
||||
stmt = (
|
||||
select(func.count())
|
||||
.select_from(TaskPipelineStateRow)
|
||||
.where(TaskPipelineStateRow.stage == str(stage))
|
||||
)
|
||||
return int((await self.session.execute(stmt)).scalar_one())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# employee_profiles
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _profile(row: EmployeeProfileRow) -> EmployeeProfile:
|
||||
return EmployeeProfile(
|
||||
employee_id=row.employee_id,
|
||||
ai_model_id=row.ai_model_id,
|
||||
effort=row.effort,
|
||||
roles=[r.strip() for r in (row.roles or "").split(",") if r.strip()],
|
||||
hired_by=row.hired_by,
|
||||
busy_task_id=row.busy_task_id,
|
||||
updated_at=row.updated_at,
|
||||
)
|
||||
|
||||
|
||||
class EmployeeProfileRepository:
|
||||
"""The engine's view of an employee: granted model, roles, current task."""
|
||||
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get(self, employee_id: str) -> EmployeeProfile | None:
|
||||
row = await self.session.get(EmployeeProfileRow, employee_id)
|
||||
return _profile(row) if row else None
|
||||
|
||||
async def list_free_with_role(self, role: str, *, limit: int = 10) -> list[str]:
|
||||
"""Employee ids that hold ``role`` and are not busy, least busy first."""
|
||||
stmt = select(EmployeeProfileRow).where(
|
||||
EmployeeProfileRow.busy_task_id.is_(None),
|
||||
EmployeeProfileRow.roles.like(f"%{role}%"),
|
||||
)
|
||||
return [r.employee_id for r in (await self.session.execute(stmt)).scalars()][:limit]
|
||||
|
||||
async def upsert(self, profile: EmployeeProfile) -> EmployeeProfile:
|
||||
row = await self.session.get(EmployeeProfileRow, profile.employee_id)
|
||||
if row is None:
|
||||
row = EmployeeProfileRow(
|
||||
employee_id=profile.employee_id,
|
||||
ai_model_id=profile.ai_model_id,
|
||||
effort=profile.effort,
|
||||
roles=",".join(profile.role_list()),
|
||||
hired_by=profile.hired_by,
|
||||
busy_task_id=profile.busy_task_id,
|
||||
)
|
||||
self.session.add(row)
|
||||
else:
|
||||
row.ai_model_id = profile.ai_model_id
|
||||
row.effort = profile.effort
|
||||
row.roles = ",".join(profile.role_list())
|
||||
row.hired_by = profile.hired_by
|
||||
row.busy_task_id = profile.busy_task_id
|
||||
await self.session.flush()
|
||||
return _profile(row)
|
||||
|
||||
async def set_busy(self, *, employee_id: str, task_id: str | None) -> None:
|
||||
row = await self.session.get(EmployeeProfileRow, employee_id)
|
||||
if row is None:
|
||||
return
|
||||
row.busy_task_id = task_id
|
||||
await self.session.flush()
|
||||
|
||||
async def set_model(self, *, employee_id: str, ai_model_id: str | None) -> None:
|
||||
row = await self.session.get(EmployeeProfileRow, employee_id)
|
||||
if row is None:
|
||||
return
|
||||
row.ai_model_id = ai_model_id
|
||||
await self.session.flush()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# engine_turns
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _turn(row: EngineTurnRow) -> EngineTurn:
|
||||
return EngineTurn(
|
||||
id=row.id,
|
||||
task_id=row.task_id,
|
||||
run_id=row.run_id,
|
||||
employee_id=row.employee_id,
|
||||
role=RoleKey(row.role),
|
||||
mode=Mode(row.mode),
|
||||
stage=Stage(row.stage) if row.stage else Stage.PENDING_IMPLEMENTATION,
|
||||
prompt_chars=row.prompt_chars,
|
||||
reply_excerpt=row.reply_excerpt,
|
||||
tags=dict(row.tags or {}),
|
||||
tool_calls=row.tool_calls,
|
||||
ok=row.ok,
|
||||
error_kind=row.error_kind,
|
||||
started_at=row.started_at,
|
||||
finished_at=row.finished_at,
|
||||
)
|
||||
|
||||
|
||||
class EngineTurnRepository:
|
||||
"""One row per agent turn: the audit trail of who thought what."""
|
||||
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def create(self, turn: EngineTurn) -> EngineTurn:
|
||||
row = EngineTurnRow(
|
||||
id=turn.id,
|
||||
task_id=turn.task_id,
|
||||
run_id=turn.run_id,
|
||||
employee_id=turn.employee_id,
|
||||
role=str(turn.role),
|
||||
mode=str(turn.mode),
|
||||
stage=str(turn.stage),
|
||||
prompt_chars=turn.prompt_chars,
|
||||
reply_excerpt=turn.reply_excerpt,
|
||||
tags=dict(turn.tags),
|
||||
tool_calls=turn.tool_calls,
|
||||
ok=turn.ok,
|
||||
error_kind=turn.error_kind,
|
||||
started_at=turn.started_at,
|
||||
finished_at=turn.finished_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _turn(row)
|
||||
|
||||
async def list_for_task(self, task_id: str, *, limit: int = 200) -> list[EngineTurn]:
|
||||
stmt = (
|
||||
select(EngineTurnRow)
|
||||
.where(EngineTurnRow.task_id == task_id)
|
||||
.order_by(EngineTurnRow.started_at.asc())
|
||||
.limit(min(limit, 1000))
|
||||
)
|
||||
return [_turn(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def list_for_run(self, run_id: str) -> list[EngineTurn]:
|
||||
stmt = (
|
||||
select(EngineTurnRow)
|
||||
.where(EngineTurnRow.run_id == run_id)
|
||||
.order_by(EngineTurnRow.started_at.asc())
|
||||
)
|
||||
return [_turn(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def finish(
|
||||
self, turn_id: str, *, ok: bool, tool_calls: int, error_kind: str = ""
|
||||
) -> None:
|
||||
row = await self.session.get(EngineTurnRow, turn_id)
|
||||
if row is None:
|
||||
return
|
||||
row.ok = ok
|
||||
row.tool_calls = tool_calls
|
||||
row.error_kind = error_kind
|
||||
row.finished_at = utcnow()
|
||||
await self.session.flush()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# gate_results
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _gate(row: GateResultRow) -> GateResult:
|
||||
return GateResult(
|
||||
id=row.id,
|
||||
task_id=row.task_id,
|
||||
gate=GateName(row.gate),
|
||||
ok=row.ok,
|
||||
detail=row.detail,
|
||||
created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
class GateResultRepository:
|
||||
"""One row per gate verdict, with the detail that justifies it."""
|
||||
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def create(self, result: GateResult) -> GateResult:
|
||||
row = GateResultRow(
|
||||
id=result.id,
|
||||
task_id=result.task_id,
|
||||
gate=str(result.gate),
|
||||
ok=result.ok,
|
||||
detail=result.detail,
|
||||
created_at=result.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _gate(row)
|
||||
|
||||
async def list_for_task(self, task_id: str, *, limit: int = 200) -> list[GateResult]:
|
||||
stmt = (
|
||||
select(GateResultRow)
|
||||
.where(GateResultRow.task_id == task_id)
|
||||
.order_by(GateResultRow.created_at.asc())
|
||||
.limit(min(limit, 1000))
|
||||
)
|
||||
return [_gate(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def latest_for(self, task_id: str, gate: str) -> GateResult | None:
|
||||
stmt = (
|
||||
select(GateResultRow)
|
||||
.where(GateResultRow.task_id == task_id, GateResultRow.gate == gate)
|
||||
.order_by(GateResultRow.created_at.desc())
|
||||
.limit(1)
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return _gate(row) if row else None
|
||||
Reference in New Issue
Block a user