AI Orchestration & Tool Governance
01 / 22

AI Orchestration &
Tool Governance

Agent Bricks, MCP Servers, Connections, Federation, and Observability
Architecture

The Orchestration Stack

Every layer has a governance boundary. Identity flows top to bottom.

UserHuman or External Application
↓
IdentityIdP + Token Path (OBO / M2M / Federation)
↓
ApplicationDatabricks App / External App
↓
OrchestratorAgent Bricks Supervisor (Model Serving)
↓
ToolsGenie · Vector Search · UC Functions · MCP Servers
↓
DataTables · Lakebase
↓
External Servicesvia UC Connections (Jira, GitHub, Slack, etc.)
Agent Bricks

Multi-Agent Orchestration

A supervisor coordinates specialized sub-agents, each with its own tools and identity context.

Supervisor Model

  • Coordinates up to 20 sub-agents
  • Each agent: Knowledge Assistant, Genie Space, or Custom Tool
  • Deployed as a Model Serving endpoint
  • Supervisor routes intent to the right agent

Identity Modes

OBO Supervisor auto-forwards the user's token to all sub-agents. current_user() = human email everywhere.

M2M Supervisor uses its own SP credentials. Sub-agents see the SP identity.

Agent Bricks

Identity Flow

OBO token propagation from browser to the deepest tool call.

UserBrowser
→
Databricks AppX-Forwarded-Access-Token
→
Agent BricksModel Serving
→
ModelServingUserCredentials()Extracts user identity
→
Sub-agentsUser context forwarded
→
Resourcescurrent_user() = human
current_user() returns the human email at every resource — from the supervisor down to the table query.
Agent Bricks

Gotchas

Runtime Limits

  • Max 20 sub-agents per supervisor
  • ModelServingUserCredentials() only works inside Model Serving — do not use in Apps code
  • Supervisor endpoint requires model-serving scope

Permissions & Token Propagation

  • Each sub-agent’s tools must have their own UC grants
  • Genie sub-agents need both genie and dashboards.genie scopes
  • Missing either scope → 403: "required scopes: genie"
  • Grant UC permissions to the SP, not the human user

Single App: X-Forwarded-Access-Token has the user’s token — pass it as Bearer to Agent Bricks. Works fine. App → App (two-proxy): User token is lost. Use X-Forwarded-Email (survives proxy hops, cannot be forged) + M2M SQL.

MCP Servers

MCP on Databricks — Three Types

Managed MCP

Databricks-hosted, pre-built connectors (e.g., Unity Catalog). No code needed.

Zero Config
  • Unity Catalog tools
  • Pre-configured auth
  • Managed lifecycle

Custom MCP

Your code, deployed as Databricks App. Full control over tools and auth.

Full Control
  • Streamlit or FastAPI
  • OBO / M2M / mixed
  • Custom tool logic

External MCP

Third-party MCP servers accessed via UC HTTP Connections.

Governed
  • Jira, GitHub, Slack
  • USE CONNECTION grants
  • Credential-free code
MCP Servers

Custom MCP Server Pattern

Deployment

  • Deploy as Databricks App (Streamlit/FastAPI)
  • App proxy injects X-Forwarded-Email (trustworthy, signed by proxy)
  • Tool registration via MCP protocol (tools/list, tools/call)

Auth Options

OBORead X-Forwarded-Access-Token, forward to downstream
M2MWorkspaceClient() no-args, SP identity
MixedM2M for data, OBO for user context

Example: Tool Registration

# tools/list response
{
  "tools": [
    {
      "name": "query_genie",
      "description": "Ask a question",
      "inputSchema": {
        "type": "object",
        "properties": {
          "question": {"type": "string"}
        }
      }
    }
  ]
}

# tools/call invocation
POST /tools/call
{
  "name": "query_genie",
  "arguments": {"question": "..."}
}
MCP Servers

External MCP via UC HTTP Connections

Access third-party MCP servers through governed connection objects.

How It Works

  • MCP servers outside Databricks (Jira, GitHub, Slack)
  • Connected through UC HTTP Connections
  • GRANT USE CONNECTION controls access
  • Platform injects auth headers automatically

IDE Integration

  • uc-mcp-proxy for Cursor / Windsurf
  • Managed OAuth providers: Glean, GitHub, Atlassian

SDK Patterns

OpenAI Agents SDK + MCP
Native MCP client integration. Each tool backed by a UC connection.

LangGraph + MCP
MCP tools as LangGraph nodes. Graph orchestration with governed tool access.

uc-mcp-proxy
Local proxy bridges IDE MCP clients to Databricks-governed connections.

UC HTTP Connections

Governed External Access

