orchestrata_ultimate.gateway.audit

Reference › orchestrata_ultimate › gateway › audit

AuditWriter records gateway events as audit nodes in the operational knowledge graph. It is built by server and called by rpc, and it takes its graph endpoint from config. The package narrative is in gateway.

connect builds an async graph repository through core.graph and does nothing but log a warning when the graph driver is not installed. log_event validates its arguments (TypeError for wrong types, ValueError for empty strings), generates a UUID entry id and an ISO-8601 UTC timestamp, and returns the id in every case. If there is no repository, or the circuit is open, it skips the write. Otherwise it runs _write_event, which executes one create-audit query carrying the event type, connection id, stringified data and the project slug from _project_slug, under a bounded timeout.

The writer protects the socket from a slow or absent graph. Each failed write increments a counter; when the counter reaches a threshold the circuit opens for a cooldown and later events are skipped without touching the graph. A successful write resets the counter. Failures are logged and never raised, so an audit outage cannot stall a connection. close closes the repository.

At a glance

Class diagram

Dependency diagram

API

Classes

AuditWriter

Line 72

Writes gateway audit events to the graph.

Functions

Source

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

"""Graph audit writer for WebSocket gateway events."""

import asyncio
import logging
import os
import time
import uuid
from datetime import UTC, datetime
from typing import Any

from orchestrata_ultimate.gateway.config import GatewayConfig
from orchestrata_ultimate.knowledge_graph.canonical_schema import (
    LABEL_PROJECT,
    LABEL_SESSION,
    REL_HAS_SESSION,
)

logger = logging.getLogger(__name__)

try:
    import neo4j  # noqa: F401  (presence check only — real import is local)

    GRAPH_DRIVER_AVAILABLE = True
except ImportError:
    GRAPH_DRIVER_AVAILABLE = False

# Audit writes are best-effort and must never stall the RPC/connection path.
# A single write is time-bounded, and after a failure the circuit opens so we
# stop hammering an unreachable Neo4j (each failed connect can cost seconds).
# The circuit half-opens after the cooldown so audit self-heals when Neo4j
# comes back.
_AUDIT_WRITE_TIMEOUT_S = 2.0
_AUDIT_FAILURE_THRESHOLD = 1
_AUDIT_CIRCUIT_COOLDOWN_S = 30.0

# The root Project slug the hub edge below MERGEs onto — same default every
# other KG writer in this codebase uses (e.g.
# factory/promotion/store.py's `_project_slug()`). Read at call time via
# `_project_slug()` (RUL-003), never cached at import time.
_DEFAULT_PROJECT_SLUG = "orchestrata-ultimate"


def _project_slug() -> str:
    """Resolve the root ``:Project`` slug the audit hub edge MERGEs onto."""
    return os.environ.get("ORCHESTRATA_PROJECT_ID", _DEFAULT_PROJECT_SLUG)


# Session-254 root-connectivity ruling: this writer previously MERGEd
# `:Session` with no edge back to the root `:Project` node, leaving every
# Session it minted (and the GatewayAuditEvent nodes hanging off it via
# DURING) permanently unreachable — 20 of the 25 orphaned Sessions
# scripts/check_kg_connectivity.py found, dragging 60 GatewayAuditEvent
# nodes with them. The Project hub MERGE below closes that gap; the
# GatewayAuditEvent's own label/DURING edge is left untouched (out of scope
# here — this fix only adds the missing hub edge to nodes already created).
_CREATE_AUDIT_QUERY = (
    "CREATE (e:GatewayAuditEvent {"
    "  entry_id: $entry_id,"
    "  event_type: $event_type,"
    "  connection_id: $connection_id,"
    "  timestamp: $timestamp,"
    "  data: $data"
    "}) "
    "WITH e "
    f"MERGE (s:{LABEL_SESSION} {{session_id: $connection_id}}) "
    "MERGE (e)-[:DURING]->(s) "
    f"MERGE (p:{LABEL_PROJECT} {{slug: $project}}) "
    f"MERGE (p)-[:{REL_HAS_SESSION}]->(s)"
)


