Know your user:
Identity-aware MCP servers
with Cosmos DB

Before MCP

Every integration with an AI agent had to be done individually:

Model Context Protocol

An open protocol that defines how AI applications get context from external tools and data sources.

MCP architecture

MCP server components

This identity-aware MCP server combines four pieces:

MCP server architecture

Cosmos DB data model

Azure Cosmos DB is a NoSQL database that stores JSON documents in containers.

{
  "id": "a1b2c3d4-...",
  "user_id": "00000000-0000-...",
  "date": "2026-03-31",
  "amount": 20.00,
  "category": "food",
  "description": "Avocado toast",
  "payment_method": "visa"
}
            

Setting up the Cosmos DB client

Use azure-cosmos for easy Python integration with Cosmos DB:

from azure.cosmos.aio import CosmosClient
from azure.identity.aio import ManagedIdentityCredential

azure_credential = ManagedIdentityCredential(
  client_id=os.environ["AZURE_CLIENT_ID"])

cosmos_client = CosmosClient(
  url=f"https://{os.environ['AZURE_COSMOSDB_ACCOUNT']}.documents.azure.com/",
  credential=azure_credential,
)
cosmos_db = cosmos_client.get_database_client(
    os.environ["AZURE_COSMOSDB_DATABASE"])
cosmos_container = cosmos_db.get_container_client(
    os.environ["AZURE_COSMOSDB_USER_CONTAINER"])
    

🔒 Cosmos DB supports keyless access via managed identity.

Tool: add_user_expense

@mcp.tool
async def add_user_expense(
    date: Annotated[date, "Date of the expense in YYYY-MM-DD format"],
    amount: Annotated[float, "Positive numeric amount of the expense"],
    category: Annotated[Category, "Category label"],
    description: Annotated[str, "Human-readable description of the expense"],
    payment_method: Annotated[PaymentMethod, "Payment method used"],
    ctx: Context,
):
    """Add a new expense to Cosmos DB."""
    user_id = await ctx.get_state("user_id")
    if not user_id:
        return "Error: Authentication required"
    expense_item = {
        "id": str(uuid.uuid4()),
        "user_id": user_id,
        "date": date.isoformat(),
        "amount": amount,
        "category": category.value,
        "description": description,
        "payment_method": payment_method.value,
    }
    await cosmos_container.create_item(body=expense_item)
    return f"Successfully added expense: ${amount} for {description}"
    

Tool: get_user_expenses

@mcp.tool
async def get_user_expenses(ctx: Context):
    """Get the authenticated user's expense data from Cosmos DB."""
    user_id = await ctx.get_state("user_id")
    if not user_id:
        return "Error: Authentication required"
    
    query = "SELECT * FROM c WHERE c.user_id = @uid ORDER BY c.date DESC"
    parameters = [{"name": "@uid", "value": user_id}]
    expenses_data = []
    async for item in cosmos_container.query_items(
        query=query, parameters=parameters, partition_key=user_id
    ):
        expenses_data.append(item)

    return json.dumps([{
        "date": e.get("date"), "amount": e.get("amount"),
        "category": e.get("category"), "description": e.get("description"),
        "payment_method": e.get("payment_method"),
    } for e in expenses_data], indent=2)
    

Passing partition_key=user_id makes this a single-partition query.

OAuth-based access flow

An MCP client can make requests to an MCP server on behalf of a signed-in user:

OAuth 2.1 roles in MCP

MCP authentication builds on the standard OAuth 2.1 roles and token flow.

Choosing a client registration path

For this demo: Microsoft Entra does not support CIMD or DCR, so we use pre-registration with VS Code's known client ID. (Alternative approach: put a DCR proxy in front of Entra.)

Pre-registered OAuth flow for MCP

Setting up OAuth with Entra in FastMCP

FastMCP sends clients directly to Entra, then validates the returned JWTs with Entra's public signing keys.

