llm-principal-broker/llm_broker/db.py

"""Database — SQLModel + aiosqlite, matching fastapi-gsap pattern."""

from datetime import datetime
from typing import Optional
from uuid import UUID, uuid4

from sqlalchemy.ext.asyncio import create_async_engine, AsyncEngine
from sqlmodel import SQLModel, Field, select
from sqlmodel.ext.asyncio.session import AsyncSession

from llm_broker.settings import settings

engine: AsyncEngine = create_async_engine(settings.database_url, echo=False)


class DelegationDB(SQLModel, table=True):
    __tablename__ = "delegations"

    delegation_id: str = Field(primary_key=True)
    status: str = Field(default="active", index=True)
    agent_type: str
    agent_model: Optional[str] = None
    agent_did: str
    agent_keycloak_client_id: Optional[str] = None
    delegator_did: str
    delegator_ac_id: str
    delegated_ac_id: Optional[str] = None
    capability_ceiling: str = "CAP_MUTATE"
    ceremony_required_for: str = ""
    max_commands: int = 500
    commands_executed: int = 0
    created_at: datetime = Field(default_factory=lambda: datetime.utcnow())
    expires_at: datetime = Field(default_factory=lambda: datetime.utcnow())
    revoked_at: Optional[datetime] = None
    revoke_reason: Optional[str] = None
    chronicle_cid: Optional[str] = None


async def init_db():
    async with engine.begin() as conn:
        await conn.run_sync(SQLModel.metadata.create_all)


async def get_session():
    async with AsyncSession(engine) as session:
        yield session


async def create_delegation(delegation: DelegationDB) -> None:
    async with AsyncSession(engine) as session:
        session.add(delegation)
        await session.commit()


async def get_delegation(delegation_id: str) -> Optional[DelegationDB]:
    async with AsyncSession(engine) as session:
        result = await session.exec(
            select(DelegationDB).where(DelegationDB.delegation_id == delegation_id)
        )
        return result.first()


async def revoke_delegation(delegation_id: str, reason: str) -> bool:
    async with AsyncSession(engine) as session:
        result = await session.exec(
            select(DelegationDB).where(
                DelegationDB.delegation_id == delegation_id,
                DelegationDB.status == "active",
            )
        )
        d = result.first()
        if not d:
            return False
        d.status = "revoked"
        d.revoked_at = datetime.utcnow()
        d.revoke_reason = reason
        session.add(d)
        await session.commit()
        return True


async def get_active_delegations() -> list[DelegationDB]:
    async with AsyncSession(engine) as session:
        result = await session.exec(
            select(DelegationDB).where(DelegationDB.status == "active")
        )
        return list(result.all())


async def increment_commands(delegation_id: str) -> int:
    async with AsyncSession(engine) as session:
        result = await session.exec(
            select(DelegationDB).where(DelegationDB.delegation_id == delegation_id)
        )
        d = result.first()
        if not d:
            return 0
        d.commands_executed += 1
        session.add(d)
        await session.commit()
        return d.commands_executed


async def expire_stale() -> list[DelegationDB]:
    """Find and expire delegations past TTL or command limit."""
    now = datetime.utcnow()
    async with AsyncSession(engine) as session:
        result = await session.exec(
            select(DelegationDB).where(DelegationDB.status == "active")
        )
        expired = []
        for d in result.all():
            if now > d.expires_at or d.commands_executed >= d.max_commands:
                d.status = "expired"
                d.revoke_reason = (
                    "command_limit" if d.commands_executed >= d.max_commands else "ttl_elapsed"
                )
                d.revoked_at = now
                session.add(d)
                expired.append(d)
        await session.commit()
        return expired


async def revoke_by_delegator_ac(delegator_ac_id: str) -> int:
    """Cascading revocation — all delegations from a specific AC."""
    now = datetime.utcnow()
    async with AsyncSession(engine) as session:
        result = await session.exec(
            select(DelegationDB).where(
                DelegationDB.status == "active",
                DelegationDB.delegator_ac_id == delegator_ac_id,
            )
        )
        count = 0
        for d in result.all():
            d.status = "revoked"
            d.revoked_at = now
            d.revoke_reason = "delegator_ac_revoked"
            session.add(d)
            count += 1
        await session.commit()
        return count