UC Connections are the governance boundary for every external API call.

  • Every external API call goes through a UC Connection
  • Credentials stored in UC, never in code
  • GRANT USE CONNECTION controls who can use each connection
  • REVOKE removes access immediately (runtime enforcement)
  • Platform injects auth headers via http_request() / SDK
Agent / Tool CodeNo credentials in source
↓
UC ConnectionCredentials + GRANT/REVOKE
↓
External APIJira / GitHub / Slack / etc.
UC HTTP Connections

Four Auth Methods

Bearer Token

Static token stored in UC. Simple but shared identity.

Static

OAuth M2M

Client credentials flow. App-level identity at external service.

App Identity

OAuth U2M Shared

One user’s OAuth token shared across callers.

Shared

OAuth U2M Per User

Each user’s own OAuth token. True per-user identity on both sides.

Per-User
Per-user OAuth is the gold standard: the external service sees the actual human identity, not a shared service account. Use bearer tokens only for services that don’t support OAuth.
UC HTTP Connections

GRANT/REVOKE as Runtime Control

-- Enable agent to use Jira
GRANT USE CONNECTION ON CONNECTION jira_mcp TO `agent-sp-uuid`;

-- Immediately disable (no redeploy needed)
REVOKE USE CONNECTION ON CONNECTION jira_mcp FROM `agent-sp-uuid`;

This is the kill switch. Revoke connection access and the agent loses the tool immediately. No code change, no redeploy.

Before Revoke

Agent
→
Jira Connection
→
Jira API

✓ Tool call succeeds

After Revoke

Agent
→
Jira Connection
→
Jira API

✗ Permission denied at UC layer

Reference Architecture

AppKit Recipe — 4 Governance Patterns in One App

One AppKit app proves all auth patterns side by side. External service is provider-agnostic.

TabPatternIdentityWhat It Proves
SQL Analytics OBO vs M2M (.obo.sql / .sql) User or SP Row filters fire per-user with OBO; shared cache with SP
Genie Chat OBO token → Genie API User current_user() = human, per-user data access
External Service UC Connection (U2M Per User) User on both sides Two-sided identity: Databricks & external service see user
Connection Gov GRANT / REVOKE USE CONNECTION Admin Runtime kill switch — revoke = tool gone immediately

External Service tab is a template. Swap connection_name + path for Salesforce, Jira, GitHub, Slack, ServiceNow — same pattern, same code.

Reference Architecture

AppKit Recipe — How It Wires Together

Request Flow

Browser → Apps Proxy
  (injects OBO token)
→ AppKit Server
  ├── analytics() → SQL Warehouse
  │   .obo.sql = user, .sql = SP
  ├── genie() → Genie API
  │   OBO auto-forwarded
  ├── external-service → UC Connection
  │   Any OAuth provider
  └── connections → GRANT/REVOKE
      Runtime governance

Critical Post-Deploy Steps

  1. Re-grant UC permissions to new SP
    Every deploy = new SP = zero grants
  2. PATCH genie scope via account API
    Not in UI — must use REST API
  3. Set connection ownership to app SP
    Required for GRANT/REVOKE
  4. Verify scopes in incognito
    Cached tokens miss new scopes

Deploy raw .mjs files — esbuild bundles break AppKit’s runtime import(). Include package.json, exclude .npmrc.

Token Federation

Token Federation

Exchange a trusted IdP JWT for a Databricks OAuth token. No Databricks secrets needed.

How It Works

Your IdPIssues JWT with user claims
↓
POST /oidc/v1/tokengrant_type = token-exchange
↓
Databricks OAuth TokenUser or SP identity

Key Properties

  • One endpoint: POST /oidc/v1/token
  • No Databricks secrets needed in your infrastructure
  • Works for user (U2M) and SP (M2M) identities
  • Standards-based: OAuth token exchange
  • Exchange trusted JWT → get Databricks token
Token Federation

Two Federation Types

Account Token Federation

  • Users & SPs from your IdP
  • SCIM required (sync users first)
  • Interactive + M2M flows
  • Refresh tokens for U2M

Workload Identity Federation (WIF)

  • CI/CD workloads → SP
  • No SCIM required
  • Per-SP policy
  • Secretless, M2M only
FeatureAccount FederationWIF
Applies toUsers & SPsWorkload → SP
SCIM requiredYesNo
Issuer limit5Unlimited
SecretsNoneNone
Refresh tokensYes (U2M)No
Token Federation

Token TTL & IdP Support

Token Lifetime

Databricks copies the IdP’s exp claim verbatim.

  • No override, no extension
  • Short-lived IdP tokens = short-lived Databricks tokens
  • Plan token refresh in your client code

Supported IdPs

Entra IDAccount + WIF
OktaAccount + WIF
PingOneAccount + WIF
Google Cloud IdentityAccount + WIF
Auth0Account + WIF
AWS IAMWIF only
GitHub Actions OIDCWIF only

