diff --git a/roboco/services/gateway/choreographer/_impl.py b/roboco/services/gateway/choreographer/_impl.py index 61bcb96d..b87122d4 100644 --- a/roboco/services/gateway/choreographer/_impl.py +++ b/roboco/services/gateway/choreographer/_impl.py @@ -1168,6 +1168,22 @@ class Choreographer: # above. These migrate into spec.extra_preconditions in a later # task; until then, keep them imperative so concurrency invariants # stay enforced. + # + # The agent-wide guards below read the agent's OTHER tasks via plain + # (unlocked) SELECTs, and claim()'s FOR UPDATE locks only the TARGET + # row — neither serializes two concurrent claims by the SAME agent on + # TWO DIFFERENT pending tasks. A transaction-scoped advisory lock keyed + # by agent_id, acquired here BEFORE the guard reads and held until the + # request transaction commits, makes the second concurrent claim's read + # see the first's committed in_progress task and get rejected. The + # in-process asyncio.Lock is lost on orchestrator-restart split-brain, + # so this is the only DB-level guarantee of the one-task-per-agent + # invariant. Coordinators (cell_pm / main_pm) are exempt — matching the + # already_active/paused guard exemption below — so a PM can still plan + + # delegate many roots in parallel (the PM coordinator concurrency + # feature). A hash collision only causes benign false serialization. + if role_str not in self._COORDINATOR_ROLES: + await self.task.acquire_claim_lock(ctx.agent_id) if guard := await self._run_claim_guards( agent_id=ctx.agent_id, task=t, diff --git a/roboco/services/task.py b/roboco/services/task.py index 50dfc559..8b9e93ce 100644 --- a/roboco/services/task.py +++ b/roboco/services/task.py @@ -11,7 +11,7 @@ from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING, Any, ClassVar, cast from uuid import UUID, uuid4 -from sqlalchemy import and_, func, or_, select, update +from sqlalchemy import and_, func, or_, select, text, update from sqlalchemy import inspect as sa_inspect from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import InstanceState @@ -2341,6 +2341,31 @@ class TaskService(BaseService): self._background_tasks.add(bg_task) bg_task.add_done_callback(self._background_tasks.discard) + async def acquire_claim_lock(self, agent_id: UUID) -> None: + """Take a per-agent transaction-scoped advisory lock. + + ``claim``'s ``SELECT ... FOR UPDATE`` serializes concurrent claims of + the SAME task, but the one-task-per-agent invariant is an agent-WIDE + guard: two concurrent claims by the SAME agent on TWO DIFFERENT pending + tasks lock different rows and both pass the already_active guard (each + reads an empty in_progress set before either commits). A + ``pg_advisory_xact_lock`` keyed by the agent serializes the WHOLE + claim+guard+start critical section per agent — held until the outer + request transaction commits, the second concurrent claim blocks until + the first commits, then its guard read sees the first's committed + in_progress task and is rejected. + + Transaction-scoped (not session-scoped) so it auto-releases on commit + or rollback and cannot outlive the request. ``hashtextextended`` maps a + UUID to a ``bigint`` key; a hash collision only causes benign false + serialization (two different agents momentarily serializing), never a + false negative — so it is not a correctness concern. + """ + await self.session.execute( + text("SELECT pg_advisory_xact_lock(hashtextextended(CAST(:aid AS text), 0))"), + {"aid": str(agent_id)}, + ) + async def claim( self, task_id: UUID, agent_id: UUID, allow_reassign: bool = False ) -> TaskTable | None: diff --git a/tests/integration/test_claim_lock_serialization.py b/tests/integration/test_claim_lock_serialization.py new file mode 100644 index 00000000..80d25c2a --- /dev/null +++ b/tests/integration/test_claim_lock_serialization.py @@ -0,0 +1,112 @@ +"""F074 — real-Postgres proof that ``TaskService.acquire_claim_lock`` serializes +concurrent claims by the SAME agent (the one-task-per-agent invariant) while NOT +serializing claims by DIFFERENT agents. + +The choreographer-level ordering + coordinator-exemption is covered by the unit +suite (``test_choreographer_claim_lock.py``); this test pins the DB-level +contract the unit suite mocks out: that ``pg_advisory_xact_lock`` keyed by +``hashtextextended(agent_id)`` actually blocks a second transaction trying to +acquire the same agent's lock until the first commits/rolls back, and that a +different agent's lock is uncontended. Skips when Postgres is unreachable. + +Each ``acquire_claim_lock`` call uses its own fresh session/engine. The +blocking call's session is single-use: ``asyncio.wait_for`` cancelling an +in-flight asyncpg query leaves the SQLAlchemy session mid-connection-checkout +("provisioning a new connection"), so a throwaway session per timed acquire +keeps the rest of the test on clean connections. +""" + +from __future__ import annotations + +import asyncio +import contextlib +from uuid import uuid4 + +import pytest +from roboco.services.task import TaskService +from sqlalchemy.ext.asyncio import ( + AsyncEngine, + AsyncSession, + async_sessionmaker, + create_async_engine, +) + + +async def _fresh_session(url: str) -> tuple[AsyncSession, AsyncEngine]: + """A session on a brand-new engine (caller disposes via ``_dispose``).""" + engine = create_async_engine(url, future=True) + factory = async_sessionmaker( + bind=engine, class_=AsyncSession, expire_on_commit=False + ) + return factory(), engine + + +async def _dispose(session: AsyncSession, engine: AsyncEngine) -> None: + with contextlib.suppress(Exception): + await session.rollback() + await engine.dispose() + + +@pytest.mark.asyncio +async def test_advisory_lock_serializes_same_agent_not_different( + _test_database_url: str, +) -> None: + """TX1 holds agent A's lock (uncommitted). A second acquire for agent A + must block past a short deadline; an acquire for a DIFFERENT agent B must + return immediately. After TX1 rolls back, agent A's lock is releasable. + Proves the SQL is valid against real PG, serializes per-agent, is + collision-free across agents, and is transaction-scoped (releases on + rollback).""" + url = _test_database_url + agent_a = uuid4() + agent_b = uuid4() + + # TX1 acquires agent A's tx-scoped advisory lock and holds it (no commit). + holder, holder_engine = await _fresh_session(url) + try: + await TaskService(holder).acquire_claim_lock(agent_a) + + # A second transaction acquiring the SAME agent's lock must block. + blocked, blocked_engine = await _fresh_session(url) + try: + with pytest.raises(TimeoutError): + await asyncio.wait_for( + TaskService(blocked).acquire_claim_lock(agent_a), timeout=0.5 + ) + finally: + await _dispose(blocked, blocked_engine) + + # A different agent's lock is uncontended — acquires immediately. + other, other_engine = await _fresh_session(url) + try: + await asyncio.wait_for( + TaskService(other).acquire_claim_lock(agent_b), timeout=2.0 + ) + finally: + await _dispose(other, other_engine) + + # Rolling back TX1 ends its transaction → releases the tx-scoped lock. + await holder.rollback() + finally: + await _dispose(holder, holder_engine) + + # Agent A's lock is now free — a fresh transaction acquires it at once. + reacquire, reacquire_engine = await _fresh_session(url) + try: + await asyncio.wait_for( + TaskService(reacquire).acquire_claim_lock(agent_a), timeout=2.0 + ) + finally: + await _dispose(reacquire, reacquire_engine) + + +@pytest.mark.asyncio +async def test_acquire_runs_without_error(_test_database_url: str) -> None: + """Sanity: the advisory-lock SQL executes against the migrated schema + (``hashtextextended`` exists, the statement is valid) — pins that the + one-liner hasn't been broken by a schema/PG-version drift.""" + session, engine = await _fresh_session(_test_database_url) + try: + await TaskService(session).acquire_claim_lock(uuid4()) + finally: + await _dispose(session, engine) diff --git a/tests/unit/gateway/test_choreographer_claim_lock.py b/tests/unit/gateway/test_choreographer_claim_lock.py new file mode 100644 index 00000000..3f462f54 --- /dev/null +++ b/tests/unit/gateway/test_choreographer_claim_lock.py @@ -0,0 +1,280 @@ +"""F074 — the one-task-per-agent invariant had no DB-level enforcement. + +``_run_claim_guards`` read the agent's other tasks via unlocked SELECTs +(``list_in_progress_for_agent`` / ``list_paused_for_agent``) BEFORE ``claim()`` +took its row lock, and ``claim()``'s ``FOR UPDATE`` locked only the TARGET row +— not the agent-wide invariant. So two concurrent ``i_will_work_on`` calls by +the SAME agent on TWO DIFFERENT pending tasks each locked their own target row +(no contention), each read an empty in_progress set, each passed +``already_active_guard``, and each claim+start succeeded → the agent ended with +two in_progress tasks. The in-process ``asyncio.Lock`` serializes container +spawns per agent but is lost on orchestrator-restart split-brain, so it is not +a DB-level guarantee. + +The fix: a PostgreSQL transaction-scoped advisory lock keyed by agent_id, +acquired in ``_claim_plan_start_gate`` BEFORE the guard reads (for non- +coordinator roles only). Held until the request transaction commits, it spans +the guard read + the savepoint + the claim write, so the second concurrent +claim's guard read sees the first's committed in_progress task and is rejected. + +CRITICAL logical-regression guard: the advisory lock is acquired ONLY for non- +coordinator roles. The PM coordinator concurrency feature (CLAUDE.md) lets a +cell_pm / main_pm plan + delegate many roots in parallel — acquiring a per- +agent advisory lock for a coordinator would serialize those claims and +REGRESS that feature. So coordinators are exempt (matching the existing +``_COORDINATOR_ROLES`` guard exemption for ``already_active`` / ``paused``). +""" + +from __future__ import annotations + +from datetime import UTC, datetime +from typing import Any +from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 + +import pytest +from roboco.services.gateway.choreographer import Choreographer, ChoreographerDeps + +# A developer fresh claim must carry a substantive step checklist + rich plan +# (parity with a PM) so the tracing-gate passes and the flow reaches claim. +_STEPS = [ + { + "title": "Implement the change", + "description": ( + "edit the target file, add tests, run them, and stage the " + "change for commit on the task branch" + ), + } +] +_GOOD_PLAN = ( + "Append the timestamp HTML comment to the very bottom of README.md without " + "touching any other line, then commit it on the task branch and open a PR. " + "Verify the diff is a single-line addition before submitting for QA." +) +_GOOD_TC = ["Use a trailing newline so the comment sits on its own line."] +_GOOD_RISKS = [ + { + "risk": "An accidental reformat of README.md balloons the diff.", + "mitigation": "Append only; assert the diff touches one line pre-commit.", + } +] + + +def _make_deps(task: AsyncMock) -> ChoreographerDeps: + task.session = MagicMock() + task.session.begin_nested = MagicMock( + return_value=MagicMock( + __aenter__=AsyncMock(return_value=None), + __aexit__=AsyncMock(return_value=False), + ) + ) + evidence_repo = AsyncMock() + for method in ( + "list_unread_a2a", + "list_unread_mentions", + "list_pending_notifications", + "task_metadata_gaps", + "recent_team_activity", + "blockers_in_lane", + "journal_highlights_for_task", + ): + getattr(evidence_repo, method).return_value = [] + journal = AsyncMock() + journal.latest_decision_at.return_value = datetime.now(UTC) + return ChoreographerDeps( + task=task, + work_session=AsyncMock(), + git=AsyncMock(), + a2a=AsyncMock(), + journal=journal, + audit=AsyncMock(), + evidence_repo=evidence_repo, + ) + + +def _dev_task_svc(agent_id: object, task_id: object) -> AsyncMock: + """TaskService mock that completes the dev claim→plan→start flow.""" + pending = MagicMock( + id=task_id, + status="pending", + plan=None, + assigned_to=None, + parent_task_id=None, + sequence=0, + task_type="code", + commits=[], + pr_number=None, + branch_name="feature/backend/abc", + quick_context=None, + ) + in_progress = MagicMock( + id=task_id, status="in_progress", plan={"text": "do x"}, assigned_to=agent_id + ) + task_svc = AsyncMock() + task_svc.get.return_value = pending + task_svc.agent_for.return_value = MagicMock( + id=agent_id, role="developer", team="backend", slug=None + ) + task_svc.list_in_progress_for_agent.return_value = [] + task_svc.list_paused_for_agent.return_value = [] + task_svc.get_subtasks.return_value = [] + task_svc.claim.return_value = MagicMock( + id=task_id, status="claimed", plan=None, assigned_to=agent_id + ) + task_svc.set_plan.return_value = MagicMock( + id=task_id, status="claimed", plan={"text": "do x"}, assigned_to=agent_id + ) + task_svc.start.return_value = in_progress + return task_svc + + +@pytest.mark.asyncio +async def test_dev_claim_acquires_lock_before_guard_read() -> None: + """A developer claim acquires the per-agent advisory lock BEFORE the + already_active guard reads the agent's other tasks — the ordering that + closes the TOCTOU (the second concurrent claim's read must see the first's + committed in_progress task, which requires the lock to be held first).""" + agent_id = uuid4() + task_id = uuid4() + task_svc = _dev_task_svc(agent_id, task_id) + + # Shared call-order recorder: the lock MUST be acquired before the guard + # read, so a shared list distinguishes the correct ordering from the bug + # (read-then-claim, where the lock is too late). + calls: list[str] = [] + + async def _lock(_aid: object) -> None: + calls.append("lock") + + async def _read_in_progress(_aid: object) -> list[Any]: + calls.append("read_in_progress") + return [] + + async def _read_paused(_aid: object) -> list[Any]: + calls.append("read_paused") + return [] + + task_svc.acquire_claim_lock = _lock + task_svc.list_in_progress_for_agent.side_effect = _read_in_progress + task_svc.list_paused_for_agent.side_effect = _read_paused + + deps = _make_deps(task_svc) + c = Choreographer(deps) + + env = await c.i_will_work_on( + agent_id, + task_id, + plan=_GOOD_PLAN, + steps=_STEPS, + technical_considerations=_GOOD_TC, + risks=_GOOD_RISKS, + ) + assert env.error is None, env.as_dict() + + # The lock was acquired once for this agent, and BEFORE the guard reads. + assert calls[:1] == ["lock"], ( + f"advisory lock must be acquired before the guard reads; order was {calls}" + ) + assert "read_in_progress" in calls and calls.index("lock") < calls.index( + "read_in_progress" + ) + + +@pytest.mark.asyncio +async def test_coordinator_claim_does_not_acquire_lock() -> None: + """A coordinator PM (cell_pm / main_pm) must NOT acquire the per-agent + advisory lock — the PM coordinator concurrency feature (CLAUDE.md) lets a + PM plan + delegate many roots in parallel. Acquiring the lock would + serialize those claims and REGRESS that feature. The lock is only for the + one-task-per-agent invariant, which coordinators are exempt from (matching + the existing ``_COORDINATOR_ROLES`` already_active/paused guard exemption).""" + pm_id = uuid4() + task_id = uuid4() + task_svc = AsyncMock() + pending = MagicMock( + id=task_id, + status="pending", + plan=None, + assigned_to=None, + task_type="planning", + parent_task_id=None, + sequence=0, + team="backend", + commits=[], + pr_number=None, + branch_name=None, + quick_context=None, + work_session_id=None, + ) + task_svc.get.return_value = pending + task_svc.agent_for.return_value = MagicMock( + id=pm_id, role="cell_pm", team="backend", slug=None + ) + task_svc.list_in_progress_for_agent.return_value = [] + task_svc.list_paused_for_agent.return_value = [] + task_svc.get_subtasks.return_value = [] + task_svc.claim.return_value = MagicMock( + id=task_id, status="claimed", plan=None, assigned_to=pm_id + ) + task_svc.set_plan.return_value = MagicMock( + id=task_id, status="claimed", plan={"text": "x"}, assigned_to=pm_id + ) + task_svc.start.return_value = MagicMock( + id=task_id, status="in_progress", plan={"text": "x"}, assigned_to=pm_id + ) + task_svc.ensure_work_session.return_value = None + + lock_calls: list[object] = [] + + async def _lock(aid: object) -> None: + lock_calls.append(aid) + + task_svc.acquire_claim_lock = _lock + + deps = _make_deps(task_svc) + c = Choreographer(deps) + + env = await c.i_will_plan( + pm_id, + task_id, + plan="decompose the feature into cell tasks with acceptance criteria", + rich_plan={ + "approach": ( + "Two-slice decomposition: backend builds the API + DB migration " + "first, UX consumes it after the endpoint merges. Strict " + "sequencing; the UX slice depends on the backend slice landing." + ), + "sub_tasks": [ + { + "title": "Backend slice", + "description": ( + "be-dev-1 implements the API endpoint and DB migration, " + "with tests, behind the existing service layer." + ), + }, + { + "title": "UX slice", + "description": ( + "fe-dev-1 wires the panel view to the new endpoint and " + "renders the result list with loading + error states." + ), + }, + ], + "risks": [ + { + "risk": "schema migration may block", + "mitigation": "rehearse the migration on a throwaway DB first", + } + ], + }, + ) + body = env.as_dict() + # The flow MUST reach the claim gate (not short-circuit at input validation) + # — otherwise ``lock_calls == []`` would pass for the wrong reason. + assert body.get("error") != "incomplete_input", body + + # The coordinator did NOT acquire the lock — parallel planning preserved. + assert lock_calls == [], ( + "coordinator PM must not acquire the per-agent claim lock (would regress " + "the PM coordinator concurrency feature)" + )