llm-principal-broker/llm_broker/main.py

"""LLM Principal Broker — GCAP-SPEC-LLM-PRINCIPAL-BROKER-0001.

AI agent identity delegation for governed shell sessions.
Companion service to the GSAP broker.
"""

import asyncio
import logging
from contextlib import asynccontextmanager
from datetime import datetime

import structlog
from fastapi import FastAPI, HTTPException, Header
from fastapi.middleware.cors import CORSMiddleware

from llm_broker import db
from llm_broker.delegation import DelegationManager
from llm_broker.models import (
    ActiveDelegation,
    AgentListResponse,
    DelegationInfo,
    DelegationRequest,
    DelegationResponse,
    DelegationStatus,
    RevokeRequest,
    RevokeResponse,
)
from llm_broker.settings import settings

logger = structlog.get_logger()
manager = DelegationManager(settings)


@asynccontextmanager
async def lifespan(app: FastAPI):
    await db.init_db()
    logger.info("llm-principal-broker started", broker_did=settings.broker_did)
    task = asyncio.create_task(_cleanup_loop())
    yield
    task.cancel()


async def _cleanup_loop():
    """Periodically expire stale delegations (30s interval)."""
    while True:
        try:
            count = await manager.cleanup_expired()
        except Exception as e:
            logger.warning("cleanup error", error=str(e))
        await asyncio.sleep(30)


app = FastAPI(
    title="LLM Principal Broker",
    description="AI agent identity delegation — GCAP-SPEC-LLM-PRINCIPAL-BROKER-0001",
    version="0.1.0",
    lifespan=lifespan,
)
app.add_middleware(
    CORSMiddleware,
    allow_origins=settings.cors_origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.post("/delegate", response_model=DelegationResponse, tags=["Delegation"])
async def create_delegation(
    request: DelegationRequest,
    x_delegator_did: str = Header(..., alias="X-Delegator-DID"),
):
    """Request delegation of authority to an AI agent — §8.1."""
    try:
        return await manager.create_delegation(request, x_delegator_did)
    except RuntimeError as e:
        raise HTTPException(status_code=502, detail=str(e))


@app.post("/delegate/{delegation_id}/revoke", response_model=RevokeResponse, tags=["Delegation"])
async def revoke_delegation(
    delegation_id: str,
    request: RevokeRequest = RevokeRequest(),
):
    """Revoke an active delegation — §8.2."""
    success = await manager.revoke_delegation(delegation_id, request.reason)
    if not success:
        raise HTTPException(status_code=404, detail="Delegation not found or not active")
    return RevokeResponse(delegation_id=delegation_id, reason=request.reason)


@app.get("/delegate/{delegation_id}", response_model=DelegationInfo, tags=["Delegation"])
async def get_delegation(delegation_id: str):
    """Query delegation status — §8.3."""
    d = await db.get_delegation(delegation_id)
    if not d:
        raise HTTPException(status_code=404, detail="Delegation not found")

    now = datetime.utcnow()
    ttl_remaining = max(0, int((d.expires_at - now).total_seconds()))

    return DelegationInfo(
        delegation_id=d.delegation_id,
        status=DelegationStatus(d.status),
        agent_did=d.agent_did,
        agent_type=d.agent_type,
        delegator_did=d.delegator_did,
        commands_executed=d.commands_executed,
        commands_remaining=max(0, d.max_commands - d.commands_executed),
        ttl_remaining_seconds=ttl_remaining,
        created_at=d.created_at.isoformat(),
        expires_at=d.expires_at.isoformat(),
    )


@app.get("/agents", response_model=AgentListResponse, tags=["Delegation"])
async def list_agents():
    """List all active agent delegations — §8.4."""
    active = await db.get_active_delegations()
    now = datetime.utcnow()

    return AgentListResponse(
        active_delegations=[
            ActiveDelegation(
                delegation_id=d.delegation_id,
                agent_type=d.agent_type,
                delegator=d.delegator_did,
                commands_executed=d.commands_executed,
                ttl_remaining_seconds=max(0, int((d.expires_at - now).total_seconds())),
                status=d.status,
            )
            for d in active
        ],
        total_active=len(active),
    )


@app.get("/health", tags=["Health"])
async def health():
    active = await db.get_active_delegations()
    return {
        "status": "healthy",
        "service": "llm-principal-broker",
        "version": "0.1.0",
        "active_delegations": len(active),
    }