mesa-memory 0.9.0__py3-none-any.whl

mesa_storage/control/connection_repo.py

@@ -0,0 +1,118 @@
+import logging
+from datetime import datetime, timezone
+from typing import Any
+
+from mesa_storage.sqlite_engine import AsyncEngine
+
+logger = logging.getLogger(__name__)
+
+
+class ConnectionRepository:
+    def __init__(self, sqlite_engine: AsyncEngine):
+        self._sql = sqlite_engine
+
+    async def register_connection(
+        self, connection_id: str, client_id: str, transport: str, **kwargs: Any
+    ) -> None:
+        now = datetime.now(timezone.utc).isoformat()
+        status = kwargs.get("status", "CONNECTED")
+
+        async with self._sql.transaction() as db:
+            await db.execute(
+                """
+                INSERT INTO mcp_connections (
+                    connection_id, client_id, transport, status, connected_at, last_seen_at,
+                    remote_address_hash, protocol_version, client_version, user_agent, session_id, project_id
+                ) VALUES (
+                    :connection_id, :client_id, :transport, :status, :now, :now,
+                    :remote_address_hash, :protocol_version, :client_version, :user_agent, :session_id, :project_id
+                )
+                """,
+                {
+                    "connection_id": connection_id,
+                    "client_id": client_id,
+                    "transport": transport,
+                    "status": status,
+                    "now": now,
+                    "remote_address_hash": kwargs.get("remote_address_hash"),
+                    "protocol_version": kwargs.get("protocol_version"),
+                    "client_version": kwargs.get("client_version"),
+                    "user_agent": kwargs.get("user_agent"),
+                    "session_id": kwargs.get("session_id"),
+                    "project_id": kwargs.get("project_id"),
+                },
+            )
+            await db.commit()
+
+    async def update_connection_status(self, connection_id: str, status: str) -> None:
+        now = datetime.now(timezone.utc).isoformat()
+        disconnected_at = (
+            now if status in ("DISCONNECTED", "REVOKED", "ERROR") else None
+        )
+
+        async with self._sql.transaction() as db:
+            await db.execute(
+                """
+                UPDATE mcp_connections
+                SET status = :status, last_seen_at = :now, disconnected_at = COALESCE(disconnected_at, :disconnected_at)
+                WHERE connection_id = :connection_id
+                """,
+                {
+                    "connection_id": connection_id,
+                    "status": status,
+                    "now": now,
+                    "disconnected_at": disconnected_at,
+                },
+            )
+            await db.commit()
+
+    async def get_connection(self, connection_id: str) -> dict[str, Any] | None:
+        async with self._sql.connection() as db:
+            db.row_factory = __import__("aiosqlite").Row
+            async with db.execute(
+                "SELECT * FROM mcp_connections WHERE connection_id = :connection_id",
+                {"connection_id": connection_id},
+            ) as cursor:
+                row = await cursor.fetchone()
+                return dict(row) if row else None
+
+    async def list_active_connections(
+        self, client_id: str | None = None
+    ) -> list[dict[str, Any]]:
+        async with self._sql.connection() as db:
+            db.row_factory = __import__("aiosqlite").Row
+            query = "SELECT * FROM mcp_connections WHERE status = 'CONNECTED'"
+            params: dict[str, Any] = {}
+            if client_id:
+                query += " AND client_id = :client_id"
+                params["client_id"] = client_id
+
+            async with db.execute(query, params) as cursor:
+                rows = await cursor.fetchall()
+                return [dict(row) for row in rows]
+
+    async def list_all_connections(
+        self, limit: int = 50, status: str | None = None
+    ) -> list[dict[str, Any]]:
+        query = "SELECT * FROM mcp_connections WHERE 1=1"
+        params: dict[str, Any] = {}
+        if status:
+            query += " AND status = :status"
+            params["status"] = status
+        query += " ORDER BY connected_at DESC LIMIT :limit"
+        params["limit"] = min(limit, 200)
+
+        async with self._sql.connection() as db:
+            db.row_factory = __import__("aiosqlite").Row
+            async with db.execute(query, params) as cursor:
+                rows = await cursor.fetchall()
+                return [dict(row) for row in rows]
+
+    async def count_by_status(self) -> dict[str, int]:
+        async with self._sql.connection() as db:
+            db.row_factory = __import__("aiosqlite").Row
+            async with db.execute(
+                "SELECT status, COUNT(*) as cnt FROM mcp_connections GROUP BY status"
+            ) as cursor:
+                rows = await cursor.fetchall()
+                return {row["status"]: row["cnt"] for row in rows}

