"""``muse coord shard`` — partition a codebase into low-coupling work zones. Divides files into N groups (shards) such that files within each shard are tightly coupled to each other (many imports between them) and loosely coupled to files in other shards. Each shard is a safe parallel-work zone for agents. Agents assigned to different shards are unlikely to create merge conflicts because they work on different parts of the dependency graph. Algorithm --------- 1. Build the import graph from the committed snapshot. 2. Compute connected components of the import graph. 3. Distribute components greedily into N shards, balancing by symbol count. 4. Report each shard with its files, symbol count, and coupling score. The coupling score is the count of cross-shard import edges (lower = better). Usage:: muse coord shard --agents 4 muse coord shard --agents 8 --commit HEAD~10 muse coord shard --agents 4 --language Python muse coord shard --agents 4 --json Output (text):: Shard plan — 4 agents, commit a1b2c3d4 ────────────────────────────────────────────────────────────── Shard 1 (12 symbols, 3 files, coupling 1): src/billing.py src/billing_utils.py src/tax.py Shard 2 (8 symbols, 2 files, coupling 0): src/auth.py src/session.py ... Cross-shard edges: 1 JSON output schema:: { "schema_version": str, "commit": str, // 8-char short ID (display only) "full_commit_id": str, // full commit ID for cross-referencing "agents": int, "shards_created": int, "total_files": int, "total_symbols": int, "cross_shard_edges": int, "shards": [ { "shard": int, "files": [str, ...], "symbol_count": int, "coupling_score": int }, ... ], "elapsed_seconds": float } Flags: ``--agents N`` Number of parallel work zones to create (1–256, default: 4). ``--commit, -c REF`` Use a historical snapshot instead of HEAD. ``--language LANG`` Restrict to files of this language. ``--format`` / ``--json`` Emit the shard plan as compact JSON. Exit codes:: 0 — success (includes no-files-found case) 1 — bad arguments (--agents out of range) or commit not found """ from __future__ import annotations import argparse import json import logging import pathlib import sys import time from muse._version import __version__ from muse.core.errors import ExitCode from muse.core.object_store import read_object from muse.core.repo import read_repo_id, require_repo from muse.core.store import Manifest, get_commit_snapshot_manifest, read_current_branch, resolve_commit_ref from muse.plugins.code._query import language_of, symbols_for_snapshot from muse.plugins.code.ast_parser import parse_symbols from muse.core.validation import sanitize_display type _AdjMap = dict[str, set[str]] type _CounterMap = dict[str, int] type _ShardMap = dict[str, int] logger = logging.getLogger(__name__) # ── Input constraints ───────────────────────────────────────────────────────── #: Allowed range for ``--agents``. _MIN_AGENTS: int = 1 _MAX_AGENTS: int = 256 # ── Error helper ────────────────────────────────────────────────────────────── def _err(msg: str, as_json: bool, status: str = "error") -> None: """Print an error and return. Caller raises SystemExit.""" if as_json: print(json.dumps({"error": msg, "status": status})) else: print(f"❌ {msg}", file=sys.stderr) # ── Internal helpers ────────────────────────────────────────────────────────── def _file_stem(fp: str) -> str: return pathlib.PurePosixPath(fp).stem def _build_import_edges( root: pathlib.Path, manifest: Manifest, language_filter: str | None, ) -> list[tuple[str, str]]: """Return ``(importer, importee)`` file-path pairs from the snapshot. Reads each file's object from the object store, parses its import symbols, then resolves each imported name to a file path in the manifest via a stem-to-path lookup table. Only files matching *language_filter* (when provided) are included. Files whose object is missing from the store are silently skipped. Args: root: Repository root. manifest: Snapshot file-path → object-ID mapping. language_filter: Language name to restrict to, or ``None`` for all. Returns: List of ``(importer_path, importee_path)`` tuples. """ stem_to_file: Manifest = {} for fp in manifest: if language_filter and language_of(fp) != language_filter: continue stem_to_file[_file_stem(fp)] = fp edges: list[tuple[str, str]] = [] for file_path, obj_id in sorted(manifest.items()): if language_filter and language_of(file_path) != language_filter: continue raw = read_object(root, obj_id) if raw is None: continue tree = parse_symbols(raw, file_path) for rec in tree.values(): if rec["kind"] != "import": continue imported = rec["qualified_name"].split(".")[-1].replace("import::", "") target = stem_to_file.get(imported) if target and target != file_path: edges.append((file_path, target)) return edges def _connected_components( files: list[str], edges: list[tuple[str, str]], ) -> list[frozenset[str]]: """Return weakly-connected components of the import graph. Uses an iterative DFS to avoid stack overflow on large graphs. Args: files: All file paths (graph nodes). edges: ``(importer, importee)`` pairs (directed edges treated as undirected for connectivity purposes). Returns: List of :class:`frozenset` — one per connected component. """ adj: _AdjMap = {f: set() for f in files} for a, b in edges: adj.setdefault(a, set()).add(b) adj.setdefault(b, set()).add(a) visited: set[str] = set() components: list[frozenset[str]] = [] for start in files: if start in visited: continue component: set[str] = set() stack = [start] while stack: node = stack.pop() if node in visited: continue visited.add(node) component.add(node) for neighbour in adj.get(node, set()): if neighbour not in visited: stack.append(neighbour) components.append(frozenset(component)) return components def _greedy_partition( components: list[frozenset[str]], sym_counts: _CounterMap, n_shards: int, ) -> list[frozenset[str]]: """Distribute connected components into *n_shards* shards. Uses a greedy largest-first strategy: components are sorted by total symbol count (descending) and each is assigned to the shard with the currently smallest symbol total. This approximates the Longest Processing Time (LPT) algorithm for makespan minimisation. Args: components: Connected components to distribute. sym_counts: File-path → symbol count mapping. n_shards: Number of shards to create. Returns: List of *n_shards* :class:`frozenset` objects (some may be empty when there are fewer components than shards). """ shards: list[set[str]] = [set() for _ in range(n_shards)] shard_sizes: list[int] = [0] * n_shards # Sort components largest-first for better balance (LPT heuristic). sorted_comps = sorted( components, key=lambda c: sum(sym_counts.get(f, 0) for f in c), reverse=True, ) for comp in sorted_comps: smallest = min(range(n_shards), key=lambda i: shard_sizes[i]) shards[smallest].update(comp) shard_sizes[smallest] += sum(sym_counts.get(f, 0) for f in comp) return [frozenset(s) for s in shards] # ── CLI registration ────────────────────────────────────────────────────────── def register(subparsers: "argparse._SubParsersAction[argparse.ArgumentParser]") -> None: """Register the ``shard`` subcommand on *subparsers* (under ``muse coord``). Wires all flags with their defaults, choices, and help text so that ``--help`` output is accurate. Sets ``func`` to :func:`run`. """ parser = subparsers.add_parser( "shard", help="Partition the codebase into N low-coupling work zones for parallel agents.", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter, ) parser.add_argument( "--agents", "-n", type=int, default=4, metavar="N", help=f"Number of work zones ({_MIN_AGENTS}–{_MAX_AGENTS}, default: 4).", ) parser.add_argument( "--commit", "-c", dest="ref", default=None, metavar="REF", help="Use this commit instead of HEAD.", ) parser.add_argument( "--language", "-l", default=None, metavar="LANG", help="Restrict to files of this language.", ) parser.add_argument( "--format", "-f", default="text", dest="fmt", choices=("text", "json"), help="Output format: text (default) or json.", ) parser.add_argument( "--json", action="store_const", const="json", dest="fmt", help="Shorthand for --format json.", ) parser.set_defaults(func=run) # ── Command implementation ──────────────────────────────────────────────────── def run(args: argparse.Namespace) -> None: """Partition the codebase into N low-coupling work zones for parallel agents. Uses the import graph connectivity to group files that are tightly coupled together and loosely coupled to other shards. Agents assigned to different shards minimise merge conflicts. Execution order --------------- 1. **Validate inputs** — ``--agents`` range checked before any file I/O. Failure exits :attr:`~muse.core.errors.ExitCode.USER_ERROR` (1) with a message on *stderr* or compact JSON on *stdout* when ``--format json``. 2. **Resolve repo** — :func:`~muse.core.repo.require_repo`. 3. **Resolve commit** — :func:`~muse.core.store.resolve_commit_ref`. 4. **Load symbol map** — :func:`~muse.plugins.code._query.symbols_for_snapshot`. 5. **Build import graph** — :func:`_build_import_edges` reads each file's object from the store and parses import symbols. 6. **Partition** — :func:`_connected_components` then :func:`_greedy_partition`. 7. **Emit output** — compact JSON or human-readable text. Security -------- * ``--agents`` is validated against :data:`_MIN_AGENTS` / :data:`_MAX_AGENTS` before any I/O. * ``--commit`` refs are sanitised by :func:`~muse.core.store.resolve_commit_ref`. * ``--language`` values are passed through :func:`~muse.core.validation.sanitize_display` before text output to prevent ANSI injection. * All file paths in text output pass through :func:`~muse.core.validation.sanitize_display`. * Error messages go to *stdout* as compact JSON when ``--format json``, or to *stderr* with an ``❌`` prefix otherwise. Performance ----------- * Import-graph build is O(F × S) where F is the number of files in the manifest and S is the average number of symbols per file. * Connected-components DFS is O(F + E) where E is the number of import edges. * Greedy partition is O(C log C + C × N) where C is the number of components and N the number of shards. * Shards are balanced by symbol count using the LPT heuristic (largest-processing-time first), which gives a ≤ 4/3 OPT approximation. Args: args: Parsed ``argparse.Namespace`` with attributes ``agents``, ``ref``, ``language``, and ``fmt``. Exit codes: 0 — success (no files found is also success). 1 — bad arguments or commit not found. """ as_json: bool = args.fmt == "json" t0 = time.monotonic() # ── Input validation (before any file I/O) ──────────────────────────────── agents_raw: int = args.agents if not (_MIN_AGENTS <= agents_raw <= _MAX_AGENTS): msg = f"--agents must be {_MIN_AGENTS}–{_MAX_AGENTS}, got {agents_raw}" _err(msg, as_json, "bad_args") raise SystemExit(ExitCode.USER_ERROR) agents: int = agents_raw ref: str | None = args.ref language: str | None = args.language root = require_repo() repo_id = read_repo_id(root) branch = read_current_branch(root) commit = resolve_commit_ref(root, repo_id, branch, ref) if commit is None: msg = f"commit '{sanitize_display(ref or 'HEAD')}' not found" _err(msg, as_json, "commit_not_found") raise SystemExit(ExitCode.USER_ERROR) manifest = get_commit_snapshot_manifest(root, commit.commit_id) or {} sym_map = symbols_for_snapshot(root, manifest, language_filter=language) sym_counts: _CounterMap = {fp: len(tree) for fp, tree in sym_map.items()} files = sorted(sym_counts.keys()) if not files: elapsed = round(time.monotonic() - t0, 4) if as_json: print(json.dumps({ "schema_version": __version__, "commit": commit.commit_id[:8], "full_commit_id": commit.commit_id, "agents": agents, "shards_created": 0, "total_files": 0, "total_symbols": 0, "cross_shard_edges": 0, "shards": [], "elapsed_seconds": elapsed, })) else: print(" (no semantic files found in snapshot)") return edges = _build_import_edges(root, manifest, language) components = _connected_components(files, edges) n = min(agents, len(components)) # Can't have more shards than components. shard_sets = _greedy_partition(components, sym_counts, n) # Compute cross-shard edges. file_to_shard: _ShardMap = {} for i, s in enumerate(shard_sets): for f in s: file_to_shard[f] = i cross_edges = sum( 1 for a, b in edges if file_to_shard.get(a, -1) != file_to_shard.get(b, -2) ) total_symbols = sum(sym_counts.values()) elapsed = round(time.monotonic() - t0, 4) if as_json: print(json.dumps({ "schema_version": __version__, "commit": commit.commit_id[:8], "full_commit_id": commit.commit_id, "agents": agents, "shards_created": n, "total_files": len(files), "total_symbols": total_symbols, "cross_shard_edges": cross_edges, "shards": [ { "shard": i + 1, "files": sorted(s), "symbol_count": sum(sym_counts.get(f, 0) for f in s), "coupling_score": sum( 1 for a, b in edges if (a in s) != (b in s) ), } for i, s in enumerate(shard_sets) if s ], "elapsed_seconds": elapsed, })) return # ── Text output ─────────────────────────────────────────────────────────── print(f"\nShard plan — {n} agent(s), commit {commit.commit_id[:8]}") if language: print(f" (language: {sanitize_display(language)})") print("─" * 62) for i, s in enumerate(shard_sets): if not s: continue sym_total = sum(sym_counts.get(f, 0) for f in s) coupling = sum(1 for a, b in edges if (a in s) != (b in s)) print(f"\n Shard {i + 1} ({sym_total} symbols, {len(s)} files, coupling {coupling}):") for fp in sorted(s): print(f" {sanitize_display(fp)}") print(f"\n Cross-shard edges: {cross_edges}") if cross_edges == 0: print(" ✅ Perfect isolation — no cross-shard dependencies") print(f"\n ({elapsed:.3f}s)")