"""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