class AuditWriter:
    """Writes gateway audit events to the graph."""

    def __init__(self, config: GatewayConfig) -> None:
        """Initialise from gateway configuration."""
        self._config = config
        self._repo: Any = None
        self._failure_count = 0
        self._circuit_open_until = 0.0

    async def connect(self) -> None:
        """Connect to the graph (no-op when driver absent)."""
        if not GRAPH_DRIVER_AVAILABLE:
            logger.warning("neo4j package not installed")
            return
        from orchestrata_ultimate.core.graph import (
            GraphRepositoryConfig,
            create_async_graph_repository,
        )

        self._repo = create_async_graph_repository(
            GraphRepositoryConfig(
                uri=self._config.graph_uri,
                auth=(
                    self._config.graph_user,
                    self._config.graph_password,
                ),
                driver_kwargs={
                    "connection_timeout": _AUDIT_WRITE_TIMEOUT_S,
                    "connection_acquisition_timeout": _AUDIT_WRITE_TIMEOUT_S,
                },
            )
        )
        logger.info("Audit writer connected to the graph")

    async def log_event(
        self,
        event_type: str,
        connection_id: str,
        data: dict[str, Any] | None = None,
    ) -> str:
        """Log an audit event and return its entry ID.

        Args:
            event_type: Category of the event.
            connection_id: Originating connection.
            data: Optional event payload.

        Returns:
            Generated UUID entry ID.

        Raises:
            TypeError: If arguments have wrong types.
            ValueError: If required strings are empty.
        """
        if not isinstance(event_type, str):
            raise TypeError("event_type must be a str")
        if not event_type:
            raise ValueError("event_type must not be empty")
        if not isinstance(connection_id, str):
            raise TypeError("connection_id must be a str")
        if not connection_id:
            raise ValueError("connection_id must not be empty")
        if data is not None and not isinstance(data, dict):
            raise TypeError("data must be a dict or None")
        entry_id = str(uuid.uuid4())
        timestamp = datetime.now(UTC).isoformat()
        if not self._repo:
            logger.debug(
                "No driver; skipping audit %s",
                entry_id,
            )
            return entry_id
        if time.monotonic() < self._circuit_open_until:
            logger.debug(
                "Audit circuit open; skipping %s",
                entry_id,
            )
            return entry_id
        try:
            await asyncio.wait_for(
                self._write_event(
                    entry_id,
                    event_type,
                    connection_id,
                    timestamp,
                    data,
                ),
                timeout=_AUDIT_WRITE_TIMEOUT_S,
            )
            self._failure_count = 0
        except Exception:
            self._failure_count += 1
            if self._failure_count >= _AUDIT_FAILURE_THRESHOLD:
                self._circuit_open_until = time.monotonic() + _AUDIT_CIRCUIT_COOLDOWN_S
                logger.warning(
                    "Audit write failed; circuit open for %.0fs "
                    "(graph unreachable/slow)",
                    _AUDIT_CIRCUIT_COOLDOWN_S,
                )
            else:
                logger.exception("Failed to write audit event")
        return entry_id

    async def _write_event(
        self,
        entry_id: str,
        event_type: str,
        connection_id: str,
        timestamp: str,
        data: dict[str, Any] | None,
    ) -> None:
        """Write a single audit event to the graph.

        Args:
            entry_id: Generated event UUID.
            event_type: Category of the event.
            connection_id: Originating connection.
            timestamp: ISO-8601 event time.
            data: Optional event payload.
        """
        await self._repo.execute_write(
            _CREATE_AUDIT_QUERY,
            {
                "entry_id": entry_id,
                "event_type": event_type,
                "connection_id": connection_id,
                "timestamp": timestamp,
                "data": str(data or {}),
                "project": _project_slug(),
            },
        )

    async def close(self) -> None:
        """Close the Neo4j driver."""
        if self._repo:
            await self._repo.close()
            self._repo = None
            logger.info("Audit writer disconnected")