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