mesa_storage/control/credential_repo.py

@@ -0,0 +1,196 @@
+"""Durable, binding-scoped bearer credentials for HTTP MCP clients."""
+
+from __future__ import annotations
+
+import hashlib
+import secrets
+import uuid
+from datetime import datetime, timedelta, timezone
+from typing import Any
+
+from mesa_storage.sqlite_engine import AsyncEngine
+
+DEFAULT_CREDENTIAL_TTL_SECONDS = 90 * 24 * 60 * 60
+
+
+class CredentialRepository:
+    """Issue and verify high-entropy credentials without retaining their secret."""
+
+    def __init__(self, sqlite_engine: AsyncEngine):
+        self._sql = sqlite_engine
+
+    async def issue(
+        self,
+        client_id: str,
+        binding_id: str,
+        *,
+        token_kind: str = "codex",
+        expires_in_seconds: int = DEFAULT_CREDENTIAL_TTL_SECONDS,
+    ) -> tuple[dict[str, Any], str]:
+        credential_id = f"cred_{uuid.uuid4().hex}"
+        secret = secrets.token_urlsafe(32)
+        if token_kind not in {"codex", "antigravity"}:
+            raise ValueError("unsupported credential token kind")
+        if expires_in_seconds <= 0:
+            raise ValueError("expires_in_seconds must be positive")
+        token = f"mesa_{token_kind}_{credential_id}_{secret}"
+        now = datetime.now(timezone.utc).isoformat()
+        expires_at = (
+            datetime.now(timezone.utc) + timedelta(seconds=expires_in_seconds)
+        ).isoformat()
+        record = {
+            "credential_id": credential_id,
+            "client_id": client_id,
+            "binding_id": binding_id,
+            "token_hash": _token_hash(token),
+            "token_prefix": token[:24],
+            "status": "ACTIVE",
+            "created_at": now,
+            "expires_at": expires_at,
+        }
+        async with self._sql.transaction() as db:
+            await db.execute(
+                """
+                INSERT INTO mcp_client_credentials
+                    (credential_id, client_id, binding_id, token_hash, token_prefix,
+                     status, created_at, updated_at, expires_at)
+                VALUES
+                    (:credential_id, :client_id, :binding_id, :token_hash, :token_prefix,
+                     :status, :created_at, :created_at, :expires_at)
+                """,
+                record,
+            )
+            await db.commit()
+        return record, token
+
+    async def rotate(
+        self,
+        credential_id: str,
+        *,
+        expires_in_seconds: int = DEFAULT_CREDENTIAL_TTL_SECONDS,
+    ) -> tuple[dict[str, Any], str]:
+        """Atomically replace one active credential within its existing binding."""
+        if expires_in_seconds <= 0:
+            raise ValueError("expires_in_seconds must be positive")
+        replacement_id = f"cred_{uuid.uuid4().hex}"
+        secret = secrets.token_urlsafe(32)
+        now = datetime.now(timezone.utc).isoformat()
+        expires_at = (
+            datetime.now(timezone.utc) + timedelta(seconds=expires_in_seconds)
+        ).isoformat()
+        async with self._sql.transaction() as db:
+            async with db.execute(
+                "SELECT client_id, binding_id, token_prefix FROM mcp_client_credentials "
+                "WHERE credential_id = ? AND status = 'ACTIVE'",
+                (credential_id,),
+            ) as cursor:
+                previous = await cursor.fetchone()
+            if previous is None:
+                raise ValueError("active credential not found")
+            token_kind = _token_kind_from_prefix(str(previous[2]))
+            token = f"mesa_{token_kind}_{replacement_id}_{secret}"
+            record: dict[str, Any] = {
+                "credential_id": replacement_id,
+                "client_id": str(previous[0]),
+                "binding_id": str(previous[1]),
+                "token_hash": _token_hash(token),
+                "token_prefix": token[:24],
+                "status": "ACTIVE",
+                "created_at": now,
+                "expires_at": expires_at,
+            }
+            await db.execute(
+                "INSERT INTO mcp_client_credentials "
+                "(credential_id, client_id, binding_id, token_hash, token_prefix, status, created_at, updated_at, expires_at) "
+                "VALUES (:credential_id, :client_id, :binding_id, :token_hash, :token_prefix, "
+                ":status, :created_at, :created_at, :expires_at)",
+                record,
+            )
+            cursor = await db.execute(
+                "UPDATE mcp_client_credentials SET status = 'REVOKED', revoked_at = :now, updated_at = :now "
+                "WHERE credential_id = :credential_id AND status = 'ACTIVE'",
+                {"credential_id": credential_id, "now": now},
+            )
+            if cursor.rowcount != 1:
+                raise RuntimeError("credential rotation fence lost")
+            await db.commit()
+        return record, token
+
+    async def authenticate(self, token: str) -> dict[str, Any] | None:
+        digest = _token_hash(token)
+        async with self._sql.connection() as db:
+            db.row_factory = __import__("aiosqlite").Row
+            async with db.execute(
+                """
+                SELECT credential.*, client.enabled AS client_enabled, binding.enabled AS binding_enabled
+                FROM mcp_client_credentials AS credential
+                JOIN mcp_clients AS client ON client.client_id = credential.client_id
+                JOIN mcp_project_bindings AS binding ON binding.binding_id = credential.binding_id
+                WHERE credential.token_hash = :token_hash AND credential.status = 'ACTIVE'
+                  AND (credential.expires_at IS NULL OR credential.expires_at > :now)
+                """,
+                {"token_hash": digest, "now": datetime.now(timezone.utc).isoformat()},
+            ) as cursor:
+                row = await cursor.fetchone()
+        if row is None:
+            return None
+        result = dict(row)
+        if not result["client_enabled"] or not result["binding_enabled"]:
+            return None
+        now = datetime.now(timezone.utc).isoformat()
+        async with self._sql.transaction() as db:
+            await db.execute(
+                "UPDATE mcp_client_credentials SET last_used_at = :now, updated_at = :now WHERE credential_id = :credential_id",
+                {"now": now, "credential_id": result["credential_id"]},
+            )
+            await db.commit()
+        return result
+
+    async def revoke(self, credential_id: str) -> bool:
+        now = datetime.now(timezone.utc).isoformat()
+        async with self._sql.transaction() as db:
+            cursor = await db.execute(
+                """UPDATE mcp_client_credentials
+                   SET status = 'REVOKED', revoked_at = :now, updated_at = :now
+                   WHERE credential_id = :credential_id AND status = 'ACTIVE'""",
+                {"credential_id": credential_id, "now": now},
+            )
+            await db.commit()
+        return bool(cursor.rowcount)
+
+    async def list_for_binding(self, binding_id: str) -> list[dict[str, Any]]:
+        """Return secret-safe credential summaries; token hashes stay private."""
+        async with self._sql.connection() as db:
+            db.row_factory = __import__("aiosqlite").Row
+            async with db.execute(
+                """SELECT credential_id, client_id, binding_id, token_prefix, status,
+                          created_at, updated_at, last_used_at, revoked_at, expires_at
+                   FROM mcp_client_credentials
+                   WHERE binding_id = ? ORDER BY created_at DESC""",
+                (binding_id,),
+            ) as cursor:
+                return [dict(row) for row in await cursor.fetchall()]
+
+    async def get_summary(self, credential_id: str) -> dict[str, Any] | None:
+        async with self._sql.connection() as db:
+            db.row_factory = __import__("aiosqlite").Row
+            async with db.execute(
+                """SELECT credential_id, client_id, binding_id, token_prefix, status,
+                          created_at, updated_at, last_used_at, revoked_at, expires_at
+                   FROM mcp_client_credentials WHERE credential_id = ?""",
+                (credential_id,),
+            ) as cursor:
+                row = await cursor.fetchone()
+        return dict(row) if row else None
+
+
+def _token_hash(token: str) -> str:
+    return hashlib.sha256(token.encode("utf-8")).hexdigest()
+
+
+def _token_kind_from_prefix(token_prefix: str) -> str:
+    if token_prefix.startswith("mesa_antigravity_"):
+        return "antigravity"
+    if token_prefix.startswith("mesa_codex_"):
+        return "codex"
+    raise ValueError("stored credential has an unsupported token kind")

mesa_storage/control/policy_repo.py

@@ -0,0 +1,81 @@
+import json
+import logging
+from datetime import datetime, timezone
+from typing import Any
+
+from mesa_storage.sqlite_engine import AsyncEngine
+
+logger = logging.getLogger(__name__)
+
+
+class PolicyRepository:
+    def __init__(self, sqlite_engine: AsyncEngine):
+        self._sql = sqlite_engine
+
+    async def create_rule(
+        self,
+        rule_id: str,
+        scope_type: str,
+        operation: str,
+        effect: str,
+        created_by: str,
+        scope_id: str | None = None,
+        priority: int = 100,
+        conditions: dict[str, Any] | None = None,
+    ) -> None:
+        now = datetime.now(timezone.utc).isoformat()
+        cond_json = json.dumps(conditions or {})
+
+        async with self._sql.transaction() as db:
+            await db.execute(
+                """
+                INSERT INTO mcp_policy_rules (
+                    rule_id, scope_type, scope_id, operation, effect,
+                    priority, conditions_json, created_by, created_at, updated_at
+                ) VALUES (
+                    :rule_id, :scope_type, :scope_id, :operation, :effect,
+                    :priority, :cond_json, :created_by, :now, :now
+                )
+                """,
+                {
+                    "rule_id": rule_id,
+                    "scope_type": scope_type,
+                    "scope_id": scope_id,
+                    "operation": operation,
+                    "effect": effect,
+                    "priority": priority,
+                    "cond_json": cond_json,
+                    "created_by": created_by,
+                    "now": now,
+                },
+            )
+            await db.commit()
+
+    async def get_rule(self, rule_id: str) -> dict[str, Any] | None:
+        async with self._sql.connection() as db:
+            db.row_factory = __import__("aiosqlite").Row
+            async with db.execute(
+                "SELECT * FROM mcp_policy_rules WHERE rule_id = :rule_id",
+                {"rule_id": rule_id},
+            ) as cursor:
+                row = await cursor.fetchone()
+                return dict(row) if row else None
+
+    async def list_rules(
+        self, scope_type: str | None = None, operation: str | None = None
+    ) -> list[dict[str, Any]]:
+        query = "SELECT * FROM mcp_policy_rules WHERE 1=1"
+        params = {}
+        if scope_type:
+            query += " AND scope_type = :scope_type"
+            params["scope_type"] = scope_type
+        if operation:
+            query += " AND operation = :operation"
+            params["operation"] = operation
+        query += " ORDER BY priority DESC"
+
+        async with self._sql.connection() as db:
+            db.row_factory = __import__("aiosqlite").Row
+            async with db.execute(query, params) as cursor:
+                rows = await cursor.fetchall()
+                return [dict(row) for row in rows]

