orchestrata_ultimate.agents.swarm.coordinator

Reference › orchestrata_ultimate › agents › swarm › coordinator

coordinator in the swarm package fans code generation out file by file. SwarmCoordinator is created with a workspace path, a retry limit (default 2), the task summary, the API spec, the original intent and optional structured requirements. The CodeDirector's implementation helpers (implementation) call it after the scaffold has written a skeleton; results are ShardResult objects from types.

execute_swarm keeps the source files of the skeleton (_is_source_file), registers each Python file's exports in the knowledge-graph registry (_sync_file_to_registry), builds one prompt per file with _build_shard_prompt, and runs all shards concurrently in an asyncio TaskGroup. It returns the shard results and a scan of every file in the workspace. Each shard runs in an isolated temporary copy of the workspace (_create_shard_workspace, _copy_workspace) and uses a junior-volk driver obtained from the driver factory. After a successful generation the file is copied back atomically (_copy_file_back) and checked in order: syntax, lint, KG registration and a deterministic review (_review_file, which covers @traces presence, interface compliance against the API spec and requirement coverage). Any failure builds a fix prompt (_build_fix_prompt) and retries until the retry limit, after which the shard returns a failed ShardResult with the error. _execute_shard_safe converts any exception into a failed result so one shard cannot cancel the others.

At a glance

Source src/orchestrata_ultimate/agents/swarm/coordinator.py
Lines of code 1030
Classes 1
Functions 2
Imports (internal) orchestrata_ultimate.agents.swarm.types, orchestrata_ultimate.drivers.driver_factory
Imported by orchestrata_ultimate.agents.swarm, orchestrata_ultimate.agents.volks.codedirector.implementation, orchestrata_ultimate.agents.volks.codedirector.volk
Tests that import it tests/agents/swarm/test_coordinator.py, tests/agents/swarm/test_coordinator_extended.py, tests/agents/swarm/test_coordinator_kg_registry_concurrency_regression.py, tests/agents/test_swarm_coordinator_real.py, tests/generated_agents/test_swarm_coordinator_generated.py, tests/generated_recovery/test_agents_swarm_coordinator.py

Class diagram

Class diagram orchestrata_ultimate.agents.swarm.coordinator

Dependency diagram

Dependency diagram orchestrata_ultimate.agents.swarm.coordinator

API

Classes

SwarmCoordinator

Line 91

Async parallel per-file code generation coordinator.

Method Signature Summary
__init__ __init__(workspace: str, max_retries: int = 2, task_summary: str = '', api_spec: str = '', original_intent: str = '', structured_requirements: list | None = None) -> None Initialize the swarm coordinator.
execute_swarm async execute_swarm(skeleton_files: list[str]) -> tuple[list[ShardResult], list[str]] Execute parallel code generation for all source files.
_build_shard_prompt _build_shard_prompt(rel_path: str) -> str Build the coding prompt for a single shard.
_get_file_requirements _get_file_requirements(rel_path: str) -> str G222: Get requirement IDs relevant to this file.
_match_requirements_for_basename _match_requirements_for_basename(basename: str) -> list[str] Select requirement lines whose text references basename.
_all_requirement_ids_section _all_requirement_ids_section() -> str Build the fallback prompt section listing all requirement IDs.
_execute_shard_safe async _execute_shard_safe(shard_id: int, rel_path: str, abs_path: str, prompt: str, total: int) -> ShardResult Run _execute_shard, converting any exception into a failed ShardResult instead of letting it propagate.
_execute_shard async _execute_shard(shard_id: int, rel_path: str, abs_path: str, prompt: str) -> ShardResult Execute a single shard with retry.
_create_shard_workspace async _create_shard_workspace(shard_id: int) -> str Create an isolated temp workspace for a shard.
_copy_file_back async _copy_file_back(shard_dir: str, rel_path: str) -> None Atomically copy a generated file from shard to main workspace.
_review_file async _review_file(rel_path: str) -> list[str] G220: Run deterministic review on a single file immediately.
_check_traceability _check_traceability(rel_path: str, source: str) -> list[str] G221: Check that all public functions/classes have @traces.
_missing_trace_decorator_issues _missing_trace_decorator_issues(rel_path: str, tree: ast.AST) -> list[str] List public module-level defs/classes missing a @traces decorator.
_has_traces_decorator _has_traces_decorator(node: object) -> bool Check if an AST node has a @traces decorator.
_check_interface_compliance _check_interface_compliance(rel_path: str, source: str) -> list[str] Check that generated code matches the API spec signatures.
_expected_names_from_spec _expected_names_from_spec(basename: str) -> list[str] Parse the API spec for def/class names declared under basename.
_check_syntax async _check_syntax(rel_path: str) -> tuple[bool, str] G227 Gate 1: Check if file parses as valid Python.
_check_lint async _check_lint(rel_path: str) -> list[str] G227 Gate 2: Run linter on the file.
_sync_file_to_registry async _sync_file_to_registry(rel_path: str) -> None G227: Parse file and register its exports in the KG registry.
_get_kg_import_context _get_kg_import_context() -> str G227: Build import context from KG registry.
_build_fix_prompt _build_fix_prompt(rel_path: str, issues_text: str) -> str Build a fix prompt with current file content and issues.
_scan_workspace _scan_workspace() -> list[str] Scan the workspace for all source files.

