Custom Scope Mappings

import os
from datetime import UTC, datetime, timedelta

import jwt
from agno.agent import Agent
from agno.db.postgres import PostgresDb
from agno.models.openai import OpenAIChat
from agno.os import AgentOS
from agno.os.middleware import JWTMiddleware
from agno.tools.duckduckgo import DuckDuckGoTools

JWT_SECRET = os.getenv("JWT_VERIFICATION_KEY", "your-secret-key-at-least-256-bits-long")

db = PostgresDb(db_url="postgresql+psycopg://ai:ai@localhost:5532/ai")

research_agent = Agent(
    id="research-agent",
    name="Research Agent",
    model=OpenAIChat(id="gpt-4o"),
    db=db,
    tools=[DuckDuckGoTools()],
    markdown=True,
)

# Define custom scope mappings
# Format: "METHOD /path": ["scope1", "scope2"]
custom_scopes = {
    # Override default agent scopes with custom ones
    "GET /agents": ["app:read"],
    "GET /agents/*": ["app:read"],
    "POST /agents/*/runs": ["app:run", "app:execute"],  # Require BOTH scopes
    
    # Session endpoints - only admins can view
    "GET /sessions": ["app:admin"],
    "GET /sessions/*": ["app:read", "sessions:read"],
    
    # Memory endpoints with custom scopes
    "GET /memories": ["memory:admin"],
    "POST /memories": ["memory:write"],
    
    # Config endpoint - system admins only
    "GET /config": ["app:admin"],
}

# Create AgentOS without built-in authorization
agent_os = AgentOS(
    id="my-agent-os",
    description="Custom Scope Mappings AgentOS",
    agents=[research_agent],
)

app = agent_os.get_app()

# Add JWT middleware with custom scope mappings
app.add_middleware(
    JWTMiddleware,
    verification_keys=[JWT_SECRET],
    algorithm="HS256",
    scope_mappings=custom_scopes,  # Custom scopes enable RBAC automatically
    admin_scope="app:superadmin",  # Custom admin scope
)


if __name__ == "__main__":
    # Basic user - can only read
    basic_user_token = jwt.encode(
        {
            "sub": "user_123",
            "scopes": ["app:read"],
            "exp": datetime.now(UTC) + timedelta(hours=24),
        },
        JWT_SECRET,
        algorithm="HS256",
    )

    # Power user - can read and execute
    power_user_token = jwt.encode(
        {
            "sub": "user_456",
            "scopes": ["app:read", "app:run", "app:execute"],
            "exp": datetime.now(UTC) + timedelta(hours=24),
        },
        JWT_SECRET,
        algorithm="HS256",
    )

    # Admin user - has admin scope
    admin_token = jwt.encode(
        {
            "sub": "admin_789",
            "scopes": ["app:admin", "app:read", "app:run", "app:execute"],
            "exp": datetime.now(UTC) + timedelta(hours=24),
        },
        JWT_SECRET,
        algorithm="HS256",
    )

    # Super admin - bypasses all checks
    superadmin_token = jwt.encode(
        {
            "sub": "superadmin",
            "scopes": ["app:superadmin"],
            "exp": datetime.now(UTC) + timedelta(hours=24),
        },
        JWT_SECRET,
        algorithm="HS256",
    )

    print("Basic User (app:read only):")
    print(basic_user_token)
    print("\nPower User (app:read, app:run, app:execute):")
    print(power_user_token)
    print("\nAdmin (app:admin + all permissions):")
    print(admin_token)
    print("\nSuper Admin (app:superadmin - bypasses all):")
    print(superadmin_token)

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