memini-ai-dev v1.0.0 — Embedded pgembed Architecture Design¶

def get_server(
    pgdata: Union[Path, str],
    cleanup_mode: Optional[str] = 'stop'  # 'stop' = stop on cleanup, None = keep running, 'delete' = stop + delete
) -> PostgresServer: ...

class PostgresServer:
    def get_uri(self) -> str: ...
    def cleanup(self) -> None: ...  # segfaults in psql subprocess on Python 3.13, use asyncpg instead

# src/memini_ai/postgres/driver.py (NEW FILE, ~470 LOC)

from __future__ import annotations

from typing import Protocol, runtime_checkable


@runtime_checkable
class DatabaseDriver(Protocol):
    """Minimal interface for backend-specific connection lifecycle."""

    async def get_uri(self) -> str:
        """Return a PostgreSQL connection URI ready for asyncpg.

        For pgembed: returns the Unix socket URI from get_server().
        For external: returns the configured MEMINI_DB_URL.

        Must be callable before initialize() — the URI is needed
        to create the asyncpg pool at PostgresDatabase.initialize():167.
        Idempotent.
        """
        ...

    async def initialize(self) -> None:
        """Perform backend-specific initialization AFTER the pool exists.

        For pgembed: starts the heartbeat + sweep background task.
        For external: no-op.
        Idempotent.
        """
        ...

    async def shutdown(self) -> None:
        """Gracefully release backend resources.

        For pgembed: stops the heartbeat task, removes this client from
        the clients map, initiates shutdown if last client. Does NOT
        stop the server unless this is the last client AND the sweep
        would also initiate shutdown.
        For external: no-op.
        """
        ...

    def is_ready(self) -> bool:
        """Synchronous check: is the backend accepting connections?

        For pgembed: checks postmaster PID is alive.
        For external: always returns True.
        """
        ...

class ExternalPGDriver:
    """Driver for external (Docker/team) Postgres — preserves v0.8.2 behavior."""

    def __init__(self, db_url: str) -> None:
        self._db_url = db_url

    async def get_uri(self) -> str:
        return self._db_url

    async def initialize(self) -> None:
        pass

    async def shutdown(self) -> None:
        pass

    def is_ready(self) -> bool:
        return True

class EmbeddedPGDriver:
    def __init__(self, data_dir: str | Path) -> None:
        self._client_id = str(uuid4())          # stable per process
        self._heartbeat_task: asyncio.Task | None = None
        self._heartbeat_stop: asyncio.Event = asyncio.Event()
        self._healthy: bool = True
        self._role: str = "passive"

    async def _heartbeat_loop(self) -> None:
        interval = self._read_state().get("heartbeat_interval_s", 60)
        consecutive_failures = 0

        while not self._heartbeat_stop.is_set():
            try:
                await asyncio.sleep(interval)
            except asyncio.CancelledError:
                break
            if self._heartbeat_stop.is_set():
                break

            try:
                self._write_heartbeat()
                consecutive_failures = 0
            except OSError as e:
                consecutive_failures += 1
                logger.warning("heartbeat_write_failed",
                    client_id=self._client_id, attempt=consecutive_failures, error=str(e))
                if consecutive_failures >= 3:
                    logger.error("heartbeat_write_failed_permanently", client_id=self._client_id)
                    self._healthy = False
                continue

            await self._sweep_stale_clients()

async def _sweep_stale_clients(self) -> None:
    state = self._read_state()
    if state is None or state.get("shutdown_initiated_by") is not None:
        return

    clients = state.get("clients", {})
    if not clients:
        return

    now = datetime.now(timezone.utc)
    timeout_s = state.get("heartbeat_timeout_s", 120)
    grace_s = state.get("grace_margin_s", 10)
    threshold = now - timedelta(seconds=timeout_s + grace_s)

    for client_id, info in clients.items():
        last_hb = datetime.fromisoformat(info["last_heartbeat"])
        if last_hb > threshold:
            return  # At least one client alive

    # ALL stale — initiate shutdown
    if not self._try_claim_shutdown(state):
        return  # Another client beat us

    await asyncio.sleep(5)  # 5s grace for late heartbeats

    # Re-check: did anyone heartbeat during grace?
    state = self._read_state()
    if state:
        for info in state.get("clients", {}).values():
            last_hb = datetime.fromisoformat(info["last_heartbeat"])
            if last_hb > threshold:
                self._clear_shutdown_token()
                return

    await self._stop_server()

def _try_claim_shutdown(self, state: dict) -> bool:
    with fasteners.InterProcessLock(str(self._state_lock)):
        current = self._read_state()
        if current is None or current.get("shutdown_initiated_by") is not None:
            return False
        current["shutdown_initiated_by"] = self._client_id
        current["shutdown_initiated_at"] = datetime.now(timezone.utc).isoformat()
        current["state"] = "draining"
        self._state_file.write_text(json.dumps(current, indent=2))
        return True

def is_healthy(self) -> bool:
    """Check if this client's connection is healthy."""
    if not self._healthy:
        return False
    state = self._read_state()
    if state is None:
        return False
    pid = state.get("pid")
    if pid is None:
        return False
    try:
        os.kill(pid, 0)
        return True
    except OSError:
        return False

def get_health_report(self) -> dict:
    """Detailed health report for memini-ai status / MCP tool."""
    # ... see architect's Section 5.1

import asyncio
import json
import os
import uuid
from datetime import datetime, timezone, timedelta
from pathlib import Path

import fasteners
import pgembed


class EmbeddedPGDriver:
    """Driver for embedded pgembed Postgres with cooperative heartbeat liveness.

    ONE embedded server, shared by ALL memini-ai processes on the same
    machine. Uses pgembed's native get_server() for process coordination
    and a server.json state file for observability + heartbeat map.
    """

    def __init__(self, data_dir: str | Path) -> None:
        self._data_dir = Path(data_dir).expanduser().resolve()
        # server.json stays in ~/.memini-ai/pgembed/ for namespace consistency
        self._state_dir = Path("~/.memini-ai/pgembed").expanduser().resolve()
        self._state_file = self._state_dir / "server.json"
        self._state_lock = self._state_dir / "server.json.lock"
        self._server = None
        self._uri = None
        self._client_id = str(uuid4())
        self._role = "passive"
        self._healthy = True
        self._heartbeat_task = None
        self._heartbeat_stop = asyncio.Event()
        self._explicit_shutdown = False

        self._state_dir.mkdir(parents=True, exist_ok=True)

    # ── Public API (DatabaseDriver Protocol) ──────────────────────

    async def get_uri(self) -> str:
        if self._uri is not None:
            return self._uri
        existing = self._read_state()
        if existing and existing.get("state") in ("running", "draining"):
            self._uri = await self._attach_to_existing(existing)
            return self._uri
        if existing and existing.get("state") == "dead":
            self._cleanup_stale_state()
        self._uri = await self._start_new_server()
        return self._uri

    async def initialize(self) -> None:
        if self._server is None:
            await self.get_uri()
        self._role = "primary"  # Now actively querying

    async def shutdown(self) -> None:
        # Stop heartbeat first
        if self._heartbeat_task is not None:
            self._heartbeat_stop.set()
            self._heartbeat_task.cancel()
            try:
                await self._heartbeat_task
            except asyncio.CancelledError:
                pass
            self._heartbeat_task = None

        if self._server is None:
            return

        # Remove this client from the map
        self._remove_client()

        # Check if we're the last
        state = self._read_state()
        remaining = state.get("clients", {}) if state else {}
        if not remaining or self._explicit_shutdown:
            await self._stop_server()

    def is_ready(self) -> bool:
        state = self._read_state()
        if state and state.get("state") == "running":
            pid = state.get("pid")
            if pid is not None:
                try:
                    os.kill(pid, 0)
                    return True
                except OSError:
                    return False
        return False

    def request_explicit_shutdown(self) -> None:
        self._explicit_shutdown = True

    # ── Server lifecycle ──────────────────────────────────────────

    async def _attach_to_existing(self, state: dict) -> str:
        self._server = pgembed.get_server(state["data_dir"], cleanup_mode=None)
        uri = self._server.get_uri()
        self._register_client()
        self._start_heartbeat()
        return uri

    async def _start_new_server(self) -> str:
        self._write_initial_state()
        try:
            self._server = pgembed.get_server(str(self._data_dir), cleanup_mode=None)
        except Exception as e:
            self._write_state("dead", refcount=0)
            raise RuntimeError(f"Failed to start embedded Postgres: {e}") from e
        uri = self._server.get_uri()
        self._register_client()
        self._start_heartbeat()
        return uri

    def _cleanup_stale_state(self) -> None:
        if self._state_file.exists():
            self._state_file.unlink()

    # ── State file management ────────────────────────────────────

    def _read_state(self) -> dict | None:
        if not self._state_file.exists():
            return None
        try:
            return json.loads(self._state_file.read_text())
        except (json.JSONDecodeError, OSError):
            return None

    def _write_initial_state(self) -> None:
        with fasteners.InterProcessLock(str(self._state_lock)):
            existing = self._read_state() or {}
            # Preserve server_started_at across restarts
            started_at = existing.get("server_started_at") or datetime.now(timezone.utc).isoformat()
            data = {
                "version": "1.0.0",
                "state": "starting",
                "pid": os.getpid(),  # will be replaced by postmaster pid
                "uri": "",
                "data_dir": str(self._data_dir),
                "server_started_at": started_at,
                "heartbeat_interval_s": 60,
                "heartbeat_timeout_s": 120,
                "grace_margin_s": 10,
                "shutdown_initiated_by": None,
                "shutdown_initiated_at": None,
                "clients": {},
            }
            self._state_file.write_text(json.dumps(data, indent=2))

    def _register_client(self) -> None:
        with fasteners.InterProcessLock(str(self._state_lock)):
            state = self._read_state() or {}
            state.setdefault("clients", {})
            state["clients"][self._client_id] = {
                "pid": os.getpid(),
                "last_heartbeat": datetime.now(timezone.utc).isoformat(),
                "role": self._role,
            }
            # Update postmaster pid if we just started
            if self._server and self._server._postmaster_info:
                state["pid"] = self._server._postmaster_info.pid
            state["state"] = "running"
            state["uri"] = self._uri or state.get("uri", "")
            self._state_file.write_text(json.dumps(state, indent=2))

    def _remove_client(self) -> None:
        with fasteners.InterProcessLock(str(self._state_lock)):
            state = self._read_state()
            if state and self._client_id in state.get("clients", {}):
                del state["clients"][self._client_id]
                self._state_file.write_text(json.dumps(state, indent=2))

    def _write_heartbeat(self) -> None:
        with fasteners.InterProcessLock(str(self._state_lock)):
            state = self._read_state() or {}
            state.setdefault("clients", {})
            state["clients"][self._client_id] = {
                "pid": os.getpid(),
                "last_heartbeat": datetime.now(timezone.utc).isoformat(),
                "role": self._role,
            }
            self._state_file.write_text(json.dumps(state, indent=2))