Functions

Function Signature Summary
_is_source_file _is_source_file(filepath: str) -> bool Check if a file is a swarm shard candidate.
_copy_workspace async _copy_workspace(src: str, dst: str) -> None Copy workspace to a temp directory for shard isolation.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/agents/swarm/coordinator.py

# mypy: ignore-errors
"""Async swarm coordinator for parallel per-file code generation.

Fans out N source files across N concurrent asyncio tasks, each
running an async code generation driver in an isolated temporary
workspace. Results are atomically merged back to the main workspace.

Under Python 3.14t (free-threaded), the asyncio tasks run with
true parallelism for the Python-side work (prompt building, result
parsing, file I/O), while subprocess drivers (Claude Code, Aider)
are always parallel regardless of GIL status.
"""

import ast
import asyncio
import logging
import os
import shutil
import subprocess
import threading

from orchestrata_ultimate.agents.swarm.types import ShardResult

_logger = logging.getLogger(__name__)

_SKIP_DIRS = frozenset(
    {
        ".git",
        ".venv",
        "__pycache__",
        "node_modules",
        ".claude",
    }
)
_SOURCE_EXTS = frozenset({".py"})
_SKIP_FILES = frozenset(
    {
        "__init__.py",
        "conftest.py",
    }
)


def _is_source_file(filepath: str) -> bool:
    """Check if a file is a swarm shard candidate.

    Args:
        filepath: Relative file path.

    Returns:
        True if the file is a source file eligible for shard gen.
    """
    basename = os.path.basename(filepath)
    ext = os.path.splitext(basename)[1]
    if ext not in _SOURCE_EXTS:
        return False
    if basename in _SKIP_FILES:
        return False
    if basename.startswith("test_"):
        return False
    return True


async def _copy_workspace(
    src: str,
    dst: str,
) -> None:
    """Copy workspace to a temp directory for shard isolation.

    Args:
        src: Source workspace directory.
        dst: Destination temp directory.
    """

    def _copy() -> None:
        shutil.copytree(
            src,
            dst,
            ignore=shutil.ignore_patterns(
                ".git",
                ".venv",
                "__pycache__",
                "node_modules",
            ),
            dirs_exist_ok=True,
        )

    await asyncio.to_thread(_copy)


