How we built an enterprise-grade multi-agent system that preserves user identity across AI agents and Databricks

Introduction

When building AI-powered applications for the enterprise, a common challenge emerges: how do you maintain user identity and access controls when an AI agent queries backend services on behalf of a user?

In many implementations, AI agents authenticate to backend systems using a shared service account or with PAT (Personal Access Token) tokens, effectively bypassing row-level security (RLS), column masking, and other data governance policies that organizations carefully configure. This creates a security gap where users can potentially access data they shouldn’t see, simply by asking an AI agent.

In this post, I’ll walk through how we solved this challenge for a current enterprise customer by implementing Microsoft Entra ID On-Behalf-Of (OBO) secure flow in a custom multi-agent LangGraph solution, enabling our Databricks Genie agent to query data and the data agent designed to modify or update delta tables, to do so as the authenticated user, while preserving all RBAC policies.

The Architecture

Our system is built on several key components:

Chainlit: Python-based web interface for LLM-driven conversational applications, integrated with OAuth 2.0–based authentication. Customizing the framework to satisfy customer UI requirements eliminated the need to develop and maintain a bespoke React front end. It fulfilled the majority of requirements while reducing maintenance overhead.
LangGraph: Opensource Multi-agent orchestration framework.
Azure Databricks Genie: Natural language to SQL agent.
Azure Cosmos DB: Long-term memory and checkpoint storage.
Microsoft Entra ID: Identity provider with OBO support.

This shows:

Genie: Read-only natural language queries, per-user OBO
Task Agent: Handles sensitive operations (SQL modifications, etc.) with HITL approval + OBO
Memory: Shared agent, no per-user auth needed

The Problem with Chainlit OAuth Provider

Chainlit was integrated with Microsoft Entra ID for OAuth authentication; however, the default implementation assumes Microsoft Graph scopes, requiring extension to support custom resource scopes. This means:

The access token you receive is scoped for Microsoft Graph API
You can’t use it for OBO flow to downstream services like Databricks
The token’s audience is graph.microsoft.com, not your application
For OBO to work, you need an access token where:

The audience is your application’s client ID
The scope includes your custom API permission (e.g., api://{client_id}/access_as_user)

Solution: Custom Entra ID OBO Provider

We created a custom OAuth provider that replaces Chainlit’s built-in one.

Key insight: By requesting api://{client_id}/access_as_user as the scope, the returned access token has the correct audience for OBO exchange.

Since we can’t call Graph API with this token (wrong audience), we extract user information from the ID token claims instead.

The OBO Token Exchange

Once we have the user’s access token (with correct audience), we exchange it for a Databricks-scoped token using MSAL.

The resulting token:

Has audience = Databricks resource ID
Contains the user’s identity (UPN, OID)
Can be used with Databricks SDK/API
Respects all Unity Catalog permissions configured for that user

Per-User Agent Creation

A critical design decision: never cache user-specific agents globally. Each user needs their own Genie agent instance.

Using the OBO Token with Databricks Genie

The key integration point is passing the OBO-acquired token to the Databricks SDK’s WorkspaceClient as indicated in the above screenshot, which the Genie agent uses internally for all API calls as shown in the following image.

Initialize Genie Agent with User’s Access Token:

Wire It Into LangGraph:

The user_access_token flows from Chainlit’s OAuth callback → session config → LangGraph config → agent creation, ensuring every Genie query runs with the authenticated user’s permissions.

Human-in-the-Loop for Destructive SQL Operations

While Databricks Genie handles natural language queries (read-only), our system also supports custom SQL execution for data modifications. Since these operations can DELETE or UPDATE data, we implement human-in-the-loop approval using LangGraph’s interrupt feature.

The OBO token ensures that even when executing user-authored SQL, the query runs with the user’s permissions: they can only modify data they’re authorized to change.

The destructive operation detector uses LLM-based intent analysis

Entra ID App Registration Requirements

Your Entra ID app registration needs:

1. API Permissions: Azure Databricks → user_impersonation (admin consent required)
2. Expose an API: Scope access_as_user on URI api://{client-id}
3. Redirect URI: {your-app-url}/auth/oauth/azure-ad/callback

Lessons Learned

1. Token audience matters: OBO fails if your initial token has the wrong audience
2. Don’t cache user-specific clients: breaks user isolation
3. ID tokens contain user info: use claims when you can’t call Graph API
4. HITL for destructive ops: even with RBAC, require explicit user confirmation

Conclusion

By implementing Entra ID OBO flow in our multi-agent system, we achieved:

1. User identity preservation across AI agents
2. RBAC enforcement at the Databricks/Unity Catalog level
3. Audit trail showing actual user making queries
4. Zero-trust architecture: the AI agent never has more access than the user
5. Human-in-the-loop for destructive SQL operations
   

This design applies beyond Databricks, any Azure service supporting OAuth 2.0 can be accessed via OBO. It also forms part of the AI governance foundation for this enterprise custom multiagent AI solution and ensures compliance with Microsoft SFI (Secure Future Initiative) principles.