from fastmcp import FastMCP
from fastmcp.server.auth import RemoteAuthProvider
from fastmcp.server.auth.providers.azure import AzureJWTVerifier

verifier = AzureJWTVerifier(
    client_id=entra_client_id,
    tenant_id=os.environ["AZURE_TENANT_ID"],
    required_scopes=["user_impersonation"],
)
auth = RemoteAuthProvider(
    token_verifier=verifier,
    authorization_servers=[f"https://login.microsoftonline.com/{os.environ['AZURE_TENANT_ID']}/v2.0"],
    base_url=base_url,
)
mcp = FastMCP("Expenses Tracker", auth=auth)
    

Extracting user identity from the token

We pull user_id from the OAuth token claims so that each tool can access the user id.

from fastmcp.server.dependencies import get_access_token
from fastmcp.server.middleware import Middleware, MiddlewareContext

class UserAuthMiddleware(Middleware):
    def _get_user_id(self):
        token = get_access_token()
        if not (token and hasattr(token, "claims")):
            return None
        return token.claims.get("oid")  # Entra Object ID

    async def on_call_tool(self, context: MiddlewareContext, call_next):
        user_id = self._get_user_id()
        if context.fastmcp_context is not None:
            await context.fastmcp_context.set_state("user_id", user_id)
        return await call_next(context)

mcp = FastMCP("Expenses Tracker", auth=auth,
              middleware=[UserAuthMiddleware()])
    

Why role-based access?

Some tools should only be available to certain users:

• Regular users: add and view their own expenses
• Admins: view aggregate stats across all users

On-Behalf-Of (OBO) flow for Graph API

Checking group membership via Graph API

The server uses MSAL for token exchange, then Graph API to check group membership:

confidential_client = ConfidentialClientApplication(
    client_id=entra_client_id,
    client_credential=client_credential,
    authority=f"https://login.microsoftonline.com/{os.environ['AZURE_TENANT_ID']}",
    token_cache=TokenCache(),
)
graph_token = confidential_client.acquire_token_on_behalf_of(
    user_assertion=ctx.token.token,
    scopes=["https://graph.microsoft.com/.default"],
)

client = httpx.AsyncClient()
response = await client.get(
    "https://graph.microsoft.com/v1.0/me/transitiveMemberOf"
    "/microsoft.graph.group"
    f"?$filter=id eq '{group_id}'&$count=true",
    headers={
        "Authorization": f"Bearer {graph_token['access_token']}",
        "ConsistencyLevel": "eventual",
    })
is_admin = response.json().get("@odata.count", 0) > 0

Enforcing role-based access in tools

FastMCP runs this auth check at both tools listing time
and tool invocation time:

async def require_admin_group(ctx: AuthContext) -> bool:
    admin_group_id = os.environ["ENTRA_ADMIN_GROUP_ID"]
    graph_token = confidential_client.acquire_token_on_behalf_of(
        user_assertion=ctx.token.token,
        scopes=["https://graph.microsoft.com/.default"],
    )
    return await check_user_in_group(
        graph_token["access_token"], admin_group_id)

@mcp.tool(auth=require_admin_group)
async def get_expense_stats(ctx: Context):
    """Get expense statistics. Only accessible to admins."""
    ...
    

Non-admin users won't even see get_expense_stats in the tools list.

Agent skills for Cosmos DB

Use the cosmosdb-best-practices agent skill to:

• Review partition key design
• Optimize query performance
• Validate data model choices
• Check SDK usage patterns

"Review my Cosmos DB data model and suggest improvements"

Install the skills:

npx skills add AzureCosmosDB/cosmosdb-agent-kit

github.com/AzureCosmosDB/cosmosdb-agent-kit

Thank you!

Slides: aka.ms/cosmos-identity-mcp/slides

Code: aka.ms/cosmos-identity-mcp/code

Learn more from Python + MCP livestream series: aka.ms/pythonmcp/rewatch

Questions? Find me online: