diff --git a/docs/a2a.md b/docs/a2a.md index be3e882c..dfe6cc02 100644 --- a/docs/a2a.md +++ b/docs/a2a.md @@ -94,3 +94,44 @@ Configure via environment variable: ```bash export A2A_MAX_CONTEXTS=1000 ``` + +## Security + +By default, the A2A agent runs without authenticationto. For production deployments, you should add authentication. + +### Adding Authentication + +The `create_a2a_app()` function accepts optional security parameters that declare authentication requirements in the agent card: + +```python +from haiku.rag.a2a import create_a2a_app + +app = create_a2a_app( + db_path, + security_schemes={ + "apiKeyAuth": { + "type": "apiKey", + "in": "header", + "name": "X-API-Key", + "description": "API key authentication", + } + }, + security=[{"apiKeyAuth": []}], +) +``` + +This populates the agent card at `/.well-known/agent-card.json` so other agents can discover your authentication requirements. + +### Security Examples + +Three working examples are provided in `examples/a2a-security/`: + +1. **API Key** (`apikey_example.py`) - Simple header-based authentication +2. **OAuth2 GitHub** (`oauth2_github.py`) - GitHub Personal Access Token authentication +3. **OAuth2 Enterprise** (`oauth2_example.py`) - Full OAuth2 with JWT verification + +Each example shows: + +- How to declare security in the agent card +- How to implement authentication middleware +- How to verify credentials diff --git a/examples/a2a-security/apikey_example.py b/examples/a2a-security/apikey_example.py new file mode 100644 index 00000000..c37981ba --- /dev/null +++ b/examples/a2a-security/apikey_example.py @@ -0,0 +1,130 @@ +"""Example: Adding API Key authentication to haiku.rag A2A agent. + +Simple header-based authentication suitable for internal services and development. +Perfect for getting started with A2A authentication. + +Setup: + # Run with default key + python apikey_example.py /path/to/database.lancedb + + # Or use your own key + export API_KEY='your-secret-key' + python apikey_example.py /path/to/database.lancedb + +Usage: + # Make authenticated request (default key is demo-key-12345) + curl -H "X-API-Key: demo-key-12345" \ + -H "Content-Type: application/json" \ + -X POST http://localhost:8000/ \ + -d '{"jsonrpc":"2.0","method":"message/send","params":{"contextId":"test","message":{"kind":"message","role":"user","messageId":"msg-1","parts":[{"kind":"text","text":"What is Python?"}]}},"id":1}' +""" + +import os +from pathlib import Path + +from starlette.exceptions import HTTPException +from starlette.responses import JSONResponse +from starlette.status import HTTP_401_UNAUTHORIZED + +from haiku.rag.a2a import create_a2a_app + +# API Key Configuration - In production, use environment variables or a secure key store +API_KEY_NAME = "X-API-Key" +VALID_API_KEY = os.getenv("API_KEY", "demo-key-12345") + + +def verify_api_key(api_key: str | None) -> str: + """Verify API key from request header. + + Args: + api_key: API key from X-API-Key header + + Returns: + The verified API key + + Raises: + HTTPException: If API key is missing or invalid + """ + if not api_key: + raise HTTPException( + status_code=HTTP_401_UNAUTHORIZED, + detail="Missing API key", + headers={"WWW-Authenticate": f'ApiKey realm="{API_KEY_NAME}"'}, + ) + + if api_key != VALID_API_KEY: + raise HTTPException( + status_code=HTTP_401_UNAUTHORIZED, + detail="Invalid API key", + headers={"WWW-Authenticate": f'ApiKey realm="{API_KEY_NAME}"'}, + ) + + return api_key + + +def create_secure_a2a_app(db_path: Path): + """Create A2A app with API key authentication. + + Args: + db_path: Path to LanceDB database + + Returns: + FastA2A application with API key security + """ + # Create app with security declared in AgentCard + app = create_a2a_app( + db_path, + security_schemes={ + "apiKeyAuth": { + "type": "apiKey", + "in": "header", + "name": API_KEY_NAME, + "description": "API key authentication", + } + }, + security=[{"apiKeyAuth": []}], + ) + + # Add authentication middleware + @app.middleware("http") + async def authenticate_request(request, call_next): + """Middleware to verify API key on all requests.""" + # Skip authentication for well-known endpoints + if request.url.path in [ + "/.well-known/agent-card.json", + "/health", + "/docs", + "/openapi.json", + ]: + return await call_next(request) + + # Verify API key + api_key = request.headers.get(API_KEY_NAME) + try: + verify_api_key(api_key) + except HTTPException as e: + return JSONResponse( + status_code=e.status_code, + content={"detail": e.detail}, + headers=e.headers or {}, + ) + + # Continue with request + return await call_next(request) + + return app + + +if __name__ == "__main__": + import sys + + import uvicorn + + if len(sys.argv) < 2: + print("Usage: python apikey_example.py ") + sys.exit(1) + + db_path = Path(sys.argv[1]) + app = create_secure_a2a_app(db_path) + + uvicorn.run(app, host="127.0.0.1", port=8000) diff --git a/examples/a2a-security/oauth2_example.py b/examples/a2a-security/oauth2_example.py new file mode 100644 index 00000000..6caa4627 --- /dev/null +++ b/examples/a2a-security/oauth2_example.py @@ -0,0 +1,222 @@ +"""Example: Adding OAuth2 authentication to haiku.rag A2A agent. + +This example demonstrates OAuth2 client credentials flow with JWT token verification. +Suitable for enterprise environments with existing OAuth2 infrastructure. + +Requirements: + uv pip install python-jose[cryptography] + +Setup: + 1. Set up an OAuth2 provider (Auth0, Okta, Azure AD, Keycloak, etc.) + 2. Create an API and a machine-to-machine application + 3. Get the token URL and public key from your provider + 4. Set environment variables: + export OAUTH2_TOKEN_URL='https://your-auth.example.com/oauth/token' + export OAUTH2_PUBLIC_KEY='-----BEGIN PUBLIC KEY-----...' + +Usage: + python oauth2_example.py /path/to/database.lancedb + + # Get access token from your OAuth2 provider: + TOKEN=$(curl -X POST $OAUTH2_TOKEN_URL \ + -d "grant_type=client_credentials" \ + -d "client_id=your-client-id" \ + -d "client_secret=your-client-secret" \ + -d "scope=read:documents query:documents" \ + | jq -r '.access_token') + + # Make authenticated request: + curl -H "Authorization: Bearer $TOKEN" \ + -H "Content-Type: application/json" \ + -X POST http://localhost:8000/ \ + -d '{"jsonrpc":"2.0","method":"message/send","params":{"contextId":"test","message":{"kind":"message","role":"user","messageId":"msg-1","parts":[{"kind":"text","text":"What is Python?"}]}},"id":1}' +""" + +import os +from pathlib import Path + +from jose import JWTError, jwt +from starlette.exceptions import HTTPException +from starlette.responses import JSONResponse +from starlette.status import ( + HTTP_401_UNAUTHORIZED, + HTTP_403_FORBIDDEN, + HTTP_500_INTERNAL_SERVER_ERROR, +) + +from haiku.rag.a2a import create_a2a_app + +# OAuth2 Configuration +OAUTH2_TOKEN_URL = os.getenv( + "OAUTH2_TOKEN_URL", "https://your-auth.example.com/oauth/token" +) +OAUTH2_AUTH_URL = os.getenv( + "OAUTH2_AUTH_URL", "https://your-auth.example.com/oauth/authorize" +) +OAUTH2_PUBLIC_KEY = os.getenv("OAUTH2_PUBLIC_KEY", "") +OAUTH2_ALGORITHM = os.getenv("OAUTH2_ALGORITHM", "RS256") + +# Define required scopes for each skill +SKILL_SCOPES = { + "document-qa": ["read:documents", "query:documents"], +} + + +def verify_token(token: str) -> dict: + """Verify JWT token from OAuth2 provider. + + Args: + token: JWT token from Authorization header + + Returns: + Dictionary with user info and scopes + + Raises: + HTTPException: If token is invalid or expired + """ + credentials_exception = HTTPException( + status_code=HTTP_401_UNAUTHORIZED, + detail="Could not validate credentials", + headers={"WWW-Authenticate": "Bearer"}, + ) + + if not OAUTH2_PUBLIC_KEY: + raise HTTPException( + status_code=HTTP_500_INTERNAL_SERVER_ERROR, + detail="OAuth2 public key not configured", + ) + + try: + payload = jwt.decode( + token, + OAUTH2_PUBLIC_KEY, + algorithms=[OAUTH2_ALGORITHM], + ) + + username: str | None = payload.get("sub") + scopes: list[str] = ( + payload.get("scope", "").split() + if isinstance(payload.get("scope"), str) + else payload.get("scope", []) + ) + + if username is None: + raise credentials_exception + + return {"username": username, "scopes": scopes} + + except JWTError as e: + raise HTTPException( + status_code=HTTP_401_UNAUTHORIZED, + detail=f"Invalid token: {str(e)}", + headers={"WWW-Authenticate": "Bearer"}, + ) from e + + +def check_skill_permissions(skill_id: str, credentials: dict) -> None: + """Verify that user has required scopes for a skill. + + Args: + skill_id: The skill being accessed + credentials: User credentials with scopes + + Raises: + HTTPException: If user lacks required permissions + """ + required_scopes = SKILL_SCOPES.get(skill_id, []) + user_scopes = credentials.get("scopes", []) + + missing_scopes = [scope for scope in required_scopes if scope not in user_scopes] + + if missing_scopes: + raise HTTPException( + status_code=HTTP_403_FORBIDDEN, + detail=f"Missing required scopes: {', '.join(missing_scopes)} for skill: {skill_id}", + ) + + +def create_secure_a2a_app(db_path: Path): + """Create A2A app with OAuth2 authentication. + + Args: + db_path: Path to LanceDB database + + Returns: + FastA2A application with OAuth2 security + """ + # Create app with security declared in AgentCard + app = create_a2a_app( + db_path, + security_schemes={ + "oauth2": { + "type": "oauth2", + "flows": { + "clientCredentials": { + "tokenUrl": OAUTH2_TOKEN_URL, + "scopes": { + "read:documents": "Read document content", + "query:documents": "Search and query documents", + }, + } + }, + "description": "OAuth2 client credentials flow", + } + }, + security=[{"oauth2": ["read:documents", "query:documents"]}], + ) + + # Add authentication middleware + @app.middleware("http") + async def authenticate_request(request, call_next): + """Middleware to verify OAuth2 token on all requests.""" + # Skip authentication for well-known endpoints + if request.url.path in [ + "/.well-known/agent-card.json", + "/health", + "/docs", + "/openapi.json", + ]: + return await call_next(request) + + # Get token from Authorization header + auth_header = request.headers.get("Authorization", "") + if not auth_header.startswith("Bearer "): + return JSONResponse( + status_code=HTTP_401_UNAUTHORIZED, + content={"detail": "Missing or invalid Authorization header"}, + headers={"WWW-Authenticate": "Bearer"}, + ) + + token = auth_header[7:] # Remove "Bearer " prefix + + # Verify token + try: + credentials = verify_token(token) + # Attach credentials to request state for use in handlers + request.state.credentials = credentials + except HTTPException as e: + return JSONResponse( + status_code=e.status_code, + content={"detail": e.detail}, + headers=e.headers or {}, + ) + + # Continue with request + return await call_next(request) + + return app + + +if __name__ == "__main__": + import sys + + import uvicorn + + if len(sys.argv) < 2: + print("Usage: python oauth2_example.py ") + sys.exit(1) + + db_path = Path(sys.argv[1]) + app = create_secure_a2a_app(db_path) + + uvicorn.run(app, host="127.0.0.1", port=8000) diff --git a/examples/a2a-security/oauth2_github.py b/examples/a2a-security/oauth2_github.py new file mode 100644 index 00000000..f2b10a41 --- /dev/null +++ b/examples/a2a-security/oauth2_github.py @@ -0,0 +1,193 @@ +"""Example: Using GitHub Personal Access Tokens for authentication. + +This is a simplified OAuth2 example that uses GitHub Personal Access Tokens. +It's much easier to set up than full OAuth2 and perfect for testing. + +Setup: + 1. Go to https://github.com/settings/tokens + 2. Click "Generate new token (classic)" + 3. Give it a name and select scopes + 4. Copy the generated token + +Usage: + export GITHUB_TOKENS="your_github_token_here" + python oauth2_github.py /path/to/database.lancedb + + # Make authenticated request: + curl -H "Authorization: Bearer ghp_your_token" \ + -H "Content-Type: application/json" \ + -X POST http://localhost:8000/ \ + -d '{"jsonrpc":"2.0","method":"message/send","params":{"contextId":"test","message":{"kind":"message","role":"user","messageId":"msg-1","parts":[{"kind":"text","text":"What is Python?"}]}},"id":1}' +""" + +import os +from pathlib import Path + +import httpx +from starlette.exceptions import HTTPException +from starlette.responses import JSONResponse +from starlette.status import HTTP_401_UNAUTHORIZED + +from haiku.rag.a2a import create_a2a_app + +# Configuration +GITHUB_API_URL = "https://api.github.com" +ALLOWED_TOKENS = ( + set(os.getenv("GITHUB_TOKENS", "").split(",")) + if os.getenv("GITHUB_TOKENS") + else set() +) + + +async def verify_github_token(token: str) -> dict: + """Verify GitHub Personal Access Token by calling GitHub API. + + Args: + token: GitHub Personal Access Token (starts with ghp_) + + Returns: + Dictionary with user info + + Raises: + HTTPException: If token is invalid + """ + if not token.startswith("ghp_") and not token.startswith("github_pat_"): + raise HTTPException( + status_code=HTTP_401_UNAUTHORIZED, + detail="Invalid GitHub token format", + headers={"WWW-Authenticate": "Bearer"}, + ) + + # If we have a list of allowed tokens, check against it + if ALLOWED_TOKENS and token not in ALLOWED_TOKENS: + raise HTTPException( + status_code=HTTP_401_UNAUTHORIZED, + detail="Token not in allowed list", + headers={"WWW-Authenticate": "Bearer"}, + ) + + # Verify token with GitHub API + async with httpx.AsyncClient() as client: + try: + response = await client.get( + f"{GITHUB_API_URL}/user", + headers={ + "Authorization": f"Bearer {token}", + "Accept": "application/vnd.github+json", + "X-GitHub-Api-Version": "2022-11-28", + }, + timeout=5.0, + ) + + if response.status_code == 401: + raise HTTPException( + status_code=HTTP_401_UNAUTHORIZED, + detail="Invalid or expired GitHub token", + headers={"WWW-Authenticate": "Bearer"}, + ) + + if response.status_code != 200: + raise HTTPException( + status_code=HTTP_401_UNAUTHORIZED, + detail=f"GitHub API error: {response.status_code}", + headers={"WWW-Authenticate": "Bearer"}, + ) + + user_data = response.json() + return { + "username": user_data.get("login"), + "email": user_data.get("email"), + "name": user_data.get("name"), + "github_id": user_data.get("id"), + } + + except httpx.TimeoutException: + raise HTTPException( + status_code=HTTP_401_UNAUTHORIZED, + detail="GitHub API timeout", + headers={"WWW-Authenticate": "Bearer"}, + ) + except httpx.RequestError as e: + raise HTTPException( + status_code=HTTP_401_UNAUTHORIZED, + detail=f"Failed to verify token with GitHub: {str(e)}", + headers={"WWW-Authenticate": "Bearer"}, + ) + + +def create_secure_a2a_app(db_path: Path): + """Create A2A app with GitHub token authentication. + + Args: + db_path: Path to LanceDB database + + Returns: + FastA2A application with GitHub authentication + """ + # Create app with security declared in AgentCard + app = create_a2a_app( + db_path, + security_schemes={ + "githubAuth": { + "type": "http", + "scheme": "bearer", + "description": "GitHub Personal Access Token authentication", + } + }, + security=[{"githubAuth": []}], + ) + + @app.middleware("http") + async def authenticate_request(request, call_next): + """Middleware to verify GitHub token on all requests.""" + # Skip authentication for well-known endpoints + if request.url.path in [ + "/.well-known/agent-card.json", + "/health", + "/docs", + "/openapi.json", + ]: + return await call_next(request) + + # Get token from Authorization header + auth_header = request.headers.get("Authorization", "") + if not auth_header.startswith("Bearer "): + return JSONResponse( + status_code=HTTP_401_UNAUTHORIZED, + content={"detail": "Missing or invalid Authorization header"}, + headers={"WWW-Authenticate": 'Bearer realm="GitHub"'}, + ) + + token = auth_header[7:] # Remove "Bearer " prefix + + # Verify token + try: + user_data = await verify_github_token(token) + # Attach user data to request state + request.state.user = user_data + except HTTPException as e: + return JSONResponse( + status_code=e.status_code, + content={"detail": e.detail}, + headers=e.headers or {}, + ) + + # Continue with request + return await call_next(request) + + return app + + +if __name__ == "__main__": + import sys + + import uvicorn + + if len(sys.argv) < 2: + print("Usage: python oauth2_github.py ") + sys.exit(1) + + db_path = Path(sys.argv[1]) + app = create_secure_a2a_app(db_path) + + uvicorn.run(app, host="127.0.0.1", port=8000) diff --git a/pyproject.toml b/pyproject.toml index b117b8d4..15caf15c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,7 +50,7 @@ requires = ["hatchling"] build-backend = "hatchling.build" [tool.hatch.build] -exclude = ["/docs", "/tests", "/.github"] +exclude = ["/docs", "/examples", "/tests", "/.github"] [tool.hatch.build.targets.wheel] packages = ["src/haiku"] diff --git a/src/haiku/rag/a2a/__init__.py b/src/haiku/rag/a2a/__init__.py index 06a5467c..0467deac 100644 --- a/src/haiku/rag/a2a/__init__.py +++ b/src/haiku/rag/a2a/__init__.py @@ -42,11 +42,17 @@ __all__ = [ ] -def create_a2a_app(db_path: Path): +def create_a2a_app( + db_path: Path, + security_schemes: dict | None = None, + security: list[dict[str, list[str]]] | None = None, +): """Create an A2A app for the conversational QA agent. Args: db_path: Path to the LanceDB database + security_schemes: Optional security scheme definitions for the AgentCard + security: Optional security requirements for the AgentCard Returns: A FastA2A ASGI application @@ -133,7 +139,7 @@ def create_a2a_app(db_path: Path): async with worker.run(): yield - return FastA2A( + app = FastA2A( storage=storage, broker=broker, name="haiku-rag", @@ -141,3 +147,44 @@ def create_a2a_app(db_path: Path): skills=get_agent_skills(), lifespan=lifespan, ) + + # Add security configuration if provided + if security_schemes or security: + # Monkey-patch the agent card endpoint to include security + async def _agent_card_endpoint_with_security(request): + from fasta2a.schema import AgentCapabilities, AgentCard, agent_card_ta + from starlette.responses import Response + + if app._agent_card_json_schema is None: + agent_card = AgentCard( + name=app.name, + description=app.description + or "An AI agent exposed as an A2A agent.", + url=app.url, + version=app.version, + protocol_version="0.3.0", + skills=app.skills, + default_input_modes=app.default_input_modes, + default_output_modes=app.default_output_modes, + capabilities=AgentCapabilities( + streaming=False, + push_notifications=False, + state_transition_history=False, + ), + ) + if app.provider is not None: + agent_card["provider"] = app.provider + if security_schemes: + agent_card["security_schemes"] = security_schemes + if security: + agent_card["security"] = security + app._agent_card_json_schema = agent_card_ta.dump_json( + agent_card, by_alias=True + ) + return Response( + content=app._agent_card_json_schema, media_type="application/json" + ) + + app._agent_card_endpoint = _agent_card_endpoint_with_security + + return app