task_queue.py
file-level
1
files
1
commits
0
hotspots
0
🧊 dead
0
💥 blast risk
| 1 | """Task queue with atomic claiming for multi-agent work distribution. |
| 2 | |
| 3 | This module implements a **file-system–based task queue** that allows an |
| 4 | orchestrator to enqueue work items and N competing agents to claim them |
| 5 | without any central coordinator or external service. |
| 6 | |
| 7 | The fundamental invariant is:: |
| 8 | |
| 9 | For any task, at most one agent holds an active claim at any time. |
| 10 | |
| 11 | Atomic claiming |
| 12 | --------------- |
| 13 | The atomic claiming primitive is ``open(path, O_CREAT | O_EXCL | O_WRONLY)``. |
| 14 | On POSIX-compliant file systems (including all common local filesystems and |
| 15 | most NFS mounts), this syscall either: |
| 16 | |
| 17 | * **Succeeds** — creating the file and returning the fd (the callee won the |
| 18 | race). |
| 19 | * **Fails with ENOENT/FileExistsError** — another agent created the file first |
| 20 | (the callee lost the race). |
| 21 | |
| 22 | No lock file, no mutex, no advisory lock, no external service is needed. |
| 23 | |
| 24 | Re-claiming timed-out tasks |
| 25 | ---------------------------- |
| 26 | When an agent crashes or is killed while holding a claim, the claim's |
| 27 | ``expires_at`` timestamp passes without a ``complete`` or ``fail`` call. |
| 28 | Such claims are called **timed-out** and the task becomes eligible for |
| 29 | re-claiming. |
| 30 | |
| 31 | Re-claiming uses an **optimistic concurrency** approach: |
| 32 | |
| 33 | 1. Write a new claim via ``write_text_atomic`` (atomic rename over the |
| 34 | existing claim file) embedding a unique ``claim_nonce``. |
| 35 | 2. Read back the file and verify the ``claim_nonce`` matches what we wrote. |
| 36 | 3. If the nonce matches → we won. If not → another agent beat us; continue |
| 37 | to the next candidate. |
| 38 | |
| 39 | This gives at-most-once semantics for the common case (local filesystem, |
| 40 | sequential renames) with a tiny window of ambiguity under extreme concurrency. |
| 41 | The window closes immediately on read-back verification. |
| 42 | |
| 43 | Directory layout |
| 44 | ---------------- |
| 45 | :: |
| 46 | |
| 47 | .muse/coordination/ |
| 48 | tasks/ — task definitions (write-once, immutable after enqueue) |
| 49 | <task-id>.json |
| 50 | claims/ — claim records (O_CREAT|O_EXCL for initial claim; |
| 51 | write_text_atomic for updates) |
| 52 | <task-id>.json |
| 53 | |
| 54 | Derived task status (from file existence + claim.status) |
| 55 | --------------------------------------------------------- |
| 56 | ``"pending"`` |
| 57 | Task file exists, no claim file. |
| 58 | |
| 59 | ``"claimed"`` |
| 60 | Claim exists, ``status == "claimed"`` and ``expires_at > now``. |
| 61 | |
| 62 | ``"timed_out"`` |
| 63 | Claim exists, ``status == "claimed"`` and ``expires_at <= now``. |
| 64 | The task is eligible for re-claiming. |
| 65 | |
| 66 | ``"completed"`` |
| 67 | Claim exists, ``status == "completed"``. |
| 68 | |
| 69 | ``"failed"`` |
| 70 | Claim exists, ``status == "failed"``. |
| 71 | |
| 72 | ``"cancelled"`` |
| 73 | Claim exists, ``status == "cancelled"``. |
| 74 | |
| 75 | Security |
| 76 | -------- |
| 77 | Every ``task_id`` is validated as a sha256: content ID before constructing any file path, |
| 78 | mirroring :func:`~muse.core.coordination._validate_reservation_id`. A valid |
| 79 | sha256: content ID cannot contain ``/`` or ``..``, which prevents directory traversal |
| 80 | attacks. |
| 81 | |
| 82 | Performance |
| 83 | ----------- |
| 84 | ``claim_next_task`` scans the ``tasks/`` directory once (O(n)) to build a |
| 85 | priority-sorted candidate list, then attempts O_EXCL claim in order. On a |
| 86 | quiet queue the first attempt almost always succeeds. On a hot queue |
| 87 | (many agents competing for few tasks) each agent tries at most *n* attempts, |
| 88 | but the total number of attempts across all agents is bounded by the number |
| 89 | of pending tasks × the number of competing agents — no spin loops. |
| 90 | """ |
| 91 | |
| 92 | import dataclasses |
| 93 | import datetime |
| 94 | import json |
| 95 | import logging |
| 96 | import os |
| 97 | import pathlib |
| 98 | import re |
| 99 | import secrets |
| 100 | |
| 101 | from muse._version import __version__ |
| 102 | from muse.core.types import content_hash, load_json_file |
| 103 | from muse.core.coordination import _coord_dir, _now_utc, _parse_dt |
| 104 | from muse.core.io import write_text_atomic |
| 105 | |
| 106 | type _ClaimMap = dict[str, "ClaimRecord"] |
| 107 | # JSON-compatible value type for arbitrary task payloads and results. |
| 108 | JsonValue = str | int | float | bool | None | list["JsonValue"] | "JsonDict" |
| 109 | JsonDict = dict[str, JsonValue] |
| 110 | |
| 111 | logger = logging.getLogger(__name__) |
| 112 | |
| 113 | # ── Constants ───────────────────────────────────────────────────────────────── |
| 114 | |
| 115 | # Valid task statuses stored in a claim record. |
| 116 | _CLAIM_STATUSES: frozenset[str] = frozenset( |
| 117 | ("claimed", "completed", "failed", "cancelled", "timed_out") |
| 118 | ) |
| 119 | |
| 120 | # Derived task status values (not stored; computed from file state). |
| 121 | _TASK_STATUSES: frozenset[str] = frozenset( |
| 122 | ("pending", "claimed", "timed_out", "completed", "failed", "cancelled") |
| 123 | ) |
| 124 | |
| 125 | # Maximum allowed tag count per task. |
| 126 | _MAX_TAGS = 32 |
| 127 | |
| 128 | # Maximum tag length in characters. |
| 129 | _MAX_TAG_LEN = 64 |
| 130 | |
| 131 | # Maximum title length in characters. |
| 132 | _MAX_TITLE_LEN = 256 |
| 133 | |
| 134 | # Maximum queue name length in characters. |
| 135 | _MAX_QUEUE_LEN = 64 |
| 136 | |
| 137 | # Content-addressed task_id regex: sha256: followed by 64 hex chars. |
| 138 | # No path separators or special characters, preventing directory traversal. |
| 139 | _TASK_ID_RE = re.compile(r"^sha256:[0-9a-f]{64}$") |
| 140 | |
| 141 | # ── Directory helpers ───────────────────────────────────────────────────────── |
| 142 | |
| 143 | def _tasks_dir(root: pathlib.Path) -> pathlib.Path: |
| 144 | """Return ``.muse/coordination/tasks/`` for *root*.""" |
| 145 | return _coord_dir(root) / "tasks" |
| 146 | |
| 147 | def _claims_dir(root: pathlib.Path) -> pathlib.Path: |
| 148 | """Return ``.muse/coordination/claims/`` for *root*.""" |
| 149 | return _coord_dir(root) / "claims" |
| 150 | |
| 151 | def ensure_task_dirs(root: pathlib.Path) -> None: |
| 152 | """Create both task-queue subdirectories if they do not exist.""" |
| 153 | _tasks_dir(root).mkdir(parents=True, exist_ok=True) |
| 154 | _claims_dir(root).mkdir(parents=True, exist_ok=True) |
| 155 | |
| 156 | # ── Input validation ────────────────────────────────────────────────────────── |
| 157 | |
| 158 | def _validate_task_id(task_id: str) -> None: |
| 159 | """Raise :exc:`ValueError` if *task_id* is not a well-formed content-addressed ID. |
| 160 | |
| 161 | Security note |
| 162 | ------------- |
| 163 | ``task_id`` is used to construct file paths inside |
| 164 | ``.muse/coordination/tasks/`` and ``.muse/coordination/claims/``. |
| 165 | The ``sha256:`` prefix and hex-only suffix cannot contain ``/`` or ``..``, |
| 166 | preventing directory traversal attacks. |
| 167 | """ |
| 168 | if not _TASK_ID_RE.match(str(task_id)): |
| 169 | raise ValueError(f"task_id must be a sha256: content-addressed ID: {task_id!r}") |
| 170 | |
| 171 | def _validate_queue_name(queue: str) -> None: |
| 172 | """Raise :exc:`ValueError` if *queue* is not a safe queue name. |
| 173 | |
| 174 | Queue names must be non-empty, at most :data:`_MAX_QUEUE_LEN` characters, |
| 175 | and contain only ASCII alphanumerics, hyphens, and underscores. |
| 176 | """ |
| 177 | if not queue: |
| 178 | raise ValueError("queue name must be non-empty") |
| 179 | if len(queue) > _MAX_QUEUE_LEN: |
| 180 | raise ValueError( |
| 181 | f"queue name too long: {len(queue)} > {_MAX_QUEUE_LEN} chars" |
| 182 | ) |
| 183 | if not re.match(r"^[a-zA-Z0-9_-]+$", queue): |
| 184 | raise ValueError( |
| 185 | f"queue name may only contain [a-zA-Z0-9_-]: {queue!r}" |
| 186 | ) |
| 187 | |
| 188 | # ── Time helpers ────────────────────────────────────────────────────────────── |
| 189 | |
| 190 | # ── Data model ──────────────────────────────────────────────────────────────── |
| 191 | |
| 192 | @dataclasses.dataclass |
| 193 | class TaskRecord: |
| 194 | """An immutable task definition stored in ``tasks/<task-id>.json``. |
| 195 | |
| 196 | A :class:`TaskRecord` is written once (by the enqueuing agent) and never |
| 197 | modified. All mutable state lives in :class:`ClaimRecord`. |
| 198 | |
| 199 | Attributes |
| 200 | ---------- |
| 201 | task_id: |
| 202 | content-addressed sha256: ID that uniquely identifies this task. |
| 203 | title: |
| 204 | Short human-readable description (≤ :data:`_MAX_TITLE_LEN` chars). |
| 205 | payload: |
| 206 | Arbitrary JSON-serialisable dict. Typically contains addresses, |
| 207 | operation type, or other agent-specific parameters. |
| 208 | priority: |
| 209 | Integer priority — **higher values are processed first**. Ties are |
| 210 | broken by ``created_at`` (FIFO: earliest first). |
| 211 | queue: |
| 212 | Logical queue name. Agents claiming with ``queue="billing"`` only see |
| 213 | billing tasks. Must match ``[a-zA-Z0-9_-]+``. |
| 214 | created_at: |
| 215 | UTC datetime when the task was enqueued. |
| 216 | created_by: |
| 217 | ``run_id`` of the agent or orchestrator that enqueued this task. |
| 218 | ttl_seconds: |
| 219 | How long (in seconds) from ``created_at`` the task may remain |
| 220 | ``pending`` before it is considered expired and skipped during |
| 221 | claiming. Completed/failed tasks are not affected. |
| 222 | tags: |
| 223 | List of string labels for filtering or analytics. |
| 224 | """ |
| 225 | |
| 226 | task_id: str |
| 227 | title: str |
| 228 | payload: JsonDict |
| 229 | priority: int |
| 230 | queue: str |
| 231 | created_at: datetime.datetime |
| 232 | created_by: str |
| 233 | ttl_seconds: int |
| 234 | tags: list[str] |
| 235 | |
| 236 | def to_dict(self) -> JsonDict: |
| 237 | """Serialise to a JSON-compatible dict.""" |
| 238 | return { |
| 239 | "schema_version": __version__, |
| 240 | "task_id": self.task_id, |
| 241 | "title": self.title, |
| 242 | "payload": self.payload, |
| 243 | "priority": self.priority, |
| 244 | "queue": self.queue, |
| 245 | "created_at": self.created_at.isoformat(), |
| 246 | "created_by": self.created_by, |
| 247 | "ttl_seconds": self.ttl_seconds, |
| 248 | "tags": self.tags, |
| 249 | } |
| 250 | |
| 251 | @classmethod |
| 252 | def from_dict(cls, d: JsonDict) -> "TaskRecord": |
| 253 | """Deserialise from a JSON-compatible dict. |
| 254 | |
| 255 | Missing fields default to safe values so records written by older |
| 256 | versions remain loadable. |
| 257 | """ |
| 258 | return cls( |
| 259 | task_id=str(d.get("task_id", "")), |
| 260 | title=str(d.get("title", ""))[:_MAX_TITLE_LEN], |
| 261 | payload=dict(d.get("payload") or {}), |
| 262 | priority=int(d.get("priority", 0)), |
| 263 | queue=str(d.get("queue", "default")), |
| 264 | created_at=_parse_dt(str(d["created_at"])) if "created_at" in d else _now_utc(), |
| 265 | created_by=str(d.get("created_by", "")), |
| 266 | ttl_seconds=int(d.get("ttl_seconds", 3600)), |
| 267 | tags=[str(t) for t in (d.get("tags") or [])], |
| 268 | ) |
| 269 | |
| 270 | def is_expired(self, now: datetime.datetime | None = None) -> bool: |
| 271 | """Return True if the task's pending TTL has passed.""" |
| 272 | if now is None: |
| 273 | now = _now_utc() |
| 274 | return now >= self.created_at + datetime.timedelta(seconds=self.ttl_seconds) |
| 275 | |
| 276 | @dataclasses.dataclass |
| 277 | class ClaimRecord: |
| 278 | """Mutable claim state for a task stored in ``claims/<task-id>.json``. |
| 279 | |
| 280 | A :class:`ClaimRecord` is **created atomically** via |
| 281 | ``open(O_CREAT | O_EXCL)`` by the claiming agent, then **updated** via |
| 282 | :func:`~muse.core.store.write_text_atomic` (atomic rename) as the task |
| 283 | progresses through its lifecycle. |
| 284 | |
| 285 | Attributes |
| 286 | ---------- |
| 287 | task_id: |
| 288 | content-addressed ID of the task being claimed. |
| 289 | claimer_run_id: |
| 290 | ``run_id`` of the agent that claimed the task. |
| 291 | claimed_at: |
| 292 | UTC datetime when the claim was created. |
| 293 | expires_at: |
| 294 | UTC datetime after which the claim is considered timed-out and the |
| 295 | task becomes eligible for re-claiming. |
| 296 | status: |
| 297 | One of ``"claimed"``, ``"completed"``, ``"failed"``, |
| 298 | ``"cancelled"``, ``"timed_out"``. |
| 299 | heartbeat_at: |
| 300 | UTC datetime of the most recent heartbeat. Agents should call |
| 301 | :func:`heartbeat_claim` every ``claim_ttl_seconds / 2`` to prevent |
| 302 | expiry. |
| 303 | claim_nonce: |
| 304 | Random nonce generated at claim time. Used for **optimistic re-claim |
| 305 | verification**: after overwriting a timed-out claim, the new claimer |
| 306 | reads back the file and checks that the nonce matches, confirming |
| 307 | they won the race. |
| 308 | result: |
| 309 | Optional JSON-serialisable dict set by :func:`complete_task`. |
| 310 | error: |
| 311 | Optional human-readable error message set by :func:`fail_task`. |
| 312 | """ |
| 313 | |
| 314 | task_id: str |
| 315 | claimer_run_id: str |
| 316 | claimed_at: datetime.datetime |
| 317 | expires_at: datetime.datetime |
| 318 | status: str |
| 319 | heartbeat_at: datetime.datetime |
| 320 | claim_nonce: str |
| 321 | result: JsonDict | None |
| 322 | error: str | None |
| 323 | |
| 324 | def to_dict(self) -> JsonDict: |
| 325 | """Serialise to a JSON-compatible dict.""" |
| 326 | return { |
| 327 | "schema_version": __version__, |
| 328 | "task_id": self.task_id, |
| 329 | "claimer_run_id": self.claimer_run_id, |
| 330 | "claimed_at": self.claimed_at.isoformat(), |
| 331 | "expires_at": self.expires_at.isoformat(), |
| 332 | "status": self.status, |
| 333 | "heartbeat_at": self.heartbeat_at.isoformat(), |
| 334 | "claim_nonce": self.claim_nonce, |
| 335 | "result": self.result, |
| 336 | "error": self.error, |
| 337 | } |
| 338 | |
| 339 | @classmethod |
| 340 | def from_dict(cls, d: JsonDict) -> "ClaimRecord": |
| 341 | """Deserialise from a JSON-compatible dict.""" |
| 342 | return cls( |
| 343 | task_id=str(d.get("task_id", "")), |
| 344 | claimer_run_id=str(d.get("claimer_run_id", "")), |
| 345 | claimed_at=_parse_dt(str(d["claimed_at"])), |
| 346 | expires_at=_parse_dt(str(d["expires_at"])), |
| 347 | status=str(d.get("status", "claimed")), |
| 348 | heartbeat_at=_parse_dt(str(d.get("heartbeat_at", d["claimed_at"]))), |
| 349 | claim_nonce=str(d.get("claim_nonce", secrets.token_hex(16))), |
| 350 | result=d.get("result"), |
| 351 | error=d.get("error"), |
| 352 | ) |
| 353 | |
| 354 | def is_expired(self, now: datetime.datetime | None = None) -> bool: |
| 355 | """Return True if this claim's TTL has passed without renewal.""" |
| 356 | if now is None: |
| 357 | now = _now_utc() |
| 358 | return now >= self.expires_at |
| 359 | |
| 360 | # ── Loading helpers ─────────────────────────────────────────────────────────── |
| 361 | |
| 362 | def load_all_tasks(root: pathlib.Path) -> list[TaskRecord]: |
| 363 | """Return all task records from ``tasks/``. |
| 364 | |
| 365 | Corrupt or unreadable files are skipped with a warning. Expired tasks |
| 366 | are included — the caller decides whether to filter them. |
| 367 | """ |
| 368 | tdir = _tasks_dir(root) |
| 369 | if not tdir.exists(): |
| 370 | return [] |
| 371 | tasks: list[TaskRecord] = [] |
| 372 | for path in tdir.glob("*.json"): |
| 373 | raw = load_json_file(path) |
| 374 | if raw is None: |
| 375 | logger.warning("Corrupt task file %s", path.name) |
| 376 | continue |
| 377 | try: |
| 378 | tasks.append(TaskRecord.from_dict(raw)) |
| 379 | except (KeyError, ValueError) as exc: |
| 380 | logger.warning("Corrupt task file %s: %s", path.name, exc) |
| 381 | return tasks |
| 382 | |
| 383 | def load_all_claims(root: pathlib.Path) -> _ClaimMap: |
| 384 | """Return a mapping of ``task_id → ClaimRecord`` from ``claims/``. |
| 385 | |
| 386 | Corrupt or unreadable files are skipped with a warning. |
| 387 | """ |
| 388 | cdir = _claims_dir(root) |
| 389 | if not cdir.exists(): |
| 390 | return {} |
| 391 | claims: _ClaimMap = {} |
| 392 | for path in cdir.glob("*.json"): |
| 393 | raw = load_json_file(path) |
| 394 | if raw is None: |
| 395 | logger.warning("Corrupt claim file %s", path.name) |
| 396 | continue |
| 397 | try: |
| 398 | claim = ClaimRecord.from_dict(raw) |
| 399 | claims[claim.task_id] = claim |
| 400 | except (KeyError, ValueError) as exc: |
| 401 | logger.warning("Corrupt claim file %s: %s", path.name, exc) |
| 402 | return claims |
| 403 | |
| 404 | def load_task(root: pathlib.Path, task_id: str) -> TaskRecord | None: |
| 405 | """Load and return the task with *task_id*, or ``None`` if not found.""" |
| 406 | _validate_task_id(task_id) |
| 407 | path = _tasks_dir(root) / f"{task_id}.json" |
| 408 | raw = load_json_file(path) |
| 409 | if raw is None: |
| 410 | return None |
| 411 | try: |
| 412 | return TaskRecord.from_dict(raw) |
| 413 | except (KeyError, ValueError): |
| 414 | return None |
| 415 | |
| 416 | def load_claim(root: pathlib.Path, task_id: str) -> ClaimRecord | None: |
| 417 | """Load and return the claim for *task_id*, or ``None`` if unclaimed.""" |
| 418 | _validate_task_id(task_id) |
| 419 | path = _claims_dir(root) / f"{task_id}.json" |
| 420 | raw = load_json_file(path) |
| 421 | if raw is None: |
| 422 | return None |
| 423 | try: |
| 424 | return ClaimRecord.from_dict(raw) |
| 425 | except (KeyError, ValueError): |
| 426 | return None |
| 427 | |
| 428 | def get_task_status( |
| 429 | task: TaskRecord, |
| 430 | claim: ClaimRecord | None, |
| 431 | now: datetime.datetime | None = None, |
| 432 | ) -> str: |
| 433 | """Derive the current status string for *task* given its *claim*. |
| 434 | |
| 435 | Parameters |
| 436 | ---------- |
| 437 | task: |
| 438 | The task definition. |
| 439 | claim: |
| 440 | The claim record if one exists, else ``None``. |
| 441 | now: |
| 442 | Current UTC time. Defaults to :func:`_now_utc`. |
| 443 | |
| 444 | Returns |
| 445 | ------- |
| 446 | str |
| 447 | One of ``"pending"``, ``"claimed"``, ``"timed_out"``, |
| 448 | ``"completed"``, ``"failed"``, ``"cancelled"``. |
| 449 | """ |
| 450 | if now is None: |
| 451 | now = _now_utc() |
| 452 | if claim is None: |
| 453 | return "pending" |
| 454 | if claim.status == "claimed": |
| 455 | return "timed_out" if claim.is_expired(now) else "claimed" |
| 456 | return claim.status # completed | failed | cancelled | timed_out |
| 457 | |
| 458 | # ── Core operations ─────────────────────────────────────────────────────────── |
| 459 | |
| 460 | def compute_task_id( |
| 461 | title: str, |
| 462 | queue: str, |
| 463 | payload: "JsonDict | None", |
| 464 | priority: int, |
| 465 | created_by: str, |
| 466 | ) -> str: |
| 467 | """Return a content-addressed ``sha256:`` ID for a task genesis record. |
| 468 | |
| 469 | The same inputs always produce the same ID, making task creation |
| 470 | naturally idempotent for identical enqueue requests. |
| 471 | """ |
| 472 | return content_hash({ |
| 473 | "created_by": created_by, |
| 474 | "payload": payload or {}, |
| 475 | "priority": int(priority), |
| 476 | "queue": queue, |
| 477 | "title": title, |
| 478 | }) |
| 479 | |
| 480 | def create_task( |
| 481 | root: pathlib.Path, |
| 482 | title: str, |
| 483 | *, |
| 484 | payload: JsonDict | None = None, |
| 485 | priority: int = 0, |
| 486 | queue: str = "default", |
| 487 | ttl_seconds: int = 86400, |
| 488 | created_by: str = "unknown", |
| 489 | tags: list[str] | None = None, |
| 490 | ) -> TaskRecord: |
| 491 | """Enqueue a new task and return its :class:`TaskRecord`. |
| 492 | |
| 493 | Parameters |
| 494 | ---------- |
| 495 | root: |
| 496 | Repository root (directory containing ``.muse/``). |
| 497 | title: |
| 498 | Short human-readable description. Truncated to |
| 499 | :data:`_MAX_TITLE_LEN` characters. |
| 500 | payload: |
| 501 | Arbitrary JSON-serialisable dict attached to the task. Agents use |
| 502 | this to carry domain-specific parameters (addresses, operation type, |
| 503 | branch names, etc.). |
| 504 | priority: |
| 505 | Integer priority. **Higher = processed first** by |
| 506 | :func:`claim_next_task`. Default: ``0``. |
| 507 | queue: |
| 508 | Logical queue name. Must match ``[a-zA-Z0-9_-]+``. Default: |
| 509 | ``"default"``. |
| 510 | ttl_seconds: |
| 511 | Seconds from ``created_at`` after which the task is no longer |
| 512 | eligible for claiming. Default: 86 400 s (24 h). |
| 513 | created_by: |
| 514 | ``run_id`` of the enqueuing agent. |
| 515 | tags: |
| 516 | Optional list of string labels. |
| 517 | |
| 518 | Returns |
| 519 | ------- |
| 520 | TaskRecord |
| 521 | The newly created task. |
| 522 | |
| 523 | Raises |
| 524 | ------ |
| 525 | ValueError |
| 526 | If *queue* is not a valid queue name or *title* is empty. |
| 527 | """ |
| 528 | if not title: |
| 529 | raise ValueError("task title must be non-empty") |
| 530 | _validate_queue_name(queue) |
| 531 | ensure_task_dirs(root) |
| 532 | |
| 533 | safe_tags = [str(t)[:_MAX_TAG_LEN] for t in (tags or [])][:_MAX_TAGS] |
| 534 | |
| 535 | task = TaskRecord( |
| 536 | task_id=compute_task_id( |
| 537 | title=title[:_MAX_TITLE_LEN], |
| 538 | queue=queue, |
| 539 | payload=dict(payload or {}), |
| 540 | priority=int(priority), |
| 541 | created_by=str(created_by), |
| 542 | ), |
| 543 | title=title[:_MAX_TITLE_LEN], |
| 544 | payload=dict(payload or {}), |
| 545 | priority=int(priority), |
| 546 | queue=queue, |
| 547 | created_at=_now_utc(), |
| 548 | created_by=str(created_by), |
| 549 | ttl_seconds=max(1, int(ttl_seconds)), |
| 550 | tags=safe_tags, |
| 551 | ) |
| 552 | path = _tasks_dir(root) / f"{task.task_id}.json" |
| 553 | write_text_atomic(path, f"{json.dumps(task.to_dict(), indent=2)}\n") |
| 554 | logger.debug( |
| 555 | "Enqueued task %s (priority=%d, queue=%s, ttl=%ds)", |
| 556 | task.task_id[:8], task.priority, task.queue, task.ttl_seconds, |
| 557 | ) |
| 558 | return task |
| 559 | |
| 560 | def claim_next_task( |
| 561 | root: pathlib.Path, |
| 562 | run_id: str, |
| 563 | *, |
| 564 | queue: str | None = None, |
| 565 | claim_ttl_seconds: int = 3600, |
| 566 | ) -> tuple[TaskRecord, ClaimRecord] | None: |
| 567 | """Atomically claim the highest-priority pending task. |
| 568 | |
| 569 | Scans all tasks in *queue* (or all queues if ``None``), sorts by priority |
| 570 | (descending) then by ``created_at`` (ascending, FIFO within same priority), |
| 571 | and tries to claim each one in order. |
| 572 | |
| 573 | **Initial claim** (no existing claim file): uses ``O_CREAT | O_EXCL`` for |
| 574 | a guaranteed-atomic, winner-takes-all claim. |
| 575 | |
| 576 | **Re-claim of timed-out task** (existing claim with ``expires_at`` in the |
| 577 | past): uses ``write_text_atomic`` (atomic rename) followed by read-back |
| 578 | nonce verification. The agent whose nonce survives the rename wins. In |
| 579 | the rare case of a tie (two agents renaming within the same kernel |
| 580 | operation), the losing agent detects the mismatch on read-back and moves |
| 581 | to the next candidate. |
| 582 | |
| 583 | Parameters |
| 584 | ---------- |
| 585 | root: |
| 586 | Repository root. |
| 587 | run_id: |
| 588 | The claiming agent's identifier. |
| 589 | queue: |
| 590 | If set, only claim tasks from this queue. If ``None``, claim from |
| 591 | any queue (sorted globally by priority + FIFO). |
| 592 | claim_ttl_seconds: |
| 593 | How long (seconds) the claim is valid before another agent may |
| 594 | re-claim the task. Agents should call :func:`heartbeat_claim` |
| 595 | every ``claim_ttl_seconds / 2`` to prevent expiry. |
| 596 | |
| 597 | Returns |
| 598 | ------- |
| 599 | (TaskRecord, ClaimRecord) | None |
| 600 | The claimed task and its claim record, or ``None`` if the queue |
| 601 | contains no pending/claimable tasks. |
| 602 | """ |
| 603 | ensure_task_dirs(root) |
| 604 | now = _now_utc() |
| 605 | all_tasks = load_all_tasks(root) |
| 606 | all_claims = load_all_claims(root) |
| 607 | |
| 608 | # Build the candidate list: unclaimed OR timed-out, correct queue, not expired TTL. |
| 609 | candidates: list[TaskRecord] = [] |
| 610 | for task in all_tasks: |
| 611 | if queue is not None and task.queue != queue: |
| 612 | continue |
| 613 | if task.is_expired(now): |
| 614 | continue # Task TTL exhausted — skip. |
| 615 | claim = all_claims.get(task.task_id) |
| 616 | if claim is None: |
| 617 | candidates.append(task) # Unclaimed. |
| 618 | elif claim.status == "claimed" and claim.is_expired(now): |
| 619 | candidates.append(task) # Timed-out — eligible for re-claim. |
| 620 | |
| 621 | # Sort: highest priority first, then FIFO (earliest created_at). |
| 622 | candidates.sort(key=lambda t: (-t.priority, t.created_at)) |
| 623 | |
| 624 | for task in candidates: |
| 625 | existing = all_claims.get(task.task_id) |
| 626 | if existing is None: |
| 627 | # ── Initial claim via O_CREAT | O_EXCL ─────────────────────────── |
| 628 | result = _try_excl_claim(root, task.task_id, run_id, now, claim_ttl_seconds) |
| 629 | else: |
| 630 | # ── Re-claim of timed-out task via optimistic write + verify ───── |
| 631 | result = _try_optimistic_reclaim( |
| 632 | root, task.task_id, run_id, now, claim_ttl_seconds |
| 633 | ) |
| 634 | if result is not None: |
| 635 | logger.debug( |
| 636 | "Agent %s claimed task %s (queue=%s, priority=%d)", |
| 637 | run_id, task.task_id[:8], task.queue, task.priority, |
| 638 | ) |
| 639 | return task, result |
| 640 | |
| 641 | return None |
| 642 | |
| 643 | def _try_excl_claim( |
| 644 | root: pathlib.Path, |
| 645 | task_id: str, |
| 646 | run_id: str, |
| 647 | now: datetime.datetime, |
| 648 | claim_ttl_seconds: int, |
| 649 | ) -> ClaimRecord | None: |
| 650 | """Attempt to claim *task_id* via ``O_CREAT | O_EXCL``. |
| 651 | |
| 652 | Returns the :class:`ClaimRecord` on success, or ``None`` if another agent |
| 653 | claimed the task first (:exc:`FileExistsError`). |
| 654 | |
| 655 | The claim file is written with an fsync before close to ensure durability |
| 656 | even if the process is killed immediately after claiming. |
| 657 | """ |
| 658 | claim_path = _claims_dir(root) / f"{task_id}.json" |
| 659 | nonce = secrets.token_hex(16) |
| 660 | claim = ClaimRecord( |
| 661 | task_id=task_id, |
| 662 | claimer_run_id=run_id, |
| 663 | claimed_at=now, |
| 664 | expires_at=now + datetime.timedelta(seconds=claim_ttl_seconds), |
| 665 | status="claimed", |
| 666 | heartbeat_at=now, |
| 667 | claim_nonce=nonce, |
| 668 | result=None, |
| 669 | error=None, |
| 670 | ) |
| 671 | content = f"{json.dumps(claim.to_dict(), indent=2)}\n".encode() |
| 672 | try: |
| 673 | fd = os.open( |
| 674 | str(claim_path), |
| 675 | os.O_CREAT | os.O_EXCL | os.O_WRONLY, |
| 676 | 0o644, |
| 677 | ) |
| 678 | try: |
| 679 | os.write(fd, content) |
| 680 | try: |
| 681 | os.fsync(fd) |
| 682 | except OSError: |
| 683 | pass # Best-effort durability; atomicity still holds. |
| 684 | finally: |
| 685 | os.close(fd) |
| 686 | return claim |
| 687 | except FileExistsError: |
| 688 | return None # Another agent got there first. |
| 689 | |
| 690 | def _try_optimistic_reclaim( |
| 691 | root: pathlib.Path, |
| 692 | task_id: str, |
| 693 | run_id: str, |
| 694 | now: datetime.datetime, |
| 695 | claim_ttl_seconds: int, |
| 696 | ) -> ClaimRecord | None: |
| 697 | """Re-claim a timed-out task using optimistic concurrency. |
| 698 | |
| 699 | Overwrites the existing (expired) claim file via ``write_text_atomic`` |
| 700 | (atomic rename), then reads it back to verify the ``claim_nonce`` matches |
| 701 | what we wrote. The agent whose write survives the atomic rename wins. |
| 702 | |
| 703 | Returns the :class:`ClaimRecord` if we won the race, or ``None`` if |
| 704 | another agent's write arrived later. |
| 705 | |
| 706 | Security note |
| 707 | ------------- |
| 708 | This function does NOT use ``O_EXCL`` — it is possible (though unlikely on |
| 709 | a local filesystem) for two agents to both believe they won in a very tight |
| 710 | race window. The read-back verification closes this window to the duration |
| 711 | of a single file stat + read, which is microseconds on local storage. |
| 712 | Callers should treat this as best-effort at-most-once (not strict exactly- |
| 713 | once) for the re-claim case. |
| 714 | """ |
| 715 | claim_path = _claims_dir(root) / f"{task_id}.json" |
| 716 | nonce = secrets.token_hex(16) |
| 717 | claim = ClaimRecord( |
| 718 | task_id=task_id, |
| 719 | claimer_run_id=run_id, |
| 720 | claimed_at=now, |
| 721 | expires_at=now + datetime.timedelta(seconds=claim_ttl_seconds), |
| 722 | status="claimed", |
| 723 | heartbeat_at=now, |
| 724 | claim_nonce=nonce, |
| 725 | result=None, |
| 726 | error=None, |
| 727 | ) |
| 728 | write_text_atomic(claim_path, f"{json.dumps(claim.to_dict(), indent=2)}\n") |
| 729 | # Read back and verify nonce. |
| 730 | actual = load_json_file(claim_path) |
| 731 | if actual is not None and actual.get("claim_nonce") == nonce: |
| 732 | return claim |
| 733 | return None |
| 734 | |
| 735 | def complete_task( |
| 736 | root: pathlib.Path, |
| 737 | task_id: str, |
| 738 | run_id: str, |
| 739 | *, |
| 740 | result: JsonDict | None = None, |
| 741 | ) -> ClaimRecord: |
| 742 | """Mark a claimed task as completed. |
| 743 | |
| 744 | Parameters |
| 745 | ---------- |
| 746 | root: |
| 747 | Repository root. |
| 748 | task_id: |
| 749 | content-addressed ID of the task to complete. |
| 750 | run_id: |
| 751 | The claiming agent's ``run_id``. Must match ``claim.claimer_run_id``. |
| 752 | result: |
| 753 | Optional JSON-serialisable dict describing the outcome (e.g. proposal URL, |
| 754 | output summary, created artefacts). |
| 755 | |
| 756 | Returns |
| 757 | ------- |
| 758 | ClaimRecord |
| 759 | The updated claim record with ``status="completed"``. |
| 760 | |
| 761 | Raises |
| 762 | ------ |
| 763 | ValueError |
| 764 | If *task_id* is not a valid sha256: content ID. |
| 765 | FileNotFoundError |
| 766 | If the task or claim does not exist. |
| 767 | PermissionError |
| 768 | If *run_id* does not match the claimer. |
| 769 | RuntimeError |
| 770 | If the task is not in a claimable state (already completed, failed, etc.). |
| 771 | """ |
| 772 | _validate_task_id(task_id) |
| 773 | return _update_claim_status( |
| 774 | root, task_id, run_id, "completed", |
| 775 | result=result, error=None, |
| 776 | ) |
| 777 | |
| 778 | def fail_task( |
| 779 | root: pathlib.Path, |
| 780 | task_id: str, |
| 781 | run_id: str, |
| 782 | *, |
| 783 | error: str = "", |
| 784 | ) -> ClaimRecord: |
| 785 | """Mark a claimed task as failed. |
| 786 | |
| 787 | Parameters |
| 788 | ---------- |
| 789 | root: |
| 790 | Repository root. |
| 791 | task_id: |
| 792 | content-addressed ID of the task to fail. |
| 793 | run_id: |
| 794 | The claiming agent's ``run_id``. Must match ``claim.claimer_run_id``. |
| 795 | error: |
| 796 | Human-readable error message describing why the task failed. |
| 797 | |
| 798 | Returns |
| 799 | ------- |
| 800 | ClaimRecord |
| 801 | The updated claim record with ``status="failed"``. |
| 802 | |
| 803 | Raises |
| 804 | ------ |
| 805 | ValueError |
| 806 | If *task_id* is not a valid sha256: content ID. |
| 807 | FileNotFoundError |
| 808 | If the task or claim does not exist. |
| 809 | PermissionError |
| 810 | If *run_id* does not match the claimer. |
| 811 | RuntimeError |
| 812 | If the task is not in ``"claimed"`` status. |
| 813 | """ |
| 814 | _validate_task_id(task_id) |
| 815 | return _update_claim_status( |
| 816 | root, task_id, run_id, "failed", |
| 817 | result=None, error=str(error), |
| 818 | ) |
| 819 | |
| 820 | def cancel_task( |
| 821 | root: pathlib.Path, |
| 822 | task_id: str, |
| 823 | run_id: str, |
| 824 | *, |
| 825 | force: bool = False, |
| 826 | ) -> ClaimRecord: |
| 827 | """Cancel a pending or claimed task. |
| 828 | |
| 829 | A **pending** task (no claim file) is cancelled by creating a new claim |
| 830 | record directly with ``status="cancelled"`` via ``O_CREAT | O_EXCL``. |
| 831 | If another agent claims the task between the pending-check and the cancel |
| 832 | write, the cancel fails with :exc:`FileExistsError` (the task is now |
| 833 | claimed and must be cancelled via the claimer or ``force=True``). |
| 834 | |
| 835 | A **claimed** task may only be cancelled by its claimer (``run_id == |
| 836 | claim.claimer_run_id``) unless ``force=True``. |
| 837 | |
| 838 | Parameters |
| 839 | ---------- |
| 840 | root: |
| 841 | Repository root. |
| 842 | task_id: |
| 843 | content-addressed ID of the task to cancel. |
| 844 | run_id: |
| 845 | The calling agent's ``run_id``. |
| 846 | force: |
| 847 | If ``True``, cancel even if the task is claimed by a different agent. |
| 848 | Use with caution — this aborts in-flight work without notifying the |
| 849 | claimer. |
| 850 | |
| 851 | Returns |
| 852 | ------- |
| 853 | ClaimRecord |
| 854 | The resulting claim record with ``status="cancelled"``. |
| 855 | |
| 856 | Raises |
| 857 | ------ |
| 858 | ValueError |
| 859 | If *task_id* is not a valid sha256: content ID. |
| 860 | FileNotFoundError |
| 861 | If the task does not exist. |
| 862 | RuntimeError |
| 863 | If the task is already in a terminal state (completed, failed, cancelled). |
| 864 | PermissionError |
| 865 | If the task is claimed by a different agent and ``force=False``. |
| 866 | """ |
| 867 | _validate_task_id(task_id) |
| 868 | ensure_task_dirs(root) |
| 869 | |
| 870 | task_path = _tasks_dir(root) / f"{task_id}.json" |
| 871 | if not task_path.exists(): |
| 872 | raise FileNotFoundError(f"task not found: {task_id}") |
| 873 | |
| 874 | claim = load_claim(root, task_id) |
| 875 | now = _now_utc() |
| 876 | |
| 877 | if claim is None: |
| 878 | # Pending task — cancel it via O_EXCL (prevent race with a claiming agent). |
| 879 | cancel_claim = ClaimRecord( |
| 880 | task_id=task_id, |
| 881 | claimer_run_id=run_id, |
| 882 | claimed_at=now, |
| 883 | expires_at=now + datetime.timedelta(hours=24), |
| 884 | status="cancelled", |
| 885 | heartbeat_at=now, |
| 886 | claim_nonce=secrets.token_hex(16), |
| 887 | result=None, |
| 888 | error="cancelled before claiming", |
| 889 | ) |
| 890 | claim_path = _claims_dir(root) / f"{task_id}.json" |
| 891 | try: |
| 892 | fd = os.open(str(claim_path), os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o644) |
| 893 | try: |
| 894 | os.write(fd, f"{json.dumps(cancel_claim.to_dict(), indent=2)}\n".encode()) |
| 895 | try: |
| 896 | os.fsync(fd) |
| 897 | except OSError: |
| 898 | pass |
| 899 | finally: |
| 900 | os.close(fd) |
| 901 | except FileExistsError as exc: |
| 902 | raise FileExistsError( |
| 903 | f"task {task_id} was claimed by another agent before cancel; " |
| 904 | f"use force=True to cancel a claimed task" |
| 905 | ) from exc |
| 906 | return cancel_claim |
| 907 | |
| 908 | # Claimed / terminal task. |
| 909 | if claim.status in ("completed", "failed", "cancelled"): |
| 910 | raise RuntimeError( |
| 911 | f"task {task_id} is already in terminal state: {claim.status!r}" |
| 912 | ) |
| 913 | if not force and claim.claimer_run_id != run_id: |
| 914 | raise PermissionError( |
| 915 | f"task {task_id} is claimed by {claim.claimer_run_id!r}; " |
| 916 | f"only the claimer or force=True can cancel it" |
| 917 | ) |
| 918 | |
| 919 | return _update_claim_status( |
| 920 | root, task_id, run_id if force else claim.claimer_run_id, |
| 921 | "cancelled", result=None, error="cancelled", |
| 922 | _skip_ownership_check=True, |
| 923 | ) |
| 924 | |
| 925 | def heartbeat_claim( |
| 926 | root: pathlib.Path, |
| 927 | task_id: str, |
| 928 | run_id: str, |
| 929 | *, |
| 930 | extension_seconds: int = 3600, |
| 931 | ) -> ClaimRecord: |
| 932 | """Extend the claim TTL by writing a fresh heartbeat. |
| 933 | |
| 934 | Agents should call this every ``claim_ttl_seconds / 2`` to prevent their |
| 935 | claim from timing out. Only the current claimer may heartbeat a claim. |
| 936 | |
| 937 | Parameters |
| 938 | ---------- |
| 939 | root: |
| 940 | Repository root. |
| 941 | task_id: |
| 942 | content-addressed ID of the claimed task. |
| 943 | run_id: |
| 944 | The claiming agent's ``run_id``. Must match ``claim.claimer_run_id``. |
| 945 | extension_seconds: |
| 946 | How far into the future to set the new ``expires_at``. |
| 947 | Default: 3600 s (1 h). |
| 948 | |
| 949 | Returns |
| 950 | ------- |
| 951 | ClaimRecord |
| 952 | The updated claim record with refreshed ``heartbeat_at`` and |
| 953 | ``expires_at``. |
| 954 | |
| 955 | Raises |
| 956 | ------ |
| 957 | ValueError |
| 958 | If *task_id* is not a valid sha256: content ID. |
| 959 | FileNotFoundError |
| 960 | If no claim exists for *task_id*. |
| 961 | PermissionError |
| 962 | If *run_id* does not match the claimer. |
| 963 | RuntimeError |
| 964 | If the task is not in ``"claimed"`` status. |
| 965 | """ |
| 966 | _validate_task_id(task_id) |
| 967 | claim = load_claim(root, task_id) |
| 968 | if claim is None: |
| 969 | raise FileNotFoundError(f"no claim found for task {task_id}") |
| 970 | if claim.claimer_run_id != run_id: |
| 971 | raise PermissionError( |
| 972 | f"task {task_id} is claimed by {claim.claimer_run_id!r}, not {run_id!r}" |
| 973 | ) |
| 974 | if claim.status != "claimed": |
| 975 | raise RuntimeError( |
| 976 | f"cannot heartbeat task {task_id} with status {claim.status!r}" |
| 977 | ) |
| 978 | now = _now_utc() |
| 979 | claim.heartbeat_at = now |
| 980 | claim.expires_at = now + datetime.timedelta(seconds=extension_seconds) |
| 981 | claim_path = _claims_dir(root) / f"{task_id}.json" |
| 982 | write_text_atomic(claim_path, f"{json.dumps(claim.to_dict(), indent=2)}\n") |
| 983 | return claim |
| 984 | |
| 985 | # ── Private helpers ─────────────────────────────────────────────────────────── |
| 986 | |
| 987 | def _update_claim_status( |
| 988 | root: pathlib.Path, |
| 989 | task_id: str, |
| 990 | run_id: str, |
| 991 | new_status: str, |
| 992 | *, |
| 993 | result: JsonDict | None, |
| 994 | error: str | None, |
| 995 | _skip_ownership_check: bool = False, |
| 996 | ) -> ClaimRecord: |
| 997 | """Update a claim's status field via atomic rename. |
| 998 | |
| 999 | Internal helper used by :func:`complete_task`, :func:`fail_task`, and |
| 1000 | :func:`cancel_task`. |
| 1001 | |
| 1002 | Raises |
| 1003 | ------ |
| 1004 | FileNotFoundError |
| 1005 | If the task file or claim file does not exist. |
| 1006 | PermissionError |
| 1007 | If *run_id* does not match the claimer (unless *_skip_ownership_check*). |
| 1008 | RuntimeError |
| 1009 | If the claim is not in a state that permits the transition. |
| 1010 | """ |
| 1011 | task_path = _tasks_dir(root) / f"{task_id}.json" |
| 1012 | if not task_path.exists(): |
| 1013 | raise FileNotFoundError(f"task not found: {task_id}") |
| 1014 | |
| 1015 | claim = load_claim(root, task_id) |
| 1016 | if claim is None: |
| 1017 | raise FileNotFoundError(f"no claim found for task {task_id}") |
| 1018 | |
| 1019 | if not _skip_ownership_check and claim.claimer_run_id != run_id: |
| 1020 | raise PermissionError( |
| 1021 | f"task {task_id} is claimed by {claim.claimer_run_id!r}, not {run_id!r}" |
| 1022 | ) |
| 1023 | |
| 1024 | # Only "claimed" tasks can transition to completed/failed/cancelled. |
| 1025 | # (A timed_out task is still status="claimed" on disk — it's just expired.) |
| 1026 | if claim.status not in ("claimed",): |
| 1027 | raise RuntimeError( |
| 1028 | f"task {task_id} cannot transition from {claim.status!r} to {new_status!r}" |
| 1029 | ) |
| 1030 | |
| 1031 | claim.status = new_status |
| 1032 | claim.result = result |
| 1033 | claim.error = error |
| 1034 | |
| 1035 | claim_path = _claims_dir(root) / f"{task_id}.json" |
| 1036 | write_text_atomic(claim_path, f"{json.dumps(claim.to_dict(), indent=2)}\n") |
| 1037 | logger.debug("Task %s → %s (by %s)", task_id[:8], new_status, run_id) |
| 1038 | return claim |