# In src/memini_ai/config.py, inside class MeminiConfig(BaseSettings):

vector_backend: Literal["pgembed", "postgres-external"] = Field(
    default="pgembed",
    alias="MEMINI_VECTOR_BACKEND",
    description=(
        "Vector database backend. 'pgembed' (default) starts an embedded "
        "PostgreSQL 17 server in-process. 'postgres-external' connects to "
        "an external PostgreSQL server via MEMINI_DB_URL (Docker, team server, etc.)."
    ),
)

pgembed_data_dir: str = Field(
    default="~/.local/share/memini-ai/pgembed/data",
    alias="MEMINI_PGEMBED_DATA_DIR",
    description=(
        "Data directory for the embedded pgembed PostgreSQL server. "
        "Persistent across restarts. Default: ~/.local/share/memini-ai/pgembed/data"
    ),
)

team_db_url: str = Field(
    default="",
    alias="MEMINI_TEAM_DB_URL",
    description=(
        "Optional team/shared PostgreSQL server URL for RRF fusion with "
        "embedded results."
    ),
)

fusion_mode: Literal["none", "rrf"] = Field(
    default="none",
    alias="MEMINI_FUSION_MODE",
    description=(
        "Fusion mode for multi-backend queries. 'none' (default) queries "
        "only the primary backend. 'rrf' fuses results from embedded + team."
    ),
)

def create_database(config: MeminiConfig | None = None) -> VectorDatabase:
    if config is None:
        config = get_config()

    # Q4: vector_backend always wins. No auto-detect, no silent compat.
    # v0.8.2 users with MEMINI_DB_URL set MUST set MEMINI_VECTOR_BACKEND=postgres-explicit
    # or run `memini-ai migrate` to copy data to embedded mode.

    from memini_ai.postgres.driver import EmbeddedPGDriver, ExternalPGDriver

    if config.vector_backend == "pgembed":
        data_dir = Path(config.pgembed_data_dir).expanduser()
        driver = EmbeddedPGDriver(data_dir)
    elif config.vector_backend == "postgres-external":
        db_url = config.db_url or os.environ.get("MEMINI_DB_URL", "")
        if not db_url:
            raise ValueError(
                "MEMINI_VECTOR_BACKEND=postgres-external but MEMINI_DB_URL is not set. "
                "Set MEMINI_DB_URL or switch to MEMINI_VECTOR_BACKEND=pgembed for embedded mode."
            )
        driver = ExternalPGDriver(db_url)
    else:
        raise ValueError(f"Unknown vector_backend: {config.vector_backend}")

    from memini_ai.postgres import PostgresDatabase
    db = PostgresDatabase(
        driver=driver,
        project_id=config.project_id,
        sslmode=config.db_sslmode,
        sslrootcert=config.db_sslrootcert,
    )

    if config.fusion_mode == "rrf" and config.team_db_url:
        from memini_ai.memory.rrf_database import RRFDatabase
        team_driver = ExternalPGDriver(config.team_db_url)
        team_db = PostgresDatabase(driver=team_driver, project_id=config.project_id)
        db = RRFDatabase(primary=db, secondary=team_db, k=config.rrf_k)

    return db

