aios/active-runtime/app/aios_v2/persistence.py

"""Durable single-process persistence for the AIOS V2 control plane.

The candidate uses an atomic SQLite snapshot so a gateway restart does not lose
tenants, agents, task state, idempotency records, or audit events.  The domain
service remains storage-agnostic; a later HA backend can replace this module
without changing API or runtime-adapter contracts.
"""
from __future__ import annotations

import json
import sqlite3
import threading
from dataclasses import asdict
from pathlib import Path
from typing import Any, Callable

from .control_plane import (
    AgentRecord,
    AuditEvent,
    ControlPlaneService,
    InMemoryControlPlaneStore,
    TaskRecord,
    TaskState,
    TenantRecord,
)


class SQLiteSnapshotRepository:
    """Store one versioned control-plane snapshot in SQLite."""

    SCHEMA_VERSION = 1

    def __init__(self, path: str | Path) -> None:
        self.path = Path(path).resolve()
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._lock = threading.RLock()
        self._initialize()

    def _connect(self) -> sqlite3.Connection:
        connection = sqlite3.connect(self.path, timeout=30.0)
        connection.execute("PRAGMA journal_mode=WAL")
        connection.execute("PRAGMA synchronous=FULL")
        connection.execute("PRAGMA foreign_keys=ON")
        return connection

    def _initialize(self) -> None:
        with self._connect() as connection:
            connection.execute(
                """
                CREATE TABLE IF NOT EXISTS control_plane_snapshot (
                    singleton_id INTEGER PRIMARY KEY CHECK (singleton_id = 1),
                    schema_version INTEGER NOT NULL,
                    payload_json TEXT NOT NULL,
                    updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
                )
                """
            )

    def load(self) -> dict[str, Any] | None:
        with self._lock, self._connect() as connection:
            row = connection.execute(
                "SELECT schema_version, payload_json FROM control_plane_snapshot WHERE singleton_id = 1"
            ).fetchone()
        if row is None:
            return None
        schema_version, payload_json = row
        if int(schema_version) != self.SCHEMA_VERSION:
            raise RuntimeError(
                f"unsupported control-plane snapshot version: {schema_version}"
            )
        decoded = json.loads(str(payload_json))
        if not isinstance(decoded, dict):
            raise RuntimeError("invalid control-plane snapshot")
        return decoded

    def save(self, payload: dict[str, Any]) -> None:
        encoded = json.dumps(
            payload,
            ensure_ascii=False,
            sort_keys=True,
            separators=(",", ":"),
            default=str,
        )
        with self._lock, self._connect() as connection:
            connection.execute("BEGIN IMMEDIATE")
            connection.execute(
                """
                INSERT INTO control_plane_snapshot (
                    singleton_id, schema_version, payload_json, updated_at
                ) VALUES (1, ?, ?, CURRENT_TIMESTAMP)
                ON CONFLICT(singleton_id) DO UPDATE SET
                    schema_version = excluded.schema_version,
                    payload_json = excluded.payload_json,
                    updated_at = CURRENT_TIMESTAMP
                """,
                (self.SCHEMA_VERSION, encoded),
            )
            connection.commit()


class DurableControlPlaneService(ControlPlaneService):
    """Control-plane service that persists after every successful mutation."""

    def __init__(
        self,
        repository: SQLiteSnapshotRepository,
        *,
        policy_evaluator=None,
    ) -> None:
        self.repository = repository
        store = InMemoryControlPlaneStore()
        self._restore(store, repository.load())
        super().__init__(store, policy_evaluator=policy_evaluator)

    @staticmethod
    def _restore(
        store: InMemoryControlPlaneStore,
        snapshot: dict[str, Any] | None,
    ) -> None:
        if not snapshot:
            return
        for item in snapshot.get("tenants", []):
            record = TenantRecord(**item)
            store.tenants[record.tenant_id] = record
        for item in snapshot.get("agents", []):
            normalized = dict(item)
            normalized["capabilities"] = tuple(normalized.get("capabilities", ()))
            record = AgentRecord(**normalized)
            store.agents[record.agent_id] = record
        for item in snapshot.get("tasks", []):
            normalized = dict(item)
            normalized["state"] = TaskState(normalized["state"])
            record = TaskRecord(**normalized)
            store.tasks[record.task_id] = record
        store.audit_events.extend(AuditEvent(**item) for item in snapshot.get("audit", []))
        for item in snapshot.get("idempotency", []):
            store.idempotency[(item["tenant_id"], item["key"])] = (
                item["task_id"],
                item["digest"],
            )

    def _snapshot(self) -> dict[str, Any]:
        with self.store.lock:
            return {
                "schema_version": self.repository.SCHEMA_VERSION,
                "tenants": [asdict(item) for item in self.store.tenants.values()],
                "agents": [asdict(item) for item in self.store.agents.values()],
                "tasks": [
                    {**asdict(item), "state": item.state.value}
                    for item in self.store.tasks.values()
                ],
                "audit": [asdict(item) for item in self.store.audit_events],
                "idempotency": [
                    {
                        "tenant_id": tenant_id,
                        "key": key,
                        "task_id": task_id,
                        "digest": digest,
                    }
                    for (tenant_id, key), (task_id, digest)
                    in self.store.idempotency.items()
                ],
            }

    def _mutate(self, operation: Callable[[], Any]) -> Any:
        with self.store.lock:
            result = operation()
            self.repository.save(self._snapshot())
            return result

    def register_tenant(self, *args, **kwargs):
        return self._mutate(lambda: super(DurableControlPlaneService, self).register_tenant(*args, **kwargs))

    def register_agent(self, *args, **kwargs):
        return self._mutate(lambda: super(DurableControlPlaneService, self).register_agent(*args, **kwargs))

    def create_task(self, *args, **kwargs):
        return self._mutate(lambda: super(DurableControlPlaneService, self).create_task(*args, **kwargs))

    def transition_task(self, *args, **kwargs):
        return self._mutate(lambda: super(DurableControlPlaneService, self).transition_task(*args, **kwargs))