88 lines
3.3 KiB
Python
Raw Permalink Normal View History

Enable Auth in AGS (#5928) <!-- Thank you for your contribution! Please review https://microsoft.github.io/autogen/docs/Contribute before opening a pull request. --> <!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. --> ## Why are these changes needed? https://github.com/user-attachments/assets/b649053b-c377-40c7-aa51-ee64af766fc2 <img width="100%" alt="image" src="https://github.com/user-attachments/assets/03ba1df5-c9a2-4734-b6a2-0eb97ec0b0e0" /> ## Authentication This PR implements an experimental authentication feature to enable personalized experiences (multiple users). Currently, only GitHub authentication is supported. You can extend the base authentication class to add support for other authentication methods. By default authenticatio is disabled and only enabled when you pass in the `--auth-config` argument when running the application. ### Enable GitHub Authentication To enable GitHub authentication, create a `auth.yaml` file in your app directory: ```yaml type: github jwt_secret: "your-secret-key" token_expiry_minutes: 60 github: client_id: "your-github-client-id" client_secret: "your-github-client-secret" callback_url: "http://localhost:8081/api/auth/callback" scopes: ["user:email"] ``` Please see the documentation on [GitHub OAuth](https://docs.github.com/en/apps/oauth-apps/building-oauth-apps/authenticating-to-the-rest-api-with-an-oauth-app) for more details on obtaining the `client_id` and `client_secret`. To pass in this configuration you can use the `--auth-config` argument when running the application: ```bash autogenstudio ui --auth-config /path/to/auth.yaml ``` Or set the environment variable: ```bash export AUTOGENSTUDIO_AUTH_CONFIG="/path/to/auth.yaml" ``` ```{note} - Authentication is currently experimental and may change in future releases - User data is stored in your configured database - When enabled, all API endpoints require authentication except for the authentication endpoints - WebSocket connections require the token to be passed as a query parameter (`?token=your-jwt-token`) ``` ## Related issue number <!-- For example: "Closes #1234" --> Closes #4350 ## Checks - [ ] I've included any doc changes needed for <https://microsoft.github.io/autogen/>. See <https://github.com/microsoft/autogen/blob/main/CONTRIBUTING.md> to build and test documentation locally. - [ ] I've added tests (if relevant) corresponding to the changes introduced in this PR. - [ ] I've made sure all auto checks have passed. --------- Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
2025-03-14 15:02:05 -07:00
import jwt
from fastapi import WebSocket, WebSocketDisconnect, status
from loguru import logger
from .manager import AuthManager
from .models import User
class WebSocketAuthHandler:
"""
Helper class for authenticating WebSocket connections.
"""
def __init__(self, auth_manager: AuthManager):
self.auth_manager = auth_manager
async def authenticate(self, websocket: WebSocket) -> tuple[bool, User | None]:
"""
Authenticate a WebSocket connection.
Returns (success, user) tuple.
"""
if self.auth_manager.config.type == "none":
# No authentication required
return True, User(id="guestuser@gmail.com", name="Default User", provider="none")
try:
# Extract token from query params or headers query_params)
token = None
if "token" in websocket.query_params:
token = websocket.query_params["token"]
elif "authorization" in websocket.headers:
auth_header = websocket.headers["authorization"]
if auth_header.startswith("Bearer "):
token = auth_header.replace("Bearer ", "")
if not token:
logger.warning("No token found for WebSocket connection")
return False, None
# Validate token
if not self.auth_manager.config.jwt_secret:
# Development mode with no JWT secret
return True, User(id="guestuser@gmail.com", name="Default User", provider="none")
try:
# Decode and validate JWT
if not self.auth_manager.config.jwt_secret:
logger.warning("Invalid token for WebSocket connection")
return False, None
payload = jwt.decode(token, self.auth_manager.config.jwt_secret, algorithms=["HS256"])
# Create User object from token payload
user = User(
id=payload.get("sub"),
name=payload.get("name", "Unknown User"),
email=payload.get("email"),
provider=payload.get("provider", "jwt"),
roles=payload.get("roles", ["user"]),
)
return True, user
except jwt.ExpiredSignatureError:
logger.warning("Expired token for WebSocket connection")
return False, None
except jwt.InvalidTokenError:
logger.warning("Invalid token for WebSocket connection")
return False, None
except Exception as e:
logger.error(f"WebSocket auth error: {str(e)}")
return False, None
async def on_connect(self, websocket: WebSocket) -> User | None:
"""
Handle WebSocket connection with authentication.
Returns authenticated user if successful, otherwise closes the connection.
"""
success, user = await self.authenticate(websocket)
if not success:
# Authentication failed, close the connection
await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="Authentication failed")
raise WebSocketDisconnect(code=status.WS_1008_POLICY_VIOLATION)
# Authentication successful, return the user
return user