orchestrata_ultimate.knowledge_graph.kg_design

Reference › orchestrata_ultimate › knowledge_graph › kg_design

KGDesignStore is the graph-backed store for design and traceability data in the spike pipeline: specs, user stories, requirements, design modules and classes, code nodes and sessions, all in the operational graph (orchkg on orchgraph). The analyst, compliance and other volks import it. It works on a synchronous graph repository, supplied or created by _connect, and _ensure_indexes creates lookup indexes. All operations are non-fatal: _run wraps each Cypher call so that when the graph is unavailable the methods return empty results and log.

Write methods are ensure_project, store_spec (a Spec linked to a Project), store_user_stories, store_requirements, store_design (DesignModule and DesignClass nodes), sync_code and sync_code_classes (code traces become CodeClass and CodeFunction nodes through the _build_* and _persist_* helpers), store_session and close_session. Read methods are get_requirements, query_existing_code for the refactoring case, get_session_history and get_rtm_trend, the coverage trend across sessions. check_compliance compares the design classes with the code classes and returns a ComplianceResult. The dataclasses SessionData and RTMSnapshot carry session lifecycle and requirements-traceability-matrix coverage over time. The dialect is taken from resolve_kg_dialect.

At a glance

Source src/orchestrata_ultimate/knowledge_graph/kg_design.py
Lines of code 924
Classes 4
Functions 1
Imports (internal) orchestrata_ultimate.core.base.traces, orchestrata_ultimate.core.graph, orchestrata_ultimate.core.graph_factory, orchestrata_ultimate.core.infra_config, orchestrata_ultimate.knowledge_graph.schema
Imported by orchestrata_ultimate.agents.volks.analyst.agent, orchestrata_ultimate.agents.volks.compliance._pipeline, orchestrata_ultimate.agents.volks.compliance.volk
Tests that import it tests/agents/volks/compliance/test_pipeline.py, tests/agents/volks/compliance/test_volk.py, tests/knowledge_graph/test_kg_design.py, tests/knowledge_graph/test_kg_design_coverage.py, tests/knowledge_graph/test_kg_design_extended.py, tests/knowledge_graph/test_knowledge_signatures.py

Class diagram

Class diagram orchestrata_ultimate.knowledge_graph.kg_design

Dependency diagram

Dependency diagram orchestrata_ultimate.knowledge_graph.kg_design

API

Classes

SessionData (dataclass)

Decorators: dataclass · Line 33

Session lifecycle data.

Attribute Type Default
id str
purpose str ''
started_at str ''
ended_at str ''
status str 'ACTIVE'
achievements list[str] field(default_factory=list)
work_log list[str] field(default_factory=list)
Method Signature Summary
__post_init__ __post_init__() -> None

RTMSnapshot (dataclass)

Decorators: dataclass · Line 50

RTM snapshot for tracking coverage over time.

Attribute Type Default
id str
timestamp str ''
coverage_pct float 0.0
verified_count int 0
total_count int 0
missing_count int 0
orphan_count int 0
Method Signature Summary
__post_init__ __post_init__() -> None

ComplianceResult (dataclass)

Decorators: dataclass · Line 94

Result of design vs code compliance check.

Attribute Type Default
missing_classes list[str] field(default_factory=list)
extra_classes list[str] field(default_factory=list)
missing_methods dict[str, list[str]] field(default_factory=dict)
extra_methods dict[str, list[str]] field(default_factory=dict)
compliance_pct float 0.0
Method Signature Summary
to_dict to_dict() -> dict[str, Any]

KGDesignStore

Line 113

Graph-backed store for design and traceability data.

