JWT-Driven Factory

"""JWT-Driven Factory -- RBAC tool grants from trusted claims.

Demonstrates the trust split: authorization decisions (which tools to grant)
come from `ctx.trusted.claims` (set by verified JWT middleware), while
non-privileged customization comes from `ctx.input` (untrusted client input).
"""

from datetime import UTC, datetime, timedelta
from typing import Literal, Optional

import jwt as pyjwt
from agno.agent import Agent, AgentFactory
from agno.db.postgres import PostgresDb
from agno.factory import FactoryPermissionError, RequestContext
from agno.models.openai import OpenAIResponses
from agno.os import AgentOS
from agno.os.middleware import JWTMiddleware
from pydantic import BaseModel

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

JWT_SECRET = "a-string-secret-at-least-256-bits-long"

db = PostgresDb(
    id="factory-jwt-db",
    db_url="postgresql+psycopg://ai:ai@localhost:5532/ai",
)


# ---------------------------------------------------------------------------
# Simulated tools (replace with real implementations)
# ---------------------------------------------------------------------------


def read_docs() -> str:
    """Read workspace documents."""
    return "Document list: [design-spec.md, roadmap.md, api-docs.md]"


def write_docs(title: str, content: str) -> str:
    """Create or update a workspace document."""
    return f"Document '{title}' saved."


def manage_members(action: str, email: str) -> str:
    """Add or remove workspace members."""
    return f"Member {email} {action}d."


# ---------------------------------------------------------------------------
# Input schema (untrusted -- cosmetic only)
# ---------------------------------------------------------------------------


class WorkspaceInput(BaseModel):
    theme: Literal["light", "dark"] = "light"


# ---------------------------------------------------------------------------
# Factory
# ---------------------------------------------------------------------------


def build_workspace_agent(ctx: RequestContext) -> Agent:
    """Build an agent whose tools depend on the caller's JWT role."""
    # Trusted: from verified JWT middleware (request.state.claims)
    claims = ctx.trusted.claims
    role = claims.get("role")
    org_id = claims.get("org_id", "unknown")

    if not role:
        raise FactoryPermissionError("JWT must contain a 'role' claim")

    # Untrusted: from client request body (factory_input)
    cfg: Optional[WorkspaceInput] = ctx.input
    theme = cfg.theme if cfg else "light"

    # Role-based tool grants
    tools = [read_docs]
    if role in ("admin", "editor"):
        tools.append(write_docs)
    if role == "admin":
        tools.append(manage_members)

    return Agent(
        model=OpenAIResponses(id="gpt-5.4"),
        db=db,
        tools=tools,
        instructions=(
            f"You are a workspace assistant for org {org_id}.\n"
            f"The caller's role is: {role}.\n"
            f"UI theme: {theme}.\n"
            "Only use the tools available to you."
        ),
        add_datetime_to_context=True,
        markdown=True,
    )


workspace_factory = AgentFactory(
    db=db,
    id="workspace-agent",
    name="Workspace Agent",
    description="RBAC workspace agent -- tools depend on JWT role",
    factory=build_workspace_agent,
    input_schema=WorkspaceInput,
)

# ---------------------------------------------------------------------------
# AgentOS
# ---------------------------------------------------------------------------

agent_os = AgentOS(
    id="factory-jwt-demo",
    description="Demo: JWT-driven agent factory with RBAC tool grants",
    agents=[workspace_factory],
)
app = agent_os.get_app()

# Standard JWTMiddleware -- sets request.state.claims automatically
app.add_middleware(
    JWTMiddleware,
    verification_keys=[JWT_SECRET],
    algorithm="HS256",
    user_id_claim="sub",
    validate=False,  # Set True in production
)

if __name__ == "__main__":

    def make_token(role: str, org_id: str = "acme", user_id: str = "user_1") -> str:
        payload = {
            "sub": user_id,
            "role": role,
            "org_id": org_id,
            "exp": datetime.now(UTC) + timedelta(hours=24),
            "iat": datetime.now(UTC),
        }
        return pyjwt.encode(payload, JWT_SECRET, algorithm="HS256")

    print("Test tokens (valid for 24h):")
    print()
    print(f"  VIEWER:  {make_token('viewer')}")
    print(f"  EDITOR:  {make_token('editor')}")
    print(f"  ADMIN:   {make_token('admin')}")
    print()

    agent_os.serve(app="03_jwt_role_factory:app", port=7777, reload=True)

# Clone and setup repo
git clone https://github.com/agno-agi/agno.git
cd agno

# Create and activate virtual environment
./scripts/demo_setup.sh
source .venvs/demo/bin/activate

# Start Postgres for session storage
./cookbook/scripts/run_pgvector.sh

python cookbook/05_agent_os/factories/agent/03_jwt_role_factory.py