JDBC/ODBC driver support: Work in progress. Not all drivers support federation yet. Check driver release notes before depending on this path.

AI Gateway

AI Gateway v2

Centralized model traffic management: rate limiting, cost tracking, guardrails, fallback routing.

Rate Limiting

Per endpoint, per user, per model. Sliding window with configurable QPM/TPM.

Usage Tracking

Token counts, latency, cost allocation per team/project/endpoint.

Guardrails

Content filtering, PII detection, topic blocking on input and output.

Fallback Routing

Primary model fails → automatically route to secondary. No client code change.

A/B Testing

Traffic splitting across model versions with built-in metrics comparison.

AI Gateway

When to Use What

ScenarioAI GatewayUC + ScopesBoth
Rate limit model access✓
Track token usage / cost✓
Content guardrails✓
Data access control✓
External tool governance✓
Production agent deployment✓

AI Gateway manages MODEL traffic. UC manages DATA and TOOL access. Both needed for production.

AI Gateway

Gateway Patterns

1. Single Model + Guardrails

Client Request
↓
AI GatewayInput filter
↓
Model Endpoint
↓
AI GatewayOutput filter
↓
Safe Response

2. Fallback Routing

Client Request
↓
AI Gateway
↓
Primary Model
↓ on failure
Secondary Model
↓
Response

3. A/B Testing

Client Request
↓
AI GatewayTraffic split
80%
Model A
20%
Model B
↓
Metrics Comparison
Observability

Observability Stack

Three pillars. Join on request_id for end-to-end agent lineage.

system.access.audit

Every API call logged. Join on SP UUID to trace agent actions.

  • Permission changes
  • Connection usage
  • Function executions
  • Token exchanges

MLflow Tracing

Every agent step, tool call, LLM invocation. Latency, tokens, errors.

  • Supervisor → sub-agent spans
  • Tool call duration
  • Token counts per step
  • Error propagation

App Logs (OTel)

Structured application logs from Databricks Apps. Auto-configured OTel env vars.

  • 10 env vars auto-set
  • OpenTelemetry export
  • Request correlation
  • App Insights (Beta)
Join system.access.audit + MLflow traces on request_id for end-to-end agent lineage.
Observability

Security Monitoring Queries

Permission Changes

-- Who changed permissions?
SELECT *
FROM system.access.audit
WHERE action_name IN (
  'grantPermission',
  'revokePermission'
)
AND event_date >=
  current_date() - 7;

OBO Token Usage

-- OBO token usage
SELECT
  request_params.user_email,
  count(*)
FROM system.access.audit
WHERE source_ip_address
  LIKE '10.%'
AND user_identity.email
  LIKE '%ServicePrincipal%'
GROUP BY 1;

Agent Tool Invocations

-- Agent tool invocations
SELECT *
FROM system.access.audit
WHERE action_name =
  'executeFunction'
AND user_identity.email =
  '<agent-sp-uuid>';
Observability

OTel in Databricks Apps

Auto-Configured Environment

OTEL_EXPORTER_*Export endpoint + protocol
OTEL_SERVICE_NAMEApp name auto-set
OTEL_RESOURCE_*Resource attributes
DATABRICKS_APP_PORTFramework port

10 env vars total, auto-configured at app startup. No manual OTel setup needed.

Runtime & Limits

OSUbuntu 22.04
Python3.11
Node22.16
P50 latency~50ms
P95 latency~200ms

App Insights (Beta): Built-in dashboard in the App UI. Metrics, traces, and logs in one view.

Production Ready

Production Checklist

✓
Agent SP per capability boundary
Not shared. Each agent gets its own SP with scoped permissions.
✓
UC grants: least privilege, explicit
No broad grants. GRANT only what is needed, REVOKE everything else.
✓
Row filters: current_user() only
Filter data to the calling user’s identity, not the SP.
✓
UC Connections: USE CONNECTION grants
Every external tool governed by GRANT/REVOKE on connection objects.
✓
OAuth scopes: minimal
Not all-apis. Use sql, genie, serving, dashboards.genie as needed.
✓
Token federation: no stored secrets
Token federation exchange eliminates secrets in CI/CD and external apps.
✓
Audit: system.access.audit enabled
Every API call logged. Retention configured. Alerting on anomalies.
✓
MLflow tracing: enabled for all agents
Every tool call, LLM invocation, and agent step traced.
✓
AI Gateway: rate limits + guardrails
QPM/TPM limits per user. Content filtering on input and output.
✓
Tokens never in logs or error messages
Bearer tokens excluded from all log output and exception payloads.

Orchestration is governance.

Every agent action traces back to an identity.
Every tool call goes through UC.
Every external service call goes through a connection.
Every step is audited.