# In create_database(), called at server startup:
if (
    (config.db_url or os.environ.get("MEMINI_DB_URL"))
    and "MEMINI_VECTOR_BACKEND" not in os.environ
):
    raise RuntimeError(
        "memini-ai v1.0.0: MEMINI_DB_URL is set but MEMINI_VECTOR_BACKEND is not.\n"
        "\n"
        "v1.0.0 defaults to 'pgembed' (embedded PostgreSQL) which IGNORES MEMINI_DB_URL.\n"
        "Your v0.8.2 setup connected to an external Postgres server via MEMINI_DB_URL.\n"
        "\n"
        "To keep using your external server:\n"
        "  export MEMINI_VECTOR_BACKEND=postgres-external\n"
        "\n"
        "To switch to embedded mode (and copy your data):\n"
        "  unset MEMINI_DB_URL\n"
        "  memini-ai migrate --from='<your old MEMINI_DB_URL>'\n"
        "\n"
        "Set MEMINI_VECTOR_BACKEND explicitly to suppress this error."
    )

@field_validator("vector_backend")
@classmethod
def _validate_vector_backend(cls, v: str) -> str:
    allowed = {"pgembed", "postgres-external"}
    if v not in allowed:
        raise ValueError(f"vector_backend must be one of {allowed}, got '{v}'")
    return v

@field_validator("fusion_mode")
@classmethod
def _validate_fusion_mode(cls, v: str) -> str:
    allowed = {"none", "rrf"}
    if v not in allowed:
        raise ValueError(f"fusion_mode must be one of {allowed}, got '{v}'")
    return v

@field_validator("vector_backend")
@classmethod
def _validate_vector_backend(cls, v: str) -> str:
    allowed = {"pgembed", "postgres-external"}
    if v not in allowed:
        raise ValueError(f"vector_backend must be one of {allowed}, got '{v}'")
    return v

@field_validator("fusion_mode")
@classmethod
def _validate_fusion_mode(cls, v: str) -> str:
    allowed = {"none", "rrf"}
    if v not in allowed:
        raise ValueError(f"fusion_mode must be one of {allowed}, got '{v}'")
    return v

# src/memini_ai/memory/rrf_database.py (NEW FILE, ~210 LOC)

class RRFDatabase(VectorDatabase):
    """Wraps two VectorDatabase instances and fuses query results via RRF.

    Writes go to primary only. Queries fan out to both backends in
    parallel and fuse ranked lists using RRF.

    Tie-breaking: same memory_id in both → primary (embedded) wins.
    Failure mode: secondary unreachable → degrade to primary-only.
    """

    def __init__(self, primary: VectorDatabase, secondary: VectorDatabase, k: int = 60) -> None:
        self._primary = primary
        self._secondary = secondary
        self._k = k
        self._initialized = False
        self._team_write_failures: int = 0

    async def initialize(self) -> None:
        if self._initialized:
            return
        await self._primary.initialize()
        try:
            await self._secondary.initialize()
        except Exception:
            logger.warning("team_server_init_failed", message="Team server unavailable at startup")
        self._initialized = True

    async def add_memory(self, entry: MemoryEntry) -> str:
        await self.initialize()
        memory_id = await self._primary.add_memory(entry)
        if self._secondary is not None:
            asyncio.create_task(self._write_to_team(entry, memory_id))  # Q3: fire-and-forget
        return memory_id

    async def _write_to_team(self, entry: MemoryEntry, memory_id: str) -> None:
        try:
            await self._secondary.add_memory(entry)
        except Exception as e:
            self._team_write_failures += 1
            logger.warning("team_write_failed", memory_id=memory_id, error=str(e)[:200])

    async def get_memory(self, memory_id: str, include_archived: bool = False) -> MemoryEntry | None:
        await self.initialize()
        result = await self._primary.get_memory(memory_id, include_archived)
        if result is not None:
            return result
        try:
            return await self._secondary.get_memory(memory_id, include_archived)
        except Exception:
            return None

    async def query_memories(self, vector, options, collection_name=None) -> list[MemoryEntry]:
        await self.initialize()
        primary_task = asyncio.create_task(self._primary.query_memories(vector, options, collection_name))
        secondary_task = asyncio.create_task(self._secondary.query_memories(vector, options, collection_name))
        results_primary = await primary_task
        try:
            results_secondary = await secondary_task
        except Exception as e:
            logger.warning("team_server_query_failed", error=str(e)[:200], fallback="embedded_only")
            return results_primary

        ranked_ids = [[e.id for e in results_primary], [e.id for e in results_secondary]]
        fused_ids = rrf_with_limit(ranked_ids, k=self._k, limit=options.top_k)
        entries_by_id = {e.id: e for e in results_primary}
        for e in results_secondary:
            entries_by_id.setdefault(e.id, e)
        return [entries_by_id[mid] for mid in fused_ids if mid in entries_by_id]

    # ... all other read methods delegate to primary, same pattern