Method Signature Summary
__init__ __init__(repo: SyncGraphRepository | None = None) -> None Initialize the design store.
_connect _connect() -> None
_ensure_indexes _ensure_indexes() -> None Create indexes for efficient querying.
_run _run(query: str, *, write: bool = True, ****params: Any) -> list[dict] Execute a Cypher query and return results.
ensure_project ensure_project(slug: str, name: str, project_type: str = '') -> None Create or update a Project node.
store_spec store_spec(project_slug: str, spec_id: str, title: str, content: str = '') -> str Store a Spec node linked to a Project.
store_user_stories store_user_stories(spec_id: str, stories: list[dict[str, Any]]) -> dict[str, str] Store UserStory nodes linked to a Spec.
store_requirements store_requirements(story_id: str, requirements: list[dict[str, Any]]) -> None Store Requirement nodes linked to a UserStory.
get_requirements get_requirements(project_slug: str) -> list[dict[str, str]] Get all requirements for a project.
store_design store_design(project_slug: str, modules: list[dict[str, Any]]) -> None Store DesignModule and DesignClass nodes.
sync_code sync_code(project_slug: str, trace_records: list[dict[str, Any]]) -> None Sync code traces to CodeClass/CodeFunction nodes.
_build_class_rows staticmethod _build_class_rows(trace_records: list[dict[str, Any]], include: Any) -> tuple[list[dict[str, Any]], list[dict[str, Any]]] Build CodeClass node rows and their TRACES_TO requirement rows.
_build_function_rows staticmethod _build_function_rows(trace_records: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]] Build CodeFunction node rows (methods) and their requirement rows.
_persist_code_classes _persist_code_classes(project_slug: str, ts: str, class_rows: list[dict[str, Any]], class_req_rows: list[dict[str, Any]]) -> None Write CodeClass nodes and their TRACES_TO requirement edges.
_persist_code_functions _persist_code_functions(project_slug: str, ts: str, func_rows: list[dict[str, Any]], func_req_rows: list[dict[str, Any]]) -> None Write CodeFunction nodes and their TRACES_TO requirement edges.
sync_code_classes sync_code_classes(project_slug: str, classes: list[dict[str, Any]]) -> None Sync code classes with their methods for compliance.
check_compliance check_compliance(project_slug: str) -> ComplianceResult Compare DesignClass nodes against CodeClass nodes.
query_existing_code query_existing_code(project_slug: str) -> dict[str, Any] Query KG for existing code in a project (refactoring case).
store_session store_session(project_slug: str, session: SessionData) -> None Store a Session node linked to a Project.
close_session close_session(session_id: str, achievements: list[str], rtm_snapshot: RTMSnapshot | None = None) -> None Close a session and optionally link an RTM snapshot.
get_session_history get_session_history(project_slug: str) -> list[dict[str, Any]] Get all sessions for a project.
get_rtm_trend get_rtm_trend(project_slug: str) -> list[dict[str, Any]] Get RTM coverage trend across sessions.
close close() -> None Close the repository.

Functions

Function Signature Summary
_require_project_slug _require_project_slug(project_slug: str, caller: str) -> None Guard every project-scoped WRITE against an empty/whitespace slug.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/knowledge_graph/kg_design.py

