danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Document securing a2a

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-10-13 15:00:47 +03:00
parent 6ea7d73eed
commit fce68a6673
No known key found for this signature in database
6 changed files with 636 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
docs/a2a.md
View file

@ -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

130
examples/a2a-security/apikey_example.py Normal file
View file

@ -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 <path-to-database.lancedb>")
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)

222
examples/a2a-security/oauth2_example.py Normal file
View file

@ -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 <path-to-database.lancedb>")
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)

193
examples/a2a-security/oauth2_github.py Normal file
View file

@ -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 <path-to-database.lancedb>")
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)

2
pyproject.toml
View file

@ -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"]

51
src/haiku/rag/a2a/__init__.py
View file

@ -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