JWT Middleware with Cookies

from datetime import UTC, datetime, timedelta

import jwt
from agno.agent import Agent
from agno.db.postgres import PostgresDb
from agno.models.openai import OpenAIResponses
from agno.os import AgentOS
from agno.run import RunContext
from agno.os.middleware.jwt import AuthMiddleware, TokenSource
from fastapi import FastAPI, Response

# JWT Secret (use environment variable in production)
JWT_SECRET = "a-string-secret-at-least-256-bits-long"

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


def get_user_profile(run_context: RunContext) -> dict:
    """
    Get the current user's profile.
    """
    dependencies = run_context.dependencies or {}
    return {
        "name": dependencies.get("name", "Unknown"),
        "email": dependencies.get("email", "Unknown"),
        "roles": dependencies.get("roles", []),
        "organization": dependencies.get("org", "Unknown"),
    }


# Create agent
profile_agent = Agent(
    id="profile-agent",
    name="Profile Agent",
    model=OpenAIResponses(id="gpt-5.2"),
    db=db,
    tools=[get_user_profile],
    instructions="You are a profile agent. You can search for information and access user profiles.",
    add_history_to_context=True,
    markdown=True,
)


app = FastAPI()


# Add a simple endpoint to set the JWT authentication cookie
@app.get("/set-auth-cookie")
async def set_auth_cookie(response: Response):
    """
    Endpoint to set the JWT authentication cookie.
    In a real application, this would be done after successful login.
    """
    # Create a test JWT token
    payload = {
        "sub": "cookie_user_789",
        "session_id": "cookie_session_123",
        "name": "Jane Smith",
        "email": "jane.smith@example.com",
        "roles": ["user", "premium"],
        "org": "Example Corp",
        "exp": datetime.now(UTC) + timedelta(hours=24),
        "iat": datetime.now(UTC),
    }

    token = jwt.encode(payload, JWT_SECRET, algorithm="HS256")

    # Set HTTP-only cookie (more secure than localStorage for JWT storage)
    response.set_cookie(
        key="auth_token",
        value=token,
        httponly=True,  # Prevents JavaScript from reading the cookie
        secure=True,  # Only send over HTTPS in production
        samesite="strict",  # Restricts cross-site cookie sending
        max_age=24 * 60 * 60,  # 24 hours
    )

    return {
        "message": "Authentication cookie set successfully",
        "cookie_name": "auth_token",
        "expires_in": "24 hours",
        "security_features": ["httponly", "secure", "samesite=strict"],
        "instructions": "Now you can make authenticated requests without Authorization headers",
    }


# Add a simple endpoint to clear the JWT authentication cookie
@app.get("/clear-auth-cookie")
async def clear_auth_cookie(response: Response):
    """Endpoint to clear the JWT authentication cookie (logout)."""
    response.delete_cookie(key="auth_token")
    return {"message": "Authentication cookie cleared successfully"}


# Add auth middleware configured for cookie-based authentication
app.add_middleware(
    AuthMiddleware,
    verification_keys=[JWT_SECRET], # or use JWT_VERIFICATION_KEY environment variable
    algorithm="HS256",
    excluded_route_paths=[
        "/set-auth-cookie",
        "/clear-auth-cookie",
    ],
    token_source=TokenSource.COOKIE,  # Extract JWT from cookies
    cookie_name="auth_token",  # Name of the cookie containing the JWT
    user_id_claim="sub",  # Extract user_id from 'sub' claim
    session_id_claim="session_id",  # Extract session_id from 'session_id' claim
    dependencies_claims=[
        "name",
        "email",
        "roles",
        "org",
    ],  # Additional claims to extract
    validate=True,  # We want to ensure the token is valid
)


agent_os = AgentOS(
    description="JWT Cookie-Based AgentOS",
    agents=[profile_agent],
    base_app=app,
)

# Get the final app
app = agent_os.get_app()


if __name__ == "__main__":
    """
    Run your AgentOS with JWT cookie authentication.
    """

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