aaronrene / gabriel-muse public
dag.py python
611 lines 21.4 KB
Raw
sha256:cf6265cea8c21d9228d90dec13ef6ec2dab5103d466db9cc4590681832de4bf8 docs(KD-STAGING): sync governance after KD-6b DONE Human 12 days ago
1 """Dependency DAG for multi-agent coordination reservations.
2
3 Agents working in parallel often have strict ordering constraints: agent B
4 cannot safely start its work until agent A completes and releases its
5 reservation. This module implements a lightweight **dependency DAG** that
6 layers on top of the existing reservation system to express and evaluate those
7 constraints without any central coordinator.
8
9 Layout
10 ------
11 ::
12
13 .muse/coordination/
14 dependencies/<reservation-id>.json — dependency record (write-once)
15
16 Dependency schema
17 -----------------
18 ::
19
20 {
21 "schema_version": "<muse package version>",
22 "reservation_id": "<uuid>", // the *dependent* reservation
23 "depends_on": ["<uuid>", ...], // must complete (release/expire) first
24 "created_at": "2026-03-30T12:00:00+00:00"
25 }
26
27 A reservation is **blocked** when at least one of its declared dependencies is
28 still *active* (not yet released and not yet expired). Agents poll
29 :func:`is_blocked` to decide whether to start work or wait.
30
31 Graph model
32 -----------
33 Edges are directed from *dependent → dependency* ("A depends on B" is the edge
34 ``A → B``). Topological order therefore lists dependencies before dependents —
35 the correct execution sequence for a pipeline.
36
37 DAG invariants
38 --------------
39 * **Acyclicity** is enforced at write time by :func:`add_dependencies`.
40 Any call that would introduce a cycle raises :exc:`ValueError` immediately,
41 so a cycle can never exist on disk.
42 * **Self-dependency** (``reservation_id in depends_on``) is rejected.
43 * **Write-once**: each reservation has at most one dependency record.
44 Calling :func:`add_dependencies` twice for the same ID raises
45 :exc:`FileExistsError`.
46 * **UUID-safe**: every ID is validated before constructing any file path,
47 preventing directory traversal attacks.
48
49 Security
50 --------
51 Every reservation ID supplied by callers is validated against a UUID4 regex
52 (same pattern as :func:`~muse.core.coordination._validate_reservation_id`).
53 A valid UUID cannot contain ``/`` or ``..``, ruling out path traversal.
54
55 Performance
56 -----------
57 * :func:`load_all_dependencies` — one directory scan: O(n).
58 * :func:`detect_cycle` — iterative DFS: O(V + E).
59 * :func:`topological_sort` — Kahn's algorithm: O(V + E).
60 * :func:`is_blocked` / :func:`get_blocking` — O(degree) set intersection.
61
62 All graph operations are pure in-memory after the initial load.
63 """
64
65 from __future__ import annotations
66
67 import datetime
68 import json
69 import logging
70 import pathlib
71 import re
72 from collections import deque
73 from typing import Sequence
74
75 from muse._version import __version__
76 from muse.core.store import write_text_atomic
77
78 # JSON-compatible value type for dependency record serialization.
79 JsonValue = str | int | float | bool | None | list["JsonValue"] | "JsonDict"
80 JsonDict = dict[str, JsonValue]
81
82 type AdjacencyMap = dict[str, set[str]] # node → set of direct dependencies
83 type ColorMap = dict[str, int] # node → DFS color (WHITE/GREY/BLACK)
84 type ParentMap = dict[str, str | None] # node → parent in DFS traversal
85 type InDegreeMap = dict[str, int] # node → in-degree count
86
87 logger = logging.getLogger(__name__)
88
89 # ── Constants ─────────────────────────────────────────────────────────────────
90
91 # UUID4 regex — identical to coordination._UUID_RE for consistency.
92 _UUID_RE = re.compile(
93 r"^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$",
94 re.IGNORECASE,
95 )
96
97 # Maximum number of direct dependencies per reservation.
98 _MAX_DEPS: int = 256
99
100 # ── Directory helpers ─────────────────────────────────────────────────────────
101
102
103 def _coord_dir(root: pathlib.Path) -> pathlib.Path:
104 """Return ``.muse/coordination/`` for *root*."""
105 return root / ".muse" / "coordination"
106
107
108 def _dependencies_dir(root: pathlib.Path) -> pathlib.Path:
109 """Return ``.muse/coordination/dependencies/`` for *root*."""
110 return _coord_dir(root) / "dependencies"
111
112
113 def ensure_dag_dirs(root: pathlib.Path) -> None:
114 """Create the dependencies directory if it does not exist."""
115 _dependencies_dir(root).mkdir(parents=True, exist_ok=True)
116
117
118 # ── Validation ────────────────────────────────────────────────────────────────
119
120
121 def _validate_id(rid: str) -> None:
122 """Raise :exc:`ValueError` if *rid* is not a well-formed UUID4.
123
124 Security note
125 -------------
126 ``rid`` is used to construct file paths inside
127 ``.muse/coordination/dependencies/``. A valid UUID4 cannot contain
128 ``/`` or ``..``, preventing directory traversal attacks.
129 """
130 if not _UUID_RE.match(str(rid)):
131 raise ValueError(f"reservation_id must be a valid UUID4: {rid!r}")
132
133
134 # ── Time helpers ──────────────────────────────────────────────────────────────
135
136
137 def _now_utc() -> datetime.datetime:
138 return datetime.datetime.now(datetime.timezone.utc)
139
140
141 def _parse_dt(s: str) -> datetime.datetime:
142 return datetime.datetime.fromisoformat(s)
143
144
145 # ── Data model ────────────────────────────────────────────────────────────────
146
147
148 class DependencyRecord:
149 """Stores the direct dependency edges for one reservation.
150
151 A :class:`DependencyRecord` is **write-once**: once created for a given
152 ``reservation_id`` it is never modified. It declares that
153 ``reservation_id`` must wait for every ID in ``depends_on`` to be released
154 or expired before starting work.
155
156 Attributes
157 ----------
158 reservation_id:
159 The *dependent* reservation — the one that must wait.
160 depends_on:
161 List of reservation IDs that must complete (be released or expire)
162 before the dependent reservation is unblocked. De-duplicated and
163 order-preserving at write time.
164 created_at:
165 UTC time when this dependency record was written to disk.
166 """
167
168 def __init__(
169 self,
170 reservation_id: str,
171 depends_on: list[str],
172 created_at: datetime.datetime,
173 ) -> None:
174 self.reservation_id = reservation_id
175 self.depends_on = depends_on
176 self.created_at = created_at
177
178 def to_dict(self) -> JsonDict:
179 """Serialise to a JSON-compatible dict."""
180 return {
181 "schema_version": __version__,
182 "reservation_id": self.reservation_id,
183 "depends_on": self.depends_on,
184 "created_at": self.created_at.isoformat(),
185 }
186
187 @classmethod
188 def from_dict(cls, d: JsonDict) -> "DependencyRecord":
189 """Deserialise from a JSON-compatible dict.
190
191 Missing fields default to safe values so records written by older
192 versions remain loadable.
193 """
194 created_raw = d.get("created_at")
195 created_at = _parse_dt(str(created_raw)) if created_raw else _now_utc()
196 deps_raw = d.get("depends_on", [])
197 depends_on = (
198 [str(x) for x in deps_raw if x]
199 if isinstance(deps_raw, list)
200 else []
201 )
202 return cls(
203 reservation_id=str(d.get("reservation_id", "")),
204 depends_on=depends_on,
205 created_at=created_at,
206 )
207
208
209 # ── Loading helpers ───────────────────────────────────────────────────────────
210
211
212 def load_all_dependencies(root: pathlib.Path) -> list[DependencyRecord]:
213 """Return all dependency records from ``.muse/coordination/dependencies/``.
214
215 Corrupt or unreadable files are skipped with a warning — this mirrors the
216 resilient loading strategy used throughout the coordination layer.
217 """
218 ddir = _dependencies_dir(root)
219 if not ddir.exists():
220 return []
221 records: list[DependencyRecord] = []
222 for path in ddir.glob("*.json"):
223 try:
224 raw = json.loads(path.read_text())
225 records.append(DependencyRecord.from_dict(raw))
226 except (json.JSONDecodeError, KeyError, ValueError) as exc:
227 logger.warning("Corrupt dependency file %s: %s", path.name, exc)
228 return records
229
230
231 def load_dependencies(
232 root: pathlib.Path,
233 reservation_id: str,
234 ) -> DependencyRecord | None:
235 """Load the dependency record for *reservation_id*, or ``None`` if absent.
236
237 Parameters
238 ----------
239 root:
240 Repository root.
241 reservation_id:
242 UUID4 of the reservation to look up.
243
244 Raises
245 ------
246 ValueError
247 If *reservation_id* is not a valid UUID4.
248 """
249 _validate_id(reservation_id)
250 path = _dependencies_dir(root) / f"{reservation_id}.json"
251 try:
252 return DependencyRecord.from_dict(json.loads(path.read_text()))
253 except (OSError, json.JSONDecodeError, KeyError, ValueError):
254 return None
255
256
257 # ── Graph construction ────────────────────────────────────────────────────────
258
259
260 def build_graph(deps: Sequence[DependencyRecord]) -> AdjacencyMap:
261 """Return an adjacency map ``node → {direct_dependencies}``.
262
263 Every node that appears (either as a dependent or as a dependency) is
264 included in the returned graph with an explicit entry, so the set of graph
265 nodes is complete. Nodes that are referenced only as dependencies (no
266 record of their own) get an empty dependency set.
267
268 Parameters
269 ----------
270 deps:
271 Sequence of :class:`DependencyRecord` objects, typically from
272 :func:`load_all_dependencies`.
273
274 Returns
275 -------
276 dict[str, set[str]]
277 Adjacency map: ``node → {ids this node directly depends on}``.
278 """
279 graph: AdjacencyMap = {}
280 for rec in deps:
281 if rec.reservation_id not in graph:
282 graph[rec.reservation_id] = set()
283 for dep in rec.depends_on:
284 graph[rec.reservation_id].add(dep)
285 if dep not in graph:
286 graph[dep] = set()
287 return graph
288
289
290 # ── Cycle detection ───────────────────────────────────────────────────────────
291
292
293 def detect_cycle(graph: AdjacencyMap) -> list[str] | None:
294 """Return a cycle path if *graph* contains a directed cycle, else ``None``.
295
296 Uses iterative DFS with three-colour marking:
297
298 * **WHITE** (0) — not yet visited.
299 * **GREY** (1) — currently on the DFS stack (back edge → cycle).
300 * **BLACK** (2) — fully explored.
301
302 Returns
303 -------
304 list[str] | None
305 A list of reservation IDs forming the cycle (first == last), or
306 ``None`` if the graph is acyclic.
307
308 Notes
309 -----
310 Nodes that appear in the ``depends_on`` lists but have no record of their
311 own are treated as leaf nodes with no outgoing edges — they cannot be part
312 of a cycle.
313 """
314 WHITE, GREY, BLACK = 0, 1, 2
315 color: ColorMap = {n: WHITE for n in graph}
316 parent: ParentMap = {n: None for n in graph}
317
318 def _reconstruct(node: str) -> list[str]:
319 """Walk parent pointers back to *node* to extract the cycle."""
320 path = [node]
321 cur = parent.get(node)
322 while cur is not None and cur != node:
323 path.append(cur)
324 cur = parent.get(cur)
325 path.append(node)
326 path.reverse()
327 return path
328
329 for start in list(graph):
330 if color.get(start, WHITE) != WHITE:
331 continue
332 stack = [(start, iter(graph.get(start, set())))]
333 color[start] = GREY
334 while stack:
335 node, children = stack[-1]
336 try:
337 child = next(children)
338 if child not in color:
339 # Leaf node (dependency with no record) — cannot cycle.
340 continue
341 if color[child] == GREY:
342 # Back edge detected — cycle found.
343 parent[child] = node
344 return _reconstruct(child)
345 if color[child] == WHITE:
346 color[child] = GREY
347 parent[child] = node
348 stack.append((child, iter(graph.get(child, set()))))
349 except StopIteration:
350 color[node] = BLACK
351 stack.pop()
352
353 return None
354
355
356 # ── Topological sort ──────────────────────────────────────────────────────────
357
358
359 def topological_sort(
360 graph: AdjacencyMap,
361 nodes: Sequence[str] | None = None,
362 ) -> list[str]:
363 """Return nodes in topological order using Kahn's algorithm.
364
365 Dependencies appear **before** their dependents in the returned list —
366 i.e. the correct execution order for a pipeline where each step must wait
367 for its declared predecessors.
368
369 Parameters
370 ----------
371 graph:
372 Adjacency map from :func:`build_graph`.
373 nodes:
374 Optional subset of node IDs to sort. When supplied, the function
375 restricts to the subgraph reachable from those nodes (including all
376 transitive dependencies). When ``None``, the full graph is sorted.
377
378 Returns
379 -------
380 list[str]
381 Node IDs in execution order (dependencies first).
382
383 Raises
384 ------
385 ValueError
386 If the graph contains a cycle. Call :func:`detect_cycle` first to
387 get a human-readable description of the cycle.
388 """
389 if nodes is not None:
390 # Restrict to the subgraph reachable (downward through deps) from
391 # the requested nodes.
392 visited: set[str] = set()
393 queue: deque[str] = deque(nodes)
394 while queue:
395 n = queue.popleft()
396 if n in visited:
397 continue
398 visited.add(n)
399 for dep in graph.get(n, set()):
400 queue.append(dep)
401 sub: AdjacencyMap = {
402 n: {d for d in graph.get(n, set()) if d in visited}
403 for n in visited
404 }
405 else:
406 sub = {n: set(deps) for n, deps in graph.items()}
407
408 # Build in-degree and reverse adjacency map (dep → {dependents}).
409 in_degree: InDegreeMap = {n: len(sub[n]) for n in sub}
410 reverse: AdjacencyMap = {n: set() for n in sub}
411 for n, deps in sub.items():
412 for dep in deps:
413 if dep in reverse:
414 reverse[dep].add(n)
415
416 # Kahn's: repeatedly dequeue nodes with no remaining prerequisites.
417 ready: deque[str] = deque(
418 sorted(n for n in sub if in_degree[n] == 0)
419 )
420 result: list[str] = []
421 while ready:
422 n = ready.popleft()
423 result.append(n)
424 for dependent in sorted(reverse.get(n, set())):
425 in_degree[dependent] -= 1
426 if in_degree[dependent] == 0:
427 ready.append(dependent)
428
429 if len(result) != len(sub):
430 raise ValueError(
431 "DAG contains a cycle — call detect_cycle() for details"
432 )
433 return result
434
435
436 # ── Blocking status ───────────────────────────────────────────────────────────
437
438
439 def is_blocked(
440 reservation_id: str,
441 graph: AdjacencyMap,
442 active_ids: frozenset[str],
443 ) -> bool:
444 """Return ``True`` if *reservation_id* has at least one active dependency.
445
446 A reservation is *blocked* when any of its direct dependencies are still
447 present in *active_ids* (i.e. not yet released and not yet expired).
448
449 Parameters
450 ----------
451 reservation_id:
452 The UUID of the reservation to check.
453 graph:
454 Adjacency map from :func:`build_graph`.
455 active_ids:
456 Set of currently active reservation IDs, typically derived from
457 :func:`~muse.core.coordination.active_reservations`.
458
459 Returns
460 -------
461 bool
462 ``True`` if any direct dependency is still active.
463 """
464 return bool(graph.get(reservation_id, set()) & active_ids)
465
466
467 def get_blocking(
468 reservation_id: str,
469 graph: AdjacencyMap,
470 active_ids: frozenset[str],
471 ) -> list[str]:
472 """Return the direct dependencies of *reservation_id* that are still active.
473
474 Parameters
475 ----------
476 reservation_id:
477 The UUID of the reservation to check.
478 graph:
479 Adjacency map from :func:`build_graph`.
480 active_ids:
481 Set of currently active reservation IDs.
482
483 Returns
484 -------
485 list[str]
486 Sorted list of active dependency IDs that are blocking
487 *reservation_id*. Empty when the reservation is unblocked.
488 """
489 return sorted(graph.get(reservation_id, set()) & active_ids)
490
491
492 # ── Write: add_dependencies ───────────────────────────────────────────────────
493
494
495 def add_dependencies(
496 root: pathlib.Path,
497 reservation_id: str,
498 depends_on: list[str],
499 ) -> DependencyRecord:
500 """Declare that *reservation_id* depends on each ID in *depends_on*.
501
502 This is the sole write function for the dependency layer. It validates
503 all IDs, enforces the acyclicity invariant against the current on-disk
504 graph, and writes the record atomically.
505
506 Parameters
507 ----------
508 root:
509 Repository root (directory containing ``.muse/``).
510 reservation_id:
511 UUID4 of the *dependent* reservation — the one that must wait.
512 depends_on:
513 List of UUID4 reservation IDs that must be released or expire before
514 *reservation_id* is unblocked. Duplicates are silently de-duplicated
515 (order preserved). An empty list is accepted and produces a no-op
516 dependency record that can serve as an explicit audit marker.
517
518 Returns
519 -------
520 DependencyRecord
521 The newly created dependency record.
522
523 Raises
524 ------
525 ValueError
526 If any ID is not a valid UUID4, if ``reservation_id in depends_on``
527 (self-dependency), if more than :data:`_MAX_DEPS` dependencies are
528 supplied, or if the new edge(s) would create a cycle in the existing
529 DAG.
530 FileExistsError
531 If a dependency record for *reservation_id* already exists.
532 Dependency declarations are write-once — use :func:`load_dependencies`
533 to check first if idempotency is needed.
534
535 Security
536 --------
537 Every ID is validated as a UUID4 before any file path is constructed.
538 The write is atomic (``write_text_atomic`` uses mkstemp → fsync → rename),
539 ensuring no partial state is visible to concurrent readers.
540 """
541 _validate_id(reservation_id)
542 for dep_id in depends_on:
543 _validate_id(dep_id)
544
545 # Deduplicate while preserving order.
546 unique_deps: list[str] = list(dict.fromkeys(depends_on))
547
548 if reservation_id in unique_deps:
549 raise ValueError(
550 f"reservation_id {reservation_id!r} cannot depend on itself"
551 )
552 if len(unique_deps) > _MAX_DEPS:
553 raise ValueError(
554 f"too many dependencies: {len(unique_deps)} > {_MAX_DEPS}"
555 )
556
557 ensure_dag_dirs(root)
558
559 # ── Cycle check ───────────────────────────────────────────────────────────
560 # Build the current on-disk graph and tentatively add the new edges before
561 # writing anything. If detect_cycle reports a cycle, reject the write.
562 existing_deps = load_all_dependencies(root)
563 graph = build_graph(existing_deps)
564
565 if reservation_id not in graph:
566 graph[reservation_id] = set()
567 for dep_id in unique_deps:
568 graph[reservation_id].add(dep_id)
569 if dep_id not in graph:
570 graph[dep_id] = set()
571
572 cycle = detect_cycle(graph)
573 if cycle is not None:
574 raise ValueError(
575 f"adding dependency would create a cycle: {' → '.join(cycle)}"
576 )
577
578 # ── Write-once guard ──────────────────────────────────────────────────────
579 dep_path = _dependencies_dir(root) / f"{reservation_id}.json"
580 if dep_path.exists():
581 raise FileExistsError(
582 f"dependency record for {reservation_id!r} already exists; "
583 "dependency declarations are write-once"
584 )
585
586 now = _now_utc()
587 record = DependencyRecord(
588 reservation_id=reservation_id,
589 depends_on=unique_deps,
590 created_at=now,
591 )
592 write_text_atomic(dep_path, json.dumps(record.to_dict(), indent=2) + "\n")
593 logger.debug(
594 "Added %d dependency edge(s) for reservation %s",
595 len(unique_deps),
596 reservation_id[:8],
597 )
598 return record
599
600
601 # ── High-level convenience helper ─────────────────────────────────────────────
602
603
604 def load_dag(root: pathlib.Path) -> AdjacencyMap:
605 """Load all dependency records and return the full adjacency graph.
606
607 Convenience wrapper around :func:`load_all_dependencies` +
608 :func:`build_graph`. Returns an empty dict when no dependency records
609 exist yet.
610 """
611 return build_graph(load_all_dependencies(root))
File History 1 commit
sha256:cf6265cea8c21d9228d90dec13ef6ec2dab5103d466db9cc4590681832de4bf8 docs(KD-STAGING): sync governance after KD-6b DONE Human 12 days ago