fastapi-gsap/gsap_broker/delegations/storage.py

"""Delegation persistence — shares the GSAP async engine.

The DelegationDB SQLModel is registered against the same SQLModel.metadata
as the rest of GSAP, so init_db() in gsap_broker.db creates this table
alongside the AuthorizationContextDB and friends.
"""

from datetime import datetime
from typing import Optional

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

from gsap_broker.db import engine


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 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:
    """Fix C-10: atomic increment with SQL-level limit check."""
    from sqlalchemy import text as sa_text
    async with AsyncSession(engine) as session:
        result = await session.execute(
            sa_text(
                "UPDATE delegations "
                "SET commands_executed = commands_executed + 1 "
                "WHERE delegation_id = :id "
                "AND commands_executed < max_commands "
                "AND status = 'active'"
            ),
            {"id": delegation_id},
        )
        await session.commit()
        if result.rowcount == 0:
            return -1  # limit reached or delegation not found/active
        # Read back the new count
        row = await session.execute(
            sa_text("SELECT commands_executed FROM delegations WHERE delegation_id = :id"),
            {"id": delegation_id},
        )
        r = row.first()
        return r[0] if r else 0


async def expire_stale() -> list[DelegationDB]:
    """Find and expire delegations past TTL or command limit.
    Fix M-23: uses atomic SQL update."""
    from sqlalchemy import text as sa_text
    from datetime import datetime, UTC
    now = datetime.now(UTC).replace(tzinfo=None)
    async with AsyncSession(engine) as session:
        # Atomically expire by TTL
        await session.execute(
            sa_text(
                "UPDATE delegations SET status = 'expired', "
                "revoke_reason = 'ttl_elapsed', revoked_at = :now "
                "WHERE status = 'active' AND expires_at < :now"
            ),
            {"now": str(now)},
        )
        # Atomically expire by command limit
        await session.execute(
            sa_text(
                "UPDATE delegations SET status = 'expired', "
                "revoke_reason = 'command_limit', revoked_at = :now "
                "WHERE status = 'active' AND commands_executed >= max_commands"
            ),
            {"now": str(now)},
        )
        await session.commit()
        # Return the expired rows for cleanup
        result = await session.exec(
            select(DelegationDB).where(
                DelegationDB.status == "expired",
                DelegationDB.revoked_at != None,
            )
        )
        return list(result.all())


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