diff --git a/python/packages/ag-ui/README.md b/python/packages/ag-ui/README.md index 1e3d6b567f..ec5602cef9 100644 --- a/python/packages/ag-ui/README.md +++ b/python/packages/ag-ui/README.md @@ -82,6 +82,53 @@ This integration supports all 7 AG-UI features: 6. **Shared State**: Bidirectional state sync between client and server 7. **Predictive State Updates**: Stream tool arguments as optimistic state updates during execution +## Security: Authentication & Authorization + +The AG-UI endpoint does not enforce authentication by default. **For production deployments, you should add authentication** using FastAPI's dependency injection system via the `dependencies` parameter. + +### API Key Authentication Example + +```python +import os +from fastapi import Depends, FastAPI, HTTPException, Security +from fastapi.security import APIKeyHeader +from agent_framework import ChatAgent +from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint + +# Configure API key authentication +API_KEY_HEADER = APIKeyHeader(name="X-API-Key", auto_error=False) +EXPECTED_API_KEY = os.environ.get("AG_UI_API_KEY") + +async def verify_api_key(api_key: str | None = Security(API_KEY_HEADER)) -> None: + """Verify the API key provided in the request header.""" + if not api_key or api_key != EXPECTED_API_KEY: + raise HTTPException(status_code=401, detail="Invalid or missing API key") + +# Create agent and app +agent = ChatAgent(name="my_agent", instructions="...", chat_client=...) +app = FastAPI() + +# Register endpoint WITH authentication +add_agent_framework_fastapi_endpoint( + app, + agent, + "/", + dependencies=[Depends(verify_api_key)], # Authentication enforced here +) +``` + +### Other Authentication Options + +The `dependencies` parameter accepts any FastAPI dependency, enabling integration with: + +- **OAuth 2.0 / OpenID Connect** - Use `fastapi.security.OAuth2PasswordBearer` +- **JWT Tokens** - Validate tokens with libraries like `python-jose` +- **Azure AD / Entra ID** - Use `azure-identity` for Microsoft identity platform +- **Rate Limiting** - Add request throttling dependencies +- **Custom Authentication** - Implement your organization's auth requirements + +For a complete authentication example, see [getting_started/server.py](getting_started/server.py). + ## Architecture The package uses a clean, orchestrator-based architecture: diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py b/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py index dda045b3fb..7948d4f935 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py @@ -4,11 +4,13 @@ import copy import logging +from collections.abc import Sequence from typing import Any from ag_ui.encoder import EventEncoder from agent_framework import AgentProtocol from fastapi import FastAPI +from fastapi.params import Depends from fastapi.responses import StreamingResponse from ._agent import AgentFrameworkAgent @@ -26,6 +28,7 @@ def add_agent_framework_fastapi_endpoint( allow_origins: list[str] | None = None, default_state: dict[str, Any] | None = None, tags: list[str] | None = None, + dependencies: Sequence[Depends] | None = None, ) -> None: """Add an AG-UI endpoint to a FastAPI app. @@ -39,6 +42,10 @@ def add_agent_framework_fastapi_endpoint( allow_origins: CORS origins (not yet implemented) default_state: Optional initial state to seed when the client does not provide state keys tags: OpenAPI tags for endpoint categorization (defaults to ["AG-UI"]) + dependencies: Optional FastAPI dependencies for authentication/authorization. + These dependencies run before the endpoint handler. Use this to add + authentication checks, rate limiting, or other middleware-like behavior. + Example: `dependencies=[Depends(verify_api_key)]` """ if isinstance(agent, AgentProtocol): wrapped_agent = AgentFrameworkAgent( @@ -49,7 +56,7 @@ def add_agent_framework_fastapi_endpoint( else: wrapped_agent = agent - @app.post(path, tags=tags or ["AG-UI"]) # type: ignore[arg-type] + @app.post(path, tags=tags or ["AG-UI"], dependencies=dependencies) # type: ignore[arg-type] async def agent_endpoint(request_body: AGUIRequest): # type: ignore[misc] """Handle AG-UI agent requests. diff --git a/python/packages/ag-ui/getting_started/server.py b/python/packages/ag-ui/getting_started/server.py index e4ed669516..c8889126e9 100644 --- a/python/packages/ag-ui/getting_started/server.py +++ b/python/packages/ag-ui/getting_started/server.py @@ -9,7 +9,8 @@ from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint from agent_framework.azure import AzureOpenAIChatClient from dotenv import load_dotenv -from fastapi import FastAPI +from fastapi import Depends, FastAPI, HTTPException, Security +from fastapi.security import APIKeyHeader load_dotenv() @@ -31,6 +32,60 @@ raise ValueError("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME environment variable is required") +# ============================================================================ +# AUTHENTICATION EXAMPLE +# ============================================================================ +# This demonstrates how to secure the AG-UI endpoint with API key authentication. +# In production, you should use a more robust authentication mechanism such as: +# - OAuth 2.0 / OpenID Connect +# - JWT tokens with proper validation +# - Azure AD / Entra ID integration +# - Your organization's identity provider +# +# The API key should be stored securely (e.g., Azure Key Vault, environment variables) +# and rotated regularly. +# ============================================================================ + +# API key header configuration +API_KEY_HEADER = APIKeyHeader(name="X-API-Key", auto_error=False) + +# Get the expected API key from environment variable +# In production, use a secrets manager like Azure Key Vault +EXPECTED_API_KEY = os.environ.get("AG_UI_API_KEY") + + +async def verify_api_key(api_key: str | None = Security(API_KEY_HEADER)) -> None: + """Verify the API key provided in the request header. + + Args: + api_key: The API key from the X-API-Key header + + Raises: + HTTPException: If the API key is missing or invalid + """ + if not EXPECTED_API_KEY: + # If no API key is configured, log a warning but allow the request + # This maintains backward compatibility but warns about the security risk + logger.warning( + "AG_UI_API_KEY environment variable not set. " + "The endpoint is accessible without authentication. " + "Set AG_UI_API_KEY to enable API key authentication." + ) + return + + if not api_key: + raise HTTPException( + status_code=401, + detail="Missing API key. Provide X-API-Key header.", + ) + + if api_key != EXPECTED_API_KEY: + raise HTTPException( + status_code=403, + detail="Invalid API key.", + ) + + # Server-side tool (executes on server) @ai_function(description="Get the time zone for a location.") def get_time_zone(location: str) -> str: @@ -72,8 +127,14 @@ def get_time_zone(location: str) -> str: # Create FastAPI app app = FastAPI(title="AG-UI Server") -# Register the AG-UI endpoint -add_agent_framework_fastapi_endpoint(app, agent, "/") +# Register the AG-UI endpoint with authentication +# The dependencies parameter accepts FastAPI Depends() objects that run before the handler +add_agent_framework_fastapi_endpoint( + app, + agent, + "/", + dependencies=[Depends(verify_api_key)], +) if __name__ == "__main__": import uvicorn diff --git a/python/packages/ag-ui/tests/test_endpoint.py b/python/packages/ag-ui/tests/test_endpoint.py index 02b23a544b..59cb884c5c 100644 --- a/python/packages/ag-ui/tests/test_endpoint.py +++ b/python/packages/ag-ui/tests/test_endpoint.py @@ -7,11 +7,12 @@ from pathlib import Path from agent_framework import ChatAgent, ChatResponseUpdate, TextContent -from fastapi import FastAPI +from fastapi import FastAPI, Header, HTTPException +from fastapi.params import Depends from fastapi.testclient import TestClient +from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint from agent_framework_ag_ui._agent import AgentFrameworkAgent -from agent_framework_ag_ui._endpoint import add_agent_framework_fastapi_endpoint sys.path.insert(0, str(Path(__file__).parent)) from utils_test_ag_ui import StreamingChatClientStub, stream_from_updates @@ -380,3 +381,88 @@ async def test_endpoint_internal_error_handling(): assert response.status_code == 200 assert response.json() == {"error": "An internal error has occurred."} + + +async def test_endpoint_with_dependencies_blocks_unauthorized(): + """Test that endpoint blocks requests when authentication dependency fails.""" + app = FastAPI() + agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client()) + + async def require_api_key(x_api_key: str | None = Header(None)): + if x_api_key != "secret-key": + raise HTTPException(status_code=401, detail="Unauthorized") + + add_agent_framework_fastapi_endpoint(app, agent, path="/protected", dependencies=[Depends(require_api_key)]) + + client = TestClient(app) + + # Request without API key should be rejected + response = client.post("/protected", json={"messages": [{"role": "user", "content": "Hello"}]}) + assert response.status_code == 401 + assert response.json()["detail"] == "Unauthorized" + + +async def test_endpoint_with_dependencies_allows_authorized(): + """Test that endpoint allows requests when authentication dependency passes.""" + app = FastAPI() + agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client()) + + async def require_api_key(x_api_key: str | None = Header(None)): + if x_api_key != "secret-key": + raise HTTPException(status_code=401, detail="Unauthorized") + + add_agent_framework_fastapi_endpoint(app, agent, path="/protected", dependencies=[Depends(require_api_key)]) + + client = TestClient(app) + + # Request with valid API key should succeed + response = client.post( + "/protected", + json={"messages": [{"role": "user", "content": "Hello"}]}, + headers={"x-api-key": "secret-key"}, + ) + assert response.status_code == 200 + assert response.headers["content-type"] == "text/event-stream; charset=utf-8" + + +async def test_endpoint_with_multiple_dependencies(): + """Test that endpoint supports multiple dependencies.""" + app = FastAPI() + agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client()) + + execution_order: list[str] = [] + + async def first_dependency(): + execution_order.append("first") + + async def second_dependency(): + execution_order.append("second") + + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/multi-deps", + dependencies=[Depends(first_dependency), Depends(second_dependency)], + ) + + client = TestClient(app) + response = client.post("/multi-deps", json={"messages": [{"role": "user", "content": "Hello"}]}) + + assert response.status_code == 200 + assert "first" in execution_order + assert "second" in execution_order + + +async def test_endpoint_without_dependencies_is_accessible(): + """Test that endpoint without dependencies remains accessible (backward compatibility).""" + app = FastAPI() + agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client()) + + # No dependencies parameter - should be accessible without auth + add_agent_framework_fastapi_endpoint(app, agent, path="/open") + + client = TestClient(app) + response = client.post("/open", json={"messages": [{"role": "user", "content": "Hello"}]}) + + assert response.status_code == 200 + assert response.headers["content-type"] == "text/event-stream; charset=utf-8"