# mypy: ignore-errors
"""Knowledge Graph design store for spike pipeline.

Reads/writes specs, user stories, requirements, design nodes, code nodes,
and session data to the graph. Provides compliance queries comparing design
vs implementation.

Non-fatal: all Neo4j operations are wrapped in try/except. If Neo4j is
unavailable, operations return empty results and log warnings/errors.
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.core.graph import BoltSyncRepository, GraphRepositoryConfig
from orchestrata_ultimate.core.graph_factory import resolve_kg_dialect
from orchestrata_ultimate.core.infra_config import GRAPH_PASSWORD, GRAPH_URI, GRAPH_USER
from orchestrata_ultimate.knowledge_graph.schema import to_memgraph_ddl

if TYPE_CHECKING:
    from orchestrata_ultimate.core.graph import SyncGraphRepository

_logger = logging.getLogger(__name__)


@dataclass
class SessionData:
    """Session lifecycle data."""

    id: str
    purpose: str = ""
    started_at: str = ""
    ended_at: str = ""
    status: str = "ACTIVE"
    achievements: list[str] = field(default_factory=list)
    work_log: list[str] = field(default_factory=list)

    def __post_init__(self) -> None:
        if not self.started_at:
            self.started_at = datetime.now(UTC).isoformat()


@dataclass
class RTMSnapshot:
    """RTM snapshot for tracking coverage over time."""

    id: str
    timestamp: str = ""
    coverage_pct: float = 0.0
    verified_count: int = 0
    total_count: int = 0
    missing_count: int = 0
    orphan_count: int = 0

    def __post_init__(self) -> None:
        if not self.timestamp:
            self.timestamp = datetime.now(UTC).isoformat()


def _require_project_slug(project_slug: str, caller: str) -> None:
    """Guard every project-scoped WRITE against an empty/whitespace slug.

    2026-07 audit: callers that slugified free-text (a task summary /
    intent) instead of resolving the canonical project id produced an
    empty-string slug and a truncated-to-50-chars slug, each MERGEing its
    own throwaway ``:Project`` hub that thousands of CodeClass/CodeFunction
    ``DEFINES`` edges then anchored onto. Fail fast instead: never write an
    anonymous Project hub. Callers must resolve a real slug via
    ``knowledge_graph.project_resolver`` (falling back to the ``"unscoped"``
    sentinel — never ``""`` — when no project directory is known).

    Args:
        project_slug: The slug the caller is about to write under.
        caller: Method name, for the error message.

    Raises:
        ValueError: If ``project_slug`` is empty or whitespace-only.
    """
    if not project_slug or not project_slug.strip():
        raise ValueError(
            f"KGDesignStore.{caller}: project_slug must be a non-empty, "
            "non-whitespace string -- refusing to write an anonymous "
            "Project hub node."
        )


@dataclass
class ComplianceResult:
    """Result of design vs code compliance check."""

    missing_classes: list[str] = field(default_factory=list)
    extra_classes: list[str] = field(default_factory=list)
    missing_methods: dict[str, list[str]] = field(default_factory=dict)
    extra_methods: dict[str, list[str]] = field(default_factory=dict)
    compliance_pct: float = 0.0

    def to_dict(self) -> dict[str, Any]:
        return {
            "missing_classes": self.missing_classes,
            "extra_classes": self.extra_classes,
            "missing_methods": self.missing_methods,
            "extra_methods": self.extra_methods,
            "compliance_pct": self.compliance_pct,
        }


class KGDesignStore:
    """Graph-backed store for design and traceability data.

    All operations are non-fatal. If the graph is unavailable, methods
    return empty results and log warnings/errors.
    """

    def __init__(self, repo: SyncGraphRepository | None = None) -> None:
        """Initialize the design store.

        Ownership (backlog #163 phase 4): when ``repo`` is None, the store
        builds and owns its own repository lazily via ``_connect()`` (on
        first ``_run()`` call) — ``close()`` on it really closes the
        underlying driver. When a ``repo`` is injected, ownership belongs to
        whoever built *that* repo (shared repos no-op on ``close()``); this
        class no longer tracks its own ``_owns_driver`` flag.
        """
        self._repo: SyncGraphRepository | None = repo
        self._available = False
        self._indexes_ensured = False
        if repo is not None:
            try:
                repo.verify_connectivity()
                self._available = True
                self._ensure_indexes()
            except Exception as e:
                self._available = False
                _logger.error(
                    "shared driver unavailable (%s) -- design store disabled", e
                )

    def _connect(self) -> None:
        if self._repo is not None:
            return
        try:
            repo = BoltSyncRepository(
                GraphRepositoryConfig(uri=GRAPH_URI, auth=(GRAPH_USER, GRAPH_PASSWORD))
            )
            self._repo = repo
            repo.verify_connectivity()
            self._available = True
            self._ensure_indexes()
            _logger.info("graph connected: %s", GRAPH_URI)
        except Exception as e:
            self._available = False
            _logger.error("graph unavailable (%s) -- design store disabled", e)

    @traces(spec="008", story="US-001", reqs=["REQ-008-037"])
    def _ensure_indexes(self) -> None:
        """Create indexes for efficient querying."""
        repo = self._repo
        if not self._available or self._indexes_ensured or repo is None:
            return
        # Project.slug and Session.session_id are canonical UNIQUE keys
        # (canonical_schema CANONICAL_MERGE_KEYS): their constraints provide
        # backing indexes for free, and a plain index created here FIRST
        # permanently blocks constraint creation (Neo.ClientError.Schema.
        # IndexAlreadyExists — hit in the 2026-07-12 suite runs). Never
        # index a canonical key here.
        indexes = [
            # Wave 6: Spec/UserStory/Requirement keys are canonical `id`
            # (backed by uniqueness constraints), so plain indexes on the old
            # `spec_id`/`story_id`/`req_id` properties are removed.
            # DesignModule/DesignClass now key on project-qualified `id`, so
            # their name indexes are also retired.
            "CREATE INDEX IF NOT EXISTS FOR (c:CodeClass) ON (c.name)",
            "CREATE INDEX IF NOT EXISTS FOR (c:CodeFunction) ON (c.name)",
            "CREATE INDEX IF NOT EXISTS FOR (r:RTMSnapshot) ON (r.snapshot_id)",
        ]
        # gate-ok: a missing index degrades query SPEED, not correctness; the failing statement is logged and re-applying is idempotent on the next connect
        try:
            for idx in indexes:
                # Rendered for the live engine, then run via execute_autocommit
                # rather than execute_write: the statements above are Neo4j
                # dialect, and orchgraph/Memgraph both spells them differently
                # ("CREATE INDEX ON :L(p)") and refuses index DDL inside a
                # managed/explicit transaction ("Index manipulation is not
                # allowed in multicommand transactions" — measured live
                # 2026-07-26). Neo4j accepts DDL either way, so this is the
                # portable form for both engines.
                stmt = idx
                if resolve_kg_dialect() == "memgraph":
                    translated = to_memgraph_ddl(idx)
                    if translated is None:
                        continue
                    stmt = translated
                try:
                    repo.execute_autocommit(stmt)
                except Exception as exc:
                    # The Memgraph dialect has no IF NOT EXISTS grammar, so
                    # re-application surfaces as "already exists" — exactly the
                    # no-op the Neo4j dialect gets for free.
                    if "already exists" not in str(exc).lower():
                        raise
            self._indexes_ensured = True
        except Exception as e:
            _logger.warning("index creation error: %s", e)

    def _run(self, query: str, *, write: bool = True, **params: Any) -> list[dict]:
        """Execute a Cypher query and return results.

        Args:
            query: The Cypher query string.
            write: Whether to route through a managed write transaction
                (default) or a managed read transaction. Most callers write
                (MERGE/SET); the handful of pure lookups pass ``write=False``.
            **params: Cypher query parameters.
        """
        if not self._available and self._repo is None:
            self._connect()
        if not self._available:
            return []
        repo = self._repo
        if repo is None:
            return []
        try:
            if write:
                return repo.execute_write(query, params)
            return repo.execute_read(query, params)
        except Exception as e:
            _logger.error("query error: %s", e)
            return []

    # -- Project Operations ----------------------------------------

    def ensure_project(self, slug: str, name: str, project_type: str = "") -> None:
        """Create or update a Project node."""
        _require_project_slug(slug, "ensure_project")
        self._run(
            "MERGE (p:Project {slug: $slug}) "
            "SET p.name = $name, "
            "    p.project_type = $type, "
            "    p.updated_at = $ts",
            slug=slug,
            name=name,
            type=project_type,
            ts=datetime.now(UTC).isoformat(),
        )

    # -- Spec Operations ---------------------------------------------

    def store_spec(
        self,
        project_slug: str,
        spec_id: str,
        title: str,
        content: str = "",
    ) -> str:
        """Store a Spec node linked to a Project.

        ``id`` is project-qualified (``"{project_slug}:{spec_id}"``): a bare
        ``spec_id`` (e.g. ``"001"``, this class's own default) recurs across
        analyst runs/projects, so it is unsound as a global MERGE key in a
        multi-project graph (2026-07-12 multi-project ruling). The bare
        ``spec_id`` is kept as a plain property for display/lookup. Label is
        ``Spec`` (canonical_schema LABEL_SPEC), not the prior ``Specification``
        — this class was the only reader/writer of that label, so it converges
        onto the canon rather than keeping its own synonym.

        Returns:
            The project-qualified Spec id, to pass forward into
            ``store_user_stories``.

        Session-254 root-connectivity ruling: the ``Project`` hub used to be
        ``MATCH``ed rather than ``MERGE``d here, so when no ``:Project``
        existed yet for ``project_slug`` the match failed OPEN -- the ``MERGE
        (p)-[:HAS_SPEC]->(s)`` line never ran, and the Spec (plus every
        UserStory/Requirement ``store_user_stories``/``store_requirements``
        hang off it) was minted as an unreachable island (183 live nodes
        found by ``scripts/check_kg_connectivity.py``). ``MERGE`` is safe here
        specifically because ``_require_project_slug`` above already fails
        closed on an empty/whitespace slug -- every other hub writer in the
        repo (``cold_store.py``, ``session_manager.py``, ``gateway/audit.py``)
        follows the same MERGE-the-hub pattern.
        """
        _require_project_slug(project_slug, "store_spec")
        qid = f"{project_slug}:{spec_id}"
        self._run(
            "MERGE (s:Spec {id: $id}) "
            "SET s.spec_id = $spec_id, "
            "    s.project_slug = $slug, "
            "    s.title = $title, "
            "    s.content = $content, "
            "    s.updated_at = $ts "
            "WITH s "
            "MERGE (p:Project {slug: $slug}) "
            "MERGE (p)-[:HAS_SPEC]->(s)",
            id=qid,
            spec_id=spec_id,
            title=title,
            content=content,
            slug=project_slug,
            ts=datetime.now(UTC).isoformat(),
        )
        return qid

    # -- User Story Operations -------------------------------------

    def store_user_stories(
        self,
        spec_id: str,
        stories: list[dict[str, Any]],
    ) -> dict[str, str]:
        """Store UserStory nodes linked to a Spec.

        Args:
            spec_id: The project-qualified Spec id returned by
                ``store_spec`` (NOT the bare spec number).
            stories: List of dicts with keys:
                id, title, as_a, i_want, so_that,
                acceptance_criteria (list of strings).

        Returns:
            Mapping of each story's bare ``id`` to its project-qualified
            UserStory id (``"{spec_id}:{story id}"``), to pass forward into
            ``store_requirements``. Empty when ``stories`` is empty.
        """
        if not stories:
            return {}
        rows = [
            {
                "id": f"{spec_id}:{s['id']}",
                "story_id": s["id"],
                "title": s.get("title", ""),
                "as_a": s.get("as_a", ""),
                "i_want": s.get("i_want", ""),
                "so_that": s.get("so_that", ""),
                "criteria": json.dumps(s.get("acceptance_criteria", [])),
            }
            for s in stories
        ]
        ts = datetime.now(UTC).isoformat()
        self._run(
            "MATCH (sp:Spec {id: $spec_id}) "
            "UNWIND $rows AS row "
            "MERGE (u:UserStory {id: row.id}) "
            "SET u.story_id = row.story_id, "
            "    u.title = row.title, "
            "    u.as_a = row.as_a, "
            "    u.i_want = row.i_want, "
            "    u.so_that = row.so_that, "
            "    u.acceptance_criteria = row.criteria, "
            "    u.updated_at = $ts "
            "MERGE (sp)-[:HAS_STORY]->(u)",
            spec_id=spec_id,
            rows=rows,
            ts=ts,
        )
        return {row["story_id"]: row["id"] for row in rows}

    # -- Requirement Operations ------------------------------------

    def store_requirements(
        self,
        story_id: str,
        requirements: list[dict[str, Any]],
    ) -> None:
        """Store Requirement nodes linked to a UserStory.

        Args:
            story_id: The project-qualified UserStory id returned by
                ``store_user_stories`` (NOT the bare story id).
            requirements: List of dicts with keys:
                id, description, priority, category.

        Note:
            ``id`` is project-qualified (``"{story_id}:{req id}"``). The bare
            ``req_id`` is kept as a plain property — ``_persist_code_classes``/
            ``_persist_code_functions`` below still MATCH Requirement by that
            bare ``req_id`` (a separate, @traces-decorator-sourced identity
            space not reconciled by this wave; W3 reconciliation ledger).
        """
        if not requirements:
            return
        rows = [
            {
                "id": f"{story_id}:{r['id']}",
                "req_id": r["id"],
                "desc": r.get("description", ""),
                "priority": r.get("priority", "NORMAL"),
                "category": r.get("category", ""),
            }
            for r in requirements
        ]
        ts = datetime.now(UTC).isoformat()
        self._run(
            "MATCH (u:UserStory {id: $story_id}) "
            "UNWIND $rows AS row "
            "MERGE (r:Requirement {id: row.id}) "
            "SET r.req_id = row.req_id, "
            "    r.description = row.desc, "
            "    r.priority = row.priority, "
            "    r.category = row.category, "
            "    r.updated_at = $ts "
            "MERGE (u)-[:HAS_REQUIREMENT]->(r)",
            story_id=story_id,
            rows=rows,
            ts=ts,
        )

    def get_requirements(self, project_slug: str) -> list[dict[str, str]]:
        """Get all requirements for a project.

        Returns:
            List of dicts with req_id, description, story_id, spec_id (all
            the bare/display ids, not the internal qualified ``id``).
        """
        results = self._run(
            "MATCH (p:Project {slug: $slug})-[:HAS_SPEC]->"
            "(s:Spec)-[:HAS_STORY]->"
            "(u:UserStory)-[:HAS_REQUIREMENT]->"
            "(r:Requirement) "
            "RETURN r.req_id AS req_id, "
            "       r.description AS description, "
            "       u.story_id AS story_id, "
            "       s.spec_id AS spec_id "
            "ORDER BY r.req_id",
            write=False,
            slug=project_slug,
        )
        return results

    # -- Design Operations -----------------------------------------

    def store_design(
        self,
        project_slug: str,
        modules: list[dict[str, Any]],
    ) -> None:
        """Store DesignModule and DesignClass nodes.

        Args:
            project_slug: Project slug.
            modules: List of module dicts with keys:
                name, description, dependencies,
                classes: [{name, methods, attributes, responsibilities}]
        """
        if not modules:
            return
        _require_project_slug(project_slug, "store_design")
        mod_rows = [
            {
                "id": f"{project_slug}:{m['name']}",
                "name": m["name"],
                "desc": m.get("description", ""),
                "deps": json.dumps(m.get("dependencies", [])),
            }
            for m in modules
        ]
        cls_rows = [
            {
                "id": f"{project_slug}:{module['name']}:{cls['name']}",
                "name": cls["name"],
                "methods": json.dumps(cls.get("methods", [])),
                "attrs": json.dumps(cls.get("attributes", [])),
                "resp": cls.get("responsibilities", ""),
                "mod_id": f"{project_slug}:{module['name']}",
            }
            for module in modules
            for cls in module.get("classes", [])
        ]
        ts = datetime.now(UTC).isoformat()
        if mod_rows:
            self._run(
                "MATCH (p:Project {slug: $slug}) "
                "UNWIND $rows AS row "
                "MERGE (m:DesignModule {id: row.id}) "
                "SET m.name = row.name, "
                "    m.description = row.desc, "
                "    m.dependencies = row.deps, "
                "    m.updated_at = $ts "
                "MERGE (p)-[:HAS_MODULE]->(m)",
                slug=project_slug,
                rows=mod_rows,
                ts=ts,
            )
        if cls_rows:
            self._run(
                "UNWIND $rows AS row "
                "MERGE (c:DesignClass {id: row.id}) "
                "SET c.name = row.name, "
                "    c.methods = row.methods, "
                "    c.attributes = row.attrs, "
                "    c.responsibilities = row.resp, "
                "    c.updated_at = $ts "
                "WITH c, row "
                "MATCH (m:DesignModule {id: row.mod_id}) "
                "MERGE (m)-[:CONTAINS_CLASS]->(c)",
                rows=cls_rows,
                ts=ts,
            )

    # -- Code Sync Operations --------------------------------------

    def sync_code(
        self,
        project_slug: str,
        trace_records: list[dict[str, Any]],
    ) -> None:
        """Sync code traces to CodeClass/CodeFunction nodes.

        Creates nodes for each traced class/function and links them
        to requirements via TRACES_TO relationships.

        Args:
            project_slug: Project slug.
            trace_records: List of TraceInfo.to_dict() dicts.
        """
        if not trace_records:
            return
        _require_project_slug(project_slug, "sync_code")
        ts = datetime.now(UTC).isoformat()

        # Batch 1: CodeClass nodes (non-method)
        class_rows, class_req_rows = self._build_class_rows(
            trace_records,
            include=lambda t: bool(t.get("class_name") and not t.get("method")),
        )
        self._persist_code_classes(project_slug, ts, class_rows, class_req_rows)

        # Batch 2: CodeFunction nodes (methods)
        func_rows, func_req_rows = self._build_function_rows(trace_records)
        self._persist_code_functions(project_slug, ts, func_rows, func_req_rows)

        # Batch 3: CodeClass nodes for class-bearing methods (class_name + method)
        class_only_rows, class_only_req_rows = self._build_class_rows(
            trace_records,
            include=lambda t: bool(t.get("class_name") and t.get("method")),
        )
        self._persist_code_classes(
            project_slug, ts, class_only_rows, class_only_req_rows
        )

    @staticmethod
    def _build_class_rows(
        trace_records: list[dict[str, Any]],
        include: Any,
    ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
        """Build CodeClass node rows and their TRACES_TO requirement rows.

        Args:
            trace_records: Trace dicts to scan.
            include: Predicate deciding whether a trace contributes a class row.
        """
        class_rows: list[dict[str, Any]] = []
        req_rows: list[dict[str, Any]] = []
        for trace in trace_records:
            if not include(trace):
                continue
            class_rows.append(
                {
                    "name": trace["class_name"],
                    "fp": trace.get("file_path", ""),
                    "ln": trace.get("line_number", 0),
                    "mod": trace.get("module", ""),
                }
            )
            req_rows.extend(
                {"name": trace["class_name"], "req_id": req_id}
                for req_id in trace.get("reqs", [])
            )
        return class_rows, req_rows

    @staticmethod
    def _build_function_rows(
        trace_records: list[dict[str, Any]],
    ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
        """Build CodeFunction node rows (methods) and their requirement rows."""
        func_rows: list[dict[str, Any]] = []
        req_rows: list[dict[str, Any]] = []
        for trace in trace_records:
            if not trace.get("method"):
                continue
            class_name = trace.get("class_name", "")
            func_name = trace["method"]
            full_name = f"{class_name}.{func_name}" if class_name else func_name
            func_rows.append(
                {
                    "name": full_name,
                    "fp": trace.get("file_path", ""),
                    "ln": trace.get("line_number", 0),
                    "cls": class_name,
                    "mod": trace.get("module", ""),
                }
            )
            req_rows.extend(
                {"name": full_name, "req_id": req_id}
                for req_id in trace.get("reqs", [])
            )
        return func_rows, req_rows

    def _persist_code_classes(
        self,
        project_slug: str,
        ts: str,
        class_rows: list[dict[str, Any]],
        class_req_rows: list[dict[str, Any]],
    ) -> None:
        """Write CodeClass nodes and their TRACES_TO requirement edges."""
        if class_rows:
            self._run(
                "MATCH (p:Project {slug: $slug}) "
                "UNWIND $rows AS row "
                "MERGE (c:CodeClass {name: row.name, project: $slug}) "
                "SET c.file_path = row.fp, "
                "    c.line_number = row.ln, "
                "    c.module = row.mod, "
                "    c.updated_at = $ts "
                "MERGE (p)-[:DEFINES]->(c)",
                slug=project_slug,
                rows=class_rows,
                ts=ts,
            )
        if class_req_rows:
            self._run(
                "UNWIND $rows AS row "
                "MATCH (c:CodeClass {name: row.name, project: $slug}) "
                "MATCH (r:Requirement {req_id: row.req_id}) "
                "MERGE (c)-[:TRACES_TO]->(r)",
                slug=project_slug,
                rows=class_req_rows,
            )

    def _persist_code_functions(
        self,
        project_slug: str,
        ts: str,
        func_rows: list[dict[str, Any]],
        func_req_rows: list[dict[str, Any]],
    ) -> None:
        """Write CodeFunction nodes and their TRACES_TO requirement edges."""
        if func_rows:
            self._run(
                "MATCH (p:Project {slug: $slug}) "
                "UNWIND $rows AS row "
                "MERGE (f:CodeFunction {name: row.name, project: $slug}) "
                "SET f.file_path = row.fp, "
                "    f.line_number = row.ln, "
                "    f.class_name = row.cls, "
                "    f.module = row.mod, "
                "    f.updated_at = $ts "
                "MERGE (p)-[:DEFINES]->(f)",
                slug=project_slug,
                rows=func_rows,
                ts=ts,
            )
        if func_req_rows:
            self._run(
                "UNWIND $rows AS row "
                "MATCH (f:CodeFunction {name: row.name, project: $slug}) "
                "MATCH (r:Requirement {req_id: row.req_id}) "
                "MERGE (f)-[:TRACES_TO]->(r)",
                slug=project_slug,
                rows=func_req_rows,
            )

    def sync_code_classes(
        self,
        project_slug: str,
        classes: list[dict[str, Any]],
    ) -> None:
        """Sync code classes with their methods for compliance.

        Args:
            project_slug: Project slug.
            classes: List of dicts with name, methods, file_path.
        """
        _require_project_slug(project_slug, "sync_code_classes")
        for cls in classes:
            methods = json.dumps(cls.get("methods", []))
            self._run(
                "MERGE (c:CodeClass {name: $name, "
                "       project: $slug}) "
                "SET c.methods = $methods, "
                "    c.file_path = $fp, "
                "    c.updated_at = $ts",
                name=cls["name"],
                slug=project_slug,
                methods=methods,
                fp=cls.get("file_path", ""),
                ts=datetime.now(UTC).isoformat(),
            )

    # -- Compliance Checking ---------------------------------------

    def check_compliance(self, project_slug: str) -> ComplianceResult:
        """Compare DesignClass nodes against CodeClass nodes.

        Returns:
            ComplianceResult with missing/extra classes and methods.
        """
        design_classes = self._run(
            "MATCH (p:Project {slug: $slug})-[:HAS_MODULE]->"
            "(:DesignModule)-[:CONTAINS_CLASS]->"
            "(d:DesignClass) "
            "RETURN d.name AS name, d.methods AS methods",
            write=False,
            slug=project_slug,
        )

        code_classes = self._run(
            "MATCH (c:CodeClass {project: $slug}) "
            "RETURN c.name AS name, c.methods AS methods",
            write=False,
            slug=project_slug,
        )

        design_names = {d["name"] for d in design_classes}
        code_names = {c["name"] for c in code_classes}

        missing_classes = sorted(design_names - code_names)
        extra_classes = sorted(code_names - design_names)

        design_methods: dict[str, list[str]] = {}
        for d in design_classes:
            methods_raw = d.get("methods", "[]")
            if isinstance(methods_raw, str):
                methods_raw = json.loads(methods_raw)
            design_methods[d["name"]] = methods_raw

        code_methods: dict[str, list[str]] = {}
        for c in code_classes:
            methods_raw = c.get("methods", "[]")
            if isinstance(methods_raw, str):
                methods_raw = json.loads(methods_raw)
            code_methods[c["name"]] = methods_raw

        missing_methods: dict[str, list[str]] = {}
        extra_methods_dict: dict[str, list[str]] = {}

        for cls_name in design_names & code_names:
            d_set = set(design_methods.get(cls_name, []))
            c_set = set(code_methods.get(cls_name, []))
            missing = sorted(d_set - c_set)
            extra = sorted(c_set - d_set)
            if missing:
                missing_methods[cls_name] = missing
            if extra:
                extra_methods_dict[cls_name] = extra

        total_designed = len(design_names)
        implemented = len(design_names & code_names)
        compliance_pct = (
            round(implemented / total_designed * 100, 1)
            if total_designed > 0
            else 100.0
        )

        return ComplianceResult(
            missing_classes=missing_classes,
            extra_classes=extra_classes,
            missing_methods=missing_methods,
            extra_methods=extra_methods_dict,
            compliance_pct=compliance_pct,
        )

    # -- Query Operations ------------------------------------------

    def query_existing_code(self, project_slug: str) -> dict[str, Any]:
        """Query KG for existing code in a project (refactoring case).

        Returns:
            Dict with classes, functions, modules lists.
        """
        classes = self._run(
            "MATCH (c:CodeClass {project: $slug}) "
            "RETURN c.name AS name, c.file_path AS file_path, "
            "       c.methods AS methods",
            write=False,
            slug=project_slug,
        )

        functions = self._run(
            "MATCH (f:CodeFunction {project: $slug}) "
            "RETURN f.name AS name, f.file_path AS file_path, "
            "       f.class_name AS class_name",
            write=False,
            slug=project_slug,
        )

        modules = self._run(
            "MATCH (p:Project {slug: $slug})-[:HAS_MODULE]->"
            "(m:DesignModule) "
            "RETURN m.name AS name, m.description AS description",
            write=False,
            slug=project_slug,
        )

        return {
            "classes": classes,
            "functions": functions,
            "modules": modules,
        }

    # -- Session Operations ----------------------------------------

    def store_session(self, project_slug: str, session: SessionData) -> None:
        """Store a Session node linked to a Project.

        Session-254 root-connectivity ruling: found while fixing the
        identical ``store_spec`` defect above -- this writer also ``MATCH``ed
        (rather than ``MERGE``d) the Project hub, so it fails open exactly
        the same way whenever no ``:Project`` exists yet for ``project_slug``,
        despite ``plans/hub-link-plan.md`` §4.6 having (incorrectly) listed
        this method as one of the five already-anchored Session writers.
        ``_require_project_slug`` above makes ``MERGE`` safe here too.
        """
        _require_project_slug(project_slug, "store_session")
        self._run(
            "MERGE (s:Session {session_id: $sid}) "
            "SET s.purpose = $purpose, "
            "    s.started_at = $started, "
            "    s.status = $status, "
            "    s.updated_at = $ts "
            "WITH s "
            "MERGE (p:Project {slug: $slug}) "
            "MERGE (p)-[:HAS_SESSION]->(s)",
            sid=session.id,
            purpose=session.purpose,
            started=session.started_at,
            status=session.status,
            slug=project_slug,
            ts=datetime.now(UTC).isoformat(),
        )

    def close_session(
        self,
        session_id: str,
        achievements: list[str],
        rtm_snapshot: RTMSnapshot | None = None,
    ) -> None:
        """Close a session and optionally link an RTM snapshot."""
        self._run(
            "MATCH (s:Session {session_id: $sid}) "
            "SET s.ended_at = $ts, "
            "    s.status = 'COMPLETED', "
            "    s.achievements = $achievements",
            sid=session_id,
            ts=datetime.now(UTC).isoformat(),
            achievements=json.dumps(achievements),
        )

        if rtm_snapshot:
            self._run(
                "MERGE (r:RTMSnapshot "
                "       {snapshot_id: $snap_id}) "
                "SET r.timestamp = $ts, "
                "    r.coverage_pct = $cov, "
                "    r.verified_count = $verified, "
                "    r.total_count = $total, "
                "    r.missing_count = $missing, "
                "    r.orphan_count = $orphan "
                "WITH r "
                "MATCH (s:Session {session_id: $sid}) "
                "MERGE (s)-[:PRODUCED_RTM]->(r)",
                snap_id=rtm_snapshot.id,
                ts=rtm_snapshot.timestamp,
                cov=rtm_snapshot.coverage_pct,
                verified=rtm_snapshot.verified_count,
                total=rtm_snapshot.total_count,
                missing=rtm_snapshot.missing_count,
                orphan=rtm_snapshot.orphan_count,
                sid=session_id,
            )

    def get_session_history(self, project_slug: str) -> list[dict[str, Any]]:
        """Get all sessions for a project."""
        return self._run(
            "MATCH (p:Project {slug: $slug})-[:HAS_SESSION]->"
            "(s:Session) "
            "RETURN s.session_id AS session_id, "
            "       s.purpose AS purpose, "
            "       s.started_at AS started_at, "
            "       s.ended_at AS ended_at, "
            "       s.status AS status, "
            "       s.achievements AS achievements "
            "ORDER BY s.started_at",
            write=False,
            slug=project_slug,
        )

    def get_rtm_trend(self, project_slug: str) -> list[dict[str, Any]]:
        """Get RTM coverage trend across sessions.

        Returns:
            List of dicts with timestamp, coverage_pct, verified,
            total for each session.
        """
        return self._run(
            "MATCH (p:Project {slug: $slug})-[:HAS_SESSION]->"
            "(s:Session)-[:PRODUCED_RTM]->(r:RTMSnapshot) "
            "RETURN r.timestamp AS timestamp, "
            "       r.coverage_pct AS coverage_pct, "
            "       r.verified_count AS verified, "
            "       r.total_count AS total "
            "ORDER BY r.timestamp",
            write=False,
            slug=project_slug,
        )

    # -- Cleanup ---------------------------------------------------

    def close(self) -> None:
        """Close the repository.

        A no-op when ``self._repo`` wraps a driver it does not own (shared
        semantics live on the repository itself — see ``__init__``).
        """
        if self._repo is not None:
            self._repo.close()
        self._repo = None
        self._available = False