class SwarmCoordinator:
    """Async parallel per-file code generation coordinator.

    Fans out source files across N concurrent tasks, each with an
    isolated temporary workspace. Files are atomically copied back
    to the main workspace on success.

    Attributes:
        workspace: Main project workspace directory.
        max_retries: Maximum retry attempts per shard.
        task_summary: High-level description of the coding task.
        api_spec: API specification text for context.
        original_intent: Original user intent text.
    """

    def __init__(
        self,
        workspace: str,
        max_retries: int = 2,
        task_summary: str = "",
        api_spec: str = "",
        original_intent: str = "",
        structured_requirements: list | None = None,
    ) -> None:
        """Initialize the swarm coordinator.

        Args:
            workspace: Main project workspace directory.
            max_retries: Max retries per shard (default 2).
            task_summary: Task summary for shard prompts.
            api_spec: API spec for context in shard prompts.
            original_intent: Original user intent.
            structured_requirements: List of requirement dicts
                from the Analyst volk, each with 'id' and 'text'.
        """
        self.workspace = workspace
        self.max_retries = max_retries
        self.task_summary = task_summary
        self.api_spec = api_spec
        self.original_intent = original_intent
        self.structured_requirements = structured_requirements or []
        # G227: Live KG registry — tracks synced modules/classes/paths
        # during swarm execution. Each shard that passes syntax+lint
        # gets registered here so subsequent shards know what exists.
        # Guarded by a threading.Lock: under 3.14t free-threading, shard
        # tasks write (_sync_file_to_registry) and read
        # (_get_kg_import_context) the dict with true parallelism.
        self._kg_registry: dict[str, dict] = {}
        self._kg_registry_lock = threading.Lock()

    async def execute_swarm(
        self,
        skeleton_files: list[str],
    ) -> tuple[list[ShardResult], list[str]]:
        """Execute parallel code generation for all source files.

        Args:
            skeleton_files: List of relative file paths from scaffold.

        Returns:
            Tuple of (shard results, all files in workspace).
        """
        source_files = [f for f in skeleton_files if _is_source_file(f)]
        total = len(source_files)

        if total == 0:
            _logger.warning("No source files for swarm")
            return [], self._scan_workspace()

        _logger.info("Starting %d parallel shards", total)

        # G227: Pre-sync scaffold files to KG registry so shards
        # know what already exists (exception classes, __init__, etc.)
        for f in skeleton_files:
            if f.endswith(".py"):
                await self._sync_file_to_registry(f)

        # Build shard prompts
        shards = []
        for i, rel_path in enumerate(source_files):
            prompt = self._build_shard_prompt(rel_path)
            abs_path = os.path.join(self.workspace, rel_path)
            shards.append((i, rel_path, abs_path, prompt))

        # Fan out with asyncio.TaskGroup. Each shard is wrapped by
        # _execute_shard_safe() so a single shard's exception cannot
        # propagate through TaskGroup.__aexit__ — an uncaught exception in
        # ANY task would cancel every sibling task and abort the whole
        # swarm run, defeating the per-file failure isolation this method
        # is documented to provide.
        results: list[ShardResult] = []
        async with asyncio.TaskGroup() as tg:
            tasks = {}
            for shard_id, rel_path, abs_path, prompt in shards:
                task = tg.create_task(
                    self._execute_shard_safe(
                        shard_id, rel_path, abs_path, prompt, total
                    )
                )
                tasks[task] = (shard_id, rel_path)

        for task, (shard_id, rel_path) in tasks.items():
            result = task.result()
            results.append(result)
            status = "OK" if result.success else "FAIL"
            _logger.info(
                "Shard %d/%d (%s): %s",
                shard_id,
                total,
                rel_path,
                status,
            )

        all_files = await asyncio.to_thread(self._scan_workspace)
        succeeded = sum(1 for r in results if r.success)
        _logger.info(
            "Complete: %d/%d shards OK, %d total files",
            succeeded,
            total,
            len(all_files),
        )

        return results, all_files

    def _build_shard_prompt(self, rel_path: str) -> str:
        """Build the coding prompt for a single shard.

        Args:
            rel_path: Relative path of the file to generate.

        Returns:
            Complete prompt string.
        """
        abs_path = os.path.join(self.workspace, rel_path)
        skeleton = ""
        if os.path.exists(abs_path):
            try:
                with open(abs_path, encoding="utf-8") as f:
                    skeleton = f.read()
            except Exception as exc:
                _logger.debug("read skeleton %s: %s", abs_path, exc)

        parts = [
            f"Generate the COMPLETE implementation for: {rel_path}",
            "\nMANDATORY CODE STANDARDS (Python 3.14t):",
            "- Type hints on ALL function parameters and return types",
            "  def func(name: str, count: int = 0) -> list[str]:",
            "- Use modern syntax: list[], dict[], X | None (not Optional)",
            "- Google-style docstrings on every public class and function",
            "- @traces decorator on every public class and function",
            "- Specific exception types (never bare except:)",
            "- Max 40 lines per function, max 5 parameters",
            "- snake_case functions, PascalCase classes, UPPER_CASE constants",
            "- No mutable default arguments (use None + if param is None)",
            "- Validate EVERY parameter in EVERY public method:",
            "  if not isinstance(x, str): raise TypeError(...)",
            "  if x is None: raise ValueError('x must not be None')",
            "  if x < 0: raise ValueError('x must be non-negative')",
        ]
        if self.task_summary:
            parts.append(f"\nTask: {self.task_summary}")
        if skeleton:
            parts.append(f"\nSkeleton:\n{skeleton}")
        if self.api_spec:
            parts.append(f"\nAPI Spec (read-only reference):\n{self.api_spec}")

        # G222: Inject requirement IDs this file implements
        req_section = self._get_file_requirements(rel_path)
        if req_section:
            parts.append(req_section)

        # G227: Inject KG import context — what's already written
        kg_context = self._get_kg_import_context()
        if kg_context:
            parts.append(kg_context)

        parts.append(
            f"\nCRITICAL: Write ONLY the file {rel_path}. "
            "Do not create or modify any other file."
        )
        return "\n".join(parts)

    def _get_file_requirements(self, rel_path: str) -> str:
        """G222: Get requirement IDs relevant to this file.

        Maps requirements to files by matching module names in
        the requirement text against the file basename.

        Args:
            rel_path: Relative file path being generated.

        Returns:
            Formatted requirement section for the prompt, or empty.
        """
        if not self.structured_requirements:
            return ""

        basename = os.path.basename(rel_path).removesuffix(".py")
        matched = self._match_requirements_for_basename(basename)

        if not matched:
            return self._all_requirement_ids_section()

        return (
            "\nRequirements this file implements:\n"
            + "\n".join(matched)
            + "\n\nAdd @traces(spec='001', story='US-001', "
            "reqs=[<matching IDs>]) to every public function "
            "and class implementing these requirements.\n"
            "If orchestrata.core.traceability is not available, "
            "define:\n"
            "  def traces(**kw): return lambda f: f"
        )

    def _match_requirements_for_basename(self, basename: str) -> list[str]:
        """Select requirement lines whose text references *basename*.

        Args:
            basename: Module basename (without .py) to match against.

        Returns:
            Formatted requirement lines that match this file.
        """
        matched: list[str] = []
        for req in self.structured_requirements:
            if isinstance(req, dict):
                rid = req.get("id", req.get("req_id", ""))
                text = req.get("text", req.get("description", ""))
                # Match if the module name appears in the requirement
                if basename in text.lower() or not text:
                    matched.append(f"  {rid}: {text[:100]}")
            elif isinstance(req, str) and basename in req.lower():
                matched.append(f"  {req}")
        return matched

    def _all_requirement_ids_section(self) -> str:
        """Build the fallback prompt section listing all requirement IDs.

        Returns:
            Formatted requirement-ID reference section, or '' if none.
        """
        all_ids = []
        for req in self.structured_requirements[:10]:
            if isinstance(req, dict):
                rid = req.get("id", req.get("req_id", ""))
                if rid:
                    all_ids.append(rid)
        if not all_ids:
            return ""
        return (
            f"\nRequirement IDs (for @traces decorator):\n"
            f"  {', '.join(all_ids)}\n"
            "Add @traces(spec='001', story='US-001', "
            f"reqs={all_ids[:3]}) to every public "
            "function and class.\n"
            "If orchestrata.core.traceability is not "
            "available, define:\n"
            "  def traces(**kw): return lambda f: f"
        )

    async def _execute_shard_safe(
        self,
        shard_id: int,
        rel_path: str,
        abs_path: str,
        prompt: str,
        total: int,
    ) -> ShardResult:
        """Run ``_execute_shard``, converting any exception into a failed
        ``ShardResult`` instead of letting it propagate.

        Required because ``asyncio.TaskGroup`` cancels every sibling task
        and re-raises on exit as soon as ANY task raises — an uncaught
        exception here would abort the entire swarm run instead of just
        this one shard, defeating per-file failure isolation.

        Args:
            shard_id: Numeric shard identifier.
            rel_path: Relative file path.
            abs_path: Absolute file path.
            prompt: The coding prompt.
            total: Total number of shards (for logging).

        Returns:
            ShardResult with success/failure info.
        """
        try:
            return await self._execute_shard(shard_id, rel_path, abs_path, prompt)
        except Exception as exc:
            _logger.error(
                "Shard %d/%d (%s) exception: %s",
                shard_id,
                total,
                rel_path,
                exc,
            )
            return ShardResult(
                file_path=rel_path,
                success=False,
                error=str(exc),
            )

    async def _execute_shard(
        self,
        shard_id: int,
        rel_path: str,
        abs_path: str,
        prompt: str,
    ) -> ShardResult:
        """Execute a single shard with retry.

        Args:
            shard_id: Numeric shard identifier.
            rel_path: Relative file path.
            abs_path: Absolute file path.
            prompt: The coding prompt.

        Returns:
            ShardResult with success/failure info.
        """
        current_prompt = prompt
        for attempt in range(self.max_retries + 1):
            shard_dir = await self._create_shard_workspace(shard_id)
            try:
                from orchestrata_ultimate.drivers.driver_factory import get_driver

                driver = get_driver(
                    volk="J:Junior",
                    workspace=shard_dir,
                )
                success, output = await driver.execute(current_prompt, files=[rel_path])
                if success:
                    await self._copy_file_back(shard_dir, rel_path)

                    # G227 Gate 1: Syntax check immediately
                    syntax_ok, syntax_err = await self._check_syntax(rel_path)
                    if not syntax_ok:
                        if attempt < self.max_retries:
                            current_prompt = self._build_fix_prompt(
                                rel_path,
                                f"SYNTAX ERROR: {syntax_err}\n"
                                "Fix the syntax error. The file must "
                                "parse as valid Python.",
                            )
                            _logger.info(
                                "Shard %d: syntax error, retrying",
                                shard_id,
                            )
                            continue
                        return ShardResult(
                            file_path=rel_path,
                            success=False,
                            error=f"SYNTAX ERROR: {syntax_err}",
                            retries=attempt,
                        )

                    # G227 Gate 2: Lint check
                    lint_issues = await self._check_lint(rel_path)
                    if lint_issues:
                        if attempt < self.max_retries:
                            current_prompt = self._build_fix_prompt(
                                rel_path,
                                "LINT ERRORS:\n"
                                + "\n".join(f"  - {i}" for i in lint_issues)
                                + "\nFix ALL lint errors.",
                            )
                            _logger.info(
                                "Shard %d: %d lint issues, retrying",
                                shard_id,
                                len(lint_issues),
                            )
                            continue
                        return ShardResult(
                            file_path=rel_path,
                            success=False,
                            error=f"LINT ERRORS: {'; '.join(lint_issues)}",
                            retries=attempt,
                        )

                    # G227: Syntax+lint passed → sync to KG registry
                    await self._sync_file_to_registry(rel_path)

                    # G220: Full review (traceability, quality)
                    issues = await self._review_file(rel_path)
                    if issues:
                        if attempt < self.max_retries:
                            issue_list = "\n".join(f"  - {i}" for i in issues)
                            current_prompt = self._build_fix_prompt(
                                rel_path,
                                f"CODE REVIEW ISSUES:\n{issue_list}\nFix ALL issues.",
                            )
                            _logger.info(
                                "Shard %d: %d review issues, retrying",
                                shard_id,
                                len(issues),
                            )
                            continue
                        return ShardResult(
                            file_path=rel_path,
                            success=False,
                            error=f"CODE REVIEW ISSUES: {'; '.join(issues)}",
                            retries=attempt,
                        )
                    return ShardResult(
                        file_path=rel_path,
                        success=True,
                        output=output,
                        retries=attempt,
                    )
            except Exception as exc:
                _logger.warning(
                    "Shard %d attempt %d failed: %s",
                    shard_id,
                    attempt + 1,
                    exc,
                )
            finally:
                await asyncio.to_thread(shutil.rmtree, shard_dir, True)

        return ShardResult(
            file_path=rel_path,
            success=False,
            error=f"Failed after {self.max_retries + 1} attempts",
            retries=self.max_retries,
        )

    async def _create_shard_workspace(self, shard_id: int) -> str:
        """Create an isolated temp workspace for a shard.

        Args:
            shard_id: Shard identifier for naming.

        Returns:
            Path to the temp workspace.
        """
        import tempfile

        shard_dir = os.path.join(
            tempfile.gettempdir(),
            f"orchestrata-shard-{shard_id}",
        )
        if os.path.exists(shard_dir):
            await asyncio.to_thread(shutil.rmtree, shard_dir, True)
        await _copy_workspace(self.workspace, shard_dir)
        return shard_dir

    async def _copy_file_back(
        self,
        shard_dir: str,
        rel_path: str,
    ) -> None:
        """Atomically copy a generated file from shard to main workspace.

        Args:
            shard_dir: Shard temp directory.
            rel_path: Relative file path to copy.
        """
        src = os.path.join(shard_dir, rel_path)
        dst = os.path.join(self.workspace, rel_path)

        if not os.path.exists(src):
            return

        def _copy() -> None:
            os.makedirs(os.path.dirname(dst), exist_ok=True)
            if os.path.exists(dst) and os.path.samefile(src, dst):
                _logger.debug("skip copy — same file: %s", rel_path)
                return
            shutil.copy2(src, dst)

        await asyncio.to_thread(_copy)

    async def _review_file(self, rel_path: str) -> list[str]:
        """G220: Run deterministic review on a single file immediately.

        Also checks interface compliance against the API spec.

        Args:
            rel_path: Relative path to the generated file.

        Returns:
            List of issue strings. Empty if file passes review.
        """
        issues: list[str] = []
        abs_path = os.path.join(self.workspace, rel_path)

        if not abs_path.endswith(".py") or not os.path.exists(abs_path):
            return issues

        # Syntax check
        try:
            with open(abs_path, encoding="utf-8") as f:
                source = f.read()
            compile(source, rel_path, "exec")
        except SyntaxError as exc:
            issues.append(f"SyntaxError: {exc}")
            return issues

        # Interface compliance: check that functions in API spec exist
        if self.api_spec:
            spec_issues = await asyncio.to_thread(
                self._check_interface_compliance, rel_path, source
            )
            issues.extend(spec_issues)

        # Import check: verify no spike/orchestrata imports
        try:
            import ast

            tree = ast.parse(source)
            for node in ast.walk(tree):
                if isinstance(node, ast.Import):
                    for alias in node.names:
                        mod = alias.name.split(".")[0]
                        if mod in {"spike", "orchestrata"}:
                            issues.append(
                                f"Import '{alias.name}' — is expected to "
                                "not import spike/orchestrata in "
                                "generated code"
                            )
        except SyntaxError as exc:
            _logger.debug("AST parse failed for %s: %s", rel_path, exc)

        # G221: Traceability enforcement — every public function
        # and class must have a @traces decorator
        trace_issues = await asyncio.to_thread(
            self._check_traceability, rel_path, source
        )
        issues.extend(trace_issues)

        return issues

    def _check_traceability(
        self,
        rel_path: str,
        source: str,
    ) -> list[str]:
        """G221: Check that all public functions/classes have @traces.

        Args:
            rel_path: File being checked.
            source: Source code content.

        Returns:
            List of traceability issues.
        """
        import ast

        basename = os.path.basename(rel_path)

        # Skip __init__.py, conftest.py, test files
        if basename in {"__init__.py", "conftest.py"}:
            return []
        if basename.startswith("test_"):
            return []

        try:
            tree = ast.parse(source)
        except SyntaxError:
            return []

        # Check if @traces is imported or defined
        if "traces" not in source:
            return [
                f"TRACEABILITY: {rel_path} has no @traces decorator "
                "import. Add: from orchestrata_ultimate.core.tracing.traceability "
                "import traces (or define: def traces(**kw): "
                "return lambda f: f)"
            ]

        return self._missing_trace_decorator_issues(rel_path, tree)

    def _missing_trace_decorator_issues(
        self,
        rel_path: str,
        tree: ast.AST,
    ) -> list[str]:
        """List public module-level defs/classes missing a @traces decorator.

        Args:
            rel_path: File being checked (for issue messages).
            tree: Parsed AST module.

        Returns:
            One issue string per public def/class without @traces.
        """
        import ast

        issues: list[str] = []
        for node in ast.iter_child_nodes(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                kind = "function"
            elif isinstance(node, ast.ClassDef):
                kind = "class"
            else:
                continue
            if node.name.startswith("_"):
                continue
            if not self._has_traces_decorator(node):
                issues.append(
                    f"TRACEABILITY: {rel_path}:{node.lineno} "
                    f"{kind} '{node.name}' missing @traces"
                )
        return issues

    def _has_traces_decorator(self, node: object) -> bool:
        """Check if an AST node has a @traces decorator.

        Args:
            node: AST FunctionDef, AsyncFunctionDef, or ClassDef.

        Returns:
            True if @traces is present.
        """
        import ast

        for dec in getattr(node, "decorator_list", []):
            if isinstance(dec, ast.Call):
                func = dec.func
                if isinstance(func, ast.Name) and func.id == "traces":
                    return True
                if isinstance(func, ast.Attribute) and func.attr == "traces":
                    return True
            elif isinstance(dec, ast.Name) and dec.id == "traces":
                return True
            elif isinstance(dec, ast.Attribute) and dec.attr == "traces":
                return True
        return False

    def _check_interface_compliance(
        self,
        rel_path: str,
        source: str,
    ) -> list[str]:
        """Check that generated code matches the API spec signatures.

        Args:
            rel_path: File being checked.
            source: Source code content.

        Returns:
            List of compliance issues.
        """
        import ast

        basename = os.path.basename(rel_path).removesuffix(".py")

        # Extract function/class names from source
        try:
            tree = ast.parse(source)
        except SyntaxError:
            return []

        defined_names = set()
        for node in ast.walk(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                defined_names.add(node.name)
            elif isinstance(node, ast.ClassDef):
                defined_names.add(node.name)

        # Check API spec for expected functions/classes in this file
        if not self.api_spec:
            return []

        expected_names = self._expected_names_from_spec(basename)
        return [
            f"Interface mismatch: {name} expected in "
            f"{rel_path} (from API spec) but not defined"
            for name in expected_names
            if name not in defined_names
        ]

    def _expected_names_from_spec(self, basename: str) -> list[str]:
        """Parse the API spec for def/class names declared under *basename*.

        Scans the spec line by line, entering a file section when *basename*
        appears in a header line and collecting the def/class names that follow
        until the section ends.

        Args:
            basename: Module basename (without .py) marking the spec section.

        Returns:
            Expected function/class names declared for this file.
        """
        in_file_section = False
        expected_names: list[str] = []
        for line in self.api_spec.splitlines():
            stripped = line.strip()
            if basename in stripped and (
                "def " in stripped
                or "class " in stripped
                or stripped.endswith(".py")
                or stripped.endswith(":")
            ):
                in_file_section = True
                continue
            if not in_file_section:
                continue
            if stripped.startswith("def "):
                expected_names.append(
                    stripped.split("(")[0].replace("def ", "").strip()
                )
            elif stripped.startswith("class "):
                expected_names.append(
                    stripped.split("(")[0].split(":")[0].replace("class ", "").strip()
                )
            elif stripped and not stripped.startswith((" ", "\t", "#", '"""', "'''")):
                in_file_section = False
        return expected_names

    # ── G227: Post-write gates + KG sync ─────────────────────────

    async def _check_syntax(
        self,
        rel_path: str,
    ) -> tuple[bool, str]:
        """G227 Gate 1: Check if file parses as valid Python.

        Args:
            rel_path: Relative file path.

        Returns:
            (True, "") if syntax OK, (False, error_msg) if not.
        """
        filepath = os.path.join(self.workspace, rel_path)
        if not filepath.endswith(".py") or not os.path.exists(filepath):
            return True, ""

        def _parse() -> tuple[bool, str]:
            try:
                with open(filepath, encoding="utf-8") as f:
                    source = f.read()
                ast.parse(source, filename=rel_path)
                return True, ""
            except SyntaxError as exc:
                return False, (
                    f"Line {exc.lineno}: {exc.msg}" if exc.lineno else str(exc)
                )

        return await asyncio.to_thread(_parse)

    async def _check_lint(
        self,
        rel_path: str,
    ) -> list[str]:
        """G227 Gate 2: Run linter on the file.

        Tries ruff, then flake8. Returns issue strings.

        Args:
            rel_path: Relative file path.

        Returns:
            List of lint issue strings. Empty = clean.
        """
        filepath = os.path.join(self.workspace, rel_path)
        if not filepath.endswith(".py") or not os.path.exists(filepath):
            return []

        def _lint() -> list[str]:
            for linter in ["ruff check", "flake8"]:
                cmd_parts = linter.split() + [filepath]
                try:
                    result = subprocess.run(
                        cmd_parts,
                        capture_output=True,
                        text=True,
                        timeout=15,
                    )
                    if result.returncode != 0 and result.stdout:
                        return [
                            line.strip()
                            for line in result.stdout.strip().splitlines()
                            if line.strip()
                        ]
                    return []
                except FileNotFoundError, subprocess.TimeoutExpired:
                    continue
            return []

        return await asyncio.to_thread(_lint)

    async def _sync_file_to_registry(
        self,
        rel_path: str,
    ) -> None:
        """G227: Parse file and register its exports in the KG registry.

        Extracts classes, functions, and their import paths so
        subsequent shards can get correct import context.

        Args:
            rel_path: Relative file path (already syntax-checked).
        """
        filepath = os.path.join(self.workspace, rel_path)
        if not filepath.endswith(".py") or not os.path.exists(filepath):
            return

        def _extract() -> dict:
            with open(filepath, encoding="utf-8") as f:
                source = f.read()
            tree = ast.parse(source)

            classes = []
            functions = []
            for node in ast.iter_child_nodes(tree):
                if isinstance(node, ast.ClassDef):
                    methods = [
                        n.name
                        for n in ast.iter_child_nodes(node)
                        if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef))
                    ]
                    classes.append(
                        {
                            "name": node.name,
                            "methods": methods,
                            "line": node.lineno,
                        }
                    )
                elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                    functions.append(
                        {
                            "name": node.name,
                            "line": node.lineno,
                        }
                    )

            # Build import path from relative path
            # src/pkg/module.py → pkg.module
            import_path = rel_path.removesuffix(".py").replace(os.sep, ".")
            if import_path.startswith("src."):
                import_path = import_path[4:]

            return {
                "rel_path": rel_path,
                "import_path": import_path,
                "classes": classes,
                "functions": functions,
            }

        info = await asyncio.to_thread(_extract)
        with self._kg_registry_lock:
            self._kg_registry[rel_path] = info
        _logger.info(
            "KG sync: %s → %d classes, %d functions",
            rel_path,
            len(info["classes"]),
            len(info["functions"]),
        )

    def _get_kg_import_context(self) -> str:
        """G227: Build import context from KG registry.

        Returns a formatted string showing what modules, classes,
        and functions already exist with their correct import paths.

        Returns:
            Import context string for injection into shard prompts.
        """
        # Snapshot under the lock so a concurrent _sync_file_to_registry
        # write cannot mutate the dict mid-iteration (3.14t parallelism).
        with self._kg_registry_lock:
            snapshot = list(self._kg_registry.items())
        if not snapshot:
            return ""

        lines = [
            "\nAVAILABLE MODULES (already written — use these exact import paths):"
        ]
        for rel_path, info in sorted(snapshot):
            imp = info["import_path"]
            lines.append(f"\n  # {rel_path}")
            for cls in info["classes"]:
                lines.append(f"  from {imp} import {cls['name']}")
            for func in info["functions"]:
                if not func["name"].startswith("_"):
                    lines.append(f"  from {imp} import {func['name']}")

        lines.append(
            "\nIMPORTANT: Use EXACTLY these import paths. "
            "Do NOT guess or invent import paths."
        )
        return "\n".join(lines)

    def _build_fix_prompt(
        self,
        rel_path: str,
        issues_text: str,
    ) -> str:
        """Build a fix prompt with current file content and issues.

        Args:
            rel_path: Relative file path.
            issues_text: Formatted issue description.

        Returns:
            Complete fix prompt string.
        """
        filepath = os.path.join(self.workspace, rel_path)
        current_source = ""
        if os.path.exists(filepath):
            try:
                with open(filepath, encoding="utf-8") as f:
                    current_source = f.read()
            except OSError as exc:
                _logger.debug("read for fix: %s", exc)

        # Include KG import context so fixes use correct paths
        kg_context = self._get_kg_import_context()

        return (
            f"The file {rel_path} has issues that MUST be fixed:\n"
            f"{issues_text}\n\n"
            "REMINDER — mandatory on every function:\n"
            "- Type hints: def f(x: str, n: int = 0) -> list[str]:\n"
            "- Use list[], dict[], X | None (Python 3.14 syntax)\n"
            "- Google docstring on every public class/function\n\n"
            f"Current file content:\n"
            f"```python\n{current_source}\n```\n"
            f"{kg_context}\n\n"
            f"Fix ALL listed issues. Output the complete "
            f"corrected file."
        )

    def _scan_workspace(self) -> list[str]:
        """Scan the workspace for all source files.

        Returns:
            List of relative file paths.
        """
        files: list[str] = []
        for root, dirs, filenames in os.walk(self.workspace):
            dirs[:] = [d for d in dirs if d not in _SKIP_DIRS]
            for fname in filenames:
                if fname.startswith("."):
                    continue
                rel = os.path.relpath(os.path.join(root, fname), self.workspace)
                files.append(rel)
        return sorted(files)