mesa_storage/control/settings_repo.py

@@ -0,0 +1,56 @@
+import json
+import logging
+from datetime import datetime, timezone
+from typing import Any
+
+from mesa_storage.sqlite_engine import AsyncEngine
+
+logger = logging.getLogger(__name__)
+
+
+class SettingsRepository:
+    def __init__(self, sqlite_engine: AsyncEngine):
+        self._sql = sqlite_engine
+
+    async def get_setting(self, setting_key: str) -> Any | None:
+        async with self._sql.connection() as db:
+            db.row_factory = __import__("aiosqlite").Row
+            async with db.execute(
+                "SELECT setting_value_json FROM control_plane_settings WHERE setting_key = :key",
+                {"key": setting_key},
+            ) as cursor:
+                row = await cursor.fetchone()
+                if not row:
+                    return None
+                return json.loads(row["setting_value_json"])
+
+    async def get_all_settings(self) -> dict[str, Any]:
+        async with self._sql.connection() as db:
+            db.row_factory = __import__("aiosqlite").Row
+            async with db.execute(
+                "SELECT setting_key, setting_value_json FROM control_plane_settings"
+            ) as cursor:
+                rows = await cursor.fetchall()
+                return {
+                    row["setting_key"]: json.loads(row["setting_value_json"])
+                    for row in rows
+                }
+
+    async def set_setting(
+        self, setting_key: str, value: Any, updated_by: str = "system"
+    ) -> None:
+        now = datetime.now(timezone.utc).isoformat()
+        value_json = json.dumps(value)
+        async with self._sql.transaction() as db:
+            await db.execute(
+                """
+                INSERT INTO control_plane_settings (setting_key, setting_value_json, updated_at, updated_by)
+                VALUES (:key, :val, :now, :by)
+                ON CONFLICT(setting_key) DO UPDATE SET
+                    setting_value_json = excluded.setting_value_json,
+                    updated_at = excluded.updated_at,
+                    updated_by = excluded.updated_by
+                """,
+                {"key": setting_key, "val": value_json, "now": now, "by": updated_by},
+            )
+            await db.commit()