[Improvement] Add concurrency guard to prevent parallel advance_dag for same team

[mvillmow]

Objective

Add a per-team lock or deduplication mechanism to prevent multiple concurrent advance_dag calls for the same team from racing each other.

Motivation

When multiple tasks in the same team complete in rapid succession, the NATS listener will trigger multiple advance_dag("team-X") calls concurrently. Since advance_dag fetches tasks, finds ready ones, and claims them in a non-atomic sequence, two concurrent calls could both see the same pending task as ready and both attempt to claim it. While ProjectAgamemnon's PUT endpoint may reject the second claim, this creates unnecessary API calls and noisy error logs.

Deliverables

• Add a dict[str, asyncio.Lock] in TaskClaimer keyed by team_id
• Acquire the per-team lock at the start of advance_dag and release it on completion
• Optionally, if a call arrives while the lock is held, skip it (coalesce) rather than queue, since the in-flight call will already see the latest state
• Log when a call is coalesced/skipped for observability

Success Criteria

• At most one advance_dag call executes per team at any given time
• No deadlocks or lock leaks (use async with for safety)
• Unit test demonstrating that concurrent advance_dag calls for the same team are serialized or coalesced

Priority

HIGH -- Race conditions in task claiming are the most likely source of production bugs in the current architecture

---

Updated per ADR-006 — ai-maestro references replaced with current architecture terminology.

[mvillmow]

Implementation Plan

Now I have a complete picture of the codebase. Here's the implementation plan:

---

Implementation Plan: Issue #131 — Per-Team Concurrency Guard for advance_dag

Objective

Add a per-team asyncio.Lock with coalesce/skip semantics to TaskClaimer.advance_dag() so that at most one call executes per team at any time. Concurrent calls for the same team are skipped (not queued) with a log message.

Approach

Use a two-layer guard:

1. _in_progress: set[str] — checked before acquiring the lock to implement coalesce/skip. If team_id is in the set, the call returns immediately instead of queuing behind the lock.
2. _team_locks: dict[str, asyncio.Lock] — per-team lock created lazily. The lock is held for the entire fetch→find→claim cycle inside advance_dag.

This avoids queuing redundant calls (since the in-flight call already sees latest state) while still providing mutual exclusion if the skip-check and lock-acquire race.

Files to Create

None.

Files to Modify

1. src/keystone/task_claimer.py

• Add import asyncio at the top.
• Add two instance attributes in __init__:
  • self._team_locks: dict[str, asyncio.Lock] = {}
  • self._advancing: set[str] = set()
• Add a private method _get_team_lock(team_id: str) -> asyncio.Lock that lazily creates and returns the lock for a given team.
• Wrap advance_dag body with the coalesce check + lock:
  • At the top: if team_id in self._advancing, log a skip message and return [].
  • Otherwise: self._advancing.add(team_id), then async with self._get_team_lock(team_id): around the existing body, with self._advancing.discard(team_id) in a finally block.

2. tests/test_task_claimer.py

Add three new test cases:

• test_concurrent_advance_dag_coalesced — Launch two advance_dag calls concurrently for the same team using asyncio.gather. Use asyncio.Event gates inside a slow mock get_tasks to ensure the first call is in-flight when the second arrives. Assert: get_tasks is called only once (second call was skipped).
• test_concurrent_advance_dag_different_teams_parallel — Launch two concurrent calls for different teams. Assert: both complete successfully and get_tasks is called twice (once per team).
• test_advance_dag_lock_released_on_error — Make get_tasks raise on the first call. Call advance_dag a second time (should not deadlock). Assert: second call proceeds normally.

Implementation Order

1. Modify task_claimer.py — Add the asyncio import, instance attributes, _get_team_lock method, and wrap advance_dag with the coalesce guard.
2. Add tests in test_task_claimer.py — Write the three new tests using asyncio.Event for deterministic scheduling.
3. Run just test — Verify all existing + new tests pass.
4. Run just lint — Verify no lint issues.

Detailed Code Design

_get_team_lock

def _get_team_lock(self, team_id: str) -> asyncio.Lock:
    if team_id not in self._team_locks:
        self._team_locks[team_id] = asyncio.Lock()
    return self._team_locks[team_id]

This is safe because all asyncio code runs on a single thread — no race creating the lock.

advance_dag guard structure

async def advance_dag(self, team_id: str) -> list[str]:
    # Coalesce: skip if already in-flight for this team
    if team_id in self._advancing:
        logger.info("advance_dag already in-flight for team %s — skipping", team_id)
        return []

    self._advancing.add(team_id)
    try:
        async with self._get_team_lock(team_id):
            # ... existing body unchanged ...
    finally:
        self._advancing.discard(team_id)

Key properties:

• _advancing set provides fast O(1) skip check without waiting on the lock.
• async with guarantees the lock is released even on exceptions.
• finally guarantees _advancing is cleaned up even on exceptions.
• The lock scope covers the entire fetch→find→claim cycle (no gap for races).

Verification

1. Unit tests — The three new tests cover:
   • Same-team concurrent calls are coalesced (core requirement)
   • Different-team concurrent calls proceed independently (no false serialization)
   • Lock/set cleanup on exceptions (no deadlocks or leaks)
2. Existing tests — All existing test_task_claimer.py tests continue to pass unchanged (the guard is transparent to single-call scenarios).
3. Lint — just lint passes.

Skills Used

Skills Invoked During Planning

• None (direct codebase reading was sufficient)

Team Knowledge Base Skills Referenced (from Prior Learnings)

Skill	How It Informed the Plan
global-semaphore-parallelism	Per-resource (per-team) concurrency control pattern; acquire in execution path not dispatcher; try/finally for guaranteed release
e2e-framework-crash-recovery-bugs	Module-level lock with try/finally pattern; PID-only keying fails in same-process concurrency
checkpoint-disk-reconciliation	Reinforced that the lock must cover the full read-write cycle (fetch→find→claim), not just a subset
checkpoint-save-guard-bug	Guard conditions must cover all affected operations; informed the _advancing set design to properly gate the entire path
fix-flaky-sleep-mock	Test with deterministic asyncio.Event gates, not wall-clock sleep assertions
logging-context-deduplication	Structured log messages with team_id context for the skip/coalesce log line
checkpoint-intermediate-state-resume	Keep complexity down — extract coalesce logic rather than bloating advance_dag

---

Generated by Claude Code Planner

[mvillmow]

Created 4 follow-up issue(s): #295, #296, #297, #298