Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 47 additions & 0 deletions python/packages/ag-ui/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
9 changes: 8 additions & 1 deletion python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.

Expand All @@ -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(
Expand All @@ -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.

Expand Down
67 changes: 64 additions & 3 deletions python/packages/ag-ui/getting_started/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand All @@ -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:
Expand Down Expand Up @@ -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
Expand Down
90 changes: 88 additions & 2 deletions python/packages/ag-ui/tests/test_endpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"
Loading