mjalmoudy94
36564e3f2b
The AI workforce engine's driver. One tick moves every ready task one step:
a turn, a gate, or a platform merge. The AI never runs git.
- scheduler.py: tick() over TURN_STAGES/MERGE_STAGES; selection caps
(per-org 3, global 8); optimistic advance so two workers cannot both win;
worktree created for code pipelines, released on completion/cancel/failure
(kept during error_review); budget pause before any turn is spent; merge
conflict back-off; transient-failure retries then the manager's error turn.
- gates.py: diff guard (empty/.git/CI/symlink/oversize), no-progress hash,
tests in the sandbox, green gate that only trusts a green tests row or an
explicit waiver. Every verdict is a gate_results row + a task.gate event.
Fixed forbidden_paths: lstrip("./") stripped all leading dots, hiding
.git and .github from the guard.
- container_engine.build_scheduler + worker drives the engine tick.
- ManagerOutcome.error_kind; TaskSpendRepository.spend_for_task.
- New settings engine_max_concurrent_turns, engine_max_turns_per_org,
engine_tests_timeout_seconds, engine_max_merge_conflicts,
engine_max_task_credits (+ .env.example).
- tests: test_engine_gates.py (13), test_engine_flow.py (11).
Checks: pytest 746 passed/1 skipped, ruff clean, mypy clean (160 files).
508 lines
18 KiB
Python
508 lines
18 KiB
Python
"""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 import RunRow, UsageEventRow
|
|
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 = "",
|
|
reply_excerpt: str | None = None,
|
|
tags: dict | None = None,
|
|
) -> None:
|
|
"""Close a turn: outcome, tool count and the bounded reply excerpt."""
|
|
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
|
|
if reply_excerpt is not None:
|
|
row.reply_excerpt = reply_excerpt
|
|
if tags is not None:
|
|
row.tags = dict(tags)
|
|
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
|
|
|
|
|
|
class TaskSpendRepository:
|
|
"""Per-task spend, read from the usage ledger (never written here).
|
|
|
|
The scheduler's budget guard needs "how much has this task cost so far",
|
|
which the ledger can answer by joining its events to the task's runs. The
|
|
ledger stays the source of truth; this only reads it.
|
|
"""
|
|
|
|
def __init__(self, session: AsyncSession) -> None:
|
|
self.session = session
|
|
|
|
async def spend_for_task(self, task_id: str) -> float:
|
|
run_ids = select(RunRow.id).where(RunRow.task_id == task_id)
|
|
stmt = select(func.coalesce(func.sum(UsageEventRow.total_price), 0.0)).where(
|
|
UsageEventRow.run_id.in_(run_ids)
|
|
)
|
|
return float((await self.session.execute(stmt)).scalar_one())
|